refactor deploy scripts.

deploy/deploy_mujoco/configs/go2.yaml

@@ -1,5 +1,10 @@
-policy_path: "{ROOT_DIR}/deploy/pretrain/go2/go2_moe_cts_1500.pt"
+policy_path: "{ROOT_DIR}/deploy/pre_train/wyh/go2_moe_cts_v1_49500.pt"
-xml_path: "{ROOT_DIR}/resources/go2/stairs.xml"
+
+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

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,
-# Types & Constants
+    display_current_command,
-# ============================================================================
+    gravity_from_quat,
-
+    infer_action,
-ROOT_DIR = str(Path(__file__).parent.parent.parent)
+    init_joystick,
-CONFIG_DIR = f"{ROOT_DIR}/deploy/deploy_mujoco/configs"
+    latest_obs_frame,
-VIDEO_DIR = Path(__file__).parent / "videos"
+    load_config,
-
+    open_video_writer,
-class CTSPolicyInputs(NamedTuple):
+    pd_control,
-    """Input format for CTS policy."""
+    push_obs_history,
-    policy: torch.Tensor
+    read_joystick_command,
-    single_obs: torch.Tensor
+    sample_delayed_targets,
-
+    set_initial_state,
-
+    setup_tracking_camera,
-@dataclass
+    MujocoRenderUtils,
-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"],
-    }
-
-    # 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 = {
+CONFIG_NAME = "go2.yaml"
-        "lin_vel": config["lin_vel_scale"],
+VIDEO_DIR = Path(__file__).with_name("videos")
-        "ang_vel": config["ang_vel_scale"],
+ACTUATOR_GROUPS = (
-        "dof_pos": config["dof_pos_scale"],
+    np.array([0, 3, 6, 9], dtype=np.int64),
-        "dof_vel": config["dof_vel_scale"],
+    np.array([1, 4, 7, 10], dtype=np.int64),
-        "action_pos": config["action_pos_scale"],
+    np.array([2, 5, 8, 11], dtype=np.int64),
-        "cmd": np.array(config["cmd_scale"], dtype=np.float32),
+)
+
+
+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],
     }
 
-    # 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
+def action_to_target(action: np.ndarray, cfg):
-    idx_model2mj = idx_mj2model = list(range(num_actions))
+    return cfg.default_angles + action * cfg.action_pos_scale
-    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
+    renderer = mujoco.Renderer(model, height=360, width=640)
-    VIDEO_DIR.mkdir(parents=True, exist_ok=True)
+    policy = torch.jit.load(str(cfg.policy_path))
-    model_name = Path(config["policy_path"]).stem
+    writer, frame_skip, video_path = open_video_writer(
-    cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
+        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(cfg.num_actions, dtype=np.float32)
-    action = np.zeros(num_actions, dtype=np.float32)
+    target_pos = cfg.default_angles.copy()
-    target_dof_pos = default_angles.copy()
+    target_vel = np.zeros(cfg.num_actions, dtype=np.float32)
-    target_dof_vel = np.zeros(num_actions, dtype=np.float32)
+    obs = np.zeros(cfg.num_obs * cfg.history_len, dtype=np.float32)
-    obs = np.zeros(num_obs * history_len, dtype=np.float32)
+    pos_history = build_delay_buffers(target_pos, delay_max=cfg.delay_max)
-    pos_history = build_delay_buffers(target_dof_pos, delay_max)
+    delay_rng = np.random.default_rng(cfg.delay_seed)
-
+    render_substeps = max(1, int((1.0 / cfg.render_fps) / cfg.dt))
-    # Build observation config dynamically
+    mujoco_render_utils = MujocoRenderUtils()
-    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)
 
+    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
+            data.ctrl[:] = pd_control(target_pos, data.qpos[7:], cfg.kps, target_vel, data.qvel[6:], cfg.kds)
-            if use_joystick and counter % sim_cfg["decimation"] == 0:
+            mujoco.mj_step(model, data)
-                cmd = get_joystick_command(joystick, max_cmd)
+            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:
+                except Exception as exc:
-                    print(f"Render error: {e}")
+                    print(f"Render error: {exc}")
 
             counter += 1
-
+            if counter % cfg.decimation == 0:
-            # Policy update at control frequency
+                push_obs_history(obs, build_features(data, action, cmd, cfg), layout, cfg.history_len)
-            if counter % sim_cfg["decimation"] == 0:
+                action = infer_action(policy, obs, latest_obs_frame(obs, layout, cfg.history_len), cfg.idx_model2mj)
-                # Extract sensor data
+                pos_history.append(action_to_target(action, cfg).copy())
-                qj = d.qpos[7:]
+                target_pos = sample_delayed_targets(pos_history, ACTUATOR_GROUPS, cfg.delay_min, cfg.delay_max, delay_rng)
-                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)
                 display_current_command(cmd)
 
-            # Sync viewer
             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
-            # Time management
-            sleep_time = sim_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__":

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