Add robogauge eval; update single obs export jit

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)