Add robogauge eval; update single obs export jit

MUJOCO_LOG.TXT

@@ -0,0 +1,6 @@
+Sat Apr 11 17:39:23 2026
+WARNING: Nan, Inf or huge value in CTRL at ACTUATOR 9. The simulation is unstable. Time = 4.4000.
+
+Sat Apr 11 17:39:34 2026
+WARNING: Nan, Inf or huge value in CTRL at ACTUATOR 9. The simulation is unstable. Time = 4.4000.
+

deploy/deploy_mujoco/deploy_go2.py

@@ -1,4 +1,16 @@
-"""Simplified CTS policy deployment for Unitree Go2 in MuJoCo."""
+"""Run CTS Go2 policy deployment in MuJoCo with joystick command input.
+
+Overview:
+This script loads a TorchScript policy, steps MuJoCo simulation with PD control, and
+queries the policy at control decimation using only current-frame observations. The
+exported policy manages internal observation history for CTS inference.
+
+Quick Start:
+    python deploy/deploy_mujoco/deploy_go2.py
+
+Notes:
+    This script currently uses CONFIG_NAME = "go2.yaml" and does not expose CLI flags.
+"""
 
 import time
 from pathlib import Path
@@ -12,13 +24,10 @@ from utils import (
     build_delay_buffers,
     display_current_command,
     gravity_from_quat,
-    infer_action,
     init_joystick,
-    latest_obs_frame,
     load_config,
     open_video_writer,
     pd_control,
-    push_obs_history,
     read_joystick_command,
     sample_delayed_targets,
     set_initial_state,
@@ -37,6 +46,17 @@ ACTUATOR_GROUPS = (
 
 
 def build_features(data, action: np.ndarray, cmd: np.ndarray, cfg):
+    """Build the current single-frame feature dictionary for policy inference.
+
+    Args:
+        data: MuJoCo runtime data object.
+        action: Latest action in MuJoCo joint order.
+        cmd: Current command vector.
+        cfg: Loaded deployment configuration.
+
+    Returns:
+        A feature dictionary keyed by observation group name.
+    """
     joint_pos = (data.qpos[7:] - cfg.default_angles) * cfg.dof_pos_scale
     joint_vel = data.qvel[6:] * cfg.dof_vel_scale
     return {
@@ -50,10 +70,33 @@ def build_features(data, action: np.ndarray, cmd: np.ndarray, cfg):
 
 
 def action_to_target(action: np.ndarray, cfg):
+    """Convert normalized action output to joint position targets.
+
+    Args:
+        action: Action in model joint order.
+        cfg: Loaded deployment configuration.
+
+    Returns:
+        Target joint positions in model joint order.
+    """
     return cfg.default_angles + action * cfg.action_pos_scale
 
 
+def build_single_obs(features: dict[str, np.ndarray], layout: list[tuple[str, int]]) -> np.ndarray:
+    """Flatten current feature groups into one single observation vector.
+
+    Args:
+        features: Current-step feature dictionary.
+        layout: Ordered feature layout specification.
+
+    Returns:
+        A concatenated single-frame observation vector.
+    """
+    return np.concatenate([features[name] for name, _ in layout], axis=0).astype(np.float32, copy=False)
+
+
 def main() -> None:
+    """Run MuJoCo simulation and deploy the CTS policy in closed-loop control."""
     cfg = load_config(CONFIG_NAME)
     layout = [
         ("ang_vel", 3),
@@ -82,7 +125,6 @@ def main() -> None:
     action = np.zeros(cfg.num_actions, dtype=np.float32)
     target_pos = cfg.default_angles.copy()
     target_vel = np.zeros(cfg.num_actions, dtype=np.float32)
-    obs = np.zeros(cfg.num_obs * cfg.history_len, dtype=np.float32)
     pos_history = build_delay_buffers(target_pos, delay_max=cfg.delay_max)
     delay_rng = np.random.default_rng(cfg.delay_seed)
     render_substeps = max(1, int((1.0 / cfg.render_fps) / cfg.dt))
@@ -112,8 +154,10 @@ def main() -> None:
 
             counter += 1
             if counter % cfg.decimation == 0:
-                push_obs_history(obs, build_features(data, action, cmd, cfg), layout, cfg.history_len)
-                action = infer_action(policy, obs, latest_obs_frame(obs, layout, cfg.history_len), cfg.idx_model2mj)
+                features = build_features(data, action, cmd, cfg)
+                single_obs = build_single_obs(features, layout)
+                action_tensor = policy(torch.from_numpy(single_obs).unsqueeze(0))
+                action = action_tensor.detach().cpu().numpy().squeeze()[cfg.idx_model2mj]
                 pos_history.append(action_to_target(action, cfg).copy())
                 target_pos = sample_delayed_targets(pos_history, ACTUATOR_GROUPS, cfg.delay_min, cfg.delay_max, delay_rng)
                 display_current_command(cmd)

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)

source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py

@@ -9,6 +9,8 @@ import os
 import time
 import torch
 import warnings
+import yaml
+import numpy as np
 from tensordict import TensorDict
 
 from rsl_rl.algorithms import MoECTS
@@ -21,6 +23,21 @@ from rsl_rl.modules import (
 from rsl_rl.storage import RolloutStorageCTS
 from rsl_rl.utils import resolve_callable, resolve_obs_groups
 from rsl_rl.utils.logger_cts import LoggerCTS
+from rsl_rl.utils.exporter_cts import export_cts_policy_as_jit
+
+
+def numpy_representer(dumper: yaml.SafeDumper, data: np.floating) -> yaml.Node:
+    return dumper.represent_float(float(data))
+
+
+def numpy_int_representer(dumper: yaml.SafeDumper, data: np.integer) -> yaml.Node:
+    return dumper.represent_int(int(data))
+
+
+yaml.add_representer(np.float32, numpy_representer, Dumper=yaml.SafeDumper)
+yaml.add_representer(np.float64, numpy_representer, Dumper=yaml.SafeDumper)
+yaml.add_representer(np.int32, numpy_int_representer, Dumper=yaml.SafeDumper)
+yaml.add_representer(np.int64, numpy_int_representer, Dumper=yaml.SafeDumper)
 
 
 class OnPolicyRunnerCTS:
@@ -58,6 +75,18 @@ class OnPolicyRunnerCTS:
 
         self.current_learning_iteration = 0
 
+        # robogauge client
+        try:
+            robogauge_cfg = train_cfg.get("robogauge", {})
+            if not robogauge_cfg.get("enabled", False):
+                raise ImportError("config disabled")
+            from robogauge.scripts.client import RoboGaugeClient
+
+            self.robogauge_client = RoboGaugeClient(f"http://127.0.0.1:{robogauge_cfg.get('port', 9973)}")
+        except Exception as e:
+            print(f"[INFO] RoboGauge client could not be initialized: {e}, disabling RoboGauge interface.")
+            self.robogauge_client = None
+
     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:
@@ -124,13 +153,17 @@ class OnPolicyRunnerCTS:
 
             # Save model
             if it % self.cfg["save_interval"] == 0:
-                self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt"))  # type: ignore
+                self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt"), it=it, last_model=False)  # 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"))
+            self.save(
+                os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"),
+                it=self.current_learning_iteration,
+                last_model=True,
+            )
 
-    def save(self, path: str, infos: dict | None = None) -> None:
+    def save(self, path: str, it: int, last_model: bool, infos: dict | None = None) -> None:
         # Save model
         saved_dict = {
             "model_state_dict": self.alg.policy.state_dict(),
@@ -148,6 +181,80 @@ class OnPolicyRunnerCTS:
 
         # Upload model to external logging services
         self.logger.save_model(path, self.current_learning_iteration)
+        self.update_robogauge(it, last_model)
+
+    def update_robogauge(self, it: int, last_model: bool) -> None:
+        if self.robogauge_client is None or self.logger.log_dir is None or self.logger.disable_logs:
+            return
+
+        try:
+            if it % 500 == 0 or last_model:
+                # export jit model
+                jit_dir = os.path.join(self.logger.log_dir, "jit_models")
+                jit_path = os.path.join(jit_dir, f"policy_jit_{it}.pt")
+                export_cts_policy_as_jit(
+                    self.alg.policy,
+                    actor_obs_normalizer=self.alg.policy.actor_obs_normalizer,
+                    single_obs_normalizer=self.alg.policy.single_obs_normalizer,
+                    path=jit_dir,
+                    filename=f"policy_jit_{it}.pt",
+                )
+                # upload to robogauge
+                self.robogauge_client.submit_task(
+                    model_path=jit_path,
+                    step=it,
+                    task_name="go2_lab",
+                    experiment_name=self.cfg["experiment_name"],
+                )
+        except Exception as e:
+            print(f"[WARN] RoboGauge submit failed at step {it}: {e}")
+            return
+
+        check_times = 1
+        if last_model:
+            check_times = int(1e9)  # keep checking until manually stopped
+
+        while check_times > 0:
+            check_times -= 1
+            try:
+                self.robogauge_client.monitor_tasks()
+            except Exception as e:
+                print(f"[WARN] RoboGauge monitor failed at step {it}: {e}")
+                break
+
+            results_dir = os.path.join(self.logger.log_dir, "robogauge_results")
+            os.makedirs(results_dir, exist_ok=True)
+            result_received = False
+
+            for task_id, resp in self.robogauge_client.response_data.items():
+                if not isinstance(resp, dict):
+                    print(f"[WARN] RoboGauge returned an invalid response for task {task_id}: {resp}")
+                    continue
+                results = resp.get("results")
+                step = resp.get("step", it)
+                if results is None:
+                    print(f"[WARN] RoboGauge returned empty results for task {task_id} at step {step}.")
+                    continue
+                scores = results.get("scores")
+                if scores is None:
+                    print(f"[WARN] RoboGauge results for task {task_id} at step {step} do not contain 'scores'.")
+                    continue
+                if step == it:
+                    result_received = True
+                if self.logger.writer is not None:
+                    for key, val in scores.items():
+                        self.logger.writer.add_scalar(f"RoboGauge/{key}", val, step)
+                results_path = os.path.join(results_dir, f"results_{step}.yaml")
+                with open(results_path, "w", encoding="utf-8") as f:
+                    yaml.dump(results, f, allow_unicode=True, sort_keys=False)
+
+            if last_model and result_received:
+                print(f"RoboGauge result for step {it} received. Exiting wait loop.")
+                break
+
+            if check_times > 0:
+                print("Sleeping for 1 minute before checking RoboGauge results again...")
+                time.sleep(60)  # wait for 1 minute before checking again
 
     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)

source/rsl_rl/rsl_rl/utils/__init__.py

@@ -14,8 +14,10 @@ from .utils import (
     split_and_pad_trajectories,
     unpad_trajectories,
 )
+from .exporter_cts import export_cts_policy_as_jit
 
 __all__ = [
+    "export_cts_policy_as_jit",
     "get_param",
     "resolve_callable",
     "resolve_nn_activation",

source/rsl_rl/rsl_rl/utils/exporter_cts.py

@@ -0,0 +1,98 @@
+import copy
+import os
+import torch
+
+
+def export_cts_policy_as_jit(
+    policy: object,
+    actor_obs_normalizer: object | None,
+    single_obs_normalizer: object | None,
+    path: str,
+    filename: str = "policy.pt",
+) -> None:
+    """Export CTS policy into a Torch JIT file with single_obs input."""
+    policy_exporter = _TorchPolicyExporter(policy, actor_obs_normalizer, single_obs_normalizer)
+    policy_exporter.export(path, filename)
+
+
+class _TorchPolicyExporter(torch.nn.Module):
+    """Exporter of CTS actor-critic into JIT file."""
+
+    def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None):
+        assert not policy.is_recurrent, "CTS policy should not be recurrent"
+        super().__init__()
+
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+        elif hasattr(policy, "student"):
+            self.actor = copy.deepcopy(policy.student)
+        else:
+            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
+        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))
+
+        if actor_obs_normalizer:
+            self.actor_obs_normalizer = copy.deepcopy(actor_obs_normalizer)
+        else:
+            self.actor_obs_normalizer = torch.nn.Identity()
+        if single_obs_normalizer:
+            self.single_obs_normalizer = copy.deepcopy(single_obs_normalizer)
+        else:
+            self.single_obs_normalizer = torch.nn.Identity()
+
+    def forward(self, single_obs: torch.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:
+            return self.actor(latent_and_obs)[..., 0, :]
+        return self.actor(latent_and_obs)
+
+    @torch.jit.export
+    def reset(self):
+        self.obs_history.zero_()
+
+    def export(self, path, filename):
+        os.makedirs(path, exist_ok=True)
+        path = os.path.join(path, filename)
+        self.to("cpu")
+        traced_script_module = torch.jit.script(self)
+        traced_script_module.save(path)