[BugFix] Complete offline-to-online trainer wiring

Author: vmoens

docs/source/reference/config.rst

@@ -480,6 +480,7 @@ Training and Optimization Configurations
     TrainerConfig
     PPOTrainerConfig
     SACTrainerConfig
+    OfflineToOnlineTrainerConfig
     DQNTrainerConfig
     DDPGTrainerConfig
     IQLTrainerConfig
@@ -599,6 +600,7 @@ TorchRL currently provides configuration-driven trainers for the following algor
 
 - **PPO** (on-policy): ``PPOTrainerConfig``, ``PPOLossConfig``
 - **SAC** (off-policy, continuous): ``SACTrainerConfig``, ``SACLossConfig``
+- **Offline-to-online SAC**: ``OfflineToOnlineTrainerConfig``, ``SACLossConfig``
 - **DQN** (off-policy, discrete): ``DQNTrainerConfig``, ``DQNLossConfig``
 - **DDPG** (off-policy, continuous): ``DDPGTrainerConfig``, ``DDPGLossConfig``
 - **IQL** (offline): ``IQLTrainerConfig``, ``IQLLossConfig``

docs/source/reference/trainers_basics.rst

@@ -26,6 +26,7 @@ Algorithm-specific trainers
 
     PPOTrainer
     SACTrainer
+    OfflineToOnlineTrainer
     DQNTrainer
     DDPGTrainer
     IQLTrainer

sota-implementations/offline_to_online/train.py

@@ -118,7 +118,7 @@ def main():
 
     # Immutable offline dataset (DoubleToFloat to match the online float32 stream)
     # paired with a growing online buffer.
-    offline = load_dataset(args.dataset)
+    offline = load_dataset(args.dataset, batch_size=args.batch_size)
     offline.append_transform(DoubleToFloat())
     replay_buffer = OfflineToOnlineReplayBuffer(
         offline_dataset=offline,

test/test_offline_to_online.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import argparse
+import inspect
 
 import pytest
 import torch
@@ -498,13 +499,21 @@ def test_state_dict_roundtrip(self):
         )
         hook = OfflineToOnlineReplayBufferHook(rb)
         hook.extend(_make_online_data(20))
+        rb.anneal(step=50, total_steps=100)
 
         rb2 = OfflineToOnlineReplayBuffer(
-            offline_dataset=_make_offline_buffer(), online_capacity=500, batch_size=16
+            offline_dataset=_make_offline_buffer(),
+            online_capacity=500,
+            offline_fraction=0.8,
+            batch_size=16,
         )
         hook2 = OfflineToOnlineReplayBufferHook(rb2)
         hook2.load_state_dict(hook.state_dict())
         assert len(rb2.online_buffer) == 20
+        assert rb2.offline_fraction == pytest.approx(0.25)
+
+        rb2.anneal(step=50, total_steps=100)
+        assert rb2.offline_fraction == pytest.approx(0.25)
 
 
 class TestOfflineToOnlineAnnealHook:
@@ -555,6 +564,32 @@ def test_requires_offline_to_online_buffer(self):
                 replay_buffer=plain,
             )
 
