From f6ec5b68cc96e33f01888e65aa859e538a2d4755 Mon Sep 17 00:00:00 2001 From: wertyuilife Date: Tue, 31 Mar 2026 12:58:24 +0800 Subject: [PATCH] refactor deploy scripts. --- deploy/deploy_mujoco/configs/go2.yaml | 11 +- deploy/deploy_mujoco/deploy_go2.py | 435 ++++++-------------------- deploy/deploy_mujoco/utils.py | 328 +++++++++++++++++++ 3 files changed, 427 insertions(+), 347 deletions(-) create mode 100644 deploy/deploy_mujoco/utils.py diff --git a/deploy/deploy_mujoco/configs/go2.yaml b/deploy/deploy_mujoco/configs/go2.yaml index 626cbc7..6b3a321 100644 --- a/deploy/deploy_mujoco/configs/go2.yaml +++ b/deploy/deploy_mujoco/configs/go2.yaml @@ -1,5 +1,10 @@ -policy_path: "{ROOT_DIR}/deploy/pretrain/go2/go2_moe_cts_1500.pt" -xml_path: "{ROOT_DIR}/resources/go2/stairs.xml" +policy_path: "{ROOT_DIR}/deploy/pre_train/wyh/go2_moe_cts_v1_49500.pt" + +xml_path: "{ROOT_DIR}/resource/stairs.xml" + +render_fps: 240 +video_fps: 60 +save_video: false # Total simulation time simulation_duration: 60000000.0 @@ -10,7 +15,7 @@ control_decimation: 10 base_init_pos: [0.0, 0.0, 0.57] base_init_quat: [1.0, 0.0, 0.0, 0.0] -# For IsaacLab DelayedPDActuatorCfg, if not used, set all delay to 0. +# Based on the implementation of IsaacLab's DelayedPDActuatorCfg. actuator_delay_min: 0 actuator_delay_max: 0 actuator_delay_seed: 0 diff --git a/deploy/deploy_mujoco/deploy_go2.py b/deploy/deploy_mujoco/deploy_go2.py index 1c345c6..af4bb42 100644 --- a/deploy/deploy_mujoco/deploy_go2.py +++ b/deploy/deploy_mujoco/deploy_go2.py @@ -1,387 +1,134 @@ -"""CTS Policy deployment for Unitree Go2 in MuJoCo.""" +"""Simplified CTS policy deployment for Unitree Go2 in MuJoCo.""" import time -from collections import deque -from dataclasses import dataclass from pathlib import Path -from typing import NamedTuple -import imageio import mujoco import mujoco.viewer import numpy as np -import pygame import torch -import yaml -from argparse import ArgumentParser + +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, + setup_tracking_camera, + MujocoRenderUtils, +) -# ============================================================================ -# Types & Constants -# ============================================================================ - -ROOT_DIR = str(Path(__file__).parent.parent.parent) -CONFIG_DIR = f"{ROOT_DIR}/deploy/deploy_mujoco/configs" -VIDEO_DIR = Path(__file__).parent / "videos" - -class CTSPolicyInputs(NamedTuple): - """Input format for CTS policy.""" - policy: torch.Tensor - single_obs: torch.Tensor +CONFIG_NAME = "go2.yaml" +VIDEO_DIR = Path(__file__).with_name("videos") +ACTUATOR_GROUPS = ( + np.array([0, 3, 6, 9], dtype=np.int64), + np.array([1, 4, 7, 10], dtype=np.int64), + np.array([2, 5, 8, 11], dtype=np.int64), +) -@dataclass -class ObsBlockCfg: - """Configuration for observation block.""" - name: str - dim: int # Single frame dimension - - -ACTUATOR_GROUPS = { - "hip": [0, 3, 6, 9], - "thigh": [1, 4, 7, 10], - "calf": [2, 5, 8, 11], -} - - -# ============================================================================ -# Helper Functions -# ============================================================================ - -def get_gravity_orientation(quaternion: np.ndarray) -> np.ndarray: - """Compute gravity vector in body frame from quaternion.""" - qw, qx, qy, qz = quaternion - - gravity = np.zeros(3) - gravity[0] = 2 * (-qz * qx + qw * qy) - gravity[1] = -2 * (qz * qy + qw * qx) - gravity[2] = 1 - 2 * (qw * qw + qz * qz) - - return gravity - - -def pd_control(target_q: np.ndarray, q: np.ndarray, kp: np.ndarray, - target_dq: np.ndarray, dq: np.ndarray, kd: np.ndarray) -> np.ndarray: - """Compute PD control torques.""" - return (target_q - q) * kp + (target_dq - dq) * kd - - -def get_joystick_command(joystick, max_cmd: np.ndarray) -> np.ndarray: - """Read command from Xbox controller.""" - pygame.event.pump() - - dead_zone = 0.1 - axes = [joystick.get_axis(i) for i in [0, 1, 3]] # LX, LY, RX - axes = [0 if abs(a) < dead_zone else a for a in axes] - - cmd = np.array([-axes[1] * max_cmd[0], -axes[0] * max_cmd[1], -axes[2] * max_cmd[2]], dtype=np.float32) - return cmd - - -def load_config(config_file: str) -> dict: - """Load and parse YAML configuration.""" - with open(f"{CONFIG_DIR}/{config_file}", "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - # Replace path placeholders - config["policy_path"] = config["policy_path"].replace("{ROOT_DIR}", ROOT_DIR) - config["xml_path"] = config["xml_path"].replace("{ROOT_DIR}", ROOT_DIR) - - return config - - -def build_observation(obs: np.ndarray, features: dict, obs_cfg: list, history_len: int) -> None: - """Update stacked observation buffer with new frame features (in-place).""" - ptr = 0 - for cfg in obs_cfg: - dim = cfg.dim - start, end = ptr, ptr + dim * history_len - - # Roll history and insert new frame - obs[start:end] = np.roll(obs[start:end], shift=-dim, axis=0) - obs[end - dim:end] = features[cfg.name] - - ptr = end - - -def extract_single_obs(obs: np.ndarray, obs_cfg: list, history_len: int) -> np.ndarray: - """Extract most recent single-frame observation from stacked buffer.""" - single = [] - ptr = 0 - for cfg in obs_cfg: - dim = cfg.dim - block = obs[ptr:ptr + dim * history_len] - single.append(block[-dim:]) # Get last dim elements (most recent frame) - ptr += dim * history_len - return np.concatenate(single, axis=0) - - -def apply_action(action: np.ndarray, default_angles: np.ndarray, - action_pos_scale: float) -> np.ndarray: - """Transform policy action to position targets.""" - return action * action_pos_scale + default_angles - - -def set_initial_state(data: mujoco.MjData, base_pos: np.ndarray, base_quat: np.ndarray, - joint_pos: np.ndarray) -> None: - """Apply the Isaac-style initial pose before the first mj_forward.""" - data.qpos[:] = 0.0 - data.qvel[:] = 0.0 - data.qpos[:3] = base_pos - data.qpos[3:7] = base_quat - data.qpos[7:] = joint_pos - - -def build_delay_buffers(default_pos: np.ndarray, delay_max: int) -> deque: - """Create history buffers for delayed control targets.""" - history_len = delay_max + 1 - pos_history = deque((default_pos.copy() for _ in range(history_len)), maxlen=history_len) - return pos_history - - -def sample_delayed_targets(pos_history: deque, delay_min: int, - delay_max: int, rng: np.random.Generator) -> np.ndarray: - """Apply per-actuator-group discrete control delays.""" - delayed_pos = pos_history[-1].copy() - - for joint_ids in ACTUATOR_GROUPS.values(): - delay_steps = int(rng.integers(delay_min, delay_max + 1)) if delay_max > 0 else 0 - pos_source = pos_history[-1 - delay_steps] - delayed_pos[joint_ids] = pos_source[joint_ids] - - return delayed_pos - - -def init_joystick() -> tuple: - """Initialize pygame and joystick if available.""" - pygame.init() - - if pygame.joystick.get_count() > 0: - joystick = pygame.joystick.Joystick(0) - joystick.init() - print(f"Detected Joystick: {joystick.get_name()}") - return joystick, True - - print("No Joystick detected. Using default commands from config.") - return None, False - - -def setup_viewer_camera(viewer) -> None: - """Configure tracking camera.""" - viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING - viewer.cam.trackbodyid = 1 - viewer.cam.distance = 3.0 - viewer.cam.elevation = -30.0 - viewer.cam.azimuth = 0.0 - -def display_current_command(cmd: np.ndarray) -> None: - """Refresh the current command on the same terminal line.""" - cmd_text = f"\rCurrent command | vx: {cmd[0]: .3f} vy: {cmd[1]: .3f} wz: {cmd[2]: .3f}" - print(cmd_text, end="", flush=True) - - -# ============================================================================ -# Main -# ============================================================================ - -def main(): - parser = ArgumentParser() - parser.add_argument("--save-video", action="store_true", help="Save video of simulation.") - args = parser.parse_args() - - # Configuration - config_file = "go2.yaml" - render_fps = 240 - - config = load_config(config_file) - - # Extract config parameters - sim_cfg = { - "duration": config["simulation_duration"], - "dt": config["simulation_dt"], - "decimation": config["control_decimation"], +def build_features(data, action: np.ndarray, cmd: np.ndarray, cfg): + joint_pos = (data.qpos[7:] - cfg.default_angles) * cfg.dof_pos_scale + joint_vel = data.qvel[6:] * cfg.dof_vel_scale + return { + "ang_vel": data.qvel[3:6] * cfg.ang_vel_scale, + "gravity": gravity_from_quat(data.qpos[3:7]), + "cmd": cmd * cfg.cmd_scale, + "joint_pos": joint_pos[cfg.idx_mj2model], + "joint_vel": joint_vel[cfg.idx_mj2model], + "last_action": action[cfg.idx_mj2model], } - # PD controller gains - kps = np.array(config["kps"], dtype=np.float32) - kds = np.array(config["kds"], dtype=np.float32) - default_angles = np.array(config["default_angles"], dtype=np.float32) - base_init_pos = np.array(config["base_init_pos"], dtype=np.float32) - base_init_quat = np.array(config["base_init_quat"], dtype=np.float32) - delay_min = int(config.get("actuator_delay_min", 0)) - delay_max = int(config.get("actuator_delay_max", 0)) - delay_seed = int(config.get("actuator_delay_seed", 0)) - if delay_min < 0 or delay_max < delay_min: - raise ValueError( - f"Invalid actuator delay range: min={delay_min}, max={delay_max}." - ) - delay_rng = np.random.default_rng(delay_seed) - # Scaling factors - scales = { - "lin_vel": config["lin_vel_scale"], - "ang_vel": config["ang_vel_scale"], - "dof_pos": config["dof_pos_scale"], - "dof_vel": config["dof_vel_scale"], - "action_pos": config["action_pos_scale"], - "cmd": np.array(config["cmd_scale"], dtype=np.float32), - } +def action_to_target(action: np.ndarray, cfg): + return cfg.default_angles + action * cfg.action_pos_scale - # Policy dimensions - num_actions = config["num_actions"] - num_obs = config["num_obs"] - history_len = config.get("history_len", 1) - max_cmd = np.array(config["max_cmd"], dtype=np.float32) - cmd = np.array(config["cmd_init"], dtype=np.float32) - # Joint name mapping - idx_model2mj = idx_mj2model = list(range(num_actions)) - if "mujoco_joint_names" in config and "model_joint_names" in config: - mj_names = config["mujoco_joint_names"] - model_names = config["model_joint_names"] - idx_model2mj = [model_names.index(j) for j in mj_names] - idx_mj2model = [mj_names.index(j) for j in model_names] - - # Initialize joystick - joystick, use_joystick = init_joystick() +def main() -> None: + cfg = load_config(CONFIG_NAME) + layout = [ + ("ang_vel", 3), + ("gravity", 3), + ("cmd", 3), + ("joint_pos", cfg.num_actions), + ("joint_vel", cfg.num_actions), + ("last_action", cfg.num_actions), + ] + joystick = init_joystick() + cmd = cfg.cmd_init.copy() display_current_command(cmd) + model = mujoco.MjModel.from_xml_path(str(cfg.xml_path)) + data = mujoco.MjData(model) + model.opt.timestep = cfg.dt + set_initial_state(data, cfg.base_init_pos, cfg.base_init_quat, cfg.default_angles) + mujoco.mj_forward(model, data) - # Prepare video output - VIDEO_DIR.mkdir(parents=True, exist_ok=True) - model_name = Path(config["policy_path"]).stem - cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}" + renderer = mujoco.Renderer(model, height=360, width=640) + policy = torch.jit.load(str(cfg.policy_path)) + writer, frame_skip, video_path = open_video_writer( + cfg.save_video, policy_path=cfg.policy_path, cmd=cmd, dt=cfg.dt, video_dir=VIDEO_DIR, video_fps=cfg.video_fps + ) - # Initialize state - action = np.zeros(num_actions, dtype=np.float32) - target_dof_pos = default_angles.copy() - target_dof_vel = np.zeros(num_actions, dtype=np.float32) - obs = np.zeros(num_obs * history_len, dtype=np.float32) - pos_history = build_delay_buffers(target_dof_pos, delay_max) - - # Build observation config dynamically - num_joints = num_actions - obs_cfg = [ - ObsBlockCfg("ang_vel", 3), - ObsBlockCfg("gravity", 3), - ObsBlockCfg("cmd", 3), - ObsBlockCfg("joint_pos", num_joints), - ObsBlockCfg("joint_vel", num_actions), - ObsBlockCfg("last_action", num_actions), - ] - - # Load MuJoCo model - m = mujoco.MjModel.from_xml_path(config["xml_path"]) - d = mujoco.MjData(m) - m.opt.timestep = sim_cfg["dt"] - set_initial_state(d, base_init_pos, base_init_quat, default_angles) - mujoco.mj_forward(m, d) - - renderer = mujoco.Renderer(m, height=360, width=640) - - # Load policy - policy = torch.jit.load(config["policy_path"]) - - # Setup video recording - writer = None - if args.save_video: - video_path = VIDEO_DIR / f"{model_name}_{cmd_str}.mp4" - video_fps = 50 - sim_fps = 1.0 / sim_cfg["dt"] - frame_skip = max(1, int(sim_fps / video_fps)) - writer = imageio.get_writer(video_path, fps=video_fps) - print(f"Recording: {video_path} (Sim FPS: {sim_fps:.1f}, Video FPS: {video_fps})") - - render_substeps = int((1.0 / render_fps) / sim_cfg["dt"]) - - # Run simulation - with mujoco.viewer.launch_passive(m, d) as viewer: - setup_viewer_camera(viewer) + 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)) + mujoco_render_utils = MujocoRenderUtils() + with mujoco.viewer.launch_passive(model, data) as viewer: + setup_tracking_camera(viewer) start_time = time.time() counter = 0 - while viewer.is_running() and time.time() - start_time < sim_cfg["duration"]: + while viewer.is_running() and time.time() - start_time < cfg.duration: step_start = time.time() + if joystick and counter % cfg.decimation == 0: + cmd = read_joystick_command(joystick, cfg.max_cmd) - # Update command from joystick - if use_joystick and counter % sim_cfg["decimation"] == 0: - cmd = get_joystick_command(joystick, max_cmd) + data.ctrl[:] = pd_control(target_pos, data.qpos[7:], cfg.kps, target_vel, data.qvel[6:], cfg.kds) + mujoco.mj_step(model, data) + mujoco_render_utils.update(cmd, data) - # Compute and apply control - tau = pd_control(target_dof_pos, d.qpos[7:], kps, target_dof_vel, d.qvel[6:], kds) - d.ctrl[:] = tau - - mujoco.mj_step(m, d) - - # Record frame if writer and counter % frame_skip == 0: try: - renderer.update_scene(d, camera=viewer.cam) + renderer.update_scene(data, camera=viewer.cam) + mujoco_render_utils.update_external_rendering(renderer, ctype='renderer') writer.append_data(renderer.render()) - except Exception as e: - print(f"Render error: {e}") + except Exception as exc: + print(f"Render error: {exc}") counter += 1 - - # Policy update at control frequency - if counter % sim_cfg["decimation"] == 0: - # Extract sensor data - qj = d.qpos[7:] - dqj = d.qvel[6:] - quat = d.qpos[3:7] - ang_vel = d.qvel[3:6] - - # Scale observations - qj = (qj - default_angles) * scales["dof_pos"] - dqj = dqj * scales["dof_vel"] - ang_vel = ang_vel * scales["ang_vel"] - - gravity = get_gravity_orientation(quat) - - # Build observation features - features = { - "ang_vel": ang_vel, - "gravity": gravity, - "cmd": cmd * scales["cmd"], - "joint_pos": qj[idx_mj2model], - "joint_vel": dqj[idx_mj2model], - "last_action": action[idx_mj2model], - } - - # Update stacked observation - build_observation(obs, features, obs_cfg, history_len) - - # Extract single-frame observation - single_obs = extract_single_obs(obs, obs_cfg, history_len) - - # Run policy - obs_tensor = torch.from_numpy(obs).unsqueeze(0) - single_tensor = torch.from_numpy(single_obs).unsqueeze(0) - - result = policy(CTSPolicyInputs(policy=obs_tensor, single_obs=single_tensor)) - action = result.detach().cpu().numpy().squeeze()[idx_model2mj] - - # Apply action - latest_target_pos = apply_action(action, default_angles, scales["action_pos"]) - pos_history.append(latest_target_pos.copy()) - target_dof_pos = sample_delayed_targets(pos_history, delay_min, delay_max, delay_rng) + 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) + 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) - - # Sync viewer - if counter % render_substeps == 0: - viewer.sync() - # Time management - sleep_time = sim_cfg["dt"] - (time.time() - step_start) - 0.1 + if counter % render_substeps == 0: + mujoco_render_utils.update_external_rendering(viewer, ctype='viewer') + viewer.sync() + sleep_time = cfg.dt - (time.time() - step_start) - 0.1 if sleep_time > 0: time.sleep(sleep_time) - print() - # Cleanup + print() if writer: - print(f"Video saved: {video_path}") writer.close() + print(f"Video saved: {video_path}") if __name__ == "__main__": diff --git a/deploy/deploy_mujoco/utils.py b/deploy/deploy_mujoco/utils.py new file mode 100644 index 0000000..c136f94 --- /dev/null +++ b/deploy/deploy_mujoco/utils.py @@ -0,0 +1,328 @@ +"""Shared helpers for MuJoCo deployment scripts.""" + +from collections import deque +from pathlib import Path +from types import SimpleNamespace +from typing import NamedTuple + +import imageio +import numpy as np +import pygame +import torch +import yaml + +from typing import Union, Literal +import mujoco +import mujoco.viewer + +ROOT_DIR = Path(__file__).resolve().parents[2] +CONFIG_DIR = Path(__file__).with_name("configs") + + +class CTSPolicyInputs(NamedTuple): + policy: torch.Tensor + single_obs: torch.Tensor + + +def load_config(config_name: str): + with (CONFIG_DIR / config_name).open("r", encoding="utf-8") as file: + raw = yaml.safe_load(file) + + def path_value(name: str) -> Path: + return Path(raw[name].replace("{ROOT_DIR}", str(ROOT_DIR))) + + data = { + "policy_path": path_value("policy_path"), + "xml_path": path_value("xml_path"), + "duration": float(raw["simulation_duration"]), + "dt": float(raw["simulation_dt"]), + "decimation": int(raw["control_decimation"]), + "history_len": int(raw.get("history_len", 1)), + "num_actions": int(raw["num_actions"]), + "num_obs": int(raw["num_obs"]), + "delay_min": int(raw.get("actuator_delay_min", 0)), + "delay_max": int(raw.get("actuator_delay_max", 0)), + "delay_seed": int(raw.get("actuator_delay_seed", 0)), + "save_video": bool(raw.get("save_video", False)), + "render_fps": int(raw.get("render_fps", 60)), + "video_fps": int(raw.get("video_fps", 50)), + } + if data["delay_min"] < 0 or data["delay_max"] < data["delay_min"]: + raise ValueError(f"Invalid actuator delay range: min={data['delay_min']}, max={data['delay_max']}.") + + for name in ( + "kps", + "kds", + "default_angles", + "torque_limit", + "base_init_pos", + "base_init_quat", + "max_cmd", + "cmd_init", + "cmd_scale", + ): + if name in raw: + data[name] = np.asarray(raw[name], dtype=np.float32) + + for name in ( + "lin_vel_scale", + "ang_vel_scale", + "dof_pos_scale", + "dof_vel_scale", + "action_pos_scale", + "action_vel_scale", + ): + if name in raw: + data[name] = float(raw[name]) + + idx_model2mj = idx_mj2model = np.arange(data["num_actions"], dtype=np.int64) + if "mujoco_joint_names" in raw and "model_joint_names" in raw: + mj_names = raw["mujoco_joint_names"] + model_names = raw["model_joint_names"] + idx_model2mj = np.asarray([model_names.index(name) for name in mj_names], dtype=np.int64) + idx_mj2model = np.asarray([mj_names.index(name) for name in model_names], dtype=np.int64) + data["idx_model2mj"] = idx_model2mj + data["idx_mj2model"] = idx_mj2model + return SimpleNamespace(**data) + + +def gravity_from_quat(quaternion: np.ndarray) -> np.ndarray: + qw, qx, qy, qz = quaternion + return np.array( + [ + 2 * (-qz * qx + qw * qy), + -2 * (qz * qy + qw * qx), + 1 - 2 * (qw * qw + qz * qz), + ], + dtype=np.float32, + ) + + +def pd_control(target_q: np.ndarray, q: np.ndarray, kp: np.ndarray, + target_dq: np.ndarray, dq: np.ndarray, kd: np.ndarray) -> np.ndarray: + return (target_q - q) * kp + (target_dq - dq) * kd + + +def init_joystick(): + pygame.init() + if pygame.joystick.get_count() == 0: + print("No Joystick detected. Using default commands from config.") + return None + joystick = pygame.joystick.Joystick(0) + joystick.init() + print(f"Detected Joystick: {joystick.get_name()}") + return joystick + + +def read_joystick_command(joystick, max_cmd: np.ndarray) -> np.ndarray: + pygame.event.pump() + axes = np.array([joystick.get_axis(i) for i in (0, 1, 3)], dtype=np.float32) + axes[np.abs(axes) < 0.1] = 0.0 + return np.array( + [-axes[1] * max_cmd[0], -axes[0] * max_cmd[1], -axes[2] * max_cmd[2]], + dtype=np.float32, + ) + + +def display_current_command(cmd: np.ndarray) -> None: + print(f"\rCurrent command | vx: {cmd[0]: .3f} vy: {cmd[1]: .3f} wz: {cmd[2]: .3f}", end="", flush=True) + + +def set_initial_state(data, base_pos: np.ndarray, base_quat: np.ndarray, joint_pos: np.ndarray) -> None: + data.qpos[:] = 0.0 + data.qvel[:] = 0.0 + data.qpos[:3] = base_pos + data.qpos[3:7] = base_quat + data.qpos[7:] = joint_pos + + +def setup_tracking_camera(viewer, *, trackbodyid: int = 1, distance: float = 3.0, + elevation: float = -30.0, azimuth: float = 0.0) -> None: + + viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING + viewer.cam.trackbodyid = trackbodyid + viewer.cam.distance = distance + viewer.cam.elevation = elevation + viewer.cam.azimuth = azimuth + + +def open_video_writer(enabled: bool, *, policy_path: Path, cmd: np.ndarray, + dt: float, video_dir: Path, video_fps: int): + if not enabled: + return None, 1, None + video_dir.mkdir(parents=True, exist_ok=True) + video_path = video_dir / f"{policy_path.stem}_cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}.mp4" + frame_skip = max(1, int((1.0 / dt) / video_fps)) + writer = imageio.get_writer(video_path, fps=video_fps) + print(f"Recording: {video_path} (Sim FPS: {1.0 / dt:.1f}, Video FPS: {video_fps})") + return writer, frame_skip, video_path + + +def push_obs_history(obs: np.ndarray, features: dict, layout: list[tuple[str, int]], history_len: int) -> None: + offset = 0 + for name, dim in layout: + block = obs[offset:offset + dim * history_len] + block[:-dim] = block[dim:] + block[-dim:] = features[name] + offset += dim * history_len + + +def latest_obs_frame(obs: np.ndarray, layout: list[tuple[str, int]], history_len: int) -> np.ndarray: + frames = [] + offset = 0 + for _, dim in layout: + end = offset + dim * history_len + frames.append(obs[end - dim:end]) + offset = end + return np.concatenate(frames, axis=0) + + +def infer_action(policy, obs: np.ndarray, single_obs: np.ndarray, idx_model2mj: np.ndarray) -> np.ndarray: + result = policy( + CTSPolicyInputs( + policy=torch.from_numpy(obs).unsqueeze(0), + single_obs=torch.from_numpy(single_obs).unsqueeze(0), + ) + ) + return result.detach().cpu().numpy().squeeze()[idx_model2mj] + + +def build_delay_buffers(*defaults: np.ndarray, delay_max: int): + size = delay_max + 1 + buffers = tuple(deque((value.copy() for _ in range(size)), maxlen=size) for value in defaults) + return buffers[0] if len(buffers) == 1 else buffers + + +def sample_delayed_targets(histories, actuator_groups, delay_min: int, delay_max: int, rng: np.random.Generator): + if isinstance(histories, deque): + histories = (histories,) + delayed = [history[-1].copy() for history in histories] + for joint_ids in actuator_groups: + delay_steps = int(rng.integers(delay_min, delay_max + 1)) if delay_max > 0 else 0 + for i, history in enumerate(histories): + delayed[i][joint_ids] = history[-1 - delay_steps][joint_ids] + return delayed[0] if len(delayed) == 1 else tuple(delayed) + + +class MujocoRenderUtils: + def __init__(self): + self.target_velocity = None + + self.vis_smooth_factor = 1.0 + self.ren_smooth_factor = 1.0 + + self.vis_cur_vel = np.zeros(3) + self.ren_cur_vel = np.zeros(3) + + self.mj_data = None + self._renderer_capacity_warned = False + self._viewer_capacity_warned = False + + def update(self, target_velocity, mj_data): + self.target_velocity = target_velocity + self.mj_data = mj_data + + def update_external_rendering(self, + handle: Union[mujoco.viewer.Handle, mujoco.Renderer], + ctype: Literal['viewer', 'renderer'], + ): + """ Update external rendering handle (viewer or renderer). """ + + def has_geom_slot(container, index: int, warned_attr: str) -> bool: + capacity = len(container.geoms) + if index < capacity: + return True + if not getattr(self, warned_attr): + print(f"Warning: {ctype} geom capacity exceeded ({capacity}); skipping velocity arrows.") + setattr(self, warned_attr, True) + return False + + def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7): + vel_norm = np.linalg.norm(vec) + display_norm = min(vel_norm * scale, 1.0) + + if display_norm < 0.10: + mujoco.mjv_initGeom( + geom_elem, + type=mujoco.mjtGeom.mjGEOM_NONE, + size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0] + ) + return + + mat = np.zeros(9) + target_quat = np.zeros(4) + vec_normalized = vec / vel_norm + mujoco.mju_quatZ2Vec(target_quat, vec_normalized) + mujoco.mju_quat2Mat(mat, target_quat) + + mat = mat.reshape(3, 3) + mat[:, 2] *= display_norm + + mujoco.mjv_initGeom( + geom_elem, + type=mujoco.mjtGeom.mjGEOM_ARROW, + size=[0.02, 0.02, display_norm], # [height, width, length] + pos=pos, + mat=mat.flatten(), + rgba=rgba + ) + + viewer_geom_idx = 0 + if ctype == 'viewer': + handle.user_scn.ngeom = 0 # reset user scene geometry + + if self.target_velocity is not None: + base_pos_world = self.mj_data.qpos[:3] + base_quat = self.mj_data.qpos[3:7] + + # rendering arrows start position + offset_body = np.array([0.0, 0.0, 0.2]) + offset_world = np.zeros(3) + mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat) + start_pos = base_pos_world + offset_world + + tgt_vel_body = np.array([self.target_velocity[0], self.target_velocity[1], 0.0]) + + raw_cur_vel_world = self.mj_data.qvel[:3] + raw_cur_vel = np.zeros(3) + neg_quat = np.zeros(4) + mujoco.mju_negQuat(neg_quat, base_quat) + mujoco.mju_rotVecQuat(raw_cur_vel, raw_cur_vel_world, neg_quat) + cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0]) + + # EMA: v_smooth = alpha * v_new + (1 - alpha) * v_old + # alpha = self.vis_smooth_factor if ctype == 'viewer' else self.ren_smooth_factor + self.vis_cur_vel = cur_vel_body + self.ren_cur_vel = cur_vel_body + + tgt_vel_world = np.zeros(3) + cur_vel_world = np.zeros(3) + mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat) + if ctype == 'viewer': + mujoco.mju_rotVecQuat(cur_vel_world, self.vis_cur_vel, base_quat) + else: + mujoco.mju_rotVecQuat(cur_vel_world, self.ren_cur_vel, base_quat) + + COLOR_CMD = [0, 1, 0, 1] # Green 0x00ff00 + COLOR_REAL = [0, 0, 1, 1] # Blue 0x0000ff + + if ctype == 'viewer': + # Cmd Arrow + if has_geom_slot(handle.user_scn, viewer_geom_idx, "_viewer_capacity_warned"): + add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD) + viewer_geom_idx += 1 + # Real Arrow + if has_geom_slot(handle.user_scn, viewer_geom_idx, "_viewer_capacity_warned"): + add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL) + viewer_geom_idx += 1 + else: + scene = handle.scene + if has_geom_slot(scene, scene.ngeom, "_renderer_capacity_warned"): + add_thick_arrow(scene.geoms[scene.ngeom], start_pos, tgt_vel_world, COLOR_CMD) + scene.ngeom += 1 + if has_geom_slot(scene, scene.ngeom, "_renderer_capacity_warned"): + add_thick_arrow(scene.geoms[scene.ngeom], start_pos, cur_vel_world, COLOR_REAL) + scene.ngeom += 1 + + if ctype == 'viewer': + handle.user_scn.ngeom = viewer_geom_idx