Add clip and constant std functionalities to Gaussian Distribution (#201)

---------

Co-authored-by: epalmaEth <epalma@student.ethz.ch>
Co-authored-by: ClemensSchwarke <clemens.schwarke@gmail.com>

Author: 3 people

CITATION.cff

@@ -3,7 +3,7 @@ title: "RSL-RL: A Learning Library for Robotics Research"
 message: "If you use this work, please cite the following paper."
 repository-code: "https://github.com/leggedrobotics/rsl_rl"
 license: BSD-3-Clause
-version: 5.1.0
+version: 5.2.0
 type: software
 authors:
   - family-names: Schwarke

CONTRIBUTORS.md

@@ -30,6 +30,7 @@ Please keep the lists sorted alphabetically.
 ## Contributors
 
 * Bikram Pandit
+* Emilio Palma
 * Eric Vollenweider
 * Fabian Jenelten
 * Lorenzo Terenzi

docs/guide/configuration.rst

@@ -52,7 +52,8 @@ Currently, RSL-RL implements two runner classes:
    * - ``obs_groups``
      - dict[str, list[str]]
      - required
-     - Mapping from observation sets to observation groups coming from the environment. See :ref:`here <observation-configuration>` for more details.
+     - Mapping from observation sets to observation groups coming from the environment. See :ref:`here
+       <observation-configuration>` for more details.
    * - ``save_interval``
      - int
      - required
@@ -290,33 +291,7 @@ MLPModel
      - dict | None
      - ``None``
      - Optional output distribution configuration. If provided, the model can output stochastic values sampled from 
-       the distribution.
-
-The  ``distribution_cfg`` dictionary contains all parameters required by a specific distribution. RSL-RL implements two
-distributions by default: A simple Gaussian distribution (:class:`~rsl_rl.modules.distribution.GaussianDistribution`) 
-and a Gaussian distribution with state-dependent standard deviation 
-(:class:`~rsl_rl.modules.distribution.HeteroscedasticGaussianDistribution`). Both require the same parameters:
-
-.. list-table::
-   :header-rows: 1
-   :class: no-wrap-type-column
-
-   * - Key
-     - Type
-     - Default
-     - Description
-   * - ``class_name``
-     - str
-     - required
-     - Distribution class name. Valid values: ``"GaussianDistribution"``, ``"HeteroscedasticGaussianDistribution"``.
-   * - ``init_std``
-     - float
-     - ``1.0``
-     - Initial standard deviation.
-   * - ``std_type``
-     - str
-     - ``"scalar"``
-     - Parameterization of the standard deviation. Valid values: ``"scalar"``, ``"log"``.
+       the specified distribution.
 
 RNNModel
 ^^^^^^^^
@@ -441,6 +416,75 @@ configuration includes the following parameters:
      - ``True``
      - Whether to flatten the output tensor.
 
+
+Distribution Configuration
+--------------------------
+
+RSL-RL implements two distributions that enable stochastic model outputs: A
+:class:`~rsl_rl.modules.distribution.GaussianDistribution` and a
+:class:`~rsl_rl.modules.distribution.HeteroscedasticGaussianDistribution` with state-dependent standard deviation, which
+may be configured as follows.
+
+GaussianDistribution
+^^^^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :class: no-wrap-type-column
+
+   * - Key
+     - Type
+     - Default
+     - Description
+   * - ``class_name``
+     - str
+     - required
+     - Distribution class name. Valid values: ``"GaussianDistribution"``.
+   * - ``init_std``
+     - float
+     - ``1.0``
+     - Initial standard deviation for all dimensions.
+   * - ``std_range``
+     - tuple[float, float]
+     - ``(1e-6, 1e6)``
+     - Minimum and maximum allowed values for the standard deviation for numerical stability.
+   * - ``std_type``
+     - str
+     - ``"scalar"``
+     - Whether the standard deviation is stored directly or in log-space. Valid values: ``"scalar"``, ``"log"``.
+   * - ``learn_std``
+     - bool
+     - ``True``
+     - Whether the standard deviation is learnable or fixed.
+
+HeteroscedasticGaussianDistribution
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :class: no-wrap-type-column
+
+   * - Key
+     - Type
+     - Default
+     - Description
+   * - ``class_name``
+     - str
+     - required
+     - Distribution class name. Valid values: ``"HeteroscedasticGaussianDistribution"``.
+   * - ``init_std``
+     - float
+     - ``1.0``
+     - Initial standard deviation (used to initialize the std head bias).
+   * - ``std_range``
+     - tuple[float, float]
+     - ``(1e-6, 1e6)``
+     - Minimum and maximum allowed values for the standard deviation for numerical stability.
+   * - ``std_type``
+     - str
+     - ``"scalar"``
+     - Whether the standard deviation is stored directly or in log-space. Valid values: ``"scalar"``, ``"log"``.
+
 Extension Configuration
 -----------------------
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rsl-rl-lib"
-version = "5.1.0"
+version = "5.2.0"
 keywords = ["reinforcement-learning", "robotics"]
 maintainers = [
   { name="Clemens Schwarke", email="cschwarke@ethz.ch" },

rsl_rl/modules/distribution.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import numpy as np
 import torch
 import torch.nn as nn
 from torch.distributions import Normal
@@ -129,37 +130,50 @@ def init_mlp_weights(self, mlp: nn.Module) -> None:
 
 
 class GaussianDistribution(Distribution):
-    """Gaussian (Normal) distribution module with state-independent standard deviation.
+    """Gaussian distribution module with state-independent standard deviation.
 
-    This distribution parameterizes actions using a multivariate Gaussian with diagonal covariance. The standard
-    deviation is a learnable parameter that is independent of the model input. It can be parameterized in either
-    "scalar" space (directly) or "log" space.
+    This distribution parameterizes stochastic outputs using a multivariate Gaussian with diagonal covariance. The
+    standard deviation can be a learnable parameter or a constant. It can be parameterized in either "scalar" space or
+    "log" space and is clamped to a specified range.
+
+    .. note::
+        If the standard deviation type is set to "log", the provided arguments are still interpreted in scalar space,
+        and converted to log space internally.
     """
 
     def __init__(
         self,
         output_dim: int,
         init_std: float = 1.0,
+        std_range: tuple[float, float] = (1e-6, 1e6),
         std_type: str = "scalar",
+        learn_std: bool = True,
     ) -> None:
         """Initialize the Gaussian distribution module.
 
         Args:
             output_dim: Dimension of the action/output space.
             init_std: Initial standard deviation.
+            std_range: Range for the standard deviation. Should be a tuple of (min, max) values for clamping.
             std_type: Parameterization of the standard deviation: "scalar" or "log".
+            learn_std: Whether the standard deviation should be learnable. If False, it will be fixed to `init_std`.
         """
         super().__init__(output_dim)
         self.std_type = std_type
 
         # Learnable std parameters
         if std_type == "scalar":
-            self.std_param = nn.Parameter(init_std * torch.ones(output_dim))
+            self.std_param = nn.Parameter(init_std * torch.ones(output_dim), requires_grad=learn_std)
         elif std_type == "log":
-            self.log_std_param = nn.Parameter(torch.log(init_std * torch.ones(output_dim)))
+            self.log_std_param = nn.Parameter(torch.log(init_std * torch.ones(output_dim)), requires_grad=learn_std)
         else:
             raise ValueError(f"Unknown standard deviation type: {std_type}. Should be 'scalar' or 'log'.")
 
+        # Clamp the std range to ensure numerical stability and store log space range if needed
+        self.std_range = list(std_range)
+        self.std_range[0] = max(self.std_range[0], 1e-6)  # Avoid zero std for numerical stability
+        self.log_std_range = [float(np.log(self.std_range[0])), float(np.log(self.std_range[1]))]
+
         # Internal torch distribution (populated by update())
         self._distribution: Normal | None = None
 
@@ -170,9 +184,10 @@ def update(self, mlp_output: torch.Tensor) -> None:
         """Update the Gaussian distribution from MLP output."""
         mean = mlp_output
         if self.std_type == "scalar":
-            std = self.std_param.expand_as(mean)
+            std = self.std_param.clamp(self.std_range[0], self.std_range[1])
         elif self.std_type == "log":
-            std = torch.exp(self.log_std_param).expand_as(mean)
+            log_std = self.log_std_param.clamp(self.log_std_range[0], self.log_std_range[1])
+            std = torch.exp(log_std)
         self._distribution = Normal(mean, std)
 
     def sample(self) -> torch.Tensor:
@@ -226,24 +241,30 @@ def kl_divergence(self, old_params: tuple[torch.Tensor, ...], new_params: tuple[
 
 
 class HeteroscedasticGaussianDistribution(GaussianDistribution):
-    """Gaussian (Normal) distribution module with state-dependent standard deviation.
+    """Gaussian distribution module with state-dependent standard deviation.
 
-    This distribution parameterizes actions using a multivariate Gaussian with diagonal covariance. The standard
-    deviation is output by the MLP alongside the mean, making it state-dependent (heteroscedastic). It can be
-    parameterized in either "scalar" space (directly) or "log" space.
+    This distribution parameterizes stochastic outputs using a multivariate Gaussian with diagonal covariance. The
+    standard deviation is output by the MLP alongside the mean, making it state-dependent. It can be parameterized in
+    either "scalar" space or "log" space, and is clamped to a specified range.
+
+    .. note::
+        If the standard deviation type is set to "log", the provided arguments are still interpreted in scalar space,
+        and converted to log space internally.
     """
 
     def __init__(
         self,
         output_dim: int,
         init_std: float = 1.0,
+        std_range: tuple[float, float] = (1e-6, 1e6),
         std_type: str = "scalar",
     ) -> None:
         """Initialize the heteroscedastic Gaussian distribution module.
 
         Args:
             output_dim: Dimension of the action/output space.
-            init_std: Initial standard deviation (used to initialize MLP std head bias).
+            init_std: Initial standard deviation (used to initialize the MLP's std head bias).
+            std_range: Range for the standard deviation. Should be a tuple of (min, max) values for clamping.
             std_type: Parameterization of the standard deviation: "scalar" or "log".
         """
         # Skip GaussianDistribution.__init__ to avoid creating unnecessary learnable std parameters.
@@ -254,6 +275,11 @@ def __init__(
         if std_type not in ("scalar", "log"):
             raise ValueError(f"Unknown standard deviation type: {std_type}. Should be 'scalar' or 'log'.")
 
+        # Clamp the std range to ensure numerical stability and store log space range if needed
+        self.std_range = list(std_range)
+        self.std_range[0] = max(self.std_range[0], 1e-6)  # Avoid zero std for numerical stability
+        self.log_std_range = [float(np.log(self.std_range[0])), float(np.log(self.std_range[1]))]
+
         # Internal torch distribution (populated by update())
         self._distribution: Normal | None = None
 
@@ -264,8 +290,10 @@ def update(self, mlp_output: torch.Tensor) -> None:
         """Update the Gaussian distribution from MLP output."""
         if self.std_type == "scalar":
             mean, std = torch.unbind(mlp_output, dim=-2)
+            std = torch.clamp(std, self.std_range[0], self.std_range[1])
         elif self.std_type == "log":
             mean, log_std = torch.unbind(mlp_output, dim=-2)
+            log_std = torch.clamp(log_std, self.log_std_range[0], self.log_std_range[1])
             std = torch.exp(log_std)
         self._distribution = Normal(mean, std)