Backup DreamWaQ rslrl stability fixes

scripts/play_amp_mujoco.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Deploy Go1 AMP model in MuJoCo."""
+import mujoco, numpy as np, os, torch, time, sys
+from mujoco import viewer
+
+# Use AMP's rsl_rl (NOT MotrixLab's)
+sys.path.insert(0, '/home/8x54zj-m/amp_go2/rsl_rl')
+sys.path.insert(0, '/home/8x54zj-m/amp_go2/legged_gym')
+# Force reload of rsl_rl to pick up AMP version
+for mod in list(sys.modules.keys()):
+    if 'rsl_rl' in mod:
+        del sys.modules[mod]
+
+XML_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                       'motrix_envs/src/motrix_envs/locomotion/go1/xmls')
+CKPT = '/home/8x54zj-m/amp_go2/legged_gym/logs/go1_amp/Jul05_01-15-41_flat/model_6000.pt'
+
+import argparse
+p = argparse.ArgumentParser()
+p.add_argument("--terrain", default="flat", choices=["flat", "flat_stairs"])
+args = p.parse_args()
+
+scene = "scene_flat_stairs.xml" if args.terrain == "flat_stairs" else "go1_motor_actuator.xml"
+os.chdir(XML_DIR)
+m = mujoco.MjModel.from_xml_string(open(scene).read())
+d = mujoco.MjData(m)
+
+# Spawn on level 1 platform if using flat_stairs
+spawn_x, spawn_y, spawn_z = 0.0, 0.0, 0.45
+if args.terrain == "flat_stairs":
+    col = np.random.randint(0, 4)
+    spawn_x = -12.0 + col * 8.0
+    spawn_y = -4.0  # level 1 center
+    # Sample terrain height at spawn
+    floor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
+    if floor_id >= 0 and m.geom_type[floor_id] == mujoco.mjtGeom.mjGEOM_HFIELD:
+        hf = m.geom_dataid[floor_id]
+        nr, nc = int(m.hfield_nrow[hf]), int(m.hfield_ncol[hf])
+        sx, sy, zt, sb = m.hfield_size[hf]
+        hd = m.hfield_data[m.hfield_adr[hf]:m.hfield_adr[hf]+nr*nc].reshape(nr, nc)
+        gp = m.geom_pos[floor_id]
+        col_idx = int(np.clip(((spawn_x-gp[0])/sx*0.5+0.5)*(nc-1), 0, nc-1))
+        row_idx = int(np.clip(((spawn_y-gp[1])/sy*0.5+0.5)*(nr-1), 0, nr-1))
+        spawn_z = float(gp[2] + sb + hd[row_idx, col_idx] * zt) + 0.45
+    print(f"[Terrain] flat_stairs level=1 spawn=({spawn_x:.1f},{spawn_y:.1f},{spawn_z:.2f})")
+
+ckpt = torch.load(CKPT, map_location='cpu', weights_only=False)
+normalizer = ckpt['amp_normalizer']
+
+actor = torch.nn.Sequential(
+    torch.nn.Linear(45, 512), torch.nn.ELU(),
+    torch.nn.Linear(512, 256), torch.nn.ELU(),
+    torch.nn.Linear(256, 128), torch.nn.ELU(),
+    torch.nn.Linear(128, 12)
+)
+sd = {k.replace('actor.',''): v for k,v in ckpt['model_state_dict'].items() if k.startswith('actor.')}
+actor.load_state_dict(sd)
+actor.eval()
+
+DEFAULT = np.array([-0.1,0.8,-1.5, 0.1,0.8,-1.5, -0.1,1.0,-1.5, 0.1,1.0,-1.5], dtype=np.float32)
+LIN_VEL_SCALE, ANG_VEL_SCALE = 2.0, 0.25
+DOF_POS_SCALE, DOF_VEL_SCALE = 1.0, 0.05
+CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
+KP, KD, ACT_SCALE = 20.0, 0.5, 0.25
+
+def get_gravity(q):
+    qw,qx,qy,qz = q[3],q[0],q[1],q[2]
+    return np.array([2*(-qz*qx+qw*qy), -2*(qz*qy+qw*qx), 1-2*(qw*qw+qz*qz)])
+
+def compute_obs(d, last_action, cmd):
+    obs = np.zeros(45, dtype=np.float32)
+    obs[0:3] = d.qvel[3:6] * ANG_VEL_SCALE
+    obs[3:6] = get_gravity(d.qpos[3:7])
+    obs[6:9] = cmd * CMD_SCALE
+    obs[9:21] = (d.qpos[7:19] - DEFAULT) * DOF_POS_SCALE
+    obs[21:33] = d.qvel[6:18] * DOF_VEL_SCALE
+    obs[33:45] = last_action
+    return obs
+
+d.qpos[7:19] = DEFAULT; d.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; d.qvel[:] = 0
+mujoco.mj_forward(m, d)
+
+view = viewer.launch_passive(m, d)
+cmd = np.array([0.8, 0.0, 0.0], dtype=np.float32)
+last_action = np.zeros(12, dtype=np.float32)
+
+print("AMP Go1 | W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
+
+while view.is_running():
+    if view.is_running():
+        obs = compute_obs(d, last_action, cmd)
+        obs_t = torch.from_numpy(obs).unsqueeze(0).float()
+        obs_norm = obs_t.clone()
+        obs_norm[:, :43] = normalizer.normalize(obs_t[:, :43])
+        with torch.no_grad():
+            act = actor(obs_norm)[0].numpy()
+        act = np.clip(act, -18, 18)
+        last_action = act.copy()
+        target = DEFAULT + act * ACT_SCALE
+        d.ctrl[:] = KP * (target - d.qpos[7:19]) - KD * d.qvel[6:18]
+    mujoco.mj_step(m, d)
+    view.sync()
+    time.sleep(0.001)