Add trainer connection status broadcasting in mujoco_server.py

- Implemented a new function to notify UI clients about the trainer's connection status, enhancing real-time feedback.
- Updated the trainer websocket handler to call the new broadcasting function upon client connection and disconnection.
- Modified the UI message handling to update the trainer connection status dynamically, improving user experience during teleoperation.

mujoco_server.py

@@ -565,6 +565,21 @@ def broadcast_state():
         # Remove dead clients
         ws_clients.difference_update(dead_clients)
 
+
+def broadcast_trainer_connection_status():
+    """Notify UI clients about trainer connection state."""
+    with trainer_ws_clients_lock:
+        connected = len(trainer_ws_clients) > 0
+    message = json.dumps({"type": "trainer_status", "data": {"connected": connected}})
+    with ws_clients_lock:
+        dead_clients = set()
+        for client in ws_clients:
+            try:
+                client.send(message)
+            except:
+                dead_clients.add(client)
+        ws_clients.difference_update(dead_clients)
+
 def broadcast_to_trainer(event_type: str, payload: dict):
     """Broadcast a structured event to connected trainer WebSocket clients."""
     message = json.dumps({"type": event_type, "data": payload})
@@ -1100,6 +1115,7 @@ def trainer_websocket_handler(ws):
     print('Trainer client connected')
     with trainer_ws_clients_lock:
         trainer_ws_clients.add(ws)
+    broadcast_trainer_connection_status()
     broadcast_state()
 
     try:
@@ -1119,6 +1135,7 @@ def trainer_websocket_handler(ws):
             trainer_ws_clients.discard(ws)
         print('Trainer client disconnected')
         broadcast_state()
+        broadcast_trainer_connection_status()
 
 
 def _handle_trainer_message(data):
@@ -2442,6 +2459,11 @@ def index():
                                 };
                                 updateRobotUI(data.robot, sceneValue);
                             }
+                        } else if (msg.type === 'trainer_status') {
+                            const payload = msg.data || {};
+                            if (typeof payload.connected === 'boolean') {
+                                updateTrainerStatus(payload.connected);
+                            }
                         } else if (msg.type === 'trainer_notification') {
                             showTrainerNotification(msg.data);
                         }