""" 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 = """ """ 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!")