init commit.

source/rsl_rl/rsl_rl/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Main module for the rsl_rl package."""

source/rsl_rl/rsl_rl/algorithms/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Implementation of different learning algorithms."""
+
+from .distillation import Distillation
+from .ppo import PPO
+from .moe_cts import MoECTS
+
+__all__ = ["PPO", "Distillation", "MoECTS"]

source/rsl_rl/rsl_rl/algorithms/distillation.py

@@ -0,0 +1,179 @@
+# 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 rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_optimizer
+
+
+class Distillation:
+    """Distillation algorithm for training a student model to mimic a teacher model."""
+
+    policy: StudentTeacher | StudentTeacherRecurrent
+    """The student teacher model."""
+
+    def __init__(
+        self,
+        policy: StudentTeacher | StudentTeacherRecurrent,
+        storage: RolloutStorage,
+        num_learning_epochs: int = 1,
+        gradient_length: int = 15,
+        learning_rate: float = 1e-3,
+        max_grad_norm: float | None = None,
+        loss_type: str = "mse",
+        optimizer: str = "adam",
+        device: str = "cpu",
+        # Distributed training parameters
+        multi_gpu_cfg: dict | None = None,
+    ) -> None:
+        # Device-related parameters
+        self.device = device
+        self.is_multi_gpu = multi_gpu_cfg is not None
+
+        # Multi-GPU parameters
+        if multi_gpu_cfg is not None:
+            self.gpu_global_rank = multi_gpu_cfg["global_rank"]
+            self.gpu_world_size = multi_gpu_cfg["world_size"]
+        else:
+            self.gpu_global_rank = 0
+            self.gpu_world_size = 1
+
+        # Distillation components
+        self.policy = policy
+        self.policy.to(self.device)
+
+        # Create the optimizer
+        self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate)
+
+        # Add storage
+        self.storage = storage
+        self.transition = RolloutStorage.Transition()
+        self.last_hidden_states = (None, None)
+
+        # Distillation parameters
+        self.num_learning_epochs = num_learning_epochs
+        self.gradient_length = gradient_length
+        self.learning_rate = learning_rate
+        self.max_grad_norm = max_grad_norm
+
+        # Initialize the loss function
+        loss_fn_dict = {
+            "mse": nn.functional.mse_loss,
+            "huber": nn.functional.huber_loss,
+        }
+        if loss_type in loss_fn_dict:
+            self.loss_fn = loss_fn_dict[loss_type]
+        else:
+            raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}")
+
+        self.num_updates = 0
+
+    def act(self, obs: TensorDict) -> torch.Tensor:
+        # Compute the actions
+        self.transition.actions = self.policy.act(obs).detach()
+        self.transition.privileged_actions = self.policy.evaluate(obs).detach()
+        # Record the observations
+        self.transition.observations = obs
+        return self.transition.actions
+
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ) -> None:
+        # Update the normalizers
+        self.policy.update_normalization(obs)
+        # Record the rewards and dones
+        self.transition.rewards = rewards
+        self.transition.dones = dones
+        # Record the transition
+        self.storage.add_transition(self.transition)
+        self.transition.clear()
+        self.policy.reset(dones)
+
+    def compute_returns(self, obs: TensorDict) -> None:
+        # Not needed for distillation
+        pass
+
+    def update(self) -> dict[str, float]:
+        self.num_updates += 1
+        mean_behavior_loss = 0
+        loss = 0
+        cnt = 0
+
+        for epoch in range(self.num_learning_epochs):
+            self.policy.reset(hidden_states=self.last_hidden_states)
+            self.policy.detach_hidden_states()
+            for obs, _, privileged_actions, dones in self.storage.generator():
+                # Inference of the student for gradient computation
+                actions = self.policy.act_inference(obs)
+
+                # Behavior cloning loss
+                behavior_loss = self.loss_fn(actions, privileged_actions)
+
+                # Total loss
+                loss = loss + behavior_loss
+                mean_behavior_loss += behavior_loss.item()
+                cnt += 1
+
+                # Gradient step
+                if cnt % self.gradient_length == 0:
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    if self.is_multi_gpu:
+                        self.reduce_parameters()
+                    if self.max_grad_norm:
+                        nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm)
+                    self.optimizer.step()
+                    self.policy.detach_hidden_states()
+                    loss = 0
+
+                # Reset dones
+                self.policy.reset(dones.view(-1))
+                self.policy.detach_hidden_states(dones.view(-1))
+
+        mean_behavior_loss /= cnt
+        self.storage.clear()
+        self.last_hidden_states = self.policy.get_hidden_states()
+        self.policy.detach_hidden_states()
+
+        # Construct the loss dictionary
+        loss_dict = {"behavior": mean_behavior_loss}
+
+        return loss_dict
+
+    def broadcast_parameters(self) -> None:
+        """Broadcast model parameters to all GPUs."""
+        # Obtain the model parameters on current GPU
+        model_params = [self.policy.state_dict()]
+        # Broadcast the model parameters
+        torch.distributed.broadcast_object_list(model_params, src=0)
+        # Load the model parameters on all GPUs from source GPU
+        self.policy.load_state_dict(model_params[0])
+
+    def reduce_parameters(self) -> None:
+        """Collect gradients from all GPUs and average them.
+
+        This function is called after the backward pass to synchronize the gradients across all GPUs.
+        """
+        # Create a tensor to store the gradients
+        grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None]
+        all_grads = torch.cat(grads)
+        # Average the gradients across all GPUs
+        torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM)
+        all_grads /= self.gpu_world_size
+        # Update the gradients for all parameters with the reduced gradients
+        offset = 0
+        for param in self.policy.parameters():
+            if param.grad is not None:
+                numel = param.numel()
+                # Copy data back from shared buffer
+                param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data))
+                # Update the offset for the next parameter
+                offset += numel

source/rsl_rl/rsl_rl/algorithms/moe_cts.py

@@ -0,0 +1,484 @@
+# 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 torch.optim as optim
+from itertools import chain
+from tensordict import TensorDict
+import itertools
+
+from rsl_rl.modules import ActorCriticMoECTS
+from rsl_rl.modules.rnd import RandomNetworkDistillation
+from rsl_rl.storage import RolloutStorageCTS
+
+
+class MoECTS:
+    """Concurrent Teacher-Student algorithm (https://arxiv.org/abs/2405.10830) with MoE."""
+
+    policy: ActorCriticMoECTS
+    """The actor critic module."""
+
+    def __init__(
+        self,
+        policy: ActorCriticMoECTS,
+        storage: RolloutStorageCTS,
+        num_envs: int,
+        num_learning_epochs: int = 5,
+        num_mini_batches: int = 4,
+        clip_param: float = 0.2,
+        gamma: float = 0.99,
+        lam: float = 0.95,
+        betas: tuple = (0.9, 0.999),
+        weight_decay: float = 0.0,
+        value_loss_coef: float = 1.0,
+        entropy_coef: float = 0.01,
+        load_balance_coef: float = 0.01,
+        learning_rate: float = 0.001,
+        student_encoder_learning_rate: float = 0.001,
+        max_grad_norm: float = 1.0,
+        use_clipped_value_loss: bool = True,
+        schedule: str = "adaptive",
+        desired_kl: float = 0.01,
+        teacher_env_ratio: float = 0.75,
+        normalize_advantage_per_mini_batch: bool = False,
+        device: str = "cpu",
+        # RND parameters
+        rnd_cfg: dict | None = None,
+        # Symmetry parameters
+        symmetry_cfg: dict | None = None,
+        # Distributed training parameters
+        multi_gpu_cfg: dict | None = None,
+    ) -> None:
+        assert isinstance(policy, ActorCriticMoECTS), "Policy must be an instance of ActorCriticMoECTS."
+        assert not policy.is_recurrent, "Recurrent policies are not supported yet for MoECTS."
+        # Device-related parameters
+        self.device = device
+        self.is_multi_gpu = multi_gpu_cfg is not None
+
+        # Multi-GPU parameters
+        if multi_gpu_cfg is not None:
+            self.gpu_global_rank = multi_gpu_cfg["global_rank"]
+            self.gpu_world_size = multi_gpu_cfg["world_size"]
+        else:
+            self.gpu_global_rank = 0
+            self.gpu_world_size = 1
+
+        # RND components
+        if rnd_cfg:
+            # Extract parameters used in ppo
+            rnd_lr = rnd_cfg.pop("learning_rate", 1e-3)
+            # Create RND module
+            self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg)
+            # Create RND optimizer
+            params = self.rnd.predictor.parameters()
+            self.rnd_optimizer = optim.Adam(params, lr=rnd_lr)
+        else:
+            self.rnd = None
+            self.rnd_optimizer = None
+
+        # Symmetry components
+        if symmetry_cfg is not None:
+            print("[WARNING] `symmetry_cfg` detected, but MoECTS does not currently support symmetry; the configuration will be ignored.")
+        self.symmetry = None
+
+        # PPO components
+        self.policy = policy
+        self.policy.to(self.device)
+
+        # Create the optimizer
+        params1 = [
+            {"params": self.policy.teacher_encoder.parameters()},
+            {"params": self.policy.critic.parameters()},
+            {"params": self.policy.actor.parameters()},
+            {"params": getattr(self.policy, 'std', getattr(self.policy, 'log_std', []))}
+        ]
+        self.optimizer = optim.Adam(params1, lr=learning_rate, betas=betas, weight_decay=weight_decay)
+        self.optimizer_stu_enc = optim.Adam(self.policy.student_moe_encoder.parameters(), lr=student_encoder_learning_rate, betas=betas, weight_decay=weight_decay)
+
+        # Add storage
+        self.storage = storage
+        self.transition = RolloutStorageCTS.Transition()
+
+        # MoECTS & PPO parameters
+        self.clip_param = clip_param
+        self.num_learning_epochs = num_learning_epochs
+        self.num_mini_batches = num_mini_batches
+        self.value_loss_coef = value_loss_coef
+        self.entropy_coef = entropy_coef
+        self.load_balance_coef = load_balance_coef
+        self.gamma = gamma
+        self.lam = lam
+        self.max_grad_norm = max_grad_norm
+        self.use_clipped_value_loss = use_clipped_value_loss
+        self.desired_kl = desired_kl
+        self.schedule = schedule
+        self.learning_rate = learning_rate
+        self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch
+        
+        # Teacher-student environment split
+        self.teacher_num_envs = max(int(num_envs * teacher_env_ratio), 1)
+        self.student_num_envs = num_envs - self.teacher_num_envs
+        student_env_ratio = 1 - teacher_env_ratio
+        self.teacher_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) != 0], device=self.device)
+        self.student_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) == 0], device=self.device)
+        assert len(self.teacher_env_idxs) == self.teacher_num_envs, f"{len(self.teacher_env_idxs)=} != {self.teacher_num_envs=}"
+        assert len(self.student_env_idxs) == self.student_num_envs, f"{len(self.student_env_idxs)=} != {self.student_num_envs=}"
+        
+    def act(self, obs: TensorDict) -> torch.Tensor:
+        # Compute the actions and values
+        def _get_results(obs, is_teacher):
+            actions = self.policy.act(obs, is_teacher)
+            return (
+                actions.detach(),
+                self.policy.evaluate(obs, is_teacher).detach(),
+                self.policy.get_actions_log_prob(actions).detach(),
+                self.policy.action_mean.detach(),
+                self.policy.action_std.detach(),
+            )
+        ti, si = self.teacher_env_idxs, self.student_env_idxs
+        teacher_results = _get_results(obs[ti], is_teacher=True)
+        student_results = _get_results(obs[si], is_teacher=False)
+        results = []
+        for x1, x2 in zip(teacher_results, student_results):
+            results.append(torch.cat([x1, x2], dim=0))
+        self.transition.actions = results[0]
+        self.transition.values = results[1]
+        self.transition.actions_log_prob = results[2]
+        self.transition.action_mean = results[3]
+        self.transition.action_sigma = results[4]
+                
+        # Record observations before env.step()
+        self.transition.observations = torch.cat([obs[ti], obs[si]], dim=0)
+        
+        # Reconstruct the actions in the original order
+        reordered_actions = torch.zeros_like(self.transition.actions)
+        reordered_actions[ti] = self.transition.actions[:self.teacher_num_envs]
+        reordered_actions[si] = self.transition.actions[self.teacher_num_envs:]
+        return reordered_actions
+
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ) -> None:
+        # Update the normalizers
+        self.policy.update_normalization(obs)
+        if self.rnd:
+            self.rnd.update_normalization(obs)
+
+        # Record the rewards and dones
+        # Note: We clone here because later on we bootstrap the rewards based on timeouts
+        ti, si = self.teacher_env_idxs, self.student_env_idxs
+        rewards = rewards.clone()
+        self.transition.rewards = torch.cat([rewards[ti], rewards[si]], dim=0)
+        self.transition.dones = torch.cat([dones[ti], dones[si]], dim=0)
+
+        # Compute the intrinsic rewards and add to extrinsic rewards
+        if self.rnd:
+            # Compute the intrinsic rewards
+            reordered_obs = torch.cat([obs[ti], obs[si]], dim=0)
+            self.intrinsic_rewards = self.rnd.get_intrinsic_reward(reordered_obs)
+            # Add intrinsic rewards to extrinsic rewards
+            self.transition.rewards += self.intrinsic_rewards
+
+        # Bootstrapping on time outs
+        if "time_outs" in extras:
+            time_outs = extras["time_outs"].to(self.device)
+            reordered_time_outs = torch.cat([time_outs[ti], time_outs[si]], dim=0)
+            self.transition.rewards += self.gamma * torch.squeeze(
+                self.transition.values * reordered_time_outs.unsqueeze(1).to(self.device), 1
+            )
+
+        # Record the transition
+        self.storage.add_transition(self.transition)
+        self.transition.clear()
+        self.policy.reset(dones)
+
+    def compute_returns(self, obs: TensorDict) -> None:
+        st = self.storage
+        # Compute value for the last step
+        ti, si = self.teacher_env_idxs, self.student_env_idxs
+        last_values = torch.cat([
+            self.policy.evaluate(obs[ti], is_teacher=True).detach(),
+            self.policy.evaluate(obs[si], is_teacher=False).detach(),
+        ], dim=0)
+        # Compute returns and advantages
+        advantage = 0
+        for step in reversed(range(st.num_transitions_per_env)):
+            # If we are at the last step, bootstrap the return value
+            next_values = last_values if step == st.num_transitions_per_env - 1 else st.values[step + 1]
+            # 1 if we are not in a terminal state, 0 otherwise
+            next_is_not_terminal = 1.0 - st.dones[step].float()
+            # TD error: r_t + gamma * V(s_{t+1}) - V(s_t)
+            delta = st.rewards[step] + next_is_not_terminal * self.gamma * next_values - st.values[step]
+            # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1})
+            advantage = delta + next_is_not_terminal * self.gamma * self.lam * advantage
+            # Return: R_t = A(s_t, a_t) + V(s_t)
+            st.returns[step] = advantage + st.values[step]
+        # Compute the advantages
+        st.advantages = st.returns - st.values
+        # Normalize the advantages if per minibatch normalization is not used
+        if not self.normalize_advantage_per_mini_batch:
+            st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std() + 1e-8)
+
+    def update(self) -> dict[str, float]:
+        mean_value_loss = 0
+        mean_surrogate_loss = 0
+        mean_entropy = 0
+        mean_latent_loss = 0
+        mean_load_balance_loss = 0
+        # RND loss
+        mean_rnd_loss = 0 if self.rnd else None
+
+        # Get mini batch generator
+        generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
+        data = list(generator)
+
+        # Iterate over batches
+        teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
+        student_samples = self.student_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
+        for (
+            obs_batch,
+            actions_batch,
+            target_values_batch,
+            advantages_batch,
+            returns_batch,
+            old_actions_log_prob_batch,
+            old_mu_batch,
+            old_sigma_batch,
+            hidden_states_batch,
+            masks_batch,
+        ) in data:
+            original_batch_size = obs_batch.batch_size[0]
+
+            # Check if we should normalize advantages per mini batch
+            if self.normalize_advantage_per_mini_batch:
+                with torch.no_grad():
+                    advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8)
+ 
+            def _get_results(start, end, is_teacher):
+                self.policy.act(obs_batch[start:end], is_teacher)
+                actions_log_prob = self.policy.get_actions_log_prob(actions_batch[start:end])
+                value = self.policy.evaluate(obs_batch[start:end], is_teacher)
+                mu = self.policy.action_mean
+                sigma = self.policy.action_std
+                entropy = self.policy.entropy
+                return actions_log_prob, value, mu, sigma, entropy
+            teacher_results = _get_results(0, teacher_samples, is_teacher=True)
+            student_results = _get_results(teacher_samples, teacher_samples + student_samples, is_teacher=False)
+            results = []
+            for x1, x2 in zip(teacher_results, student_results):
+                results.append(torch.cat([x1, x2], dim=0))
+            actions_log_prob_batch, value_batch, mu_batch, sigma_batch, entropy_batch = results
+
+            # Compute KL divergence and adapt the learning rate
+            if self.desired_kl is not None and self.schedule == "adaptive":
+                with torch.inference_mode():
+                    kl = torch.sum(
+                        torch.log(sigma_batch / old_sigma_batch + 1.0e-5)
+                        + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch))
+                        / (2.0 * torch.square(sigma_batch))
+                        - 0.5,
+                        axis=-1,
+                    )
+                    kl_mean = torch.mean(kl)
+
+                    # Reduce the KL divergence across all GPUs
+                    if self.is_multi_gpu:
+                        torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM)
+                        kl_mean /= self.gpu_world_size
+
+                    # Update the learning rate only on the main process
+                    # TODO: Is this needed? If KL-divergence is the "same" across all GPUs,
+                    #       then the learning rate should be the same across all GPUs.
+                    if self.gpu_global_rank == 0:
+                        if kl_mean > self.desired_kl * 2.0:
+                            self.learning_rate = max(1e-5, self.learning_rate / 1.5)
+                        elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
+                            self.learning_rate = min(1e-2, self.learning_rate * 1.5)
+
+                    # Update the learning rate for all GPUs
+                    if self.is_multi_gpu:
+                        lr_tensor = torch.tensor(self.learning_rate, device=self.device)
+                        torch.distributed.broadcast(lr_tensor, src=0)
+                        self.learning_rate = lr_tensor.item()
+
+                    # Update the learning rate for all parameter groups
+                    for param_group in self.optimizer.param_groups:
+                        param_group["lr"] = self.learning_rate
+
+            # Surrogate loss
+            ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
+            surrogate = -torch.squeeze(advantages_batch) * ratio
+            surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
+                ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
+            )
+            surrogate_losses = torch.max(surrogate, surrogate_clipped)
+            teacher_surrogate_loss = surrogate_losses[:teacher_samples].mean()
+            student_surrogate_loss = surrogate_losses[teacher_samples:].mean()
+            surrogate_loss = teacher_surrogate_loss + student_surrogate_loss
+
+            # Value function loss
+            if self.use_clipped_value_loss:
+                value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(
+                    -self.clip_param, self.clip_param
+                )
+                value_losses = (value_batch - returns_batch).pow(2)
+                value_losses_clipped = (value_clipped - returns_batch).pow(2)
+                value_loss = torch.max(value_losses, value_losses_clipped).mean()
+            else:
+                value_loss = (returns_batch - value_batch).pow(2).mean()
+
+            loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
+
+            # RND loss
+            # TODO: Move this processing to inside RND module.
+            if self.rnd:
+                # Extract the rnd_state
+                # TODO: Check if we still need torch no grad. It is just an affine transformation.
+                with torch.no_grad():
+                    rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size])
+                    rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch)
+                # Predict the embedding and the target
+                predicted_embedding = self.rnd.predictor(rnd_state_batch)
+                target_embedding = self.rnd.target(rnd_state_batch).detach()
+                # Compute the loss as the mean squared error
+                mseloss = torch.nn.MSELoss()
+                rnd_loss = mseloss(predicted_embedding, target_embedding)
+
+            # Compute the gradients for PPO
+            self.optimizer.zero_grad()
+            loss.backward()
+            # Compute the gradients for RND
+            if self.rnd:
+                self.rnd_optimizer.zero_grad()
+                rnd_loss.backward()
+
+            # Collect gradients from all GPUs
+            if self.is_multi_gpu:
+                self.reduce_parameters()
+
+            # Apply the gradients for PPO
+            params_to_clip = itertools.chain.from_iterable(g['params'] for g in self.optimizer.param_groups)
+            nn.utils.clip_grad_norm_(params_to_clip, self.max_grad_norm)
+            self.optimizer.step()
+            # Apply the gradients for RND
+            if self.rnd_optimizer:
+                self.rnd_optimizer.step()
+
+            # Store the losses
+            mean_value_loss += value_loss.item()
+            mean_surrogate_loss += surrogate_loss.item()
+            mean_entropy += entropy_batch.mean().item()
+            # RND loss
+            if mean_rnd_loss is not None:
+                mean_rnd_loss += rnd_loss.item()
+
+        for (
+            obs_batch,
+            actions_batch,
+            target_values_batch,
+            advantages_batch,
+            returns_batch,
+            old_actions_log_prob_batch,
+            old_mu_batch,
+            old_sigma_batch,
+            hidden_states_batch,
+            masks_batch,
+        ) in data:
+            # Student encoder loss
+            obs_a_batch = self.policy.get_actor_obs(obs_batch)
+            obs_a_batch = self.policy.actor_obs_normalizer(obs_a_batch)
+            student_latent, gating_weights = self.policy.student_moe_encoder(obs_a_batch[teacher_samples:])
+            with torch.no_grad():
+                obs_c_batch = self.policy.get_critic_obs(obs_batch)
+                obs_c_batch = self.policy.critic_obs_normalizer(obs_c_batch)
+                teacher_latent = self.policy.teacher_encoder(obs_c_batch[teacher_samples:])
+            latent_loss = (teacher_latent - student_latent).pow(2).mean()
+
+            # Load balance loss
+            mean_usage = torch.mean(gating_weights, dim=0)
+            target_usage = torch.full_like(mean_usage, 1.0 / gating_weights.shape[1])
+            load_balance_loss = torch.mean((mean_usage - target_usage).pow(2))
+            # load_balance_loss = torch.sum(mean_usage.pow(2)) * gating_weights.shape[1]  # Switch Transformer style
+            student_loss = latent_loss + self.load_balance_coef * load_balance_loss
+            
+            self.optimizer_stu_enc.zero_grad()
+            student_loss.backward()
+            nn.utils.clip_grad_norm_(self.policy.student_moe_encoder.parameters(), self.max_grad_norm)
+            self.optimizer_stu_enc.step()
+
+            mean_latent_loss += latent_loss.item()
+            mean_load_balance_loss += load_balance_loss.item()
+
+        # Divide the losses by the number of updates
+        num_updates = self.num_learning_epochs * self.num_mini_batches
+        mean_value_loss /= num_updates
+        mean_surrogate_loss /= num_updates
+        mean_entropy /= num_updates
+        mean_latent_loss /= num_updates
+        mean_load_balance_loss /= num_updates
+        if mean_rnd_loss is not None:
+            mean_rnd_loss /= num_updates
+
+        # Clear the storage
+        self.storage.clear()
+
+        # Construct the loss dictionary
+        loss_dict = {
+            "value": mean_value_loss,
+            "surrogate": mean_surrogate_loss,
+            "entropy": mean_entropy,
+            "mean_latent_loss": mean_latent_loss,
+            "mean_load_balance_loss": mean_load_balance_loss
+        }
+        if self.rnd:
+            loss_dict["rnd"] = mean_rnd_loss
+
+        return loss_dict
+
+    def broadcast_parameters(self) -> None:
+        """Broadcast model parameters to all GPUs."""
+        # Obtain the model parameters on current GPU
+        model_params = [self.policy.state_dict()]
+        if self.rnd:
+            model_params.append(self.rnd.predictor.state_dict())
+        # Broadcast the model parameters
+        torch.distributed.broadcast_object_list(model_params, src=0)
+        # Load the model parameters on all GPUs from source GPU
+        self.policy.load_state_dict(model_params[0])
+        if self.rnd:
+            self.rnd.predictor.load_state_dict(model_params[1])
+
+    def reduce_parameters(self) -> None:
+        """Collect gradients from all GPUs and average them.
+
+        This function is called after the backward pass to synchronize the gradients across all GPUs.
+        """
+        # Create a tensor to store the gradients
+        grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None]
+        if self.rnd:
+            grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None]
+        all_grads = torch.cat(grads)
+
+        # Average the gradients across all GPUs
+        torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM)
+        all_grads /= self.gpu_world_size
+
+        # Get all parameters
+        all_params = self.policy.parameters()
+        if self.rnd:
+            all_params = chain(all_params, self.rnd.parameters())
+
+        # Update the gradients for all parameters with the reduced gradients
+        offset = 0
+        for param in all_params:
+            if param.grad is not None:
+                numel = param.numel()
+                # Copy data back from shared buffer
+                param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data))
+                # Update the offset for the next parameter
+                offset += numel

