File size: 4,253 Bytes
1851ea6 | 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 105 106 107 108 109 | """eval_sprint_speed.py — 在 warp 训练环境内直接测 checkpoint 的真实速度。
绕开本地 CPU harness 的所有差异项（BAM/延迟/噪声/归一化烘焙），用训练同款
环境 + rsl_rl 推理策略（含 obs normalizer）回答：策略在原生环境里跑多快？
用法（实例上，工作目录 = repo 根）：
uv run --no-sync python scripts/eval_sprint_speed.py \
--checkpoint logs/rsl_rl/sprint/<run>/model_1999.pt --vx 0.4 1.2 2.0
"""
from __future__ import annotations
import argparse
from dataclasses import asdict
import numpy as np
import torch
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--task", default="Mjlab-Sprint-Flat-MicroDuck")
ap.add_argument("--checkpoint", default=None)
ap.add_argument("--onnx", default=None, help="测 ONNX（归一化已烘焙）而非 .pt checkpoint")
ap.add_argument("--vx", type=float, nargs="+", default=[0.4, 0.8, 1.2, 1.6, 2.0])
ap.add_argument("--num-envs", type=int, default=64)
ap.add_argument("--seconds", type=float, default=8.0)
ap.add_argument("--device", default="cuda:0")
args = ap.parse_args()
import mjlab.tasks # noqa: F401 — populate the registry
import mjlab_microduck.tasks # noqa: F401 — microduck registrations
from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.envs import ManagerBasedRlEnv
from mjlab.rl import RslRlVecEnvWrapper, MjlabOnPolicyRunner
env_cfg = load_env_cfg(args.task)
agent_cfg = load_rl_cfg(args.task)
env_cfg.scene.num_envs = args.num_envs
env = ManagerBasedRlEnv(cfg=env_cfg, device=args.device)
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
if args.onnx:
import onnxruntime as ort
session = ort.InferenceSession(args.onnx)
def policy(obs):
vec = obs["actor"].detach().cpu().numpy().astype(np.float32)
# 导出的 ONNX batch 维固定为 1，逐 env 推理
acts = [session.run(None, {session.get_inputs()[0].name: vec[i : i + 1]})[0]
for i in range(vec.shape[0])]
return torch.tensor(np.concatenate(acts, axis=0),
device=args.device, dtype=torch.float32)
else:
assert args.checkpoint, "--checkpoint or --onnx required"
runner_cls = load_runner_cls(args.task) or MjlabOnPolicyRunner
runner = runner_cls(env, 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)
uw = env.unwrapped
robot = uw.scene["robot"]
cmd_term = uw.command_manager.get_term("twist")
dt = uw.step_dt
print(f"policy: {args.onnx or args.checkpoint}")
print(f"control dt: {dt:.4f}s envs: {args.num_envs} task: {args.task}")
def force_cmd(vx):
cmd_term.vel_command_b[:, 0] = vx
cmd_term.vel_command_b[:, 1] = 0.0
cmd_term.vel_command_b[:, 2] = 0.0
cmd_term.is_standing_env[:] = False
obs, _ = env.reset()
for vx in args.vx:
obs, _ = env.reset()
force_cmd(vx)
# 2s 加速/稳定段
for _ in range(int(2.0 / dt)):
force_cmd(vx)
actions = policy(obs)
obs, _, _, _ = env.step(actions)
x0 = robot.data.root_link_pos_w[:, 0].clone()
y0 = robot.data.root_link_pos_w[:, 1].clone()
n_meas = int(args.seconds / dt)
dones_total = 0
for _ in range(n_meas):
force_cmd(vx)
actions = policy(obs)
obs, _, dones, _ = env.step(actions)
dones_total += int(dones.sum())
x1 = robot.data.root_link_pos_w[:, 0]
y1 = robot.data.root_link_pos_w[:, 1]
dist = torch.hypot(x1 - x0, y1 - y0)
speed = dist / args.seconds
alive = (dones_total == 0)
print(f"cmd={vx:4.1f} 实测 speed mean={speed.mean():.3f} "
f"p10={speed.quantile(0.1):.3f} p90={speed.quantile(0.9):.3f} "
f"max={speed.max():.3f} m/s resets={dones_total}")
env.close()
if __name__ == "__main__":
main()
|