Blame for scripts/render_all_readme_media.py

#!/usr/bin/env python3
"""Batch renders all showcase videos directly to docs/media/ at natural 1x speed:
1. 5-Duck Olympic Stadium Race (1x original playback, 10s competition)
2. 100g Barbell Physical Clamping & Dynamic Walk
3. 200g Barbell Physical Clamping & Dynamic Walk
4. 300g Barbell Physical Clamping & Dynamic Walk
5. Table Tennis Racket Carry (with Steadicam leveling)
6. Badminton Racket Carry (BWF official scaled spec)
7. Pickleball Paddle Carry
8. Lemon & Tableware Spoon Dynamic Balance
"""

import os
import sys
import math
import numpy as np
import mujoco
import imageio

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_MANIP_DIR = os.path.dirname(_SCRIPT_DIR)
_DOCS_MEDIA = os.path.join(_MANIP_DIR, "docs", "media")
os.makedirs(_DOCS_MEDIA, exist_ok=True)

sys.path.insert(0, _SCRIPT_DIR)

from infer_policy import PolicyInference
from test_racket_drop_and_walk import (
    create_scaled_scene_xml,
    TableTennisPolicyInference,
    TableTennisRacketCarryPolicyInference,
    _DEFAULT_POLICY,
    _TT_POLICY,
    _TT_CARRY_POLICY,
)
from test_physical_clamp_policy import setup_physical_clamp_sim, PHYSICAL_CLAMP_XML
import test_duck_racer_lemon_balance as racer
import test_spoon_lemon_balance as spoon_sim


def render_to_mp4_and_webp(frames, output_base, fps=30):
    """Saves frame list as MP4 and WebP in docs/media/."""
    mp4_path = os.path.join(_DOCS_MEDIA, f"{output_base}.mp4")
    webp_path = os.path.join(_DOCS_MEDIA, f"{output_base}.webp")

    writer = imageio.get_writer(mp4_path, fps=fps, codec="libx264", quality=8)
    for frame in frames:
        writer.append_data(frame)
    writer.close()
    print(f"  [SAVED] {mp4_path} ({len(frames)} frames @ {fps}fps)")

    # WebP compression for GitHub display
    cmd = f'ffmpeg -y -i "{mp4_path}" -vf "fps=20,scale=540:-1:flags=lanczos" -vcodec libwebp -lossless 0 -compression_level 4 -q:v 70 -loop 0 "{webp_path}" >/dev/null 2>&1'
    os.system(cmd)
    print(f"  [SAVED] {webp_path}")


def render_duck_race_1x():
    """Renders 5-Duck Olympic Race at natural 1x speed."""
    print("\n[1/8] Rendering 5-Duck Olympic Race (1x natural playback)...")
    xml_str = racer.create_racer_scene_xml()
    model = mujoco.MjModel.from_xml_string(xml_str)
    model.opt.timestep = 0.002
    data = mujoco.MjData(model)

    session = racer.ort.InferenceSession(racer._DEFAULT_POLICY, providers=["CPUExecutionProvider"])
    ducks = [racer.DuckRacer(cfg["id"], cfg, model, data, session) for cfg in racer.DUCK_CONFIGS]
    racer.setup_racer_episode(ducks, model, data)

    # Physical pacing:
    # 1st: Emerald Green (0.175), 2nd: Ruby Coral Pink (0.165), 3rd: Duck Gold (0.155),
    # 4th: Amethyst Purple (0.148), 5th: Riviera Sky Blue (0.140)
    pacing = [0.165, 0.140, 0.155, 0.175, 0.148]
    for i, d in enumerate(ducks):
        d.target_vx = pacing[i]

    renderer = mujoco.Renderer(model, width=1280, height=720)
    vopt = mujoco.MjvOption()
    vopt.geomgroup[3] = 0

    cam = mujoco.MjvCamera()
    cam.type = mujoco.mjtCamera.mjCAMERA_FREE
    cam.distance = 2.15
    cam.elevation = -22.0
    cam.azimuth = 68.0

    duration_sec = 10.0
    fps = 30
    total_steps = int(duration_sec / model.opt.timestep)
    steps_per_frame = int(round(1.0 / (fps * model.opt.timestep)))
    decimation = 10
    smooth_x = 0.45

    frames = []
    for step in range(total_steps):
        t_sim = step * model.opt.timestep
        for d in ducks:
            d.step_clamp()

        if step % decimation == 0:
            for d in ducks:
                d.infer_and_apply(decimation, model.opt.timestep)
                d.update_telemetry(t_sim)

        mujoco.mj_step(model, data)

        if step % steps_per_frame == 0:
            xs = [d.distance for d in ducks]
            pack_mean_x = float(np.mean(xs))
            target_x = pack_mean_x + 0.45
            smooth_x = 0.95 * smooth_x + 0.05 * target_x
            cam.lookat = [smooth_x, 0.0, 0.16]

            renderer.update_scene(data, camera=cam, scene_option=vopt)
            frames.append(renderer.render())

    render_to_mp4_and_webp(frames, "duck_racer_competition_1x", fps=fps)
    render_to_mp4_and_webp(frames, "duck_racer_competition_2x", fps=fps * 2)


