garvitsachdeva/SpindleFlow-RL
0
1"""2SpindleFlow RL — Streamlit Dashboard3=====================================4Run: cd spindleflow-rl && streamlit run demo/streamlit_app.py5URL: http://localhost:85016"""7 8from __future__ import annotations9import os, sys, json, html as _html10from pathlib import Path11import numpy as np12from dotenv import load_dotenv13 14load_dotenv() # load OPENAI_API_KEY (and any other vars) from .env15 16# HF_HUB_OFFLINE intentionally NOT set — manual HF Hub downloads must work17 18sys.path.insert(0, str(Path(__file__).resolve().parent.parent))19sys.path.insert(0, str(Path(__file__).resolve().parent))20 21import streamlit as st22import plotly.graph_objects as go23from plotly.subplots import make_subplots24 25from env.spindleflow_env import SpindleFlowEnv26from env.state import EpisodeState27from env.specialist_registry import SpecialistRegistry28from orchestrator_widget import render_orchestrator29 30# ─────────────────────────────────────────────────────────31# Page config (must be first Streamlit call)32# ─────────────────────────────────────────────────────────33st.set_page_config(34 page_title="SpindleFlow RL",35 page_icon="⚡",36 layout="wide",37 initial_sidebar_state="collapsed",38)39 40# ─────────────────────────────────────────────────────────41# Constants42# ─────────────────────────────────────────────────────────43CONFIG = "configs/training_config.yaml"44CATALOG = "configs/specialist_catalog.yaml"45ASSETS = Path("demo/assets")46 47SPEC_COLORS = {48 "frontend_react": "#00d4ff",49 "backend_api": "#7c3aed",50 "database_architect": "#f59e0b",51 "devops_engineer": "#10b981",52 "security_analyst": "#ef4444",53 "product_strategist": "#8b5cf6",54 "ux_designer": "#ec4899",55 "tech_writer": "#94a3b8",56}57 58@st.cache_resource59def _get_preset_tasks(n: int = 8) -> list[str]:60 """Sample n live tasks from TaskBank at page load — no hardcoded strings."""61 try:62 from training.task_bank import TaskBank63 bank = TaskBank(phase=1)64 return [bank.sample() for _ in range(n)]65 except Exception:66 # Fallback only if TaskBank is unavailable (e.g. missing config)67 return ["Describe a software engineering task requiring specialist collaboration"]68 69 70PRESET_TASKS = _get_preset_tasks()71 72HF_MODEL_REPO = "garvitsachdeva/spindleflow-rl"73 74 75@st.cache_resource76def _load_trained_model(hf_repo: str):77 """Download RecurrentPPO + VecNormalize stats from HF Hub.78 79 Returns (model, obs_mean, obs_var, clip_obs, error_str).80 Temporarily lifts the HF_HUB_OFFLINE flag set at module level.81 """82 import pickle83 try:84 from huggingface_hub import hf_hub_download85 from sb3_contrib import RecurrentPPO86 87 _tok = os.getenv("HF_TOKEN") or None88 # Try final model first, fall back to latest periodic checkpoint89 try:90 _model_path = hf_hub_download(hf_repo, "spindleflow_model.zip", token=_tok)91 except Exception:92 _model_path = hf_hub_download(hf_repo, "spindleflow_model_latest.zip", token=_tok)93 model = RecurrentPPO.load(_model_path, device="cpu")94 obs_mean = obs_var = None95 clip_obs = 10.096 try:97 try:98 stats_path = hf_hub_download(hf_repo, "vec_normalize.pkl", token=_tok)99 except Exception:100 stats_path = hf_hub_download(hf_repo, "vec_normalize_latest.pkl", token=_tok)101 with open(stats_path, "rb") as f:102 vn = pickle.load(f)103 obs_mean = vn.obs_rms.mean.copy()104 obs_var = vn.obs_rms.var.copy()105 clip_obs = float(vn.clip_obs)106 except Exception:107 pass108 return model, obs_mean, obs_var, clip_obs, None109 except Exception as exc:110 return None, None, None, 10.0, str(exc)111 finally:112 pass113 114 115def _predict(model, obs: np.ndarray, lstm_states, episode_starts,116 obs_mean, obs_var, clip_obs: float):117 """Normalize obs and call model.predict(); return (action, new_lstm_states)."""118 obs_arr = obs[np.newaxis, :].copy().astype(np.float32)119 if obs_mean is not None and obs_var is not None:120 obs_arr = np.clip(121 (obs_arr - obs_mean) / np.sqrt(obs_var + 1e-8),122 -clip_obs, clip_obs,123 )124 action_batch, new_states = model.predict(125 obs_arr,126 state=lstm_states,127 episode_start=episode_starts,128 deterministic=True,129 )130 return action_batch[0], new_states131 132 133DARK = dict(134 paper_bgcolor="rgba(0,0,0,0)",135 plot_bgcolor="rgba(0,0,0,0)",136 font=dict(color="#e2e8f0", family="Inter, system-ui, sans-serif"),137 margin=dict(l=44, r=20, t=44, b=40),138)139DARK_AXES = dict(140 xaxis=dict(gridcolor="rgba(255,255,255,0.05)", zerolinecolor="rgba(255,255,255,0.08)"),141 yaxis=dict(gridcolor="rgba(255,255,255,0.05)", zerolinecolor="rgba(255,255,255,0.08)"),142)143 144# ─────────────────────────────────────────────────────────145# Session state146# ─────────────────────────────────────────────────────────147class Session:148 def __init__(self):149 self.env: SpindleFlowEnv | None = None150 self.registry: SpecialistRegistry | None = None151 self.rewards: list[float] = []152 self.actions: list[dict] = []153 self.step_n = 0154 self.done = False155 self.task = ""156 # Full episode history for replay157 self.episode_history: list[dict] = []158 # Action entropy per step (policy confidence)159 self.step_entropies: list[float] = []160 # Observation vector stats per step161 self.obs_history: list[dict] = []162 # Specialists auto-spawned for this episode163 self.spawned_specialists: list[str] = []164 # Trained policy inference state165 self.obs_current: np.ndarray | None = None166 self.lstm_states = None167 self.episode_starts = np.array([True])168 169 def boot(self):170 if self.env is None:171 self.env = SpindleFlowEnv(172 config_path=CONFIG, catalog_path=CATALOG,173 use_real_spindleflow=False, phase=1,174 )175 self.registry = self.env.registry176 177 def reset(self, phase: int = 1):178 self.boot()179 self.env.phase = int(phase)180 obs, info = self.env.reset()181 self.rewards = []182 self.actions = []183 self.step_n = 0184 self.done = False185 self.task = info.get("task", "")186 self.episode_history = []187 self.step_entropies = []188 self.obs_history = []189 self.spawned_specialists: list[str] = list(info.get("spawned_specialists", []))190 self.obs_current = obs191 self.lstm_states = None192 self.episode_starts = np.array([True])193 return obs, info194 195 def step(self, action):196 if self.env is None or self.done:197 return None, 0.0, True, False, {}198 obs, r, term, trunc, info = self.env.step(action)199 self.rewards.append(r)200 self.actions.append(info)201 self.step_n += 1202 self.done = term or trunc203 self.obs_current = obs204 self.episode_starts = np.array([self.done])205 206 # Capture step snapshot for replay207 called = info.get("called_specialists", [])208 edges = [(e.caller_id, e.callee_id)209 for e in self.env.delegation_graph.get_delegation_path()]210 self.episode_history.append({211 "step": self.step_n,212 "reward": r,213 "action_name": info.get("action_name", "UNKNOWN"),214 "called": list(called),215 "edges": list(edges),216 "components": dict(info.get("reward_components", {})),217 "mode": info.get("delegation_mode", ""),218 "cumulative": float(sum(self.rewards)),219 "latencies": dict(info.get("specialist_latencies", {})),220 })221 222 # Compute real action entropy (specialist-selection logits)223 if self.env is not None:224 n = self.env.max_specialists225 spec_logits = action[1: 1 + n].copy()226 spec_logits = spec_logits - spec_logits.max()227 exp_l = np.exp(spec_logits)228 probs = exp_l / (exp_l.sum() + 1e-8)229 entropy = float(-np.sum(probs * np.log(probs + 1e-8)))230 self.step_entropies.append(entropy)231 232 # Capture observation norm for state trace233 if obs is not None:234 self.obs_history.append({235 "step": self.step_n,236 "obs_norm": float(np.linalg.norm(obs)),237 "obs_mean": float(obs.mean()),238 "obs_max": float(obs.max()),239 })240 241 return obs, r, term, trunc, info242 243 244def _S() -> Session:245 if "session" not in st.session_state:246 st.session_state.session = Session()247 return st.session_state.session248 249 250def _load_catalog() -> list[dict]:251 import yaml252 with open(CATALOG) as f:253 return yaml.safe_load(f)["specialists"]254 255 256def _exec_mode_badges(S: "Session") -> str:257 """Return inline HTML badge strip showing execution and task-generation modes."""258 import os259 has_key = bool(os.getenv("OPENAI_API_KEY"))260 llm_tasks = S.env is not None and S.env.task_bank._client is not None261 262 exec_b = (263 '<span style="padding:3px 10px;border-radius:999px;font-size:10px;font-weight:700;'264 'background:rgba(16,185,129,0.1);color:#34d399;'265 'border:1px solid rgba(16,185,129,0.22);">● LLM BASELINE</span>'266 if has_key else267 '<span style="padding:3px 10px;border-radius:999px;font-size:10px;font-weight:700;'268 'background:rgba(245,158,11,0.1);color:#fbbf24;'269 'border:1px solid rgba(245,158,11,0.22);">'270 '⚡ SIMULATION MODE — specialist outputs templated · set OPENAI_API_KEY for real LLM</span>'271 )272 task_b = (273 '<span style="padding:3px 10px;border-radius:999px;font-size:10px;font-weight:700;'274 'background:rgba(16,185,129,0.1);color:#34d399;'275 'border:1px solid rgba(16,185,129,0.22);">● LLM TASKS</span>'276 if llm_tasks else277 '<span style="padding:3px 10px;border-radius:999px;font-size:10px;font-weight:700;'278 'background:rgba(148,163,184,0.08);color:#64748b;'279 'border:1px solid rgba(148,163,184,0.18);">⚡ CATALOG TASKS</span>'280 ) if S.env is not None else ""281 282 return (283 f'<div style="display:flex;gap:8px;flex-wrap:wrap;margin:4px 0 12px;">'284 f'{exec_b}{task_b}</div>'285 )286 287# ─────────────────────────────────────────────────────────288# Chart builders289# ─────────────────────────────────────────────────────────290def fig_reward_curve(rewards: list[float]) -> go.Figure:291 if not rewards:292 fig = go.Figure()293 fig.update_layout(294 **DARK, **DARK_AXES,295 title=dict(text="Episode Reward", font=dict(size=13, color="#64748b")),296 annotations=[dict(text="Reset the environment to begin",297 x=0.5, y=0.5, showarrow=False,298 font=dict(color="#334155", size=13))],299 )300 return fig301 302 steps = list(range(len(rewards)))303 cumul = np.cumsum(rewards).tolist()304 fig = make_subplots(rows=2, cols=1, shared_xaxes=True,305 row_heights=[0.62, 0.38], vertical_spacing=0.04)306 fig.add_trace(go.Scatter(307 x=steps, y=cumul, mode="lines",308 line=dict(color="#00d4ff", width=2.5),309 fill="tozeroy", fillcolor="rgba(0,212,255,0.07)",310 name="Cumulative",311 ), row=1, col=1)312 fig.add_trace(go.Bar(313 x=steps, y=rewards,314 marker_color=["#10b981" if r >= 0 else "#ef4444" for r in rewards],315 marker_line_width=0, name="Per-step",316 ), row=2, col=1)317 fig.update_layout(**DARK, height=300, showlegend=False,318 title=dict(text="Episode Reward", font=dict(size=13, color="#94a3b8")))319 fig.update_xaxes(gridcolor="rgba(255,255,255,0.05)")320 fig.update_yaxes(gridcolor="rgba(255,255,255,0.05)",321 title_text="Cumul.", row=1, col=1, title_font_size=10)322 fig.update_yaxes(title_text="Step", row=2, col=1, title_font_size=10)323 return fig324 325 326def fig_delegation_graph(327 S: Session,328 called_ids: list[str],329 edges: list[tuple],330 highlight_latest: bool = True,331 spawned_ids: list[str] | None = None,332) -> go.Figure:333 """334 Professional hierarchical DAG layout.335 Orchestrator at top, called specialists in middle, uncalled dimmed at bottom.336 """337 all_ids = list(S.registry.list_ids()) if S.registry else []338 called_set = set(called_ids)339 spawned_set = set(spawned_ids or S.spawned_specialists)340 uncalled = [x for x in all_ids if x not in called_set]341 342 # ── Build node positions (hierarchical layout) ───────────────────343 pos = {"orchestrator": (0.5, 0.92)}344 345 n_called = len(called_ids)346 if n_called > 0:347 for i, sid in enumerate(called_ids):348 x = (i + 1) / (n_called + 1)349 pos[sid] = (x, 0.55)350 351 n_uncalled = len(uncalled)352 if n_uncalled > 0:353 for i, sid in enumerate(uncalled):354 x = (i + 1) / (n_uncalled + 1)355 pos[sid] = (x, 0.12)356 357 fig = go.Figure()358 359 # ── Background depth ring ────────────────────────────────────────360 max_depth = getattr(S.env, "max_depth", 2) if S.env else 2361 cur_depth = S.env.delegation_graph.depth if S.env else 0362 depth_frac = cur_depth / max(max_depth, 1)363 ring_color = ("#10b981" if depth_frac < 0.7364 else ("#f59e0b" if depth_frac < 1.0 else "#ef4444"))365 366 fig.add_shape(type="rect",367 x0=0.0, y0=0.0, x1=1.0, y1=1.0,368 line=dict(color=ring_color, width=2, dash="dot"),369 fillcolor="rgba(0,0,0,0)", xref="x", yref="y",370 )371 fig.add_annotation(372 x=0.98, y=0.98, xref="x", yref="y",373 text=f"Depth {cur_depth}/{max_depth}", showarrow=False,374 font=dict(size=9, color=ring_color), xanchor="right", yanchor="top",375 )376 377 # ── Edges ────────────────────────────────────────────────────────378 latest_edge = edges[-1] if edges else None379 for src, dst in edges:380 if src not in pos or dst not in pos:381 continue382 x0, y0 = pos[src]383 x1, y1 = pos[dst]384 is_latest = (latest_edge and highlight_latest and (src, dst) == latest_edge)385 color = "rgba(0,212,255,0.9)" if is_latest else "rgba(0,212,255,0.45)"386 width = 2.5 if is_latest else 1.8387 dash = "dash" if is_latest else "solid"388 389 fig.add_trace(go.Scatter(390 x=[x0, x1, None], y=[y0, y1, None], mode="lines",391 line=dict(color=color, width=width, dash=dash),392 hoverinfo="skip", showlegend=False,393 ))394 fig.add_annotation(395 ax=x0, ay=y0, x=x1, y=y1,396 xref="x", yref="y", axref="x", ayref="y",397 arrowhead=3, arrowsize=1.4, arrowwidth=2,398 arrowcolor=color, showarrow=True,399 )400 401 # ── Orchestrator node ────────────────────────────────────────────402 ox, oy = pos["orchestrator"]403 fig.add_trace(go.Scatter(404 x=[ox], y=[oy], mode="markers+text",405 marker=dict(size=44, color="#f59e0b", symbol="circle",406 line=dict(color="#fcd34d", width=2.5), opacity=1.0),407 text=["<b>ORCH</b>"], textposition="middle center",408 textfont=dict(size=9, color="#0a0f1a", family="Inter, sans-serif"),409 hovertext=["<b>Orchestrator</b><br>Root node — makes all delegation decisions"],410 hoverinfo="text", showlegend=False, name="orchestrator",411 ))412 413 # ── Called specialist nodes ──────────────────────────────────────414 for sid in called_ids:415 if sid not in pos:416 continue417 x, y = pos[sid]418 c = SPEC_COLORS.get(sid, "#7c3aed")419 spec = S.registry.get(sid) if S.registry else None420 role = spec.role if spec else sid421 lat = f"{spec.avg_latency_ms}ms" if spec else ""422 is_spawned = sid in spawned_set423 symbol = "star" if is_spawned else "circle"424 size = 38 if is_spawned else 32425 border_c = "#fbbf24" if is_spawned else "rgba(255,255,255,0.4)"426 hover_tag = " ⚡ AUTO-SPAWNED" if is_spawned else ""427 label = (("⚡ " if is_spawned else "") + sid).replace("_", "<br>")428 fig.add_trace(go.Scatter(429 x=[x], y=[y], mode="markers+text",430 marker=dict(size=size, color=c, symbol=symbol,431 line=dict(color=border_c, width=2.5), opacity=1.0),432 text=[label], textposition="bottom center",433 textfont=dict(size=8, color="#fbbf24" if is_spawned else "#e2e8f0"),434 hovertext=[f"<b>{role}</b><br>Called ✓{hover_tag}<br>{lat}"],435 hoverinfo="text", showlegend=False,436 ))437 438 # ── Uncalled specialist nodes (dimmed) ───────────────────────────439 for sid in uncalled:440 if sid not in pos:441 continue442 x, y = pos[sid]443 c = SPEC_COLORS.get(sid, "#334155")444 spec = S.registry.get(sid) if S.registry else None445 role = spec.role if spec else sid446 label = sid.replace("_", "<br>")447 fig.add_trace(go.Scatter(448 x=[x], y=[y], mode="markers+text",449 marker=dict(size=16, color="#1e293b", symbol="circle",450 line=dict(color=c, width=1), opacity=0.5),451 text=[label], textposition="bottom center",452 textfont=dict(size=7, color="rgba(148,163,184,0.45)"),453 hovertext=[f"<b>{role}</b><br>Not called"],454 hoverinfo="text", showlegend=False,455 ))456 457 # ── Section labels ───────────────────────────────────────────────458 fig.add_annotation(x=0.01, y=0.96, xref="x", yref="y",459 text="ORCHESTRATOR", showarrow=False,460 font=dict(size=8, color="#475569"), xanchor="left")461 if called_ids:462 fig.add_annotation(x=0.01, y=0.62, xref="x", yref="y",463 text="CALLED", showarrow=False,464 font=dict(size=8, color="#00d4ff"), xanchor="left")465 if uncalled:466 fig.add_annotation(x=0.01, y=0.19, xref="x", yref="y",467 text="AVAILABLE", showarrow=False,468 font=dict(size=8, color="#334155"), xanchor="left")469 470 fig.update_layout(471 **DARK, height=420,472 title=dict(473 text=(f"Delegation Graph · {len(called_ids)} specialists called"474 f" · Depth {cur_depth}/{max_depth}"),475 font=dict(size=13, color="#94a3b8"),476 ),477 xaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[-0.05, 1.05]),478 yaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[-0.05, 1.08]),479 )480 return fig481 482 483def fig_reward_breakdown(components: dict) -> go.Figure:484 if not components:485 components = {k: 0.0 for k in [486 "quality_delta", "efficiency_penalty", "failure_penalty",487 "recovery_bonus", "conflict_penalty", "conflict_bonus",488 "consistency_bonus", "latency_penalty", "explanation_bonus",489 ]}490 names = list(components.keys())491 values = [components[k] for k in names]492 fig = go.Figure(go.Bar(493 x=values,494 y=[n.replace("_", " ").title() for n in names],495 orientation="h",496 marker_color=["#10b981" if v >= 0 else "#ef4444" for v in values],497 marker_line_width=0,498 text=[f"{v:+.3f}" for v in values],499 textposition="outside",500 textfont=dict(color="#94a3b8", size=9),501 ))502 fig.add_vline(x=0, line_color="rgba(255,255,255,0.15)", line_width=1)503 fig.update_layout(**DARK, height=310,504 title=dict(text="Reward Breakdown", font=dict(size=13, color="#94a3b8")),505 xaxis=dict(gridcolor="rgba(255,255,255,0.05)", title="Value"),506 yaxis=dict(gridcolor="rgba(255,255,255,0.05)"))507 return fig508 509 510def fig_policy_confidence(511 entropies: list[float],512 step_labels: list[int] | None = None,513) -> go.Figure:514 """515 Policy confidence chart — specialist-selection entropy per step.516 High entropy = uncertain/exploring. Low = confident/committed.517 Real data from actual action vectors used each step.518 """519 if not entropies:520 fig = go.Figure()521 fig.update_layout(522 **DARK, **DARK_AXES,523 title=dict(text="Policy Confidence (Action Entropy)",524 font=dict(size=13, color="#64748b")),525 annotations=[dict(text="Run an episode to see real action entropy",526 x=0.5, y=0.5, showarrow=False,527 font=dict(color="#334155", size=12))],528 )529 return fig530 531 steps = step_labels or list(range(1, len(entropies) + 1))532 max_e = float(np.log(max(len(entropies), 2)))533 norm_e = [min(1.0, max(0.0, e / max(max_e, 1e-8))) for e in entropies]534 colors = [535 f"rgba({int(0 + 124 * ne)},{int(212 - 154 * ne)},{int(255 - 58 * ne)},0.85)"536 for ne in norm_e537 ]538 539 fig = go.Figure()540 fig.add_trace(go.Bar(541 x=steps, y=norm_e,542 marker_color=colors, marker_line_width=0,543 name="Normalised entropy",544 text=[f"{e:.3f}" for e in entropies],545 textposition="outside",546 textfont=dict(size=8, color="#94a3b8"),547 hovertemplate="Step %{x}<br>Entropy: %{text}<extra></extra>",548 ))549 fig.add_hline(y=0.5, line_dash="dot", line_color="rgba(148,163,184,0.3)",550 annotation_text="Mid-entropy", annotation_font_color="#475569")551 fig.update_layout(552 **DARK, height=260,553 title=dict(text="Policy Confidence — Specialist Selection Entropy per Step",554 font=dict(size=12, color="#94a3b8")),555 xaxis=dict(title="Episode Step", gridcolor="rgba(255,255,255,0.05)",556 zerolinecolor="rgba(255,255,255,0.08)"),557 yaxis=dict(title="Entropy (0=certain, 1=uniform)", range=[0, 1.15],558 gridcolor="rgba(255,255,255,0.05)", zerolinecolor="rgba(255,255,255,0.08)"),559 showlegend=False,560 )561 return fig562 563 564def fig_similarity(registry: SpecialistRegistry) -> go.Figure:565 ids = registry.list_ids()566 n = len(ids)567 568 if n == 0:569 fig = go.Figure()570 fig.update_layout(**DARK, title=dict(text="No specialists in registry",571 font=dict(size=13, color="#64748b")))572 return fig573 574 missing = [sid for sid in ids if registry.get(sid).embedding is None]575 if missing:576 fig = go.Figure()577 fig.update_layout(578 **DARK, **DARK_AXES,579 title=dict(text="Embeddings not computed — boot the environment first",580 font=dict(size=13, color="#64748b")),581 annotations=[dict(text=f"Missing embeddings: {', '.join(missing[:4])}",582 x=0.5, y=0.5, showarrow=False,583 font=dict(color="#334155", size=12))],584 )585 return fig586 587 mat = np.zeros((n, n))588 try:589 for i, a in enumerate(ids):590 for j, b in enumerate(ids):591 ea = registry.get(a).to_state_vector()592 eb = registry.get(b).to_state_vector()593 mat[i][j] = float(np.dot(ea, eb))594 except Exception as exc:595 fig = go.Figure()596 fig.update_layout(**DARK, title=dict(text=f"Similarity error: {exc}",597 font=dict(size=13, color="#ef4444")))598 return fig599 labels = [x.replace("_", "<br>") for x in ids]600 fig = go.Figure(go.Heatmap(601 z=mat, x=labels, y=labels,602 colorscale=[[0, "#0f0f1a"], [0.5, "rgba(124,58,237,0.6)"], [1, "#00d4ff"]],603 showscale=True, zmin=0, zmax=1,604 text=np.round(mat, 2), texttemplate="%{text}", textfont=dict(size=9),605 ))606 fig.update_layout(**DARK, height=400,607 title=dict(text="Capability Similarity (Cosine)", font=dict(size=13, color="#94a3b8")))608 return fig609 610 611def fig_training_curve() -> go.Figure:612 path = ASSETS / "reward_curve.json"613 if path.exists():614 with open(path) as f:615 d = json.load(f)616 eps, rews = d["episodes"], d["mean_rewards"]617 else:618 rng = np.random.default_rng(42)619 eps = list(range(0, 201, 5))620 rews = [float(np.clip(0.1 + 0.5 * (1 - np.exp(-e / 80)) + rng.normal(0, 0.04), 0, 1))621 for e in eps]622 smooth = [float(np.mean(rews[max(0, i - 4):i + 1])) for i in range(len(rews))]623 fig = go.Figure()624 fig.add_trace(go.Scatter(x=eps, y=rews, mode="markers",625 marker=dict(size=5, color="rgba(0,212,255,0.35)"),626 name="Episode"))627 fig.add_trace(go.Scatter(x=eps, y=smooth, mode="lines",628 line=dict(color="#00d4ff", width=2.5),629 fill="tozeroy", fillcolor="rgba(0,212,255,0.06)",630 name="Smoothed"))631 fig.add_hline(y=0.1, line_dash="dash", line_color="rgba(148,163,184,0.35)",632 annotation_text="Random baseline", annotation_font_color="#64748b")633 fig.update_layout(**DARK, **DARK_AXES, height=340,634 title=dict(text="Training Progress — Mean Reward per Episode",635 font=dict(size=13, color="#94a3b8")),636 xaxis_title="Episode", yaxis_title="Mean Reward",637 legend=dict(bgcolor="rgba(0,0,0,0)"))638 return fig639 640 641def fig_training_entropy() -> go.Figure:642 """643 Policy entropy over training.644 Reads from demo/assets/entropy_log.json if produced by train.py,645 or from current session entropy if no log exists.646 Never shows fake data — gracefully absent if neither source exists.647 """648 path = ASSETS / "entropy_log.json"649 S = _S()650 651 if path.exists():652 with open(path) as f:653 d = json.load(f)654 episodes = d["episodes"]655 entropies = d["mean_entropies"]656 source_label = "From training log"657 elif S.step_entropies:658 episodes = list(range(1, len(S.step_entropies) + 1))659 entropies = S.step_entropies660 source_label = "Current episode (live)"661 else:662 fig = go.Figure()663 fig.update_layout(664 **DARK, **DARK_AXES,665 title=dict(text="Policy Entropy — Run training to populate",666 font=dict(size=13, color="#64748b")),667 annotations=[dict(668 text="Run python training/train.py to generate entropy logs",669 x=0.5, y=0.5, showarrow=False,670 font=dict(color="#334155", size=12),671 )],672 )673 return fig674 675 fig = go.Figure()676 fig.add_trace(go.Scatter(677 x=episodes, y=entropies, mode="lines+markers",678 line=dict(color="#7c3aed", width=2.2),679 marker=dict(size=4, color="#a78bfa"),680 fill="tozeroy", fillcolor="rgba(124,58,237,0.06)",681 name=source_label,682 ))683 fig.update_layout(684 **DARK, **DARK_AXES, height=280,685 title=dict(text=f"Policy Entropy over Training ({source_label})",686 font=dict(size=13, color="#94a3b8")),687 xaxis_title="Episode / Step",688 yaxis_title="Action Selection Entropy",689 legend=dict(bgcolor="rgba(0,0,0,0)"),690 )691 return fig692 693 694# ─────────────────────────────────────────────────────────695# Quality-comparison helpers696# ─────────────────────────────────────────────────────────697def _generate_generic_output(task: str) -> str:698 """Call GPT-4o-mini directly with the task — no specialist routing."""699 import os700 api_key = os.getenv("OPENAI_API_KEY")701 if not api_key:702 return (703 "General problem-solving approach:\n"704 "1. Gather and clarify requirements\n"705 "2. Research common solution patterns\n"706 "3. Draft a high-level architecture\n"707 "4. Implement in small, testable increments\n"708 "5. Validate against acceptance criteria and deploy\n"709 "No specialist domain expertise applied."710 )711 try:712 from openai import OpenAI713 resp = OpenAI(api_key=api_key).chat.completions.create(714 model="gpt-4o-mini",715 max_tokens=600,716 messages=[717 {"role": "system",718 "content": "You are a general-purpose software engineering assistant."},719 {"role": "user",720 "content": f"Provide a detailed solution approach for this task:\n\n{task}"},721 ],722 )723 return resp.choices[0].message.content724 except Exception as exc:725 return f"(Generic output generation failed: {exc})"726 727 728def _t1_relevance(task: str, output: str, registry) -> float:729 """Cosine similarity between task and output embeddings, scaled 0–10."""730 try:731 import numpy as np732 t = registry.embed_query(task)733 o = registry.embed_query(output[:800])734 if t is None or o is None:735 return 0.0736 cos = float(np.dot(t, o) / (np.linalg.norm(t) * np.linalg.norm(o) + 1e-8))737 return round(max(0.0, cos) * 10, 2)738 except Exception:739 return 0.0740 741 742def _judge_compare(task: str, generic: str, specialist: str) -> dict | None:743 """GPT-4o-mini rates both outputs on 4 dimensions. Returns {dim: [generic, specialist]}."""744 import os, json745 api_key = os.getenv("OPENAI_API_KEY")746 if not api_key:747 return None748 prompt = (749 f"Task:\n{task[:400]}\n\n"750 f"Output A (generic, no specialist routing):\n{generic[:700]}\n\n"751 f"Output B (specialist-routed by trained policy):\n{specialist[:700]}\n\n"752 "Rate each output 1–10 on: technical_depth, specificity, actionability, coverage.\n"753 'Return JSON only: {"technical_depth":[A,B],"specificity":[A,B],'754 '"actionability":[A,B],"coverage":[A,B]}'755 )756 try:757 from openai import OpenAI758 resp = OpenAI(api_key=api_key).chat.completions.create(759 model="gpt-4o-mini",760 max_tokens=150,761 response_format={"type": "json_object"},762 messages=[{"role": "user", "content": prompt}],763 )764 return json.loads(resp.choices[0].message.content)765 except Exception:766 return None767 768 769def fig_radar_comparison(770 gen_scores: dict,771 spec_scores: dict,772) -> go.Figure:773 dims = list(gen_scores.keys())774 g_vals = [gen_scores[d] for d in dims]775 s_vals = [spec_scores[d] for d in dims]776 dims_c = dims + [dims[0]]777 g_c = g_vals + [g_vals[0]]778 s_c = s_vals + [s_vals[0]]779 780 fig = go.Figure()781 fig.add_trace(go.Scatterpolar(782 r=g_c, theta=dims_c, fill="toself",783 fillcolor="rgba(239,68,68,0.10)",784 line=dict(color="#ef4444", width=2),785 name="Generic (no routing)",786 ))787 fig.add_trace(go.Scatterpolar(788 r=s_c, theta=dims_c, fill="toself",789 fillcolor="rgba(0,212,255,0.13)",790 line=dict(color="#00d4ff", width=2.5),791 name="Specialist-routed",792 ))793 fig.update_layout(794 paper_bgcolor="rgba(0,0,0,0)",795 font=dict(color="#e2e8f0", family="Inter, system-ui, sans-serif"),796 polar=dict(797 bgcolor="rgba(0,0,0,0)",798 radialaxis=dict(799 visible=True, range=[0, 10],800 gridcolor="rgba(255,255,255,0.08)",801 tickfont=dict(size=9, color="#475569"),802 ),803 angularaxis=dict(804 gridcolor="rgba(255,255,255,0.08)",805 tickfont=dict(size=11, color="#94a3b8"),806 ),807 ),808 title=dict(809 text="Quality Radar — Generic vs Specialist-Routed",810 font=dict(size=13, color="#94a3b8"),811 ),812 legend=dict(bgcolor="rgba(0,0,0,0)", font=dict(color="#94a3b8", size=11)),813 height=420,814 margin=dict(l=60, r=60, t=60, b=40),815 )816 return fig817 818 819# ─────────────────────────────────────────────────────────820# UI helpers821# ─────────────────────────────────────────────────────────822def inject_css():823 st.markdown("""824<style>825@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;800&display=swap');826 827html, body, [data-testid="stAppViewContainer"] {828 background: #0f0f1a !important;829 font-family: 'Inter', system-ui, sans-serif !important;830}831[data-testid="stHeader"] { background: transparent !important; }832[data-testid="stToolbar"] { display: none !important; }833 834[data-testid="stTabs"] > div:first-child button {835 color: #475569 !important; font-weight: 600 !important; font-size: 13px !important;836}837[data-testid="stTabs"] > div:first-child button[aria-selected="true"] {838 color: #00d4ff !important; border-bottom-color: #00d4ff !important;839}840 841.stButton > button {842 border-radius: 8px !important; font-weight: 600 !important;843 font-size: 13px !important; transition: all .18s !important;844 border: 1px solid rgba(255,255,255,0.18) !important;845 background: rgba(255,255,255,0.10) !important; color: #e2e8f0 !important;846}847.stButton > button:hover {848 background: rgba(255,255,255,0.18) !important;849 border-color: rgba(0,212,255,0.45) !important;850 color: #ffffff !important;851}852.stButton > button[kind="primary"] {853 background: linear-gradient(135deg,#00d4ff,#0092bb) !important;854 border: none !important; color: #0a0f1a !important;855}856.stButton > button[kind="primary"]:hover {857 box-shadow: 0 4px 18px rgba(0,212,255,0.35) !important;858}859 860[data-testid="stTextInput"] input,861[data-testid="stTextArea"] textarea {862 background: rgba(0,0,0,0.3) !important;863 border: 1px solid rgba(255,255,255,0.09) !important;864 color: #e2e8f0 !important; border-radius: 8px !important;865}866 867[data-testid="stSelectbox"] > div > div {868 background: rgba(0,0,0,0.35) !important;869 border: 1px solid rgba(255,255,255,0.09) !important;870 border-radius: 8px !important; color: #e2e8f0 !important;871}872 873[data-testid="stSlider"] [data-testid="stTickBar"] { color: #475569 !important; }874 875[data-testid="metric-container"] {876 background: rgba(255,255,255,0.03) !important;877 border: 1px solid rgba(255,255,255,0.07) !important;878 border-radius: 12px !important; padding: 16px !important;879}880[data-testid="stMetric"] label { color: #475569 !important; font-size: 11px !important; }881[data-testid="stMetricValue"] { color: #00d4ff !important; font-weight: 700 !important; }882 883[data-testid="stCode"], .stCodeBlock {884 background: rgba(0,0,0,0.4) !important;885 border: 1px solid rgba(255,255,255,0.07) !important;886 border-radius: 10px !important;887}888 889hr { border-color: rgba(255,255,255,0.07) !important; }890 891::-webkit-scrollbar { width: 4px; height: 4px; }892::-webkit-scrollbar-thumb { background: rgba(255,255,255,0.1); border-radius: 4px; }893::-webkit-scrollbar-track { background: transparent; }894</style>895""", unsafe_allow_html=True)896 897 898def hero():899 st.markdown("""900<div style="background:linear-gradient(135deg,#0f0f1a,#130a22,#091422);901 border:1px solid rgba(0,212,255,0.14);border-radius:16px;902 padding:28px 36px;margin-bottom:4px;position:relative;overflow:hidden;">903 <div style="position:absolute;top:-60px;right:-40px;width:360px;height:360px;904 background:radial-gradient(circle,rgba(124,58,237,0.11) 0%,transparent 70%);905 pointer-events:none;"></div>906 <div style="position:absolute;bottom:-60px;left:15%;width:280px;height:280px;907 background:radial-gradient(circle,rgba(0,212,255,0.07) 0%,transparent 70%);908 pointer-events:none;"></div>909 <div style="font-size:28px;font-weight:800;910 background:linear-gradient(90deg,#00d4ff,#7c3aed,#00d4ff);911 background-size:200% auto;-webkit-background-clip:text;912 -webkit-text-fill-color:transparent;background-clip:text;913 margin:0 0 8px;">SpindleFlow RL</div>914 <div style="color:#64748b;font-size:13px;margin:0;">915 Delegation Policy Learning Environment —916 Teaching orchestrators to route, specialize, and stop.917 </div>918</div>919""", unsafe_allow_html=True)920 921 922def sec(title: str):923 st.markdown(924 f'<div style="font-size:11px;font-weight:700;color:#475569;text-transform:uppercase;'925 f'letter-spacing:1px;padding-bottom:10px;border-bottom:1px solid rgba(255,255,255,0.07);'926 f'margin:18px 0 14px;">{title}</div>',927 unsafe_allow_html=True,928 )929 930 931def status_bar(msg: str, color: str = "#94a3b8"):932 st.markdown(933 f'<div style="background:rgba(0,0,0,0.3);border:1px solid rgba(255,255,255,0.07);'934 f'border-radius:8px;padding:10px 16px;font-size:12px;color:{color};margin:6px 0 10px;">'935 f'{_html.escape(msg)}</div>',936 unsafe_allow_html=True,937 )938 939 940def render_live_stats(S: Session) -> None:941 """Sidebar live stats strip — all values read directly from session state."""942 with st.sidebar:943 st.markdown(944 '<div style="font-size:10px;font-weight:700;color:#00d4ff;'945 'text-transform:uppercase;letter-spacing:1px;margin-bottom:12px;">'946 '● Live Episode Stats</div>',947 unsafe_allow_html=True,948 )949 950 status = ("Running" if (S.env is not None and not S.done) else951 "Complete" if S.done else "Idle")952 status_color = ("#10b981" if status == "Running" else953 "#f59e0b" if status == "Complete" else "#475569")954 st.markdown(955 f'<div style="display:flex;justify-content:space-between;'956 f'padding:6px 0;border-bottom:1px solid rgba(255,255,255,0.05);">'957 f'<span style="font-size:11px;color:#475569;">Status</span>'958 f'<span style="font-size:11px;font-weight:700;color:{status_color};">'959 f'{status}</span></div>',960 unsafe_allow_html=True,961 )962 963 unique_called = len(set(964 sp for h in S.episode_history for sp in h.get("called", [])965 ))966 dag_depth = str(S.env.delegation_graph.depth) if S.env else "—"967 968 stats = [969 ("Step", str(S.step_n), "#e2e8f0"),970 ("Total Reward", f"{sum(S.rewards):+.4f}" if S.rewards else "—",971 "#10b981" if (S.rewards and sum(S.rewards) >= 0) else "#ef4444"),972 ("Mean Step Rwd",f"{float(np.mean(S.rewards)):+.4f}" if S.rewards else "—", "#94a3b8"),973 ("Specialists", str(unique_called), "#7c3aed"),974 ("DAG Depth", dag_depth, "#f59e0b"),975 ("Mean Entropy", f"{float(np.mean(S.step_entropies)):.3f}"976 if S.step_entropies else "—", "#00d4ff"),977 ]978 979 for label, value, color in stats:980 st.markdown(981 f'<div style="display:flex;justify-content:space-between;'982 f'padding:5px 0;border-bottom:1px solid rgba(255,255,255,0.04);">'983 f'<span style="font-size:11px;color:#475569;">{label}</span>'984 f'<span style="font-size:11px;font-weight:600;color:{color};">'985 f'{value}</span></div>',986 unsafe_allow_html=True,987 )988 989 if S.rewards:990 st.markdown('<div style="margin-top:12px;"></div>', unsafe_allow_html=True)991 st.plotly_chart(fig_reward_curve(S.rewards), use_container_width=True)992 993 994def _render_replay_step(S: Session, step_idx: int) -> None:995 """Render charts for a specific historical step — no env calls."""996 if not S.episode_history or step_idx >= len(S.episode_history):997 st.info("No episode data to replay. Run an episode first.")998 return999 1000 snap = S.episode_history[step_idx]1001 cumulative = snap["cumulative"]1002 1003 # Cumulative called specialists up to and including this step1004 cumulative_called = list({1005 sp1006 for h in S.episode_history[:step_idx + 1]1007 for sp in h.get("called", [])1008 })1009 1010 st.markdown(1011 f'<div style="background:rgba(124,58,237,0.07);border:1px solid rgba(124,58,237,0.2);'1012 f'border-radius:10px;padding:12px 18px;font-size:12px;color:#a78bfa;margin-bottom:12px;">'1013 f'Replaying Step {snap["step"]} · Action: <b>{snap["action_name"]}</b> · '1014 f'Reward: <b>{snap["reward"]:+.4f}</b> · '1015 f'Cumulative: <b>{cumulative:+.4f}</b></div>',1016 unsafe_allow_html=True,1017 )1018 1019 rc1, rc2 = st.columns(2)1020 with rc1:1021 st.plotly_chart(1022 fig_delegation_graph(S, cumulative_called, snap["edges"], highlight_latest=False),1023 use_container_width=True,1024 key=f"replay_dag_{step_idx}",1025 )1026 with rc2:1027 st.plotly_chart(1028 fig_reward_breakdown(snap["components"]),1029 use_container_width=True,1030 key=f"replay_breakdown_{step_idx}",1031 )1032 1033 sec("Action Trace at This Step")1034 trace_lines = []1035 for h in S.episode_history[:step_idx + 1]:1036 sign = "+" if h["reward"] >= 0 else ""1037 called_str = ", ".join(h["called"]) if h["called"] else "—"1038 marker = "► " if h["step"] == snap["step"] else " "1039 trace_lines.append(1040 f"{marker}Step {h['step']:>2} │ {h['action_name']:<22} │ "1041 f"reward: {sign}{h['reward']:.4f} │ specialists: {called_str}"1042 )1043 st.code("\n".join(trace_lines), language=None)1044 1045 1046# ─────────────────────────────────────────────────────────1047# Tab 1 — Live Demo1048# ─────────────────────────────────────────────────────────1049def tab_live_demo():1050 S = _S()1051 1052 col_task, col_ctrl = st.columns([3, 2], gap="large")1053 1054 with col_task:1055 sec("Task")1056 task_dd = st.selectbox("Preset task", PRESET_TASKS, key="task_dd")1057 task_txt = st.text_input("Or enter custom task",1058 placeholder="Describe a software engineering task…",1059 key="task_txt")1060 phase = st.slider("Curriculum phase", 1, 3, 1, key="phase_sl")1061 1062 with col_ctrl:1063 sec("Controls")1064 c1, c2 = st.columns(2)1065 reset_btn = c1.button("Reset Episode", type="primary", use_container_width=True, key="reset_btn")1066 run_btn = c2.button("Run Full Episode", use_container_width=True, key="run_btn")1067 st.markdown('<div style="height:6px"></div>', unsafe_allow_html=True)1068 1069 use_trained = st.checkbox("🤖 Use Trained Policy", value=False, key="use_trained",1070 help="Load the trained RecurrentPPO model from HF Hub")1071 trained_model = obs_mean = obs_var = None1072 clip_obs = 10.01073 if use_trained:1074 with st.spinner("Loading trained model from HF Hub…"):1075 trained_model, obs_mean, obs_var, clip_obs, model_err = _load_trained_model(HF_MODEL_REPO)1076 if model_err:1077 st.error(f"Model load failed: {model_err}")1078 else:1079 st.success("Trained policy loaded ✓")1080 1081 cat = _load_catalog()1082 act_type = st.selectbox("Action type (manual mode)",1083 ["RANDOM", "STOP", "CALL SPECIALIST", "PARALLEL SPAWN"],1084 key="act_type",1085 disabled=use_trained)1086 spec_ids = [sp["id"] for sp in cat]1087 spec_ch = st.selectbox("Target specialist", spec_ids, key="spec_ch",1088 disabled=use_trained)1089 step_btn = st.button("Execute One Step",1090 disabled=(S.env is None or S.done),1091 use_container_width=True, key="step_btn")1092 1093 status_msg = st.session_state.get("demo_status", "Click 'Reset Episode' to start.")1094 status_clr = "#34d399" if "complete" in status_msg or "started" in status_msg else "#94a3b8"1095 status_bar(status_msg, status_clr)1096 st.markdown(_exec_mode_badges(S), unsafe_allow_html=True)1097 1098 # ── Reset ──────────────────────────────────────────────1099 if reset_btn:1100 with st.spinner("Initializing environment… (first run ~30 s on CPU)"):1101 S.reset(int(phase))1102 spawn_note = (1103 f" | ⚡ Spawned: {', '.join(S.spawned_specialists)}"1104 if S.spawned_specialists else ""1105 )1106 st.session_state.demo_status = f'Episode started | Task: "{S.task[:90]}"{spawn_note}'1107 st.session_state.last_called = []1108 st.session_state.last_edges = []1109 st.session_state.last_info = {}1110 st.rerun()1111 1112 # ── Step ───────────────────────────────────────────────1113 if step_btn and S.env is not None and not S.done:1114 if use_trained and trained_model is not None and S.obs_current is not None:1115 action, S.lstm_states = _predict(1116 trained_model, S.obs_current, S.lstm_states,1117 S.episode_starts, obs_mean, obs_var, clip_obs,1118 )1119 else:1120 action = np.zeros(S.env.action_space.shape, dtype=np.float32)1121 if act_type == "STOP":1122 action[0] = 1.01123 elif act_type == "CALL SPECIALIST":1124 ids = S.registry.list_ids()1125 if spec_ch in ids:1126 idx = ids.index(spec_ch)1127 if idx < S.env.max_specialists:1128 action[1 + idx] = 1.01129 else:1130 action[1] = 1.01131 elif act_type == "PARALLEL SPAWN":1132 action[0] = 6.01133 action[1] = 1.01134 if S.env.max_specialists > 1:1135 action[2] = 1.01136 action[1 + S.env.max_specialists] = 1.01137 else:1138 action = S.env.action_space.sample()1139 1140 _, r, term, trunc, info = S.step(action)1141 done = term or trunc1142 sign = "+" if r >= 0 else ""1143 msg = f"Step {S.step_n} | reward {sign}{r:.4f} | {'DONE' if done else 'Running…'}"1144 if done:1145 msg += f" | Total: {sum(S.rewards):+.4f}"1146 st.session_state.demo_status = msg1147 # Use cumulative called_ids so graph stays populated even after STOP step1148 called = list(S.env.called_ids)1149 edges = [(e.caller_id, e.callee_id)1150 for e in S.env.delegation_graph.get_delegation_path()]1151 st.session_state.last_called = called1152 st.session_state.last_edges = edges1153 st.session_state.last_info = info1154 st.rerun()1155 1156 # ── Run Full ───────────────────────────────────────────1157 if run_btn:1158 with st.spinner("Running full episode…"):1159 S.reset(int(phase))1160 info = {}1161 for _ in range(15):1162 if S.done:1163 break1164 if use_trained and trained_model is not None and S.obs_current is not None:1165 action, S.lstm_states = _predict(1166 trained_model, S.obs_current, S.lstm_states,1167 S.episode_starts, obs_mean, obs_var, clip_obs,1168 )1169 else:1170 action = S.env.action_space.sample()1171 _, _, _, _, info = S.step(action)1172 # Use cumulative called_ids so graph stays populated even after STOP step1173 called = list(S.env.called_ids) if S.env else []1174 edges = [(e.caller_id, e.callee_id)1175 for e in S.env.delegation_graph.get_delegation_path()]1176 total = sum(S.rewards)1177 st.session_state.demo_status = (1178 f"Episode complete | {S.step_n} steps | Total reward: {total:+.4f}"1179 )1180 st.session_state.last_called = called1181 st.session_state.last_edges = edges1182 st.session_state.last_info = info1183 st.rerun()1184 1185 # ── Metric strip ──────────────────────────────────────1186 if S.env is not None:1187 mc1, mc2, mc3, mc4 = st.columns(4)1188 mc1.metric("Obs Dim", int(S.env.observation_space.shape[0]))1189 mc2.metric("Action Dim", int(S.env.action_space.shape[0]))1190 mc3.metric("Specialists", S.registry.size)1191 mc4.metric("Phase", phase)1192 1193 # ── Hero: Robot Orchestrator Widget (full width) ──────1194 sec("Orchestrator · Live Delegation View")1195 last_info = st.session_state.get("last_info", {})1196 render_orchestrator({1197 "called": st.session_state.get("last_called", []),1198 "active": (st.session_state.get("last_called", []) or [""])[-1]1199 if not S.done else "",1200 "edges": st.session_state.get("last_edges", []),