Home

rl-triton

High-performance Triton GPU kernels for reinforcement-learning credit assignment and return estimation.

Kernels

Function	Algorithm	Description
`compute_gae`	GAE	Generalized Advantage Estimation – backward scan over `δ + γλ·A`
`compute_vtrace`	V-Trace	IS-weighted targets and advantages – fused single-kernel for seq_len ≤ 131072
`compute_retrace`	Retrace(λ)	Off-policy return estimate with truncated IS ratios
`compute_lambda_returns`	TD(λ)	λ-return targets mixing one-step TD and Monte Carlo
`compute_discounted_returns`	Returns	Discounted reward-to-go
`compute_eligibility_traces`	Elig. traces	Accumulating forward traces `e[t] = x[t] + γλ(1-d[t-1])e[t-1]`
`compute_episodic_prefix_sum`	Prefix sum	Episodic cumulative sum with done-mask resets

For the seq_len > 131072 fallback behavior of each kernel, see GPU Concepts and the per-kernel pages under Kernels.

Installation

Use it in your project:

pip install git+https://github.com/simonsays1980/rl-triton

From source, editable (for modifying the kernels):

git clone https://github.com/simonsays1980/rl-triton
cd rl-triton
pip install -e .

Contributors (adds test/dev tooling):

pip install -e ".[dev]"

Usage

import torch
from rl_triton import compute_gae

rewards     = torch.randn(64, 512, device="cuda")
values      = torch.randn(64, 512, device="cuda")
terminateds = torch.zeros(64, 512, device="cuda")
advantages  = compute_gae(rewards, values, terminateds, gamma=0.99, lambda_=0.95)

Testing

# Correctness tests
pytest tests/ -v

# PR performance safeguard (one config per algorithm, requires CUDA)
pytest -m perf -v

# Full slow benchmark suite (all configs, requires CUDA)
pytest -m slow -v

Benchmarking

Run the full release benchmark suite with:

python tests/bench_release.py --parent-sweep --gpu "RTX 2000 Ada"

See Benchmarks for methodology, full results, and benchmark details.

Performance

Full sweep, methodology, and truncation-path results: Benchmarks.

Representative benchmark results at num_envs=4096, seq_len=128. Numbers report full-call speedup over torch.compile.

NVIDIA H100 80GB HBM3

algorithm	speedup vs torch.compile (full-call)
GAE	2.36×
V-Trace	2.78×
Retrace	1.91×
lambda-returns	2.85×
discounted-returns	2.70×
eligibility-traces	2.48×
prefix-sum	2.39×

With truncations, same configuration:

algorithm	speedup vs torch.compile, with truncations (full-call)
GAE	1.7×
V-Trace	2.7×
Retrace	1.9×
lambda-returns	2.6×
discounted-returns	2.6×

NVIDIA RTX 2000 Ada Generation

algorithm	speedup vs torch.compile (full-call)
GAE	2.57×
V-Trace	3.46×
Retrace	1.62×
lambda-returns	5.14×
discounted-returns	5.70×
eligibility-traces	2.28×
prefix-sum	2.25×

With truncations, same configuration:

algorithm	speedup vs torch.compile, with truncations (full-call)
GAE	1.6×
V-Trace	3.3×
Retrace	1.6×
lambda-returns	4.4×
discounted-returns	4.6×