File size: 4,108 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 105 106 107 108 109 | """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)
# Render context (glfw) -- needs DISPLAY, so run under xvfb-run.
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 # [envs, slots]
up = robot.data.projected_gravity_b[:, 2] # [envs]
# env 0 only (single-env record)
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 (fall) step
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()
|