#!/usr/bin/env python3 """ sim2sim test for WTW deploy_wtw_pro_sdk.py using MuJoCo Go1 XML. Parameters verified against go1_walk_these_ways_inference.py reference. Usage: conda activate free_dog_sdk mjpython sim2sim_wtw_test.py Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 1-4=步态 Esc=退出 """ import os, signal, time, queue, threading import mujoco, numpy as np, torch from mujoco import viewer HERE = os.path.dirname(os.path.abspath(__file__)) # ─── Paths ─── XML_PATH = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml") MESH_DIR = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "meshes") BODY_JIT = os.path.join(HERE, "body_latest.jit") ADAPT_JIT = os.path.join(HERE, "adaptation_module_latest.jit") # ─── WTW constants (verified against reference) ─── NUM_OBS = 70 NUM_ACTIONS = 12 NUM_COMMANDS = 15 NUM_OBS_HISTORY = 30 OBS_BUFFER_SIZE = 2100 ACTION_SCALE = 0.25 HIP_SCALE_REDUCTION = 0.5 CLIP_ACTIONS = 10.0 CLIP_OBS = 100.0 # Joint orders: # MuJoCo/SDK: [FR_hip,FR_thigh,FR_calf, FL_hip,FL_thigh,FL_calf, RR_hip,RR_thigh,RR_calf, RL_hip,RL_thigh,RL_calf] # WTW/Deploy: [FL_hip,FL_thigh,FL_calf, FR_hip,FR_thigh,FR_calf, RL_hip,RL_thigh,RL_calf, RR_hip,RR_thigh,RR_calf] DEPLOY_TO_MUJOCO = np.array([3,4,5, 0,1,2, 9,10,11, 6,7,8], dtype=np.int64) # Default angles in WTW order (from reference) DEFAULT_WTW = np.array([ 0.1, 0.8, -1.5, # FL -0.1, 0.8, -1.5, # FR 0.1, 1.0, -1.5, # RL -0.1, 1.0, -1.5, # RR ], dtype=np.float32) DEFAULT_MUJOCO = np.array([ -0.1, 0.8, -1.5, # FR 0.1, 0.8, -1.5, # FL -0.1, 1.0, -1.5, # RR 0.1, 1.0, -1.5, # RL ], dtype=np.float32) # Commands scale (verified) COMMANDS_SCALE = np.array([ 2.0, 2.0, 0.25, # vx, vy, wz 2.0, # body_height 1, 1, 1, 1, 1, # freq, phase, offset, bound, duration 0.15, # footswing_height 0.3, 0.3, # body_pitch, body_roll 1.0, 1.0, # stance_width, stance_length 1.0, # aux_reward ], dtype=np.float32)[:15] OBS_SCALES = {"dof_pos":1.0, "dof_vel":0.05} # PD gains (from reference) KP = 20.0; KD = 0.1 # matches reference (XML passive damping=1.0, total≈1.1) # Friction (training had zero floor friction) FLOOR_FRICTION = [0.0, 0.0, 0.0] BODY_FRICTION = [0.6, 0.3, 0.3] # Gait presets GAITS = { '1': ('Trot', 0.5, 0.0, 0.0), '2': ('Pace', 0.0, 0.0, 0.5), '3': ('Bound', 0.0, 0.5, 0.0), '4': ('Pronk', 0.0, 0.0, 0.0), } EXIT = False def _sig(s, f): global EXIT; EXIT = True signal.signal(signal.SIGINT, _sig) class Keyboard: def __init__(self): self.held = set(); self._l = None def _p(self, k): try: self.held.add(k.char.lower()) except: self.held.add(str(k)) def _r(self, k): try: self.held.discard(k.char.lower()) except: self.held.discard(str(k)) def init(self): from pynput import keyboard self._l = keyboard.Listener(on_press=self._p, on_release=self._r); self._l.start() def keys(self): return self.held.copy() def stop(self): if self._l: self._l.stop() def main(): if not os.path.exists(BODY_JIT): print(f"[ERROR] body not found: {BODY_JIT}"); return if not os.path.exists(ADAPT_JIT): print(f"[ERROR] adapt not found: {ADAPT_JIT}"); return # Load MuJoCo with mesh path fix with open(XML_PATH) as f: xml = f.read() xml = xml.replace('meshdir="../meshes/"', f'meshdir="{MESH_DIR}"') model = mujoco.MjModel.from_xml_string(xml) data = mujoco.MjData(model) # Set friction for i in range(model.ngeom): model.geom_friction[i] = BODY_FRICTION print(f"[INFO] MuJoCo: {model.nbody} bodies, {model.nq} DoF, KP={KP}, KD(active)={KD}") # Load models body = torch.jit.load(BODY_JIT, map_location='cpu').eval() adapt = torch.jit.load(ADAPT_JIT, map_location='cpu').eval() print(f"[INFO] WTW body+adapt loaded") # Init data.qpos[0:3] = [0, 0, 0.35] data.qpos[3:7] = [1, 0, 0, 0] data.qpos[7:19] = DEFAULT_MUJOCO mujoco.mj_forward(model, data) obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32) prev_action = np.zeros(12, dtype=np.float32) last_action = np.zeros(12, dtype=np.float32) gait_idx = 0.0 step, vx, vy, wz = 0, 0.0, 0.0, 0.0 fh = 0.15 # footswing height bp, br = 0.0, 0.0 # body pitch/roll gait_phase, gait_offset, gait_bound, gait_dur = 0.5, 0.0, 0.0, 0.5 current_gait = '1' ctrl_dt = 0.02 kb = Keyboard(); kb.init() view = viewer.launch_passive(model, data) print("[INFO] W/S=前后 Q/E=左右 A/D=旋转 1-4=步态 R=重置 Esc=退出") t0 = time.perf_counter() while view.is_running() and not EXIT: keys = kb.keys() if 'key.esc' in keys: break if 'r' in keys: data.qpos[0:3]=[0,0,0.35]; data.qpos[3:7]=[1,0,0,0] data.qpos[7:19]=DEFAULT_MUJOCO; data.qvel[:]=0 obs_buffer[:]=0; prev_action[:]=0; last_action[:]=0; gait_idx=0 mujoco.mj_forward(model, data) # Gait switch for k, (name, ph, off, bd) in GAITS.items(): if k in keys and k != current_gait: current_gait = k gait_phase, gait_offset, gait_bound = ph, off, bd print(f"[INFO] Gait: {name}") 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=3.0 if 'a' in keys else (-3.0 if 'd' in keys else 0.0) if ' ' in keys: vx=vy=wz=0.0 # Inference at 50Hz (every 10 sim steps at dt=0.002) if step % 10 == 0: # Commands raw_cmd = np.zeros(15, dtype=np.float32) raw_cmd[0]=vx; raw_cmd[1]=vy; raw_cmd[2]=wz raw_cmd[3]=0.0; raw_cmd[4]=3.0 raw_cmd[5]=gait_phase; raw_cmd[6]=gait_offset; raw_cmd[7]=gait_bound; raw_cmd[8]=gait_dur raw_cmd[9]=0.15; raw_cmd[10]=bp; raw_cmd[11]=br raw_cmd[12]=0.25; raw_cmd[13]=0.4 commands = raw_cmd * COMMANDS_SCALE # Gait index & clock gait_idx += 0.02 * 3.0 if gait_idx > 1.0: gait_idx -= 1.0 p, o, b = gait_phase, gait_offset, gait_bound fi = [gait_idx+p+o+b, gait_idx+o, gait_idx+b, gait_idx+p] clock = np.array([np.sin(2*np.pi*f) for f in fi], dtype=np.float32) # Observation obs = np.zeros(NUM_OBS, dtype=np.float32) base_rot = data.xmat[1].reshape(3, 3) obs[0:3] = (base_rot.T @ np.array([0., 0., -1.], dtype=np.float64)).astype(np.float32) obs[3:18] = commands dof_wtw = data.qpos[7:19][DEPLOY_TO_MUJOCO] obs[18:30] = (dof_wtw - DEFAULT_WTW) * 1.0 obs[30:42] = data.qvel[6:18][DEPLOY_TO_MUJOCO] * 0.05 obs[42:54] = np.clip(prev_action, -CLIP_ACTIONS, CLIP_ACTIONS) obs[54:66] = np.clip(last_action, -CLIP_ACTIONS, CLIP_ACTIONS) obs[66:70] = clock obs = np.clip(obs, -CLIP_OBS, CLIP_OBS) # History buffer obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs]) # Inference obs_hist = torch.from_numpy(obs_buffer).float().unsqueeze(0) with torch.inference_mode(): latent = adapt(obs_hist) action = body(torch.cat([obs_hist, latent], dim=1)).numpy().flatten() action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) last_action = prev_action.copy() prev_action = action.copy() # PD control action_scaled = prev_action * ACTION_SCALE for i in [0,3,6,9]: action_scaled[i] *= HIP_SCALE_REDUCTION targets_mujoco = action_scaled[DEPLOY_TO_MUJOCO] + DEFAULT_MUJOCO torques = KP*(targets_mujoco - data.qpos[7:19]) - KD*data.qvel[6:18] data.ctrl[:] = np.clip(torques, -23.7, 23.7) mujoco.mj_step(model, data) view.sync() if step % 200 == 0: print(f"[STEP {step}] z={data.qpos[2]:.3f} cmd=[{vx:.1f},{vy:.1f},{wz:.1f}] " f"pos=[{data.qpos[0]:.2f},{data.qpos[1]:.2f}]") step += 1 expected = (step+1)*model.opt.timestep sleep = expected - (time.perf_counter()-t0) if sleep > 0: time.sleep(sleep) kb.stop(); view.close() print("[INFO] Done.") if __name__ == "__main__": main()