Extending UniRL

Use existing boundaries rather than adding cross-cutting glue. Most extensions need one implementation package, one config dataclass, one recipe, and one focused test or compose check.

Extension Map

Adding a Model

1. Add unirl/models/<model_name>/.
2. Define the model's config @dataclass next to it.
3. Implement a bundle that exposes the trainable and frozen modules needed by training and rollout.
4. Implement a pipeline that consumes that bundle and exposes the trainable stages and sampling path.
5. Wire both bundle: and pipeline: blocks in at least one examples/<path-under-examples>.yaml recipe.
6. Document required external checkpoints through YAML env interpolation or launcher docs.

Adding a Rollout Engine

1. Add a typed config under unirl/rollout/engine/<engine>/.
2. Reference it by _target_ under the recipe's rollout: block.
3. Implement the engine contract from unirl/rollout/engine/base.py.
4. Implement generate(sample: Sample) -> Sample, filling the appropriate generated Part in Sample.parts.
5. If the engine is dedicated, define trainer-to-rollout weight sync.

Adding a Train-Side Algorithm

The training loss for a generated Part is a StageAlgorithm. Add a new class only when the loss math changes (recipe compositions like DanceGRPO / MixGRPO reuse FlowGRPO; autoregressive GRPO uses GRPO).

1. Subclass StageAlgorithm (unirl/algorithms/base.py).
2. Define a config subclassing BaseAlgorithmConfig (a plain @dataclass) next to it.
3. Implement compute_loss_and_backward(...) — replay the stage, compute the loss, call backward().
4. Set supports_multi_update to match whether the loss is valid across multiple optimizer steps on one rollout shard, and requires_ema_rollout only for off-policy losses that need EMA sampling (NFT).
5. Bind the class under the track's algorithm: node in a recipe.

See the generated Algorithms Package README for the full contract.

Adding a Reward Scorer

Add the spec and scorer near the reward implementation under unirl/reward/local/. Prefer LocalRewardBackend for in-process model scorers because it provides device resolution, eager load, offload(), and onload(). Define the spec as a plain @dataclass and reference it by _target_ under reward.backend.config in the recipe. For out-of-process scoring, point the reward backend at the remote service. See Rewards.

Training Backends and Parallelism

New training-backend or parallelism work (a new FSDP / SP / EP plan, or a VeOmni-style backend) targets the backend contract under unirl/train/backend/: implement the same Remote surface as FSDPBackend and map TrainTopology onto a parallel plan rather than hardcoding shard sizes. See Trainer & Training Stack and the generated Train Stack README.

Agent Tip

When an agent is asked to implement a new feature, first map the request to the table above, then read the closest package README before editing code.