+    def test_constructor_exposes_sac_key_and_logging_kwargs(self):
+        from torchrl.trainers.algorithms.offline_to_online import OfflineToOnlineTrainer
+
+        parameters = inspect.signature(OfflineToOnlineTrainer).parameters
+        for name in (
+            "log_rewards",
+            "log_actions",
+            "log_observations",
+            "log_timings",
+            "auto_log_optim_steps",
+            "done_key",
+            "terminated_key",
+            "reward_key",
+            "episode_reward_key",
+            "action_key",
+            "observation_key",
+        ):
+            assert name in parameters
+
+    def test_config_registered(self):
+        from torchrl.trainers.algorithms.configs import OfflineToOnlineTrainerConfig
+
+        assert OfflineToOnlineTrainerConfig._target_.endswith(
+            "_make_offline_to_online_trainer"
+        )
+
     def test_hooks_drive_offline_online_flow(self):
         """The three hooks together grow the online buffer, keep the mixed batch
         flat, and anneal the offline fraction -- the data path the trainer runs,

torchrl/trainers/algorithms/configs/__init__.py

@@ -123,6 +123,7 @@
     DDPGTrainerConfig,
     DQNTrainerConfig,
     IQLTrainerConfig,
+    OfflineToOnlineTrainerConfig,
     PPOTrainerConfig,
     SACTrainerConfig,
     TD3TrainerConfig,
@@ -397,6 +398,7 @@
     "DDPGTrainerConfig",
     "DQNTrainerConfig",
     "IQLTrainerConfig",
+    "OfflineToOnlineTrainerConfig",
     "PPOTrainerConfig",
     "SACTrainerConfig",
     "TD3TrainerConfig",
@@ -671,6 +673,11 @@ def _register_configs():
     cs.store(group="trainer", name="ddpg", node=DDPGTrainerConfig)
     cs.store(group="trainer", name="dqn", node=DQNTrainerConfig)
     cs.store(group="trainer", name="iql", node=IQLTrainerConfig)
+    cs.store(
+        group="trainer",
+        name="offline_to_online",
+        node=OfflineToOnlineTrainerConfig,
+    )
     cs.store(group="trainer", name="ppo", node=PPOTrainerConfig)
     cs.store(group="trainer", name="sac", node=SACTrainerConfig)
     cs.store(group="trainer", name="td3", node=TD3TrainerConfig)

torchrl/trainers/algorithms/configs/trainers.py

@@ -22,6 +22,7 @@
 from torchrl.trainers.algorithms.ddpg import DDPGTrainer
 from torchrl.trainers.algorithms.dqn import DQNTrainer
 from torchrl.trainers.algorithms.iql import IQLTrainer
+from torchrl.trainers.algorithms.offline_to_online import OfflineToOnlineTrainer
 from torchrl.trainers.algorithms.ppo import PPOTrainer
 from torchrl.trainers.algorithms.sac import SACTrainer
 from torchrl.trainers.algorithms.td3 import TD3Trainer
@@ -218,6 +219,147 @@ def _make_sac_trainer(*args, **kwargs) -> SACTrainer:
     return trainer
 
 
+@dataclass
+class OfflineToOnlineTrainerConfig(SACTrainerConfig):
+    """Hydra configuration for
+    :class:`~torchrl.trainers.algorithms.OfflineToOnlineTrainer`.
+
+    Every kwarg accepted by ``OfflineToOnlineTrainer.__init__`` is exposed as a
+    field here, with SAC network-construction helper fields inherited from
+    :class:`SACTrainerConfig`.
+    """
+
+    anneal_frames: int | None = None
+
+    _target_: str = (
+        "torchrl.trainers.algorithms.configs.trainers."
+        "_make_offline_to_online_trainer"
+    )
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for offline-to-online trainer configuration."""
+        super().__post_init__()
+
+
+def _make_offline_to_online_trainer(*args, **kwargs) -> OfflineToOnlineTrainer:
+    from torchrl.trainers.trainers import Logger
+
+    collector = kwargs.pop("collector")
+    total_frames = kwargs.pop("total_frames")
+    if total_frames is None:
+        total_frames = collector.total_frames
+    frame_skip = kwargs.pop("frame_skip", 1)
+    optim_steps_per_batch = kwargs.pop("optim_steps_per_batch", 1)
+    loss_module = kwargs.pop("loss_module")
+    optimizer = kwargs.pop("optimizer")
+    logger = kwargs.pop("logger")
+    clip_grad_norm = kwargs.pop("clip_grad_norm", True)
+    clip_norm = kwargs.pop("clip_norm")
+    progress_bar = kwargs.pop("progress_bar", True)
+    replay_buffer = kwargs.pop("replay_buffer")
+    save_trainer_interval = kwargs.pop("save_trainer_interval", 10000)
+    log_interval = kwargs.pop("log_interval", 10000)
+    save_trainer_file = kwargs.pop("save_trainer_file")
+    seed = kwargs.pop("seed")
+    actor_network = kwargs.pop("actor_network")
+    critic_network = kwargs.pop("critic_network")
+    kwargs.pop("create_env_fn")
+    target_net_updater = kwargs.pop("target_net_updater")
+    async_collection = kwargs.pop("async_collection", False)
+    if async_collection:
+        raise ValueError(
+            "OfflineToOnlineTrainer does not support async_collection."
+        )
+    log_timings = kwargs.pop("log_timings", False)
+    auto_log_optim_steps = kwargs.pop("auto_log_optim_steps", True)
+    batch_size = kwargs.pop("batch_size", None)
+    anneal_frames = kwargs.pop("anneal_frames", None)
+    enable_logging = kwargs.pop("enable_logging", True)
+    log_rewards = kwargs.pop("log_rewards", True)
+    log_actions = kwargs.pop("log_actions", True)
+    log_observations = kwargs.pop("log_observations", False)
+    done_key = _normalize_hydra_key(kwargs.pop("done_key", "done"))
+    terminated_key = _normalize_hydra_key(kwargs.pop("terminated_key", "terminated"))
+    reward_key = _normalize_hydra_key(kwargs.pop("reward_key", "reward"))
+    episode_reward_key = _normalize_hydra_key(
+        kwargs.pop("episode_reward_key", "reward_sum")
+    )
+    action_key = _normalize_hydra_key(kwargs.pop("action_key", "action"))
+    observation_key = _normalize_hydra_key(kwargs.pop("observation_key", "observation"))
+    hooks = kwargs.pop("hooks", None)
+
+    # Instantiate networks first
+    if actor_network is not None and not isinstance(actor_network, torch.nn.Module):
+        actor_network = actor_network()
+    if critic_network is not None and not isinstance(critic_network, torch.nn.Module):
+        critic_network = critic_network()
+
+    if not isinstance(collector, BaseCollector):
+        collector = collector()
+
+    if not isinstance(loss_module, LossModule):
+        # then it's a partial config
+        loss_module = loss_module(
+            actor_network=actor_network, critic_network=critic_network
+        )
+    if target_net_updater is not None and not isinstance(
+        target_net_updater, TargetNetUpdater
+    ):
+        # target_net_updater must be a partial taking the loss as input
+        target_net_updater = target_net_updater(loss_module)
+    if not isinstance(optimizer, torch.optim.Optimizer):
+        # then it's a partial config
+        optimizer = optimizer(params=loss_module.parameters())
+
+    # Quick instance checks
+    if not isinstance(collector, BaseCollector):
+        raise ValueError(f"collector must be a BaseCollector, got {type(collector)}")
+    if not isinstance(loss_module, LossModule):
+        raise ValueError(f"loss_module must be a LossModule, got {type(loss_module)}")
+    if not isinstance(optimizer, torch.optim.Optimizer):
+        raise ValueError(
+            f"optimizer must be a torch.optim.Optimizer, got {type(optimizer)}"
+        )
+    if not isinstance(logger, Logger) and logger is not None:
+        raise ValueError(f"logger must be a Logger, got {type(logger)}")
+
+    trainer = OfflineToOnlineTrainer(
+        collector=collector,
+        total_frames=total_frames,
+        frame_skip=frame_skip,
+        optim_steps_per_batch=optim_steps_per_batch,
+        loss_module=loss_module,
+        replay_buffer=replay_buffer,
+        anneal_frames=anneal_frames,
+        batch_size=batch_size,
+        optimizer=optimizer,
+        logger=logger,
+        clip_grad_norm=clip_grad_norm,
+        clip_norm=clip_norm,
+        progress_bar=progress_bar,
+        seed=seed,
+        save_trainer_interval=save_trainer_interval,
+        log_interval=log_interval,
+        save_trainer_file=save_trainer_file,
+        enable_logging=enable_logging,
+        log_rewards=log_rewards,
+        log_actions=log_actions,
+        log_observations=log_observations,
+        target_net_updater=target_net_updater,
+        async_collection=async_collection,
+        log_timings=log_timings,
+        auto_log_optim_steps=auto_log_optim_steps,
+        done_key=done_key,
+        terminated_key=terminated_key,
+        reward_key=reward_key,
+        episode_reward_key=episode_reward_key,
+        action_key=action_key,
+        observation_key=observation_key,
+    )
+    _register_trainer_hooks(trainer, hooks)
+    return trainer
+
+
 @dataclass
 class PPOTrainerConfig(TrainerConfig):
     """Hydra configuration for :class:`~torchrl.trainers.algorithms.PPOTrainer`.

torchrl/trainers/algorithms/offline_to_online.py

@@ -10,6 +10,8 @@
 from collections.abc import Callable
 
 from tensordict import TensorDictBase
+from tensordict.utils import NestedKey
+from torch import optim
 
 from torchrl.collectors import BaseCollector
 from torchrl.data.replay_buffers.offline_to_online import OfflineToOnlineReplayBuffer
@@ -89,10 +91,20 @@ def sample(self, batch: TensorDictBase) -> TensorDictBase:
         return sample.to(self.device) if self.device is not None else sample
 
     def state_dict(self) -> dict:
-        return {"online_buffer": self.replay_buffer.online_buffer.state_dict()}
+        return {
+            "online_buffer": self.replay_buffer.online_buffer.state_dict(),
+            "offline_fraction": self.replay_buffer._offline_fraction,
+            "base_offline_fraction": self.replay_buffer._base_offline_fraction,
+        }
 
     def load_state_dict(self, state_dict: dict) -> None:
         self.replay_buffer.online_buffer.load_state_dict(state_dict["online_buffer"])
+        self.replay_buffer._offline_fraction = state_dict.get(
+            "offline_fraction", self.replay_buffer._offline_fraction
+        )
+        self.replay_buffer._base_offline_fraction = state_dict.get(
+            "base_offline_fraction", self.replay_buffer._base_offline_fraction
+        )
 
     def register(self, trainer, name: str = "replay_buffer") -> None:
         trainer.register_op("pre_epoch", self.extend)
@@ -136,6 +148,10 @@ def register(self, trainer, name: str = "offline_to_online_anneal") -> None:
 class OfflineToOnlineTrainer(SACTrainer):
     """A SAC trainer for the offline-pretrain -> online-finetune transition.
 
