# 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 collections.abc import Sequence
from dataclasses import dataclass, field
from functools import partial
from typing import Any
import torch
import torch.nn as nn
from hydra.utils import instantiate
from omegaconf import MISSING
from tensordict.nn import TensorDictModule, TensorDictSequential
from torchrl.modules import (
AdditiveGaussianModule,
DreamerV3DiscreteActor,
LowLevelController,
MLP,
QValueActor,
RSSMStateEstimatorV3,
SimplicialNormalization,
TanhModule,
TdMpc2Planner,
TdMpc2QEnsemble,
ValueOperator,
WorldModel,
)
from torchrl.modules.models.tdmpc2 import _TdMpc2PolicyPrior
from torchrl.trainers.algorithms.configs.common import (
_normalize_hydra_key,
_normalize_hydra_keys,
ConfigBase,
)
@dataclass
class ActivationConfig(ConfigBase):
"""A class to configure an activation function.
Defaults to :class:`torch.nn.Tanh`.
.. seealso:: :class:`torch.nn.Tanh`
"""
_target_: str = "torch.nn.Tanh"
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for activation configurations."""
@dataclass
class LayerConfig(ConfigBase):
"""A class to configure a layer.
Defaults to :class:`torch.nn.Linear`.
.. seealso:: :class:`torch.nn.Linear`
"""
_target_: str = "torch.nn.Linear"
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for layer configurations."""
[docs]
@dataclass
class NetworkConfig(ConfigBase):
"""Parent class to configure a network."""
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for network configurations."""
@dataclass
class LowLevelControllerConfig(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.LowLevelController`.
Pass the pretrained policy to Hydra's instantiate and declare the unbatched
Composite decision spec with a Hydra target. Adapter and nested keys can
also be configured in YAML.
"""
policy: Any = MISSING
decision_spec: Any = MISSING
adapter: Any = None
group_key: Any = None
state_key: Any = "controller"
policy_action_key: Any = "action"
action_key: Any = "action"
reset_key: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_low_level_controller"
)
def _make_low_level_controller(**kwargs: Any) -> LowLevelController:
"""Convert Hydra lists to nested TensorDict keys before construction."""
return LowLevelController(
policy=kwargs.pop("policy"),
decision_spec=kwargs.pop("decision_spec"),
adapter=kwargs.pop("adapter", None),
group_key=_normalize_hydra_key(kwargs.pop("group_key", None)),
state_key=_normalize_hydra_key(kwargs.pop("state_key", "controller")),
policy_action_key=_normalize_hydra_key(
kwargs.pop("policy_action_key", "action")
),
action_key=_normalize_hydra_key(kwargs.pop("action_key", "action")),
reset_key=_normalize_hydra_key(kwargs.pop("reset_key", None)),
**kwargs,
)
[docs]
@dataclass
class MLPConfig(NetworkConfig):
"""A class to configure a multi-layer perceptron.
Example:
>>> cfg = MLPConfig(in_features=10, out_features=5, depth=2, num_cells=32)
>>> net = instantiate(cfg)
>>> y = net(torch.randn(1, 10))
>>> assert y.shape == (1, 5)
.. seealso:: :class:`torchrl.modules.MLP`
"""
in_features: int | None = None
out_features: Any = None
depth: int | None = None
num_cells: Any = None
activation_class: ActivationConfig = field(
default_factory=partial(
ActivationConfig, _target_="torch.nn.Tanh", _partial_=True
)
)
activation_kwargs: Any = None
norm_class: Any = None
norm_kwargs: Any = None
dropout: float | None = None
bias_last_layer: bool = True
single_bias_last_layer: bool = False
layer_class: LayerConfig = field(
default_factory=partial(LayerConfig, _target_="torch.nn.Linear", _partial_=True)
)
layer_kwargs: dict | None = None
activate_last_layer: bool = False
device: Any = None
_target_: str = "torchrl.modules.MLP"
def __post_init__(self):
if isinstance(self.activation_class, str):
self.activation_class = ActivationConfig(
_target_=self.activation_class, _partial_=True
)
if isinstance(self.layer_class, str):
self.layer_class = LayerConfig(_target_=self.layer_class, _partial_=True)
[docs]
@dataclass
class DreamerV3MLPConfig(NetworkConfig):
"""A class to configure a DreamerV3 multilayer perceptron.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import DreamerV3MLPConfig
>>> cfg = DreamerV3MLPConfig(
... in_features=6, out_features=4, depth=2, num_cells=8
... )
>>> net = instantiate(cfg)
>>> y = net(torch.randn(3, 2), torch.randn(3, 4))
>>> assert y.shape == (3, 4)
.. seealso:: :class:`~torchrl.modules.DreamerV3MLP`
"""
in_features: int = MISSING
out_features: int | None = MISSING
depth: int = 3
num_cells: int = 1024
outscale: float = 1.0
norm_eps: float = 1e-4
device: Any = None
_target_: str = "torchrl.modules.DreamerV3MLP"
[docs]
@dataclass
class TdMpc2MLPConfig(NetworkConfig):
"""A class to configure a TDMPC2 multilayer perceptron.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import TdMpc2MLPConfig
>>> cfg = TdMpc2MLPConfig(
... in_features=6, out_features=4, depth=2, num_cells=8
... )
>>> net = instantiate(cfg)
>>> y = net(torch.randn(3, 6))
>>> assert y.shape == (3, 4)
.. seealso:: :class:`~torchrl.modules.MLP`
"""
in_features: int = MISSING
out_features: int = MISSING
depth: int | None = None
num_cells: Any = MISSING
output_activation: Any = None
dropout: float = 0.0
device: Any = None
_target_: str = "torchrl.trainers.algorithms.configs.modules._make_tdmpc2_mlp"
def _make_tdmpc2_mlp(
*,
in_features: int,
out_features: int,
depth: int | None,
num_cells: int | Sequence[int],
output_activation: nn.Module | None = None,
dropout: float = 0.0,
device: Any = None,
) -> MLP:
"""Helper function to create a TDMPC2 MLP."""
if isinstance(num_cells, int):
hidden_sizes = [num_cells] * (depth or 0)
else:
hidden_sizes = list(num_cells)
network = MLP(
in_features=in_features,
out_features=out_features,
depth=depth,
num_cells=num_cells,
activation_class=nn.Mish,
activation_kwargs={"inplace": False},
norm_class=nn.LayerNorm,
norm_kwargs=[{"normalized_shape": size} for size in hidden_sizes],
dropout=dropout or None,
activate_last_layer=False,
device=device,
)
# TDMPC2 only applies dropout after the first hidden layer.
# MLP applies it to all layers, so we disable all after the first.
if dropout:
dropout_layers = [
module for module in network.children() if isinstance(module, nn.Dropout)
]
for module in dropout_layers[1:]:
module.p = 0.0
if output_activation is not None:
network.append(nn.LayerNorm(out_features, device=device))
network.append(
output_activation.to(device=device)
if device is not None
else output_activation
)
for module in network.modules():
if isinstance(module, nn.Linear):
nn.init.trunc_normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
return network
[docs]
@dataclass
class DreamerV3ImageEncoderConfig(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.DreamerV3ImageEncoder`.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import DreamerV3ImageEncoderConfig
>>> cfg = DreamerV3ImageEncoderConfig(depth=8, mults=[1, 2])
>>> net = instantiate(cfg)
>>> image = torch.randint(0, 256, (4, 3, 16, 16), dtype=torch.uint8)
>>> assert net(image).shape == (4, 256)
.. seealso:: :class:`~torchrl.modules.DreamerV3ImageEncoder`
"""
in_channels: int = 3
depth: int = 64
mults: list[int] = field(default_factory=partial(list, (2, 3, 4, 4)))
kernel_size: int = 5
norm_eps: float = 1e-4
device: Any = None
_target_: str = "torchrl.modules.DreamerV3ImageEncoder"
[docs]
@dataclass
class DreamerV3ImageDecoderConfig(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.DreamerV3ImageDecoder`.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import DreamerV3ImageDecoderConfig
>>> cfg = DreamerV3ImageDecoderConfig(
... in_features=12, image_shape=[3, 16, 16], depth=8, mults=[1, 2], num_blocks=2
... )
>>> net = instantiate(cfg)
>>> assert net(torch.randn(4, 12)).shape == (4, 3, 16, 16)
.. seealso:: :class:`~torchrl.modules.DreamerV3ImageDecoder`
"""
in_features: int = MISSING
image_shape: list[int] = field(default_factory=partial(list, (3, 64, 64)))
depth: int = 64
mults: list[int] = field(default_factory=partial(list, (2, 3, 4, 4)))
kernel_size: int = 5
num_blocks: int = 8
norm_eps: float = 1e-4
device: Any = None
_target_: str = "torchrl.modules.DreamerV3ImageDecoder"
[docs]
@dataclass
class DreamerV3DiscreteActorConfig(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.DreamerV3DiscreteActor`.
Examples:
>>> import torch
>>> from hydra.utils import instantiate
>>> from tensordict import TensorDict
>>> from torchrl.trainers.algorithms.configs import DreamerV3DiscreteActorConfig
>>> actor = instantiate(DreamerV3DiscreteActorConfig(in_features=12, out_features=3))
>>> data = TensorDict({"state": torch.randn(4, 8), "belief": torch.randn(4, 4)}, [4])
>>> actor(data)["action"].shape
torch.Size([4, 3])
"""
in_features: int = MISSING
out_features: int = MISSING
depth: int = 3
num_cells: int = 1024
norm_eps: float = 1e-4
unimix: float = 0.01
in_keys: Any = None
action_key: Any = "action"
logits_key: Any = "logits"
log_prob_key: Any = "action_log_prob"
device: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_dreamer_v3_discrete_actor"
)
def _make_dreamer_v3_discrete_actor(**kwargs) -> DreamerV3DiscreteActor:
"""Normalize Hydra's nested key lists before constructing the actor."""
in_keys = _normalize_hydra_keys(kwargs.pop("in_keys", None))
for key in ("action_key", "logits_key", "log_prob_key"):
if key in kwargs:
kwargs[key] = _normalize_hydra_key(kwargs[key])
return DreamerV3DiscreteActor(in_keys=in_keys, **kwargs)
[docs]
@dataclass
class DreamerV3SeededPolicyConfig(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.DreamerV3SeededPolicy`.
Examples:
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import DreamerV3DiscreteActorConfig, DreamerV3SeededPolicyConfig
>>> config = DreamerV3SeededPolicyConfig(
... module=DreamerV3DiscreteActorConfig(in_features=6, out_features=3), seed=7,
... )
>>> policy = instantiate(config)
>>> policy.get_extra_state()
{'seed': 7, 'counter': 0}
"""
module: Any = MISSING
seed: int = MISSING
_target_: str = "torchrl.modules.DreamerV3SeededPolicy"
@dataclass
class RSSMStateEstimatorV3Config(NetworkConfig):
"""Hydra configuration for :class:`~torchrl.modules.RSSMStateEstimatorV3`.
Examples:
Given the shared prior and posterior in the estimator's example:
>>> from hydra.utils import instantiate
>>> from torchrl.trainers.algorithms.configs import RSSMStateEstimatorV3Config
>>> estimator = instantiate( # doctest: +SKIP
... RSSMStateEstimatorV3Config(), prior=prior, posterior=posterior,
... )
"""
prior: Any = MISSING
posterior: Any = MISSING
in_keys: Any = None
out_keys: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_rssm_state_estimator_v3"
)
def _make_rssm_state_estimator_v3(**kwargs) -> RSSMStateEstimatorV3:
"""Normalize configured nested keys before constructing the estimator."""
in_keys = _normalize_hydra_keys(kwargs.pop("in_keys", None))
out_keys = _normalize_hydra_keys(kwargs.pop("out_keys", None))
return RSSMStateEstimatorV3(in_keys=in_keys, out_keys=out_keys, **kwargs)
@dataclass
class NormConfig(ConfigBase):
"""A class to configure a normalization layer.
Defaults to :class:`torch.nn.BatchNorm1d`.
.. seealso:: :class:`torch.nn.BatchNorm1d`
"""
_target_: str = "torch.nn.BatchNorm1d"
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for normalization configurations."""
@dataclass
class AggregatorConfig(ConfigBase):
"""A class to configure an aggregator layer.
Defaults to :class:`torchrl.modules.models.utils.SquashDims`.
.. seealso:: :class:`torchrl.modules.models.utils.SquashDims`
"""
_target_: str = "torchrl.modules.models.utils.SquashDims"
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for aggregator configurations."""
[docs]
@dataclass
class ConvNetConfig(NetworkConfig):
"""A class to configure a convolutional network.
Defaults to :class:`torchrl.modules.ConvNet`.
Example:
>>> cfg = ConvNetConfig(in_features=3, depth=2, num_cells=[32, 64], kernel_sizes=[3, 5], strides=[1, 2], paddings=[1, 2])
>>> net = instantiate(cfg)
>>> y = net(torch.randn(1, 3, 32, 32))
>>> assert y.shape == (1, 64)
.. seealso:: :class:`torchrl.modules.ConvNet`
"""
in_features: int | None = None
depth: int | None = None
num_cells: Any = None
kernel_sizes: Any = 3
strides: Any = 1
paddings: Any = 0
activation_class: ActivationConfig = field(
default_factory=partial(
ActivationConfig, _target_="torch.nn.ELU", _partial_=True
)
)
activation_kwargs: Any = None
norm_class: NormConfig | None = None
norm_kwargs: Any = None
bias_last_layer: bool = True
aggregator_class: AggregatorConfig = field(
default_factory=partial(
AggregatorConfig,
_target_="torchrl.modules.models.utils.SquashDims",
_partial_=True,
)
)
aggregator_kwargs: dict | None = None
squeeze_output: bool = False
device: Any = None
_target_: str = "torchrl.modules.ConvNet"
def __post_init__(self):
if self.activation_class is None and isinstance(self.activation_class, str):
self.activation_class = ActivationConfig(
_target_=self.activation_class, _partial_=True
)
if self.norm_class is None and isinstance(self.norm_class, str):
self.norm_class = NormConfig(_target_=self.norm_class, _partial_=True)
if self.aggregator_class is None and isinstance(self.aggregator_class, str):
self.aggregator_class = AggregatorConfig(
_target_=self.aggregator_class, _partial_=True
)
@dataclass
class QMixerNetworkConfig(NetworkConfig):
"""A class to configure a QMIX mixer network.
.. seealso:: :class:`torchrl.modules.models.multiagent.QMixer`
"""
state_shape: Any = MISSING
mixing_embed_dim: int = 32
n_agents: int = MISSING
device: Any = None
_target_: str = "torchrl.modules.models.multiagent.QMixer"
def __post_init__(self) -> None:
super().__post_init__()
@dataclass
class VDNMixerNetworkConfig(NetworkConfig):
"""A class to configure a VDN mixer network.
.. seealso:: :class:`torchrl.modules.models.multiagent.VDNMixer`
"""
n_agents: int = MISSING
device: Any = None
_target_: str = "torchrl.modules.models.multiagent.VDNMixer"
def __post_init__(self) -> None:
super().__post_init__()
[docs]
@dataclass
class ModelConfig(ConfigBase):
"""Parent class to configure a model.
A model can be made of several networks. It is always a :class:`~tensordict.nn.TensorDictModuleBase` instance.
.. seealso:: :class:`TanhNormalModelConfig`, :class:`ValueModelConfig`
"""
_partial_: bool = False
in_keys: Any = None
out_keys: Any = None
shared: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for model configurations."""
[docs]
@dataclass
class TdMpc2WorldModelConfig(ConfigBase):
"""Configuration for a TD-MPC2 world model.
The resulting :class:`~torchrl.modules.WorldModel` encodes observations,
predicts the next latent state from the current latent state and action,
and predicts distributional rewards from the current latent state and
action. TensorDict routing is configured with the specialized key fields
below. Set ``shared=True`` to place the constructed module in shared
memory.
.. seealso:: :class:`~torchrl.modules.WorldModel`
"""
observation_dim: int = MISSING
action_dim: int = MISSING
latent_dim: int = MISSING
encoder_dim: int = 256
encoder_depth: int = 1
mlp_dim: int = 512
simnorm_dim: int = 8
num_bins: int = 101
observation_key: Any = "observation"
action_key: Any = "action"
latent_key: Any = "latent"
reward_logits_key: Any = "reward_logits"
device: Any = None
shared: bool = False
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tdmpc2_world_model"
)
def __post_init__(self) -> None:
"""Post-initialization hook for the TD-MPC2 world model config."""
def _next_key(key: Any) -> Any:
if isinstance(key, tuple):
return ("next", *key)
return ("next", key)
def _make_tdmpc2_world_model(
*,
observation_dim: int,
action_dim: int,
latent_dim: int,
encoder_dim: int = 256,
encoder_depth: int = 1,
mlp_dim: int = 512,
simnorm_dim: int = 8,
num_bins: int = 101,
observation_key: Any = "observation",
action_key: Any = "action",
latent_key: Any = "latent",
reward_logits_key: Any = "reward_logits",
device: Any = None,
shared: bool = False,
) -> WorldModel:
"""Build the state-based encoder, dynamics, and reward components of TD-MPC2."""
if observation_dim <= 0 or action_dim <= 0 or latent_dim <= 0:
raise ValueError("observation_dim, action_dim, and latent_dim must be positive")
if encoder_dim <= 0 or mlp_dim <= 0:
raise ValueError("encoder_dim and mlp_dim must be positive")
if simnorm_dim <= 0 or latent_dim % simnorm_dim:
raise ValueError("simnorm_dim must be a positive divisor of latent_dim")
if num_bins <= 1:
raise ValueError("num_bins must be greater than 1")
observation_key = _normalize_hydra_key(observation_key)
action_key = _normalize_hydra_key(action_key)
latent_key = _normalize_hydra_key(latent_key)
reward_logits_key = _normalize_hydra_key(reward_logits_key)
encoder = TensorDictModule(
_make_tdmpc2_mlp(
in_features=observation_dim,
out_features=latent_dim,
depth=encoder_depth,
num_cells=encoder_dim,
output_activation=SimplicialNormalization(simnorm_dim),
device=device,
),
in_keys=[observation_key],
out_keys=[latent_key],
)
dynamics = TensorDictModule(
_make_tdmpc2_mlp(
in_features=latent_dim + action_dim,
out_features=latent_dim,
depth=2,
num_cells=mlp_dim,
output_activation=SimplicialNormalization(simnorm_dim),
device=device,
),
in_keys=[latent_key, action_key],
out_keys=[_next_key(latent_key)],
)
reward = TensorDictModule(
_make_tdmpc2_mlp(
in_features=latent_dim + action_dim,
out_features=max(num_bins, 1),
depth=2,
num_cells=mlp_dim,
device=device,
),
in_keys=[latent_key, action_key],
out_keys=[_next_key(reward_logits_key)],
)
with torch.no_grad():
reward.module[-1].weight.zero_()
world_model = WorldModel(encoder, dynamics, reward)
if shared:
world_model = world_model.share_memory()
return world_model
[docs]
@dataclass
class TdMpc2QEnsembleConfig(ModelConfig):
"""Configuration for a TD-MPC2 Q-function ensemble.
The resulting
:class:`~torchrl.modules.TdMpc2QEnsemble` maps a latent
state and action to distributional Q-function logits. Its :meth:`forward`
method writes all ensemble logits, while
:meth:`~torchrl.modules.TdMpc2QEnsemble.reduce` decodes two randomly
selected Q-functions and writes a reduced value.
Args:
latent_dim: Size of the latent state input.
action_dim: Number of action dimensions.
mlp_dim: Width of each hidden layer in every Q-function.
num_q: Number of Q-functions in the ensemble.
num_bins: Number of categorical bins for the Q-function output. Must be
greater than 1.
vmin: Minimum value of the symlog-space categorical support.
vmax: Maximum value of the symlog-space categorical support.
dropout: Dropout probability applied in the Q-functions.
q_value_key: TensorDict key written by the ensemble reduction.
device: Device on which to construct the Q-functions.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from tensordict import TensorDict
>>> from torchrl.trainers.algorithms.configs import TdMpc2QEnsembleConfig
>>> cfg = TdMpc2QEnsembleConfig(latent_dim=8, action_dim=2, mlp_dim=16)
>>> q_ensemble = instantiate(cfg)
>>> data = TensorDict(
... {"latent": torch.randn(3, 8), "action": torch.randn(3, 2)},
... batch_size=[3],
... )
>>> q_ensemble(data)["q_logits"].shape
torch.Size([3, 5, 101])
.. seealso:: :class:`~torchrl.modules.TdMpc2QEnsemble`
"""
latent_dim: int = MISSING
action_dim: int = MISSING
mlp_dim: int = 512
num_q: int = 5
num_bins: int = 101
vmin: float = -10.0
vmax: float = 10.0
dropout: float = 0.01
q_value_key: Any = "q_value"
device: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tdmpc2_q_ensemble"
)
def __post_init__(self) -> None:
super().__post_init__()
if self.in_keys is None:
self.in_keys = ["latent", "action"]
if self.out_keys is None:
self.out_keys = ["q_logits"]
def _make_tdmpc2_q_ensemble(
*,
latent_dim: int,
action_dim: int,
mlp_dim: int = 512,
num_q: int = 5,
num_bins: int = 101,
vmin: float = -10.0,
vmax: float = 10.0,
dropout: float = 0.01,
q_value_key: Any = "q_value",
device: Any = None,
in_keys: Any = None,
out_keys: Any = None,
shared: bool = False,
) -> TdMpc2QEnsemble:
"""Build the TD-MPC2 Q-function ensemble."""
if latent_dim <= 0 or action_dim <= 0 or mlp_dim <= 0:
raise ValueError("latent_dim, action_dim, and mlp_dim must be positive")
if num_bins <= 1:
raise ValueError("num_bins must be greater than 1")
if not 0 <= dropout <= 1:
raise ValueError("dropout must be in the interval [0, 1]")
if in_keys is None:
in_keys = ["latent", "action"]
if out_keys is None:
out_keys = ["q_logits"]
in_keys = _normalize_hydra_keys(in_keys)
out_keys = _normalize_hydra_keys(out_keys)
q_value_key = _normalize_hydra_key(q_value_key)
q_networks = [
_make_tdmpc2_mlp(
in_features=latent_dim + action_dim,
out_features=num_bins,
depth=2,
num_cells=mlp_dim,
dropout=dropout,
device=device,
)
for _ in range(num_q)
]
with torch.no_grad():
for q_network in q_networks:
q_network[-1].weight.zero_()
q_ensemble = TdMpc2QEnsemble(
q_networks,
num_bins=num_bins,
vmin=vmin,
vmax=vmax,
in_keys=in_keys,
out_keys=out_keys,
q_value_key=q_value_key,
)
if shared:
q_ensemble = q_ensemble.share_memory()
return q_ensemble
[docs]
@dataclass
class TdMpc2PolicyPriorConfig(ModelConfig):
"""Configuration for a TD-MPC2 policy prior.
The policy prior maps a latent state to a sampled action and the
distribution statistics used by the TD-MPC2 objective. By default, the
module reads ``"latent"`` and writes ``"action"``, ``"mean"`,
``"log_std"``, ``"entropy"``, and ``"scaled_entropy"``.
Example:
>>> import torch
>>> from hydra.utils import instantiate
>>> from tensordict import TensorDict
>>> from torchrl.trainers.algorithms.configs import TdMpc2PolicyPriorConfig
>>> cfg = TdMpc2PolicyPriorConfig(latent_dim=8, action_dim=2, mlp_dim=16)
>>> policy = instantiate(cfg)
>>> td = TensorDict({"latent": torch.randn(3, 8)}, batch_size=[3])
>>> assert policy(td)["action"].shape == (3, 2)
"""
latent_dim: int = MISSING
action_dim: int = MISSING
mlp_dim: int = 512
log_std_min: float = -10.0
log_std_max: float = 2.0
device: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tdmpc2_policy_prior"
)
def __post_init__(self) -> None:
super().__post_init__()
if self.in_keys is None:
self.in_keys = ["latent"]
if self.out_keys is None:
self.out_keys = [
"action",
"mean",
"log_std",
"entropy",
"scaled_entropy",
]
def _make_tdmpc2_policy_prior(
*,
latent_dim: int,
action_dim: int,
mlp_dim: int = 512,
log_std_min: float = -10.0,
log_std_max: float = 2.0,
device: Any = None,
in_keys: Any = None,
out_keys: Any = None,
shared: bool = False,
) -> TensorDictModule:
"""Build a TD-MPC2 policy prior module."""
if latent_dim <= 0 or action_dim <= 0 or mlp_dim <= 0:
raise ValueError("latent_dim, action_dim, and mlp_dim must be positive")
if log_std_max <= log_std_min:
raise ValueError("log_std_max must be greater than log_std_min")
if in_keys is None:
in_keys = ["latent"]
if out_keys is None:
out_keys = [
"action",
"mean",
"log_std",
"entropy",
"scaled_entropy",
]
in_keys = _normalize_hydra_keys(in_keys)
out_keys = _normalize_hydra_keys(out_keys)
if len(in_keys) != 1:
raise ValueError("TdMpc2PolicyPrior requires exactly one input key")
if len(out_keys) != 5:
raise ValueError(
"TdMpc2PolicyPrior requires five output keys: action, mean, "
"log_std, entropy, and scaled_entropy"
)
network = _make_tdmpc2_mlp(
in_features=latent_dim,
out_features=2 * action_dim,
depth=2,
num_cells=mlp_dim,
device=device,
)
policy_prior = TensorDictModule(
_TdMpc2PolicyPrior(
network,
log_std_min=log_std_min,
log_std_max=log_std_max,
device=device,
),
in_keys=in_keys,
out_keys=out_keys,
)
if shared:
policy_prior = policy_prior.share_memory()
return policy_prior
[docs]
@dataclass
class TdMpc2PlannerConfig(ConfigBase):
"""Configuration for the TD-MPC2 action planner."""
world_model: Any = None
policy_prior: Any = None
q_ensemble: Any = None
horizon: int | None = None
discount: float | None = None
num_samples: int = 512
num_elites: int = 64
num_pi_trajs: int = 24
iterations: int = 6
min_std: float = 0.05
max_std: float = 2.0
temperature: float = 0.5
observation_key: Any = None
action_key: Any = None
is_init_key: Any = "is_init"
prev_mean_key: Any = "_tdmpc2_prev_mean"
action_dim: int | None = None
_partial_: bool = False
_target_: str = "torchrl.trainers.algorithms.configs.modules._make_tdmpc2_planner"
def __post_init__(self) -> None:
if self.horizon is not None and self.horizon <= 0:
raise ValueError(f"horizon must be positive, got {self.horizon}.")
if self.discount is not None and not 0 <= self.discount <= 1:
raise ValueError(f"discount must be in [0, 1], got {self.discount}.")
def _make_tdmpc2_planner(
*,
world_model: Any,
policy_prior: Any,
q_ensemble: Any,
horizon: int | None = None,
discount: float | None = None,
num_samples: int = 512,
num_elites: int = 64,
num_pi_trajs: int = 24,
iterations: int = 6,
min_std: float = 0.05,
max_std: float = 2.0,
temperature: float = 0.5,
observation_key: Any = None,
action_key: Any = None,
is_init_key: Any = "is_init",
prev_mean_key: Any = "_tdmpc2_prev_mean",
action_dim: int | None = None,
) -> TdMpc2Planner:
"""Build a TD-MPC2 planner from shared live learner components."""
components = {
"world_model": world_model,
"policy_prior": policy_prior,
"q_ensemble": q_ensemble,
}
for name, component in components.items():
if not isinstance(component, torch.nn.Module) and hasattr(
component, "_target_"
):
components[name] = instantiate(component)
missing = [name for name, component in components.items() if component is None]
if missing:
raise TypeError(
"TdMpc2PlannerConfig requires live component modules for "
+ ", ".join(missing)
+ "; the trainer factory should inject these automatically."
)
if horizon is None:
horizon = 3
if discount is None:
discount = 0.99
return TdMpc2Planner(
**components,
horizon=horizon,
discount=discount,
num_samples=num_samples,
num_elites=num_elites,
num_pi_trajs=num_pi_trajs,
iterations=iterations,
min_std=min_std,
max_std=max_std,
temperature=temperature,
observation_key=(
None if observation_key is None else _normalize_hydra_key(observation_key)
),
action_key=(None if action_key is None else _normalize_hydra_key(action_key)),
is_init_key=_normalize_hydra_key(is_init_key),
prev_mean_key=_normalize_hydra_key(prev_mean_key),
action_dim=action_dim,
)
[docs]
@dataclass
class TensorDictModuleConfig(ModelConfig):
"""A class to configure a TensorDictModule.
Example:
>>> cfg = TensorDictModuleConfig(module=MLPConfig(in_features=10, out_features=10, depth=2, num_cells=32), in_keys=["observation"], out_keys=["action"])
>>> module = instantiate(cfg)
>>> assert isinstance(module, TensorDictModule)
>>> assert module(observation=torch.randn(10, 10)).shape == (10, 10)
.. seealso:: :class:`tensordict.nn.TensorDictModule`
"""
module: MLPConfig = MISSING
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tensordict_module"
)
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for TensorDict module configurations."""
return super().__post_init__()
[docs]
@dataclass
class TensorDictSequentialConfig(ModelConfig):
"""A class to configure a TensorDictSequential.
Example:
>>> cfg = TensorDictSequentialConfig(
... modules=[
... TensorDictModuleConfig(module=MLPConfig(in_features=10, out_features=10, depth=2, num_cells=32), in_keys=["observation"], out_keys=["hidden"]),
... TensorDictModuleConfig(module=MLPConfig(in_features=10, out_features=5, depth=2, num_cells=32), in_keys=["hidden"], out_keys=["action"])
... ]
... )
>>> seq = instantiate(cfg)
>>> assert isinstance(seq, TensorDictSequential)
.. seealso:: :class:`tensordict.nn.TensorDictSequential`
"""
modules: Any | None = None
partial_tolerant: bool = False
selected_out_keys: Any | None = None
inplace: bool | str | None = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tensordict_sequential"
)
_partial_: bool = False
def __post_init__(self) -> None:
return super().__post_init__()
[docs]
@dataclass
class TanhNormalModelConfig(ModelConfig):
"""A class to configure a TanhNormal model.
Example:
>>> cfg = TanhNormalModelConfig(network=MLPConfig(in_features=10, out_features=5, depth=2, num_cells=32))
>>> net = instantiate(cfg)
>>> y = net(torch.randn(1, 10))
>>> assert y.shape == (1, 5)
Args:
low: lower bound of the action support handed to
:class:`~torchrl.modules.TanhNormal` (a scalar or a per-dimension
sequence). Defaults to ``None``, i.e. the distribution default of ``-1``.
high: upper bound of the action support. Defaults to ``None``, i.e. ``1``.
tanh_loc: if ``True``, the location is squashed to ``[-upscale, upscale]``
before the tanh transform, which keeps the log-probability of actions
at the bounds finite (see :class:`~torchrl.modules.TanhNormal`).
Defaults to ``False``.
.. seealso:: :class:`torchrl.modules.TanhNormal`
"""
network: MLPConfig = MISSING
eval_mode: bool = False
extract_normal_params: bool = True
scale_mapping: str = "biased_softplus_1.0"
scale_lb: float = 1e-4
low: Any = None
high: Any = None
tanh_loc: bool = False
param_keys: Any = None
exploration_type: Any = "RANDOM"
return_log_prob: bool = False
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_tanh_normal_model"
)
def __post_init__(self):
"""Post-initialization hook for TanhNormal model configurations."""
super().__post_init__()
if self.in_keys is None:
self.in_keys = ["observation"]
if self.param_keys is None:
self.param_keys = ["loc", "scale"]
if self.out_keys is None:
self.out_keys = ["action"]
[docs]
@dataclass
class ValueModelConfig(ModelConfig):
"""A class to configure a Value model.
Example:
>>> cfg = ValueModelConfig(network=MLPConfig(in_features=10, out_features=5, depth=2, num_cells=32))
>>> net = instantiate(cfg)
>>> y = net(torch.randn(1, 10))
>>> assert y.shape == (1, 5)
.. seealso:: :class:`torchrl.modules.ValueOperator`
"""
_target_: str = "torchrl.trainers.algorithms.configs.modules._make_value_model"
network: NetworkConfig = MISSING
def __post_init__(self) -> None:
"""Post-initialization hook for value model configurations."""
super().__post_init__()
[docs]
@dataclass
class TanhModuleConfig(ModelConfig):
"""A class to configure a TanhModule.
Example:
>>> cfg = TanhModuleConfig(in_keys=["action"], out_keys=["action"], low=-1.0, high=1.0)
>>> module = instantiate(cfg)
>>> assert isinstance(module, TanhModule)
.. seealso:: :class:`torchrl.modules.TanhModule`
"""
spec: Any = None
low: Any = None
high: Any = None
clamp: bool = False
_target_: str = "torchrl.trainers.algorithms.configs.modules._make_tanh_module"
def __post_init__(self) -> None:
"""Post-initialization hook for TanhModule configurations."""
super().__post_init__()
[docs]
@dataclass
class AdditiveGaussianModuleConfig(ModelConfig):
"""A class to configure an AdditiveGaussianModule.
Example:
>>> cfg = AdditiveGaussianModuleConfig(
... spec=None,
... sigma_init=1.0,
... sigma_end=0.1,
... mean=0.0,
... std=1.0,
... action_key="action",
... )
>>> module = instantiate(cfg)
>>> assert isinstance(module, AdditiveGaussianModule)
.. seealso:: :class:`torchrl.modules.AdditiveGaussianModule`
"""
spec: Any = None
sigma_init: float = 1.0
sigma_end: float = 0.1
annealing_num_steps: int = 1000
mean: float = 0.0
std: float = 1.0
action_key: Any = "action"
safe: bool = False
device: Any = None
_target_: str = (
"torchrl.trainers.algorithms.configs.modules._make_additive_gaussian_module"
)
_partial_: bool = False
def __post_init__(self) -> None:
super().__post_init__()
def _make_tensordict_module(*args, **kwargs) -> TensorDictModule:
"""Helper function to create a TensorDictModule."""
from hydra.utils import instantiate
module = kwargs.pop("module")
shared = kwargs.pop("shared", False)
for key in ("in_keys", "out_keys"):
if key in kwargs:
kwargs[key] = _normalize_hydra_keys(kwargs[key])
# Instantiate the module if it's a config
if hasattr(module, "_target_"):
module = instantiate(module)
elif callable(module) and hasattr(module, "func"): # partial function
module = module()
# Create the TensorDictModule
tensordict_module = TensorDictModule(module, **kwargs)
# Apply share_memory if needed
if shared:
tensordict_module = tensordict_module.share_memory()
return tensordict_module
def _make_tensordict_sequential(*args, **kwargs) -> TensorDictSequential:
"""Helper function to create a TensorDictSequential."""
from hydra.utils import instantiate
from omegaconf import DictConfig, ListConfig
modules = kwargs.pop("modules")
shared = kwargs.pop("shared", False)
partial_tolerant = kwargs.pop("partial_tolerant", False)
selected_out_keys = _normalize_hydra_keys(kwargs.pop("selected_out_keys", None))
inplace = kwargs.pop("inplace", None)
def _instantiate_module(module):
if hasattr(module, "_target_"):
return instantiate(module)
elif callable(module) and hasattr(module, "func"):
return module()
else:
return module
if isinstance(modules, (dict, DictConfig)):
instantiated_modules = {
key: _instantiate_module(module) for key, module in modules.items()
}
elif isinstance(modules, (list, ListConfig)):
instantiated_modules = [_instantiate_module(module) for module in modules]
else:
raise ValueError(
f"modules must be a dict or list, got {type(modules).__name__}"
)
tensordict_sequential = TensorDictSequential(
instantiated_modules,
partial_tolerant=partial_tolerant,
selected_out_keys=selected_out_keys,
inplace=inplace,
)
if shared:
tensordict_sequential = tensordict_sequential.share_memory()
return tensordict_sequential
def _make_tanh_normal_model(*args, **kwargs):
"""Helper function to create a TanhNormal model with ProbabilisticTensorDictSequential."""
from hydra.utils import instantiate
from tensordict.nn import (
ProbabilisticTensorDictModule,
ProbabilisticTensorDictSequential,
TensorDictModule,
)
from torchrl.modules import NormalParamExtractor, TanhNormal
# Extract parameters
network = kwargs.pop("network")
in_keys = _normalize_hydra_keys(kwargs.pop("in_keys", ["observation"]))
param_keys = _normalize_hydra_keys(kwargs.pop("param_keys", ["loc", "scale"]))
out_keys = _normalize_hydra_keys(kwargs.pop("out_keys", ["action"]))
extract_normal_params = kwargs.pop("extract_normal_params", True)
scale_mapping = kwargs.pop("scale_mapping", "biased_softplus_1.0")
scale_lb = kwargs.pop("scale_lb", 1e-4)
return_log_prob = kwargs.pop("return_log_prob", False)
eval_mode = kwargs.pop("eval_mode", False)
exploration_type = kwargs.pop("exploration_type", "RANDOM")
shared = kwargs.pop("shared", False)
distribution_kwargs = dict(kwargs.pop("distribution_kwargs", None) or {})
for bound in ("low", "high"):
value = kwargs.pop(bound, None)
if value is None:
continue
# omegaconf hands sequences over as ListConfig
if not isinstance(value, (int, float)):
value = torch.as_tensor(list(value), dtype=torch.get_default_dtype())
distribution_kwargs[bound] = value
if kwargs.pop("tanh_loc", False):
distribution_kwargs["tanh_loc"] = True
if distribution_kwargs:
kwargs["distribution_kwargs"] = distribution_kwargs
# Now instantiate the network
if hasattr(network, "_target_"):
network = instantiate(network)
elif callable(network) and hasattr(network, "func"): # partial function
network = network()
# Create the sequential
if extract_normal_params:
# Add NormalParamExtractor to split the output
network = torch.nn.Sequential(
network,
NormalParamExtractor(scale_mapping=scale_mapping, scale_lb=scale_lb),
)
module = TensorDictModule(network, in_keys=in_keys, out_keys=param_keys)
if shared:
module = module.share_memory()
# Create ProbabilisticTensorDictModule
prob_module = ProbabilisticTensorDictModule(
in_keys=param_keys,
out_keys=out_keys,
distribution_class=TanhNormal,
return_log_prob=return_log_prob,
default_interaction_type=exploration_type,
**kwargs,
)
result = ProbabilisticTensorDictSequential(module, prob_module)
if eval_mode:
result.eval()
return result
def _make_value_model(*args, **kwargs) -> ValueOperator:
"""Helper function to create a ValueOperator with the given network."""
from hydra.utils import instantiate
network = kwargs.pop("network")
shared = kwargs.pop("shared", False)
for key in ("in_keys", "out_keys"):
if key in kwargs:
kwargs[key] = _normalize_hydra_keys(kwargs[key])
# Instantiate the network if it's a config
if hasattr(network, "_target_"):
network = instantiate(network)
elif callable(network) and hasattr(network, "func"): # partial function
network = network()
# Create the ValueOperator
value_operator = ValueOperator(network, **kwargs)
# Apply share_memory if needed
if shared:
value_operator = value_operator.share_memory()
return value_operator
def _make_tanh_module(*args, **kwargs) -> TanhModule:
"""Helper function to create a TanhModule."""
kwargs.pop("shared", False)
if "in_keys" in kwargs:
kwargs["in_keys"] = _normalize_hydra_keys(kwargs["in_keys"])
if "out_keys" in kwargs:
kwargs["out_keys"] = _normalize_hydra_keys(kwargs["out_keys"])
return TanhModule(**kwargs)
def _make_additive_gaussian_module(*args, **kwargs) -> AdditiveGaussianModule:
"""Helper function to create an AdditiveGaussianModule."""
kwargs.pop("shared", False)
kwargs.pop("in_keys", None)
kwargs.pop("out_keys", None)
if "action_key" in kwargs:
kwargs["action_key"] = _normalize_hydra_key(kwargs["action_key"])
return AdditiveGaussianModule(**kwargs)
[docs]
@dataclass
class QValueModelConfig(ModelConfig):
"""A class to configure a QValueActor model.
.. seealso:: :class:`torchrl.modules.QValueActor`
"""
_target_: str = "torchrl.trainers.algorithms.configs.modules._make_qvalue_model"
network: NetworkConfig = MISSING
action_space: Any = None
action_key: Any = None
action_value_key: Any = None
chosen_action_value_key: Any = None
action_mask_key: Any = None
def __post_init__(self) -> None:
super().__post_init__()
def _make_qvalue_model(*args, **kwargs) -> QValueActor:
"""Helper function to create a QValueActor with the given network."""
from hydra.utils import instantiate
network = kwargs.pop("network")
shared = kwargs.pop("shared", False)
kwargs.pop("out_keys", None)
if "in_keys" in kwargs:
kwargs["in_keys"] = _normalize_hydra_keys(kwargs["in_keys"])
for key in (
"action_key",
"action_value_key",
"chosen_action_value_key",
"action_mask_key",
):
if key in kwargs:
kwargs[key] = _normalize_hydra_key(kwargs[key])
if hasattr(network, "_target_"):
network = instantiate(network)
elif callable(network) and hasattr(network, "func"):
network = network()
qvalue_actor = QValueActor(network, **kwargs)
if shared:
qvalue_actor = qvalue_actor.share_memory()
return qvalue_actor
