Visually render the coordinate axes (X, Y, Z) in the MuJoCo simulation

tray_sim.py

@@ -13,7 +13,7 @@ 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.
+IMPACT_STEP = 60  # is a good starting point, just after the pusher activates.
 
 
 def classify_stability(data, model):
@@ -133,12 +133,51 @@ def format_as_natural_language_prompt(scene_description, task_description=None):
 
     return "\n".join(lines)
 
+def draw_world_axes(scene):
+    """
+    Adds red (X), green (Y), and blue (Z) arrows at the origin to visualize coordinate axes.
+    """
+    mujoco.mjv_addGeom(
+        scene,
+        type=mujoco.mjtGeom.mjGEOM_ARROW,
+        data=None,
+        pos=np.array([0, 0, 0]),
+        mat=np.eye(3).flatten(),  # Identity rotation
+        size=np.array([0.01, 0.01, 0.2]),  # shaft_radius, head_radius, length
+        rgba=np.array([1.0, 0.0, 0.0, 1.0]),  # Red for X
+        label=-1,
+        textureid=-1
+    )
+    mujoco.mjv_addGeom(
+        scene,
+        type=mujoco.mjtGeom.mjGEOM_ARROW,
+        data=None,
+        pos=np.array([0, 0, 0]),
+        mat=np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]]).flatten(),  # 90 deg about Z to point Y
+        size=np.array([0.01, 0.01, 0.2]),
+        rgba=np.array([0.0, 1.0, 0.0, 1.0]),  # Green for Y
+        label=-1,
+        textureid=-1
+    )
+    mujoco.mjv_addGeom(
+        scene,
+        type=mujoco.mjtGeom.mjGEOM_ARROW,
+        data=None,
+        pos=np.array([0, 0, 0]),
+        mat=np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]).flatten(),  # Pointing up
+        size=np.array([0.01, 0.01, 0.2]),
+        rgba=np.array([0.0, 0.0, 1.0, 1.0]),  # Blue for Z
+        label=-1,
+        textureid=-1
+    )
+
 
 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")
+    draw_world_axes(renderer.scene)
 
     # Apply a transient force to the tray
     tray_force_applied = True # False   # One-time trigger
@@ -253,8 +292,6 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
         json.dump(physics_state, f, indent=2)
 
     # LLM-friendly prompt output 
-    #prompt = format_llm_prompt(physics_state)
-
     formatted = format_llm_prompt(physics_state)
     prompt = format_as_natural_language_prompt(
         formatted,
@@ -265,7 +302,4 @@ def run_tray_simulation(seed=0, num_objects=N_OBJECTS, azimuth=45, elevation=-25
     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, physics_state  # optionally also return json_path
-    return gif_path, stability_flags, physics_state, prompt, json_path # optionally also return json_path
+    return gif_path, stability_flags, physics_state, prompt, json_path # optionally also return json_path