Enhance mujoco_server.py and documentation with camera and scene updates

- Updated the UR5 overlay camera configuration to include a new look target and adjusted offsets for improved visualization.
- Enhanced the simulation loop to support dynamic camera positioning based on site IDs and offsets.
- Added a new endpoint for scheduling robot and scene switches, allowing for automated testing and integration.
- Updated README.md to reflect changes in camera functionality and added instructions for the new switch endpoint.
- Introduced tests for API endpoints, camera streams, and WebSocket interactions to ensure reliability and performance.

README.md

@@ -247,7 +247,7 @@ docker run --gpus all -p 3004:3004 \
 
 ### UR5 Scene & Camera Hints
 - The UI now selects between exactly two UR5 options: the gripper-ready scene and the T-push scene (both enumerated via `/nova-sim/api/v1/metadata`). That metadata entry also returns `overlay_camera_presets`, so the trainer CLI can build dashboards based on the same stream list the UI exposes.
-- UR5 T-push streams three auxiliary MJPEG tiles (`aux_top`, `aux_side`, and the new `aux_flange` first-person camera) via `/nova-sim/api/v1/camera/<name>/video_feed`. Those overlays appear below the state panel only when the scene actually provides them, and the flange view mirrors a tool-mounted perspective so a trainer-side visualization can know what the robot “sees” without interfering with your primary viewport.
+- UR5 T-push streams three auxiliary MJPEG tiles (`aux_top`, `aux_side`, and the new `aux_flange` first-person camera) via `/nova-sim/api/v1/camera/<name>/video_feed`. Those overlays appear below the state panel only when the scene actually provides them; the flange camera now sits below the tool, offset slightly along the flange +X axis, and always looks toward the stick tip so the trainer knows what the robot is actually pressing even if the stick partially hides the table view.
 - The T-shape target stays anchored at its configured pose across resets, which keeps the training objective consistent even when you hit Reset from the UI.
 
 ## Architecture
@@ -552,16 +552,27 @@ The server responds with `gym_reset`, `gym_step`, `gym_spaces`, or `gym_configur
 |----------|--------|-------------|
 | `/nova-sim/api/v1` | GET | Web interface (HTML/JS) |
 | `/nova-sim/api/v1/video_feed` | GET | MJPEG video stream |
+| `/nova-sim/api/v1/switch` | POST | Schedule a robot/scene switch without the UI |
 
-**Video stream usage:**
-```html
-<img src="http://localhost:3004/nova-sim/api/v1/video_feed" />
-```
+ **Video stream usage:**
+ ```html
+ <img src="http://localhost:3004/nova-sim/api/v1/video_feed" />
+ ```
+
+  **Robot switching (HTTP)**
+  ```
+  POST /nova-sim/api/v1/switch
+  Content-Type: application/json
+
+  {"robot": "ur5_t_push", "scene": "scene_t_push"}
+  ```
+  Returns `202 Accepted` while the server asynchronously performs the switch. This endpoint enables tests or automation frameworks to select robots without opening the browser.
 
 ### Metadata & Camera Feeds
 
 - `GET /nova-sim/api/v1/metadata` returns JSON describing every available robot/scene pair, the configured defaults (UR5 T-push), the supported commands (including teleop commands), and camera feed definitions so RL trainers can build dynamic dashboards.
 - `GET /nova-sim/api/v1/camera/<name>/video_feed` streams MJPEG for the primary viewport (`main`) or additional overlays such as `aux_top` / `aux_side`.
+- `pytest tests/` exercises the HTTP metadata/video endpoints, the `/ws` control socket, and every camera feed. Keep Nova-Sim running at `http://localhost:3004` when you run it so the suite can talk to the live server.
 - The metadata payload also exposes `overlay_camera_presets`, which lists any auxiliary video streams (e.g., `aux_top`, `aux_side`, `aux_flange`) that the UI will render below the state panel. Those same stream names are the ones your RL trainer should subscribe to whenever they are present.
 
 ## Wandelbots Nova API Integration
