Implement ML-Agents WebSocket bridge for web environments.

Author: lukehollis

api/.python-version

@@ -0,0 +1 @@
+3.11

api/README.md

@@ -0,0 +1,64 @@
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+import asyncio
+import json
+import logging
+import queue
+from typing import Dict
+
+logger = logging.getLogger("ThreeJSServer")
+
+def create_server(env_map: Dict[str, "ThreeJSEnv"]):
+    """
+    Creates a FastAPI app with a websocket endpoint that connects to ThreeJSEnvs.
+    env_map: dict of channel_id -> ThreeJSEnv
+    """
+    app = FastAPI()
+
+    @app.websocket("/ws/mlagents")
+    async def websocket_endpoint(ws: WebSocket):
+        await ws.accept()
+        logger.info("Client connected to /ws/mlagents")
+        
+        # We assume single channel/single environment for now, or use query param?
+        # Let's assume default channel
+        env = env_map.get("default")
+        
+        if not env:
+            await ws.close(code=1000, reason="No environment found")
+            return
+
+        # Connect queues
+        # We need async wrappers or run in thread
+        
+        try:
+            while True:
+                # 1. Check for outbound messages from Python (Poll non-blocking?)
+                # Since we are in async loop, we can't easily wait on a blocking Queue.
+                # But we can use run_in_executor or just poll with small sleep.
+                
+                # Check outbound (Python -> Browser)
+                try:
+                    while True:
+                         msg = env.outbound_queue.get_nowait()
+                         await ws.send_json(msg)
+                except queue.Empty:
+                    pass
+                
+                # 2. Check for inbound messages from Browser
+                # receive_json is awaitable
+                try:
+                    # We use wait_for to allow polling both directions
+                    data = await asyncio.wait_for(ws.receive_json(), timeout=0.01)
+                    # Put into inbound queue for Python to read
+                    env.inbound_queue.put(data)
+                except asyncio.TimeoutError:
+                    pass # continue loop
+                
+                await asyncio.sleep(0.001)
+
+        except WebSocketDisconnect:
+            logger.info("Client disconnected")
+            env.close()
+
+    return app

api/mlagents_bridge/__init__.py

@@ -0,0 +1,19 @@
+[project]
+name = "api"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "fastapi>=0.128.0",
+    "numpy>=2.4.0",
+    "onnx>=1.20.0",
+    "onnxruntime>=1.23.2",
+    "python-multipart>=0.0.21",
+    "shimmy[gym-v21]>=2.0.0",
+    "stable-baselines3[extra]>=2.7.1",
+    "tensorboard>=2.20.0",
+    "torch>=2.9.1",
+    "uvicorn>=0.40.0",
+    "websockets>=15.0.1",
+]

api/mlagents_bridge/env.py

@@ -0,0 +1,91 @@
+
+import threading
+import uvicorn
+import time
+import queue
+import logging
+import numpy as np
+import uuid
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("MigrationDemo")
+
+# Updated Imports
+from mlagents_bridge.env import ThreeJSEnv
+from mlagents_bridge.server import create_server
+from mlagents_envs.side_channel.environment_parameters_channel import EnvironmentParametersChannel
+
+def run_server(env):
+    """Starts the uvicorn server in a separate thread."""
+    env_map = {"default": env}
+    app = create_server(env_map)
+    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="warning")
+
+def main():
+    logger.info("Starting Migration Demo...")
+    
+    # 1. Create Side Channels
+    env_params = EnvironmentParametersChannel()
+    
+    # 2. Create Environment
+    # We pass the side channel list to the env
+    env = ThreeJSEnv(channel_id="default", side_channels=[env_params])
+    
+    # 3. Start Server in Thread
+    server_thread = threading.Thread(target=run_server, args=(env,), daemon=True)
+    server_thread.start()
+    
+    logger.info("Server started on port 8000. Waiting for client to connect...")
+    
+    # 4. Wait for handshake and Reset
+    try:
+        env.reset()
+        logger.info("Environment reset complete. Handshake successful.")
+    except Exception as e:
+        logger.error(f"Failed to reset: {e}")
+        return
+
+    # Set Initial Gravity
+    env_params.set_float_parameter("gravity", 9.81)
+
+    # 5. Run Loop
+    behavior_name = "Ball3D"  # Must match JS
+    
+    try:
+        for i in range(10000): # Run longer
+            # Get decision steps
+            decision_steps, terminal_steps = env.get_steps(behavior_name)
+            
+            # Log Obs Shapes occasionally
+            if i % 100 == 0:
+                n_agents = len(decision_steps)
+                logger.info(f"Step {i}: Agents={n_agents}")
+                for idx, obs in enumerate(decision_steps.obs):
+                    logger.info(f"  Obs[{idx}] Shape: {obs.shape}")
+
+            # Create random actions
+            if behavior_name in env.behavior_specs:
+                spec = env.behavior_specs[behavior_name]
+                n_agents = len(decision_steps)
+                
+                if n_agents > 0:
+                    action = spec.action_spec.random_action(n_agents)
+                    env.set_actions(behavior_name, action)
+            
+            # Step env
+            env.step()
+            
+            # Test Side Channel
+            if i % 100 == 0:
+                new_g = 5.0 + np.random.rand() * 10.0
+                logger.info(f"Step {i}: Setting Gravity to {new_g:.2f}")
+                env_params.set_float_parameter("gravity", float(new_g))
+                
+    except KeyboardInterrupt:
+        logger.info("Stopping...")
+    finally:
+        env.close()
+
+if __name__ == "__main__":
+    main()