Rollout Engine

1 Position in the System

Everything in rl_engine/rollout/ answers one question each trainer iteration: given the current student weights, produce the trajectories to train on. The entry point is SchedulerRolloutManager — a Ray actor subclassing slime's RolloutManager that routes generation and evaluation through a TaskScheduler.

SchedulerRolloutManager

Ray forbids inheriting from a @ray.remote actor class directly, so rollout_manager.py unwraps the decorator, subclasses the plain base, and re-applies it:

# rl_engine/rollout/rollout_manager.py
from slime.ray.rollout import RolloutManager as _RolloutManagerActor

# Unwrap the @ray.remote decorator to get the plain Python class.
_RolloutManagerBase = _RolloutManagerActor.__ray_metadata__.modified_class

@ray.remote
class SchedulerRolloutManager(_RolloutManagerBase):
    # overrides: __init__, dispose, _get_rollout_data, eval (+ helpers)
    ...

It overrides only what it must, inheriting the rest from slime:

Two construction paths

placement_group.py exposes two factories — the default scheduler-backed path, and a legacy bypass for ablation runs:

Top-level architecture

One rollout step flows left to right — data source to scheduler, a short-lived Evaluator per batch, one episode handler per sample, and back out as list[Box]:

2 Directory Tour

Paths are relative to rl_engine/rollout/.

Top level

scheduler/ — the control plane deep dive: Task Scheduler

agents/ — episode handlers

The in-house SWE / DAgger handlers (SWEAgent, DAggerSWEAgent) live with their lane in rl_engine/examples/swe_agent/agents/, not here — rollout/agents/ holds only the base contract.

backends/ · proxy/ · sandbox/ · utils/

3 The Episode Abstraction

The scheduler knows nothing about SWE-Bench, tool calls, or teachers — all of that lives behind one interface, BaseAgent (agents/base.py). A handler owns the entire episode for one sample, so both black-box scaffolds (OpenHands, Terminus-2) and native token-level agents (DAgger) fit one contract.

The three data types

# agents/base.py — the contract every scaffold implements
class BaseAgent(abc.ABC):
    async def init(self, sample, context)          -> None:  # lightweight, no placement yet
    async def run(self, sample, context)           -> AgentResult:  # the full multi-turn loop
    async def eval(self, sample, context, result)  -> float | dict:  # reward
    async def cleanup(self, sample, context)       -> None:  # finally-block; stop containers

How the Evaluator runs one sample

scheduler/evaluator.py drives the lifecycle in _execute_sample. Three guarantees hold:

• The semaphore spans the whole episode. async with self._semaphore wraps init → run → eval → cleanup, capping live samples at Resource.concurrency — not turns.
• One placement token spans run and eval. Acquired before run, released in the finally after cleanup, so the patch and verification containers never double-count against a node's budget.
• Failures degrade to a masked dummy. Any exception marks the sample ABORTED with 2-token dummy data ([0, 0], remove_sample=True, zero gradient); classify_abort buckets it and _annotate_abort attaches diagnostics (phase, elapsed vs. window, node IP).

After eval, _populate_token_data lifts the token-level fields off AgentResult onto the Sample — what makes a trajectory trainable, whether captured natively or through the recording proxy.

The multi-turn loop

Inside run(), a native handler like DAggerSWEAgent (overriding _agent_loop) alternates policy generation with environment steps, tagging every token with a loss mask.

4 The Data Path

Each prompt must reach the right handler with the right stopping criteria. Three pieces cooperate: MultiDataSource serves prompts, DatasetEntry / EvalDatasetConfig carry per-dataset overrides, and task_type routes each sample to an EvaluatorSpec.

MultiDataSource

multi_data_source.py replaces slime's single-file DataSource so one rollout step can span heterogeneous task types. It serves datasets in FIFO order (dataset A exhausted before B). Per sample it deep-copies the prompt, stamps metadata["task_type"] and metadata["_dataset_cfg"] (the full DatasetEntry), and expands n_samples_per_prompt into a group.

parse_rollout_datasets_config builds the DatasetEntry list from YAML; resolve_dataset_path anchors relative paths at the repo root (or STACX_ROOT).

# --rollout-config YAML — parse_rollout_datasets_config()
rollout:
  defaults:                  # shared across datasets
    task_type: swe
    max_context_len: 65536
  datasets:
    swe_gym_train:
      path: data/swe/swe_gym_train_minus_val100.jsonl
      input_key: prompt
      label_key: instance_id
    retool_math:
      path: data/retool_math.jsonl
      task_type: retool     # overrides the default
      max_context_len: 32768

The 3-level cascade

Generation parameters resolve from coarsest to finest, the per-dataset value winning:

CLI args      →      EvaluatorSpec.default_stopping      →      DatasetEntry / EvalDatasetConfig
(coarsest)                                                                        (finest, wins)

Resolved by the scheduler's _resolve_stopping (for StoppingCriteria) and the sampling-param builder; the Task Scheduler page documents it in full.

