[Feature] OfflineToOnlineTrainer + sota-implementation for offline->online RL

Follow-up to the OfflineToOnlineReplayBuffer PR: a config-driven SAC trainer
that drives the offline-pretrain -> online-finetune transition, plus a
sota-implementations script demonstrating it.

- OfflineToOnlineTrainer (subclasses SACTrainer): routes collected experience
  to the online buffer (pre_epoch), samples a mixed offline/online batch
  (process_optim_batch), and anneals the offline fraction to zero over
  anneal_frames (post_steps). Two reusable hooks back it:
  OfflineToOnlineReplayBufferHook (projects online transitions onto the offline
  dataset schema so the mixed-batch concat stays valid) and
  OfflineToOnlineAnnealHook.
- Hydra configs: OfflineToOnlineReplayBufferConfig and
  OfflineToOnlineTrainerConfig, registered under the replay_buffer /
  trainer groups (name "offline_to_online").
- sota-implementations/offline_to_online_trainer/: train.py + config.yaml
  (SAC on HalfCheetah, offline dataset via d4rl:/minari: string).
- Tests: hook unit + flow tests, a gated functional train() run on Pendulum,
  and config-instantiation tests; OfflineToOnlineTrainer added to the
  auto_log_optim_steps invariant.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: theap06

sota-implementations/offline_to_online_trainer/config/config.yaml

@@ -0,0 +1,146 @@
+# Offline-to-online SAC fine-tuning for HalfCheetah.
+#
+# Warm-starts SAC on an offline dataset and fine-tunes online, annealing the
+# offline sampling fraction to zero over `trainer.anneal_frames` frames. Built
+# on the configurable trainer system (see sota-implementations/sac_trainer).
+
+defaults:
+
+  - transform@transform0: step_counter
+  - transform@transform1: double_to_float
+  - transform@transform2: reward_sum
+
+  - env@training_env: batched_env
+  - env@training_env.create_env_fn: transformed_env
+  - env@training_env.create_env_fn.base_env: gym
+  - transform@training_env.create_env_fn.transform: compose
+
+  - model@models.policy_model: tanh_normal
+  - model@models.qvalue_model: value
+
+  - network@networks.policy_network: mlp
+  - network@networks.qvalue_network: mlp
+
+  - collector@collector: sync
+
+  - replay_buffer@replay_buffer: offline_to_online
+  - trainer@trainer: offline_to_online
+  - optimizer@optimizer: adam
+  - loss@loss: sac
+  - target_net_updater@target_net_updater: soft
+  - logger@logger: wandb
+  - _self_
+
+# Network configurations
+networks:
+  policy_network:
+    out_features: 12  # HalfCheetah action space is 6-dimensional (loc + scale)
+    in_features: 17   # HalfCheetah observation space is 17-dimensional
+    num_cells: [256, 256]
+
+  qvalue_network:
+    out_features: 1  # Q-value output
+    in_features: 23   # HalfCheetah observation space (17) + action space (6)
+    num_cells: [256, 256]
+
+# Model configurations
+models:
+  policy_model:
+    return_log_prob: true
+    in_keys: ["observation"]
+    param_keys: ["loc", "scale"]
+    out_keys: ["action"]
+    network: ${networks.policy_network}
+
+  qvalue_model:
+    in_keys: ["observation", "action"]
+    out_keys: ["state_action_value"]
+    network: ${networks.qvalue_network}
+
+transform0:
+  max_steps: 1000
+  step_count_key: "step_count"
+
+transform1:
+  # DoubleToFloatTransform - converts double precision to float to fix dtype mismatch
+  in_keys: null
+  out_keys: null
+
+transform2:
+  # RewardSumTransform - sums up the rewards
+  in_keys: ["reward"]
+  out_keys: ["reward_sum"]
+
+training_env:
+  num_workers: 1
+  create_env_fn:
+    base_env:
+      env_name: HalfCheetah-v4
+    transform:
+      transforms:
+        - ${transform0}
+        - ${transform1}
+        - ${transform2}
+    _partial_: true
+
+# Loss configuration
+loss:
+  actor_network: ${models.policy_model}
+  qvalue_network: ${models.qvalue_model}
+  target_entropy: "auto"
+  loss_function: l2
+  alpha_init: 1.0
+  delay_qvalue: true
+  num_qvalue_nets: 2
+
+target_net_updater:
+  tau: 0.001
+
+# Optimizer configuration
+optimizer:
+  lr: 3.0e-4
+
+# Collector configuration (sync: offline-to-online does not support async collection)
+collector:
+  create_env_fn: ${training_env}
+  policy: ${models.policy_model}
+  total_frames: 1_000_000
+  frames_per_batch: 1000
+  init_random_frames: 0  # the offline dataset already warm-starts learning
+  track_policy_version: true
+  _partial_: true
+
+# Replay buffer: immutable offline dataset + growing online buffer.
+# Swap offline_dataset for e.g. "minari:mujoco/halfcheetah/expert-v0".
+replay_buffer:
+  offline_dataset: "d4rl:halfcheetah-medium-v2"
+  online_capacity: 1_000_000
+  offline_fraction: 0.5
+  batch_size: 256
+
+logger:
+  exp_name: offline_to_online_halfcheetah_v4
+  offline: false
+  project: torchrl-sota-implementations
+
+# Trainer configuration
+trainer:
+  collector: ${collector}
+  optimizer: ${optimizer}
+  replay_buffer: ${replay_buffer}
+  target_net_updater: ${target_net_updater}
+  loss_module: ${loss}
+  logger: ${logger}
+  total_frames: ${collector.total_frames}
+  # Anneal the offline sampling fraction to 0 over the first half of training.
+  anneal_frames: 500_000
+  batch_size: 256
+  frame_skip: 1
+  clip_grad_norm: false  # SAC typically doesn't use gradient clipping
+  clip_norm: null
+  progress_bar: true
+  seed: 42
+  save_trainer_interval: 25000
+  log_interval: 25000
+  save_trainer_file: null
+  optim_steps_per_batch: 64