source/rsl_rl/rsl_rl/algorithms/ppo.py

@@ -0,0 +1,458 @@
+# 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 torch.optim as optim
+from itertools import chain
+from tensordict import TensorDict
+
+from rsl_rl.modules import ActorCritic, ActorCriticCNN, ActorCriticRecurrent
+from rsl_rl.modules.rnd import RandomNetworkDistillation
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_callable
+
+
+class PPO:
+    """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347)."""
+
+    policy: ActorCritic | ActorCriticRecurrent | ActorCriticCNN
+    """The actor critic module."""
+
+    def __init__(
+        self,
+        policy: ActorCritic | ActorCriticRecurrent | ActorCriticCNN,
+        storage: RolloutStorage,
+        num_learning_epochs: int = 5,
+        num_mini_batches: int = 4,
+        clip_param: float = 0.2,
+        gamma: float = 0.99,
+        lam: float = 0.95,
+        value_loss_coef: float = 1.0,
+        entropy_coef: float = 0.01,
+        learning_rate: float = 0.001,
+        max_grad_norm: float = 1.0,
+        use_clipped_value_loss: bool = True,
+        schedule: str = "adaptive",
+        desired_kl: float = 0.01,
+        normalize_advantage_per_mini_batch: bool = False,
+        device: str = "cpu",
+        # RND parameters
+        rnd_cfg: dict | None = None,
+        # Symmetry parameters
+        symmetry_cfg: dict | None = None,
+        # Distributed training parameters
+        multi_gpu_cfg: dict | None = None,
+    ) -> None:
+        # Device-related parameters
+        self.device = device
+        self.is_multi_gpu = multi_gpu_cfg is not None
+
+        # Multi-GPU parameters
+        if multi_gpu_cfg is not None:
+            self.gpu_global_rank = multi_gpu_cfg["global_rank"]
+            self.gpu_world_size = multi_gpu_cfg["world_size"]
+        else:
+            self.gpu_global_rank = 0
+            self.gpu_world_size = 1
+
+        # RND components
+        if rnd_cfg:
+            # Extract parameters used in ppo
+            rnd_lr = rnd_cfg.pop("learning_rate", 1e-3)
+            # Create RND module
+            self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg)
+            # Create RND optimizer
+            params = self.rnd.predictor.parameters()
+            self.rnd_optimizer = optim.Adam(params, lr=rnd_lr)
+        else:
+            self.rnd = None
+            self.rnd_optimizer = None
+
+        # Symmetry components
+        if symmetry_cfg is not None:
+            # Check if symmetry is enabled
+            use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"]
+            # Print that we are not using symmetry
+            if not use_symmetry:
+                print("Symmetry not used for learning. We will use it for logging instead.")
+            # Resolve the data augmentation function (supports string names or direct callables)
+            symmetry_cfg["data_augmentation_func"] = resolve_callable(symmetry_cfg["data_augmentation_func"])
+            # Check valid configuration
+            if not callable(symmetry_cfg["data_augmentation_func"]):
+                raise ValueError(
+                    f"Symmetry configuration exists but the function is not callable: "
+                    f"{symmetry_cfg['data_augmentation_func']}"
+                )
+            # Check if the policy is compatible with symmetry
+            if isinstance(policy, ActorCriticRecurrent):
+                raise ValueError("Symmetry augmentation is not supported for recurrent policies.")
+            # Store symmetry configuration
+            self.symmetry = symmetry_cfg
+        else:
+            self.symmetry = None
+
+        # PPO components
+        self.policy = policy
+        self.policy.to(self.device)
+
+        # Create the optimizer
+        self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate)
+
+        # Add storage
+        self.storage = storage
+        self.transition = RolloutStorage.Transition()
+
+        # PPO parameters
+        self.clip_param = clip_param
+        self.num_learning_epochs = num_learning_epochs
+        self.num_mini_batches = num_mini_batches
+        self.value_loss_coef = value_loss_coef
+        self.entropy_coef = entropy_coef
+        self.gamma = gamma
+        self.lam = lam
+        self.max_grad_norm = max_grad_norm
+        self.use_clipped_value_loss = use_clipped_value_loss
+        self.desired_kl = desired_kl
+        self.schedule = schedule
+        self.learning_rate = learning_rate
+        self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch
+
+    def act(self, obs: TensorDict) -> torch.Tensor:
+        if self.policy.is_recurrent:
+            self.transition.hidden_states = self.policy.get_hidden_states()
+        # Compute the actions and values
+        self.transition.actions = self.policy.act(obs).detach()
+        self.transition.values = self.policy.evaluate(obs).detach()
+        self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach()
+        self.transition.action_mean = self.policy.action_mean.detach()
+        self.transition.action_sigma = self.policy.action_std.detach()
+        # Record observations before env.step()
+        self.transition.observations = obs
+        return self.transition.actions
+
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ) -> None:
+        # Update the normalizers
+        self.policy.update_normalization(obs)
+        if self.rnd:
+            self.rnd.update_normalization(obs)
+
+        # Record the rewards and dones
+        # Note: We clone here because later on we bootstrap the rewards based on timeouts
+        self.transition.rewards = rewards.clone()
+        self.transition.dones = dones
+
+        # Compute the intrinsic rewards and add to extrinsic rewards
+        if self.rnd:
+            # Compute the intrinsic rewards
+            self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs)
+            # Add intrinsic rewards to extrinsic rewards
+            self.transition.rewards += self.intrinsic_rewards
+
+        # Bootstrapping on time outs
+        if "time_outs" in extras:
+            self.transition.rewards += self.gamma * torch.squeeze(
+                self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1
+            )
+
+        # Record the transition
+        self.storage.add_transition(self.transition)
+        self.transition.clear()
+        self.policy.reset(dones)
+
+    def compute_returns(self, obs: TensorDict) -> None:
+        st = self.storage
+        # Compute value for the last step
+        last_values = self.policy.evaluate(obs).detach()
+        # Compute returns and advantages
+        advantage = 0
+        for step in reversed(range(st.num_transitions_per_env)):
+            # If we are at the last step, bootstrap the return value
+            next_values = last_values if step == st.num_transitions_per_env - 1 else st.values[step + 1]
+            # 1 if we are not in a terminal state, 0 otherwise
+            next_is_not_terminal = 1.0 - st.dones[step].float()
+            # TD error: r_t + gamma * V(s_{t+1}) - V(s_t)
+            delta = st.rewards[step] + next_is_not_terminal * self.gamma * next_values - st.values[step]
+            # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1})
+            advantage = delta + next_is_not_terminal * self.gamma * self.lam * advantage
+            # Return: R_t = A(s_t, a_t) + V(s_t)
+            st.returns[step] = advantage + st.values[step]
+        # Compute the advantages
+        st.advantages = st.returns - st.values
+        # Normalize the advantages if per minibatch normalization is not used
+        if not self.normalize_advantage_per_mini_batch:
+            st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std() + 1e-8)
+
+    def update(self) -> dict[str, float]:
+        mean_value_loss = 0
+        mean_surrogate_loss = 0
+        mean_entropy = 0
+        # RND loss
+        mean_rnd_loss = 0 if self.rnd else None
+        # Symmetry loss
+        mean_symmetry_loss = 0 if self.symmetry else None
+
+        # Get mini batch generator
+        if self.policy.is_recurrent:
+            generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
+        else:
+            generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
+
+        # Iterate over batches
+        for (
+            obs_batch,
+            actions_batch,
+            target_values_batch,
+            advantages_batch,
+            returns_batch,
+            old_actions_log_prob_batch,
+            old_mu_batch,
+            old_sigma_batch,
+            hidden_states_batch,
+            masks_batch,
+        ) in generator:
+            num_aug = 1  # Number of augmentations per sample. Starts at 1 for no augmentation.
+            original_batch_size = obs_batch.batch_size[0]
+
+            # Check if we should normalize advantages per mini batch
+            if self.normalize_advantage_per_mini_batch:
+                with torch.no_grad():
+                    advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8)
+
+            # Perform symmetric augmentation
+            if self.symmetry and self.symmetry["use_data_augmentation"]:
+                # Augmentation using symmetry
+                data_augmentation_func = self.symmetry["data_augmentation_func"]
+                # Returned shape: [batch_size * num_aug, ...]
+                obs_batch, actions_batch = data_augmentation_func(
+                    obs=obs_batch,
+                    actions=actions_batch,
+                    env=self.symmetry["_env"],
+                )
+                # Compute number of augmentations per sample
+                num_aug = int(obs_batch.batch_size[0] / original_batch_size)
+                # Repeat the rest of the batch
+                old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1)
+                target_values_batch = target_values_batch.repeat(num_aug, 1)
+                advantages_batch = advantages_batch.repeat(num_aug, 1)
+                returns_batch = returns_batch.repeat(num_aug, 1)
+
+            # Recompute actions log prob and entropy for current batch of transitions
+            # Note: We need to do this because we updated the policy with the new parameters
+            self.policy.act(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[0])
+            actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch)
+            value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[1])
+            # Note: We only keep the entropy of the first augmentation (the original one)
+            mu_batch = self.policy.action_mean[:original_batch_size]
+            sigma_batch = self.policy.action_std[:original_batch_size]
+            entropy_batch = self.policy.entropy[:original_batch_size]
+
+            # Compute KL divergence and adapt the learning rate
+            if self.desired_kl is not None and self.schedule == "adaptive":
+                with torch.inference_mode():
+                    kl = torch.sum(
+                        torch.log(sigma_batch / old_sigma_batch + 1.0e-5)
+                        + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch))
+                        / (2.0 * torch.square(sigma_batch))
+                        - 0.5,
+                        axis=-1,
+                    )
+                    kl_mean = torch.mean(kl)
+
+                    # Reduce the KL divergence across all GPUs
+                    if self.is_multi_gpu:
+                        torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM)
+                        kl_mean /= self.gpu_world_size
+
+                    # Update the learning rate only on the main process
+                    # TODO: Is this needed? If KL-divergence is the "same" across all GPUs,
+                    #       then the learning rate should be the same across all GPUs.
+                    if self.gpu_global_rank == 0:
+                        if kl_mean > self.desired_kl * 2.0:
+                            self.learning_rate = max(1e-5, self.learning_rate / 1.5)
+                        elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
+                            self.learning_rate = min(1e-2, self.learning_rate * 1.5)
+
+                    # Update the learning rate for all GPUs
+                    if self.is_multi_gpu:
+                        lr_tensor = torch.tensor(self.learning_rate, device=self.device)
+                        torch.distributed.broadcast(lr_tensor, src=0)
+                        self.learning_rate = lr_tensor.item()
+
+                    # Update the learning rate for all parameter groups
+                    for param_group in self.optimizer.param_groups:
+                        param_group["lr"] = self.learning_rate
+
+            # Surrogate loss
+            ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
+            surrogate = -torch.squeeze(advantages_batch) * ratio
+            surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
+                ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
+            )
+            surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
+
+            # Value function loss
+            if self.use_clipped_value_loss:
+                value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(
+                    -self.clip_param, self.clip_param
+                )
+                value_losses = (value_batch - returns_batch).pow(2)
+                value_losses_clipped = (value_clipped - returns_batch).pow(2)
+                value_loss = torch.max(value_losses, value_losses_clipped).mean()
+            else:
+                value_loss = (returns_batch - value_batch).pow(2).mean()
+
+            loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
+
+            # Symmetry loss
+            if self.symmetry:
+                # Obtain the symmetric actions
+                # Note: If we did augmentation before then we don't need to augment again
+                if not self.symmetry["use_data_augmentation"]:
+                    data_augmentation_func = self.symmetry["data_augmentation_func"]
+                    obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"])
+                    # Compute number of augmentations per sample
+                    num_aug = int(obs_batch.shape[0] / original_batch_size)
+
+                # Actions predicted by the actor for symmetrically-augmented observations
+                mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone())
+
+                # Compute the symmetrically augmented actions
+                # Note: We are assuming the first augmentation is the original one. We do not use the action_batch from
+                # earlier since that action was sampled from the distribution. However, the symmetry loss is computed
+                # using the mean of the distribution.
+                action_mean_orig = mean_actions_batch[:original_batch_size]
+                _, actions_mean_symm_batch = data_augmentation_func(
+                    obs=None, actions=action_mean_orig, env=self.symmetry["_env"]
+                )
+
+                # Compute the loss
+                mse_loss = torch.nn.MSELoss()
+                symmetry_loss = mse_loss(
+                    mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:]
+                )
+                # Add the loss to the total loss
+                if self.symmetry["use_mirror_loss"]:
+                    loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss
+                else:
+                    symmetry_loss = symmetry_loss.detach()
+
+            # RND loss
+            # TODO: Move this processing to inside RND module.
+            if self.rnd:
+                # Extract the rnd_state
+                # TODO: Check if we still need torch no grad. It is just an affine transformation.
+                with torch.no_grad():
+                    rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size])
+                    rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch)
+                # Predict the embedding and the target
+                predicted_embedding = self.rnd.predictor(rnd_state_batch)
+                target_embedding = self.rnd.target(rnd_state_batch).detach()
+                # Compute the loss as the mean squared error
+                mseloss = torch.nn.MSELoss()
+                rnd_loss = mseloss(predicted_embedding, target_embedding)
+
+            # Compute the gradients for PPO
+            self.optimizer.zero_grad()
+            loss.backward()
+            # Compute the gradients for RND
+            if self.rnd:
+                self.rnd_optimizer.zero_grad()
+                rnd_loss.backward()
+
+            # Collect gradients from all GPUs
+            if self.is_multi_gpu:
+                self.reduce_parameters()
+
+            # Apply the gradients for PPO
+            nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
+            self.optimizer.step()
+            # Apply the gradients for RND
+            if self.rnd_optimizer:
+                self.rnd_optimizer.step()
+
+            # Store the losses
+            mean_value_loss += value_loss.item()
+            mean_surrogate_loss += surrogate_loss.item()
+            mean_entropy += entropy_batch.mean().item()
+            # RND loss
+            if mean_rnd_loss is not None:
+                mean_rnd_loss += rnd_loss.item()
+            # Symmetry loss
+            if mean_symmetry_loss is not None:
+                mean_symmetry_loss += symmetry_loss.item()
+
+        # Divide the losses by the number of updates
+        num_updates = self.num_learning_epochs * self.num_mini_batches
+        mean_value_loss /= num_updates
+        mean_surrogate_loss /= num_updates
+        mean_entropy /= num_updates
+        if mean_rnd_loss is not None:
+            mean_rnd_loss /= num_updates
+        if mean_symmetry_loss is not None:
+            mean_symmetry_loss /= num_updates
+
+        # Clear the storage
+        self.storage.clear()
+
+        # Construct the loss dictionary
+        loss_dict = {
+            "value": mean_value_loss,
+            "surrogate": mean_surrogate_loss,
+            "entropy": mean_entropy,
+        }
+        if self.rnd:
+            loss_dict["rnd"] = mean_rnd_loss
+        if self.symmetry:
+            loss_dict["symmetry"] = mean_symmetry_loss
+
+        return loss_dict
+
+    def broadcast_parameters(self) -> None:
+        """Broadcast model parameters to all GPUs."""
+        # Obtain the model parameters on current GPU
+        model_params = [self.policy.state_dict()]
+        if self.rnd:
+            model_params.append(self.rnd.predictor.state_dict())
+        # Broadcast the model parameters
+        torch.distributed.broadcast_object_list(model_params, src=0)
+        # Load the model parameters on all GPUs from source GPU
+        self.policy.load_state_dict(model_params[0])
+        if self.rnd:
+            self.rnd.predictor.load_state_dict(model_params[1])
+
+    def reduce_parameters(self) -> None:
+        """Collect gradients from all GPUs and average them.
+
+        This function is called after the backward pass to synchronize the gradients across all GPUs.
+        """
+        # Create a tensor to store the gradients
+        grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None]
+        if self.rnd:
+            grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None]
+        all_grads = torch.cat(grads)
+
+        # Average the gradients across all GPUs
+        torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM)
+        all_grads /= self.gpu_world_size
+
+        # Get all parameters
+        all_params = self.policy.parameters()
+        if self.rnd:
+            all_params = chain(all_params, self.rnd.parameters())
+
+        # Update the gradients for all parameters with the reduced gradients
+        offset = 0
+        for param in all_params:
+            if param.grad is not None:
+                numel = param.numel()
+                # Copy data back from shared buffer
+                param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data))
+                # Update the offset for the next parameter
+                offset += numel

