134 lines
4.7 KiB
Python
134 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Sim2sim test for 57-dim no-linevel policy (better3.onnx).
|
|
obs: gyro(3) + gravity(3) + dof_pos(12) + dof_vel(12) + last_actions(12) + commands(3) + contacts(12) = 57
|
|
Contacts set to zero.
|
|
|
|
Usage:
|
|
conda activate free_dog_sdk
|
|
mjpython sim2sim_57dim_test.py
|
|
mjpython sim2sim_57dim_test.py --action-scale 0.07 --action-ema-alpha 0.3
|
|
|
|
Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 Esc=退出
|
|
"""
|
|
|
|
import argparse, os, signal, time
|
|
import mujoco, numpy as np, onnxruntime as ort
|
|
from mujoco import viewer
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
XML = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml")
|
|
ONNX = os.path.join(HERE, "better3.onnx")
|
|
|
|
NUM_OBS, NUM_ACTIONS = 57, 12
|
|
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)
|
|
|
|
EXIT = False
|
|
def _sig(s, f): global EXIT; EXIT = True
|
|
signal.signal(signal.SIGINT, _sig)
|
|
|
|
class KB:
|
|
def __init__(s): s.h = set(); s._l = None
|
|
def _p(s,k):
|
|
try: s.h.add(k.char.lower())
|
|
except: s.h.add(str(k))
|
|
def _r(s,k):
|
|
try: s.h.discard(k.char.lower())
|
|
except: s.h.discard(str(k))
|
|
def init(s):
|
|
from pynput import keyboard
|
|
s._l = keyboard.Listener(on_press=s._p, on_release=s._r); s._l.start()
|
|
def keys(s): return s.h.copy()
|
|
def stop(s):
|
|
if s._l: s._l.stop()
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--action-scale", type=float, default=0.05)
|
|
p.add_argument("--action-ema-alpha", type=float, default=0.0)
|
|
p.add_argument("--kp", type=float, default=80.0)
|
|
p.add_argument("--kd", type=float, default=1.0)
|
|
args = p.parse_args()
|
|
|
|
model = mujoco.MjModel.from_xml_path(XML)
|
|
data = mujoco.MjData(model)
|
|
|
|
print(f"[INFO] 57-dim no-linevel ONNX: {ONNX}")
|
|
print(f"[INFO] action_scale={args.action_scale} ema_alpha={args.action_ema_alpha}")
|
|
print(f"[INFO] KP={args.kp} KD={args.kd} passive_damping={model.dof_damping[6]}")
|
|
|
|
sess = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"])
|
|
input_name = sess.get_inputs()[0].name
|
|
|
|
kb = KB(); kb.init()
|
|
|
|
data.qpos[:3] = [0,0,0.42]; data.qpos[3:7] = [1,0,0,0]; data.qpos[7:19] = DEFAULT_ANGLES
|
|
mujoco.mj_forward(model, data)
|
|
|
|
view = viewer.launch_passive(model, data)
|
|
step = 0; si = int(0.01 / model.opt.timestep) # 100Hz
|
|
last_a = np.zeros(12, dtype=np.float32); action = np.zeros(12, dtype=np.float32)
|
|
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[: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 % si == 0:
|
|
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
|
|
|
# gyro
|
|
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Gyro")
|
|
obs[0:3] = data.sensordata[model.sensor_adr[sid]:model.sensor_adr[sid]+3] * 0.25
|
|
|
|
# gravity
|
|
R = data.xmat[1].reshape(3, 3)
|
|
obs[3:6] = (R.T @ [0., 0., -1.]).astype(np.float32)
|
|
|
|
# dof_pos
|
|
obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
|
|
# dof_vel
|
|
obs[18:30] = data.qvel[6:18] * 0.05
|
|
# last_actions
|
|
obs[30:42] = last_a
|
|
# commands
|
|
obs[42:45] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25])
|
|
# contact forces = 0
|
|
obs[45:57] = 0.0
|
|
|
|
obs = np.clip(obs, -100., 100.)
|
|
action = sess.run(None, {input_name: obs.reshape(1, -1).astype(np.float32)})[0][0]
|
|
action = np.clip(action, -23.7, 23.7)
|
|
if 0 < args.action_ema_alpha < 1:
|
|
action = args.action_ema_alpha * action + (1 - args.action_ema_alpha) * last_a
|
|
last_a = action.copy()
|
|
|
|
targets = DEFAULT_ANGLES + action * args.action_scale
|
|
data.ctrl[:] = np.clip(args.kp*(targets-data.qpos[7:19])-args.kd*data.qvel[6:18], -23.7, 23.7)
|
|
|
|
mujoco.mj_step(model, data)
|
|
view.sync()
|
|
|
|
if step % 200 == 0:
|
|
print(f"[{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()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|