Refactor mujoco_server.py to use WebSocket for communication and update README.md for API changes

- Replaced Flask-SocketIO with Flask-Sock for WebSocket handling in mujoco_server.py.
- Implemented WebSocket message handling for commands, resets, and camera controls.
- Updated state broadcasting to use JSON format for WebSocket clients.
- Modified the README.md to reflect the new WebSocket API structure and usage instructions.
- Changed HTTP endpoints to include a versioned API prefix for better organization.

README.md

@@ -31,7 +31,7 @@ A unified MuJoCo-based robot simulation platform with web interface for multiple
 
 ```bash
 # Install dependencies
-pip install mujoco gymnasium flask flask-socketio opencv-python torch numpy
+pip install mujoco gymnasium flask flask-sock opencv-python torch numpy
 
 # Optional: For PyMPC gait controller
 pip install jax jaxlib quadruped-pympc gym-quadruped
@@ -39,7 +39,7 @@ pip install jax jaxlib quadruped-pympc gym-quadruped
 # Start the server
 python mujoco_server.py
 
-# Open browser at http://localhost:3004
+# Open browser at http://localhost:3004/nova-sim/api/v1
 ```
 
 ### Docker
@@ -169,7 +169,7 @@ docker run --gpus all -p 3004:3004 \
 ┌─────────────────────────────────────────────────────────────────────────┐
 │                         mujoco_server.py                                │
 │  ┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐     │
-│  │  Flask + WSGI   │    │SocketIO Handler │    │  Render Thread  │     │
+│  │  Flask + WSGI   │    │   WebSocket     │    │  Render Thread  │     │
 │  │  HTTP endpoints │    │  cmd, reset,    │    │  60 FPS loop    │     │
 │  │                 │    │  switch_robot   │    │  MJPEG encode   │     │
 │  └─────────────────┘    └────────┬────────┘    └────────┬────────┘     │
@@ -289,80 +289,93 @@ env = SpotEnv(controller_type='trot')
 
 ## API
 
-### WebSocket Events
+All endpoints are prefixed with `/nova-sim/api/v1`.
 
-Connect to `ws://localhost:3004/socket.io/` using Socket.IO client.
+### WebSocket
 
-#### Client → Server
+Connect using standard WebSocket:
 
-| Event | Payload | Description |
-|-------|---------|-------------|
+```javascript
+const ws = new WebSocket('ws://localhost:3004/nova-sim/api/v1/ws');
+
+// Send message
+ws.send(JSON.stringify({type: 'command', data: {vx: 0.5, vy: 0, vyaw: 0}}));
+
+// Receive messages
+ws.onmessage = (event) => {
+  const msg = JSON.parse(event.data);
+  if (msg.type === 'state') {
+    console.log(msg.data);
+  }
+};
+```
+
+#### Client → Server Messages
+
+All messages are JSON with `{type, data}` structure:
+
+| Type | Data | Description |
+|------|------|-------------|
 | `command` | `{vx, vy, vyaw}` | Set velocity command |
-| `reset` | *(none)* | Reset robot to standing pose |
+| `reset` | `{}` | Reset robot to standing pose |
 | `switch_robot` | `{robot}` | Switch active robot |
-| `camera` | `{type, ...}` | Camera control |
+| `camera` | `{action, ...}` | Camera control |
 | `camera_follow` | `{follow}` | Toggle camera follow mode |
 
-**`command` payload:**
+**`command`:**
 ```json
-{
-  "vx": 0.5,    // Forward/backward velocity [-1, 1]
-  "vy": 0.0,    // Left/right strafe velocity [-1, 1]
-  "vyaw": 0.0   // Turn rate [-1, 1]
-}
+{"type": "command", "data": {"vx": 0.5, "vy": 0.0, "vyaw": 0.0}}
 ```
+- `vx`: Forward/backward velocity [-1, 1]
+- `vy`: Left/right strafe velocity [-1, 1]
+- `vyaw`: Turn rate [-1, 1]
 
