XYZ-AI-Lab/axrl

解决的问题

AxisRL 专为代理型强化学习（RL）后训练设计，模型需在多轮交互中与环境进行交互、调用工具，并根据长期结果获得奖励。它解决了管理多轮回放、环境状态、工具调用和权重同步的协调挑战，同时确保训练过程稳定且可观察。

工作原理

AxisRL 作为一个系统层，连接高吞吐量回放引擎（SGLang）与大规模分布式训练引擎（Megatron）。它管理一个循环：回放代理执行代理工作流，SGLang 工作者生成文本，环境产生奖励，Megatron 工作者执行策略优化（如 PPO 或 GRPO）。

关键技术机制包括：

• 白盒与黑盒集成：支持内部代理循环（白盒）或通过 OpenAI 兼容代理捕获外部框架的交互（黑盒）。
• 部分回放：允许已完成的样本立即发送给训练器，以减少因长尾轨迹延迟导致的空闲时间。
• 基于句柄的数据平面：通过句柄而非中央驱动程序传输重型数据（如 MoE 路由数据），以避免瓶颈。
• 上下文管理：使用前缀树合并和 MagiAttention 优化多轮上下文中的重复注意力计算。
• 一致性工具：包含不匹配分析和「脉冲重播」，用于调试回放与训练器执行路径之间的偏差。

适用人群

适用于训练大规模代理型 LLM（数百亿参数）的研究人员和工程师，这些模型需要复杂的多轮交互和高吞吐量 RL 后训练。

主要亮点

• 基于 SGLang 实现回放，基于 Megatron 实现分布式训练。
• 支持多种优化目标，包括 PPO、GRPO/GRPO2、GSPO、TOPR 和 TIS。
• 可处理超过 300 轮的轨迹。
• 提供 Rollout Routing Replay（R3）以在训练期间稳定 MoE 专家路由。
• 提供可复现的调试套件，用于分析损失突增和 token 级别不匹配。