WorldRepresentation / render_scripts /render_exp2_estimation.py
Nirav Madhani
Import World Representation Space
5af572e
Raw
History Blame Contribute Delete
5.4 kB
"""
Render Experiment 2: Material Estimation via Sequential Interactions
Shows a robot finger pushing an object with different forces.
Each push provides information about mass and friction.
"""
import mujoco
import numpy as np
import imageio
WIDTH, HEIGHT = 1280, 720
FPS = 60
OUTPUT_DIR = "."
XML_TEMPLATE = """
<mujoco model="material_estimation">
<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"/>
<material name="wood_mat" rgba="0.45 0.35 0.25 1" specular="0.1"/>
<material name="obj_mat" rgba="{r} {g} {b} 1" specular="0.4" shininess="0.6"/>
<material name="pusher_mat" rgba="0.25 0.30 0.65 0.9" specular="0.5" shininess="0.8"/>
</asset>
<worldbody>
<light pos="0 -0.4 1.5" dir="0 0.3 -1" castshadow="true" diffuse="0.8 0.8 0.8"/>
<light pos="0.5 0.5 1.0" 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"/>
<body name="table" pos="0 0 0.38">
<geom type="box" size="0.5 0.3 0.02" material="wood_mat"/>
<geom type="cylinder" size="0.015 0.18" pos="-0.45 -0.25 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos=" 0.45 -0.25 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos="-0.45 0.25 -0.2" rgba="0.4 0.3 0.2 1"/>
<geom type="cylinder" size="0.015 0.18" pos=" 0.45 0.25 -0.2" rgba="0.4 0.3 0.2 1"/>
</body>
<!-- Object under test -->
<body name="object" pos="0 0 0.435">
<joint type="free"/>
<geom type="box" size="0.035 0.035 0.035" mass="{mass}" material="obj_mat"
friction="{friction} 0.005 0.0001"/>
</body>
<!-- Pusher finger -->
<body name="pusher" pos="-0.15 0 0.435">
<geom type="capsule" size="0.012 0.02" euler="0 90 0" material="pusher_mat"/>
</body>
<camera name="main" pos="0.05 -0.55 0.7" xyaxes="1 0 0 0 0.55 0.84" fovy="45"/>
<camera name="close" pos="0.0 -0.3 0.5" xyaxes="1 0 0 0 0.4 0.92" fovy="40"/>
</worldbody>
</mujoco>
"""
MATERIALS = {
"wood": {"mass": 0.5, "friction": 0.4, "color": (0.65, 0.45, 0.25)},
"rubber": {"mass": 1.0, "friction": 0.7, "color": (0.15, 0.15, 0.15)},
"metal": {"mass": 3.0, "friction": 0.2, "color": (0.7, 0.72, 0.75)},
"plastic": {"mass": 0.2, "friction": 0.35, "color": (0.2, 0.6, 0.85)},
}
def render_interaction_sequence(material_name, forces, filename, camera="main"):
"""Render sequential push interactions on a single object."""
props = MATERIALS[material_name]
xml = XML_TEMPLATE.format(
w=WIDTH, h=HEIGHT,
mass=props["mass"], friction=props["friction"],
r=props["color"][0], g=props["color"][1], b=props["color"][2],
)
model = mujoco.MjModel.from_xml_string(xml)
renderer = mujoco.Renderer(model, height=HEIGHT, width=WIDTH)
cam_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera)
all_frames = []
for i, force in enumerate(forces):
# Fresh sim for each push (object resets to center)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)
obj_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "object")
# Settle
for _ in range(200):
mujoco.mj_step(model, data)
frames_per_step = max(1, int(1.0 / (FPS * model.opt.timestep)))
force_steps = 40
# Brief pause
for _ in range(int(0.2 * FPS)):
renderer.update_scene(data, camera=cam_id)
all_frames.append(renderer.render().copy())
# Push and coast
total = int(1.5 / model.opt.timestep)
for step in range(total):
if step < force_steps:
data.xfrc_applied[obj_id, :3] = [force, 0, 0]
else:
data.xfrc_applied[obj_id, :3] = [0, 0, 0]
mujoco.mj_step(model, data)
if step % frames_per_step == 0:
renderer.update_scene(data, camera=cam_id)
all_frames.append(renderer.render().copy())
print(f" Push {i+1}/{len(forces)}: F={force}N, disp={data.xpos[obj_id][0]:.4f}m")
# 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 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 2: Material Estimation ===\n")
forces = [2.0, 5.0, 10.0, 15.0, 8.0, 3.0]
for mat_name in ["wood", "metal"]:
print(f"\n--- {mat_name} ---")
render_interaction_sequence(
mat_name, forces,
f"exp2_{mat_name}_interactions.mp4"
)
print("\nDone!")