Backup DreamWaQ rslrl stability fixes

scripts/test_abrupt_stop.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""无头 MuJoCo 测试：前进 → 突然停止 → 观察姿态变化。"""
+import numpy as np
+import mujoco
+import onnxruntime as ort
+import os, sys, time
+
+_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
+ONNX = os.path.join(_PROJECT, "runs/go1-dreamwaq-walk/rslrl/26-07-02_20-03-23-_36054_PPO/policy.onnx")
+
+NUM_OBS = 45
+NUM_ACTIONS = 12
+HISTORY_LEN = 5
+ACTION_SCALE = 0.25
+KP = 28.0
+KD = 0.7
+DECIMATION = 4
+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)
+
+
+def compute_obs(model, data, commands, last_action):
+    obs = np.zeros(NUM_OBS, dtype=np.float32)
+    g = None  # no gyro sensor in headless
+    obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25
+    grav_world = model.opt.gravity.copy()
+    grav_world = grav_world / np.linalg.norm(grav_world)
+    R = data.xmat[1].reshape(3, 3)
+    obs[3:6] = (R.T @ grav_world).astype(np.float32)
+    obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=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_action
+    return obs
+
+
+def get_base_pose(data):
+    """返回 base 的 z 高度和 roll/pitch（度）。"""
+    quat = data.xquat[1]  # base body quaternion
+    w, x, y, z = quat[0], quat[1], quat[2], quat[3]
+    # roll, pitch from quaternion
+    sinr_cosp = 2 * (w * x + y * z)
+    cosr_cosp = 1 - 2 * (x * x + y * y)
+    roll = np.arctan2(sinr_cosp, cosr_cosp)
+    sinp = 2 * (w * y - z * x)
+    pitch = np.arcsin(np.clip(sinp, -1, 1))
+    return float(data.xpos[1, 2]), np.degrees(roll), np.degrees(pitch)
+
+
+def main():
+    session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
+    xml_file = os.path.join(XML_DIR, "scene_dreamwaq_flat.xml")
+    model = mujoco.MjModel.from_xml_path(xml_file)
+    data = mujoco.MjData(model)
+
+    # 初始化
+    data.qpos[7:19] = DEFAULT_ANGLES
+    data.qpos[2] = 0.35
+    mujoco.mj_forward(model, data)
+
+    history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
+    last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
+
+    print(f"{'Step':>5s}  {'Cmd_vx':>7s}  {'base_z':>8s}  {'roll':>8s}  {'pitch':>8s}  {'speed_xy':>8s}")
+    print("-" * 60)
+
+    for step in range(3000):
+        # 命令：前 1500 步前进，之后突然停止
+        if step < 1500:
+            vx, vy, wz = 0.5, 0.0, 0.0  # 前进
+        else:
+            vx, vy, wz = 0.0, 0.0, 0.0  # 突然停止！
+
+        for _ in range(DECIMATION):
+            mujoco.mj_step(model, data)
+
+        if step % DECIMATION == 0:
+            cmd = np.array([vx, vy, wz], dtype=np.float32)
+            obs = compute_obs(model, data, cmd, last_action)
+            history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
+            outputs = session.run(None, {
+                'obs': obs.reshape(1, -1).astype(np.float32),
+                'obs_history': history.reshape(1, -1).astype(np.float32),
+            })
+            action = outputs[0][0]
+            last_action = action
+
+        target = DEFAULT_ANGLES + action * ACTION_SCALE
+        data.ctrl[:] = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
+
+        # 每 100 步打印
+        if step % 100 == 0:
+            bz, roll, pitch = get_base_pose(data)
+            speed_xy = np.linalg.norm(data.qvel[0:2])
+            print(f"{step:5d}  {vx:7.1f}  {bz:8.3f}  {roll:+8.1f}  {pitch:+8.1f}  {speed_xy:8.3f}")
+
+    # 最终状态
+    bz, roll, pitch = get_base_pose(data)
+    print(f"\n最终: base_z={bz:.3f} roll={roll:.1f}° pitch={pitch:.1f}°")
+    if abs(roll) > 30 or abs(pitch) > 30:
+        print("⚠ 机器人倾覆！突然停止导致翻跟头")
+    else:
+        print("✅ 机器人在突然停止后保持稳定")
+
+
+if __name__ == "__main__":
+    main()