diff --git a/deploy_45dim/better3.onnx b/deploy_45dim/better3.onnx new file mode 100644 index 0000000..fdd3380 Binary files /dev/null and b/deploy_45dim/better3.onnx differ diff --git a/deploy_45dim/go1_no_linevel_sim2sim_mujoco.py b/deploy_45dim/go1_no_linevel_sim2sim_mujoco.py new file mode 100644 index 0000000..194b629 --- /dev/null +++ b/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()