torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/trainers/algorithms/configs/objectives.py

@@ -0,0 +1,177 @@
+# 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.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from torchrl.objectives import ClipPPOLoss, KLPENPPOLoss, PPOLoss, SACLoss
+from torchrl.objectives.sac import DiscreteSACLoss
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@dataclass
+class LossConfig(ConfigBase):
+    """A class to configure a loss.
+
+    Args:
+        loss_type: The type of loss to use.
+    """
+
+    _partial_: bool = False
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for loss configurations."""
+
+
+@dataclass
+class SACLossConfig(LossConfig):
+    """A class to configure a SAC loss."""
+
+    actor_network: Any = None
+    qvalue_network: Any = None
+    value_network: Any = None
+    discrete: bool = False
+    num_qvalue_nets: int = 2
+    loss_function: str = "smooth_l1"
+    alpha_init: float = 1.0
+    min_alpha: float | None = None
+    max_alpha: float | None = None
+    action_spec: Any = None
+    fixed_alpha: bool = False
+    target_entropy: str | float = "auto"
+    delay_actor: bool = False
+    delay_qvalue: bool = True
+    delay_value: bool = True
+    gamma: float | None = None
+    priority_key: str | None = None
+    separate_losses: bool = False
+    reduction: str | None = None
+    skip_done_states: bool = False
+    deactivate_vmap: bool = False
+    _target_: str = "torchrl.trainers.algorithms.configs.objectives._make_sac_loss"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for SAC loss configurations."""
+        super().__post_init__()
+
+
+def _make_sac_loss(*args, **kwargs) -> SACLoss:
+    discrete_loss_type = kwargs.pop("discrete", False)
+
+    # Instantiate networks if they are config objects
+    actor_network = kwargs.get("actor_network")
+    qvalue_network = kwargs.get("qvalue_network")
+    value_network = kwargs.get("value_network")
+
+    if actor_network is not None and hasattr(actor_network, "_target_"):
+        kwargs["actor_network"] = actor_network()
+    if qvalue_network is not None and hasattr(qvalue_network, "_target_"):
+        kwargs["qvalue_network"] = qvalue_network()
+    if value_network is not None and hasattr(value_network, "_target_"):
+        kwargs["value_network"] = value_network()
+
+    if discrete_loss_type:
+        return DiscreteSACLoss(*args, **kwargs)
+    else:
+        return SACLoss(*args, **kwargs)
+
+
+@dataclass
+class PPOLossConfig(LossConfig):
+    """A class to configure a PPO loss."""
+
+    actor_network: Any = None
+    critic_network: Any = None
+    loss_type: str = "clip"
+    entropy_bonus: bool = True
+    samples_mc_entropy: int = 1
+    entropy_coeff: float | None = None
+    log_explained_variance: bool = True
+    critic_coeff: float = 0.25
+    loss_critic_type: str = "smooth_l1"
+    normalize_advantage: bool = True
+    normalize_advantage_exclude_dims: tuple = ()
+    gamma: float | None = None
+    separate_losses: bool = False
+    advantage_key: str | None = None
+    value_target_key: str | None = None
+    value_key: str | None = None
+    functional: bool = True
+    actor: Any = None
+    critic: Any = None
+    reduction: str | None = None
+    clip_value: float | None = None
+    device: Any = None
+    _target_: str = "torchrl.trainers.algorithms.configs.objectives._make_ppo_loss"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for PPO loss configurations."""
+        super().__post_init__()
+
+
+def _make_ppo_loss(*args, **kwargs) -> PPOLoss:
+    loss_type = kwargs.pop("loss_type", "clip")
+    if loss_type == "clip":
+        return ClipPPOLoss(*args, **kwargs)
+    elif loss_type == "kl":
+        return KLPENPPOLoss(*args, **kwargs)
+    elif loss_type == "ppo":
+        return PPOLoss(*args, **kwargs)
+    else:
+        raise ValueError(f"Invalid loss type: {loss_type}")
+
+
+@dataclass
+class TargetNetUpdaterConfig:
+    """An abstract class to configure target net updaters."""
+
+    loss_module: Any
+    _partial_: bool = True
+
+
+@dataclass
+class SoftUpdateConfig(TargetNetUpdaterConfig):
+    """A class for soft update instantiation."""
+
+    _target_: str = "torchrl.objectives.utils.SoftUpdate"
+    eps: float | None = None  # noqa # type-ignore
+    tau: float | None = 0.001  # noqa # type-ignore
+
+
+@dataclass
+class HardUpdateConfig(TargetNetUpdaterConfig):
+    """A class for hard update instantiation."""
+
+    _target_: str = "torchrl.objectives.utils.HardUpdate."
+    value_network_update_interval: int = 1000
+
+
+@dataclass
+class GAEConfig(LossConfig):
+    """A class to configure a GAELoss."""
+
+    gamma: float | None = None
+    lmbda: float | None = None
+    value_network: Any = None
+    average_gae: bool = True
+    differentiable: bool = False
+    vectorized: bool | None = None
+    skip_existing: bool | None = None
+    advantage_key: str | None = None
+    value_target_key: str | None = None
+    value_key: str | None = None
+    shifted: bool = False
+    device: Any = None
+    time_dim: int | None = None
+    auto_reset_env: bool = False
+    deactivate_vmap: bool = False
+    _target_: str = "torchrl.objectives.value.GAE"
+    _partial_: bool = False
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for GAELoss configurations."""
+        super().__post_init__()

