Add GPU support with docker-compose.gpu.yml for enhanced performance

- Introduced a new docker-compose.gpu.yml file to enable GPU-accelerated rendering using NVIDIA GPUs and the nvidia-container-toolkit.
- Updated environment variables for improved rendering performance, including higher resolution and frame rate settings.
- Enhanced mujoco_server.py to dynamically adjust rendering settings based on the environment (Docker vs. native).
- Updated README.md with instructions for using the new GPU configuration and performance optimization details.

Dockerfile

@@ -28,9 +28,14 @@ RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
 
 # Environment variables for headless rendering
+# Options: osmesa (software, slow), egl (GPU if available), glx (X11 GPU)
 ENV MUJOCO_GL=osmesa
 ENV PYOPENGL_PLATFORM=osmesa
 
+# Performance tuning
+ENV OMP_NUM_THREADS=4
+ENV MKL_NUM_THREADS=4
+
 # Expose port
 EXPOSE 3004
 

README.md

@@ -11,8 +11,10 @@ A unified MuJoCo-based robot simulation platform with web interface for multiple
 
 ### Boston Dynamics Spot (Quadruped)
 - 12 degrees of freedom (3 per leg)
-- Classical trot gait controller
-- Diagonal leg pairing locomotion
+- Multiple gait controllers:
+  - **MPC Gait**: Feedback-based balance control with trot pattern (default)
+  - **PyMPC Gait**: Uses Quadruped-PyMPC gait generator (requires extra dependencies)
+  - **Trot Gait**: Simple open-loop trot pattern
 
 ## Features
 
@@ -23,68 +25,203 @@ A unified MuJoCo-based robot simulation platform with web interface for multiple
 - Robot switching without restart
 - Keyboard and button controls for locomotion
 
-## Installation
+## Quick Start
 
-```bash
-# Clone the repository
-git clone git@github.com:wandelbotsgmbh/nova-sim.git
-cd nova-sim
+### Native (Recommended for Development)
 
+```bash
 # Install dependencies
 pip install mujoco gymnasium flask flask-socketio opencv-python torch numpy
 
-# Clone MuJoCo Menagerie for Spot model (in parent directory)
-cd ..
-git clone https://github.com/google-deepmind/mujoco_menagerie.git
+# Optional: For PyMPC gait controller
+pip install jax jaxlib quadruped-pympc gym-quadruped
+
+# Start the server
+python mujoco_server.py
 
-# For G1 RL policy, clone unitree_rl_gym (in parent directory)
-git clone https://github.com/unitreerobotics/unitree_rl_gym.git
+# Open browser at http://localhost:3004
 ```
 
-## Usage
+### Docker
 
 ```bash
-# Start the server
-cd nova-sim
-python mujoco_server.py
+# Build and run (CPU/software rendering)
+docker-compose up --build
 
-# Open browser at http://localhost:3004
+# Or with docker run
+docker run -p 3004:3004 nova-sim
+```
+
+## Docker Configuration
+
+### Performance Optimization
+
+Docker uses software rendering (OSMesa) which is slower than native GPU rendering. Configure these environment variables to optimize performance:
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `RENDER_WIDTH` | 640 | Render width in pixels |
+| `RENDER_HEIGHT` | 360 | Render height in pixels |
+| `TARGET_FPS` | 30 | Target frame rate |
+| `SIM_STEPS_PER_FRAME` | 10 | Physics steps per rendered frame |
+| `OMP_NUM_THREADS` | 4 | OpenMP thread count |
+| `MKL_NUM_THREADS` | 4 | MKL thread count |
+
+### docker-compose.yml (CPU - Default)
+
+```yaml
+services:
+  nova-sim:
+    build: .
+    ports:
+      - "3004:3004"
+    environment:
+      - MUJOCO_GL=osmesa
+      - PYOPENGL_PLATFORM=osmesa
+      - RENDER_WIDTH=640      # Lower resolution for speed
+      - RENDER_HEIGHT=360
+      - TARGET_FPS=30
+```
+
+### docker-compose.gpu.yml (NVIDIA GPU)
+
+For significantly better performance, use GPU acceleration:
+
+```bash
+# Requires: NVIDIA GPU + nvidia-container-toolkit
+docker-compose -f docker-compose.gpu.yml up
+```
+
+```yaml
+services:
+  nova-sim:
+    build: .
+    ports:
+      - "3004:3004"
+    environment:
+      - MUJOCO_GL=egl         # GPU rendering
+      - PYOPENGL_PLATFORM=egl
+      - RENDER_WIDTH=1280     # Full resolution
+      - RENDER_HEIGHT=720
+      - TARGET_FPS=60
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+```
+
+### Performance Comparison
+
+| Mode | Resolution | Expected FPS | Notes |
+|------|------------|--------------|-------|
+| Native (macOS/Linux) | 1280x720 | 60+ | Best performance |
+| Docker + GPU (EGL) | 1280x720 | 60 | Requires NVIDIA GPU |
+| Docker + CPU (OSMesa) | 640x360 | 15-20 | Slower, but compatible |
+| Docker + CPU (OSMesa) | 320x180 | 30+ | Very low quality |
+
+### Custom Docker Run
+
+```bash
+# Ultra-low resolution for maximum speed
+docker run -p 3004:3004 \
+  -e RENDER_WIDTH=320 \
+  -e RENDER_HEIGHT=180 \
+  -e MUJOCO_GL=osmesa \
+  nova-sim
+
+# High quality with GPU
+docker run --gpus all -p 3004:3004 \
+  -e RENDER_WIDTH=1920 \
+  -e RENDER_HEIGHT=1080 \
+  -e MUJOCO_GL=egl \
+  nova-sim
 ```
 
 ## Controls
 
-- **WASD / Arrow Keys**: Walk forward/backward, turn left/right
-- **Q/E**: Strafe left/right
+- **W / Arrow Up**: Walk forward
+- **S / Arrow Down**: Walk backward
+- **A / Arrow Left**: Turn left
+- **D / Arrow Right**: Turn right
+- **Q**: Strafe left
+- **E**: Strafe right
 - **Mouse drag**: Rotate camera
 - **Scroll wheel**: Zoom camera
 - **Robot selector**: Switch between G1 and Spot
+- **Reset button**: Reset robot to standing pose
 
 ## Project Structure
 
 ```
 nova_sim/