-**`reset`:** No payload - emit the event name only:
-```javascript
-socket.emit('reset');  // No second argument
+**`reset`:**
+```json
+{"type": "reset", "data": {}}
 ```
-Resets the robot to its initial standing pose and resets the camera position.
 
-**`switch_robot` payload:**
+**`switch_robot`:**
 ```json
-{
-  "robot": "g1"  // "g1" or "spot"
-}
+{"type": "switch_robot", "data": {"robot": "spot"}}
 ```
+- `robot`: `"g1"` or `"spot"`
 
-**`camera` payload:**
+**`camera`:**
 ```json
 // Rotate camera
-{"type": "rotate", "dx": 10, "dy": 5}
+{"type": "camera", "data": {"action": "rotate", "dx": 10, "dy": 5}}
 
 // Zoom camera
-{"type": "zoom", "dz": -50}
+{"type": "camera", "data": {"action": "zoom", "dz": -50}}
 
 // Pan camera
-{"type": "pan", "dx": 10, "dy": 5}
+{"type": "camera", "data": {"action": "pan", "dx": 10, "dy": 5}}
 
 // Set absolute distance
-{"type": "set_distance", "distance": 3.0}
+{"type": "camera", "data": {"action": "set_distance", "distance": 3.0}}
 ```
 
-**`camera_follow` payload:**
+**`camera_follow`:**
 ```json
-{
-  "follow": true  // true to follow robot, false for fixed camera
-}
+{"type": "camera_follow", "data": {"follow": true}}
 ```
 
-#### Server → Client
-
-| Event | Payload | Description |
-|-------|---------|-------------|
-| `state` | `{robot, base_height, ...}` | Robot state (broadcast at render rate) |
+#### Server → Client Messages
 
-**`state` payload:**
+**`state`** (broadcast at ~10 Hz):
 ```json
 {
-  "robot": "spot",       // Current robot type
-  "base_height": 0.46,   // Robot base height in meters
-  "upright": 0.98,       // Uprightness metric [0-1] (1 = perfectly upright)
-  "steps": 1234,         // Simulation step count
-  "vx": 0.5,             // Current forward velocity command
-  "vy": 0.0,             // Current strafe velocity command
-  "vyaw": 0.0            // Current turn rate command
+  "type": "state",
+  "data": {
+    "robot": "spot",
+    "base_height": 0.46,
+    "upright": 0.98,
+    "steps": 1234,
+    "vx": 0.5,
+    "vy": 0.0,
+    "vyaw": 0.0
+  }
 }
 ```
 
@@ -370,12 +383,12 @@ Resets the robot to its initial standing pose and resets the camera position.
 
 | Endpoint | Method | Description |
 |----------|--------|-------------|
-| `/` | GET | Web interface (HTML/JS) |
-| `/video_feed` | GET | MJPEG video stream |
+| `/nova-sim/api/v1` | GET | Web interface (HTML/JS) |
+| `/nova-sim/api/v1/video_feed` | GET | MJPEG video stream |
 
 **Video stream usage:**
 ```html
