Add scene object extraction and update state broadcasting in mujoco_server.py

- Introduced a new function `_get_scene_objects()` to extract position and orientation data for scene objects, specifically for UR5 scenes.
- Enhanced the `broadcast_state()` function to include scene object information in the state broadcast to connected clients.
- Updated `protocol_types.py` to define a new `SceneObject` type and included a list of scene objects in the `StateData` structure.
- Modified `README.md` to document the new `scene_objects` field, detailing its structure and usage.

README.md

@@ -646,6 +646,10 @@ For robot arm (UR5):
   - Locomotion robots: Use `vx`, `vy`, `vyaw` (other fields are 0)
   - UR5: Use `vx`/`vy`/`vz` for Cartesian translation, `vrx`/`vry`/`vrz` for rotation, `j1`-`j6` for joint velocities, and `gripper` for gripper control
 - `connected_clients`: List of connected external client IDs (e.g., `["trainer_v1", "monitor"]`)
+- `scene_objects`: List of scene objects with their positions and orientations. Each object has:
+  - `name`: Object identifier (e.g., "t_object", "t_target", "box")
+  - `position`: Object position in meters (x, y, z)
+  - `orientation`: Object orientation as quaternion (w, x, y, z)
 
 *Locomotion observation fields (inside `observation`):*
 - `position`: Robot base position in world coordinates (x, y, z) in meters

mujoco_server.py

@@ -554,6 +554,54 @@ def switch_robot(robot_type, scene_name=None):
 env = init_g1()
 
 
+def _get_scene_objects():
+    """Extract scene objects (position and orientation) from the environment."""
+    if env is None:
+        return []
+
+    scene_objects = []
+
+    # For UR5 scenes, extract T-shape objects if they exist
+    if current_robot in ("ur5", "ur5_t_push"):
+        # Get t_object body if it exists
+        t_object_id = getattr(env, 't_object_body_id', -1)
+        if t_object_id >= 0:
+            pos = env.data.xpos[t_object_id]
+            quat = env._get_body_quat(t_object_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+            scene_objects.append({
+                'name': 't_object',
+                'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
+                'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
+            })
+
+        # Get t_target body if it exists
+        t_target_id = getattr(env, 't_target_body_id', -1)
+        if t_target_id >= 0:
+            pos = env.data.xpos[t_target_id]
+            quat = env._get_body_quat(t_target_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+            scene_objects.append({
+                'name': 't_target',
+                'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
+                'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
+            })
+
+        # Get box body if it exists
+        try:
+            box_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, "box")
+            if box_id >= 0:
+                pos = env.data.xpos[box_id]
+                quat = env._get_body_quat(box_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+                scene_objects.append({
+                    'name': 'box',
+                    'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
+                    'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
+                })
+        except:
+            pass
+
+    return scene_objects
+
+
 def broadcast_state():
     """Broadcast robot state to all connected WebSocket clients."""
     with mujoco_lock:
@@ -577,6 +625,9 @@ def broadcast_state():
         with external_ws_clients_lock:
             connected_clients = [info.get("identity", "client") for ws, info in list(client_metadata.items()) if ws in external_ws_clients]
 
+        # Get scene objects
+        scene_objects = _get_scene_objects()
+
         # UR5 has different state structure
         if current_robot in ("ur5", "ur5_t_push"):
             ee_pos = env.get_end_effector_pos()
@@ -592,8 +643,13 @@ def broadcast_state():
             nova_client = getattr(env, '_nova_client', None)
             nova_state_streaming = getattr(env, '_use_nova_state_stream', False)
             nova_ik = getattr(env, '_use_nova_ik', False)
-            # Show connected if client exists (even if first state not received yet)
-            nova_connected = nova_client is not None
+            # Check if actually connected (WebSocket active and receiving data)
+            nova_connected = False
+            if nova_client is not None and hasattr(nova_client, 'is_state_stream_connected'):
+                nova_connected = nova_client.is_state_stream_connected()
+            elif nova_client is not None and not nova_state_streaming:
+                # If not using state streaming, consider connected if client exists
+                nova_connected = True
 
             nova_available = False
             nova_enabled_pref = False
@@ -623,6 +679,7 @@ def broadcast_state():
                     'steps': int(steps),
                     'reward': reward_value,
                     'teleop_action': teleop_snapshot,
+                    'scene_objects': scene_objects,
                     'nova_api': {
                         'connected': nova_connected,
                         'state_streaming': nova_state_streaming,
@@ -650,6 +707,7 @@ def broadcast_state():
                     'steps': int(steps),
                     'reward': reward_value,
                     'teleop_action': teleop_snapshot,
+                    'scene_objects': scene_objects,
                     'connected_clients': connected_clients
                 }
             })

protocol_types.py

@@ -325,6 +325,13 @@ class EulerAngles(TypedDict):
     yaw: float
 
 
+class SceneObject(TypedDict):
+    """Scene object information (position and orientation)."""
+    name: str
+    position: Position
+    orientation: Quaternion
+
+
 class LocomotionObservation(TypedDict):
     """Observation data for locomotion robots (G1, Spot)."""
     position: Position
@@ -359,6 +366,7 @@ class StateData(TypedDict, total=False):
     reward: float
     teleop_action: ActionData  # Current action/velocity commands
     connected_clients: List[str]  # List of connected client IDs
+    scene_objects: List[SceneObject]  # Scene objects (position and orientation)
 
     # UR5-specific fields
     control_mode: ControlMode