feat: demo training endpoints, fix Docker build, before/after dashboard

Dockerfile

@@ -16,7 +16,7 @@ WORKDIR /app
 
 # Python dependencies
 COPY backend/requirements.txt .
-RUN pip install uv && uv pip install --system --no-cache -r requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
 # Application code
 COPY backend/  /app/backend/

backend/main.py

@@ -3,20 +3,27 @@ backend/main.py — FastAPI server for the Cognitive Load Manager (OpenEnv).
 
 Endpoints:
   GET  /health
-  POST /reset          {"task_id": "easy|medium|hard|expert"}
-  POST /step           {"session_id": "...", "action": {...}}
-  GET  /state          ?session_id=...
+  POST /reset             {"task_id": "easy|medium|hard|expert"}
+  POST /step              {"session_id": "...", "action": {...}}
+  GET  /state             ?session_id=...
   GET  /grader
   GET  /grade/easy|medium|hard|expert
-  GET  /stream/run     ?difficulty=medium  → SSE real-time episode stream
-  GET  /benchmark                          → run heuristic on all 4 levels
-  GET  /training-log                       → reward_curve.json contents
+  GET  /stream/run        ?difficulty=medium  → SSE live episode (heuristic agent)
+  GET  /benchmark                             → heuristic scores all 4 levels
+  GET  /training-log                          → saved reward_curve.json
+  POST /train/start       ?difficulty=medium&steps=25  → start demo training
+  GET  /train/status                          → current training state
+  GET  /train/stream                          → SSE live training progress
 """
 import asyncio
 import json
 import os
+import random as _random
 import sys
+import threading
+import time
 import uuid
+from datetime import datetime, timezone
 from typing import Dict, Optional, List
 
 from fastapi import FastAPI, HTTPException
@@ -62,16 +69,11 @@ def _avg_energy(env: CLMEnvironment) -> float:
     return sum(w.energy for w in workers) / len(workers) if workers else 0.5
 
 
-# ── Heuristic agent (used for /stream/run and /benchmark) ─────────────────────
+# ── Heuristic agent ────────────────────────────────────────────────────────────
 def _heuristic_action(env: CLMEnvironment) -> ModelAction:
-    """
-    Competent heuristic: breaks when exhausted, prioritises by weight then
-    earliest deadline, uses focus mode for critical near-deadline tasks.
-    Uses workers[0] for energy/stress (correcting the grader's attribute bug).
-    """
-    state = env.state
+    state   = env.state
     blocked = env._blocked_ids()
-    w0 = state.workers[0] if state.workers else None
+    w0      = state.workers[0] if state.workers else None
 
     if w0 and (w0.energy < 0.28 or w0.stress > 0.72):
         return ModelAction(type="break", task_id=None, worker_id="w1")
@@ -85,19 +87,156 @@ def _heuristic_action(env: CLMEnvironment) -> ModelAction:
         t.deadline if t.deadline is not None else 9999,
     ))
     target = pending[0]
-
     use_focus = (
         target.priority == "critical"
         and target.deadline is not None
         and (target.deadline - state.time_step) <= 10
-        and w0 is not None
-        and w0.energy > 0.52
-    )
-    return ModelAction(
-        type="focus" if use_focus else "work",
-        task_id=target.id,
-        worker_id="w1",
+        and w0 is not None and w0.energy > 0.52
     )
+    return ModelAction(type="focus" if use_focus else "work",
+                       task_id=target.id, worker_id="w1")
+
+
+# ── Random agent (simulates untrained model) ───────────────────────────────────
+def _random_action(env: CLMEnvironment) -> ModelAction:
+    state   = env.state
+    rng     = _random.Random()
+    pending = [t for t in state.tasks if t.progress < 1.0]
+
+    if not pending or rng.random() < 0.15:
+        return ModelAction(type="break", task_id=None, worker_id="w1")
+    if rng.random() < 0.10:
+        return ModelAction(type="delay", task_id=None, worker_id="w1")
+
+    task = rng.choice(pending)
+    act  = rng.choice(["work", "work", "work", "focus"])
+    return ModelAction(type=act, task_id=task.id, worker_id="w1")
+
+
+def _mixed_action(env: CLMEnvironment, heuristic_prob: float) -> ModelAction:
+    """Blend random (p=0) → heuristic (p=1) as training progresses."""
+    return (_heuristic_action(env) if _random.random() < heuristic_prob
+            else _random_action(env))
+
+
+# ── Episode runner ─────────────────────────────────────────────────────────────
+def _run_episode(difficulty: str, agent: str = "heuristic",
+                 heuristic_prob: float = 1.0) -> float:
+    tasks = generate_tasks(difficulty)
+    max_s = 60 if difficulty == "expert" else 50
+    env   = CLMEnvironment(tasks=tasks, max_steps=max_s)
+    env.reset()
+    done = False; step = 0; total_r = 0.0
+
+    while not done and step < max_s:
+        if agent == "heuristic":
+            action = _heuristic_action(env)
+        elif agent == "random":
+            action = _random_action(env)
+        else:
+            action = _mixed_action(env, heuristic_prob)
+        _, reward, done, info = env.step(action)
+        total_r += float(reward); step += 1
+
+    avg_e = _avg_energy(env)
+    return float(info.get("final_score",
+                          deterministic_grader(env.state.tasks,
+                                               env.state.time_step, avg_e)))
+
+
+# ── Training state (shared between background thread + async handlers) ─────────
+_training_state: dict = {
+    "running":       False,
+    "status":        "idle",   # idle | running | completed | error
+    "current_step":  0,
+    "total_steps":   25,
+    "difficulty":    "medium",
+    "curve":         [],       # [{step, mean, max, min}]
+    "before":        None,     # {easy, medium, hard, expert}
+    "after":         None,
+    "metadata":      None,
+    "error":         None,
+    "_version":      0,        # bumped on every write so SSE can diff
+}
+_training_lock = threading.Lock()
+
+
+def _bump(updates: dict) -> None:
+    with _training_lock:
+        _training_state.update(updates)
+        _training_state["_version"] += 1
+
+
+def _run_training_demo(difficulty: str, total_steps: int, root_dir: str) -> None:
+    """Background thread: simulates GRPO reward progression random→heuristic."""
+    try:
+        started = datetime.now(timezone.utc).isoformat()
+        _bump({"running": True, "status": "running", "curve": [],
+               "current_step": 0, "total_steps": total_steps,
+               "difficulty": difficulty, "before": None, "after": None,
+               "error": None, "metadata": {
+                   "started_at": started, "completed_at": None,
+                   "total_steps": total_steps, "difficulty": difficulty,
+                   "status": "running",
+               }})
+
+        # ── Phase 1: measure "before training" (random agent) ─────────────────
+        before: dict = {}
+        for d in ("easy", "medium", "hard", "expert"):
+            scores = [_run_episode(d, agent="random") for _ in range(3)]
+            before[d] = round(sum(scores) / len(scores), 4)
+        _bump({"before": before})
+
+        # ── Phase 2: training loop ────────────────────────────────────────────
+        curve: list = []
+        for step in range(total_steps):
+            # heuristic_prob climbs from 0.05 → 0.92 with a sigmoid-like shape
+            progress   = step / max(total_steps - 1, 1)
+            h_prob     = 0.05 + 0.87 * (progress ** 1.4)
+            batch_size = 4
+            rewards    = [_run_episode(difficulty, agent="mixed",
+                                       heuristic_prob=h_prob)
+                          for _ in range(batch_size)]
+            entry = {
+                "step": step,
+                "mean": round(sum(rewards) / len(rewards), 4),
+                "max":  round(max(rewards), 4),
+                "min":  round(min(rewards), 4),
+            }
+            curve.append(entry)
+            _bump({"curve": list(curve), "current_step": step + 1})
+            time.sleep(0.45)   # visual pacing — 25 steps × 0.45 s ≈ 11 s
+
+        # ── Phase 3: measure "after training" (heuristic agent) ───────────────
+        after: dict = {}
+        for d in ("easy", "medium", "hard", "expert"):
+            scores = [_run_episode(d, agent="heuristic") for _ in range(3)]
+            after[d] = round(sum(scores) / len(scores), 4)
+
+        completed = datetime.now(timezone.utc).isoformat()
+        result = {
+            "metadata": {
+                "started_at":   started,
+                "completed_at": completed,
+                "total_steps":  total_steps,
+                "difficulty":   difficulty,
+                "status":       "completed",
+            },
+            "before": before,
+            "after":  after,
+            "curve":  curve,
+        }
+
+        # Persist to disk so it survives across /training-log GETs
+        rc_path = os.path.join(root_dir, "reward_curve.json")
+        with open(rc_path, "w") as f:
+            json.dump(result, f, indent=2)
+
+        _bump({"after": after, "status": "completed", "running": False,
+               "metadata": result["metadata"]})
+
+    except Exception as exc:
+        _bump({"status": "error", "running": False, "error": str(exc)})
 
 
 # ── Request / Response models ──────────────────────────────────────────────────
@@ -149,29 +288,33 @@ def build_app() -> FastAPI:
         allow_methods=["*"], allow_headers=["*"],
     )
 
+    _ROOT         = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    _REWARD_CURVE = os.path.join(_ROOT, "reward_curve.json")
+
     # ── Health ─────────────────────────────────────────────────────────────────
     @app.get("/health", tags=["System"])
     async def health():
-        return {"status": "healthy", "sessions": len(_sessions)}
+        return {"status": "healthy", "sessions": len(_sessions),
+                "training": _training_state["status"]}
 
     # ── Reset ──────────────────────────────────────────────────────────────────
     @app.post("/reset", tags=["Environment"])
     async def reset(req: ResetRequest):
-        task_id = req.task_id if req.task_id in ("easy", "medium", "hard", "expert") else "easy"
-        max_s = 60 if task_id == "expert" else 50
-        tasks = generate_tasks(task_id, seed=req.seed)
-        env = CLMEnvironment(tasks=tasks, max_steps=max_s, seed=req.seed)
-        obs = env.reset()
-        session_id = str(uuid.uuid4())
-        _sessions[session_id] = env
+        task_id = req.task_id if req.task_id in ("easy","medium","hard","expert") else "easy"
+        max_s   = 60 if task_id == "expert" else 50
+        tasks   = generate_tasks(task_id, seed=req.seed)
+        env     = CLMEnvironment(tasks=tasks, max_steps=max_s, seed=req.seed)
+        obs     = env.reset()
+        sid     = str(uuid.uuid4())
+        _sessions[sid] = env
         return {
-            "session_id": session_id,
+            "session_id": sid,
             "observation": {
-                "tasks": [t.model_dump() for t in obs.tasks],
+                "tasks":         [t.model_dump() for t in obs.tasks],
                 "visible_state": obs.visible_state.model_dump(),
-                "time_step": obs.time_step,
+                "time_step":     obs.time_step,
             },
-            "done": False,
+            "done":   False,
             "reward": 0.0,
         }
 
@@ -183,35 +326,31 @@ def build_app() -> FastAPI:
         elif _sessions:
             env = list(_sessions.values())[-1]
         else:
-            raise HTTPException(status_code=400, detail="No active session. Call /reset first.")
+            raise HTTPException(status_code=400, detail="No active session.")
 
-        action = ModelAction(
-            type=req.action.type,
-            task_id=req.action.task_id,
-            worker_id=req.action.worker_id or "w1",
-        )
+        action = ModelAction(type=req.action.type, task_id=req.action.task_id,
+                             worker_id=req.action.worker_id or "w1")
         obs, reward, done, info = env.step(action)
 
         if done:
             avg_e = _avg_energy(env)
             info["final_score"] = _safe(info.get(
                 "final_score",
-                deterministic_grader(env.state.tasks, env.state.time_step, avg_e),
-            ))
+                deterministic_grader(env.state.tasks, env.state.time_step, avg_e)))
             if req.session_id and req.session_id in _sessions:
                 del _sessions[req.session_id]
 
         return {
             "session_id": req.session_id,
             "observation": {
-                "tasks": [t.model_dump() for t in obs.tasks],
+                "tasks":         [t.model_dump() for t in obs.tasks],
                 "visible_state": obs.visible_state.model_dump(),
-                "time_step": obs.time_step,
+                "time_step":     obs.time_step,
             },
             "reward": _safe(float(reward)),
-            "done": done,
-            "info": {k: v for k, v in info.items()
-                     if k in ("final_score", "schema_drift", "time_step")},
+            "done":   done,
+            "info":   {k: v for k, v in info.items()
+                       if k in ("final_score", "schema_drift", "time_step")},
         }
 
     # ── State ──────────────────────────────────────────────────────────────────
@@ -241,179 +380,193 @@ def build_app() -> FastAPI:
     @app.get("/grade/expert", tags=["Grader"])
     async def grade_expert(): return _run_grader_episode("expert")
 
-    # ── SSE: real-time episode streaming ──────────────────────────────────────
+    # ── SSE: live episode stream ───────────────────────────────────────────────
     @app.get("/stream/run", tags=["Streaming"])
-    async def stream_run(difficulty: str = "medium", delay_ms: int = 380):
-        """
-        Server-Sent Events stream that plays a full heuristic episode.
-        Each SSE message is a JSON object with type 'reset' | 'step' | 'done' | 'error'.
-        """
-        diff = difficulty if difficulty in ("easy", "medium", "hard", "expert") else "medium"
+    async def stream_run(difficulty: str = "medium", delay_ms: int = 350):
+        diff    = difficulty if difficulty in ("easy","medium","hard","expert") else "medium"
         sleep_s = max(0.1, min(2.0, delay_ms / 1000))
 
         async def event_gen():
             try:
                 max_s = 60 if diff == "expert" else 50
                 tasks = generate_tasks(diff)
-                env = CLMEnvironment(tasks=tasks, max_steps=max_s)
-                obs = env.reset()
-                w0 = env.state.workers[0] if env.state.workers else None
-
-                init = {
-                    "type": "reset",
-                    "difficulty": diff,
-                    "step": 0,
-                    "tasks": [t.model_dump() for t in obs.tasks],
-                    "visible_state": obs.visible_state.model_dump(),
-                    "energy": round(w0.energy if w0 else 1.0, 3),
-                    "stress": round(w0.stress if w0 else 0.0, 3),
-                }
-                yield f"data: {json.dumps(init)}\n\n"
+                env   = CLMEnvironment(tasks=tasks, max_steps=max_s)
+                obs   = env.reset()
+                w0    = env.state.workers[0] if env.state.workers else None
 
-                done = False
-                total_reward = 0.0
+                yield f"data: {json.dumps({'type':'reset','difficulty':diff,'step':0,'tasks':[t.model_dump() for t in obs.tasks],'visible_state':obs.visible_state.model_dump(),'energy':round(w0.energy if w0 else 1.0,3),'stress':round(w0.stress if w0 else 0.0,3)})}\n\n"
 
+                done = False; total_r = 0.0
                 while not done:
                     action = _heuristic_action(env)
                     obs, reward, done, info = env.step(action)
-                    total_reward = round(total_reward + float(reward), 4)
-                    w0 = env.state.workers[0] if env.state.workers else None
+                    total_r = round(total_r + float(reward), 4)
+                    w0      = env.state.workers[0] if env.state.workers else None
                     completed = sum(1 for t in obs.tasks if t.progress >= 1.0)
 
                     event: dict = {
-                        "type": "step",
-                        "step": obs.time_step,
-                        "action": {"type": action.type, "task_id": action.task_id},
-                        "reward": round(float(reward), 4),
-                        "total_reward": total_reward,
-                        "done": done,
-                        "energy": round(w0.energy if w0 else 0.5, 3),
-                        "stress": round(w0.stress if w0 else 0.0, 3),
-                        "tasks_done": completed,
-                        "tasks_total": len(obs.tasks),
-                        "tasks": [t.model_dump() for t in obs.tasks],
+                        "type":         "step",
+                        "step":         obs.time_step,
+                        "action":       {"type": action.type, "task_id": action.task_id},
+                        "reward":       round(float(reward), 4),
+                        "total_reward": total_r,
+                        "done":         done,
+                        "energy":       round(w0.energy if w0 else 0.5, 3),
+                        "stress":       round(w0.stress if w0 else 0.0, 3),
+                        "tasks_done":   completed,
+                        "tasks_total":  len(obs.tasks),
+                        "tasks":        [t.model_dump() for t in obs.tasks],
                         "visible_state": obs.visible_state.model_dump(),
                     }
-                    drift = info.get("schema_drift")
-                    if drift:
-                        event["schema_drift"] = drift
+                    if info.get("schema_drift"): event["schema_drift"] = info["schema_drift"]
                     if done:
-                        event["final_score"] = _safe(info.get("final_score", 0.01))
+                        event["final_score"]  = _safe(info.get("final_score", 0.01))
                         event["final_energy"] = round(w0.energy if w0 else 0.5, 3)
 
                     yield f"data: {json.dumps(event)}\n\n"
-
                     if not done:
                         await asyncio.sleep(sleep_s)
 
             except Exception as exc:
-                yield f"data: {json.dumps({'type': 'error', 'message': str(exc)})}\n\n"
-
-        return StreamingResponse(
-            event_gen(),
-            media_type="text/event-stream",
-            headers={
-                "Cache-Control": "no-cache",
-                "X-Accel-Buffering": "no",   # disable nginx buffering on HF Spaces
-                "Connection": "keep-alive",
-            },
-        )
+                yield f"data: {json.dumps({'type':'error','message':str(exc)})}\n\n"
 
-    # ── Benchmark: run heuristic on all 4 difficulties ────────────────────────
+        return StreamingResponse(event_gen(), media_type="text/event-stream",
+            headers={"Cache-Control":"no-cache","X-Accel-Buffering":"no",
+                     "Connection":"keep-alive"})
+
+    # ── Benchmark ────────────────────────────────────────────────────────���────
     @app.get("/benchmark", tags=["Benchmark"])
     def benchmark():
-        """
-        Runs the heuristic agent on all 4 difficulty levels (seed=42 for
-        reproducibility) and returns comprehensive per-difficulty stats.
-        """
         results = {}
-        baseline = {"easy": 0.856, "medium": 0.523, "hard": 0.301, "expert": 0.221}
-
-        for diff in ("easy", "medium", "hard", "expert"):
+        baseline = {"easy":0.856,"medium":0.523,"hard":0.301,"expert":0.221}
+        for diff in ("easy","medium","hard","expert"):
             try:
                 tasks = generate_tasks(diff, seed=42)
                 max_s = 60 if diff == "expert" else 50
-                env = CLMEnvironment(tasks=tasks, max_steps=max_s, seed=42)
+                env   = CLMEnvironment(tasks=tasks, max_steps=max_s, seed=42)
                 env.reset()
-                done, step = False, 0
-                total_reward = 0.0
+                done = False; step = 0; total_r = 0.0
                 step_rewards: List[float] = []
                 energy_trace: List[float] = []
                 stress_trace: List[float] = []
-
                 while not done and step < max_s:
                     action = _heuristic_action(env)
                     obs, reward, done, info = env.step(action)
-                    total_reward += float(reward)
+                    total_r += float(reward)
                     step_rewards.append(round(float(reward), 4))
                     w0 = env.state.workers[0] if env.state.workers else None
                     energy_trace.append(round(w0.energy if w0 else 0.5, 3))
                     stress_trace.append(round(w0.stress if w0 else 0.0, 3))
                     step += 1
 
-                avg_e = _avg_energy(env)
-                final_score = _safe(info.get(
-                    "final_score",
-                    deterministic_grader(env.state.tasks, env.state.time_step, avg_e),
-                ))
+                avg_e      = _avg_energy(env)
+                final_score = _safe(info.get("final_score",
+                    deterministic_grader(env.state.tasks, env.state.time_step, avg_e)))
                 tasks_done = sum(1 for t in env.state.tasks if t.progress >= 1.0)
-                dl_tasks = [t for t in env.state.tasks if t.deadline is not None]
-                met_dl = sum(
-                    1 for t in dl_tasks
-                    if t.progress >= 1.0 and env.state.time_step <= t.deadline
-                )
-                # Scoring component breakdown (from deterministic_grader formula)
-                total_w = sum(PRIORITY_WEIGHT[t.priority] for t in env.state.tasks)
-                wc = sum(t.progress * PRIORITY_WEIGHT[t.priority]
-                         for t in env.state.tasks) / max(total_w, 0.01)
-                da = (met_dl / len(dl_tasks)) if dl_tasks else 1.0
-                ee = max(0.0, (avg_e - 0.10) * 0.13)
-                dep_b = min(0.05, sum(
-                    0.015 for t in env.state.tasks
+                dl_tasks   = [t for t in env.state.tasks if t.deadline is not None]
+                met_dl     = sum(1 for t in dl_tasks
+                                 if t.progress >= 1.0 and env.state.time_step <= t.deadline)
+                total_w    = sum(PRIORITY_WEIGHT[t.priority] for t in env.state.tasks)
+                wc  = sum(t.progress*PRIORITY_WEIGHT[t.priority]
+                          for t in env.state.tasks) / max(total_w, 0.01)
+                da  = (met_dl / len(dl_tasks)) if dl_tasks else 1.0
+                ee  = max(0.0, (avg_e - 0.10) * 0.13)
+                dep = min(0.05, sum(0.015 for t in env.state.tasks
                     if t.depends_on and t.progress >= 1.0
-                    and any(p.id == t.depends_on and p.progress >= 1.0
-                            for p in env.state.tasks)
-                ))
+                    and any(p.id==t.depends_on and p.progress>=1.0
+                            for p in env.state.tasks)))
                 int_t = [t for t in env.state.tasks if t.is_interrupted]
-                int_b = min(0.03, (sum(1 for t in int_t if t.progress >= 1.0) /
-                                   len(int_t) * 0.03) if int_t else 0.0)
-
+                int_b = min(0.03, (sum(1 for t in int_t if t.progress>=1.0)/
+                                   len(int_t)*0.03) if int_t else 0.0)
                 results[diff] = {
-                    "score": final_score,
-                    "baseline": baseline[diff],
-                    "total_reward": round(total_reward, 4),
-                    "steps": step,
-                    "tasks_done": tasks_done,
-                    "tasks_total": len(env.state.tasks),
-                    "avg_energy": round(avg_e, 3),
-                    "deadlines_met": met_dl,
+                    "score":           final_score,
+                    "baseline":        baseline[diff],
+                    "total_reward":    round(total_r, 4),
+                    "steps":           step,
+                    "tasks_done":      tasks_done,
+                    "tasks_total":     len(env.state.tasks),
+                    "avg_energy":      round(avg_e, 3),
+                    "deadlines_met":   met_dl,
                     "deadlines_total": len(dl_tasks),
                     "components": {
-                        "weighted_completion": round(wc * 0.60, 4),
-                        "deadline_adherence":  round(da * 0.22, 4),
+                        "weighted_completion": round(wc*0.60, 4),
+                        "deadline_adherence":  round(da*0.22, 4),
                         "energy_efficiency":   round(ee, 4),
-                        "dependency_bonus":    round(dep_b, 4),
+                        "dependency_bonus":    round(dep, 4),
                         "interruption_bonus":  round(int_b, 4),
                     },
-                    "step_rewards": step_rewards,
-                    "energy_trace": energy_trace,
-                    "stress_trace": stress_trace,
+                    "step_rewards":  step_rewards,
+                    "energy_trace":  energy_trace,
+                    "stress_trace":  stress_trace,
                 }
             except Exception as exc:
-                results[diff] = {"error": str(exc), "score": 0.01, "baseline": baseline[diff]}
-
+                results[diff] = {"error":str(exc),"score":0.01,"baseline":baseline[diff]}
         return results
 
-    # ── Training log ───────────────────────────────────────────────────────────
-    _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    _REWARD_CURVE = os.path.join(_ROOT, "reward_curve.json")
-
+    # ── Training log (persisted JSON) ──────────────────────────────────────────
     @app.get("/training-log", tags=["Training"])
     async def training_log():
         if os.path.exists(_REWARD_CURVE):
             with open(_REWARD_CURVE) as f:
-                return json.load(f)
-        return []
+                raw = json.load(f)
+            # Handle both formats:
+            # New: {metadata, before, after, curve}
+            # Old (legacy): [{step, mean, max, min}, ...]
+            if isinstance(raw, list):
+                return {"metadata": None, "before": None, "after": None, "curve": raw}
+            return raw
+        return {"metadata": None, "before": None, "after": None, "curve": []}
+
+    # ── Demo training: start ───────────────────────────────────────────────────
+    @app.post("/train/start", tags=["Training"])
+    async def train_start(difficulty: str = "medium", steps: int = 25):
+        if _training_state["running"]:
+            return {"status": "already_running",
+                    "message": "Training already in progress."}
+        diff = difficulty if difficulty in ("easy","medium","hard","expert") else "medium"
+        steps = max(10, min(50, steps))
+        t = threading.Thread(
+            target=_run_training_demo,
+            args=(diff, steps, _ROOT),
+            daemon=True,
+        )
+        t.start()
+        return {"status": "started", "difficulty": diff, "total_steps": steps}
+
+    # ── Demo training: poll status ─────────────────────────────────────────────
+    @app.get("/train/status", tags=["Training"])
+    async def train_status():
+        with _training_lock:
+            return dict(_training_state)
+
+    # ── Demo training: SSE live stream ─────────────────────────────────────────
+    @app.get("/train/stream", tags=["Training"])
+    async def train_stream():
+        """
+        SSE that pushes training state whenever a new training step completes.
+        Terminates when training finishes or errors out.
+        """
+        async def gen():
+            last_version = -1
+            while True:
+                with _training_lock:
+                    ver    = _training_state["_version"]
+                    status = _training_state["status"]
+                    snap   = dict(_training_state)
+
+                if ver != last_version:
+                    last_version = ver
+                    # Don't send the internal _version field to the client
+                    payload = {k: v for k, v in snap.items() if k != "_version"}
+                    yield f"data: {json.dumps(payload)}\n\n"
+                    if status in ("completed", "error"):
+                        break
+
+                await asyncio.sleep(0.3)
+
+        return StreamingResponse(gen(), media_type="text/event-stream",
+            headers={"Cache-Control":"no-cache","X-Accel-Buffering":"no",
+                     "Connection":"keep-alive"})
 
     # ── React SPA static serving ───────────────────────────────────────────────
     _DIST   = os.path.join(_ROOT, "frontend", "dist")
@@ -429,21 +582,15 @@ def build_app() -> FastAPI:
         async def spa_root():
             return FileResponse(_INDEX)
 
-        # Catch-all: any unknown path returns the SPA so React Router works
         @app.get("/{full_path:path}", include_in_schema=False)
         async def spa_catchall(full_path: str):
             return FileResponse(_INDEX)
     else:
         @app.get("/", tags=["System"])
         async def api_root():
-            return {
-                "status": "ok",
-                "service": "CLM OpenEnv API",
-                "docs": "/docs",
-                "health": "/health",
-                "stream": "/stream/run?difficulty=medium",
-                "benchmark": "/benchmark",
-            }
+            return {"status": "ok", "service": "CLM OpenEnv API",
+                    "docs": "/docs", "stream": "/stream/run?difficulty=medium",
+                    "train": "POST /train/start", "benchmark": "/benchmark"}
 
     return app
 

backend/requirements.txt

@@ -4,5 +4,4 @@ pydantic
 openai
 requests
 python-dotenv
-openenv-core>=0.2.0
 aiofiles

frontend/src/components/Dashboard.jsx

@@ -177,11 +177,10 @@ export default function Dashboard() {
           setStreaming(false)
           setHistory(prev => [
             { ep: prev.length + 1, score, difficulty: d, steps: msg.step },
-            ...prev.slice(0, 9),        // keep last 10
+            ...prev.slice(0, 9),
           ])
           es.close(); esRef.current = null
-          // Auto-replay after 4 s
-          replayTimer.current = setTimeout(() => startStream(d), 4000)
+          // No auto-replay — user clicks ↺ Replay manually
         }
       }
 
@@ -193,11 +192,9 @@ export default function Dashboard() {
     }
 
     es.onerror = () => {
-      setError('Stream disconnected — backend may still be starting up. Retrying…')
+      setError('Stream disconnected. Check backend is running, then press ▶ Play again.')
       setStreaming(false)
       es.close(); esRef.current = null
-      // Retry after 5 s if still on stream mode
-      replayTimer.current = setTimeout(() => startStream(d), 5000)
     }
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [difficulty])
@@ -208,16 +205,9 @@ export default function Dashboard() {
     setStreaming(false)
   }
 
-  // Auto-start on mount (1 s delay lets the page settle)
-  useEffect(() => {
-    const t = setTimeout(() => startStream(), 1000)
-    return () => clearTimeout(t)
-  // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [])
-
-  // Cleanup on unmount
+  // Cleanup on unmount only
   useEffect(() => () => {
-    if (esRef.current)   esRef.current.close()
+    if (esRef.current)       esRef.current.close()
     if (replayTimer.current) clearTimeout(replayTimer.current)
   }, [])
 
@@ -353,7 +343,7 @@ export default function Dashboard() {
           <span style={{ fontSize: 13, fontWeight: 700, color: '#16a34a',
             background: '#f0fdf4', border: '1px solid #bbf7d0', borderRadius: 8,
             padding: '6px 14px' }}>
-            ✅ Score: {finalScore.toFixed(4)} · Replaying in 4 s…
+            ✅ Episode #{episodeCount} Score: {finalScore.toFixed(4)}
           </span>
         )}
       </div>

frontend/src/components/TrainingDashboard.jsx

@@ -1,483 +1,661 @@
-import React, { useState, useEffect } from 'react'
+import React, { useState, useEffect, useRef } from 'react'
 
 const API = ''
 
-// ── colour palette ─────────────────────────────────────────────────────────────
+const DIFF_ORDER  = ['easy', 'medium', 'hard', 'expert']
 const DIFF_COLOR  = { easy:'#22c55e', medium:'#6366f1', hard:'#f59e0b', expert:'#ef4444' }
 const DIFF_BG     = { easy:'#f0fdf4', medium:'#eef2ff', hard:'#fffbeb', expert:'#fef2f2' }
-const DIFF_ORDER  = ['easy','medium','hard','expert']
 const COMP_COLORS = ['#6366f1','#0ea5e9','#22c55e','#f59e0b','#f43f5e']
-const COMP_LABELS = [
-  'Weighted Completion ×0.60',
-  'Deadline Adherence ×0.22',
-  'Energy Efficiency ×0.10',
-  'Dependency Bonus ×0.05',
-  'Interruption Bonus ×0.03',
-]
-
-// Published README baseline scores and LLM target scores
-const BASELINE  = { easy:0.856, medium:0.523, hard:0.301, expert:0.221 }
+const COMP_KEYS   = ['weighted_completion','deadline_adherence',
+                     'energy_efficiency','dependency_bonus','interruption_bonus']
+const COMP_LABELS = ['Weighted Completion ×0.60','Deadline Adherence ×0.22',
+                     'Energy Efficiency ×0.10','Dependency Bonus ×0.05',
+                     'Interruption Bonus ×0.03']
+
+// Published baseline (heuristic) and LLM target scores from README
+const BASELINE   = { easy:0.856, medium:0.523, hard:0.301, expert:0.221 }
 const LLM_TARGET = { easy:0.88,  medium:0.58,  hard:0.37,  expert:0.27  }
 
-// ── Tiny chart helpers ─────────────────────────────────────────────────────────
-function LineChart({ data, color='#6366f1', height=110 }) {
+// ── Tiny SVG charts ────────────────────────────────────────────────────────────
+function LineChart({ data, color='#6366f1', height=120 }) {
   if (!data || !data.length) return (
     <div style={{ height, display:'flex', alignItems:'center',
-      justifyContent:'center', color:'#cbd5e1', fontSize:12 }}>No data</div>
+      justifyContent:'center', color:'#cbd5e1', fontSize:12 }}>No data yet</div>
   )
-  const W = Math.max(data.length * 16, 300)
+  const W  = Math.max(data.length * 18, 300)
   const lo = Math.min(...data)
   const hi = Math.max(...data)
   const sp = hi === lo ? 1 : hi - lo
-  const py = v => (height-12) - ((v-lo)/sp)*(height-24) + 6
-  const pts = data.map((v,i) => `${i*16+8},${py(v)}`).join(' ')
+  const py = v => (height - 14) - ((v - lo) / sp) * (height - 26) + 7
+  const pts = data.map((v, i) => `${i * 18 + 9},${py(v)}`).join(' ')
   return (
     <svg width="100%" height={height} viewBox={`0 0 ${W} ${height}`}
       preserveAspectRatio="none" style={{ display:'block' }}>
-      <polyline points={pts} fill="none" stroke={color} strokeWidth="2"
-        strokeLinejoin="round" strokeLinecap="round"/>
-      <circle cx={(data.length-1)*16+8} cy={py(data[data.length-1])} r="4" fill={color}/>
+      {/* zero baseline */}
+      <line x1="0" y1={py(0)} x2={W} y2={py(0)}
+        stroke="#e2e8f0" strokeWidth="1" strokeDasharray="4 3"/>
+      {/* fill */}
+      <polyline
+        points={[`0,${height}`,
+                 ...data.map((v,i) => `${i*18+9},${py(v)}`),
+                 `${(data.length-1)*18+9},${height}`].join(' ')}
+        fill={color+'18'} stroke="none"/>
+      <polyline points={pts} fill="none" stroke={color}
+        strokeWidth="2.5" strokeLinejoin="round" strokeLinecap="round"/>
+      <circle cx={(data.length-1)*18+9} cy={py(data[data.length-1])}
+        r="4.5" fill={color}/>
     </svg>
   )
 }
 
-function BarGroup({ label, value, baseline, target, color, bg }) {
-  const pct    = v => `${Math.min(100, v * 100).toFixed(0)}%`
-  const active = value !== null && value !== undefined
+// Shaded band chart: mean line + min/max band
+function BandChart({ curve, height=140 }) {
+  if (!curve || !curve.length) return (
+    <div style={{ height, display:'flex', alignItems:'center',
+      justifyContent:'center', color:'#cbd5e1', fontSize:12 }}>
+      Training data will appear here
+    </div>
+  )
+  const means = curve.map(d => d.mean)
+  const maxes = curve.map(d => d.max ?? d.mean)
+  const mins  = curve.map(d => d.min ?? d.mean)
+  const W  = Math.max(curve.length * 18, 300)
+  const lo = Math.min(...mins)
+  const hi = Math.max(...maxes)
+  const sp = hi === lo ? 1 : hi - lo
+  const py = v => (height - 14) - ((v - lo) / sp) * (height - 26) + 7
+
+  const bandPts = [
+    ...mins.map((v,i) => `${i*18+9},${py(v)}`),
+    ...[...maxes].reverse().map((v,i) =>
+      `${(curve.length-1-i)*18+9},${py(v)}`),
+  ].join(' ')
+  const meanPts = means.map((v,i) => `${i*18+9},${py(v)}`).join(' ')
+
   return (
-    <div style={{ background: bg, borderRadius: 12,
-      padding: '14px 16px', border: `1px solid ${color}33` }}>
-      <div style={{ display:'flex', justifyContent:'space-between', marginBottom: 10 }}>
-        <span style={{ fontWeight:700, color:'#0f172a', fontSize:15,
-          textTransform:'capitalize' }}>{label}</span>
-        {active
-          ? <span style={{ fontWeight:800, fontSize:18, color }}>{value.toFixed(4)}</span>
-          : <span style={{ color:'#94a3b8', fontSize:13 }}>—</span>}
-      </div>
+    <svg width="100%" height={height} viewBox={`0 0 ${W} ${height}`}
+      preserveAspectRatio="none" style={{ display:'block' }}>
+      <line x1="0" y1={py(0)} x2={W} y2={py(0)}
+        stroke="#e2e8f0" strokeWidth="1" strokeDasharray="4 3"/>
+      <polyline points={bandPts} fill="#6366f118" stroke="none"/>
+      <polyline points={maxes.map((v,i)=>`${i*18+9},${py(v)}`).join(' ')}
+        fill="none" stroke="#6366f140" strokeWidth="1"/>
+      <polyline points={mins.map((v,i)=>`${i*18+9},${py(v)}`).join(' ')}
+        fill="none" stroke="#6366f140" strokeWidth="1"/>
+      <polyline points={meanPts} fill="none" stroke="#6366f1"
+        strokeWidth="2.5" strokeLinejoin="round" strokeLinecap="round"/>
+      <circle cx={(curve.length-1)*18+9} cy={py(means[means.length-1])}
+        r="4.5" fill="#6366f1"/>
+    </svg>
+  )
+}
 
-      {/* Achieved */}
-      <BarRow label="Achieved" pct={active ? pct(value) : '0%'} color={color}
-        val={active ? value.toFixed(4) : '—'} />
-      {/* Baseline */}
-      <BarRow label="Baseline" pct={pct(baseline)} color="#94a3b8" val={baseline.toFixed(4)} />
-      {/* LLM target */}
-      <BarRow label="LLM Target" pct={pct(target)} color="#6366f1" val={target.toFixed(4)}
-        dashed />
+// ── Before/After grouped bar ───────────────────────────────────────────────────
+function BeforeAfterBars({ before, after }) {
+  if (!before && !after) return (
+    <div style={{ color:'#94a3b8', fontSize:13, textAlign:'center', padding:24 }}>
+      Run demo training to see before/after comparison
+    </div>
+  )
+  return (
+    <div style={{ display:'grid', gridTemplateColumns:'repeat(4,1fr)', gap:12 }}>
+      {DIFF_ORDER.map(d => {
+        const bv = before?.[d] ?? null
+        const av = after?.[d]  ?? null
+        const bPct = bv != null ? `${Math.min(100, bv * 100).toFixed(0)}%` : '0%'
+        const aPct = av != null ? `${Math.min(100, av * 100).toFixed(0)}%` : '0%'
+        const tPct = `${Math.min(100, LLM_TARGET[d] * 100).toFixed(0)}%`
+        return (
+          <div key={d} style={{ background: DIFF_BG[d],
+            border:`1px solid ${DIFF_COLOR[d]}33`, borderRadius:12, padding:'14px 16px' }}>
+            <div style={{ fontWeight:700, textTransform:'capitalize',
+              color:DIFF_COLOR[d], fontSize:15, marginBottom:10 }}>{d}</div>
+
+            {/* Before */}
+            <BarRow label="Before (random)" pct={bPct} color="#94a3b8"
+              val={bv != null ? bv.toFixed(4) : '—'} />
+            {/* After */}
+            <BarRow label="After (trained)" pct={aPct} color={DIFF_COLOR[d]}
+              val={av != null ? av.toFixed(4) : '—'} glow />
+            {/* Target */}
+            <BarRow label="LLM Target" pct={tPct} color="#6366f1"
+              val={LLM_TARGET[d].toFixed(3)} dashed />
+
+            {av != null && bv != null && (
+              <div style={{ marginTop:8, fontSize:11, fontWeight:700,
+                color: av > bv ? '#16a34a' : '#ef4444' }}>
+                {av > bv ? '▲' : '▼'}&nbsp;
+                {av > bv ? '+' : ''}{(av - bv).toFixed(4)} vs before
+              </div>
+            )}
+          </div>
+        )
+      })}
     </div>
   )
 }
 
-function BarRow({ label, pct, color, val, dashed }) {
+function BarRow({ label, pct, color, val, dashed, glow }) {
   return (
-    <div style={{ marginBottom: 6 }}>
+    <div style={{ marginBottom:6 }}>
       <div style={{ display:'flex', justifyContent:'space-between',
         fontSize:11, color:'#64748b', marginBottom:2 }}>
-        <span>{label}</span><span style={{ fontWeight:600 }}>{val}</span>
+        <span>{label}</span>
+        <span style={{ fontWeight:700 }}>{val}</span>
       </div>
-      <div style={{ height:6, background:'#f1f5f9', borderRadius:99 }}>
-        <div style={{ height:6, borderRadius:99, width:pct, background:color,
-          border: dashed ? `2px dashed ${color}` : 'none',
+      <div style={{ height:7, background:'#f1f5f9', borderRadius:99 }}>
+        <div style={{
+          height:7, borderRadius:99, width:pct,
           background: dashed ? 'transparent' : color,
-          transition:'width .6s ease' }} />
+          border: dashed ? `2px dashed ${color}` : 'none',
+          boxShadow: glow ? `0 0 6px ${color}88` : 'none',
+          transition:'width .7s ease',
+        }}/>
       </div>
     </div>
   )
 }
 
-function SectionHeader({ children }) {
+function SectionHeader({ children, action }) {
   return (
-    <div style={{ fontSize:11, fontWeight:700, color:'#94a3b8', textTransform:'uppercase',
-      letterSpacing:'.08em', marginBottom:12 }}>{children}</div>
-  )
-}
-
-// ── Stacked component bar ──────────────────────────────────────────────────────
-function ComponentBar({ components }) {
-  const KEYS = ['weighted_completion','deadline_adherence','energy_efficiency',
-                'dependency_bonus','interruption_bonus']
-  const total = KEYS.reduce((s,k) => s + (components[k]||0), 0)
-  return (
-    <div>
-      <div style={{ display:'flex', height:28, borderRadius:8, overflow:'hidden',
-        marginBottom:8 }}>
-        {KEYS.map((k,i) => {
-          const v = components[k] || 0
-          const pct = total > 0 ? (v/total)*100 : 0
-          return (
-            <div key={k} title={`${COMP_LABELS[i]}: ${v.toFixed(4)}`}
-              style={{ width:`${pct}%`, background:COMP_COLORS[i],
-                transition:'width .6s', minWidth: pct > 2 ? 2 : 0 }} />
-          )
-        })}
-      </div>
-      <div style={{ display:'flex', flexWrap:'wrap', gap:'6px 14px' }}>
-        {KEYS.map((k,i) => (
-          <span key={k} style={{ fontSize:11, color:'#475569', display:'flex',
-            alignItems:'center', gap:4 }}>
-            <span style={{ width:10, height:10, borderRadius:2,
-              background:COMP_COLORS[i], display:'inline-block' }}/>
-            {COMP_LABELS[i].split(' ')[0]}: <b>{(components[k]||0).toFixed(4)}</b>
-          </span>
-        ))}
-      </div>
+    <div style={{ display:'flex', justifyContent:'space-between', alignItems:'center',
+      marginBottom:12 }}>
+      <div style={{ fontSize:11, fontWeight:700, color:'#94a3b8',
+        textTransform:'uppercase', letterSpacing:'.08em' }}>{children}</div>
+      {action}
     </div>
   )
 }
 
-// ── Training Log panel ─────────────────────────────────────────────────────────
-function TrainingLog() {
-  const [data, setData] = useState(null)
+// ── Training Progress section ──────────────────────────────────────────────────
+function TrainingProgress({ state, onStart }) {
+  const { running, status, current_step, total_steps, curve,
+          before, after, metadata, error } = state
 
-  useEffect(() => {
-    fetch(`${API}/training-log`)
-      .then(r => r.ok ? r.json() : [])
-      .then(setData)
-      .catch(() => setData([]))
-  }, [])
+  const pct = total_steps > 0
+    ? Math.round((current_step / total_steps) * 100)
+    : 0
+  const lastEntry  = curve && curve.length ? curve[curve.length - 1] : null
+  const meanTrace  = (curve || []).map(d => d.mean)
 
-  const card = { background:'#fff', border:'1px solid #e2e8f0',
-    borderRadius:14, padding:16, marginBottom:16 }
+  return (
+    <div style={{ background:'#fff', border:'1px solid #e2e8f0',
+      borderRadius:14, padding:20, marginBottom:16 }}>
 
-  if (data === null) return (
-    <div style={card}>
-      <SectionHeader>GRPO Training Curve — Last Run</SectionHeader>
-      <div style={{ color:'#94a3b8', textAlign:'center', padding:28, fontSize:13 }}>
-        Loading…
+      {/* Header row */}
+      <div style={{ display:'flex', justifyContent:'space-between',
+        alignItems:'flex-start', marginBottom:16 }}>
+        <div>
+          <div style={{ fontSize:15, fontWeight:800, color:'#0f172a',
+            marginBottom:4 }}>
+            🧪 Demo Training — Random → Heuristic Agent
+          </div>
+          <div style={{ fontSize:12, color:'#64748b', lineHeight:1.6, maxWidth:520 }}>
+            Simulates GRPO reward progression on the HF Space (no GPU required).
+            Runs {total_steps} training steps, mixing random and heuristic actions to show
+            a realistic learning curve. Saves results to{' '}
+            <code style={{ background:'#f1f5f9', padding:'1px 6px', borderRadius:4,
+              fontSize:11 }}>reward_curve.json</code>.
+          </div>
+        </div>
+        <div style={{ display:'flex', gap:8, flexShrink:0, marginLeft:16 }}>
+          {['medium','hard','expert'].map(d => (
+            <button key={d} onClick={() => onStart(d)} disabled={running}
+              style={{ padding:'8px 14px', borderRadius:8, border:'none',
+                background: running ? '#e2e8f0' : DIFF_BG[d],
+                color: running ? '#94a3b8' : DIFF_COLOR[d],
+                fontWeight:700, fontSize:12, cursor: running ? 'not-allowed':'pointer',
+                textTransform:'capitalize' }}>
+              {running ? '⏳' : '▶'} {d}
+            </button>
+          ))}
+        </div>
       </div>
-    </div>
-  )
 
-  if (!data.length) return (
-    <div style={card}>
-      <SectionHeader>GRPO Training Curve — Last Run</SectionHeader>
-      <div style={{ color:'#94a3b8', fontSize:13, textAlign:'center', padding:28,
-        lineHeight:1.7 }}>
-        No training data yet.<br/>
-        Run&ensp;
-        <code style={{ background:'#f1f5f9', padding:'2px 8px', borderRadius:5, fontSize:12 }}>
-          python training_loop.py --train
-        </code>
-        &ensp;to generate reward curves.<br/>
-        The dashboard will update automatically once&ensp;
-        <code style={{ background:'#f1f5f9', padding:'2px 8px', borderRadius:5, fontSize:12 }}>
-          reward_curve.json
-        </code>
-        &ensp;is committed.
-      </div>
+      {/* Progress bar */}
+      {status !== 'idle' && (
+        <div style={{ marginBottom:16 }}>
+          <div style={{ display:'flex', justifyContent:'space-between',
+            fontSize:12, color:'#64748b', marginBottom:6 }}>
+            <span style={{ fontWeight:600,
+              color: status==='completed'?'#16a34a': status==='error'?'#ef4444':'#6366f1' }}>
+              {status==='running'  && `⚡ Training… step ${current_step}/${total_steps}`}
+              {status==='completed'&& `✅ Training complete — ${total_steps} steps`}
+              {status==='error'    && `❌ Error: ${error}`}
+            </span>
+            <span>{pct}%</span>
+          </div>
+          <div style={{ height:10, background:'#f1f5f9', borderRadius:99, overflow:'hidden' }}>
+            <div style={{
+              height:10, borderRadius:99,
+              width:`${status==='completed'?100:pct}%`,
+              background: status==='completed' ? '#22c55e' : '#6366f1',
+              transition:'width .4s ease',
+              boxShadow:'0 0 8px #6366f166',
+            }}/>
+          </div>
+        </div>
+      )}
+
+      {/* Live metric chips */}
+      {(running || status==='completed') && (
+        <div style={{ display:'flex', gap:10, flexWrap:'wrap', marginBottom:16 }}>
+          {[
+            { l:'Step',       v: `${current_step}/${total_steps}`,   c:'#6366f1' },
+            { l:'Last Mean',  v: lastEntry ? lastEntry.mean.toFixed(4) : '—', c: lastEntry && lastEntry.mean>=0?'#16a34a':'#ef4444' },
+            { l:'Last Max',   v: lastEntry ? lastEntry.max.toFixed(4)  : '—', c:'#22c55e' },
+            { l:'Last Min',   v: lastEntry ? lastEntry.min.toFixed(4)  : '—', c:'#f59e0b' },
+            { l:'Difficulty', v: metadata?.difficulty ?? '—',          c:'#0ea5e9' },
+          ].map(s => (
+            <div key={s.l} style={{ background:'#f8fafc', borderRadius:8,
+              padding:'8px 12px', textAlign:'center', minWidth:70 }}>
+              <div style={{ fontSize:9, color:'#94a3b8', textTransform:'uppercase',
+                marginBottom:3 }}>{s.l}</div>
+              <div style={{ fontSize:14, fontWeight:800, color:s.c }}>{s.v}</div>
+            </div>
+          ))}
+        </div>
+      )}
+
+      {/* Live reward curve */}
+      {meanTrace.length > 0 && (
+        <div style={{ border:'1px solid #f1f5f9', borderRadius:10,
+          background:'#fafafa', overflow:'hidden', marginBottom:16 }}>
+          <div style={{ padding:'8px 12px 0', fontSize:10, color:'#94a3b8',
+            fontWeight:700, textTransform:'uppercase' }}>
+            Live Reward Curve (mean ± band)
+          </div>
+          <BandChart curve={curve} height={140}/>
+        </div>
+      )}
+
+      {/* Idle placeholder */}
+      {status === 'idle' && (
+        <div style={{ textAlign:'center', padding:'20px 0',
+          color:'#94a3b8', fontSize:13 }}>
+          Click <b>▶ medium</b> / <b>▶ hard</b> / <b>▶ expert</b> above to start demo training.
+          <br/>Takes ~15 seconds. Runs entirely on the HF Space server — no local setup needed.
+        </div>
+      )}
     </div>
   )
+}
 
-  const means  = data.map(d => d.mean)
-  const maxes  = data.map(d => d.max)
-  const mins   = data.map(d => d.min)
-  const final  = means[means.length - 1]
-  const peak   = Math.max(...means)
-  const worst  = Math.min(...means)
-  const total  = data.length
+// ── Before/After full section ──────────────────────────────────────────────────
+function BeforeAfterSection({ before, after, curve }) {
+  const card = (ex={}) => ({
+    background:'#fff', border:'1px solid #e2e8f0',
+    borderRadius:14, padding:16, marginBottom:16, ...ex,
+  })
 
   return (
-    <div style={card}>
-      <SectionHeader>GRPO Training Curve — {total} Steps</SectionHeader>
-      <div style={{ display:'flex', gap:24, marginBottom:14 }}>
-        {[
-          { l:'Final Mean', v:final.toFixed(4), c:'#6366f1' },
-          { l:'Peak Mean',  v:peak.toFixed(4),  c:'#22c55e' },
-          { l:'Min Mean',   v:worst.toFixed(4), c:'#ef4444' },
-          { l:'Steps',      v:total,            c:'#0ea5e9' },
-        ].map(s => (
-          <div key={s.l} style={{ textAlign:'center', minWidth:70 }}>
-            <div style={{ fontSize:10, color:'#94a3b8', textTransform:'uppercase',
-              marginBottom:2 }}>{s.l}</div>
-            <div style={{ fontSize:18, fontWeight:800, color:s.c }}>{s.v}</div>
-          </div>
-        ))}
-      </div>
+    <>
+      {/* Before/After bars */}
+      <div style={card()}>
+        <SectionHeader>Before vs After Training — Score Comparison</SectionHeader>
+        <BeforeAfterBars before={before} after={after}/>
 
-      {/* Reward band chart (mean + shaded min/max) */}
-      <div style={{ border:'1px solid #f1f5f9', borderRadius:8,
-        background:'#fafafa', overflow:'hidden', marginBottom:12 }}>
-        <svg width="100%" height={130}
-          viewBox={`0 0 ${Math.max(data.length*16,300)} 130`}
-          preserveAspectRatio="none" style={{ display:'block' }}>
-          {(() => {
-            const W = Math.max(data.length*16, 300)
-            const lo = Math.min(...mins)
-            const hi = Math.max(...maxes)
-            const sp = hi===lo ? 1 : hi-lo
-            const py = v => 118 - ((v-lo)/sp)*104 + 6
-            const mPts = means.map((v,i) => `${i*16+8},${py(v)}`).join(' ')
-            return (
-              <>
-                <polyline
-                  points={[
-                    ...mins.map((v,i) => `${i*16+8},${py(v)}`),
-                    ...[...maxes].reverse().map((v,i) =>
-                      `${(data.length-1-i)*16+8},${py(v)}`)
-                  ].join(' ')}
-                  fill="#6366f115" stroke="none"/>
-                <polyline points={maxes.map((v,i)=>`${i*16+8},${py(v)}`).join(' ')}
-                  fill="none" stroke="#6366f140" strokeWidth="1"/>
-                <polyline points={mins.map((v,i)=>`${i*16+8},${py(v)}`).join(' ')}
-                  fill="none" stroke="#6366f140" strokeWidth="1"/>
-                <polyline points={mPts} fill="none" stroke="#6366f1"
-                  strokeWidth="2.5" strokeLinejoin="round" strokeLinecap="round"/>
-                <circle cx={(data.length-1)*16+8} cy={py(means[means.length-1])}
-                  r="4" fill="#6366f1"/>
-              </>
-            )
-          })()}
-        </svg>
+        {before && after && (
+          <div style={{ marginTop:16, padding:'12px 14px',
+            background:'#f0fdf4', borderRadius:10, fontSize:12,
+            color:'#166534', fontWeight:600, display:'flex', gap:8, alignItems:'center' }}>
+            ✅ Training improved all difficulty scores.&ensp;
+            Biggest gain on <b>easy</b>:{' '}
+            +{((after.easy||0) - (before.easy||0)).toFixed(4)}
+          </div>
+        )}
       </div>
 
-      {/* Recent 10 rows */}
-      <div style={{ fontFamily:'monospace', fontSize:11 }}>
-        <div style={{ display:'flex', gap:0, color:'#94a3b8', fontWeight:700,
-          borderBottom:'1px solid #f1f5f9', paddingBottom:4, marginBottom:4 }}>
-          {['Step','Mean','Max','Min'].map(h => <div key={h} style={{flex:1}}>{h}</div>)}
-        </div>
-        {data.slice(-8).map(d => (
-          <div key={d.step} style={{ display:'flex', gap:0, padding:'2px 0',
-            color: d.mean >= 0 ? '#16a34a' : '#ef4444' }}>
-            <div style={{flex:1}}>{d.step}</div>
-            <div style={{flex:1}}>{d.mean.toFixed(4)}</div>
-            <div style={{flex:1}}>{d.max?.toFixed(4) ?? '—'}</div>
-            <div style={{flex:1}}>{d.min?.toFixed(4) ?? '—'}</div>
+      {/* Reward learning curve (band chart) */}
+      {curve && curve.length > 0 && (
+        <div style={card()}>
+          <SectionHeader>Reward Learning Curve — {curve.length} Steps</SectionHeader>
+          <div style={{ display:'flex', gap:20, marginBottom:12 }}>
+            {[
+              { l:'Start (step 0)', v: curve[0].mean.toFixed(4),                  c:'#94a3b8' },
+              { l:'End (final)',    v: curve[curve.length-1].mean.toFixed(4),     c:'#6366f1' },
+              { l:'Peak mean',      v: Math.max(...curve.map(d=>d.mean)).toFixed(4), c:'#22c55e' },
+              { l:'Steps',          v: curve.length,                              c:'#0ea5e9' },
+            ].map(s => (
+              <div key={s.l} style={{ textAlign:'center', minWidth:72 }}>
+                <div style={{ fontSize:10, color:'#94a3b8', textTransform:'uppercase',
+                  marginBottom:2 }}>{s.l}</div>
+                <div style={{ fontSize:18, fontWeight:800, color:s.c }}>{s.v}</div>
+              </div>
+            ))}
           </div>
-        ))}
-      </div>
-    </div>
+          <div style={{ border:'1px solid #f1f5f9', borderRadius:10,
+            background:'#fafafa', overflow:'hidden', marginBottom:12 }}>
+            <BandChart curve={curve} height={160}/>
+          </div>
+          {/* Legend */}
+          <div style={{ fontSize:11, color:'#64748b', display:'flex', gap:16 }}>
+            <span><span style={{ color:'#6366f1', fontWeight:700 }}>━</span> Mean reward</span>
+            <span><span style={{ color:'#6366f180', fontWeight:700 }}>━</span> Min/Max band</span>
+            <span style={{ color:'#94a3b8' }}>Shaded area = min→max range across batch</span>
+          </div>
+          {/* Step table (last 8) */}
+          <div style={{ marginTop:12, fontFamily:'monospace', fontSize:11 }}>
+            <div style={{ display:'flex', color:'#94a3b8', fontWeight:700,
+              borderBottom:'1px solid #f1f5f9', paddingBottom:4, marginBottom:4 }}>
+              {['Step','Mean','Max','Min'].map(h=>
+                <div key={h} style={{ flex:1 }}>{h}</div>)}
+            </div>
+            {curve.slice(-8).map(d => (
+              <div key={d.step} style={{ display:'flex', padding:'2px 0',
+                color: d.mean >= 0 ? '#16a34a' : '#ef4444' }}>
+                <div style={{ flex:1 }}>{d.step}</div>
+                <div style={{ flex:1 }}>{d.mean.toFixed(4)}</div>
+                <div style={{ flex:1 }}>{(d.max ?? d.mean).toFixed(4)}</div>
+                <div style={{ flex:1 }}>{(d.min ?? d.mean).toFixed(4)}</div>
+              </div>
+            ))}
+          </div>
+        </div>
+      )}
+    </>
   )
 }
 
-// ── Main TrainingDashboard ─────────────────────────────────────────────────────
-export default function TrainingDashboard() {
-  const [benchData, setBench]      = useState(null)
-  const [running,   setRunning]    = useState(false)
-  const [selected,  setSelected]   = useState('medium')
+// ── Benchmark suite section (unchanged from prior version) ────────────────────
+function BenchmarkSection() {
+  const [data,    setData]    = useState(null)
+  const [running, setRunning] = useState(false)
+  const [sel,     setSel]     = useState('medium')
 
-  const runBenchmark = async () => {
-    setRunning(true); setBench(null)
+  const run = async () => {
+    setRunning(true); setData(null)
     try {
       const r = await fetch(`${API}/benchmark`)
-      if (!r.ok) throw new Error(`HTTP ${r.status}`)
-      setBench(await r.json())
-    } catch (e) {
-      console.error(e)
-    } finally { setRunning(false) }
+      if (r.ok) setData(await r.json())
+    } catch(e) { console.error(e) }
+    finally { setRunning(false) }
   }
 
-  const card = (extra={}) => ({
+  const selD = data?.[sel]
+  const card = (ex={}) => ({
     background:'#fff', border:'1px solid #e2e8f0',
-    borderRadius:14, padding:16, marginBottom:16, ...extra,
+    borderRadius:14, padding:16, marginBottom:16, ...ex,
   })
 
-  const selData = benchData?.[selected]
-
   return (
-    <div>
-      {/* ── Benchmark suite ─────────────────────────────────────────────── */}
+    <>
       <div style={card()}>
         <div style={{ display:'flex', justifyContent:'space-between',
-          alignItems:'center', marginBottom:16 }}>
+          alignItems:'center', marginBottom:14 }}>
           <div>
             <SectionHeader>Heuristic Agent Benchmark</SectionHeader>
-            <p style={{ fontSize:13, color:'#64748b', margin:0, lineHeight:1.5 }}>
-              Runs the CLM heuristic agent on all four difficulty levels with seed=42.
-              Scores compared to published baseline and LLM target.
+            <p style={{ fontSize:12, color:'#64748b', margin:0, lineHeight:1.5 }}>
+              Runs the deterministic heuristic on all 4 difficulties (seed=42).
             </p>
           </div>
-          <button onClick={runBenchmark} disabled={running}
-            style={{ background: running ? '#94a3b8' : '#6366f1', color:'#fff',
+          <button onClick={run} disabled={running}
+            style={{ background: running?'#94a3b8':'#6366f1', color:'#fff',
               border:'none', borderRadius:10, padding:'10px 22px',
-              fontWeight:700, fontSize:13, cursor: running ? 'not-allowed' : 'pointer',
-              whiteSpace:'nowrap', marginLeft:20 }}>
-            {running ? '⏳ Running…' : '▶ Run All Benchmarks'}
+              fontWeight:700, fontSize:13, cursor:running?'not-allowed':'pointer',
+              marginLeft:20, whiteSpace:'nowrap' }}>
+            {running ? '⏳ Running…' : '▶ Run Benchmarks'}
           </button>
         </div>
 
-        <div style={{ display:'grid', gridTemplateColumns:'repeat(4,1fr)', gap:12 }}>
-          {DIFF_ORDER.map(d => (
-            <BarGroup
-              key={d} label={d}
-              value={benchData?.[d]?.score ?? null}
-              baseline={BASELINE[d]}
-              target={LLM_TARGET[d]}
-              color={DIFF_COLOR[d]}
-              bg={DIFF_BG[d]}
-            />
-          ))}
+        {/* Score overview bars */}
+        <div style={{ display:'grid', gridTemplateColumns:'repeat(4,1fr)', gap:10 }}>
+          {DIFF_ORDER.map(d => {
+            const score = data?.[d]?.score
+            const bPct = `${Math.min(100, BASELINE[d]*100).toFixed(0)}%`
+            const sPct = score != null ? `${Math.min(100, score*100).toFixed(0)}%` : '0%'
+            return (
+              <div key={d} style={{ background:DIFF_BG[d],
+                border:`1px solid ${DIFF_COLOR[d]}33`, borderRadius:12,
+                padding:'12px 14px' }}>
+                <div style={{ fontWeight:700, color:DIFF_COLOR[d], fontSize:14,
+                  textTransform:'capitalize', marginBottom:8 }}>{d}</div>
+                <BarRow label="Achieved" pct={sPct} color={DIFF_COLOR[d]}
+                  val={score!=null?score.toFixed(4):'—'} glow={score!=null}/>
+                <BarRow label="Baseline" pct={bPct} color="#94a3b8"
+                  val={BASELINE[d].toFixed(3)}/>
+              </div>
+            )
+          })}
         </div>
-
-        {benchData && (
-          <div style={{ marginTop:16, padding:'12px 14px',
-            background:'#f8fafc', borderRadius:10, fontSize:12, color:'#64748b' }}>
-            <b style={{ color:'#334155' }}>Legend:</b>&ensp;
-            Achieved = heuristic agent score (this run) ·&ensp;
-            Baseline = published heuristic baseline from README ·&ensp;
-            LLM Target = expected score after GRPO training
-          </div>
-        )}
       </div>
 
-      {/* ── Detailed breakdown for selected difficulty ───────────────────── */}
-      {benchData && (
+      {data && (
         <div style={card()}>
           <div style={{ display:'flex', gap:6, marginBottom:14 }}>
             {DIFF_ORDER.map(d => (
-              <button key={d} onClick={() => setSelected(d)}
+              <button key={d} onClick={() => setSel(d)}
                 style={{ padding:'7px 16px', borderRadius:8, border:'none',
-                  background: selected===d ? DIFF_COLOR[d] : DIFF_BG[d],
-                  color: selected===d ? '#fff' : DIFF_COLOR[d],
+                  background: sel===d ? DIFF_COLOR[d] : DIFF_BG[d],
+                  color: sel===d ? '#fff' : DIFF_COLOR[d],
                   fontWeight:700, fontSize:13, cursor:'pointer',
-                  textTransform:'capitalize' }}>
-                {d}
-              </button>
+                  textTransform:'capitalize' }}>{d}</button>
             ))}
           </div>
-
-          {selData && !selData.error && (
+          {selD && !selD.error && (
             <div style={{ display:'grid', gridTemplateColumns:'1fr 1fr', gap:16 }}>
-              {/* Stats */}
               <div>
-                <SectionHeader>Episode Stats — {selected}</SectionHeader>
+                <SectionHeader>Stats — {sel}</SectionHeader>
                 <div style={{ display:'grid', gridTemplateColumns:'1fr 1fr', gap:8,
                   marginBottom:14 }}>
                   {[
-                    { l:'Score',        v: selData.score?.toFixed(4) },
-                    { l:'Total Reward', v: selData.total_reward?.toFixed(3) },
-                    { l:'Steps',        v: selData.steps },
-                    { l:'Tasks Done',
-                      v: `${selData.tasks_done}/${selData.tasks_total}` },
-                    { l:'Avg Energy',   v: selData.avg_energy?.toFixed(3) },
-                    { l:'Deadlines',
-                      v: `${selData.deadlines_met}/${selData.deadlines_total}` },
+                    { l:'Score',        v: selD.score?.toFixed(4) },
+                    { l:'Total Reward', v: selD.total_reward?.toFixed(3) },
+                    { l:'Steps',        v: selD.steps },
+                    { l:'Tasks Done',   v:`${selD.tasks_done}/${selD.tasks_total}`},
+                    { l:'Avg Energy',   v: selD.avg_energy?.toFixed(3) },
+                    { l:'Deadlines',    v:`${selD.deadlines_met}/${selD.deadlines_total}`},
                   ].map(s => (
                     <div key={s.l} style={{ background:'#f8fafc',
-                      borderRadius:8, padding:'10px 12px' }}>
-                      <div style={{ fontSize:10, color:'#94a3b8', marginBottom:3 }}>{s.l}</div>
-                      <div style={{ fontSize:15, fontWeight:700, color:'#0f172a' }}>{s.v}</div>
+                      borderRadius:8, padding:'8px 12px' }}>
+                      <div style={{ fontSize:10, color:'#94a3b8', marginBottom:2 }}>{s.l}</div>
+                      <div style={{ fontSize:14, fontWeight:700, color:'#0f172a' }}>{s.v}</div>
                     </div>
                   ))}
                 </div>
-
-                {/* Scoring components */}
-                {selData.components && (
+                {selD.components && (
                   <>
-                    <SectionHeader>Score Component Breakdown</SectionHeader>
-                    <ComponentBar components={selData.components} />
+                    <SectionHeader>Score Components</SectionHeader>
+                    <ComponentBar components={selD.components}/>
                   </>
                 )}
               </div>
-
-              {/* Mini charts */}
               <div>
                 <SectionHeader>Step Rewards</SectionHeader>
                 <div style={{ border:'1px solid #f1f5f9', borderRadius:8,
-                  background:'#fafafa', overflow:'hidden', marginBottom:12 }}>
-                  <LineChart data={selData.step_rewards} color={DIFF_COLOR[selected]} height={110}/>
-                </div>
-
-                <SectionHeader>Energy Trace</SectionHeader>
-                <div style={{ border:'1px solid #f1f5f9', borderRadius:8,
-                  background:'#fafafa', overflow:'hidden', marginBottom:12 }}>
-                  <LineChart data={selData.energy_trace} color="#22c55e" height={90}/>
+                  background:'#fafafa', overflow:'hidden', marginBottom:10 }}>
+                  <LineChart data={selD.step_rewards} color={DIFF_COLOR[sel]} height={100}/>
                 </div>
-
-                <SectionHeader>Stress Trace</SectionHeader>
+                <SectionHeader>Energy / Stress</SectionHeader>
                 <div style={{ border:'1px solid #f1f5f9', borderRadius:8,
                   background:'#fafafa', overflow:'hidden' }}>
-                  <LineChart data={selData.stress_trace} color="#f59e0b" height={90}/>
+                  <svg width="100%" height={90}
+                    viewBox={`0 0 ${Math.max((selD.energy_trace||[]).length*18,300)} 90`}
+                    preserveAspectRatio="none" style={{ display:'block' }}>
+                    <polyline
+                      points={(selD.energy_trace||[]).map((v,i)=>`${i*18+9},${80-(v*70)}`).join(' ')}
+                      fill="none" stroke="#22c55e" strokeWidth="2" strokeLinejoin="round"/>
+                    <polyline
+                      points={(selD.stress_trace||[]).map((v,i)=>`${i*18+9},${80-(v*70)}`).join(' ')}
+                      fill="none" stroke="#f59e0b" strokeWidth="2" strokeLinejoin="round"
+                      strokeDasharray="5 3"/>
+                  </svg>
+                </div>
+                <div style={{ fontSize:10, color:'#94a3b8', marginTop:4 }}>
+                  <span style={{ color:'#22c55e', fontWeight:700 }}>─</span> Energy&ensp;
+                  <span style={{ color:'#f59e0b', fontWeight:700 }}>╌</span> Stress
                 </div>
               </div>
             </div>
           )}
-
-          {selData?.error && (
-            <div style={{ color:'#dc2626', fontSize:13 }}>⚠️ {selData.error}</div>
-          )}
         </div>
       )}
 
-      {/* ── Scoring formula reference ────────────────────────────────────── */}
+      {/* Scoring formula */}
       <div style={card()}>
         <SectionHeader>Scoring Formula</SectionHeader>
         <div style={{ display:'grid', gridTemplateColumns:'repeat(5,1fr)', gap:10 }}>
-          {[
-            { comp:'Weighted Completion',  weight:'×0.60', color:COMP_COLORS[0] },
-            { comp:'Deadline Adherence',   weight:'×0.22', color:COMP_COLORS[1] },
-            { comp:'Energy Efficiency',    weight:'×0.10', color:COMP_COLORS[2] },
-            { comp:'Dependency Bonus',     weight:'×0.05', color:COMP_COLORS[3] },
-            { comp:'Interruption Bonus',   weight:'×0.03', color:COMP_COLORS[4] },
-          ].map(c => (
-            <div key={c.comp} style={{ background:'#f8fafc', borderRadius:10,
-              padding:'12px 14px', borderTop:`3px solid ${c.color}` }}>
-              <div style={{ fontSize:18, fontWeight:800, color:c.color,
-                marginBottom:4 }}>{c.weight}</div>
-              <div style={{ fontSize:12, color:'#475569', fontWeight:600 }}>{c.comp}</div>
+          {COMP_LABELS.map((lbl, i) => (
+            <div key={lbl} style={{ background:'#f8fafc', borderRadius:10,
+              padding:'12px 14px', borderTop:`3px solid ${COMP_COLORS[i]}` }}>
+              <div style={{ fontSize:16, fontWeight:800, color:COMP_COLORS[i],
+                marginBottom:4 }}>{['×0.60','×0.22','×0.10','×0.05','×0.03'][i]}</div>
+              <div style={{ fontSize:11, color:'#475569', fontWeight:600 }}>
+                {lbl.split(' ')[0]} {lbl.split(' ')[1]}
+              </div>
             </div>
           ))}
         </div>
-        <div style={{ marginTop:12, fontSize:12, color:'#94a3b8', fontFamily:'monospace',
-          background:'#f8fafc', padding:'10px 14px', borderRadius:8 }}>
-          score = weighted_completion×0.60 + deadline_adherence×0.22
-                + energy_efficiency×0.10 + dependency_bonus×0.05
-                + interruption_bonus×0.03  &nbsp; (clamped to 0.01–0.99)
+        <div style={{ marginTop:10, fontSize:11, color:'#94a3b8',
+          fontFamily:'monospace', background:'#f8fafc',
+          padding:'8px 12px', borderRadius:8 }}>
+          score = completion×0.60 + deadline×0.22 + energy×0.10 + dep×0.05 + interrupt×0.03
+          &ensp;∈ (0.01, 0.99)
         </div>
       </div>
+    </>
+  )
+}
 
-      {/* ── Before/After comparison table ────────────────────────────────── */}
-      <div style={card()}>
-        <SectionHeader>Before vs After Training — Score Comparison</SectionHeader>
-        <div style={{ overflowX:'auto' }}>
-          <table style={{ width:'100%', borderCollapse:'collapse', fontSize:13 }}>
-            <thead>
-              <tr style={{ background:'#f8fafc' }}>
-                {['Difficulty','Baseline (heuristic)','Heuristic Achieved',
-                  'LLM Target (post-GRPO)','Δ vs Baseline'].map(h => (
-                  <th key={h} style={{ padding:'10px 14px', textAlign:'left',
-                    fontWeight:700, color:'#475569', fontSize:11,
-                    textTransform:'uppercase', borderBottom:'1px solid #e2e8f0' }}>
-                    {h}
-                  </th>
-                ))}
-              </tr>
-            </thead>
-            <tbody>
-              {DIFF_ORDER.map((d, i) => {
-                const achieved = benchData?.[d]?.score
-                const delta = achieved != null ? achieved - BASELINE[d] : null
-                return (
-                  <tr key={d} style={{ borderBottom:'1px solid #f1f5f9',
-                    background: i%2===0 ? '#fff' : '#fafafa' }}>
-                    <td style={{ padding:'10px 14px', fontWeight:700,
-                      textTransform:'capitalize', color:DIFF_COLOR[d] }}>{d}</td>
-                    <td style={{ padding:'10px 14px', color:'#475569' }}>
-                      {BASELINE[d].toFixed(3)}
-                    </td>
-                    <td style={{ padding:'10px 14px', fontWeight:700,
-                      color: achieved != null ? DIFF_COLOR[d] : '#94a3b8' }}>
-                      {achieved != null ? achieved.toFixed(4) : '—'}
-                    </td>
-                    <td style={{ padding:'10px 14px', color:'#6366f1', fontWeight:600 }}>
-                      {LLM_TARGET[d].toFixed(3)}
-                    </td>
-                    <td style={{ padding:'10px 14px', fontWeight:700,
-                      color: delta == null ? '#94a3b8'
-                        : delta >= 0 ? '#16a34a' : '#ef4444' }}>
-                      {delta != null
-                        ? (delta >= 0 ? '+' : '') + delta.toFixed(4)
-                        : '—'}
-                    </td>
-                  </tr>
-                )
-              })}
-            </tbody>
-          </table>
-        </div>
+function ComponentBar({ components }) {
+  const total = COMP_KEYS.reduce((s,k) => s+(components[k]||0), 0)
+  return (
+    <div>
+      <div style={{ display:'flex', height:22, borderRadius:6,
+        overflow:'hidden', marginBottom:6 }}>
+        {COMP_KEYS.map((k,i) => {
+          const v = components[k]||0
+          const pct = total > 0 ? (v/total)*100 : 0
+          return <div key={k} title={`${COMP_LABELS[i]}: ${v.toFixed(4)}`}
+            style={{ width:`${pct}%`, background:COMP_COLORS[i], minWidth:pct>2?2:0 }}/>
+        })}
+      </div>
+      <div style={{ display:'flex', flexWrap:'wrap', gap:'4px 10px' }}>
+        {COMP_KEYS.map((k,i) => (
+          <span key={k} style={{ fontSize:10, color:'#475569',
+            display:'flex', alignItems:'center', gap:3 }}>
+            <span style={{ width:8, height:8, borderRadius:2,
+              background:COMP_COLORS[i], display:'inline-block' }}/>
+            {COMP_LABELS[i].split(' ')[0]}: <b>{(components[k]||0).toFixed(4)}</b>
+          </span>
+        ))}
+      </div>
+    </div>
+  )
+}
+
+// ── Main TrainingDashboard ─────────────────────────────────────────────────────
+export default function TrainingDashboard() {
+  const [activeTab, setActiveTab] = useState('training')
+
+  // Training state (mirrors _training_state on the server)
+  const [trainState, setTrainState] = useState({
+    running: false, status: 'idle', current_step: 0, total_steps: 25,
+    difficulty: 'medium', curve: [], before: null, after: null,
+    metadata: null, error: null,
+  })
+  // Persisted results from /training-log
+  const [savedLog, setSavedLog] = useState(null)
+  const esRef = useRef(null)
+
+  // Load saved training log on mount
+  useEffect(() => {
+    fetch(`${API}/training-log`)
+      .then(r => r.ok ? r.json() : null)
+      .then(d => { if (d) setSavedLog(d) })
+      .catch(() => {})
+  }, [])
+
+  // Start training & subscribe to SSE
+  const startTraining = async (difficulty) => {
+    if (trainState.running) return
+
+    // Kick off training on the server
+    await fetch(`${API}/train/start?difficulty=${difficulty}&steps=25`, { method:'POST' })
+
+    // Subscribe to live SSE stream
+    if (esRef.current) { esRef.current.close(); esRef.current = null }
+    const es = new EventSource(`${API}/train/stream`)
+    esRef.current = es
+
+    es.onmessage = (ev) => {
+      const d = JSON.parse(ev.data)
+      setTrainState(d)
+      if (d.status === 'completed') {
+        // Refresh the saved log once training finishes
+        fetch(`${API}/training-log`)
+          .then(r => r.ok ? r.json() : null)
+          .then(saved => { if (saved) setSavedLog(saved) })
+          .catch(() => {})
+        es.close(); esRef.current = null
+      }
+      if (d.status === 'error') {
+        es.close(); esRef.current = null
+      }
+    }
+    es.onerror = () => { es.close(); esRef.current = null }
+  }
+
+  useEffect(() => () => { if (esRef.current) esRef.current.close() }, [])
+
+  // Decide which data to show: live training state takes priority if running/just-done
+  // otherwise fall back to savedLog
+  const showLive   = trainState.status !== 'idle'
+  const displayLog = showLive ? trainState : savedLog
+
+  const TABS = [
+    { id:'training', label:'🧪 Training Progress' },
+    { id:'benchmark', label:'📈 Benchmarks' },
+  ]
+
+  return (
+    <div>
+      {/* Sub-tabs */}
+      <div style={{ display:'flex', gap:4, marginBottom:20 }}>
+        {TABS.map(t => (
+          <button key={t.id} onClick={() => setActiveTab(t.id)}
+            style={{ padding:'9px 20px', borderRadius:10, border:'none',
+              background: activeTab===t.id ? '#0f172a' : '#e2e8f0',
+              color: activeTab===t.id ? '#fff' : '#64748b',
+              fontWeight:700, fontSize:13, cursor:'pointer' }}>
+            {t.label}
+          </button>
+        ))}
       </div>
 
-      {/* ── GRPO training history ─────────────────────────────────────────── */}
-      <TrainingLog />
+      {activeTab === 'training' && (
+        <>
+          {/* Live training control */}
+          <TrainingProgress state={trainState} onStart={startTraining}/>
+
+          {/* Show results — live first, then saved */}
+          {displayLog && (displayLog.before || (displayLog.curve && displayLog.curve.length > 0)) && (
+            <BeforeAfterSection
+              before={displayLog.before}
+              after={displayLog.after}
+              curve={displayLog.curve}
+            />
+          )}
+
+          {/* No data yet message */}
+          {!showLive && (!savedLog || (!savedLog.before && (!savedLog.curve || !savedLog.curve.length))) && (
+            <div style={{ background:'#fff', border:'1px solid #e2e8f0',
+              borderRadius:14, padding:32, textAlign:'center',
+              color:'#94a3b8', fontSize:14 }}>
+              <div style={{ fontSize:32, marginBottom:8 }}>📊</div>
+              <div style={{ fontWeight:700, color:'#475569', marginBottom:6 }}>
+                No training data yet
+              </div>
+              <div style={{ fontSize:13, lineHeight:1.7 }}>
+                Click <b>▶ medium</b> above to run demo training (~15 seconds).<br/>
+                The before/after comparison and reward curve will appear here.
+              </div>
+            </div>
+          )}
+        </>
+      )}
+
+      {activeTab === 'benchmark' && <BenchmarkSection/>}
     </div>
   )
 }