Using LLMPhy paper's definition of stability to implement classify orientation stability

tray_sim.py

@@ -16,6 +16,57 @@ SIM_STEPS = 200
 IMPACT_STEP = 60  # is a good starting point, just after the pusher activates.
 
 
+def quat_to_rotmat(q):
+    """Convert quaternion to rotation matrix."""
+    w, x, y, z = q
+    return np.array([
+        [1 - 2*y**2 - 2*z**2,     2*x*y - 2*z*w,     2*x*z + 2*y*w],
+        [2*x*y + 2*z*w,     1 - 2*x**2 - 2*z**2,     2*y*z - 2*x*w],
+        [2*x*z - 2*y*w,         2*y*z + 2*x*w, 1 - 2*x**2 - 2*y**2]
+    ])
+
+def classify_orientation_stability(data, model, alpha_deg=45):
+    """
+    Check object stability using orientation.
+    Returns a list of booleans, True if the object is stable.
+    """
+    stable_objects = []
+    cos_alpha = np.cos(np.deg2rad(alpha_deg))
+    g = np.array([0, 0, -1])
+
+    for i in range(N_OBJECTS):
+        q = data.qpos[i*7 + 3 : i*7 + 7]
+        rotmat = quat_to_rotmat(q)
+        obj_z_world = rotmat[:, 2]  # local z-axis in world coords
+        cos_theta = -np.dot(g, obj_z_world)  # because gravity is downward
+        stable = cos_theta > cos_alpha
+        stable_objects.append(stable)
+
+    return stable_objects
+
+def classify_orientation_stability_verbose(data, model, alpha_deg=45):
+    results = []
+    cos_alpha = np.cos(np.deg2rad(alpha_deg))
+    g = np.array([0, 0, -1])
+
+    for i in range(N_OBJECTS):
+        q = data.qpos[i*7 + 3 : i*7 + 7]
+        rotmat = quat_to_rotmat(q)
+        obj_z_world = rotmat[:, 2]
+        cos_theta = -np.dot(g, obj_z_world)
+
+        if cos_theta > cos_alpha:
+            results.append({"stable": True, "reason": "upright (tilt within limit)"})
+        else:
+            angle_deg = np.rad2deg(np.arccos(cos_theta))
+            results.append({
+                "stable": False,
+                "reason": f"tilt angle {angle_deg:.2f}° exceeds {alpha_deg}°"
+            })
+
+    return results
+
+    
 def classify_stability(data, model):
     """
     Classify object stability based on position data.