File size: 4,368 Bytes
24fe873 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Measure what a TwoLeg checkpoint actually does, in metres.
Reward and curriculum terms are not speed. This loads a checkpoint, holds a
fixed forward command, rolls out, and reports the base's world displacement and
mean velocity. Same posture as microduck_rl's scripts/eval_sprint_speed.py.
uv run python training/scripts/eval_walk.py --checkpoint <path/to/model.pt>
uv run python training/scripts/eval_walk.py --checkpoint <pt> --command 0.3 --steps 500
"""
import argparse
import os
from dataclasses import asdict
import torch
TASK = "Mjlab-Velocity-Flat-TwoLeg"
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--checkpoint", required=True)
ap.add_argument("--command", type=float, default=0.3, help="forward m/s to hold")
ap.add_argument("--steps", type=int, default=500, help="control steps (50 Hz)")
ap.add_argument("--envs", type=int, default=16)
ap.add_argument("--device", default="cuda:0")
args = ap.parse_args()
from mjlab.envs import ManagerBasedRlEnv
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls
env_cfg = load_env_cfg(TASK, play=True)
env_cfg.scene.num_envs = args.envs
agent_cfg = load_rl_cfg(TASK)
env = ManagerBasedRlEnv(cfg=env_cfg, device=args.device, render_mode=None)
venv = RslRlVecEnvWrapper(env)
runner_cls = load_runner_cls(TASK) or MjlabOnPolicyRunner
runner = runner_cls(venv, asdict(agent_cfg), device=args.device)
runner.load(args.checkpoint, load_cfg={"actor": True}, strict=True, map_location=args.device)
policy = runner.get_inference_policy(device=args.device)
robot = env.scene["robot"]
# Global heading metric: the model's forward is body +y here, so reset
# mjlab's +x default to avoid measuring side-sway as forward travel.
robot.data.forward_vec_b[:] = 0.0
robot.data.forward_vec_b[:, 1] = -1.0
twist = env.command_manager.get_term("twist")
import math
obs, _ = venv.reset()
start = robot.data.root_link_pos_w[:, :2].clone()
# NOTE: each env resets to its own heading, so world-frame x averages out
# across envs and a signed world-x travel reads near zero even when every
# env walks forward. Integrate each step's displacement against that step's
# heading instead. (Verified: matches the body-velocity integral to 1e-3.)
forward, P, H = [], [], []
for _ in range(args.steps):
cmd = torch.zeros_like(twist.command)
cmd[:, 1] = -args.command
twist.command[:] = cmd
with torch.no_grad():
actions = policy(obs)
obs, _, _, _ = venv.step(actions)
forward.append(robot.data.root_com_lin_vel_w[:, :2].norm(dim=1).mean().item())
P.append(robot.data.root_link_pos_w[:, :2].clone())
H.append(robot.data.heading_w.clone())
P = torch.stack(P)
H = torch.stack(H)
d = P[1:] - P[:-1]
h = H[:-1]
fwd = torch.stack([torch.cos(h), torch.sin(h)], -1)
travel = (d * fwd).sum(-1).sum(0) # per env, metres along its own axis
elapsed = (P.shape[0] - 1) * env.step_dt
fwd_speed = (travel / elapsed).mean().item()
end = robot.data.root_link_pos_w[:, :2]
displacement = (end - start).norm(dim=1)
mean_speed = displacement.mean().item() / elapsed
print(f"checkpoint : {os.path.basename(args.checkpoint)}")
print(f"held command : {args.command:.2f} m/s forward")
print(f"control steps : {args.steps} ({elapsed:.1f} s at {1/env.step_dt:.0f} Hz)")
print(f"envs : {args.envs}")
print()
print(f"mean |v| xy : {sum(forward)/len(forward):.3f} m/s (speed, direction ignored)")
print(f"net displacement: {displacement.mean().item():.3f} m "
f"(p10 {displacement.quantile(0.1).item():.3f}, p90 {displacement.quantile(0.9).item():.3f})")
print(f"net speed : {mean_speed:.3f} m/s (displacement / time)")
print(f"forward travel : {travel.mean().item():+.3f} m along own axis "
f"(p10 {travel.quantile(0.1).item():+.3f}, p90 {travel.quantile(0.9).item():+.3f})")
print(f"forward speed : {fwd_speed:+.3f} m/s")
print()
verdict = "WALKS" if fwd_speed > 0.05 else "does not translate forward"
print(f"verdict : {verdict}")
if __name__ == "__main__":
main()
|