torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/modules/tensordict_module/sequence.py

@@ -0,0 +1,132 @@
+# 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 tensordict.nn import TensorDictModule, TensorDictSequential
+from torch import nn
+
+from torchrl.data.tensor_specs import Composite
+from torchrl.modules.tensordict_module.common import SafeModule
+
+
+class SafeSequential(TensorDictSequential, SafeModule):
+    """A safe sequence of TensorDictModules.
+
+    Similarly to :obj:`nn.Sequence` which passes a tensor through a chain of mappings that read and write a single tensor
+    each, this module will read and write over a tensordict by querying each of the input modules.
+    When calling a :obj:`TensorDictSequential` instance with a functional module, it is expected that the parameter lists (and
+    buffers) will be concatenated in a single list.
+
+    Args:
+         modules (iterable of TensorDictModules): ordered sequence of TensorDictModule instances to be run sequentially.
+         partial_tolerant (bool, optional): if ``True``, the input tensordict can miss some of the input keys.
+            If so, the only modules that will be executed are those which can be executed given the keys that
+            are present.
+            Also, if the input tensordict is a lazy stack of tensordicts AND if partial_tolerant is ``True`` AND if the
+            stack does not have the required keys, then SafeSequential will scan through the sub-tensordicts
+            looking for those that have the required keys, if any.
+        inplace (bool or str, optional): if `True`, the input tensordict is modified in-place. If `False`, a new empty
+            :class:`~tensordict.TensorDict` instance is created. If `"empty"`, `input.empty()` is used instead (ie, the
+            output preserves type, device and batch-size). Defaults to `None` (relies on sub-modules).
+
+    TensorDictSequence supports functional, modular and vmap coding:
+    Examples:
+        >>> import torch
+        >>> from tensordict import TensorDict
+        >>> from torchrl.data import Composite, Unbounded
+        >>> from torchrl.modules import TanhNormal, SafeSequential, TensorDictModule, NormalParamExtractor
+        >>> from torchrl.modules.tensordict_module import SafeProbabilisticModule
+        >>> td = TensorDict({"input": torch.randn(3, 4)}, [3,])
+        >>> spec1 = Composite(hidden=Unbounded(4), loc=None, scale=None)
+        >>> net1 = nn.Sequential(torch.nn.Linear(4, 8), NormalParamExtractor())
+        >>> module1 = TensorDictModule(net1, in_keys=["input"], out_keys=["loc", "scale"])
+        >>> td_module1 = SafeProbabilisticModule(
+        ...     module=module1,
+        ...     spec=spec1,
+        ...     in_keys=["loc", "scale"],
+        ...     out_keys=["hidden"],
+        ...     distribution_class=TanhNormal,
+        ...     return_log_prob=True,
+        ... )
+        >>> spec2 = Unbounded(8)
+        >>> module2 = torch.nn.Linear(4, 8)
+        >>> td_module2 = TensorDictModule(
+        ...    module=module2,
+        ...    spec=spec2,
+        ...    in_keys=["hidden"],
+        ...    out_keys=["output"],
+        ...    )
+        >>> td_module = SafeSequential(td_module1, td_module2)
+        >>> params = TensorDict.from_module(td_module)
+        >>> with params.to_module(td_module):
+        ...     td_module(td)
+        >>> print(td)
+        TensorDict(
+            fields={
+                hidden: Tensor(torch.Size([3, 4]), dtype=torch.float32),
+                input: Tensor(torch.Size([3, 4]), dtype=torch.float32),
+                loc: Tensor(torch.Size([3, 4]), dtype=torch.float32),
+                output: Tensor(torch.Size([3, 8]), dtype=torch.float32),
+                sample_log_prob: Tensor(torch.Size([3, 1]), dtype=torch.float32),
+                scale: Tensor(torch.Size([3, 4]), dtype=torch.float32)},
+            batch_size=torch.Size([3]),
+            device=None,
+            is_shared=False)
+        >>> # The module spec aggregates all the input specs:
+        >>> print(td_module.spec)
+        Composite(
+            hidden: UnboundedContinuous(
+                shape=torch.Size([4]), space=None, device=cpu, dtype=torch.float32, domain=continuous),
+            loc: None,
+            scale: None,
+            output: UnboundedContinuous(
+                shape=torch.Size([8]), space=None, device=cpu, dtype=torch.float32, domain=continuous))
+
+    In the vmap case:
+        >>> from torch import vmap
+        >>> params = params.expand(4, *params.shape)
+        >>> td_vmap = vmap(td_module, (None, 0))(td, params)
+        >>> print(td_vmap)
+        TensorDict(
+            fields={
+                hidden: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32),
+                input: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32),
+                loc: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32),
+                output: Tensor(torch.Size([4, 3, 8]), dtype=torch.float32),
+                sample_log_prob: Tensor(torch.Size([4, 3, 1]), dtype=torch.float32),
+                scale: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32)},
+            batch_size=torch.Size([4, 3]),
+            device=None,
+            is_shared=False)
+
+    """
+
+    module: nn.ModuleList
+
+    def __init__(
+        self,
+        *modules: TensorDictModule,
+        partial_tolerant: bool = False,
+        inplace: bool | str | None = None,
+    ):
+        self.partial_tolerant = partial_tolerant
+
+        in_keys, out_keys = self._compute_in_and_out_keys(modules)
+
+        spec = Composite()
+        for module in modules:
+            try:
+                spec.update(module.spec)
+            except AttributeError:
+                spec.update(Composite({key: None for key in module.out_keys}))
+
+        super(TensorDictSequential, self).__init__(
+            spec=spec,
+            module=nn.ModuleList(list(modules)),
+            in_keys=in_keys,
+            out_keys=out_keys,
+            inplace=inplace,
+        )

