feat(server): add training broadcast server and Colab training FastAPI app

- TrainingBroadcastServer: fire-and-forget WS hub, stores full step history
in episode registry for /episode/replay, fixes publish() to use stored
event loop (asyncio.run_coroutine_threadsafe from training threads)
- server/app.py: new Colab training server with /ws/training, /targets,
/episode/demo, /episode/replay/:ep_id; mounts React build + viewer

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

server/app.py

@@ -0,0 +1,162 @@
+"""
+server/app.py — Training WebSocket server for Colab environment.
+
+Provides /ws/training for live streaming of RL training episodes to browsers.
+Mount at a publicly accessible URL in Colab (e.g., via ngrok or Colab's proxy).
+
+Usage in training:
+    from server.app import broadcast
+    broadcast.publish(episode_id, {"type": "episode_update", ...})
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import uvicorn
+from fastapi import FastAPI, HTTPException, WebSocket
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+
+from server.training_broadcast import TrainingBroadcastServer
+
+app = FastAPI(title="Optigami Training Server", version="1.0")
+
+# Allow cross-origin connections (Colab public URL → browser)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global broadcast server — import and use from training code
+broadcast = TrainingBroadcastServer()
+
+
+@app.on_event("startup")
+async def _store_loop() -> None:
+    """Capture the asyncio event loop so training threads can schedule coroutines."""
+    import asyncio
+    broadcast._loop = asyncio.get_running_loop()
+
+
+@app.websocket("/ws/training")
+async def training_ws(websocket: WebSocket) -> None:
+    """Spectator WebSocket endpoint. Viewers connect here to watch training."""
+    await broadcast.connect_spectator(websocket)
+
+
+@app.get("/health")
+def health() -> dict:
+    return {
+        "status": "ok",
+        "spectators": broadcast.spectator_count,
+        "active_episodes": broadcast.active_episodes,
+    }
+
+
+# ── Demo endpoints (same as openenv_server/app.py so the React UI works) ──
+
+@app.get("/targets")
+def get_targets() -> dict:
+    from server.tasks import available_task_names, get_task_by_name
+    return {
+        name: {
+            "name": name,
+            "level": t["difficulty"],
+            "description": t.get("description", ""),
+            "n_creases": t.get("max_folds", 3),
+            "difficulty": t["difficulty"],
+            "material": t.get("material", "paper"),
+        }
+        for name in available_task_names()
+        if (t := get_task_by_name(name))
+    }
+
+
+_DEMO_SEQUENCES: dict[str, list[dict]] = {
+    "half_fold":    [{"type": "valley",   "line": {"start": [0.0, 0.5],   "end": [1.0, 0.5]},   "angle": 180.0}],
+    "quarter_fold": [{"type": "valley",   "line": {"start": [0.0, 0.5],   "end": [1.0, 0.5]},   "angle": 180.0},
+                     {"type": "valley",   "line": {"start": [0.5, 0.0],   "end": [0.5, 1.0]},   "angle": 180.0}],
+    "letter_fold":  [{"type": "valley",   "line": {"start": [0.0, 0.333], "end": [1.0, 0.333]}, "angle": 180.0},
+                     {"type": "mountain", "line": {"start": [0.0, 0.667], "end": [1.0, 0.667]}, "angle": 180.0}],
+    "map_fold":     [{"type": "valley",   "line": {"start": [0.0, 0.5],   "end": [1.0, 0.5]},   "angle": 180.0},
+                     {"type": "mountain", "line": {"start": [0.5, 0.0],   "end": [0.5, 1.0]},   "angle": 180.0}],
+    "solar_panel":  [{"type": "valley",   "line": {"start": [0.0, 0.25],  "end": [1.0, 0.25]},  "angle": 180.0},
+                     {"type": "mountain", "line": {"start": [0.0, 0.5],   "end": [1.0, 0.5]},   "angle": 180.0},
+                     {"type": "valley",   "line": {"start": [0.0, 0.75],  "end": [1.0, 0.75]},  "angle": 180.0}],
+}
+
+
+@app.get("/episode/demo")
+def demo_episode(target: str = "half_fold") -> dict:
+    from server.origami_environment import OrigamiEnvironment
+    from server.models import OrigamiAction as NewAction
+    from server.tasks import get_task_by_name
+
+    folds = _DEMO_SEQUENCES.get(target, _DEMO_SEQUENCES["half_fold"])
+    env = OrigamiEnvironment()
+    obs = env.reset(task_name=target)
+    steps: list[dict] = []
+
+    for i, fold_dict in enumerate(folds):
+        action = NewAction(
+            fold_type=fold_dict["type"],
+            fold_line=fold_dict["line"],
+            fold_angle=float(fold_dict.get("angle", 180.0)),
+        )
+        obs = env.step(action)
+        steps.append({"step": i + 1, "fold": fold_dict,
+                       "paper_state": obs.paper_state, "metrics": obs.metrics,
+                       "done": obs.done})
+        if obs.done:
+            break
+
+    return {"task_name": target, "task": get_task_by_name(target) or {},
+            "steps": steps, "final_metrics": obs.metrics if steps else {}}
+
+
+@app.get("/episode/replay/{ep_id}")
+def replay_episode(ep_id: str) -> dict:
+    """Return a stored training episode in the same format as /episode/demo."""
+    from server.tasks import get_task_by_name
+    ep = broadcast._registry.get(ep_id)
+    if not ep:
+        raise HTTPException(status_code=404, detail=f"Episode '{ep_id}' not found in registry")
+    return {
+        "task_name": ep.task_name,
+        "task": get_task_by_name(ep.task_name) or {},
+        "steps": ep.steps,
+        "final_metrics": ep.final_metrics or (ep.steps[-1]["metrics"] if ep.steps else {}),
+    }
+
+
+# ── Static files — viewer first, then React app (LAST, catch-all) ──
+
+_VIEWER_DIR = Path(__file__).resolve().parent.parent / "viewer"
+_BUILD_DIR  = Path(__file__).resolve().parent.parent / "build"
+
+if _VIEWER_DIR.exists():
+    app.mount("/viewer", StaticFiles(directory=str(_VIEWER_DIR), html=True), name="viewer")
+
+
+if _BUILD_DIR.exists():
+    app.mount("/", StaticFiles(directory=str(_BUILD_DIR), html=True), name="react")
+else:
+    @app.get("/", include_in_schema=False)
+    def _no_build() -> HTMLResponse:
+        return HTMLResponse(
+            "<p>React build not found. Run <code>npm run build</code> in the frontend directory.</p>"
+            "<p>Training viewer: <a href='/viewer/training.html'>/viewer/training.html</a></p>"
+        )
+
+
+def run(host: str = "0.0.0.0", port: int = 9001) -> None:
+    """Start the training server. Call from Colab notebook."""
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    run()

server/training_broadcast.py

@@ -27,6 +27,7 @@ class EpisodeInfo:
     observation: dict = field(default_factory=dict)
     metrics: dict = field(default_factory=dict)
     fold_history: list = field(default_factory=list)
+    steps: list = field(default_factory=list)   # full step history for replay
     score: Optional[float] = None
     final_metrics: Optional[dict] = None
 
@@ -50,17 +51,13 @@ class TrainingBroadcastServer:
     def publish(self, episode_id: str, data: dict) -> None:
         """Fire-and-forget: push an update from the training process.
 
