Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -173,11 +173,13 @@ _TOOL_CALL_RE = re.compile(
 
 
 def _env_tool_callables(env: Any) -> list[Callable[..., str]]:
-    return [
-        getattr(env, name)
+    # Bound methods so held-out eval can call fn(**kwargs); TRL uses class functions itself.
+    names = [
+        name
         for name, _ in inspect.getmembers(type(env), predicate=inspect.isfunction)
         if name not in ("reset", "get_reward") and not name.startswith("_")
     ]
+    return [getattr(env, name) for name in names]
 
 
 def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
@@ -227,7 +229,8 @@ def run_heldout_agent_eval(
                 messages = [{"role": "user", "content": str(prompt or "")}]
 
             final_text = ""
-            for _turn in range(max_turns):
+            lf = getattr(env, "_lf", None)
+            for turn in range(max_turns):
                 try:
                     prompt_text = tokenizer.apply_chat_template(
                         messages,
@@ -242,19 +245,40 @@ def run_heldout_agent_eval(
                 except Exception:
                     prompt_text = messages[-1]["content"] if messages else ""
 
-                inputs = tokenizer(prompt_text, return_tensors="pt")
-                inputs = {k: v.to(model.device) for k, v in inputs.items()}
-                with torch.no_grad():
-                    out = model.generate(
-                        **inputs,
-                        max_new_tokens=max_new_tokens,
-                        do_sample=False,
-                        pad_token_id=tokenizer.eos_token_id,
-                    )
-                gen = out[0][inputs["input_ids"].shape[-1] :]
-                completion = tokenizer.decode(gen, skip_special_tokens=True)
+                def _generate_once() -> str:
+                    inputs_local = tokenizer(prompt_text, return_tensors="pt")
+                    inputs_local = {k: v.to(model.device) for k, v in inputs_local.items()}
+                    with torch.no_grad():
+                        out_ids = model.generate(
+                            **inputs_local,
+                            max_new_tokens=max_new_tokens,
+                            do_sample=False,
+                            pad_token_id=tokenizer.eos_token_id,
+                        )
+                    gen_ids = out_ids[0][inputs_local["input_ids"].shape[-1] :]
+                    return tokenizer.decode(gen_ids, skip_special_tokens=True)
+
+                if lf is not None and getattr(env, "_root", None) is not None:
+                    with lf.start_as_current_observation(
+                        name=f"llm-turn-{turn}",
+                        as_type="generation",
+                        input={"messages_tail": messages[-2:], "turn": turn},
+                        metadata={"phase": "ckpt_eval", "turn": turn},
+                    ) as gen_span:
+                        completion = _generate_once()
+                        calls = _parse_tool_calls(completion)
+                        gen_span.update(
+                            output={
+                                "completion": completion[:4000],
+                                "n_tool_calls": len(calls),
+                                "tool_names": [c["name"] for c in calls],
+                            }
+                        )
+                else:
+                    completion = _generate_once()
+                    calls = _parse_tool_calls(completion)
+
                 final_text = completion
-                calls = _parse_tool_calls(completion)
                 messages.append({"role": "assistant", "content": completion})
                 if not calls:
                     break
@@ -295,6 +319,11 @@ def run_heldout_agent_eval(
                     "reward_total": reward,
                     "resolved": bool(report.get("resolved")),
                     "resolution": report.get("resolution"),
+                    "n_tool_calls": len(getattr(env, "_tool_calls", []) or []),
+                    "tools_used": [
+                        c.get("name") for c in (getattr(env, "_tool_calls", []) or [])
+                    ],
+                    "langfuse_trace_id": getattr(env, "_trace_id", "") or "",
                 }
             )
         except Exception as e:  # noqa: BLE001
@@ -311,6 +340,10 @@ def run_heldout_agent_eval(
                 }
             )
         finally:
+            try:
+                env._end_trace()
+            except Exception:
+                pass
             try:
                 env._cleanup()
             except Exception:
@@ -373,6 +406,9 @@ def langfuse_log_ckpt_eval(
                         "pass_frac": r.get("pass_frac"),
                         "reward_total": r.get("reward_total"),
                         "resolved": r.get("resolved"),
+                        "n_tool_calls": r.get("n_tool_calls"),
+                        "tools_used": r.get("tools_used"),
+                        "rollout_trace_id": r.get("langfuse_trace_id"),
                     }
                 )
         summary = {