Move variable definition and initialization up.

tray_sim.py

@@ -15,6 +15,11 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
     model = mujoco.MjModel.from_xml_path(MODEL_PATH)
     data = mujoco.MjData(model)
 
+    # Locate pusher DOF index (last one in qpos if you have 10 objects)
+    # Compute DOF offsets
+    pusher_idx = num_objects * 7
+    pusher_vel_idx = num_objects * 6
+
     # Pusher setup
     pusher_pos = np.array([0.35, -0.35, 0.08])  # Slightly outside tray corner
     pusher_quat = [1, 0, 0, 0]
@@ -37,10 +42,6 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
         data.qpos[obj_start+3:obj_start+7] = quat
         data.qvel[i*6:i*6+6] = 0
 
-    # Locate pusher DOF index (last one in qpos if you have 10 objects)
-    pusher_idx = num_objects * 7
-    pusher_vel_idx = num_objects * 6
-
     # Renderer setup
     renderer = mujoco.Renderer(model, width=480, height=480) # width=640, height=640