arXiv:2608.17641cs.LGcs.DC2026-08

用统一框架加速强化学习信用分配,速度提升1.6到5.7倍。

rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment

  • 将七种信用分配算法统一为一个可并行计算的递推结构。
  • 在数千环境并行场景下,相比基线提速1.6至5.7倍。
  • 适合需要高效训练的强化学习研究者和工程团队使用。

我们提出rl-triton,一个基于Triton的开源GPU内核库,用于强化学习信用分配。核心贡献是一个统一的关联扫描框架,将七种不同算法——广义优势估计(GAE)、V-Trace、Retrace(λ)、TD(λ)回报、折扣回报、资格迹和回合前缀和——全部重构成一个一阶线性递推关系,可在O(log T)并行步骤中求解。所有算法共享同一关联扫描算子,特定算法的融合Triton内核在芯片上构造其递推系数。我们从代数上验证了关联算子性质,并明确定义了终止与截断回合的处理方式。基准测试显示,在大规模并行模拟场景(数千环境,短轨迹)下,相比向量化torch-compile基线,速度提升1.6至5.70倍。该范围涵盖所有七种算法及两种GPU,含与不含每步截断处理。多数算法在序列更长时提速更高,因基线需更多扫描阶段,每阶段增加一次HBM往返开销。代码已开源:https://github.com/simonsays1980/rl-triton。

原文摘要 · Abstract (English)

We present rl-triton, an open-source library of high-performance GPU kernels for reinforcement learning credit assignment, implemented in Triton. The core contribution is a unified associative scan framework that recasts seven distinct RL estimation algorithms - Generalized Advantage Estimation (GAE), V-Trace, Retrace($λ$), TD($λ$) returns, discounted returns, eligibility traces, and episodic prefix sums - as instances of a single first-order linear recurrence solved in $O(\log T)$ parallel steps. All algorithms share the same associative scan operator, with algorithm-specific fused Triton kernels constructing their recurrence coefficients on-chip. We verify the associative operator algebraically and define the treatment of terminated and truncated episodes explicitly. Benchmarks show a 1.6-5.70$\times$ full-call speedup over a vectorized torch-compile baseline in the massively parallel simulation regime (thousands of environments, short rollouts). The reported range covers all seven algorithms on both GPUs, both with and without per-step truncation handling. For most algorithms, speedups increase at longer sequence lengths, as the baseline requires more scan stages as $\log T$ grows, each adding an intermediate HBM round-trip. The library is available at https://github.com/simonsays1980/rl-triton.

强化学习GPU优化Triton信用分配

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。