sota-implementations/offline_to_online_trainer/train.py

@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Offline-to-online SAC fine-tuning via the configurable trainer system.
+
+Warm-starts SAC on an offline dataset (D4RL/Minari) and fine-tunes it online,
+sampling a mixed offline/online batch whose offline fraction is annealed to zero
+over ``trainer.anneal_frames`` collected frames.
+
+Run with::
+
+    python train.py
+    python train.py replay_buffer.offline_dataset=minari:mujoco/halfcheetah/expert-v0
+
+Requires the offline dataset's backend (``pip install d4rl`` or
+``pip install minari``) and the matching MuJoCo environment.
+"""
+
+import hydra
+from torchrl.trainers.algorithms.configs import *  # noqa: F401, F403
+
+
+@hydra.main(config_path="config", config_name="config", version_base="1.1")
+def main(cfg):
+    trainer = hydra.utils.instantiate(cfg.trainer)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()

test/test_configs.py

@@ -1130,7 +1130,6 @@ def test_collector_config(self, factory, collector):
 
         # Define cfg_cls and kwargs based on collector type
         if collector == "async":
-
             cfg_cls = AsyncCollectorConfig
             kwargs = {"create_env_fn": env_cfg, "frames_per_batch": 10}
         elif collector == "multi_sync":
@@ -1869,6 +1868,59 @@ def test_td3_trainer_config(self):
         assert cfg.policy_update_delay == 2
         assert cfg.clip_grad_norm is True
 
+    @pytest.mark.skipif(not _has_gymnasium, reason="Gymnasium is not installed")
+    def test_offline_to_online_trainer_config(self):
+        from torchrl.trainers.algorithms.configs.trainers import (
+            OfflineToOnlineTrainerConfig,
+        )
+
+        cfg = OfflineToOnlineTrainerConfig(
+            collector=None,
+            total_frames=600,
+            optim_steps_per_batch=2,
+            loss_module=None,
+            optimizer=None,
+            logger=None,
+            save_trainer_file=None,
+            replay_buffer=None,
+            anneal_frames=300,
+        )
+
+        assert (
+            cfg._target_
+            == "torchrl.trainers.algorithms.configs.trainers._make_offline_to_online_trainer"
+        )
+        assert cfg.total_frames == 600
+        assert cfg.anneal_frames == 300
+        # defaults inherited from the SAC-style config
+        assert cfg.frame_skip == 1
+        assert cfg.auto_log_optim_steps is True
+        assert cfg.actor_network is None
+        assert cfg.critic_network is None
+
+    @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed")
+    def test_offline_to_online_replay_buffer_config_instantiation(self):
+        from hydra.utils import instantiate
+
+        from torchrl.data import OfflineToOnlineReplayBuffer
+        from torchrl.trainers.algorithms.configs.data import (
+            LazyTensorStorageConfig,
+            OfflineToOnlineReplayBufferConfig,
+            ReplayBufferConfig,
+        )
+
+        cfg = OfflineToOnlineReplayBufferConfig(
+            offline_dataset=ReplayBufferConfig(
+                storage=LazyTensorStorageConfig(max_size=100)
+            ),
+            online_capacity=200,
+            offline_fraction=0.25,
+            batch_size=32,
+        )
+        rb = instantiate(cfg)
+        assert isinstance(rb, OfflineToOnlineReplayBuffer)
+        assert rb.offline_fraction == 0.25
+
 
 @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed")
 @pytest.mark.skipif(not _has_gymnasium, reason="Gymnasium is not installed")