| |
| """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)") |
|
|
| |
| 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) |
|
|
| |
| |
| |
| 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 |
| 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() |
|
|
| |
| 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/!") |
|
|
