| """Rollout recorder: render a real video of a policy + dump a per-step |
| duty/altitude CSV, so we can SEE how the robot falls before changing rewards. |
| |
| Unlike verify_walk.py (which deadlocks on headless GL) this renders the same |
| way training does -- under xvfb-run with MUJOCO_GL=glfw. Run it wrapped: |
| |
| xvfb-run -a uv run --no-sync python scripts/record_rollout.py \ |
| --checkpoint <pt> --out /tmp/opencode/rollout.mp4 |
| """ |
| import argparse, csv, json, os |
| import torch |
| from dataclasses import asdict |
|
|
| import sys |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) |
|
|
| 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 |
|
|
| TASK = "Mjlab-Velocity-Flat-TwoLeg" |
| DUTY_FLOOR = 0.20 |
| FRAME_ALT_CEIL = -0.20 |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--checkpoint", required=True) |
| ap.add_argument("--out", default="/tmp/opencode/rollout.mp4") |
| ap.add_argument("--command", type=float, default=0.1) |
| ap.add_argument("--steps", type=int, default=300) |
| ap.add_argument("--envs", type=int, default=1) |
| ap.add_argument("--device", default="cuda:0") |
| args = ap.parse_args() |
|
|
| 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="rgb_array") |
| 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"] |
| robot.data.forward_vec_b[:] = 0.0 |
| robot.data.forward_vec_b[:, 1] = -1.0 |
| twist = env.command_manager.get_term("twist") |
| step_dt = env.step_dt |
| obs, _ = venv.reset() |
| for _ in range(50): |
| cmd = torch.zeros_like(twist.command) |
| cmd[:, 1] = -args.command |
| twist.command[:] = cmd |
| with torch.no_grad(): |
| obs, _, _, _ = venv.step(policy(obs)) |
|
|
| import imageio.v2 as imageio |
| writer = imageio.get_writer(args.out, fps=int(1 / step_dt), |
| macro_block_size=1) |
| rows = [] |
| for s in range(args.steps): |
| cmd = torch.zeros_like(twist.command) |
| cmd[:, 1] = -args.command |
| twist.command[:] = cmd |
| with torch.no_grad(): |
| obs, _, _, _ = venv.step(policy(obs)) |
| img = env.render() |
| frame = img[0] if img.ndim == 4 else img |
| writer.append_data(frame) |
| found = env.scene["feet_ground_contact"].data.found[:] |
| found = found.squeeze(-1) if found.ndim == 4 else found |
| up = robot.data.projected_gravity_b[:, 2] |
| |
| rows.append({ |
| "step": s, |
| "left_contact": float(found[0, 0].item()), |
| "right_contact": float(found[0, 1].item()), |
| "upright": float(up[0].item()), |
| "torso_z": float(robot.data.root_link_pos_w[0, 2].item()), |
| }) |
| writer.close() |
|
|
| csv_path = os.path.splitext(args.out)[0] + ".csv" |
| with open(csv_path, "w", newline="") as fh: |
| w = csv.DictWriter(fh, fieldnames=["step", "left_contact", |
| "right_contact", "upright", "torso_z"]) |
| w.writeheader() |
| w.writerows(rows) |
| |
| first_fail = next((r["step"] for r in rows |
| if r["upright"] <= FRAME_ALT_CEIL), None) |
| print(json.dumps({ |
| "video": args.out, "csv": csv_path, |
| "first_fall_step": first_fail, |
| "torso_z_start": rows[0]["torso_z"], |
| "torso_z_end": rows[-1]["torso_z"], |
| "upright_end": rows[-1]["upright"], |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
