Initial commit

demo/export_skrl_policy_to_onnx.py

@@ -0,0 +1,185 @@
+"""
+Export SKRL PyTorch policy to ONNX format.
+
+Usage:
+    python demo/export_skrl_policy_to_onnx.py --checkpoint <path_to_checkpoint.pt> --output <output_dir>
+
+Example:
+    python demo/export_skrl_policy_to_onnx.py \
+        --checkpoint runs/go1-rough-terrain-walk/skrl/[time]_PPO/checkpoints/best_agent.pt \
+        --output exports_go1_rough
+"""
+
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# ============== PyTorch Checkpoint Export ==============
+
+class SKRLPolicyTorch(nn.Module):
+    """PyTorch policy matching SKRL GaussianMixin architecture.
+
+    The architecture is inferred from the checkpoint keys:
+    - net.N: Hidden layers (Linear -> ELU)
+    - mean_layer: Output layer (no activation)
+    """
+
+    def __init__(self, obs_dim, action_dim, hidden_dims=[512, 256, 128]):
+        super().__init__()
+        self.obs_dim = obs_dim
+        self.action_dim = action_dim
+        self.hidden_dims = hidden_dims
+
+        layers = []
+        in_dim = obs_dim
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(in_dim, hidden_dim),
+                nn.ELU(),
+            ])
+            in_dim = hidden_dim
+
+        self.net = nn.Sequential(*layers)
+        self.mean_layer = nn.Linear(in_dim, action_dim)
+        self.log_std = nn.Parameter(torch.zeros(action_dim))
+
+    def forward(self, x):
+        return self.mean_layer(self.net(x))
+
+
+class ONNXPolicyExporterTorch(nn.Module):
+    """Exporter wrapper that applies normalization and policy.
+
+    Matches SKRL's RunningStandardScaler behavior:
+    - Normalize: (x - mean) / sqrt(var + eps)
+    - Clip to [-5.0, 5.0] (clip_threshold)
+    """
+
+    def __init__(self, policy, normalizer_mean, normalizer_std):
+        super().__init__()
+        self.policy = policy
+        # Register normalizer buffers as float32 (matching RunningStandardScaler)
+        self.register_buffer('mean', normalizer_mean.float())
+        self.register_buffer('std', normalizer_std.float())
+        self.clip_threshold = 5.0
+
+    def forward(self, x):
+        # Normalize (RunningStandardScaler._compute)
+        x = (x - self.mean) / (self.std + 1e-8)
+        # Clip to RunningStandardScaler clip_threshold
+        x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
+        return self.policy(x)
+
+
+def infer_architecture_from_state_dict(state_dict):
+    """Infer the policy architecture from state dict keys.
+
+    Args:
+        state_dict: Policy state dict
+
+    Returns:
+        tuple: (obs_dim, action_dim, hidden_dims)
+    """
+    # Find observation dimension from first layer
+    obs_dim = state_dict['net.0.weight'].shape[1]
+
+    # Find action dimension from mean_layer
+    action_dim = state_dict['mean_layer.weight'].shape[0]
+
+    # Infer hidden dimensions from net layers
+    hidden_dims = []
+    net_keys = sorted([k for k in state_dict.keys() if k.startswith('net.') and k.endswith('.weight')])
+    for key in net_keys:
+        if 'mean_layer' not in key:  # Skip output layer
+            layer_idx = int(key.split('.')[1])
+            if layer_idx % 2 == 0:  # Only Linear layers, not activation layers
+                hidden_dims.append(state_dict[key].shape[0])
+
+    return obs_dim, action_dim, hidden_dims
+
+
+def export_torch_policy_to_onnx(checkpoint_path, output_path, verbose=False):
+    """Export PyTorch SKRL policy to ONNX.
+
+    Args:
+        checkpoint_path: Path to .pt checkpoint file
+        output_path: Directory to save ONNX file
+        verbose: Whether to print verbose output
+    """
+    print(f"Loading PyTorch checkpoint from: {checkpoint_path}")
+    state = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+
+    policy_state = state['policy']
+    normalizer_state = state['state_preprocessor']
+
+    # Infer architecture
+    obs_dim, action_dim, hidden_dims = infer_architecture_from_state_dict(policy_state)
+    print(f"Inferred architecture: obs={obs_dim}, action={action_dim}, hidden={hidden_dims}")
+
+    # Create model
+    print("Creating PyTorch model...")
+    policy = SKRLPolicyTorch(obs_dim, action_dim, hidden_dims)
+
+    # Remove value layers from state dict (keep only policy)
+    policy_state_policy_only = {k: v for k, v in policy_state.items()
+                                if not k.startswith('value') and 'value_layer' not in k}
+    policy.load_state_dict(policy_state_policy_only, strict=False)
+    policy.eval()
+
+    # Load normalizer
+    mean = normalizer_state['running_mean']
+    std = torch.sqrt(normalizer_state['running_variance'])
+
+    # Create exporter
+    exporter = ONNXPolicyExporterTorch(policy, mean, std)
+    exporter.eval()
+
+    # Create output directory
+    os.makedirs(output_path, exist_ok=True)
+    output_file = os.path.join(output_path, "policy.onnx")
+
+    # Export to ONNX
+    print(f"Exporting policy to: {output_file}")
+    obs = torch.zeros(1, obs_dim, dtype=torch.float32)
+
+    torch.onnx.export(
+        exporter,
+        obs,
+        output_file,
+        export_params=True,
+        opset_version=11,
+        verbose=verbose,
+        input_names=["observations"],
+        output_names=["actions"],
+        dynamic_axes={},
+    )
+
+    print(f"Successfully exported policy to ONNX: {output_file}")
+
+    # Save PyTorch model and normalizer
+    torch.save(policy.state_dict(), os.path.join(output_path, "policy.pt"))
+    np.savez(os.path.join(output_path, "normalizer.npz"),
+             mean=mean.numpy(), std=std.numpy())
+
+    return output_file
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Export SKRL PyTorch policy to ONNX")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint file (.pt)")
+    parser.add_argument("--output", type=str, default="exports", help="Output directory")
+    parser.add_argument("--verbose", action="store_true", help="Verbose output")
+
+    args = parser.parse_args()
+
+    output_file = export_torch_policy_to_onnx(
+        checkpoint_path=args.checkpoint,
+        output_path=args.output,
+        verbose=args.verbose
+    )
+
+    print(f"\nExported to: {output_file}")

