Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -1,89 +1,58 @@
 """
-Kaggle GRPO trainer for deepagents-rl coding.
+Kaggle agent GRPO: Qwen + deepagents SWE harness + remote Docker reward.
 
 Hard requirements:
   - W&B metrics every step
-  - BEFORE every checkpoint: held-out eval -> Langfuse scores/trace -> THEN save/push HF
+  - BEFORE every checkpoint: held-out agent eval -> Langfuse -> THEN save/push HF
   - HF adapter revisions tagged grpo-step-{n}
+  - SWE_GRADER_URL must reach local scripts/swe_grader_server.py (tunnel)
 
-Secrets (Kaggle User Secrets OR env):
+Secrets (dataset / User Secrets / env):
   HF_TOKEN, WANDB_API_KEY, LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY
-  optional LANGFUSE_HOST / LANGFUSE_BASE_URL
+  SWE_GRADER_URL, optional SWE_GRADER_TOKEN
 """
 from __future__ import annotations
 
+import inspect
 import json
 import os
 import re
-import subprocess
 import sys
 import traceback
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 
 # ---------------------------------------------------------------------------
-# Inline reward helpers (also shipped in package; duplicated for notebook isolation)
+# Bootstrap package (notebook may only download this file)
 # ---------------------------------------------------------------------------
 
 
-def extract_python_code(completion: str) -> str:
-    if "```python" in completion:
-        start = completion.find("```python") + len("```python")
-        end = completion.find("```", start)
-        return completion[start:end if end != -1 else None].strip()
-    if "```" in completion:
-        start = completion.find("```") + 3
-        end = completion.find("```", start)
-        if end != -1:
-            return completion[start:end].strip()
-    return completion.strip()
-
-
-def run_hidden_tests(solution_code: str, hidden_tests: list[str], timeout_s: float = 5.0) -> dict[str, Any]:
-    per_test = []
-    passed = 0
-    for i, assertion in enumerate(hidden_tests):
-        script = f"{solution_code}\n{assertion}\n"
-        try:
-            proc = subprocess.run(
-                [sys.executable, "-c", script],
-                capture_output=True,
-                text=True,
-                timeout=timeout_s,
-            )
-            ok = proc.returncode == 0
-            err = (proc.stderr or "")[:300]
-        except Exception as e:  # noqa: BLE001
-            ok = False
-            err = f"{type(e).__name__}: {e}"
-        per_test.append({"i": i, "ok": ok, "stderr": err})
-        if ok:
-            passed += 1
-    total = len(hidden_tests)
-    return {
-        "pass_frac": (passed / total) if total else 0.0,
-        "passed": passed,
-        "total": total,
-        "per_test": per_test,
-    }
-
-
-def grade_completion(completion: str, hidden_tests: list[str]) -> dict[str, Any]:
-    code = extract_python_code(completion)
-    g = run_hidden_tests(code, hidden_tests)
-    g["extracted_code"] = code
-    return g
-
-
-def coding_reward(completions: list[str], hidden_tests: list[Any], **kwargs: Any) -> list[float]:
-    """TRL GRPO reward: fraction of hidden unit tests passed."""
-    rewards = []
-    for completion, tests in zip(completions, hidden_tests):
-        if isinstance(tests, str):
-            tests = json.loads(tests)
-        g = grade_completion(completion, list(tests))
-        rewards.append(float(g["pass_frac"]))
-    return rewards
+def _ensure_package() -> None:
+    try:
+        import deepagents_rl  # noqa: F401
+
+        return
+    except ImportError:
+        pass
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise RuntimeError("HF_TOKEN required to download deepagents_rl sources")
+
+    from huggingface_hub import snapshot_download
+
+    root = Path("/kaggle/working/deepagents_rl_src")
+    snapshot_download(
+        repo_id=os.getenv("HF_REPO_ID", "vanishingradient/deepagents-rl-coding"),
+        repo_type="model",
+        token=token,
+        allow_patterns=["src/deepagents_rl/**"],
+        local_dir=str(root),
+    )
+    pkg_parent = root / "src"
+    if str(pkg_parent) not in sys.path:
+        sys.path.insert(0, str(pkg_parent))
+    import deepagents_rl  # noqa: F401
 
 
 # ---------------------------------------------------------------------------