-<img src="http://localhost:3004/video_feed" />
+<img src="http://localhost:3004/nova-sim/api/v1/video_feed" />
 ```
 
 ## License

mujoco_server.py

@@ -2,11 +2,12 @@ import os
 import sys
 import time
 import threading
+import json
 import cv2
 import numpy as np
 import mujoco
 from flask import Flask, Response, render_template_string, request, jsonify
-from flask_socketio import SocketIO, emit
+from flask_sock import Sock
 
 # Add robots directory to path for imports
 _nova_sim_dir = os.path.dirname(os.path.abspath(__file__))
@@ -14,9 +15,12 @@ sys.path.insert(0, os.path.join(_nova_sim_dir, 'robots', 'g1'))
 from g1_env import G1Env
 sys.path.pop(0)
 
+# API prefix for all endpoints
+API_PREFIX = '/nova-sim/api/v1'
+
 app = Flask(__name__)
 app.config['SECRET_KEY'] = 'robotsim-secret'
-socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')
+sock = Sock(app)
 
 # Detect if running in Docker (check for /.dockerenv or cgroup)
 def is_running_in_docker():
@@ -54,6 +58,10 @@ mujoco_lock = threading.Lock()
 renderer = None
 needs_robot_switch = None  # Robot to switch to
 
+# WebSocket clients
+ws_clients = set()
+ws_clients_lock = threading.Lock()
+
 # Camera state for orbit controls
 cam = mujoco.MjvCamera()
 cam.azimuth = 135
@@ -128,7 +136,7 @@ env = init_g1()
 
 
 def broadcast_state():
-    """Broadcast robot state to all connected clients."""
+    """Broadcast robot state to all connected WebSocket clients."""
     with mujoco_lock:
         if env is None:
             return
@@ -140,16 +148,30 @@ def broadcast_state():
     base_quat = obs[3:7]
     upright = float(base_quat[0] ** 2)
 
-    socketio.emit('state', {
-        'robot': current_robot,
-        'base_height': base_height,
-        'upright': upright,
-        'steps': int(steps),
-        'vx': float(cmd[0]),
-        'vy': float(cmd[1]),
-        'vyaw': float(cmd[2])
+    state_msg = json.dumps({
+        'type': 'state',
+        'data': {
+            'robot': current_robot,
+            'base_height': base_height,
+            'upright': upright,
+            'steps': int(steps),
+            'vx': float(cmd[0]),
+            'vy': float(cmd[1]),
+            'vyaw': float(cmd[2])
+        }
     })
 
+    # Send to all connected clients
+    with ws_clients_lock:
+        dead_clients = set()
+        for ws in ws_clients:
+            try:
+                ws.send(state_msg)
+            except:
+                dead_clients.add(ws)
+        # Remove dead clients
+        ws_clients.difference_update(dead_clients)
+
 
 def simulation_loop():
     global running, latest_frame, camera_follow, renderer, needs_robot_switch, env
@@ -238,14 +260,120 @@ def generate_frames():
         time.sleep(0.04)
 
 
+def handle_ws_message(data):
+    """Handle incoming WebSocket message."""
+    global needs_robot_switch, camera_follow
+
+    msg_type = data.get('type')
+
+    if msg_type == 'command':
+        payload = data.get('data', {})
+        vx = payload.get('vx', 0.0)
+        vy = payload.get('vy', 0.0)
+        vyaw = payload.get('vyaw', 0.0)
+        with mujoco_lock:
+            if env is not None:
+                env.set_command(vx, vy, vyaw)
+
+    elif msg_type == 'reset':
+        with mujoco_lock:
+            if env is not None:
+                env.reset()
+        cam.azimuth = 135
+        cam.elevation = -20
+        if current_robot == "g1":
+            cam.distance = 3.0
+            cam.lookat = np.array([0.0, 0.0, 0.8])
+        else:
+            cam.distance = 2.5
+            cam.lookat = np.array([0.0, 0.0, 0.4])
+
+    elif msg_type == 'switch_robot':
+        payload = data.get('data', {})
+        robot = payload.get('robot', 'g1')
+        print(f"Robot switch requested: {robot}")
+        needs_robot_switch = robot
+
+    elif msg_type == 'camera':
+        payload = data.get('data', {})
+        a = cam.azimuth * np.pi / 180.0
+        e = cam.elevation * np.pi / 180.0
+
+        cam_type = payload.get('action')
+        if cam_type == 'rotate':
+            cam.azimuth -= payload.get('dx', 0) * 0.5
+            cam.elevation -= payload.get('dy', 0) * 0.5
+            cam.elevation = np.clip(cam.elevation, -89, 89)
+        elif cam_type == 'zoom':
+            cam.distance += payload.get('dz', 0) * 0.01
+            cam.distance = max(0.5, min(20.0, cam.distance))
+        elif cam_type == 'set_distance':
+            cam.distance = payload.get('distance', 3.0)
+        elif cam_type == 'pan':
+            right = np.array([np.sin(a), -np.cos(a), 0])
+            up = np.array([-np.cos(a) * np.sin(e), -np.sin(a) * np.sin(e), np.cos(e)])
+            scale = cam.distance * 0.002
+            cam.lookat -= (payload.get('dx', 0) * right - payload.get('dy', 0) * up) * scale
+
+    elif msg_type == 'camera_follow':
+        payload = data.get('data', {})
+        camera_follow = payload.get('follow', True)
+
+
+@sock.route(f'{API_PREFIX}/ws')
+def websocket_handler(ws):
+    """Handle WebSocket connections."""
+    print('WebSocket client connected')
+
+    # Register client
+    with ws_clients_lock:
+        ws_clients.add(ws)
+
+    try:
+        while True:
+            message = ws.receive()
+            if message is None:
+                break
+            try:
+                data = json.loads(message)
+                handle_ws_message(data)
+            except json.JSONDecodeError:
+                print(f"Invalid JSON received: {message}")
+            except Exception as e:
+                print(f"Error handling message: {e}")
+    except:
+        pass
+    finally:
+        # Unregister client
+        with ws_clients_lock:
+            ws_clients.discard(ws)
+        print('WebSocket client disconnected')
+
+
+# Redirect root to UI
 @app.route('/')
+def root_redirect():
+    return render_template_string("""
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <meta http-equiv="refresh" content="0; url='""" + API_PREFIX + """'" />
+    </head>
+    <body>
+        <p>Redirecting to <a href=\"""" + API_PREFIX + """\">""" + API_PREFIX + """</a>...</p>
+    </body>
+    </html>
+    """)
+
+
+@app.route(API_PREFIX)
+@app.route(f'{API_PREFIX}/')
 def index():
     return render_template_string("""
     <!DOCTYPE html>
     <html>
     <head>
-        <title>Robot Simulator - MuJoCo</title>
-        <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.6.1/socket.io.min.js"></script>
+        <title>Nova Sim - Robot Simulator</title>
         <style>
             body, html {
                 margin: 0; padding: 0; width: 100%; height: 100%;
@@ -331,10 +459,10 @@ def index():
     </head>
     <body>
         <div class="video-container" id="viewport">
-            <img src="/video_feed" draggable="false">
+            <img src=\"""" + API_PREFIX + """/video_feed" draggable="false">
         </div>
 
-        <div class="connection-status" id="conn_status">Connected</div>
+        <div class="connection-status" id="conn_status">Connecting...</div>
 
         <div class="overlay">
             <h2 id="robot_title">Unitree G1 Humanoid</h2>
@@ -404,7 +532,12 @@ def index():
         </div>
 
         <script>
-            const socket = io();
+            const API_PREFIX = '""" + API_PREFIX + """';
+            const WS_URL = (window.location.protocol === 'https:' ? 'wss://' : 'ws://') +
+                           window.location.host + API_PREFIX + '/ws';
+
+            let ws = null;
+            let reconnectTimer = null;
             const connStatus = document.getElementById('conn_status');
             const robotSelect = document.getElementById('robot_select');
             const robotTitle = document.getElementById('robot_title');
@@ -420,29 +553,63 @@ def index():
                 'spot': 'Boston Dynamics Spot'
             };
 
-            socket.on('connect', () => {
-                connStatus.textContent = 'Connected';
-                connStatus.classList.remove('disconnected');
-            });
-
-            socket.on('disconnect', () => {
-                connStatus.textContent = 'Disconnected';
-                connStatus.classList.add('disconnected');
-            });
+            function connect() {
+                ws = new WebSocket(WS_URL);
+
+                ws.onopen = () => {
+                    connStatus.textContent = 'Connected';
+                    connStatus.classList.remove('disconnected');
+                    if (reconnectTimer) {
+                        clearInterval(reconnectTimer);
+                        reconnectTimer = null;
+                    }
+                };
+
+                ws.onclose = () => {
+                    connStatus.textContent = 'Disconnected';
+                    connStatus.classList.add('disconnected');
+                    // Auto-reconnect
+                    if (!reconnectTimer) {
+                        reconnectTimer = setInterval(() => {
+                            if (ws.readyState === WebSocket.CLOSED) {
+                                connect();
+                            }
+                        }, 2000);
+                    }
+                };
+
+                ws.onerror = (err) => {
+                    console.error('WebSocket error:', err);
+                };
+
+                ws.onmessage = (event) => {
+                    try {
+                        const msg = JSON.parse(event.data);
+                        if (msg.type === 'state') {
+                            const data = msg.data;
+                            heightVal.innerText = data.base_height.toFixed(2);
+                            uprightVal.innerText = data.upright.toFixed(2);
+                            stepVal.innerText = data.steps;
+                            cmdVx.innerText = data.vx.toFixed(1);
+                            cmdVy.innerText = data.vy.toFixed(1);
+                            cmdYaw.innerText = data.vyaw.toFixed(1);
+
+                            if (data.robot && robotSelect.value !== data.robot) {
+                                robotSelect.value = data.robot;
+                                updateRobotUI(data.robot);
+                            }
+                        }
+                    } catch (e) {
+                        console.error('Error parsing message:', e);
+                    }
+                };
+            }
 
-            socket.on('state', (data) => {
-                heightVal.innerText = data.base_height.toFixed(2);
-                uprightVal.innerText = data.upright.toFixed(2);
-                stepVal.innerText = data.steps;
-                cmdVx.innerText = data.vx.toFixed(1);
-                cmdVy.innerText = data.vy.toFixed(1);
-                cmdYaw.innerText = data.vyaw.toFixed(1);
-
-                if (data.robot && robotSelect.value !== data.robot) {
-                    robotSelect.value = data.robot;
-                    updateRobotUI(data.robot);
+            function send(type, data = {}) {
+                if (ws && ws.readyState === WebSocket.OPEN) {
+                    ws.send(JSON.stringify({type, data}));
                 }
-            });
+            }
 
             function updateRobotUI(robot) {
                 robotTitle.innerText = robotTitles[robot] || robot;
@@ -451,7 +618,7 @@ def index():
 
             function switchRobot() {
                 const robot = robotSelect.value;
-                socket.emit('switch_robot', {robot});
+                send('switch_robot', {robot});
                 updateRobotUI(robot);
             }
 
@@ -468,16 +635,16 @@ def index():
             let keysPressed = new Set();
 
             function resetEnv() {
-                socket.emit('reset');
+                send('reset');
             }
 
             function setCmd(vx, vy, vyaw) {
-                socket.emit('command', {vx, vy, vyaw});
+                send('command', {vx, vy, vyaw});
             }
 
             function setCameraFollow() {
                 const follow = document.getElementById('cam_follow').checked;
-                socket.emit('camera_follow', {follow});
+                send('camera_follow', {follow});
             }
 
             function updateCmdFromKeys() {
@@ -510,7 +677,7 @@ def index():
 
             camDist.oninput = () => {
                 camDistVal.innerText = parseFloat(camDist.value).toFixed(1);
-                socket.emit('camera', {type: 'set_distance', distance: parseFloat(camDist.value)});
+                send('camera', {action: 'set_distance', distance: parseFloat(camDist.value)});
             };
 
             let isDragging = false;
@@ -534,101 +701,28 @@ def index():
                     const dy = e.clientY - lastY;
                     lastX = e.clientX;
                     lastY = e.clientY;
-                    socket.emit('camera', {type: 'rotate', dx, dy});
+                    send('camera', {action: 'rotate', dx, dy});
                 }
             };
 
             viewport.onwheel = (e) => {
                 e.preventDefault();
-                socket.emit('camera', {type: 'zoom', dz: e.deltaY});
+                send('camera', {action: 'zoom', dz: e.deltaY});
             };
+
+            // Connect on load
+            connect();
         </script>
     </body>
     </html>
     """)
 
 
-@app.route('/video_feed')
+@app.route(f'{API_PREFIX}/video_feed')
 def video_feed():
     return Response(generate_frames(),
                     mimetype='multipart/x-mixed-replace; boundary=frame')
 
 
-@socketio.on('connect')
-def handle_connect():
-    print('Client connected')
-
-
-@socketio.on('disconnect')
-def handle_disconnect():
-    print('Client disconnected')
-
-
-@socketio.on('switch_robot')
-def handle_switch_robot(data):
-    """Handle robot switch request from client."""
-    global needs_robot_switch
-    robot = data.get('robot', 'g1')
-    print(f"Robot switch requested: {robot}")
-    needs_robot_switch = robot
-
-
-@socketio.on('command')
-def handle_command(data):
-    """Handle velocity command from client."""
-    vx = data.get('vx', 0.0)
-    vy = data.get('vy', 0.0)
-    vyaw = data.get('vyaw', 0.0)
-    with mujoco_lock:
-        if env is not None:
-            env.set_command(vx, vy, vyaw)
-
-
-@socketio.on('reset')
-def handle_reset():
-    """Handle reset request from client."""
-    global camera_follow
-    with mujoco_lock:
-        if env is not None:
-            env.reset()
-    cam.azimuth = 135
-    cam.elevation = -20
-    if current_robot == "g1":
-        cam.distance = 3.0
-        cam.lookat = np.array([0.0, 0.0, 0.8])
-    else:
-        cam.distance = 2.5
-        cam.lookat = np.array([0.0, 0.0, 0.4])
-
-
-@socketio.on('camera')
-def handle_camera(data):
-    """Handle camera control from client."""
-    a = cam.azimuth * np.pi / 180.0
-    e = cam.elevation * np.pi / 180.0
-
-    if data['type'] == 'rotate':
-        cam.azimuth -= data['dx'] * 0.5
-        cam.elevation -= data['dy'] * 0.5
-        cam.elevation = np.clip(cam.elevation, -89, 89)
-    elif data['type'] == 'zoom':
-        cam.distance += data['dz'] * 0.01
-        cam.distance = max(0.5, min(20.0, cam.distance))
-    elif data['type'] == 'set_distance':
-        cam.distance = data['distance']
-    elif data['type'] == 'pan':
-        right = np.array([np.sin(a), -np.cos(a), 0])
-        up = np.array([-np.cos(a) * np.sin(e), -np.sin(a) * np.sin(e), np.cos(e)])
-        scale = cam.distance * 0.002
-        cam.lookat -= (data['dx'] * right - data['dy'] * up) * scale
-
-
-@socketio.on('camera_follow')
-def handle_camera_follow(data):
-    """Handle camera follow toggle from client."""
-    global camera_follow
-    camera_follow = data.get('follow', True)
-
-
 if __name__ == '__main__':
-    socketio.run(app, host='0.0.0.0', port=3004, debug=False, allow_unsafe_werkzeug=True)
+    app.run(host='0.0.0.0', port=3004, debug=False, threaded=True)

requirements.txt

@@ -1,7 +1,7 @@
 mujoco>=3.1.3
 gymnasium>=0.29.0
 flask>=3.0.0
-flask-socketio>=5.3.0
+flask-sock>=0.7.0
 opencv-python>=4.8.0
 torch>=2.0.0
 numpy>=1.24.0