demo/go1_sim2sim_mujoco.py

@@ -0,0 +1,371 @@
+#!/usr/bin/env python3
+"""
+    MOTRIXLAB_UNTRIEE_GO1_SIM2SIM
+    source /opt/mujoco/venv/bin/activate
+    cd /opt/unitree_mujoco
+    python demo/go1_sim2sim_mujoco.py
+
+"""
+
+import numpy as np
+import mujoco
+from mujoco import viewer
+import os
+import threading
+import signal
+import queue
+import argparse
+import time
+
+g_exit_requested = False
+
+def signal_handler(signum, frame):
+    global g_exit_requested
+    g_exit_requested = True
+
+signal.signal(signal.SIGINT, signal_handler)
+
+# ============================================================
+# 配置
+# ============================================================
+DEFAULT_ONNX_PATH = "/opt/motrixlab/repo/export_flat/policy.onnx"
+MOTRIX_XML_DIR = "/opt/motrixlab/repo/motrix_envs/src/motrix_envs/locomotion/go1/xmls"
+XML_PATH = f"{MOTRIX_XML_DIR}/go1_motor_actuator.xml"
+
+TERRAIN = "none"
+
+# ============================================================
+# MotrixLab 参数 (来自 cfg.py)
+# ============================================================
+NUM_OBS = 48
+NUM_ACTIONS = 12
+OBS_SCALES = {'lin_vel': 2.0, 'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
+ACTION_SCALE = 0.05
+KP = 80.0
+KD = 1.0
+CLIP_ACTIONS = 23.7
+CLIP_OBSERVATIONS = 100.0
+MAX_LIN_VEL_X = 1.0
+MAX_LIN_VEL_Y = 1.0
+MAX_ANG_VEL = 0.5
+
+# ============================================================
+# 关节名称和顺序
+# ============================================================
+POLICY_JOINT_NAMES = [
+    "FR_hip", "FR_thigh", "FR_calf",
+    "FL_hip", "FL_thigh", "FL_calf",
+    "RR_hip", "RR_thigh", "RR_calf",
+    "RL_hip", "RL_thigh", "RL_calf",
+]
+
+DEFAULT_JOINT_ANGLES = np.array([
+    -0.0,  0.9, -1.8,   # FR_hip, FR_thigh, FR_calf
+     0.0,  0.9, -1.8,   # FL_hip, FL_thigh, FL_calf
+    -0.0,  0.9, -1.8,   # RR_hip, RR_thigh, RR_calf
+     0.0,  0.9, -1.8,   # RL_hip, RL_thigh, RL_calf
+], dtype=np.float32)
+
+MUJOCO_TO_POLICY = np.arange(12, dtype=np.int64)
+POLICY_TO_MUJOCO = np.arange(12, dtype=np.int64)
+
+# ============================================================
+# 键盘输入
+# ============================================================
+from pynput import keyboard
+
+class KeyboardReader:
+    def __init__(self):
+        self._event_queue = queue.Queue()
+        self.running = True
+        self.shared_keys_held = set()
+        self.shared_one_shot = set()
+        self._reader_thread = None
+        self._listener = None
+
+    def _normalize_key(self, key):
+        try:
+            if hasattr(key, 'char') and key.char is not None:
+                return key.char.lower()
+        except:
+            pass
+        key_str = str(key)
+        if key_str == 'Key.esc':
+            return 'escape'
+        elif key_str == 'Key.space':
+            return 'space'
+        elif key_str.startswith('Key.'):
+            return key_str.lower()
+        return key_str.lower()
+
+    def _reader_worker(self):
+        while self.running:
+            try:
+                event_type, key = self._event_queue.get(timeout=0.05)
+                k = self._normalize_key(key)
+                if event_type == 'press':
+                    self.shared_keys_held.add(k)
+                    self.shared_one_shot.discard(k)
+                elif event_type == 'release':
+                    self.shared_keys_held.discard(k)
+                    self.shared_one_shot.discard(k)
+            except queue.Empty:
+                pass
+
+    def init(self):
+        def on_press(key):
+            self._event_queue.put(('press', key))
+        def on_release(key):
+            self._event_queue.put(('release', key))
+        try:
+            self._listener = keyboard.Listener(on_press=on_press, on_release=on_release)
+            self._listener.start()
+            self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True)
+            self._reader_thread.start()
+            print("[INFO] 键盘监听已启动")
+        except Exception as e:
+            print(f"[WARN] 无法初始化键盘监听: {e}")
+
+    def is_key_pressed(self, key):
+        k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
+        if k not in self.shared_keys_held or k in self.shared_one_shot:
+            return False
+        self.shared_one_shot.add(k)
+        return True
+
+    def is_key_held(self, key):
+        k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
+        return k in self.shared_keys_held
+
+    def restore(self):
+        self.running = False
+        if self._listener:
+            self._listener.stop()
+
+
+# ============================================================
+# Sensor 读取 
+# ============================================================
+def get_sensor(model, data, name):
+    sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
+    if sid < 0:
+        return None
+    adr = model.sensor_adr[sid]
+    dim = model.sensor_dim[sid]
+    return data.sensordata[adr:adr + dim].copy()
+
+
+def compute_observations_motrix(model, data, commands, last_actions):
+    """计算 48 维观测"""
+    obs = np.zeros(NUM_OBS, dtype=np.float32)
+
+    # 局部线速度 (来自 sensor)
+    local_vel = get_sensor(model, data, "local_linvel")
+    if local_vel is not None:
+        obs[0:3] = local_vel * OBS_SCALES['lin_vel']
+    else:
+        base_rot = data.xmat[1].reshape(3, 3)
+        world_linvel = data.qvel[0:3]
+        local_linvel = base_rot.T @ world_linvel
+        obs[0:3] = local_linvel * OBS_SCALES['lin_vel']
+
+    # 陀螺仪
+    gyro = get_sensor(model, data, "gyro")
+    if gyro is not None:
+        obs[3:6] = gyro * OBS_SCALES['ang_vel']
+    else:
+        obs[3:6] = data.qvel[3:6] * OBS_SCALES['ang_vel']
+
+    # 重力向量 (躯干坐标系
+    base_rot = data.xmat[1].reshape(3, 3)
+    gravity_world = np.array([0., 0., -1.], dtype=np.float64)
+    local_gravity = base_rot.T @ gravity_world
+    obs[6:9] = local_gravity.astype(np.float32)
+
+    # 关节位置偏差
+    joint_pos = data.qpos[7:19]
+    dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES['dof_pos']
+    obs[9:21] = dof_pos_rel
+
+    # 关节速度
+    joint_vel = data.qvel[6:18]
+    obs[21:33] = joint_vel * OBS_SCALES['dof_vel']
+
+    # 上一步动作
+    obs[33:45] = last_actions
+
+    # 命令
+    obs[45:48] = commands * np.array([OBS_SCALES['lin_vel'], OBS_SCALES['lin_vel'], OBS_SCALES['ang_vel']], dtype=np.float32)
+
+    # 限幅
+    obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
+    return obs
+
+
+def main():
+    import re
+    import onnxruntime as ort
+
+    parser = argparse.ArgumentParser(description="MotrixLab Go1 Policy Inference in MuJoCo")
+    parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
+    parser.add_argument("--terrain", type=str, default=TERRAIN, choices=["none", "rough", "stairs"])
+    args = parser.parse_args()
+
+    os.chdir(MOTRIX_XML_DIR)
+
+    if args.terrain == "rough":
+        xml_file = f"{MOTRIX_XML_DIR}/scene_rough_terrain.xml"
+    elif args.terrain == "stairs":
+        xml_file = f"{MOTRIX_XML_DIR}/scene_stairs_terrain.xml"
+    else:
+        xml_file = f"{MOTRIX_XML_DIR}/scene_motor_actuator.xml"  # flat floor
+
+    with open(xml_file, 'r') as f:
+        xml_content = f.read()
+
+    model = mujoco.MjModel.from_xml_string(xml_content)
+    data = mujoco.MjData(model)
+
+    print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
+    print(f"[INFO] MuJoCo timestep: {model.opt.timestep}")
+    print(f"[INFO] 关节顺序 (qpos[7:19]): {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
+
+    for sensor_name in ["gyro", "local_linvel"]:
+        sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name)
+        print(f"[SENSOR] {sensor_name}: {'存在' if sid >= 0 else '不存在'}")
+
+    # 初始化
+    data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
+    data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])  # 默认四元数
+    data.qpos[7:19] = DEFAULT_JOINT_ANGLES
+    data.qvel[:] = 0.0
+    data.ctrl[:] = 0.0
+    mujoco.mj_forward(model, data)
+
+    print(f"[INIT] qpos[2]={data.qpos[2]:.3f}")
+    print(f"[INIT] qpos[7:19]={data.qpos[7:19]}")
+    print(f"[INIT] DEFAULT_JOINT_ANGLES={DEFAULT_JOINT_ANGLES}")
+
+    # 加载onnx
+    session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
+    print(f"[INFO] loaded")
+
+    # 主循环
+    ctrl_dt = 0.01  # 100Hz
+    num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
+    print(f"[INFO] 每 {num_steps_per_inference} 步推理一次 ")
+
+    step_count = 0
+    inference_step = 0
+
+    x_vel_cmd = 0.0
+    y_vel_cmd = 0.0
+    yaw_vel_cmd = 0.0
+    commands = np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32)
+    last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
+    action = np.zeros(NUM_ACTIONS, dtype=np.float32)
+
+    # 键盘
+    keyboard_reader = KeyboardReader()
+    keyboard_reader.init()
+
+    view = viewer.launch_passive(model, data)
+    print("[INFO] 已启动!")
+
+    loop_start_time = time.time()
+
+    while view.is_running() and not g_exit_requested:
+        # 键盘命令
+        if keyboard_reader.is_key_pressed(' '):
+            x_vel_cmd = 0.0
+            y_vel_cmd = 0.0
+            yaw_vel_cmd = 0.0
+
+        if keyboard_reader.is_key_held('w'):
+            x_vel_cmd = MAX_LIN_VEL_X
+        elif keyboard_reader.is_key_held('s'):
+            x_vel_cmd = -MAX_LIN_VEL_X
+        else:
+            x_vel_cmd = 0.0
+
+        if keyboard_reader.is_key_held('q'):
+            y_vel_cmd = MAX_LIN_VEL_Y
+        elif keyboard_reader.is_key_held('e'):
+            y_vel_cmd = -MAX_LIN_VEL_Y
+        else:
+            y_vel_cmd = 0.0
+
+        if keyboard_reader.is_key_held('a'):
+            yaw_vel_cmd = MAX_ANG_VEL
+        elif keyboard_reader.is_key_held('d'):
+            yaw_vel_cmd = -MAX_ANG_VEL
+        else:
+            yaw_vel_cmd = 0.0
+
+        if keyboard_reader.is_key_pressed('r'):
+            # 重置机器人到初始位置
+            data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
+            data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
+            data.qpos[7:19] = DEFAULT_JOINT_ANGLES
+            data.qvel[:] = 0.0
+            data.ctrl[:] = 0.0
+            last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
+            mujoco.mj_forward(model, data)
+            print("[RESET] 机器人已重置")
+
+        if keyboard_reader.is_key_pressed('escape'):
+            break
+
+        # 推理 (每 N 步一次)
+        if inference_step == 0:
+            commands[0] = x_vel_cmd
+            commands[1] = y_vel_cmd
+            commands[2] = yaw_vel_cmd
+
+            obs = compute_observations_motrix(model, data, commands, last_actions)
+
+            # 推理
+            action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
+            action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
+            last_actions = action.copy()
+
+        # PD 控制
+        # joint_targets = action * action_scale + default_angles
+        joint_targets = DEFAULT_JOINT_ANGLES + action * ACTION_SCALE
+
+        current_pos = data.qpos[7:19]
+        current_vel = data.qvel[6:18]
+        torques = KP * (joint_targets - current_pos) - KD * current_vel
+        torques = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
+        data.ctrl[:] = torques
+
+        mujoco.mj_step(model, data)
+        view.sync()
+
+        expected_time = step_count * ctrl_dt
+        elapsed = time.time() - loop_start_time
+        sleep_time = expected_time - elapsed
+        if sleep_time > 0:
+            time.sleep(sleep_time)
+
+        step_count += 1
+        inference_step = (inference_step + 1) % num_steps_per_inference
+
+        if step_count % 200 == 0:
+            trunk_z = data.qpos[2]
+            lin_vel = np.linalg.norm(data.qvel[0:3])
+            print(f"\n========== Step {step_count} ==========")
+            print(f"[CMD] x={x_vel_cmd:.2f}, y={y_vel_cmd:.2f}, yaw={yaw_vel_cmd:.2f}")
+            print(f"[OBS] lin_vel={obs[0:3]}, ang_vel={obs[3:6]}, grav={obs[6:9]}")
+            print(f"[ACTION] raw={action[:4]}... scaled={action[:4]*ACTION_SCALE}...")
+            print(f"[TARGET] {joint_targets[:4]}...")
+            print(f"[TORQUE] {torques[:4]}...")
+            print(f"[STATE] z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
+            print(f"==========================================\n")
+
+    keyboard_reader.restore()
+    view.close()
+
+
+if __name__ == "__main__":
+    main()