feat(fireworks): fold ECHO env loss into the client loss path (single pass)

cookbooks/swe-rl/README.md

@@ -0,0 +1,183 @@
+# SWE-RL
+
+End-to-end agentic-RL recipe for software engineering: train on [R2E-Gym](https://huggingface.co/datasets/R2E-Gym/R2E-Gym-Subset) (R2E-Gym Subset, 4,578 bug-fix tasks across 12 Python repos), validate on [SWE-bench Verified](https://www.swebench.com/) (500 real GitHub issues). The agent harness is [mini-swe-agent](https://github.com/SWE-agent/mini-swe-agent); the base model is `Qwen/Qwen3.5-9B`.
+
+This cookbook deliberately ships **no custom AgentFlow and no custom evaluator** — it's a thin wrapper around primitives that already live in `rllm/`. The flow is `mini-swe-agent` running as a CLI inside per-task sandboxes, and the evaluator is each task's own `tests/test.sh`. Both `r2egym` and `harbor:swebench-verified` ship that verifier with the dataset.
+
+## Architecture
+
+```
+AgentTrainer.train()
+  │
+  ├── for each task: launch a sandbox (Modal / Daytona / Docker)
+  │       │
+  │       └── mini-swe-agent --yolo --task=<problem_statement>
+  │             │   (multi-turn shell loop; each LLM call → gateway)
+  │             │
+  │             └── rLLM gateway routes to the trainer-hosted policy,
+  │                  capturing the full trajectory (prompt + response
+  │                  tokens + sampling params per turn).
+  │
+  └── verifier: tests/test.sh inside the sandbox
+        │   run_tests.sh for r2egym;  SWE-bench harness for verified.
+        │
+        └── writes /logs/verifier/reward.txt  →  RL reward signal
+```
+
+The trainer never parses tool calls or model outputs directly. The agent harness owns the action loop; the gateway owns the trajectory; the in-sandbox verifier owns the reward. This is the same setup as `examples/harbor_swe/` — the cookbook packages it as a versioned, installable recipe.
+
+## Installation
+
+```bash
+uv pip install -e ".[tinker]"                       # rllm + tinker backend
+uv pip install --no-deps -e cookbooks/swe-rl        # this cookbook (registers prepare_data)
+```
+
+`mini-swe-agent` is installed automatically into each task sandbox on first run — no host-side install required.
+
+## Datasets
+
+```bash
+python cookbooks/swe-rl/prepare_data.py
+# or, faster smoke run:
+python cookbooks/swe-rl/prepare_data.py --train-limit 50 --val-limit 20
+```
+
+This pulls:
+
+| Dataset | Role | Source | Verifier |
+|---|---|---|---|
+| `r2egym` | train (4,578) | `R2E-Gym/R2E-Gym-Subset` | in-sandbox `run_tests.sh`, pytest-output equality (`tests/test.sh`) |
+| `harbor:swebench-verified` | eval (500) | `princeton-nlp/SWE-bench_Verified` | official SWE-bench harness (F2P / P2P) |
+
+Both materialize as Harbor-format task directories (`task.toml`, `instruction.md`, `environment/Dockerfile`, `tests/test.sh`).
+
+## Training
+
+### Tinker (single-machine, LoRA)
+
+```bash
+bash cookbooks/swe-rl/train_tinker.sh
+```
+
+Defaults: Qwen/Qwen3.5-9B + LoRA rank 32, GRPO with compact filtering, 64 parallel Modal sandboxes, async rollout/training. Override anything via Hydra:
+
+```bash
+SWE_SANDBOX_BACKEND=docker bash cookbooks/swe-rl/train_tinker.sh \
+    model.name=Qwen/Qwen3-8B \
+    rllm.workflow.n_parallel_tasks=32
+```
+
+For a simpler on-policy loop (generate a full batch, then one optimizer step — easier to debug), use the synchronous variant:
+
+```bash
+bash cookbooks/swe-rl/train_tinker_sync.sh
+```
+
+It drops `async_training` and uses a real `data.train_batch_size` (default 4; effective batch = `train_batch_size × group_size`).
+
+### Verl (distributed GPU)
+
+```bash
+uv pip install -e ".[verl]"
+bash scripts/install_megatron.sh <cu128|cu129|...>
+bash cookbooks/swe-rl/train_verl.sh
+```
+
+vLLM rollouts + FSDP/Megatron training. Sandboxes still run mini-swe-agent — only the trainer hosting changes.
+
+### Fireworks (managed, LoRA)
+
+```bash
+uv pip install -e ".[fireworks]"
+export FIREWORKS_API_KEY=...
+bash cookbooks/swe-rl/train_fireworks.sh
+```
+
+Same async GRPO + compact-filtering recipe as `train_tinker.sh`, but the trainer
+job and inference deployment are provisioned on Fireworks at startup and torn
+down on shutdown. Defaults to `accounts/fireworks/models/qwen3p5-9b` + LoRA rank
+32 on the `qwen3p5-9b-256k-lora` training shape (Fireworks ships a 3.5-9B LoRA
+shape but no 3.5-4B; swap `model.name` / `model.tokenizer_model` /
+`fireworks_config.policy_trainer_shape_id` together to change it — see
+[`docs/backends/fireworks.mdx`](../../docs/backends/fireworks.mdx) for the
+model/shape catalog). The synchronous (on-policy) variant is
+`train_fireworks_sync.sh`.
+
+### ECHO (train on environment feedback)
+
+[ECHO](https://arxiv.org/abs/2605.24517) adds a cross-entropy loss on the
+environment-observation tokens (tool/terminal output) that the policy already
+conditions on but GRPO never trains. For a terminal/SWE agent this turns every
+rollout — including the many that fail — into dense supervision, at no extra
+rollout or forward-pass cost. It uses GRPO's advantages unchanged; the only
+difference is the extra loss term.
+
+Flip GRPO → ECHO with one override on any backend (verl / tinker / fireworks):
+
+```bash
+# tinker (async or sync), verl, or fireworks — same switch:
+bash cookbooks/swe-rl/train_tinker.sh    rllm.algorithm.adv_estimator=echo
+bash cookbooks/swe-rl/train_verl.sh      algorithm.adv_estimator=echo
+bash cookbooks/swe-rl/train_fireworks.sh rllm.algorithm.adv_estimator=echo
+```
+
+`adv_estimator=echo` defaults the env-loss weight λ to the paper's 0.05. Tune it
+explicitly with `rllm.algorithm.env_loss_coef=<λ>` (productive range 0.01–0.05;
+`0.0` reproduces plain GRPO). It is implemented as an `env_prediction`
+[auxiliary loss](../../design/auxiliary-losses.md); watch
+`actor/aux_env_prediction_loss` (verl) / `train/aux_*` (tinker, fireworks) to
+confirm the environment-prediction loss is falling.
+
+> On verl the env term shares GRPO's single forward pass (free, exact). On
+> tinker/fireworks (managed training services with fixed server-side loss
+> kernels) it is a second, gradient-accumulated `cross_entropy` pass over the
+> same rollouts — no extra rollouts, but one extra backward. λ may need
+> per-backend retuning since loss normalization differs across services.
+
+## Evaluation (no training)
+
+```bash
+rllm eval harbor:swebench-verified \
+    --agent mini-swe-agent \
+    --model Qwen/Qwen3.5-9B \
+    --base-url http://localhost:8000/v1 \
+    --max-examples 20
+```
+
+Per-task results land in `~/.rllm/eval_results/`; aggregated resolve rate is printed at the end. `rllm view` opens the per-task trajectory UI.
+
+## Sandbox backend
+
+Training uses rLLM's own `SandboxedAgentFlow` path (`AgentFlowEngine`) — not the
+remote Harbor runtime. `mini-swe-agent` runs inside one sandbox per task, created
+by `SandboxTaskHooks`. Pick a backend via the `SWE_SANDBOX_BACKEND` env var:
+
+| Backend | Setup | Notes |
+|---|---|---|
+| `modal` | `pip install modal` + `modal token new` | Default for training — per-task billing, scales to many parallel sandboxes. |
+| `daytona` | `pip install daytona` + `DAYTONA_API_KEY` | Cloud sandboxes; scales to thousands in parallel. |
+| `docker` | local | Fastest iteration; needs the Docker daemon and ~20 GB free disk. |
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `prepare_data.py` | Pulls `r2egym` (train) and `harbor:swebench-verified` (eval) |
+| `train.py` | Loads the two datasets, hands them to `AgentTrainer` |
+| `train_tinker.sh` | Tinker backend — Qwen3.5-9B LoRA, GRPO + async, Modal sandboxes |
+| `train_tinker_sync.sh` | Tinker backend — synchronous (on-policy) variant, simpler for testing |
+| `train_fireworks.sh` | Fireworks backend — Qwen3.5-9B LoRA, GRPO + async, managed trainer/deployment |
+| `train_fireworks_sync.sh` | Fireworks backend — synchronous (on-policy) variant, simpler for testing |
+| `train_verl.sh` | Verl backend — same recipe with vLLM + FSDP |
+| `test.py` | Catalog wiring + harness import smoke tests |
+| `pyproject.toml` | Cookbook metadata (no entry points — the harness is in-tree) |
+
+## Why no custom flow or evaluator?
+
+Other cookbooks in this repo (`finqa`, `math`, `deepcoder`, …) ship a custom AgentFlow because their workloads either fit in a single LLM turn or need bespoke tool wiring. SWE doesn't — the existing in-tree primitives already cover it:
+
+- **`rllm.harnesses.mini_swe_agent`** is the agent. It exposes the `mini-swe-agent` CLI as an rLLM harness (installs in-sandbox on first run; reads the gateway URL from the env; logs to `/tmp/mini-swe-agent.log`).
+- **Per-task `tests/test.sh`** is the evaluator. The sandbox-shell verifier kind (`rllm.eval.script_evaluator`) reads `/logs/verifier/reward.txt` and returns it as the RL reward. For `harbor:swebench-verified`, that script invokes the official SWE-bench harness; for `r2egym`, it runs the image's own `/testbed/run_tests.sh` and scores 1.0 iff the pytest output matches the row's expected output.
+
+The only thing this cookbook adds on top is the recipe: dataset pairing, sampling/optimizer hyperparams, and the `mini-swe-agent` harness selection. Forking `train_tinker.sh` is the place to start customizing.

cookbooks/swe-rl/prepare_data.py

@@ -0,0 +1,66 @@
+"""Pull the train + eval datasets for the swe-rl cookbook.
+
+Both are sandbox-format benchmarks (per-task ``environment/Dockerfile`` +
+``tests/test.sh`` verifier). The training set is rLLM's native
+``r2egym`` (R2E-Gym Subset, 4,578 bug-fix tasks across 12 Python repos,
+each shipped as a per-instance Docker image graded by the image's own
+``/testbed/run_tests.sh``). The eval set is ``harbor:swebench-verified``
+(500 real-world GitHub issues, evaluated against the official SWE-bench
+harness inside the sandbox).
+
+This script is a thin wrapper around ``rllm dataset pull`` so the
+cookbook can be used end-to-end with a single command. Re-runs are
+no-ops once the on-disk benchmark directory exists.
+
+Usage::
+
+    python cookbooks/swe-rl/prepare_data.py
+    # or, smoke run with a small training cap (rebuilds the benchmark dir):
+    python cookbooks/swe-rl/prepare_data.py --train-limit 50
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+
+TRAIN_DATASET = "r2egym"
+VAL_DATASET = "harbor:swebench-verified"
+
+
+def _pull(name: str, limit: int | None = None) -> None:
+    cmd = [sys.executable, "-m", "rllm.cli.main", "dataset", "pull", name]
+    if limit is not None:
+        cmd += ["--limit", str(limit)]
+    print(f"[swe-rl] $ {' '.join(cmd)}", flush=True)
+    subprocess.run(cmd, check=True)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument(
+        "--train-limit",
+        type=int,
+        default=None,
+        help="Cap training tasks (default: full ~4.6K). Useful for smoke runs.",
+    )
+    ap.add_argument(
+        "--val-limit",
+        type=int,
+        default=None,
+        help="Cap eval tasks (default: full 500).",
+    )
+    args = ap.parse_args()
+
+    _pull(TRAIN_DATASET, limit=args.train_limit)
+    _pull(VAL_DATASET, limit=args.val_limit)
+
+    print(
+        f"\n[swe-rl] Done. Train: {TRAIN_DATASET}   Eval: {VAL_DATASET}\n        Run `bash cookbooks/swe-rl/train_tinker.sh` to train.",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

cookbooks/swe-rl/pyproject.toml

@@ -0,0 +1,13 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "swe-rl"
+version = "0.1.0"
+description = "End-to-end SWE-RL recipe: train on R2E-Gym, eval on SWE-bench Verified"
+requires-python = ">=3.10"
+dependencies = ["rllm"]
+
+[tool.setuptools]
+py-modules = ["prepare_data"]

cookbooks/swe-rl/test.py

@@ -0,0 +1,67 @@
+"""Smoke tests for the swe-rl cookbook.
+
+These tests don't boot sandboxes, run agents, or train. They only
+verify the wiring: the dataset names resolve in the catalog, the
+``mini-swe-agent`` harness is importable, and ``train.py`` can be
+imported without side effects.
+
+Run::
+
+    pytest cookbooks/swe-rl/test.py -v
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+from pathlib import Path
+
+# -- Dataset catalog entries --------------------------------------------------
+
+
+def _load_catalog() -> dict:
+    registry_path = Path(__file__).resolve().parents[2] / "rllm" / "registry" / "datasets.json"
+    return json.loads(registry_path.read_text())
+
+
+def test_train_dataset_in_catalog():
+    """r2egym must be registered as a sandbox dataset using mini-swe-agent."""
+    catalog = _load_catalog()
+    entry = catalog["datasets"]["r2egym"]
+    assert entry["category"] == "code"
+    assert entry["default_agent"] == "mini-swe-agent"
+    assert "train" in entry["splits"]
+
+
+def test_val_dataset_resolvable():
+    """SWE-bench Verified must be reachable either as the native row dataset
+    (``swebench_verified``) or as the Harbor sandbox dataset
+    (``harbor:swebench-verified``)."""
+    catalog = _load_catalog()
+    assert "swebench_verified" in catalog["datasets"], "native swebench_verified entry missing from registry"
+
+
+# -- Harness wiring -----------------------------------------------------------
+
+
+def test_mini_swe_agent_harness_importable():
+    """``mini-swe-agent`` is the harness this cookbook drives; it must import."""
+    mod = importlib.import_module("rllm.harnesses.mini_swe_agent")
+    assert hasattr(mod, "MiniSweAgentHarness")
+    assert mod.MiniSweAgentHarness.name == "mini-swe-agent"
+
+
+# -- Train script -------------------------------------------------------------
+
+
+def test_train_module_imports():
+    """``train.py`` must be importable without triggering Hydra or starting training."""
+    import sys
+
+    script_dir = Path(__file__).resolve().parent
+    if str(script_dir) not in sys.path:
+        sys.path.insert(0, str(script_dir))
+    mod = importlib.import_module("train")
+    assert mod.TRAIN_DATASET == "r2egym"
+    assert mod.VAL_DATASET == "swebench-verified"
+    assert callable(mod.main)