Code Architecture

unirl/ is the framework package — the layered, composable core that turns one Hydra recipe into a multi-GPU RL training run. This README is the map of the source tree: what the layers are, how they depend on each other, and which module README to read next. (For what UniRL is and how to launch it, see the top-level README; for how tensors move across GPUs at runtime, see distributed/README.md.)

UniRL is a thin entrypoint that hands a recipe to the trainer (the driver), which places a Worker pool and wires each loop component as a Remote from the recipe's _target_ — not by import. (Deliberately not "stage": in this repo a stage is a trainable unit of a model pipeline — DiffusionStage / ARStage, see models/README.md — never a position in the training loop.)

As source, the package falls into four groups:

• Entrypoints (train_diffusion.py, train_ar.py, train_pe.py, train_unified_model.py, plus train_agentic.py for multi-turn training) — each composes and validates a Hydra recipe, then hands off to its trainer.
• Orchestration (trainer/) — the per-domain <Domain>Trainer owns GPU placement, builds the rollout and train workers, and runs the rollout→reward→advantage→train loop.
• Training loop (rollout/, reward/, algorithms/, train/) — the four pluggable components of one rollout, plus what they share: models/ (per-model bundles), sde/ (step kernels / σ schedule), and data/ (sources).
• Foundation (distributed/, config/, types/) — the cross-cutting runtime and contracts every layer rests on: Ray worker/dispatch/transport, config validation, and shared typed data. utils/ contains small supporting mechanisms; it is not an ownership layer or a default destination for code.

Module Map

Ownership and dependency direction

The source tree is not one rigid linear stack: rollout adapters may translate concrete model wire formats, and train/ consumes algorithm and model contracts. The stable direction is nevertheless simple:

entrypoints / trainer                  compose and orchestrate
          ↓
rollout · reward · algorithms · train  own loop policy
          ↓
models · data · sde                    own domain policy and transformations
          ↓
types · config · distributed           own stable contracts and mechanisms

Imports should point downward. Integration code that must understand both a backend and a model belongs at that integration seam (normally a rollout or train-backend adapter), not in generic infrastructure. In particular:

• concrete model packages do not import trainers, algorithms, rewards, or rollout engines;
• algorithms depend on stage/type contracts, not concrete model families;
• data sources do not depend on trainers, models, or training backends;
• distributed mechanisms do not depend on concrete models;
• trainers may compose everything, but model-specific semantics stay behind a model-owned contract or hook;
• utils/ is only for stable, domain-agnostic leaves. A function used by one model, dataset, backend, or algorithm stays with that owner.

The already-clean directions are locked by lint/check_core_dependencies.py (TYPE_CHECKING-only imports exempt): types/, config/, utils/, sde/, data/, and models/types/ never import the loop tier or concrete models; models/ never imports algorithms, data, rewards, rollout, or trainers; algorithms/ sees model contracts, never concrete model families; train/ sees model contracts and algorithms/base, never concrete algorithms; distributed/ never imports algorithms, data, models, rewards, rollout, the train stack, or trainers.

Directions the tree does not yet satisfy are left unlocked rather than hidden behind allowlists. The open debt, named so it is not copied as precedent: models/ (hunyuan_image3, qwen3_5, qwen3_moe) reaches into train/backend/veomni (EP fusing / _compat), models/cosmos3 subclasses the train-side SFT track builder, and algorithms/ flips train/lora adapters. That is migration work, not a convention to copy.

Where new code goes

Prefer a small amount of policy duplication over a shared base class that makes one model change touch several unrelated modules. Share mechanisms; keep model, algorithm, reward, and backend policy local.

Deployment modes

The rollout engine and optional sync: section define how GPUs are used:

See distributed/README.md for how each mode places workers and moves rollout data between them.

Runtime Data Flow

The layers above turn one recipe into a repeating loop. A single training step flows through them like this:

1. An entrypoint composes the chosen examples/<domain>/<recipe>.yaml and runs validators.
2. The <Domain>Trainer (e.g. trainer/diffusion.py) acquires a Ray DevicePool and builds the rollout and train workers.
3. The trainer builds a typed Sample lineage and dispatches it to the rollout engine.
4. The engine returns the lineage with generated Parts filled with conditions, segments, primitive maps, and media previews.
5. RewardService.score_and_attach attaches rewards; Part.compute_advantages z-scores them into advantages.
6. TrainStack.train_track(...) shards the generated Part across train workers and runs the mini-batch optimizer loop.
7. Each train worker owns a model Bundle, an FSDPBackend, and one loss algorithm.
8. Dedicated-rollout modes (separate / colocate) sync trainer weights back to the rollout workers.

The agentic workflow extends step 3 into trajectories: a driver-side manager dispatches one task per remote engine slot while each engine executes model and environment turns. AgenticTrainer waits for complete sibling groups, scores their terminal answers through RewardService, and assigns each trajectory's group-normalized advantage to all of its generated turns before training.

Deeper Module Docs

• trainer/README.md: the orchestration hub — how a <Domain>Trainer places workers and drives the loop.
• types/README.md: the Sample / Part contract and migration from the retired request/response API.
• config/README.md: flat-recipe config — require/precision validators, _target_ instantiation, cross-component contracts.
• rollout/README.md: rollout modes, engines, and the Sample / Part generation flow.
• rollout/env/README.md: agentic environments, tools, trajectories, and manager quiescence behavior.
• train/readme.md: train stack, FSDP backend, injection, EMA shadow.
• algorithms/README.md: per-track loss algorithms.
• reward/README.md: reward backends and custom scorers.
• models/README.md: model bundle and per-model package contracts.
• sde/README.md: SDE kernels, σ schedule, initial noise.
• distributed/README.md: distributed runtime — workers, dispatch, placement, and the rollout→train data plane.
• distributed/weight_sync/README.md: trainer→rollout weight-sync backends.

Source: unirl/README.md — edit the README next to the code, then run npm run sync:readmes from docs/.