+    See also
+    :class:`~torchrl.trainers.algorithms.configs.OfflineToOnlineTrainerConfig`
+    for the Hydra configuration counterpart.
+
     Builds on :class:`~torchrl.trainers.algorithms.SACTrainer`, swapping the
     plain replay buffer for an :class:`~torchrl.data.OfflineToOnlineReplayBuffer`.
     Each collected batch is routed to the online buffer while optimization
@@ -175,7 +191,7 @@ def __init__(
         replay_buffer: OfflineToOnlineReplayBuffer,
         anneal_frames: int | None = None,
         batch_size: int | None = None,
-        optimizer=None,
+        optimizer: optim.Optimizer | None = None,
         logger: Logger | None = None,
         clip_grad_norm: bool = True,
         clip_norm: float | None = None,
@@ -185,13 +201,29 @@ def __init__(
         log_interval: int = 10000,
         save_trainer_file: str | pathlib.Path | None = None,
         enable_logging: bool = True,
+        log_rewards: bool = True,
+        log_actions: bool = True,
+        log_observations: bool = False,
         target_net_updater: TargetNetUpdater | None = None,
+        async_collection: bool = False,
+        log_timings: bool = False,
+        auto_log_optim_steps: bool = True,
+        done_key: NestedKey = "done",
+        terminated_key: NestedKey = "terminated",
+        reward_key: NestedKey = "reward",
+        episode_reward_key: NestedKey = "reward_sum",
+        action_key: NestedKey = "action",
+        observation_key: NestedKey = "observation",
     ) -> None:
         if not isinstance(replay_buffer, OfflineToOnlineReplayBuffer):
             raise TypeError(
                 "OfflineToOnlineTrainer requires an OfflineToOnlineReplayBuffer, "
                 f"got {type(replay_buffer).__name__}."
             )
+        if async_collection:
+            raise ValueError(
+                "OfflineToOnlineTrainer does not support async_collection."
+            )
 
         # Let SACTrainer wire up everything except the replay buffer (its
         # ReplayBufferTrainer assumes a sampler/priority API the offline-to-online
@@ -213,8 +245,19 @@ def __init__(
             save_trainer_file=save_trainer_file,
             replay_buffer=None,
             enable_logging=enable_logging,
+            log_rewards=log_rewards,
+            log_actions=log_actions,
+            log_observations=log_observations,
             target_net_updater=target_net_updater,
             async_collection=False,
+            log_timings=log_timings,
+            auto_log_optim_steps=auto_log_optim_steps,
+            done_key=done_key,
+            terminated_key=terminated_key,
+            reward_key=reward_key,
+            episode_reward_key=episode_reward_key,
+            action_key=action_key,
+            observation_key=observation_key,
         )
 
         self.replay_buffer = replay_buffer