def render_barbell(mass_g, torque_nm, vel_x, output_name):
    """Renders Barbell carry at specified mass."""
    print(f"\nRendering {mass_g:.0f}g Barbell Carry...")
    model, data, policy, meta = setup_physical_clamp_sim(
        xml_path=PHYSICAL_CLAMP_XML,
        policy_path=os.path.join(_MANIP_DIR, "policies/onnx/payload_clamping_locomotion.onnx"),
        payload_mass_g=mass_g,
        vel_x=vel_x,
        clamp_torque_nm=torque_nm,
    )

    renderer = mujoco.Renderer(model, height=720, width=1280)
    camera = mujoco.MjvCamera()
    camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING
    camera.trackbodyid = 1
    camera.distance = 0.62
    camera.elevation = -8.0
    camera.azimuth = 125.0

    frames = []
    beak_dof = meta["beak_dof"]
    num_steps = 300  # 300 policy steps * 0.02s = 6.0s walk (150 frames @ 30fps = 5.0s video, 100% longer)
    fps = 30

    for step in range(num_steps):
        action = policy.infer()
        policy.apply_action(action)

        for _ in range(4):
            data.qfrc_applied[beak_dof] = -torque_nm
            mujoco.mj_step(model, data)

        if step % 2 == 0:
            renderer.update_scene(data, camera=camera)
            frames.append(renderer.render())

    render_to_mp4_and_webp(frames, output_name, fps=fps)


def render_racket(racket_type, output_name):
    """Renders racket carry for table tennis, badminton, or pickleball."""
    print(f"\nRendering {racket_type} Carry...")
    is_tt = (racket_type == "table_tennis")
    xml_str = create_scaled_scene_xml(racket_type=racket_type, include_ball=False)
    model = mujoco.MjModel.from_xml_string(xml_str)
    model.opt.timestep = 0.002
    data = mujoco.MjData(model)

    policy_path = _TT_CARRY_POLICY if is_tt and os.path.exists(_TT_CARRY_POLICY) else _DEFAULT_POLICY
    import onnxruntime as ort
    sess = ort.InferenceSession(policy_path, providers=["CPUExecutionProvider"])
    input_dim = sess.get_inputs()[0].shape[1]

    if input_dim == 64 and is_tt:
        policy = TableTennisRacketCarryPolicyInference(model, data, onnx_path=policy_path)
    else:
        policy = PolicyInference(
            model, data,
            walking_onnx_path=policy_path,
            action_scale=1.0,
            new_cmd_obs=True,
            use_projected_gravity=True,
        )

    policy.set_vel_cmd(0.18, 0.0, 0.0)
    policy.head_offset[:] = 0.0
    policy._update_command()

    # Reset
    mujoco.mj_resetData(model, data)
    data.qpos[0:3] = [0.0, 0.0, 0.125]
    data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
    for i, idx in enumerate(policy.joint_qpos_indices):
        data.qpos[idx] = policy.default_pose[i]

    racket_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "racket")
    racket_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "racket_free")
    beak_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "beak")
    beak_dof = model.jnt_dofadr[beak_jid]
    beak_qpos = model.jnt_qposadr[beak_jid]
    jaw_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "jaw_soft")
    trunk_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk_base")

    target_clamp_angle = 1.637 if is_tt else (1.654 if racket_type == "badminton" else 1.666)
    clamp_torque = 0.50
    data.qpos[beak_qpos] = target_clamp_angle
    data.qfrc_applied[beak_dof] = -clamp_torque
    if model.nu > 14:
        data.ctrl[14] = 1.0

    mujoco.mj_forward(model, data)
    p_jaw = data.xpos[jaw_bid]
    R_jaw = data.xmat[jaw_bid].reshape(3, 3)
    HOVER_LOCAL = np.array([-0.0120, 0.0, -0.0680])
    CYL_QUAT = np.array([0.0, 0.7071068, 0.7071068, 0.0]) if is_tt else np.array([0.7071068, 0.7071068, 0.0, 0.0])

    r_adr = model.jnt_qposadr[racket_jid]
    data.qpos[r_adr:r_adr+3] = p_jaw + R_jaw @ HOVER_LOCAL
    q_jaw = np.empty(4)
    mujoco.mju_mat2Quat(q_jaw, data.xmat[jaw_bid])
    q_r = np.empty(4)
    mujoco.mju_mulQuat(q_r, q_jaw, CYL_QUAT)
    data.qpos[r_adr+3:r_adr+7] = q_r

    for _ in range(60):
        data.qfrc_applied[beak_dof] = -clamp_torque
        mujoco.mj_step(model, data)

    renderer = mujoco.Renderer(model, width=1280, height=720)
    vopt = mujoco.MjvOption()
    vopt.geomgroup[3] = 0

    cam = mujoco.MjvCamera()
    cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
    cam.trackbodyid = 1
    cam.distance = 0.72
    cam.elevation = -14.0
    cam.azimuth = 138.0

    duration_s = 3.44 if racket_type in ["badminton", "pickleball"] else 6.87
    total_steps = int(duration_s / model.opt.timestep)
    fps = 30
    steps_per_frame = int(round(1.0 / (fps * model.opt.timestep)))
    decimation = 10

    frames = []
    steadicam_pos = 0.0

    for step in range(total_steps):
        beak_err = target_clamp_angle - data.qpos[beak_qpos]
        data.qfrc_applied[beak_dof] = float(np.clip(-clamp_torque + 12.0 * beak_err - 0.5 * data.qvel[beak_dof], -1.5, 0.2))

        if step % decimation == 0:
            action = policy.infer()

            if is_tt:
                R_trunk = data.xmat[trunk_bid].reshape(3, 3)
                h_fwd = np.array([R_trunk[0, 0], R_trunk[1, 0], 0.0])
                h_fwd /= (np.linalg.norm(h_fwd) + 1e-6)
                R_r = data.xmat[racket_bid].reshape(3, 3)
                b_norm = R_r @ np.array([0.0, 1.0, 0.0])
                racket_pitch = math.asin(np.clip(float(np.dot(b_norm, h_fwd)), -0.999, 0.999))

                rate_cmd = np.clip(1.20 * racket_pitch, -0.30, 0.30)
                steadicam_pos = np.clip(steadicam_pos + rate_cmd * 0.02, -0.20, 0.20)
                action[5] += steadicam_pos
                action[6] -= steadicam_pos

            policy.apply_action(action)
            if model.nu > 14:
                data.ctrl[14] = 1.0

        mujoco.mj_step(model, data)

        if step % steps_per_frame == 0:
            renderer.update_scene(data, camera=cam, scene_option=vopt)
            frames.append(renderer.render())

    render_to_mp4_and_webp(frames, output_name, fps=fps)