source/rsl_rl/rsl_rl/env/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Submodule defining the environment definitions."""
+
+from .vec_env import VecEnv
+
+__all__ = ["VecEnv"]

source/rsl_rl/rsl_rl/env/vec_env.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+from abc import ABC, abstractmethod
+from tensordict import TensorDict
+
+
+class VecEnv(ABC):
+    """Abstract class for a vectorized environment.
+
+    The vectorized environment is a collection of environments that are synchronized. This means that the same type of
+    action is applied to all environments and the same type of observation is returned from all environments.
+    """
+
+    num_envs: int
+    """Number of environments."""
+
+    num_actions: int
+    """Number of actions."""
+
+    max_episode_length: int | torch.Tensor
+    """Maximum episode length.
+
+    The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments.
+    If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode
+    lengths.
+    """
+
+    episode_length_buf: torch.Tensor
+    """Buffer for current episode lengths."""
+
+    device: torch.device | str
+    """Device to use."""
+
+    cfg: dict | object
+    """Configuration object."""
+
+    @abstractmethod
+    def get_observations(self) -> TensorDict:
+        """Return the current observations.
+
+        Returns:
+            The observations from the environment.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
+        """Apply input action to the environment.
+
+        Args:
+            actions: Input actions to apply. Shape: (num_envs, num_actions)
+
+        Returns:
+            observations: Observations from the environment.
+            rewards: Rewards from the environment. Shape: (num_envs,)
+            dones: Done flags from the environment. Shape: (num_envs,)
+            extras: Extra information from the environment.
+
+        Observations:
+            The observations TensorDict usually contains multiple observation groups. The `obs_groups`
+            dictionary of the runner configuration specifies which observation groups are used for which
+            purpose, i.e., it maps the available observation groups to observation sets. The observation sets
+            (keys of the `obs_groups` dictionary) currently used by rsl_rl are:
+
+            - "policy": Specified observation groups are used as input to the actor/student network.
+            - "critic": Specified observation groups are used as input to the critic network.
+            - "teacher": Specified observation groups are used as input to the teacher network.
+            - "rnd_state": Specified observation groups are used as input to the RND network.
+
+            Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in
+            `rsl_rl/utils/utils.py`.
+
+        Extras:
+            The extras dictionary includes metrics such as the episode reward, episode length, etc. The following
+            dictionary keys are used by rsl_rl:
+
+            - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that
+               happen due to time limits and not due to the environment reaching a terminal state. This is useful
+               for environments that have a fixed episode length.
+
+            - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes.
+               The key should be a string and start with "/" for namespacing. The value can be a scalar or a
+               tensor. If it is a tensor, the mean of the tensor is used for logging.
+        """
+        raise NotImplementedError

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

source/rsl_rl/rsl_rl/networks/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Definitions for components of modules."""
+
+from .cnn import CNN
+from .memory import HiddenState, Memory
+from .mlp import MLP
+from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
+from .moe import L2Norm, SimNorm, MoE
+__all__ = [
+    "CNN",
+    "MLP",
+    "EmpiricalDiscountedVariationNormalization",
+    "EmpiricalNormalization",
+    "HiddenState",
+    "Memory",
+    "L2Norm",
+    "SimNorm",
+    "MoE",
+]

source/rsl_rl/rsl_rl/networks/cnn.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import math
+import torch
+from torch import nn as nn
+
+from rsl_rl.utils import get_param, resolve_nn_activation
+
+
+class CNN(nn.Sequential):
+    """Convolutional Neural Network (CNN).
+
+    The CNN network is a sequence of convolutional layers, optional normalization layers, optional activation functions,
+    and optional pooling. The final output can be flattened.
+    """
+
+    def __init__(
+        self,
+        input_dim: tuple[int, int],
+        input_channels: int,
+        output_channels: tuple[int] | list[int],
+        kernel_size: int | tuple[int] | list[int],
+        stride: int | tuple[int] | list[int] = 1,
+        dilation: int | tuple[int] | list[int] = 1,
+        padding: str = "none",
+        norm: str | tuple[str] | list[str] = "none",
+        activation: str = "elu",
+        max_pool: bool | tuple[bool] | list[bool] = False,
+        global_pool: str = "none",
+        flatten: bool = True,
+    ) -> None:
+        """Initialize the CNN.
+
+        Args:
+            input_dim: Height and width of the input.
+            input_channels: Number of input channels.
+            output_channels: List of output channels for each convolutional layer.
+            kernel_size: List of kernel sizes for each convolutional layer or a single kernel size for all layers.
+            stride: List of strides for each convolutional layer or a single stride for all layers.
+            dilation: List of dilations for each convolutional layer or a single dilation for all layers.
+            padding: Padding type to use. Either 'none', 'zeros', 'reflect', 'replicate', or 'circular'.
+            norm: List of normalization types for each convolutional layer or a single type for all layers. Either
+                'none', 'batch', or 'layer'.
+            activation: Activation function to use.
+            max_pool: List of booleans indicating whether to apply max pooling after each convolutional layer or a
+                single boolean for all layers.
+            global_pool: Global pooling type to apply at the end. Either 'none', 'max', or 'avg'.
+            flatten: Whether to flatten the output tensor.
+        """
+        super().__init__()
+
+        # Resolve activation function
+        activation_function = resolve_nn_activation(activation)
+
+        # Create layers sequentially
+        layers = []
+        last_channels = input_channels
+        last_dim = input_dim
+        for idx in range(len(output_channels)):
+            # Get parameters for the current layer
+            k = get_param(kernel_size, idx)
+            s = get_param(stride, idx)
+            d = get_param(dilation, idx)
+            p = (
+                _compute_padding(last_dim, k, s, d)
+                if padding in ["zeros", "reflect", "replicate", "circular"]
+                else (0, 0)
+            )
+
+            # Append convolutional layer
+            layers.append(
+                nn.Conv2d(
+                    in_channels=last_channels,
+                    out_channels=output_channels[idx],
+                    kernel_size=k,
+                    stride=s,
+                    padding=p,
+                    dilation=d,
+                    padding_mode=padding if padding in ["zeros", "reflect", "replicate", "circular"] else "zeros",
+                )
+            )
+
+            # Append normalization layer if specified
+            n = get_param(norm, idx)
+            if n == "none":
+                pass
+            elif n == "batch":
+                layers.append(nn.BatchNorm2d(output_channels[idx]))
+            elif n == "layer":
+                norm_input_dim = _compute_output_dim(last_dim, k, s, d, p)
+                layers.append(nn.LayerNorm([output_channels[idx], norm_input_dim[0], norm_input_dim[1]]))
+            else:
+                raise ValueError(
+                    f"Unsupported normalization type: {n}. Supported types are 'none', 'batch', and 'layer'."
+                )
+
+            # Append activation function
+            layers.append(activation_function)
+
+            # Apply max pooling if specified
+            if get_param(max_pool, idx):
+                layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+
+            # Update last channels and dimensions
+            last_channels = output_channels[idx]
+            last_dim = _compute_output_dim(last_dim, k, s, d, p, is_max_pool=get_param(max_pool, idx))
+
+        # Apply global pooling if specified
+        if global_pool == "none":
+            pass
+        elif global_pool == "max":
+            layers.append(nn.AdaptiveMaxPool2d((1, 1)))
+            last_dim = (1, 1)
+        elif global_pool == "avg":
+            layers.append(nn.AdaptiveAvgPool2d((1, 1)))
+            last_dim = (1, 1)
+        else:
+            raise ValueError(
+                f"Unsupported global pooling type: {global_pool}. Supported types are 'none', 'max', and 'avg'."
+            )
+
+        # Apply flattening if specified
+        if flatten:
+            layers.append(nn.Flatten(start_dim=1))
+
+        # Store final output dimension
+        self._output_channels = last_channels if not flatten else None
+        self._output_dim = last_dim if not flatten else last_channels * last_dim[0] * last_dim[1]
+
+        # Register the layers
+        for idx, layer in enumerate(layers):
+            self.add_module(f"{idx}", layer)
+
+    @property
+    def output_channels(self) -> int | None:
+        """Get the number of output channels or None if output is flattened."""
+        return self._output_channels
+
+    @property
+    def output_dim(self) -> tuple[int, int] | int:
+        """Get the output height and width or total output dimension if output is flattened."""
+        return self._output_dim
+
+    def init_weights(self) -> None:
+        """Initialize the weights of the CNN with Kaiming initialization."""
+        for idx, module in enumerate(self):
+            if isinstance(module, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(module.weight)
+                torch.nn.init.zeros_(module.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of the CNN."""
+        for layer in self:
+            x = layer(x)
+        return x
+
+
+def _compute_padding(input_hw: tuple[int, int], kernel: int, stride: int, dilation: int) -> tuple[int, int]:
+    """Compute the optimal padding for the current layer.
+
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
+    """
+    h = math.ceil((stride * math.floor(input_hw[0] / stride) - input_hw[0] - stride + dilation * (kernel - 1) + 1) / 2)
+    w = math.ceil((stride * math.floor(input_hw[1] / stride) - input_hw[1] - stride + dilation * (kernel - 1) + 1) / 2)
+    return (h, w)
+
+
+def _compute_output_dim(
+    input_hw: tuple[int, int],
+    kernel: int,
+    stride: int,
+    dilation: int,
+    padding: tuple[int, int],
+    is_max_pool: bool = False,
+) -> tuple[int, int]:
+    """Compute the output height and width of the current layer.
+
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
+    """
+    h = math.floor((input_hw[0] + 2 * padding[0] - dilation * (kernel - 1) - 1) / stride + 1)
+    w = math.floor((input_hw[1] + 2 * padding[1] - dilation * (kernel - 1) - 1) / stride + 1)
+
+    if is_max_pool:
+        h = math.ceil(h / 2)
+        w = math.ceil(w / 2)
+
+    return (h, w)

source/rsl_rl/rsl_rl/networks/memory.py

@@ -0,0 +1,81 @@
+# 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 typing import Union
+
+from rsl_rl.utils import unpad_trajectories
+
+HiddenState = Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor], None]  # Using Union due to Python <3.10
+"""Type alias for the hidden state of RNNs (GRU/LSTM).
+
+For GRUs, this is a single tensor while for LSTMs, this is a tuple of two tensors (hidden state and cell state).
+"""
+
+
+class Memory(nn.Module):
+    """Memory network for recurrent architectures.
+
+    This network is used to store the hidden state of the policy. It currently supports GRU and LSTM.
+    """
+
+    def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None:
+        super().__init__()
+        rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM
+        self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_dim, num_layers=num_layers)
+        self.hidden_state = None
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        masks: torch.Tensor | None = None,
+        hidden_state: HiddenState = None,
+    ) -> torch.Tensor:
+        batch_mode = masks is not None
+        if batch_mode:
+            # Batch mode needs saved hidden states
+            if hidden_state is None:
+                raise ValueError("Hidden states not passed to memory module during policy update")
+            out, _ = self.rnn(input, hidden_state)
+            out = unpad_trajectories(out, masks)
+        else:
+            # Inference/distillation mode uses hidden state of last step
+            out, self.hidden_state = self.rnn(input.unsqueeze(0), self.hidden_state)
+        return out
+
+    def reset(self, dones: torch.Tensor | None = None, hidden_state: HiddenState = None) -> None:
+        if dones is None:  # Reset hidden state
+            if hidden_state is None:
+                self.hidden_state = None
+            else:
+                self.hidden_state = hidden_state
+        elif self.hidden_state is not None:  # Reset hidden state of done environments
+            if hidden_state is None:
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    for hidden_state in self.hidden_state:
+                        hidden_state[..., dones == 1, :] = 0.0
+                else:
+                    self.hidden_state[..., dones == 1, :] = 0.0
+            else:
+                NotImplementedError(
+                    "Resetting the hidden state of done environments with a custom hidden state is not implemented"
+                )
+
+    def detach_hidden_state(self, dones: torch.Tensor | None = None) -> None:
+        if self.hidden_state is not None:
+            if dones is None:  # Detach hidden state
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    self.hidden_state = tuple(hidden_state.detach() for hidden_state in self.hidden_state)
+                else:
+                    self.hidden_state = self.hidden_state.detach()
+            else:  # Detach hidden state of done environments
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    for hidden_state in self.hidden_state:
+                        hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach()
+                else:
+                    self.hidden_state[..., dones == 1, :] = self.hidden_state[..., dones == 1, :].detach()

