[Draft] core: harness-rollout contract + runnable example for online RL with TRL

examples/trl_harness_rollout/README.md

@@ -0,0 +1,109 @@
+# Harness rollout worker for online RL with TRL (draft)
+
+A runnable draft example of training an OpenEnv-style **agentic harness** (an agent that owns its own
+loop) with TRL's online RL (AsyncGRPO). It runs today with **no GPU** (fake mode) and with **real
+vLLM**.
+
+This example ships in the same PR as the minimal contract it builds on (`src/openenv/core/harness/`:
+`rollout.py` + `interception.py`), so the whole dynamic is reviewable in one place. The contract is the
+OpenEnv contribution. The **worker is the trainer side**: for a real integration it gets **vendored
+into TRL** (the worker conforms to TRL's `RolloutWorkerProtocol`, so it belongs with the trainer). It
+is shown here so the contract is not just interfaces but a thing you can read and run.
+
+The example vendors `rollout.py` + `interception.py` locally so it stays self-contained (runs with just
+`aiohttp` + `requests`, no heavy imports) and is copy-pasteable into TRL as-is. Those two files mirror
+the canonical contract in `src/openenv/core/harness/`. The richer interception/sandbox stack lives in
+OpenEnv PR #694.
+
+To pick it up: implement the two seams on the trainer side and wire it into AsyncGRPO. See
+`trl_adapter.py` for the exact spots.
+
+## The idea
+
+A harness is an agent that owns its loop (it calls the model, uses tools, decides when to stop). To
+train it on-policy you must run its LLM calls through the policy being trained and capture the exact
+tokens. This worker drives the harness through an interception proxy, scores the episode with the
+env's `verify()`, and emits a **message-level rollout**. It does NOT tokenize. Tokenization (the
+prefix-preserving, token-in/token-out work) is the trainer's job, behind the `generate` seam.
+
+```
+harness (owns its loop) --HTTP--> interception proxy --> HarnessRolloutWorker
+   ^                                                       |- generate(...)   [Seam 1: trainer]
+   |---------------- completion ---------------------------|- verify() -> reward
+                                                           |- emit RolloutMessages [Seam 2: trainer consumes]
+```
+
+## The two seams (what the trainer implements)
+
+- **Seam 1:** `generate(rollout_id, turn, messages, tools, sampling) -> completion_text`. The trainer
+  generates with its inference engine and records token_ids + logprobs keyed by `(rollout_id, turn)`.
+  This is also where messages-to-tokens / TITO lives.
+- **Seam 2:** the worker emits `RolloutMessages{rollout_id, messages, reward}`. The trainer stitches
+  the captured per-turn tokens into a training sample, applies the reward, computes the advantage.
+
+## What is done vs to do
+
+**Done (this example):**
+- `HarnessRolloutWorker`: drives the harness via interception, `verify()`, emits message-level rollouts.
+- A minimal interception proxy (the richer one is OpenEnv PR #694).
+- A harness with the right dynamic (a ReAct agent that owns its loop, multi-turn, interception-gated).
+- Runnable with no GPU (fake) and with real vLLM (real token_ids + logprobs capture).
+
+**To do (trainer side):** see `trl_adapter.py`. Implement `generate` (Seam 1, with TITO) and wrap the
+worker into TRL's `RolloutWorkerProtocol` (Seam 2 + wiring into `AsyncGRPOTrainer(rollout_worker=...)`).
+
+> Weight sync is a trainer-side concern. TRL handles it on the default path (`environment_factory` +
+> the built-in `AsyncRolloutWorker`). When you inject a custom `rollout_worker`, TRL sets
+> `weight_transfer = None` and `_sync_weight()` becomes a no-op, so the trainer side must wire weight
+> sync itself (worker-driven, the way OpenEnv PR #695 does it with `send_weights` +
+> `update_model_version`, or by extending `_sync_weight`). It does not touch OpenEnv. See
+> `trl_adapter.py`. This draft validates the rollout dynamic and the seam, not a converged training run.
+
+## Run
+
+```bash
+pip install aiohttp requests           # fake mode deps
+# transformers + vllm are only needed for vllm mode
+
+# no GPU: scripted generator driving the real harness dynamic
+python run.py --mode fake
+
+# GPU box, with `vllm serve <model> --port 8000` already running:
+python run.py --mode vllm --vllm-url http://localhost:8000 --model Qwen/Qwen2.5-3B-Instruct --dump captures.json
+
+# tests (no GPU)
+pytest test_rollout_worker.py
+```
+
+`--dump captures.json` (vllm mode) writes the real per-turn `(prompt_ids, completion_ids, logprobs)` to
+JSON: the exact data a TITO step would stitch into a training sample.
+
+## Files
+
+Vendored minimal contract + transport (mirror the canonical contract in `src/openenv/core/harness/`):
+
+| File | Role |
+|------|------|
+| `rollout.py` | The contract: `AgentSession`, `AgentSessionFactory`, `RolloutMessages`, `GenerateAPI`. |
+| `interception.py` | `InterceptionServer`: the OpenAI-compatible gating proxy. |
+
+The worker + the example:
+
+| File | Role |
+|------|------|
+| `rollout_worker.py` | `HarnessRolloutWorker` (the trainer side): drives the harness, emits message-level rollouts. |
+| `harness.py` | The harness (ReAct agent owning its loop) + session + arithmetic task + verifier. |
+| `generate.py` | `FakeGenerate` (no GPU) and `VLLMGenerate` (real vLLM + token capture), both Seam 1. |
+| `trl_adapter.py` | Skeleton of the trainer-side integration (where Seam 1/2 + AsyncGRPO wiring go). |
+| `run.py` | Entry point, `--mode fake|vllm`. |
+| `test_rollout_worker.py` | Unit tests, no GPU. |
+
+## Notes
+
+- `rollout.py` (contract) and `interception.py` are vendored here so the example is self-contained and
+  directly portable into TRL. They mirror the canonical contract in `src/openenv/core/harness/`.
+- The worker is harness-agnostic: it does not change when you swap the harness. The ReAct agent here
+  gives the right dynamic without a full coding agent. A follow-up can swap it for a real harness
+  (Pi / OpenCode) over OpenEnv's interception + sandbox (PR #694).
+- The core (worker + interception + message-level rollout) is framework-neutral. TRL is the first
+  integration, not the only one.

examples/trl_harness_rollout/generate.py

@@ -0,0 +1,122 @@
+"""Seam 1 (`generate`) implementations.
+
+`FakeGenerate` runs the whole harness dynamic with NO GPU and NO model (a scripted "model" that drives
+the ReAct agent deterministically). `VLLMGenerate` is the real path: it generates with vLLM and
+captures the engine's exact token_ids + logprobs.
+
+`VLLMGenerate` does naive tokenization here so the example runs. In a real trainer integration this is
+where the prefix-preserving, token-in/token-out tokenization (TITO) lives, on the trainer side. See
+`trl_adapter.py`.
+"""
+
+from __future__ import annotations
+
+import re
+import threading
+from typing import Any
+
+
+# ── No-GPU scripted model: drives the ReAct harness deterministically ──
+class FakeGenerate:
+    """Lets the full harness dynamic (interception, multi-turn, verify) run with no GPU and no model.
+
+    Turn logic mirrors the ReAct protocol: if the last message is a calc result, answer with it,
+    otherwise emit ACTION: calc(<expression from the question>).
+    """
+
+    def generate(self, *, rollout_id, turn, messages, tools, sampling) -> str:
+        last = messages[-1]["content"] if messages else ""
+        m = re.search(r"calc result:\s*(-?\d+)", last)
+        if m:
+            return f"ANSWER: {m.group(1)}"
+        question = next(
+            (
+                mm["content"]
+                for mm in messages
+                if mm.get("role") == "user" and "What is" in mm["content"]
+            ),
+            "",
+        )
+        e = re.search(r"What is (.+?)\?", question)
+        return f"ACTION: calc({e.group(1) if e else '0'})"
+
+
+# ── Real path: vLLM generation + token_ids/logprobs capture ──
+class VLLMGenerate:
+    def __init__(self, *, vllm_base_url: str, model: str, api_key: str = "token"):
+        from transformers import AutoTokenizer
+
+        self._url = vllm_base_url.rstrip("/")
+        self._model = model
+        self._key = api_key
+        self._tok = AutoTokenizer.from_pretrained(model)
+        if self._tok.pad_token is None:
+            self._tok.pad_token = self._tok.eos_token
+        self.captures: dict[tuple[str, int], dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+    def generate(self, *, rollout_id, turn, messages, tools, sampling) -> str:
+        import requests
+
+        # NAIVE tokenization. The trainer replaces this with the TITO-safe, prefix-preserving path.
+        encoded = self._tok.apply_chat_template(
+            messages,
+            tools=tools,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=False,
+        )
+        prompt_ids = _to_id_list(encoded)
+        r = requests.post(
+            f"{self._url}/v1/completions",
+            headers={"Authorization": f"Bearer {self._key}"},
+            json={
+                "model": self._model,
+                "prompt": prompt_ids,
+                "max_tokens": sampling.get("max_tokens", 256),
+                "temperature": sampling.get("temperature", 1.0),
+                "logprobs": 0,
+                "return_token_ids": True,  # real engine ids: the capture a TITO step would stitch
+            },
+            timeout=600,
+        )
+        r.raise_for_status()
+        choice = r.json()["choices"][0]
+        with self._lock:
+            self.captures[(rollout_id, turn)] = {
+                "prompt_ids": prompt_ids,
+                "completion_ids": choice["token_ids"],
+                "logprobs": choice["logprobs"]["token_logprobs"],
+                "text": choice.get("text", ""),
+            }
+        return choice.get("text", "")
+
+    def capture_summary(self) -> str:
+        with self._lock:
+            n = len(self.captures)
+            toks = sum(len(c["completion_ids"]) for c in self.captures.values())
+        return f"{n} captured turns, {toks} real completion tokens (with logprobs)"
+
+    def dump(self, path: str) -> int:
+        import json
+
+        with self._lock:
+            records = [
+                {"rollout_id": rid, "turn": t, **cap}
+                for (rid, t), cap in sorted(self.captures.items())
+            ]
+        with open(path, "w") as f:
+            json.dump(records, f, indent=2)
+        return len(records)
+
+
+def _to_id_list(encoded: Any) -> list[int]:
+    """Normalize apply_chat_template output to a flat list[int] across transformers versions."""
+    out = encoded
+    if hasattr(out, "input_ids"):  # BatchEncoding
+        out = out["input_ids"]
+    if hasattr(out, "tolist"):  # tensor
+        out = out.tolist()
+    if out and isinstance(out[0], list):  # batched [[...]] -> [...]
+        out = out[0]
+    return [int(t) for t in out]

examples/trl_harness_rollout/harness.py

@@ -0,0 +1,131 @@
+"""A harness (an agent that owns its own loop) plus its OpenEnv session, task, and verifier.
+
+The harness here is a small ReAct agent with a calculator tool, talking OpenAI over HTTP to the
+interception proxy. It owns its loop: it decides when to call the tool and when to answer. This gives
+the right dynamic (agent-owned loop, multi-turn, interception-gated) without a full coding agent. To
+raise fidelity, swap this for a real harness (Pi / OpenCode over OpenEnv's InterceptionServer, PR #694)
+in a follow-up: the rollout worker does not change.
+
+Exposes `ProxyAgentSessionFactory`, which satisfies the worker's `AgentSessionFactory` protocol.
+"""
+
+from __future__ import annotations
+
+import re
+import threading
+from typing import Any
+
+import requests
+from interception import InterceptionServer
+
+
+# ── Task + verifier (the env): verifiable arithmetic ──
+TASKS = [
+    {"id": "t1", "prompt": "What is 347 * 29?", "answer": 10063},
+    {"id": "t2", "prompt": "What is 9876 / 12?", "answer": 823},
+    {"id": "t3", "prompt": "What is 1234 + 5678 - 999?", "answer": 5913},
+    {"id": "t4", "prompt": "What is (84 + 16) * 47?", "answer": 4700},
+]
+
+_SYS = (
+    "You solve an arithmetic question. You are bad at mental math, so you MUST use the calculator.\n"
+    "To use it, output exactly one line and nothing else:\n"
+    "ACTION: calc(<expression>)\n"
+    "Only after you have seen a calc result, output exactly one line:\n"
+    "ANSWER: <number>\n"
+    "Never output ANSWER before using ACTION: calc(...) at least once."
+)
+_SAFE = re.compile(r"^[0-9+\-*/(). ]+$")
+
+
+def calc(expr: str) -> str:
+    expr = expr.strip()
+    if not _SAFE.match(expr):
+        return "error: invalid expression"
+    try:
+        val = eval(expr, {"__builtins__": {}}, {})  # noqa: S307 (regex-restricted input)
+        if isinstance(val, float) and val.is_integer():
+            val = int(val)
+        return str(val)
+    except Exception as e:  # noqa: BLE001
+        return f"error: {e}"
+
+
+# ── The agent: owns its loop, talks OpenAI over HTTP to the interception proxy ──
+def react_agent(
+    base_url: str, rollout_id: str, prompt: str, max_turns: int = 6
+) -> None:
+    url = f"{base_url}/rollout/{rollout_id}/v1/chat/completions"
+    messages = [
+        {"role": "system", "content": _SYS},
+        {"role": "user", "content": prompt},
+    ]
+    used_calc = False
+    for _ in range(max_turns):
+        content = requests.post(url, json={"messages": messages}, timeout=600).json()
+        content = content["choices"][0]["message"]["content"]
+        messages.append({"role": "assistant", "content": content})
+        m = re.search(
+            r"ACTION:\s*calc\((.+)\)", content
+        )  # greedy: capture up to the last ) on the line
+        if m:
+            used_calc = True
+            messages.append(
+                {"role": "user", "content": f"calc result: {calc(m.group(1))}"}
+            )
+            continue
+        if "ANSWER:" in content:
+            if used_calc:
+                break
+            messages.append(
+                {"role": "user", "content": "Use ACTION: calc(...) first, then ANSWER."}
+            )
+            continue
+        messages.append(
+            {"role": "user", "content": "Output ACTION: calc(...) or ANSWER: <number>."}
+        )
+    requests.post(f"{base_url}/rollout/{rollout_id}/exit", json={}, timeout=30)
+
+
+# ── OpenEnv session over the proxy (satisfies AgentSession) ──
+class ProxyAgentSession:
+    def __init__(self, proxy: InterceptionServer, rollout_id: str, q, answer: int):
+        self._proxy, self._rid, self._q, self._answer = proxy, rollout_id, q, answer
+        self._last = ""
+
+    def next_request(self) -> dict | None:
+        try:
+            req_id = self._q.get(timeout=120)
+        except Exception:
+            return None
+        return None if req_id is None else self._proxy.get_intercept(req_id)
+
+    def deliver(self, intercept: dict, completion_text: str) -> None:
+        self._last = completion_text
+        self._proxy.deliver(intercept["request_id"], completion_text)
+
+    def verify(self) -> Any:
+        m = re.search(r"ANSWER:\s*(-?\d+)", self._last)
+        ok = (
+            m is not None and int(m.group(1)) == self._answer
+        )  # reward stays in the env
+        return type("V", (), {"env_reward": 1.0 if ok else 0.0})()
+
+    def close(self) -> None:
+        self._proxy.unregister_rollout(
+            self._rid
+        )  # drop the per-rollout queue, no leak on long runs
+
+
+class ProxyAgentSessionFactory:
+    def __init__(self, proxy: InterceptionServer):
+        self._proxy = proxy
+
+    def create(self, *, task: Any, rollout_id: str) -> ProxyAgentSession:
+        q = self._proxy.register_rollout(rollout_id)
+        threading.Thread(
+            target=react_agent,
+            args=(self._proxy.base_url, rollout_id, task["prompt"]),
+            daemon=True,
+        ).start()
+        return ProxyAgentSession(self._proxy, rollout_id, q, task["answer"])