244 lines
8.5 KiB
Python
244 lines
8.5 KiB
Python
#!/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()
|