@@ -786,6 +797,18 @@ The Nova API integration is implemented in:
 - [robots/ur5/nova_api.py](robots/ur5/nova_api.py) - API client and configuration
 - [robots/ur5/ur5_env.py](robots/ur5/ur5_env.py) - Environment integration (lines 123-203, 499-520)
 
+## Testing
+
+1. Start the Nova Sim server (e.g. `python nova-sim/mujoco_server.py` or via `docker-compose`).
+2. Keep it running at `http://localhost:3004` so the HTTP/websocket endpoints stay reachable.
+3. Run `pytest nova-sim/tests` to exercise:
+   - API endpoints (`/metadata`, `/camera/<name>/video_feed`, `/video_feed`)
+   - WebSocket control (`/ws`)
+   - Gym-style websocket (`/gym/ws`)
+   - Auxiliary MJPEG overlays after switching to the T-push UR5 scene
+
+The tests assume the server is accessible via `http://localhost:3004/nova-sim/api/v1` and will skip automatically if the API is unreachable.
+
 ## License
 
 This project uses models from:

mujoco_server.py

@@ -166,11 +166,13 @@ UR5_T_PUSH_OVERLAY_PRESETS = [
         "name": "aux_flange",
         "label": "Flange View",
         "follow_site": "ee_site",
+        "look_target_site": "stick_tip",
         "track_orientation": True,
-        "offset": [0.0, 0.0, 0.02],
-        "distance": 0.2,
-        "azimuth": 0,
-        "elevation": -5,
+        "offset": [0.03, 0.0, -0.04],
+        "forward_offset": 0.03,
+        "distance": 0.08,
+        "azimuth": 10,
+        "elevation": 20,
     },
 ]
 
@@ -183,7 +185,7 @@ CAMERA_FEEDS = [
     {"name": "main", "label": "Main", "description": "Primary viewport"},
     {"name": "aux_top", "label": "Aux Top", "description": "Additional top perspective"},
     {"name": "aux_side", "label": "Aux Side", "description": "Side perspective"},
-    {"name": "aux_flange", "label": "Flange", "description": "Tool-mounted first-person view"},
+    {"name": "aux_flange", "label": "Flange", "description": "Tool view mounted below the flange"},
 ]
 
 overlay_camera_states: dict[str, dict[str, Any]] = {}
@@ -201,15 +203,25 @@ def _make_mjv_camera(config: dict[str, float]) -> mujoco.MjvCamera:
 
 def _get_site_forward(env_obj, site_id: int) -> np.ndarray:
     """Compute the forward (X axis) vector for a site."""
+    default = np.array([0.0, 0.0, 1.0], dtype=np.float32)
     if env_obj is None or site_id < 0:
-        return np.array([0.0, 0.0, 1.0], dtype=np.float32)
-    xmat = env_obj.data.site_xmat
+        return default
+    if site_id >= getattr(env_obj.model, "nsite", 0):
+        return default
+    xmat = getattr(env_obj.data, "site_xmat", None)
+    if xmat is None or xmat.size < (site_id + 1) * 9:
+        return default
     start = site_id * 9
     mat = xmat[start:start + 9]
-    forward = np.array([mat[0], mat[3], mat[6]], dtype=np.float32)
+    if mat.size < 9:
+        return default
+    try:
+        forward = np.array([mat[0], mat[3], mat[6]], dtype=np.float32)
+    except IndexError:
+        return default
     norm = np.linalg.norm(forward)
     if norm < 1e-6:
-        return np.array([0.0, 0.0, 1.0], dtype=np.float32)
+        return default
     return forward / norm
 
 def _close_overlay_renderers():
@@ -250,8 +262,15 @@ def prepare_overlay_renderers(robot_type: str, scene_name: str | None):
             "follow_site": follow_site,
             "site_id": site_id,
             "offset": offset,
+            "forward_offset": float(config.get("forward_offset", 0.0)),
             "track_orientation": bool(config.get("track_orientation")),
