Implement camera and environment reset functions in mujoco_server.py

- Added _reset_camera_for_current_robot to configure camera settings based on the active robot.
- Introduced _reset_active_environment to reset the simulation environment and camera together.
- Updated _handle_gym_ws_message to call _reset_active_environment on specific message types.
- Refactored reset logic in handle_ws_message to utilize the new reset functions, improving code organization and maintainability.

mujoco_server.py

@@ -122,6 +122,36 @@ cam.lookat = np.array([0, 0, 0.8])
 # Camera follow mode
 camera_follow = True
 
+def _reset_camera_for_current_robot() -> None:
+    """Reset the orbit camera to the defaults for the active robot."""
+    cam.azimuth = 135
+    cam.elevation = -20
+    if current_robot == "g1":
+        cam.distance = 3.0
+        cam.lookat = np.array([0.0, 0.0, 0.8])
+    elif current_robot in ("ur5", "ur5_t_push"):
+        cam.distance = 1.8
+        cam.lookat = np.array([0.3, 0.0, 0.6])
+        cam.azimuth = 150
+        cam.elevation = -25
+    else:
+        cam.distance = 2.5
+        cam.lookat = np.array([0.0, 0.0, 0.4])
+
+
+def _reset_active_environment() -> None:
+    """Reset the primary simulation environment and associated camera."""
+    with mujoco_lock:
+        if env is not None:
+            env.reset()
+    _reset_camera_for_current_robot()
+
+
+def _schedule_robot_switch(robot: str, scene: Optional[str]) -> None:
+    """Schedule a robot/scene switch for the simulation loop."""
+    global needs_robot_switch
+    needs_robot_switch = {"robot": robot, "scene": scene}
+
 MODEL_ROOT = Path(_nova_sim_dir) / "robots"
 
 def discover_scene_files(robot: str) -> list[str]:
@@ -736,6 +766,7 @@ def _handle_gym_ws_message(ws, data):
                     "info": _serialize_value(info),
                 },
             }
+            _reset_active_environment()
         elif msg_type == "step":
             action = payload.get("action", [])
             render = bool(payload.get("render", False))
@@ -755,9 +786,12 @@ def _handle_gym_ws_message(ws, data):
                 if frame_jpeg:
                     response["data"]["frame_jpeg"] = frame_jpeg
         elif msg_type == "configure":
-            robot = payload.get("robot", "ur5")
+            robot = payload.get("robot")
             scene = payload.get("scene")
             session.configure(robot, scene)
+            if robot:
+                print(f"Configure requested (gym): {robot} / scene: {scene}")
+                _schedule_robot_switch(robot, scene)
             response = {
                 "type": "gym_configured",
                 "data": {"robot": session.robot, "scene": session.scene},
@@ -990,29 +1024,14 @@ def handle_ws_message(ws, data):
                 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])
-        elif current_robot in ("ur5", "ur5_t_push"):
-            cam.distance = 1.8
-            cam.lookat = np.array([0.3, 0.0, 0.6])
-            cam.azimuth = 150
-            cam.elevation = -25
-        else:
-            cam.distance = 2.5
-            cam.lookat = np.array([0.0, 0.0, 0.4])
+        _reset_active_environment()
 
     elif msg_type == 'switch_robot':
         payload = data.get('data', {})
         robot = payload.get('robot', 'g1')
         scene = payload.get('scene')
         print(f"Robot switch requested: {robot} / scene: {scene}")
-        needs_robot_switch = {"robot": robot, "scene": scene}
+        _schedule_robot_switch(robot, scene)
 
     elif msg_type == 'configure':
         payload = data.get('data', {})
@@ -1020,7 +1039,7 @@ def handle_ws_message(ws, data):
         scene = payload.get('scene')
         if robot:
             print(f"Configure requested: {robot} / scene: {scene}")
-            needs_robot_switch = {"robot": robot, "scene": scene}
+            _schedule_robot_switch(robot, scene)
 
     elif msg_type == 'camera':
         payload = data.get('data', {})
@@ -3207,8 +3226,7 @@ def switch_robot_endpoint():
     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}
+    _schedule_robot_switch(robot, scene)
     return jsonify({"status": "pending", "robot": robot, "scene": scene}), 202