Add reset Data Keyframe

tray_sim.py

@@ -34,7 +34,8 @@ def run_tray_simulation(
     # Explicitly reset the simulation state
     model = mujoco.MjModel.from_xml_path(MODEL_PATH)
     data = mujoco.MjData(model)
-    mujoco.mj_resetData(model, data)
+    #mujoco.mj_resetData(model, data)
+    mujoco.mj_resetDataKeyframe(model, data, 0)  # Reset the state to keyframe 0
 
     # Object naming convention: puck, obj0, obj1, ...
     object_names = ["puck"] + [f"obj{i}" for i in range(num_objects - 1)]