source/rsl_rl/rsl_rl/networks/mlp.py

@@ -0,0 +1,94 @@
+# 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 functools import reduce
+
+from rsl_rl.utils import get_param, resolve_nn_activation
+
+
+class MLP(nn.Sequential):
+    """Multi-layer perceptron.
+
+    The MLP network is a sequence of linear layers and activation functions. The last layer is a linear layer that
+    outputs the desired dimension unless the last activation function is specified.
+
+    It provides additional conveniences:
+    - If the hidden dimensions have a value of ``-1``, the dimension is inferred from the input dimension.
+    - If the output dimension is a tuple, the output is reshaped to the desired shape.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int | tuple[int] | list[int],
+        hidden_dims: tuple[int] | list[int],
+        activation: str = "elu",
+        last_activation: str | None = None,
+    ) -> None:
+        """Initialize the MLP.
+
+        Args:
+            input_dim: Dimension of the input.
+            output_dim: Dimension of the output.
+            hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates that the dimension should be
+                inferred from the input dimension.
+            activation: Activation function.
+            last_activation: Activation function of the last layer. None results in a linear last layer.
+        """
+        super().__init__()
+
+        # Resolve activation functions
+        activation_mod = resolve_nn_activation(activation)
+        last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None
+        # Resolve number of hidden dims if they are -1
+        hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims]
+
+        # Create layers sequentially
+        layers = []
+        layers.append(nn.Linear(input_dim, hidden_dims_processed[0]))
+        layers.append(activation_mod)
+
+        for layer_index in range(len(hidden_dims_processed) - 1):
+            layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1]))
+            layers.append(activation_mod)
+
+        # Add last layer
+        if isinstance(output_dim, int):
+            layers.append(nn.Linear(hidden_dims_processed[-1], output_dim))
+        else:
+            # Compute the total output dimension
+            total_out_dim = reduce(lambda x, y: x * y, output_dim)
+            # Add a layer to reshape the output to the desired shape
+            layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim))
+            layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim))
+
+        # Add last activation function if specified
+        if last_activation_mod is not None:
+            layers.append(last_activation_mod)
+
+        # Register the layers
+        for idx, layer in enumerate(layers):
+            self.add_module(f"{idx}", layer)
+
+    def init_weights(self, scales: float | tuple[float]) -> None:
+        """Initialize the weights of the MLP.
+
+        Args:
+            scales: Scale factor for the weights.
+        """
+        for idx, module in enumerate(self):
+            if isinstance(module, nn.Linear):
+                nn.init.orthogonal_(module.weight, gain=get_param(scales, idx))
+                nn.init.zeros_(module.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of the MLP."""
+        for layer in self:
+            x = layer(x)
+        return x

source/rsl_rl/rsl_rl/networks/moe.py

@@ -0,0 +1,170 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class CatELU(nn.Module):
+    """
+    CatELU activation (feature-doubling version).
+
+    Applies ELU to the input and its negation,
+    doubling the feature dimension.
+
+    Output: [..., 2 * D] given input [..., D]
+
+    NOTE:
+        This is a structural activation and NOT element-wise.
+        Assumes the last dimension is the feature dimension.
+    """
+    def __init__(self, inplace: bool = False):
+        super().__init__()
+        self.elu = nn.ELU(inplace=inplace)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.dim() >= 2, \
+            f"CatELU expects feature dimension in the last axis, got shape {x.shape}"
+
+        y1 = self.elu(x)
+        y2 = self.elu(-x)
+        return torch.cat((y1, y2), dim=-1)
+
+def get_activation(act_name):
+    if act_name == "elu":
+        return nn.ELU()
+    elif act_name == "selu":
+        return nn.SELU()
+    elif act_name == "relu":
+        return nn.ReLU()
+    elif act_name == "crelu":
+        return nn.ReLU()
+    elif act_name == "lrelu":
+        return nn.LeakyReLU()
+    elif act_name == "tanh":
+        return nn.Tanh()
+    elif act_name == "sigmoid":
+        return nn.Sigmoid()
+    elif act_name == "cat_elu":
+        return CatELU()
+    else:
+        print("invalid activation function!")
+        return None
+
+class L2Norm(nn.Module):
+    
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return F.normalize(x, p=2.0, dim=-1)
+
+class SimNorm(nn.Module):
+    """
+    Simplicial normalization.
+    Adapted from https://arxiv.org/abs/2204.00616.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.dim = 8  # for latent dim 512
+
+    def forward(self, x):
+        shp = x.shape
+        x = x.view(*shp[:-1], -1, self.dim)
+        x = F.softmax(x, dim=-1)
+        return x.view(*shp)
+
+    def __repr__(self):
+        return f"SimNorm(dim={self.dim})"
+
+# MLP implementation for MoE
+class MLP(nn.Module):
+    def __init__(self, input_dim, output_dim, hidden_dims, activation='elu', last_activation: str | None = None):
+        super().__init__()
+
+        dims = [input_dim] + hidden_dims
+        act_func = get_activation(activation)
+        layers = []
+        last_dim = dims[0]
+        for h_dim in dims[1:]:
+            layers.append(nn.Linear(last_dim, h_dim))
+            layers.append(act_func)
+            if activation == 'cat_elu':
+                last_dim = h_dim * 2
+            else:
+                last_dim = h_dim
+            
+        if isinstance(output_dim, int):
+            layers.append(nn.Linear(last_dim, output_dim))
+        elif isinstance(output_dim, tuple) or isinstance(output_dim, list):
+            layers.append(nn.Linear(last_dim, np.prod(output_dim)))
+            layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim))
+        else:
+            raise ValueError("output_dim must be int, tuple or list")        
+        
+        if last_activation is not None:
+            last_act_func = get_activation(last_activation)
+            layers.append(last_act_func)
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.network(x)
+
+class Experts(nn.Module):
+    def __init__(self,
+                 expert_num,
+                 input_dim,
+                 backbone_hidden_dims,
+                 expert_hidden_dim,
+                 output_dim,
+                 activation='elu',
+    ):
+        super().__init__()
+        self.expert_num = expert_num
+        self.output_dim = output_dim
+
+        self.backbone = MLP(input_dim, expert_num * expert_hidden_dim, backbone_hidden_dims, activation, last_activation=activation)
+        self.experts = nn.Conv1d(
+            in_channels=expert_num*expert_hidden_dim if activation != 'cat_elu' else expert_num*expert_hidden_dim*2,
+            out_channels=expert_num*output_dim,
+            kernel_size=1,
+            groups=expert_num,
+        )
+    
+    def forward(self, x):
+        shared_features = self.backbone(x).unsqueeze(-1)  # (B, expert_num * expert_hidden_dim, 1)
+        expert_outs = self.experts(shared_features).squeeze(-1)  # (B, expert_num * output_dim)
+        expert_outs = expert_outs.reshape(-1, self.expert_num, self.output_dim)
+        return expert_outs
+
+class MoE(nn.Module):
+    def __init__(self,
+                 expert_num,
+                 input_dim,
+                 hidden_dims,
+                 output_dim,
+                 activation='elu',
+    ):
+        super().__init__()
+
+        # Expert networks
+        self.experts = Experts(
+            expert_num=expert_num,
+            input_dim=input_dim,
+            backbone_hidden_dims=hidden_dims[:-1],
+            expert_hidden_dim=hidden_dims[-1],
+            output_dim=output_dim,
+            activation=activation,
+        )
+        
+        # Gating network
+        self.gating_network = nn.Sequential(
+            MLP(input_dim, expert_num, hidden_dims, activation),
+            nn.Softmax(dim=-1)
+        )
+
+    def forward(self, x):
+        weights = self.gating_network(x)  # (B, expert_num)
+        expert_outs = self.experts(x)  # (B, expert_num, output_dim)
+        output = torch.sum(weights.unsqueeze(-1) * expert_outs, dim=1)  # (B, output_dim)
+        return output, weights
+

source/rsl_rl/rsl_rl/networks/normalization.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Copyright (c) 2020 Preferred Networks, Inc.
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class EmpiricalNormalization(nn.Module):
+    """Normalize mean and variance of values based on empirical values."""
+
+    def __init__(self, shape: int | tuple[int] | list[int], eps: float = 1e-2, until: int | None = None) -> None:
+        """Initialize EmpiricalNormalization module.
+
+        .. note:: The normalization parameters are computed over the whole batch, not for each environment separately.
+
+        Args:
+            shape: Shape of input values except batch axis.
+            eps: Small value for stability.
+            until: If this arg is specified, the module learns input values until the sum of batch sizes exceeds it.
+        """
+        super().__init__()
+        self.eps = eps
+        self.until = until
+        self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0))
+        self.register_buffer("_var", torch.ones(shape).unsqueeze(0))
+        self.register_buffer("_std", torch.ones(shape).unsqueeze(0))
+        self.register_buffer("count", torch.tensor(0, dtype=torch.long))
+
+    @property
+    def mean(self) -> torch.Tensor:
+        return self._mean.squeeze(0).clone()
+
+    @property
+    def std(self) -> torch.Tensor:
+        return self._std.squeeze(0).clone()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Normalize mean and variance of values based on empirical values."""
+        return (x - self._mean) / (self._std + self.eps)
+
+    @torch.jit.unused
+    def update(self, x: torch.Tensor) -> None:
+        """Learn input values without computing the output values of them."""
+        if not self.training:
+            return
+        if self.until is not None and self.count >= self.until:
+            return
+
+        count_x = x.shape[0]
+        self.count += count_x
+        rate = count_x / self.count
+        var_x = torch.var(x, dim=0, unbiased=False, keepdim=True)
+        mean_x = torch.mean(x, dim=0, keepdim=True)
+        delta_mean = mean_x - self._mean
+        self._mean += rate * delta_mean
+        self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean))
+        self._std = torch.sqrt(self._var)
+
+    @torch.jit.unused
+    def inverse(self, y: torch.Tensor) -> torch.Tensor:
+        """De-normalize values based on empirical values."""
+        return y * (self._std + self.eps) + self._mean
+
+
+class EmpiricalDiscountedVariationNormalization(nn.Module):
+    """Reward normalization from Pathak's large scale study on PPO.
+
+    Reward normalization. Since the reward function is non-stationary, it is useful to normalize the scale of the
+    rewards so that the value function can learn quickly. We did this by dividing the rewards by a running estimate of
+    the standard deviation of the sum of discounted rewards.
+    """
+
+    def __init__(
+        self, shape: int | tuple[int] | list[int], eps: float = 1e-2, gamma: float = 0.99, until: int | None = None
+    ) -> None:
+        super().__init__()
+
+        self.emp_norm = EmpiricalNormalization(shape, eps, until)
+        self.disc_avg = _DiscountedAverage(gamma)
+
+    def forward(self, rew: torch.Tensor) -> torch.Tensor:
+        if self.training:
+            # Update discounted rewards
+            avg = self.disc_avg.update(rew)
+            # Update moments from discounted rewards
+            self.emp_norm.update(avg)
+
+        # Normalize rewards with the empirical std
+        if self.emp_norm._std > 0:
+            return rew / self.emp_norm._std
+        else:
+            return rew
+
+
+class _DiscountedAverage:
+    r"""Discounted average of rewards.
+
+    The discounted average is defined as:
+
+    .. math::
+
+        \bar{R}_t = \gamma \bar{R}_{t-1} + r_t
+    """
+
+    def __init__(self, gamma: float) -> None:
+        self.avg = None
+        self.gamma = gamma
+
+    def update(self, rew: torch.Tensor) -> torch.Tensor:
+        if self.avg is None:
+            self.avg = rew
+        else:
+            self.avg = self.avg * self.gamma + rew
+        return self.avg

