142 lines
6.3 KiB
Python
142 lines
6.3 KiB
Python
#!/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()
|