Routing: task_type is the evaluator_id

An EvaluatorSpec (scheduler/evaluator.py) is a blueprint: a handler (a BaseAgent), its resource model (resource / run_resource / eval_resource), an optional batch_size, and default_stopping. The scheduler selects one via batch.task_type, so a dataset's task_type must equal some spec's evaluator_id or the batch has nowhere to run.

Specs come from a build_evaluators(args) factory (list[EvaluatorSpec]), loaded via --evaluator-config-path and wired into a TaskScheduler by build_task_scheduler.

# rl_engine/examples/swe_agent/evaluator_config_dagger.py
def build_evaluators(args):
    teacher = TeacherBackend(sglang_url=..., model_path=..., tools=SWE_TOOL_SPECS)
    handler = DAggerSWEAgent(teacher_backend=teacher, max_turns=100, ...)
    return [EvaluatorSpec(
        evaluator_id="swe",                     # == dataset task_type
        handler=handler,
        run_resource=Resource(cpu=4, memory_gb=16),   # agent container
        eval_resource=Resource(cpu=4, memory_gb=16),  # verification container
        default_stopping=StoppingCriteria(max_turns=100, max_context_len=65536),
    )]

5 Trajectory Post-Processing Lanes

Once execute_rollout() returns raw trajectories, _get_rollout_data shapes them for the loss. The order is fixed — filter → route by algorithm → shape — with the lane chosen entirely by args.algorithm.

Filters

Two independent filters run before shaping, on the raw one-per-task trajectories:

They compose: rejection first, finish-only on the survivors. If a filter empties the batch, _get_rollout_data returns one masked dummy so logging never trips on an empty rollout.

The algorithm fork

_maybe_apply_per_token_capture is the switch: per-token-loss algorithms capture teacher logprobs and return immediately (no chunk expansion); everything else falls through to CE-only expansion.

Teacher force-decode capture (per-token lanes)

Both per-token lanes call capture_teacher_logprobs_for_samples (distill_capture.py): kept trajectories are force-decoded through the teacher SGLang server with top_logprobs_num=K, yielding the teacher's chosen-token logprob at each response position. That pins sample.teacher_log_probs and overwrites sample.loss_mask with the action mask, so the loss reduces only over student action tokens.

The mixed CE/K3 lane snapshots the agent's per-token is-teacher mask, then overwrites teacher_log_probs with a +1.0 sentinel at teacher positions. Real SGLang logprobs are always ≤ 0, so the kernel reads the sign — > 0 ⇒ CE, ≤ 0 ⇒ K3 — meaning β=1.0 is pure CE (≡ legacy online-DAgger) and β=0.0 pure K3 (≡ OPD). The loss kernel lives in rl_engine/trainer/specialization/mixed_loss_dagger.py — see Algorithm Recipes.

DAgger trajectory-SFT sub-samples (CE-only lane)

The chunk-expansion lane (_build_trajectory_sft_samples) is the older CE-only path: each trajectory with M student-executed turns becomes up to M+1 plain-CE SFT sub-samples, driven by the _student_cut_points metadata.

A per-trajectory buffer makes the drop atomic — if any sub-sample exceeds max_seq_len the whole trajectory is discarded (semantics pinned by tests/rollout/test_trajectory_sft_samples.py). A deterministic per-iter shuffle (STACX_SHUFFLE_SEED_BASE) de-clusters sibling chunks; optional dynamic GBS scaling (STACX_TASKS_PER_GRAD_STEP) holds tasks-per-step constant as chunks grow.

Termination taxonomy & per-iter summary

termination_summary.py classifies outcomes. classify_termination(reward, termination_reason) maps every trajectory into one of eight classes; write_summary aggregates a pass into a summary.json with per-dataset counts and a per-task list.

Analysis dumps

Three optional, best-effort dumps capture token-level data for offline analysis. All are atomic (tmp + os.replace) and swallow failures so they can never abort training:

• OPD discrepancy dump (STACX_DISTILL_DUMP_DIR) — per-iter .pt with the exact (teacher, student) logprob pair the K3 estimator consumes, plus the action_mask.
• Mixed-loss dump (STACX_MIXED_LOSS_DAGGER_DUMP_DIR) — the sentinel'd teacher_log_probs and the original is-teacher mask, so tools can split CE vs. K3 per position.
• Per-iter student dump (STACX_KL_DUMP_DIR, kl_eval_cache.py) — dump_samples_from_eval writes one .pt per finished eval task (tokens, loss_mask, rollout logprobs, reward) plus a manifest.json; armed by default in the launch scripts.

6 Rollout-Layer Env Knobs

The STACX_* variables read inside rollout_manager.py and scheduler/task_scheduler.py — the rollout-layer slice of the full catalog (with launch-script defaults) in scripts/train/swe/ENV_REFERENCE.md.

Data path & filters

Teacher force-decode capture

Analysis dumps

7 Where To Go Next