diff --git a/configs/debug/replay/README.md b/configs/debug/replay/README.md new file mode 100644 index 0000000000..88b0530199 --- /dev/null +++ b/configs/debug/replay/README.md @@ -0,0 +1,28 @@ +# Replay — Debug Configs + +Minimal end-to-end configs for the replay derivation tasksets (`environments/replay_*_v1`, +thin subclasses of `verifiers.v1.tasksets.replay`), which turn +saved rollouts (`/rollouts/step_*/train_rollouts.jsonl`) back into training tasks. Each config +mixes a fresh `reverse-text-v1` env with one replay env via ratios (ratios are all-or-none across +train envs), using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the policy and the `null` +harness with the subprocess runtime for both envs. + +| Config | Taskset | Buffer | Notes | +|---|---|---|---| +| `offline_recheck.toml` | `replay-recheck-v1` | offline, a prior run's `rollouts` dir | edit `buffer_dir` to a real prior run first | +| `online_recheck.toml` | `replay-recheck-v1` | `buffer_dir = "self"` — this run's own rollouts | forces whole-group dispatch; buffer fills as the run trains | + +Every derivation is scored by its `inner` taskset, which must reproduce the source run's taskset +config. The replay env gets extra `max_input_tokens` headroom because its seeds carry whole source +conversations. `replay-continue-v1` (`anchor = "compaction" | "tool-call"`) has no debug config here: +reverse-text rollouts neither compact nor call tools — point one at an agentic run's buffer instead. + +## Run the debug configs + +```bash +# Offline recheck (edit buffer_dir in the TOML to a real prior run's rollouts dir first) +uv run rl @ configs/debug/replay/offline_recheck.toml --output-dir outputs/replay_recheck + +# Online self-recheck (no prior run needed — the buffer fills as the run trains) +uv run rl @ configs/debug/replay/online_recheck.toml --output-dir outputs/replay_recheck_online +``` diff --git a/configs/debug/replay/offline_recheck.toml b/configs/debug/replay/offline_recheck.toml new file mode 100644 index 0000000000..03fbf172aa --- /dev/null +++ b/configs/debug/replay/offline_recheck.toml @@ -0,0 +1,55 @@ +# Debug run mixing fresh reverse-text rollouts with replayed "recheck" tasks. The +# replay env re-serves finished conversations from a prior run's saved rollouts +# (/rollouts/step_*/train_rollouts.jsonl) with an appended check-your-work +# turn, scored by the original (inner) taskset. Edit `buffer_dir` below to point +# at a real prior run before launching (CLI can't index into env lists; TOML +# overlays replace the env list wholesale): +# uv run rl @ configs/debug/replay/offline_recheck.toml + +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "replay-debug" +name = "debug-offline-recheck" + +[orchestrator] +batch_size = 128 +group_size = 16 + +[orchestrator.renderer] +name = "qwen3" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +# Fresh reverse-text rollouts (3/4 of the batch). +[[orchestrator.train.env]] +name = "reverse-text" +ratio = 3.0 +taskset = { id = "reverse-text-v1" } +harness = { id = "null", runtime = { type = "subprocess" } } + +# Replayed recheck tasks (1/4). Ratios are all-or-none: every train env sets one. +[[orchestrator.train.env]] +name = "replay-recheck" +ratio = 1.0 +harness = { id = "null", runtime = { type = "subprocess" } } +[orchestrator.train.env.taskset] +id = "replay-recheck-v1" +# Placeholder: a prior run's saved-rollout dir (or the run dir containing it). +buffer_dir = "outputs/prior_run/rollouts" + +# Replay rollouts are scored by the original taskset, so `inner` must reproduce the +# source run's taskset config (here the base env's defaults). +[orchestrator.train.env.taskset.inner] +id = "reverse-text-v1" + +[trainer.optim] +lr = 3e-6 + +[inference] +gpu_memory_utilization = 0.5 diff --git a/configs/debug/replay/online_recheck.toml b/configs/debug/replay/online_recheck.toml new file mode 100644 index 0000000000..a73b7073fa --- /dev/null +++ b/configs/debug/replay/online_recheck.toml @@ -0,0 +1,57 @@ +# Debug run where the model re-checks its own rollouts as it trains. The recheck env +# replays this run's freshly written rollouts (buffer_dir = "self" resolves to +# /rollouts before the env servers spawn): each task re-serves a finished +# conversation with an appended check-your-work turn, scored by the original (inner) +# taskset. The online buffer samples a fresh source per request — forcing whole-group +# dispatch so GRPO groups share one source. +# uv run rl @ configs/debug/replay/online_recheck.toml + +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "replay-debug" +name = "debug-online-recheck" + +[orchestrator] +batch_size = 128 +group_size = 16 + +[orchestrator.renderer] +name = "qwen3" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +# Fresh reverse-text rollouts (3/4 of the batch) — also what fills the buffer. +[[orchestrator.train.env]] +name = "reverse-text" +ratio = 3.0 +taskset = { id = "reverse-text-v1" } +harness = { id = "null", runtime = { type = "subprocess" } } + +# Recheck tasks over this run's own rollouts (1/4). Ratios are all-or-none. +[[orchestrator.train.env]] +name = "replay-recheck" +ratio = 1.0 +harness = { id = "null", runtime = { type = "subprocess" } } +[orchestrator.train.env.taskset] +id = "replay-recheck-v1" +buffer_dir = "self" # implies online +# Recheck only the fresh env's rollouts. (Unset, source_envs replays every env except +# replay envs; listing a replay env by name opts into chained derivations.) +source_envs = ["reverse-text"] + +# Rechecks are scored by the original taskset, so `inner` must reproduce the source +# run's taskset config (here the base env's defaults). +[orchestrator.train.env.taskset.inner] +id = "reverse-text-v1" + +[trainer.optim] +lr = 3e-6 + +[inference] +gpu_memory_utilization = 0.5 diff --git a/deps/verifiers b/deps/verifiers index 22d6333a00..cd08b8dcce 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 22d6333a00fad300bd4b1c1fcc947b58a2605280 +Subproject commit cd08b8dccec0f89cd68272cffc6bc313c72dd350 diff --git a/docs/configuration.md b/docs/configuration.md index 42bdce744c..5e108d5a24 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -151,10 +151,11 @@ Training environments are an array of tables — set one per env, optionally wit id = "math-env" name = "gsm8k" args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } +ratio = 0.75 # 75% of the batch [[orchestrator.train.env]] id = "reverse-text" -ratio = 0.25 # 25% of batches; remaining 75% goes to math-env +ratio = 0.25 # ratios are relative weights; setting one means setting all [[orchestrator.eval.env]] id = "math-env" @@ -166,6 +167,10 @@ args = { dataset_name = "openai/gsm8k", dataset_subset = "main" } The same `id` can appear multiple times across train and eval (or with different `args`) — useful for evaluating on a held-out split of the env you're training on, or comparing two configurations of the same env side by side. When `id` is reused, set a distinct `name` on each entry; `name` defaults to `id` and must be unique across all envs in the same group. +`ratio` is all-or-none: once any train env sets one, every train env must. + +One special environment ships with `prime-rl` itself: the `replay-v1` taskset replays a run's saved rollouts (its own or a prior run's) as derived training tasks — see [Replay](replay.md). + ### Environment Variables OS environment variables exported into launched component process(es). In `rl` configs, top-level `[env_vars]` applies to trainer, inference, and orchestrator: diff --git a/docs/mint.json b/docs/mint.json index 6fc234f12b..6fd886e723 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -10,6 +10,7 @@ "inference", "scaling", "algorithms", + "replay", "advanced", "development" ] diff --git a/docs/replay.md b/docs/replay.md new file mode 100644 index 0000000000..a5ebce3440 --- /dev/null +++ b/docs/replay.md @@ -0,0 +1,95 @@ +# Replay + +The replay tasksets turn a run's saved rollouts back into training tasks. Every RL run writes its train rollouts to `/rollouts/step_N/train_rollouts.jsonl` (one WireTrace per line, stamped with `info.prime_rl` metadata); a replay taskset indexes those files as a buffer and serves tasks derived from them — resume a conversation from a compaction point or a tool call, or ask the model to re-check a finished attempt. A replay env is a normal `[[orchestrator.train.env]]` entry mixed into training with a `ratio`, over either a previous run's rollouts (offline) or the current run's own (online, `buffer_dir = "self"`). + +The machinery is split like verifiers' `harbor`/`textarena` tasksets: the shared base (buffer, trace surgery, scoring delegation) lives in verifiers as `verifiers.v1.tasksets.replay`, and each derivation is a thin subclass package under `environments/`: + +| Taskset id | Derived task | +|---|---| +| `replay-continue-v1` | Resume a source rollout mid-way — from a context-compaction point (`anchor = "compaction"`: the seed is the post-compaction prompt the original harness restarted from; one task per compaction) or right after a tool result (`anchor = "tool-call"`: one deterministically sampled resumable point per source; only points where every issued call has its result). | +| `replay-recheck-v1` | The finished conversation plus an appended check-your-work user turn (`instruction`); one task per rollout, seeded from its final state. | + +Every derivation is scored by the original taskset: `inner` must reproduce the source run's taskset config, and the inner rewards land on the trace with their original names and weights. The replay base also delegates tools, setup, and finalize to the inner taskset, so tool-using tasks replay with their tools intact. Each drawn source spawns a fresh GRPO group. + +## Table of Contents + +- [Offline vs Online Buffers](#offline-vs-online-buffers) +- [Configuration](#configuration) +- [Metrics](#metrics) +- [Writing a New Derivation](#writing-a-new-derivation) +- [Constraints and Caveats](#constraints-and-caveats) + +## Offline vs Online Buffers + +**Offline** (the default): point `buffer_dir` at a finished run's rollouts. The buffer is indexed once at env-server startup — steps are scanned newest-first up to the buffer's capacity (the newest few thousand candidates are retained) — and each task index maps deterministically to one candidate. GRPO group members dispatched as independent rollouts therefore still bind the same source rollout. + +**Online** (`buffer_dir = "self"`, which implies `online = true`): the buffer is the run's own growing rollout dir. The orchestrator resolves the `"self"` sentinel to `/rollouts` before the env servers spawn. Four consequences: + +- **Barrier semantics.** The jsonl is written non-atomically, so an online buffer only reads steps whose sibling `train_rollouts.bin` exists — the orchestrator writes and closes the jsonl strictly before the atomic rename that creates the `.bin`, so the barrier file marks the jsonl complete. This means online replay requires the filesystem rollout transport (a run with a non-filesystem transport never writes the `.bin`). Offline buffers skip the barrier: finished runs' files are complete by definition. +- **Group dispatch.** Every request samples a fresh source rollout, so the whole GRPO group must share one draw — the taskset forces whole-group dispatch (`REQUIRES_GROUP_ROLLOUTS`), and all group members arrive at the env server as a single request. +- **Startup.** Until the run has produced its first replayable step, replay requests wait briefly, then fail (with a logged warning) so their dispatch permits free up for the source envs — the orchestrator retries replay groups naturally, and they start succeeding once the first step's barrier appears. The buffer rescans for new steps during training, and its capacity doubles as recency eviction: the newest candidates are retained, keeping the buffer on recent policy behavior. +- **Choose your sources.** A `"self"` buffer sees every train env's saved rollouts — including the replay envs' own. By default (`source_envs` unset) replay-derived records are skipped: rechecking your own rechecks is a feedback loop unless chosen deliberately. Listing env names replays exactly those — and naming a replay env is the deliberate opt-in for chained derivations (a recheck env sourcing another recheck env is depth-2 self-correction; the env topology *is* the depth control, and scoring/tools/provisioning always resolve to the innermost original task). + +## Configuration + +The derivation packages are installed via the `envs` extra. A replay env's taskset config is flat — the base fields plus the package's own: + +```toml +[orchestrator] +batch_size = 128 +group_size = 16 + +# Fresh rollouts (3/4 of the batch) — also what fills the buffer. +[[orchestrator.train.env]] +name = "reverse-text" +ratio = 3.0 +taskset = { id = "reverse-text-v1" } +harness = { id = "null", runtime = { type = "subprocess" } } + +# Recheck tasks over this run's own rollouts (1/4). Ratios are all-or-none. +[[orchestrator.train.env]] +name = "replay-recheck" +ratio = 1.0 +harness = { id = "null", runtime = { type = "subprocess" } } + +[orchestrator.train.env.taskset] +id = "replay-recheck-v1" +buffer_dir = "self" +source_envs = ["reverse-text"] + +[orchestrator.train.env.taskset.inner] +id = "reverse-text-v1" +``` + +Base fields (every replay taskset, from `verifiers.v1.tasksets.replay`): + +| Field | Default | What it does | +|---|---|---| +| `buffer_dir` | *(required)* | The saved-rollout dir to replay: a run's `rollouts` dir (or the run dir containing it). Under prime-rl the literal `"self"` resolves to this run's own rollout dir (an online buffer over the run's freshly written rollouts). | +| `inner` | *(required)* | The original taskset's config — it scores the derived rollouts and provides their tools, so it must reproduce the source run's taskset config. | +| `source_envs` | `None` | Which envs' rollouts to replay, by their stamped name (`info.prime_rl.env_name`). Unset: every env except replay envs. An explicit list replays exactly those — naming a replay env opts into chained derivations (recheck a recheck). With a list set, records without the stamp never match. | +| `allow_container` | `false` | Allow sources whose task ran in a container image. The container state the transcript references is gone — a fresh container is provisioned from the same image, so the model resumes in a reset world. | +| `online` | *(inferred)* | Set automatically for `buffer_dir = "self"`. Set it yourself only to treat another still-running run's dir as a growing buffer. | + +Per-package fields: `replay-continue-v1` adds `anchor = "compaction" | "tool-call"`; `replay-recheck-v1` adds `instruction` (the appended check-your-work turn, with a built-in default). + +Runnable debug configs live at `configs/debug/replay/` (`online_recheck.toml`, `offline_recheck.toml`). + +## Metrics + +Every replay rollout logs `replay/source_reward` (the reward the source rollout received) and `replay/source_step` (the trainer step that wrote it) alongside the inner taskset's own rewards and metrics. + +## Writing a New Derivation + +A derivation is a thin package over the base, exactly like `swebench-pro-v1` over verifiers' `harbor` taskset: subclass `ReplayTaskset` (binding your narrowed config in the generic) and implement `build_prompt` (the seeded conversation for one anchor); derivations that resume mid-rollout also override `record_anchors` (the resume points one saved rollout offers; each becomes one task — the default is one task per rollout, anchored at its final state). The base owns the buffer, online semantics, lazy binding, lineage, and inner-taskset delegation; `verifiers.v1.tasksets.replay.surgery` provides the graph enumerators (`compaction_forks`, `tool_call_anchors`, `recheck_seed`, ...). See `environments/replay_recheck_v1` for the ~30-line reference. Register the package in the root `pyproject.toml` (`envs` extra + `[tool.uv.sources]`) and `uv sync`. + +## Constraints and Caveats + +- **One derivation per env entry.** Env-level settings (harness runtime, token budgets, `ratio`) are per-env — mix derivations with multiple env entries. +- **Harness: match the source.** Replay resumes conversations — run the replay env under the same harness the source env used, so the resumed rollout has the same scaffold and harness-side tools as the original. The harness must support message-prompt seeding (`SUPPORTS_MESSAGE_PROMPT`; the built-in `default` and `null` harnesses do — string-prompt CLI harnesses can't be seeded with a prior conversation). +- **Replay envs need explicit ratios.** Without ratios, envs are weighted by task count — an online replay env advertises a single virtual task and would effectively never be drawn. Setting `ratio` on the replay env forces ratios on every train env (the all-or-none rule), which is the intended, explicit state. +- **Token budgets default from `seq_len`.** Every v1 train env's `max_total_tokens` defaults to the run's `seq_len` at config-resolve time, so replayed seeds stop instead of producing samples that can't fit training. Override `max_total_tokens` (or set it to `"None"`) per env entry to opt out. +- **Long continue tasks age off-policy.** A continue task resumes deep into a long trajectory and keeps going, so its rollout can span many trainer steps and be discarded by `orchestrator.max_off_policy_steps` (default 8). If a continue env shows sustained `errored_rollouts`, raise `max_off_policy_steps` or shorten the tasks. +- **`source_envs` needs the stamp.** The filter matches `info.prime_rl.env_name`, which prime-rl's rollout writer stamps into every saved record. Records without the stamp (rollout files from other producers) never match an explicit list — leave `source_envs` unset to replay them. +- **`allow_container` is off for a reason.** The trace records the conversation, not the container: for a containerized source task, replay provisions a fresh container from the same image, so any filesystem state the transcript references is gone and the model resumes in a reset world — especially acute for `anchor = "tool-call"`, which resumes mid-trajectory. Enable it only when the task's image is self-contained enough for that to be fair — and give the replay env's harness a container runtime (docker/prime): with a subprocess runtime every imaged request fails at episode construction. +- **Inner-taskset limits.** Replay delegates scoring to `inner` but cannot delegate group scoring (`@group_reward` tasksets), custom `State` classes, or user simulators — those are rejected loudly at env construction. Shared-placement inner toolsets are also unsupported, but fail per-rollout (the serving layer inspects toolsets on a stub task and never starts them). diff --git a/environments/replay_continue_v1/pyproject.toml b/environments/replay_continue_v1/pyproject.toml new file mode 100644 index 0000000000..4672b37b92 --- /dev/null +++ b/environments/replay_continue_v1/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "replay-continue-v1" +version = "0.1.0" +description = "replay-continue-v1 — a derivation over verifiers' replay taskset base." +requires-python = ">=3.10" +dependencies = ["verifiers"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["replay_continue_v1"] diff --git a/environments/replay_continue_v1/replay_continue_v1/__init__.py b/environments/replay_continue_v1/replay_continue_v1/__init__.py new file mode 100644 index 0000000000..231c447759 --- /dev/null +++ b/environments/replay_continue_v1/replay_continue_v1/__init__.py @@ -0,0 +1,4 @@ +from replay_continue_v1.taskset import ReplayContinueTaskset + +# The loader requires exactly one Taskset subclass exported via __all__. +__all__ = ["ReplayContinueTaskset"] diff --git a/environments/replay_continue_v1/replay_continue_v1/taskset.py b/environments/replay_continue_v1/replay_continue_v1/taskset.py new file mode 100644 index 0000000000..5053c6c2f2 --- /dev/null +++ b/environments/replay_continue_v1/replay_continue_v1/taskset.py @@ -0,0 +1,44 @@ +"""replay-continue-v1 — resume saved rollouts mid-way and let the model finish the task. + +A thin derivation over the `verifiers.v1.tasksets.replay` base: anchors a saved rollout +at a context-compaction point (the seed is the post-compaction prompt the original +harness restarted from) or right after a tool result, and the new rollout is scored by +the original (`inner`) taskset. +""" + +import random +from typing import Literal + +import verifiers.v1 as vf +from verifiers.v1.tasksets.replay import ( + ReplayTask, + ReplayTaskset, + ReplayTasksetConfig, + compaction_forks, + continue_seed, + tool_call_anchors, +) + + +class ReplayContinueConfig(ReplayTasksetConfig): + anchor: Literal["compaction", "tool-call"] = "compaction" + """Where to resume: `"compaction"` seeds the post-compaction prompt of each + compaction point (one task per compaction; only compacted rollouts are sources). + `"tool-call"` seeds the conversation up to a tool result and resumes right after it + (one deterministically sampled resume point per source rollout).""" + + +class ReplayContinueTaskset(ReplayTaskset, vf.Taskset[ReplayTask, ReplayContinueConfig]): + def record_anchors(self, record, children, roots, tree) -> list[int | None]: + nodes = record["nodes"] + if self.config.anchor == "compaction": + # One task per compaction point. + return compaction_forks(nodes, children, tree) + # Tool calls are numerous (~dozens per rollout); one deterministically sampled + # resume point per source keeps the index one-candidate-per-rollout and + # identical across pool workers. + anchors = tool_call_anchors(nodes, children, tree) + return [random.Random(record.get("id", "")).choice(anchors)] if anchors else [] + + def build_prompt(self, record: dict, anchor: int | None) -> list[dict]: + return continue_seed(record["nodes"], anchor) diff --git a/environments/replay_recheck_v1/pyproject.toml b/environments/replay_recheck_v1/pyproject.toml new file mode 100644 index 0000000000..c3dab27adf --- /dev/null +++ b/environments/replay_recheck_v1/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "replay-recheck-v1" +version = "0.1.0" +description = "replay-recheck-v1 — a derivation over verifiers' replay taskset base." +requires-python = ">=3.10" +dependencies = ["verifiers"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["replay_recheck_v1"] diff --git a/environments/replay_recheck_v1/replay_recheck_v1/__init__.py b/environments/replay_recheck_v1/replay_recheck_v1/__init__.py new file mode 100644 index 0000000000..5b27eee1ef --- /dev/null +++ b/environments/replay_recheck_v1/replay_recheck_v1/__init__.py @@ -0,0 +1,4 @@ +from replay_recheck_v1.taskset import ReplayRecheckTaskset + +# The loader requires exactly one Taskset subclass exported via __all__. +__all__ = ["ReplayRecheckTaskset"] diff --git a/environments/replay_recheck_v1/replay_recheck_v1/taskset.py b/environments/replay_recheck_v1/replay_recheck_v1/taskset.py new file mode 100644 index 0000000000..080751bfda --- /dev/null +++ b/environments/replay_recheck_v1/replay_recheck_v1/taskset.py @@ -0,0 +1,33 @@ +"""replay-recheck-v1 — re-serve finished rollouts with a check-your-work turn. + +A thin derivation over the `verifiers.v1.tasksets.replay` base: seeds a saved rollout's +final conversation plus an appended instruction to re-verify and fix the answer, and +the new rollout is scored by the original (`inner`) taskset. +""" + +import verifiers.v1 as vf +from verifiers.v1.tasksets.replay import ( + ReplayTask, + ReplayTaskset, + ReplayTasksetConfig, + build_children, + main_tree, + recheck_seed, +) + +RECHECK_INSTRUCTION = ( + "Carefully check your work above. Re-verify the reasoning and the final answer, and " + "fix any mistakes you find. Then state your final answer." +) + + +class ReplayRecheckConfig(ReplayTasksetConfig): + instruction: str = RECHECK_INSTRUCTION + """The user turn appended to the finished conversation.""" + + +class ReplayRecheckTaskset(ReplayTaskset, vf.Taskset[ReplayTask, ReplayRecheckConfig]): + def build_prompt(self, record: dict, anchor: int | None) -> list[dict]: + nodes = record["nodes"] + children, _ = build_children(nodes) + return recheck_seed(nodes, children, main_tree(children), self.config.instruction) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index d296178c05..2aa15f6e74 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -748,7 +748,7 @@ def auto_setup_bench(self): @model_validator(mode="after") def resolve_env_config(self): - """Set vLLM sampling defaults + legacy env kwargs on each train env from top-level fields.""" + """Set vLLM sampling defaults + token budgets on each train env from top-level fields.""" for env in self.train.env: # Policy-sourced rollouts hit our vLLM server; frozen-sourced # rollouts may hit external OAI endpoints that reject these knobs. @@ -761,4 +761,10 @@ def resolve_env_config(self): # v0 env: cap per-turn response tokens to the training budget (the legacy # bridge applies extra_env_kwargs via env.set_kwargs). env.extra_env_kwargs["max_seq_len"] = self.seq_len + elif "max_total_tokens" not in env.model_fields_set: + # v1 env: cap each rollout's total (prompt + completion) tokens to the + # training budget, so long-prompt tasks (e.g. replayed conversations) + # stop instead of producing samples that can't fit seq_len. Explicit + # values (including "None" = unlimited) are respected. + env.max_total_tokens = self.seq_len return self diff --git a/pyproject.toml b/pyproject.toml index 0c69c94e31..9b67e7dd25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,6 +110,8 @@ envs = [ "prolog-v1", "r2e-gym-v1", "redsearcher-v1", + "replay-continue-v1", + "replay-recheck-v1", "reverse-text", "reverse-text-v1", "rlm-search", @@ -283,6 +285,8 @@ patterned-needle-in-haystack-v1 = { path = "deps/research-environments/environme prolog-v1 = { path = "deps/research-environments/environments/prolog_v1", editable = true } r2e-gym-v1 = { path = "deps/research-environments/environments/r2e_gym_v1", editable = true } redsearcher-v1 = { path = "deps/research-environments/environments/redsearcher_v1", editable = true } +replay-continue-v1 = { path = "environments/replay_continue_v1", editable = true } +replay-recheck-v1 = { path = "environments/replay_recheck_v1", editable = true } reverse-text = { path = "deps/verifiers/environments/reverse_text", editable = true } reverse-text-v1 = { path = "deps/verifiers/environments/reverse_text_v1", editable = true } rlm-search = { path = "deps/research-environments/environments/rlm_search", editable = true } diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index 7eebe87f86..0f2cfb220c 100644 --- a/skills/configs/SKILL.md +++ b/skills/configs/SKILL.md @@ -53,6 +53,8 @@ CLI: `--env.0.id reverse-text --env.1.id math-env`. **Algorithms** — `[orchestrator.algo] type = "grpo" | "max_rl" | "opd" | "opsd" | "sft" | "echo"` — the type names the algorithm (credit assignment + loss routing, fused), and each type's class defaults are its vetted setting; any other key you set is your own assembly (e.g. `[orchestrator.algo.roles.user] alpha = 0.1` for echo — setting any echo role replaces the whole role table). There is no preset layer, and no config hook that points at user code — a new algorithm is a named class in the repo (subclass `Algorithm`, register it). Per-env override: `[orchestrator.train.env.algo] type = "opd"` (the env assembles its own algorithm). prime-rl only hosts the trainable policy; frozen models are inline external endpoints on the algorithm, named where the model is used — `[orchestrator.algo.teacher]` for opd (the frozen model scored against), `[orchestrator.algo.sampling.source]` for sft (the model it samples from), each with `name` + `base_url`. There is no shared `teacher` slot. opsd declares no model — it self-distills against the live policy. See `docs/algorithms.md`. +**Replay envs** — taskset ids `replay-continue-v1` (`anchor = "compaction" | "tool-call"`) and `replay-recheck-v1` (`instruction`) replay saved rollouts (`/rollouts/step_*/train_rollouts.jsonl`) as training tasks; thin packages over `verifiers.v1.tasksets.replay` (like swebench over harbor), configured as normal train env entries with a `ratio` (ratios are all-or-none, and replay envs need one — they advertise a single virtual task). Flat taskset config, base fields: `buffer_dir` (a prior run's rollouts dir, or `"self"` for this run's own growing buffer — implies online + whole-group dispatch), required `inner` (the source run's taskset config, which scores the derived rollouts), `source_envs` (allowlist over the `info.prime_rl.env_name` stamp; unset = every env except replay envs, listing a replay env opts into chained derivations), `allow_container`. Run replay envs under the source env's harness (must support message prompts); token budgets default from `seq_len` (every v1 train env's unset `max_total_tokens` resolves to it). Debug configs: `configs/debug/replay/`. See `docs/replay.md`. + **`BaseModel | None` fields** — bare flag enables defaults; nested override enables and sets: ```bash diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index c0bcc346bc..a027aec9d0 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -141,8 +141,11 @@ curl -s http://localhost:8000/metrics | grep -E "num_requests|gpu_cache_usage" wc -l {output_dir}/rollouts/step_42/train_rollouts.jsonl head -1 {output_dir}/rollouts/step_42/train_rollouts.jsonl | python -m json.tool jq '.reward' {output_dir}/rollouts/step_42/train_rollouts.jsonl +jq '.info.prime_rl' {output_dir}/rollouts/step_42/train_rollouts.jsonl # attribution stamp ``` +Each jsonl line carries an `info.prime_rl` stamp (`kind`, `env_name`, `group_id`, `policy_version`, `is_filtered`) attributing the record to its env, GRPO group, and producing policy version — consumed by e.g. the `replay-v1` env (`docs/replay.md`). The `.bin` is renamed into place only after the jsonl is complete, so it doubles as a completeness barrier for readers of a live run. + ### Common failure modes A few warnings are normal. Escalate when errors are persistent, growing, or hit a large fraction of rollouts. diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index e18952ea72..c341f6d2fb 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -236,6 +236,16 @@ async def setup(self) -> None: pre_filters = setup_filters(config.pre_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="pre-batch") post_filters = setup_filters(config.post_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="post-batch") + # Resolve the replay taskset's "self" buffer sentinel to this run's rollout dir. + # Env-server children never learn the orchestrator's output_dir, and the configs + # package doesn't know the rollouts-dir layout — so the resolved path is written + # into the taskset config here, just before the servers spawn (the config + # persisted to control/orch.toml above keeps the relocatable sentinel). + eval_envs = config.eval.env if config.eval is not None else [] + for env_config in [*config.train.env, *eval_envs]: + if getattr(env_config.taskset, "buffer_dir", None) == "self": + env_config.taskset.buffer_dir = str(get_rollout_dir(config.output_dir)) + get_logger().info("Loading training environments") self.train_envs = TrainEnvs( config.train.env, policy_pool=self.policy_inference, renderer_config=config.renderer diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 17eb1113f8..473ac97dc3 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -73,8 +73,9 @@ class Rollout(vf.Trace[TaskT], Generic[TaskT]): """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout — prime-rl's orchestration metadata lives on it directly (set by the dispatcher once the rollout returns), so there's no wrapper. Train vs eval is the ``kind`` discriminator. All metadata - fields are ``exclude=True``, so dumping a Rollout yields a plain trace — the on-disk - ``results.jsonl`` is unchanged. + fields are ``exclude=True``, so dumping a Rollout yields a plain trace; :meth:`to_record` + additionally stamps the metadata into ``info["prime_rl"]`` so saved rollouts stay + attributable (e.g. to their env) after the plain-trace round trip. It is also the single currency the scoring hooks receive: a hook reads the trace directly (``rollout.reward``, ``rollout.nodes``, ``rollout.num_turns``) and writes @@ -98,6 +99,22 @@ class Rollout(vf.Trace[TaskT], Generic[TaskT]): filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) eval_step: int | None = Field(default=None, exclude=True) + def to_record(self) -> dict: + """The plain-trace dump, with the orchestration metadata stamped into + ``info["prime_rl"]`` — ``info`` round-trips through ``vf.WireTrace.model_validate``, + so consumers of saved rollouts (e.g. replay buffers) can attribute each record to + its env, GRPO group, and producing policy version (== the trainer step whose + weights generated it). The step itself is encoded in the ``step_N`` dirname.""" + record = super().to_record() + record.setdefault("info", {})["prime_rl"] = { + "kind": self.kind, + "env_name": self.env_name, + "group_id": str(self.group_id), + "policy_version": self.policy_version, + "is_filtered": self.is_filtered, + } + return record + def assign_advantages(self, values: float | list[float]) -> None: """Write the rl advantage stream: a scalar broadcast over the rollout's trainable (mask-True) tokens (0.0 elsewhere), or a per-token diff --git a/tests/unit/test_replay_v1.py b/tests/unit/test_replay_v1.py new file mode 100644 index 0000000000..eb682a4087 --- /dev/null +++ b/tests/unit/test_replay_v1.py @@ -0,0 +1,94 @@ +"""Unit tests for the replay derivation packages (replay-continue-v1 / replay-recheck-v1): +config validation and anchor/prompt binding. The shared base's surgery and buffer logic +is tested in verifiers (tests/v1/test_replay.py).""" + +import json +from pathlib import Path + +import pytest +from pydantic import ValidationError +from replay_continue_v1.taskset import ReplayContinueConfig, ReplayContinueTaskset +from replay_recheck_v1.taskset import ReplayRecheckConfig +from reverse_text_v1.taskset import ReverseTextConfig +from verifiers.v1.loaders import task_type, taskset_class, taskset_config_type + +INNER = {"id": "reverse-text-v1"} + + +def test_loader_resolves_both_packages(): + assert taskset_class("replay-continue-v1") is ReplayContinueTaskset + assert taskset_config_type("replay-continue-v1") is ReplayContinueConfig + assert taskset_config_type("replay-recheck-v1") is ReplayRecheckConfig + assert task_type("replay-recheck-v1").__name__ == "ReplayTask" + + +def test_config_requires_inner(): + # `inner` is a required base field, so leaving it out is a plain missing-field error. + with pytest.raises(ValidationError, match="inner"): + ReplayRecheckConfig(buffer_dir="/tmp/buf") + + +def test_config_rejects_unknown_anchor(): + with pytest.raises(ValidationError, match="anchor"): + ReplayContinueConfig(buffer_dir="/tmp/buf", anchor="turn", inner=INNER) + + +def test_config_requires_buffer_dir(): + with pytest.raises(ValidationError, match="buffer_dir"): + ReplayRecheckConfig(inner=INNER) + + +def test_config_self_buffer_implies_online(): + cfg = ReplayRecheckConfig(buffer_dir="self", inner=INNER) + assert cfg.online is True # pinned before the orchestrator rewrites the sentinel + + +def test_config_narrows_inner(): + cfg = ReplayContinueConfig(buffer_dir="/tmp/buf", anchor="tool-call", inner=INNER) + assert isinstance(cfg.inner, ReverseTextConfig) # narrowed to the inner taskset's type + + +def _node(parent: int | None, sampled: bool, message: dict) -> dict: + return {"parent": parent, "sampled": sampled, "message": message} + + +@pytest.fixture +def tool_buffer(tmp_path: Path) -> Path: + call = {"name": "t", "arguments": "{}"} + record = { + "id": "ttt", + "errors": None, + "stop_condition": "agent_completed", + "rewards": {"r": 1.0}, + "task": {"idx": 0, "prompt": "Reverse: abc", "answer": "cba"}, + "info": {}, + "nodes": [ + _node(None, False, {"role": "system", "content": "s"}), + _node(0, False, {"role": "user", "content": "task"}), + _node(1, True, {"role": "assistant", "content": "", "tool_calls": [call | {"id": "a"}]}), + _node(2, False, {"role": "tool", "tool_call_id": "a", "content": "ra"}), + _node(3, True, {"role": "assistant", "content": "", "tool_calls": [call | {"id": "b"}]}), + _node(4, False, {"role": "tool", "tool_call_id": "b", "content": "rb"}), + _node(5, True, {"role": "assistant", "content": "done"}), + ], + } + step = tmp_path / "step_1" + step.mkdir() + (step / "train_rollouts.jsonl").write_text(json.dumps(record) + "\n") + return tmp_path + + +def test_continue_tool_call_anchor_is_sampled_deterministically(tool_buffer): + def scan(taskset): + return taskset.buffer.scan() + + make = lambda: ReplayContinueTaskset( # noqa: E731 + ReplayContinueConfig(buffer_dir=str(tool_buffer), anchor="tool-call", inner=INNER) + ) + first, second = scan(make()), scan(make()) + # One resume point per source rollout, identical across taskset instances (pool workers). + assert len(first) == 1 and first[0].anchor_node in (3, 5) + assert first[0].anchor_node == second[0].anchor_node + # Compaction anchoring finds nothing in this rollout (it never compacted). + compaction = ReplayContinueTaskset(ReplayContinueConfig(buffer_dir=str(tool_buffer), inner=INNER)) + assert compaction.buffer.scan() == [] diff --git a/uv.lock b/uv.lock index fe29eac05d..d5062bbb2e 100644 --- a/uv.lock +++ b/uv.lock @@ -4025,6 +4025,8 @@ envs = [ { name = "prolog-v1" }, { name = "r2e-gym-v1" }, { name = "redsearcher-v1" }, + { name = "replay-continue-v1" }, + { name = "replay-recheck-v1" }, { name = "reverse-text" }, { name = "reverse-text-v1" }, { name = "rlm-search" }, @@ -4162,6 +4164,8 @@ requires-dist = [ { name = "r2e-gym-v1", marker = "extra == 'envs'", editable = "deps/research-environments/environments/r2e_gym_v1" }, { name = "redsearcher-v1", marker = "extra == 'envs'", editable = "deps/research-environments/environments/redsearcher_v1" }, { name = "renderers", editable = "deps/renderers" }, + { name = "replay-continue-v1", marker = "extra == 'envs'", editable = "environments/replay_continue_v1" }, + { name = "replay-recheck-v1", marker = "extra == 'envs'", editable = "environments/replay_recheck_v1" }, { name = "reverse-text", marker = "extra == 'envs'", editable = "deps/verifiers/environments/reverse_text" }, { name = "reverse-text-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/reverse_text_v1" }, { name = "rich", specifier = ">=14.0.0" }, @@ -4851,6 +4855,28 @@ dev = [ { name = "ty", specifier = ">=0.0.1a29,<0.0.22" }, ] +[[package]] +name = "replay-continue-v1" +version = "0.1.0" +source = { editable = "environments/replay_continue_v1" } +dependencies = [ + { name = "verifiers" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + +[[package]] +name = "replay-recheck-v1" +version = "0.1.0" +source = { editable = "environments/replay_recheck_v1" } +dependencies = [ + { name = "verifiers" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + [[package]] name = "requests" version = "2.34.2"