Apply a Force to the Tray

tray_sim.py

@@ -8,13 +8,46 @@ MODEL_PATH = "assets/tray.xml"
 N_OBJECTS = 5  # number of dynamic blocks to randomize
 PUSH_START_STEP = 50
 SIM_STEPS = 200
-
+IMPACT_STEP = 60  #  is a good starting point, just after pusher activates.
+
+def classify_stability(data, model):
+    """
+    Classify object stability based on position data.
+
+    For each object in the scene:
+    - Check the object's position using data.qpos[z].
+    - Determine whether the object is upright, within the tray bounds, or has fallen/rolled over.
+    - Define thresholds for "stable" as:
+        * z-position > 0.01
+        * x and y within tray bounds of ±0.3
+
+    Args:
+        data: MuJoCo data object containing simulation state.
+        model: MuJoCo model (unused in this function but typically passed alongside data).
+
+    Returns:
+        List of booleans indicating whether each object is stable.
+    """
+    
+    stable_objects = []
+    tray_bounds = 0.3
+    for i in range(N_OBJECTS):
+        pos = data.qpos[i*7 : i*7 + 3]
+        if abs(pos[0]) < tray_bounds and abs(pos[1]) < tray_bounds and pos[2] > 0.01:
+            stable_objects.append(True)
+        else:
+            stable_objects.append(False)
+    return stable_objects
 
 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")
 
+    # Apply a transient force to the tray
+    tray_force_applied = False   # One-time trigger
+    
     # Locate pusher DOF index (last one in qpos if you have 10 objects)
     # Compute DOF offsets
     pusher_idx = num_objects * 7
@@ -70,6 +103,17 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
             #data.qvel[pusher_vel_idx:pusher_vel_idx+3] = [-0.05, 0.05, 0.0]  # toward center
             data.qvel[pusher_vel_idx:pusher_vel_idx+3] = push_velocity  # toward center
 
+        if t == IMPACT_STEP and not tray_force_applied:
+            mujoco.mj_applyFT(
+                model, data,
+                tray_id,
+                force=[5.0, 0.0, 0.0],   # Impulse to the right; adjust magnitude
+                torque=[0.0, 0.0, 0.0],
+                point=[0.0, 0.0, 0.0],   # Apply at body origin
+                frame=0                  # World frame
+            )
+            tray_force_applied = True
+        
         mujoco.mj_step(model, data)
         renderer.update_scene(data, camera=cam, scene_option=scene_option)
         img = renderer.render()