fix: clamp std before distribution, lower init_noise to 0.5, NaN guard
This commit is contained in:
8x54zj-m
236
deploy_orig/go1_sim2sim.py
Normal file
236
deploy_orig/go1_sim2sim.py
Normal file
@@ -0,0 +1,236 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Go1 sim2sim MuJoCo viewer — Original 30k flat training.
|
||||
|
||||
Usage: python go1_sim2sim.py
|
||||
|
||||
Requires: mujoco, onnxruntime, pynput
|
||||
Install: pip install mujoco onnxruntime pynput
|
||||
|
||||
Controls:
|
||||
W/S: forward/back Q/E: strafe left/right
|
||||
A/D: rotate Space: stop R: reset Esc: quit
|
||||
"""
|
||||
import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
|
||||
from mujoco import viewer
|
||||
from pynput import keyboard
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
ONNX = os.path.join(HERE, "policy.onnx")
|
||||
|
||||
# ── Parameters (original MotrixLab Go1 config) ──
|
||||
NUM_OBS = 48
|
||||
KP, KD = 80.0, 0.5 # KD=0.5 + MuJoCo joint_damping(0.5) = 1.0 = training kd
|
||||
ACTION_SCALE = 0.05
|
||||
CLIP = 23.7
|
||||
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)
|
||||
|
||||
# ── Keyboard ──
|
||||
class KB:
|
||||
def __init__(s):
|
||||
s._q = queue.Queue(); s.running = True; s.held = set()
|
||||
def _n(s, k):
|
||||
try:
|
||||
if hasattr(k, 'char') and k.char: return k.char.lower()
|
||||
except: pass
|
||||
return str(k).lower()
|
||||
def _w(s):
|
||||
while s.running:
|
||||
try:
|
||||
et, k = s._q.get(timeout=0.05)
|
||||
n = s._n(k)
|
||||
if et == 'press': s.held.add(n)
|
||||
elif et == 'release': s.held.discard(n)
|
||||
except queue.Empty: pass
|
||||
def init(s):
|
||||
s._l = keyboard.Listener(
|
||||
on_press=lambda k: s._q.put(('press', k)),
|
||||
on_release=lambda k: s._q.put(('release', k)))
|
||||
s._l.start()
|
||||
s._t = threading.Thread(target=s._w, daemon=True); s._t.start()
|
||||
def keys(s): return s.held.copy()
|
||||
def stop(s): s.running = False; s._l.stop()
|
||||
|
||||
# ── Main ──
|
||||
def main():
|
||||
# The model XML is embedded below
|
||||
xml = '''<mujoco model="go1 scene">
|
||||
<compiler angle="radian" autolimits="true"/>
|
||||
<option timestep="0.005" integrator="Euler" iterations="60">
|
||||
<flag eulerdamp="disable"/>
|
||||
</option>
|
||||
<custom>
|
||||
<numeric data="30" name="max_contact_points"/>
|
||||
<numeric data="12" name="max_geom_pairs"/>
|
||||
</custom>
|
||||
<default>
|
||||
<default class="go1">
|
||||
<geom condim="1"/>
|
||||
<joint axis="0 1 0" armature="0.005" damping="0.5"/>
|
||||
<default class="abduction">
|
||||
<joint axis="1 0 0" range="-0.863 0.863" frictionloss="0.3"/>
|
||||
</default>
|
||||
<default class="hip">
|
||||
<joint range="-0.686 4.501" frictionloss="0.3"/>
|
||||
</default>
|
||||
<default class="knee">
|
||||
<joint range="-2.818 -0.888" frictionloss="1.0"/>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
<asset>
|
||||
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
|
||||
<texture name="ground" type="2d" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
||||
<material name="ground" texture="ground" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
|
||||
<geom name="floor" size="0 0 0.01" type="plane" material="ground" contype="1" conaffinity="0" priority="1" friction="0.6" condim="3"/>
|
||||
<body name="trunk" pos="0 0 0.4" childclass="go1">
|
||||
<freejoint/>
|
||||
<inertial pos="0.0223 0.002 -0.0005" quat="-0.00342088 0.705204 0.000106698 0.708996" mass="5.204" diaginertia="0.0716565 0.0630105 0.0168101"/>
|
||||
<geom name="trunk_geom" contype="0" conaffinity="0" group="2" type="box" size="0.35 0.12 0.08" rgba="0.4 0.4 0.4 1"/>
|
||||
<geom name="trunk_col" contype="1" conaffinity="1" group="3" pos="0.24 0 0" size="0.05 0.05 0.05" type="box"/>
|
||||
<site name="imu" pos="-0.01592 -0.06659 -0.00617" group="5"/>
|
||||
<!-- FR leg -->
|
||||
<body name="FR_hip" pos="0.1881 -0.04675 0">
|
||||
<joint class="abduction" name="FR_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FR_thigh" pos="0 -0.08 0">
|
||||
<joint class="hip" name="FR_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="FR_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="FR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- FL leg -->
|
||||
<body name="FL_hip" pos="0.1881 0.04675 0">
|
||||
<joint class="abduction" name="FL_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FL_thigh" pos="0 0.08 0">
|
||||
<joint class="hip" name="FL_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="FL_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="FL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- RR leg -->
|
||||
<body name="RR_hip" pos="-0.1881 -0.04675 0">
|
||||
<joint class="abduction" name="RR_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RR_thigh" pos="0 -0.08 0">
|
||||
<joint class="hip" name="RR_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="RR_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="RR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- RL leg -->
|
||||
<body name="RL_hip" pos="-0.1881 0.04675 0">
|
||||
<joint class="abduction" name="RL_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RL_thigh" pos="0 0.08 0">
|
||||
<joint class="hip" name="RL_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="RL_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="RL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="FR_calf" joint="FR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="FL_calf" joint="FL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="RR_calf" joint="RR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="RL_calf" joint="RL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
</actuator>
|
||||
<sensor>
|
||||
<gyro site="imu" name="gyro"/>
|
||||
<velocimeter site="imu" name="local_linvel"/>
|
||||
</sensor>
|
||||
</mujoco>'''
|
||||
|
||||
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, providers=['CPUExecutionProvider'])
|
||||
print(f"[Go1 sim2sim] ONNX={ONNX}")
|
||||
print(f"[Go1 sim2sim] PD kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} obs={NUM_OBS}-dim")
|
||||
print(f"[Go1 sim2sim] 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_a = np.zeros(12, dtype=np.float32)
|
||||
|
||||
while view.is_running():
|
||||
keys = kb.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_a[:] = 0
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
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 = 1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.0)
|
||||
if ' ' in keys: vx = vy = wz = 0.0
|
||||
|
||||
if step % 2 == 0: # 100Hz control (MuJoCo dt=0.005)
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro")
|
||||
adr = model.sensor_adr[sid]
|
||||
obs[0:3] = data.sensordata[adr:adr+3] * 0.25
|
||||
R = data.xmat[1].reshape(3, 3)
|
||||
obs[6:9] = (R.T @ np.array([0., 0., -1.])).astype(np.float32)
|
||||
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
|
||||
obs[21:33] = data.qvel[6:18] * 0.05
|
||||
obs[33:45] = last_a
|
||||
obs[45:48] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25])
|
||||
obs = np.clip(obs, -100., 100.)
|
||||
action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
|
||||
action = np.clip(action, -CLIP, CLIP)
|
||||
last_a = 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, CLIP)
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
step += 1
|
||||
time.sleep(0.001)
|
||||
|
||||
kb.stop(); view.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user