TraySim module

tray_sim.py

@@ -0,0 +1,29 @@
+import mujoco
+import numpy as np
+import imageio
+import os
+import tempfile
+from mujoco import viewer
+
+MODEL_PATH = "assets/tray.xml"
+
+def run_tray_simulation(seed=0):
+    np.random.seed(seed)
+    model = mujoco.MjModel.from_xml_path(MODEL_PATH)
+    data = mujoco.MjData(model)
+
+    # Randomize initial object position
+    data.qpos[:3] = [0.0, 0.0, 0.3 + 0.2 * np.random.rand()]  # x, y, z
+    data.qvel[:] = 0
+
+    # Simulate
+    frames = []
+    for _ in range(200):
+        mujoco.mj_step(model, data)
+        img = mujoco.mj_render(model, data, width=480, height=480, camera=None)
+        frames.append(img)
+
+    # Save as GIF to temp file
+    gif_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{seed}.gif")
+    imageio.mimsave(gif_path, frames, fps=20)
+    return gif_path