-├── mujoco_server.py      # Main Flask server with WebSocket
+├── mujoco_server.py          # Main Flask server with WebSocket
+├── Dockerfile                # Docker build configuration
+├── docker-compose.yml        # CPU/OSMesa configuration
+├── docker-compose.gpu.yml    # GPU/EGL configuration
+├── requirements.txt          # Python dependencies
 ├── robots/
-│   ├── g1/               # Unitree G1 humanoid
-│   │   ├── g1_env.py     # Gymnasium environment
-│   │   ├── scene.xml     # MuJoCo scene
-│   │   ├── g1_29dof.xml  # Robot model
-│   │   ├── meshes/       # 3D mesh files
-│   │   └── controllers/  # G1 controllers
-│   │       ├── rl_policy.py    # RL walking policy
-│   │       ├── pd_standing.py  # Standing controller
-│   │       └── keyframe.py     # Keyframe controller
-│   └── spot/             # Boston Dynamics Spot quadruped
-│       ├── spot_env.py   # Gymnasium environment
-│       └── controllers/  # Spot controllers
-│           ├── trot_gait.py    # Trot gait locomotion
-│           └── pd_standing.py  # Standing controller
+│   ├── g1/                   # Unitree G1 humanoid
+│   │   ├── g1_env.py         # Gymnasium environment
+│   │   ├── scene.xml         # MuJoCo scene
+│   │   ├── g1_29dof.xml      # Robot model
+│   │   ├── meshes/           # 3D mesh files
+│   │   ├── policy/           # RL policy weights
+│   │   └── controllers/      # G1 controllers
+│   │       ├── rl_policy.py      # RL walking policy
+│   │       ├── pd_standing.py    # Standing controller
+│   │       └── keyframe.py       # Keyframe controller
+│   └── spot/                 # Boston Dynamics Spot quadruped
+│       ├── spot_env.py       # Gymnasium environment
+│       ├── model/            # MuJoCo model files
+│       └── controllers/      # Spot controllers
+│           ├── mpc_gait.py           # MPC-inspired gait (default)
+│           ├── quadruped_pympc_controller.py  # PyMPC gait
+│           ├── trot_gait.py          # Simple trot gait
+│           └── pd_standing.py        # Standing controller
 └── README.md
 ```
 
+## Spot Controller Options
+
+The Spot robot supports multiple controller types via the `controller_type` parameter:
+
+```python
+from spot_env import SpotEnv
+
+# MPC Gait - Feedback-based balance control (default, no extra deps)
+env = SpotEnv(controller_type='mpc_gait')
+
+# PyMPC Gait - Uses Quadruped-PyMPC gait generator (requires jax, gym-quadruped)
+env = SpotEnv(controller_type='pympc')
+
+# Trot Gait - Simple open-loop trot pattern
+env = SpotEnv(controller_type='trot')
+```
+
+## API
+
+### WebSocket Events
+
+**Client to Server:**
+- `command`: `{vx, vy, vyaw}` - Velocity command
+- `reset`: Reset robot to initial pose
+- `switch_robot`: `{robot: "g1" | "spot"}` - Switch robot type
+
+**Server to Client:**
+- `state`: Robot state data (position, orientation, joint angles)
+
+### HTTP Endpoints
+
+- `GET /` - Web interface
+- `GET /video_feed` - MJPEG video stream
+
 ## License
 
 This project uses models from:
 - [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie) (BSD-3-Clause)
 - [unitree_mujoco](https://github.com/unitreerobotics/unitree_mujoco) (BSD-3-Clause)
 - [unitree_rl_gym](https://github.com/unitreerobotics/unitree_rl_gym)
+- [Quadruped-PyMPC](https://github.com/iit-DLSLab/Quadruped-PyMPC) (BSD-3-Clause)

docker-compose.gpu.yml

@@ -0,0 +1,31 @@
+version: '3.8'
+
+# GPU-accelerated version - requires NVIDIA GPU and nvidia-container-toolkit
+# Usage: docker-compose -f docker-compose.gpu.yml up
+
+services:
+  nova-sim:
+    build: .
+    ports:
+      - "3004:3004"
+    environment:
+      # Use EGL for GPU rendering (much faster than OSMesa)
+      - MUJOCO_GL=egl
+      - PYOPENGL_PLATFORM=egl
+      # Higher resolution since GPU can handle it
+      - RENDER_WIDTH=1280
+      - RENDER_HEIGHT=720
+      - TARGET_FPS=60
+      - SIM_STEPS_PER_FRAME=5
+      # NVIDIA driver
+      - NVIDIA_VISIBLE_DEVICES=all
+      - NVIDIA_DRIVER_CAPABILITIES=graphics,utility,compute
+    restart: unless-stopped
+    # Enable GPU access
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

docker-compose.yml

@@ -6,6 +6,21 @@ services:
     ports:
       - "3004:3004"
     environment:
+      # Rendering backend (osmesa = software, slow but compatible)
       - MUJOCO_GL=osmesa
       - PYOPENGL_PLATFORM=osmesa
+      # Performance settings - reduce resolution for faster software rendering
+      - RENDER_WIDTH=640
+      - RENDER_HEIGHT=360
+      - TARGET_FPS=30
+      - SIM_STEPS_PER_FRAME=10
+      # Threading optimization
+      - OMP_NUM_THREADS=4
+      - MKL_NUM_THREADS=4
     restart: unless-stopped
+    # Optional: limit CPU to prevent runaway usage
+    # deploy:
+    #   resources:
+    #     limits:
+    #       cpus: '2.0'
+    #       memory: 2G

mujoco_server.py

@@ -18,6 +18,26 @@ app = Flask(__name__)
 app.config['SECRET_KEY'] = 'robotsim-secret'
 socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')
 
+# Detect if running in Docker (check for /.dockerenv or cgroup)
+def is_running_in_docker():
+    if os.path.exists('/.dockerenv'):
+        return True
+    try:
+        with open('/proc/1/cgroup', 'r') as f:
+            return 'docker' in f.read()
+    except:
+        return False
+
+IN_DOCKER = is_running_in_docker()
+
+# Performance settings from environment
+# Native defaults: 1280x720 @ 60 FPS (GPU accelerated)
+# Docker defaults: 640x360 @ 30 FPS (software rendering)
+RENDER_WIDTH = int(os.environ.get('RENDER_WIDTH', 640 if IN_DOCKER else 1280))
+RENDER_HEIGHT = int(os.environ.get('RENDER_HEIGHT', 360 if IN_DOCKER else 720))
+TARGET_FPS = int(os.environ.get('TARGET_FPS', 30 if IN_DOCKER else 60))
+SIM_STEPS_PER_FRAME = int(os.environ.get('SIM_STEPS_PER_FRAME', 10 if IN_DOCKER else 5))
+
 # Current robot type
 current_robot = "g1"  # or "spot"
 
@@ -49,7 +69,7 @@ def init_g1():
     """Initialize G1 environment."""
     global env_g1
     if env_g1 is None:
-        env_g1 = G1Env(render_mode="rgb_array")
+        env_g1 = G1Env(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT)
         env_g1.reset()
     return env_g1
 
@@ -64,7 +84,7 @@ def init_spot():
         from spot_env import SpotEnv
         sys.path.pop(0)
         # Use pympc controller if available, otherwise fallback to mpc_gait
-        env_spot = SpotEnv(render_mode="rgb_array", controller_type='pympc')
+        env_spot = SpotEnv(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT, controller_type='pympc')
         env_spot.reset()
     return env_spot