Shaped rewards + root listing in reset

src/deepagents_rl/train/swe_env.py

@@ -447,6 +447,12 @@ class DeepAgentSWEEnv:
                 span.update(output={"sandbox": str(self._sandbox)})
         else:
             _materialize()
+        # Pre-list root so the model sees real paths even before its first tool call.
+        root_listing = ""
+        try:
+            root_listing = self._root_listing()
+        except Exception as e:  # noqa: BLE001
+            root_listing = f"(root listing failed: {e})"
         _vlog(
             "env.reset.done",
             instance_id=instance_id,
@@ -454,16 +460,18 @@ class DeepAgentSWEEnv:
             n_f2p=len(fail_to_pass),
             n_p2p=len(pass_to_pass),
             langfuse_trace_id=self._trace_id,
+            root_listing_chars=len(root_listing),
         )
         # TRL appends this string to the last user message — force tool use.
         return (
-            "\n\nCRITICAL REMINDER (edge cases):\n"
-            "- Use structured tool/function calls only (not markdown ```json``` blobs).\n"
-            "- FIRST tool call MUST be ls with path='.' (repo root). Never ls 'src'/'lib'/'app'.\n"
-            "- Only read_file/edit_file paths that appeared in ls/glob output.\n"
-            "- Prefer edit_file with unique old_string; keep the diff minimal.\n"
-            "- If a path is missing, ls '.' again and pivot — do not retry the same missing path.\n"
-            "- When done editing, stop tools and reply: DONE <summary>."
+            "\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:
@@ -664,11 +672,45 @@ class DeepAgentSWEEnv:
         """
         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:
-        """Export git diff and grade via remote Docker SWE evaluator."""
+        """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"),
@@ -677,11 +719,17 @@ class DeepAgentSWEEnv:
         )
 
         def _grade() -> float:
-            nonlocal report
+            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")
-                return 0.0
+                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
 
@@ -697,31 +745,45 @@ class DeepAgentSWEEnv:
             except Exception as e:  # noqa: BLE001
                 report = {"error": f"git_diff_failed: {e}"}
                 _vlog("env.get_reward.git_diff_failed", error=str(e))
-                return 0.0
+                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")
-                return 0.0
-
-            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,
-                )
-            except Exception as e:  # noqa: BLE001
-                report = {
-                    "error": f"grader_failed: {e}",
-                    "traceback": traceback.format_exc()[-1500:],
-                }
-                _vlog("env.get_reward.grader_failed", error=str(e))
-                return 0.0
+            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:  # noqa: BLE001
+                    report = {
+                        "error": f"grader_failed: {e}",
+                        "traceback": traceback.format_exc()[-1500:],
+                    }
+                    _vlog("env.get_reward.grader_failed", error=str(e))
+                    docker_reward = 0.0
 
-            return float(report.get("reward_total") or 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:
@@ -740,6 +802,8 @@ class DeepAgentSWEEnv:
                     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"),
@@ -759,6 +823,9 @@ class DeepAgentSWEEnv:
                     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:
@@ -769,6 +836,7 @@ class DeepAgentSWEEnv:
                                 "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:
@@ -787,6 +855,7 @@ class DeepAgentSWEEnv:
             resolved=report.get("resolved"),
             error=report.get("error"),
             n_tool_calls=len(self._tool_calls),
+            shaped=report.get("shaped"),
         )
 
         # W&B: one point per completed rollout (train + ckpt eval)
@@ -795,9 +864,11 @@ class DeepAgentSWEEnv:
 
             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)),
@@ -805,6 +876,9 @@ class DeepAgentSWEEnv:
                         "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")