torchrl/trainers/algorithms/configs/trainers.py

@@ -0,0 +1,340 @@
+# 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.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from tensordict.nn import TensorDictModuleBase
+
+from torchrl.collectors import BaseCollector
+from torchrl.objectives.common import LossModule
+from torchrl.objectives.utils import TargetNetUpdater
+from torchrl.objectives.value.advantages import GAE
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+from torchrl.trainers.algorithms.ppo import PPOTrainer
+from torchrl.trainers.algorithms.sac import SACTrainer
+
+
+@dataclass
+class TrainerConfig(ConfigBase):
+    """Base configuration class for trainers."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for trainer configurations."""
+
+
+@dataclass
+class SACTrainerConfig(TrainerConfig):
+    """Configuration class for SAC (Soft Actor Critic) trainer.
+
+    This class defines the configuration parameters for creating a SAC trainer,
+    including both required and optional fields with sensible defaults.
+    """
+
+    collector: Any
+    total_frames: int
+    optim_steps_per_batch: int | None
+    loss_module: Any
+    optimizer: Any
+    logger: Any
+    save_trainer_file: Any
+    replay_buffer: Any
+    frame_skip: int = 1
+    clip_grad_norm: bool = True
+    clip_norm: float | None = None
+    progress_bar: bool = True
+    seed: int | None = None
+    save_trainer_interval: int = 10000
+    log_interval: int = 10000
+    create_env_fn: Any = None
+    actor_network: Any = None
+    critic_network: Any = None
+    target_net_updater: Any = None
+    async_collection: bool = False
+    log_timings: bool = False
+
+    _target_: str = "torchrl.trainers.algorithms.configs.trainers._make_sac_trainer"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for SAC trainer configuration."""
+        super().__post_init__()
+
+
+def _make_sac_trainer(*args, **kwargs) -> SACTrainer:
+    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)
+    log_timings = kwargs.pop("log_timings", False)
+
+    # Instantiate networks first
+    if actor_network is not None:
+        actor_network = actor_network()
+    if critic_network is not None:
+        critic_network = critic_network()
+
+    if not isinstance(collector, BaseCollector):
+        # then it's a partial config
+        if not async_collection:
+            collector = collector()
+        elif replay_buffer is not None:
+            collector = collector(replay_buffer=replay_buffer)
+    elif getattr(collector, "replay_buffer", None) is None:
+        if async_collection and (
+            collector.replay_buffer is None or replay_buffer is None
+        ):
+            raise ValueError(
+                "replay_buffer must be provided when async_collection is True"
+            )
+
+    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 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)}")
+
+    return SACTrainer(
+        collector=collector,
+        total_frames=total_frames,
+        frame_skip=frame_skip,
+        optim_steps_per_batch=optim_steps_per_batch,
+        loss_module=loss_module,
+        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,
+        replay_buffer=replay_buffer,
+        target_net_updater=target_net_updater,
+        async_collection=async_collection,
+        log_timings=log_timings,
+    )
+
+
+@dataclass
+class PPOTrainerConfig(TrainerConfig):
+    """Configuration class for PPO (Proximal Policy Optimization) trainer.
+
+    This class defines the configuration parameters for creating a PPO trainer,
+    including both required and optional fields with sensible defaults.
+
+    Args:
+        collector: The data collector for gathering training data.
+        total_frames: Total number of frames to train for.
+        optim_steps_per_batch: Number of optimization steps per batch.
+        loss_module: The loss module for computing policy and value losses.
+        optimizer: The optimizer for training.
+        logger: Logger for tracking training metrics.
+        save_trainer_file: File path for saving trainer state.
+        replay_buffer: Replay buffer for storing data.
+        frame_skip: Frame skip value for the environment. Default: 1.
+        clip_grad_norm: Whether to clip gradient norms. Default: True.
+        clip_norm: Maximum gradient norm value.
+        progress_bar: Whether to show a progress bar. Default: True.
+        seed: Random seed for reproducibility.
+        save_trainer_interval: Interval for saving trainer state. Default: 10000.
+        log_interval: Interval for logging metrics. Default: 10000.
+        create_env_fn: Environment creation function.
+        actor_network: Actor network configuration.
+        critic_network: Critic network configuration.
+        num_epochs: Number of epochs per batch. Default: 4.
+        async_collection: Whether to use async collection. Default: False.
+        add_gae: Whether to add GAE computation. Default: True.
+        gae: Custom GAE module configuration.
+        weight_update_map: Mapping from collector destination paths to trainer source paths.
+            Required if collector has weight_sync_schemes configured.
+            Example: ``{"policy": "loss_module.actor_network", "replay_buffer.transforms[0]": "loss_module.critic_network"}``.
+        log_timings: Whether to automatically log timing information for all hooks.
+            If True, timing metrics will be logged to the logger (e.g., wandb, tensorboard)
+            with prefix "time/" (e.g., "time/hook/UpdateWeights"). Default: False.
+    """
+
+    collector: Any
+    total_frames: int
+    optim_steps_per_batch: int | None
+    loss_module: Any
+    optimizer: Any
+    logger: Any
+    save_trainer_file: Any
+    replay_buffer: Any
+    frame_skip: int = 1
+    clip_grad_norm: bool = True
+    clip_norm: float | None = None
+    progress_bar: bool = True
+    seed: int | None = None
+    save_trainer_interval: int = 10000
+    log_interval: int = 10000
+    create_env_fn: Any = None
+    actor_network: Any = None
+    critic_network: Any = None
+    num_epochs: int = 4
+    async_collection: bool = False
+    add_gae: bool = True
+    gae: Any = None
+    weight_update_map: dict[str, str] | None = None
+    log_timings: bool = False
+
+    _target_: str = "torchrl.trainers.algorithms.configs.trainers._make_ppo_trainer"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for PPO trainer configuration."""
+        super().__post_init__()
+
+
+def _make_ppo_trainer(*args, **kwargs) -> PPOTrainer:
+    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")
+    add_gae = kwargs.pop("add_gae", True)
+    gae = kwargs.pop("gae")
+    create_env_fn = kwargs.pop("create_env_fn")
+    weight_update_map = kwargs.pop("weight_update_map", None)
+    log_timings = kwargs.pop("log_timings", False)
+
+    if create_env_fn is not None:
+        # could be referenced somewhere else, no need to raise an error
+        pass
+    num_epochs = kwargs.pop("num_epochs", 4)
+    async_collection = kwargs.pop("async_collection", False)
+
+    # Instantiate networks first
+    if actor_network is not None:
+        actor_network = actor_network()
+    if critic_network is not None:
+        critic_network = critic_network()
+    else:
+        critic_network = loss_module.critic_network
+
+    # Ensure GAE in replay buffer uses the same value network instance as loss module
+    # This fixes the issue where Hydra instantiates separate instances of value_model
+    if (
+        replay_buffer is not None
+        and hasattr(replay_buffer, "_transform")
+        and len(replay_buffer._transform) > 1
+        and hasattr(replay_buffer._transform[1], "module")
+        and hasattr(replay_buffer._transform[1].module, "value_network")
+    ):
+        replay_buffer._transform[1].module.value_network = critic_network
+
+    if not isinstance(collector, BaseCollector):
+        # then it's a partial config
+        if not async_collection:
+            collector = collector()
+        else:
+            collector = collector(replay_buffer=replay_buffer)
+    elif async_collection and getattr(collector, "replay_buffer", None) is None:
+        raise RuntimeError(
+            "replay_buffer must be provided when async_collection is True"
+        )
+    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 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)}")
+    # instantiate gae if it is a partial config
+    if not isinstance(gae, (GAE, TensorDictModuleBase)) and gae is not None:
+        gae = gae()
+
+    return PPOTrainer(
+        collector=collector,
+        total_frames=total_frames,
+        frame_skip=frame_skip,
+        optim_steps_per_batch=optim_steps_per_batch,
+        loss_module=loss_module,
+        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,
+        replay_buffer=replay_buffer,
+        num_epochs=num_epochs,
+        async_collection=async_collection,
+        add_gae=add_gae,
+        gae=gae,
+        weight_update_map=weight_update_map,
+        log_timings=log_timings,
+    )