Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -48,6 +48,7 @@ def _ensure_package() -> None:
         token=token,
         allow_patterns=["src/deepagents_rl/**"],
         local_dir=str(root),
+        force_download=True,
     )
     pkg_parent = root / "src"
     if str(pkg_parent) not in sys.path:
@@ -498,13 +499,45 @@ def make_eval_before_save_callback(
                 if wandb.run is not None:
                     wandb_run_id = wandb.run.id
                     mean_r = sum(float(r["reward_total"]) for r in results) / max(len(results), 1)
+                    mean_tools = sum(float(r.get("n_tool_calls") or 0) for r in results) / max(
+                        len(results), 1
+                    )
+                    table = wandb.Table(
+                        columns=[
+                            "task_id",
+                            "reward_total",
+                            "pass_frac",
+                            "resolved",
+                            "resolution",
+                            "n_tool_calls",
+                            "tools_used",
+                            "langfuse_trace_id",
+                            "completion",
+                        ]
+                    )
+                    for r in results:
+                        table.add_data(
+                            str(r.get("task_id") or ""),
+                            float(r.get("reward_total") or 0),
+                            float(r.get("pass_frac") or 0),
+                            bool(r.get("resolved")),
+                            str(r.get("resolution") or ""),
+                            int(r.get("n_tool_calls") or 0),
+                            ",".join(r.get("tools_used") or []),
+                            str(r.get("langfuse_trace_id") or ""),
+                            str(r.get("completion") or "")[:1500],
+                        )
                     wandb.log(
                         {
                             "eval/mean_reward_total": mean_r,
-                            "eval/mean_pass_frac": sum(float(r["pass_frac"]) for r in results)
+                            "eval/mean_pass_frac": sum(
+                                float(r["pass_frac"]) for r in results
+                            )
                             / max(len(results), 1),
                             "eval/n_resolved": sum(1 for r in results if r.get("resolved")),
+                            "eval/mean_n_tool_calls": mean_tools,
                             "eval/checkpoint_step": step,
+                            "eval/episodes": table,
                         },
                         step=step,
                     )
@@ -742,12 +775,40 @@ def main(cfg: Cfg | None = None) -> None:
     Path(output_dir, "ckpt_eval_step_0.json").write_text(
         json.dumps(base_summary | {"results": baseline}, indent=2), encoding="utf-8"
     )
+    base_table = wandb.Table(
+        columns=[
+            "task_id",
+            "reward_total",
+            "pass_frac",
+            "resolved",
+            "resolution",
+            "n_tool_calls",
+            "tools_used",
+            "langfuse_trace_id",
+            "completion",
+        ]
+    )
+    for r in baseline:
+        base_table.add_data(
+            str(r.get("task_id") or ""),
+            float(r.get("reward_total") or 0),
+            float(r.get("pass_frac") or 0),
+            bool(r.get("resolved")),
+            str(r.get("resolution") or ""),
+            int(r.get("n_tool_calls") or 0),
+            ",".join(r.get("tools_used") or []),
+            str(r.get("langfuse_trace_id") or ""),
+            str(r.get("completion") or "")[:1500],
+        )
     wandb.log(
         {
             "eval/mean_pass_frac": base_summary["mean_pass_frac"],
             "eval/mean_reward_total": base_summary["mean_reward_total"],
             "eval/n_resolved": base_summary.get("n_resolved", 0),
+            "eval/mean_n_tool_calls": sum(float(r.get("n_tool_calls") or 0) for r in baseline)
+            / max(len(baseline), 1),
             "eval/checkpoint_step": 0,
+            "eval/episodes": base_table,
             "langfuse/ckpt_trace_id": base_summary.get("langfuse_trace_id"),
         },
         step=0,