def render_lemon_spoon():
    """Renders Lemon & Tableware Spoon dynamic balancing."""
    print("\nRendering Lemon & Spoon Balance...")
    mp4_path = os.path.join(_DOCS_MEDIA, "lemon_spoon_balance.mp4")
    webp_path = os.path.join(_DOCS_MEDIA, "lemon_spoon_balance.webp")
    spoon_sim.render_walking_video(output_path=mp4_path, duration_sec=7.0, vx=0.16)
    cmd = f'ffmpeg -y -i "{mp4_path}" -vf "fps=20,scale=540:-1:flags=lanczos" -vcodec libwebp -lossless 0 -compression_level 4 -q:v 70 -loop 0 "{webp_path}" >/dev/null 2>&1'
    os.system(cmd)
    print(f"  [SAVED] {webp_path}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Render all showcase videos")
    parser.add_argument("--force", action="store_true", help="Force re-rendering all videos")
    args = parser.parse_args()

    print(f"Target media directory: {_DOCS_MEDIA}")
    def exists(name):
        return (not args.force) and os.path.exists(os.path.join(_DOCS_MEDIA, f"{name}.webp"))

    if not exists("duck_racer_competition_1x"):
        render_duck_race_1x()
    else:
        print("  [EXISTS] duck_racer_competition_1x")

    if not exists("barbell_100g_carry"):
        render_barbell(100.0, 0.10, 0.25, "barbell_100g_carry")
    else:
        print("  [EXISTS] barbell_100g_carry")

    if not exists("barbell_200g_carry"):
        render_barbell(200.0, 0.25, 0.20, "barbell_200g_carry")
    else:
        print("  [EXISTS] barbell_200g_carry")

    if not exists("barbell_300g_carry"):
        render_barbell(300.0, 0.50, 0.15, "barbell_300g_carry")
    else:
        print("  [EXISTS] barbell_300g_carry")

    if not exists("table_tennis_carry"):
        render_racket("table_tennis", "table_tennis_carry")
    else:
        print("  [EXISTS] table_tennis_carry")

    if not exists("badminton_carry"):
        render_racket("badminton", "badminton_carry")
    else:
        print("  [EXISTS] badminton_carry")

    if not exists("pickleball_carry"):
        render_racket("pickleball", "pickleball_carry")
    else:
        print("  [EXISTS] pickleball_carry")

    if not exists("lemon_spoon_balance"):
        render_lemon_spoon()
    else:
        print("  [EXISTS] lemon_spoon_balance")

    print("\nAll README media rendered successfully into docs/media/!")