WorldRepresentation / render_scripts /render_exp1_push.py
Nirav Madhani
Import World Representation Space
5af572e
Raw
History Blame Contribute Delete
8.05 kB
"""
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 = """
<mujoco model="indistinguishable_pair">
<option timestep="0.002" gravity="0 0 -9.81"/>
<visual>
<global offwidth="{w}" offheight="{h}"/>
<quality shadowsize="4096"/>
<headlight ambient="0.3 0.3 0.3" diffuse="0.6 0.6 0.6"/>
</visual>
<asset>
<texture name="grid" type="2d" builtin="checker" rgb1="0.35 0.35 0.38" rgb2="0.30 0.30 0.33"
width="256" height="256"/>
<material name="grid_mat" texture="grid" texrepeat="8 8" reflectance="0.1"/>
<texture name="wood" type="2d" builtin="flat" rgb1="0.45 0.35 0.25" width="1" height="1"/>
<material name="wood_mat" texture="wood" specular="0.1"/>
<material name="red_mat" rgba="0.85 0.18 0.18 1" specular="0.4" shininess="0.6"/>
<material name="pusher_mat" rgba="0.3 0.35 0.7 0.8" specular="0.3"/>
</asset>
<worldbody>
<light pos="0 -0.5 1.5" dir="0 0.3 -1" castshadow="true" diffuse="0.8 0.8 0.8"/>
<light pos="0.5 0.5 1.2" dir="-0.3 -0.3 -1" castshadow="false" diffuse="0.3 0.3 0.35"/>
<geom name="floor" type="plane" size="2 2 0.1" material="grid_mat"/>
<!-- Table -->
<body name="table" pos="0 0 0.38">
<geom type="box" size="0.45 0.25 0.02" material="wood_mat"/>
<!-- Legs -->
<geom type="cylinder" size="0.015 0.18" pos="-0.40 -0.20 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos=" 0.40 -0.20 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos="-0.40 0.20 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos=" 0.40 0.20 -0.2" rgba="0.4 0.3 0.2 1"/>
</body>
<!-- Label markers on table surface -->
<body name="label_heavy" pos="-0.15 0.15 0.405">
<geom type="box" size="0.06 0.008 0.001" rgba="0.2 0.3 0.7 0.6"/>
</body>
<body name="label_light" pos="0.15 0.15 0.405">
<geom type="box" size="0.06 0.008 0.001" rgba="0.7 0.3 0.1 0.6"/>
</body>
<!-- Heavy cube: steel, 2kg -->
<body name="cube_heavy" pos="-0.15 0 0.435">
<joint type="free"/>
<geom type="box" size="0.03 0.03 0.03" mass="2.0" material="red_mat"
friction="0.5 0.005 0.0001"/>
</body>
<!-- Light cube: foam, 0.05kg — VISUALLY IDENTICAL -->
<body name="cube_light" pos="0.15 0 0.435">
<joint type="free"/>
<geom type="box" size="0.03 0.03 0.03" mass="0.05" material="red_mat"
friction="0.3 0.005 0.0001"/>
</body>
<!-- Camera -->
<camera name="main" pos="0 -0.65 0.75" xyaxes="1 0 0 0 0.55 0.84" fovy="50"/>
<camera name="top" pos="0 0 1.3" xyaxes="1 0 0 0 1 0" fovy="50"/>
<camera name="side" pos="0.8 0 0.55" xyaxes="0 1 0 -0.5 0 0.87" fovy="50"/>
</worldbody>
</mujoco>
""".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!")