| """TRL GRPO `environment_factory` for deepagents SWE-bench sandboxes.
|
|
|
| Kaggle has no Docker for instance images — `get_reward` POSTs the model patch
|
| to `SWE_GRADER_URL` (local `scripts/swe_grader_server.py` + tunnel).
|
|
|
| Each rollout opens a Langfuse agent trace; every tool call is a nested `tool`
|
| span; grading is an `evaluator` span with reward scores.
|
| """
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import shutil
|
| import tempfile
|
| import traceback
|
| import uuid
|
| from functools import wraps
|
| from pathlib import Path
|
| from typing import Any, Callable
|
|
|
|
|
| def _vlog(msg: str, **fields: Any) -> None:
|
| """Always-on verbose progress (Kaggle logs + local stderr)."""
|
| extra = ""
|
| if fields:
|
| try:
|
| extra = " | " + json.dumps(fields, default=str)[:1200]
|
| except Exception:
|
| extra = f" | {fields!r}"[:1200]
|
| line = f"[darl] {msg}{extra}"
|
| print(line, flush=True)
|
| try:
|
| from deepagents_rl.observability.logging import step
|
|
|
| step("swe_env", msg, **fields)
|
| except Exception:
|
| pass
|
|
|
|
|
| def _result_to_str(obj: Any) -> str:
|
| if obj is None:
|
| return ""
|
| if isinstance(obj, str):
|
| return obj[:8000]
|
| if hasattr(obj, "output") and isinstance(obj.output, str):
|
| return obj.output[:8000]
|
|
|
| entries = getattr(obj, "entries", None)
|
| if entries is not None:
|
| lines: list[str] = []
|
| for e in entries:
|
| if isinstance(e, dict):
|
| lines.append(str(e.get("path") or e))
|
| else:
|
| lines.append(str(e))
|
| return ("\n".join(lines) if lines else "(empty directory)")[:8000]
|
| matches = getattr(obj, "matches", None)
|
| if matches is not None and not getattr(obj, "error", None):
|
| try:
|
| return json.dumps(matches, default=str)[:8000]
|
| except Exception:
|
| return str(matches)[:8000]
|
| if hasattr(obj, "error") and obj.error:
|
| return f"ERROR: {obj.error}"[:8000]
|
| try:
|
| return json.dumps(obj, default=str)[:8000]
|
| except Exception:
|
| return str(obj)[:8000]
|
|
|
|
|
|
|
| _PHANTOM_DIRS = frozenset(
|
| {
|
| "src",
|
| "lib",
|
| "app",
|
| "source",
|
| "sources",
|
| "code",
|
| "project",
|
| "workspace",
|
| }
|
| )
|
|
|
|
|
| def _as_str_list(val: Any) -> list[str]:
|
| if val is None:
|
| return []
|
| if isinstance(val, list):
|
| return [str(x) for x in val]
|
| if isinstance(val, str):
|
| s = val.strip()
|
| if not s:
|
| return []
|
| try:
|
| parsed = json.loads(s)
|
| if isinstance(parsed, list):
|
| return [str(x) for x in parsed]
|
| except json.JSONDecodeError:
|
| pass
|
| return [s]
|
| return [str(val)]
|
|
|
|
|
| def _truncate(obj: Any, n: int = 4000) -> Any:
|
| if isinstance(obj, str):
|
| return obj[:n]
|
| if isinstance(obj, dict):
|
| return {k: _truncate(v, n) for k, v in list(obj.items())[:40]}
|
| if isinstance(obj, list):
|
| return [_truncate(v, n) for v in obj[:40]]
|
| return obj
|
|
|
|
|
| def grade_patch_remote(
|
| *,
|
| instance_id: str,
|
| docker_image: str,
|
| eval_script: str,
|
| fail_to_pass: list[str],
|
| pass_to_pass: list[str],
|
| model_patch: str,
|
| timeout_s: int = 1800,
|
| ) -> dict[str, Any]:
|
| """POST patch to the local Docker grader service."""
|
| import requests
|
|
|
| base = (os.environ.get("SWE_GRADER_URL") or "").rstrip("/")
|
| if not base:
|
| raise RuntimeError("SWE_GRADER_URL is not set (point at the local grader tunnel)")
|
|
|
| headers = {
|
| "Content-Type": "application/json",
|
| "ngrok-skip-browser-warning": "1",
|
| }
|
| token = os.environ.get("SWE_GRADER_TOKEN") or ""
|
| if token:
|
| headers["Authorization"] = f"Bearer {token}"
|
|
|
| payload = {
|
| "instance_id": instance_id,
|
| "image": docker_image,
|
| "eval_script": eval_script,
|
| "fail_to_pass": fail_to_pass,
|
| "pass_to_pass": pass_to_pass,
|
| "model_patch": model_patch or "",
|
| "timeout_s": timeout_s,
|
| }
|
| _vlog(
|
| "grader.POST /grade",
|
| url=base,
|
| instance_id=instance_id,
|
| patch_bytes=len((model_patch or "").encode("utf-8")),
|
| timeout_s=timeout_s,
|
| )
|
| resp = requests.post(f"{base}/grade", headers=headers, json=payload, timeout=timeout_s + 60)
|
| _vlog("grader.response", status=resp.status_code, body_preview=(resp.text or "")[:400])
|
| resp.raise_for_status()
|
| data = resp.json()
|
| _vlog(
|
| "grader.result",
|
| instance_id=instance_id,
|
| reward_total=data.get("reward_total"),
|
| resolution=data.get("resolution"),
|
| resolved=data.get("resolved"),
|
| error=data.get("error"),
|
| )
|
| return data
|
|
|
|
|
| def _traced_tool(method: Callable[..., str]) -> Callable[..., str]:
|
| """Wrap a public tool method so each call becomes a Langfuse tool span."""
|
|
|
| @wraps(method)
|
| def wrapper(self: "DeepAgentSWEEnv", *args: Any, **kwargs: Any) -> str:
|
| return self._call_tool(method.__name__, method, args, kwargs)
|
|
|
| return wrapper
|
|
|
|
|
| class DeepAgentSWEEnv:
|
| """Per-rollout TRL environment: LocalShellBackend tools + remote Docker reward."""
|
|
|
| def __init__(self) -> None:
|
| self._sandbox_parent: Path | None = None
|
| self._sandbox: Path | None = None
|
| self._backend: Any = None
|
| self._task: dict[str, Any] = {}
|
| self._last_reward: float = 0.0
|
| self._last_report: dict[str, Any] = {}
|
| self._rollout_id: str = ""
|
| self._tool_calls: list[dict[str, Any]] = []
|
| self._lf: Any = None
|
| self._root_cm: Any = None
|
| self._root: Any = None
|
| self._trace_id: str = ""
|
| self._root_span_id: str = ""
|
|
|
| def _langfuse_enabled(self) -> bool:
|
| if os.getenv("DARL_LANGFUSE", "1").strip() in ("0", "false", "no"):
|
| return False
|
| return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))
|
|
|
| def _ensure_langfuse(self) -> Any | None:
|
| if not self._langfuse_enabled():
|
| return None
|
| if self._lf is not None:
|
| return self._lf
|
| try:
|
| from langfuse import Langfuse, get_client
|
|
|
| Langfuse(
|
| public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
|
| secret_key=os.environ["LANGFUSE_SECRET_KEY"],
|
| host=os.getenv("LANGFUSE_HOST") or os.getenv("LANGFUSE_BASE_URL"),
|
| )
|
| self._lf = get_client()
|
| return self._lf
|
| except Exception as e:
|
| print(f"[langfuse] init failed: {e}")
|
| return None
|
|
|
| def _child_trace_context(self) -> dict[str, str] | None:
|
| """Attach children to THIS rollout's root — never ambient OTEL context.
|
|
|
| Passing only ``trace_id`` (no parent) makes the Langfuse SDK invent a
|
| NonRecordingSpan parent id → orphan parents + mega-trace pollution when
|
| parallel GRPO rollouts share get_client().
|
| """
|
| if not self._trace_id or not self._root_span_id:
|
| return None
|
| return {"trace_id": self._trace_id, "parent_span_id": self._root_span_id}
|
|
|
| def _end_trace(self) -> None:
|
| if self._root_cm is not None:
|
| try:
|
| self._root_cm.__exit__(None, None, None)
|
| except Exception:
|
| pass
|
| self._root_cm = None
|
| self._root = None
|
| self._root_span_id = ""
|
|
|
| if self._lf is not None:
|
| try:
|
| self._lf.flush()
|
| except Exception:
|
| pass
|
|
|
| def _call_tool(
|
| self,
|
| name: str,
|
| method: Callable[..., str],
|
| args: tuple[Any, ...],
|
| kwargs: dict[str, Any],
|
| ) -> str:
|
| payload = _truncate({"args": list(args), "kwargs": kwargs})
|
| n = len(self._tool_calls) + 1
|
| _vlog(
|
| f"tool[{n}].begin {name}",
|
| instance_id=self._task.get("instance_id"),
|
| rollout_id=self._rollout_id,
|
| input=payload,
|
| )
|
| lf = self._lf
|
| ctx = self._child_trace_context()
|
| if lf is None or ctx is None:
|
| try:
|
| out = method(self, *args, **kwargs)
|
| except Exception as e:
|
| out = f"TOOL_ERROR: {type(e).__name__}: {e}"
|
| self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
|
| _vlog(
|
| f"tool[{n}].end {name}",
|
| output_preview=str(out)[:500],
|
| output_chars=len(str(out)),
|
| )
|
| return out
|
|
|
| with lf.start_as_current_observation(
|
| name=name,
|
| as_type="tool",
|
| trace_context=ctx,
|
| input=payload,
|
| metadata={
|
| "rollout_id": self._rollout_id,
|
| "instance_id": self._task.get("instance_id"),
|
| "harness": "deepagents-swe-grpo",
|
| "tool_index": n,
|
| "trace_id": self._trace_id,
|
| },
|
| ) as span:
|
| try:
|
| out = method(self, *args, **kwargs)
|
| span.update(output=_truncate(out))
|
| except Exception as e:
|
| out = f"TOOL_ERROR: {type(e).__name__}: {e}"
|
| span.update(output=out, level="ERROR")
|
| self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
|
| _vlog(
|
| f"tool[{n}].end {name}",
|
| output_preview=str(out)[:500],
|
| output_chars=len(str(out)),
|
| )
|
| return out
|
|
|
| def reset(self, **kwargs: Any) -> str:
|
| """Materialize repo sandbox; return tool-forcing text TRL appends to the user message."""
|
| self._cleanup()
|
| self._end_trace()
|
|
|
| instance_id = str(kwargs.get("instance_id") or kwargs.get("task_id") or "")
|
| repo = str(kwargs.get("repo") or "")
|
| base_commit = str(kwargs.get("base_commit") or "")
|
| docker_image = str(kwargs.get("docker_image") or kwargs.get("image") or "")
|
| eval_script = str(kwargs.get("eval_script") or "")
|
| fail_to_pass = _as_str_list(
|
| kwargs.get("fail_to_pass")
|
| if kwargs.get("fail_to_pass") is not None
|
| else kwargs.get("fail_to_pass_json")
|
| )
|
| pass_to_pass = _as_str_list(
|
| kwargs.get("pass_to_pass")
|
| if kwargs.get("pass_to_pass") is not None
|
| else kwargs.get("pass_to_pass_json")
|
| )
|
|
|
| _vlog(
|
| "env.reset.begin",
|
| instance_id=instance_id,
|
| repo=repo,
|
| base_commit=(base_commit[:12] if base_commit else ""),
|
| keys=list(kwargs.keys())[:30],
|
| )
|
|
|
| if not instance_id or not repo or not base_commit:
|
| raise ValueError(
|
| f"DeepAgentSWEEnv.reset requires instance_id/repo/base_commit; got keys={list(kwargs)}"
|
| )
|
| if not docker_image or not eval_script.strip():
|
| raise ValueError(f"{instance_id}: missing docker_image/eval_script")
|
|
|
| from deepagents.backends import LocalShellBackend
|
|
|
| from deepagents_rl.harness.repo import materialize_repo_at_commit
|
|
|
| self._rollout_id = str(uuid.uuid4())
|
| self._tool_calls = []
|
| self._task = {
|
| "instance_id": instance_id,
|
| "task_id": instance_id,
|
| "repo": repo,
|
| "base_commit": base_commit,
|
| "image": docker_image,
|
| "docker_image": docker_image,
|
| "eval_script": eval_script,
|
| "fail_to_pass": fail_to_pass,
|
| "pass_to_pass": pass_to_pass,
|
| }
|
| self._last_reward = 0.0
|
| self._last_report = {}
|
| self._trace_id = ""
|
|
|
| lf = self._ensure_langfuse()
|
| _vlog("env.reset.langfuse", enabled=lf is not None)
|
| if lf is not None:
|
| try:
|
|
|
|
|
|
|
|
|
| self._root_cm = lf.start_as_current_observation(
|
| name=f"grpo-swe-rollout:{instance_id}",
|
| as_type="agent",
|
| input={
|
| "instance_id": instance_id,
|
| "repo": repo,
|
| "base_commit": base_commit,
|
| "rollout_id": self._rollout_id,
|
| "prompt_preview": _truncate(
|
| kwargs.get("prompt")
|
| if isinstance(kwargs.get("prompt"), str)
|
| else (
|
| kwargs.get("problem_statement")
|
| or (kwargs.get("prompt") or "")
|
| ),
|
| 2000,
|
| ),
|
| },
|
| metadata={
|
| "rollout_id": self._rollout_id,
|
| "instance_id": instance_id,
|
| "docker_image": docker_image,
|
| "phase": "grpo_train_or_eval",
|
| "harness": "deepagents-swe-grpo",
|
| "session_id": self._rollout_id,
|
| "tags": ["deepagents-rl", "swe-bench", "grpo", "tools"],
|
| },
|
| end_on_exit=True,
|
| )
|
| self._root = self._root_cm.__enter__()
|
| self._trace_id = str(
|
| getattr(self._root, "trace_id", None)
|
| or getattr(self._root, "traceId", None)
|
| or ""
|
| )
|
| self._root_span_id = str(
|
| getattr(self._root, "id", None)
|
| or getattr(self._root, "observation_id", None)
|
| or ""
|
| )
|
| _vlog(
|
| "env.reset.langfuse_trace",
|
| trace_id=self._trace_id,
|
| root_span_id=self._root_span_id,
|
| )
|
| except Exception as e:
|
| _vlog("env.reset.langfuse_failed", error=str(e))
|
| self._root_cm = None
|
| self._root = None
|
| self._trace_id = ""
|
| self._root_span_id = ""
|
|
|
| def _materialize() -> None:
|
| _vlog("env.reset.materialize", repo=repo, base_commit=base_commit[:12])
|
| self._sandbox_parent = Path(
|
| tempfile.mkdtemp(prefix=f"darl_grpo_{instance_id[:40]}_")
|
| )
|
| self._sandbox = self._sandbox_parent / "repo"
|
| materialize_repo_at_commit(repo=repo, base_commit=base_commit, dest=self._sandbox)
|
| self._backend = LocalShellBackend(
|
| root_dir=str(self._sandbox),
|
| virtual_mode=True,
|
| timeout=60,
|
| )
|
| _vlog("env.reset.sandbox_ready", sandbox=str(self._sandbox))
|
|
|
| ctx = self._child_trace_context()
|
| if lf is not None and ctx is not None:
|
| with lf.start_as_current_observation(
|
| name="repo_materialize",
|
| as_type="span",
|
| trace_context=ctx,
|
| input={"repo": repo, "base_commit": base_commit},
|
| ) as span:
|
| _materialize()
|
| span.update(output={"sandbox": str(self._sandbox)})
|
| else:
|
| _materialize()
|
|
|
| root_listing = ""
|
| try:
|
| root_listing = self._root_listing()
|
| except Exception as e:
|
| root_listing = f"(root listing failed: {e})"
|
| _vlog(
|
| "env.reset.done",
|
| instance_id=instance_id,
|
| rollout_id=self._rollout_id,
|
| n_f2p=len(fail_to_pass),
|
| n_p2p=len(pass_to_pass),
|
| langfuse_trace_id=self._trace_id,
|
| root_listing_chars=len(root_listing),
|
| )
|
|
|
| return (
|
| "\n\nREPO ROOT LISTING (already gathered for you):\n"
|
| f"{root_listing[:3500]}\n\n"
|
| "CRITICAL REMINDER:\n"
|
| "- You MUST use structured tool/function calls (not markdown ```json```).\n"
|
| "- FIRST tool call: ls path='.' OR grep/glob/read_file on a REAL path from the listing above.\n"
|
| "- Never invent src/lib/app. Prefer edit_file with a unique old_string.\n"
|
| "- Text-only DONE with zero tools scores NEGATIVE reward. Call tools.\n"
|
| "- After a minimal fix, stop tools and reply: DONE <summary>."
|
| )
|
|
|
| def _cleanup(self) -> None:
|
| self._backend = None
|
| if self._sandbox_parent is not None and self._sandbox_parent.exists():
|
| shutil.rmtree(self._sandbox_parent, ignore_errors=True)
|
| self._sandbox_parent = None
|
| self._sandbox = None
|
|
|
| def _require_backend(self) -> Any:
|
| if self._backend is None:
|
| raise RuntimeError("DeepAgentSWEEnv tool called before reset()")
|
| return self._backend
|
|
|
| @staticmethod
|
| def _norm_path(path: str | None) -> str:
|
| p = (path or ".").strip().strip("'\"")
|
| if p in ("", "/", "./", ".\\", "~"):
|
| return "."
|
|
|
| if p.startswith("./"):
|
| p = p[2:]
|
| if p.startswith("/"):
|
| p = p.lstrip("/") or "."
|
|
|
| p = p.rstrip("/") or "."
|
| return p or "."
|
|
|
| @classmethod
|
| def _resolve_ls_path(cls, path: str | None) -> tuple[str, str | None]:
|
| """Map invented dirs (src/lib/…) to repo root before hitting the backend."""
|
| p = cls._norm_path(path)
|
| head = p.split("/", 1)[0].lower()
|
| if head in _PHANTOM_DIRS:
|
| return ".", p
|
| return p, None
|
|
|
| def _is_missing(self, result_obj: Any, text: str) -> bool:
|
| err = getattr(result_obj, "error", None)
|
| if err and "path_not_found" in str(err).lower():
|
| return True
|
| return "path_not_found" in (text or "").lower()
|
|
|
| def _root_listing(self) -> str:
|
| backend = self._require_backend()
|
|
|
| for candidate in ("/", "."):
|
| result = backend.ls(candidate)
|
| out = _result_to_str(result)
|
| if not self._is_missing(result, out):
|
| return out
|
| return "(unable to list repo root)"
|
|
|
| @_traced_tool
|
| def ls(self, path: str = ".") -> str:
|
| """List files in a directory inside the sandbox repo root.
|
|
|
| Args:
|
| path: Directory relative to the repo root. Use "." for the root.
|
| Do NOT assume a "src/" folder exists — many repos have none.
|
|
|
| Returns:
|
| Directory listing as text (falls back to repo root if path missing).
|
| """
|
| backend = self._require_backend()
|
| path, aliased_from = self._resolve_ls_path(path)
|
| result = backend.ls(path)
|
| out = _result_to_str(result)
|
| if self._is_missing(result, out):
|
| root = self._root_listing()
|
| _vlog("tool.ls.fallback_root", requested=aliased_from or path)
|
| return (
|
| f"Path '{aliased_from or path}' does not exist here.\n"
|
| f"Listing repo root instead (use these real paths):\n{root}"
|
| )
|
| if aliased_from and aliased_from != path:
|
| _vlog("tool.ls.alias", requested=aliased_from, used=path)
|
| return (
|
| f"Note: '{aliased_from}' is not used in this repo; showing repo root:\n{out}"
|
| )
|
| return out
|
|
|
| @_traced_tool
|
| def glob(self, pattern: str, path: str = ".") -> str:
|
| """Glob for files matching pattern under path.
|
|
|
| Args:
|
| pattern: Glob pattern (e.g. '**/*.py').
|
| path: Root path for the search (use '.' for repo root).
|
|
|
| Returns:
|
| Matching paths as text.
|
| """
|
| backend = self._require_backend()
|
| path, aliased_from = self._resolve_ls_path(path)
|
| result = backend.glob(pattern, path)
|
| out = _result_to_str(result)
|
| if self._is_missing(result, out):
|
| return (
|
| f"Path '{aliased_from or path}' not found for glob; retrying at repo root.\n"
|
| f"{_result_to_str(backend.glob(pattern, '.'))}"
|
| )
|
| return out
|
|
|
| @_traced_tool
|
| def grep(self, pattern: str, path: str = ".") -> str:
|
| """Search file contents for a regex pattern.
|
|
|
| Args:
|
| pattern: Regex pattern to search for.
|
| path: Root path for the search (use '.' for repo root).
|
|
|
| Returns:
|
| Matching lines as text.
|
| """
|
| backend = self._require_backend()
|
| path, aliased_from = self._resolve_ls_path(path)
|
| result = backend.grep(pattern, path)
|
| out = _result_to_str(result)
|
| if self._is_missing(result, out):
|
| return (
|
| f"Path '{aliased_from or path}' not found for grep; retrying at repo root.\n"
|
| f"{_result_to_str(backend.grep(pattern, '.'))}"
|
| )
|
| return out
|
|
|
| @_traced_tool
|
| def read_file(self, file_path: str, offset: int = 0, limit: int = 200) -> str:
|
| """Read a text file from the sandbox.
|
|
|
| Args:
|
| file_path: Path relative to the repo root (discover via ls/glob first).
|
| offset: Starting line offset.
|
| limit: Max number of lines to return.
|
|
|
| Returns:
|
| File contents as text.
|
| """
|
| backend = self._require_backend()
|
| file_path = self._norm_path(file_path)
|
|
|
| head = file_path.split("/", 1)[0].lower()
|
| candidates = [file_path]
|
| if head in _PHANTOM_DIRS and "/" in file_path:
|
| candidates.append(file_path.split("/", 1)[1])
|
| last_out = ""
|
| for cand in candidates:
|
| result = backend.read(cand, offset=offset, limit=limit)
|
| out = _result_to_str(result)
|
| last_out = out
|
| if not self._is_missing(result, out):
|
| if cand != file_path:
|
| return f"(read via '{cand}' — '{file_path}' missing)\n{out}"
|
| return out
|
| return (
|
| f"File '{file_path}' not found. Use ls path='.' or glob to discover real paths.\n"
|
| f"Repo root listing:\n{self._root_listing()}\nLast error: {last_out}"
|
| )
|
|
|
| @_traced_tool
|
| def write_file(self, file_path: str, content: str) -> str:
|
| """Write or create a text file in the sandbox.
|
|
|
| Args:
|
| file_path: Path relative to the repo root.
|
| content: Full file contents to write.
|
|
|
| Returns:
|
| Write result as text.
|
| """
|
| return _result_to_str(self._require_backend().write(file_path, content))
|
|
|
| @_traced_tool
|
| def edit_file(self, file_path: str, old_string: str, new_string: str) -> str:
|
| """Replace old_string with new_string in an existing file.
|
|
|
| Args:
|
| file_path: Path relative to the repo root.
|
| old_string: Exact text to find.
|
| new_string: Replacement text.
|
|
|
| Returns:
|
| Edit result as text.
|
| """
|
| return _result_to_str(
|
| self._require_backend().edit(file_path, old_string, new_string)
|
| )
|
|
|
| @_traced_tool
|
| def execute(self, command: str) -> str:
|
| """Run a shell command in the sandbox.
|
|
|
| Args:
|
| command: Shell command (e.g. python -c '...').
|
|
|
| Returns:
|
| Command stdout/stderr as text.
|
| """
|
| return _result_to_str(self._require_backend().execute(command))
|
|
|
| def _shaped_reward(
|
| self,
|
| *,
|
| docker_reward: float,
|
| model_patch: str,
|
| report: dict[str, Any],
|
| ) -> tuple[float, dict[str, float]]:
|
| """Dense process reward so GRPO sees non-zero variance before resolve.
|
|
|
| Components (additive, then clamped):
|
| - no tool calls: -0.15
|
| - tool calls: +0.05 each (cap +0.25)
|
| - used edit_file/write_file: +0.15
|
| - non-empty git patch: +0.30
|
| - docker/pass reward (0..1 typically): +1.0 * docker_reward
|
| """
|
| names = [str(c.get("name") or "") for c in self._tool_calls]
|
| n = len(names)
|
| parts = {
|
| "tool_penalty": -0.15 if n == 0 else 0.0,
|
| "tool_bonus": 0.0 if n == 0 else float(min(0.25, 0.05 * n)),
|
| "edit_bonus": 0.15
|
| if any(x in names for x in ("edit_file", "write_file"))
|
| else 0.0,
|
| "patch_bonus": 0.30 if (model_patch or "").strip() else 0.0,
|
| "docker": float(docker_reward or 0.0),
|
| }
|
| total = sum(parts.values())
|
| total = max(-0.2, min(1.6, total))
|
| parts["total"] = total
|
| return total, parts
|
|
|
| def get_reward(self) -> float:
|
| """Shaped reward = process bonuses + remote Docker SWE score (if patch)."""
|
| lf = self._lf
|
| reward = 0.0
|
| report: dict[str, Any] = {}
|
| shape_parts: dict[str, float] = {}
|
| model_patch = ""
|
| _vlog(
|
| "env.get_reward.begin",
|
| instance_id=self._task.get("instance_id"),
|
| n_tool_calls=len(self._tool_calls),
|
| tools_used=[c.get("name") for c in self._tool_calls],
|
| )
|
|
|
| def _grade() -> float:
|
| nonlocal report, model_patch, shape_parts
|
| docker_reward = 0.0
|
| if self._sandbox is None:
|
| report = {"error": "no_sandbox", "reward_total": 0.0}
|
| _vlog("env.get_reward.no_sandbox")
|
| shaped, shape_parts = self._shaped_reward(
|
| docker_reward=0.0, model_patch="", report=report
|
| )
|
| report["shaped"] = shape_parts
|
| report["reward_total"] = shaped
|
| return shaped
|
|
|
| from deepagents_rl.harness.repo import git_diff_patch
|
|
|
| try:
|
| _vlog("env.get_reward.git_diff")
|
| model_patch = git_diff_patch(self._sandbox)
|
| _vlog(
|
| "env.get_reward.patch",
|
| patch_bytes=len(model_patch.encode("utf-8")),
|
| empty=not bool(model_patch.strip()),
|
| preview=model_patch[:300],
|
| )
|
| except Exception as e:
|
| report = {"error": f"git_diff_failed: {e}"}
|
| _vlog("env.get_reward.git_diff_failed", error=str(e))
|
| shaped, shape_parts = self._shaped_reward(
|
| docker_reward=0.0, model_patch="", report=report
|
| )
|
| report["shaped"] = shape_parts
|
| report["reward_total"] = shaped
|
| return shaped
|
|
|
| if not model_patch.strip():
|
| report = {"error": "empty_model_patch", "reward_total": 0.0}
|
| _vlog("env.get_reward.empty_patch")
|
| else:
|
| try:
|
| report = grade_patch_remote(
|
| instance_id=self._task["instance_id"],
|
| docker_image=self._task["docker_image"],
|
| eval_script=self._task["eval_script"],
|
| fail_to_pass=list(self._task["fail_to_pass"]),
|
| pass_to_pass=list(self._task["pass_to_pass"]),
|
| model_patch=model_patch,
|
| )
|
| docker_reward = float(report.get("reward_total") or 0.0)
|
| except Exception as e:
|
| report = {
|
| "error": f"grader_failed: {e}",
|
| "traceback": traceback.format_exc()[-1500:],
|
| }
|
| _vlog("env.get_reward.grader_failed", error=str(e))
|
| docker_reward = 0.0
|
|
|
| shaped, shape_parts = self._shaped_reward(
|
| docker_reward=docker_reward,
|
| model_patch=model_patch,
|
| report=report,
|
| )
|
| report["docker_reward"] = docker_reward
|
| report["shaped"] = shape_parts
|
| report["reward_total"] = shaped
|
| _vlog("env.get_reward.shaped", **shape_parts)
|
| return shaped
|
|
|
| ctx = self._child_trace_context()
|
| if lf is not None and ctx is not None:
|
| with lf.start_as_current_observation(
|
| name="reward_grader",
|
| as_type="evaluator",
|
| trace_context=ctx,
|
| input={
|
| "instance_id": self._task.get("instance_id"),
|
| "n_tool_calls": len(self._tool_calls),
|
| "tools_used": [c["name"] for c in self._tool_calls],
|
| },
|
| ) as span:
|
| reward = _grade()
|
| span.update(
|
| output=_truncate(
|
| {
|
| "reward_total": reward,
|
| "docker_reward": report.get("docker_reward"),
|
| "shaped": report.get("shaped"),
|
| "pass_frac": report.get("pass_frac"),
|
| "resolved": report.get("resolved"),
|
| "resolution": report.get("resolution"),
|
| "error": report.get("error"),
|
| "n_tool_calls": len(self._tool_calls),
|
| }
|
| )
|
| )
|
| try:
|
| lf.score_current_trace(name="reward_total", value=float(reward))
|
| lf.score_current_trace(
|
| name="pass_frac", value=float(report.get("pass_frac") or 0.0)
|
| )
|
| lf.score_current_trace(
|
| name="resolved", value=1.0 if report.get("resolved") else 0.0
|
| )
|
| lf.score_current_trace(
|
| name="n_tool_calls", value=float(len(self._tool_calls))
|
| )
|
| if isinstance(report.get("shaped"), dict):
|
| for sk, sv in report["shaped"].items():
|
| lf.score_current_trace(name=f"shaped_{sk}", value=float(sv))
|
| except Exception:
|
| pass
|
| if self._root is not None:
|
| try:
|
| self._root.update(
|
| output={
|
| "reward_total": reward,
|
| "n_tool_calls": len(self._tool_calls),
|
| "tools_used": [c["name"] for c in self._tool_calls],
|
| "resolution": report.get("resolution"),
|
| "shaped": report.get("shaped"),
|
| }
|
| )
|
| except Exception:
|
| pass
|
| else:
|
| reward = _grade()
|
|
|
| self._last_report = report
|
| self._last_reward = reward
|
| _vlog(
|
| "env.get_reward.done",
|
| instance_id=self._task.get("instance_id"),
|
| reward_total=reward,
|
| pass_frac=report.get("pass_frac"),
|
| resolution=report.get("resolution"),
|
| resolved=report.get("resolved"),
|
| error=report.get("error"),
|
| n_tool_calls=len(self._tool_calls),
|
| shaped=report.get("shaped"),
|
| )
|
|
|
|
|
| try:
|
| import wandb
|
|
|
| if wandb.run is not None:
|
| tools_used = [c.get("name") for c in self._tool_calls]
|
| shaped = report.get("shaped") if isinstance(report.get("shaped"), dict) else {}
|
| wandb.log(
|
| {
|
| "rollout/reward_total": float(reward),
|
| "rollout/docker_reward": float(report.get("docker_reward") or 0.0),
|
| "rollout/pass_frac": float(report.get("pass_frac") or 0.0),
|
| "rollout/resolved": 1.0 if report.get("resolved") else 0.0,
|
| "rollout/n_tool_calls": float(len(self._tool_calls)),
|
| "rollout/instance_id": self._task.get("instance_id") or "",
|
| "rollout/tools_used": ",".join(tools_used),
|
| "rollout/resolution": str(report.get("resolution") or ""),
|
| "rollout/grader_error": str(report.get("error") or ""),
|
| "rollout/shaped_tool_bonus": float(shaped.get("tool_bonus") or 0.0),
|
| "rollout/shaped_edit_bonus": float(shaped.get("edit_bonus") or 0.0),
|
| "rollout/shaped_patch_bonus": float(shaped.get("patch_bonus") or 0.0),
|
| }
|
| )
|
| _vlog("env.get_reward.wandb_logged")
|
| except Exception as e:
|
| _vlog("env.get_reward.wandb_failed", error=str(e))
|
|
|
| self._end_trace()
|
| _vlog("env.get_reward.trace_closed", trace_id=self._trace_id)
|
| return reward
|
|
|
|
|
| def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
| """Normalize HF SWE Lite rows for TRL GRPO + DeepAgentSWEEnv.
|
|
|
| Renames `image` -> `docker_image` so TRL does not treat it as multimodal.
|
| Converts `prompt` to conversational chat messages (required for tool calling).
|
| """
|
| from deepagents_rl.data.swe_bench_lite import build_agent_prompt, from_hf_record
|
|
|
| rows: list[dict[str, Any]] = []
|
| for rec in records:
|
| if "fail_to_pass_json" in rec or (
|
| "prompt" in rec and isinstance(rec.get("prompt"), str) and "image" in rec
|
| ):
|
| task = from_hf_record(rec)
|
| else:
|
| task = dict(rec)
|
|
|
|
|
| messages = [
|
| {
|
| "role": "user",
|
| "content": build_agent_prompt(
|
| str(task.get("problem_statement") or ""),
|
| instance_id=str(task.get("instance_id") or task.get("task_id") or ""),
|
| ),
|
| }
|
| ]
|
|
|
| docker_image = str(task.get("docker_image") or task.get("image") or "")
|
| rows.append(
|
| {
|
| "prompt": messages,
|
| "task_id": task["task_id"],
|
| "instance_id": task["instance_id"],
|
| "repo": task["repo"],
|
| "base_commit": task["base_commit"],
|
| "docker_image": docker_image,
|
| "eval_script": task.get("eval_script") or "",
|
| "fail_to_pass": list(task.get("fail_to_pass") or []),
|
| "pass_to_pass": list(task.get("pass_to_pass") or []),
|
| "problem_statement": task.get("problem_statement") or "",
|
| }
|
| )
|
| return rows
|
|
|