source/rsl_rl/rsl_rl/runners/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Implementation of runners for environment-agent interaction."""
+
+from .on_policy_runner import OnPolicyRunner  # noqa: I001
+from .distillation_runner import DistillationRunner
+from .on_policy_runner_cts import OnPolicyRunnerCTS
+
+__all__ = ["DistillationRunner", "OnPolicyRunner", "OnPolicyRunnerCTS"]

source/rsl_rl/rsl_rl/runners/distillation_runner.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from tensordict import TensorDict
+
+from rsl_rl.algorithms import Distillation
+from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent
+from rsl_rl.runners import OnPolicyRunner
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_callable
+
+
+class DistillationRunner(OnPolicyRunner):
+    """Distillation runner for training and evaluation of teacher-student methods."""
+
+    def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
+        # Check if teacher is loaded
+        if not self.alg.policy.loaded_teacher:
+            raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.")
+
+        super().learn(num_learning_iterations, init_at_random_ep_len)
+
+    def _get_default_obs_sets(self) -> list[str]:
+        """Get the the default observation sets required for the algorithm.
+
+        .. note::
+            See :func:`resolve_obs_groups` for more details on the handling of observation sets.
+        """
+        return ["teacher"]
+
+    def _construct_algorithm(self, obs: TensorDict) -> Distillation:
+        """Construct the distillation algorithm."""
+        # Initialize the policy
+        student_teacher_class = resolve_callable(self.policy_cfg.pop("class_name"))
+        student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class(
+            obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
+        ).to(self.device)
+
+        # Initialize the storage
+        storage = RolloutStorage(
+            "distillation", self.env.num_envs, self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
+        )
+
+        # Initialize the algorithm
+        alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
+        alg: Distillation = alg_class(
+            student_teacher, storage, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
+        )
+
+        # Set RND configuration to None as it does not apply to distillation
+        self.cfg["algorithm"]["rnd_cfg"] = None
+
+        return alg

source/rsl_rl/rsl_rl/runners/on_policy_runner.py

@@ -0,0 +1,286 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import os
+import time
+import torch
+import warnings
+from tensordict import TensorDict
+
+from rsl_rl.algorithms import PPO
+from rsl_rl.env import VecEnv
+from rsl_rl.modules import (
+    ActorCritic,
+    ActorCriticCNN,
+    ActorCriticRecurrent,
+    resolve_rnd_config,
+    resolve_symmetry_config,
+)
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_callable, resolve_obs_groups
+from rsl_rl.utils.logger import Logger
+
+
+class OnPolicyRunner:
+    """On-policy runner for training and evaluation of actor-critic methods."""
+
+    def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device: str = "cpu") -> None:
+        self.cfg = train_cfg
+        self.policy_cfg = train_cfg["policy"]
+        self.alg_cfg = train_cfg["algorithm"]
+        self.device = device
+        self.env = env
+
+        # Setup multi-GPU training if enabled
+        self._configure_multi_gpu()
+
+        # Query observations from environment for algorithm construction
+        obs = self.env.get_observations()
+        self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], self._get_default_obs_sets())
+
+        # Create the algorithm
+        self.alg = self._construct_algorithm(obs)
+
+        # Create the logger
+        self.logger = Logger(
+            log_dir=log_dir,
+            cfg=self.cfg,
+            env_cfg=self.env.cfg,
+            num_envs=self.env.num_envs,
+            is_distributed=self.is_distributed,
+            gpu_world_size=self.gpu_world_size,
+            gpu_global_rank=self.gpu_global_rank,
+            device=self.device,
+        )
+
+        self.current_learning_iteration = 0
+
+    def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
+        # Randomize initial episode lengths (for exploration)
+        if init_at_random_ep_len:
+            self.env.episode_length_buf = torch.randint_like(
+                self.env.episode_length_buf, high=int(self.env.max_episode_length)
+            )
+
+        # Start learning
+        obs = self.env.get_observations().to(self.device)
+        self.train_mode()  # switch to train mode (for dropout for example)
+
+        # Ensure all parameters are in-synced
+        if self.is_distributed:
+            print(f"Synchronizing parameters for rank {self.gpu_global_rank}...")
+            self.alg.broadcast_parameters()
+
+        # Start training
+        start_it = self.current_learning_iteration
+        total_it = start_it + num_learning_iterations
+        for it in range(start_it, total_it):
+            start = time.time()
+            # Rollout
+            with torch.inference_mode():
+                for _ in range(self.cfg["num_steps_per_env"]):
+                    # Sample actions
+                    actions = self.alg.act(obs)
+                    # Step the environment
+                    obs, rewards, dones, extras = self.env.step(actions.to(self.env.device))
+                    # Move to device
+                    obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device))
+                    # Process the step
+                    self.alg.process_env_step(obs, rewards, dones, extras)
+                    # Extract intrinsic rewards (only for logging)
+                    intrinsic_rewards = self.alg.intrinsic_rewards if self.alg_cfg["rnd_cfg"] else None
+                    # Book keeping
+                    self.logger.process_env_step(rewards, dones, extras, intrinsic_rewards)
+
+                stop = time.time()
+                collect_time = stop - start
+                start = stop
+
+                # Compute returns
+                self.alg.compute_returns(obs)
+
+            # Update policy
+            loss_dict = self.alg.update()
+
+            stop = time.time()
+            learn_time = stop - start
+            self.current_learning_iteration = it
+
+            # Log information
+            self.logger.log(
+                it=it,
+                start_it=start_it,
+                total_it=total_it,
+                collect_time=collect_time,
+                learn_time=learn_time,
+                loss_dict=loss_dict,
+                learning_rate=self.alg.learning_rate,
+                action_std=self.alg.policy.action_std,
+                rnd_weight=self.alg.rnd.weight if self.alg_cfg["rnd_cfg"] else None,
+            )
+
+            # Save model
+            if it % self.cfg["save_interval"] == 0:
+                self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt"))  # type: ignore
+
+        # Save the final model after training
+        if self.logger.log_dir is not None and not self.logger.disable_logs:
+            self.save(os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"))
+
+    def save(self, path: str, infos: dict | None = None) -> None:
+        # Save model
+        saved_dict = {
+            "model_state_dict": self.alg.policy.state_dict(),
+            "optimizer_state_dict": self.alg.optimizer.state_dict(),
+            "iter": self.current_learning_iteration,
+            "infos": infos,
+        }
+        # Save RND model if used
+        if self.alg_cfg["rnd_cfg"]:
+            saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict()
+            if self.alg.rnd_optimizer:
+                saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict()
+        torch.save(saved_dict, path)
+
+        # Upload model to external logging services
+        self.logger.save_model(path, self.current_learning_iteration)
+
+    def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict:
+        loaded_dict = torch.load(path, weights_only=False, map_location=map_location)
+        # Load model
+        resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"])
+        # Load RND model if used
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"])
+        # Load optimizer if used
+        if load_optimizer and resumed_training:
+            # Algorithm optimizer
+            self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
+            # RND optimizer if used
+            if self.alg_cfg["rnd_cfg"]:
+                self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"])
+        # Load current learning iteration
+        if resumed_training:
+            self.current_learning_iteration = loaded_dict["iter"]
+        return loaded_dict["infos"]
+
+    def get_inference_policy(self, device: str | None = None) -> callable:
+        self.eval_mode()  # Switch to evaluation mode (e.g. for dropout)
+        if device is not None:
+            self.alg.policy.to(device)
+        return self.alg.policy.act_inference
+
+    def train_mode(self) -> None:
+        # PPO
+        self.alg.policy.train()
+        # RND
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.train()
+
+    def eval_mode(self) -> None:
+        # PPO
+        self.alg.policy.eval()
+        # RND
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.eval()
+
+    def add_git_repo_to_log(self, repo_file_path: str) -> None:
+        self.logger.git_status_repos.append(repo_file_path)
+
+    def _get_default_obs_sets(self) -> list[str]:
+        """Get the the default observation sets required for the algorithm.
+
+        .. note::
+            See :func:`resolve_obs_groups` for more details on the handling of observation sets.
+        """
+        default_sets = ["critic"]
+        if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None:
+            default_sets.append("rnd_state")
+        return default_sets
+
+    def _configure_multi_gpu(self) -> None:
+        """Configure multi-gpu training."""
+        # Check if distributed training is enabled
+        self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1"))
+        self.is_distributed = self.gpu_world_size > 1
+
+        # If not distributed training, set local and global rank to 0 and return
+        if not self.is_distributed:
+            self.gpu_local_rank = 0
+            self.gpu_global_rank = 0
+            self.multi_gpu_cfg = None
+            return
+
+        # Get rank and world size
+        self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        self.gpu_global_rank = int(os.getenv("RANK", "0"))
+
+        # Make a configuration dictionary
+        self.multi_gpu_cfg = {
+            "global_rank": self.gpu_global_rank,  # Rank of the main process
+            "local_rank": self.gpu_local_rank,  # Rank of the current process
+            "world_size": self.gpu_world_size,  # Total number of processes
+        }
+
+        # Check if user has device specified for local rank
+        if self.device != f"cuda:{self.gpu_local_rank}":
+            raise ValueError(
+                f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'."
+            )
+        # Validate multi-GPU configuration
+        if self.gpu_local_rank >= self.gpu_world_size:
+            raise ValueError(
+                f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'."
+            )
+        if self.gpu_global_rank >= self.gpu_world_size:
+            raise ValueError(
+                f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'."
+            )
+
+        # Initialize torch distributed
+        torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size)
+        # Set device to the local rank
+        torch.cuda.set_device(self.gpu_local_rank)
+
+    def _construct_algorithm(self, obs: TensorDict) -> PPO:
+        """Construct the actor-critic algorithm."""
+        # Resolve RND config if used
+        self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env)
+
+        # Resolve symmetry config if used
+        self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env)
+
+        # Resolve deprecated normalization config
+        if self.cfg.get("empirical_normalization") is not None:
+            warnings.warn(
+                "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and "
+                "`critic_obs_normalization` as part of the `policy` configuration instead.",
+                DeprecationWarning,
+            )
+            if self.policy_cfg.get("actor_obs_normalization") is None:
+                self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"]
+            if self.policy_cfg.get("critic_obs_normalization") is None:
+                self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"]
+
+        # Initialize the policy
+        actor_critic_class = resolve_callable(self.policy_cfg.pop("class_name"))
+        actor_critic: ActorCritic | ActorCriticRecurrent | ActorCriticCNN = actor_critic_class(
+            obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
+        ).to(self.device)
+
+        # Initialize the storage
+        storage = RolloutStorage(
+            "rl", self.env.num_envs, self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
+        )
+
+        # Initialize the algorithm
+        alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
+        alg: PPO = alg_class(
+            actor_critic, storage, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
+        )
+
+        return alg

source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import os
+import time
+import torch
+import warnings
+from tensordict import TensorDict
+
+from rsl_rl.algorithms import MoECTS
+from rsl_rl.env import VecEnv
+from rsl_rl.modules import (
+    ActorCriticMoECTS,
+    resolve_rnd_config,
+    resolve_symmetry_config,
+)
+from rsl_rl.storage import RolloutStorageCTS
+from rsl_rl.utils import resolve_callable, resolve_obs_groups
+from rsl_rl.utils.logger_cts import LoggerCTS
+
+
+class OnPolicyRunnerCTS:
+    """On-policy runner for training and evaluation of actor-critic methods."""
+
+    def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device: str = "cpu") -> None:
+        self.cfg = train_cfg
+        self.policy_cfg = train_cfg["policy"]
+        self.alg_cfg = train_cfg["algorithm"]
+        self.device = device
+        self.env = env
+
+        # Setup multi-GPU training if enabled
+        self._configure_multi_gpu()
+
+        # Query observations from environment for algorithm construction
+        obs = self.env.get_observations()
+        self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], self._get_default_obs_sets())
+
+        # Create the algorithm
+        self.alg = self._construct_algorithm(obs)
+        
+        # Create the logger
+        self.logger = LoggerCTS(
+            log_dir=log_dir,
+            cfg=self.cfg,
+            env_cfg=self.env.cfg,
+            num_envs=self.env.num_envs,
+            is_distributed=self.is_distributed,
+            gpu_world_size=self.gpu_world_size,
+            gpu_global_rank=self.gpu_global_rank,
+            teacher_env_idxs=self.alg.teacher_env_idxs,
+            device=self.device,
+        )
+
+        self.current_learning_iteration = 0
+
+    def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
+        # Randomize initial episode lengths (for exploration)
+        if init_at_random_ep_len:
+            self.env.episode_length_buf = torch.randint_like(
+                self.env.episode_length_buf, high=int(self.env.max_episode_length)
+            )
+
+        # Start learning
+        obs = self.env.get_observations().to(self.device)
+        self.train_mode()  # switch to train mode (for dropout for example)
+
+        # Ensure all parameters are in-synced
+        if self.is_distributed:
+            print(f"Synchronizing parameters for rank {self.gpu_global_rank}...")
+            self.alg.broadcast_parameters()
+
+        # Start training
+        start_it = self.current_learning_iteration
+        total_it = start_it + num_learning_iterations
+        for it in range(start_it, total_it):
+            start = time.time()
+            # Rollout
+            with torch.inference_mode():
+                for _ in range(self.cfg["num_steps_per_env"]):
+                    # Sample actions
+                    actions = self.alg.act(obs)
+                    # Step the environment
+                    obs, rewards, dones, extras = self.env.step(actions.to(self.env.device))
+                    # Move to device
+                    obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device))
+                    # Process the step
+                    self.alg.process_env_step(obs, rewards, dones, extras)
+                    # Extract intrinsic rewards (only for logging)
+                    intrinsic_rewards = self.alg.intrinsic_rewards if self.alg_cfg["rnd_cfg"] else None
+                    # Book keeping
+                    self.logger.process_env_step(rewards, dones, extras, intrinsic_rewards)
+
+                stop = time.time()
+                collect_time = stop - start
+                start = stop
+
+                # Compute returns
+                self.alg.compute_returns(obs)
+
+            # Update policy
+            loss_dict = self.alg.update()
+
+            stop = time.time()
+            learn_time = stop - start
+            self.current_learning_iteration = it
+
+            # Log information
+            self.logger.log(
+                it=it,
+                start_it=start_it,
+                total_it=total_it,
+                collect_time=collect_time,
+                learn_time=learn_time,
+                loss_dict=loss_dict,
+                learning_rate=self.alg.learning_rate,
+                action_std=self.alg.policy.action_std,
+                rnd_weight=self.alg.rnd.weight if self.alg_cfg["rnd_cfg"] else None,
+            )
+
+            # Save model
+            if it % self.cfg["save_interval"] == 0:
+                self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt"))  # type: ignore
+
+        # Save the final model after training
+        if self.logger.log_dir is not None and not self.logger.disable_logs:
+            self.save(os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"))
+
+    def save(self, path: str, infos: dict | None = None) -> None:
+        # Save model
+        saved_dict = {
+            "model_state_dict": self.alg.policy.state_dict(),
+            "optimizer_state_dict": self.alg.optimizer.state_dict(),
+            "optimizer_stu_enc_state_dict": self.alg.optimizer_stu_enc.state_dict(),
+            "iter": self.current_learning_iteration,
+            "infos": infos,
+        }
+        # Save RND model if used
+        if self.alg_cfg["rnd_cfg"]:
+            saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict()
+            if self.alg.rnd_optimizer:
+                saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict()
+        torch.save(saved_dict, path)
+
+        # Upload model to external logging services
+        self.logger.save_model(path, self.current_learning_iteration)
+
+    def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict:
+        loaded_dict = torch.load(path, weights_only=False, map_location=map_location)
+        # Load model
+        resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"])
+        # Load RND model if used
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"])
+        # Load optimizer if used
+        if load_optimizer and resumed_training:
+            # Algorithm optimizer
+            self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
+            # Student encoder optimizer
+            self.alg.optimizer_stu_enc.load_state_dict(loaded_dict["optimizer_stu_enc_state_dict"])
+            # RND optimizer if used
+            if self.alg_cfg["rnd_cfg"]:
+                self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"])
+        # Load current learning iteration
+        if resumed_training:
+            self.current_learning_iteration = loaded_dict["iter"]
+        return loaded_dict["infos"]
+
+    def get_inference_policy(self, device: str | None = None) -> callable:
+        self.eval_mode()  # Switch to evaluation mode (e.g. for dropout)
+        if device is not None:
+            self.alg.policy.to(device)
+        return self.alg.policy.act_inference
+
+    def train_mode(self) -> None:
+        # PPO
+        self.alg.policy.train()
+        # RND
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.train()
+
+    def eval_mode(self) -> None:
+        # PPO
+        self.alg.policy.eval()
+        # RND
+        if self.alg_cfg["rnd_cfg"]:
+            self.alg.rnd.eval()
+
+    def add_git_repo_to_log(self, repo_file_path: str) -> None:
+        self.logger.git_status_repos.append(repo_file_path)
+
+    def _get_default_obs_sets(self) -> list[str]:
+        """Get the the default observation sets required for the algorithm.
+
+        .. note::
+            See :func:`resolve_obs_groups` for more details on the handling of observation sets.
+        """
+        default_sets = ["critic"]
+        if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None:
+            default_sets.append("rnd_state")
+        return default_sets
+
+    def _configure_multi_gpu(self) -> None:
+        """Configure multi-gpu training."""
+        # Check if distributed training is enabled
+        self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1"))
+        self.is_distributed = self.gpu_world_size > 1
+
+        # If not distributed training, set local and global rank to 0 and return
+        if not self.is_distributed:
+            self.gpu_local_rank = 0
+            self.gpu_global_rank = 0
+            self.multi_gpu_cfg = None
+            return
+
+        # Get rank and world size
+        self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        self.gpu_global_rank = int(os.getenv("RANK", "0"))
+
+        # Make a configuration dictionary
+        self.multi_gpu_cfg = {
+            "global_rank": self.gpu_global_rank,  # Rank of the main process
+            "local_rank": self.gpu_local_rank,  # Rank of the current process
+            "world_size": self.gpu_world_size,  # Total number of processes
+        }
+
+        # Check if user has device specified for local rank
+        if self.device != f"cuda:{self.gpu_local_rank}":
+            raise ValueError(
+                f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'."
+            )
+        # Validate multi-GPU configuration
+        if self.gpu_local_rank >= self.gpu_world_size:
+            raise ValueError(
+                f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'."
+            )
+        if self.gpu_global_rank >= self.gpu_world_size:
+            raise ValueError(
+                f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'."
+            )
+
+        # Initialize torch distributed
+        torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size)
+        # Set device to the local rank
+        torch.cuda.set_device(self.gpu_local_rank)
+
+    def _construct_algorithm(self, obs: TensorDict) -> MoECTS:
+        """Construct the actor-critic algorithm."""
+        # Resolve RND config if used
+        self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env)
+
+        # Resolve symmetry config if used
+        self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env)
+
+        # Resolve deprecated normalization config
+        if self.cfg.get("empirical_normalization") is not None:
+            warnings.warn(
+                "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and "
+                "`critic_obs_normalization` as part of the `policy` configuration instead.",
+                DeprecationWarning,
+            )
+            if self.policy_cfg.get("actor_obs_normalization") is None:
+                self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"]
+            if self.policy_cfg.get("critic_obs_normalization") is None:
+                self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"]
+
+        # Initialize the policy
+        # actor_critic_class = resolve_callable(self.policy_cfg.pop("class_name"))
+        actor_critic_class = eval(self.policy_cfg.pop("class_name")) # temporally use eval to avoid import bugs
+        actor_critic: ActorCriticMoECTS = actor_critic_class(
+            obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
+        ).to(self.device)
+
+        # Initialize the storage
+        storage = RolloutStorageCTS(
+            "rl", self.env.num_envs, max(int(self.env.num_envs*self.alg_cfg["teacher_env_ratio"]), 1), self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
+        )
+
+        # Initialize the algorithm
+        # alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
+        alg_class = eval(self.alg_cfg.pop("class_name")) # temporally use eval to avoid import bugs
+        alg: MoECTS = alg_class(
+            actor_critic, storage, self.env.num_envs, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
+        )
+
+        return alg

source/rsl_rl/rsl_rl/storage/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Implementation of transitions storage for RL-agent."""
+
+from .rollout_storage import RolloutStorage
+from .rollout_storage_cts import RolloutStorageCTS
+
+__all__ = ["RolloutStorage", "RolloutStorageCTS"]

source/rsl_rl/rsl_rl/storage/rollout_storage.py

