Algorithm Recipes

1 recipes.py — TrainingConfig

All algorithms live in TRAINING_CONFIGS (rl_engine/trainer/recipes.py) as frozen TrainingConfig dataclasses (rl_engine/trainer/config.py). Three dataclass types cover the entire authoring surface:

TrainingConfig.__post_init__ validates at import time: unique role names, sync endpoints resolve, non-empty keys, pool-sharing acyclic — a malformed recipe fails at import, not mid-run.

The @recipe decorator

# rl_engine/trainer/recipes.py
TRAINING_CONFIGS: dict[str, TrainingConfig] = {}

def recipe(name: str):
    """Decorator to register a training config."""
    def decorator(fn):
        TRAINING_CONFIGS[name] = fn()   # build + register at import
        return fn
    return decorator

factory.create_trainer() (rl_engine/trainer/factory.py) builds a RoleWorkers per role, connects exchange syncs over NCCL, and hands the graph to StacxTrainer.

2 Recipe catalog

22 recipes ship in recipes.py — 16 LLM plus 6 experimental diffusion.

LLM recipes

Diffusion recipes experimental

Flow-matching kernels for DiT/MMDiT models; loss math in rl_engine/trainer/objectives/diffusion_loss.py.

Recipe selection — infer_algorithm()

The runtime key args.algorithm (rl_engine/trainer/factory.py) resolves as:

• --algorithm set explicitly → wins verbatim (DAgger / distillation / diffusion always use this).
• use_critic set → ppo.
• Otherwise --advantage-estimator maps to the matching recipe (e.g. reinforce_plus_plus → reinforce_pp).

3 Algorithm objectives

Stateless loss functions in rl_engine/trainer/objectives/; specializations in rl_engine/trainer/specialization/ compose them inside compute_loss().

GRPO — PPO-clipped PG on group-normalized rewards

ActorSpecialization.compute_loss (specialization/agent.py). compute_grpo_advantages() (objectives/advantage.py) broadcasts the scalar reward to every response token and whitens across the group. Optional KL anchor: --use-kl-loss --kl-loss-type low_var_kl --kl-loss-coef 0.001 + --ref-load (in-Megatron reference, no separate worker). gspo and reinforce_pp[_baseline] reuse this specialization with a different estimator.

SFT & iter-SFT — supervised cross-entropy

CausalDecoderSpecialization.compute_loss (specialization/causal_decoder.py): per-token CE on the response, masked to loss_masks. The iter-SFT baseline uses --algorithm online_dagger with STACX_DAGGER_SFT_MODE=trajectory to replay a pre-saved teacher corpus instead of live rollout.

On-policy distillation (OPD) — K3 reverse-KL vs. a frozen teacher

DistillSpecialization (specialization/distill.py): reverse KL against a frozen teacher SGLang server, whose per-token logprobs arrive as rollout fields. Mode chosen by STACX_DISTILL_KL_MODE:

• k3_sample (default) — single-sample MC estimate on the chosen token; TP-aware, higher variance.
• topk (TP=1 only) — RKL over the union of student and teacher top-K supports; lower variance.

Mixed-loss DAgger-OPD — per-token CE / K3 dispatch

MixedLossDaggerOPDSpecialization (specialization/mixed_loss_dagger.py) routes each action token by executor — teacher-executed → CE, student-executed → K3 reverse-KL — via a sign sentinel packed into teacher_log_probs (SGLang logprobs are ≤ 0; rollout writes +1.0 at teacher positions). WandB tracks ce_mean, k3_mean, teacher_action_frac.

Both endpoints asserted in tests/trainer/test_mixed_loss_dagger.py.

Rollin protocols — β vs. κ

Both recipes share the loss kernel but differ in how the teacher/student split is produced (STACX_DAGGER_SAMPLING_MODE):

• mixture_dagger_opd — per-turn β-mixture: each turn flips Bernoulli(β), β decaying over training (linear / cosine / exponential / step; launcher default linear 1.0→0.0 over 5 iters).
• mixture_aggrevate_opd — trajectory-level AggreVaTe: κ ~ U{0, T_max(i)} per task; student leads turns 0..κ−1, teacher force-decodes the tail. T_max(i) = min(start + step·i, cap), launcher default 0/20/100.

DPO — contrastive preference loss

4 Data syncs: inject vs. exchange

Single-role recipes declare syncs=(); the external SGLang teacher server needs no sync entry.

5 Adding an algorithm

A new algorithm is a specialization (only if no existing loss fits) plus one @recipe. No backend, Ray, or orchestrator changes needed.

Step 1 — Subclass a specialization

# rl_engine/trainer/specialization/my_algo.py
from collections.abc import Callable
from typing import Any
import torch
from rl_engine.trainer.specialization.causal_decoder import CausalDecoderSpecialization
from rl_engine.trainer.objectives.policy_loss import sum_of_sample_mean

class MyAlgoSpecialization(CausalDecoderSpecialization):
    def compute_loss(
        self,
        model_output: torch.Tensor,          # [seq_len, vocab] (TP-local shard if TP>1)
        batch: dict[str, Any],
        reducer: Callable[[torch.Tensor], torch.Tensor],
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        response_lengths = batch["response_lengths"]
        loss_masks = batch["loss_masks"]
        # ... derive per-token student logprobs from model_output, then:
        per_token = my_kernel(...)                       # shape [N_response]
        loss = sum_of_sample_mean(per_token, response_lengths, loss_masks)
        loss = reducer(loss)                             # slime CP-aware reduction
        return loss, {"my_loss": loss.detach()}

Step 2 — Register the recipe

# rl_engine/trainer/recipes.py
@recipe("my_algo")
def my_algo_config():
    from rl_engine.trainer.specialization.my_algo import MyAlgoSpecialization
    return TrainingConfig(
        name="my_algo",
        roles=(
            ModelRole(
                name="policy",
                specialization_cls=MyAlgoSpecialization,
                sync_to_rollout=True,          # push weights to SGLang each step
            ),
        ),
        syncs=(),                                # single role → no cross-model sync
        description="My custom reverse-KL algorithm",
    )

Step 3 — Launch it

python -m rl_engine.train --algorithm my_algo \
    --hf-checkpoint /root/models/<student> \
    --ref-load /root/models/<student>_torch_dist \
    ...   # the SWE launchers under scripts/train/swe/ wrap this

6 Objectives reference

7 Algorithm knobs

Algorithm-level flags and env vars. Full launcher catalog: scripts/train/swe/ENV_REFERENCE.md.