# 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, field
from functools import partial
from typing import Any
import torch
from omegaconf import MISSING
from tensordict.nn import TensorDictModule, TensorDictSequential
from torchrl.modules import (
AdditiveGaussianModule,
DreamerV3DiscreteActor,
LowLevelController,
QValueActor,
RSSMStateEstimatorV3,
TanhModule,
ValueOperator,
)
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 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 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
