Shortcuts

Source code for torchrl.modules.tensordict_module.sequence

# 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 CompositeSpec
from torchrl.modules.tensordict_module.common import SafeModule


[docs]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:`TensorDictSequencial` 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 module that will be executed are those who 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. TensorDictSequence supports functional, modular and vmap coding: Examples: >>> import torch >>> from tensordict import TensorDict >>> from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec >>> from torchrl.modules import TanhNormal, SafeSequential, TensorDictModule, NormalParamWrapper >>> from torchrl.modules.tensordict_module import SafeProbabilisticModule >>> td = TensorDict({"input": torch.randn(3, 4)}, [3,]) >>> spec1 = CompositeSpec(hidden=UnboundedContinuousTensorSpec(4), loc=None, scale=None) >>> net1 = NormalParamWrapper(torch.nn.Linear(4, 8)) >>> 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 = UnboundedContinuousTensorSpec(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) CompositeSpec( hidden: UnboundedContinuousTensorSpec( shape=torch.Size([4]), space=None, device=cpu, dtype=torch.float32, domain=continuous), loc: None, scale: None, output: UnboundedContinuousTensorSpec( 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, ): self.partial_tolerant = partial_tolerant in_keys, out_keys = self._compute_in_and_out_keys(modules) spec = CompositeSpec() for module in modules: try: spec.update(module.spec) except AttributeError: spec.update(CompositeSpec({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, )

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources