#!/usr/bin/env python3 """DreamWaQ rsl_rl play — render the trained ActorCritic_DWAQ policy in NATIVE MotrixSim. Loads the PyTorch checkpoint DIRECTLY (no ONNX). ONNX is only for cross-sim deployment (e.g. MuJoCo sim2sim); the native MotrixSim env runs the torch policy. Deterministic inference: mean CENet code + actor. Usage: uv run scripts/play_dreamwaq_rsl.py # auto-find latest, walk forward uv run scripts/play_dreamwaq_rsl.py --checkpoint runs/.../model_1100.pt --vx 0.5 uv run scripts/play_dreamwaq_rsl.py --vx 0 --num-envs 1 # stand still, single robot """ import argparse, glob, os, sys, time # avoid JAX grabbing GPU memory and starving the MotrixSim (Vulkan) renderer os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") os.environ.setdefault("JAX_PLATFORMS", "cpu") # --terrain / --level / --flat-stairs / --stairs: pick hfield scene (before import). terrain_type = "pyramid" if "--flat-stairs" in sys.argv: terrain_type = "flat_stairs" elif "--stairs" in sys.argv: terrain_type = "stairs" if "--terrain" in sys.argv or "--level" in sys.argv or "--flat-stairs" in sys.argv or "--stairs" in sys.argv: os.environ["DREAMWAQ_TERRAIN"] = terrain_type sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import numpy as np import torch import motrix_envs.locomotion.go1.dreamwaq # noqa: F401 register env from motrix_envs import registry as env_registry from motrix_envs.np.renderer import NpRenderer from motrix_rl.dwaq_rsl.actor_critic_dwaq import ActorCritic_DWAQ PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) NUM_OBS, NUM_PRIV, NUM_HIST, NUM_ACT, CENET_OUT = 45, 235, 5, 12, 19 CLIP_ACT = 23.7 def _iter_of(path): try: return int(os.path.basename(path).split("_")[1].split(".")[0]) except Exception: return -1 def find_latest(): models = glob.glob(os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "rsl_dwaq", "*", "model_*.pt")) if not models: print("[ERROR] no rsl_dwaq checkpoints found"); sys.exit(1) return max(models, key=_iter_of) # highest iteration (flat model_1100 > terrain early models) def main(): p = argparse.ArgumentParser() p.add_argument("--checkpoint", default=None) p.add_argument("--num-envs", type=int, default=4) p.add_argument("--vx", type=float, default=0.5, help="forward velocity command [m/s]") p.add_argument("--vy", type=float, default=0.0, help="lateral velocity command [m/s]") p.add_argument("--wz", type=float, default=0.0, help="yaw rate command [rad/s]") p.add_argument("--terrain", action="store_true", help="view the training pyramid hfield (default: flat plane)") p.add_argument("--flat-stairs", action="store_true", help="view the 2-level flat+stairs terrain (implies terrain)") p.add_argument("--stairs", action="store_true", help="view the stairs terrain scene (implies terrain)") p.add_argument("--level", type=int, default=None, help="force ALL spawns at this terrain level (implies --terrain)") p.add_argument("--spawn-height", type=float, default=None, help="spawn clearance above terrain in meters (default 0.45; try 1-2 to experiment)") args = p.parse_args() ckpt = args.checkpoint or find_latest() ac = ActorCritic_DWAQ(NUM_OBS + CENET_OUT, NUM_PRIV, NUM_ACT, NUM_HIST * NUM_OBS, CENET_OUT) ac.load_state_dict(torch.load(ckpt, map_location="cpu")["model_state_dict"]) ac.eval() print(f"[Play-rsl] policy (native torch): {ckpt}") env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs) if args.level is not None: env._force_level = args.level # pin all spawns to this level (read in reset) print(f"[Play-rsl] forcing ALL spawns at terrain level {args.level}") if args.spawn_height is not None: env._spawn_absolute = args.spawn_height # absolute world z, no offset print(f"[Play-rsl] spawn absolute z = {args.spawn_height}m") renderer = NpRenderer(env) env.init_state() n = env._num_envs cmd = np.array([args.vx, args.vy, args.wz], dtype=np.float32) print(f"[Play-rsl] {n} envs | cmd=(vx={args.vx}, vy={args.vy}, wz={args.wz}) | Ctrl+C to stop") @torch.no_grad() def act_fn(obs, hist): obs_t = torch.from_numpy(obs) h = ac.encoder(torch.from_numpy(hist).reshape(obs.shape[0], -1)) # (n,225)->(n,64) code = torch.cat([ac.encode_mean_vel(h), ac.encode_mean_latent(h)], dim=-1) # (n,19) return ac.actor(torch.cat([code, obs_t], dim=-1)).numpy() # (n,12) from motrixsim.render import RenderClosedError show_heights = False try: while True: if renderer._render.input.is_key_just_pressed("r"): env.init_state() print("[R] Reset all envs") if renderer._render.input.is_key_just_pressed("h"): show_heights = not show_heights print(f"[H] Height points: {'ON' if show_heights else 'OFF'}") env._state.info["commands"][:] = cmd obs = env._state.obs.astype(np.float32) hist = env._state.info.get("obs_history", np.zeros((n, NUM_HIST, NUM_OBS), np.float32)).astype(np.float32) act = act_fn(obs, hist) env.step(np.clip(act, -CLIP_ACT, CLIP_ACT).astype(np.float32)) if show_heights: from motrix_envs.math import quaternion pose = env._body.get_pose(env._state.data) bp = pose[0, :3] yaw = quaternion.get_yaw(pose[0:1, 3:7])[0] cos_y, sin_y = np.cos(yaw), np.sin(yaw) for gy in env._hy: for gx in env._hx: wx = bp[0] + cos_y*gx - sin_y*gy wy = bp[1] + sin_y*gx + cos_y*gy wz = float(env._sample_terrain_height(np.array([[wx,wy]]))[0]) g = renderer._render.gizmos g.draw_sphere(0.02, (np.float32(wx), np.float32(wy), np.float32(wz))) renderer.render() time.sleep(0.01) except (KeyboardInterrupt, RenderClosedError): pass try: renderer.close() except Exception: pass print("[Play-rsl] done") if __name__ == "__main__": main()