Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -437,22 +437,26 @@ def langfuse_log_ckpt_eval(
     n_perfect = sum(1 for r in results if float(r.get("pass_frac") or 0) >= 1.0)
     n_resolved = sum(1 for r in results if r.get("resolved"))
 
+    # Own root trace (do not nest under / swallow per-rollout traces).
+    agg_trace_id = lf.create_trace_id()
     with lf.start_as_current_observation(
-        name="kaggle-ckpt-eval",
+        name=f"kaggle-ckpt-eval:step-{step}",
         as_type="evaluator",
+        trace_context={"trace_id": agg_trace_id},
         metadata={
             "step": step,
             "hf_revision": hf_revision,
             "train_model_id": train_model_id,
             "wandb_run_id": wandb_run_id,
-            "phase": "ckpt_eval",
+            "phase": "ckpt_eval_aggregate",
             "harness": "deepagents-swe",
         },
         input={"step": step, "hf_revision": hf_revision, "n": n},
+        end_on_exit=True,
     ) as root:
         for r in results:
             with lf.start_as_current_observation(
-                name=f"task:{r['task_id']}",
+                name=f"ckpt-summary:{r['task_id']}",
                 as_type="span",
                 input={"prompt": str(r.get("prompt", ""))[:2000]},
             ) as span:
@@ -482,13 +486,59 @@ def langfuse_log_ckpt_eval(
         lf.score_current_trace(name="n_perfect", value=float(n_perfect))
         lf.score_current_trace(name="n_resolved", value=float(n_resolved))
         lf.score_current_trace(name="checkpoint_step", value=float(step))
-        trace_id = lf.get_current_trace_id() or ""
     lf.flush()
-    print(
-        f"[langfuse] ckpt-eval step={step} rev={hf_revision} "
-        f"mean_reward={mean_reward:.3f} resolved={n_resolved}/{n} trace={trace_id}"
+    _vlog(
+        "langfuse.ckpt_eval",
+        step=step,
+        rev=hf_revision,
+        mean_reward=mean_reward,
+        resolved=n_resolved,
+        n=n,
+        trace_id=agg_trace_id,
     )
-    return {"langfuse_trace_id": trace_id, **summary}
+    return {"langfuse_trace_id": agg_trace_id, **summary}
+
+
+def log_completions_reward(completions: list[Any], **kwargs: Any) -> list[float]:
+    """Side-channel: log raw model completions to W&B/console (reward always 0)."""
+    previews = []
+    for i, c in enumerate(completions):
+        if isinstance(c, list):
+            # conversational messages
+            bits = []
+            for m in c[-4:]:
+                if isinstance(m, dict):
+                    role = m.get("role")
+                    content = str(m.get("content") or "")[:400]
+                    tcalls = m.get("tool_calls")
+                    bits.append(f"{role}: {content}" + (f" tool_calls={tcalls}" if tcalls else ""))
+                else:
+                    bits.append(str(m)[:400])
+            preview = " | ".join(bits)
+        else:
+            preview = str(c)[:800]
+        previews.append(preview)
+        _vlog("train.completion", i=i, preview=preview[:500], has_tool_calls=("tool_calls" in preview))
+
+    try:
+        import wandb
+
+        if wandb.run is not None:
+            n_with_tools = sum(1 for p in previews if "tool_calls" in p)
+            table = wandb.Table(columns=["i", "preview", "has_tool_calls"])
+            for i, p in enumerate(previews[:16]):
+                table.add_data(i, p[:1500], "tool_calls" in p)
+            wandb.log(
+                {
+                    "train/completion_previews": table,
+                    "train/n_completions": len(previews),
+                    "train/n_completions_with_tool_calls_str": n_with_tools,
+                }
+            )
+    except Exception as e:  # noqa: BLE001
+        _vlog("train.completion.wandb_failed", error=str(e))
+
+    return [0.0] * len(completions)
 
 
 # ---------------------------------------------------------------------------
@@ -944,6 +994,8 @@ def main(cfg: Cfg | None = None) -> None:
         bf16=use_bf16,
         fp16=use_fp16,
         remove_unused_columns=False,
+        # log_completions_reward (0) + DeepAgentSWEEnv.get_reward (1)
+        reward_weights=[0.0, 1.0],
         hub_model_id=cfg.hub_model_id,
         hub_token=os.environ["HF_TOKEN"],
         push_to_hub=False,
@@ -967,6 +1019,7 @@ def main(cfg: Cfg | None = None) -> None:
     )
     trainer = GRPOTrainer(
         model=model,
+        reward_funcs=[log_completions_reward],  # logs only; env get_reward still scores SWE
         args=args,
         train_dataset=train,
         processing_class=tokenizer,