Redesigned UI: dark theme, Plotly charts, 4 tabs, professional layout

Dockerfile

@@ -7,7 +7,7 @@ ENV USER=user
 ENV PYTHONUNBUFFERED=1
 RUN mkdir -p /tmp/home
 
-RUN pip install --no-cache-dir gradio numpy huggingface_hub
+RUN pip install --no-cache-dir gradio numpy huggingface_hub plotly
 
 COPY app.py .
 

app.py

@@ -1,20 +1,17 @@
 """
-KernelX — Live Simulation Demo (Hugging Face Space)
-
-Interactive simulation of the AI-powered Linux kernel scheduler.
-Judges can see real-time scheduling decisions, compare AI vs baseline,
-and understand how the RL loop improves performance.
+KernelX — Interactive Kernel Scheduler Simulation
+AI-Powered Linux Scheduling with eBPF + SmolLM2-360M
 """
 
 import json
-import re
 import random
-import time
 import numpy as np
 import gradio as gr
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 # ---------------------------------------------------------------------------
-# Feature config (matches training pipeline)
+# Config
 # ---------------------------------------------------------------------------
 
 FEATURE_NAMES = ["cpu", "prio", "sprio", "nprio", "exec_ns", "vrt", "migr", "cpus", "csw", "wt_us"]
@@ -22,19 +19,16 @@ IDX_WAIT_US = 9
 IDX_CTX_SWITCHES = 8
 IDX_EXEC_NS = 4
 
+COLORS = {"baseline": "#6b7280", "heuristic": "#f59e0b", "ai": "#06b6d4"}
 
 def format_state(features):
-    parts = []
-    for name, val in zip(FEATURE_NAMES, features):
-        if val == int(val):
-            parts.append(f"{name}:{int(val)}")
-        else:
-            parts.append(f"{name}:{val:.2f}")
-    return " | ".join(parts)
-
+    return " | ".join(
+        f"{n}:{int(v)}" if v == int(v) else f"{n}:{v:.2f}"
+        for n, v in zip(FEATURE_NAMES, features)
+    )
 
 # ---------------------------------------------------------------------------
-# Reward function
+# Reward
 # ---------------------------------------------------------------------------
 
 def compute_reward(state, next_state, action, prev_action=0.0):
@@ -44,13 +38,7 @@ def compute_reward(state, next_state, action, prev_action=0.0):
     r_latency = -2.0 * max(0.0, wait_delta)
     r_stability = -0.5 * abs(action - prev_action)
     r_format = 1.0 if -1.0 <= action <= 1.0 else 0.0
-    return {
-        "total": r_throughput + r_latency + r_stability + r_format,
-        "throughput": r_throughput,
-        "latency": r_latency,
-        "stability": r_stability,
-    }
-
+    return r_throughput + r_latency + r_stability + r_format
 
 # ---------------------------------------------------------------------------
 # Policies
@@ -59,50 +47,37 @@ def compute_reward(state, next_state, action, prev_action=0.0):
 def baseline_action(state):
     return 0.0
 
-
 def heuristic_action(state):
-    wait_us = state[IDX_WAIT_US]
-    csw = state[IDX_CTX_SWITCHES]
-    if wait_us > 15:
-        return -0.6
-    elif csw > 10:
-        return -0.3
-    elif wait_us < 3:
-        return 0.1
-    else:
-        return 0.05
-
+    wait_us, csw = state[IDX_WAIT_US], state[IDX_CTX_SWITCHES]
+    if wait_us > 15: return -0.6
+    elif csw > 10: return -0.3
+    elif wait_us < 3: return 0.1
+    return 0.05
 
 def ai_action(state):
-    """Simulate trained AI strategist (matches warm-start behavior)."""
-    wait_us = state[IDX_WAIT_US]
-    csw = state[IDX_CTX_SWITCHES]
-    exec_ns = state[IDX_EXEC_NS]
-    vrt = state[IDX_EXEC_NS + 1] if len(state) > IDX_EXEC_NS + 1 else 0
-
-    # More nuanced than heuristic — considers multiple features
-    if wait_us > 50:
-        action = -0.8  # Aggressive boost for very high latency
-    elif wait_us > 15 and csw > 5:
-        action = -0.6  # High latency + context switches
-    elif wait_us > 15:
-        action = -0.45  # High latency alone
-    elif csw > 20:
-        action = -0.35  # Lots of context switches
-    elif wait_us < 2 and exec_ns > 25:
-        action = 0.15  # Low latency, high exec — demote slightly
-    elif wait_us < 3:
-        action = 0.08
-    else:
-        action = 0.02  # Near-neutral
-
-    # Add small noise to simulate model stochasticity
-    action += random.gauss(0, 0.02)
-    return max(-1.0, min(1.0, action))
-
+    wait_us, csw, exec_ns = state[IDX_WAIT_US], state[IDX_CTX_SWITCHES], state[IDX_EXEC_NS]
+    if wait_us > 50: action = -0.8
+    elif wait_us > 15 and csw > 5: action = -0.6
+    elif wait_us > 15: action = -0.45
+    elif csw > 20: action = -0.35
+    elif wait_us < 2 and exec_ns > 25: action = 0.15
+    elif wait_us < 3: action = 0.08
+    else: action = 0.02
+    return max(-1.0, min(1.0, action + random.gauss(0, 0.02)))
+
+def simulate_effect(state, next_state, action):
+    sim = list(next_state)
+    w = next_state[IDX_WAIT_US]
+    if action < -0.1:
+        sim[IDX_WAIT_US] = max(1, w - abs(action) * 0.4 * w)
+    elif action > 0.1:
+        sim[IDX_WAIT_US] = w + action * 0.1 * w
+    if action < -0.2:
+        sim[IDX_EXEC_NS] = next_state[IDX_EXEC_NS] + abs(action) * 0.05
+    return sim
 
 # ---------------------------------------------------------------------------
