fix: clamp std before distribution, lower init_noise to 0.5, NaN guard

scripts/view_orig.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""MuJoCo viewer for original Go1 (45-dim, PD 80/1.0, action_scale=0.05)"""
+import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
+from mujoco import viewer
+from pynput import keyboard
+
+PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ONNX = os.path.join(PROJECT, "exports_go1_orig", "policy.onnx")
+XML_DIR = os.path.join(PROJECT, "motrix_envs/src/motrix_envs/locomotion/go1/xmls")
+
+# Original params
+NUM_OBS = 57
+KP, KD = 80.0, 0.5   # KD=0.5 + joint_damping(0.5) = 1.0 = training kd
+ACTION_SCALE = 0.05
+CLIP = 23.7
+DEFAULT = 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)
+
+class KB:
+    def __init__(self):
+        self._q=queue.Queue(); self.running=True; self.held=set()
+    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):
+        self._l=keyboard.Listener(on_press=lambda k:self._q.put(('press',k)), on_release=lambda k:self._q.put(('release',k)))
+        self._l.start(); self._t=threading.Thread(target=self._w,daemon=True); self._t.start()
+    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]; return d.sensordata[adr:adr+m.sensor_dim[sid]].copy()
+
+def compute_obs(model,data,cmd,last_a):
+    obs=np.zeros(NUM_OBS,dtype=np.float32)
+    g=get_sensor(model,data,"gyro")
+    obs[0:3]=(g if g is not None else data.qvel[3:6])*0.25
+    R=data.xmat[1].reshape(3,3)
+    obs[3:6]=(R.T@np.array([0.,0.,-1.])).astype(np.float32)
+    obs[6:18]=(data.qpos[7:19]-DEFAULT)*1.0
+    obs[18:30]=data.qvel[6:18]*0.05
+    obs[30:42]=last_a
+    obs[42:45]=cmd; obs[45:57]=0.0  # contact_force*np.array([2.,2.,0.25],dtype=np.float32)
+    return np.clip(obs,-100.,100.)
+
+def main():
+    os.chdir(XML_DIR)
+    model=mujoco.MjModel.from_xml_string(open("scene_motor_actuator.xml").read())
+    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
+    mujoco.mj_forward(model,data)
+    session=ort.InferenceSession(ONNX,providers=['CPUExecutionProvider'])
+    print(f"[ORIG] kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} | W/S前后 Q/E左右 A/D旋转")
+    kb=KB(); kb.init()
+    view=viewer.launch_passive(model,data)
+    step,vx,vy,wz=0,0.,0.,0.
+    last_a=np.zeros(12,dtype=np.float32); action=np.zeros(12,dtype=np.float32)
+    dec=2
+    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
+            data.qvel[:]=0; last_a[:]=0; mujoco.mj_forward(model,data)
+        if ' ' in keys: vx=vy=wz=0.
+        vx=1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.)
+        vy=1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.)
+        wz=1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.)
+        if step%dec==0:
+            obs=compute_obs(model,data,np.array([vx,vy,wz],dtype=np.float32),last_a)
+            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+action*ACTION_SCALE
+        t=KP*(target-data.qpos[7:19])-KD*data.qvel[6:18]
+        data.ctrl[:]=np.clip(t,-CLIP,CLIP)
+        mujoco.mj_step(model,data); view.sync(); step+=1; time.sleep(0.001)
+    kb.stop(); view.close()
+
+if __name__=="__main__": main()