+            "look_target_site": config.get("look_target_site"),
+            "look_target_id": -1,
         }
+        look_target_site = config.get("look_target_site")
+        if look_target_site and env is not None:
+            look_target_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_SITE, look_target_site)
+            overlay_camera_states[config["name"]]["look_target_id"] = look_target_id
         with overlay_frame_lock:
             overlay_frames[config["name"]] = None
 
@@ -598,10 +617,24 @@ def simulation_loop():
                     continue
                 if state.get("follow_site") and env is not None:
                     site_id = state.get("site_id", -1)
-                    if site_id >= 0:
-                        lookat = env.data.site_xpos[site_id]
+                    if site_id >= 0 and site_id < getattr(env.model, "nsite", 0):
+                        site_xpos = getattr(env.data, "site_xpos", None)
+                        if site_xpos is not None and site_xpos.size >= (site_id + 1) * 3:
+                            lookat_point = site_xpos[site_id]
+                        else:
+                            continue
                         offset = state.get("offset")
-                        cam_obj.lookat = lookat + offset if offset is not None else lookat
+                        if offset is not None:
+                            lookat_point = lookat_point + offset
+                        forward_offset = state.get("forward_offset", 0.0)
+                        if forward_offset:
+                            forward = _get_site_forward(env, site_id)
+                            lookat_point = lookat_point + forward * forward_offset
+                        look_target_id = state.get("look_target_id", -1)
+                        target_point = None
+                        if look_target_id >= 0 and site_xpos.size >= (look_target_id + 1) * 3:
+                            target_point = site_xpos[look_target_id]
+                        cam_obj.lookat = target_point if target_point is not None else lookat_point
                         if state.get("track_orientation"):
                             forward = _get_site_forward(env, site_id)
                             yaw = math.degrees(math.atan2(forward[1], forward[0]))
@@ -1225,12 +1258,12 @@ def index():
             .state-panel {
                 position: absolute; top: 20px; right: 20px;
                 background: rgba(33, 28, 68, 0.85);
-                backdrop-filter: blur(15px);
-                padding: 15px 20px; border-radius: 10px;
-                box-shadow: 0 4px 20px rgba(1, 4, 15, 0.5);
+                backdrop-filter: blur(12px);
+                padding: 8px 12px; border-radius: 8px;
+                box-shadow: 0 3px 14px rgba(1, 4, 15, 0.35);
                 color: var(--wb-secondary); border: 1px solid rgba(188, 190, 236, 0.15);
-                z-index: 100; min-width: 200px;
-                font-size: 0.8em; line-height: 1.5;
+                z-index: 100; min-width: 180px; max-width: 220px;
+                font-size: 0.75em; line-height: 1.3;
             }
             .state-panel strong { color: #fff; }
             /* Camera controls - bottom left */
@@ -1687,8 +1720,8 @@ def index():
                 'KeyW', 'KeyA', 'KeyS', 'KeyD', 'KeyQ', 'KeyE',
                 'ArrowUp', 'ArrowDown', 'ArrowLeft', 'ArrowRight'
             ]);