@@ -103,6 +72,9 @@ def load_secrets() -> None:
             "LANGFUSE_SECRET_KEY",
             "LANGFUSE_HOST",
             "LANGFUSE_BASE_URL",
+            "SWE_GRADER_URL",
+            "SWE_GRADER_TOKEN",
+            "WANDB_PROJECT",
         ):
             try:
                 val = us.get_secret(key)
@@ -113,6 +85,22 @@ def load_secrets() -> None:
     except Exception:
         pass
 
+    # Private dataset secrets.json (API-friendly path)
+    secret_candidates = [Path("/kaggle/input/deepagents-rl-secrets/secrets.json")]
+    input_root = Path("/kaggle/input")
+    if input_root.exists():
+        secret_candidates.extend(input_root.rglob("secrets.json"))
+    for p in secret_candidates:
+        if p.exists():
+            try:
+                data = json.loads(p.read_text(encoding="utf-8"))
+                for k, v in data.items():
+                    if v and not os.getenv(k):
+                        os.environ[k] = str(v)
+            except Exception:
+                pass
+            break
+
     missing = [
         k
         for k in (
@@ -120,18 +108,20 @@ def load_secrets() -> None:
             "WANDB_API_KEY",
             "LANGFUSE_PUBLIC_KEY",
             "LANGFUSE_SECRET_KEY",
+            "SWE_GRADER_URL",
         )
         if not os.getenv(k)
     ]
     if missing:
         raise RuntimeError(
             f"Missing required secrets/env: {missing}. "
-            "Add them in Kaggle Notebook Secrets (Add-ons -> Secrets)."
+            "Need SWE_GRADER_URL pointing at the local Docker grader tunnel."
         )
     os.environ.setdefault("LANGFUSE_HOST", "https://cloud.langfuse.com")
     os.environ.setdefault("LANGFUSE_BASE_URL", os.environ["LANGFUSE_HOST"])
     os.environ.setdefault("WANDB_PROJECT", "deepagents-rl-coding")
     os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")
+    os.environ.setdefault("TRL_EXPERIMENTAL_SILENCE", "1")
 
 
 # ---------------------------------------------------------------------------
@@ -160,7 +150,9 @@ def resolve_model() -> str:
         try:
             api.model_info(repo)
             print(f"[resolve] using {repo}")