@@ -0,0 +1,276 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+from collections.abc import Generator
+from tensordict import TensorDict
+
+from rsl_rl.networks import HiddenState
+from rsl_rl.utils import split_and_pad_trajectories
+
+
+class RolloutStorage:
+    """Storage for the data collected during a rollout.
+
+    The rollout storage is populated by adding transitions during the rollout phase. It then returns a generator for
+    learning, depending on the algorithm and the policy architecture.
+    """
+
+    class Transition:
+        """Storage for a single state transition."""
+
+        def __init__(self) -> None:
+            self.observations: TensorDict | None = None
+            self.actions: torch.Tensor | None = None
+            self.privileged_actions: torch.Tensor | None = None
+            self.rewards: torch.Tensor | None = None
+            self.dones: torch.Tensor | None = None
+            self.values: torch.Tensor | None = None
+            self.actions_log_prob: torch.Tensor
+            self.action_mean: torch.Tensor | None = None
+            self.action_sigma: torch.Tensor | None = None
+            self.hidden_states: tuple[HiddenState, HiddenState] = (None, None)
+
+        def clear(self) -> None:
+            self.__init__()
+
+    def __init__(
+        self,
+        training_type: str,
+        num_envs: int,
+        num_transitions_per_env: int,
+        obs: TensorDict,
+        actions_shape: tuple[int] | list[int],
+        device: str = "cpu",
+    ) -> None:
+        self.training_type = training_type
+        self.device = device
+        self.num_transitions_per_env = num_transitions_per_env
+        self.num_envs = num_envs
+        self.actions_shape = actions_shape
+
+        # Core
+        self.observations = TensorDict(
+            {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()},
+            batch_size=[num_transitions_per_env, num_envs],
+            device=self.device,
+        )
+        self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+        self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+        self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
+
+        # For distillation
+        if training_type == "distillation":
+            self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+
+        # For reinforcement learning
+        if training_type == "rl":
+            self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+            self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+            self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+
+        # For RNN networks
+        self.saved_hidden_state_a = None
+        self.saved_hidden_state_c = None
+
+        # Counter for the number of transitions stored
+        self.step = 0
+
+    def add_transition(self, transition: Transition) -> None:
+        # Check if the transition is valid
+        if self.step >= self.num_transitions_per_env:
+            raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.")
+
+        # Core
+        self.observations[self.step].copy_(transition.observations)
+        self.actions[self.step].copy_(transition.actions)
+        self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
+        self.dones[self.step].copy_(transition.dones.view(-1, 1))
+
+        # For distillation
+        if self.training_type == "distillation":
+            self.privileged_actions[self.step].copy_(transition.privileged_actions)
+
+        # For reinforcement learning
+        if self.training_type == "rl":
+            self.values[self.step].copy_(transition.values)
+            self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
+            self.mu[self.step].copy_(transition.action_mean)
+            self.sigma[self.step].copy_(transition.action_sigma)
+
+        # For RNN networks
+        self._save_hidden_states(transition.hidden_states)
+
+        # Increment the counter
+        self.step += 1
+
+    def clear(self) -> None:
+        self.step = 0
+
+    # For distillation
+    def generator(self) -> Generator:
+        if self.training_type != "distillation":
+            raise ValueError("This function is only available for distillation training.")
+
+        for i in range(self.num_transitions_per_env):
+            yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i]
+
+    # For reinforcement learning with feedforward networks
+    def mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
+        if self.training_type != "rl":
+            raise ValueError("This function is only available for reinforcement learning training.")
+        batch_size = self.num_envs * self.num_transitions_per_env
+        mini_batch_size = batch_size // num_mini_batches
+        indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device)
+
+        # Core
+        observations = self.observations.flatten(0, 1)
+        actions = self.actions.flatten(0, 1)
+        values = self.values.flatten(0, 1)
+        returns = self.returns.flatten(0, 1)
+
+        # For PPO
+        old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
+        advantages = self.advantages.flatten(0, 1)
+        old_mu = self.mu.flatten(0, 1)
+        old_sigma = self.sigma.flatten(0, 1)
+
+        for epoch in range(num_epochs):
+            for i in range(num_mini_batches):
+                # Select the indices for the mini-batch
+                start = i * mini_batch_size
+                stop = (i + 1) * mini_batch_size
+                batch_idx = indices[start:stop]
+
+                # Create the mini-batch
+                obs_batch = observations[batch_idx]
+                actions_batch = actions[batch_idx]
+                target_values_batch = values[batch_idx]
+                returns_batch = returns[batch_idx]
+                old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
+                advantages_batch = advantages[batch_idx]
+                old_mu_batch = old_mu[batch_idx]
+                old_sigma_batch = old_sigma[batch_idx]
+
+                hidden_state_a_batch = None
+                hidden_state_c_batch = None
+                masks_batch = None
+
+                # Yield the mini-batch
+                yield (
+                    obs_batch,
+                    actions_batch,
+                    target_values_batch,
+                    advantages_batch,
+                    returns_batch,
+                    old_actions_log_prob_batch,
+                    old_mu_batch,
+                    old_sigma_batch,
+                    (
+                        hidden_state_a_batch,
+                        hidden_state_c_batch,
+                    ),
+                    masks_batch,
+                )
+
+    # For reinforcement learning with recurrent networks
+    def recurrent_mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
+        if self.training_type != "rl":
+            raise ValueError("This function is only available for reinforcement learning training.")
+        padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones)
+
+        mini_batch_size = self.num_envs // num_mini_batches
+        for ep in range(num_epochs):
+            first_traj = 0
+            for i in range(num_mini_batches):
+                start = i * mini_batch_size
+                stop = (i + 1) * mini_batch_size
+
+                dones = self.dones.squeeze(-1)
+                last_was_done = torch.zeros_like(dones, dtype=torch.bool)
+                last_was_done[1:] = dones[:-1]
+                last_was_done[0] = True
+                trajectories_batch_size = torch.sum(last_was_done[:, start:stop])
+                last_traj = first_traj + trajectories_batch_size
+
+                masks_batch = trajectory_masks[:, first_traj:last_traj]
+                obs_batch = padded_obs_trajectories[:, first_traj:last_traj]
+                actions_batch = self.actions[:, start:stop]
+                old_mu_batch = self.mu[:, start:stop]
+                old_sigma_batch = self.sigma[:, start:stop]
+                returns_batch = self.returns[:, start:stop]
+                advantages_batch = self.advantages[:, start:stop]
+                values_batch = self.values[:, start:stop]
+                old_actions_log_prob_batch = self.actions_log_prob[:, start:stop]
+
+                # Reshape to [num_envs, time, num layers, hidden dim]
+                # Original shape: [time, num_layers, num_envs, hidden_dim])
+                last_was_done = last_was_done.permute(1, 0)
+                # Take only time steps after dones (flattens num envs and time dimensions),
+                # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim]
+                hidden_state_a_batch = [
+                    saved_hidden_state.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj]
+                    .transpose(1, 0)
+                    .contiguous()
+                    for saved_hidden_state in self.saved_hidden_state_a
+                ]
+                hidden_state_c_batch = [
+                    saved_hidden_state.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj]
+                    .transpose(1, 0)
+                    .contiguous()
+                    for saved_hidden_state in self.saved_hidden_state_c
+                ]
+                # Remove the tuple for GRU
+                hidden_state_a_batch = (
+                    hidden_state_a_batch[0] if len(hidden_state_a_batch) == 1 else hidden_state_a_batch
+                )
+                hidden_state_c_batch = (
+                    hidden_state_c_batch[0] if len(hidden_state_c_batch) == 1 else hidden_state_c_batch
+                )
+
+                # Yield the mini-batch
+                yield (
+                    obs_batch,
+                    actions_batch,
+                    values_batch,
+                    advantages_batch,
+                    returns_batch,
+                    old_actions_log_prob_batch,
+                    old_mu_batch,
+                    old_sigma_batch,
+                    (
+                        hidden_state_a_batch,
+                        hidden_state_c_batch,
+                    ),
+                    masks_batch,
+                )
+
+                first_traj = last_traj
+
+    def _save_hidden_states(self, hidden_states: tuple[HiddenState, HiddenState]) -> None:
+        if hidden_states == (None, None):
+            return
+        # Make a tuple out of GRU hidden states to match the LSTM format
+        hidden_state_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
+        hidden_state_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
+        # Initialize hidden states if needed
+        if self.saved_hidden_state_a is None:
+            self.saved_hidden_state_a = [
+                torch.zeros(self.observations.shape[0], *hidden_state_a[i].shape, device=self.device)
+                for i in range(len(hidden_state_a))
+            ]
+            self.saved_hidden_state_c = [
+                torch.zeros(self.observations.shape[0], *hidden_state_c[i].shape, device=self.device)
+                for i in range(len(hidden_state_c))
+            ]
+        # Copy the states
+        for i in range(len(hidden_state_a)):
+            self.saved_hidden_state_a[i][self.step].copy_(hidden_state_a[i])
+            self.saved_hidden_state_c[i][self.step].copy_(hidden_state_c[i])

source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py

@@ -0,0 +1,207 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+from collections.abc import Generator
+from tensordict import TensorDict
+
+from rsl_rl.networks import HiddenState
+from rsl_rl.utils import split_and_pad_trajectories
+from functools import partial
+
+class RolloutStorageCTS:
+    """Storage for the data collected during a rollout.
+
+    The rollout storage is populated by adding transitions during the rollout phase. It then returns a generator for
+    learning, depending on the algorithm and the policy architecture.
+    """
+
+    class Transition:
+        """Storage for a single state transition."""
+
+        def __init__(self) -> None:
+            self.observations: TensorDict | None = None
+            self.actions: torch.Tensor | None = None
+            self.privileged_actions: torch.Tensor | None = None
+            self.rewards: torch.Tensor | None = None
+            self.dones: torch.Tensor | None = None
+            self.values: torch.Tensor | None = None
+            self.actions_log_prob: torch.Tensor
+            self.action_mean: torch.Tensor | None = None
+            self.action_sigma: torch.Tensor | None = None
+            self.hidden_states: tuple[HiddenState, HiddenState] = (None, None)
+
+        def clear(self) -> None:
+            self.__init__()
+
+    def __init__(
+        self,
+        training_type: str,
+        num_envs: int,
+        teacher_num_envs: int, 
+        num_transitions_per_env: int,
+        obs: TensorDict,
+        actions_shape: tuple[int] | list[int],
+        device: str = "cpu",
+    ) -> None:
+        self.training_type = training_type
+        self.device = device
+        self.num_transitions_per_env = num_transitions_per_env
+        self.num_envs = num_envs
+        self.actions_shape = actions_shape
+        self.teacher_num_envs = teacher_num_envs
+        self.student_num_envs = num_envs - teacher_num_envs
+
+        # Core
+        self.observations = TensorDict(
+            {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()},
+            batch_size=[num_transitions_per_env, num_envs],
+            device=self.device,
+        )
+
+        self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+        self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+        self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
+        
+        # For distillation
+        if training_type == "distillation":
+            self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+
+        # For reinforcement learning
+        if training_type == "rl":
+            self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+            self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
+            self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+            self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
+
+        # For RNN networks
+        self.saved_hidden_state_a = None
+        self.saved_hidden_state_c = None
+
+        # Counter for the number of transitions stored
+        self.step = 0
+
+    def add_transition(self, transition: Transition) -> None:
+        # Check if the transition is valid
+        if self.step >= self.num_transitions_per_env:
+            raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.")
+
+        # Core
+        self.observations[self.step].copy_(transition.observations)
+        self.actions[self.step].copy_(transition.actions)
+        self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
+        self.dones[self.step].copy_(transition.dones.view(-1, 1))
+
+        # For distillation
+        if self.training_type == "distillation":
+            self.privileged_actions[self.step].copy_(transition.privileged_actions)
+
+        # For reinforcement learning
+        if self.training_type == "rl":
+            self.values[self.step].copy_(transition.values)
+            self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
+            self.mu[self.step].copy_(transition.action_mean)
+            self.sigma[self.step].copy_(transition.action_sigma)
+
+        # For RNN networks
+        self._save_hidden_states(transition.hidden_states)
+
+        # Increment the counter
+        self.step += 1
+
+    def clear(self) -> None:
+        self.step = 0
+
+    # For distillation
+    def generator(self) -> Generator:
+        if self.training_type != "distillation":
+            raise ValueError("This function is only available for distillation training.")
+
+        for i in range(self.num_transitions_per_env):
+            yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i]
+
+    # For reinforcement learning with feedforward networks
+    def mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
+        if self.training_type != "rl":
+            raise ValueError("This function is only available for reinforcement learning training.")
+        
+        # Prepare indices
+        teacher_samples_num = self.teacher_num_envs * self.num_transitions_per_env
+        student_samples_num = self.student_num_envs * self.num_transitions_per_env
+        teacher_mini_batch_size = teacher_samples_num // num_mini_batches
+        student_mini_batch_size = student_samples_num // num_mini_batches
+        teacher_indices = torch.randperm(teacher_samples_num, requires_grad=False, device=self.device)
+        student_indices = teacher_samples_num + torch.randperm(student_samples_num, requires_grad=False, device=self.device)
+        
+        # Core
+        observations = self.observations.transpose(0, 1).flatten(0, 1)
+        actions = self.actions.transpose(0, 1).flatten(0, 1)
+        values = self.values.transpose(0, 1).flatten(0, 1)
+        returns = self.returns.transpose(0, 1).flatten(0, 1)
+
+        # For PPO
+        old_actions_log_prob = self.actions_log_prob.transpose(0, 1).flatten(0, 1)
+        advantages = self.advantages.transpose(0, 1).flatten(0, 1)
+        old_mu = self.mu.transpose(0, 1).flatten(0, 1)
+        old_sigma = self.sigma.transpose(0, 1).flatten(0, 1)
+        
+        def _get_teacher_student_samples(data, slice):
+            (i1, i2), (j1, j2) = slice
+            return torch.cat([data[teacher_indices[i1:i2]], data[student_indices[j1:j2]]], 0).detach()
+
+        for epoch in range(num_epochs):
+            for i in range(num_mini_batches):
+                # Select the indices for the mini-batch
+                slice = (
+                    (i * teacher_mini_batch_size, (i+1) * teacher_mini_batch_size),
+                    (i * student_mini_batch_size, (i+1) * student_mini_batch_size),
+                )
+                
+                # Create the mini-batch
+                get_batch = partial(_get_teacher_student_samples, slice=slice)
+                obs_batch, actions_batch, target_values_batch, returns_batch, \
+                old_actions_log_prob_batch, advantages_batch, old_mu_batch, \
+                old_sigma_batch = map(get_batch, [
+                    observations,
+                    actions,
+                    values,
+                    returns,
+                    old_actions_log_prob,
+                    advantages,
+                    old_mu,
+                    old_sigma
+                ])
+
+                hidden_state_a_batch = None
+                hidden_state_c_batch = None
+                masks_batch = None
+
+                # Yield the mini-batch
+                yield (
+                    obs_batch,
+                    actions_batch,
+                    target_values_batch,
+                    advantages_batch,
+                    returns_batch,
+                    old_actions_log_prob_batch,
+                    old_mu_batch,
+                    old_sigma_batch,
+                    (
+                        hidden_state_a_batch,
+                        hidden_state_c_batch,
+                    ),
+                    masks_batch,
+                )
+
+    # For reinforcement learning with recurrent networks
+    def recurrent_mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
+        return NotImplementedError("CTS rollout storage does not support RNNs yet.")
+
+    def _save_hidden_states(self, hidden_states: tuple[HiddenState, HiddenState]) -> None:
+        return NotImplementedError("CTS rollout storage does not support RNNs yet.")