-            const TELEOP_STEP = 0.015;
-            const TELEOP_Z_STEP = 0.01;
+            let teleopTranslationStep = 0.05; // meters per keyboard nudge
+            let teleopVerticalStep = 0.01;
             let lastTeleopCommand = {dx: 0, dy: 0, dz: 0};
 
             function humanizeScene(scene) {
@@ -2194,6 +2227,7 @@ def index():
             function updateTransVelocity() {
                 transVelocity = parseFloat(document.getElementById('trans_velocity').value);
                 document.getElementById('trans_vel_val').innerText = transVelocity.toFixed(0);
+                teleopTranslationStep = transVelocity / 1000;
             }
 
             function updateRotVelocity() {
@@ -2285,12 +2319,12 @@ def index():
 
             function updateArmTeleopFromKeys() {
                 let dx = 0, dy = 0, dz = 0;
-                if (keysPressed.has('KeyW')) dx += TELEOP_STEP;
-                if (keysPressed.has('KeyS')) dx -= TELEOP_STEP;
-                if (keysPressed.has('KeyA')) dy += TELEOP_STEP;
-                if (keysPressed.has('KeyD')) dy -= TELEOP_STEP;
-                if (keysPressed.has('KeyR')) dz += TELEOP_Z_STEP;
-                if (keysPressed.has('KeyF')) dz -= TELEOP_Z_STEP;
+                if (keysPressed.has('KeyW')) dx += teleopTranslationStep;
+                if (keysPressed.has('KeyS')) dx -= teleopTranslationStep;
+                if (keysPressed.has('KeyA')) dy += teleopTranslationStep;
+                if (keysPressed.has('KeyD')) dy -= teleopTranslationStep;
+                if (keysPressed.has('KeyR')) dz += teleopVerticalStep;
+                if (keysPressed.has('KeyF')) dz -= teleopVerticalStep;
 
                 const unchanged =
                     Math.abs(dx - lastTeleopCommand.dx) < 1e-6 &&
@@ -2465,5 +2499,17 @@ def metadata():
     })
 
 
+@app.route(f'{API_PREFIX}/switch', methods=['POST'])
+def switch_robot_endpoint():
+    payload = request.get_json(silent=True) or {}
+    robot = payload.get('robot')
+    scene = payload.get('scene')
+    if not robot:
+        return jsonify({"error": "robot is required"}), 400
+    global needs_robot_switch
+    needs_robot_switch = {"robot": robot, "scene": scene}
+    return jsonify({"status": "pending", "robot": robot, "scene": scene}), 202
+
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=3004, debug=False, threaded=True)

requirements.txt

@@ -5,3 +5,7 @@ flask-sock>=0.7.0
 opencv-python>=4.8.0
 torch>=2.0.0
 numpy>=1.24.0
+pytest>=7.4.0
+pytest-asyncio>=0.23.0
+requests>=2.30.0
+websockets>=11.0

robots/ur5/model/scene_t_push.xml

@@ -15,7 +15,7 @@
     <global azimuth="120" elevation="-20"/>
   </visual>
 
-  <asset>
+    <asset>
     <!-- Wandelbots gradient skybox - deep purple to near black -->
     <texture type="skybox" builtin="gradient" rgb1="0.13 0.11 0.27" rgb2="0.004 0.016 0.059" width="512" height="3072"/>
     <!-- Ground with Wandelbots purple accent -->
@@ -31,6 +31,9 @@
     <material name="t_target_mat" rgba="0.2 0.7 0.35 0.25" specular="0.2" shininess="0.2"/>
     <material name="t_object_mat" rgba="0.55 0.65 0.98 1" specular="0.3" shininess="0.2"/>
     <material name="stick_mat" rgba="0.6 0.6 0.62 1" specular="0.4" shininess="0.3"/>
+    <material name="axis_x" rgba="1 0.2 0.2 1" specular="0.4" shininess="0.3"/>
+    <material name="axis_y" rgba="0.2 1 0.1 1" specular="0.4" shininess="0.3"/>
+    <material name="axis_z" rgba="0.25 0.4 1 1" specular="0.4" shininess="0.3"/>
 
     <!-- UR5e materials - with Wandelbots accent colors -->
     <material name="black" rgba="0.02 0.02 0.04 1" specular="0.5" shininess="0.25"/>
@@ -105,7 +108,7 @@
     </body>
 
     <!-- Target visualization sphere (for IK target) -->
-    <body name="target" pos="0.4 0.0 0.6" mocap="true">
+    <body name="target" pos="0.4 0.0 0.72" mocap="true">
       <geom name="target_vis" type="sphere" size="0.03" material="target_mat" contype="0" conaffinity="0"/>
     </body>
 
@@ -167,6 +170,7 @@
                     <!-- In flange frame: +Z points away from robot (downward when wrist is horizontal) -->
                     <geom name="push_stick" type="cylinder" fromto="0 0 0 0 0 0.18" size="0.008"
                           material="stick_mat" mass="0.02" friction="1 0.01 0.01"/>
