Trainer Architecture

Reference for rl_engine/trainer/ — scripts, classes, and options.

1 factory.py — create_trainer()

# rl_engine/train.py — outer loop (simplified)
for rollout_id in range(start_rollout_id, args.num_rollout):
    rollout_data_ref = rollout_manager.generate.remote(rollout_id)   # upstream: Rollout
    for pass_idx in range(num_train_passes):                       # NUM_TRAIN_PASSES_PER_ROLLOUT
        pass_data_ref = rollout_data_ref if pass_idx == 0 \
            else _reshuffle_rollout_data(rollout_data_ref, pass_seed)  # prime-strided seed
        train_metrics = trainer.train_step(rollout_id, pass_data_ref)   # StacxTrainer
    trainer.update_weights()          # push new weights to SGLang (sync_to_rollout roles)
    trainer.save_model(rollout_id)      # checkpoint trainable roles

• Reshuffle seed = base + rollout_id + pass_idx·100003. The global pass default is 1; shipped DAgger/AggreVaTe/OPD scripts set 3.

2 Two-branch design

Algorithm-specific code lives entirely in the role branch. The shared backend branch — TrainBackend, realized by Megatron-LM or FSDP v2 — is identical for every algorithm: distributed setup, forward/backward/optimizer, weight sync, checkpointing, sleep/wake offload.

Trainer architecture

3 config.py — Recipes as data

An algorithm is a frozen TrainingConfig in trainer/config.py. Three dataclasses:

Recipes register in TRAINING_CONFIGS via @recipe in trainer/recipes.py — 22 built-in; the full catalog is on the Algorithm Recipes page.

# trainer/recipes.py
@recipe("grpo")
def grpo_config():
    return TrainingConfig(
        name="grpo",
        roles=(ModelRole(name="policy", specialization_cls=ActorSpecialization,
                         sync_to_rollout=True, init_kwargs={"with_ref": True}),),
        syncs=(),                                   # reward-only advantages, no cross-role data
    )

@recipe("dpo_distributed")
def dpo_distributed_config():
    return TrainingConfig(
        name="dpo_distributed",
        roles=(ModelRole(name="policy",    specialization_cls=DPOSpecialization),
               ModelRole(name="reference", specialization_cls=DPOSpecialization, trainable=False)),
        syncs=(DataSync(source="reference", target="policy", keys=("log_probs",),
                        target_keys=("ref_log_probs",), stage="inject"),),
    )

4 role_workers.py — RoleWorkers & WorkerPool

• Role names use RL terminology (policy, value, reference, reward); the classes behind them are function-named (ActorSpecialization, CriticSpecialization), selected by specialization_cls rather than a role-string switch. build_gpu_allocations (factory.py) maps slime placement-group keys onto role names.

5 specialization/base.py — ModelSpecialization

Specialization class hierarchy

The three-method interface

Optional: compute_advantages(rollout_data) — no-op on base; overridden by the RL-actor and diffusion specializations.

Adding an algorithm

# 1. A specialization — implement the interface (specialization/base.py contract)
class MySpecialization(ModelSpecialization):
    def prepare_batch(self, rollout_data):      return [...]        # micro-batch dicts
    def extract_model_output(self, model, batch): return model(...)   # logits / velocity
    def compute_loss(self, model_output, batch, reducer):
        loss = ...
        return reducer(loss), {"my_loss": loss.detach()}

# 2. Register it as data (trainer/recipes.py)
@recipe("my_algorithm")
def my_algorithm_config():
    return TrainingConfig(name="my_algorithm",
        roles=(ModelRole(name="policy", specialization_cls=MySpecialization, sync_to_rollout=True),))

# 3. Use it
python -m rl_engine.train --algorithm my_algorithm ...

• MixedLossDaggerOPDSpecialization (specialization/mixed_loss_dagger.py) blends per-token CE (teacher-executed tokens) and K3 reverse-KL (student-executed) on a β schedule: β=1 → pure CE, β=0 → pure K3.

6 StacxTrainer step loop

StacxTrainer.train_step(): run inject syncs → dispatch train() to each trainable role, async for independent pools and serialized for shared.

One backend training step

Sync mechanisms

7 Directory reference

See Backends for infra/ internals and Algorithm Recipes for the diffusion objectives.

8 Trainer knobs

Defined in rl_engine/train.py; parallelism/optimizer flags from slime's argument parser via scripts/train/swe/. Full STACX_* catalog: scripts/train/swe/ENV_REFERENCE.md.