Add demo script, API endpoints, and training notebook baselines

- Add demo.py: standalone heuristic agent for full 6-quarter episodes
(CLI + JSON API, no LLM key required)
- Add /demo, /hyperparameters, /scenarios API endpoints to HF Space
- Update Dockerfile to include demo.py and agent/ directory
- Enhance training notebook with baseline comparisons and reward curves:
- Random baseline (5 seeds, ~0.15-0.25 floor)
- Heuristic baseline via /demo API (~0.50-0.65 target)
- Reward + loss curve plotting after GRPO training
- Before/after evaluation on same seeds
- Results summary with bar charts and JSON export

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

Dockerfile

@@ -5,7 +5,8 @@ WORKDIR /app
 COPY pyproject.toml README.md openenv.yaml ./
 COPY hr_env/ hr_env/
 COPY server/ server/
-COPY __init__.py client.py models.py ./
+COPY agent/ agent/
+COPY demo.py __init__.py client.py models.py ./
 
 RUN pip install --no-cache-dir -e ".[server]"
 

demo.py

@@ -0,0 +1,826 @@
+"""Standalone demo: run a full 6-quarter HR simulation with a heuristic strategy.
+
+No server, no LLM API key required. Seeds the environment with sample data and
+executes data-driven HR decisions through all HCM:21 phases each quarter.
+
+Usage:
+    python demo.py                                  # Default: seed=42, size=300
+    python demo.py --seed 123 --size 250 --scenario budget_cuts
+    python demo.py --scenario high_eng_turnover
+    python demo.py --all-scenarios                   # Run all 4 task scenarios
+
+API usage (from server):
+    from demo import run_demo_json
+    result = run_demo_json(seed=42, size=300, scenario="high_eng_turnover")
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Any, Dict, List, Optional
+
+from hr_env.models import HRAction
+from hr_env.server.environment import HRProductivityEnvironment
+
+DEPARTMENTS = ["Engineering", "Sales", "Operations", "HR", "Finance"]
+
+SCENARIOS = [
+    {"name": "high_eng_turnover", "seed": 42, "size": 300},
+    {"name": "budget_cuts", "seed": 123, "size": 250},
+    {"name": "rapid_growth", "seed": 456, "size": 350},
+    {"name": "balanced_optimization", "seed": 789, "size": 300},
+]
+
+
+# ── Formatting Helpers ──────────────────────────────────────────────
+
+
+def fmt_dollar(v: float) -> str:
+    if abs(v) >= 1_000_000:
+        return f"${v / 1_000_000:,.1f}M"
+    if abs(v) >= 1_000:
+        return f"${v:,.0f}"
+    return f"${v:.2f}"
+
+
+def fmt_pct(curr: float, prev: float) -> str:
+    if prev == 0:
+        return "  N/A"
+    change = (curr - prev) / abs(prev) * 100
+    return f"{change:+.1f}%"
+
+
+def print_header(seed: int, size: int, scenario: Optional[str]) -> None:
+    print()
+    print("=" * 64)
+    print("  HCM:21 HR Productivity Environment - Demo Playthrough")
+    print("=" * 64)
+    scn = scenario or "default (no scenario)"
+    print(f"  Seed: {seed} | Company Size: {size} | Scenario: {scn}")
+    print()
+
+
+def print_baseline(obs_data: dict) -> None:
+    summary = obs_data.get("company_summary", {})
+    metrics = obs_data.get("baseline_metrics", {})
+    print("--- COMPANY BASELINE ---")
+    print(f"  Employees: {summary.get('total_headcount', '?')} across 5 departments")
+    print(f"  Revenue: {fmt_dollar(summary.get('revenue', 0))} | Profit: {fmt_dollar(summary.get('profit', 0))}")
+    hcva = metrics.get("hcva", 0)
+    hcroi = metrics.get("hcroi", 0)
+    qips_c = metrics.get("qips", {}).get("composite", 0)
+    print(f"  HCVA: {fmt_dollar(hcva)}/FTE | HCROI: {hcroi:.2f}x | QIPS: {qips_c:.4f}")
+    # Department breakdown
+    depts = summary.get("departments", {})
+    for name in DEPARTMENTS:
+        d = depts.get(name, {})
+        print(
+            f"    {name:12s}: {d.get('headcount', '?'):>3} employees, "
+            f"perf {d.get('avg_performance', 0):.1f}, "
+            f"engage {d.get('avg_engagement', 0):.0f}, "
+            f"risk {d.get('avg_flight_risk', 0):.2f}"
+        )
+    print()
+
+
+def print_quarter_header(q: int) -> None:
+    print(f"{'=' * 64}")
+    print(f"  QUARTER {q}")
+    print(f"{'=' * 64}")
+
+
+def print_phase(phase: str) -> None:
+    print(f"\n  --- {phase.upper()} ---")
+
+
+def print_metrics_table(
+    metrics: dict, prev_metrics: Optional[dict], reward: Optional[float]
+) -> None:
+    prev = prev_metrics or {}
+    hcva = metrics.get("hcva", 0)
+    hcroi = metrics.get("hcroi", 0)
+    qips = metrics.get("qips", {}).get("composite", 0)
+    ev = metrics.get("employee_value", 0)
+    snap = metrics.get("snapshot", {})
+    headcount = snap.get("headcount", 0)
+    engagement = snap.get("avg_engagement", 0)
+
+    p_hcva = prev.get("hcva", 0)
+    p_hcroi = prev.get("hcroi", 0)
+    p_qips = prev.get("qips", {}).get("composite", 0)
+    p_snap = prev.get("snapshot", {})
+    p_hc = p_snap.get("headcount", 0)
+    p_eng = p_snap.get("avg_engagement", 0)
+
+    print()
+    print("  Metrics:")
+    print(f"    {'Metric':<14s} {'Value':>12s} {'Change':>8s}")
+    print(f"    {'-'*14} {'-'*12} {'-'*8}")
+    print(f"    {'HCVA':<14s} {fmt_dollar(hcva):>12s} {fmt_pct(hcva, p_hcva):>8s}")
+    print(f"    {'HCROI':<14s} {hcroi:>12.2f}x {fmt_pct(hcroi, p_hcroi):>8s}")
+    print(f"    {'QIPS':<14s} {qips:>12.4f} {fmt_pct(qips, p_qips):>8s}")
+    print(f"    {'Empl. Value':<14s} {ev:>12.4f}")
+    print(f"    {'Headcount':<14s} {headcount:>12d} {fmt_pct(headcount, p_hc):>8s}")
+    print(f"    {'Engagement':<14s} {engagement:>12.1f} {fmt_pct(engagement, p_eng):>8s}")
+
+    if reward is not None:
+        print(f"\n  Quarterly Reward: {reward:+.4f}")
+
+
+def print_final_summary(
+    total_steps: int,
+    final_score: float,
+    metric_history: List[dict],
+    quarterly_rewards: List[float],
+    all_events: List[str],
+) -> None:
+    print()
+    print("=" * 64)
+    print("                   FINAL EPISODE SUMMARY")
+    print("=" * 64)
+    print(f"  Total Steps: {total_steps}")
+    print(f"  Final Score: {final_score:.4f}")
+
+    if len(metric_history) >= 2:
+        first = metric_history[0]
+        last = metric_history[-1]
+        print(f"\n  Metric Trajectory (Baseline -> Q6):")
+        h0, h1 = first.get("hcva", 0), last.get("hcva", 0)
+        r0, r1 = first.get("hcroi", 0), last.get("hcroi", 0)
+        q0, q1 = first.get("qips", {}).get("composite", 0), last.get("qips", {}).get("composite", 0)
+        print(f"    HCVA:  {fmt_dollar(h0):>10s} -> {fmt_dollar(h1):>10s}  ({fmt_pct(h1, h0)})")
+        print(f"    HCROI: {r0:>10.2f}x -> {r1:>10.2f}x  ({fmt_pct(r1, r0)})")
+        print(f"    QIPS:  {q0:>10.4f} -> {q1:>10.4f}  ({fmt_pct(q1, q0)})")
+
+    if quarterly_rewards:
+        formatted = ", ".join(f"{r:+.4f}" for r in quarterly_rewards)
+        print(f"\n  Quarterly Rewards: [{formatted}]")
+
+    if all_events:
+        print(f"\n  Events Log ({len(all_events)} total):")
+        for ev in all_events:
+            print(f"    - {ev}")
+
+    print("=" * 64)
+    print()
+
+
+# ── Phase Execution ─────────────────────────────────────────────────
+
+
+def run_scanning(env: HRProductivityEnvironment) -> dict:
+    """Execute scanning phase: query all departments, employees, metrics, financials."""
+    print_phase("scanning")
+    dept_data = {}
+
+    # Query each department
+    for dept_name in DEPARTMENTS:
+        obs = env.step(HRAction(action_type="query_department", department=dept_name))
+        d = obs.data or {}
+        dept_data[dept_name] = d
+        print(
+            f"    [SCAN] {dept_name:12s}: {d.get('headcount', '?'):>3} employees, "
+            f"perf {d.get('avg_performance', 0):.1f}, "
+            f"engage {d.get('avg_engagement', 0):.0f}, "
+            f"risk {d.get('avg_flight_risk', 0):.2f}"
+        )
+
+    # Query high performers for promotion candidates
+    obs = env.step(HRAction(
+        action_type="query_employees", parameters={"min_performance": 4.0}
+    ))
+    high_performers = (obs.data or {}).get("employees", [])
+    promote_ids = [
+        e["id"] for e in high_performers if e.get("level", 5) < 5
+    ][:3]
+
+    # Query low performers for termination candidates
+    obs = env.step(HRAction(
+        action_type="query_employees", parameters={"min_flight_risk": 0.0}
+    ))
+    all_emps = (obs.data or {}).get("employees", [])
+    terminate_ids = [
+        e["id"] for e in all_emps if e.get("performance_score", 5.0) < 1.8
+    ][:2]
+
+    # Calculate all metrics
+    obs = env.step(HRAction(action_type="calculate_metric", metric_name="all"))
+    all_metrics = obs.data or {}
+    hcva = all_metrics.get("hcva", 0)
+    hcroi = all_metrics.get("hcroi", 0)
+    qips_c = all_metrics.get("qips", {}).get("composite", 0)
+    print(f"    [METRIC] HCVA: {fmt_dollar(hcva)} | HCROI: {hcroi:.2f}x | QIPS: {qips_c:.4f}")
+
+    # Review financials
+    obs = env.step(HRAction(action_type="review_financials"))
+    fin = obs.data or {}
+    hr_budget = fin.get("hr_budget_remaining", 0)
+    print(f"    [FINANCE] Revenue: {fmt_dollar(fin.get('revenue', 0))} | "
+          f"Profit: {fmt_dollar(fin.get('profit', 0))} | "
+          f"HR Budget: {fmt_dollar(hr_budget)}")
+
+    # Advance to planning
+    env.step(HRAction(action_type="advance_phase"))
+
+    return {
+        "depts": dept_data,
+        "metrics": all_metrics,
+        "financials": fin,
+        "hr_budget": hr_budget,
+        "promote_ids": promote_ids,
+        "terminate_ids": terminate_ids,
+    }
+
+
+def plan_quarter(
+    env: HRProductivityEnvironment,
+    scan_data: dict,
+    quarter: int,
+    prev_events: List[str],
+) -> dict:
+    """Execute planning phase with data-driven decisions."""
+    print_phase("planning")
+    depts = scan_data["depts"]
+    hr_budget = scan_data["hr_budget"]
+    plan = {}
+
+    # Rank departments by flight risk (descending) and performance (ascending)
+    dept_list = [
+        {
+            "name": name,
+            "headcount": d.get("headcount", 0),
+            "avg_performance": d.get("avg_performance", 3.0),
+            "avg_engagement": d.get("avg_engagement", 70),
+            "avg_flight_risk": d.get("avg_flight_risk", 0.1),
+        }
+        for name, d in depts.items()
+    ]
+
+    by_risk = sorted(dept_list, key=lambda d: d["avg_flight_risk"], reverse=True)
+    by_perf = sorted(dept_list, key=lambda d: d["avg_performance"])
+    by_headcount = sorted(
+        [d for d in dept_list if d["name"] != "HR"],
+        key=lambda d: d["headcount"],
+    )
+
+    # 1. Hiring: spread across the 2 smallest non-HR depts to offset turnover
+    hiring_dept = by_headcount[0]["name"]
+    # React to competitor poaching events
+    for ev in prev_events:
+        if "poaching" in ev.lower() or "competitor" in ev.lower():
+            for d in dept_list:
+                if d["name"].lower() in ev.lower():
+                    hiring_dept = d["name"]
+                    break
+
+    # Hire aggressively to offset ~20% quarterly turnover
+    total_hc = sum(d["headcount"] for d in dept_list)
+    hire_count = max(3, min(15, int(total_hc * 0.12)))
+    obs = env.step(HRAction(
+        action_type="set_hiring_target", department=hiring_dept, count=hire_count
+    ))
+    plan["hiring"] = {"dept": hiring_dept, "count": hire_count}
+    print(f"    [PLAN] Hiring target: {hiring_dept} +{hire_count}")
+
+    # Also set hiring target for second-smallest dept
+    hiring_dept2 = by_headcount[1]["name"] if len(by_headcount) > 1 else hiring_dept
+    hire_count2 = max(2, min(10, int(total_hc * 0.08)))
+    if hiring_dept2 != hiring_dept:
+        obs = env.step(HRAction(
+            action_type="set_hiring_target", department=hiring_dept2, count=hire_count2
+        ))
+        plan["hiring2"] = {"dept": hiring_dept2, "count": hire_count2}
+        print(f"    [PLAN] Hiring target: {hiring_dept2} +{hire_count2}")
+
+    # 2. Training: lowest-performing department
+    training_dept = by_perf[0]["name"]
+    training_amount = min(hr_budget * 0.30, 100_000)
+    if training_amount > 0:
+        obs = env.step(HRAction(
+            action_type="set_training_budget", department=training_dept, amount=training_amount
+        ))
+        # Update budget tracking from observation
+        if obs.data and "hr_budget_remaining" in obs.data:
+            hr_budget = obs.data["hr_budget_remaining"]
+        else:
+            hr_budget -= training_amount
+        plan["training"] = {"dept": training_dept, "amount": training_amount}
+        print(f"    [PLAN] Training budget: {fmt_dollar(training_amount)} -> {training_dept}")
+
+    # 3. Compensation: highest flight-risk department, +3%
+    comp_dept = by_risk[0]["name"]
+    comp_pct = 3.0
+    obs = env.step(HRAction(
+        action_type="set_compensation_policy", department=comp_dept, amount=comp_pct
+    ))
+    plan["compensation"] = {"dept": comp_dept, "pct": comp_pct}
+    print(f"    [PLAN] Compensation: +{comp_pct:.0f}% for {comp_dept}")
+
+    # 4. Retention: highest flight-risk dept (if budget allows)
+    retention_dept = by_risk[0]["name"]
+    retention_amount = min(hr_budget * 0.20, 50_000)
+    if retention_amount > 1000:
+        obs = env.step(HRAction(
+            action_type="set_retention_program", department=retention_dept, amount=retention_amount
+        ))
+        plan["retention"] = {"dept": retention_dept, "amount": retention_amount}
+        print(f"    [PLAN] Retention program: {retention_dept} ({fmt_dollar(retention_amount)})")
+
+    # Advance to producing
+    env.step(HRAction(action_type="advance_phase"))
+    return plan
+
+
+def execute_quarter(
+    env: HRProductivityEnvironment,
+    plan: dict,
+    scan_data: dict,
+) -> None:
+    """Execute producing phase: hire, train, promote, optionally terminate."""
+    print_phase("producing")
+
+    # 1. Execute hiring (primary dept)
+    hiring = plan.get("hiring", {})
+    if hiring:
+        obs = env.step(HRAction(
+            action_type="execute_hiring",
+            department=hiring["dept"],
+            count=hiring["count"],
+        ))
+        cost = (obs.data or {}).get("cost", 0)
+        print(f"    [EXEC] Hired {hiring['count']} in {hiring['dept']} (cost: {fmt_dollar(cost)})")
+
+    # 1b. Execute hiring (secondary dept)
+    hiring2 = plan.get("hiring2", {})
+    if hiring2:
+        obs = env.step(HRAction(
+            action_type="execute_hiring",
+            department=hiring2["dept"],
+            count=hiring2["count"],
+        ))
+        cost = (obs.data or {}).get("cost", 0)
+        print(f"    [EXEC] Hired {hiring2['count']} in {hiring2['dept']} (cost: {fmt_dollar(cost)})")
+
+    # 2. Execute training
+    training = plan.get("training", {})
+    if training:
+        obs = env.step(HRAction(
+            action_type="execute_training",
+            department=training["dept"],
+            amount=20,  # 20 hours per employee
+        ))
+        cost = (obs.data or {}).get("cost", 0)
+        hc = (obs.data or {}).get("headcount", "?")
+        print(f"    [EXEC] Training: {training['dept']}, 20 hrs x {hc} employees (cost: {fmt_dollar(cost)})")
+
+    # 3. Promote top performers (identified during scanning)
+    promote_ids = scan_data.get("promote_ids", [])
+    if promote_ids:
+        obs = env.step(HRAction(action_type="execute_promotion", employee_ids=promote_ids))
+        promoted = sum(
+            1 for p in (obs.data or {}).get("promotions", []) if p.get("success")
+        )
+        print(f"    [EXEC] Promoted {promoted} high-performing employees")
+    else:
+        print("    [EXEC] No promotion candidates found")
+
+    # 4. Terminate underperformers (identified during scanning)
+    terminate_ids = scan_data.get("terminate_ids", [])
+    if terminate_ids:
+        obs = env.step(HRAction(action_type="execute_termination", employee_ids=terminate_ids))
+        terminated = sum(
+            1 for t in (obs.data or {}).get("terminations", []) if t.get("success")
+        )
+        print(f"    [EXEC] Terminated {terminated} underperformers")
+    else:
+        print("    [EXEC] No underperformers to terminate")
+
+    # Advance to controlling
+    env.step(HRAction(action_type="advance_phase"))
+
+
+def run_controlling(env: HRProductivityEnvironment, q: int) -> dict:
+    """Execute controlling phase: submit report and advance quarter."""
+    print_phase("controlling")
+
+    # Submit report
+    env.step(HRAction(action_type="submit_report"))
+    print(f"    [REPORT] Q{q} report submitted")
+
+    # Advance quarter
+    obs = env.step(HRAction(action_type="advance_quarter"))
+    return {
+        "obs": obs,
+        "data": obs.data or {},
+        "reward": obs.reward,
+        "done": obs.done,
+    }
+
+
+# ── Main Demo Runner ────────────────────────────────────────────────
+
+
+def run_demo(seed: int = 42, size: int = 300, scenario: Optional[str] = None) -> float:
+    """Run a full 6-quarter demo episode with a heuristic strategy."""
+    print_header(seed, size, scenario)
+
+    env = HRProductivityEnvironment()
+    kwargs: Dict[str, Any] = {"size": size}
+    if scenario:
+        kwargs["scenario"] = scenario
+    obs = env.reset(seed=seed, **kwargs)
+
+    print_baseline(obs.data or {})
+
+    # Track state across quarters
+    baseline_metrics = (obs.data or {}).get("baseline_metrics", {})
+    metric_history = [baseline_metrics]
+    quarterly_rewards: List[float] = []
+    all_events: List[str] = []
+    prev_events: List[str] = []
+    total_steps = 0
+
+    for q in range(1, 7):
+        print_quarter_header(q)
+
+        # Scanning
+        scan_data = run_scanning(env)
+
+        # Planning
+        plan = plan_quarter(env, scan_data, q, prev_events)
+
+        # Producing
+        execute_quarter(env, plan, scan_data)
+
+        # Controlling
+        result = run_controlling(env, q)
+        data = result["data"]
+
+        # Collect results
+        reward = result["reward"]
+        events = data.get("events", [])
+        turnover = data.get("turnover", {})
+        departed = turnover.get("departed_count", 0)
+        metrics = data.get("metrics") or data.get("last_quarter_metrics", {})
+
+        all_events.extend(events)
+        prev_events = events
+
+        if reward is not None and not result["done"]:
+            quarterly_rewards.append(reward)
+
+        if metrics:
+            metric_history.append(metrics)
+
+        # Print quarter results
+        if events:
+            print(f"\n  Events:")
+            for ev in events:
+                print(f"    - {ev}")
+        if departed:
+            print(f"  Turnover: {departed} employees departed")
+
+        prev_m = metric_history[-2] if len(metric_history) >= 2 else None
+        if metrics:
+            print_metrics_table(metrics, prev_m, reward)
+
+        total_steps = data.get("total_steps", total_steps)
+
+        if result["done"]:
+            final_score = data.get("final_score", reward or 0.0)
+            total_steps = data.get("total_steps", total_steps)
+            # Collect all quarterly rewards from final data
+            qr = data.get("quarterly_rewards", quarterly_rewards)
+            if qr:
+                quarterly_rewards = qr
+            print_final_summary(
+                total_steps, final_score, metric_history, quarterly_rewards, all_events
+            )
+            return final_score
+
+    # Should not reach here, but just in case
+    return 0.0
+
+
+def _pct_change(curr: float, prev: float) -> Optional[float]:
+    if prev == 0:
+        return None
+    return round((curr - prev) / abs(prev) * 100, 2)
+
+
+def run_demo_json(
+    seed: int = 42, size: int = 300, scenario: Optional[str] = None
+) -> Dict[str, Any]:
+    """Run a full 6-quarter demo and return structured JSON results.
+
+    This is the API-friendly version of run_demo() — no print output,
+    returns all data as a dict suitable for JSON serialization.
+    """
+    env = HRProductivityEnvironment()
+    kwargs: Dict[str, Any] = {"size": size}
+    if scenario:
+        kwargs["scenario"] = scenario
+    obs = env.reset(seed=seed, **kwargs)
+
+    obs_data = obs.data or {}
+    baseline_metrics = obs_data.get("baseline_metrics", {})
+    baseline_summary = obs_data.get("company_summary", {})
+    metric_history = [baseline_metrics]
+    quarterly_rewards: List[float] = []
+    all_events: List[str] = []
+    prev_events: List[str] = []
+    quarters: List[Dict[str, Any]] = []
+
+    for q in range(1, 7):
+        scan_data = run_scanning.__wrapped__(env) if hasattr(run_scanning, '__wrapped__') else _scan_quiet(env)
+        plan = _plan_quiet(env, scan_data, q, prev_events)
+        _exec_quiet(env, plan, scan_data)
+        result = _control_quiet(env)
+        data = result["data"]
+
+        reward = result["reward"]
+        events = data.get("events", [])
+        turnover = data.get("turnover", {})
+        departed = turnover.get("departed_count", 0)
+        metrics = data.get("metrics") or data.get("last_quarter_metrics", {})
+
+        all_events.extend(events)
+        prev_events = events
+
+        prev_m = metric_history[-1] if metric_history else {}
+        quarter_result: Dict[str, Any] = {
+            "quarter": q,
+            "actions": {
+                "hiring": plan.get("hiring"),
+                "hiring2": plan.get("hiring2"),
+                "training": plan.get("training"),
+                "compensation": plan.get("compensation"),
+                "retention": plan.get("retention"),
+            },
+            "events": events,
+            "turnover_count": departed,
+            "quarterly_reward": reward,
+        }
+
+        if metrics:
+            metric_history.append(metrics)
+            quarter_result["metrics"] = {
+                "hcva": metrics.get("hcva", 0),
+                "hcroi": metrics.get("hcroi", 0),
+                "qips": metrics.get("qips", {}).get("composite", 0),
+                "employee_value": metrics.get("employee_value", 0),
+                "headcount": metrics.get("snapshot", {}).get("headcount", 0),
+                "avg_engagement": round(metrics.get("snapshot", {}).get("avg_engagement", 0), 1),
+            }
+            quarter_result["changes"] = {
+                "hcva_pct": _pct_change(metrics.get("hcva", 0), prev_m.get("hcva", 0)),
+                "hcroi_pct": _pct_change(metrics.get("hcroi", 0), prev_m.get("hcroi", 0)),
+                "qips_pct": _pct_change(
+                    metrics.get("qips", {}).get("composite", 0),
+                    prev_m.get("qips", {}).get("composite", 0),
+                ),
+            }
+
+        if reward is not None and not result["done"]:
+            quarterly_rewards.append(reward)
+
+        quarters.append(quarter_result)
+
+        if result["done"]:
+            final_score = data.get("final_score", reward or 0.0)
+            qr = data.get("quarterly_rewards", quarterly_rewards)
+            if qr:
+                quarterly_rewards = qr
+            break
+
+    first_m = metric_history[0] if metric_history else {}
+    last_m = metric_history[-1] if metric_history else {}
+
+    return {
+        "seed": seed,
+        "size": size,
+        "scenario": scenario,
+        "baseline": {
+            "headcount": baseline_summary.get("total_headcount", 0),
+            "revenue": baseline_summary.get("revenue", 0),
+            "profit": baseline_summary.get("profit", 0),
+            "hcva": first_m.get("hcva", 0),
+            "hcroi": first_m.get("hcroi", 0),
+            "qips": first_m.get("qips", {}).get("composite", 0),
+        },
+        "final": {
+            "score": final_score,
+            "total_steps": data.get("total_steps", 0),
+            "hcva": last_m.get("hcva", 0),
+            "hcroi": last_m.get("hcroi", 0),
+            "qips": last_m.get("qips", {}).get("composite", 0),
+            "hcva_change_pct": _pct_change(last_m.get("hcva", 0), first_m.get("hcva", 0)),
+            "hcroi_change_pct": _pct_change(last_m.get("hcroi", 0), first_m.get("hcroi", 0)),
+            "qips_change_pct": _pct_change(
+                last_m.get("qips", {}).get("composite", 0),
+                first_m.get("qips", {}).get("composite", 0),
+            ),
+        },
+        "quarterly_rewards": quarterly_rewards,
+        "events": all_events,
+        "quarters": quarters,
+    }
+
+
+def _scan_quiet(env: HRProductivityEnvironment) -> dict:
+    """Scanning phase — no print output."""
+    dept_data = {}
+    for dept_name in DEPARTMENTS:
+        obs = env.step(HRAction(action_type="query_department", department=dept_name))
+        dept_data[dept_name] = obs.data or {}
+
+    obs = env.step(HRAction(
+        action_type="query_employees", parameters={"min_performance": 4.0}
+    ))
+    high_performers = (obs.data or {}).get("employees", [])
+    promote_ids = [e["id"] for e in high_performers if e.get("level", 5) < 5][:3]
+
+    obs = env.step(HRAction(
+        action_type="query_employees", parameters={"min_flight_risk": 0.0}
+    ))
+    all_emps = (obs.data or {}).get("employees", [])
+    terminate_ids = [
+        e["id"] for e in all_emps if e.get("performance_score", 5.0) < 1.8
+    ][:2]
+
+    obs = env.step(HRAction(action_type="calculate_metric", metric_name="all"))
+    all_metrics = obs.data or {}
+
+    obs = env.step(HRAction(action_type="review_financials"))
+    fin = obs.data or {}
+
+    env.step(HRAction(action_type="advance_phase"))
+
+    return {
+        "depts": dept_data,
+        "metrics": all_metrics,
+        "financials": fin,
+        "hr_budget": fin.get("hr_budget_remaining", 0),
+        "promote_ids": promote_ids,
+        "terminate_ids": terminate_ids,
+    }
+
+
+def _plan_quiet(
+    env: HRProductivityEnvironment, scan_data: dict, quarter: int, prev_events: List[str]
+) -> dict:
+    """Planning phase — no print output."""
+    depts = scan_data["depts"]
+    hr_budget = scan_data["hr_budget"]
+    plan: Dict[str, Any] = {}
+
+    dept_list = [
+        {
+            "name": name,
+            "headcount": d.get("headcount", 0),
+            "avg_performance": d.get("avg_performance", 3.0),
+            "avg_engagement": d.get("avg_engagement", 70),
+            "avg_flight_risk": d.get("avg_flight_risk", 0.1),
+        }
+        for name, d in depts.items()
+    ]
+
+    by_risk = sorted(dept_list, key=lambda d: d["avg_flight_risk"], reverse=True)
+    by_perf = sorted(dept_list, key=lambda d: d["avg_performance"])
+    by_headcount = sorted(
+        [d for d in dept_list if d["name"] != "HR"], key=lambda d: d["headcount"]
+    )
+
+    hiring_dept = by_headcount[0]["name"]
+    for ev in prev_events:
+        if "poaching" in ev.lower() or "competitor" in ev.lower():
+            for d in dept_list:
+                if d["name"].lower() in ev.lower():
+                    hiring_dept = d["name"]
+                    break
+
+    total_hc = sum(d["headcount"] for d in dept_list)
+    hire_count = max(3, min(15, int(total_hc * 0.12)))
+    env.step(HRAction(action_type="set_hiring_target", department=hiring_dept, count=hire_count))
+    plan["hiring"] = {"dept": hiring_dept, "count": hire_count}
+
+    hiring_dept2 = by_headcount[1]["name"] if len(by_headcount) > 1 else hiring_dept
+    hire_count2 = max(2, min(10, int(total_hc * 0.08)))
+    if hiring_dept2 != hiring_dept:
+        env.step(HRAction(action_type="set_hiring_target", department=hiring_dept2, count=hire_count2))
+        plan["hiring2"] = {"dept": hiring_dept2, "count": hire_count2}
+
+    training_dept = by_perf[0]["name"]
+    training_amount = min(hr_budget * 0.30, 100_000)
+    if training_amount > 0:
+        obs = env.step(HRAction(
+            action_type="set_training_budget", department=training_dept, amount=training_amount
+        ))
+        if obs.data and "hr_budget_remaining" in obs.data:
+            hr_budget = obs.data["hr_budget_remaining"]
+        else:
+            hr_budget -= training_amount
+        plan["training"] = {"dept": training_dept, "amount": training_amount}
+
+    comp_dept = by_risk[0]["name"]
+    env.step(HRAction(action_type="set_compensation_policy", department=comp_dept, amount=3.0))
+    plan["compensation"] = {"dept": comp_dept, "pct": 3.0}
+
+    retention_dept = by_risk[0]["name"]
+    retention_amount = min(hr_budget * 0.20, 50_000)
+    if retention_amount > 1000:
+        env.step(HRAction(
+            action_type="set_retention_program", department=retention_dept, amount=retention_amount
+        ))
+        plan["retention"] = {"dept": retention_dept, "amount": retention_amount}
+
+    env.step(HRAction(action_type="advance_phase"))
+    return plan
+
+
+def _exec_quiet(
+    env: HRProductivityEnvironment, plan: dict, scan_data: dict
+) -> None:
+    """Producing phase — no print output."""
+    hiring = plan.get("hiring", {})
+    if hiring:
+        env.step(HRAction(action_type="execute_hiring", department=hiring["dept"], count=hiring["count"]))
+
+    hiring2 = plan.get("hiring2", {})
+    if hiring2:
+        env.step(HRAction(action_type="execute_hiring", department=hiring2["dept"], count=hiring2["count"]))
+
+    training = plan.get("training", {})
+    if training:
+        env.step(HRAction(action_type="execute_training", department=training["dept"], amount=20))
+
+    promote_ids = scan_data.get("promote_ids", [])
+    if promote_ids:
+        env.step(HRAction(action_type="execute_promotion", employee_ids=promote_ids))
+
+    terminate_ids = scan_data.get("terminate_ids", [])
+    if terminate_ids:
+        env.step(HRAction(action_type="execute_termination", employee_ids=terminate_ids))
+
+    env.step(HRAction(action_type="advance_phase"))
+
+
+def _control_quiet(env: HRProductivityEnvironment) -> dict:
+    """Controlling phase — no print output."""
+    env.step(HRAction(action_type="submit_report"))
+    obs = env.step(HRAction(action_type="advance_quarter"))
+    return {
+        "obs": obs,
+        "data": obs.data or {},
+        "reward": obs.reward,
+        "done": obs.done,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Run a full 6-quarter HR simulation demo with heuristic strategy"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
+    parser.add_argument("--size", type=int, default=300, help="Company size (default: 300)")
+    parser.add_argument(
+        "--scenario",
+        choices=["high_eng_turnover", "budget_cuts", "rapid_growth", "balanced_optimization"],
+        help="Scenario variant to run",
+    )
+    parser.add_argument(
+        "--all-scenarios", action="store_true",
+        help="Run all 4 task scenarios and compare results",
+    )
+    args = parser.parse_args()
+
+    if args.all_scenarios:
+        print()
+        print("=" * 64)
+        print("  Running all 4 scenarios...")
+        print("=" * 64)
+        results = []
+        for scn in SCENARIOS:
+            score = run_demo(seed=scn["seed"], size=scn["size"], scenario=scn["name"])
+            results.append((scn["name"], score))
+
+        print()
+        print("=" * 64)
+        print("              SCENARIO COMPARISON")
+        print("=" * 64)
+        print(f"  {'Scenario':<28s} {'Score':>8s}")
+        print(f"  {'-'*28} {'-'*8}")
+        for name, score in results:
+            print(f"  {name:<28s} {score:>8.4f}")
+        avg = sum(s for _, s in results) / len(results)
+        print(f"  {'-'*28} {'-'*8}")
+        print(f"  {'Average':<28s} {avg:>8.4f}")
+        print("=" * 64)
+    else:
+        score = run_demo(seed=args.seed, size=args.size, scenario=args.scenario)
+        print(f"  Exit score: {score:.4f}")
+
+
+if __name__ == "__main__":
+    main()

hr_env/server/app.py

@@ -1,5 +1,11 @@
 """FastAPI application entry point for the HR Productivity Environment."""
 
+import json
+from pathlib import Path
+from typing import Optional
+
+from fastapi import Query
+from fastapi.responses import RedirectResponse
 from openenv.core.env_server.http_server import create_app
 
 from hr_env.models import HRAction, HRObservation
@@ -12,3 +18,215 @@ app = create_app(
     env_name="hr-productivity-env",
     max_concurrent_envs=1,
 )
+
+
+@app.get("/")
+async def root():
+    return RedirectResponse(url="/web")
+
+
+# ── Hyperparameters endpoint ────────────────────────────────────────
+
+
+@app.get("/hyperparameters")
+async def hyperparameters():
+    """Return all environment hyperparameters and configuration constants."""
+    from hr_env.server.company import (
+        BENEFITS_MULTIPLIER,
+        DEPT_CONFIGS,
+        RECRUITING_COST_PER_HIRE,
+    )
+    from hr_env.server.data_gen import (
+        DEPT_SIZE_FRACTIONS,
+        DEPT_SKILLS,
+        ROLES,
+        SALARY_BASE,
+        SALARY_SIGMA,
+    )
+    from hr_env.server.events import ALL_EVENTS
+    from hr_env.server.phases import PHASE_ACTIONS, PHASE_MIN_ACTIONS, PHASES
+    from hr_env.server.scoring import compute_final_score  # noqa: F401
+
+    return {
+        "environment": {
+            "max_quarters": 6,
+            "max_steps": 300,
+            "max_concurrent_envs": 1,
+            "default_company_size": 300,
+            "recommended_size_range": [200, 500],
+            "departments": ["Engineering", "Sales", "Operations", "HR", "Finance"],
+        },
+        "company_financials": {
+            "base_revenue_formula": "size * 150,000",
+            "hr_budget_formula": "size * 6,000",
+            "quarterly_hr_budget": "hr_budget / 4",
+            "non_employment_cost_ratio": 0.55,
+            "benefits_multiplier": BENEFITS_MULTIPLIER,
+            "recruiting_cost_per_hire": RECRUITING_COST_PER_HIRE,
+            "training_cost_per_hour_per_employee": 50,
+        },
+        "cobb_douglas_production": {
+            dept: {
+                "A": cfg["A"],
+                "alpha": cfg["alpha"],
+                "beta": cfg["beta"],
+                "base_revenue_share": cfg["base_revenue_share"],
+            }
+            for dept, cfg in DEPT_CONFIGS.items()
+        },
+        "employee_generation": {
+            "salary_base_by_level": SALARY_BASE,
+            "salary_lognormal_sigma": SALARY_SIGMA,
+            "performance_distribution": {"type": "normal", "mean": 3.2, "std": 0.8, "clip": [1.0, 5.0]},
+            "tenure_distribution": {"type": "exponential", "lambda_months": 24},
+            "engagement_distribution": {"type": "beta", "alpha": 7, "beta": 3, "scale": 100},
+            "level_distribution": {
+                "type": "categorical",
+                "probabilities": {"L1": 0.35, "L2": 0.30, "L3": 0.20, "L4": 0.10, "L5": 0.05},
+            },
+            "dept_size_fractions": DEPT_SIZE_FRACTIONS,
+            "skills_per_employee": {"min": 2, "max": 4},
+        },
+        "employee_lifecycle": {
+            "promotion_salary_increase": 0.15,
+            "promotion_engagement_boost": 10,
+            "promotion_flight_risk_reduction": 0.15,
+            "quarterly_engagement_decay": -1.5,
+            "long_tenure_no_promotion_penalty": -3.0,
+            "long_tenure_threshold_quarters": 4,
+            "transfer_engagement_penalty": -5,
+            "termination_colleague_engagement_penalty": -3,
+            "new_hire_engagement_distribution": {"type": "beta", "alpha": 8, "beta": 2, "scale": 100},
+        },
+        "flight_risk": {
+            "model": "logistic",
+            "formula": "sigmoid(-1 + 2*pay_factor + 1.5*engagement_factor + 1*tenure_factor + 0.8*market_demand)",
+            "turnover_probability": "flight_risk * 0.4 per quarter",
+            "bounds": [0.02, 0.95],
+        },
+        "phase_system": {
+            "phases": PHASES,
+            "actions_per_phase": PHASE_ACTIONS,
+            "minimum_actions_per_phase": PHASE_MIN_ACTIONS,
+        },
+        "stochastic_events": [
+            {
+                "name": ev.name,
+                "description": ev.description,
+                "probability": ev.probability,
+            }
+            for ev in ALL_EVENTS
+        ],
+        "scoring": {
+            "quarterly_reward_weights": {
+                "hcva_improvement": 0.40,
+                "hcroi_improvement": 0.20,
+                "qips_improvement": 0.20,
+                "five_indexes": 0.20,
+            },
+            "final_score_weights": {
+                "hcva_trajectory": 0.25,
+                "hcroi_final_vs_baseline": 0.20,
+                "employee_value": 0.20,
+                "qips_consistency": 0.15,
+                "five_indexes_cumulative": 0.10,
+                "financial_health": 0.10,
+            },
+            "calibration_targets": {
+                "random_agent": "0.15-0.25",
+                "reasonable_heuristic": "~0.50",
+                "strong_agent": "0.70+",
+            },
+        },
+        "heuristic_strategy": {
+            "description": "Data-driven heuristic used in the demo endpoint",
+            "hiring_rate": "12% of current headcount for primary dept, 8% for secondary",
+            "training_budget_fraction": 0.30,
+            "training_hours_per_session": 20,
+            "compensation_adjustment_pct": 3.0,
+            "retention_budget_fraction": 0.20,
+            "promotion_criteria": {"min_performance": 4.0, "max_level": 4, "max_per_quarter": 3},
+            "termination_criteria": {"max_performance": 1.8, "max_per_quarter": 2},
+        },
+        "roles_by_department": {
+            dept: [{"role": role, "level": level} for role, level in roles]
+            for dept, roles in ROLES.items()
+        },
+        "skills_by_department": DEPT_SKILLS,
+    }
+
+
+# ── Scenarios endpoint ──────────────────────────────────────────────
+
+
+@app.get("/scenarios")
+async def scenarios():
+    """Return all available task scenarios with their configurations."""
+    tasks_dir = Path(__file__).resolve().parent.parent / "tasks"
+    task_list = []
+
+    if tasks_dir.exists():
+        for f in sorted(tasks_dir.glob("*.json")):
+            try:
+                task_list.append(json.loads(f.read_text()))
+            except (json.JSONDecodeError, OSError):
+                continue
+
+    # Also include scenario effect descriptions
+    scenario_effects = {
+        "high_eng_turnover": {
+            "modifications": "Engineering flight_risk += 0.20, engagement -= 10",
+            "challenge": "Stabilize retention in a bleeding engineering department",
+        },
+        "budget_cuts": {
+            "modifications": "HR budget halved (50% reduction)",
+            "challenge": "Optimize workforce with severely constrained resources",
+        },
+        "rapid_growth": {
+            "modifications": "Revenue +30%, 15% employees deactivated (understaffing)",
+            "challenge": "Scale the workforce while maintaining quality and culture",
+        },
+        "balanced_optimization": {
+            "modifications": "All employees: performance -0.3, engagement -8",
+            "challenge": "Improve below-average metrics across the board",
+        },
+    }
+
+    return {
+        "scenarios": [
+            {
+                **task,
+                "effects": scenario_effects.get(task.get("params", {}).get("scenario"), {}),
+            }
+            for task in task_list
+        ],
+        "default": {
+            "seed": 42,
+            "size": 300,
+            "scenario": None,
+            "description": "Standard company with no scenario modifications",
+        },
+    }
+
+
+# ── Demo endpoint ───────────────────────────────────────────────────
+
+
+@app.get("/demo")
+async def demo(
+    seed: int = Query(42, description="Random seed for reproducibility"),
+    size: int = Query(300, ge=50, le=1000, description="Company size (employees)"),
+    scenario: Optional[str] = Query(
+        None,
+        description="Scenario variant",
+        enum=["high_eng_turnover", "budget_cuts", "rapid_growth", "balanced_optimization"],
+    ),
+):
+    """Run a full 6-quarter demo with a heuristic strategy and return results.
+
+    Executes data-driven HR decisions through all HCM:21 phases for 6 quarters.
+    No LLM API key required — uses a built-in heuristic agent.
+    """
+    from demo import run_demo_json
+
+    return run_demo_json(seed=seed, size=size, scenario=scenario)

notebooks/hcm21_trl_training.ipynb

@@ -1,440 +1,695 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "accelerator": "GPU"
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
   },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "# HCM:21 Training with TRL + Unsloth\n",
-        "\n",
-        "This notebook trains a small LLM to play as Chief Human Capital Officer in the\n",
-        "[HCM:21 environment](https://huggingface.co/spaces/ParetoOptimal/hcm21) using\n",
-        "**GRPO** (Group Relative Policy Optimization) from HuggingFace TRL, with\n",
-        "**Unsloth** for efficient fine-tuning on a free Colab T4 GPU.\n",
-        "\n",
-        "The agent interacts with the HCM:21 OpenEnv environment deployed on HF Spaces\n",
-        "via WebSocket, making HR decisions (hiring, training, compensation, retention)\n",
-        "across 6 simulated quarters and receiving sparse rewards based on Fitz-enz\n",
-        "HR analytics metrics.\n",
-        "\n",
-        "**Environment:** https://huggingface.co/spaces/ParetoOptimal/hcm21  \n",
-        "**GitHub:** https://github.com/n8mauer/HCM-21"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 1. Install Dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "%%capture\n",
-        "!pip install unsloth trl openenv-core pydantic datasets\n",
-        "# Unsloth automatically patches transformers for 2x faster training on T4"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 2. Load Model with Unsloth"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from unsloth import FastLanguageModel\n",
-        "\n",
-        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-        "    model_name=\"unsloth/Qwen3-0.6B\",\n",
-        "    max_seq_length=4096,\n",
-        "    load_in_4bit=True,\n",
-        ")\n",
-        "\n",
-        "model = FastLanguageModel.get_peft_model(\n",
-        "    model,\n",
-        "    r=16,\n",
-        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
-        "                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
-        "    lora_alpha=16,\n",
-        "    lora_dropout=0,\n",
-        "    use_gradient_checkpointing=\"unsloth\",\n",
-        ")\n",
-        "print(f\"Model loaded. Trainable params: {model.print_trainable_parameters()}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 3. Define the HCM:21 Environment Wrapper\n",
-        "\n",
-        "We wrap the OpenEnv WebSocket client as a TRL `environment_factory` that\n",
-        "exposes `reset()`, `step()`, and `get_reward()` methods. The agent calls\n",
-        "`step()` as a tool, and the reward comes from quarterly HR metric improvements.\n",
-        "\n",
-        "The OpenEnv WebSocket protocol maintains stateful sessions — each connection\n",
-        "gets its own Environment instance on the server."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import json\n",
-        "from openenv.core import GenericEnvClient\n",
-        "\n",
-        "HCM21_URL = \"https://paretooptimal-hcm21.hf.space\"\n",
-        "\n",
-        "\n",
-        "class HCM21Env:\n",
-        "    \"\"\"TRL-compatible wrapper for the HCM:21 OpenEnv environment.\n",
-        "\n",
-        "    Uses the OpenEnv GenericEnvClient (WebSocket) for stateful sessions.\n",
-        "    Each instance gets its own environment on the server.\n",
-        "    \"\"\"\n",
-        "\n",
-        "    def __init__(self):\n",
-        "        self._client = GenericEnvClient(base_url=HCM21_URL)\n",
-        "        self._client.connect()\n",
-        "        self._reward = 0.0\n",
-        "        self._done = False\n",
-        "\n",
-        "    def reset(self, seed: int = 42) -> str:\n",
-        "        \"\"\"Reset the environment for a new episode.\"\"\"\n",
-        "        result = self._client.reset(seed=seed)\n",
-        "        obs = result.observation\n",
-        "        self._reward = result.reward or 0.0\n",
-        "        self._done = result.done\n",
-        "        return self._format_obs(obs)\n",
-        "\n",
-        "    def step(self, action_json: str) -> str:\n",
-        "        \"\"\"Take an action in the environment.\n",
-        "\n",
-        "        Args:\n",
-        "            action_json: JSON string with action_type and parameters.\n",
-        "                Example: {\\\"action_type\\\": \\\"query_department\\\", \\\"department\\\": \\\"Engineering\\\"}\n",
-        "\n",
-        "        Returns:\n",
-        "            Observation string describing the result.\n",
-        "        \"\"\"\n",
-        "        try:\n",
-        "            action = json.loads(action_json)\n",
-        "        except json.JSONDecodeError:\n",
-        "            return \"Error: Invalid JSON. Provide a valid action JSON object.\"\n",
-        "\n",
-        "        from openenv.core import GenericAction\n",
-        "        generic_action = GenericAction(**action)\n",
-        "        result = self._client.step(generic_action)\n",
-        "        obs = result.observation\n",
-        "        if result.reward is not None:\n",
-        "            self._reward = float(result.reward)\n",
-        "        self._done = result.done\n",
-        "        return self._format_obs(obs)\n",
-        "\n",
-        "    def get_reward(self) -> float:\n",
-        "        \"\"\"Return the reward from the last step.\"\"\"\n",
-        "        return self._reward\n",
-        "\n",
-        "    def close(self):\n",
-        "        \"\"\"Disconnect from the environment.\"\"\"\n",
-        "        try:\n",
-        "            self._client.disconnect()\n",
-        "        except Exception:\n",
-        "            pass\n",
-        "\n",
-        "    def _format_obs(self, obs) -> str:\n",
-        "        \"\"\"Format observation as a concise string for the LLM.\"\"\"\n",
-        "        if isinstance(obs, dict):\n",
-        "            data = obs\n",
-        "        else:\n",
-        "            data = obs if isinstance(obs, dict) else vars(obs) if hasattr(obs, '__dict__') else {\"raw\": str(obs)}\n",
-        "        parts = []\n",
-        "        parts.append(f\"Q{data.get('current_quarter', '?')} | {data.get('current_phase', '?')} phase\")\n",
-        "        parts.append(str(data.get('message', '')))\n",
-        "        warnings = data.get('warnings', [])\n",
-        "        if warnings:\n",
-        "            parts.append(f\"Warnings: {', '.join(str(w) for w in warnings)}\")\n",
-        "        actions = data.get('available_actions', [])\n",
-        "        if actions:\n",
-        "            parts.append(f\"Available: {', '.join(str(a) for a in actions)}\")\n",
-        "        obs_data = data.get('data')\n",
-        "        if obs_data:\n",
-        "            data_str = json.dumps(obs_data, indent=2, default=str)\n",
-        "            if len(data_str) > 1500:\n",
-        "                data_str = data_str[:1500] + \"...\"\n",
-        "            parts.append(f\"Data: {data_str}\")\n",
-        "        return \"\\n\".join(parts)\n",
-        "\n",
-        "\n",
-        "# Quick connectivity test\n",
-        "import requests\n",
-        "try:\n",
-        "    health = requests.get(f\"{HCM21_URL}/health\", timeout=10)\n",
-        "    print(f\"Environment health: {health.json()}\")\n",
-        "    env_test = HCM21Env()\n",
-        "    obs = env_test.reset(seed=42)\n",
-        "    print(f\"\\nReset OK. Observation preview:\")\n",
-        "    print(obs[:300])\n",
-        "    env_test.close()\n",
-        "except Exception as e:\n",
-        "    print(f\"Environment not reachable: {e}\")\n",
-        "    print(\"If the Space is still building, wait a few minutes and re-run.\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 4. Build Training Dataset\n",
-        "\n",
-        "For GRPO, we need a dataset of prompts that start each episode.\n",
-        "Each prompt describes the CHCO role and the current scenario."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from datasets import Dataset\n",
-        "\n",
-        "SYSTEM_PROMPT = \"\"\"You are the Chief Human Capital Officer (CHCO) of a simulated company.\n",
-        "You manage 300 employees across 5 departments over 6 quarters.\n",
-        "Each quarter has 4 phases: Scanning, Planning, Producing, Controlling.\n",
-        "Maximize HCVA, HCROI, and QIPS metrics by making strategic HR decisions.\n",
-        "\n",
-        "When calling the step tool, provide a JSON action object like:\n",
-        "{\"action_type\": \"query_department\", \"department\": \"Engineering\"}\n",
-        "\n",
-        "Available action types vary by phase:\n",
-        "- Scanning: query_department, query_employees, calculate_metric, review_financials\n",
-        "- Planning: set_hiring_target, set_training_budget, set_compensation_policy, set_retention_program\n",
-        "- Producing: execute_hiring, execute_promotion, execute_transfer, execute_training, execute_termination\n",
-        "- Controlling: submit_report, calculate_metric, advance_quarter\n",
-        "- advance_phase: move to next phase (after minimum actions taken)\n",
-        "\n",
-        "Think step by step. Start by scanning each department.\"\"\"\n",
-        "\n",
-        "# Create prompts with different seeds for diversity\n",
-        "prompts = []\n",
-        "for seed in range(20):\n",
-        "    prompts.append({\n",
-        "        \"prompt\": [\n",
-        "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-        "            {\"role\": \"user\", \"content\": f\"Begin managing the company. Episode seed: {seed}. Start by scanning your organization to understand its current state.\"},\n",
-        "        ]\n",
-        "    })\n",
-        "\n",
-        "dataset = Dataset.from_list(prompts)\n",
-        "print(f\"Training dataset: {len(dataset)} episodes\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 5. Define Reward Function\n",
-        "\n",
-        "The reward function uses the environment's built-in scoring.\n",
-        "Quarterly rewards are based on HCVA/HCROI/QIPS improvement.\n",
-        "The final episode score (0-1) is the terminal reward."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "def reward_func(completions, environments, **kwargs):\n",
-        "    \"\"\"Extract rewards from the HCM:21 environment after each rollout.\"\"\"\n",
-        "    rewards = []\n",
-        "    for env in environments:\n",
-        "        r = env.get_reward()\n",
-        "        rewards.append(r)\n",
-        "    return rewards\n",
-        "\n",
-        "print(\"Reward function defined.\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 6. Configure and Run GRPO Training\n",
-        "\n",
-        "We use TRL's `GRPOTrainer` with `environment_factory=HCM21Env`.\n",
-        "The model generates tool calls to interact with the environment,\n",
-        "and GRPO optimizes the policy based on the environment rewards.\n",
-        "\n",
-        "On a free T4 GPU, we use:\n",
-        "- Small batch size (2) with gradient accumulation\n",
-        "- 4-bit quantization via Unsloth\n",
-        "- Short max_completion_length to fit in memory"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from trl import GRPOTrainer, GRPOConfig\n",
-        "\n",
-        "training_args = GRPOConfig(\n",
-        "    output_dir=\"./hcm21-grpo-output\",\n",
-        "    num_train_epochs=1,\n",
-        "    per_device_train_batch_size=2,\n",
-        "    gradient_accumulation_steps=4,\n",
-        "    num_generations=4,\n",
-        "    max_completion_length=2048,\n",
-        "    learning_rate=5e-6,\n",
-        "    logging_steps=1,\n",
-        "    save_steps=10,\n",
-        "    bf16=True,\n",
-        "    chat_template_kwargs={\"enable_thinking\": False},\n",
-        ")\n",
-        "\n",
-        "trainer = GRPOTrainer(\n",
-        "    model=model,\n",
-        "    processing_class=tokenizer,\n",
-        "    train_dataset=dataset,\n",
-        "    reward_funcs=reward_func,\n",
-        "    environment_factory=HCM21Env,\n",
-        "    args=training_args,\n",
-        ")\n",
-        "\n",
-        "print(\"Trainer initialized. Starting training...\")\n",
-        "trainer.train()\n",
-        "print(\"Training complete!\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 7. Save & Push Model"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Save locally\n",
-        "model.save_pretrained(\"./hcm21-agent-lora\")\n",
-        "tokenizer.save_pretrained(\"./hcm21-agent-lora\")\n",
-        "\n",
-        "# Optional: push to HF Hub\n",
-        "# model.push_to_hub(\"your-username/hcm21-agent-lora\")\n",
-        "# tokenizer.push_to_hub(\"your-username/hcm21-agent-lora\")\n",
-        "\n",
-        "print(\"Model saved to ./hcm21-agent-lora\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 8. Evaluate the Trained Agent\n",
-        "\n",
-        "Run a full 6-quarter episode with the trained model to see its score."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import json\n",
-        "\n",
-        "FastLanguageModel.for_inference(model)\n",
-        "\n",
-        "env = HCM21Env()\n",
-        "obs_text = env.reset(seed=99)\n",
-        "print(\"Episode started.\")\n",
-        "\n",
-        "messages = [\n",
-        "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-        "    {\"role\": \"user\", \"content\": f\"Begin. Here is your initial observation:\\n{obs_text}\"},\n",
-        "]\n",
-        "\n",
-        "step_num = 0\n",
-        "max_steps = 200\n",
-        "\n",
-        "while not env._done and step_num < max_steps:\n",
-        "    # Generate action\n",
-        "    inputs = tokenizer.apply_chat_template(\n",
-        "        messages, return_tensors=\"pt\", add_generation_prompt=True\n",
-        "    ).to(model.device)\n",
-        "\n",
-        "    outputs = model.generate(\n",
-        "        input_ids=inputs,\n",
-        "        max_new_tokens=256,\n",
-        "        temperature=0.7,\n",
-        "        do_sample=True,\n",
-        "    )\n",
-        "    response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
-        "\n",
-        "    # Extract JSON action from response\n",
-        "    try:\n",
-        "        start = response.index(\"{\")\n",
-        "        depth = 0\n",
-        "        for i, c in enumerate(response[start:], start):\n",
-        "            if c == \"{\": depth += 1\n",
-        "            elif c == \"}\": depth -= 1\n",
-        "            if depth == 0:\n",
-        "                action_json = response[start:i+1]\n",
-        "                break\n",
-        "        obs_text = env.step(action_json)\n",
-        "        action = json.loads(action_json)\n",
-        "        print(f\"Step {step_num} | {action.get('action_type', '?')} | Reward: {env._reward:.4f}\")\n",
-        "    except (ValueError, json.JSONDecodeError):\n",
-        "        obs_text = env.step(json.dumps({\"action_type\": \"advance_phase\"}))\n",
-        "        print(f\"Step {step_num} | fallback advance_phase\")\n",
-        "\n",
-        "    # Rolling context window\n",
-        "    messages.append({\"role\": \"assistant\", \"content\": response})\n",
-        "    messages.append({\"role\": \"user\", \"content\": obs_text})\n",
-        "    if len(messages) > 12:\n",
-        "        messages = messages[:2] + messages[-10:]\n",
-        "\n",
-        "    step_num += 1\n",
-        "\n",
-        "env.close()\n",
-        "print(f\"\\nEpisode complete after {step_num} steps.\")\n",
-        "print(f\"Final score: {env._reward:.4f}\")"
-      ]
-    }
-  ]
-}
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# HCM:21 Training with TRL + Unsloth\n",
+    "\n",
+    "This notebook trains a small LLM to play as Chief Human Capital Officer in the\n",
+    "[HCM:21 environment](https://huggingface.co/spaces/ParetoOptimal/hcm21) using\n",
+    "**GRPO** (Group Relative Policy Optimization) from HuggingFace TRL, with\n",
+    "**Unsloth** for efficient fine-tuning on a free Colab T4 GPU.\n",
+    "\n",
+    "The agent interacts with the HCM:21 OpenEnv environment deployed on HF Spaces\n",
+    "via WebSocket, making HR decisions (hiring, training, compensation, retention)\n",
+    "across 6 simulated quarters and receiving sparse rewards based on Fitz-enz\n",
+    "HR analytics metrics.\n",
+    "\n",
+    "**Environment:** https://huggingface.co/spaces/ParetoOptimal/hcm21  \n",
+    "**GitHub:** https://github.com/n8mauer/HCM-21"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Install Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "!pip install unsloth trl openenv-core pydantic datasets\n",
+    "# Unsloth automatically patches transformers for 2x faster training on T4"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Load Model with Unsloth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name=\"unsloth/Qwen3-0.6B\",\n",
+    "    max_seq_length=4096,\n",
+    "    load_in_4bit=True,\n",
+    ")\n",
+    "\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r=16,\n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
+    "                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    lora_alpha=16,\n",
+    "    lora_dropout=0,\n",
+    "    use_gradient_checkpointing=\"unsloth\",\n",
+    ")\n",
+    "print(f\"Model loaded. Trainable params: {model.print_trainable_parameters()}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Define the HCM:21 Environment Wrapper\n",
+    "\n",
+    "We wrap the OpenEnv WebSocket client as a TRL `environment_factory` that\n",
+    "exposes `reset()`, `step()`, and `get_reward()` methods. The agent calls\n",
+    "`step()` as a tool, and the reward comes from quarterly HR metric improvements.\n",
+    "\n",
+    "The OpenEnv WebSocket protocol maintains stateful sessions \u00e2\u20ac\u201d each connection\n",
+    "gets its own Environment instance on the server."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from openenv.core import GenericEnvClient\n",
+    "\n",
+    "HCM21_URL = \"https://paretooptimal-hcm21.hf.space\"\n",
+    "\n",
+    "\n",
+    "class HCM21Env:\n",
+    "    \"\"\"TRL-compatible wrapper for the HCM:21 OpenEnv environment.\n",
+    "\n",
+    "    Uses the OpenEnv GenericEnvClient (WebSocket) for stateful sessions.\n",
+    "    Each instance gets its own environment on the server.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self):\n",
+    "        self._client = GenericEnvClient(base_url=HCM21_URL)\n",
+    "        self._client.connect()\n",
+    "        self._reward = 0.0\n",
+    "        self._done = False\n",
+    "\n",
+    "    def reset(self, seed: int = 42) -> str:\n",
+    "        \"\"\"Reset the environment for a new episode.\"\"\"\n",
+    "        result = self._client.reset(seed=seed)\n",
+    "        obs = result.observation\n",
+    "        self._reward = result.reward or 0.0\n",
+    "        self._done = result.done\n",
+    "        return self._format_obs(obs)\n",
+    "\n",
+    "    def step(self, action_json: str) -> str:\n",
+    "        \"\"\"Take an action in the environment.\n",
+    "\n",
+    "        Args:\n",
+    "            action_json: JSON string with action_type and parameters.\n",
+    "                Example: {\\\"action_type\\\": \\\"query_department\\\", \\\"department\\\": \\\"Engineering\\\"}\n",
+    "\n",
+    "        Returns:\n",
+    "            Observation string describing the result.\n",
+    "        \"\"\"\n",
+    "        try:\n",
+    "            action = json.loads(action_json)\n",
+    "        except json.JSONDecodeError:\n",
+    "            return \"Error: Invalid JSON. Provide a valid action JSON object.\"\n",
+    "\n",
+    "        from openenv.core import GenericAction\n",
+    "        generic_action = GenericAction(**action)\n",
+    "        result = self._client.step(generic_action)\n",
+    "        obs = result.observation\n",
+    "        if result.reward is not None:\n",
+    "            self._reward = float(result.reward)\n",
+    "        self._done = result.done\n",
+    "        return self._format_obs(obs)\n",
+    "\n",
+    "    def get_reward(self) -> float:\n",
+    "        \"\"\"Return the reward from the last step.\"\"\"\n",
+    "        return self._reward\n",
+    "\n",
+    "    def close(self):\n",
+    "        \"\"\"Disconnect from the environment.\"\"\"\n",
+    "        try:\n",
+    "            self._client.disconnect()\n",
+    "        except Exception:\n",
+    "            pass\n",
+    "\n",
+    "    def _format_obs(self, obs) -> str:\n",
+    "        \"\"\"Format observation as a concise string for the LLM.\"\"\"\n",
+    "        if isinstance(obs, dict):\n",
+    "            data = obs\n",
+    "        else:\n",
+    "            data = obs if isinstance(obs, dict) else vars(obs) if hasattr(obs, '__dict__') else {\"raw\": str(obs)}\n",
+    "        parts = []\n",
+    "        parts.append(f\"Q{data.get('current_quarter', '?')} | {data.get('current_phase', '?')} phase\")\n",
+    "        parts.append(str(data.get('message', '')))\n",
+    "        warnings = data.get('warnings', [])\n",
+    "        if warnings:\n",
+    "            parts.append(f\"Warnings: {', '.join(str(w) for w in warnings)}\")\n",
+    "        actions = data.get('available_actions', [])\n",
+    "        if actions:\n",
+    "            parts.append(f\"Available: {', '.join(str(a) for a in actions)}\")\n",
+    "        obs_data = data.get('data')\n",
+    "        if obs_data:\n",
+    "            data_str = json.dumps(obs_data, indent=2, default=str)\n",
+    "            if len(data_str) > 1500:\n",
+    "                data_str = data_str[:1500] + \"...\"\n",
+    "            parts.append(f\"Data: {data_str}\")\n",
+    "        return \"\\n\".join(parts)\n",
+    "\n",
+    "\n",
+    "# Quick connectivity test\n",
+    "import requests\n",
+    "try:\n",
+    "    health = requests.get(f\"{HCM21_URL}/health\", timeout=10)\n",
+    "    print(f\"Environment health: {health.json()}\")\n",
+    "    env_test = HCM21Env()\n",
+    "    obs = env_test.reset(seed=42)\n",
+    "    print(f\"\\nReset OK. Observation preview:\")\n",
+    "    print(obs[:300])\n",
+    "    env_test.close()\n",
+    "except Exception as e:\n",
+    "    print(f\"Environment not reachable: {e}\")\n",
+    "    print(\"If the Space is still building, wait a few minutes and re-run.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3a. Random Baseline\n",
+    "\n",
+    "Before training, establish the **floor** by running episodes with random actions.\n",
+    "A random agent scores ~0.15-0.25, showing how hard the environment is without strategy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import random\n",
+    "import numpy as np\n",
+    "\n",
+    "BASELINE_SEEDS = [42, 99, 123, 456, 789]\n",
+    "\n",
+    "# Valid actions by phase for a random agent\n",
+    "RANDOM_ACTIONS = {\n",
+    "    \"scanning\": [\n",
+    "        {\"action_type\": \"query_department\", \"department\": \"Engineering\"},\n",
+    "        {\"action_type\": \"query_department\", \"department\": \"Sales\"},\n",
+    "        {\"action_type\": \"query_department\", \"department\": \"Operations\"},\n",
+    "        {\"action_type\": \"review_financials\"},\n",
+    "        {\"action_type\": \"calculate_metric\", \"metric_name\": \"hcva\"},\n",
+    "    ],\n",
+    "    \"planning\": [\n",
+    "        {\"action_type\": \"set_hiring_target\", \"department\": \"Engineering\", \"count\": 5},\n",
+    "        {\"action_type\": \"set_training_budget\", \"department\": \"Sales\", \"amount\": 10000},\n",
+    "        {\"action_type\": \"set_compensation_policy\", \"department\": \"Operations\", \"parameters\": {\"adjustment_pct\": 2.0}},\n",
+    "    ],\n",
+    "    \"producing\": [\n",
+    "        {\"action_type\": \"execute_hiring\", \"department\": \"Engineering\", \"count\": 3},\n",
+    "        {\"action_type\": \"execute_training\", \"department\": \"Sales\", \"parameters\": {\"hours\": 10}},\n",
+    "    ],\n",
+    "    \"controlling\": [\n",
+    "        {\"action_type\": \"calculate_metric\", \"metric_name\": \"hcva\"},\n",
+    "        {\"action_type\": \"submit_report\", \"rationale\": \"Random agent report\"},\n",
+    "    ],\n",
+    "}\n",
+    "\n",
+    "def run_random_episode(seed):\n",
+    "    \"\"\"Run one episode with random actions, return final score.\"\"\"\n",
+    "    env = HCM21Env()\n",
+    "    try:\n",
+    "        env.reset(seed=seed)\n",
+    "        rng = random.Random(seed)\n",
+    "        steps = 0\n",
+    "        while not env._done and steps < 300:\n",
+    "            if steps % 5 == 4:\n",
+    "                env.step(json.dumps({\"action_type\": \"advance_phase\"}))\n",
+    "            elif steps % 8 == 7:\n",
+    "                env.step(json.dumps({\"action_type\": \"advance_quarter\"}))\n",
+    "            else:\n",
+    "                phase = rng.choice(list(RANDOM_ACTIONS.keys()))\n",
+    "                action = rng.choice(RANDOM_ACTIONS[phase])\n",
+    "                env.step(json.dumps(action))\n",
+    "            steps += 1\n",
+    "        return env.get_reward()\n",
+    "    finally:\n",
+    "        env.close()\n",
+    "\n",
+    "print(\"Running random baseline episodes...\")\n",
+    "random_scores = []\n",
+    "for seed in BASELINE_SEEDS:\n",
+    "    score = run_random_episode(seed)\n",
+    "    random_scores.append(score)\n",
+    "    print(f\"  Seed {seed:>3d}: {score:.4f}\")\n",
+    "\n",
+    "print(f\"\\n{'='*40}\")\n",
+    "print(f\"Random Baseline: {np.mean(random_scores):.4f} +/- {np.std(random_scores):.4f}\")\n",
+    "print(f\"{'='*40}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Build Training Dataset\n",
+    "\n",
+    "For GRPO, we need a dataset of prompts that start each episode.\n",
+    "Each prompt describes the CHCO role and the current scenario."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import Dataset\n",
+    "\n",
+    "SYSTEM_PROMPT = \"\"\"You are the Chief Human Capital Officer (CHCO) of a simulated company.\n",
+    "You manage 300 employees across 5 departments over 6 quarters.\n",
+    "Each quarter has 4 phases: Scanning, Planning, Producing, Controlling.\n",
+    "Maximize HCVA, HCROI, and QIPS metrics by making strategic HR decisions.\n",
+    "\n",
+    "When calling the step tool, provide a JSON action object like:\n",
+    "{\"action_type\": \"query_department\", \"department\": \"Engineering\"}\n",
+    "\n",
+    "Available action types vary by phase:\n",
+    "- Scanning: query_department, query_employees, calculate_metric, review_financials\n",
+    "- Planning: set_hiring_target, set_training_budget, set_compensation_policy, set_retention_program\n",
+    "- Producing: execute_hiring, execute_promotion, execute_transfer, execute_training, execute_termination\n",
+    "- Controlling: submit_report, calculate_metric, advance_quarter\n",
+    "- advance_phase: move to next phase (after minimum actions taken)\n",
+    "\n",
+    "Think step by step. Start by scanning each department.\"\"\"\n",
+    "\n",
+    "# Create prompts with different seeds for diversity\n",
+    "prompts = []\n",
+    "for seed in range(20):\n",
+    "    prompts.append({\n",
+    "        \"prompt\": [\n",
+    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "            {\"role\": \"user\", \"content\": f\"Begin managing the company. Episode seed: {seed}. Start by scanning your organization to understand its current state.\"},\n",
+    "        ]\n",
+    "    })\n",
+    "\n",
+    "dataset = Dataset.from_list(prompts)\n",
+    "print(f\"Training dataset: {len(dataset)} episodes\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4a. Heuristic Baseline\n",
+    "\n",
+    "Run the built-in heuristic agent (from `demo.py`) on the same seeds.\n",
+    "This data-driven strategy scores ~0.50-0.65, representing what a\n",
+    "well-designed non-LLM approach achieves. The trained model should aim\n",
+    "to match or exceed this."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import requests\n",
+    "\n",
+    "def run_heuristic_episode(seed, size=300):\n",
+    "    \"\"\"Run the built-in heuristic agent via the /demo API endpoint.\"\"\"\n",
+    "    try:\n",
+    "        resp = requests.get(\n",
+    "            f\"{HCM21_URL}/demo\",\n",
+    "            params={\"seed\": seed, \"size\": size},\n",
+    "            timeout=120,\n",
+    "        )\n",
+    "        result = resp.json()\n",
+    "        return result.get(\"final_score\", 0.0)\n",
+    "    except Exception as e:\n",
+    "        print(f\"  Error on seed {seed}: {e}\")\n",
+    "        return 0.0\n",
+    "\n",
+    "print(\"Running heuristic baseline episodes...\")\n",
+    "heuristic_scores = []\n",
+    "for seed in BASELINE_SEEDS:\n",
+    "    score = run_heuristic_episode(seed)\n",
+    "    heuristic_scores.append(score)\n",
+    "    print(f\"  Seed {seed:>3d}: {score:.4f}\")\n",
+    "\n",
+    "print(f\"\\n{'='*40}\")\n",
+    "print(f\"Heuristic Baseline: {np.mean(heuristic_scores):.4f} +/- {np.std(heuristic_scores):.4f}\")\n",
+    "print(f\"{'='*40}\")\n",
+    "\n",
+    "print(f\"\\nBaseline Summary:\")\n",
+    "print(f\"  Random:    {np.mean(random_scores):.4f} (floor)\")\n",
+    "print(f\"  Heuristic: {np.mean(heuristic_scores):.4f} (target)\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Define Reward Function\n",
+    "\n",
+    "The reward function uses the environment's built-in scoring.\n",
+    "Quarterly rewards are based on HCVA/HCROI/QIPS improvement.\n",
+    "The final episode score (0-1) is the terminal reward."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def reward_func(completions, environments, **kwargs):\n",
+    "    \"\"\"Extract rewards from the HCM:21 environment after each rollout.\"\"\"\n",
+    "    rewards = []\n",
+    "    for env in environments:\n",
+    "        r = env.get_reward()\n",
+    "        rewards.append(r)\n",
+    "    return rewards\n",
+    "\n",
+    "print(\"Reward function defined.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Configure and Run GRPO Training\n",
+    "\n",
+    "We use TRL's `GRPOTrainer` with `environment_factory=HCM21Env`.\n",
+    "The model generates tool calls to interact with the environment,\n",
+    "and GRPO optimizes the policy based on the environment rewards.\n",
+    "\n",
+    "On a free T4 GPU, we use:\n",
+    "- Small batch size (2) with gradient accumulation\n",
+    "- 4-bit quantization via Unsloth\n",
+    "- Short max_completion_length to fit in memory\n",
+    "\n",
+    "After training, we plot the reward trajectory to visualize learning progress."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from trl import GRPOTrainer, GRPOConfig\n",
+    "\n",
+    "training_args = GRPOConfig(\n",
+    "    output_dir=\"./hcm21-grpo-output\",\n",
+    "    logging_dir=\"./hcm21-grpo-logs\",\n",
+    "    num_train_epochs=1,\n",
+    "    per_device_train_batch_size=2,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    num_generations=4,\n",
+    "    max_completion_length=2048,\n",
+    "    learning_rate=5e-6,\n",
+    "    logging_steps=1,\n",
+    "    save_steps=10,\n",
+    "    bf16=True,\n",
+    "    chat_template_kwargs={\"enable_thinking\": False},\n",
+    ")\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=model,\n",
+    "    processing_class=tokenizer,\n",
+    "    train_dataset=dataset,\n",
+    "    reward_funcs=reward_func,\n",
+    "    environment_factory=HCM21Env,\n",
+    "    args=training_args,\n",
+    ")\n",
+    "\n",
+    "print(\"Trainer initialized. Starting training...\")\n",
+    "trainer.train()\n",
+    "print(\"Training complete!\")\n",
+    "\n",
+    "# Extract and plot reward curve from training logs\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "log_history = trainer.state.log_history\n",
+    "rewards = [entry[\"reward\"] for entry in log_history if \"reward\" in entry]\n",
+    "losses = [entry[\"loss\"] for entry in log_history if \"loss\" in entry]\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "# Reward curve\n",
+    "if rewards:\n",
+    "    axes[0].plot(rewards, alpha=0.4, label=\"Per-step reward\", color=\"steelblue\")\n",
+    "    if len(rewards) >= 3:\n",
+    "        window = min(5, len(rewards))\n",
+    "        rolling = np.convolve(rewards, np.ones(window)/window, mode=\"valid\")\n",
+    "        axes[0].plot(range(window-1, len(rewards)), rolling, linewidth=2,\n",
+    "                     label=f\"{window}-step rolling avg\", color=\"darkblue\")\n",
+    "    axes[0].axhline(y=np.mean(random_scores), color=\"red\", linestyle=\"--\",\n",
+    "                    label=f\"Random baseline ({np.mean(random_scores):.3f})\")\n",
+    "    axes[0].axhline(y=np.mean(heuristic_scores), color=\"green\", linestyle=\"--\",\n",
+    "                    label=f\"Heuristic baseline ({np.mean(heuristic_scores):.3f})\")\n",
+    "    axes[0].set_xlabel(\"Training Step\")\n",
+    "    axes[0].set_ylabel(\"Episode Reward\")\n",
+    "    axes[0].set_title(\"GRPO Training Reward Curve\")\n",
+    "    axes[0].legend()\n",
+    "    axes[0].grid(True, alpha=0.3)\n",
+    "\n",
+    "# Loss curve\n",
+    "if losses:\n",
+    "    axes[1].plot(losses, color=\"orange\", alpha=0.7)\n",
+    "    axes[1].set_xlabel(\"Training Step\")\n",
+    "    axes[1].set_ylabel(\"Loss\")\n",
+    "    axes[1].set_title(\"Training Loss\")\n",
+    "    axes[1].grid(True, alpha=0.3)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"./hcm21-grpo-output/training_curves.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(f\"\\nTraining stats: {len(rewards)} reward entries, mean={np.mean(rewards):.4f}\" if rewards else \"No reward data logged.\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Save & Push Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save locally\n",
+    "model.save_pretrained(\"./hcm21-agent-lora\")\n",
+    "tokenizer.save_pretrained(\"./hcm21-agent-lora\")\n",
+    "\n",
+    "# Optional: push to HF Hub\n",
+    "# model.push_to_hub(\"your-username/hcm21-agent-lora\")\n",
+    "# tokenizer.push_to_hub(\"your-username/hcm21-agent-lora\")\n",
+    "\n",
+    "print(\"Model saved to ./hcm21-agent-lora\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Evaluate the Trained Agent & Compare to Baselines\n",
+    "\n",
+    "Run the trained model on the **same seeds** as the baselines.\n",
+    "This produces a direct before/after comparison showing whether\n",
+    "GRPO training improved the agent's HR decision-making."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import json\n",
+    "\n",
+    "FastLanguageModel.for_inference(model)\n",
+    "\n",
+    "def run_trained_episode(seed, max_steps=200):\n",
+    "    \"\"\"Run one episode with the trained model, return final score.\"\"\"\n",
+    "    env = HCM21Env()\n",
+    "    try:\n",
+    "        obs_text = env.reset(seed=seed)\n",
+    "        messages = [\n",
+    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "            {\"role\": \"user\", \"content\": f\"Begin. Observation:\\n{obs_text}\"},\n",
+    "        ]\n",
+    "        step_num = 0\n",
+    "        while not env._done and step_num < max_steps:\n",
+    "            inputs = tokenizer.apply_chat_template(\n",
+    "                messages, return_tensors=\"pt\", add_generation_prompt=True\n",
+    "            ).to(model.device)\n",
+    "            outputs = model.generate(\n",
+    "                input_ids=inputs, max_new_tokens=256,\n",
+    "                temperature=0.7, do_sample=True,\n",
+    "            )\n",
+    "            response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
+    "            try:\n",
+    "                start = response.index(\"{\")\n",
+    "                depth = 0\n",
+    "                for i, c in enumerate(response[start:], start):\n",
+    "                    if c == \"{\": depth += 1\n",
+    "                    elif c == \"}\": depth -= 1\n",
+    "                    if depth == 0:\n",
+    "                        action_json = response[start:i+1]\n",
+    "                        break\n",
+    "                obs_text = env.step(action_json)\n",
+    "            except (ValueError, json.JSONDecodeError):\n",
+    "                obs_text = env.step(json.dumps({\"action_type\": \"advance_phase\"}))\n",
+    "            messages.append({\"role\": \"assistant\", \"content\": response})\n",
+    "            messages.append({\"role\": \"user\", \"content\": obs_text})\n",
+    "            if len(messages) > 12:\n",
+    "                messages = messages[:2] + messages[-10:]\n",
+    "            step_num += 1\n",
+    "        return env.get_reward(), step_num\n",
+    "    finally:\n",
+    "        env.close()\n",
+    "\n",
+    "print(\"Evaluating trained agent on baseline seeds...\")\n",
+    "trained_scores = []\n",
+    "for seed in BASELINE_SEEDS:\n",
+    "    score, steps = run_trained_episode(seed)\n",
+    "    trained_scores.append(score)\n",
+    "    print(f\"  Seed {seed:>3d}: {score:.4f} ({steps} steps)\")\n",
+    "\n",
+    "sep = \"=\" * 50\n",
+    "print(f\"\\n{sep}\")\n",
+    "print(f\"         RESULTS COMPARISON\")\n",
+    "print(f\"{sep}\")\n",
+    "print(f\"  Random Agent:    {np.mean(random_scores):.4f} +/- {np.std(random_scores):.4f}\")\n",
+    "print(f\"  Heuristic Agent: {np.mean(heuristic_scores):.4f} +/- {np.std(heuristic_scores):.4f}\")\n",
+    "print(f\"  Trained Agent:   {np.mean(trained_scores):.4f} +/- {np.std(trained_scores):.4f}\")\n",
+    "print(f\"{sep}\")\n",
+    "\n",
+    "improvement = np.mean(trained_scores) - np.mean(random_scores)\n",
+    "print(f\"\\n  Improvement over random: +{improvement:.4f} ({improvement/max(np.mean(random_scores), 0.01)*100:.1f}%)\")\n",
+    "if np.mean(trained_scores) > np.mean(random_scores):\n",
+    "    print(\"  >>> Training shows positive improvement! <<<\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Results Summary & Visualization\n",
+    "\n",
+    "Final comparison across all agent types with visualization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import os\n",
+    "\n",
+    "# === Bar chart comparison ===\n",
+    "agents = [\"Random\", \"Heuristic\", \"Trained (GRPO)\"]\n",
+    "means = [np.mean(random_scores), np.mean(heuristic_scores), np.mean(trained_scores)]\n",
+    "stds = [np.std(random_scores), np.std(heuristic_scores), np.std(trained_scores)]\n",
+    "colors = [\"#e74c3c\", \"#f39c12\", \"#2ecc71\"]\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "# Bar chart\n",
+    "bars = axes[0].bar(agents, means, yerr=stds, capsize=8, color=colors, edgecolor=\"black\", alpha=0.85)\n",
+    "axes[0].set_ylabel(\"Final Score (0-1)\")\n",
+    "axes[0].set_title(\"Agent Performance Comparison\")\n",
+    "axes[0].set_ylim(0, 1.0)\n",
+    "axes[0].grid(True, axis=\"y\", alpha=0.3)\n",
+    "for bar, mean in zip(bars, means):\n",
+    "    axes[0].text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.02,\n",
+    "                 f\"{mean:.3f}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\")\n",
+    "\n",
+    "# Per-seed breakdown\n",
+    "x = np.arange(len(BASELINE_SEEDS))\n",
+    "width = 0.25\n",
+    "axes[1].bar(x - width, random_scores, width, label=\"Random\", color=\"#e74c3c\", alpha=0.8)\n",
+    "axes[1].bar(x, heuristic_scores, width, label=\"Heuristic\", color=\"#f39c12\", alpha=0.8)\n",
+    "axes[1].bar(x + width, trained_scores, width, label=\"Trained\", color=\"#2ecc71\", alpha=0.8)\n",
+    "axes[1].set_xlabel(\"Seed\")\n",
+    "axes[1].set_ylabel(\"Final Score\")\n",
+    "axes[1].set_title(\"Per-Seed Score Breakdown\")\n",
+    "axes[1].set_xticks(x)\n",
+    "axes[1].set_xticklabels([str(s) for s in BASELINE_SEEDS])\n",
+    "axes[1].legend()\n",
+    "axes[1].grid(True, axis=\"y\", alpha=0.3)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"./hcm21-grpo-output/results_comparison.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "\n",
+    "# Save results to JSON\n",
+    "os.makedirs(\"./hcm21-grpo-output\", exist_ok=True)\n",
+    "results = {\n",
+    "    \"baselines\": {\n",
+    "        \"random\": {\"mean\": float(np.mean(random_scores)), \"std\": float(np.std(random_scores)),\n",
+    "                    \"per_seed\": {str(s): float(v) for s, v in zip(BASELINE_SEEDS, random_scores)}},\n",
+    "        \"heuristic\": {\"mean\": float(np.mean(heuristic_scores)), \"std\": float(np.std(heuristic_scores)),\n",
+    "                      \"per_seed\": {str(s): float(v) for s, v in zip(BASELINE_SEEDS, heuristic_scores)}},\n",
+    "    },\n",
+    "    \"trained\": {\"mean\": float(np.mean(trained_scores)), \"std\": float(np.std(trained_scores)),\n",
+    "                \"per_seed\": {str(s): float(v) for s, v in zip(BASELINE_SEEDS, trained_scores)}},\n",
+    "    \"improvement_over_random\": float(np.mean(trained_scores) - np.mean(random_scores)),\n",
+    "    \"model\": \"unsloth/Qwen3-0.6B\",\n",
+    "    \"method\": \"GRPO\",\n",
+    "    \"lora_rank\": 16,\n",
+    "}\n",
+    "with open(\"./hcm21-grpo-output/training_results.json\", \"w\") as f:\n",
+    "    json.dump(results, f, indent=2)\n",
+    "print(\"Results saved to ./hcm21-grpo-output/training_results.json\")\n",
+    "print(json.dumps(results, indent=2))"
+   ],
+   "outputs": [],
+   "execution_count": null
+  }
+ ]
+}