+                    <site name="stick_tip" pos="0 0 0.18" type="sphere" size="0.005" rgba="0.6 0.9 0.3 0.2"/>
                   </body>
 
                   <!-- End-effector site for IK (at flange) -->
@@ -180,17 +184,27 @@
     </body>
 
     <!-- T-shaped target (visual marker) -->
-    <body name="t_target" pos="0.62 -0.18 0.425">
+    <body name="t_target" pos="0.62 -0.18 0.42">
       <geom name="t_target_stem" type="box" pos="0 -0.05 0" size="0.02 0.07 0.002" material="t_target_mat" contype="0" conaffinity="0"/>
       <geom name="t_target_cap" type="box" pos="0 0.03 0" size="0.08 0.02 0.002" material="t_target_mat" contype="0" conaffinity="0"/>
     </body>
 
     <!-- Movable T-shaped object to push into target -->
-    <body name="t_object" pos="0.45 0.2 0.43">
+    <body name="t_object" pos="0.45 0.25 0.46">
       <freejoint name="t_object_joint"/>
       <geom name="t_object_stem" type="box" pos="0 -0.05 0" size="0.02 0.07 0.03" material="t_object_mat" mass="3.0" friction="0.3 0.005 0.005"/>
       <geom name="t_object_cap" type="box" pos="0 0.03 0" size="0.08 0.02 0.03" material="t_object_mat" mass="2.0" friction="0.3 0.005 0.005"/>
     </body>
+
+    <!-- Axis arrows (XYZ) for orientation -->
+    <body name="axis_arrows" pos="0.35 0.6 0.43">
+      <geom type="capsule" fromto="0 0 0 0.12 0 0" size="0.006" material="axis_x" contype="0" conaffinity="0"/>
+      <geom type="sphere" pos="0.13 0 0" size="0.008" material="axis_x" contype="0" conaffinity="0"/>
+      <geom type="capsule" fromto="0 0 0 0 0.12 0" size="0.006" material="axis_y" contype="0" conaffinity="0"/>
+      <geom type="sphere" pos="0 0.13 0" size="0.008" material="axis_y" contype="0" conaffinity="0"/>
+      <geom type="capsule" fromto="0 0 0 0 0 0.12" size="0.006" material="axis_z" contype="0" conaffinity="0"/>
+      <geom type="sphere" pos="0 0 0.13" size="0.008" material="axis_z" contype="0" conaffinity="0"/>
+    </body>
   </worldbody>
 
   <actuator>

tests/conftest.py

@@ -0,0 +1,27 @@
+import os
+import requests
+import pytest
+
+API_BASE = os.getenv("NOVA_SIM_API_URL", "http://localhost:3004/nova-sim/api/v1")
+WS_BASE = os.getenv("NOVA_SIM_WS_URL", "ws://localhost:3004/nova-sim/api/v1/ws")
+
+
+def _ping_api() -> None:
+    try:
+        resp = requests.get(API_BASE, timeout=3)
+        resp.raise_for_status()
+    except (requests.RequestException, ValueError) as exc:
+        pytest.skip(f"Nova-Sim API not reachable at {API_BASE}: {exc}")
+
+
+@pytest.fixture(scope="session")
+def api_base():
+    _ping_api()
+    return API_BASE
+
+
+@pytest.fixture(scope="session")
+def ws_url():
+    # Ensure API is accessible before touching the websocket
+    _ping_api()
+    return WS_BASE

tests/test_api.py

