init commit.

source/rsl_rl/rsl_rl/modules/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Definitions for neural-network components for RL-agents."""
+
+from .actor_critic import ActorCritic
+from .actor_critic_cnn import ActorCriticCNN
+from .actor_critic_recurrent import ActorCriticRecurrent
+from .actor_critic_moe_cts import ActorCriticMoECTS
+from .rnd import RandomNetworkDistillation, resolve_rnd_config
+from .student_teacher import StudentTeacher
+from .student_teacher_recurrent import StudentTeacherRecurrent
+from .symmetry import resolve_symmetry_config
+
+__all__ = [
+    "ActorCritic",
+    "ActorCriticCNN",
+    "ActorCriticRecurrent",
+    "ActorCriticMoECTS",
+    "RandomNetworkDistillation",
+    "StudentTeacher",
+    "StudentTeacherRecurrent",
+    "resolve_rnd_config",
+    "resolve_symmetry_config",
+]

source/rsl_rl/rsl_rl/modules/actor_critic.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any, NoReturn
+
+from rsl_rl.networks import MLP, EmpiricalNormalization
+
+
+class ActorCritic(nn.Module):
+    is_recurrent: bool = False
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        actor_obs_normalization: bool = False,
+        critic_obs_normalization: bool = False,
+        actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        activation: str = "elu",
+        init_noise_std: float = 1.0,
+        noise_std_type: str = "scalar",
+        state_dependent_std: bool = False,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if kwargs:
+            print(
+                "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs])
+            )
+        super().__init__()
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_actor_obs = 0
+        for obs_group in obs_groups["policy"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations."
+            num_actor_obs += obs[obs_group].shape[-1]
+        num_critic_obs = 0
+        for obs_group in obs_groups["critic"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations."
+            num_critic_obs += obs[obs_group].shape[-1]
+
+        # Actor
+        self.state_dependent_std = state_dependent_std
+        if self.state_dependent_std:
+            self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation)
+        else:
+            self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation)
+        print(f"Actor MLP: {self.actor}")
+
+        # Actor observation normalization
+        self.actor_obs_normalization = actor_obs_normalization
+        if actor_obs_normalization:
+            self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs)
+        else:
+            self.actor_obs_normalizer = torch.nn.Identity()
+
+        # Critic
+        self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation)
+        print(f"Critic MLP: {self.critic}")
+
+        # Critic observation normalization
+        self.critic_obs_normalization = critic_obs_normalization
+        if critic_obs_normalization:
+            self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs)
+        else:
+            self.critic_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.state_dependent_std:
+            torch.nn.init.zeros_(self.actor[-2].weight[num_actions:])
+            if self.noise_std_type == "scalar":
+                torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std)
+            elif self.noise_std_type == "log":
+                torch.nn.init.constant_(
+                    self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7))
+                )
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            if self.noise_std_type == "scalar":
+                self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+            elif self.noise_std_type == "log":
+                self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    def reset(self, dones: torch.Tensor | None = None) -> None:
+        pass
+
+    def forward(self) -> NoReturn:
+        raise NotImplementedError
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.distribution.mean
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.distribution.stddev
+
+    @property
+    def entropy(self) -> torch.Tensor:
+        return self.distribution.entropy().sum(dim=-1)
+
+    def _update_distribution(self, obs: torch.Tensor) -> None:
+        if self.state_dependent_std:
+            # Compute mean and standard deviation
+            mean_and_std = self.actor(obs)
+            if self.noise_std_type == "scalar":
+                mean, std = torch.unbind(mean_and_std, dim=-2)
+            elif self.noise_std_type == "log":
+                mean, log_std = torch.unbind(mean_and_std, dim=-2)
+                std = torch.exp(log_std)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            # Compute mean
+            mean = self.actor(obs)
+            # Compute standard deviation
+            if self.noise_std_type == "scalar":
+                std = self.std.expand_as(mean)
+            elif self.noise_std_type == "log":
+                std = torch.exp(self.log_std).expand_as(mean)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # Create distribution
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
+        obs = self.get_actor_obs(obs)
+        obs = self.actor_obs_normalizer(obs)
+        self._update_distribution(obs)
+        return self.distribution.sample()
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_actor_obs(obs)
+        obs = self.actor_obs_normalizer(obs)
+        if self.state_dependent_std:
+            return self.actor(obs)[..., 0, :]
+        else:
+            return self.actor(obs)
+
+    def evaluate(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
+        obs = self.get_critic_obs(obs)
+        obs = self.critic_obs_normalizer(obs)
+        return self.critic(obs)
+
+    def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.actor_obs_normalization:
+            actor_obs = self.get_actor_obs(obs)
+            self.actor_obs_normalizer.update(actor_obs)
+        if self.critic_obs_normalization:
+            critic_obs = self.get_critic_obs(obs)
+            self.critic_obs_normalizer.update(critic_obs)
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
+        """Load the parameters of the actor-critic model.
+
+        Args:
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
+                :meth:`state_dict` function.
+
+        Returns:
+            Whether this training resumes a previous training. This flag is used by the :func:`load` function of
+                :class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
+        """
+        super().load_state_dict(state_dict, strict=strict)
+        return True

source/rsl_rl/rsl_rl/modules/actor_critic_cnn.py

@@ -0,0 +1,262 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any
+
+from rsl_rl.networks import CNN, MLP, EmpiricalNormalization
+
+from .actor_critic import ActorCritic
+
+
+class ActorCriticCNN(ActorCritic):
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        actor_obs_normalization: bool = False,
+        critic_obs_normalization: bool = False,
+        actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        actor_cnn_cfg: dict[str, dict] | dict | None = None,
+        critic_cnn_cfg: dict[str, dict] | dict | None = None,
+        activation: str = "elu",
+        init_noise_std: float = 1.0,
+        noise_std_type: str = "scalar",
+        state_dependent_std: bool = False,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if kwargs:
+            print(
+                "ActorCriticCNN.__init__ got unexpected arguments, which will be ignored: "
+                + str([key for key in kwargs])
+            )
+        super(ActorCritic, self).__init__()
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_actor_obs_1d = 0
+        self.actor_obs_groups_1d = []
+        actor_in_dims_2d = []
+        actor_in_channels_2d = []
+        self.actor_obs_groups_2d = []
+        for obs_group in obs_groups["policy"]:
+            if len(obs[obs_group].shape) == 4:  # B, C, H, W
+                self.actor_obs_groups_2d.append(obs_group)
+                actor_in_dims_2d.append(obs[obs_group].shape[2:4])
+                actor_in_channels_2d.append(obs[obs_group].shape[1])
+            elif len(obs[obs_group].shape) == 2:  # B, C
+                self.actor_obs_groups_1d.append(obs_group)
+                num_actor_obs_1d += obs[obs_group].shape[-1]
+            else:
+                raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
+        num_critic_obs_1d = 0
+        self.critic_obs_groups_1d = []
+        critic_in_dims_2d = []
+        critic_in_channels_2d = []
+        self.critic_obs_groups_2d = []
+        for obs_group in obs_groups["critic"]:
+            if len(obs[obs_group].shape) == 4:  # B, C, H, W
+                self.critic_obs_groups_2d.append(obs_group)
+                critic_in_dims_2d.append(obs[obs_group].shape[2:4])
+                critic_in_channels_2d.append(obs[obs_group].shape[1])
+            elif len(obs[obs_group].shape) == 2:  # B, C
+                self.critic_obs_groups_1d.append(obs_group)
+                num_critic_obs_1d += obs[obs_group].shape[-1]
+            else:
+                raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
+
+        # Assert that there are 2D observations
+        assert self.actor_obs_groups_2d or self.critic_obs_groups_2d, (
+            "No 2D observations are provided. If this is intentional, use the ActorCritic module instead."
+        )
+
+        # Actor CNN
+        if self.actor_obs_groups_2d:
+            # Resolve the actor CNN configuration
+            assert actor_cnn_cfg is not None, "An actor CNN configuration is required for 2D actor observations."
+            # If a single configuration dictionary is provided, create a dictionary for each 2D observation group
+            if not all(isinstance(v, dict) for v in actor_cnn_cfg.values()):
+                actor_cnn_cfg = {group: actor_cnn_cfg for group in self.actor_obs_groups_2d}
+            # Check that the number of configs matches the number of observation groups
+            assert len(actor_cnn_cfg) == len(self.actor_obs_groups_2d), (
+                "The number of CNN configurations must match the number of 2D actor observations."
+            )
+
+            # Create CNNs for each 2D actor observation
+            self.actor_cnns = nn.ModuleDict()
+            encoding_dim = 0
+            for idx, obs_group in enumerate(self.actor_obs_groups_2d):
+                self.actor_cnns[obs_group] = CNN(
+                    input_dim=actor_in_dims_2d[idx],
+                    input_channels=actor_in_channels_2d[idx],
+                    **actor_cnn_cfg[obs_group],
+                )
+                print(f"Actor CNN for {obs_group}: {self.actor_cnns[obs_group]}")
+                # Get the output dimension of the CNN
+                if self.actor_cnns[obs_group].output_channels is None:
+                    encoding_dim += int(self.actor_cnns[obs_group].output_dim)
+                else:
+                    raise ValueError("The output of the actor CNN must be flattened before passing it to the MLP.")
+        else:
+            self.actor_cnns = None
+            encoding_dim = 0
+
+        # Actor MLP
+        self.state_dependent_std = state_dependent_std
+        if self.state_dependent_std:
+            self.actor = MLP(num_actor_obs_1d + encoding_dim, [2, num_actions], actor_hidden_dims, activation)
+        else:
+            self.actor = MLP(num_actor_obs_1d + encoding_dim, num_actions, actor_hidden_dims, activation)
+        print(f"Actor MLP: {self.actor}")
+
+        # Actor observation normalization (only for 1D actor observations)
+        self.actor_obs_normalization = actor_obs_normalization
+        if actor_obs_normalization:
+            self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs_1d)
+        else:
+            self.actor_obs_normalizer = torch.nn.Identity()
+
+        # Critic CNN
+        if self.critic_obs_groups_2d:
+            # Resolve the critic CNN configuration
+            assert critic_cnn_cfg is not None, "A critic CNN configuration is required for 2D critic observations."
+            # If a single configuration dictionary is provided, create a dictionary for each 2D observation group
+            if not all(isinstance(v, dict) for v in critic_cnn_cfg.values()):
+                critic_cnn_cfg = {group: critic_cnn_cfg for group in self.critic_obs_groups_2d}
+            # Check that the number of configs matches the number of observation groups
+            assert len(critic_cnn_cfg) == len(self.critic_obs_groups_2d), (
+                "The number of CNN configurations must match the number of 2D critic observations."
+            )
+
+            # Create CNNs for each 2D critic observation
+            self.critic_cnns = nn.ModuleDict()
+            encoding_dim = 0
+            for idx, obs_group in enumerate(self.critic_obs_groups_2d):
+                self.critic_cnns[obs_group] = CNN(
+                    input_dim=critic_in_dims_2d[idx],
+                    input_channels=critic_in_channels_2d[idx],
+                    **critic_cnn_cfg[obs_group],
+                )
+                print(f"Critic CNN for {obs_group}: {self.critic_cnns[obs_group]}")
+                # Get the output dimension of the CNN
+                if self.critic_cnns[obs_group].output_channels is None:
+                    encoding_dim += int(self.critic_cnns[obs_group].output_dim)
+                else:
+                    raise ValueError("The output of the critic CNN must be flattened before passing it to the MLP.")
+        else:
+            self.critic_cnns = None
+            encoding_dim = 0
+
+        # Critic MLP
+        self.critic = MLP(num_critic_obs_1d + encoding_dim, 1, critic_hidden_dims, activation)
+        print(f"Critic MLP: {self.critic}")
+
+        # Critic observation normalization (only for 1D critic observations)
+        self.critic_obs_normalization = critic_obs_normalization
+        if critic_obs_normalization:
+            self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs_1d)
+        else:
+            self.critic_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.state_dependent_std:
+            torch.nn.init.zeros_(self.actor[-2].weight[num_actions:])
+            if self.noise_std_type == "scalar":
+                torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std)
+            elif self.noise_std_type == "log":
+                torch.nn.init.constant_(
+                    self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7))
+                )
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            if self.noise_std_type == "scalar":
+                self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+            elif self.noise_std_type == "log":
+                self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    def _update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Tensor]) -> None:
+        if self.actor_cnns is not None:
+            # Encode the 2D actor observations
+            cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_groups_2d]
+            cnn_enc = torch.cat(cnn_enc_list, dim=-1)
+            # Concatenate to the MLP observations
+            mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
+
+        super()._update_distribution(mlp_obs)
+
+    def act(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
+        mlp_obs, cnn_obs = self.get_actor_obs(obs)
+        mlp_obs = self.actor_obs_normalizer(mlp_obs)
+        self._update_distribution(mlp_obs, cnn_obs)
+        return self.distribution.sample()  # type: ignore
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        mlp_obs, cnn_obs = self.get_actor_obs(obs)
+        mlp_obs = self.actor_obs_normalizer(mlp_obs)
+
+        if self.actor_cnns is not None:
+            # Encode the 2D actor observations
+            cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_groups_2d]
+            cnn_enc = torch.cat(cnn_enc_list, dim=-1)
+            # Concatenate to the MLP observations
+            mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
+
+        if self.state_dependent_std:
+            return self.actor(mlp_obs)[..., 0, :]
+        else:
+            return self.actor(mlp_obs)
+
+    def evaluate(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
+        mlp_obs, cnn_obs = self.get_critic_obs(obs)
+        mlp_obs = self.critic_obs_normalizer(mlp_obs)
+
+        if self.critic_cnns is not None:
+            # Encode the 2D critic observations
+            cnn_enc_list = [self.critic_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.critic_obs_groups_2d]
+            cnn_enc = torch.cat(cnn_enc_list, dim=-1)
+            # Concatenate to the MLP observations
+            mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
+
+        return self.critic(mlp_obs)
+
+    def get_actor_obs(self, obs: TensorDict) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        obs_list_1d = [obs[obs_group] for obs_group in self.actor_obs_groups_1d]
+        obs_dict_2d = {}
+        for obs_group in self.actor_obs_groups_2d:
+            obs_dict_2d[obs_group] = obs[obs_group]
+        return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
+
+    def get_critic_obs(self, obs: TensorDict) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        obs_list_1d = [obs[obs_group] for obs_group in self.critic_obs_groups_1d]
+        obs_dict_2d = {}
+        for obs_group in self.critic_obs_groups_2d:
+            obs_dict_2d[obs_group] = obs[obs_group]
+        return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.actor_obs_normalization:
+            actor_obs, _ = self.get_actor_obs(obs)
+            self.actor_obs_normalizer.update(actor_obs)
+        if self.critic_obs_normalization:
+            critic_obs, _ = self.get_critic_obs(obs)
+            self.critic_obs_normalizer.update(critic_obs)

source/rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py

@@ -0,0 +1,281 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any, NoReturn
+from rsl_rl.networks.moe import MLP
+
+from rsl_rl.networks import EmpiricalNormalization, L2Norm, SimNorm, MoE
+
+class StudentMoEEncoder(nn.Module):
+    def __init__(
+        self,
+        expert_num,
+        input_dim,
+        hidden_dims,
+        output_dim,
+        activation='elu',
+        norm_type='l2norm',
+    ):
+        super().__init__()
+        self.norm_layer = L2Norm() if norm_type == 'l2norm' else SimNorm()
+        self.moe = MoE(
+            expert_num=expert_num,
+            input_dim=input_dim,
+            hidden_dims=hidden_dims,
+            output_dim=output_dim,
+            activation=activation,
+        )
+    
+    def forward(self, obs):
+        latent, weights = self.moe(obs)
+        latent = self.norm_layer(latent)
+        return latent, weights
+
+class ActorCriticMoECTS(nn.Module):
+    is_recurrent: bool = False
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        actor_obs_normalization: bool = False,
+        critic_obs_normalization: bool = False,
+        actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        teacher_encoder_hidden_dims: tuple[int] | list[int] = [512, 256],
+        student_encoder_hidden_dims: tuple[int] | list[int] = [512, 256, 128],
+        expert_num: int = 8,
+        activation: str = "elu",
+        init_noise_std: float = 1.0,
+        noise_std_type: str = "scalar",
+        state_dependent_std: bool = False,
+        latent_dim: int = 32,
+        norm_type: str = 'l2norm',
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if kwargs:
+            print(
+                "ActorCriticMoECTS.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs])
+            )
+        assert norm_type in ['l2norm', 'simnorm'], f"Normalization type {norm_type} not supported!"
+        assert "policy" in obs.keys() and "critic" in obs.keys() and "single_obs" in obs.keys(), \
+            "obs must contain 'policy', 'critic' and 'single_obs' keys for ActorCriticMoECTS."
+        super().__init__()
+        
+        self.num_actions = num_actions
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_actor_obs = 0
+        for obs_group in obs_groups["policy"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCriticMoECTS module only supports 1D observations."
+            num_actor_obs += obs[obs_group].shape[-1]
+        num_critic_obs = 0
+        for obs_group in obs_groups["critic"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCriticMoECTS module only supports 1D observations."
+            num_critic_obs += obs[obs_group].shape[-1]
+        
+        # MLP input dimensions (teacher, student, actor, critic)
+        self.num_actor_obs = num_actor_obs
+        self.num_single_obs = obs['single_obs'].shape[-1]
+        mlp_input_dim_t = num_critic_obs
+        mlp_input_dim_s = num_actor_obs
+        mlp_input_dim_a = latent_dim + self.num_single_obs
+        mlp_input_dim_c = latent_dim + num_critic_obs
+
+        # Teacher encoder
+        self.teacher_encoder = nn.Sequential(
+            MLP(mlp_input_dim_t, latent_dim, teacher_encoder_hidden_dims, activation=activation),
+            L2Norm() if norm_type == 'l2norm' else SimNorm()
+        )
+        print(f"Teacher Encoder: {self.teacher_encoder}")
+        
+        # Student MoE encoder
+        self.student_moe_encoder = StudentMoEEncoder(
+            expert_num=expert_num,
+            input_dim=mlp_input_dim_s,
+            hidden_dims=student_encoder_hidden_dims,
+            output_dim=latent_dim,
+            activation=activation,
+            norm_type=norm_type,
+        )
+        print(f"Student MoE Encoder: {self.student_moe_encoder}")
+        
+        # Actor
+        self.state_dependent_std = state_dependent_std
+        if self.state_dependent_std:
+            self.actor = MLP(mlp_input_dim_a, [2, num_actions], actor_hidden_dims, activation)
+        else:
+            self.actor = MLP(mlp_input_dim_a, num_actions, actor_hidden_dims, activation)
+        print(f"Actor MLP: {self.actor}")
+
+        # Actor observation normalization
+        self.actor_obs_normalization = actor_obs_normalization
+        if actor_obs_normalization:
+            self.actor_obs_normalizer = EmpiricalNormalization(self.num_actor_obs)
+            self.single_obs_normalizer = EmpiricalNormalization(self.num_single_obs)
+        else:
+            self.actor_obs_normalizer = torch.nn.Identity()
+            self.single_obs_normalizer = torch.nn.Identity()
+
+        # Critic
+        self.critic = MLP(mlp_input_dim_c, 1, critic_hidden_dims, activation)
+        print(f"Critic MLP: {self.critic}")
+
+        # Critic observation normalization
+        self.critic_obs_normalization = critic_obs_normalization
+        if critic_obs_normalization:
+            self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs)
+        else:
+            self.critic_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.state_dependent_std:
+            torch.nn.init.zeros_(self.actor[-2].weight[num_actions:])
+            if self.noise_std_type == "scalar":
+                torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std)
+            elif self.noise_std_type == "log":
+                torch.nn.init.constant_(
+                    self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7))
+                )
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            if self.noise_std_type == "scalar":
+                self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+            elif self.noise_std_type == "log":
+                self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    def reset(self, dones: torch.Tensor | None = None) -> None:
+        pass
+
+    def forward(self) -> NoReturn:
+        raise NotImplementedError
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.distribution.mean
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.distribution.stddev
+
+    @property
+    def entropy(self) -> torch.Tensor:
+        return self.distribution.entropy().sum(dim=-1)
+
+    def _update_distribution(self, latent_and_obs: torch.Tensor) -> None:
+        if self.state_dependent_std:
+            # Compute mean and standard deviation
+            mean_and_std = self.actor(latent_and_obs)
+            if self.noise_std_type == "scalar":
+                mean, std = torch.unbind(mean_and_std, dim=-2)
+            elif self.noise_std_type == "log":
+                mean, log_std = torch.unbind(mean_and_std, dim=-2)
+                std = torch.exp(log_std)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            # Compute mean
+            mean = self.actor(latent_and_obs)
+            # Compute standard deviation
+            if self.noise_std_type == "scalar":
+                std = self.std.expand_as(mean)
+            elif self.noise_std_type == "log":
+                std = torch.exp(self.log_std).expand_as(mean)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # Create distribution
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs: TensorDict, is_teacher: bool, **kwargs: dict[str, Any]) -> torch.Tensor:
+        single_obs = self.single_obs_normalizer(obs['single_obs'])
+        if is_teacher:
+            obs_c = self.get_critic_obs(obs)
+            obs_c = self.critic_obs_normalizer(obs_c)
+            latent = self.teacher_encoder(obs_c)
+        else:
+            with torch.no_grad():
+                obs_a = self.get_actor_obs(obs)
+                obs_a = self.actor_obs_normalizer(obs_a)
+                latent, _ = self.student_moe_encoder(obs_a)
+        latent_and_obs = torch.cat([latent, single_obs], dim=-1)
+        self._update_distribution(latent_and_obs)
+        return self.distribution.sample()
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        single_obs = self.single_obs_normalizer(obs['single_obs'])
+        obs_a = self.get_actor_obs(obs)
+        obs_a = self.actor_obs_normalizer(obs_a)
+        latent, _ = self.student_moe_encoder(obs_a)
+        latent_and_obs = torch.cat([latent, single_obs], dim=-1)
+        if self.state_dependent_std:
+            return self.actor(latent_and_obs)[..., 0, :]
+        else:
+            return self.actor(latent_and_obs)
+
+    def evaluate(self, obs: TensorDict, is_teacher: bool, **kwargs: dict[str, Any]) -> torch.Tensor:
+        obs_c = self.get_critic_obs(obs)
+        obs_c = self.critic_obs_normalizer(obs_c)
+        if is_teacher:
+            latent = self.teacher_encoder(obs_c)
+        else:
+            obs_a = self.get_actor_obs(obs)
+            obs_a = self.actor_obs_normalizer(obs_a)
+            latent, _ = self.student_moe_encoder(obs_a)
+        latent_and_obs = torch.cat([latent.detach(), obs_c], dim=-1)
+        return self.critic(latent_and_obs)
+
+    def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.actor_obs_normalization:
+            actor_obs = self.get_actor_obs(obs)
+            self.actor_obs_normalizer.update(actor_obs)
+            self.single_obs_normalizer.update(obs['single_obs'])
+        if self.critic_obs_normalization:
+            critic_obs = self.get_critic_obs(obs)
+            self.critic_obs_normalizer.update(critic_obs)
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
+        """Load the parameters of the actor-critic model.
+
+        Args:
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
+                :meth:`state_dict` function.
+
+        Returns:
+            Whether this training resumes a previous training. This flag is used by the :func:`load` function of
+                :class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
+        """
+        super().load_state_dict(state_dict, strict=strict)
+        return True

source/rsl_rl/rsl_rl/modules/actor_critic_recurrent.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import warnings
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any, NoReturn
+
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
+
+
+class ActorCriticRecurrent(nn.Module):
+    is_recurrent: bool = True
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        actor_obs_normalization: bool = False,
+        critic_obs_normalization: bool = False,
+        actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        activation: str = "elu",
+        init_noise_std: float = 1.0,
+        noise_std_type: str = "scalar",
+        state_dependent_std: bool = False,
+        rnn_type: str = "lstm",
+        rnn_hidden_dim: int = 256,
+        rnn_num_layers: int = 1,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if "rnn_hidden_size" in kwargs:
+            warnings.warn(
+                "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. "
+                "Please use `rnn_hidden_dim` instead.",
+                DeprecationWarning,
+            )
+            if rnn_hidden_dim == 256:  # Only override if the new argument is at its default
+                rnn_hidden_dim = kwargs.pop("rnn_hidden_size")
+        if kwargs:
+            print(
+                "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()),
+            )
+        super().__init__()
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_actor_obs = 0
+        for obs_group in obs_groups["policy"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations."
+            num_actor_obs += obs[obs_group].shape[-1]
+        num_critic_obs = 0
+        for obs_group in obs_groups["critic"]:
+            assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations."
+            num_critic_obs += obs[obs_group].shape[-1]
+
+        # Actor
+        self.state_dependent_std = state_dependent_std
+        self.memory_a = Memory(num_actor_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
+        if self.state_dependent_std:
+            self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation)
+        else:
+            self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation)
+        print(f"Actor RNN: {self.memory_a}")
+        print(f"Actor MLP: {self.actor}")
+
+        # Actor observation normalization
+        self.actor_obs_normalization = actor_obs_normalization
+        if actor_obs_normalization:
+            self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs)
+        else:
+            self.actor_obs_normalizer = torch.nn.Identity()
+
+        # Critic
+        self.memory_c = Memory(num_critic_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
+        self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation)
+        print(f"Critic RNN: {self.memory_c}")
+        print(f"Critic MLP: {self.critic}")
+
+        # Critic observation normalization
+        self.critic_obs_normalization = critic_obs_normalization
+        if critic_obs_normalization:
+            self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs)
+        else:
+            self.critic_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.state_dependent_std:
+            torch.nn.init.zeros_(self.actor[-2].weight[num_actions:])
+            if self.noise_std_type == "scalar":
+                torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std)
+            elif self.noise_std_type == "log":
+                torch.nn.init.constant_(
+                    self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7))
+                )
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            if self.noise_std_type == "scalar":
+                self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+            elif self.noise_std_type == "log":
+                self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.distribution.mean
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.distribution.stddev
+
+    @property
+    def entropy(self) -> torch.Tensor:
+        return self.distribution.entropy().sum(dim=-1)
+
+    def reset(self, dones: torch.Tensor | None = None) -> None:
+        self.memory_a.reset(dones)
+        self.memory_c.reset(dones)
+
+    def forward(self) -> NoReturn:
+        raise NotImplementedError
+
+    def _update_distribution(self, obs: torch.Tensor) -> None:
+        if self.state_dependent_std:
+            # Compute mean and standard deviation
+            mean_and_std = self.actor(obs)
+            if self.noise_std_type == "scalar":
+                mean, std = torch.unbind(mean_and_std, dim=-2)
+            elif self.noise_std_type == "log":
+                mean, log_std = torch.unbind(mean_and_std, dim=-2)
+                std = torch.exp(log_std)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        else:
+            # Compute mean
+            mean = self.actor(obs)
+            # Compute standard deviation
+            if self.noise_std_type == "scalar":
+                std = self.std.expand_as(mean)
+            elif self.noise_std_type == "log":
+                std = torch.exp(self.log_std).expand_as(mean)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # Create distribution
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None) -> torch.Tensor:
+        obs = self.get_actor_obs(obs)
+        obs = self.actor_obs_normalizer(obs)
+        out_mem = self.memory_a(obs, masks, hidden_state).squeeze(0)
+        self._update_distribution(out_mem)
+        return self.distribution.sample()
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_actor_obs(obs)
+        obs = self.actor_obs_normalizer(obs)
+        out_mem = self.memory_a(obs).squeeze(0)
+        if self.state_dependent_std:
+            return self.actor(out_mem)[..., 0, :]
+        else:
+            return self.actor(out_mem)
+
+    def evaluate(
+        self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None
+    ) -> torch.Tensor:
+        obs = self.get_critic_obs(obs)
+        obs = self.critic_obs_normalizer(obs)
+        out_mem = self.memory_c(obs, masks, hidden_state).squeeze(0)
+        return self.critic(out_mem)
+
+    def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
+        return self.memory_a.hidden_state, self.memory_c.hidden_state
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.actor_obs_normalization:
+            actor_obs = self.get_actor_obs(obs)
+            self.actor_obs_normalizer.update(actor_obs)
+        if self.critic_obs_normalization:
+            critic_obs = self.get_critic_obs(obs)
+            self.critic_obs_normalizer.update(critic_obs)
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
+        """Load the parameters of the actor-critic model.
+
+        Args:
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
+                :meth:`state_dict` function.
+
+        Returns:
+            Whether this training resumes a previous training. This flag is used by the :func:`load` function of
+                :class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
+        """
+        super().load_state_dict(state_dict, strict=strict)
+        return True

source/rsl_rl/rsl_rl/modules/rnd.py

@@ -0,0 +1,210 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+from typing import Any, NoReturn
+
+from rsl_rl.env import VecEnv
+from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
+
+
+class RandomNetworkDistillation(nn.Module):
+    """Implementation of Random Network Distillation (RND) [1].
+
+    References:
+        .. [1] Burda, Yuri, et al. "Exploration by Random Network Distillation." arXiv preprint arXiv:1810.12894 (2018).
+    """
+
+    def __init__(
+        self,
+        num_states: int,
+        obs_groups: dict,
+        num_outputs: int,
+        predictor_hidden_dims: tuple[int] | list[int],
+        target_hidden_dims: tuple[int] | list[int],
+        activation: str = "elu",
+        weight: float = 0.0,
+        state_normalization: bool = False,
+        reward_normalization: bool = False,
+        device: str = "cpu",
+        weight_schedule: dict | None = None,
+    ) -> None:
+        """Initialize the RND module.
+
+        - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization
+          layer.
+        - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted
+          Variation Normalization layer.
+        - If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states
+          is used as the hidden dimension.
+
+        Args:
+            num_states: Number of states/inputs to the predictor and target networks.
+            obs_groups: Dictionary of observation groups.
+            num_outputs: Number of outputs (embedding size) of the predictor and target networks.
+            predictor_hidden_dims: List of hidden dimensions of the predictor network.
+            target_hidden_dims: List of hidden dimensions of the target network.
+            activation: Activation function.
+            weight: Scaling factor of the intrinsic reward.
+            state_normalization: Whether to normalize the input state.
+            reward_normalization: Whether to normalize the intrinsic reward.
+            device: Device to use.
+            weight_schedule: Type of schedule to use for the RND weight parameter.
+                It is a dictionary with the following keys:
+
+                - "mode": Type of schedule to use for the RND weight parameter.
+                    - "constant": Constant weight schedule.
+                    - "step": Step weight schedule.
+                    - "linear": Linear weight schedule.
+
+                For the "step" weight schedule, the following parameters are required:
+
+                - "final_step": Step at which the weight parameter is set to the final value.
+                - "final_value": Final value of the weight parameter.
+
+                For the "linear" weight schedule, the following parameters are required:
+                - "initial_step": Step at which the weight parameter is set to the initial value.
+                - "final_step": Step at which the weight parameter is set to the final value.
+                - "final_value": Final value of the weight parameter.
+        """
+        # Initialize parent class
+        super().__init__()
+
+        # Store parameters
+        self.num_states = num_states
+        self.obs_groups = obs_groups
+        self.num_outputs = num_outputs
+        self.initial_weight = weight
+        self.device = device
+        self.state_normalization = state_normalization
+        self.reward_normalization = reward_normalization
+
+        # Normalization of input gates
+        if state_normalization:
+            self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device)
+        else:
+            self.state_normalizer = torch.nn.Identity()
+
+        # Normalization of intrinsic reward
+        if reward_normalization:
+            self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device)
+        else:
+            self.reward_normalizer = torch.nn.Identity()
+
+        # Counter for the number of updates
+        self.update_counter = 0
+
+        # Resolve weight schedule
+        if weight_schedule is not None:
+            self.weight_scheduler_params = weight_schedule
+            self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule")
+        else:
+            self.weight_scheduler = None
+
+        # Create network architecture
+        self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device)
+        self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device)
+
+        # Make target network not trainable
+        self.target.eval()
+
+    def get_intrinsic_reward(self, obs: TensorDict) -> torch.Tensor:
+        # Note: The counter is updated number of env steps per learning iteration
+        self.update_counter += 1
+        # Extract the rnd state from the observation
+        rnd_state = self.get_rnd_state(obs)
+        rnd_state = self.state_normalizer(rnd_state)
+        # Obtain the embedding of the rnd state from the target and predictor networks
+        target_embedding = self.target(rnd_state).detach()
+        predictor_embedding = self.predictor(rnd_state).detach()
+        # Compute the intrinsic reward as the distance between the embeddings
+        intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1)
+        # Normalize intrinsic reward
+        intrinsic_reward = self.reward_normalizer(intrinsic_reward)
+        # Check the weight schedule
+        if self.weight_scheduler is not None:
+            self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params)
+        else:
+            self.weight = self.initial_weight
+        # Scale intrinsic reward
+        intrinsic_reward *= self.weight
+
+        return intrinsic_reward
+
+    def forward(self, *args: Any, **kwargs: dict[str, Any]) -> NoReturn:
+        raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.")
+
+    def train(self, mode: bool = True) -> RandomNetworkDistillation:
+        # Set module into training mode
+        self.predictor.train(mode)
+        if self.state_normalization:
+            self.state_normalizer.train(mode)
+        if self.reward_normalization:
+            self.reward_normalizer.train(mode)
+        return self
+
+    def eval(self) -> RandomNetworkDistillation:
+        return self.train(False)
+
+    def get_rnd_state(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["rnd_state"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        # Normalize the state
+        if self.state_normalization:
+            rnd_state = self.get_rnd_state(obs)
+            self.state_normalizer.update(rnd_state)
+
+    def _constant_weight_schedule(self, step: int, **kwargs: dict[str, Any]) -> float:
+        return self.initial_weight
+
+    def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs: dict[str, Any]) -> float:
+        return self.initial_weight if step < final_step else final_value
+
+    def _linear_weight_schedule(
+        self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs: dict[str, Any]
+    ) -> float:
+        if step < initial_step:
+            return self.initial_weight
+        elif step > final_step:
+            return final_value
+        else:
+            return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / (
+                final_step - initial_step
+            )
+
+
+def resolve_rnd_config(alg_cfg: dict, obs: TensorDict, obs_groups: dict[str, list[str]], env: VecEnv) -> dict:
+    """Resolve the RND configuration.
+
+    Args:
+        alg_cfg: Algorithm configuration dictionary.
+        obs: Observation dictionary.
+        obs_groups: Observation groups dictionary.
+        env: Environment object.
+
+    Returns:
+        The resolved algorithm configuration dictionary.
+    """
+    # Resolve dimension of rnd gated state
+    if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None:
+        # Get dimension of rnd gated state
+        num_rnd_state = 0
+        for obs_group in obs_groups["rnd_state"]:
+            assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations."
+            num_rnd_state += obs[obs_group].shape[-1]
+        # Add rnd gated state to config
+        alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state
+        alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups
+        # Scale down the rnd weight with timestep
+        alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt
+    else:
+        alg_cfg["rnd_cfg"] = None
+    return alg_cfg

source/rsl_rl/rsl_rl/modules/student_teacher.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any, NoReturn
+
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState
+
+
+class StudentTeacher(nn.Module):
+    is_recurrent: bool = False
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        student_obs_normalization: bool = False,
+        teacher_obs_normalization: bool = False,
+        student_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        teacher_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        activation: str = "elu",
+        init_noise_std: float = 0.1,
+        noise_std_type: str = "scalar",
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if kwargs:
+            print(
+                "StudentTeacher.__init__ got unexpected arguments, which will be ignored: "
+                + str([key for key in kwargs])
+            )
+        super().__init__()
+
+        self.loaded_teacher = False  # Indicates if teacher has been loaded
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_student_obs = 0
+        for obs_group in obs_groups["policy"]:
+            assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
+            num_student_obs += obs[obs_group].shape[-1]
+        num_teacher_obs = 0
+        for obs_group in obs_groups["teacher"]:
+            assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
+            num_teacher_obs += obs[obs_group].shape[-1]
+
+        # Student
+        self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation)
+        print(f"Student MLP: {self.student}")
+
+        # Student observation normalization
+        self.student_obs_normalization = student_obs_normalization
+        if student_obs_normalization:
+            self.student_obs_normalizer = EmpiricalNormalization(num_student_obs)
+        else:
+            self.student_obs_normalizer = torch.nn.Identity()
+
+        # Teacher
+        self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation)
+        print(f"Teacher MLP: {self.teacher}")
+
+        # Teacher observation normalization
+        self.teacher_obs_normalization = teacher_obs_normalization
+        if teacher_obs_normalization:
+            self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs)
+        else:
+            self.teacher_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.noise_std_type == "scalar":
+            self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+        elif self.noise_std_type == "log":
+            self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    def reset(
+        self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, None)
+    ) -> None:
+        pass
+
+    def forward(self) -> NoReturn:
+        raise NotImplementedError
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.distribution.mean
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.distribution.stddev
+
+    @property
+    def entropy(self) -> torch.Tensor:
+        return self.distribution.entropy().sum(dim=-1)
+
+    def _update_distribution(self, obs: TensorDict) -> None:
+        # Compute mean
+        mean = self.student(obs)
+        # Compute standard deviation
+        if self.noise_std_type == "scalar":
+            std = self.std.expand_as(mean)
+        elif self.noise_std_type == "log":
+            std = torch.exp(self.log_std).expand_as(mean)
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # Create distribution
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_student_obs(obs)
+        obs = self.student_obs_normalizer(obs)
+        self._update_distribution(obs)
+        return self.distribution.sample()
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_student_obs(obs)
+        obs = self.student_obs_normalizer(obs)
+        return self.student(obs)
+
+    def evaluate(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_teacher_obs(obs)
+        obs = self.teacher_obs_normalizer(obs)
+        with torch.no_grad():
+            return self.teacher(obs)
+
+    def get_student_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
+        return None, None
+
+    def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
+        pass
+
+    def train(self, mode: bool = True) -> None:
+        super().train(mode)
+        # Make sure teacher is in eval mode
+        self.teacher.eval()
+        self.teacher_obs_normalizer.eval()
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.student_obs_normalization:
+            student_obs = self.get_student_obs(obs)
+            self.student_obs_normalizer.update(student_obs)
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
+        """Load the parameters of the student and teacher networks.
+
+        Args:
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
+                :meth:`state_dict` function.
+
+        Returns:
+            Whether this training resumes a previous training. This flag is used by the :func:`load` function of
+                :class:`OnPolicyRunner` to determine how to load further parameters.
+        """
+        # Check if state_dict contains teacher and student or just teacher parameters
+        if any("actor" in key for key in state_dict):  # Load parameters from rl training
+            # Rename keys to match teacher and remove critic parameters
+            teacher_state_dict = {}
+            teacher_obs_normalizer_state_dict = {}
+            for key, value in state_dict.items():
+                if "actor." in key:
+                    teacher_state_dict[key.replace("actor.", "")] = value
+                if "actor_obs_normalizer." in key:
+                    teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value
+            self.teacher.load_state_dict(teacher_state_dict, strict=strict)
+            self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict)
+            # Set flag for successfully loading the parameters
+            self.loaded_teacher = True
+            self.teacher.eval()
+            self.teacher_obs_normalizer.eval()
+            return False  # Training does not resume
+        elif any("student" in key for key in state_dict):  # Load parameters from distillation training
+            super().load_state_dict(state_dict, strict=strict)
+            # Set flag for successfully loading the parameters
+            self.loaded_teacher = True
+            self.teacher.eval()
+            self.teacher_obs_normalizer.eval()
+            return True  # Training resumes
+        else:
+            raise ValueError("state_dict does not contain student or teacher parameters")

source/rsl_rl/rsl_rl/modules/student_teacher_recurrent.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import warnings
+from tensordict import TensorDict
+from torch.distributions import Normal
+from typing import Any, NoReturn
+
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
+
+
+class StudentTeacherRecurrent(nn.Module):
+    is_recurrent: bool = True
+
+    def __init__(
+        self,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        student_obs_normalization: bool = False,
+        teacher_obs_normalization: bool = False,
+        student_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        teacher_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        activation: str = "elu",
+        init_noise_std: float = 0.1,
+        noise_std_type: str = "scalar",
+        rnn_type: str = "lstm",
+        rnn_hidden_dim: int = 256,
+        rnn_num_layers: int = 1,
+        teacher_recurrent: bool = False,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        if "rnn_hidden_size" in kwargs:
+            warnings.warn(
+                "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. "
+                "Please use `rnn_hidden_dim` instead.",
+                DeprecationWarning,
+            )
+            if rnn_hidden_dim == 256:  # Only override if the new argument is at its default
+                rnn_hidden_dim = kwargs.pop("rnn_hidden_size")
+        if kwargs:
+            print(
+                "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: "
+                + str(kwargs.keys()),
+            )
+        super().__init__()
+
+        self.loaded_teacher = False  # Indicates if teacher has been loaded
+        self.teacher_recurrent = teacher_recurrent  # Indicates if teacher is recurrent too
+
+        # Get the observation dimensions
+        self.obs_groups = obs_groups
+        num_student_obs = 0
+        for obs_group in obs_groups["policy"]:
+            assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
+            num_student_obs += obs[obs_group].shape[-1]
+        num_teacher_obs = 0
+        for obs_group in obs_groups["teacher"]:
+            assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
+            num_teacher_obs += obs[obs_group].shape[-1]
+
+        # Student
+        self.memory_s = Memory(num_student_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
+        self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation)
+        print(f"Student RNN: {self.memory_s}")
+        print(f"Student MLP: {self.student}")
+
+        # Student observation normalization
+        self.student_obs_normalization = student_obs_normalization
+        if student_obs_normalization:
+            self.student_obs_normalizer = EmpiricalNormalization(num_student_obs)
+        else:
+            self.student_obs_normalizer = torch.nn.Identity()
+
+        # Teacher
+        if self.teacher_recurrent:
+            self.memory_t = Memory(num_teacher_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
+        teacher_input_dim = rnn_hidden_dim if self.teacher_recurrent else num_teacher_obs
+        self.teacher = MLP(teacher_input_dim, num_actions, teacher_hidden_dims, activation)
+        if self.teacher_recurrent:
+            print(f"Teacher RNN: {self.memory_t}")
+        print(f"Teacher MLP: {self.teacher}")
+
+        # Teacher observation normalization
+        self.teacher_obs_normalization = teacher_obs_normalization
+        if teacher_obs_normalization:
+            self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs)
+        else:
+            self.teacher_obs_normalizer = torch.nn.Identity()
+
+        # Action noise
+        self.noise_std_type = noise_std_type
+        if self.noise_std_type == "scalar":
+            self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+        elif self.noise_std_type == "log":
+            self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution
+        # Note: Populated in update_distribution
+        self.distribution = None
+
+        # Disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+    def reset(
+        self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, None)
+    ) -> None:
+        self.memory_s.reset(dones, hidden_states[0])
+        if self.teacher_recurrent:
+            self.memory_t.reset(dones, hidden_states[1])
+
+    def forward(self) -> NoReturn:
+        raise NotImplementedError
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.distribution.mean
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.distribution.stddev
+
+    @property
+    def entropy(self) -> torch.Tensor:
+        return self.distribution.entropy().sum(dim=-1)
+
+    def _update_distribution(self, obs: TensorDict) -> None:
+        # Compute mean
+        mean = self.student(obs)
+        # Compute standard deviation
+        if self.noise_std_type == "scalar":
+            std = self.std.expand_as(mean)
+        elif self.noise_std_type == "log":
+            std = torch.exp(self.log_std).expand_as(mean)
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # Create distribution
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_student_obs(obs)
+        obs = self.student_obs_normalizer(obs)
+        out_mem = self.memory_s(obs).squeeze(0)
+        self._update_distribution(out_mem)
+        return self.distribution.sample()
+
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_student_obs(obs)
+        obs = self.student_obs_normalizer(obs)
+        out_mem = self.memory_s(obs).squeeze(0)
+        return self.student(out_mem)
+
+    def evaluate(self, obs: TensorDict) -> torch.Tensor:
+        obs = self.get_teacher_obs(obs)
+        obs = self.teacher_obs_normalizer(obs)
+        with torch.no_grad():
+            if self.teacher_recurrent:
+                self.memory_t.eval()
+                obs = self.memory_t(obs).squeeze(0)
+            return self.teacher(obs)
+
+    def get_student_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
+        obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
+        return torch.cat(obs_list, dim=-1)
+
+    def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
+        if self.teacher_recurrent:
+            return self.memory_s.hidden_state, self.memory_t.hidden_state
+        else:
+            return self.memory_s.hidden_state, None
+
+    def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
+        self.memory_s.detach_hidden_state(dones)
+        if self.teacher_recurrent:
+            self.memory_t.detach_hidden_state(dones)
+
+    def train(self, mode: bool = True) -> None:
+        super().train(mode)
+        # Make sure teacher is in eval mode
+        self.teacher.eval()
+        self.teacher_obs_normalizer.eval()
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        if self.student_obs_normalization:
+            student_obs = self.get_student_obs(obs)
+            self.student_obs_normalizer.update(student_obs)
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
+        """Load the parameters of the student and teacher networks.
+
+        Args:
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
+                :meth:`state_dict` function.
+
+        Returns:
+            Whether this training resumes a previous training. This flag is used by the :func:`load` function of
+                :class:`OnPolicyRunner` to determine how to load further parameters.
+        """
+        # Check if state_dict contains teacher and student or just teacher parameters
+        if any("actor" in key for key in state_dict):  # Load parameters from rl training
+            # Rename keys to match teacher and remove critic parameters
+            teacher_state_dict = {}
+            teacher_obs_normalizer_state_dict = {}
+            for key, value in state_dict.items():
+                if "actor." in key:
+                    teacher_state_dict[key.replace("actor.", "")] = value
+                if "actor_obs_normalizer." in key:
+                    teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value
+            self.teacher.load_state_dict(teacher_state_dict, strict=strict)
+            self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict)
+            # Also load recurrent memory if teacher is recurrent
+            if self.teacher_recurrent:
+                memory_t_state_dict = {}
+                for key, value in state_dict.items():
+                    if "memory_a." in key:
+                        memory_t_state_dict[key.replace("memory_a.", "")] = value
+                self.memory_t.load_state_dict(memory_t_state_dict, strict=strict)
+            # Set flag for successfully loading the parameters
+            self.loaded_teacher = True
+            self.teacher.eval()
+            self.teacher_obs_normalizer.eval()
+            return False  # Training does not resume
+        elif any("student" in key for key in state_dict):  # Load parameters from distillation training
+            super().load_state_dict(state_dict, strict=strict)
+            # Set flag for successfully loading the parameters
+            self.loaded_teacher = True
+            self.teacher.eval()
+            self.teacher_obs_normalizer.eval()
+            return True  # Training resumes
+        else:
+            raise ValueError("state_dict does not contain student or teacher parameters")

source/rsl_rl/rsl_rl/modules/symmetry.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from rsl_rl.env import VecEnv
+
+
+def resolve_symmetry_config(alg_cfg: dict, env: VecEnv) -> dict:
+    """Resolve the symmetry configuration.
+
+    Args:
+        alg_cfg: Algorithm configuration dictionary.
+        env: Environment object.
+
+    Returns:
+        The resolved algorithm configuration dictionary.
+    """
+    # If using symmetry then pass the environment config object
+    # Note: This is used by the symmetry function for handling different observation terms
+    if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None:
+        alg_cfg["symmetry_cfg"]["_env"] = env
+    else:
+        alg_cfg["symmetry_cfg"] = None
+    return alg_cfg