Return a list of dictionaries (preferred for LLM-friendly JSON) python Copy Edit

tray_sim.py

@@ -1,8 +1,13 @@
-import mujoco
-import numpy as np
-import imageio
 import os
+import json
 import tempfile
+import imageio
+
+import numpy as np
+import mujoco
+
+# Extract final object rotations in Euler angles for interpretability (optional)
+from scipy.spatial.transform import Rotation as R
 
 MODEL_PATH = "assets/tray.xml"
 N_OBJECTS = 5  # number of dynamic blocks to randomize
@@ -10,6 +15,7 @@ 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.
@@ -39,6 +45,95 @@ def classify_stability(data, model):
             stable_objects.append(False)
     return stable_objects
 
+# Extend the classifier to explain why something is unstable:
+def classify_stability_verbose(data, model):
+    flags = []
+    tray_bounds = 0.3
+    for i in range(N_OBJECTS):
+        pos = data.qpos[i*7 : i*7 + 3]
+        reasons = []
+        if pos[2] < 0.01:
+            reasons.append("z too low (possibly toppled)")
+        if abs(pos[0]) > tray_bounds or abs(pos[1]) > tray_bounds:
+            reasons.append("outside tray bounds")
+        if not reasons:
+            flags.append({"stable": True, "reason": "object upright and within tray"})
+        else:
+            flags.append({"stable": False, "reason": ", ".join(reasons)})
+    return flags
+
+# Create LLM-Friendly Prompt Template
+def format_llm_prompt(physics_state):
+    """
+    Converts a physics state dictionary of a simulated scene into a JSON-serializable list of objects
+    (human-readable), where each object contains position, orientation, velocity, angular velocity,
+    and stability information suitable for prompting a large language model (LLM).
+
+    Each object is described by its:
+      - 3D position
+      - orientation in ZYX Euler angles (in degrees)
+      - linear velocity
+      - angular velocity
+      - stability status (Stable or Unstable)
+
+    Parameters:
+        physics_state (dict): A dictionary containing the physics properties of each object,
+            with the following keys:
+            - "positions_xyz": List of shape (N, 3), each entry a position vector [x, y, z]
+            - "orientations_euler": List of shape (N, 3), each entry an orientation vector [z, y, x] in degrees
+            - "velocities_linear": List of shape (N, 3), linear velocity vectors
+            - "velocities_angular": List of shape (N, 3), angular velocity vectors
+            - "stable_flags": List of bools indicating whether each object is stable
+
+    Returns:
+        List[dict]: Each dictionary represents an object with formatted fields
+                    suitable for JSON serialization.
+    """
+    formatted = []
+    for i in range(len(physics_state["positions_xyz"])):        
+        obj = {
+            "id": i,
+            "position": np.round(physics_state["positions_xyz"][i], 3).tolist(),
+            "orientation_euler_zyx_deg": np.round(physics_state["orientations_euler"][i], 1).tolist(),
+            "linear_velocity": np.round(physics_state["velocities_linear"][i], 3).tolist(),
+            "angular_velocity": np.round(physics_state["velocities_angular"][i], 3).tolist(),
+            "status": "Stable" if physics_state["stable_flags"][i] else "Unstable"
+            #"status": "Stable" if verbose["stable"] else f"Unstable ({verbose['reason']})"
+        }
+        formatted.append(obj)
+    return formatted
+
+def format_as_natural_language_prompt(scene_description, task_description=None):
+    """
+    Convert structured physics state into a natural language prompt for LLMs.
+
+    Args:
+        scene_description (List[dict]): Output of format_llm_prompt().
+        task_description (str, optional): Additional instruction to prepend, such as
+            "Describe which blocks are unstable and why" or
+            "Predict which objects might move further if the tray tilts".
+
+    Returns:
+        str: Natural language prompt for an LLM.
+    """
+    lines = []
+    if task_description:
+        lines.append(task_description.strip())
+        lines.append("")  # blank line for separation
+
+    lines.append("Here is the scene summary:")
+
+    for obj in scene_description:
+        line = (
+            f"Object {obj['id']} is at position {obj['position']} with orientation {obj['orientation_euler_zyx_deg']} degrees. "
+            f"It has linear velocity {obj['linear_velocity']} and angular velocity {obj['angular_velocity']}. "
+            f"Status: {obj['status']}."
+        )
+        lines.append(line)
+
+    return "\n".join(lines)
+
+
 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)
@@ -130,9 +225,47 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
 
     # Optional: print or return
     print("Stability:", stability_flags)
+    
+    euler_angles = []
+    for i in range(num_objects):
+        quat = data.qpos[i*7+3 : i*7+7]  # qw, qx, qy, qz
+        # Convert to Euler (yaw-pitch-roll)
+        rot = R.from_quat([quat[1], quat[2], quat[3], quat[0]])  # MuJoCo uses [qw, qx, qy, qz]
+        euler = rot.as_euler('zyx', degrees=True)
+        euler_angles.append(euler.tolist())   
+
+    # Fully structured snapshot of: Positions, 
+    # Orientations (in both quaternion and Euler angle forms),
+    # Linear and angular velocities, and Stability flags
+    physics_state = {
+        "positions_xyz": data.qpos[:num_objects * 7].reshape(num_objects, 7)[:, :3].tolist(),  # [x,y,z,qw,qx,qy,qz]
+        "orientations_quat": data.qpos[:num_objects * 7].reshape(num_objects, 7)[:, 3:].tolist(),
+        "orientations_euler": euler_angles,
+        "velocities_linear": data.qvel[:num_objects * 6].reshape(num_objects, 6)[:, :3].tolist(),  # [vx,vy,vz,wx,wy,wz]
+        "velocities_angular": data.qvel[:num_objects * 6].reshape(num_objects, 6)[:, 3:].tolist(),
+        "stable_flags": stability_flags,
+        #"stable_verbose": classify_stability_verbose(data, model), more interpretable to LLMs
+    }
+
+    # Save JSON Snapshot (optional, for LLM input/debugging)
+    json_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{seed}_state.json")
+    with open(json_path, "w") as f:
+        json.dump(physics_state, f, indent=2)
+
+    llm_friendly_output = format_llm_prompt(physics_state)
+
+    #formatted = format_llm_prompt(physics_state)
+    #prompt = format_as_natural_language_prompt(
+    #    formatted,
+    #    task_description="Explain which objects are likely to fall if the tray is tilted slightly to the right."
+    #)
+    print(prompt)
+
 
     # Save to GIF
     gif_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{seed}.gif")
     imageio.mimsave(gif_path, frames, fps=20)
     #return gif_path
-    return gif_path, stability_flags
+    #return gif_path, stability_flags
+    #return gif_path, stability_flags, physics_state  # optionally also return json_path
+    return gif_path, stability_flags, physics_state, llm_friendly_output, json_path # optionally also return json_path