torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/modules/models/__init__.py

@@ -0,0 +1,86 @@
+# 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.modules.tensordict_module.common import DistributionalDQNnet
+
+from .batchrenorm import BatchRenorm1d
+
+from .decision_transformer import DecisionTransformer
+from .exploration import (
+    ConsistentDropout,
+    ConsistentDropoutModule,
+    NoisyLazyLinear,
+    NoisyLinear,
+    reset_noise,
+)
+from .llm import GPT2RewardModel
+from .model_based import (
+    DreamerActor,
+    ObsDecoder,
+    ObsEncoder,
+    RSSMPosterior,
+    RSSMPrior,
+    RSSMRollout,
+)
+from .models import (
+    Conv2dNet,
+    Conv3dNet,
+    ConvNet,
+    DdpgCnnActor,
+    DdpgCnnQNet,
+    DdpgMlpActor,
+    DdpgMlpQNet,
+    DTActor,
+    DuelingCnnDQNet,
+    DuelingMlpDQNet,
+    MLP,
+    OnlineDTActor,
+)
+from .multiagent import (
+    MultiAgentConvNet,
+    MultiAgentMLP,
+    MultiAgentNetBase,
+    QMixer,
+    VDNMixer,
+)
+from .utils import Squeeze2dLayer, SqueezeLayer
+
+__all__ = [
+    "DistributionalDQNnet",
+    "BatchRenorm1d",
+    "DecisionTransformer",
+    "GPT2RewardModel",
+    "ConsistentDropout",
+    "ConsistentDropoutModule",
+    "NoisyLazyLinear",
+    "NoisyLinear",
+    "reset_noise",
+    "DreamerActor",
+    "ObsDecoder",
+    "ObsEncoder",
+    "RSSMPosterior",
+    "RSSMPrior",
+    "RSSMRollout",
+    "Conv2dNet",
+    "Conv3dNet",
+    "ConvNet",
+    "DdpgCnnActor",
+    "DdpgCnnQNet",
+    "DdpgMlpActor",
+    "DdpgMlpQNet",
+    "DTActor",
+    "DuelingCnnDQNet",
+    "DuelingMlpDQNet",
+    "MLP",
+    "OnlineDTActor",
+    "MultiAgentConvNet",
+    "MultiAgentMLP",
+    "MultiAgentNetBase",
+    "QMixer",
+    "VDNMixer",
+    "Squeeze2dLayer",
+    "SqueezeLayer",
+]

torchrl/modules/models/batchrenorm.py