-            Path("/kaggle/working/resolved_model.json").write_text(
+            out = Path("/kaggle/working/resolved_model.json")
+            out.parent.mkdir(parents=True, exist_ok=True)
+            out.write_text(
                 json.dumps({"train_model_id": repo, "tried": tried}, indent=2),
                 encoding="utf-8",
             )
@@ -170,6 +162,162 @@ def resolve_model() -> str:
     raise RuntimeError(f"No model resolved; tried {tried}")
 
 
+# ---------------------------------------------------------------------------
+# Lightweight HF agent rollout (checkpoint eval)
+# ---------------------------------------------------------------------------
+
+_TOOL_CALL_RE = re.compile(
+    r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
+    re.DOTALL,
+)
+
+
+def _env_tool_callables(env: Any) -> list[Callable[..., str]]:
+    return [
+        getattr(env, name)
+        for name, _ in inspect.getmembers(type(env), predicate=inspect.isfunction)
+        if name not in ("reset", "get_reward") and not name.startswith("_")
+    ]
+
+
+def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
+    calls: list[dict[str, Any]] = []
+    for m in _TOOL_CALL_RE.finditer(text or ""):
+        try:
+            obj = json.loads(m.group(1))
+            name = obj.get("name") or obj.get("function", {}).get("name")
+            args = obj.get("arguments") or obj.get("parameters") or obj.get("function", {}).get("arguments") or {}
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args)
+                except json.JSONDecodeError:
+                    args = {"input": args}
+            if name:
+                calls.append({"name": str(name), "arguments": args if isinstance(args, dict) else {}})
+        except json.JSONDecodeError:
+            continue
+    return calls
+
+
+def run_heldout_agent_eval(
+    model,
+    tokenizer,
+    problems: list[dict],
+    *,
+    max_turns: int = 8,
+    max_new_tokens: int = 512,
+) -> list[dict]:
+    """Greedy multi-turn tool loop + remote Docker reward (checkpoint metric)."""
+    import torch
+
+    from deepagents_rl.train.swe_env import DeepAgentSWEEnv
+
+    rows = []
+    model.eval()
+    for p in problems:
+        env = DeepAgentSWEEnv()
+        try:
+            env.reset(**{k: v for k, v in p.items() if k != "prompt"})
+            tools = _env_tool_callables(env)
+            tool_map = {t.__name__: t for t in tools}
+            prompt = p.get("prompt")
+            if isinstance(prompt, list):
+                messages = list(prompt)
+            else:
+                messages = [{"role": "user", "content": str(prompt or "")}]
+
+            final_text = ""
+            for _turn in range(max_turns):
+                try:
+                    prompt_text = tokenizer.apply_chat_template(
+                        messages,
+                        tools=tools,
+                        tokenize=False,
+                        add_generation_prompt=True,
+                    )
+                except TypeError:
+                    prompt_text = tokenizer.apply_chat_template(
+                        messages, tokenize=False, add_generation_prompt=True
+                    )
+                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)
+                final_text = completion
+                calls = _parse_tool_calls(completion)
+                messages.append({"role": "assistant", "content": completion})
+                if not calls:
+                    break
+                for call in calls:
+                    fn = tool_map.get(call["name"])
+                    if fn is None:
+                        result = f"Unknown tool: {call['name']}"
+                    else:
+                        try:
+                            result = fn(**call["arguments"])
+                        except TypeError:
+                            try:
+                                result = fn(*call["arguments"].values())
+                            except Exception as e:  # noqa: BLE001
+                                result = f"TOOL_ERROR: {e}"
+                        except Exception as e:  # noqa: BLE001
+                            result = f"TOOL_ERROR: {e}"
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "name": call["name"],
+                            "content": str(result)[:4000],
+                        }
+                    )
+
+            reward = float(env.get_reward())
+            report = getattr(env, "_last_report", {}) or {}
+            rows.append(
+                {
+                    "task_id": p.get("task_id") or p.get("instance_id"),
+                    "prompt": (
+                        prompt
+                        if isinstance(prompt, str)
+                        else json.dumps(prompt)[:2000]
+                    ),
+                    "completion": final_text[:4000],
+                    "pass_frac": float(report.get("pass_frac") or reward),
+                    "reward_total": reward,
+                    "resolved": bool(report.get("resolved")),
+                    "resolution": report.get("resolution"),
+                }
+            )
+        except Exception as e:  # noqa: BLE001
+            traceback.print_exc()
+            rows.append(
+                {
+                    "task_id": p.get("task_id") or p.get("instance_id"),
+                    "prompt": str(p.get("prompt"))[:2000],
+                    "completion": f"EVAL_ERROR: {e}",
+                    "pass_frac": 0.0,
+                    "reward_total": 0.0,
+                    "resolved": False,
+                    "resolution": "ERROR",
+                }
+            )
+        finally:
+            try:
+                env._cleanup()
+            except Exception:
+                pass
+    return rows
+
+
 # ---------------------------------------------------------------------------
 # Langfuse checkpoint eval
 # ---------------------------------------------------------------------------