-        Safe to call from any thread. If no event loop is running, logs and returns.
+        Safe to call from any thread. Schedules onto the stored event loop
+        (set by the FastAPI startup handler). No-op if no loop is available.
         """
-        try:
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                asyncio.ensure_future(self._async_publish(episode_id, data), loop=loop)
-            else:
-                loop.run_until_complete(self._async_publish(episode_id, data))
-        except RuntimeError:
-            # No event loop — training without server
-            pass
+        loop = self._loop
+        if loop is None or loop.is_closed():
+            return
+        asyncio.run_coroutine_threadsafe(self._async_publish(episode_id, data), loop)
 
     async def _async_publish(self, episode_id: str, data: dict) -> None:
         msg_type = data.get("type", "episode_update")
@@ -91,12 +88,24 @@ class TrainingBroadcastServer:
                 ep.score = data.get("score")
                 ep.final_metrics = data.get("final_metrics")
             else:
-                ep.step = data.get("step", ep.step)
+                step_num = data.get("step", ep.step)
+                ep.step = step_num
                 ep.status = "running"
                 obs = data.get("observation", {})
                 ep.observation = obs
                 ep.metrics = obs.get("metrics", {})
                 ep.fold_history = obs.get("fold_history", ep.fold_history)
+                # Accumulate full step history for /episode/replay
+                if step_num > 0:
+                    fold_hist = obs.get("fold_history", [])
+                    latest_fold = fold_hist[-1] if fold_hist else {}
+                    ep.steps.append({
+                        "step": step_num,
+                        "fold": latest_fold,
+                        "paper_state": obs.get("paper_state", {}),
+                        "metrics": obs.get("metrics", {}),
+                        "done": obs.get("done", False),
+                    })
 
         await self._broadcast({"episode_id": episode_id, **data})