no_linevel_no_foot

deploy_45dim/go1_no_linevel_sim2sim_mujoco.py

@@ -0,0 +1,403 @@
+#!/usr/bin/env python3
+"""MuJoCo sim2sim for go1-stairs-terrain-walk-no-linevel (57-dim obs, no linvel).
+
+Loads the ONNX policy exported by export_go1_no_linevel_onnx.py and runs
+inference in MuJoCo with PD control.
+
+Usage:
+    # Default: combined flat+rough+stairs terrain, random spawn
+    uv run scripts/go1_no_linevel_sim2sim_mujoco.py
+
+    # Specific terrain
+    uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain flat
+    uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain rough
+    uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain stairs
+
+    # Custom ONNX path
+    uv run scripts/go1_no_linevel_sim2sim_mujoco.py --onnx ./exports_go1_no_linevel/policy.onnx
+
+Keyboard controls:
+    W/S     - forward/backward
+    A/D     - turn left/right
+    Q/E     - strafe left/right
+    Space   - stop
+    R       - reset robot
+    1/2/3   - switch terrain (flat/rough/stairs)
+    Esc     - quit
+"""
+
+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)
+
+# ============================================================
+# Paths
+# ============================================================
+HERE = os.path.dirname(os.path.abspath(__file__))
+DEFAULT_ONNX_PATH = os.path.join(HERE, "better3.onnx")
+GO1_XML = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml")
+
+# ============================================================
+# MotrixLab parameters (matching cfg.py + walk_stairs_terrain_no_linevel.py)
+# ============================================================
+NUM_OBS = 57  # 57-dim: NO linear velocity, WITH contact forces
+NUM_ACTIONS = 12
+OBS_SCALES = {"ang_vel": 0.25, "dof_pos": 1.0, "dof_vel": 0.05}
+ACTION_SCALE = 0.05
+KP, KD = 80.0, 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 = 1.0
+
+# ============================================================
+# Joint names and order
+# ============================================================
+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,
+    0.0, 0.9, -1.8,
+    -0.0, 0.9, -1.8,
+    0.0, 0.9, -1.8,
+], dtype=np.float32)
+
+FEET = ["FR", "FL", "RR", "RL"]
+
+# Terrain spawn positions (world Y)
+TERRAIN_SPAWN = {
+    "flat": np.array([0.0, 54.0, 0.42], dtype=np.float64),
+    "rough": np.array([0.0, 32.0, 0.42], dtype=np.float64),
+    "stairs": np.array([0.0, 0.0, 0.42], dtype=np.float64),
+}
+
+# ============================================================
+# Keyboard input
+# ============================================================
+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 Exception:
+            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] Keyboard listener started")
+        except Exception as e:
+            print(f"[WARN] Cannot init keyboard: {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 reading
+# ============================================================
+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 read_contact_forces(model, data, base_rot):
+    """Read foot contact forces (12-dim, body frame) from MuJoCo contact sensors.
+
+    Tries _stairs, _rough, _flat suffixes for each foot, picking the first
+    sensor that returns non-zero data. In MuJoCo, `data="force"` returns a
+    scalar (normal force). We construct a 3D force vector by projecting onto
+    the body-frame Z axis as an approximation.
+    """
+    forces = np.zeros(12, dtype=np.float32)
+    for i, foot in enumerate(FEET):
+        f_scalar = 0.0
+        for suffix in ["_stairs", "_rough", "_flat"]:
+            name = f"{foot}_foot_contact{suffix}"
+            v = get_sensor(model, data, name)
+            if v is not None and np.abs(v[0]) > 1e-6:
+                f_scalar = v[0]
+                break
+        # Assume contact force is approximately vertical (world Z),
+        # rotate into body frame
+        force_world = np.array([0.0, 0.0, f_scalar], dtype=np.float64)
+        force_body = base_rot.T @ force_world
+        forces[i * 3 : i * 3 + 3] = force_body.astype(np.float32)
+    return forces
+
+
+def compute_observations(model, data, commands, last_actions, base_rot):
+    """Compute 57-dim observation matching go1-stairs-terrain-walk-no-linevel.
+
+    Layout (57 dims, NO linear velocity):
+      [0:3]   gyro (ang_vel * 0.25)
+      [3:6]   gravity vector (body frame)
+      [6:18]  joint angle deviation (dof_pos * 1.0)
+      [18:30] joint velocity (dof_vel * 0.05)
+      [30:42] last actions (raw)
+      [42:45] commands [vx*2.0, vy*2.0, wz*0.25]
+      [45:57] foot contact forces (body frame, raw)
+    """
+    obs = np.zeros(NUM_OBS, dtype=np.float32)
+
+    # Gyro
+    gyro = get_sensor(model, data, "gyro")
+    if gyro is not None:
+        obs[0:3] = gyro * OBS_SCALES["ang_vel"]
+    else:
+        obs[0:3] = data.qvel[3:6] * OBS_SCALES["ang_vel"]
+
+    # Gravity vector (body frame)
+    gravity_world = np.array([0.0, 0.0, -1.0], dtype=np.float64)
+    local_gravity = base_rot.T @ gravity_world
+    obs[3:6] = local_gravity.astype(np.float32)
+
+    # Joint position deviation
+    joint_pos = data.qpos[7:19]
+    dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES["dof_pos"]
+    obs[6:18] = dof_pos_rel
+
+    # Joint velocity
+    joint_vel = data.qvel[6:18]
+    obs[18:30] = joint_vel * OBS_SCALES["dof_vel"]
+
+    # Last actions
+    obs[30:42] = last_actions
+
+    # Commands (scale matching MotrixLab: [2.0, 2.0, 0.25])
+    obs[42] = commands[0] * 2.0
+    obs[43] = commands[1] * 2.0
+    obs[44] = commands[2] * 0.25
+
+    # Contact forces
+    # obs[45:57] = read_contact_forces(model, data, base_rot)  # disabled: test with zeros
+    obs[45:57] = np.zeros(12, dtype=np.float32)
+
+    obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
+    return obs
+
+
+# ============================================================
+# Main
+# ============================================================
+def main():
+    import onnxruntime as ort
+
+    parser = argparse.ArgumentParser(description="MotrixLab Go1 No-Linevel Policy Inference in MuJoCo")
+    parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
+    parser.add_argument(
+        "--terrain", type=str, default="combined",
+        choices=["flat", "rough", "stairs", "combined"],
+        help="Terrain type (combined = flat+rough+stairs in one scene)",
+    )
+    args = parser.parse_args()
+
+    # Use local Go1 XML (terrain switching disabled on macOS)
+    os.chdir(os.path.dirname(GO1_XML))  # mesh paths are relative
+    with open(GO1_XML, "r") as f:
+        xml_content = f.read()
+
+    model = mujoco.MjModel.from_xml_string(xml_content)
+    data = mujoco.MjData(model)
+
+    print(f"[INFO] Terrain: {args.terrain}")
+    print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
+    print(f"[INFO] Timestep: {model.opt.timestep}")
+
+    # Use simple spawn
+    spawn_xyz = np.array([0.0, 0.0, 0.42], dtype=np.float64)
+
+    data.qpos[0:3] = spawn_xyz
+    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)
+
+    # Load ONNX
+    session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
+    print(f"[INFO] ONNX loaded: {args.onnx}")
+
+    # Main loop
+    ctrl_dt = 0.01
+    num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
+    print(f"[INFO] Inference every {num_steps_per_inference} sim steps")
+
+    step_count = 0
+    inference_step = 0
+    commands = np.zeros(3, 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()
+
+    viewer_handle = viewer.launch_passive(model, data)
+    print("[INFO] Viewer launched!")
+    print("[KEYS] WASD=move, QE=strafe, Space=stop, R=reset, 1/2/3=terrain, Esc=quit")
+
+    loop_start_time = time.time()
+
+    while viewer_handle.is_running() and not g_exit_requested:
+        # --- Keyboard input ---
+        x_vel, y_vel, yaw_vel = 0.0, 0.0, 0.0
+
+        if keyboard_reader.is_key_held("w"):
+            x_vel = MAX_LIN_VEL_X
+        elif keyboard_reader.is_key_held("s"):
+            x_vel = -MAX_LIN_VEL_X
+
+        if keyboard_reader.is_key_held("q"):
+            y_vel = MAX_LIN_VEL_Y
+        elif keyboard_reader.is_key_held("e"):
+            y_vel = -MAX_LIN_VEL_Y
+
+        if keyboard_reader.is_key_held("a"):
+            yaw_vel = MAX_ANG_VEL
+        elif keyboard_reader.is_key_held("d"):
+            yaw_vel = -MAX_ANG_VEL
+
+        if keyboard_reader.is_key_pressed("space"):
+            x_vel = y_vel = yaw_vel = 0.0
+
+        # Reset
+        if keyboard_reader.is_key_pressed("r"):
+            spawn_xyz[:] = [0.0, 0.0, 0.42]
+
+            data.qpos[0:3] = spawn_xyz
+            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(f"[RESET] pos={spawn_xyz}")
+
+        if keyboard_reader.is_key_pressed("escape"):
+            break
+
+        # --- Inference ---
+        if inference_step == 0:
+            commands[0] = x_vel
+            commands[1] = y_vel
+            commands[2] = yaw_vel
+
+            base_rot = data.xmat[1].reshape(3, 3)
+            obs = compute_observations(model, data, commands, last_actions, base_rot)
+
+            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 control ---
+        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)
+        viewer_handle.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 % 500 == 0:
+            trunk_z = data.qpos[2]
+            print(f"[{step_count}] cmd=({x_vel:.1f},{y_vel:.1f},{yaw_vel:.1f}) "
+                  f"z={trunk_z:.3f}m")
+
+    keyboard_reader.restore()
+    viewer_handle.close()
+
+
+if __name__ == "__main__":
+    main()