torchrl/modules/tensordict_module/world_models.py

@@ -0,0 +1,34 @@
+# 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 tensordict.nn import TensorDictModule, TensorDictSequential
+
+
+class WorldModelWrapper(TensorDictSequential):
+    """World model wrapper.
+
+    This module wraps together a transition model and a reward model.
+    The transition model is used to predict an imaginary world state.
+    The reward model is used to predict the reward of the imagined transition.
+
+    Args:
+        transition_model (TensorDictModule): a transition model that generates a new world states.
+        reward_model (TensorDictModule): a reward model, that reads the world state and returns a reward.
+
+    """
+
+    def __init__(
+        self, transition_model: TensorDictModule, reward_model: TensorDictModule
+    ):
+        super().__init__(transition_model, reward_model)
+
+    def get_transition_model_operator(self) -> TensorDictModule:
+        """Returns a transition operator that maps either an observation to a world state or a world state to the next world state."""
+        return self.module[0]
+
+    def get_reward_operator(self) -> TensorDictModule:
+        """Returns a reward operator that maps a world state to a reward."""
+        return self.module[1]

torchrl/modules/utils/__init__.py

@@ -0,0 +1,38 @@
+# 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 collections import OrderedDict
+
+import torch
+from packaging import version
+
+
+if version.parse(torch.__version__) >= version.parse("1.12.0"):
+    from torch.nn.parameter import _ParameterMeta
+else:
+    pass
+
+    # Metaclass to combine _TensorMeta and the instance check override for Parameter.
+    class _ParameterMeta(torch._C._TensorMeta):
+        # Make `isinstance(t, Parameter)` return True for custom tensor instances that have the _is_param flag.
+        def __instancecheck__(self, instance):
+            return super().__instancecheck__(instance) or (
+                isinstance(instance, torch.Tensor)
+                and getattr(instance, "_is_param", False)
+            )
+
+
+from .mappings import biased_softplus, inv_softplus, mappings
+from .utils import get_primers_from_module
+
+__all__ = [
+    "OrderedDict",
+    "torch",
+    "version",
+    "biased_softplus",
+    "inv_softplus",
+    "mappings",
+    "get_primers_from_module",
+]

torchrl/modules/utils/mappings.py

@@ -0,0 +1,9 @@
+# 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 tensordict.nn.utils import biased_softplus, expln, inv_softplus, mappings
+
+__all__ = ["biased_softplus", "expln", "inv_softplus", "mappings"]

torchrl/modules/utils/utils.py

