扩展框架

对于希望进行更深层次定制的高级用户，本章演示如何通过集成自定义环境和新的模型架构来扩展 RLinf。

你将学习如何：

• 将一个 新环境 集成到 RLinf 的任务系统中

• 添加一个使用 FSDP + HuggingFace 后端的 新模型

• 添加一个使用 Megatron + SGLang 后端的 新模型

• 参考一条完整的 Reward Model 工作流 （仿真与真机）

• 使用 Megatron-Bridge 集成 Megatron-LM 训练与 HuggingFace checkpoint

• 通过 权重同步 优化 actor 到 rollout 的权重传输（patch 与 bucket 模式）

RLinf 支持多种模型训练后端，每种后端都有自己的初始化逻辑和执行流程。 本指南提供了逐步说明，帮助你完成以下任务：

• 在 RLinf 中注册并加载自定义模型

• 配置 YAML 文件以引用你的新模型或环境

• 如果你的模型类型尚未被支持，扩展特定后端的代码

• 调整环境封装器和接口以集成新的模拟器或 API

从其他仓库添加自定义模型

如果你的项目把 RLinf 当作依赖库使用，现在可以直接注册自定义模型， 而不需要修改 RLinf 源码。

推荐做法如下：

• 在你自己的仓库中实现模型构建函数

• 在调用 build_config(...) 或启动训练之前执行 register_model(...)

• 在 YAML 配置中直接使用注册后的 model_type

from rlinf.models import register_model

def build_my_model(cfg, torch_dtype):
    from my_repo.models.custom_policy import CustomPolicy

    return CustomPolicy(cfg, torch_dtype)

register_model("my_custom_model", build_my_model, category="embodied")

注册后，RLinf 会自动：

• 在配置校验阶段接受 model.model_type: my_custom_model

• 在 get_model(cfg) 时路由到你注册的 builder

这也是当前推荐的扩展方式，适合在 RLinf 主仓库之外维护自定义具身模型。

无论你是要训练一种新的模型架构，还是要在自定义 RL 环境中进行实验， 本节都将提供工具，帮助你直接接入 RLinf 的模块化设计。