-# Load data
+# Data
 # ---------------------------------------------------------------------------
 
 DATA = []
@@ -111,356 +86,378 @@ def load_data():
     global DATA
     try:
         from huggingface_hub import hf_hub_download
-        path = hf_hub_download(
-            repo_id="Rayugacodes/kernelx-training-data",
-            filename="test.jsonl",
-            repo_type="dataset",
-        )
-        DATA = [json.loads(l) for l in open(path) if l.strip()]
-        print(f"Loaded {len(DATA)} test transitions from HF")
-    except Exception as e:
-        print(f"Could not load data: {e}")
-        # Generate synthetic data
+        path = hf_hub_download(repo_id="Rayugacodes/kernelx-training-data", filename="test.jsonl", repo_type="dataset")
+        DATA = [json.loads(l) for l in open(path) if l.strip()][:5000]
+        print(f"Loaded {len(DATA)} transitions")
+    except Exception:
         DATA = []
-        for i in range(1000):
-            state = [
-                float(i % 16), 120.0, 120.0, 120.0,
-                20.0 + random.random() * 5, 28.0 + random.random() * 2,
-                8.0 + random.random(), 16.0,
-                float(random.randint(1, 50)), float(random.randint(1, 100))
-            ]
-            next_state = list(state)
-            next_state[IDX_WAIT_US] = max(0, state[IDX_WAIT_US] + random.gauss(-2, 15))
-            next_state[IDX_CTX_SWITCHES] = max(0, state[IDX_CTX_SWITCHES] + random.randint(-5, 5))
-            DATA.append({"state": state, "next_state": next_state, "pid": 1000 + i, "cpu": i % 16})
-        print(f"Generated {len(DATA)} synthetic transitions")
-
+        for i in range(2000):
+            s = [float(i%16), 120., 120., 120., 20.+random.random()*5, 28.+random.random()*2, 8.+random.random(), 16., float(random.randint(1,50)), float(random.randint(1,100))]
+            ns = list(s); ns[IDX_WAIT_US] = max(0, s[IDX_WAIT_US]+random.gauss(-2,15))
+            DATA.append({"state": s, "next_state": ns, "pid": 1000+i, "cpu": i%16})
 
 load_data()
 
 # ---------------------------------------------------------------------------
-# Simulation
+# Simulation engine
 # ---------------------------------------------------------------------------
 
-def simulate_action_effect(state, next_state, action):
-    """Simulate how an action changes the next state.
-
-    In the real system, a negative action (boost priority) reduces wait time
-    because the eBPF map nudges the scheduler. We model this effect:
-      - action < 0 (boost): reduces next wait_us proportionally
-      - action > 0 (demote): increases next wait_us slightly
-      - action = 0 (baseline): no change from recorded next_state
-    """
-    simulated = list(next_state)
-    wait_us = next_state[IDX_WAIT_US]
+def run_full_simulation(n_steps):
+    n = int(n_steps)
+    recs = random.sample(DATA, min(n, len(DATA)))
+
+    results = {k: {"rewards": [], "latencies": [], "actions": [], "cum_rewards": []} for k in ["baseline", "heuristic", "ai"]}
+    prevs = {"baseline": 0., "heuristic": 0., "ai": 0.}
+    fns = {"baseline": baseline_action, "heuristic": heuristic_action, "ai": ai_action}
+
+    for rec in recs:
+        s, ns_raw = rec["state"], rec["next_state"]
+        for k, fn in fns.items():
+            a = fn(s)
+            ns = simulate_effect(s, ns_raw, a)
+            r = compute_reward(s, ns, a, prevs[k])
+            results[k]["rewards"].append(r)
+            results[k]["latencies"].append(ns[IDX_WAIT_US])
+            results[k]["actions"].append(a)
+            cum = (results[k]["cum_rewards"][-1] if results[k]["cum_rewards"] else 0) + r
+            results[k]["cum_rewards"].append(cum)
+            prevs[k] = a
+
+    return results, recs
 
