136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""MuJoCo sim2sim visualization for go2style policy (45-dim obs, no linvel).
|
|
|
|
Controls:
|
|
W/S: forward/back Q/E: left/right A/D: rotate Space: stop R: reset
|
|
"""
|
|
import numpy as np
|
|
import mujoco
|
|
from mujoco import viewer
|
|
import onnxruntime as ort
|
|
import os, sys, time, threading, queue
|
|
|
|
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
ONNX_PATH = os.path.join(PROJECT, "exports_go1_go2style", "policy.onnx")
|
|
XML_DIR = os.path.join(PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
|
|
|
|
# go2style params
|
|
NUM_OBS = 45
|
|
NUM_ACTIONS = 12
|
|
OBS_SCALES = {'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
|
|
ACTION_SCALE = 0.25
|
|
KP = 20.0
|
|
KD = 0.0 # MuJoCo joint自带damping=0.5, PD kd=0避免过阻尼
|
|
CLIP_ACTIONS = 23.7
|
|
CLIP_OBS = 100.0
|
|
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
|
|
|
|
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)
|
|
|
|
from pynput import keyboard
|
|
|
|
class KB:
|
|
def __init__(self):
|
|
self._q = queue.Queue(); self.running = True
|
|
self.held = set(); self._t = None; self._l = None
|
|
def _n(self, k):
|
|
try:
|
|
if hasattr(k,'char') and k.char: return k.char.lower()
|
|
except: pass
|
|
return str(k).lower()
|
|
def _w(self):
|
|
while self.running:
|
|
try:
|
|
et, k = self._q.get(timeout=0.05)
|
|
n = self._n(k)
|
|
if et == 'press': self.held.add(n)
|
|
elif et == 'release': self.held.discard(n)
|
|
except queue.Empty: pass
|
|
def init(self):
|
|
def op(k): self._q.put(('press',k))
|
|
def or_(k): self._q.put(('release',k))
|
|
self._l = keyboard.Listener(on_press=op, on_release=or_)
|
|
self._l.start()
|
|
self._t = threading.Thread(target=self._w, daemon=True); self._t.start()
|
|
print("[KB] 键盘就绪")
|
|
def held_keys(self): return self.held.copy()
|
|
def stop(self): self.running = False; self._l.stop()
|
|
|
|
def get_sensor(m, d, name):
|
|
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name)
|
|
if sid < 0: return None
|
|
adr = m.sensor_adr[sid]; dim = m.sensor_dim[sid]
|
|
return d.sensordata[adr:adr+dim].copy()
|
|
|
|
def compute_obs(model, data, commands, last_action):
|
|
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
|
# gyro [0:3]
|
|
g = get_sensor(model, data, "gyro")
|
|
obs[0:3] = (g if g is not None else data.qvel[3:6]) * OBS_SCALES['ang_vel']
|
|
# gravity [3:6]
|
|
R = data.xmat[1].reshape(3,3)
|
|
obs[3:6] = (R.T @ np.array([0.,0.,-1.])).astype(np.float32)
|
|
# joint pos [6:18]
|
|
obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * OBS_SCALES['dof_pos']
|
|
# joint vel [18:30]
|
|
obs[18:30] = data.qvel[6:18] * OBS_SCALES['dof_vel']
|
|
# last action [30:42]
|
|
obs[30:42] = last_action
|
|
# commands [42:45]
|
|
obs[42:45] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
|
|
return np.clip(obs, -CLIP_OBS, CLIP_OBS)
|
|
|
|
def main():
|
|
os.chdir(XML_DIR)
|
|
xml = open("scene_motor_actuator.xml").read()
|
|
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_PATH, providers=['CPUExecutionProvider'])
|
|
print(f"[ONNX] {ONNX_PATH}")
|
|
print(f"[CTRL] 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_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
|
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
|
decimation = 2 # MuJoCo dt=0.005, policy dt=0.01 → 2 steps per inference
|
|
|
|
while view.is_running():
|
|
keys = kb.held_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_action[:]=0; mujoco.mj_forward(model,data); print("[R] 重置")
|
|
if ' ' in keys: vx=vy=wz=0.0
|
|
|
|
vx = MAX_VX if 'w' in keys else (-MAX_VX if 's' in keys else 0.0)
|
|
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0)
|
|
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' in keys else 0.0)
|
|
|
|
if step % decimation == 0:
|
|
cmd = np.array([vx, vy, wz], dtype=np.float32)
|
|
obs = compute_obs(model, data, cmd, last_action)
|
|
action = session.run(None, {'observations': obs.reshape(1,-1).astype(np.float32)})[0][0]
|
|
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
|
last_action = 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_ACTIONS, CLIP_ACTIONS)
|
|
mujoco.mj_step(model, data)
|
|
view.sync()
|
|
step += 1
|
|
time.sleep(0.001)
|
|
|
|
kb.stop(); view.close()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|