Add robogauge eval; update single obs export jit

scripts/rsl_rl/cli_args.py

@@ -40,6 +40,12 @@ def add_rsl_rl_args(parser: argparse.ArgumentParser):
     arg_group.add_argument(
         "--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune."
     )
+    arg_group.add_argument(
+        "--robogauge", action="store_true", default=False, help="Enable robogauge evaluation interface."
+    )
+    arg_group.add_argument(
+        "--robogauge_port", type=int, default=9973, help="Port for robogauge evaluation interface."
+    )
 
 
 def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlBaseRunnerCfg:

scripts/rsl_rl/train.py

@@ -108,6 +108,11 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
     """Train with RSL-RL agent."""
     # override configurations with non-hydra CLI arguments
     agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
+    agent_cfg_dict = agent_cfg.to_dict()
+    agent_cfg_dict["robogauge"] = {
+        "enabled": args_cli.robogauge,
+        "port": args_cli.robogauge_port,
+    }
     env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
     agent_cfg.max_iterations = (
         args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
@@ -185,11 +190,11 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
 
     # create runner from rsl-rl
     if agent_cfg.class_name == "OnPolicyRunner":
-        runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
+        runner = OnPolicyRunner(env, agent_cfg_dict, log_dir=log_dir, device=agent_cfg.device)
     elif agent_cfg.class_name == "DistillationRunner":
-        runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
+        runner = DistillationRunner(env, agent_cfg_dict, log_dir=log_dir, device=agent_cfg.device)
     elif agent_cfg.class_name == "OnPolicyRunnerCTS":
-        runner = OnPolicyRunnerCTS(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
+        runner = OnPolicyRunnerCTS(env, agent_cfg_dict, log_dir=log_dir, device=agent_cfg.device)
     else:
         raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
     # write git state to logs

scripts/rsl_rl/utils.py

@@ -6,7 +6,6 @@ import torch
 import re
 import os
 import sys
-from typing import NamedTuple
 
 "Script to log terminal output to a file, stripping ANSI escape codes."
 class Logger:
@@ -28,11 +27,6 @@ class Logger:
         self.terminal.flush()
         self.log.flush()
 
-# Inputs of CTS Policy is a TensorDict with 'policy' and 'single_obs' keys, we simulate this with a NamedTuple.
-class CTSPolicyInputs(NamedTuple):
-    policy: torch.Tensor
-    single_obs: torch.Tensor
-
 def export_cts_policy_as_jit(policy: object, actor_obs_normalizer: object | None, single_obs_normalizer: object | None, path: str, filename="policy.pt"):
     """Export CTS policy into a Torch JIT file.
 
@@ -75,6 +69,16 @@ class _TorchPolicyExporter(torch.nn.Module):
     """Exporter of actor-critic into JIT file."""
 
     def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None):
+        """Initialize a TorchScript exporter for CTS policy inference.
+
+        The exported model consumes only the current `single_obs` frame and maintains
+        internal stacked history to reconstruct actor observations expected by the student encoder.
+
+        Args:
+            policy: Source CTS policy module to export.
+            actor_obs_normalizer: Normalizer applied to stacked actor observations.
+            single_obs_normalizer: Normalizer applied to current single-frame observations.
+        """
         assert not policy.is_recurrent, "CTS policy should not be recurrent"
         super().__init__()
         
@@ -87,6 +91,22 @@ class _TorchPolicyExporter(torch.nn.Module):
             raise ValueError("Policy does not have an actor/student module.")
         self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
         self.state_dependent_std = policy.state_dependent_std
+        self.num_actions = int(policy.num_actions)
+        self.num_single_obs = int(policy.num_single_obs)
+        self.num_actor_obs = int(policy.num_actor_obs)
+        if self.num_actor_obs % self.num_single_obs != 0:
+            raise ValueError(
+                f"num_actor_obs ({self.num_actor_obs}) must be divisible by num_single_obs ({self.num_single_obs})."
+            )
+        self.history_len = self.num_actor_obs // self.num_single_obs
+        # Keep the same per-term history layout as deploy-side push_obs_history:
+        # [ang_vel(3), gravity(3), cmd(3), joint_pos(A), joint_vel(A), last_action(A)].
+        self.feature_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        if sum(self.feature_dims) != self.num_single_obs:
+            raise ValueError(
+                "Unsupported single_obs layout: expected 3+3+3+3*num_actions to match num_single_obs."
+            )
+        self.register_buffer("obs_history", torch.zeros(1, self.num_actor_obs, dtype=torch.float32))
         
         # copy normalizer if exists
         if actor_obs_normalizer:
@@ -98,9 +118,43 @@ class _TorchPolicyExporter(torch.nn.Module):
         else:
             self.single_obs_normalizer = torch.nn.Identity()
 
-    def forward(self, x: CTSPolicyInputs):
-        single_obs = self.single_obs_normalizer(x.single_obs)
-        obs_a = self.actor_obs_normalizer(x.policy)
+    def forward(self, single_obs: torch.Tensor):
+        """Compute policy action from one current observation frame.
+
+        The exporter keeps an internal FIFO history buffer and shifts it by one frame
+        on each forward call before appending the latest observation.
+
+        Args:
+            single_obs: Current-step observation tensor with shape `[B, num_single_obs]`.
+
+        Returns:
+            The policy action tensor.
+        """
+        if single_obs.dim() == 1:
+            single_obs = single_obs.unsqueeze(0)
+        if single_obs.shape[-1] != self.num_single_obs:
+            raise ValueError(
+                f"Expected single_obs last dimension {self.num_single_obs}, got {single_obs.shape[-1]}."
+            )
+        if single_obs.shape[0] != 1:
+            raise ValueError("TorchScript CTS deployment currently supports batch size 1 only.")
+
+        next_history = self.obs_history.clone()
+        history_offset = 0
+        single_offset = 0
+        for dim in self.feature_dims:
+            block_size = dim * self.history_len
+            block_end = history_offset + block_size
+            single_end = single_offset + dim
+            block = self.obs_history[:, history_offset:block_end]
+            shifted_block = torch.cat([block[:, dim:], single_obs[:, single_offset:single_end]], dim=-1)
+            next_history[:, history_offset:block_end] = shifted_block
+            history_offset = block_end
+            single_offset = single_end
+        self.obs_history.copy_(next_history)
+
+        single_obs = self.single_obs_normalizer(single_obs)
+        obs_a = self.actor_obs_normalizer(self.obs_history)
         latent, _ = self.student_moe_encoder(obs_a)
         latent_and_obs = torch.cat([latent, single_obs], dim=-1)
         if self.state_dependent_std:
@@ -110,7 +164,8 @@ class _TorchPolicyExporter(torch.nn.Module):
         
     @torch.jit.export
     def reset(self):
-        pass
+        """Reset internal observation history state."""
+        self.obs_history.zero_()
 
     def export(self, path, filename):
         os.makedirs(path, exist_ok=True)