Remove superfluous code

tray_sim.py

@@ -25,164 +25,6 @@ def quat_to_rotmat(q):
         [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.
-
-    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
-
-# 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,
@@ -226,30 +68,12 @@ def run_tray_simulation(
     direction /= np.linalg.norm(direction)
     push_velocity = 0.15 * direction  # Adjust magnitude
 
-    # Object initialization
-    for i in range(num_objects):
-        obj_start = i * 7  # 7 DOF per free joint body
-        pos = np.random.uniform(low=[-0.1, -0.1, 0.05], high=[0.1, 0.1, 0.15])
-        quat = [1, 0, 0, 0]  # no rotation
-        data.qpos[obj_start:obj_start+3] = pos
-        data.qpos[obj_start+3:obj_start+7] = quat
-        data.qvel[i*6:i*6+6] = 0
-        print(f"Object {i} initial position: {pos}")  # Debug print
-
     # Renderer setup
     renderer = mujoco.Renderer(model, width=480, height=480) # width=640, height=640
 
     # Custom camera parameters
     scene_option = mujoco.MjvOption()
 
-    # Enable contact visualization in the renderer to debug collisions
-    # Check if contact points appear when the puck hits the tray. If not,
-    # collisions are not being detected.
-    #scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True
-    #scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True
-    #scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True # if you want to see joint frames or anchors to understand the articulated structure.
-    #scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True #  if the model uses transparent geoms and you want them shown correctly.
-
     cam = mujoco.MjvCamera()
     cam.type = mujoco.mjtCamera.mjCAMERA_FREE
     cam.lookat[:] = np.array([0.0, 0.0, 0.02])     # Center of the tray
@@ -263,20 +87,8 @@ def run_tray_simulation(
     for t in range(SIM_STEPS):
         if t == PUSH_START_STEP:
             # Activate pusher
-            #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,
-        #        np.array([5.0, 0.0, 0.0]),       # force (e.g. impulse to the right; adjust magnitude)
-        #        np.array([0.0, 0.0, 0.0]),       # torque
-        #        np.array([0.0, 0.0, 0.0]),       # point in global frame (apply at body origin)
-        #        tray_id,                         # body ID
-        #        data.qfrc_applied                # force accumulator to apply into
-        #    )
-        #    tray_force_applied = True
             
         mujoco.mj_step(model, data)
     
@@ -288,17 +100,11 @@ def run_tray_simulation(
         img = renderer.render()
             
         if t in keyframes:
-            #imageio.imwrite(f"/tmp/frame_{t}.png", img)
             imageio.imwrite(f"/tmp/frame_{t}_{unique_id}.png", img)
     
         frames.append(img)
 
     renderer.close()
-
-    stability_flags = classify_stability(data, model)
-
-    # Optional: print or return
-    print("Stability:", stability_flags)
     
     euler_angles = []
     for i in range(num_objects):
@@ -310,35 +116,22 @@ def run_tray_simulation(
 
     # Fully structured snapshot of: Positions, 
     # Orientations (in both quaternion and Euler angle forms),
-    # Linear and angular velocities, and Stability flags
+    # Linear and angular velocities
     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")
-    # Save JSON with unique ID
+    # Save JSON Snapshot with unique ID (optional, for LLM input/debugging)
     json_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{unique_id}_state.json")
     with open(json_path, "w") as f:
         json.dump(physics_state, f, indent=2)
 
-    # LLM-friendly prompt output 
-    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."
-    )
-
-    # Save to GIF
-    #gif_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{seed}.gif")
     # Save GIF with unique ID
     gif_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{unique_id}.gif")
     imageio.mimsave(gif_path, frames, fps=20)
 
-    return gif_path, stability_flags, physics_state, prompt, json_path # optionally also return json_path
+    return gif_path, json_path # optionally also return json_path