diff --git a/MUJOCO_LOG.TXT b/MUJOCO_LOG.TXT new file mode 100644 index 0000000..cf1c87a --- /dev/null +++ b/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. + diff --git a/deploy/deploy_mujoco/deploy_go2.py b/deploy/deploy_mujoco/deploy_go2.py index af4bb42..1f682bc 100644 --- a/deploy/deploy_mujoco/deploy_go2.py +++ b/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) diff --git a/scripts/rsl_rl/cli_args.py b/scripts/rsl_rl/cli_args.py index fa6f942..161a0c9 100644 --- a/scripts/rsl_rl/cli_args.py +++ b/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: diff --git a/scripts/rsl_rl/train.py b/scripts/rsl_rl/train.py index 4527218..fe3dac3 100644 --- a/scripts/rsl_rl/train.py +++ b/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 diff --git a/scripts/rsl_rl/utils.py b/scripts/rsl_rl/utils.py index d018d1a..2902ef3 100644 --- a/scripts/rsl_rl/utils.py +++ b/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) diff --git a/source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py b/source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py index 43bbeae..da5d551 100644 --- a/source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py +++ b/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) diff --git a/source/rsl_rl/rsl_rl/utils/__init__.py b/source/rsl_rl/rsl_rl/utils/__init__.py index 83c2bf7..eef2fe1 100644 --- a/source/rsl_rl/rsl_rl/utils/__init__.py +++ b/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", diff --git a/source/rsl_rl/rsl_rl/utils/exporter_cts.py b/source/rsl_rl/rsl_rl/utils/exporter_cts.py new file mode 100644 index 0000000..4fb073d --- /dev/null +++ b/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)