[Deprecation] Remove NormalParamWrapper

ghstack-source-id: 4a70178
Pull Request resolved: #2747

Author: Vincent Moens

docs/source/reference/modules.rst

@@ -439,7 +439,6 @@ Some distributions are typically used in RL scripts.
 
     Delta
     IndependentNormal
-    NormalParamWrapper
     TanhNormal
     TruncatedNormal
     TanhDelta

test/test_distributions.py

@@ -13,10 +13,10 @@
 import torch.nn.functional as F
 
 from tensordict import TensorDictBase
+from tensordict.nn import NormalParamExtractor
 from torch import autograd, nn
 from torch.utils._pytree import tree_map
 from torchrl.modules import (
-    NormalParamWrapper,
     OneHotCategorical,
     OneHotOrdinal,
     Ordinal,
@@ -310,14 +310,19 @@ def test_normal_mapping(batch_size, device, scale_mapping, action_dim=11, state_
     torch.manual_seed(0)
     for _ in range(100):
         module = nn.LazyLinear(2 * action_dim).to(device)
-        module = NormalParamWrapper(module, scale_mapping=scale_mapping).to(device)
         if scale_mapping != "raise_error":
+            module = nn.Sequential(
+                module, NormalParamExtractor(scale_mapping=scale_mapping)
+            ).to(device)
             loc, scale = module(torch.randn(*batch_size, state_dim, device=device))
             assert (scale > 0).all()
         else:
             with pytest.raises(
                 NotImplementedError, match="Unknown mapping " "raise_error"
             ):
+                module = nn.Sequential(
+                    module, NormalParamExtractor(scale_mapping=scale_mapping)
+                ).to(device)
                 loc, scale = module(torch.randn(*batch_size, state_dim, device=device))
 
 

torchrl/modules/distributions/continuous.py

@@ -4,10 +4,9 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-import warnings
 import weakref
 from numbers import Number
-from typing import Dict, Optional, Sequence, Tuple, Union
+from typing import Dict, Optional, Sequence, Union
 
 import numpy as np
 import torch
@@ -27,7 +26,6 @@
     safeatanh_noeps,
     safetanh_noeps,
 )
-from torchrl.modules.utils import mappings
 
 # speeds up distribution construction
 D.Distribution.set_default_validate_args(False)
@@ -126,60 +124,16 @@ def inv(self):
         return inv
 
 
-class NormalParamWrapper(nn.Module):
-    """A wrapper for normal distribution parameters.
-
-    Args:
-        operator (nn.Module): operator whose output will be transformed_in in location and scale parameters
-        scale_mapping (str, optional): positive mapping function to be used with the std.
-            default = "biased_softplus_1.0" (i.e. softplus map with bias such that fn(0.0) = 1.0)
-            choices: "softplus", "exp", "relu", "biased_softplus_1";
-        scale_lb (Number, optional): The minimum value that the variance can take. Default is 1e-4.
-
-    Examples:
-        >>> from torch import nn
-        >>> import torch
-        >>> module = nn.Linear(3, 4)
-        >>> module_normal = NormalParamWrapper(module)
-        >>> tensor = torch.randn(3)
-        >>> loc, scale = module_normal(tensor)
-        >>> print(loc.shape, scale.shape)
-        torch.Size([2]) torch.Size([2])
-        >>> assert (scale > 0).all()
-        >>> # with modules that return more than one tensor
-        >>> module = nn.LSTM(3, 4)
-        >>> module_normal = NormalParamWrapper(module)
-        >>> tensor = torch.randn(4, 2, 3)
-        >>> loc, scale, others = module_normal(tensor)
-        >>> print(loc.shape, scale.shape)
-        torch.Size([4, 2, 2]) torch.Size([4, 2, 2])
-        >>> assert (scale > 0).all()
-
-    """
-
+class NormalParamWrapper(nn.Module):  # noqa: D101
     def __init__(
         self,
         operator: nn.Module,
         scale_mapping: str = "biased_softplus_1.0",
         scale_lb: Number = 1e-4,
     ) -> None:
-        warnings.warn(
-            "The NormalParamWrapper class will be deprecated in v0.7 in favor of :class:`~tensordict.nn.NormalParamExtractor`.",
-            category=DeprecationWarning,
+        raise RuntimeError(
+            "NormalParamWrapper has been deprecated in favor of `tensordict.nn.NormalParamExtractor`. Use this class instead."
         )
-        super().__init__()
-        self.operator = operator
-        self.scale_mapping = scale_mapping
-        self.scale_lb = scale_lb
-
-    def forward(self, *tensors: torch.Tensor) -> Tuple[torch.Tensor]:
-        net_output = self.operator(*tensors)
-        others = ()
-        if not isinstance(net_output, torch.Tensor):
-            net_output, *others = net_output
-        loc, scale = net_output.chunk(2, -1)
-        scale = mappings(self.scale_mapping)(scale).clamp_min(self.scale_lb)
-        return (loc, scale, *others)
 
 
 class TruncatedNormal(D.Independent):

torchrl/modules/utils/mappings.py

@@ -3,62 +3,6 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Callable
-
-import torch
-from tensordict.nn.utils import biased_softplus, inv_softplus
+from tensordict.nn.utils import biased_softplus, expln, inv_softplus, mappings
 
 __all__ = ["biased_softplus", "expln", "inv_softplus", "mappings"]
-
-
-def expln(x):
-    """A smooth, continuous positive mapping presented in "State-Dependent Exploration for Policy Gradient Methods".
-
-    https://people.idsia.ch/~juergen/ecml2008rueckstiess.pdf
-
-    """
-    out = torch.empty_like(x)
-    idx_neg = x <= 0
-    out[idx_neg] = x[idx_neg].exp()
-    out[~idx_neg] = x[~idx_neg].log1p() + 1
-    return out
-
-
-def mappings(key: str) -> Callable:
-    """Given an input string, returns a surjective function f(x): R -> R^+.
-
-    Args:
-        key (str): one of "softplus", "exp", "relu", "expln",
-            or "biased_softplus". If the key beggins with "biased_softplus",
-            then it needs to take the following form:
-            ```"biased_softplus_{bias}"``` where ```bias``` can be converted to a floating point number that will be used to bias the softplus function.
-            Alternatively, the ```"biased_softplus_{bias}_{min_val}"``` syntax can be used. In that case, the additional ```min_val``` term is a floating point
-            number that will be used to encode the minimum value of the softplus transform.
-            In practice, the equation used is softplus(x + bias) + min_val, where bias and min_val are values computed such that the conditions above are met.
-
-    Returns:
-         a Callable
-
-    """
-    _mappings = {
-        "softplus": torch.nn.functional.softplus,
-        "exp": torch.exp,
-        "relu": torch.relu,
-        "biased_softplus": biased_softplus(1.0),
-        "expln": expln,
-    }
-    if key in _mappings:
-        return _mappings[key]
-    elif key.startswith("biased_softplus"):
-        stripped_key = key.split("_")
-        if len(stripped_key) == 3:
-            return biased_softplus(float(stripped_key[-1]))
-        elif len(stripped_key) == 4:
-            return biased_softplus(
-                float(stripped_key[-2]), min_val=float(stripped_key[-1])
-            )
-        else:
-            raise ValueError(f"Invalid number of args in  {key}")
-
-    else:
-        raise NotImplementedError(f"Unknown mapping {key}")