wtw，部署
This commit is contained in:
cyy_mac
243
deploy_wtw/sim2sim_wtw_test.py
Normal file
243
deploy_wtw/sim2sim_wtw_test.py
Normal file
@@ -0,0 +1,243 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user