@@ -195,9 +343,10 @@ def langfuse_log_ckpt_eval(
         raise RuntimeError("Langfuse auth failed during ckpt eval")
 
     n = len(results)
-    mean_pass = sum(r["pass_frac"] for r in results) / n if n else 0.0
-    mean_reward = sum(r["reward_total"] for r in results) / n if n else 0.0
-    n_perfect = sum(1 for r in results if r["pass_frac"] >= 1.0)
+    mean_pass = sum(float(r.get("pass_frac") or 0) for r in results) / n if n else 0.0
+    mean_reward = sum(float(r.get("reward_total") or 0) for r in results) / n if n else 0.0
+    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"))
 
     with lf.start_as_current_observation(
         name="kaggle-ckpt-eval",
@@ -208,6 +357,7 @@ def langfuse_log_ckpt_eval(
             "train_model_id": train_model_id,
             "wandb_run_id": wandb_run_id,
             "phase": "ckpt_eval",
+            "harness": "deepagents-swe",
         },
         input={"step": step, "hf_revision": hf_revision, "n": n},
     ) as root:
@@ -215,19 +365,21 @@ def langfuse_log_ckpt_eval(
             with lf.start_as_current_observation(
                 name=f"task:{r['task_id']}",
                 as_type="span",
-                input={"prompt": r["prompt"][:2000]},
+                input={"prompt": str(r.get("prompt", ""))[:2000]},
             ) as span:
                 span.update(
                     output={
-                        "completion": r["completion"][:2000],
-                        "pass_frac": r["pass_frac"],
-                        "reward_total": r["reward_total"],
+                        "completion": str(r.get("completion", ""))[:2000],
+                        "pass_frac": r.get("pass_frac"),
+                        "reward_total": r.get("reward_total"),
+                        "resolved": r.get("resolved"),
                     }
                 )
         summary = {
             "mean_pass_frac": mean_pass,
             "mean_reward_total": mean_reward,
             "n_perfect": n_perfect,
+            "n_resolved": n_resolved,
             "n_problems": n,
             "step": step,
             "hf_revision": hf_revision,
@@ -236,49 +388,17 @@ def langfuse_log_ckpt_eval(
         lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
         lf.score_current_trace(name="mean_reward_total", value=mean_reward)
         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} mean_pass={mean_pass:.3f} trace={trace_id}")
+    print(
+        f"[langfuse] ckpt-eval step={step} rev={hf_revision} "
+        f"mean_reward={mean_reward:.3f} resolved={n_resolved}/{n} trace={trace_id}"
+    )
     return {"langfuse_trace_id": trace_id, **summary}
 
 
-def run_heldout_eval(model, tokenizer, problems: list[dict], max_new_tokens: int = 512) -> list[dict]:
-    import torch
-
-    rows = []
-    model.eval()
-    for p in problems:
-        messages = [{"role": "user", "content": p["prompt"]}]
-        try:
-            prompt_text = tokenizer.apply_chat_template(
-                messages, tokenize=False, add_generation_prompt=True
-            )
-        except Exception:
-            prompt_text = p["prompt"]
-        inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
-        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)
-        g = grade_completion(completion, list(p["hidden_tests"]))
-        rows.append(
-            {
-                "task_id": p["task_id"],
-                "prompt": p["prompt"],
-                "completion": completion,
-                "pass_frac": float(g["pass_frac"]),
-                "reward_total": float(g["pass_frac"]),
-            }
-        )
-    return rows
-
-
 # ---------------------------------------------------------------------------
 # Callback: eval + langfuse THEN save
 # ---------------------------------------------------------------------------
@@ -291,13 +411,14 @@ def make_eval_before_save_callback(
     hub_model_id: str,
     save_steps: int,
     n_eval: int,
+    max_eval_turns: int,
 ):
     from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
 
     class EvalBeforeCheckpointCallback(TrainerCallback):
         def __init__(self) -> None:
             self._last_eval_step = -1
-            self.trainer = None  # set after GRPOTrainer construction
+            self.trainer = None
 
         def on_step_end(
             self,
@@ -326,23 +447,27 @@ def make_eval_before_save_callback(
                     self.trainer, "tokenizer", None
                 )
             if model is None or tokenizer is None:
-                print("[ckpt-eval] missing model/tokenizer; aborting save pipeline")
                 raise RuntimeError("ckpt-eval requires model+tokenizer")
+
             rev = f"grpo-step-{step}"
             problems = eval_problems[:n_eval]
-            print(f"\n=== EVAL BEFORE CHECKPOINT step={step} rev={rev} n={len(problems)} ===")
+            print(f"\n=== AGENT EVAL BEFORE CHECKPOINT step={step} rev={rev} n={len(problems)} ===")
             try:
-                results = run_heldout_eval(model, tokenizer, problems)
+                results = run_heldout_agent_eval(
+                    model, tokenizer, problems, max_turns=max_eval_turns
+                )
                 import wandb
 
                 wandb_run_id = None
                 if wandb.run is not None:
                     wandb_run_id = wandb.run.id
-                    mean_pass = sum(r["pass_frac"] for r in results) / max(len(results), 1)
+                    mean_r = sum(float(r["reward_total"]) for r in results) / max(len(results), 1)
                     wandb.log(
                         {
-                            "eval/mean_pass_frac": mean_pass,
-                            "eval/n_perfect": sum(1 for r in results if r["pass_frac"] >= 1.0),
+                            "eval/mean_reward_total": mean_r,
+                            "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/checkpoint_step": step,
                         },
                         step=step,
@@ -355,17 +480,19 @@ def make_eval_before_save_callback(
                     results=results,
                     wandb_run_id=wandb_run_id,
                 )
-                # Persist summary next to checkpoint
                 out = Path(args.output_dir) / f"ckpt_eval_step_{step}.json"
-                out.write_text(json.dumps(summary | {"results": results}, indent=2), encoding="utf-8")
+                out.write_text(
+                    json.dumps(summary | {"results": results}, indent=2), encoding="utf-8"
+                )
 
-                # Push adapter AFTER langfuse success
                 try:
                     model.push_to_hub(
                         hub_model_id,
                         token=os.environ["HF_TOKEN"],
                         revision=rev,
-                        commit_message=f"GRPO adapter {rev}; mean_pass={summary['mean_pass_frac']:.3f}",
+                        commit_message=(
+                            f"Agent GRPO {rev}; mean_reward={summary['mean_reward_total']:.3f}"
+                        ),
                         private=False,
                     )
                     print(f"[hf] pushed {hub_model_id}@{rev}")
@@ -373,7 +500,7 @@ def make_eval_before_save_callback(
                         wandb.log(
                             {
                                 "hf/revision": rev,
-                                "hf/mean_pass_at_push": summary["mean_pass_frac"],
+                                "hf/mean_reward_at_push": summary["mean_reward_total"],
                                 "langfuse/ckpt_trace_id": summary.get("langfuse_trace_id"),
                             },
                             step=step,
@@ -382,7 +509,7 @@ def make_eval_before_save_callback(
                     print(f"[hf] push failed (local ckpt still ok): {e}")
                     traceback.print_exc()
             except Exception:
-                print("[ckpt-eval] FAILED — blocking is policy; re-raising so bad ckpts aren't silent")
+                print("[ckpt-eval] FAILED — re-raising")
                 traceback.print_exc()
                 raise
             return control
@@ -400,65 +527,105 @@ class Cfg:
 
     def __init__(
         self,
-        dataset_id: str = "vanishingradient/deepagents-rl-mbpp",
+        dataset_id: str = "vanishingradient/deepagents-rl-swe-lite",
         hub_model_id: str = "vanishingradient/deepagents-rl-coding",
-        max_steps: int = 40,
+        train_split: str = "dev",
+        eval_split: str = "test",
+        max_steps: int = 20,
         save_steps: int = 5,
-        num_generations: int = 4,
+        num_generations: int = 2,
         learning_rate: float = 1e-6,
-        max_completion_length: int = 512,
-        max_seq_length: int = 2048,
-        n_eval: int = 16,
+        max_completion_length: int = 1024,
+        max_seq_length: int = 4096,
+        max_tool_calling_iterations: int = 12,
+        n_eval: int = 2,
+        max_eval_turns: int = 8,
+        max_train_instances: int = 16,
         lora_r: int = 16,
     ) -> None:
         self.dataset_id = dataset_id
         self.hub_model_id = hub_model_id
+        self.train_split = train_split
+        self.eval_split = eval_split
         self.max_steps = max_steps
         self.save_steps = save_steps
         self.num_generations = num_generations
         self.learning_rate = learning_rate
         self.max_completion_length = max_completion_length
         self.max_seq_length = max_seq_length
+        self.max_tool_calling_iterations = max_tool_calling_iterations
         self.n_eval = n_eval
+        self.max_eval_turns = max_eval_turns
+        self.max_train_instances = max_train_instances
         self.lora_r = lora_r
 
 
+def _load_swe_splits(cfg: Cfg):
+    from datasets import Dataset, load_dataset
+
+    from deepagents_rl.train.swe_env import prepare_swe_grpo_rows
+
+    ds = load_dataset(cfg.dataset_id, token=os.environ["HF_TOKEN"])
+    train_split = cfg.train_split if cfg.train_split in ds else "train"
+    eval_split = cfg.eval_split if cfg.eval_split in ds else (
+        "test" if "test" in ds else train_split
+    )
+
+    train_recs = [dict(r) for r in ds[train_split]]
+    eval_recs = [dict(r) for r in ds[eval_split]]
+
+    # Keep only rows with docker image + eval script
+    def _ok(r: dict) -> bool:
+        img = r.get("image") or r.get("docker_image") or ""
+        return bool(img) and bool(r.get("eval_script"))
+
+    train_recs = [r for r in train_recs if _ok(r)][: cfg.max_train_instances]
+    eval_recs = [r for r in eval_recs if _ok(r)]
+
+    train_rows = prepare_swe_grpo_rows(train_recs)
+    eval_rows = prepare_swe_grpo_rows(eval_recs)
+    if not train_rows:
+        raise RuntimeError(f"No usable SWE train rows in {cfg.dataset_id}/{train_split}")
+    if not eval_rows:
+        eval_rows = train_rows[: max(1, cfg.n_eval)]
+
+    print(
+        f"[data] train={len(train_rows)} eval={len(eval_rows)} "
+        f"splits={train_split}/{eval_split} dataset={cfg.dataset_id}"
+    )
+    return Dataset.from_list(train_rows), eval_rows
+
+
 def main(cfg: Cfg | None = None) -> None:
     global _TOKENIZER
     cfg = cfg or Cfg()
     load_secrets()
+    _ensure_package()
 
     import torch
     import wandb
-    from datasets import load_dataset
     from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
     from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
     from trl import GRPOConfig, GRPOTrainer
 
+    from deepagents_rl.train.swe_env import DeepAgentSWEEnv
+
+    # Probe grader
+    import requests
+
+    grader = os.environ["SWE_GRADER_URL"].rstrip("/")
+    headers = {}
+    if os.environ.get("SWE_GRADER_TOKEN"):
+        headers["Authorization"] = f"Bearer {os.environ['SWE_GRADER_TOKEN']}"
+    health = requests.get(f"{grader}/health", headers=headers, timeout=30)
+    health.raise_for_status()
+    print("[grader] health", health.json())
+
     os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT", "deepagents-rl-coding")
     wandb.login(key=os.environ["WANDB_API_KEY"])
 
     model_id = resolve_model()
-    ds = load_dataset(cfg.dataset_id, token=os.environ["HF_TOKEN"])
-    train = ds["train"]
-    ev = ds["eval"]
-    # Ensure list columns serializable for reward
-    def _ensure(example):
-        ht = example["hidden_tests"]
-        if not isinstance(ht, list):
-            ht = list(ht)
-        example["hidden_tests"] = ht
-        return example
-
-    train = train.map(_ensure)
-    eval_problems = [
-        {
-            "task_id": row["task_id"],
-            "prompt": row["prompt"],
-            "hidden_tests": list(row["hidden_tests"]),
-        }
-        for row in ev
-    ]
+    train, eval_problems = _load_swe_splits(cfg)
 
     bnb = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -466,7 +633,9 @@ def main(cfg: Cfg | None = None) -> None:
         bnb_4bit_quant_type="nf4",
         bnb_4bit_use_double_quant=True,
     )
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ["HF_TOKEN"], trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id, token=os.environ["HF_TOKEN"], trust_remote_code=True
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     _TOKENIZER = tokenizer
@@ -485,7 +654,15 @@ def main(cfg: Cfg | None = None) -> None:
         lora_dropout=0.0,
         bias="none",
         task_type="CAUSAL_LM",
-        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
     )
     model = get_peft_model(model, lora)
 
@@ -495,22 +672,30 @@ def main(cfg: Cfg | None = None) -> None:
 
     wandb.init(
         project=os.environ["WANDB_PROJECT"],
-        name="kaggle-grpo-mbpp",
+        name="kaggle-grpo-swe",
         job_type="train",
         config={
             "train_model_id": model_id,
             "dataset_id": cfg.dataset_id,
+            "harness": "deepagents-swe",
             "max_steps": cfg.max_steps,
             "save_steps": cfg.save_steps,
             "n_eval": cfg.n_eval,
+            "num_generations": cfg.num_generations,
+            "max_tool_calling_iterations": cfg.max_tool_calling_iterations,
             "hub_model_id": cfg.hub_model_id,
+            "swe_grader_url": grader,
             "eval_before_every_checkpoint": True,
         },
     )
 
-    # Baseline eval @ step 0 BEFORE any training / checkpoint
-    print("=== BASELINE EVAL (step 0) before any training ===")
-    baseline = run_heldout_eval(model, tokenizer, eval_problems[: cfg.n_eval])
+    print("=== BASELINE AGENT EVAL (step 0) ===")
+    baseline = run_heldout_agent_eval(
+        model,
+        tokenizer,
+        eval_problems[: cfg.n_eval],
+        max_turns=cfg.max_eval_turns,
+    )
     base_summary = langfuse_log_ckpt_eval(
         step=0,
         hf_revision="grpo-step-0-baseline",
@@ -524,7 +709,8 @@ def main(cfg: Cfg | None = None) -> None:
     wandb.log(
         {
             "eval/mean_pass_frac": base_summary["mean_pass_frac"],
-            "eval/n_perfect": base_summary["n_perfect"],
+            "eval/mean_reward_total": base_summary["mean_reward_total"],
+            "eval/n_resolved": base_summary.get("n_resolved", 0),
             "eval/checkpoint_step": 0,
             "langfuse/ckpt_trace_id": base_summary.get("langfuse_trace_id"),
         },
@@ -533,12 +719,13 @@ def main(cfg: Cfg | None = None) -> None:
 
     args = GRPOConfig(
         output_dir=output_dir,
-        run_name="kaggle-grpo-mbpp",
+        run_name="kaggle-grpo-swe",
         learning_rate=cfg.learning_rate,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=4,
         num_generations=cfg.num_generations,
         max_completion_length=cfg.max_completion_length,
+        max_tool_calling_iterations=cfg.max_tool_calling_iterations,
         max_steps=cfg.max_steps,
         logging_steps=1,
         save_steps=cfg.save_steps,
@@ -549,7 +736,7 @@ def main(cfg: Cfg | None = None) -> None:
         remove_unused_columns=False,
         hub_model_id=cfg.hub_model_id,
         hub_token=os.environ["HF_TOKEN"],
-        push_to_hub=False,  # push ONLY after Langfuse ckpt-eval succeeds
+        push_to_hub=False,
     )
 
     cb = make_eval_before_save_callback(
@@ -558,23 +745,30 @@ def main(cfg: Cfg | None = None) -> None:
         hub_model_id=cfg.hub_model_id,
         save_steps=cfg.save_steps,
         n_eval=cfg.n_eval,
+        max_eval_turns=cfg.max_eval_turns,
     )
 
     trainer = GRPOTrainer(
         model=model,
-        reward_funcs=coding_reward,
         args=args,
         train_dataset=train,
         processing_class=tokenizer,
         callbacks=[cb],
+        environment_factory=DeepAgentSWEEnv,
     )
     cb.trainer = trainer
 
-    trainer.train()    # Final eval + push
+    trainer.train()
+
     final_step = int(trainer.state.global_step)
     rev = f"grpo-step-{final_step}-final"
-    print(f"=== FINAL EVAL step={final_step} ===")
-    final_results = run_heldout_eval(model, tokenizer, eval_problems[: cfg.n_eval])
+    print(f"=== FINAL AGENT EVAL step={final_step} ===")
+    final_results = run_heldout_agent_eval(
+        model,
+        tokenizer,
+        eval_problems[: cfg.n_eval],
+        max_turns=cfg.max_eval_turns,
+    )
     final_summary = langfuse_log_ckpt_eval(
         step=final_step,
         hf_revision=rev,
@@ -586,17 +780,20 @@ def main(cfg: Cfg | None = None) -> None:
         cfg.hub_model_id,
         token=os.environ["HF_TOKEN"],
         revision=rev,
-        commit_message=f"Final GRPO {rev}; mean_pass={final_summary['mean_pass_frac']:.3f}",
+        commit_message=(
+            f"Final agent GRPO {rev}; mean_reward={final_summary['mean_reward_total']:.3f}"
+        ),
     )
-    # also push main
     model.push_to_hub(
         cfg.hub_model_id,
         token=os.environ["HF_TOKEN"],
-        commit_message=f"Final GRPO main; mean_pass={final_summary['mean_pass_frac']:.3f}",
+        commit_message=(
+            f"Final agent GRPO main; mean_reward={final_summary['mean_reward_total']:.3f}"
+        ),
     )
     wandb.log(
         {
-            "eval/final_mean_pass_frac": final_summary["mean_pass_frac"],
+            "eval/final_mean_reward_total": final_summary["mean_reward_total"],
             "hf/final_revision": rev,
         }
     )