@@ -0,0 +1,89 @@
+# 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
+
+import warnings
+
+import torch
+from tensordict.utils import expand_as_right
+
+
+def get_primers_from_module(module):
+    """Get all tensordict primers from all submodules of a module.
+
+    This method is useful for retrieving primers from modules that are contained within a
+    parent module.
+
+    Args:
+        module (torch.nn.Module): The parent module.
+
+    Returns:
+        TensorDictPrimer: A TensorDictPrimer Transform.
+
+    Example:
+        >>> from torchrl.modules.utils import get_primers_from_module
+        >>> from torchrl.modules import GRUModule, MLP
+        >>> from tensordict.nn import TensorDictModule, TensorDictSequential
+        >>> # Define a GRU module
+        >>> gru_module = GRUModule(
+        ...     input_size=10,
+        ...     hidden_size=10,
+        ...     num_layers=1,
+        ...     in_keys=["input", "recurrent_state", "is_init"],
+        ...     out_keys=["features", ("next", "recurrent_state")],
+        ... )
+        >>> # Define a head module
+        >>> head = TensorDictModule(
+        ...     MLP(
+        ...         in_features=10,
+        ...         out_features=10,
+        ...         num_cells=[],
+        ...     ),
+        ...     in_keys=["features"],
+        ...     out_keys=["output"],
+        ... )
+        >>> # Create a sequential model
+        >>> model = TensorDictSequential(gru_module, head)
+        >>> # Retrieve primers from the model
+        >>> primers = get_primers_from_module(model)
+        >>> print(primers)
+
+        TensorDictPrimer(primers=Composite(
+            recurrent_state: UnboundedContinuous(
+                shape=torch.Size([1, 10]),
+                space=None,
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous), device=None, shape=torch.Size([])), default_value={'recurrent_state': 0.0}, random=None)
+
+    """
+    primers = []
+
+    def make_primers(submodule):
+        if hasattr(submodule, "make_tensordict_primer"):
+            primers.append(submodule.make_tensordict_primer())
+
+    module.apply(make_primers)
+    if not primers:
+        warnings.warn("No primers found in the module.")
+        return
+    elif len(primers) == 1:
+        return primers[0]
+    else:
+        from torchrl.envs.transforms import Compose
+
+        return Compose(*primers)
+
+
+def _unpad_tensors(tensors, mask, as_nested: bool = True) -> torch.Tensor:
+    shape = tensors.shape[2:]
+    mask = expand_as_right(mask.bool(), tensors)
+    nelts = mask.sum(-1)
+    while nelts.dim() > 1:
+        nelts = nelts.sum(-1)
+    vals = [t.view(-1, *shape) for t in tensors[mask].split(nelts.tolist(), dim=0)]
+    if as_nested:
+        return torch.nested.as_nested_tensor(vals)
+    return vals

torchrl/objectives/__init__.py

@@ -0,0 +1,78 @@
+# 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 torchrl.objectives.a2c import A2CLoss
+from torchrl.objectives.common import add_random_module, LossModule
+from torchrl.objectives.cql import CQLLoss, DiscreteCQLLoss
+from torchrl.objectives.crossq import CrossQLoss
+from torchrl.objectives.ddpg import DDPGLoss
+from torchrl.objectives.decision_transformer import DTLoss, OnlineDTLoss
+from torchrl.objectives.dqn import DistributionalDQNLoss, DQNLoss
+from torchrl.objectives.dreamer import (
+    DreamerActorLoss,
+    DreamerModelLoss,
+    DreamerValueLoss,
+)
+from torchrl.objectives.gail import GAILLoss
+from torchrl.objectives.iql import DiscreteIQLLoss, IQLLoss
+from torchrl.objectives.multiagent import QMixerLoss
+from torchrl.objectives.ppo import ClipPPOLoss, KLPENPPOLoss, PPOLoss
+from torchrl.objectives.redq import REDQLoss
+from torchrl.objectives.reinforce import ReinforceLoss
+from torchrl.objectives.sac import DiscreteSACLoss, SACLoss
+from torchrl.objectives.td3 import TD3Loss
+from torchrl.objectives.td3_bc import TD3BCLoss
+from torchrl.objectives.utils import (
+    default_value_kwargs,
+    distance_loss,
+    group_optimizers,
+    HardUpdate,
+    hold_out_net,
+    hold_out_params,
+    next_state_value,
+    SoftUpdate,
+    TargetNetUpdater,
+    ValueEstimators,
+)
+
+__all__ = [
+    "A2CLoss",
+    "CQLLoss",
+    "ClipPPOLoss",
+    "CrossQLoss",
+    "DDPGLoss",
+    "DQNLoss",
+    "DTLoss",
+    "DiscreteCQLLoss",
+    "DiscreteIQLLoss",
+    "DiscreteSACLoss",
+    "DistributionalDQNLoss",
+    "DreamerActorLoss",
+    "DreamerModelLoss",
+    "DreamerValueLoss",
+    "GAILLoss",
+    "HardUpdate",
+    "IQLLoss",
+    "KLPENPPOLoss",
+    "LossModule",
+    "OnlineDTLoss",
+    "PPOLoss",
+    "QMixerLoss",
+    "REDQLoss",
+    "ReinforceLoss",
+    "SACLoss",
+    "SoftUpdate",
+    "TD3BCLoss",
+    "TD3Loss",
+    "TargetNetUpdater",
+    "ValueEstimators",
+    "add_random_module",
+    "default_value_kwargs",
+    "distance_loss",
+    "group_optimizers",
+    "hold_out_net",
+    "hold_out_params",
+    "next_state_value",
+]