UniRL 文档
指南

扩展 UniRL

新增模型、rollout engine、train-side algorithm、reward、training backend 和 recipe 的位置。

扩展时优先沿用已有边界:一个实现包、一个 config dataclass、一个 recipe,以及一个 focused check。

扩展位置

目标主要位置recipe 如何连接
新增模型unirl/models/<model_name>/@dataclass 定义 config,在 recipe 的 model: 块用 _target_ 引用
新增 rollout engineunirl/rollout/engine/<engine>/@dataclass 定义 config,在 rollout: 下用 _target_ 引用
新增 train-side algorithmunirl/algorithms/一个 BaseAlgorithmConfig 子类,用 _target_ 引用并绑定到某 track 的 algorithm: 节点
新增 reward scorerunirl/reward/local/一个 spec @dataclass,在 reward.backend.config 下用 _target_ 引用
新增 training backendunirl/train/backend/Remote backend 契约(与 FSDPBackend 并列)
新增 recipeexamples/<domain>/<recipe>.yaml--config-name=<domain>/<recipe>

新增模型

  1. 添加 unirl/models/<model_name>/
  2. 在旁边用 @dataclass 定义模型的 config(recipe 用 _target_ 引用)。
  3. 实现 bundle,暴露训练和 rollout 需要的 stage。
  4. 按需添加 condition、text/vision、diffusion、VAE helper。
  5. 添加至少一个 examples/<domain>/<recipe>.yaml recipe。

新增 Rollout Engine

  1. unirl/rollout/engine/<engine>/ 下添加 typed config。
  2. 在 recipe 的 rollout: 块用 _target_ 引用它。
  3. 实现 unirl/rollout/engine/base.py 中的 contract。
  4. 返回 canonical RolloutResp(填充 tracks[name])。
  5. 如果是 dedicated engine,还要定义 trainer-to-rollout weight sync。

新增 Train-Side Algorithm

训练 loss 是每 track 一个 StageAlgorithm。只有当 loss 数学改变时才新增类(DanceGRPO / MixGRPO 这类 recipe 组合复用 DiffusionGRPO)。

  1. 继承 StageAlgorithmunirl/algorithms/base.py)。
  2. 在旁边定义继承 BaseAlgorithmConfig 的 config(普通 @dataclass)。
  3. 实现 compute_loss_and_backward(...)——replay stage、算 loss、backward()
  4. 设置 supports_multi_update(loss 是否在一个 rollout shard 上多次 optimizer step 仍有效),requires_ema_rollout 只给需要 EMA 采样的 off-policy loss(NFT)。
  5. 在 recipe 中把该类绑定到 track 的 algorithm: 节点。

完整契约见生成的 Algorithms Package README

新增 Reward Scorer

unirl/reward/local/ 旁实现 spec 和 scorer。进程内模型 scorer 推荐从 LocalRewardBackend 开始,因为它提供 device 解析、eager load、offload()onload()。用普通 @dataclass 定义 spec,在 recipe 的 reward.backend.config 下用 _target_ 引用它(无需注册步骤)。进程外 scoring 则把 reward backend 指向远程 service。见 Rewards

训练 Backend 与并行

新的训练 backend并行工作(新的 FSDP / SP / EP plan,或 VeOmni 式 backend)面向 unirl/train/backend/ 下的 backend 契约:实现与 FSDPBackend 相同的 Remote 表面,并把 TrainTopology 映射到并行 plan,而不是硬编码 shard size。见 Trainer 与训练栈 与生成的 Train Stack README

Agent 提示

Agent 修改代码前,应先把任务映射到上表,再读对应 package README。

目录