source/rsl_rl/rsl_rl/utils/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Helper functions."""
+
+from .utils import (
+    get_param,
+    resolve_callable,
+    resolve_nn_activation,
+    resolve_obs_groups,
+    resolve_optimizer,
+    split_and_pad_trajectories,
+    unpad_trajectories,
+)
+
+__all__ = [
+    "get_param",
+    "resolve_callable",
+    "resolve_nn_activation",
+    "resolve_obs_groups",
+    "resolve_optimizer",
+    "split_and_pad_trajectories",
+    "unpad_trajectories",
+]

source/rsl_rl/rsl_rl/utils/logger.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import git
+import os
+import pathlib
+import statistics
+import time
+import torch
+from collections import deque
+
+import rsl_rl
+
+
+class Logger:
+    """Logger to save the learning metrics to different logging services."""
+
+    def __init__(
+        self,
+        log_dir: str | None,
+        cfg: dict,
+        env_cfg: dict | object,
+        num_envs: int,
+        is_distributed: bool,
+        gpu_world_size: int,
+        gpu_global_rank: int,
+        device: str,
+    ) -> None:
+        self.log_dir = log_dir
+        self.cfg = cfg
+        self.num_envs = num_envs
+        self.gpu_world_size = gpu_world_size
+        self.device = device
+        self.git_status_repos = [rsl_rl.__file__]
+        self.tot_timesteps = 0
+        self.tot_time = 0
+
+        # Create buffers
+        self.ep_extras = []
+        self.rewbuffer = deque(maxlen=100)
+        self.lenbuffer = deque(maxlen=100)
+        self.cur_reward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+        self.cur_episode_length = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+
+        # Create RND buffers
+        if self.cfg["algorithm"]["rnd_cfg"]:
+            self.erewbuffer = deque(maxlen=100)
+            self.irewbuffer = deque(maxlen=100)
+            self.cur_ereward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+            self.cur_ireward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+
+        # Decide whether to disable logging
+        # Note: We only log from the process with rank 0 (main process)
+        self.disable_logs = is_distributed and gpu_global_rank != 0
+
+        # Initialize the writer
+        self._prepare_logging_writer()
+
+        # Log code state
+        self._store_code_state()
+
+        # Log configuration
+        if self.writer and not self.disable_logs and self.logger_type in ["wandb", "neptune"]:
+            self.writer.store_config(env_cfg, self.cfg)
+
+    def process_env_step(
+        self,
+        rewards: torch.Tensor,
+        dones: torch.Tensor,
+        extras: dict,
+        intrinsic_rewards: torch.Tensor | None = None,
+    ) -> None:
+        """Add metrics from the environment step to the buffers."""
+        if self.log_dir is not None:
+            if "episode" in extras:
+                self.ep_extras.append(extras["episode"])
+            elif "log" in extras:
+                self.ep_extras.append(extras["log"])
+
+            # Update rewards and episode length
+            if intrinsic_rewards is not None:
+                self.cur_ereward_sum += rewards
+                self.cur_ireward_sum += intrinsic_rewards
+                self.cur_reward_sum += rewards + intrinsic_rewards
+            else:
+                self.cur_reward_sum += rewards
+            self.cur_episode_length += 1
+
+            # Clear data for completed episodes
+            new_ids = (dones > 0).nonzero(as_tuple=False)
+            self.rewbuffer.extend(self.cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
+            self.lenbuffer.extend(self.cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
+            self.cur_reward_sum[new_ids] = 0
+            self.cur_episode_length[new_ids] = 0
+            if intrinsic_rewards is not None:
+                self.erewbuffer.extend(self.cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist())
+                self.irewbuffer.extend(self.cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist())
+                self.cur_ereward_sum[new_ids] = 0
+                self.cur_ireward_sum[new_ids] = 0
+
+    def log(
+        self,
+        it: int,
+        start_it: int,
+        total_it: int,
+        collect_time: float,
+        learn_time: float,
+        loss_dict: dict,
+        learning_rate: float,
+        action_std: torch.Tensor,
+        rnd_weight: float | None,
+        print_minimal: bool = False,
+        width: int = 80,
+        pad: int = 40,
+    ) -> None:
+        """Log the training metrics to the logging service and print them to the console."""
+        if self.log_dir is not None and not self.disable_logs:
+            collection_size = self.cfg["num_steps_per_env"] * self.num_envs * self.gpu_world_size
+            iteration_time = collect_time + learn_time
+            self.tot_timesteps += collection_size
+            self.tot_time += iteration_time
+
+            # Log episode extras
+            extras_string = ""
+            if self.ep_extras:
+                # Iterate over all keys in the episode info dictionary
+                for key in self.ep_extras[0]:
+                    infotensor = torch.tensor([], device=self.device)
+                    # Iterate over all steps
+                    for ep_info in self.ep_extras:
+                        # Handle missing, scalar, and zero dimensional tensors
+                        if key not in ep_info:
+                            continue
+                        if not isinstance(ep_info[key], torch.Tensor):
+                            ep_info[key] = torch.Tensor([ep_info[key]])
+                        if len(ep_info[key].shape) == 0:
+                            ep_info[key] = ep_info[key].unsqueeze(0)
+                        infotensor = torch.cat((infotensor, ep_info[key].to(self.device)))
+                    value = torch.mean(infotensor)
+                    if "/" in key:
+                        self.writer.add_scalar(key, value, it)
+                        extras_string += f"""{f"{key}:":>{pad}} {value:.4f}\n"""
+                    else:
+                        self.writer.add_scalar("Episode/" + key, value, it)
+                        extras_string += f"""{f"Mean episode {key}:":>{pad}} {value:.4f}\n"""
+
+            # Log losses
+            for key, value in loss_dict.items():
+                self.writer.add_scalar(f"Loss/{key}", value, it)
+            self.writer.add_scalar("Loss/learning_rate", learning_rate, it)
+
+            # Log noise std
+            self.writer.add_scalar("Policy/mean_noise_std", action_std.mean().item(), it)
+
+            # Log performance
+            fps = int(collection_size / (collect_time + learn_time))
+            self.writer.add_scalar("Perf/total_fps", fps, it)
+            self.writer.add_scalar("Perf/collection_time", collect_time, it)
+            self.writer.add_scalar("Perf/learning_time", learn_time, it)
+
+            # Log rewards and episode length
+            if len(self.rewbuffer) > 0:
+                if self.cfg["algorithm"]["rnd_cfg"]:
+                    self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(self.erewbuffer), it)
+                    self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(self.irewbuffer), it)
+                    self.writer.add_scalar("Rnd/weight", rnd_weight, it)
+                self.writer.add_scalar("Train/mean_reward", statistics.mean(self.rewbuffer), it)
+                self.writer.add_scalar("Train/mean_episode_length", statistics.mean(self.lenbuffer), it)
+                if self.logger_type != "wandb":
+                    self.writer.add_scalar(
+                        "Train/mean_reward/time", statistics.mean(self.rewbuffer), int(self.tot_time)
+                    )
+                    self.writer.add_scalar(
+                        "Train/mean_episode_length/time", statistics.mean(self.lenbuffer), int(self.tot_time)
+                    )
+
+            # Print to console
+            log_string = f"""{"#" * width}\n"""
+            log_string += f"""\033[1m{f" Learning iteration {it}/{total_it} ".center(width)}\033[0m \n\n"""
+
+            # Print run name if provided
+            run_name = self.cfg.get("run_name")
+            log_string += f"""{"Run name:":>{pad}} {run_name}\n""" if run_name else ""
+
+            # Print performance
+            log_string += (
+                f"""{"Total steps:":>{pad}} {self.tot_timesteps} \n"""
+                f"""{"Steps per second:":>{pad}} {fps:.0f} \n"""
+                f"""{"Collection time:":>{pad}} {collect_time:.3f}s \n"""
+                f"""{"Learning time:":>{pad}} {learn_time:.3f}s \n"""
+            )
+
+            # Print losses
+            for key, value in loss_dict.items():
+                log_string += f"""{f"Mean {key} loss:":>{pad}} {value:.4f}\n"""
+
+            # Print rewards and episode length
+            if len(self.rewbuffer) > 0:
+                if self.cfg["algorithm"]["rnd_cfg"]:
+                    log_string += f"""{"Mean extrinsic reward:":>{pad}} {statistics.mean(self.erewbuffer):.2f}\n"""
+                    log_string += f"""{"Mean intrinsic reward:":>{pad}} {statistics.mean(self.irewbuffer):.2f}\n"""
+                log_string += f"""{"Mean reward:":>{pad}} {statistics.mean(self.rewbuffer):.2f}\n"""
+                log_string += f"""{"Mean episode length:":>{pad}} {statistics.mean(self.lenbuffer):.2f}\n"""
+
+            # Print noise std
+            log_string += f"""{"Mean action noise std:":>{pad}} {action_std.mean().item():.2f}\n"""
+
+            # Print episode extras
+            if not print_minimal:
+                log_string += extras_string
+
+            # Print footer
+            done_it = it + 1 - start_it
+            remaining_it = total_it - start_it - done_it
+            eta = self.tot_time / done_it * remaining_it
+            log_string += (
+                f"""{"-" * width}\n"""
+                f"""{"Iteration time:":>{pad}} {iteration_time:.2f}s\n"""
+                f"""{"Time elapsed:":>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n"""
+                f"""{"ETA:":>{pad}} {time.strftime("%H:%M:%S", time.gmtime(eta))}\n"""
+            )
+            print(log_string)
+
+            # Clear extras buffer
+            self.ep_extras.clear()
+
+    def save_model(self, path: str, it: int) -> None:
+        """Save the model to external logging services if specified."""
+        if self.writer and not self.disable_logs and self.logger_type in ["neptune", "wandb"]:
+            self.writer.save_model(path, it)
+
+    def _prepare_logging_writer(self) -> None:
+        """Prepare the logging writer, which can be either Tensorboard, W&B or Neptune."""
+        if self.log_dir is not None and not self.disable_logs:
+            self.logger_type = self.cfg.get("logger", "tensorboard")
+            self.logger_type = self.logger_type.lower()
+
+            if self.logger_type == "neptune":
+                from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter
+
+                self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg)
+            elif self.logger_type == "wandb":
+                from rsl_rl.utils.wandb_utils import WandbSummaryWriter
+
+                self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg)
+            elif self.logger_type == "tensorboard":
+                from torch.utils.tensorboard import SummaryWriter
+
+                self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
+            else:
+                raise ValueError("Logger type not found. Please choose 'wandb', 'neptune', or 'tensorboard'.")
+        else:
+            self.writer = None
+
+    def _store_code_state(self) -> None:
+        """Store the current git diff of the code repositories involved in the experiment."""
+        if self.log_dir is not None and not self.disable_logs:
+            git_log_dir = os.path.join(self.log_dir, "git")
+            os.makedirs(git_log_dir, exist_ok=True)
+            file_paths = []
+            # Iterate over all repositories to log
+            for repository_file_path in self.git_status_repos:
+                try:
+                    repo = git.Repo(repository_file_path, search_parent_directories=True)
+                    t = repo.head.commit.tree
+                except Exception:
+                    print(f"Could not find git repository in {repository_file_path}. Skipping.")
+                    continue
+                # Get the name of the repository
+                repo_name = pathlib.Path(repo.working_dir).name
+                diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff")
+                # Check if the diff file already exists
+                if os.path.isfile(diff_file_name):
+                    continue
+                # Write the diff file
+                print(f"Storing git diff for '{repo_name}' in: {diff_file_name}")
+                with open(diff_file_name, "x", encoding="utf-8") as f:
+                    content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}"
+                    f.write(content)
+                # Add the file path to the list of files to be uploaded
+                file_paths.append(diff_file_name)
+
+            # Upload diff files to external logging services
+            if self.writer and self.logger_type in ["wandb", "neptune"] and file_paths:
+                for path in file_paths:
+                    self.writer.save_file(path)

source/rsl_rl/rsl_rl/utils/logger_cts.py

@@ -0,0 +1,311 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import git
+import os
+import pathlib
+import statistics
+import time
+import torch
+from collections import deque
+
+import rsl_rl
+
+
+class LoggerCTS:
+    """Logger to save the learning metrics to different logging services."""
+
+    def __init__(
+        self,
+        log_dir: str | None,
+        cfg: dict,
+        env_cfg: dict | object,
+        num_envs: int,
+        is_distributed: bool,
+        gpu_world_size: int,
+        gpu_global_rank: int,
+        teacher_env_idxs: torch.Tensor,
+        device: str,
+    ) -> None:
+        self.log_dir = log_dir
+        self.cfg = cfg
+        self.num_envs = num_envs
+        self.gpu_world_size = gpu_world_size
+        self.device = device
+        self.git_status_repos = [rsl_rl.__file__]
+        self.tot_timesteps = 0
+        self.tot_time = 0
+        self.teacher_env_idxs = teacher_env_idxs # Indices of teacher envs in the total envs
+
+        # Create buffers
+        self.ep_extras = []
+        self.teacher_rewbuffer = deque(maxlen=100)
+        self.teacher_lenbuffer = deque(maxlen=100)
+        self.student_rewbuffer = deque(maxlen=100)
+        self.student_lenbuffer = deque(maxlen=100)
+        self.cur_reward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+        self.cur_episode_length = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+
+        # Create RND buffers
+        if self.cfg["algorithm"]["rnd_cfg"]:
+            self.erewbuffer = deque(maxlen=100)
+            self.irewbuffer = deque(maxlen=100)
+            self.cur_ereward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+            self.cur_ireward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
+
+        # Decide whether to disable logging
+        # Note: We only log from the process with rank 0 (main process)
+        self.disable_logs = is_distributed and gpu_global_rank != 0
+
+        # Initialize the writer
+        self._prepare_logging_writer()
+
+        # Log code state
+        self._store_code_state()
+
+        # Log configuration
+        if self.writer and not self.disable_logs and self.logger_type in ["wandb", "neptune"]:
+            self.writer.store_config(env_cfg, self.cfg)
+
+    def process_env_step(
+        self,
+        rewards: torch.Tensor,
+        dones: torch.Tensor,
+        extras: dict,
+        intrinsic_rewards: torch.Tensor | None = None,
+    ) -> None:
+        """Add metrics from the environment step to the buffers."""
+        if self.log_dir is not None:
+            if "episode" in extras:
+                self.ep_extras.append(extras["episode"])
+            elif "log" in extras:
+                self.ep_extras.append(extras["log"])
+
+            # Update rewards and episode length
+            if intrinsic_rewards is not None:
+                self.cur_ereward_sum += rewards
+                self.cur_ireward_sum += intrinsic_rewards
+                self.cur_reward_sum += rewards + intrinsic_rewards
+            else:
+                self.cur_reward_sum += rewards
+            self.cur_episode_length += 1
+
+            # Clear data for completed episodes
+            new_ids = (dones > 0).nonzero(as_tuple=False)
+            if new_ids.shape[0]:
+                teacher_ids = new_ids[torch.isin(new_ids, self.teacher_env_idxs)]
+                student_ids = new_ids[~torch.isin(new_ids, self.teacher_env_idxs)]
+                self.teacher_rewbuffer.extend(self.cur_reward_sum[teacher_ids].cpu().numpy().tolist())
+                self.teacher_lenbuffer.extend(self.cur_episode_length[teacher_ids].cpu().numpy().tolist())
+                self.student_rewbuffer.extend(self.cur_reward_sum[student_ids].cpu().numpy().tolist())
+                self.student_lenbuffer.extend(self.cur_episode_length[student_ids].cpu().numpy().tolist())
+                self.cur_reward_sum[new_ids] = 0
+                self.cur_episode_length[new_ids] = 0
+            
+                if intrinsic_rewards is not None:
+                    self.erewbuffer.extend(self.cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist())
+                    self.irewbuffer.extend(self.cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist())
+                    self.cur_ereward_sum[new_ids] = 0
+                    self.cur_ireward_sum[new_ids] = 0
+
+    def log(
+        self,
+        it: int,
+        start_it: int,
+        total_it: int,
+        collect_time: float,
+        learn_time: float,
+        loss_dict: dict,
+        learning_rate: float,
+        action_std: torch.Tensor,
+        rnd_weight: float | None,
+        print_minimal: bool = False,
+        width: int = 80,
+        pad: int = 40,
+    ) -> None:
+        """Log the training metrics to the logging service and print them to the console."""
+        if self.log_dir is not None and not self.disable_logs:
+            collection_size = self.cfg["num_steps_per_env"] * self.num_envs * self.gpu_world_size
+            iteration_time = collect_time + learn_time
+            self.tot_timesteps += collection_size
+            self.tot_time += iteration_time
+
+            # Log episode extras
+            extras_string = ""
+            if self.ep_extras:
+                # Iterate over all keys in the episode info dictionary
+                for key in self.ep_extras[0]:
+                    infotensor = torch.tensor([], device=self.device)
+                    # Iterate over all steps
+                    for ep_info in self.ep_extras:
+                        # Handle missing, scalar, and zero dimensional tensors
+                        if key not in ep_info:
+                            continue
+                        if not isinstance(ep_info[key], torch.Tensor):
+                            ep_info[key] = torch.Tensor([ep_info[key]])
+                        if len(ep_info[key].shape) == 0:
+                            ep_info[key] = ep_info[key].unsqueeze(0)
+                        infotensor = torch.cat((infotensor, ep_info[key].to(self.device)))
+                    value = torch.mean(infotensor)
+                    if "/" in key:
+                        self.writer.add_scalar(key, value, it)
+                        extras_string += f"""{f"{key}:":>{pad}} {value:.4f}\n"""
+                    else:
+                        self.writer.add_scalar("Episode/" + key, value, it)
+                        extras_string += f"""{f"Mean episode {key}:":>{pad}} {value:.4f}\n"""
+
+            # Log losses
+            for key, value in loss_dict.items():
+                self.writer.add_scalar(f"Loss/{key}", value, it)
+            self.writer.add_scalar("Loss/learning_rate", learning_rate, it)
+
+            # Log noise std
+            self.writer.add_scalar("Policy/mean_noise_std", action_std.mean().item(), it)
+
+            # Log performance
+            fps = int(collection_size / (collect_time + learn_time))
+            self.writer.add_scalar("Perf/total_fps", fps, it)
+            self.writer.add_scalar("Perf/collection_time", collect_time, it)
+            self.writer.add_scalar("Perf/learning_time", learn_time, it)
+
+            # Log rewards and episode length
+            if self.cfg["algorithm"]["rnd_cfg"] and len(self.erewbuffer) > 0:
+                self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(self.erewbuffer), it)
+                self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(self.irewbuffer), it)
+                self.writer.add_scalar("Rnd/weight", rnd_weight, it)
+            if len(self.teacher_rewbuffer) > 0:
+                self.writer.add_scalar("Train/mean_teacher_reward", statistics.mean(self.teacher_rewbuffer), it)
+                self.writer.add_scalar("Train/mean_teacher_episode_length", statistics.mean(self.teacher_lenbuffer), it)
+                if self.logger_type != "wandb":
+                    self.writer.add_scalar(
+                        "Train/mean_teacher_reward/time", statistics.mean(self.teacher_rewbuffer), int(self.tot_time)
+                    )
+                    self.writer.add_scalar(
+                        "Train/mean_teacher_episode_length/time", statistics.mean(self.teacher_lenbuffer), int(self.tot_time)
+                    )
+            if len(self.student_rewbuffer) > 0:
+                self.writer.add_scalar("Train/mean_student_reward", statistics.mean(self.student_rewbuffer), it)
+                self.writer.add_scalar("Train/mean_student_episode_length", statistics.mean(self.student_lenbuffer), it)
+                if self.logger_type != "wandb":
+                    self.writer.add_scalar(
+                        "Train/mean_student_reward/time", statistics.mean(self.student_rewbuffer), int(self.tot_time)
+                    )
+                    self.writer.add_scalar(
+                        "Train/mean_student_episode_length/time", statistics.mean(self.student_lenbuffer), int(self.tot_time)
+                    )
+            # Print to console
+            log_string = f"""{"#" * width}\n"""
+            log_string += f"""\033[1m{f" Learning iteration {it}/{total_it} ".center(width)}\033[0m \n\n"""
+
+            # Print run name if provided
+            run_name = self.cfg.get("run_name")
+            log_string += f"""{"Run name:":>{pad}} {run_name}\n""" if run_name else ""
+
+            # Print performance
+            log_string += (
+                f"""{"Total steps:":>{pad}} {self.tot_timesteps} \n"""
+                f"""{"Steps per second:":>{pad}} {fps:.0f} \n"""
+                f"""{"Collection time:":>{pad}} {collect_time:.3f}s \n"""
+                f"""{"Learning time:":>{pad}} {learn_time:.3f}s \n"""
+            )
+
+            # Print losses
+            for key, value in loss_dict.items():
+                log_string += f"""{f"Mean {key} loss:":>{pad}} {value:.4f}\n"""
+
+            # Print rewards and episode length
+            if self.cfg["algorithm"]["rnd_cfg"] and len(self.erewbuffer) > 0:
+                log_string += f"""{"Mean extrinsic reward:":>{pad}} {statistics.mean(self.erewbuffer):.2f}\n"""
+                log_string += f"""{"Mean intrinsic reward:":>{pad}} {statistics.mean(self.irewbuffer):.2f}\n"""
+            if len(self.teacher_rewbuffer) > 0:
+                log_string += f"""{"Mean teacher reward:":>{pad}} {statistics.mean(self.teacher_rewbuffer):.2f}\n"""
+                log_string += f"""{"Mean teacher episode length:":>{pad}} {statistics.mean(self.teacher_lenbuffer):.2f}\n"""
+            if len(self.student_rewbuffer) > 0:
+                log_string += f"""{"Mean student reward:":>{pad}} {statistics.mean(self.student_rewbuffer):.2f}\n"""
+                log_string += f"""{"Mean student episode length:":>{pad}} {statistics.mean(self.student_lenbuffer):.2f}\n"""
+            # Print noise std
+            log_string += f"""{"Mean action noise std:":>{pad}} {action_std.mean().item():.2f}\n"""
+
+            # Print episode extras
+            if not print_minimal:
+                log_string += extras_string
+
+            # Print footer
+            done_it = it + 1 - start_it
+            remaining_it = total_it - start_it - done_it
+            eta = self.tot_time / done_it * remaining_it
+            log_string += (
+                f"""{"-" * width}\n"""
+                f"""{"Iteration time:":>{pad}} {iteration_time:.2f}s\n"""
+                f"""{"Time elapsed:":>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n"""
+                f"""{"ETA:":>{pad}} {time.strftime("%H:%M:%S", time.gmtime(eta))}\n"""
+            )
+            print(log_string)
+
+            # Clear extras buffer
+            self.ep_extras.clear()
+
+    def save_model(self, path: str, it: int) -> None:
+        """Save the model to external logging services if specified."""
+        if self.writer and not self.disable_logs and self.logger_type in ["neptune", "wandb"]:
+            self.writer.save_model(path, it)
+
+    def _prepare_logging_writer(self) -> None:
+        """Prepare the logging writer, which can be either Tensorboard, W&B or Neptune."""
+        if self.log_dir is not None and not self.disable_logs:
+            self.logger_type = self.cfg.get("logger", "tensorboard")
+            self.logger_type = self.logger_type.lower()
+
+            if self.logger_type == "neptune":
+                from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter
+
+                self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg)
+            elif self.logger_type == "wandb":
+                from rsl_rl.utils.wandb_utils import WandbSummaryWriter
+
+                self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg)
+            elif self.logger_type == "tensorboard":
+                from torch.utils.tensorboard import SummaryWriter
+
+                self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
+            else:
+                raise ValueError("Logger type not found. Please choose 'wandb', 'neptune', or 'tensorboard'.")
+        else:
+            self.writer = None
+
+    def _store_code_state(self) -> None:
+        """Store the current git diff of the code repositories involved in the experiment."""
+        if self.log_dir is not None and not self.disable_logs:
+            git_log_dir = os.path.join(self.log_dir, "git")
+            os.makedirs(git_log_dir, exist_ok=True)
+            file_paths = []
+            # Iterate over all repositories to log
+            for repository_file_path in self.git_status_repos:
+                try:
+                    repo = git.Repo(repository_file_path, search_parent_directories=True)
+                    t = repo.head.commit.tree
+                except Exception:
+                    print(f"Could not find git repository in {repository_file_path}. Skipping.")
+                    continue
+                # Get the name of the repository
+                repo_name = pathlib.Path(repo.working_dir).name
+                diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff")
+                # Check if the diff file already exists
+                if os.path.isfile(diff_file_name):
+                    continue
+                # Write the diff file
+                print(f"Storing git diff for '{repo_name}' in: {diff_file_name}")
+                with open(diff_file_name, "x", encoding="utf-8") as f:
+                    content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}"
+                    f.write(content)
+                # Add the file path to the list of files to be uploaded
+                file_paths.append(diff_file_name)
+
+            # Upload diff files to external logging services
+            if self.writer and self.logger_type in ["wandb", "neptune"] and file_paths:
+                for path in file_paths:
+                    self.writer.save_file(path)

source/rsl_rl/rsl_rl/utils/neptune_utils.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import os
+from dataclasses import asdict
+from torch.utils.tensorboard import SummaryWriter
+
+try:
+    import neptune
+except ModuleNotFoundError:
+    raise ModuleNotFoundError("neptune-client is required to log to Neptune.") from None
+
+
+class NeptuneSummaryWriter(SummaryWriter):
+    """Summary writer for Neptune."""
+
+    def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
+        super().__init__(log_dir, flush_secs)
+
+        # Get the run name
+        run_name = os.path.split(log_dir)[-1]
+
+        # Get neptune project and entity
+        try:
+            project = cfg["neptune_project"]
+        except KeyError:
+            raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") from None
+        try:
+            token = os.environ["NEPTUNE_API_TOKEN"]
+        except KeyError:
+            raise KeyError(
+                "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN"
+            ) from None
+        try:
+            entity = os.environ["NEPTUNE_USERNAME"]
+        except KeyError:
+            raise KeyError(
+                "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME"
+            ) from None
+
+        # Initialize neptune
+        neptune_project = entity + "/" + project
+        self.run = neptune.init_run(project=neptune_project, api_token=token)
+        self.run["log_dir"].log(run_name)
+
+        # Name mapping for incompatible characters
+        self.name_map = {
+            "Train/mean_reward/time": "Train/mean_reward_time",
+            "Train/mean_episode_length/time": "Train/mean_episode_length_time",
+        }
+
+    def store_config(self, env_cfg: dict | object, train_cfg: dict) -> None:
+        self.run["runner_cfg"] = train_cfg
+        self.run["policy_cfg"] = train_cfg["policy"]
+        self.run["alg_cfg"] = train_cfg["algorithm"]
+        try:
+            self.run["env_cfg"] = env_cfg.to_dict()
+        except Exception:
+            self.run["env_cfg"] = asdict(env_cfg)
+
+    def add_scalar(
+        self,
+        tag: str,
+        scalar_value: float,
+        global_step: int | None = None,
+        walltime: float | None = None,
+        new_style: bool = False,
+    ) -> None:
+        super().add_scalar(
+            tag,
+            scalar_value,
+            global_step=global_step,
+            walltime=walltime,
+            new_style=new_style,
+        )
+        self.run[self._map_path(tag)].log(scalar_value, step=global_step)
+
+    def stop(self) -> None:
+        self.run.stop()
+
+    def save_model(self, model_path: str, it: int) -> None:
+        self.run["model/saved_model_" + str(it)].upload(model_path)
+
+    def save_file(self, path: str) -> None:
+        name = path.rsplit("/", 1)[-1].split(".")[0]
+        self.run["git_diff/" + name].upload(path)
+
+    def _map_path(self, path: str) -> str:
+        if path in self.name_map:
+            return self.name_map[path]
+        else:
+            return path

source/rsl_rl/rsl_rl/utils/utils.py

@@ -0,0 +1,340 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import importlib
+import pkgutil
+import torch
+import warnings
+from tensordict import TensorDict
+from typing import Any, Callable
+
+import rsl_rl
+
+
+def get_param(param: Any, idx: int) -> Any:
+    """Get a parameter for the given index.
+
+    Args:
+        param: Parameter or list/tuple of parameters.
+        idx: Index to get the parameter for.
+    """
+    if isinstance(param, (tuple, list)):
+        return param[idx]
+    else:
+        return param
+
+
+def resolve_nn_activation(act_name: str) -> torch.nn.Module:
+    """Resolve the activation function from the name.
+
+    Args:
+        act_name: Name of the activation function.
+
+    Returns:
+        The activation function.
+
+    Raises:
+        ValueError: If the activation function is not found.
+    """
+    act_dict = {
+        "elu": torch.nn.ELU(),
+        "selu": torch.nn.SELU(),
+        "relu": torch.nn.ReLU(),
+        "crelu": torch.nn.CELU(),
+        "lrelu": torch.nn.LeakyReLU(),
+        "tanh": torch.nn.Tanh(),
+        "sigmoid": torch.nn.Sigmoid(),
+        "softplus": torch.nn.Softplus(),
+        "gelu": torch.nn.GELU(),
+        "swish": torch.nn.SiLU(),
+        "mish": torch.nn.Mish(),
+        "identity": torch.nn.Identity(),
+    }
+
+    act_name = act_name.lower()
+    if act_name in act_dict:
+        return act_dict[act_name]
+    else:
+        raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}")
+
+
+def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer:
+    """Resolve the optimizer from the name.
+
+    Args:
+        optimizer_name: Name of the optimizer.
+
+    Returns:
+        The optimizer.
+
+    Raises:
+        ValueError: If the optimizer is not found.
+    """
+    optimizer_dict = {
+        "adam": torch.optim.Adam,
+        "adamw": torch.optim.AdamW,
+        "sgd": torch.optim.SGD,
+        "rmsprop": torch.optim.RMSprop,
+    }
+
+    optimizer_name = optimizer_name.lower()
+    if optimizer_name in optimizer_dict:
+        return optimizer_dict[optimizer_name]
+    else:
+        raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}")
+
+
+def split_and_pad_trajectories(
+    tensor: torch.Tensor | TensorDict, dones: torch.Tensor
+) -> tuple[torch.Tensor | TensorDict, torch.Tensor]:
+    """Split trajectories at done indices.
+
+    Split trajectories, concatenate them and pad with zeros up to the length of the longest trajectory. Return masks
+    corresponding to valid parts of the trajectories.
+
+    Example (transposed for readability):
+        Input: [[a1, a2, a3, a4 | a5, a6],
+                 [b1, b2 | b3, b4, b5 | b6]]
+
+        Output:[[a1, a2, a3, a4], | [[True, True, True, True],
+                [a5, a6, 0, 0],   |  [True, True, False, False],
+                [b1, b2, 0, 0],   |  [True, True, False, False],
+                [b3, b4, b5, 0],  |  [True, True, True, False],
+                [b6, 0, 0, 0]]    |  [True, False, False, False]]
+
+    Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions]
+    """
+    dones = dones.clone()
+    dones[-1] = 1
+    # Permute the buffers to have the order (num_envs, num_transitions_per_env, ...) for correct reshaping
+    flat_dones = dones.transpose(1, 0).reshape(-1, 1)
+    # Get length of trajectory by counting the number of successive not done elements
+    done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
+    trajectory_lengths = done_indices[1:] - done_indices[:-1]
+    trajectory_lengths_list = trajectory_lengths.tolist()
+    # Extract the individual trajectories
+    if isinstance(tensor, TensorDict):
+        padded_trajectories = {}
+        for k, v in tensor.items():
+            # Split the tensor into trajectories
+            trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list)
+            # Add at least one full length trajectory
+            trajectories = (*trajectories, torch.zeros(v.shape[0], *v.shape[2:], device=v.device))
+            # Pad the trajectories to the length of the longest trajectory
+            padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories)  # type: ignore
+            # Remove the added trajectory
+            padded_trajectories[k] = padded_trajectories[k][:, :-1]
+        padded_trajectories = TensorDict(
+            padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)], device=tensor.device
+        )
+    else:
+        # Split the tensor into trajectories
+        trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list)
+        # Add at least one full length trajectory
+        trajectories = (*trajectories, torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device))
+        # Pad the trajectories to the length of the longest trajectory
+        padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)  # type: ignore
+        # Remove the added trajectory
+        padded_trajectories = padded_trajectories[:, :-1]
+    # Create masks for the valid parts of the trajectories
+    trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1)
+    return padded_trajectories, trajectory_masks
+
+
+def unpad_trajectories(trajectories: torch.Tensor | TensorDict, masks: torch.Tensor) -> torch.Tensor | TensorDict:
+    """Do the inverse operation of `split_and_pad_trajectories()`."""
+    # Need to transpose before and after the masking to have proper reshaping
+    return (
+        trajectories.transpose(1, 0)[masks.transpose(1, 0)]
+        .view(-1, trajectories.shape[0], trajectories.shape[-1])
+        .transpose(1, 0)
+    )
+
+
+def resolve_callable(callable_or_name: type | Callable | str) -> Callable:
+    """Resolve a callable from a string, type, or return callable directly.
+
+    This function enables passing custom classes or functions directly or as strings. The following formats are
+    supported:
+        - Direct callable: Pass a type or function directly (e.g., MyClass, my_func)
+        - Qualified name with colon: "module.path:Attr.Nested" (explicit, recommended)
+        - Qualified name with dot: "module.path.ClassName" (implicit)
+        - Simple name: e.g. "PPO", "ActorCritic", ... (looks for callable in rsl_rl)
+
+    Args:
+        callable_or_name: A callable (type/function) or string name.
+
+    Returns:
+        The resolved callable.
+
+    Raises:
+        TypeError: If input is neither a callable nor a string.
+        ImportError: If the module cannot be imported.
+        AttributeError: If the attribute cannot be found in the module.
+        ValueError: If a simple name cannot be found in rsl_rl packages.
+    """
+    # Already a callable - return directly
+    if callable(callable_or_name):
+        return callable_or_name
+
+    # Must be a string at this point
+    if not isinstance(callable_or_name, str):
+        raise TypeError(f"Expected callable or string, got {type(callable_or_name)}")
+
+    # Handle qualified name with colon separator (e.g., "module.path:Attr.Nested")
+    if ":" in callable_or_name:
+        module_path, attr_path = callable_or_name.rsplit(":", 1)
+        # Try to import the module
+        module = importlib.import_module(module_path)
+        # Try to get the attribute
+        obj = module
+        for attr in attr_path.split("."):
+            obj = getattr(obj, attr)
+        return obj  # type: ignore
+
+    # Handle qualified name with dot separator (e.g., "module.path.ClassName")
+    if "." in callable_or_name:
+        parts = callable_or_name.split(".")
+        module_found = False
+        for i in range(len(parts) - 1, 0, -1):
+            # Try to import the module with the first i parts
+            module_path = ".".join(parts[:i])
+            attr_parts = parts[i:]
+            try:
+                module = importlib.import_module(module_path)
+            except ModuleNotFoundError:
+                continue
+            module_found = True
+            # Once a module is found, try to get the attribute
+            obj = module
+            try:
+                for attr in attr_parts:
+                    obj = getattr(obj, attr)
+                return obj  # type: ignore
+            except AttributeError:
+                continue
+        if module_found:
+            raise AttributeError(f"Could not resolve '{callable_or_name}': attribute not found in module")
+        else:
+            raise ImportError(f"Could not resolve '{callable_or_name}': no valid module.attr split found")
+
+    # Simple name - look for it in rsl_rl
+    for _, module_name, _ in pkgutil.iter_modules(rsl_rl.__path__, "rsl_rl."):
+        module = importlib.import_module(module_name)
+        if hasattr(module, callable_or_name):
+            return getattr(module, callable_or_name)
+
+    # Raise error if no approach worked
+    raise ValueError(
+        f"Could not resolve '{callable_or_name}'. Use qualified name like 'module.path:ClassName' "
+        f"or pass the class directly."
+    )
+
+
+def resolve_obs_groups(
+    obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str]
+) -> dict[str, list[str]]:
+    """Validate the observation configuration and resolve missing observation sets.
+
+    The input is an observation dictionary `obs` containing observation groups and a configuration dictionary
+    `obs_groups` where the keys are the observation sets and the values are lists of observation groups.
+
+    The configuration dictionary could for example look like:
+        {
+            "policy": ["group_1", "group_2"],
+            "critic": ["group_1", "group_3"]
+        }
+
+    This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the 'critic'
+    observation set will contain the observations "group_1" and "group_3". This function will check that all the
+    observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the
+    environment.
+
+    Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, this
+    function will:
+
+    1. Check if a group with the same name exists in the observations and assign this group to the observation set.
+    2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set.
+
+    Args:
+        obs: Observations from the environment in the form of a dictionary.
+        obs_groups: Observation sets configuration.
+        default_sets: Reserved observation set names used by the algorithm (besides 'policy'). If not provided in
+            'obs_groups', a default behavior gets triggered.
+
+    Returns:
+        The resolved observation groups.
+
+    Raises:
+        ValueError: If any observation set is an empty list.
+        ValueError: If any observation set contains an observation term that is not present in the observations.
+    """
+    # Check if policy observation set exists
+    if "policy" not in obs_groups:
+        if "policy" in obs:
+            obs_groups["policy"] = ["policy"]
+            warnings.warn(
+                "The observation configuration dictionary 'obs_groups' must contain the 'policy' key."
+                " As an observation group with the name 'policy' was found, this is assumed to be the observation set."
+                " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity."
+                " This behavior will be removed in a future version."
+            )
+        else:
+            raise ValueError(
+                "The observation configuration dictionary 'obs_groups' must contain the 'policy' key."
+                f" Found keys: {list(obs_groups.keys())}"
+            )
+
+    # Check all observation sets for valid observation groups
+    for set_name, groups in obs_groups.items():
+        # Check if the list is empty
+        if len(groups) == 0:
+            msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list."
+            if set_name in default_sets:
+                if set_name not in obs:
+                    msg += " Consider removing the key to default to the observations used for the 'policy' set."
+                else:
+                    msg += (
+                        f" Consider removing the key to default to the observation '{set_name}' from the environment."
+                    )
+            raise ValueError(msg)
+        # Check groups exist inside the observations from the environment
+        for group in groups:
+            if group not in obs:
+                raise ValueError(
+                    f"Observation '{group}' in observation set '{set_name}' not found in the observations from the"
+                    f" environment. Available observations from the environment: {list(obs.keys())}"
+                )
+
+    # Fill missing observation sets
+    for default_set_name in default_sets:
+        if default_set_name not in obs_groups:
+            if default_set_name in obs:
+                obs_groups[default_set_name] = [default_set_name]
+                warnings.warn(
+                    f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key."
+                    f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the"
+                    f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for"
+                    " clarity. This behavior will be removed in a future version."
+                )
+            else:
+                obs_groups[default_set_name] = obs_groups["policy"].copy()
+                warnings.warn(
+                    f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key."
+                    f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set"
+                    f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for"
+                    " clarity. This behavior will be removed in a future version."
+                )
+
+    # Print the final parsed observation sets
+    print("-" * 80)
+    print("Resolved observation sets: ")
+    for set_name, groups in obs_groups.items():
+        print("\t", set_name, ": ", groups)
+    print("-" * 80)
+
+    return obs_groups

source/rsl_rl/rsl_rl/utils/wandb_utils.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import os
+from dataclasses import asdict
+from torch.utils.tensorboard import SummaryWriter
+
+try:
+    import wandb
+except ModuleNotFoundError:
+    raise ModuleNotFoundError("wandb package is required to log to Weights and Biases.") from None
+
+
+class WandbSummaryWriter(SummaryWriter):
+    """Summary writer for Weights and Biases."""
+
+    def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
+        super().__init__(log_dir, flush_secs)
+
+        # Get the run name
+        run_name = os.path.split(log_dir)[-1]
+
+        # Get wandb project and entity
+        try:
+            project = cfg["wandb_project"]
+        except KeyError:
+            raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") from None
+        try:
+            entity = os.environ["WANDB_USERNAME"]
+        except KeyError:
+            entity = None
+
+        # Initialize wandb
+        wandb.init(project=project, entity=entity, name=run_name)
+        wandb.config.update({"log_dir": log_dir})
+
+    def store_config(self, env_cfg: dict | object, train_cfg: dict) -> None:
+        wandb.config.update({"runner_cfg": train_cfg})
+        wandb.config.update({"policy_cfg": train_cfg["policy"]})
+        wandb.config.update({"alg_cfg": train_cfg["algorithm"]})
+        try:
+            wandb.config.update({"env_cfg": env_cfg.to_dict()})
+        except Exception:
+            wandb.config.update({"env_cfg": asdict(env_cfg)})
+
+    def add_scalar(
+        self,
+        tag: str,
+        scalar_value: float,
+        global_step: int | None = None,
+        walltime: float | None = None,
+        new_style: bool = False,
+    ) -> None:
+        super().add_scalar(
+            tag,
+            scalar_value,
+            global_step=global_step,
+            walltime=walltime,
+            new_style=new_style,
+        )
+        wandb.log({tag: scalar_value}, step=global_step)
+
+    def stop(self) -> None:
+        wandb.finish()
+
+    def save_model(self, model_path: str, it: int) -> None:
+        wandb.save(model_path, base_path=os.path.dirname(model_path))
+
+    def save_file(self, path: str) -> None:
+        wandb.save(path, base_path=os.path.dirname(path))