Explicitly Reset the Simulation State

tray_sim.py

@@ -133,11 +133,25 @@ def format_as_natural_language_prompt(scene_description, task_description=None):
 
     return "\n".join(lines)
 
-def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25, distance=0.6):
+def run_tray_simulation(
+    seed=0,
+    num_objects=N_OBJECTS,
+    azimuth=45,
+    elevation=-25,
+    distance=0.6
+):
     np.random.seed(seed)
     model = mujoco.MjModel.from_xml_path(MODEL_PATH)
     data = mujoco.MjData(model)
-    tray_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "tray")
+
+    # Explicitly reset the simulation state
+    mujoco.mj_resetData(model, data)
+    
+    tray_id = mujoco.mj_name2id(
+        model,
+        mujoco.mjtObj.mjOBJ_BODY,
+        "tray"
+    )
 
     # Apply a transient force to the tray
     tray_force_applied = True # False   # One-time trigger
@@ -168,6 +182,7 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
         data.qpos[obj_start:obj_start+3] = pos
         data.qpos[obj_start+3:obj_start+7] = quat
         data.qvel[i*6:i*6+6] = 0
+        print(f"Object {i} initial position: {pos}")  # Debug print
 
     # Renderer setup
     renderer = mujoco.Renderer(model, width=480, height=480) # width=640, height=640