@@ -0,0 +1,119 @@
+# 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 torch
+import torch.nn as nn
+
+
+class BatchRenorm1d(nn.Module):
+    """BatchRenorm Module (https://arxiv.org/abs/1702.03275).
+
+    The code is adapted from https://github.com/google-research/corenet
+
+    BatchRenorm is an enhanced version of the standard BatchNorm. Unlike BatchNorm,
+    it utilizes running statistics to normalize batches after an initial warmup phase.
+    This approach reduces the impact of "outlier" batches that may occur during
+    extended training periods, making BatchRenorm more robust for long training runs.
+
+    During the warmup phase, BatchRenorm functions identically to a BatchNorm layer.
+
+    Args:
+        num_features (int): Number of features in the input tensor.
+
+    Keyword Args:
+        momentum (:obj:`float`, optional): Momentum factor for computing the running mean and variance.
+            Defaults to ``0.01``.
+        eps (:obj:`float`, optional): Small value added to the variance to avoid division by zero.
+            Defaults to ``1e-5``.
+        max_r (:obj:`float`, optional): Maximum value for the scaling factor r.
+            Defaults to ``3.0``.
+        max_d (:obj:`float`, optional): Maximum value for the bias factor d.
+            Defaults to ``5.0``.
+        warmup_steps (int, optional): Number of warm-up steps for the running mean and variance.
+            Defaults to ``10000``.
+        smooth (bool, optional): if ``True``, the behavior smoothly transitions from regular
+            batch-norm (when ``iter=0``) to batch-renorm (when ``iter=warmup_steps``).
+            Otherwise, the behavior will transition from batch-norm to batch-renorm when
+            ``iter=warmup_steps``. Defaults to ``False``.
+    """
+
+    def __init__(
+        self,
+        num_features: int,
+        *,
+        momentum: float = 0.01,
+        eps: float = 1e-5,
+        max_r: float = 3.0,
+        max_d: float = 5.0,
+        warmup_steps: int = 10000,
+        smooth: bool = False,
+    ):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.max_r = max_r
+        self.max_d = max_d
+        self.warmup_steps = warmup_steps
+        self.smooth = smooth
+
+        self.register_buffer(
+            "running_mean", torch.zeros(num_features, dtype=torch.float32)
+        )
+        self.register_buffer(
+            "running_var", torch.ones(num_features, dtype=torch.float32)
+        )
+        self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.int64))
+        self.weight = nn.Parameter(torch.ones(num_features, dtype=torch.float32))
+        self.bias = nn.Parameter(torch.zeros(num_features, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.dim() >= 2:
+            raise ValueError(
+                f"The {type(self).__name__} expects a 2D (or more) tensor, got {x.dim()}."
+            )
+
+        view_dims = [1, x.shape[1]] + [1] * (x.dim() - 2)
+
+        def _v(v):
+            return v.view(view_dims)
+
+        running_std = (self.running_var + self.eps).sqrt_()
+
+        if self.training:
+            reduce_dims = [i for i in range(x.dim()) if i != 1]
+            b_mean = x.mean(reduce_dims)
+            b_var = x.var(reduce_dims, unbiased=False)
+            b_std = (b_var + self.eps).sqrt_()
+
+            r = torch.clamp((b_std.detach() / running_std), 1 / self.max_r, self.max_r)
+            d = torch.clamp(
+                (b_mean.detach() - self.running_mean) / running_std,
+                -self.max_d,
+                self.max_d,
+            )
+
+            # Compute warmup factor (0 during warmup, 1 after warmup)
+            if self.warmup_steps > 0:
+                if self.smooth:
+                    warmup_factor = self.num_batches_tracked / self.warmup_steps
+                else:
+                    warmup_factor = self.num_batches_tracked // self.warmup_steps
+                r = 1.0 + (r - 1.0) * warmup_factor
+                d = d * warmup_factor
+
+            x = (x - _v(b_mean)) / _v(b_std) * _v(r) + _v(d)
+
+            unbiased_var = b_var.detach() * x.shape[0] / (x.shape[0] - 1)
+            self.running_var += self.momentum * (unbiased_var - self.running_var)
+            self.running_mean += self.momentum * (b_mean.detach() - self.running_mean)
+            self.num_batches_tracked += 1
+            self.num_batches_tracked.clamp_max(self.warmup_steps)
+        else:
+            x = (x - _v(self.running_mean)) / _v(running_std)
+
+        x = _v(self.weight) * x + _v(self.bias)
+        return x

torchrl/modules/models/decision_transformer.py

@@ -0,0 +1,179 @@
+# 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 dataclasses
+import importlib
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+_has_transformers = importlib.util.find_spec("transformers") is not None
+
+
+class DecisionTransformer(nn.Module):
+    """Online Decision Transformer.
+
+    Desdescribed in https://arxiv.org/abs/2202.05607 .
+
+    The transformer utilizes a default config to create the GPT2 model if the user does not provide a specific config.
+    default_config = {
+    ... "n_embd": 256,
+    ... "n_layer": 4,
+    ... "n_head": 4,
+    ... "n_inner": 1024,
+    ... "activation": "relu",
+    ... "n_positions": 1024,
+    ... "resid_pdrop": 0.1,
+    ... "attn_pdrop": 0.1,
+    }
+
+    Args:
+        state_dim (int): dimension of the state space
+        action_dim (int): dimension of the action space
+        config (:obj:`~.DTConfig` or dict, optional): transformer architecture configuration,
+            used to create the GPT2Config from transformers.
+            Defaults to ``default_config``.
+
+
+    Example:
+        >>> config = DecisionTransformer.default_config()
+        >>> config.n_embd = 128
+        >>> print(config)
+        DTConfig(n_embd: 128, n_layer: 4, n_head: 4, n_inner: 1024, activation: relu, n_positions: 1024, resid_pdrop: 0.1, attn_pdrop: 0.1)
+        >>> # alternatively
+        >>> config = DecisionTransformer.DTConfig(n_embd=128)
+        >>> model = DecisionTransformer(state_dim=4, action_dim=2, config=config)
+        >>> batch_size = [3, 32]
+        >>> length = 10
+        >>> observation = torch.randn(*batch_size, length, 4)
+        >>> action = torch.randn(*batch_size, length, 2)
+        >>> return_to_go = torch.randn(*batch_size, length, 1)
+        >>> output = model(observation, action, return_to_go)
+        >>> output.shape
+        torch.Size([3, 32, 10, 128])
+
+    """
+
+    @dataclass
+    class DTConfig:
+        """Default configuration for DecisionTransformer."""
+
+        n_embd: Any = 256
+        n_layer: Any = 4
+        n_head: Any = 4
+        n_inner: Any = 1024
+        activation: Any = "relu"
+        n_positions: Any = 1024
+        resid_pdrop: Any = 0.1
+        attn_pdrop: Any = 0.1
+
+        def __repr__(self):
+            fields = []
+            for f in dataclasses.fields(self):
+                value = getattr(self, f.name)
+                fields.append(f"{f.name}: {value}")
+            fields = ", ".join(fields)
+            return f"{self.__class__.__name__}({fields})"
+
+    @classmethod
+    def default_config(cls):
+        return cls.DTConfig()
+
+    def __init__(
+        self,
+        state_dim,
+        action_dim,
+        config: dict | DTConfig = None,
+        device: torch.device | None = None,
+    ):
+
+        if not _has_transformers:
+            raise ImportError(
+                "transformers is not installed. Please install it with `pip install transformers`."
+            )
+        import transformers
+        from transformers.models.gpt2.modeling_gpt2 import GPT2Model
+
+        if config is None:
+            config = self.default_config()
+        if isinstance(config, self.DTConfig):
+            config = dataclasses.asdict(config)
+        if not isinstance(config, dict):
+            try:
+                config = dict(config)
+            except Exception as err:
+                raise TypeError(
+                    f"Config of type {type(config)} is not supported."
+                ) from err
+
+        super().__init__()
+
+        with torch.device(device) if device is not None else nullcontext():
+            gpt_config = transformers.GPT2Config(
+                n_embd=config["n_embd"],
+                n_layer=config["n_layer"],
+                n_head=config["n_head"],
+                n_inner=config["n_inner"],
+                activation_function=config["activation"],
+                n_positions=config["n_positions"],
+                resid_pdrop=config["resid_pdrop"],
+                attn_pdrop=config["attn_pdrop"],
+                vocab_size=1,
+            )
+            self.state_dim = state_dim
+            self.action_dim = action_dim
+            self.hidden_size = config["n_embd"]
+
+            self.transformer = GPT2Model(config=gpt_config)
+
+            self.embed_return = torch.nn.Linear(1, self.hidden_size)
+            self.embed_state = torch.nn.Linear(self.state_dim, self.hidden_size)
+            self.embed_action = torch.nn.Linear(self.action_dim, self.hidden_size)
+
+            self.embed_ln = nn.LayerNorm(self.hidden_size)
+
+    def forward(
+        self,
+        observation: torch.Tensor,
+        action: torch.Tensor,
+        return_to_go: torch.Tensor,
+    ):
+        batch_size, seq_length = observation.shape[:-2], observation.shape[-2]
+        batch_size_orig = batch_size
+        if len(batch_size) != 1:
+            # TODO: vmap over transformer once this is possible
+            observation = observation.view(-1, *observation.shape[-2:])
+            action = action.view(-1, *action.shape[-2:])
+            return_to_go = return_to_go.view(-1, *return_to_go.shape[-2:])
+            batch_size = torch.Size([batch_size.numel()])
+
+        # embed each modality with a different head
+        state_embeddings = self.embed_state(observation)
+        action_embeddings = self.embed_action(action)
+        returns_embeddings = self.embed_return(return_to_go)
+
+        # this makes the sequence look like (R_1, s_1, a_1, R_2, s_2, a_2, ...)
+        # which works nice in an autoregressive sense since states predict actions
+        stacked_inputs = torch.stack(
+            (returns_embeddings, state_embeddings, action_embeddings), dim=-2
+        ).reshape(*batch_size, 3 * seq_length, self.hidden_size)
+        stacked_inputs = self.embed_ln(stacked_inputs)
+
+        # we feed in the input embeddings (not word indices as in NLP) to the model
+        transformer_outputs = self.transformer(
+            inputs_embeds=stacked_inputs,
+        )
+        x = transformer_outputs["last_hidden_state"]
+
+        # reshape x so that the second dimension corresponds to the original
+        # returns (0), states (1), or actions (2); i.e. x[:,1,t] is the token for s_t
+        x = x.reshape(*batch_size, seq_length, 3, self.hidden_size).transpose(-3, -2)
+        if batch_size_orig is batch_size:
+            return x[..., 1, :, :]  # only state tokens
+        return x[..., 1, :, :].reshape(*batch_size_orig, *x.shape[-2:])