#!/usr/bin/env python3 """Go1 sim2sim MuJoCo viewer — Original 30k flat training. Usage: python go1_sim2sim.py Requires: mujoco, onnxruntime, pynput Install: pip install mujoco onnxruntime pynput Controls: W/S: forward/back Q/E: strafe left/right A/D: rotate Space: stop R: reset Esc: quit """ import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue from mujoco import viewer from pynput import keyboard HERE = os.path.dirname(os.path.abspath(__file__)) ONNX = os.path.join(HERE, "90k_45.onnx") # ── Parameters (original MotrixLab Go1 config) ── NUM_OBS = 45 KP, KD = 80.0, 0.5 # KD=0.5 + MuJoCo joint_damping(0.5) = 1.0 = training kd ACTION_SCALE = 0.05 CLIP = 23.7 DEFAULT_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) # ── Keyboard ── class KB: def __init__(s): s._q = queue.Queue(); s.running = True; s.held = set() def _n(s, k): try: if hasattr(k, 'char') and k.char: return k.char.lower() except: pass return str(k).lower() def _w(s): while s.running: try: et, k = s._q.get(timeout=0.05) n = s._n(k) if et == 'press': s.held.add(n) elif et == 'release': s.held.discard(n) except queue.Empty: pass def init(s): s._l = keyboard.Listener( on_press=lambda k: s._q.put(('press', k)), on_release=lambda k: s._q.put(('release', k))) s._l.start() s._t = threading.Thread(target=s._w, daemon=True); s._t.start() def keys(s): return s.held.copy() def stop(s): s.running = False; s._l.stop() # ── Main ── def main(): # The model XML is embedded below xml = ''' ''' model = mujoco.MjModel.from_xml_string(xml) data = mujoco.MjData(model) data.qpos[0:3] = [0, 0, 0.42] data.qpos[3:7] = [1, 0, 0, 0] data.qpos[7:19] = DEFAULT_ANGLES mujoco.mj_forward(model, data) session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider']) print(f"[Go1 sim2sim] ONNX={ONNX}") print(f"[Go1 sim2sim] PD kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} obs={NUM_OBS}-dim") print(f"[Go1 sim2sim] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出") kb = KB(); kb.init() view = viewer.launch_passive(model, data) step, vx, vy, wz = 0, 0.0, 0.0, 0.0 last_a = np.zeros(12, dtype=np.float32) while view.is_running(): keys = kb.keys() if 'escape' in keys: break if 'r' in keys: data.qpos[0:3] = [0, 0, 0.42] data.qpos[3:7] = [1, 0, 0, 0] data.qpos[7:19] = DEFAULT_ANGLES data.qvel[:] = 0; last_a[:] = 0 mujoco.mj_forward(model, data) vx = 1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.0) vy = 1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.0) wz = 1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.0) if ' ' in keys: vx = vy = wz = 0.0 if step % 2 == 0: # 100Hz control (MuJoCo dt=0.005) obs = np.zeros(NUM_OBS, dtype=np.float32) sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro") adr = model.sensor_adr[sid] obs[0:3] = data.sensordata[adr:adr+3] * 0.25 R = data.xmat[1].reshape(3, 3) obs[3:6] = (R.T @ np.array([0., 0., -1.])).astype(np.float32) obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0 obs[18:30] = data.qvel[6:18] * 0.05 obs[30:42] = last_a obs[42:45] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25]) obs = np.clip(obs, -100., 100.) action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0] action = np.clip(action, -CLIP, CLIP) last_a = action.copy() target = DEFAULT_ANGLES + action * ACTION_SCALE torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18] data.ctrl[:] = np.clip(torques, -CLIP, CLIP) mujoco.mj_step(model, data) view.sync() step += 1 time.sleep(0.001) kb.stop(); view.close() if __name__ == "__main__": main()