jfl 0.8.1 → 0.9.1

package/scripts/train/v2/benchmark.py

@@ -0,0 +1,661 @@
+"""
+Benchmark Scenarios for PolicyHead Evaluation.
+
+Creates synthetic evaluation environments with known "right answers"
+for rapid PolicyHead iteration. Test in minutes, not weeks.
+
+Drew's recommendation (section 8.6):
+  - Create benchmark agent scenarios for rapid iteration
+  - Use World Model to create synthetic state snapshots
+  - Counterfactual Engine generates ground truth outcomes
+  - PolicyHead evaluated on these scenarios before deploying to real agents
+
+Scenarios:
+  1. Fix failing test — agent sees failing tests, should select fix_bug
+  2. Refactor messy code — high code churn, should select refactor_code
+  3. Optimize slow build — build failing/slow, should select optimize_performance
+  4. Handle hub crash — hub down, agents stranded, should select fix_bug
+  5. Improve coverage — low test counts, should select add_tests
+  6. Stale dependencies — security issues, should select dependency_update
+  7. Missing docs — new features without docs, should select update_docs
+  8. Data pipeline broken — ETL failures, should select data_pipeline
+
+Each scenario defines:
+  - A synthetic WorldState snapshot (state text for embedding)
+  - A goal string
+  - The expected correct tool
+  - Difficulty level (easy/medium/hard)
+  - Variations (to test generalization)
+
+Usage:
+  python benchmark.py --checkpoint .jfl/checkpoints/best_policy_head.pt
+  python benchmark.py --checkpoint .jfl/checkpoints/best_policy_head.pt --json
+  python benchmark.py --generate  # Write benchmark.jsonl for offline eval
+"""
+
+import json
+import os
+import sys
+import argparse
+import time
+from dataclasses import dataclass, asdict
+
+# ============================================================================
+# Scenario Definitions
+# ============================================================================
+
+@dataclass
+class BenchmarkScenario:
+    id: str
+    name: str
+    description: str
+    difficulty: str  # easy, medium, hard
+    state_text: str
+    goal: str
+    correct_tool: str
+    category: str  # diagnostic, scheduling, recovery, optimization
+    tags: list[str]
+
+
+def generate_scenarios() -> list[BenchmarkScenario]:
+    """Generate all benchmark scenarios with variations."""
+    scenarios = []
+
+    # ── Scenario 1: Fix Failing Test (Easy) ──────────────────────────
+    for variant, (tests_failing, composite, goal_text) in enumerate([
+        (3, 0.72, "Fix 3 failing tests in auth module"),
+        (1, 0.91, "Fix flaky test in session-manager.test.ts"),
+        (7, 0.45, "Multiple test failures after refactor — fix regressions"),
+        (1, 0.88, "TypeError in user-service.test.ts line 42"),
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"fix-failing-test-v{variant}",
+            name=f"Fix Failing Test (variant {variant})",
+            description="Agent sees failing tests, should identify fix_bug as the right action",
+            difficulty="easy" if variant < 2 else "medium",
+            state_text="\n".join([
+                f"Agent: test-fixer",
+                f"Composite: {composite:.4f}",
+                f"Tests: {20 - tests_failing}/{20}",
+                f"Trajectory: {variant + 1}",
+                f"Dimensions: test_pass_rate={1 - tests_failing/20:.4f}, build_health=0.9000, code_quality=0.8500",
+                f"Recent deltas: -0.0300, -0.0100",
+            ]),
+            goal=goal_text,
+            correct_tool="fix_bug",
+            category="diagnostic",
+            tags=["tests", "regression", "bug"],
+        ))
+
+    # ── Scenario 2: Refactor Messy Code (Medium) ─────────────────────
+    for variant, (churn, complexity, goal_text) in enumerate([
+        (450, "high", "Reduce complexity in orchestrator module — too many responsibilities"),
+        (200, "medium", "Extract shared logic from peter.ts and eval.ts into utils"),
+        (800, "critical", "Module has 2400 lines, 15 functions over 100 lines — needs decomposition"),
+        (100, "low", "Clean up unused imports and dead code paths in training pipeline"),
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"refactor-messy-code-v{variant}",
+            name=f"Refactor Messy Code (variant {variant})",
+            description="High complexity or code churn indicates refactoring needed",
+            difficulty="medium",
+            state_text="\n".join([
+                f"Agent: code-quality",
+                f"Composite: 0.7200",
+                f"Tests: 45/45",
+                f"Trajectory: {variant + 3}",
+                f"Dimensions: code_quality=0.5500, test_pass_rate=1.0000, build_health=1.0000, code_churn={churn}",
+                f"Recent deltas: +0.0050, +0.0020, -0.0010",
+            ]),
+            goal=goal_text,
+            correct_tool="refactor_code",
+            category="optimization",
+            tags=["complexity", "maintainability", "code-quality"],
+        ))
+
+    # ── Scenario 3: Optimize Slow Build (Medium) ─────────────────────
+    for variant, (build_time, goal_text) in enumerate([
+        ("45s", "Build taking 45s — optimize TypeScript compilation"),
+        ("120s", "CI pipeline runs 2 minutes — find bottleneck"),
+        ("30s", "Bundle size 2.4MB — tree-shake unused dependencies"),
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"optimize-build-v{variant}",
+            name=f"Optimize Slow Build (variant {variant})",
+            description="Build performance degraded, needs optimization",
+            difficulty="medium",
+            state_text="\n".join([
+                f"Agent: build-optimizer",
+                f"Composite: 0.6800",
+                f"Tests: 30/30",
+                f"Trajectory: {variant + 2}",
+                f"Dimensions: build_health=0.4000, test_pass_rate=1.0000, code_quality=0.8000",
+                f"Recent deltas: -0.0200, -0.0150",
+            ]),
+            goal=goal_text,
+            correct_tool="optimize_performance",
+            category="optimization",
+            tags=["build", "performance", "ci"],
+        ))
+
+    # ── Scenario 4: Handle Hub Crash (Hard) ──────────────────────────
+    for variant, (agents_stranded, goal_text) in enumerate([
+        (3, "Hub crashed — 3 agents stranded, need to restore connectivity"),
+        (6, "100% agent stranding — hub process died, restart and recover"),
+        (1, "Hub WebSocket connection dropped — single agent lost progress"),
+        (4, "Hub OOM killed — reduce memory usage and restart with safeguards"),
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"hub-crash-v{variant}",
+            name=f"Handle Hub Crash (variant {variant})",
+            description="Hub crash causing agent stranding — critical infrastructure fix",
+            difficulty="hard",
+            state_text="\n".join([
+                f"Agent: hub-sentinel",
+                f"Composite: 0.3000",
+                f"Tests: 10/15",
+                f"Trajectory: {variant + 5}",
+                f"Dimensions: hub_health=0.0000, agent_stranding={agents_stranded/6:.4f}, build_health=0.5000",
+                f"Recent deltas: -0.1500, -0.2000, -0.0500",
+            ]),
+            goal=goal_text,
+            correct_tool="fix_bug",
+            category="recovery",
+            tags=["hub", "crash", "infrastructure", "critical"],
+        ))
+
+    # ── Scenario 5: Improve Test Coverage (Easy) ─────────────────────
+    for variant, (coverage, goal_text) in enumerate([
+        (0.45, "Test coverage at 45% — add tests for uncovered modules"),
+        (0.60, "New feature shipped without tests — add unit tests"),
+        (0.72, "Coverage dropped after refactor — restore test coverage"),
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"improve-coverage-v{variant}",
+            name=f"Improve Test Coverage (variant {variant})",
+            description="Low test coverage, should add tests",
+            difficulty="easy",
+            state_text="\n".join([
+                f"Agent: test-coverage",
+                f"Composite: {coverage:.4f}",
+                f"Tests: {int(coverage * 50)}/{50}",
+                f"Trajectory: {variant + 1}",
+                f"Dimensions: test_coverage={coverage:.4f}, test_pass_rate=1.0000, code_quality=0.8000",
+                f"Recent deltas: +0.0100, +0.0050",
+            ]),
+            goal=goal_text,
+            correct_tool="add_tests",
+            category="diagnostic",
+            tags=["tests", "coverage", "quality"],
+        ))
+
+    # ── Scenario 6: Stale Dependencies (Medium) ──────────────────────
+    for variant, goal_text in enumerate([
+        "3 critical CVEs in dependencies — update packages",
+        "Node.js version EOL next month — migrate dependencies",
+        "Lockfile conflicts — resolve dependency tree issues",
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"stale-deps-v{variant}",
+            name=f"Stale Dependencies (variant {variant})",
+            description="Security or compatibility issues from outdated dependencies",
+            difficulty="medium",
+            state_text="\n".join([
+                f"Agent: dependency-updater",
+                f"Composite: 0.7500",
+                f"Tests: 40/40",
+                f"Trajectory: {variant + 1}",
+                f"Dimensions: security_score=0.4000, test_pass_rate=1.0000, build_health=0.9000",
+                f"Recent deltas: -0.0050, +0.0020",
+            ]),
+            goal=goal_text,
+            correct_tool="dependency_update",
+            category="optimization",
+            tags=["security", "dependencies", "maintenance"],
+        ))
+
+    # ── Scenario 7: Missing Documentation (Easy) ─────────────────────
+    for variant, goal_text in enumerate([
+        "3 new CLI commands have no documentation — write usage docs",
+        "API changed but README is stale — update docs",
+        "SPEC.md doesn't reflect recent architecture decisions — update",
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"missing-docs-v{variant}",
+            name=f"Missing Documentation (variant {variant})",
+            description="Documentation gaps need filling",
+            difficulty="easy",
+            state_text="\n".join([
+                f"Agent: docs-updater",
+                f"Composite: 0.8200",
+                f"Tests: 35/35",
+                f"Trajectory: {variant + 1}",
+                f"Dimensions: doc_coverage=0.4000, test_pass_rate=1.0000, code_quality=0.9000",
+                f"Recent deltas: +0.0100, +0.0200",
+            ]),
+            goal=goal_text,
+            correct_tool="update_docs",
+            category="diagnostic",
+            tags=["docs", "documentation", "communication"],
+        ))
+
+    # ── Scenario 8: Data Pipeline Broken (Hard) ──────────────────────
+    for variant, goal_text in enumerate([
+        "Training buffer transform failing — v1 to v2 conversion errors",
+        "Embedding precompute script OOM on large dataset — fix pipeline",
+        "Counterfactual generation producing invalid scenarios — debug pipeline",
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"data-pipeline-v{variant}",
+            name=f"Data Pipeline Broken (variant {variant})",
+            description="Data processing pipeline needs fixing",
+            difficulty="hard",
+            state_text="\n".join([
+                f"Agent: data-engineer",
+                f"Composite: 0.5500",
+                f"Tests: 20/25",
+                f"Trajectory: {variant + 4}",
+                f"Dimensions: pipeline_health=0.2000, data_quality=0.4000, test_pass_rate=0.8000",
+                f"Recent deltas: -0.0500, -0.0300, -0.0200",
+            ]),
+            goal=goal_text,
+            correct_tool="data_pipeline",
+            category="recovery",
+            tags=["data", "pipeline", "etl", "training"],
+        ))
+
+    # ── Scenario 9: Add Monitoring (Medium) ──────────────────────────
+    for variant, goal_text in enumerate([
+        "No visibility into agent performance — add telemetry",
+        "Hub crashes go undetected — add health check endpoint",
+        "Training pipeline has no metrics — add loss/accuracy logging",
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"add-monitoring-v{variant}",
+            name=f"Add Monitoring (variant {variant})",
+            description="System lacks observability, needs monitoring",
+            difficulty="medium",
+            state_text="\n".join([
+                f"Agent: observability",
+                f"Composite: 0.7000",
+                f"Tests: 30/30",
+                f"Trajectory: {variant + 2}",
+                f"Dimensions: observability=0.2000, test_pass_rate=1.0000, build_health=0.8500",
+                f"Recent deltas: +0.0050, +0.0030",
+            ]),
+            goal=goal_text,
+            correct_tool="add_monitoring",
+            category="optimization",
+            tags=["monitoring", "observability", "telemetry"],
+        ))
+
+    # ── Scenario 10: New Feature Request (Hard — ambiguous) ──────────
+    for variant, goal_text in enumerate([
+        "Add planning loop that connects PolicyHead to DynamicsModel",
+        "Implement experience replay buffer for online learning",
+        "Build counterfactual training bridge for synthetic data generation",
+    ]):
+        scenarios.append(BenchmarkScenario(
+            id=f"new-feature-v{variant}",
+            name=f"New Feature Request (variant {variant})",
+            description="Clear feature request — should select add_feature",
+            difficulty="hard",
+            state_text="\n".join([
+                f"Agent: feature-builder",
+                f"Composite: 0.8000",
+                f"Tests: 40/40",
+                f"Trajectory: {variant + 1}",
+                f"Dimensions: feature_completeness=0.6000, test_pass_rate=1.0000, code_quality=0.8500",
+                f"Recent deltas: +0.0200, +0.0100, +0.0150",
+            ]),
+            goal=goal_text,
+            correct_tool="add_feature",
+            category="diagnostic",
+            tags=["feature", "implementation", "new"],
+        ))
+
+    return scenarios
+
+
+# ============================================================================
+# Benchmark Runner
+# ============================================================================
+
+def run_benchmark(checkpoint_path: str, scenarios: list[BenchmarkScenario], device: str = "cpu") -> dict:
+    """Run all scenarios against a PolicyHead checkpoint."""
+    import torch
+    from model import PolicyHead
+    from dataset import load_embedding_cache
+
+    # Load model
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    config = ckpt["config"]
+    tool_to_index = ckpt["tool_to_index"]
+    index_to_tool = {v: k for k, v in tool_to_index.items()}
+
+    model = PolicyHead(
+        embedding_dim=config["embedding_dim"],
+        hidden_dim=config["hidden_dim"],
+        num_tools=ckpt["num_tools"],
+        num_layers=config["num_layers"],
+        num_heads=config["num_heads"],
+        dropout=config.get("dropout", 0.1),
+    ).to(device)
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+
+    # Load embeddings
+    data_dir = os.path.dirname(checkpoint_path).replace("checkpoints", "v2-data")
+    embeddings_matrix, text_to_idx = load_embedding_cache(data_dir)
+
+    if embeddings_matrix is None:
+        print("WARNING: No embedding cache — using zero vectors")
+        embedding_dim = config["embedding_dim"]
+    else:
+        embedding_dim = embeddings_matrix.shape[1]
+
+    # Run each scenario
+    results = []
+    for scenario in scenarios:
+        result = evaluate_scenario(
+            model, scenario, tool_to_index, index_to_tool,
+            embeddings_matrix, text_to_idx, embedding_dim, device
+        )
+        results.append(result)
+
+    return aggregate_results(results, scenarios)
+
+
+def evaluate_scenario(
+    model, scenario, tool_to_index, index_to_tool,
+    embeddings_matrix, text_to_idx, embedding_dim, device
+) -> dict:
+    """Evaluate a single scenario."""
+    import torch
+
+    # Get embeddings for state and goal
+    state_emb = get_embedding(scenario.state_text, embeddings_matrix, text_to_idx, embedding_dim)
+    goal_emb = get_embedding(scenario.goal, embeddings_matrix, text_to_idx, embedding_dim)
+
+    state_tensor = torch.tensor(state_emb, dtype=torch.float32).unsqueeze(0).to(device)
+    goal_tensor = torch.tensor(goal_emb, dtype=torch.float32).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        result = model.predict(state_tensor, goal_tensor, top_k=5)
+
+    top_indices = result["top_k_indices"][0].cpu().tolist()
+    top_probs = result["top_k_probs"][0].cpu().tolist()
+
+    predicted_tool = index_to_tool.get(top_indices[0], "unknown")
+    correct = predicted_tool == scenario.correct_tool
+    correct_tool_index = tool_to_index.get(scenario.correct_tool)
+
+    # Check if correct tool is in top-3
+    in_top3 = correct_tool_index in top_indices[:3] if correct_tool_index is not None else False
+    in_top5 = correct_tool_index in top_indices[:5] if correct_tool_index is not None else False
+
+    # Get rank of correct tool
+    correct_rank = -1
+    if correct_tool_index is not None:
+        all_probs = result["all_probs"][0].cpu().tolist()
+        sorted_indices = sorted(range(len(all_probs)), key=lambda i: -all_probs[i])
+        for rank, idx in enumerate(sorted_indices):
+            if idx == correct_tool_index:
+                correct_rank = rank + 1
+                break
+
+    return {
+        "scenario_id": scenario.id,
+        "correct_tool": scenario.correct_tool,
+        "predicted_tool": predicted_tool,
+        "correct": correct,
+        "in_top3": in_top3,
+        "in_top5": in_top5,
+        "correct_rank": correct_rank,
+        "confidence": top_probs[0],
+        "difficulty": scenario.difficulty,
+        "category": scenario.category,
+        "tags": scenario.tags,
+        "top5": [
+            {"tool": index_to_tool.get(idx, "?"), "prob": prob}
+            for idx, prob in zip(top_indices, top_probs)
+        ],
+    }
+
+
+def get_embedding(text: str, embeddings_matrix, text_to_idx, embedding_dim: int):
+    """Look up or generate zero embedding for text."""
+    import numpy as np
+
+    if text_to_idx and text in text_to_idx:
+        idx = text_to_idx[text]
+        return embeddings_matrix[idx].tolist()
+
+    # Fallback: zero vector (will degrade accuracy but won't crash)
+    return [0.0] * embedding_dim
+
+
+def aggregate_results(results: list[dict], scenarios: list[BenchmarkScenario]) -> dict:
+    """Aggregate individual scenario results into a report."""
+    total = len(results)
+    correct = sum(1 for r in results if r["correct"])
+    in_top3 = sum(1 for r in results if r["in_top3"])
+    in_top5 = sum(1 for r in results if r["in_top5"])
+
+    # By difficulty
+    by_difficulty = {}
+    for diff in ["easy", "medium", "hard"]:
+        subset = [r for r in results if r["difficulty"] == diff]
+        if subset:
+            by_difficulty[diff] = {
+                "total": len(subset),
+                "correct": sum(1 for r in subset if r["correct"]),
+                "accuracy": sum(1 for r in subset if r["correct"]) / len(subset),
+                "top3": sum(1 for r in subset if r["in_top3"]) / len(subset),
+            }
+
+    # By category
+    by_category = {}
+    for cat in set(r["category"] for r in results):
+        subset = [r for r in results if r["category"] == cat]
+        by_category[cat] = {
+            "total": len(subset),
+            "correct": sum(1 for r in subset if r["correct"]),
+            "accuracy": sum(1 for r in subset if r["correct"]) / len(subset),
+        }
+
+    # By correct_tool
+    by_tool = {}
+    for tool in set(r["correct_tool"] for r in results):
+        subset = [r for r in results if r["correct_tool"] == tool]
+        by_tool[tool] = {
+            "total": len(subset),
+            "correct": sum(1 for r in subset if r["correct"]),
+            "accuracy": sum(1 for r in subset if r["correct"]) / len(subset),
+        }
+
+    # Failures for debugging
+    failures = [r for r in results if not r["correct"]]
+
+    return {
+        "summary": {
+            "total_scenarios": total,
+            "top1_accuracy": correct / total if total > 0 else 0,
+            "top3_accuracy": in_top3 / total if total > 0 else 0,
+            "top5_accuracy": in_top5 / total if total > 0 else 0,
+            "correct": correct,
+        },
+        "by_difficulty": by_difficulty,
+        "by_category": by_category,
+        "by_tool": by_tool,
+        "failures": [
+            {
+                "id": f["scenario_id"],
+                "expected": f["correct_tool"],
+                "got": f["predicted_tool"],
+                "rank": f["correct_rank"],
+                "confidence": f["confidence"],
+                "top5": f["top5"],
+            }
+            for f in failures
+        ],
+        "all_results": results,
+    }
+
+
+# ============================================================================
+# Output
+# ============================================================================
+
+def print_report(report: dict):
+    """Print a human-readable benchmark report."""
+    s = report["summary"]
+
+    print("=" * 70)
+    print("  BENCHMARK EVALUATION REPORT")
+    print("=" * 70)
+    print(f"\n  Scenarios: {s['total_scenarios']}")
+    print(f"  Top-1 Accuracy: {s['top1_accuracy']:.1%} ({s['correct']}/{s['total_scenarios']})")
+    print(f"  Top-3 Accuracy: {s['top3_accuracy']:.1%}")
+    print(f"  Top-5 Accuracy: {s['top5_accuracy']:.1%}")
+
+    print(f"\n  {'Difficulty':<12} {'Acc':>8} {'Top-3':>8} {'N':>5}")
+    print("  " + "-" * 35)
+    for diff in ["easy", "medium", "hard"]:
+        if diff in report["by_difficulty"]:
+            d = report["by_difficulty"][diff]
+            print(f"  {diff:<12} {d['accuracy']:>7.1%} {d['top3']:>7.1%} {d['total']:>5}")
+
+    print(f"\n  {'Category':<18} {'Acc':>8} {'N':>5}")
+    print("  " + "-" * 33)
+    for cat, stats in sorted(report["by_category"].items()):
+        print(f"  {cat:<18} {stats['accuracy']:>7.1%} {stats['total']:>5}")
+
+    print(f"\n  {'Tool':<25} {'Acc':>8} {'N':>5}")
+    print("  " + "-" * 40)
+    for tool, stats in sorted(report["by_tool"].items(), key=lambda x: -x[1]["accuracy"]):
+        bar = "█" * int(stats["accuracy"] * 10) + "░" * (10 - int(stats["accuracy"] * 10))
+        print(f"  {tool:<25} {stats['accuracy']:>7.1%} {stats['total']:>5}  {bar}")
+
+    if report["failures"]:
+        print(f"\n  Failures ({len(report['failures'])}):")
+        print(f"  {'ID':<30} {'Expected':<20} {'Got':<20} {'Rank':>5}")
+        print("  " + "-" * 78)
+        for f in report["failures"]:
+            print(f"  {f['id']:<30} {f['expected']:<20} {f['got']:<20} {f['rank']:>5}")
+
+    # Overall grade
+    acc = s["top1_accuracy"]
+    if acc >= 0.90:
+        grade = "A — Ready for production"
+    elif acc >= 0.80:
+        grade = "B — Good, minor gaps"
+    elif acc >= 0.70:
+        grade = "C — Needs improvement"
+    elif acc >= 0.50:
+        grade = "D — Significant gaps"
+    else:
+        grade = "F — Major retraining needed"
+
+    print(f"\n  Grade: {grade}")
+    print("=" * 70)
+
+
+def generate_benchmark_jsonl(scenarios: list[BenchmarkScenario], output_path: str):
+    """Write benchmark scenarios as JSONL for use with eval.py."""
+    with open(output_path, "w") as f:
+        for s in scenarios:
+            entry = {
+                "current_state": s.state_text,
+                "goal": s.goal,
+                "correct_tool": s.correct_tool,
+                "source": "benchmark",
+                "scenario_id": s.id,
+                "difficulty": s.difficulty,
+                "category": s.category,
+            }
+            f.write(json.dumps(entry) + "\n")
+    print(f"Wrote {len(scenarios)} scenarios to {output_path}")
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark scenarios for PolicyHead evaluation")
+    parser.add_argument("--checkpoint", default=None, help="Path to PolicyHead checkpoint")
+    parser.add_argument("--generate", action="store_true", help="Generate benchmark.jsonl for offline eval")
+    parser.add_argument("--output", default=".jfl/v2-data/benchmark.jsonl", help="Output path for generated JSONL")
+    parser.add_argument("--json", action="store_true", help="Output results as JSON")
+    parser.add_argument("--difficulty", default=None, help="Filter by difficulty: easy/medium/hard")
+    parser.add_argument("--category", default=None, help="Filter by category: diagnostic/scheduling/recovery/optimization")
+    args = parser.parse_args()
+
+    scenarios = generate_scenarios()
+
+    # Apply filters
+    if args.difficulty:
+        scenarios = [s for s in scenarios if s.difficulty == args.difficulty]
+    if args.category:
+        scenarios = [s for s in scenarios if s.category == args.category]
+
+    print(f"Loaded {len(scenarios)} benchmark scenarios")
+
+    if args.generate:
+        generate_benchmark_jsonl(scenarios, args.output)
+        return
+
+    if args.checkpoint is None:
+        # Try default
+        default_ckpt = ".jfl/checkpoints/best_policy_head.pt"
+        if os.path.exists(default_ckpt):
+            args.checkpoint = default_ckpt
+        else:
+            print("No checkpoint specified. Use --checkpoint or --generate")
+            print("\n  Available commands:")
+            print("    python benchmark.py --generate                     # Create benchmark.jsonl")
+            print("    python benchmark.py --checkpoint path/to/model.pt  # Run evaluation")
+            sys.exit(1)
+
+    if not os.path.exists(args.checkpoint):
+        print(f"Checkpoint not found: {args.checkpoint}")
+        sys.exit(1)
+
+    # Detect device
+    import torch
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    print(f"Device: {device}")
+    print(f"Checkpoint: {args.checkpoint}")
+    print()
+
+    t0 = time.time()
+    report = run_benchmark(args.checkpoint, scenarios, device)
+    elapsed = time.time() - t0
+
+    if args.json:
+        # Remove all_results for cleaner JSON output
+        output = {k: v for k, v in report.items() if k != "all_results"}
+        print(json.dumps(output, indent=2))
+    else:
+        print_report(report)
+        print(f"\n  Benchmark completed in {elapsed:.1f}s")
+
+    # Save results
+    results_path = os.path.join(
+        os.path.dirname(args.checkpoint),
+        "benchmark-results.json"
+    )
+    with open(results_path, "w") as f:
+        json.dump(report, f, indent=2)
+    print(f"  Results saved to {results_path}")
+
+
+if __name__ == "__main__":
+    main()