File size: 4,166 Bytes
5a27506 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #!/usr/bin/env python
"""Closed-loop MicroDuck: obs.rs-раскладка 61D из MuJoCo, учитель ONNX или VLA.
Валидация конвенций через учителя: если alpha_walking в этой петле идёт —
obs-билдер и применение действий верны, и той же петлёй можно гнать VLA.
obs.rs layout:
0..3 gyro (trunk frame, rad/s)
3..6 projected gravity (trunk frame, unit)
6..20 joint pos - home (14, без рта)
20..34 joint vel (14)
34..48 previous action (14)
48..61 command: vx,vy,vyaw | 4 head targets | body x,y(=0) | z | roll | pitch | yaw(=0)
Действие: 14 позиционных таргетов; применение ctrl = home + action (STAND offset,
как в mjlab JointPositionAction; teacher выход в тех же единицах).
"""
import sys, json
from pathlib import Path
import numpy as np
import mujoco
SCENE='/root/microduck_rl/src/mjlab_microduck/robot/microduck/scene_walk.xml'
model=mujoco.MjModel.from_xml_path(SCENE)
data=mujoco.MjData(model)
key=model.keyframe('STAND')
HOME=key.ctrl.copy() # home-поза 14 суставов
CTRL_DT=1/50
SPC=max(1,int(round(CTRL_DT/model.opt.timestep)))
def quat_rotate_inv(q, v):
w,x,y,z=q
# R^T v
R=np.array([[1-2*(y*y+z*z),2*(x*y-w*z),2*(x*z+w*y)],
[2*(x*y+w*z),1-2*(x*x+z*z),2*(y*z-w*x)],
[2*(x*z-w*y),2*(y*z+w*x),1-2*(x*x+y*y)]])
return R.T@v
def build_obs(prev_action, command):
q=data.qpos[3:7] # trunk quat (w,x,y,z)
gyro=data.qvel[3:6].copy() # угловая скорость в ЛОКАЛЬНОЙ раме (mujoco freejoint)
grav=quat_rotate_inv(q, np.array([0,0,-1.0]))
jpos=data.qpos[7:7+14]-HOME
jvel=data.qvel[6:6+14]
return np.concatenate([gyro,grav,jpos,jvel,prev_action,command]).astype(np.float32)
def run(policy_step, command, seconds=8, tag="run", render_cam=None, vla_hooks=None):
import imageio.v2 as imageio
renderer=mujoco.Renderer(model,height=320,width=320)
cam=mujoco.MjvCamera(); cam.distance,cam.azimuth,cam.elevation=0.65,140,-12
mujoco.mj_resetDataKeyframe(model,data,key.id); mujoco.mj_forward(model,data)
prev=np.zeros(14,np.float32)
frames=[]; fell_at=None
n_ctrl=int(seconds*50)
for i in range(n_ctrl):
if i%5==0: # 10 Гц перепланирование
obs=build_obs(prev,command)
plan=policy_step(obs, i) # (k,14) минимум 5 шагов
a=plan[i%5] if plan.ndim==2 else plan
data.ctrl[:]=HOME+a
prev=a.astype(np.float32)
for _ in range(SPC): mujoco.mj_step(model,data)
if data.qpos[2]<0.05 and fell_at is None: fell_at=i/50
if i%3==0:
cam.lookat[:]=[float(data.qpos[0]),float(data.qpos[1]),0.12]
renderer.update_scene(data,camera=cam)
frames.append(renderer.render().copy())
imageio.mimwrite(f"/root/cl_{tag}.mp4",frames,fps=16,quality=7)
print(f"{tag}: {seconds}s, высота в конце {data.qpos[2]:.3f}, "
f"{'УПАЛ на '+str(round(fell_at,1))+'s' if fell_at else 'НЕ УПАЛ'}, "
f"пройдено xy {np.hypot(data.qpos[0],data.qpos[1]):.2f} м", flush=True)
return fell_at
if __name__=="__main__":
import onnxruntime as ort
which=sys.argv[1] if len(sys.argv)>1 else "walking"
sess=ort.InferenceSession(f"/root/.cache/huggingface/hub/models--pollen-robotics--microduck-policies/snapshots/{'*'}/alpha_{which}.onnx".replace('*','') if False else __import__('huggingface_hub').hf_hub_download("pollen-robotics/microduck-policies", f"alpha_{which}.onnx"))
def teacher(obs,i):
a=sess.run(None,{"obs":obs[None].astype(np.float32)})[0][0]
return np.tile(a,(5,1)) # учитель даёт 1 шаг — держим до реплана
cmd=np.zeros(13,np.float32)
if which=="walking": cmd[0]=0.15 # vx м/с
cmd[9]=0.115 # body z: стоячая высота
run(teacher,cmd,seconds=8,tag=f"teacher_{which}")
|