diff --git a/rfcs/013-task-sampling-and-training-evals.md b/rfcs/013-task-sampling-and-training-evals.md new file mode 100644 index 000000000..30a308926 --- /dev/null +++ b/rfcs/013-task-sampling-and-training-evals.md @@ -0,0 +1,533 @@ +# RFC: Task Sampling and In-Training Evaluation + +**Status**: In Review +**Created**: 2026-09-26 +**Updated**: 2026-10-22 +**Revised**: 2026-10-12 — removed a CLI limitation fixed by #1237, updated the RFC 008 reference to #1321, and corrected four small inaccuracies; 2026-10-16 — added results on three more seeds, replaced the token-weighted formula, scoped the first evaluation to environments whose turns the trainer drives, listed the gaps in the collect pipeline, offered the curriculum as an example, linked the trial's code and logs, and stated how the token weights treat unseen keys and many keys; 2026-10-22 — after review, removed weighting by cost and the `cost` argument from the curriculum, stated the two hooks a sampler needs in a trainer and their limits in TRL, added `metrics()` to the protocol, left shares across sources to the trainer, reported how far the signal formula over-predicts, cited prior work on rollout allocation and filtering, narrowed the reference experiment to uniform sampling against the curriculum on five new seeds, answered open question 6 and proposed a direction for open question 5 +**Authors**: @Matheus-F-Scatolin, @tiagoperrupato +**RFC ID**: 013 + +## Summary + +This RFC adds two trainer-side capabilities on top of the existing Task API: **task sampling**, which decides which task each episode runs, and **in-training evaluation**, which measures a policy checkpoint on a held-out split in a reproducible way. Together they turn "train on an environment" into a closed loop: evaluate a baseline, train, and evaluate the same fixed test set at regular steps, with the same environment and the same tool interface throughout. + +It introduces a small `TaskSampler` protocol with two implementations in core (`SequentialSampler`, `UniformSampler`), a generic Task API client helper, an `EnvEvalHarness` that implements the existing `EvalHarness` base class with the same pieces as the `openenv collect` rollout pipeline, and an `openenv eval` CLI command. A third sampler, `CurriculumSampler`, is offered as an example; whether it belongs in core is a question for the maintainers. The environment boundary does not change: environments keep computing rewards, and samplers only consume them. + +The protocol and the evaluation harness are the core of the proposal and are useful with any sampler. The curriculum is one pluggable sampler, and we measured it on a laptop (see [Preliminary Results](#preliminary-results)): a single-seed pilot, then a trial on three more seeds. A curriculum that targets GRPO signal reached higher test accuracy than uniform sampling per environment episode on all four seeds, and per generated token on the three new ones. The pilot's seed had pointed the other way per token, because its curriculum run happened to produce much longer answers; on the other seeds, answer length did not follow the method. Dividing the curriculum's weights by the token cost of each tier, as the first version of this RFC proposed, made it worse, and a corrected form matched the curriculum that leaves cost out per generated token, the unit it was meant to improve, while being slightly ahead of it per episode. So the `CurriculumSampler` weights by signal only; the samplers proposed here count the cost the trainer reports and report signal per episode and per unit of that cost. + +## Motivation + +### Problem Statement + +**1. Task selection has no shared mechanism.** The Task API (`list_splits`, `list_tasks`, `num_tasks`, `get_task`, `get_task_range`, and `reset(split=..., index=...)`) lets a dataset-backed environment publish its tasks, but it is discovery only. Nothing in OpenEnv decides which task an episode should use. Harbor and the wrappers generated by `openenv import` implement it; no hand-written environment under `envs/` does, including the dataset-backed `finqa_env` and the procedurally generated `reasoning_gym_env`. The core clients deliberately ship no Task API methods, because task specs are environment-specific (Task API guide), so each environment client writes its own HTTP helpers; `HarborEnv` has its own, and the guide shows the same pattern for a LaTeX OCR client. + +Today, selecting tasks well is manual. The Harbor harness takes an explicit `indices` list, and its code explains why: a task that every generation solves and one that none solves both give `reward_std` 0, and "the band that splits has to be chosen per model" (`envs/harbor_env/harness.py`). This RFC turns that manual choice into a sampler that measures the band as training runs. + +**2. Small models waste most of their GRPO rollouts.** GRPO computes each completion's advantage relative to the other completions in its group. When all `G` completions for a prompt receive the same reward, every advantage is zero and the group contributes no gradient. For a binary reward with success probability `p`, the chance that a group of size `G` carries any signal is + +``` +P(signal) = 1 - p^G - (1 - p)^G +``` + +which is maximal at `p = 0.5` and collapses as `p` approaches 0 or 1. In our pilot, 80% of GRPO groups carried no gradient under uniform sampling (Llama-3.2-1B, `G = 8`). DAPO's dynamic sampling filters these groups after generation; a sampler that targets tasks near `p = 0.5` generates fewer of them in the first place. What a wasted group costs depends on the environment: in a single-turn reasoning task it is mostly generated tokens, while in a sandboxed or tool-using environment it is mostly the episode itself (container time, tool calls, API calls). A trainer has to see that cost next to the signal, so every sampler here reports both. + +**3. There is no shared way to evaluate a checkpoint on an environment.** `openenv.core.evals` provides the abstract `EvalHarness` and an Inspect AI wrapper. There is no shared way to run a checkpoint over an environment's `test` split with fixed seeds, aggregate metrics, and compare two checkpoints: `examples/claude_code_harness_eval/run_eval.py` and the SFT warm-up tutorial (`evaluate_model`) each write their own loop. Reproducibility is also fragile: environments have accepted `seed` in `reset()` and discarded it (`textarena_env` in #1102, and `finqa_env` until #1252 fixed it), so two evaluation runs of the same checkpoint can see different tasks. + +RFC 000 schedules evals for phase three (0.6 to 0.9), starting with "simple evals that can be scored by a script". This RFC proposes that starting point, built so that training and evaluation share one interface (design principle 1: minimize lifecycle deltas). + +### Goals + +1. A trainer can choose tasks for each episode through one small protocol, independent of the training framework. +2. A reproducible evaluation of any policy reachable through an OpenAI-compatible endpoint, returning the existing `EvalResult` type. The first version covers environments that implement the Task API and whose episodes the trainer drives turn by turn; environments whose server runs a whole agent in one call, such as Harbor, are open question 8. +3. A way to run that evaluation every N training steps and compare checkpoints task by task. +4. A curriculum sampler, offered as an example, that raises the number of GRPO groups with non-zero advantage per episode, with no environment changes beyond publishing tasks. Every sampler proposed here reports those groups per episode and per unit of the cost the trainer reports (for example generated tokens). +5. A reference experiment that runs on a single laptop-class machine and reports results per episode and per generated token. + +### Non-goals + +- Owning GRPO or any trainer implementation (consistent with the non-goals of the draft recipes RFC, #924). +- LLM-judge or multi-agent evals; this RFC covers evals scored by the environment's own reward. +- Changing `reset`, `step`, or `state` signatures, or the wire types. +- Exposing task selection to the agent. Samplers run in orchestration code only. +- Setting shares across environments or data sources, or filtering groups after generation. Both belong to the trainer and its rollout scheduler; a sampler works inside one source's share (see Trainer integration). +- Claiming that a curriculum beats uniform sampling in general. The samplers make the comparison easy to run; the answer depends on the environment and the cost unit. + +## Design + +### Architecture Overview + +``` + trainer process (orchestration side) environment container + ┌─────────────────────────────────────────────────────┐ ┌───────────────────────┐ + │ │ HTTP │ Task API routes │ + │ TaskAPIClient ──── list_splits / get_task_range ───┼──────▶│ (metadata only) │ + │ │ │ │ │ + │ ▼ │ │ │ + │ TaskSampler ── sample(n) ──▶ task specs │ │ │ + │ ▲ │ │ WS │ │ + │ │ ▼ │ │ reset(split, index, │ + │ update(task, rewards, session factory + ├──────▶│ seed) │ + │ costs) harness adapter │ │ step(action) │ + │ └──── rewards, costs ◀─────┘ │◀──────┤ reward (computed │ + │ │ │ inside the env) │ + │ EnvEvalHarness = SequentialSampler + rollouts │ └───────────────────────┘ + │ + metric aggregation → EvalResult │ + └─────────────────────────────────────────────────────┘ +``` + +Training uses a `UniformSampler` or `CurriculumSampler` over the `train` split. Evaluation uses a `SequentialSampler` over the `test` split, with fixed seeds. Both feed the same rollout path, so the only difference between a training episode and an evaluation episode is which task is chosen and whether a gradient step follows. The sampler hands each task to the session factory, which turns it into what the environment needs: `reset(split=..., index=..., seed=...)` for an environment that follows the Task API guide. + +### Core Abstractions + +#### 1. Task spec conventions + +Task specs stay untyped (`Any`), as in the Task API guide. This RFC reserves a few optional keys so samplers and evaluators can work across environments: + +| Key | Type | Meaning | +|-----|------|---------| +| `id` | `str` | Stable task identifier. Already the convention. Samplers take the key's name as a parameter (`id_key`, default `"id"`), because Harbor's task specs use `task_id`. | +| `split` | `str` | Split name. Already the convention. | +| `index` | `int` | Position in the split. Already the convention. | +| `difficulty` | `int` or `float` | Optional. Higher means harder. Tiers (`0..K`) or a continuous score. | +| `group` | `str` | Optional. A family of related tasks (for example a reasoning_gym dataset name). | + +An environment that publishes neither `difficulty` nor `group` still works with every sampler; the curriculum then keys its statistics by task `id`. + +Tiers work best when they are close together. The curriculum can only concentrate on a tier whose success rate is near 0.5; in the pilot, once the middle tier reached about 95%, the next tier was still below 40% and no tier sat in between. The Task API guide should recommend enough tiers that neighbouring tiers differ by a modest step in difficulty. Two more points for the guide: train and test seeds that never collide do not make the questions disjoint, because a tier with fewer possible questions than train tasks repeats its test questions in train; and a test set should take no questions from such a tier, or report it apart. These keys and points are an addition to the guide, not a requirement: the Task API checklist and existing environments are unaffected. + +#### 2. `TaskAPIClient` + +A generic helper in `openenv.core` that calls the standard Task API routes for any environment, deriving the HTTP base URL from a `ws://` client URL: + +```python +class TaskAPIClient: + def __init__(self, base_url: str, env_name: str, + headers: dict[str, str] | None = None): ... + @classmethod + def from_env_client(cls, env: EnvClient, env_name: str | None = None, + headers: dict[str, str] | None = None) -> "TaskAPIClient": ... + def list_splits(self) -> list[dict[str, str]]: ... + def list_tasks(self, split: str) -> list[Any]: ... + def num_tasks(self, split: str) -> int: ... + def get_task(self, split: str, index: int) -> Any: ... + def get_task_range(self, split: str, start: int | None = None, + stop: int | None = None) -> list[Any]: ... + def iter_tasks(self, split: str, page_size: int = 256) -> Iterator[Any]: ... +``` + +It returns task specs as untyped JSON, exactly as the routes serve them, and maps `501` to `NotImplementedError` and `400` to `IndexError`, mirroring the server. `iter_tasks` pages with `get_task_range` and relies on `num_tasks` for the true total, since `list_tasks` may return a bounded preview. `headers` are sent with every request, as `EnvClient` and `MCPToolClient` send theirs since #1403 (for example a token for a private Space). Environment clients can keep their own typed helpers (for example `HarborEnv.get_task` returning a `HarborTaskRef`) and build them on this client. See D6 for why this belongs in core although core clients currently leave Task API methods out. + +#### 3. `TaskSampler` + +```python +class TaskSampler(Protocol): + def sample(self, n: int) -> list[Any]: + """Return n task specs for the next n prompts (GRPO groups).""" + + def update(self, task: Any, rewards: Sequence[float | None], + costs: Sequence[float] | None = None) -> None: + """Report the rewards of every rollout of one task (one GRPO group). + + A reward of None is a rollout the environment did not score. Samplers + leave it out of their estimates and counters: a missing reward is not + zero. costs, if given, has one entry per rollout in the trainer's cost + unit (for example the number of generated tokens). + """ + + def state_dict(self) -> dict[str, Any]: ... + def load_state_dict(self, state: dict[str, Any]) -> None: ... + + def metrics(self) -> dict[str, float]: + """Return numbers for the trainer's log; the samplers proposed here return the counters below.""" +``` + +The rewards passed to `update` are the environment's own, one per rollout, not a total that the trainer forms with reward terms of its own. A trainer calls `update` for every scored group, including the groups it then leaves out of the gradient because they carry no signal (see Trainer integration). + +`state_dict` and `load_state_dict` let a trainer checkpoint and resume the sampler together with the model, which a long run on one machine needs. (The pilot lost its intermediate checkpoints because they were not saved separately; resuming requires the model, the optimizer and the sampler state together.) + +The samplers proposed here keep the same counters and return them from `metrics()`, so their training logs report efficiency the same way: + +- `groups`, `signal_groups` (groups with non-zero reward variance), `episodes`, and `cost` (the sum of reported costs) +- derived: `signal_groups_per_episode`, `signal_groups_per_cost`, and `zero_signal_fraction` + +A trainer reports costs for every group or for none: these samplers raise `ValueError` on a group without costs after one with costs, and the reverse, and with none they leave `cost` and `signal_groups_per_cost` out of `metrics()`. Another `TaskSampler` returns from `metrics()` whatever it tracks, as numbers for the log. The zero-signal fraction on its own is misleading: in the pilot the curriculum had the lower zero-signal fraction and the worse result per token. Logs should report signal groups per episode and per unit of cost. + +The protocol has the shape of prime-rl's [`TaskSampler`](https://github.com/PrimeIntellect-ai/prime-rl/blob/855a2b815b804857835c635b3fd2b1ae5400f96c/src/prime_rl/orchestrator/curriculum/samplers/base.py) (`__next__`, `observe`, `state_dict`, `load_state_dict`, `metrics`), with two differences: `sample(n)` draws the tasks of a whole step at once, and `update` takes rewards and costs instead of that framework's episode objects, so core depends on no training framework. Running one of these samplers inside prime-rl takes a short wrapper, whose `__next__` returns the prime-rl task for `sample(1)[0]` and whose `observe` reads the rewards and generated token counts of the group's episodes and calls `update`, plus a change to prime-rl itself (a config class and a branch in `Curriculum.__init__`), because its sampler configs form a closed set. The other direction needs that framework's episode objects. + +Two implementations ship in core: + +- **`SequentialSampler(tasks, seeds)`**: deterministic enumeration of a split, optionally crossed with several seeds. Used by evaluation. `update` only updates the counters. +- **`UniformSampler(tasks, rng_seed)`**: shuffled sampling without replacement per epoch. The baseline for training. + +A third, **`CurriculumSampler(tasks, group_size, key="difficulty", ...)`**, described below, is offered as an example in `examples/`, and moves to core if the maintainers prefer (open question 9). + +The tasks a sampler draws, for example `sampler.sample(num_episodes)`, are a valid `tasks` iterable for the existing `CollectRunner`, so `openenv collect` and the evaluation harness gain task selection without changes to that interface. Selection does not work end to end yet, though, for four reasons. Inside `CollectRunner`: it passes no seed to the session, and it writes no record for an episode that fails. Outside it: the `reasoning_gym_env` session factory receives the task and does not use it, and the `openenv collect` command never passes tasks. Rollout steps 3 and 5 close these gaps. An adaptive sampler also needs each episode's reward, which `CollectRunner` does not report back to its `tasks` iterable today; trainers call `update` from their own loop (see the example), and open question 7 asks whether `CollectRunner` should take a per-episode callback. + +#### 4. `CurriculumSampler` + +The sampler keeps, for each key `k` (a difficulty tier, a group, or a task id), decayed success counts and an estimate `p̂_k`: + +``` +s_k ← γ · s_k + successes +f_k ← γ · f_k + failures +p̂_k = (s_k + 1) / (s_k + f_k + 2) # Beta(1, 1) prior +``` + +A rollout counts as a success when its reward is at least `success_threshold` (default `1.0` for binary rewards; open question 5 asks about graded rewards). The decay `γ` (default `0.9` per update) lets the estimate follow a policy that is improving. + +Each key is sampled with probability proportional to the chance that a group there carries signal: + +``` +w_k = 1 - p̂_k^G - (1 - p̂_k)^G + ε +``` + +where `G` is the trainer's group size and `ε` (default `0.05`) keeps every key reachable so the sampler can notice when a hard tier becomes learnable. Unseen keys start at `p̂ = 0.5`, the maximum signal, so every tier is tried early. Within a key, tasks are drawn uniformly without replacement. + +The weight ranks keys; it is not a calibrated probability. It treats the `G` rollouts of a group as independent draws with one success rate, but the questions of one tier differ in difficulty, so the outcomes inside a group are correlated and groups in which every rollout fails are far more common than the formula assumes. In the pilot, at the success rates of the uniform run's last quarter, it predicted about twice as many groups with signal as a regression on the log's per-step tier counts estimates for the two hardest tiers (0.98 and 0.92 against 0.46 and 0.44). In the trial the curriculum's own estimates did the same on tiers 2 to 4: 0.96 and 0.89 against 0.52 and 0.46 on the two hardest, in the last quarter of each run, with a correlation of about 0.5 between the outcomes of one group (pooled over the three seeds, so part of it may come from differences between seeds and from learning within the window). The factor is similar on those tiers, so it moves the sampling probabilities little. On a solved key the weight never reaches `ε`: with the default decay and `G = 8`, `p̂` approaches 81/82, so the key keeps a signal term of about 0.09, and in the trial the curriculum gave each of tiers 0 and 1, both solved from the start, about 7% of its episodes. + +The `CurriculumSampler` counts and reports costs but does not weight by them; trainers use them for budgets, concurrency and rollout allocation. The first version of this RFC divided the weights by the token cost of each key, and on our trial that did not help (Alternatives Considered). Costs come from the trainer, which knows them: the number of tokens each completion generated (the length of its generated ids, never a count made by tokenizing the text again), the token records that RFC 006's harness already carries, or a spend that the environment reports with each episode, as issue #1382 proposes (`metadata["service_spend"]`). The sampler never sees rollouts, so it takes no cost function. + +These weights are a heuristic, not an optimum: they move probability toward keys with more signal without abandoning the others. [Knapsack RL](https://arxiv.org/abs/2509.25849) uses the same expression, with a prompt's number of rollouts in place of `G`, to share a fixed rollout budget among the prompts of a batch. It multiplies that expression by an information-gain term that favours harder prompts, and warns that selecting prompts for a high share of groups with signal can concentrate on easy ones; in our trial the curriculum gave tiers 2 to 4 85% to 87% of its episodes. The reference experiment measures these weights on finer tiers. + +#### 5. `EnvEvalHarness` + +An implementation of the existing `EvalHarness` base class: + +```python +class EnvEvalHarness(EvalHarness): + def __init__(self, session_factory: ResourceSessionFactory, + harness_adapter: HarnessAdapter, + tasks: Sequence[Any], + output_dir: Path | None = None): ... + + def run(self, harness_version, library_versions, dataset, eval_parameters) -> dict: + ... +``` + +`tasks` is the list of task specs to evaluate, for example `TaskAPIClient.get_task_range("test")`; the harness does not need the client itself. `eval_parameters` accepts: + +| Key | Default | Meaning | +|-----|---------|---------| +| `limit` | `None` | Evaluate only the first `limit` tasks (for cheap in-training evals). | +| `seeds` | `[0]` | Seeds crossed with each task. | +| `samples_per_task` | `1` | `n` completions per task, for pass@k. | +| `temperature` | `0.0` | Sampling temperature; must be > 0 when `samples_per_task > 1`. | +| `success_threshold` | `1.0` | Reward at or above which an episode counts as solved. | +| `model`, `llm_endpoint`, `max_tokens`, `max_turns` | as in `openenv collect` | Policy and rollout limits. | + +It enumerates `tasks` with a `SequentialSampler`, starts each episode through the session factory with its task and seed, runs it through the harness adapter as `CollectRunner` does, writes every episode as JSONL through the existing `RolloutSerializer`, failed episodes included, and returns scores: + +- `success_rate` with its standard error, `mean_reward`, `mean_turns`, `mean_completion_tokens` +- `pass@k` for each `k ≤ samples_per_task`, with the unbiased estimator `1 - C(n - c, k) / C(n, k)` +- `error_rate`: episodes that ended in a harness or parse failure. They count as not solved in `success_rate`, so a policy cannot raise its score by failing to answer, and `error_rate` shows how many of the failures were not wrong answers +- `by_difficulty` and `by_group` breakdowns when task specs carry those keys +- `per_task`: task id → outcome, which is what checkpoint comparison needs + +`mean_completion_tokens` needs the generated token count of each episode. The pipeline has it and drops it today: the model client does not keep the provider's usage numbers, and the episode record does not copy the generated ids that the harness collects. Rollout step 5 keeps them. + +In a single-turn text environment such as `reasoning_gym_env`, the harness expects the answer as a tool call: a plain-text answer ends the episode without a reward, and it counts as failed. The harness adapter for such an environment needs a model step that turns the text into that call (rollout step 5). + +A helper compares two results: + +```python +diff = compare_eval_results(before, after) +diff.fixed # task ids that went from fail to pass +diff.regressed # task ids that went from pass to fail +``` + +#### 6. `openenv eval` + +A CLI command with the same flags as `openenv collect` for environment, endpoint and model, plus `--split` and the `eval_parameters` above. It reads the tasks through `TaskAPIClient`, writes episodes and an `EvalResult` JSON to `--output-dir`, and `--compare ` prints the fixed and regressed tasks. + +#### 7. Trainer integration + +Core provides a framework-neutral schedule object; framework adapters live in examples, following the recipes RFC: + +```python +class EvalSchedule: + def __init__(self, harness: EnvEvalHarness, config: EvalConfig, + every_n_steps: int, on_result: Callable[[int, EvalResult], None]): ... + def maybe_run(self, step: int) -> EvalResult | None: ... +``` + +A plain Python loop, an MLX loop, or a TRL callback calls `maybe_run(step)` after each optimizer step. The policy under evaluation is served through an OpenAI-compatible endpoint (vLLM, `mlx_lm.server`, or any other server), so the harness does not depend on how the trainer holds the weights. `on_result` receives the step, and the trainer should log the sampler's `metrics()` next to it, so learning curves can be drawn against episodes and against cost, not only against steps or wall-clock time. + +A sampler needs two hooks in a trainer: the trainer draws its prompts from `sample(n)`, and it calls `update` on every scored group before anything drops groups. Trainers increasingly drop the groups whose rewards are all equal: DAPO's dynamic sampling, the zero-advantage gate of prime-rl, which updates its sampler before the gate ([`Curriculum.on_result`](https://github.com/PrimeIntellect-ai/prime-rl/blob/855a2b815b804857835c635b3fd2b1ae5400f96c/src/prime_rl/orchestrator/curriculum/base.py)), and an open pull request for TRL's asynchronous GRPO trainer ([huggingface/trl#7609](https://github.com/huggingface/trl/pull/7609)). A sampler updated only with the groups that survive such a filter never sees that a key is solved or out of reach. The trial's trainer calls `update` on every group, then leaves the groups without signal out of the gradient. + +TRL has no sampler API today. In `GRPOTrainer` the training prompts come from the dataset, and the reward functions receive the dataset columns and the generated ids, so a `datasets.IterableDataset` that draws its rows from `sample(n)` (a map-style dataset goes through TRL's `RepeatSampler` instead) and a reward function that calls `update` would connect a sampler without changes to TRL. We have read that path in TRL's code, not run it, and it has limits. The trainer draws its prompts one or two rounds before they are generated, so an update first affects the tasks drawn one or two rounds later. `shuffle_dataset` must be off. On several processes, each process draws from its own copy of the sampler, the copies must draw the same tasks, and each reward function sees only its own process's rollouts, so the path as described works on one process. And the asynchronous trainer reads the dataset and scores the groups in a separate rollout process, so the sampler lives in that process and a checkpoint has to fetch its state from there. + +Asynchronous trainers add two limits of their own. Updates arrive late and out of order, and the decay weighs groups in the order they arrive, so keys whose rollouts finish first are refreshed sooner; large runs report this bias toward short rollouts, mostly after a start or a restart ([MiMo-V2.6](https://arxiv.org/abs/2610.11959), section 5.5), and [MiniMax-M2](https://arxiv.org/abs/2605.26494) bounds it with a windowed FIFO over the generation queue. And a seed reproduces a run only in a synchronous trainer such as the trial's: in an asynchronous one the draw order depends on completion times. + +Shares across environments or data sources are also the trainer's. Large runs set a share per source and keep it in the rollout scheduler, giving slow sources more concurrent rollouts (MiMo-V2.6) or a larger rollout quota ([LongCat-Flash-Thinking-2601](https://arxiv.org/abs/2601.16725)), and prime-rl draws the source by a fixed ratio before asking that source's sampler for a task. A sampler works inside one source's share, keyed by `difficulty` or `group` within that source; the trainer already knows which environment each task belongs to. + +### Key Design Decisions + +#### D1. Sampling lives on the trainer side, not in the environment + +- **Chosen approach**: Samplers run in orchestration code and hand each task to the session factory, which turns it into what the environment needs: `reset(split=..., index=...)` for an environment that follows the Task API guide. Environments only publish task metadata. +- **Rationale**: This keeps the invariants intact. Rewards are still computed inside the environment; the sampler only reads them. The agent has no path to influence which task it gets, because selection goes through the session factory and `reset()`, which stay on the infrastructure side of the dual API boundary. The environment server stays free of cross-episode state, which also keeps "one env = one trajectory" and stacking of environment instances simple. It is also where task selection already happens today: the Harbor harness picks `indices` on the trainer side. +- **Trade-offs**: Each trainer process holds its own statistics. A distributed trainer with many rollout workers must aggregate them; `state_dict` gives a merge point, but a shared sampler service is out of scope for this RFC. + +#### D2. The curriculum targets GRPO signal, and reports cost + +- **Chosen approach**: Weight keys by the probability that a group of size `G` has non-zero reward variance. Count cost in a unit the trainer chooses and report signal per unit of cost, without weighting by it. +- **Rationale**: The probability of signal is the quantity that decides whether a GRPO group produces a gradient, and it needs one parameter the trainer already knows (`G`). On the pilot's seed, the signal-only curriculum produced 39% more gradient-bearing groups per episode than uniform sampling but 42% fewer per generated token, and the first version of this RFC read that as a reason to divide by cost. Three more seeds changed the reading: there the signal-only curriculum got more groups with signal per generated token than uniform sampling on every seed, and higher test success per token, while dividing by token cost as first proposed lowered both, and a corrected division matched the signal-only weights per token without beating them (Alternatives Considered). Learning-progress curricula (weighting by the change in success rate) remain a reasonable alternative and can be added as another `TaskSampler`, as can a sampler that accounts for cost. +- **Trade-offs**: Gradient-bearing groups are a proxy for learning, not the objective. In the pilot, uniform sampling trained mostly on an easier tier and still improved the two hardest tiers, from 7% and 0% to 46% and 24%. The evaluation harness, not the sampler's counters, is what decides whether a sampler helps. The estimates also need repeated observations per key: for large datasets without `difficulty` or `group`, per-task statistics are sparse and the sampler behaves close to uniform until tasks repeat. Publishing closely spaced difficulty tiers is the intended remedy. + +#### D3. Evaluation reuses the collect pipeline's pieces + +- **Chosen approach**: `EnvEvalHarness` is built from the same pieces as `CollectRunner` (the session factory, the harness adapter and `RolloutSerializer`), plus a `SequentialSampler` and aggregation. +- **Rationale**: Training rollouts, collected datasets, and evaluation episodes then go through one code path and one serialization format (design principle 1). Evaluation episodes can be pushed to the Hub with the existing `push_to_hf_hub` for inspection. +- **Trade-offs**: `CollectRunner` itself passes no seed to the session and writes nothing for a failed episode, so until rollout step 5 adds both, the harness drives the pieces itself instead of calling `CollectRunner`. Evaluation also inherits the limits of the collect pipeline. The `openenv collect` CLI resolves only `openspiel:` and `reasoning_gym:` environment specs today, so `openenv eval` would start with the same environment coverage and grow with `collect`; the Python `EnvEvalHarness` accepts any `ResourceSessionFactory` from the start. And the pipeline drives every turn of an episode, while a Harbor server runs a whole agent in one call (`run_rollout`) and never goes through `reset`, so the first version does not cover Harbor (open question 8). + +#### D4. Evaluation checks reproducibility instead of assuming it + +- **Chosen approach**: Before a run, `EnvEvalHarness` makes two checks with fresh sessions. It resets two sessions with the same `(split, index, seed)` and compares a hash of the initial observations; then it resets two sessions with different tasks and checks that their initial observations differ. If either check fails, it warns and records `reproducible: false` in the result. +- **Rationale**: An eval that silently sees different tasks across checkpoints produces curves that cannot be trusted. Environments that ignore `seed` have existed (#1102, #1250), and RFC 008's `runtime.seed_control` and `runtime.episode_determinism` checks (pending in #1321) address the same concern at validation time. The second check catches an environment that ignores the task selection, which the first cannot see: the server passes `reset()` only the arguments its signature names, so an environment whose `reset()` does not name `split` and `index` returns the same default task twice and looks reproducible. `reasoning_gym_env` behaves this way today. +- **Trade-offs**: Four extra resets per run. The checks cover the initial observation only, not later nondeterminism in the environment. + +#### D5. Task metadata stays on HTTP for now + +- **Chosen approach**: `TaskAPIClient` uses the existing HTTP Task API routes. This RFC adds no routes and no wire types. +- **Rationale**: The Task API guide documents these routes as HTTP-only, with no WebSocket message types, so that the `/ws` session protocol stays focused on `reset`, `step`, `state` and `close`. The routes serve metadata, never the step loop. +- **Trade-offs**: This sits in tension with the communication-patterns invariant, which lists WebSocket for all environment communication, "Gym-like API + metadata", while HTTP is being deprecated (#252, #1156). This RFC does not resolve that tension; it builds on the Task API as documented. Because every sampler and the evaluation harness go through `TaskAPIClient`, moving the Task API to WebSocket later changes only that class. See open question 1. + +#### D6. A generic Task API client in core + +- **Chosen approach**: Add `TaskAPIClient` to `openenv.core`, separate from `EnvClient`, returning untyped task specs. +- **Rationale**: The Task API guide leaves Task API methods out of the core clients because task specs are environment-specific. `TaskAPIClient` keeps specs untyped and adds no methods to `EnvClient`, so that reasoning still holds; what it removes is the duplicated plumbing that is not environment-specific (route names, HTTP base derivation from a `ws://` URL, headers, paging, error mapping). Samplers and the evaluation harness need to list tasks for any environment, which is impossible if each client names its helpers differently. +- **Trade-offs**: One more public class in core. Typed task specs remain each environment client's job. + +## Examples + +### Server side: publishing difficulty tiers + +A procedurally generated environment exposes tiers as a finite split. For `reasoning_gym_env`, a task maps to a dataset name, a difficulty configuration, and a seed. The tiers below come from reasoning-gym's own curriculum for `chain_sum`, whose levels set the number of terms (2 to 12) and of digits (1 to 10): + +```python +from reasoning_gym.arithmetic.chain_sum import ChainSumCurriculum +from reasoning_gym.coaching.base_curriculum import DefaultCurriculumContext, RangeAttributeMode + +# (level of num_terms, level of num_digits): 2x1, 3x1, 3x2, 4x2, ... 7x5 terms x digits +LEVELS = [(0, 0), (1, 0), (1, 1), (2, 1), (2, 2), (3, 2), (3, 3), (4, 3), (4, 4), (5, 4)] +TASKS_PER_TIER = {"train": 2000, "test": 100} +SPLIT_SEED_OFFSET = {"train": 0, "test": 500_000_000} # train and test seeds never collide + +def tier_config(tier: int): + curriculum = ChainSumCurriculum() + terms, digits = LEVELS[tier] + curriculum.set_attr_level("num_terms", terms) + curriculum.set_attr_level("num_digits", digits) + return curriculum.generate_configuration( + context=DefaultCurriculumContext(mode=RangeAttributeMode.UPPER_BOUND)) + +class ReasoningGymEnvironment(Environment): + def list_splits(self) -> list[str]: + return ["train", "test"] + + def num_tasks(self, split: str) -> int: + return len(LEVELS) * TASKS_PER_TIER[split] + + def get_task(self, split: str, index: int) -> dict: + if not 0 <= index < self.num_tasks(split): + raise IndexError(index) + tier = index % len(LEVELS) # tiers interleave: the first N tasks cover every tier + return {"id": f"chain_sum-{split}-{index}", "split": split, "index": index, + "group": "chain_sum", "difficulty": tier} + + def reset(self, split: str | None = None, index: int | None = None, + seed: int | None = None, episode_id: str | None = None, **kwargs): + if index is None: + # No selection: keep today's seed-driven behaviour, as the Task API checklist asks. + return self._reset_from_kwargs(seed=seed, episode_id=episode_id, **kwargs) + split = split or "train" + task = self.get_task(split, index) + # The seed picks a variant of the task, so tasks crossed with seeds differ. + question_seed = (seed or 0) * 1_000_000_000 + SPLIT_SEED_OFFSET[split] + index + return self._reset_dataset("chain_sum", tier_config(task["difficulty"]), + seed=question_seed) +``` + +`_reset_from_kwargs` and `_reset_dataset` stand for the environment's existing reset paths. + +The library has no level with negative numbers, so these tiers have none, unlike the pilot's two hardest tiers. On them, the untrained Llama-3.2-1B-Instruct scored 100, 99, 99, 100, 67, 52, 15, 11, 2 and 4% (100 test questions per tier): the steps between tiers 4 and 8 are the ones a curriculum can use for that model. The three easiest tiers have few possible questions, so 100, 38 and 1 of their 100 test questions also appear among 2,000 train tasks; the guidance of section 1 applies. The pilot numbered its tasks inside each tier and derived each seed from the tier and the position, so the same index names different questions there and in this example. + +### Client side: a plain training loop + +```python +client = TaskAPIClient.from_env_client(env) +tasks = client.get_task_range("train") +sampler = CurriculumSampler(tasks, group_size=8, key="difficulty") + +evaluator = EnvEvalHarness(session_factory, harness_adapter, client.get_task_range("test")) +schedule = EvalSchedule( + evaluator, + EvalConfig(harness_name="EnvEvalHarness", harness_version="0.1", + library_versions={}, dataset="reasoning_gym:chain_sum", + eval_parameters={"llm_endpoint": POLICY_URL, + "model": "policy", "temperature": 0.0}), + every_n_steps=50, + on_result=lambda step, r: logger.log({"step": step, **sampler.metrics(), **r.scores}), +) + +schedule.maybe_run(step=0) # baseline before any training +for step in range(1, num_steps + 1): + batch = sampler.sample(n=prompts_per_step) + groups = [rollout_group(task, G=8) for task in batch] # trainer-specific + for task, group in zip(batch, groups): + sampler.update(task, [ep.reward for ep in group], + costs=[len(ep.completion_ids) for ep in group]) + trainer_update(groups) # trainer-specific; may skip groups without signal + logger.log({"step": step, **sampler.metrics()}) + schedule.maybe_run(step) +``` + +### CLI: comparing checkpoints + +The environment spec is positional, as in `openenv collect`: + +```bash +openenv eval reasoning_gym:chain_sum --base-url http://localhost:8001 --split test \ + --llm-endpoint http://localhost:8080/v1 --model base \ + --output-dir runs/eval-step0000 + +openenv eval reasoning_gym:chain_sum --base-url http://localhost:8001 --split test \ + --llm-endpoint http://localhost:8080/v1 --model ckpt-1000 \ + --output-dir runs/eval-step1000 --compare runs/eval-step0000/result.json +``` + +## Validation Plan + +The claims in this RFC are measured, not assumed. The reference experiment is designed to fit one laptop-class machine. + +### Preliminary Results + +We ran the samplers before proposing the API: first a pilot on one seed, then a trial on three more seeds, with predictions and the rules for reading the results written down before it ran. The code is a standalone trainer, not the proposed API: a GRPO loop in MLX with the samplers of section 4 and of "Weighting by cost" in Alternatives Considered. The pilot's code, its per-step logs and a figure of both learning curves are at [Matheus-F-Scatolin/OpenEnv@rfc-013-pilot](https://github.com/Matheus-F-Scatolin/OpenEnv/tree/rfc-013-pilot/pilot); the figure is also in the PR description. The trial's code, its 28 run logs, the predictions and a script that regenerates every table and figure from those logs are at [Matheus-F-Scatolin/OpenEnv@rfc-013-trial](https://github.com/Matheus-F-Scatolin/OpenEnv/tree/rfc-013-trial). + +Setup, the same for both: + +- **Model**: Llama-3.2-1B-Instruct, LoRA rank 16 on every linear layer (11.3M trainable parameters). +- **Environment**: `reasoning_gym_env` vendored from commit `9e0fd9b`, run in-process, with the pilot's five `chain_sum` tiers: 2×1, 3×2, 4×3, 5×4 and 6×5 terms × digits, negative numbers in the last two. 2,000 train tasks and 100 test tasks per tier, with disjoint seeds. The reward is whatever the environment's `step()` returns. +- **Training**: 8 prompts × 8 completions per step (64 episodes), temperature 1.0, at most 384 tokens per completion, on-policy GRPO with one update per batch, no KL term, token-level loss normalization. The sampler is updated with every group; groups with no reward variance are then left out of the gradient and not replaced. +- **Evaluation**: greedy decoding on all 500 test tasks. The standard error of one evaluation is about 2 points. +- **Reproducibility**: the model's chat template writes today's date into every prompt, so the date is fixed to the pilot's day; with it fixed, the trial's trainer reproduced the pilot's seed-0 logs step for step, for the first 370 steps of uniform sampling and 158 of the curriculum. +- **Hardware**: a MacBook Pro (M5 Pro, 24 GB). + +#### The pilot (seed 0) + +150 training minutes per sampler, one seed, an evaluation every 100 steps. At the same number of environment episodes, the curriculum was ahead at 9 of 10 evaluation points: + +| Training episodes (steps) | Uniform | Curriculum | Difference | +|---|---|---|---| +| 6,400 (100) | 50.6% | 58.0% | +7.4 | +| 19,200 (300) | 62.2% | 65.2% | +3.0 | +| 32,000 (500) | 65.4% | 68.4% | +3.0 | +| 44,800 (700) | 62.8% | 63.6% | +0.8 | +| 51,200 (800) | 62.4% | 69.0% | +6.6 | +| 64,000 (1,000) | 67.4% | 71.4% | +4.0 | + +Over the whole run: + +| | Uniform | Curriculum | +|---|---|---| +| Steps in 150 minutes | 2,196 | 1,081 | +| Tokens per completion | 44 | 106 | +| Groups with no gradient | 79.5% | 71.5% | +| Gradient-bearing groups per 1,000 episodes | 25.7 | **35.6** | +| Gradient-bearing groups per million generated tokens | **578** | 337 | +| Test success, start → end | 52.0% → 73.2% | 52.0% → 72.4% | +| Test success at ~6 M generated tokens | **73.6%** | 67.8% | +| Tiers 3 and 4, start → end | 7%, 0% → 46%, 24% | 7%, 0% → 43%, 22% | + +Per episode, the curriculum learned faster: it reached 72.4% after 69,184 episodes, which uniform sampling reached only after about 100,000. Per generated token, uniform sampling was ahead from about 3.3 M tokens on. The curriculum's completions were 2.4 times longer than uniform's, but choosing harder tiers explains a factor of only about 1.2 of that: priced at the uniform run's own completion lengths per tier, the curriculum's mix of tiers costs 54 tokens per completion, against 44. The rest came from the two trained models moving apart: on the same 500 test questions, the untrained model answered with 73 tokens on average, the uniform-trained model with 45 and the curriculum-trained model with 86, both read at 5 million generated tokens. + +#### Three more seeds + +Seeds 1, 2 and 3, each run stopped at 5 million generated tokens, with an evaluation every 100 steps and every 250,000 tokens. Four samplers: uniform; the curriculum of section 4; the token-weighted curriculum as this RFC first proposed it; and a corrected token-weighted form. Both token-weighted forms are described under "Weighting by cost" in Alternatives Considered. Test success is the mean of the evaluations at 4.0 to 5.0 million generated tokens ("per token") or at steps 100 to 700 ("per episode"). Difference to uniform sampling, in points: + +| | Seed 0 (pilot) | Seed 1 | Seed 2 | Seed 3 | +|---|---|---|---|---| +| Curriculum, per token | −7.7 | +5.5 | +18.3 | +14.4 | +| Curriculum, per episode | +2.3 | +4.3 | +9.7 | +6.5 | +| Token-weighted, first form, per token | | +1.2 | −1.6 | −1.7 | +| Token-weighted, corrected, per token | | +5.6 | +15.2 | +16.1 | +| Token-weighted, corrected, per episode | | +10.3 | +10.1 | +9.2 | + +Uniform sampling reached 66.0%, 55.6% and 59.9% per token on seeds 1 to 3 (72.2% on seed 0); the curriculum 71.6%, 73.9% and 74.4% (64.5%). Per episode, seed 2 is read at steps 100 to 700 for each pair of runs that reached step 700. The rule we wrote before the run read every run of a seed only up to the step all of them reached: one run of seed 2 stopped at step 223, and that reading gives +3.9 for the curriculum, with the same conclusion. + +At 5 million generated tokens: + +| | Uniform | Curriculum | Token-weighted, first form | Token-weighted, corrected | +|---|---|---|---|---| +| Groups with signal per million tokens, seeds 1 to 3 | 285, 287, 706 | 526, 730, 835 | 492, 119, 750 | 345, 476, 850 | +| Share of episodes on tier 0, solved from the start | 20% | 7%, 7%, 7% | 17%, 14%, 22% | 8%, 7%, 5% | +| Tokens per training completion, seeds 1 to 3 | 80, 79, 33 | 73, 59, 49 | 64, 351, 45 | 111, 80, 47 | + +What we take from both: + +1. **The mechanism works as designed.** The curriculum concentrated on the tiers nearest `p = 0.5` and raised the groups with signal per episode on every seed. +2. **The pilot's per-token result belonged to its seed.** The curriculum's completions were longer in the pilot; in a 30-step check on four new seeds they were longer by 20 tokens or more in only one, and on seeds 1 and 2 of the trial the curriculum's completions were the shorter ones. Across the four seeds the curriculum was ahead per episode four times and per token three times. +3. **Dividing by token cost did not help.** The comparisons with the signal-only curriculum were not among the questions we wrote down before the trial; we read them with the thresholds we wrote there. The first form sent 14% to 22% of its episodes to tier 0, which the model solved from the start, against about 7% for the signal-only curriculum, and learned less on every new seed: −4.3, −19.9 and −16.2 points per token. The corrected form fixed the allocation and was neither better nor worse on the 500 test questions: +0.1, −3.1 and +1.6 points per token, and +6.1, +0.4 and +2.7 per episode, below the +4 that "better" needs. On tiers 2 to 4 alone it was ahead per episode by +9.6, +0.7 and +4.3 (+4.9 on average), which passes that bar when each pair of runs is read to step 700, but not in the window we wrote before the trial, which reads seed 2 to step 200, where the corrected form is 0.7 points behind. Per token, the unit that weighting by token cost exists to improve, it did not help, and it got fewer groups with signal per generated token than the signal-only curriculum on two of three seeds. This version of the RFC therefore leaves cost out of the weights. +4. **A reward that scores a cut answer lets a run degenerate.** In one of the twelve runs (the first token-weighted form, seed 2), the policy learned to repeat its tagged answer until the length limit. From about step 50 on, every completion was cut at 384 tokens, and test success stayed at 54% (52% untrained), so the run spent its 5 million tokens by step 223. The environment scores the answer a cut completion contains, so nothing discouraged the repetition. This is the environment's reward, not the sampler, and the run stays in every number above. +5. **The tiers were too coarse.** After tier 2 reached about 95%, no tier was near 50%, which limited what any curriculum could do. + +Caveats: three new seeds per sampler are a pattern, not a significance test; ahead on three seeds of three happens one time in eight by chance alone. One model, one task family, one budget. The trainer is synchronous, so the trial says nothing about asynchronous trainers, and the rewards of `chain_sum` are almost binary (about 0.1% of training answers got partial credit), so it says nothing about graded rewards. All 100 test questions of tier 0 also appear in train, because that tier has about 200 possible questions; the model solves it from the start, so the comparisons do not change, and on tiers 2 to 4 alone they are larger. In the pilot the two arms ran one after the other on the same machine, and the curriculum arm was slowed for its last 40 minutes by an underpowered charger; this affects the comparison by wall-clock time but not the per-episode or per-token comparisons, which do not depend on machine speed. + +### Reference experiment + +- **Arms**: `UniformSampler` and `CurriculumSampler`, with the same group size, on five seeds not used so far. Before the runs we write down the test, the reading windows, how the tiers whose test questions also appear in train are read, and what happens to a run that degenerates. +- **Environment**: `reasoning_gym_env` with the tiers of the example above. In-process today, as in the trial; through the Task API once rollout step 3 exists. +- **Hardware**: the same laptop class. In the trial a run of 5 million generated tokens spent 1.7 to 2.6 hours in training and in evaluations of 500 test questions (2.1 on average; loading the model and saving adapters not counted). With the 1,000 test questions of the ten tiers above the evaluations take twice as long, so the ten runs of uniform sampling and the curriculum need about 24 hours: three overnight sessions. A CUDA machine with vLLM reproduces the same experiment faster. +- **Measurements**: test success rate with standard error, overall and per tier, against training episodes and against generated tokens; gradient-bearing groups per episode and per token; the share of completions cut at the length limit; wall-clock time per step, reported but not used for the comparison. +- **What would support the proposal**: the curriculum at least as good as uniform sampling per episode and per token, across seeds. If the curriculum does not beat uniform sampling, we report that as is; the sampler protocol and the evaluation harness do not depend on the curriculum winning. + +## Rollout Plan + +Each step is a separate PR: + +1. This RFC. +2. Reserved task spec keys in the Task API guide, including guidance on tier spacing and on train and test overlap, and `TaskAPIClient` (with `headers`) with tests against a stub server. +3. Task API support in `reasoning_gym_env`, with difficulty tiers from reasoning-gym's curriculum levels and a session factory that turns a task into `reset(split, index, seed)`; then `finqa_env`, whose `reset(seed)` fix (#1252) is merged. +4. `SequentialSampler` and `UniformSampler` in core, and `CurriculumSampler` as an example (or in core, if the maintainers prefer), with unit tests on synthetic reward and cost streams (for example, a simulated policy whose success rate per tier improves over time and whose completion length grows with difficulty). +5. `EnvEvalHarness`, `compare_eval_results`, and `openenv eval`. In the collect pipeline: pass the seed to the session, write a record for a failed episode, keep the generated token count of each episode (the provider's usage numbers in the model client, the generated ids in the episode record), let `openenv collect` pass tasks, and a model step that turns a text answer into a single-turn environment's tool call. +6. `EvalSchedule`, with a TRL `TrainerCallback` that calls it, and the reference experiment as an example in `examples/` (a plain loop and an MLX loop for Apple Silicon). A TRL example that draws its prompts from a sampler is left for later (Trainer integration). + +## Alternatives Considered + +- **Curriculum inside the environment.** The server would track success rates and pick the next task itself. Rejected: it makes the environment stateful across episodes, mixes orchestration into the environment boundary, and breaks when many trainer workers share or stack environment instances. +- **Only filtering zero-signal groups after generation (DAPO-style).** This is complementary, not a replacement: the filter still pays for generating the groups it drops, and where they are not replaced, the step trains on fewer groups. In our trial, which does not replace them, uniform sampling made no update at all on 19% to 21% of its steps. DAPO resamples until the batch is full, prime-rl keeps dispatching new tasks while its zero-advantage gate, when configured, rejects groups, and TRL's open pull request ([huggingface/trl#7609](https://github.com/huggingface/trl/pull/7609)) drops them in its asynchronous rollout worker. A trainer can use both, as long as the sampler sees each group before the filter. +- **Weighting by cost.** The first version of this RFC divided the curriculum's weights by the decayed mean token cost of each key, `ĉ_k`, as `(signal + ε) / ĉ_k`, and used it in its example. On three more seeds it learned less than the signal-only weights per token (−4.3, −19.9 and −16.2 points; the run on seed 2 degenerated, see result 4). A solved key keeps a signal term of about 0.09 under the fixed prior, its answers are the shortest, and `ε` divided by a small cost is large, so tier 0, solved from the start, got 14% to 22% of the episodes. A corrected form, `w_k = signal_k / ĉ_k` with every key kept at a share of at least `ε` outside the division and a prior that decays with the data, gave that tier 5% to 8% of the episodes, matched the signal-only weights per token (+0.1, −3.1 and +1.6 points) and was slightly ahead per episode (+6.1, +0.4 and +2.7). It also had fewer groups with signal per generated token on two of three seeds, and longer training answers on two (111 and 80 tokens against 73 and 59): the policy changes its answer length, so dividing by a decayed mean of past lengths does not control what the next answers cost. Both forms moved episodes away from the longest, hardest tiers: tiers 3 and 4 got 25% to 34% of the episodes under the first form and 47% to 56% under the corrected one, against 52% to 60% under the signal-only weights and 40% under uniform sampling. No lab report we read divides a task's sampling weight by its cost; large RL runs keep a share per source and give slow sources more rollouts instead (MiMo-V2.6, LongCat-Flash-Thinking-2601). Outside RL for language models, [Graves et al. (2017)](https://arxiv.org/abs/1704.03003) divide a task's learning progress by the time it takes to process. Two 2026 preprints that account for cost in GRPO divide a usefulness term by the square root of the cost, to share rollouts among the prompts of a batch ([DUET](https://arxiv.org/abs/2605.08441)) or to sample the minibatches of an update ([Cost-Aware Learning](https://arxiv.org/abs/2604.28020)), not to choose tasks. A sampler that accounts for cost can be added later as another `TaskSampler`. +- **reasoning-gym's own curricula.** The library defines difficulty levels for each dataset, and its curriculum experiment lets the training code raise or lower a dataset's level. That schedule lives inside one library and one dataset at a time; the `CurriculumSampler` works on the tasks any environment publishes through the Task API, and keeps the environment unaware of the schedule (D1). The example above uses the library's levels to define the tiers, so the two fit together. +- **Relying on the Inspect AI harness for evaluation.** Inspect AI remains supported, but it runs its own task and solver definitions. An evaluation that goes through the same environment, tools and rollout path as training is what design principle 1 asks for. The two can meet at the results: an `EvalResult` could later be exported in Inspect's log format, which this RFC does not attempt. + +## Open Questions + +1. Should the Task API gain WebSocket message types, given #1156 and the WebSocket invariant, or remain an HTTP metadata surface? +2. Should `difficulty` be standardized further (for example, normalized to `[0, 1]`), or remain environment-defined and only compared within one environment? +3. Where should shared sampler state live for distributed trainers with many rollout workers? +4. Should `openenv eval` results be pushable to the Hub in a schema that the environments datasets RFC (#795) can reference? +5. What is the right default for `success_threshold` on environments whose rewards are graded rather than binary? One direction, suggested in review and not yet tested: estimate each key's signal directly, as the decayed share of its groups whose rewards vary. That is the quantity a zero-variance filter keeps, it is calibrated by construction (section 4), and it needs no `success_threshold`, only a tolerance on the spread of the rewards. With continuous rewards, where almost every group varies a little, that tolerance would have to follow the scale of the rewards, or the weight would have to grow with the size of the variance. Our task can test it only on binary rewards: about 0.1% of its training answers got partial credit. There it would mainly lower the weight of solved keys: in the trial's last quarter, 0.2% to 0.4% of the curriculum's groups on tiers 0 and 1 had signal, against the formula's 0.10. We have not run it. +6. Answered in review: should environments declare their cost? Environments declare cost meters, such as #1382's `service_spend`, for reporting, and trainers use them for budgets, concurrency and rollout allocation. The `CurriculumSampler` counts and reports them and does not weight by them. +7. Should `CollectRunner` accept a per-episode callback (task, `EpisodeRecord`) so adaptive samplers can be driven by `openenv collect` and not only by a trainer loop? +8. Should the first version of the evaluation cover environments whose server runs a whole agent per call, such as Harbor's `run_rollout`, which never goes through `reset`? +9. Should `CurriculumSampler` ship in core, next to the two simple samplers, or stay an example? A review suggested keeping it as an example, which is how this RFC offers it. + +## Related RFCs and Work + +- RFC 000: phase three (evals) of the roadmap. +- RFC 004 (Rubrics): the source of the rewards that samplers consume. +- RFC 005 (Agentic harnesses) and the collect pipeline: the rollout pieces reused by evaluation. +- RFC 006 (Agentic RL harness interception): how an agent's model calls are captured; its token records are one source of per-rollout cost. +- RFC 008 (Environment auto-validation) and #1321: repeatability checks at validation time; D4 applies the same idea at evaluation time. +- RFC 012: how Harbor's agents are captured, relevant to open question 8. +- Draft recipes RFC (#924): its open question 2 asks which environment is simple enough for smoke validation. `reasoning_gym_env` answers the "simple and cheap" half, but it is single-turn and does not exercise harness-mediated tool use. +- Environments datasets RFC (#795): possible home for published evaluation results. +- Task API guide (`docs/source/guides/task-api.md`) and the Harbor harness (`envs/harbor_env/harness.py`): the current Task API contract and the manual task selection this RFC generalizes. +- Training guide (`docs/source/guides/training.md`): training stays in the user's framework, and core depends on no trainer; the samplers and `EvalSchedule` follow that. +- Issue #1382: environments that report a per-episode spend, which a trainer can count as cost and use for budgets and concurrency. +- [Knapsack RL](https://arxiv.org/abs/2509.25849): the signal term of section 4, used to share a rollout budget among the prompts of a batch. +- prime-rl's [`TaskSampler`](https://github.com/PrimeIntellect-ai/prime-rl/blob/855a2b815b804857835c635b3fd2b1ae5400f96c/src/prime_rl/orchestrator/curriculum/samplers/base.py) and `Curriculum`: the same two hooks, with the sampler updated before the zero-advantage gate. +- [huggingface/trl#7609](https://github.com/huggingface/trl/pull/7609) (open): a zero-variance filter in TRL's asynchronous GRPO trainer. +- [MiMo-V2.6](https://arxiv.org/abs/2610.11959), [LongCat-Flash-Thinking-2601](https://arxiv.org/abs/2601.16725), [MiniMax-M2](https://arxiv.org/abs/2605.26494) and [INTELLECT-3](https://arxiv.org/abs/2512.16144): shares per source, rollout scheduling and difficulty filtering in large RL runs. +- [DUET](https://arxiv.org/abs/2605.08441) and [Cost-Aware Learning](https://arxiv.org/abs/2604.28020): choices in GRPO that divide by the square root of the cost, for the rollouts of the prompts in a batch and for the minibatches of an update. +- #1250 and its fix #1252 (by one of the authors), merged on 2026-10-13: the `finqa_env` seed issue named in rollout step 3. diff --git a/rfcs/README.md b/rfcs/README.md index 8c2824b2a..59a26b85b 100644 --- a/rfcs/README.md +++ b/rfcs/README.md @@ -89,6 +89,7 @@ Each RFC should include the following sections: ### Reward & Evaluation - [004-rubrics.md](./004-rubrics.md) - Rubric System for Reward Computation +- [013-task-sampling-and-training-evals.md](./013-task-sampling-and-training-evals.md) - Task Sampling and In-Training Evaluation: `TaskSampler` (sequential, uniform; a curriculum sampler as an example), `EnvEvalHarness` and `openenv eval` ### Agentic Harnesses - [005-agentic-harnesses.md](./005-agentic-harnesses.md) - Agentic Harness Integration (OpenClaw, Claude Code, etc.)