@@ -0,0 +1,44 @@
+import requests
+from requests import Response
+
+def test_root_ui(api_base: str):
+    resp = requests.get(api_base, timeout=5)
+    assert resp.status_code == 200
+    assert '<title>Nova Sim' in resp.text
+
+
+def test_metadata(api_base: str):
+    resp = requests.get(f"{api_base}/metadata", timeout=5)
+    assert resp.status_code == 200
+    payload = resp.json()
+    assert 'robots' in payload and isinstance(payload['robots'], dict)
+    assert 'commands' in payload
+    assert 'camera_feeds' in payload
+
+
+def test_camera_feeds_structure(api_base: str):
+    resp = requests.get(f"{api_base}/metadata", timeout=5)
+    data = resp.json()
+    feeds = data.get('camera_feeds') or []
+    assert any(feed.get('name') == 'main' for feed in feeds)
+    assert any(feed.get('name') == 'aux_flange' for feed in feeds)
+
+
+def test_overlay_camera_presets(api_base: str):
+    resp = requests.get(f"{api_base}/metadata", timeout=5)
+    data = resp.json()
+    presets = data.get('overlay_camera_presets', {})
+    assert 'ur5_t_push' in presets or 'scene_t_push' in presets
+    target_presets = presets.get('ur5_t_push') or presets.get('scene_t_push') or []
+    names = {item.get('name') for item in target_presets}
+    assert {'aux_top', 'aux_side', 'aux_flange'}.issubset(names)
+
+
+def test_switch_endpoint(api_base: str):
+    # Missing robot is rejected
+    resp = requests.post(f"{api_base}/switch", json={"scene": "scene_t_push"}, timeout=5)
+    assert resp.status_code == 400
+
+    # Valid request schedules the switch
+    resp = requests.post(f"{api_base}/switch", json={"robot": "ur5_t_push", "scene": "scene_t_push"}, timeout=5)
+    assert resp.status_code == 202

tests/test_video.py

@@ -0,0 +1,34 @@
+import time
+import requests
+import pytest
+
+def _assert_mjpeg_header(resp: requests.Response) -> None:
+    content_type = resp.headers.get('Content-Type', '')
+    assert 'multipart/x-mixed-replace' in content_type
+
+
+def test_main_video_stream(api_base: str):
+    with requests.get(f"{api_base}/video_feed", stream=True, timeout=5) as resp:
+        assert resp.status_code == 200
+        _assert_mjpeg_header(resp)
+        chunks = resp.iter_content(chunk_size=1024)
+        first_chunk = next(chunks)
+        assert first_chunk, 'Expected non-empty MJPEG chunk'
+
+
+def _validate_mjpeg_stream(api_base: str, camera_name: str) -> None:
+    with requests.get(f"{api_base}/camera/{camera_name}/video_feed", stream=True, timeout=5) as resp:
+        assert resp.status_code in (200, 304), f'Stream {camera_name} returned {resp.status_code}'
+        _assert_mjpeg_header(resp)
+        chunk = next(resp.iter_content(chunk_size=1024))
+        assert chunk, f'MJPEG stream {camera_name} yielded empty chunk'
+
+
+def test_aux_camera_streams_after_switch(api_base: str):
+    resp = requests.post(f"{api_base}/switch", json={"robot": "ur5_t_push", "scene": "scene_t_push"}, timeout=5)
+    assert resp.status_code == 202, f"Switch request failed: {resp.status_code}"
+    time.sleep(3)
+
+    camera_names = ['aux_top', 'aux_side', 'aux_flange']
+    for name in camera_names:
+        _validate_mjpeg_stream(api_base, name)

tests/test_ws.py

@@ -0,0 +1,17 @@
+import asyncio
+import json
+import pytest
+import websockets
+
+@pytest.mark.asyncio
+async def test_ws_state_message(ws_url: str):
+    async with websockets.connect(ws_url, ping_interval=None) as ws:
+        raw = await asyncio.wait_for(ws.recv(), timeout=5)
+        payload = json.loads(raw)
+        assert payload.get('type') == 'state'
+        data = payload.get('data', {})
+        assert 'robot' in data
+        # send a no-op command to prove WS accepts messages
+        await ws.send(json.dumps({'type': 'command', 'data': {'vx': 0.0, 'vy': 0.0, 'vyaw': 0.0}}))
+        # read an acknowledgement (state) again
+        await asyncio.wait_for(ws.recv(), timeout=5)