""" Render Experiment 1: The Indistinguishable Pair Two visually identical cubes pushed with the same force — wildly different outcomes. """ import mujoco import numpy as np import imageio # ============ CONFIGURATION ============ WIDTH, HEIGHT = 1280, 720 FPS = 60 OUTPUT_DIR = "." # Change to desired output path # ============ SCENE ============ XML = """ """.format(w=WIDTH, h=HEIGHT) def render_push(force_magnitude, filename, camera="main", duration_sec=3.0): """Render a push experiment to video.""" model = mujoco.MjModel.from_xml_string(XML) data = mujoco.MjData(model) renderer = mujoco.Renderer(model, height=HEIGHT, width=WIDTH) heavy_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "cube_heavy") light_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "cube_light") cam_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera) mujoco.mj_forward(model, data) # Settle for _ in range(300): mujoco.mj_step(model, data) total_steps = int(duration_sec / model.opt.timestep) frames_per_step = max(1, int(1.0 / (FPS * model.opt.timestep))) force_steps = 40 # ~80ms impulse frames = [] # 0.5s of calm before push calm_steps = int(0.5 / model.opt.timestep) for step in range(calm_steps): mujoco.mj_step(model, data) if step % frames_per_step == 0: renderer.update_scene(data, camera=cam_id) frame = renderer.render() frames.append(frame.copy()) # Push for step in range(total_steps): if step < force_steps: data.xfrc_applied[heavy_id, :3] = [force_magnitude, 0, 0] data.xfrc_applied[light_id, :3] = [force_magnitude, 0, 0] else: data.xfrc_applied[heavy_id, :3] = [0, 0, 0] data.xfrc_applied[light_id, :3] = [0, 0, 0] mujoco.mj_step(model, data) if step % frames_per_step == 0: renderer.update_scene(data, camera=cam_id) frame = renderer.render() frames.append(frame.copy()) # Write video output_path = f"{OUTPUT_DIR}/{filename}" writer = imageio.get_writer(output_path, fps=FPS, codec='libx264', output_params=['-crf', '20', '-preset', 'medium']) for frame in frames: writer.append_data(frame) writer.close() print(f"Saved: {output_path} ({len(frames)} frames, {len(frames)/FPS:.1f}s)") renderer.close() return frames def render_force_sweep(filename, forces=[0.5, 1.0, 2.0, 3.0, 5.0]): """Render multiple forces sequentially in one video with title cards.""" model = mujoco.MjModel.from_xml_string(XML) renderer = mujoco.Renderer(model, height=HEIGHT, width=WIDTH) all_frames = [] for force in forces: data = mujoco.MjData(model) mujoco.mj_forward(model, data) heavy_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "cube_heavy") light_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "cube_light") cam_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, "main") # Settle for _ in range(300): mujoco.mj_step(model, data) frames_per_step = max(1, int(1.0 / (FPS * model.opt.timestep))) force_steps = 40 total_steps = int(2.5 / model.opt.timestep) # Brief pause showing initial state for _ in range(int(0.3 * FPS)): renderer.update_scene(data, camera=cam_id) frame = renderer.render() all_frames.append(frame.copy()) # Push and record for step in range(total_steps): if step < force_steps: data.xfrc_applied[heavy_id, :3] = [force, 0, 0] data.xfrc_applied[light_id, :3] = [force, 0, 0] else: data.xfrc_applied[heavy_id, :3] = [0, 0, 0] data.xfrc_applied[light_id, :3] = [0, 0, 0] mujoco.mj_step(model, data) if step % frames_per_step == 0: renderer.update_scene(data, camera=cam_id) frame = renderer.render() all_frames.append(frame.copy()) print(f" Force {force}N: {len(all_frames)} total frames so far") output_path = f"{OUTPUT_DIR}/{filename}" writer = imageio.get_writer(output_path, fps=FPS, codec='libx264', output_params=['-crf', '20', '-preset', 'medium']) for frame in all_frames: writer.append_data(frame) writer.close() print(f"Saved: {output_path} ({len(all_frames)} frames, {len(all_frames)/FPS:.1f}s)") renderer.close() if __name__ == "__main__": print("=== Rendering Experiment 1: The Indistinguishable Pair ===\n") print("--- Single push at 2N ---") render_push(2.0, "exp1_push_2N.mp4", camera="main") print("\n--- Single push at 5N ---") render_push(5.0, "exp1_push_5N.mp4", camera="main") print("\n--- Top-down view at 5N ---") render_push(5.0, "exp1_push_5N_top.mp4", camera="top") print("\n--- Force sweep ---") render_force_sweep("exp1_force_sweep.mp4") print("\nDone!")