cyy/go1_pro_deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wtw，部署

Browse Source
This commit is contained in:
cyy_mac
2026-07-24 13:19:29 +08:00
parent 04eaa8c916
commit 554d4649c5
22 changed files with 1960 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
deploy_wtw/__pycache__/deploy_wtw_pro_sdk.cpython-310.pyc Normal file
View File

Binary file not shown.

BIN
deploy_wtw/__pycache__/sim2sim_wtw_test.cpython-310.pyc Normal file
View File

Binary file not shown.

BIN
deploy_wtw/adaptation_module_latest.jit Executable file
View File

Binary file not shown.

BIN
deploy_wtw/body_latest.jit Executable file
View File

Binary file not shown.

541
deploy_wtw/deploy_isaaclab_onnx_no_torch_wtw_r2.py Normal file
View File

@@ -0,0 +1,541 @@
#!/usr/bin/env python3
"""
deploy_isaaclab_onnx_no_torch_wtw_r2.py
No-torch ONNX deployment script for Unitree Go1 using WTW LCM bridge.
R2 logic matches original Walk-These-Ways style:
startup: press R2 -> move to default pose -> press R2 -> start policy
runtime: press R2 -> return to default pose and pause -> press R2 -> resume policy
It is NOT hold-R2-to-run. R2 is treated as a rising-edge event, using
StateEstimator.right_lower_right_switch_pressed by default.
Logs model input/output and targets as JSONL when --log-dir is provided.
"""
import argparse
import json
import signal
import time
from datetime import datetime
from pathlib import Path
import lcm
import numpy as np
import onnxruntime as ort
from go1_gym_deploy.lcm_types.pd_tau_targets_lcmt import pd_tau_targets_lcmt
# ---------------- IsaacLab policy constants ----------------
NUM_OBS = 48
NUM_ACTIONS = 12
OBS_SCALES = {"lin_vel": 2.0, "ang_vel": 0.5, "dof_pos": 1.0, "dof_vel": 0.05}
ACTION_SCALE = 0.25
CLIP_OBSERVATIONS = 100.0
# IsaacLab order:
# [FL_hip, FR_hip, RL_hip, RR_hip,
# FL_thigh, FR_thigh, RL_thigh, RR_thigh,
# FL_calf, FR_calf, RL_calf, RR_calf]
DEFAULT_JOINT_ANGLES_ISAAC = np.array([
0.1, -0.1, 0.1, -0.1,
0.8, 0.8, 1.0, 1.0,
-1.5, -1.5, -1.5, -1.5,
], dtype=np.float32)
# WTW internal order:
# [FL_hip, FL_thigh, FL_calf,
# FR_hip, FR_thigh, FR_calf,
# RL_hip, RL_thigh, RL_calf,
# RR_hip, RR_thigh, RR_calf]
ISAAC_TO_WTW = np.array([0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11], dtype=np.int64)
WTW_TO_ISAAC = np.array([0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11], dtype=np.int64)
# WTW internal -> Unitree/lcm_position order. Same idea as WTW joint_idxs.
WTW_TO_UNITREE = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int64)
NAMES_ISAAC = [
"FL_hip", "FR_hip", "RL_hip", "RR_hip",
"FL_thigh", "FR_thigh", "RL_thigh", "RR_thigh",
"FL_calf", "FR_calf", "RL_calf", "RR_calf",
]
NAMES_WTW = [
"FL_hip", "FL_thigh", "FL_calf",
"FR_hip", "FR_thigh", "FR_calf",
"RL_hip", "RL_thigh", "RL_calf",
"RR_hip", "RR_thigh", "RR_calf",
]
NAMES_UNITREE = [
"FR_hip", "FR_thigh", "FR_calf",
"FL_hip", "FL_thigh", "FL_calf",
"RR_hip", "RR_thigh", "RR_calf",
"RL_hip", "RL_thigh", "RL_calf",
]
EXIT = False
def _sig_handler(signum, frame):
global EXIT
EXIT = True
signal.signal(signal.SIGINT, _sig_handler)
signal.signal(signal.SIGTERM, _sig_handler)
class FakeStateEstimator:
def __init__(self, auto_r2=False):
self.default_wtw = DEFAULT_JOINT_ANGLES_ISAAC[ISAAC_TO_WTW].copy()
self.right_lower_right_switch_pressed = bool(auto_r2)
self.right_lower_right_switch = int(auto_r2)
self.left_stick = [0.0, 0.0]
self.right_stick = [0.0, 0.0]
def get_body_linear_vel(self):
return np.zeros(3, dtype=np.float32)
def get_body_angular_vel(self):
return np.zeros(3, dtype=np.float32)
def get_gravity_vector(self):
return np.array([0.0, 0.0, -1.0], dtype=np.float32)
def get_dof_pos(self):
return self.default_wtw.copy()
def get_dof_vel(self):
return np.zeros(12, dtype=np.float32)
def get_command(self):
return np.zeros(19, dtype=np.float32)
class OnnxPolicy:
def __init__(self, onnx_path):
self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
self.input_name = self.session.get_inputs()[0].name
print("[INFO] ONNX loaded:", onnx_path)
print("[INFO] Inputs :", [(i.name, i.shape, i.type) for i in self.session.get_inputs()])
print("[INFO] Outputs:", [(o.name, o.shape, o.type) for o in self.session.get_outputs()])
def __call__(self, obs):
out = self.session.run(None, {self.input_name: obs.reshape(1, -1).astype(np.float32)})[0][0]
return np.asarray(out, dtype=np.float32)
class JsonlLogger:
def __init__(self, log_dir, args):
self.enabled = bool(log_dir)
self.fp = None
self.log_every = max(1, int(args.log_every))
self.flush_every = max(1, int(args.log_flush_every))
self.run_dir = None
if not self.enabled:
return
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
self.run_dir = Path(log_dir).expanduser().resolve() / f"run_{ts}"
self.run_dir.mkdir(parents=True, exist_ok=True)
meta = {
"created_at": ts,
"args": vars(args),
"num_obs": NUM_OBS,
"num_actions": NUM_ACTIONS,
"obs_scales": OBS_SCALES,
"action_scale": ACTION_SCALE,
"default_joint_angles_isaac": DEFAULT_JOINT_ANGLES_ISAAC.tolist(),
"isaac_to_wtw": ISAAC_TO_WTW.tolist(),
"wtw_to_isaac": WTW_TO_ISAAC.tolist(),
"wtw_to_unitree": WTW_TO_UNITREE.tolist(),
"names_isaac": NAMES_ISAAC,
"names_wtw": NAMES_WTW,
"names_unitree": NAMES_UNITREE,
}
(self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8")
print("[INFO] Runtime log dir:", self.run_dir)
def arr(self, x):
return np.asarray(x, dtype=np.float32).reshape(-1).tolist()
def log(self, step, **kw):
if not self.enabled or step % self.log_every != 0:
return
rec = {"step": int(step), "time_wall": time.time()}
for k, v in kw.items():
if isinstance(v, np.ndarray):
rec[k] = self.arr(v)
elif isinstance(v, (np.float32, np.float64)):
rec[k] = float(v)
elif isinstance(v, (np.int32, np.int64)):
rec[k] = int(v)
else:
rec[k] = v
self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
if step % self.flush_every == 0:
self.fp.flush()
def close(self):
if self.fp:
self.fp.flush()
self.fp.close()
print("[INFO] Runtime log saved:", self.run_dir)
def safe_array(x, n, name):
arr = np.asarray(x, dtype=np.float32).reshape(-1)
if arr.shape[0] != n:
raise ValueError(f"{name} should have length {n}, got {arr.shape[0]}")
return np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
def get_commands(args, se):
if args.use_rc:
if not args.real_state:
return np.zeros(3, dtype=np.float32)
cmd = np.asarray(se.get_command(), dtype=np.float32).reshape(-1)
cmd = np.nan_to_num(cmd, nan=0.0, posinf=0.0, neginf=0.0)
if cmd.shape[0] < 3:
raise RuntimeError(f"RC command should have at least 3 values, got {cmd.shape}")
return np.array([cmd[0] * args.rc_x_scale, cmd[1] * args.rc_y_scale, cmd[2] * args.rc_yaw_scale], dtype=np.float32)
return np.array([args.cmd_x, args.cmd_y, args.cmd_yaw], dtype=np.float32)
def compute_obs(se, commands, last_actions):
obs = np.zeros(NUM_OBS, dtype=np.float32)
base_lin_vel = safe_array(se.get_body_linear_vel(), 3, "body_linear_vel")
base_ang_vel = safe_array(se.get_body_angular_vel(), 3, "body_angular_vel")
gravity = safe_array(se.get_gravity_vector(), 3, "gravity_vector")
dof_pos_wtw = safe_array(se.get_dof_pos(), 12, "dof_pos")
dof_vel_wtw = safe_array(se.get_dof_vel(), 12, "dof_vel")
dof_pos_isaac = dof_pos_wtw[WTW_TO_ISAAC]
dof_vel_isaac = dof_vel_wtw[WTW_TO_ISAAC]
obs[0:3] = base_lin_vel * OBS_SCALES["lin_vel"]
obs[3:6] = base_ang_vel * OBS_SCALES["ang_vel"]
obs[6:9] = gravity
obs[9:12] = commands * np.array([OBS_SCALES["lin_vel"], OBS_SCALES["lin_vel"], OBS_SCALES["ang_vel"]], dtype=np.float32)
obs[12:24] = (dof_pos_isaac - DEFAULT_JOINT_ANGLES_ISAAC) * OBS_SCALES["dof_pos"]
obs[24:36] = dof_vel_isaac * OBS_SCALES["dof_vel"]
obs[36:48] = last_actions
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
info = {
"base_lin_vel": base_lin_vel,
"base_ang_vel": base_ang_vel,
"projected_gravity": gravity,
"dof_pos_wtw": dof_pos_wtw,
"dof_vel_wtw": dof_vel_wtw,
"dof_pos_isaac": dof_pos_isaac,
"dof_vel_isaac": dof_vel_isaac,
}
return obs.astype(np.float32), info
def state_ok(info):
dof_pos = info["dof_pos_wtw"]
grav = info["projected_gravity"]
if not np.all(np.isfinite(dof_pos)):
return False, "dof_pos has NaN/Inf"
if not np.all(np.isfinite(grav)):
return False, "gravity has NaN/Inf"
if np.linalg.norm(dof_pos) < 1e-6:
return False, "dof_pos is all zeros; likely no real leg_control_data received"
gn = float(np.linalg.norm(grav))
if gn < 0.5 or gn > 1.5:
return False, f"gravity norm suspicious: {gn:.3f}"
return True, "ok"
def action_to_targets(action, action_clip):
action_safe = np.clip(action, -action_clip, action_clip).astype(np.float32)
q_isaac = DEFAULT_JOINT_ANGLES_ISAAC + action_safe * ACTION_SCALE
q_wtw = q_isaac[ISAAC_TO_WTW]
q_unitree = q_wtw[WTW_TO_UNITREE]
return action_safe, q_isaac.astype(np.float32), q_wtw.astype(np.float32), q_unitree.astype(np.float32)
def default_targets():
q_isaac = DEFAULT_JOINT_ANGLES_ISAAC.copy()
q_wtw = q_isaac[ISAAC_TO_WTW]
q_unitree = q_wtw[WTW_TO_UNITREE]
return q_isaac, q_wtw, q_unitree
def make_msg(q_unitree, args, msg_id):
msg = pd_tau_targets_lcmt()
msg.q_des = np.asarray(q_unitree, dtype=np.float64).reshape(12).tolist()
msg.qd_des = [0.0] * 12
msg.tau_ff = [0.0] * 12
msg.kp = [float(args.kp)] * 12
msg.kd = [float(args.kd)] * 12
msg.timestamp_us = int(time.time() * 1e6)
msg.id = int(msg_id)
msg.robot_id = 0
msg.se_contactState = [0.0] * 4
return msg
def publish(lc, q_unitree, args, msg_id):
lc.publish("pd_plustau_targets", make_msg(q_unitree, args, msg_id).encode())
def consume_r2_event(args, se):
if args.fake_auto_r2 and not args.real_state:
return True
val = bool(getattr(se, args.r2_pressed_field, False))
if val:
try:
setattr(se, args.r2_pressed_field, False)
except Exception:
pass
return val
def wait_r2(args, se, text):
print(text)
if args.fake_auto_r2 and not args.real_state:
print("[INFO] fake_auto_r2 enabled: continuing immediately.")
return
while not EXIT:
if consume_r2_event(args, se):
print("[INFO] R2 press detected.")
return
time.sleep(0.01)
def calibrate_default(args, lc, se, msg_id):
if not args.publish:
print("[INFO] calibration skipped because publish=False")
return msg_id
print("[INFO] Moving slowly to default pose...")
default_wtw = DEFAULT_JOINT_ANGLES_ISAAC[ISAAC_TO_WTW].copy()
current_wtw = safe_array(se.get_dof_pos(), 12, "dof_pos")
if args.real_state and not args.no_state_safety_check and np.linalg.norm(current_wtw) < 1e-6:
print("[WARN] current dof_pos is all zeros; skip calibration publish")
return msg_id
rel = current_wtw - default_wtw
for _ in range(300):
if EXIT:
break
if np.max(np.abs(rel)) <= 0.01:
break
rel -= np.clip(rel, -args.cal_step_rad, args.cal_step_rad)
q_wtw = default_wtw + rel
publish(lc, q_wtw[WTW_TO_UNITREE], args, msg_id)
msg_id += 1
time.sleep(args.cal_dt)
_, _, q_default_unitree = default_targets()
for _ in range(10):
if EXIT:
break
publish(lc, q_default_unitree, args, msg_id)
msg_id += 1
time.sleep(args.cal_dt)
print("[INFO] Default pose commanded.")
return msg_id
def hold_default_until_r2(args, lc, se, msg_id):
print("[WTW R2] Paused at default pose. Press R2 again to resume policy.")
_, _, q_default_unitree = default_targets()
dt = 1.0 / max(1.0, args.hold_default_rate_hz)
next_t = time.perf_counter()
while not EXIT:
if consume_r2_event(args, se):
print("[WTW R2] Resume press detected.")
return msg_id
if args.publish:
publish(lc, q_default_unitree, args, msg_id)
msg_id += 1
next_t += dt
sleep = next_t - time.perf_counter()
if sleep > 0:
time.sleep(sleep)
else:
next_t = time.perf_counter()
return msg_id
def print_vec(title, names, values):
print(title)
for n, v in zip(names, values):
print(f" {n:10s}: {float(v): .4f}")
def print_debug(step, mode, args, commands, obs, info, action_raw, action_safe, q_isaac, q_wtw, q_unitree, ok, reason, sent, r2):
print("\n" + "=" * 80)
print(f"[STEP {step}] mode={mode}, publish={args.publish}, publish_sent={sent}, real_state={args.real_state}, use_rc={args.use_rc}, wtw_r2_logic={args.wtw_r2_logic}, r2_pressed={r2}")
print(f"[STATE CHECK] ok={ok}, reason={reason}")
print(f"[CMD] x={commands[0]:.3f}, y={commands[1]:.3f}, yaw={commands[2]:.3f}")
print("[OBS] lin:", np.round(info["base_lin_vel"], 4), "ang:", np.round(info["base_ang_vel"], 4), "grav:", np.round(info["projected_gravity"], 4))
print("[OBS] obs[0:12]:", np.round(obs[0:12], 4))
print("[ACTION raw ]", np.round(action_raw, 4))
print("[ACTION safe]", np.round(action_safe, 4))
print("[TARGET offset max rad]:", float(np.max(np.abs(action_safe * ACTION_SCALE))))
print_vec("[TARGET IsaacLab order]", NAMES_ISAAC, q_isaac)
print_vec("[TARGET WTW internal order]", NAMES_WTW, q_wtw)
print_vec("[TARGET Unitree/lcm_position order]", NAMES_UNITREE, q_unitree)
print("=" * 80)
def create_se(args, lc):
if not args.real_state:
print("[INFO] Using FakeStateEstimator. No real robot state will be read.")
return FakeStateEstimator(auto_r2=args.fake_auto_r2)
from go1_gym_deploy.utils.cheetah_state_estimator import StateEstimator
print("[INFO] Using real WTW StateEstimator.")
se = StateEstimator(lc)
se.spin()
return se
def parse_args():
p = argparse.ArgumentParser(description="No-torch ONNX Go1 deployment with original WTW R2 logic")
p.add_argument("--onnx", default="policy.onnx")
p.add_argument("--lcm-url", default="udpm://239.255.76.67:7667?ttl=255")
g = p.add_mutually_exclusive_group()
g.add_argument("--real-state", action="store_true")
g.add_argument("--fake-state", action="store_true")
p.add_argument("--publish", action="store_true")
p.add_argument("--allow-fake-publish", action="store_true")
p.add_argument("--no-state-safety-check", action="store_true")
p.add_argument("--rate-hz", type=float, default=50.0)
p.add_argument("--action-clip", type=float, default=0.3)
p.add_argument("--kp", type=float, default=20.0)
p.add_argument("--kd", type=float, default=0.5)
p.add_argument("--cmd-x", type=float, default=0.0)
p.add_argument("--cmd-y", type=float, default=0.0)
p.add_argument("--cmd-yaw", type=float, default=0.0)
p.add_argument("--use-rc", action="store_true")
p.add_argument("--rc-x-scale", type=float, default=0.5)
p.add_argument("--rc-y-scale", type=float, default=0.5)
p.add_argument("--rc-yaw-scale", type=float, default=0.5)
p.add_argument("--wtw-r2-logic", action="store_true", help="press R2 to calibrate/start; press R2 to pause/resume")
p.add_argument("--r2-pressed-field", default="right_lower_right_switch_pressed")
p.add_argument("--fake-auto-r2", action="store_true", help="auto-accept R2 waits in fake-state tests")
p.add_argument("--cal-step-rad", type=float, default=0.05)
p.add_argument("--cal-dt", type=float, default=0.05)
p.add_argument("--hold-default-rate-hz", type=float, default=20.0)
p.add_argument("--log-dir", default="")
p.add_argument("--log-every", type=int, default=1)
p.add_argument("--log-flush-every", type=int, default=10)
p.add_argument("--print-every", type=int, default=50)
p.add_argument("--warmup-steps", type=int, default=10)
p.add_argument("--max-steps", type=int, default=0)
return p.parse_args()
def main():
args = parse_args()
if args.publish and not args.real_state and not args.allow_fake_publish:
raise RuntimeError("Refusing to publish with fake state. Add --allow-fake-publish only for offline tests.")
if args.rate_hz <= 0:
raise ValueError("--rate-hz must be positive")
print("[INFO] Args:", vars(args))
lc = lcm.LCM(args.lcm_url)
se = create_se(args, lc)
policy = OnnxPolicy(args.onnx)
logger = JsonlLogger(args.log_dir, args)
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
step = 0
msg_id = 0
dt = 1.0 / args.rate_hz
next_t = time.perf_counter()
if not args.publish:
print("[INFO] DRY-RUN: no pd_plustau_targets will be published.")
else:
print("[WARN] PUBLISH mode enabled. Hang up the robot for first tests.")
try:
if args.wtw_r2_logic and args.publish:
wait_r2(args, se, "[WTW R2] About to calibrate; robot will stand. Press R2 to calibrate.")
msg_id = calibrate_default(args, lc, se, msg_id)
wait_r2(args, se, "[WTW R2] Starting pose calibrated. Press R2 to start ONNX controller.")
elif args.wtw_r2_logic and not args.publish:
print("[INFO] --wtw-r2-logic enabled but publish=False; startup R2 waiting/calibration skipped in dry-run.")
print("[INFO] Starting loop. Ctrl+C to exit.")
while not EXIT:
t0 = time.perf_counter()
mode = "policy"
commands = get_commands(args, se)
obs, info = compute_obs(se, commands, last_actions)
ok, reason = (True, "ok")
if args.real_state and not args.no_state_safety_check:
ok, reason = state_ok(info)
r2 = consume_r2_event(args, se) if args.wtw_r2_logic else False
if args.wtw_r2_logic and r2 and args.publish:
print("[WTW R2] R2 pressed during policy: pause and return to default pose.")
msg_id = calibrate_default(args, lc, se, msg_id)
msg_id = hold_default_until_r2(args, lc, se, msg_id)
last_actions[:] = 0.0
step += 1
continue
action_raw = policy(obs)
action_safe, q_isaac, q_wtw, q_unitree = action_to_targets(action_raw, args.action_clip)
last_actions = action_safe.copy()
publish_sent = False
should_publish = args.publish and step >= args.warmup_steps
if should_publish and args.real_state and not args.no_state_safety_check and not ok:
if step % max(1, args.print_every) == 0:
print(f"[WARN] Not publishing because state check failed: {reason}")
elif should_publish:
publish(lc, q_unitree, args, msg_id)
msg_id += 1
publish_sent = True
elif args.publish and step < args.warmup_steps and step % max(1, args.print_every) == 0:
print(f"[INFO] Warmup step {step}/{args.warmup_steps}: not publishing yet.")
loop_ms = (time.perf_counter() - t0) * 1000.0
logger.log(
step,
mode=mode,
loop_ms=loop_ms,
commands_isaac=commands,
obs_isaac=obs,
action_raw_isaac=action_raw,
action_safe_isaac=action_safe,
joint_targets_isaac=q_isaac,
joint_targets_wtw=q_wtw,
joint_targets_unitree=q_unitree,
base_lin_vel=info["base_lin_vel"],
base_ang_vel=info["base_ang_vel"],
projected_gravity=info["projected_gravity"],
dof_pos_wtw=info["dof_pos_wtw"],
dof_vel_wtw=info["dof_vel_wtw"],
dof_pos_isaac=info["dof_pos_isaac"],
dof_vel_isaac=info["dof_vel_isaac"],
state_ok=ok,
state_reason=reason,
r2_pressed_event=r2,
publish_sent=publish_sent,
)
if step % max(1, args.print_every) == 0:
print_debug(step, mode, args, commands, obs, info, action_raw, action_safe, q_isaac, q_wtw, q_unitree, ok, reason, publish_sent, r2)
step += 1
if args.max_steps > 0 and step >= args.max_steps:
print("[INFO] max_steps reached. Exiting.")
break
next_t += dt
sleep = next_t - time.perf_counter()
if sleep > 0:
time.sleep(sleep)
else:
next_t = time.perf_counter()
finally:
logger.close()
print("[INFO] Exiting.")
if __name__ == "__main__":
main()

650
deploy_wtw/deploy_wtw_pro_sdk.py Normal file
View File

@@ -0,0 +1,650 @@
#!/usr/bin/env python3
"""
deploy_wtw_pro_sdk.py
Walk-These-Ways (RMA) deployment on Unitree Go1 PRO via go1_pro_sdk.
Direct MCU control — no LCM, no official Unitree SDK.
Model: body_latest.jit (2102 → 12) + adaptation_module_latest.jit (2100 → 2)
Observation: 70-dim, 30-step history (2100 dims total)
Setup:
cd /path/to/go1_pro_sdk && pip install -e .
pip install torch onnxruntime numpy
Before running:
ssh pi@192.168.123.161
sudo pkill -9 -f keep_sport_alive; sudo pkill -9 -f Legged_sport; sudo pkill -9 -f appTransit
Usage:
python deploy_wtw_pro_sdk.py --kill-sport --log-dir ../logs
"""
import argparse
import json
import signal
import subprocess
import time
from datetime import datetime
from enum import Enum
from pathlib import Path
import numpy as np
import torch
from go1_pro_sdk import (
MCUClient, LowCmd, MotorCmd, MotorMode,
apply_safety, JOINT_NAMES as SDK_JOINT_NAMES,
)
HERE = Path(__file__).parent.resolve()
# ─── WTW policy constants ───
NUM_OBS = 70
NUM_ACTIONS = 12
NUM_COMMANDS = 15
NUM_OBS_HISTORY = 30 # 30 steps of history
OBS_HISTORY_DIM = NUM_OBS * NUM_OBS_HISTORY # 2100
BODY_INPUT_DIM = 2102 # 2100 (history) + 2 (latent)
ADAPT_INPUT_DIM = 2100
LATENT_DIM = 2
# WTW joint order: FL→FR→RL→RR (per leg: hip/thigh/calf)
# This is different from SDK order: FR→FL→RR→RL
WTW_JOINT_NAMES = [
"FL_hip", "FL_thigh", "FL_calf",
"FR_hip", "FR_thigh", "FR_calf",
"RL_hip", "RL_thigh", "RL_calf",
"RR_hip", "RR_thigh", "RR_calf",
]
# SDK joint order: FR→FL→RR→RL (per leg: hip/thigh/calf)
# Map: SDK index → WTW index
SDK_TO_WTW = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int64)
# Map: WTW index → SDK index
WTW_TO_SDK = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int64)
# Default joint angles in WTW order
DEFAULT_ANGLES_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 joint angles in SDK order
DEFAULT_ANGLES_SDK = DEFAULT_ANGLES_WTW[WTW_TO_SDK]
# Observation scales (WTW standard)
OBS_SCALES = {
"lin_vel": 2.0, "ang_vel": 0.25,
"dof_pos": 1.0, "dof_vel": 0.05,
"body_height_cmd": 2.0, "footswing_height_cmd": 0.15,
"body_pitch_cmd": 0.3, "body_roll_cmd": 0.3,
"stance_width_cmd": 1.0, "stance_length_cmd": 1.0,
"aux_reward_cmd": 1.0,
}
COMMANDS_SCALE = np.array([
OBS_SCALES["lin_vel"], OBS_SCALES["lin_vel"], OBS_SCALES["ang_vel"],
OBS_SCALES["body_height_cmd"], 1.0, 1.0, 1.0, 1.0, 1.0,
OBS_SCALES["footswing_height_cmd"],
OBS_SCALES["body_pitch_cmd"], OBS_SCALES["body_roll_cmd"],
OBS_SCALES["stance_width_cmd"], OBS_SCALES["stance_length_cmd"],
OBS_SCALES["aux_reward_cmd"],
], dtype=np.float32)[:NUM_COMMANDS]
ACTION_SCALE = 0.25
HIP_SCALE_REDUCTION = 0.5 # hip joints get half action
CLIP_ACTIONS = 10.0
CLIP_OBS = 100.0
EXIT = False
def _sig_handler(signum, frame):
global EXIT
EXIT = True
signal.signal(signal.SIGINT, _sig_handler)
signal.signal(signal.SIGTERM, _sig_handler)
# ─── State machine ───
class State(Enum):
IDLE = "IDLE"
CALIBRATE = "CALIBRATE"
HOLD = "HOLD"
RL = "RL"
# ─── Quaternion math ───
def quat_to_rot_matrix(q):
w, x, y, z = q
return np.array([
[1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y],
[ 2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x],
[ 2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y],
], dtype=np.float32)
def get_projected_gravity(quaternion):
R = quat_to_rot_matrix(quaternion)
return (R.T @ np.array([0., 0., -1.], dtype=np.float32)).astype(np.float32)
# ─── Observation ───
def build_commands_default():
"""Build default command vector for trotting gait."""
cmd = np.zeros(NUM_COMMANDS, dtype=np.float32)
cmd[0:3] = [0.0, 0.0, 0.0] # vx, vy, wz
cmd[3] = 0.0 # height command (zero = nominal)
cmd[4] = 3.0 # frequency (Hz)
cmd[5] = 0.5 # phase (trot = 0.5 offset)
cmd[6] = 0.0 # offset
cmd[7] = 0.0 # bound
cmd[8] = 0.5 # duration (stance ratio)
cmd[9] = 0.15 # swing_height (matches reference footswing_height_cmd)
cmd[10] = 0.0 # body_pitch
cmd[11] = 0.0 # body_roll
cmd[12] = 0.25 # stance_width
cmd[13] = 0.4 # stance_length
cmd[14] = 0.0 # aux_reward
return cmd
class ClockState:
"""Track gait indices and compute clock_inputs (4-dim sin per foot)."""
def __init__(self):
self.gait_indices = 0.0
self.dt = 0.01 # 100Hz
def step(self, commands, dt=None):
if dt is not None:
self.dt = dt
freq = commands[4]
phase = commands[5]
offset = commands[6]
bound = commands[7] if NUM_COMMANDS > 8 else 0.0
self.gait_indices = (self.gait_indices + self.dt * freq) % 1.0
foot_indices = [
self.gait_indices + phase + offset + bound, # FL
self.gait_indices + offset, # FR
self.gait_indices + bound, # RL
self.gait_indices + phase, # RR
]
clock = np.array([np.sin(2 * np.pi * fi) for fi in foot_indices], dtype=np.float32)
return clock
def reset(self):
self.gait_indices = 0.0
def compute_obs_wtw(imu, motor_states, commands, actions, last_actions, clock_inputs):
"""Build 70-dim observation matching WTW LCM agent layout."""
obs = np.zeros(NUM_OBS, dtype=np.float32)
# 1. projected_gravity (3)
obs[0:3] = get_projected_gravity(imu.quaternion)
# 2. commands * scale (15)
offset = 3
obs[offset:offset+NUM_COMMANDS] = commands * COMMANDS_SCALE
offset += NUM_COMMANDS
# 3. dof_pos_rel in WTW order (12)
dof_pos_sdk = np.array([motor_states[i].q for i in range(12)], dtype=np.float32)
dof_pos_wtw = dof_pos_sdk[SDK_TO_WTW]
obs[offset:offset+12] = (dof_pos_wtw - DEFAULT_ANGLES_WTW) * OBS_SCALES["dof_pos"]
offset += 12
# 4. dof_vel in WTW order (12)
dof_vel_sdk = np.array([motor_states[i].dq for i in range(12)], dtype=np.float32)
dof_vel_wtw = dof_vel_sdk[SDK_TO_WTW]
obs[offset:offset+12] = dof_vel_wtw * OBS_SCALES["dof_vel"]
offset += 12
# 5. actions clipped (12)
obs[offset:offset+12] = np.clip(actions, -CLIP_ACTIONS, CLIP_ACTIONS)
offset += 12
# 6. last_actions (12)
obs[offset:offset+12] = last_actions
offset += 12
# 7. clock_inputs (4)
obs[offset:offset+4] = clock_inputs
obs = np.clip(obs, -CLIP_OBS, CLIP_OBS)
obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
return obs
# ─── Model ───
class WTWPolicy:
def __init__(self, body_path, adapt_path):
self.body = torch.jit.load(str(body_path), map_location='cpu')
self.adapt = torch.jit.load(str(adapt_path), map_location='cpu')
self.body.eval()
self.adapt.eval()
self.obs_history = torch.zeros(1, OBS_HISTORY_DIM, dtype=torch.float)
self.latent = torch.zeros(1, LATENT_DIM, dtype=torch.float)
print(f"[INFO] WTW body: {body_path}")
print(f"[INFO] WTW adapt: {adapt_path}")
print(f"[INFO] History: {NUM_OBS} obs × {NUM_OBS_HISTORY} steps = {OBS_HISTORY_DIM}")
def reset(self):
self.obs_history.zero_()
self.latent.zero_()
def __call__(self, obs):
obs_t = torch.from_numpy(obs.reshape(1, -1)).float()
# Update history: shift left, append new obs
self.obs_history = torch.cat(
(self.obs_history[:, NUM_OBS:], obs_t), dim=-1)
# Adaptation module: history → latent
with torch.no_grad():
self.latent = self.adapt(self.obs_history)
# Body: [history, latent] → action
body_input = torch.cat((self.obs_history, self.latent), dim=-1)
with torch.no_grad():
action = self.body(body_input)
return action.numpy().flatten().astype(np.float32)
# ─── Remote controller ───
def get_rc_commands(state, base_cmd, args):
r = state.remote
base_cmd[0] = r.ly * args.rc_vx_scale
base_cmd[1] = -r.lx * args.rc_vy_scale
base_cmd[2] = -r.rx * args.rc_wz_scale
return base_cmd
class RCEdgeDetector:
def __init__(self):
self._prev = set()
def update(self, state):
current = set(state.remote.pressed)
rising = current - self._prev
falling = self._prev - current
self._prev = current
return rising, falling
# ─── Safety wrappers ───
def send_hold_cmd(client, state, args):
cmd = LowCmd()
for j in range(12):
cmd.set_motor(j, MotorCmd(
mode=MotorMode.Servo,
q=float(DEFAULT_ANGLES_SDK[j]), dq=0.0, tau=0.0,
Kp=args.kp, Kd=args.kd,
))
apply_safety(cmd, state, power_factor=args.power_factor,
position_limit_on=True, position_protect_limit=None)
client.send(cmd)
def send_rl_cmd(client, state, action_wtw, args):
"""Convert WTW action to SDK targets and send."""
# Scale action → position offset in WTW order
offset_wtw = action_wtw * ACTION_SCALE
# Apply hip scale reduction
for i in [0, 3, 6, 9]: # hip indices in WTW order
offset_wtw[i] *= HIP_SCALE_REDUCTION
# Target in WTW order
targets_wtw = DEFAULT_ANGLES_WTW + offset_wtw
# Convert to SDK order
targets_sdk = targets_wtw[WTW_TO_SDK]
cmd = LowCmd()
for j in range(12):
cmd.set_motor(j, MotorCmd(
mode=MotorMode.Servo,
q=float(targets_sdk[j]), dq=0.0, tau=0.0,
Kp=args.kp, Kd=args.kd,
))
pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None
apply_safety(cmd, state, power_factor=args.power_factor,
position_limit_on=True, position_protect_limit=pp_limit)
client.send(cmd)
def kill_sport_processes(host, user):
cmds = [
"sudo pkill -9 -f keep_sport_alive",
"sudo pkill -9 -f Legged_sport",
"sudo pkill -9 -f appTransit",
]
ssh_target = f"{user}@{host}"
print(f"[INFO] Killing sport processes on {ssh_target}...")
try:
result = subprocess.run(
["ssh", ssh_target, " && ".join(cmds)],
capture_output=True, text=True, timeout=15)
if result.returncode == 0 or "no process" in result.stderr.lower():
print("[INFO] Sport processes killed.")
return True
print(f"[WARN] SSH returned {result.returncode}: {result.stderr.strip()}")
return False
except Exception as e:
print(f"[WARN] Failed: {e}")
return False
# ─── JSONL logger ───
class JsonlLogger:
def __init__(self, log_dir, args):
self.enabled = bool(log_dir)
self.fp = None
self.run_dir = None
self.flush_every = 50
if not self.enabled:
return
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
self.run_dir = Path(log_dir).expanduser().resolve() / f"wtw_deploy_{ts}"
self.run_dir.mkdir(parents=True, exist_ok=True)
meta = {
"created_at": ts, "num_obs": NUM_OBS, "num_actions": NUM_ACTIONS,
"num_commands": NUM_COMMANDS, "num_obs_history": NUM_OBS_HISTORY,
"action_scale": ACTION_SCALE,
"default_angles_wtw": DEFAULT_ANGLES_WTW.tolist(),
"default_angles_sdk": DEFAULT_ANGLES_SDK.tolist(),
"joint_names_wtw": WTW_JOINT_NAMES,
"joint_names_sdk": list(SDK_JOINT_NAMES),
}
for k, v in vars(args).items():
if isinstance(v, (str, int, float, bool, type(None))):
meta[k] = v
(self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1)
print(f"[INFO] Log dir: {self.run_dir}")
def log(self, step, **kw):
if not self.enabled:
return
rec = {"step": int(step), "time_wall": time.time()}
for k, v in kw.items():
if isinstance(v, np.ndarray):
rec[k] = np.asarray(v, dtype=np.float32).reshape(-1).tolist()
elif isinstance(v, (np.float32, np.float64)):
rec[k] = float(v)
elif isinstance(v, (np.int32, np.int64)):
rec[k] = int(v)
else:
rec[k] = v
self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
if step % self.flush_every == 0:
self.fp.flush()
def close(self):
if self.fp:
self.fp.flush(); self.fp.close()
print(f"[INFO] Log saved: {self.run_dir}")
# ─── Ramp to default ───
def ramp_to_default(client, args, state):
print("[INFO] Ramping to default pose (~2s)...")
current_sdk = np.array([state.motorState[i].q for i in range(12)], dtype=np.float32)
error = current_sdk - DEFAULT_ANGLES_SDK
if np.max(np.abs(error)) < 0.05:
print("[INFO] Already near default pose.")
return state
ramp_steps = 200
step_err = error / ramp_steps
for i in range(ramp_steps):
if EXIT: return state
new_state = client.recv_latest()
if new_state is not None:
state = new_state
targets = DEFAULT_ANGLES_SDK + (error - step_err * min(i + 1, ramp_steps))
cmd = LowCmd()
for j in range(12):
cmd.set_motor(j, MotorCmd(
mode=MotorMode.Servo,
q=float(targets[j]), dq=0.0, tau=0.0,
Kp=args.kp_cal, Kd=args.kd_cal,
))
apply_safety(cmd, state, power_factor=args.power_factor,
position_limit_on=True, position_protect_limit=None)
client.send(cmd)
time.sleep(0.01)
if i % 50 == 0:
actual = np.array([state.motorState[j].q for j in range(12)])
print(f" ramp {i}/{ramp_steps} target_err={np.max(np.abs(targets-DEFAULT_ANGLES_SDK)):.3f} actual_err={np.max(np.abs(actual-DEFAULT_ANGLES_SDK)):.3f}")
for _ in range(50):
if EXIT: return state
new_state = client.recv_latest()
if new_state is not None:
state = new_state
send_hold_cmd(client, state, args)
time.sleep(0.01)
print("[INFO] Default pose reached.")
return state
# ─── Main ───
def main():
parser = argparse.ArgumentParser(description="WTW RMA deployment on Go1 PRO via go1_pro_sdk")
parser.add_argument("--body", default=str(HERE / "body_latest.jit"))
parser.add_argument("--adapt", default=str(HERE / "adaptation_module_latest.jit"))
parser.add_argument("--kill-sport", action="store_true")
parser.add_argument("--pi-host", default="192.168.123.161")
parser.add_argument("--pi-user", default="pi")
parser.add_argument("--kp", type=float, default=20.0)
parser.add_argument("--kd", type=float, default=0.5)
parser.add_argument("--kp-cal", type=float, default=15.0)
parser.add_argument("--kd-cal", type=float, default=0.5)
parser.add_argument("--power-factor", type=int, default=7)
parser.add_argument("--position-protect-limit", type=float, default=1.0)
parser.add_argument("--rc-vx-scale", type=float, default=1.0)
parser.add_argument("--rc-vy-scale", type=float, default=1.0)
parser.add_argument("--rc-wz-scale", type=float, default=1.0)
parser.add_argument("--cmd-x", type=float, default=0.0)
parser.add_argument("--cmd-y", type=float, default=0.0)
parser.add_argument("--cmd-yaw", type=float, default=0.0)
parser.add_argument("--rate-hz", type=float, default=50.0)
parser.add_argument("--warmup-steps", type=int, default=50)
parser.add_argument("--max-steps", type=int, default=0)
parser.add_argument("--print-every", type=int, default=50)
parser.add_argument("--log-dir", default="")
parser.add_argument("--log-flush-every", type=int, default=50)
args = parser.parse_args()
print("""
╔══════════════════════════════════════════════════════════════╗
║ WTW RMA Deployment (go1_pro_sdk direct MCU) ║
║ 1. Robot SUSPENDED ║
║ 2. Use --kill-sport to auto-kill Pi processes ║
║ 3. R2=go/stop, L2=estop, Left-stick=move, Right-stick=turn ║
╚══════════════════════════════════════════════════════════════╝
""")
input("Press Enter when ready...")
if args.kill_sport:
kill_sport_processes(args.pi_host, args.pi_user)
print("[INFO] Loading WTW models...")
policy = WTWPolicy(args.body, args.adapt)
print("[INFO] Connecting to MCU...")
client = MCUClient()
logger = JsonlLogger(args.log_dir, args)
try:
print("[INFO] Waking MCU...")
client.wake_mcu(n_frames=50, dt=0.01)
state = client.recv_state(timeout=2.0)
if state is None:
print("[ERROR] No state received.")
return 1
print(f"[INFO] Connected. Battery={state.bms.SOC}%")
print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg")
sm_state = State.IDLE
edge = RCEdgeDetector()
edge.update(state)
clock = ClockState()
base_cmd = build_commands_default()
actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
step = 0
dt = 1.0 / args.rate_hz
next_t = time.perf_counter()
print(f"[INFO] Rate: {args.rate_hz}Hz. Ctrl+C to exit.")
print(f"[INFO] State machine: IDLE → (R2) → CALIBRATE → HOLD → (R2) → RL")
while not EXIT:
new_state = client.recv_latest()
if new_state is not None:
state = new_state
if state is None:
time.sleep(0.001)
continue
rising, falling = edge.update(state)
r2_rose = "R2" in rising
l2_rose = "L2" in rising
# Emergency stop
if l2_rose and sm_state != State.IDLE:
print(f"\n[L2 EMERGENCY] {sm_state.value} → IDLE")
sm_state = State.IDLE
actions[:] = 0.0
last_actions[:] = 0.0
policy.reset()
clock.reset()
client.send(LowCmd())
# State machine
if sm_state == State.IDLE:
if step % 10 == 0:
client.send(LowCmd())
if r2_rose:
print("\n[R2] IDLE → CALIBRATE")
sm_state = State.CALIBRATE
state = ramp_to_default(client, args, state)
if EXIT: break
sm_state = State.HOLD
print("[STATE] → HOLD")
elif sm_state == State.HOLD:
send_hold_cmd(client, state, args)
if r2_rose:
print("\n[R2] HOLD → RL")
sm_state = State.RL
actions[:] = 0.0
last_actions[:] = 0.0
policy.reset()
clock.reset()
elif sm_state == State.RL:
if r2_rose:
print("\n[R2] RL → HOLD")
sm_state = State.HOLD
actions[:] = 0.0
last_actions[:] = 0.0
state = ramp_to_default(client, args, state)
if EXIT: break
continue
# Update commands from RC
base_cmd = get_rc_commands(state, base_cmd, args)
# Clock inputs
clock_inputs = clock.step(base_cmd, dt)
# Observation
obs = compute_obs_wtw(state.imu, state.motorState,
base_cmd, actions, last_actions, clock_inputs)
# Inference
action_raw = policy(obs)
actions = np.clip(action_raw, -CLIP_ACTIONS, CLIP_ACTIONS).astype(np.float32)
last_actions = actions.copy()
if step >= args.warmup_steps:
send_rl_cmd(client, state, action_raw, args)
logger.log(
step, mode="RL",
commands=base_cmd,
obs_wtw=obs,
action_raw=action_raw,
action_safe=actions,
dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]),
dof_vel_sdk=np.array([state.motorState[i].dq for i in range(12)]),
clock_inputs=clock_inputs,
imu_rpy_deg=np.degrees(state.imu.rpy),
rc_buttons=state.remote.pressed,
)
else:
logger.log(
step, mode=sm_state.value,
dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]),
imu_rpy_deg=np.degrees(state.imu.rpy),
rc_buttons=state.remote.pressed,
)
# Status print
if step % args.print_every == 0:
dof_pos = np.array([state.motorState[i].q for i in range(12)])
print(f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% "
f"rpy={np.round(np.degrees(state.imu.rpy), 1)}")
print(f" RC: lx={state.remote.lx:+.2f} ly={state.remote.ly:+.2f} "
f"btns={state.remote.pressed}")
print(f" joint: {np.round(dof_pos, 2)}")
if sm_state == State.RL:
print(f" action max: {np.max(np.abs(actions)):.2f}")
step += 1
if args.max_steps > 0 and step >= args.max_steps:
print("[INFO] max_steps reached.")
break
next_t += dt
sleep = next_t - time.perf_counter()
if sleep > 0:
time.sleep(sleep)
else:
next_t = time.perf_counter()
finally:
if logger is not None:
logger.close()
print("[INFO] Safe stopping...")
client.safe_stop(n_frames=50, dt=0.002)
client.close()
print("[INFO] Done.")
if __name__ == "__main__":
main()

243
deploy_wtw/sim2sim_wtw_test.py Normal file
View 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()