Fix format_comparison_metrics_html to accept run_comparison() dict directly

- metrics.py: format_comparison_metrics_html() now accepts either two
pre-computed metrics dicts (original API) or the single combined dict
returned by run_comparison(). When called with one arg, it auto-computes
metrics from each episode's log via compute_episode_metrics().
- colab_training.ipynb: Fix broken import paths in cells 5 and 11.
NexusEnv.train does not exist as a package — train.py lives at the repo
root, which is on sys.path after the git clone + sys.path.insert(). Both
cells now use `from train import ...` so they resolve correctly in Colab.

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

sentinelops_arena/metrics.py

@@ -488,16 +488,33 @@ def format_metrics_html(metrics: dict[str, Any]) -> str:
 
 def format_comparison_metrics_html(
     untrained_metrics: dict[str, Any],
-    trained_metrics: dict[str, Any],
+    trained_metrics: dict[str, Any] | None = None,
 ) -> str:
     """Render untrained vs. trained metrics side-by-side with diff indicators.
 
+    Accepts two calling conventions:
+
+    1. **Two metrics dicts** (original API)::
+
+           format_comparison_metrics_html(untrained_metrics, trained_metrics)
+
+    2. **Combined results dict** (output of ``run_comparison()``)::
+
+           format_comparison_metrics_html(run_comparison(seed=42))
+
+       The combined dict must have the shape
+       ``{"untrained": {"log": [...], ...}, "trained": {"log": [...], ...}}``.
+       Episode metrics are computed automatically from each log.
+
     Parameters
     ----------
     untrained_metrics : dict
-        Metrics from the untrained (baseline) episode.
-    trained_metrics : dict
-        Metrics from the trained episode.
+        Either a pre-computed metrics dict (original API) **or** the combined
+        comparison results dict returned by ``run_comparison()``.
+    trained_metrics : dict, optional
+        Pre-computed metrics for the trained episode.  Required when
+        ``untrained_metrics`` is a plain metrics dict.  Omit when passing the
+        combined results dict.
 
     Returns
     -------
@@ -505,6 +522,16 @@ def format_comparison_metrics_html(
         Self-contained HTML snippet showing both metric sets with arrows
         indicating improvement (green) or regression (red).
     """
+    # --- Handle combined results dict from run_comparison() ---
+    if trained_metrics is None:
+        combined = untrained_metrics
+        if "untrained" not in combined or "trained" not in combined:
+            raise ValueError(
+                "When called with a single argument, the dict must contain "
+                "'untrained' and 'trained' keys (output of run_comparison())."
+            )
+        untrained_metrics = compute_episode_metrics(combined["untrained"]["log"])
+        trained_metrics = compute_episode_metrics(combined["trained"]["log"])
 
     def _diff_indicator(
         before: float,

training/colab_training.ipynb

@@ -61,7 +61,7 @@
     "id": "collect-data"
    },
    "outputs": [],
-   "source": "from NexusEnv.train import collect_multi_agent_data, build_training_dataset\nfrom NexusEnv.train import WORKER_SYSTEM_PROMPT, ATTACKER_SYSTEM_PROMPT, OVERSIGHT_SYSTEM_PROMPT\nfrom NexusEnv.train import AGENT_CONFIGS\n\n# Run a single episode and show stats for each agent\nfor role in [\"worker\", \"attacker\", \"oversight\"]:\n    data = collect_multi_agent_data(seed=42, target_agent=role)\n    avg_r = sum(d[\"reward\"] for d in data) / max(len(data), 1)\n    print(f\"{role:>10}: {len(data)} turns, avg_reward={avg_r:.3f}\")"
+   "source": "from train import collect_multi_agent_data, build_training_dataset\nfrom train import WORKER_SYSTEM_PROMPT, ATTACKER_SYSTEM_PROMPT, OVERSIGHT_SYSTEM_PROMPT\nfrom train import AGENT_CONFIGS\n\n# Run a single episode and show stats for each agent\nfor role in [\"worker\", \"attacker\", \"oversight\"]:\n    data = collect_multi_agent_data(seed=42, target_agent=role)\n    avg_r = sum(d[\"reward\"] for d in data) / max(len(data), 1)\n    print(f\"{role:>10}: {len(data)} turns, avg_reward={avg_r:.3f}\")"
   },
   {
    "cell_type": "markdown",
@@ -109,7 +109,7 @@
     "id": "save"
    },
    "outputs": [],
-   "source": "from trl import GRPOConfig, GRPOTrainer\nfrom NexusEnv.train import make_reward_function\n\n# Environment-executing reward function\nreward_fn = make_reward_function(TARGET_AGENT)\n\ngrpo_config = GRPOConfig(\n    output_dir=f\"./sentinelops-grpo-{TARGET_AGENT}\",\n    max_steps=300,                      # Official recommendation\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_generations=2,                   # GRPO group size\n    max_completion_length=256,\n    max_prompt_length=512,\n    learning_rate=5e-5,                  # Official reference: 5e-5\n    temperature=1.0,                     # Official reference: 1.0\n    logging_steps=1,\n    save_steps=50,\n    report_to=\"none\",\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=[reward_fn],\n    args=grpo_config,\n    train_dataset=train_dataset,\n)\n\nprint(f\"Starting GRPO training for {TARGET_AGENT}...\")\nprint(f\"  max_steps={grpo_config.max_steps}, lr={grpo_config.learning_rate}\")\nprint(f\"  num_generations={grpo_config.num_generations}, temp={grpo_config.temperature}\")\ntrainer.train()"
+   "source": "from trl import GRPOConfig, GRPOTrainer\nfrom train import make_reward_function\n\n# Environment-executing reward function\nreward_fn = make_reward_function(TARGET_AGENT)\n\ngrpo_config = GRPOConfig(\n    output_dir=f\"./sentinelops-grpo-{TARGET_AGENT}\",\n    max_steps=300,                      # Official recommendation\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_generations=2,                   # GRPO group size\n    max_completion_length=256,\n    max_prompt_length=512,\n    learning_rate=5e-5,                  # Official reference: 5e-5\n    temperature=1.0,                     # Official reference: 1.0\n    logging_steps=1,\n    save_steps=50,\n    report_to=\"none\",\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=[reward_fn],\n    args=grpo_config,\n    train_dataset=train_dataset,\n)\n\nprint(f\"Starting GRPO training for {TARGET_AGENT}...\")\nprint(f\"  max_steps={grpo_config.max_steps}, lr={grpo_config.learning_rate}\")\nprint(f\"  num_generations={grpo_config.num_generations}, temp={grpo_config.temperature}\")\ntrainer.train()"
   },
   {
    "cell_type": "markdown",