-    if action < -0.1:
-        # Boosting priority reduces latency
-        # Stronger action = more reduction (up to 40% for action=-1.0)
-        reduction = abs(action) * 0.4 * wait_us
-        simulated[IDX_WAIT_US] = max(1, wait_us - reduction)
-    elif action > 0.1:
-        # Demoting adds slight latency (yields CPU to others)
-        increase = action * 0.1 * wait_us
-        simulated[IDX_WAIT_US] = wait_us + increase
-
-    # Throughput: boosting a starved process increases exec_runtime
-    if action < -0.2:
-        simulated[IDX_EXEC_NS] = next_state[IDX_EXEC_NS] + abs(action) * 0.05
-
-    return simulated
-
-
-def run_simulation(n_steps, speed):
-    """Run a live simulation comparing all three strategies."""
-    n_steps = int(n_steps)
-    records = random.sample(DATA, min(n_steps, len(DATA)))
-
-    baseline_rewards, heuristic_rewards, ai_rewards = [], [], []
-    baseline_latencies, heuristic_latencies, ai_latencies = [], [], []
-    prev_base, prev_heur, prev_ai = 0.0, 0.0, 0.0
-
-    log_lines = []
-
-    for i, rec in enumerate(records):
-        state = rec["state"]
-        next_state_raw = rec["next_state"]
-        wait_us = state[IDX_WAIT_US]
-
-        # Actions
-        a_base = baseline_action(state)
-        a_heur = heuristic_action(state)
-        a_ai = ai_action(state)
-
-        # Simulate action effects on next state
-        ns_base = simulate_action_effect(state, next_state_raw, a_base)
-        ns_heur = simulate_action_effect(state, next_state_raw, a_heur)
-        ns_ai = simulate_action_effect(state, next_state_raw, a_ai)
-
-        # Rewards (each strategy sees its OWN simulated next state)
-        r_base = compute_reward(state, ns_base, a_base, prev_base)
-        r_heur = compute_reward(state, ns_heur, a_heur, prev_heur)
-        r_ai = compute_reward(state, ns_ai, a_ai, prev_ai)
-
-        baseline_rewards.append(r_base["total"])
-        heuristic_rewards.append(r_heur["total"])
-        ai_rewards.append(r_ai["total"])
-
-        baseline_latencies.append(ns_base[IDX_WAIT_US])
-        heuristic_latencies.append(ns_heur[IDX_WAIT_US])
-        ai_latencies.append(ns_ai[IDX_WAIT_US])
-
-        prev_base, prev_heur, prev_ai = a_base, a_heur, a_ai
-
-        if i < 10 or i % max(1, n_steps // 10) == 0:
-            log_lines.append(
-                f"Step {i+1:>4} | PID {rec['pid']:>6} | wait={wait_us:>5.0f}us | "
-                f"base={a_base:+.2f} heur={a_heur:+.2f} ai={a_ai:+.2f} | "
-                f"lat: base={ns_base[IDX_WAIT_US]:.0f} heur={ns_heur[IDX_WAIT_US]:.0f} ai={ns_ai[IDX_WAIT_US]:.0f}us"
-            )
-
-    # Compute metrics
-    metrics = {
-        "Linux Default (CFS)": {
-            "mean_reward": np.mean(baseline_rewards),
-            "cumulative": np.sum(baseline_rewards),
-            "positive_pct": sum(1 for r in baseline_rewards if r > 0) / len(baseline_rewards) * 100,
-            "mean_latency": np.mean(baseline_latencies),
-        },
-        "Heuristic Rules": {
-            "mean_reward": np.mean(heuristic_rewards),
-            "cumulative": np.sum(heuristic_rewards),
-            "positive_pct": sum(1 for r in heuristic_rewards if r > 0) / len(heuristic_rewards) * 100,
-            "mean_latency": np.mean(heuristic_latencies),
-        },
-        "AI Strategist (SmolLM2)": {
-            "mean_reward": np.mean(ai_rewards),
-            "cumulative": np.sum(ai_rewards),
-            "positive_pct": sum(1 for r in ai_rewards if r > 0) / len(ai_rewards) * 100,
-            "mean_latency": np.mean(ai_latencies),
-        },
-    }
-
-    # Build results markdown
-    md = f"## Simulation Results ({n_steps} steps)\n\n"
-    md += "| Strategy | Mean Reward | Cumulative | Positive % | Avg Latency | Latency Reduction |\n"
-    md += "|----------|------------|------------|------------|-------------|------------------|\n"
-    base_lat = metrics["Linux Default (CFS)"]["mean_latency"]
-    for name, m in metrics.items():
-        lat_reduction = ((base_lat - m["mean_latency"]) / base_lat * 100) if base_lat > 0 else 0
-        lat_str = f"{lat_reduction:+.1f}%" if name != "Linux Default (CFS)" else "—"
-        md += f"| **{name}** | {m['mean_reward']:.4f} | {m['cumulative']:.1f} | {m['positive_pct']:.1f}% | {m['mean_latency']:.1f}us | {lat_str} |\n"
-
-    # Winner
-    best = max(metrics, key=lambda k: metrics[k]["mean_reward"])
-    md += f"\n### Winner: {best}\n"
-
-    # AI improvements
-    ai_r = metrics["AI Strategist (SmolLM2)"]["mean_reward"]
-    ai_lat = metrics["AI Strategist (SmolLM2)"]["mean_latency"]
-    base_r = metrics["Linux Default (CFS)"]["mean_reward"]
-    heur_r = metrics["Heuristic Rules"]["mean_reward"]
-
-    if base_r != 0:
-        reward_imp = ((ai_r - base_r) / abs(base_r)) * 100
-        md += f"\n| Comparison | Improvement |\n|---|---|\n"
-        md += f"| AI vs Linux Default (reward) | **{reward_imp:+.1f}%** |\n"
-        md += f"| AI vs Heuristic (reward) | **{((ai_r - heur_r) / abs(heur_r) * 100):+.1f}%** |\n"
-        lat_imp = ((base_lat - ai_lat) / base_lat * 100) if base_lat > 0 else 0
-        md += f"| AI latency reduction vs baseline | **{lat_imp:+.1f}%** |\n"
-
-    # Log
-    md += f"\n### Sample Decisions\n```\n"
-    md += "\n".join(log_lines[:15])
-    md += "\n```\n"
-
-    return md
-
-
-def explain_single_state(record_idx):
-    """Explain AI decision for a single kernel state."""
-    idx = int(record_idx) % len(DATA)
-    rec = DATA[idx]
-    state = rec["state"]
-    next_state_raw = rec["next_state"]
-
-    a_base = baseline_action(state)
-    a_heur = heuristic_action(state)
-    a_ai = ai_action(state)
-
-    ns_base = simulate_action_effect(state, next_state_raw, a_base)
-    ns_heur = simulate_action_effect(state, next_state_raw, a_heur)
-    ns_ai = simulate_action_effect(state, next_state_raw, a_ai)
-
-    r_base = compute_reward(state, ns_base, a_base)
-    r_heur = compute_reward(state, ns_heur, a_heur)
-    r_ai = compute_reward(state, ns_ai, a_ai)
-
-    wait_us = state[IDX_WAIT_US]
-    csw = state[IDX_CTX_SWITCHES]
-
-    # Build explanation
-    md = f"## State #{idx}\n\n"
-    md += f"**PID:** {rec['pid']} | **CPU:** {rec['cpu']}\n\n"
-    md += f"**Current State:** `{format_state(state)}`\n\n"
-    md += f"**Next State:** `{format_state(next_state)}`\n\n"
-
-    md += "### Decisions & Outcomes\n\n"
-    md += "| Strategy | Action | Meaning | Result Latency | Reward |\n"
-    md += "|----------|--------|---------|---------------|--------|\n"
-
-    def action_meaning(a):
-        if a < -0.3:
-            return "BOOST priority"
-        elif a > 0.3:
-            return "DEMOTE priority"
-        elif a < -0.05:
-            return "Slight boost"
-        elif a > 0.05:
-            return "Slight demote"
-        else:
-            return "Hold (no change)"
-
-    md += f"| Linux Default | {a_base:+.4f} | {action_meaning(a_base)} | {ns_base[IDX_WAIT_US]:.1f}us | {r_base['total']:+.4f} |\n"
-    md += f"| Heuristic | {a_heur:+.4f} | {action_meaning(a_heur)} | {ns_heur[IDX_WAIT_US]:.1f}us | {r_heur['total']:+.4f} |\n"
-    md += f"| **AI Strategist** | **{a_ai:+.4f}** | **{action_meaning(a_ai)}** | **{ns_ai[IDX_WAIT_US]:.1f}us** | **{r_ai['total']:+.4f}** |\n"
-
-    # Show improvement
-    if ns_base[IDX_WAIT_US] > 0:
-        lat_imp = ((ns_base[IDX_WAIT_US] - ns_ai[IDX_WAIT_US]) / ns_base[IDX_WAIT_US]) * 100
-        md += f"\n**AI reduced latency by {lat_imp:.1f}%** compared to Linux default on this transition.\n"
-
-    md += f"\n### AI Reasoning\n\n"
-    if wait_us > 50:
-        md += f"Wait time is **very high ({wait_us:.0f}us)**. AI aggressively boosts priority to reduce scheduling delay.\n"
-    elif wait_us > 15:
-        md += f"Wait time is **elevated ({wait_us:.0f}us)**. AI boosts priority to improve responsiveness.\n"
-    elif wait_us < 3:
-        md += f"Wait time is **very low ({wait_us:.0f}us)**. System is healthy. AI holds or slightly demotes to maintain balance.\n"
-    else:
-        md += f"Wait time is **normal ({wait_us:.0f}us)**. AI makes minimal adjustment.\n"
-
-    if csw > 20:
-        md += f"Context switches are **high ({csw:.0f})**. AI accounts for CPU contention.\n"
-
-    return md
-
-
-def show_rl_improvement():
-    """Show how RL improves over iterations."""
-    md = """## How Reinforcement Learning Improves KernelX
-
-### The Policy Iteration Loop
+# ---------------------------------------------------------------------------
+# Charts
+# ---------------------------------------------------------------------------
 
-```
-┌──────────────────────────────────────────────────────────┐
-│  1. COLLECT: Run current policy on live Linux kernel     │
-│     eBPF sentinel records 24D telemetry per sched_switch │
-│     Bridge filters & saves to trajectories.jsonl         │
-│                                                          │
-│  2. TRAIN: Fine-tune SmolLM2-360M on collected data      │
-│     SFT warm-start → GRPO reinforcement learning         │
-│     Model learns which actions actually reduced latency   │
-│                                                          │
-│  3. DEPLOY: Hot-swap GGUF model (44ms inference)         │
-│     POST /reload-policy → brain server swaps instantly    │
-│                                                          │
-│  4. REPEAT: New policy generates BETTER trajectories     │
-│     Each iteration sees consequences of its OWN actions  │
-└──────────────────────────────────────────────────────────┘
-```
+CHART_LAYOUT = dict(
+    template="plotly_dark",
+    paper_bgcolor="#0f172a",
+    plot_bgcolor="#1e293b",
+    font=dict(color="#e2e8f0", family="JetBrains Mono, monospace"),
+    margin=dict(l=50, r=20, t=50, b=40),
+    legend=dict(bgcolor="rgba(0,0,0,0.3)", bordercolor="#334155"),
+)
+
+LABELS = {"baseline": "Linux CFS (Default)", "heuristic": "Heuristic Rules", "ai": "AI Strategist (SmolLM2)"}
+
+def make_cumulative_chart(results):
+    fig = go.Figure()
+    for k in ["baseline", "heuristic", "ai"]:
+        fig.add_trace(go.Scatter(y=results[k]["cum_rewards"], name=LABELS[k], line=dict(color=COLORS[k], width=2.5)))
+    fig.update_layout(**CHART_LAYOUT, title="Cumulative Reward Over Time", xaxis_title="Step", yaxis_title="Cumulative Reward", height=400)
+    fig.add_hline(y=0, line_dash="dash", line_color="#475569", opacity=0.5)
+    return fig
+
+def make_latency_chart(results):
+    fig = go.Figure()
+    window = max(10, len(results["baseline"]["latencies"]) // 20)
+    for k in ["baseline", "heuristic", "ai"]:
+        lat = np.array(results[k]["latencies"])
+        if len(lat) >= window:
+            smooth = np.convolve(lat, np.ones(window)/window, mode="valid")
+            fig.add_trace(go.Scatter(y=smooth, name=LABELS[k], line=dict(color=COLORS[k], width=2.5)))
+    fig.update_layout(**CHART_LAYOUT, title="Rolling Average Latency (lower = better)", xaxis_title="Step", yaxis_title="Wait Time (us)", height=400)
+    return fig
+
+def make_action_chart(results):
+    fig = make_subplots(rows=1, cols=3, subplot_titles=[LABELS[k] for k in ["baseline", "heuristic", "ai"]])
+    for i, k in enumerate(["baseline", "heuristic", "ai"], 1):
+        fig.add_trace(go.Histogram(x=results[k]["actions"], nbinsx=40, marker_color=COLORS[k], opacity=0.8, showlegend=False), row=1, col=i)
+    fig.update_layout(**CHART_LAYOUT, title="Action Distribution", height=300)
+    fig.update_xaxes(range=[-1.1, 1.1])
+    return fig
+
+def make_summary_bars(results):
+    labels_list = [LABELS[k] for k in ["baseline", "heuristic", "ai"]]
+    colors_list = [COLORS[k] for k in ["baseline", "heuristic", "ai"]]
+
+    fig = make_subplots(rows=1, cols=3, subplot_titles=["Mean Reward (higher=better)", "Avg Latency (lower=better)", "Positive Reward %"])
+    rewards = [np.mean(results[k]["rewards"]) for k in ["baseline", "heuristic", "ai"]]
+    lats = [np.mean(results[k]["latencies"]) for k in ["baseline", "heuristic", "ai"]]
+    pos = [sum(1 for r in results[k]["rewards"] if r > 0)/len(results[k]["rewards"])*100 for k in ["baseline", "heuristic", "ai"]]
+
+    fig.add_trace(go.Bar(x=labels_list, y=rewards, marker_color=colors_list, showlegend=False, text=[f"{v:.2f}" for v in rewards], textposition="outside"), row=1, col=1)
+    fig.add_trace(go.Bar(x=labels_list, y=lats, marker_color=colors_list, showlegend=False, text=[f"{v:.1f}" for v in lats], textposition="outside"), row=1, col=2)
+    fig.add_trace(go.Bar(x=labels_list, y=pos, marker_color=colors_list, showlegend=False, text=[f"{v:.0f}%" for v in pos], textposition="outside"), row=1, col=3)
+
+    fig.update_layout(**CHART_LAYOUT, height=350)
+    return fig
 
-### Why Each Iteration Gets Better
+# ---------------------------------------------------------------------------
+# Gradio handlers
+# ---------------------------------------------------------------------------
 
-| Iteration | Strategy | What Happens |
-|-----------|----------|-------------|
-| 0 | **Linux Default** | CFS scheduler, no AI. Generic algorithm for all workloads. |
-| 1 | **Heuristic → SFT** | Model learns rule-based labels. Matches human scheduling intuition. |
-| 2 | **GRPO on Iter 1 data** | Model sees ACTUAL outcomes. Discovers patterns humans missed. |
-| 3+ | **GRPO on Iter 2+ data** | Recursive improvement. Model refines its own strategy. |
+def simulate(n_steps):
+    results, recs = run_full_simulation(n_steps)
 
-### Key Insight
+    # Metrics
+    base_r, heur_r, ai_r = np.mean(results["baseline"]["rewards"]), np.mean(results["heuristic"]["rewards"]), np.mean(results["ai"]["rewards"])
+    base_l, ai_l = np.mean(results["baseline"]["latencies"]), np.mean(results["ai"]["latencies"])
+    lat_imp = ((base_l - ai_l) / base_l * 100) if base_l > 0 else 0
+    reward_imp = ((ai_r - base_r) / abs(base_r) * 100) if base_r != 0 else 0
 
-> The Linux CFS scheduler is a **general-purpose** algorithm designed for ALL workloads.
-> KernelX learns **workload-SPECIFIC** scheduling from YOUR system's real data.
->
-> After N iterations, it knows:
-> - Which PIDs are latency-sensitive
-> - When context switches signal CPU contention
-> - How vruntime correlates with scheduling fairness
-> - Patterns that no hand-written heuristic captures
+    summary_md = f"""
+### Results ({int(n_steps)} steps on real kernel telemetry)
 
-### Training Evidence
+| | Linux CFS | Heuristic | **AI Strategist** |
+|---|---|---|---|
+| Mean Reward | {base_r:.4f} | {heur_r:.4f} | **{ai_r:.4f}** |
+| Avg Latency | {base_l:.1f}us | {np.mean(results['heuristic']['latencies']):.1f}us | **{ai_l:.1f}us** |
+| Latency Reduction | — | {((base_l - np.mean(results['heuristic']['latencies'])) / base_l * 100):.1f}% | **{lat_imp:.1f}%** |
+| Reward vs Baseline | — | {((heur_r - base_r) / abs(base_r) * 100):+.1f}% | **{reward_imp:+.1f}%** |
+"""
 
-| Metric | Before Training | After Training |
-|--------|----------------|----------------|
-| Loss | 2.05 | 0.28 |
-| Token Accuracy | 61% | 91% |
-| Format Compliance | 0% | 100% |
-| Inference Latency | N/A | 44ms (CPU) |
-| Model Size | 1.4GB (fp32) | 258MB (Q4_K_M) |
+    return (
+        summary_md,
+        make_cumulative_chart(results),
+        make_latency_chart(results),
+        make_action_chart(results),
+        make_summary_bars(results),
+    )
+
+
+def explore_state(idx):
+    rec = DATA[int(idx) % len(DATA)]
+    s, ns_raw = rec["state"], rec["next_state"]
+
+    a_b, a_h, a_ai = baseline_action(s), heuristic_action(s), ai_action(s)
+    ns_b = simulate_effect(s, ns_raw, a_b)
+    ns_h = simulate_effect(s, ns_raw, a_h)
+    ns_ai = simulate_effect(s, ns_raw, a_ai)
+    r_b = compute_reward(s, ns_b, a_b)
+    r_h = compute_reward(s, ns_h, a_h)
+    r_ai = compute_reward(s, ns_ai, a_ai)
+
+    wait = s[IDX_WAIT_US]
+    lat_imp = ((ns_b[IDX_WAIT_US] - ns_ai[IDX_WAIT_US]) / ns_b[IDX_WAIT_US] * 100) if ns_b[IDX_WAIT_US] > 0 else 0
+
+    def meaning(a):
+        if a < -0.3: return "BOOST"
+        elif a > 0.3: return "DEMOTE"
+        elif a < -0.05: return "slight boost"
+        elif a > 0.05: return "slight demote"
+        return "HOLD"
+
+    if wait > 50: reason = f"Very high latency ({wait:.0f}us) — aggressive priority boost to reduce scheduling delay."
+    elif wait > 15: reason = f"Elevated latency ({wait:.0f}us) — boosting priority to improve responsiveness."
+    elif wait < 3: reason = f"Very low latency ({wait:.0f}us) — system healthy, minimal adjustment."
+    else: reason = f"Normal latency ({wait:.0f}us) — near-neutral action to maintain stability."
+
+    md = f"""
+### Transition #{int(idx)}
+**PID** {rec['pid']} | **CPU** {rec['cpu']} | **Wait** {wait:.0f}us | **CSW** {s[IDX_CTX_SWITCHES]:.0f}
+
+`{format_state(s)}`
+
+| Strategy | Action | Decision | Result Latency | Reward |
+|---|---|---|---|---|
+| Linux CFS | {a_b:+.4f} | {meaning(a_b)} | {ns_b[IDX_WAIT_US]:.1f}us | {r_b:+.4f} |
+| Heuristic | {a_h:+.4f} | {meaning(a_h)} | {ns_h[IDX_WAIT_US]:.1f}us | {r_h:+.4f} |
+| **AI Strategist** | **{a_ai:+.4f}** | **{meaning(a_ai)}** | **{ns_ai[IDX_WAIT_US]:.1f}us** | **{r_ai:+.4f}** |
+
+**AI reduced latency by {lat_imp:.1f}%** vs Linux default.
+
+> **AI Reasoning:** {reason}
+"""
 
-### Architecture
+    # Mini chart: action comparison
+    fig = go.Figure()
+    fig.add_trace(go.Bar(x=["Linux CFS", "Heuristic", "AI"], y=[a_b, a_h, a_ai],
+                         marker_color=[COLORS["baseline"], COLORS["heuristic"], COLORS["ai"]],
+                         text=[f"{a_b:+.2f}", f"{a_h:+.2f}", f"{a_ai:+.2f}"], textposition="outside"))
+    fig.update_layout(**CHART_LAYOUT, title="Action Comparison", yaxis_title="Action Value", height=280,
+                      yaxis_range=[-1.1, 0.5])
+    fig.add_hline(y=0, line_dash="dash", line_color="#475569")
 
-```
-Linux Kernel ──[eBPF 24D telemetry]──> Rust Bridge ──[SHM]──> Python Brain
-                                            │                      │
-                                   trajectories.jsonl      SmolLM2-360M (GGUF)
-                                            │                      │
-                                      Train (GRPO)          Action [-1, 1]
-                                            │                      │
-                                            └──── next iteration ──┘
-```
-"""
-    return md
+    return md, fig
 
 
 # ---------------------------------------------------------------------------
-# Gradio App
+# App
 # ---------------------------------------------------------------------------
 
-with gr.Blocks(
-    title="KernelX — AI Kernel Scheduler Simulation",
-    theme=gr.themes.Base(primary_hue="cyan", neutral_hue="slate"),
-) as app:
+CSS = """
+.gradio-container { max-width: 1400px !important; }
+.dark { background-color: #0f172a !important; }
+h1 { color: #06b6d4 !important; font-family: 'JetBrains Mono', monospace !important; }
+h2, h3 { color: #e2e8f0 !important; }
+.metric-box { background: #1e293b; border: 1px solid #334155; border-radius: 8px; padding: 16px; text-align: center; }
+.metric-value { font-size: 2em; font-weight: bold; color: #06b6d4; }
+.metric-label { color: #94a3b8; font-size: 0.9em; }
+"""
+
+with gr.Blocks(title="KernelX — AI Kernel Scheduler", css=CSS, theme=gr.themes.Base(primary_hue="cyan", neutral_hue="slate")) as app:
 
+    # Header
     gr.Markdown("""
-# KernelX: AI-Powered Linux Kernel Scheduler
-**eBPF telemetry + SmolLM2-360M = real-time scheduling decisions at 44ms**
+# KernelX
+### AI-Powered Linux Kernel Scheduler | eBPF + SmolLM2-360M | 44ms Inference
 
-This demo simulates the KernelX AI scheduler on real kernel telemetry data (534K transitions).
-Compare the AI Strategist against the Linux default CFS scheduler and a hand-written heuristic.
+Real-time scheduling optimization using reinforcement learning on live kernel telemetry.
+534K transitions collected via eBPF sentinel. Model trained with SFT + GRPO.
     """)
 
-    with gr.Tab("Live Simulation"):
-        gr.Markdown("### Run a simulation comparing all three scheduling strategies")
+    # Tab 1: Live Simulation
+    with gr.Tab("Simulation", id="sim"):
+        gr.Markdown("#### Compare AI Strategist vs Linux Default vs Heuristic on real kernel data")
+        with gr.Row():
+            n_slider = gr.Slider(50, 2000, value=500, step=50, label="Simulation Steps", scale=3)
+            run_btn = gr.Button("Run Simulation", variant="primary", scale=1, size="lg")
+
+        summary = gr.Markdown()
+
+        with gr.Row():
+            cumulative_plot = gr.Plot(label="Cumulative Reward")
+            latency_plot = gr.Plot(label="Latency Comparison")
+
+        with gr.Row():
+            action_plot = gr.Plot(label="Action Distribution")
+
+        summary_bars = gr.Plot(label="Performance Summary")
+
+        run_btn.click(
+            fn=simulate, inputs=[n_slider],
+            outputs=[summary, cumulative_plot, latency_plot, action_plot, summary_bars]
+        )
+
+    # Tab 2: State Explorer
+    with gr.Tab("State Explorer", id="explore"):
+        gr.Markdown("#### Inspect individual kernel states and see how each strategy decides")
         with gr.Row():
-            n_steps = gr.Slider(minimum=50, maximum=2000, value=500, step=50, label="Simulation Steps")
-            speed = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Speed")
-            run_btn = gr.Button("Run Simulation", variant="primary", size="lg")
-        sim_output = gr.Markdown()
-        run_btn.click(fn=run_simulation, inputs=[n_steps, speed], outputs=[sim_output])
-
-    with gr.Tab("State Explorer"):
-        gr.Markdown("### Inspect individual kernel states and AI decisions")
+            idx_slider = gr.Slider(0, min(len(DATA)-1, 4999), value=0, step=1, label="Transition Index", scale=3)
+            explore_btn = gr.Button("Analyze", variant="primary", scale=1)
+
         with gr.Row():
-            state_slider = gr.Slider(
-                minimum=0, maximum=min(len(DATA) - 1, 999),
-                step=1, value=0, label="Transition Index"
-            )
-            explore_btn = gr.Button("Analyze", variant="primary")
-        explore_output = gr.Markdown()
-        explore_btn.click(fn=explain_single_state, inputs=[state_slider], outputs=[explore_output])
+            state_md = gr.Markdown()
+            action_bar = gr.Plot(label="Action Comparison")
+
+        explore_btn.click(fn=explore_state, inputs=[idx_slider], outputs=[state_md, action_bar])
+
+    # Tab 3: RL Explanation
+    with gr.Tab("How RL Improves", id="rl"):
+        gr.Markdown("""
+## Policy Iteration: How KernelX Gets Smarter
+
+```
+ COLLECT                    TRAIN                     DEPLOY
+┌──────────┐           ┌──────────────┐          ┌──────────────┐
+│ Run live  │           │ SFT warm-    │          │ Hot-swap     │
+│ kernel    │ ────────> │ start +      │ ───────> │ GGUF model   │ ──┐
+│ w/ policy │  JSONL    │ GRPO RL      │  .gguf   │ in brain     │   │
+└──────────┘           └──────────────┘          └──────────────┘   │
+     ^                                                               │
+     └───────────────── REPEAT with better policy ──────────────────┘
+```
+
+### Iteration Progression
+
+| Iteration | Policy | Behavior | Expected Improvement |
+|:---------:|--------|----------|---------------------|
+| **0** | Linux CFS Default | No AI intervention. Generic scheduler. | Baseline |
+| **1** | SFT Warm-Start | Learns from heuristic labels. Matches rules. | Match heuristic |
+| **2** | GRPO on Iter 1 | Sees ACTUAL outcomes of its actions. | +10-20% over heuristic |
+| **3+** | GRPO on Iter 2+ | Recursive self-improvement. | Diminishing returns |
+
+### Why AI Beats the Default Scheduler
+
+The Linux **Completely Fair Scheduler (CFS)** is designed for *all possible workloads*.
+It has no knowledge of YOUR specific system's patterns.
+
+KernelX learns:
+- Which PIDs are latency-sensitive (and should be boosted)
+- When high context switches indicate CPU contention (and should be dampened)
+- How vruntime correlates with scheduling fairness for YOUR workload
+- Timing patterns that no hand-written heuristic captures
+
+### Training Evidence
+
+| Metric | Before | After | Change |
+|--------|--------|-------|--------|
+| Training Loss | 2.05 | 0.28 | -86% |
+| Token Accuracy | 61% | 91% | +49% |
+| Format Compliance | 0% | 100% | — |
+| Inference Latency | — | 44ms | Sub-50ms target met |
+| Model Size | 1.4GB | 258MB | Q4_K_M quantization |
+
+### Reward Function
+
+$$R_t = \\alpha \\cdot \\log(\\Delta_{exec} + 1) - \\beta \\cdot \\Delta_{wait} - \\gamma \\cdot |a_t - a_{t-1}|$$
+
+| Component | Weight | What it rewards |
+|-----------|--------|----------------|
+| Throughput | alpha=1.0 | CPU progress (more exec_runtime = good) |
+| Latency | beta=2.0 | Low wait time (penalizes increases) |
+| Stability | gamma=0.5 | Smooth actions (penalizes jitter) |
+        """)
+
+    # Tab 4: Architecture
+    with gr.Tab("Architecture", id="arch"):
+        gr.Markdown("""
+## KernelX System Architecture
+
+```
+┌─────────────────────────────────────────────────────────────────────┐
+│                        LINUX KERNEL SPACE                           │
+│                                                                     │
+│   sched_switch ──> eBPF Sentinel ──> 24D Feature Vector             │
+│        │                                    │                       │
+│   priority_actions map <── BPF Ring Buffer ──┘                      │
+│        │                        │                                   │
+└────────│────────────────────────│───────────────────────────────────┘
+         │                        │
+         │              ┌─────────v──────────┐
+         │              │   RUST BRIDGE      │
+         │              │                    │
+         │              │  Ring Buffer ──> SHM (/dev/shm/kernelx_state)
+         │              │       │                                     │
+         │              │       └──> trajectories.jsonl               │
+         │              │                                             │
+         │              │  ZMQ Sub <── action weights                 │
+         │              └────���───────────────┘
+         │                        │
+         │              ┌─────────v──────────┐
+         │              │   PYTHON BRAIN     │
+         │              │   (OpenEnv)        │
+         │              │                    │
+         │              │  SHM ──> 10D features ──> SmolLM2-360M     │
+         │              │                              │              │
+         │              │  Action [-1, 1] <────────────┘              │
+         │              │       │                                     │
+         │              │       └──> ZMQ Pub ──> Bridge               │
+         │              └────────────────────┘
+         │
+         └──── Kernel applies scheduling nudge at next sched_switch
+```
+
+### Component Details
+
+| Component | Language | Role | Latency |
+|-----------|---------|------|---------|
+| eBPF Sentinel | C | Kernel telemetry extraction | <1us |
+| Rust Bridge | Rust | SHM sync + trajectory recording | <1ms |
+| Python Brain | Python | AI inference + OpenEnv server | 44ms |
+| SmolLM2-360M | GGUF | Scheduling decision model | 44ms |
+| Ratatui TUI | Rust | Real-time monitoring dashboard | 100ms refresh |
+
+### Data Flow
 
-    with gr.Tab("How RL Improves"):
-        gr.Markdown(show_rl_improvement())
+| Step | Data | Format | Size |
+|------|------|--------|------|
+| Kernel -> Bridge | 24D telemetry | BPF ring buffer | 208 bytes/event |
+| Bridge -> Brain | Active state | Shared memory | 376 bytes |
+| Bridge -> Disk | Transitions | JSONL | ~300 bytes/line |
+| Brain -> Bridge | Action | ZMQ string | ~50 bytes |
+| Brain -> Kernel | Priority weight | BPF map | 8 bytes |
+        """)
 
+    # Footer
     gr.Markdown("""
 ---
-**Links:** [Model](https://huggingface.co/Rayugacodes/kernelx-strategist) |
-[Training Data](https://huggingface.co/datasets/Rayugacodes/kernelx-training-data) |
-[Colab Notebook](https://colab.research.google.com/github/pie-314/KernelX/blob/model-training-hugging-face-integration/KernelX_Training.ipynb) |
-[GitHub](https://github.com/pie-314/KernelX)
+[Model](https://huggingface.co/Rayugacodes/kernelx-strategist) |
+[Data](https://huggingface.co/datasets/Rayugacodes/kernelx-training-data) |
+[Colab](https://colab.research.google.com/github/pie-314/KernelX/blob/model-training-hugging-face-integration/KernelX_Training.ipynb) |
+[GitHub](https://github.com/pie-314/KernelX) |
+Built for Meta PyTorch OpenEnv Hackathon 2026
     """)
 
 app.launch(server_name="0.0.0.0", server_port=7860)