dew.objectives.rl

Online RL objectives: rollouts, preference losses and group updates.

These compose the array math in dew.rl; the import gate in tests/test_rl_imports.py keeps that arrow one way. dew.rl may read dew, and nothing under dew outside these two packages may read dew.rl.

Name	Summary
ADVANTAGES_KEY
IDS_KEY
OLD_LOG_PROBS_KEY	The proximal policy’s log-probabilities, recorded or rescored before the update.
PREFERENCE_IDS_KEY	Batch key holding the [B, 2, S] pair ids, chosen at index 0. Documented in dew.data.
PREFERENCE_MASK_KEY	Batch key holding the [B, 2, S] completion marks, 1 on completion tokens. Documented in dew.data.
RESPONSE_MASK_KEY	1 on every sampled id of a trainable session, the tokens that carry loss mass.
REWARDS_KEY	Per image, the reward its trajectory scored.
Action	Record one actual model call, including EOS and both likelihood distributions.
Call	One model call as the engine served it.
CodeReward	Run the completion’s program on every test case; reward the fraction that pass.
ContainerRunner	Run each program in a fresh, network-less container of image.
DPOObjective	Train a policy on preference pairs with the DPO loss (arXiv:2305.18290, eq.
Environment	The caller owns execution and releases resources on context exit.
EnvironmentFactory
EnvironmentSource	Run Environment sessions against a RolloutServer, one worker thread each.
Episode
EpisodeCancelled
EpisodeFailure	Report a failed collection, keeping the partial episode available to the caller.
EpisodeId	Identify one sample within an attempt, reproducibly from checkpointed work.
EpisodeInference	Draw actions from a bound policy, returning the actual sampling likelihoods.
EpisodeJournal	Persist completed turns before the next tool call or trainer update.
EpisodeRecorder
EpisodeRollout	Collect complete episode groups under one policy snapshot, then train on their actions.
EpisodeStatus
FlowGRPOObjective	Train a rectified-flow policy on clipped, coordinate-normalized gradients.
FlowReward	Score decoded [-1, 1] samples and repeated source rows, one scalar per sample.
FlowRollout	Collect complete prompt groups and the first train_steps transitions.
GRPOObjective	Train a policy on sampled rollouts with the GRPO loss (arXiv:2402.03300, eq.
MathReward	Score one when the final answer equals the reference as a rational number.
Observation	Carry the exact next model context, or a terminal environment result.
Outcome
PPOObjective	Train a policy and a critic together on one token mass.
PPORollout	Collect episodes, then add critic baselines and verl’s masked GAE.
ProcessRunner	Run each program as a resource-limited Linux process in a temporary directory.
Program	Files to write, the argv to run beside them, and its stdin.
PromptSource	Draw one completion per sample of a prompt_tasks task and score its decoded text.
Publisher	Where the trainer’s weights go: load serves them under version. Documented in dew.objectives.rl.scheduler.
RecoverableEnvironment	Restore tool state from an opaque snapshot, without replaying completed calls.
Reward	Score (data_source, completion, ground_truth, extra_info).
RolloutScheduler	Train on complete rollout groups from source, ahead task batches early. Documented in dew.objectives.rl.scheduler.
Runner	Run one program to its end under limits; raise only when it cannot start.
SampledRollout	Draw G completions per prompt and pack them as one-call sessions.
SandboxFleet	Run programs on workers concurrent sandboxed workers.
SandboxLimits	Bound one worker by RLIMIT_CPU and RLIMIT_AS, plus parent-enforced session and IO limits.
SchedulerRecord	What one trainer call consumed and what it cost. Documented in dew.objectives.rl.scheduler.
Session	One harness session: its calls in submission order and how it ended.
SessionSource	Anything that turns tasks into sessions.
Status	How a session ended, which decides whether it trains.
SubprocessEnvironment	A user-selected JSON-lines worker implementing reset and step.
Task	One unit of work a session source runs: an identity and its source-specific payload.
Transition
ValueHead	Project a decoder’s hidden states to one float32 value per position.
Verdict	How a program ended.
Verifier
code_block	The last fenced block tagged language, else the last untagged one, else None.
pack	Strictly merge each trained session’s calls, then pack the chains into [rows, width].
prompt_tasks	One single-turn task per prompt row: its ids and the three reward strings.
session_of	Read an episode as an engine-style Session of the advantage group group.

ADVANTAGES_KEY

attribute source

ADVANTAGES_KEY = 'advantages'

IDS_KEY

attribute source

IDS_KEY = 'input_ids'

OLD_LOG_PROBS_KEY

attribute source

OLD_LOG_PROBS_KEY = 'old_log_probs'

The proximal policy’s log-probabilities, recorded or rescored before the update.

GRPO’s PPO ratio compares the current raw policy with this one. When a batch carries none, the behavior log-probabilities stand in.

RESPONSE_MASK_KEY

attribute source

RESPONSE_MASK_KEY = 'response_mask'

1 on every sampled id of a trainable session, the tokens that carry loss mass.

REWARDS_KEY

attribute source

REWARDS_KEY = 'rewards'

Per image, the reward its trajectory scored.

Action

dataclass source

class Action(
    context: tuple[int, ...],
    tokens: tuple[int, ...],
    raw_log_probs: tuple[float, ...],
    behavior_log_probs: tuple[float, ...],
    terminated: bool,
    policy_step: int,
    sampling: Sampling,
    *,
    _binding_id: str = '',
)

Record one actual model call, including EOS and both likelihood distributions.

Raw probabilities belong to the unmodified model; behavior probabilities include Sampling controls. EOS ends the model turn, not the episode. Context and action ids are nonnegative integers. Actions stop at the first configured EOS, and terminated agrees with that final token. Vocabulary upper bounds belong to the model-aware caller.

Call

dataclass source

class Call(
    prompt_ids: tuple[int, ...],
    sampled_ids: tuple[int, ...],
    behavior_log_probs: tuple[float, ...],
    finish_reason: str,
    version: int,
    routed_experts: np.ndarray | None = None,
    support: tuple[tuple[int, ...], ...] | None = None,
)

One model call as the engine served it.

sampled_ids includes EOS when finish_reason is a natural stop, and behavior_log_probs holds the engine-reported log-probability of each sampled id. version is the served policy version when the request was submitted, the oldest policy that may have produced any of its ids.

Two engine records are optional. routed_experts is the mixture’s routing for every id the engine forwarded, [len(prompt_ids) + len(sampled_ids) - 1, layers, top_k] expert ids (vLLM’s routed_experts, SGLang’s meta_info.routed_experts; the last sampled id is never forwarded), kept in the dtype the engine shipped, which the trainer replays (dew.nn.moe.Routes). support is, per sampled id, the token ids the sampler’s top-k/top-p filters kept (vLLM’s sampling_mask); the behavior log-probability is then the filtered one, and the trainer renormalizes over the same support. None means no filter.

CodeReward

dataclass source

class CodeReward(
    fleet: SandboxFleet,
    interpreter: tuple[str, ...] | None = None,
    language: str = 'python',
    filename: str = 'main.py',
    all_or_nothing: bool = False,
)

Run the completion’s program on every test case; reward the fraction that pass.

The ground truth is a JSON list of {"stdin", "stdout"} cases. A case passes when the program completes and its stdout matches, line by line up to trailing whitespace. A completion without a code block scores zero, as do timeouts, crashes and oversized output. all_or_nothing scores one only when every case passes. interpreter is the argv the program file is appended to; None takes the fleet runner’s python, so a container runs the image’s interpreter and a process runs this one.

ContainerRunner

dataclass source

class ContainerRunner(
    image: str,
    runtime: str = 'docker',
    cpus: float = 1.0,
    pids: int = 64,
    user: str = '65534:65534',
    python: tuple[str, ...] = ('python', '-I', '-S'),
)

Run each program in a fresh, network-less container of image.

runtime is the Docker-compatible CLI (docker or podman). The job directory is mounted read-only at /work, the working directory; /tmp is a small writable tmpfs. Memory is capped at memory_bytes with no swap, CPU time at cpu_seconds (SIGXCPU, then SIGKILL a second later), the processor share at cpus and the process count at pids. At the wall deadline the client that runs it is killed, then the container is force-removed by name. A runtime that exits without creating the container (no daemon, no permission, no image) raises.

python: tuple[str, ...]

The image’s own interpreter; the host’s path does not exist inside it.

ContainerRunner.command

def command(
    program: Program,
    limits: SandboxLimits,
    directory: str,
    name: str,
    cidfile: str,
) -> list[str]

The runtime argv that runs program from directory in a container called name.

The runtime writes the container’s id to cidfile once it creates one.

DPOObjective

class source

class DPOObjective(model, seq_len: int, beta: float = 0.1, **kwargs={})

Train a policy on preference pairs with the DPO loss (arXiv:2305.18290, eq. 7).

The reference is frozen.

beta is the KL strength; model and seq_len are the LMObjective’s, with seq_len one below the row width. The reference never moves, so an ema_decay argument is refused, and loss_role is refused with it: the completion mask already says which targets count. Validation scores the chosen responses’ perplexity under the policy.

DPOObjective.loss

def loss(params, batch, step)

Score the preference term over each pair’s completion tokens.

DPOObjective.evaluate

def evaluate(params, batch, step)

Score the chosen responses’ perplexity under the policy.

The per-token cross entropies carry the shifted completion mask as weights.

Environment

class source

class Environment(Protocol)

The caller owns execution and releases resources on context exit.

Environment.reset

def reset() -> Observation

Environment.step

def step(action: Action) -> Observation

EnvironmentFactory

attribute source

EnvironmentFactory = Callable[[EpisodeId], AbstractContextManager[Environment]]

EnvironmentSource

class source

class EnvironmentSource(
    server: RolloutServer,
    environment: Callable[[Task, EpisodeId], AbstractContextManager[Environment]],
    verifier: Callable[[Task, Episode], float],
    *,
    max_prompt_tokens: int,
    max_new_tokens: int,
    max_turns: int,
    workers: int = 64,
    seed: int = 0,
)

Run Environment sessions against a RolloutServer, one worker thread each.

environment(task, identity) enters one environment per session and verifier(task, episode) scores completed and truncated episodes; both read the task’s own payload, and identity.task is the submission’s serial number. workers bounds concurrent sessions; the server batches their calls. cancel stops a session at its next turn or mid-draw, never inside environment.step: an environment must bound its own step time, or a hung step holds its worker until it returns. Environments see the draw’s raw likelihoods when the server reports them, and its behavior likelihoods otherwise, which are the same distribution only for a sampling policy without transforms.

EnvironmentSource.submit

def submit(task: Task, samples: int, *, version: int) -> list[Future[Session]]

EnvironmentSource.cancel

def cancel(futures: Sequence[Future[Session]]) -> None

Stop sessions at their next turn; queued ones never start.

EnvironmentSource.close

def close() -> None

Cancel every session and wait for the running ones to release their environments.

Episode

dataclass source

class Episode(
    identity: EpisodeId,
    policy_step: int,
    initial: Observation | None,
    transitions: tuple[Transition, ...],
    status: EpisodeStatus,
    detail: str = '',
    reward: float | None = None,
    *,
    _binding_id: str = '',
)

EpisodeCancelled

class source

class EpisodeCancelled(episode: Episode)

EpisodeFailure

class source

class EpisodeFailure(episode: Episode)

Report a failed collection, keeping the partial episode available to the caller.

EpisodeId

dataclass source

class EpisodeId(task: int, attempt: int, sample: int, seed: tuple[int, ...])

Identify one sample within an attempt, reproducibly from checkpointed work.

The harness can use this identity for its own idempotency records. Dew does not guarantee exactly-once external effects across process failure.

EpisodeInference

class source

class EpisodeInference(Protocol)

Draw actions from a bound policy, returning the actual sampling likelihoods.

EpisodeInference.bind

def bind(variables: Variables, /) -> EpisodeInference

EpisodeJournal

dataclass source

class EpisodeJournal(directory: str)

Persist completed turns before the next tool call or trainer update.

The directory is dedicated to one run. Each rank owns one SQLite file, protected against concurrent writers by flock. WAL commits use FULL synchronization. Recovery requires the same cohort layout, controls, policy shards and key. Environment snapshots must include all state needed to continue; external effects in a pending call need idempotency from the environment. A completed, committed turn is never executed again.

EpisodeJournal.open

def open(cohort: str, signature: str, binding: str) -> Iterator[JournalRun]

Open this rank’s journal file and yield the run for one cohort.

The file is locked for the caller alone. A cohort that is already recorded has to present the same signature, and its stored binding wins over the caller’s.

EpisodeRecorder

attribute source

EpisodeRecorder = Callable[[Episode], None]

EpisodeRollout

dataclass source

class EpisodeRollout(
    policy: EpisodeInference,
    environment: EnvironmentFactory,
    verifier: Verifier,
    max_prompt_tokens: int,
    max_new_tokens: int,
    max_turns: int,
    sampling: Sampling,
    groups: int = 2,
    record: EpisodeRecorder | None = None,
    journal: EpisodeJournal | None = None,
)

Collect complete episode groups under one policy snapshot, then train on their actions.

Input batches contain integer task_id rows. The environment factory resolves each task and owns its tools, timeouts and isolation. A finite verifier reward is required for completed and truncated episodes. Errors and cancellation abort the whole group before a Trainer update; record receives the partial episode before the exception propagates.

Episodes train through sessions.pack: a call whose context extends the previous call’s context and actions merges into its chain, and chains share rows of max_prompt_tokens + max_new_tokens ids, so set the objective’s seq_len one below that. The batch keeps one row per possible call, which always fits and keeps shapes fixed. The terminal group advantage is shared by the episode’s actions; no per-turn credit rule is inferred, and truncated episodes are masked. Host records and numeric rows carry the trainer’s committed update clock. Raw and behavior likelihoods come from actual draws.

The policy binds one immutable variables snapshot for the whole collection. It must use that binding, not a mutable serving default. The Trainer cannot update or donate the tree until this call returns. EpisodeJournal adds durable turn boundaries for environments exposing get_state/set_state. Pending external effects require environment-owned idempotency; Trainer checkpoints remain the optimizer’s recovery boundary.

EpisodeRollout.collect

def collect(state: TrainState, batch: Batch, key: jax.Array) -> tuple[Episode, ...]

Collect fixed cohorts, agreeing host phases before every generation.

EpisodeRollout.project

def project(episodes: Sequence[Episode]) -> dict[str, np.ndarray]

Pack one collection’s episodes into GRPO rows through sessions.pack.

Each episode becomes a Session (session_of), so its calls merge into one chain wherever the environment’s next context extends the previous one, and the chains share [rows, width] rows with segment ids. old_log_probs carries the sampler’s raw likelihoods, recorded under the same snapshot. A truncated episode is masked, not scored.

Every episode and action must retain that collection’s private binding origin. Equal training clocks do not establish equal weight snapshots.

EpisodeStatus

class source

class EpisodeStatus(IntEnum)

FlowGRPOObjective

class source

class FlowGRPOObjective(
    model: nn.Module,
    process: Process,
    inputs: InputSpec,
    *,
    sde: FlowSDE = FlowSDE(),
    beta: float = 0.0,
    clip_range: float = 0.0001,
    adv_clip_max: float = 5.0,
    autoencoder: AutoEncoder | None = None,
    guidance: CFG | None = CFG(3.0),
    sampler: Solver = Euler(),
    steps: int = 41,
    pretrained: Variables | None = None,
)

Train a rectified-flow policy on clipped, coordinate-normalized gradients.

A conditional transition KL regularizes it.

Batches carry latents/next_latents [N, K, …], timesteps/next_timesteps, joint old_log_probs and transition_mask [N, K], and advantages [N] or [N, K]. K is the selected transition count. The denominator counts kept stochastic transitions. Deterministic intervals contribute no policy loss.

beta > 0 freezes the initial denoiser in the existing EMA slot. Evaluation and previews always use the live policy. sampler and steps configure evaluation; sde specifies both rollout and rescoring. pretrained is a model variables dict, as returned by model.init; encoders and an optional autoencoder are supplied through the existing diffusion input contract.

FlowGRPOObjective.held_variables

def held_variables() -> Variables

FlowGRPOObjective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

FlowGRPOObjective.log_probs

def log_probs(params: Variables, batch: Batch) -> jax.Array

Rescore joint transition log densities with the rollout’s guidance.

FlowGRPOObjective.loss

def loss(params: Variables, batch: Batch, step: Step) -> tuple[Mean, Aux]

Score the clipped policy gradient over the recorded transitions.

The scan carries nothing between transitions; each one contributes its surrogate, its KL to the frozen reference, whether it counted, and whether the ratio was clipped.

FlowGRPOObjective.evaluate

def evaluate(params: Variables, batch: Batch, step: Step)

Generate one live-policy sample per source row, including prompt-only batches.

FlowGRPOObjective.preview

def preview(params: Variables, batch: Batch, step: Step, *, scored=None)

Draw on all ranks; materialize before root-only caption decoding.

FlowReward

attribute source

Score decoded [-1, 1] samples and repeated source rows, one scalar per sample.

FlowRollout

dataclass source

class FlowRollout(
    objective: FlowGRPOObjective,
    reward: FlowReward,
    groups: int = 4,
    steps: int = 11,
    train_steps: int | None = None,
)

Collect complete prompt groups and the first train_steps transitions.

steps counts time points, so steps=11 draws ten transitions. None selects all transitions. Rewards use population group standard deviation and epsilon 1e-4, as Flow-GRPO’s PerPromptStatTracker does. Each prompt row defines a group; equal prompt text in other rows does not merge groups. Zero-advantage rows are masked, as the reference training loop filters them. Callback collection, JSON/byte transport, and population statistics retain float64 values. Host rewards remain float64; training advantages are float32 after normalization. The reward metric is a float32 diagnostic, and JAX device transfer also narrows the reward column when x64 is off.

The trainer supplies global arrays on every process. Generation remains collective, rewards run once on rank zero, and the result contains only this process’s owned rows for the trainer’s shard_batch boundary. Host materialization currently uses collective_host and replicates the complete trajectory on every process before selecting local rows.

GRPOObjective

class source

class GRPOObjective(
    model,
    seq_len: int,
    beta: float = 0.0,
    epsilon_low: float = 0.2,
    epsilon_high: float = 0.2,
    dual_clip: float = 3.0,
    *,
    policy_loss: str = 'ppo',
    aggregation: str = 'token-mean',
    behavior_importance: float | tuple[float, float] | None = None,
    sequence_mask: tuple[float, float] | None = None,
    geometric_mask: tuple[float, float] | None = None,
    sampling_temperature: float = 1.0,
    **kwargs={},
)

Train a policy on sampled rollouts with the GRPO loss (arXiv:2402.03300, eq. 4).

The default composition is verl’s: the dual-clipped surrogate, token-meaned over the response mask, plus beta times the token-mean k3 KL to the frozen reference (verl/trainer/ppo/core_algos.py, compute_policy_loss_vanilla with token-mean and kl_penalty_forward with k3).

beta is the KL strength; 0.0 allocates no frozen reference. epsilon_low, epsilon_high and dual_clip are the clip points; model and seq_len are the LMObjective’s, with seq_len one below the row width. An ema_decay argument is refused, and loss_role is refused with it: the response mask already says which targets count.

policy_loss picks the surrogate, each checked against verl 12ebe0c: "ppo" (compute_policy_loss_vanilla), "gspo" (compute_policy_loss_gspo: the sequence ratio of sequence_log_ratio, clipped, no dual clip) and "cispo" (compute_policy_loss_cispo). A sequence is a packed chain.

aggregation is "token-mean" (verl’s default) or "session-mean": each session’s token mean, averaged over sessions, so a long session or one split over several rows weighs as one (Agent Lightning’s per_rollout_mean, verl’s seq-mean-token-mean when a session is one row). The batch carries the weights in session_weights.

Behavior corrections read behavior_log_probs against the proximal policy (old_log_probs), all detached, from verl’s rollout_corr_helper. behavior_importance is one threshold, as verl’s rollout_is_threshold is: a number caps the token ratio (TIS), a (low, high) pair zeroes it outside the band instead (IcePop); sequence_mask and geometric_mask reject every token of a sequence whose summed (seq_sum_k1) or mean (seq_mean_k1) k1 statistic lies outside (log low, log high). Metrics add mismatch/kl, mismatch/k3_kl and mismatch/ess whenever behavior likelihoods are present, and the fraction of trainable tokens each correction masked. Without old_log_probs the corrections follow verl’s bypass mode: they compare the detached current policy with behavior, the band only masks, and a TIS cap is refused, since the ratio already is current over behavior.

Engine records on a packed batch are replayed when present: routed_experts/routed make every router use the experts the engine used (R3), and support_ids/support_columns renormalize each sampled id over the ids its top-k/top-p sampler kept, at sampling_temperature (applied after any final softcap), so the policy likelihood compares with a filtered behavior likelihood (DeepSeek-V3.2 section 3.1). sampling_temperature is the engine’s when its reported likelihoods are processed ones; raw ones need 1.0.

GRPOObjective.packed_log_probs

def packed_log_probs(params: Variables, batch) -> jax.Array

Score each packed id given its own chain’s prefix, [rows, width].

Entry t is log pi(input_ids[t] | chain prefix), aligned with input_ids; it is zero where response_mask is zero, which covers every chain start and all padding. The loss and a proximal rescoring read this one function.

GRPOObjective.loss

def loss(params, batch, step)

Score the policy surrogate over the trainable tokens, plus the KL to the reference.

The policy is rescored from the rollout’s own ids, so every term reads the tokens that were actually drawn.

GRPOObjective.evaluate

def evaluate(params, batch, step)

Score the prompts’ perplexity under the policy.

Each row is its shifted cross entropy with the real suffix as weights, taken off the row’s prompt_length. Pads predict nothing and count nothing.

MathReward

dataclass source

class MathReward(require_boxed: bool = True)

Score one when the final answer equals the reference as a rational number.

The answer is the last \boxed{} in the completion; with require_boxed=False a completion without one falls back to its last number. Integers, decimals, a/b and \frac{a}{b} compare exactly, so 0.5, 1/2 and \frac{1}{2} agree. A reference that is not a number compares as trimmed text.

Observation

dataclass source

class Observation(
    context: tuple[int, ...],
    status: EpisodeStatus = EpisodeStatus.RUNNING,
    detail: str = '',
)

Carry the exact next model context, or a terminal environment result.

The harness owns chat formatting, tool-call parsing, and context compaction. Terminal contexts may be empty. Detail can hold a verifier result or an artifact reference without encoding it into device arrays.

Outcome

dataclass source

class Outcome(
    verdict: Verdict,
    exit_code: int | None,
    stdout: str,
    stderr: str,
    seconds: float,
)

PPOObjective

class source

class PPOObjective(
    model,
    seq_len: int,
    *,
    critic: nn.Module,
    value_coefficient: float = 0.5,
    value_clip: float = 0.2,
    **policy_options={},
)

Train a policy and a critic together on one token mass.

The params collection holds policy and critic subtrees, both optimized by the ordinary Trainer. The unit-decay reference selects only policy leaves. Rollout targets are detached. beta and policy clip controls are GRPO’s existing composition; value_coefficient weights verl’s half-squared, clipped value error. A critic consumes packed token rows with their segment_ids and positions and returns [B, T] values; ValueHead supplies that interface for a decoder.

PPOObjective.held_variables

def held_variables() -> Variables | None

Return whatever the actor starts from: a loaded policy checkpoint.

The critic is drawn from the key, so the actor’s tree is the only held data here, and it reaches the trainer’s state JIT as the initializer’s argument rather than as a captured constant.

PPOObjective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

PPOObjective.policy

def policy(variables: Variables) -> EpisodeInference

Bind the policy subtree when an episode collector supplies the full tree.

PPOObjective.pipeline

def pipeline(
    state: TrainState,
    *,
    ema: bool = True,
    processor: Processor | None = None,
) -> TextGeneration

Publish the trained actor, without the critic or the frozen KL reference.

PPOObjective.values

def values(variables: Variables, batch: Mapping[str, object]) -> jax.Array

Score the state before each packed id, [rows, width] aligned with input_ids.

Entry t is the critic’s value of the chain prefix that predicts id t, the state its action was taken from; chain starts and padding, which no action follows, are zero.

PPOObjective.loss

def loss(params: Variables, batch, step: Step) -> tuple[Mean, Aux[Variables]]

Add the actor’s policy loss to the clipped value error on the same mass.

PPOObjective.evaluate

def evaluate(params: Variables, batch, step: Step)

PPOObjective.preview

def preview(params: Variables, batch, step: Step, *, scored=None)

PPORollout

dataclass source

class PPORollout(
    objective: PPOObjective,
    episodes: EpisodeRollout,
    gamma: float = 1.0,
    lam: float = 0.95,
)

Collect episodes, then add critic baselines and verl’s masked GAE.

GAE continues across the action tokens of all turns in one episode, wherever the packer placed them. Tool observations and padding have no support. The terminal verifier reward lands on the last action with zero tail bootstrap, matching the pinned verl GAE input convention; truncated episodes are masked by the packer and take no targets; a cohort with no trainable token at all returns zero targets and zero mass, while a single trainable token, whose whitening is undefined, is refused.

ProcessRunner

dataclass source

class ProcessRunner(python: tuple[str, ...] = (sys.executable, '-I', '-S'))

Run each program as a resource-limited Linux process in a temporary directory.

The whole process group is killed at the wall deadline, on excess output, and after a normal exit, so no child it forked outlives it unless it left the group on purpose. python is this interpreter, isolated from the environment, site-packages and user site.

Program

dataclass source

class Program(files: Mapping[str, str], command: tuple[str, ...], stdin: str = '')

Files to write, the argv to run beside them, and its stdin.

PromptSource

class source

class PromptSource(
    server: RolloutServer,
    reward: Reward,
    *,
    decode: Callable[[Sequence[int]], str],
    max_new_tokens: int,
    scorers: int = 16,
    seed: int = 0,
)

Draw one completion per sample of a prompt_tasks task and score its decoded text.

reward scores the completion with EOS excluded, on scorers threads. A draw that ends on its token budget is TRUNCATED and still scored; a failed draw or reward is an infrastructure failure. Anything else that fails resolves the rollout’s future with the exception, so no future is left pending.

PromptSource.submit

def submit(task: Task, samples: int, *, version: int) -> list[Future[Session]]

PromptSource.cancel

def cancel(futures: Sequence[Future[Session]]) -> None

Forget the rollouts; their draws finish on the server and are not scored.

PromptSource.close

def close() -> None

Stop the reward threads; the server belongs to the caller.

RecoverableEnvironment

class source

class RecoverableEnvironment(Environment, Protocol)

Restore tool state from an opaque snapshot, without replaying completed calls.

RecoverableEnvironment.get_state

def get_state() -> bytes

RecoverableEnvironment.set_state

def set_state(state: bytes) -> None

Reward

attribute source

Score (data_source, completion, ground_truth, extra_info).

Runner

class source

class Runner(Protocol)

Run one program to its end under limits; raise only when it cannot start.

python is the argv that runs a Python file where this runner runs programs.

SampledRollout

dataclass source

class SampledRollout(
    objective: LMObjective,
    reward: Reward,
    decode: Callable[[Sequence[int]], str] = lambda ids: ' '.join(str(token) for token in ids),
    groups: int = 4,
    max_new_tokens: int = 32,
    estimator: str = 'group',
    truncation: str = 'score',
    sampling: Sampling = Sampling(),
)

Draw G completions per prompt and pack them as one-call sessions.

EOS is a valid action in the response mask but excluded from reward text. The batch is pack’s layout, prompts * groups rows of the prompt width plus the response budget. Every completion is scored; one that ran out of max_new_tokens is TRUNCATED, and truncation (default score, train it on its reward) decides whether it trains. old_log_probs holds the raw likelihoods the cached model recorded at each sampled action, and behavior_log_probs the sampling ones.

SandboxFleet

class source

class SandboxFleet(
    runner: Runner = ProcessRunner(),
    *,
    limits: SandboxLimits = SandboxLimits(),
    workers: int = os.cpu_count() or 1,
)

Run programs on workers concurrent sandboxed workers.

Use it as a context manager, or call close. Programs queue when every worker is busy; submit returns at once.

SandboxFleet.submit

def submit(program: Program) -> Future[Outcome]

SandboxFleet.run

def run(programs: Iterable[Program]) -> list[Outcome]

Run every program, concurrently, and return their outcomes in order.

SandboxFleet.close

def close() -> None

SandboxLimits

dataclass source

class SandboxLimits(
    wall_seconds: float = 30.0,
    cpu_seconds: int = 10,
    memory_bytes: int = 256 * 1024 ** 2,
    message_bytes: int = 1024 ** 2,
)

Bound one worker by RLIMIT_CPU and RLIMIT_AS, plus parent-enforced session and IO limits.

Forked children inherit the resource limits. Process-group cleanup handles descendants that stay in that group; this is not a cgroup aggregate limit.

Session

dataclass source

class Session(
    task: str,
    group: str,
    sample: int,
    attempt: int,
    calls: tuple[Call, ...],
    status: Status,
    reward: float | None,
    components: Mapping[str, float] = (lambda: MappingProxyType({}))(),
    detail: str = '',
)

One harness session: its calls in submission order and how it ended.

(task, group) names the advantage group; sample and attempt tell members and retries apart. reward is the verifier’s score, required when the status is trainable. components holds verifier sub-scores for logging and detail the verifier or failure provenance.

SessionSource

class source

class SessionSource(Protocol)

Anything that turns tasks into sessions.

submit starts samples sessions of one task under the served policy version and returns one future per session; cancel stops sessions whose results are no longer wanted.

SessionSource.submit

def submit(task: Task, samples: int, *, version: int) -> Sequence[Future[Session]]

SessionSource.cancel

def cancel(futures: Sequence[Future[Session]]) -> None

Status

class source

class Status(Enum)

How a session ended, which decides whether it trains.

trainable: bool

Whether a session with this status is scored and carries loss mass.

SubprocessEnvironment

dataclass source

class SubprocessEnvironment(
    command: tuple[str, ...],
    limits: SandboxLimits = SandboxLimits(),
)

A user-selected JSON-lines worker implementing reset and step.

command is an argv tuple, executed without a shell in a temporary working directory. Each request has an operation and either an episode identity or an action record. Replies contain context (integer ids), status (running/completed/truncated/cancelled/error) and optional string detail.

Session exit kills the process group on success, error or cancellation. The direct worker also receives SIGKILL if its parent dies. Neither mechanism replaces filesystem/network isolation or controls descendants that deliberately leave the group.

Task

dataclass source

class Task(id: str, data: Mapping[str, object] = (lambda: MappingProxyType({}))())

One unit of work a session source runs: an identity and its source-specific payload.

Transition

dataclass source

class Transition(action: Action, observation: Observation)

ValueHead

Flax module source

class ValueHead()

Project a decoder’s hidden states to one float32 value per position.

Packed rows pass their chains’ segment_ids and positions, so no state reads another chain.

Verdict

class source

class Verdict(Enum)

How a program ended.

COMPLETED

Exited with status zero.

FAILED

Exited with a nonzero status, an uncaught exception or MemoryError included.

TIMEOUT

Ran past the wall deadline or its CPU-time limit and was killed.

CRASHED

Killed by a signal it did not ask for.

OUTPUT_LIMIT

Wrote more than message_bytes to stdout and stderr together and was killed.

Verifier

attribute source

Verifier = Callable[[Episode], float]

code_block

function source

def code_block(completion: str, language: str = 'python') -> str | None

The last fenced block tagged language, else the last untagged one, else None.

pack

function source

def pack(
    sessions: Sequence[Session],
    width: int,
    *,
    rows: int | None = None,
    estimator: str = 'group',
    truncation: str = 'mask',
    support_capacity: int | None = None,
) -> dict[str, np.ndarray]

Strictly merge each trained session’s calls, then pack the chains into [rows, width].

Every array is [rows, width] and aligned with input_ids: entry t describes id t. response_mask is 1 on sampled ids alone; behavior_log_probs, versions and call_index are set on them; advantages repeats the session’s advantage over its chain tokens; session_weights is described at SESSION_WEIGHTS_KEY. Chains are placed first-fit in decreasing length, a stable order, and rows pads the batch to a fixed count, refusing chains that need more (rows_needed over chain_lengths counts them). truncation decides whether TRUNCATED sessions train, as the module docstring describes.

Calls that recorded routed_experts or support add the arrays _engine_records describes; support_capacity, required when any call recorded a support, fixes the per-row length of the support arrays so every batch has one shape and the step compiles once.

prompt_tasks

function source

def prompt_tasks(batch: Batch) -> list[Task]

One single-turn task per prompt row: its ids and the three reward strings.

The id is a digest of the prompt ids; groups, not ids, keep repeats apart.

session_of

function source

def session_of(episode: Episode, *, group: str) -> Session

Read an episode as an engine-style Session of the advantage group group.

Each transition’s action is one call: its context is the prompt, its tokens the sampled ids with their behavior likelihoods, stop when it ended on EOS and length otherwise. An environment-reported error is an infrastructure failure; how well the agent did is the verifier’s reward.