rlmodule 0.1.7__tar.gz → 0.2.0__tar.gz

{rlmodule-0.1.7/rlmodule.egg-info → rlmodule-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rlmodule
-Version: 0.1.7
+Version: 0.2.0
 Summary: Flexible reinforcement learning models instantiators library
 Author: Lopatovsky
 Author-email: lopatovsky@gmail.com

{rlmodule-0.1.7 → rlmodule-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rlmodule"
-version = "0.1.7"
+version = "0.2.0"
 description = "Flexible reinforcement learning models instantiators library"
 readme = "README.md"
 requires-python = ">=3.10"

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/skrl/torch/network/__init__.py

@@ -9,6 +9,7 @@ __all__ = [
     "RnnMlpWithForwardedInput",
     "ResidualMlp",
     "CnnBackbone",
+    "Cnn3dBackbone",
     "CnnMlp",
     "CnnProprioMlp",
     "GraphNetwork",
@@ -23,6 +24,7 @@ __all__ = [
     "RnnMlpCfg",
     "CnnCfg",
     "CnnBackboneCfg",
+    "Cnn3dBackboneCfg",
     "CnnMlpCfg",
     "CnnProprioMlpCfg",
     "GraphNetworkCfg",
@@ -33,6 +35,7 @@ from rlmodule.source.network import (  # noqa: F401
     LSTM,
     MLP,
     RNN,
+    Cnn3dBackbone,
     CnnBackbone,
     CnnMlp,
     CnnProprioMlp,
@@ -43,6 +46,7 @@ from rlmodule.source.network import (  # noqa: F401
     RnnMlpWithForwardedInput,
 )
 from rlmodule.source.network_cfg import (  # noqa: F401
+    Cnn3dBackboneCfg,
     CnnBackboneCfg,
     CnnCfg,
     CnnMlpCfg,

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/log_std.py

@@ -102,18 +102,24 @@ class CombinedLogStd(LogStd):
     def __init__(self, device: Union[str, torch.device], input_size: int, output_size: int, cfg):
         super().__init__(device, input_size, output_size, cfg)
 
-        self._modules = [
-            module_cfg.class_type(device, input_size, output_size, module_cfg) for module_cfg in cfg.combined_modules
-        ]
+        # NOTE: must NOT be ``self._modules`` — that name is ``nn.Module``'s
+        # internal child-module registry. Assigning a plain list there
+        # silently breaks parameter registration (``.parameters()`` /
+        # ``.to()`` / ``state_dict()``), so the sub-LogStd params never reach
+        # the optimizer and their std stays frozen. Use a registered
+        # ``nn.ModuleList`` under a distinct attribute instead.
+        self._sub_log_stds = nn.ModuleList(
+            [module_cfg.class_type(device, input_size, output_size, module_cfg) for module_cfg in cfg.combined_modules]
+        )
 
         if cfg.combination_constants is None:
-            self._combination_constants = list(repeat(1.0, len(self._modules)))
+            self._combination_constants = list(repeat(1.0, len(self._sub_log_stds)))
         else:
             self._combination_constants = cfg.combination_constants
-            if len(self._modules) != len(self._combination_constants):
+            if len(self._sub_log_stds) != len(self._combination_constants):
                 raise ValueError(
                     f"Number of combination constants {len(self._combination_constants)} does not match number of"
-                    f" modules {len(self._modules)}"
+                    f" modules {len(self._sub_log_stds)}"
                 )
 
         if cfg.combination_method == "max":
@@ -128,6 +134,7 @@ class CombinedLogStd(LogStd):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """Compute standard deviation as a combination from multiple Std modules."""
         collected_stds = [
-            self._combination_constants[module_id] * module(input) for module_id, module in enumerate(self._modules)
+            self._combination_constants[module_id] * module(input)
+            for module_id, module in enumerate(self._sub_log_stds)
         ]
         return self.clip_std(reduce(self._combination_method, collected_stds))

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/model.py

@@ -3,14 +3,17 @@ from typing import Any, Mapping, Optional, Tuple, Union
 import collections
 import gym
 import gymnasium
+import logging
 
 import torch
 import torch.nn as nn
 
-from rlmodule import logger
 from rlmodule.source.utils import contains_rnn_module
 
 
+logger = logging.getLogger("rlmodule")
+
+
 class Model(torch.nn.Module):
     def __init__(self, device: Union[str, torch.device], network: nn.Module) -> None:
         """Base class representing a function approximator"""

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/network.py

@@ -1,12 +1,16 @@
 from typing import Any, Sequence, Union
 
+import logging
+
 import torch
 import torch.nn as nn
 
-from rlmodule import logger
 from rlmodule.source.utils import get_cnn_shape, get_output_size, get_space_size
 
 
+logger = logging.getLogger("rlmodule")
+
+
 LN_POSITION_PRE = "pre"
 LN_POSITION_POST = "post"
 _VALID_LN_POSITIONS = (LN_POSITION_PRE, LN_POSITION_POST)
@@ -707,14 +711,66 @@ def _build_impala_cnn(in_channels: int, widths: Sequence[int] = (16, 32, 32)) ->
     return nn.Sequential(*stages)
 
 
+class SpatialSoftmax(nn.Module):
+    """Spatial soft-argmax keypoints (Levine et al. 2016, "End-to-End Training of
+    Deep Visuomotor Policies").
+
+    For a ``(B, C, H, W)`` feature map, each channel is turned into a softmax
+    distribution over its ``H * W`` positions and reduced to that distribution's
+    expected image coordinates. The output is ``(B, 2 * C)``, flat:
+    ``[x_0, y_0, x_1, y_1, ...]``, with ``x`` (columns) and ``y`` (rows) in
+    ``[-1, 1]`` (pixel centres, -1 = left / top). It gives *where* a feature is,
+    which average pooling blurs away for features only a few pixels wide.
+
+    ``temperature`` divides the logits (lower = sharper, closer to a hard
+    argmax). With ``learnable_temperature`` it is one trained scalar, stored as
+    its log so it stays positive. The reduction runs in float32 whatever the
+    input dtype, since a bf16 softmax over hundreds of positions loses the
+    sub-pixel precision the keypoints are for.
+    """
+
+    def __init__(self, temperature: float = 1.0, learnable_temperature: bool = True):
+        super().__init__()
+        if temperature <= 0:
+            raise ValueError(f"SpatialSoftmax temperature must be > 0; got {temperature!r}.")
+        log_t = torch.tensor(float(temperature)).log()
+        if learnable_temperature:
+            self.log_temperature = nn.Parameter(log_t)
+        else:
+            self.register_buffer("log_temperature", log_t)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 4:
+            raise ValueError(f"SpatialSoftmax expects (B, C, H, W); got shape {tuple(x.shape)}.")
+        b, c, h, w = x.shape
+        logits = x.float().reshape(b, c, h * w) / self.log_temperature.float().exp()
+        attention = torch.softmax(logits, dim=-1)
+        ys = torch.linspace(-1.0, 1.0, h, device=x.device) if h > 1 else torch.zeros(1, device=x.device)
+        xs = torch.linspace(-1.0, 1.0, w, device=x.device) if w > 1 else torch.zeros(1, device=x.device)
+        grid_y, grid_x = torch.meshgrid(ys, xs, indexing="ij")
+        ex = (attention * grid_x.reshape(1, 1, -1)).sum(-1)
+        ey = (attention * grid_y.reshape(1, 1, -1)).sum(-1)
+        return torch.stack([ex, ey], dim=-1).reshape(b, 2 * c).to(x.dtype)
+
+
 _NN_LAYER_REGISTRY: dict[str, type] = {
     "Conv2d": nn.Conv2d,
     "ELU": nn.ELU,
     "ReLU": nn.ReLU,
     "MaxPool2d": nn.MaxPool2d,
     "AdaptiveAvgPool2d": nn.AdaptiveAvgPool2d,
+    # Keeps a small, strong response (a box a few pixels wide) that average
+    # pooling would dilute with its neighbourhood.
+    "AdaptiveMaxPool2d": nn.AdaptiveMaxPool2d,
+    # Per-channel expected (x, y) keypoints; its output is already flat.
+    "SpatialSoftmax": SpatialSoftmax,
     "Flatten": nn.Flatten,
     "BatchNorm2d": nn.BatchNorm2d,
+    # 3D variants for volumetric (e.g. occupancy-grid) backbones.
+    "Conv3d": nn.Conv3d,
+    "MaxPool3d": nn.MaxPool3d,
+    "AdaptiveAvgPool3d": nn.AdaptiveAvgPool3d,
+    "BatchNorm3d": nn.BatchNorm3d,
 }
 
 
@@ -796,6 +852,66 @@ class CnnBackbone(nn.Module):
         return self.cnn(x)
 
 
+class Cnn3dBackbone(nn.Module):
+    """Flat-vector **3D** CNN backbone (the volumetric analogue of :class:`CnnBackbone`).
+
+    Accepts a flat ``(B, D*H*W*C)`` observation, reshapes it internally to
+    ``(B, C, D, H, W)`` (using the DHWC-ordered ``input_states`` shape — same
+    channels-last convention as :class:`CnnBackbone`, extended by a leading
+    depth axis), runs the chosen ``Conv3d`` stack, and returns a flat feature
+    vector suitable for an MLP head.
+
+    Unlike the 2D backbone there are no canned ``nature`` / ``impala`` presets
+    (those reference 2D image architectures), so only ``preset="custom"`` is
+    supported: build the conv stack from an explicit ``cfg.layers`` sequence
+    (e.g. ``Conv3d``/``MaxPool3d``/``AdaptiveAvgPool3d``/``Flatten``).
+
+    This mirrors the flat-vector contract of :class:`CnnBackbone` exactly, so it
+    drops straight into ``SharedRLModel`` / ``SeparatedRLModel`` with no
+    SKRL-side changes. It is the right primitive for occupancy grids and other
+    small volumetric observations, where treating the grid as a stack of 2D
+    channels would discard 3D locality.
+    """
+
+    def __init__(self, cfg):
+        super().__init__()
+        self._input_shape = get_cnn_shape(cfg.input_states)  # (D, H, W, C)
+        if len(self._input_shape) != 4:
+            raise ValueError(
+                "Cnn3dBackboneCfg.input_states must be a 4-tuple in DHWC order "
+                f"(D, H, W, C); got {self._input_shape!r}."
+            )
+        d, h, w, c = self._input_shape
+        self._flat_size = int(d) * int(h) * int(w) * int(c)
+
+        preset = getattr(cfg, "preset", "custom")
+        if preset == "custom":
+            if cfg.layers is None or len(cfg.layers) == 0:
+                raise ValueError(
+                    "Cnn3dBackboneCfg.preset='custom' requires an explicit 'layers' sequence."
+                )
+            self.cnn = nn.Sequential(*list(cfg.layers))
+        else:
+            raise ValueError(
+                f"Unknown 3D CNN preset {preset!r}; only 'custom' is supported "
+                "(2D 'nature'/'impala' presets do not apply to volumetric input)."
+            )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        # Reshape a flat (B, D*H*W*C) input to (B, C, D, H, W). DHWC->CDHW permute
+        # mirrors the 2D backbone's HWC->CHW convention. Validate the flat size up
+        # front so a layout/offset mismatch fails with a clear message instead of
+        # an opaque ``.view()`` "shape is invalid for input of size" error.
+        if input.shape[-1] != self._flat_size:
+            raise ValueError(
+                f"Cnn3dBackbone expected a flat input of size {self._flat_size} "
+                f"(D*H*W*C for input_states={self._input_shape}), got "
+                f"{input.shape[-1]}."
+            )
+        x = input.view(-1, *self._input_shape).permute(0, 4, 1, 2, 3).contiguous()
+        return self.cnn(x)
+
+
 class CnnMlp(nn.Module):
     """Convolutional backbone followed by an MLP (or ResidualMlp) trunk.
 
@@ -1016,6 +1132,7 @@ class GraphNetwork(nn.Module):
         "concat",
         "residual_add",
         "cnn_backbone",
+        "cnn3d_backbone",
         "rnn",
         "gru",
         "lstm",
@@ -1236,6 +1353,7 @@ class GraphNetwork(nn.Module):
         action_input_size: int,
     ):
         from rlmodule.source.network_cfg import (  # local import to avoid cycle
+            Cnn3dBackboneCfg,
             CnnBackboneCfg,
             GruCfg,
             LstmCfg,
@@ -1386,6 +1504,43 @@ class GraphNetwork(nn.Module):
             cnn_out = get_output_size(module, in_dim)
             return module, int(cnn_out)
 
+        if node_type == "cnn3d_backbone":
+            in_dim = self._single_input_dim(node_id)
+            # Accept either ``volume_shape`` or ``image_shape`` for symmetry with
+            # the 2D node; both must be a 4-dim DHWC tuple ``(D, H, W, C)``.
+            volume_shape = params.get("volume_shape")
+            if volume_shape is None:
+                volume_shape = params.get("image_shape")
+            if volume_shape is None:
+                raise ValueError(
+                    f"GraphNetwork cnn3d_backbone {node_id!r} requires "
+                    "params['volume_shape'] = (D, H, W, C)."
+                )
+            volume_shape = tuple(volume_shape)
+            if len(volume_shape) != 4:
+                raise ValueError(
+                    f"GraphNetwork cnn3d_backbone {node_id!r} volume_shape must "
+                    f"have 4 dims (D, H, W, C); got {volume_shape!r}."
+                )
+            # 3D backbone only supports the 'custom' preset; require explicit layers.
+            layers_spec = params.get("layers")
+            if not layers_spec:
+                raise ValueError(
+                    f"GraphNetwork cnn3d_backbone {node_id!r} requires a 'layers' list "
+                    "(only the 'custom' preset is supported for 3D)."
+                )
+            custom_layers = _build_nn_layers_from_spec(layers_spec)
+            sub_cfg = Cnn3dBackboneCfg(
+                input_size=in_dim,
+                input_states=volume_shape,
+                preset="custom",
+                layers=custom_layers,
+            )
+            module = Cnn3dBackbone(sub_cfg)
+            # Probe the conv stack's flat output size.
+            cnn_out = get_output_size(module, in_dim)
+            return module, int(cnn_out)
+
         if node_type in ("rnn", "gru", "lstm"):
             in_dim = self._single_input_dim(node_id)
             # RnnBaseCfg marks all four of these as required (MISSING), so the

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/network_cfg.py

@@ -4,17 +4,18 @@ from collections.abc import Callable
 from dataclasses import MISSING, field
 import gym
 import gymnasium
+import logging
 
 import torch
 import torch.nn as nn
 
-from rlmodule import logger
 from rlmodule.source.network import (
     CNN,
     GRU,
     LSTM,
     MLP,
     RNN,
+    Cnn3dBackbone,
     CnnBackbone,
     CnnMlp,
     CnnProprioMlp,
@@ -25,6 +26,9 @@ from rlmodule.source.network import (
 )
 
 
+logger = logging.getLogger("rlmodule")
+
+
 # use isaac-lab native configclass if available to avoid it being declared twice
 try:
     from omni.isaac.lab.utils import configclass
@@ -200,6 +204,31 @@ class CnnBackboneCfg(NetworkCfg):
     """Per-stage channel widths for the IMPALA preset."""
 
 
+@configclass
+class Cnn3dBackboneCfg(NetworkCfg):
+    """Configuration for :class:`rlmodule.source.network.Cnn3dBackbone`.
+
+    The volumetric analogue of :class:`CnnBackboneCfg`. ``input_states`` must be
+    a **DHWC**-ordered shape (``(D, H, W, C)``): a leading depth axis prepended
+    to the 2D backbone's HWC convention. Internally the module reshapes the flat
+    observation to ``(B, C, D, H, W)``, runs the ``Conv3d`` stack, and returns a
+    flat feature vector suitable for an MLP head.
+
+    Only ``preset="custom"`` is supported (the 2D ``"nature"`` / ``"impala"``
+    presets reference image architectures that have no 3D counterpart here), so
+    ``layers`` must be provided — e.g. a stack of ``Conv3d`` / ``MaxPool3d`` /
+    ``AdaptiveAvgPool3d`` / ``Flatten`` modules.
+    """
+
+    module: type[Cnn3dBackbone] = Cnn3dBackbone
+
+    preset: str = "custom"
+    """Backbone preset. Only ``"custom"`` is supported for 3D."""
+
+    layers: Sequence[nn.Module] = None
+    """Explicit ``Conv3d`` layer stack; required for ``preset="custom"``."""
+
+
 @configclass
 class CnnMlpCfg(NetworkCfg):
     """Configuration for :class:`rlmodule.source.network.CnnMlp` (CNN trunk + MLP head).

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/output_layer.py

@@ -1,3 +1,4 @@
+import math
 from typing import Union
 
 import gym
@@ -5,6 +6,7 @@ import gymnasium
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.distributions import Normal
 
 from rlmodule.source.utils import get_output_size, get_space_size
@@ -80,6 +82,16 @@ class GaussianLayer(OutputLayer):
             else torch.sum if cfg.reduction == "sum" else torch.prod if cfg.reduction == "prod" else None
         )
 
+        self._squash = cfg.squash
+        if self._squash and cfg.output_activation is not nn.Identity:
+            raise ValueError(
+                "squash=True bounds the sampled action with tanh, so the mean must be the pre-squash value: "
+                f"set output_activation=nn.Identity (got {cfg.output_activation})"
+            )
+        if cfg.entropy_reduction not in ["mean", "sum"]:
+            raise ValueError("entropy_reduction must be one of 'mean' or 'sum'")
+        self._entropy_reduction = cfg.entropy_reduction
+
         self._net = nn.Sequential(nn.Linear(self._projection_in, self._output_size), cfg.output_activation())
 
         self._log_std = cfg.log_std.class_type(device, self._projection_in, self._output_size, cfg.log_std)
@@ -121,6 +133,9 @@ class GaussianLayer(OutputLayer):
         # distribution
         self._distribution = Normal(mean_actions, self._cached_log_std.exp())
 
+        if self._squash:
+            return self._squashed_forward(mean_actions, taken_actions, outputs_dict)
+
         # sample using the reparametrization trick
         actions = self._distribution.rsample()
 
@@ -162,6 +177,42 @@ class GaussianLayer(OutputLayer):
         # return output * self.instantiator_output_scale, self._log_std, {}
         # return output, self._log_std, output_dict
 
+    # Keeps atanh finite for actions stored at the bound (float32 tanh saturates to 1).
+    _SQUASH_EPS = 1e-6
+
+    def _squashed_forward(self, mean_actions, taken_actions, outputs_dict):
+        """Tanh-squashed Gaussian: ``u ~ N(mean, std)``, ``a = output_scale * tanh(u)``.
+
+        ``self._distribution`` stays the pre-squash Normal over ``u``. ``mean_actions``
+        here is ``output_scale * net(x)`` (``output_activation`` is Identity), so
+        ``mu = mean_actions / output_scale`` in ``u`` space.
+
+        ``log_prob`` applies the change of variables per action dimension before the
+        reduction: ``log p(a) = log N(u) - log(output_scale) - log(1 - tanh(u)^2)``,
+        with ``log(1 - tanh(u)^2) = 2 * (log 2 - u - softplus(-2u))`` for stability.
+        """
+        scale = self._output_scale
+        self._distribution = Normal(mean_actions / scale, self._cached_log_std.exp())
+
+        if taken_actions is None:
+            pre_squash = self._distribution.rsample()
+            actions = scale * torch.tanh(pre_squash)
+        else:
+            actions = taken_actions
+            bound = 1.0 - self._SQUASH_EPS
+            pre_squash = torch.atanh(torch.clamp(taken_actions / scale, -bound, bound))
+
+        log_det = 2.0 * (math.log(2.0) - pre_squash - F.softplus(-2.0 * pre_squash)) + math.log(scale)
+        log_prob = self._distribution.log_prob(pre_squash) - log_det
+        if self._reduction is not None:
+            log_prob = self._reduction(log_prob, dim=-1)
+        if log_prob.dim() != actions.dim():
+            log_prob = log_prob.unsqueeze(-1)
+
+        outputs_dict["mean_actions"] = scale * torch.tanh(mean_actions / scale)
+        outputs_dict["log_prob"] = log_prob
+        return actions, outputs_dict
+
     def get_entropy(self, role: str = "") -> torch.Tensor:
         """Compute and return the entropy of the model
 
@@ -174,11 +225,17 @@ class GaussianLayer(OutputLayer):
 
             >>> entropy = model.get_entropy()
             >>> print(entropy.shape)
-            torch.Size([4096, 8])
+            torch.Size([4096, 8])  # entropy_reduction="mean"; [4096, 1] with "sum"
+
+        With ``squash=True`` this is the entropy of the pre-squash Gaussian: the
+        squashed distribution has no closed form (SAC uses ``-log_prob`` instead).
         """
         if self._distribution is None:
             return torch.tensor(0.0, device=self.device)
-        return self._distribution.entropy().to(self.device)
+        entropy = self._distribution.entropy().to(self.device)
+        if self._entropy_reduction == "sum":
+            entropy = entropy.sum(dim=-1, keepdim=True)
+        return entropy
 
     def get_log_std(self, role: str = "") -> torch.Tensor:
         """Return the log standard deviation of the model

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/output_layer_cfg.py

@@ -74,6 +74,24 @@ class GaussianLayerCfg(OutputLayerCfg):
     log_std: LogStdCfg = ParameterLogStdCfg()
     """Module to specify how standard deviation is computed."""
 
+    squash: bool = False
+    """Tanh-squash the sampled action: ``a = output_scale * tanh(u)``, ``u ~ N(mean, std)``.
+
+    ``log_prob`` includes the tanh change-of-variables correction, so it is the density of
+    the bounded action the env executes. This is the SAC actor parameterisation: without it
+    an unbounded Gaussian's entropy term can grow ``std`` without limit while the env clamps
+    every sample to its bound. Requires ``output_activation=nn.Identity``; ``clip_actions``
+    is not applied (the action is already within ``[-output_scale, output_scale]``).
+    """
+
+    entropy_reduction: str = "mean"
+    """How ``get_entropy`` reduces over action dimensions: ``"mean"`` keeps the per-dimension
+    ``(num_samples, num_actions)`` entropy, which skrl's PPO/A2C/RPO then ``.mean()`` over, so
+    the bonus is the average per-dimension entropy (skrl's convention). ``"sum"`` returns the
+    joint entropy ``(num_samples, 1)`` (rsl_rl's convention), making the same
+    ``entropy_loss_scale`` ``num_actions`` times stronger.
+    """
+
 
 @configclass
 class DeterministicLayerCfg(OutputLayerCfg):

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule/source/rlmodel.py

@@ -1,13 +1,16 @@
 from typing import Any, Mapping, Optional, Tuple, Union
 
+import logging
+
 import torch
 import torch.nn as nn
 
-from rlmodule import logger
-
 from .model import Model
 
 
+logger = logging.getLogger("rlmodule")
+
+
 # SKRL 2.x roles whose canonical trunk input is the actor's observation. Every
 # other role (``value``, ``critic_*``, ``target_critic_*``) is a critic and
 # reads ``states``.

{rlmodule-0.1.7 → rlmodule-0.2.0/rlmodule.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rlmodule
-Version: 0.1.7
+Version: 0.2.0
 Summary: Flexible reinforcement learning models instantiators library
 Author: Lopatovsky
 Author-email: lopatovsky@gmail.com

{rlmodule-0.1.7 → rlmodule-0.2.0}/rlmodule.egg-info/SOURCES.txt

@@ -40,10 +40,12 @@ rlmodule/source/nvidia_utils/dict.py
 rlmodule/source/nvidia_utils/string.py
 tests/test_build_model_integration.py
 tests/test_cnn_backbone.py
+tests/test_combined_log_std.py
 tests/test_graph_network.py
 tests/test_head_mlp.py
 tests/test_mlp_layer_norm.py
 tests/test_residual_mlp.py
 tests/test_rnn_memory_and_correctness.py
 tests/test_skrl_startup.py
-tests/test_skrl_train_smoke.py
+tests/test_skrl_train_smoke.py
+tests/test_squashed_gaussian.py

{rlmodule-0.1.7 → rlmodule-0.2.0}/tests/test_cnn_backbone.py

@@ -11,8 +11,9 @@ import pytest
 import torch
 import torch.nn as nn
 
-from rlmodule.source.network import CnnBackbone, CnnMlp, CnnProprioMlp
+from rlmodule.source.network import Cnn3dBackbone, CnnBackbone, CnnMlp, CnnProprioMlp
 from rlmodule.source.network_cfg import (
+    Cnn3dBackboneCfg,
     CnnBackboneCfg,
     CnnMlpCfg,
     CnnProprioMlpCfg,
@@ -58,6 +59,121 @@ def test_cnn_backbone_custom_preset_uses_user_layers():
     assert out.shape[0] == 2 and out.dim() == 2
 
 
+# ---------------------------------------------------------------------------
+# Cnn3dBackbone (volumetric / occupancy-grid CNN, flat-vector contract)
+# ---------------------------------------------------------------------------
+
+
+def _occ_3d_layers():
+    """A small Conv3d stack that flattens to a fixed-size embedding regardless of
+    the (possibly anisotropic) input volume dims."""
+    return [
+        nn.Conv3d(1, 8, kernel_size=3, padding=1),
+        nn.ELU(),
+        nn.MaxPool3d(kernel_size=2, ceil_mode=True),
+        nn.Conv3d(8, 16, kernel_size=3, padding=1),
+        nn.ELU(),
+        nn.AdaptiveAvgPool3d((2, 2, 2)),
+        nn.Flatten(),
+    ]
+
+
+@pytest.mark.parametrize("volume", [(5, 5, 5, 1), (8, 8, 4, 1), (6, 4, 3, 1)])
+def test_cnn3d_backbone_custom_forward(volume):
+    cfg = Cnn3dBackboneCfg(input_states=volume, preset="custom", layers=_occ_3d_layers())
+    net = Cnn3dBackbone(cfg)
+    d, h, w, c = volume
+    flat = torch.randn(3, d * h * w * c)
+    out = net(flat)
+    assert out.dim() == 2
+    assert out.shape[0] == 3
+    # AdaptiveAvgPool3d((2,2,2)) over 16 channels -> 16*2*2*2 = 128 features.
+    assert out.shape[1] == 128
+
+
+def test_cnn3d_backbone_custom_preset_requires_layers():
+    with pytest.raises(ValueError):
+        Cnn3dBackbone(Cnn3dBackboneCfg(input_states=(5, 5, 5, 1), preset="custom"))
+
+
+def test_cnn3d_backbone_rejects_2d_presets():
+    for preset in ("nature", "impala", "does-not-exist"):
+        with pytest.raises(ValueError):
+            Cnn3dBackbone(
+                Cnn3dBackboneCfg(
+                    input_states=(5, 5, 5, 1), preset=preset, layers=_occ_3d_layers()
+                )
+            )
+
+
+def test_cnn3d_backbone_rejects_wrong_flat_size():
+    cfg = Cnn3dBackboneCfg(input_states=(5, 5, 5, 1), preset="custom", layers=_occ_3d_layers())
+    net = Cnn3dBackbone(cfg)
+    bad = torch.randn(2, 5 * 5 * 5 * 1 + 7)  # off-by-7 flat slice
+    with pytest.raises(ValueError, match="expected a flat input of size 125"):
+        net(bad)
+
+
+def test_cnn3d_backbone_requires_4d_input_states():
+    with pytest.raises(ValueError, match="4-tuple in DHWC"):
+        Cnn3dBackbone(
+            Cnn3dBackboneCfg(input_states=(5, 5, 1), preset="custom", layers=_occ_3d_layers())
+        )
+
+
+def test_cnn3d_backbone_dhwc_reshape_is_correct():
+    """A flat DHWC vector must map to (B, C, D, H, W) so 3D locality is preserved.
+
+    We use an identity-ish Conv3d (Flatten only) and check the permutation maps
+    the right elements: element at flat index for (d, h, w) with C=1 must land at
+    spatial position (d, h, w) of the volume the conv sees.
+    """
+    d, h, w, c = 2, 3, 4, 1
+    net = Cnn3dBackbone(
+        Cnn3dBackboneCfg(
+            input_states=(d, h, w, c),
+            preset="custom",
+            # AvgPool over the H,W plane per depth slice so each output element is
+            # the mean of one (h,w) plane -> lets us verify the depth axis ordering.
+            layers=[nn.AdaptiveAvgPool3d((d, 1, 1)), nn.Flatten()],
+        )
+    )
+    # Build a volume where every voxel in depth-slice k has value k.
+    vol = torch.zeros(1, d, h, w, c)
+    for k in range(d):
+        vol[:, k] = float(k)
+    flat = vol.reshape(1, -1)
+    out = net(flat)  # shape (1, d)
+    assert out.shape == (1, d)
+    assert torch.allclose(out[0], torch.arange(d, dtype=torch.float32))
+
+
+def test_cnn3d_backbone_backward_reaches_input_and_params():
+    torch.manual_seed(0)
+    volume = (5, 5, 5, 1)
+    cfg = Cnn3dBackboneCfg(input_states=volume, preset="custom", layers=_occ_3d_layers())
+    net = Cnn3dBackbone(cfg)
+    d, h, w, c = volume
+    flat = torch.randn(2, d * h * w * c, requires_grad=True)
+    out = net(flat)
+    out.sum().backward()
+
+    assert flat.grad is not None
+    assert torch.isfinite(flat.grad).all()
+    conv_params = [
+        (name, p) for name, p in net.named_parameters() if isinstance_conv3d_param(name)
+    ]
+    assert conv_params, "expected Conv3d parameters in the backbone"
+    for name, p in net.named_parameters():
+        assert p.grad is not None, f"{name} has no grad"
+        assert torch.isfinite(p.grad).all()
+
+
+def isinstance_conv3d_param(name: str) -> bool:
+    # Conv3d submodules live at indices 0 and 3 of the sequential stack here.
+    return name.startswith("cnn.0.") or name.startswith("cnn.3.")
+
+
 def test_cnn_mlp_with_layer_norm_trunk():
     cfg = CnnMlpCfg(
         input_states=(84, 84, 4),
@@ -309,3 +425,91 @@ def test_visual_amp_dtype_fp16_logs_warning_on_nonfinite_output():
     assert net._nonfinite_visual_warned is True
     _ = net(x)
     assert net._nonfinite_visual_warned is True
+
+
+# ---------------------------------------------------------------------------
+# AdaptiveMaxPool2d / SpatialSoftmax custom layers
+# ---------------------------------------------------------------------------
+
+
+def test_spatial_softmax_finds_a_single_hot_pixel():
+    from rlmodule.source.network import SpatialSoftmax
+
+    layer = SpatialSoftmax(temperature=0.01, learnable_temperature=False)
+    x = torch.zeros(2, 3, 5, 9)
+    x[0, 0, 0, 0] = 10.0  # top-left
+    x[0, 1, 4, 8] = 10.0  # bottom-right
+    x[0, 2, 2, 4] = 10.0  # centre
+    x[1, 0, 1, 6] = 10.0  # row 1 of 5 -> y -0.5; column 6 of 9 -> x 0.5
+    out = layer(x)
+    assert out.shape == (2, 6)
+    torch.testing.assert_close(out[0], torch.tensor([-1.0, -1.0, 1.0, 1.0, 0.0, 0.0]), atol=1e-4, rtol=0)
+    torch.testing.assert_close(out[1, :2], torch.tensor([0.5, -0.5]), atol=1e-4, rtol=0)
+    # A flat map spreads the attention uniformly: the keypoint is the centre.
+    torch.testing.assert_close(out[1, 2:], torch.zeros(4), atol=1e-6, rtol=0)
+
+
+def test_spatial_softmax_temperature_is_trainable_and_positive():
+    from rlmodule.source.network import SpatialSoftmax
+
+    layer = SpatialSoftmax(temperature=0.5)
+    assert float(layer.log_temperature.detach().exp()) == pytest.approx(0.5)
+    x = torch.randn(4, 2, 6, 6, requires_grad=True)
+    layer(x).pow(2).sum().backward()
+    assert x.grad is not None and torch.isfinite(x.grad).all()
+    assert layer.log_temperature.grad is not None
+    fixed = SpatialSoftmax(learnable_temperature=False)
+    assert not list(fixed.parameters())
+    with pytest.raises(ValueError):
+        SpatialSoftmax(temperature=0.0)
+
+
+def test_spatial_softmax_keeps_the_input_dtype():
+    from rlmodule.source.network import SpatialSoftmax
+
+    out = SpatialSoftmax()(torch.randn(2, 3, 4, 4, dtype=torch.bfloat16))
+    assert out.dtype == torch.bfloat16
+
+
+def test_adaptive_max_pool_keeps_a_small_peak_that_avg_pool_dilutes():
+    from rlmodule.source.network import _build_nn_layers_from_spec
+
+    mx, = _build_nn_layers_from_spec([{"type": "AdaptiveMaxPool2d", "output_size": [2, 2]}])
+    av, = _build_nn_layers_from_spec([{"type": "AdaptiveAvgPool2d", "output_size": [2, 2]}])
+    x = torch.zeros(1, 1, 8, 8)
+    x[0, 0, 1, 1] = 1.0
+    assert float(mx(x)[0, 0, 0, 0]) == 1.0
+    assert float(av(x)[0, 0, 0, 0]) == pytest.approx(1 / 16)
+
+
+def test_cnn_backbone_custom_spec_with_spatial_softmax_and_max_pool_branches():
+    """Two branches over one 32x48 depth image, as a GraphNetwork builds them."""
+    from rlmodule.source.network import GraphNetwork
+    from rlmodule.source.network_cfg import GraphNetworkCfg
+
+    h, w = 32, 48
+    conv = {"type": "Conv2d", "in_channels": 1, "out_channels": 8, "kernel_size": 3, "padding": 1}
+    cfg = GraphNetworkCfg(
+        nodes=[
+            {"id": "img", "type": "observation_input",
+             "params": {"indices": list(range(h * w)), "dim": h * w, "shape": [h, w, 1], "modality": "image"}},
+            {"id": "keypoints", "type": "cnn_backbone",
+             "params": {"preset": "custom", "image_shape": [h, w, 1],
+                        "layers": [conv, {"type": "ELU"}, {"type": "SpatialSoftmax", "temperature": 1.0}]}},
+            {"id": "maxpool", "type": "cnn_backbone",
+             "params": {"preset": "custom", "image_shape": [h, w, 1],
+                        "layers": [conv, {"type": "ELU"}, {"type": "AdaptiveMaxPool2d", "output_size": [4, 6]},
+                                   {"type": "Flatten"}]}},
+            {"id": "cat", "type": "concat", "params": {}},
+        ],
+        edges=[{"from": "img", "to": "keypoints"}, {"from": "img", "to": "maxpool"},
+               {"from": "keypoints", "to": "cat"}, {"from": "maxpool", "to": "cat"}],
+        output_node_id="cat",
+        flat_input_size=h * w,
+    )
+    net = GraphNetwork(cfg)
+    out = net(torch.rand(3, h * w))
+    out = out[0] if isinstance(out, tuple) else out
+    assert out.shape == (3, 8 * 2 + 8 * 4 * 6)
+    assert torch.isfinite(out).all()
+    assert out[:, :16].abs().max() <= 1.0

rlmodule-0.2.0/tests/test_combined_log_std.py

@@ -0,0 +1,42 @@
+"""Regression tests for ``CombinedLogStd`` parameter registration.
+
+``CombinedLogStd`` previously stored its sub-modules in a plain list assigned
+to ``self._modules`` — which is ``nn.Module``'s *internal* child-module
+registry. That silently prevented the sub-LogStd parameters from being
+registered, so they never reached the optimizer and the policy std stayed
+frozen. These tests pin the fix: sub-module parameters must be discoverable via
+``named_parameters()`` and must receive gradients.
+"""
+
+import torch
+
+from rlmodule.source.log_std import CombinedLogStd
+from rlmodule.source.log_std_cfg import CombinedLogStdCfg, ParameterLogStdCfg
+
+
+def _combined_cfg() -> CombinedLogStdCfg:
+    return CombinedLogStdCfg(
+        combination_method="max",
+        combined_modules=[
+            ParameterLogStdCfg(initial_log_std=0.5),
+            ParameterLogStdCfg(initial_log_std=-0.5),
+        ],
+    )
+
+
+def test_combined_log_std_registers_submodule_parameters() -> None:
+    module = CombinedLogStd("cpu", input_size=8, output_size=4, cfg=_combined_cfg())
+    log_params = [name for name, _ in module.named_parameters() if "_log_parameter" in name]
+    # Two ParameterLogStd sub-modules -> two registered learnable tensors.
+    assert len(log_params) == 2, f"expected 2 registered log_std params, got {log_params}"
+
+
+def test_combined_log_std_submodule_params_receive_grad() -> None:
+    module = CombinedLogStd("cpu", input_size=8, output_size=4, cfg=_combined_cfg())
+    out = module(torch.zeros(3, 8))
+    out.sum().backward()
+    # The element-wise ``max`` routes gradient to the dominant sub-module; the
+    # point is that the registered parameters are on the autograd graph at all
+    # (a plain-list assignment detached them entirely).
+    graphed = [p for name, p in module.named_parameters() if "_log_parameter" in name and p.grad is not None]
+    assert graphed, "no CombinedLogStd sub-parameter received a gradient"

{rlmodule-0.1.7 → rlmodule-0.2.0}/tests/test_graph_network.py

@@ -309,6 +309,75 @@ def test_graph_network_cnn_backbone_accepts_image_slice():
     assert torch.isfinite(out).all()
 
 
+def test_graph_network_cnn3d_backbone_accepts_volume_slice():
+    d, h, w, c = 5, 5, 5, 1
+    obs_dim = d * h * w * c
+    cfg = GraphNetworkCfg(
+        nodes=[
+            {
+                "id": "obs",
+                "type": "observation_input",
+                "params": {
+                    "indices": list(range(obs_dim)),
+                    "dim": obs_dim,
+                    "shape": [d, h, w, c],
+                    "modality": "volume",
+                },
+            },
+            {
+                "id": "cnn3d",
+                "type": "cnn3d_backbone",
+                "params": {
+                    "volume_shape": [d, h, w, c],
+                    "layers": [
+                        {"type": "Conv3d", "in_channels": 1, "out_channels": 8,
+                         "kernel_size": 3, "padding": 1},
+                        {"type": "ELU"},
+                        {"type": "AdaptiveAvgPool3d", "output_size": [2, 2, 2]},
+                        {"type": "Flatten"},
+                    ],
+                },
+            },
+        ],
+        edges=[{"from": "obs", "to": "cnn3d"}],
+        output_node_id="cnn3d",
+        flat_input_size=obs_dim,
+    )
+    net = GraphNetwork(cfg)
+    # 8 channels * 2*2*2 = 64 features.
+    assert net.output_dim == 64
+    flat = torch.randn(2, obs_dim)
+    out, _ = _run_forward(net, flat)
+    assert out.shape == (2, 64)
+    assert torch.isfinite(out).all()
+
+
+def test_graph_network_cnn3d_backbone_requires_4d_volume_shape():
+    obs_dim = 25
+    cfg = GraphNetworkCfg(
+        nodes=[
+            {
+                "id": "obs",
+                "type": "observation_input",
+                "params": {"indices": list(range(obs_dim)), "dim": obs_dim},
+            },
+            {
+                "id": "cnn3d",
+                "type": "cnn3d_backbone",
+                "params": {
+                    "volume_shape": [5, 5, 1],  # only 3 dims -> invalid
+                    "layers": [{"type": "Flatten"}],
+                },
+            },
+        ],
+        edges=[{"from": "obs", "to": "cnn3d"}],
+        output_node_id="cnn3d",
+        flat_input_size=obs_dim,
+    )
+    with pytest.raises(ValueError, match="4 dims"):
+        GraphNetwork(cfg)
+
+
 def test_graph_network_output_dim_is_exposed_after_construction():
     cfg = GraphNetworkCfg(
         nodes=[

rlmodule-0.2.0/tests/test_squashed_gaussian.py

@@ -0,0 +1,118 @@
+"""Tests for ``GaussianLayerCfg.squash`` (tanh-squashed Gaussian) and
+``GaussianLayerCfg.entropy_reduction``.
+
+An unbounded Gaussian head trained with SAC can grow ``std`` without limit
+while the env clamps every sample to its bound: the entropy term keeps
+rewarding a spread the env never executes. ``squash=True`` bounds the action
+with ``tanh`` and makes ``log_prob`` the density of that bounded action.
+"""
+
+import math
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.distributions import AffineTransform, Normal, TanhTransform, TransformedDistribution
+
+from rlmodule.source.log_std_cfg import ParameterLogStdCfg
+from rlmodule.source.output_layer import GaussianLayer
+from rlmodule.source.output_layer_cfg import GaussianLayerCfg
+
+N_IN, N_ACT = 6, 4
+
+
+def _layer(**kw) -> GaussianLayer:
+    torch.manual_seed(0)
+    kw.setdefault("log_std", ParameterLogStdCfg(initial_log_std=1.5))
+    cfg = GaussianLayerCfg(output_size=N_ACT, **kw)
+    return GaussianLayer("cpu", N_IN, cfg)
+
+
+def _squashed(**kw) -> GaussianLayer:
+    return _layer(squash=True, output_activation=nn.Identity, **kw)
+
+
+def test_squashed_actions_stay_within_output_scale() -> None:
+    layer = _squashed(output_scale=2.0)
+    actions, _ = layer(torch.randn(4096, N_IN) * 10, None, {})
+    assert actions.abs().max() <= 2.0
+
+
+@pytest.mark.parametrize("scale", [1.0, 2.0])
+def test_squashed_log_prob_matches_torch_transformed_distribution(scale: float) -> None:
+    layer = _squashed(output_scale=scale)
+    x = torch.randn(256, N_IN)
+    actions, out = layer(x, None, {})
+    base = Normal(layer._net(x), layer._cached_log_std.exp())
+    ref = TransformedDistribution(base, [TanhTransform(), AffineTransform(0.0, scale)])
+    # Keep away from the bound, where the reference's atanh loses precision.
+    ok = (actions.abs() < 0.99 * scale).all(dim=-1)
+    expected = ref.log_prob(actions).sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(out["log_prob"][ok], expected[ok], atol=1e-4, rtol=1e-4)
+
+
+def test_squashed_log_prob_of_taken_actions_matches_the_sampled_one() -> None:
+    # PPO-style re-evaluation: feeding the sampled action back must give the same density.
+    layer = _squashed()
+    x = torch.randn(512, N_IN)
+    actions, out = layer(x, None, {})
+    _, again = layer(x, actions, {})
+    ok = (actions.abs() < 0.999).all(dim=-1)
+    torch.testing.assert_close(again["log_prob"][ok], out["log_prob"][ok], atol=1e-3, rtol=1e-4)
+
+
+def test_squashed_log_prob_is_finite_at_the_bound() -> None:
+    layer = _squashed()
+    x = torch.randn(3, N_IN)
+    _, out = layer(x, torch.ones(3, N_ACT), {})
+    assert torch.isfinite(out["log_prob"]).all()
+
+
+def test_squashed_mean_action_is_tanh_of_the_mean() -> None:
+    layer = _squashed(output_scale=2.0)
+    x = torch.randn(8, N_IN)
+    _, out = layer(x, None, {})
+    torch.testing.assert_close(out["mean_actions"], 2.0 * torch.tanh(layer._net(x)))
+
+
+def test_saturating_the_bound_lowers_the_entropy_estimate() -> None:
+    # The SAC failure: without squashing, a huge std reads as huge entropy even though
+    # the env executes +-1. With squashing, a very wide std piles samples at the bound,
+    # where the density is high, so -log_prob falls instead of rising.
+    x = torch.zeros(20000, N_IN)
+    ent = {}
+    for log_std in (0.0, 3.0):
+        layer = _squashed(log_std=ParameterLogStdCfg(initial_log_std=log_std, max_log_std=4.0))
+        _, out = layer(x, None, {})
+        ent[log_std] = -out["log_prob"].mean().item()
+    assert ent[3.0] < ent[0.0]
+
+
+def test_squash_requires_identity_output_activation() -> None:
+    with pytest.raises(ValueError, match="output_activation=nn.Identity"):
+        _layer(squash=True)
+
+
+def test_unsquashed_layer_is_unchanged() -> None:
+    layer = _layer()
+    x = torch.randn(64, N_IN)
+    torch.manual_seed(1)
+    actions, out = layer(x, None, {})
+    mean = layer._net(x)
+    ref = Normal(mean, layer._cached_log_std.exp()).log_prob(actions).sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(out["log_prob"], ref)
+    torch.testing.assert_close(out["mean_actions"], mean)
+
+
+def test_entropy_reduction_defaults_to_per_dimension() -> None:
+    layer = _layer()
+    layer(torch.randn(5, N_IN), None, {})
+    assert layer.get_entropy().shape == (5, N_ACT)
+
+
+def test_entropy_reduction_sum_is_the_joint_entropy() -> None:
+    layer = _layer(entropy_reduction="sum", log_std=ParameterLogStdCfg(initial_log_std=0.0))
+    layer(torch.randn(5, N_IN), None, {})
+    entropy = layer.get_entropy()
+    assert entropy.shape == (5, 1)
+    torch.testing.assert_close(entropy, torch.full((5, 1), N_ACT * 0.5 * math.log(2 * math.pi * math.e)))