rlox

Rust-accelerated reinforcement learning — the Polars architecture pattern applied to RL.

Documentation  |  Learning Path  |  RL Introduction  |  22 Algorithms  |  Examples

New to Reinforcement Learning? Start with our RL Introduction to learn the fundamentals, then follow the Learning Path from beginner to production.

Why rlox?

RL frameworks like Stable-Baselines3 and TorchRL do everything in Python — environment stepping, buffer storage, advantage computation. This works, but Python interpreter overhead becomes the bottleneck long before your GPU does.

rlox applies the Polars architecture pattern to RL: a Rust data plane handles the compute-heavy, latency-sensitive work (env stepping, buffers, GAE) while a Python control plane stays in charge of training logic, configs, and neural networks via PyTorch. The two connect through PyO3 with zero-copy where possible.

The result: 3-50x faster than SB3/TorchRL on data-plane operations, with the same Python training API you're used to.

Quick Start

Prerequisites

• Rust 1.75+ -- install via rustup:

  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
  

• Python 3.10-3.13
• Optional: pip install gymnasium[mujoco] for MuJoCo environments
• Optional: pip install pettingzoo for multi-agent environments

Installation

pip install rlox

Or build from source:

python3 -m venv .venv && source .venv/bin/activate
pip install maturin numpy gymnasium torch
maturin develop --release

Train PPO on CartPole in 3 lines:

from rlox import Trainer

trainer = Trainer("ppo", env="CartPole-v1", seed=42)
metrics = trainer.train(total_timesteps=50_000)
print(f"Mean reward: {metrics['mean_reward']:.1f}")

Train SAC on Pendulum:

from rlox import Trainer

trainer = Trainer("sac", env="Pendulum-v1", config={"learning_starts": 500})
metrics = trainer.train(total_timesteps=20_000)

Note: Per-algorithm trainers (PPOTrainer, SACTrainer, etc.) are deprecated. Use the unified Trainer("algo", ...) API instead.

Config-driven training (YAML):

python -m rlox train --config config.yaml

from rlox import TrainingConfig, train_from_config

config = TrainingConfig.from_yaml("config.yaml")
metrics = train_from_config(config)

Use Rust primitives directly:

import rlox

# 140x faster GAE than Python loops
advantages, returns = rlox.compute_gae(rewards, values, dones, last_value, gamma=0.99, lam=0.95)

# 35x faster GRPO advantages
advantages = rlox.compute_batch_group_advantages(rewards, group_size=4)

# Parallel env stepping (2.7M steps/s at 512 envs)
env = rlox.VecEnv(n=256, seed=42, env_id="CartPole-v1")
result = env.step_all(actions)

More examples in examples/ — PPO, SAC, GRPO custom rewards, fast GAE, VecEnv throughput.

Documentation

Resource	Link
Full Documentation	wojciechkpl.github.io/rlox
Getting Started	Tutorial
Python API Guide	User Guide
Examples	Code Examples
Rust API	cargo doc
Migrating from SB3	Migration Guide
API Reference	Autodoc

Architecture

┌──────────────────────────────────────────────────┐
│  Python (control plane)                          │
│  PPO, SAC, DQN, TD3, A2C, MAPPO, DreamerV3,     │
│  IMPALA, GRPO, DPO                               │
│  GymVecEnv, VecNormalize, callbacks,             │
│  YAML/TOML configs, trainers, checkpointing,     │
│  diagnostics dashboard                           │
│  vLLM/TGI/SGLang backends, multi-GPU (DDP)       │
├────────────── PyO3 boundary ─────────────────────┤
│  Rust (data plane)                               │
│  rlox-core:   envs (CartPole, Pendulum,           │
│               NonStationaryCartPole),            │
│               Rayon parallel stepping,           │
│               buffers (ring, mmap, priority),    │
│               GAE, V-trace, GRPO, pipeline,      │
│               EMA stats, CUSUM change detection  │
│  rlox-nn:     RL algorithm traits                │
│  rlox-burn:   Burn backend (NdArray)             │
│  rlox-candle: Candle backend (CPU)               │
│  rlox-python: PyO3 bindings                      │
└──────────────────────────────────────────────────┘

Multi-crate workspace (crates.io):

• rlox-core — pure Rust: environments, buffers (ring, mmap, priority), GAE, V-trace, GRPO, pipeline
• rlox-nn — RL algorithm traits (ActorCritic, QFunction, StochasticPolicy, etc.)
• rlox-burn — Burn Autodiff<NdArray> implementations
• rlox-candle — Candle CPU implementations
• rlox-python — PyO3 bindings exposing rlox-core to Python

For a deep-dive into the architecture, module relationships, and API reference, see the DeepWiki.

Benchmark Highlights

Measured 2026-04-08 on Apple M3 Pro (ARM, 12-core), CPU only, against stable-baselines3 2.7.1 and torchrl 0.11.1. All results statistically significant (bootstrap 95% CI lower bound > 1.0).

Component	vs NumPy / SB3	vs TorchRL	Details
GAE (32K steps)	135x vs NumPy	1,588x	docs/benchmark/gae.md
Replay buffer push (obs_dim=4, 10K)	4.6x vs SB3	70x	docs/benchmark/buffer-ops.md
Replay buffer sample (batch=32)	9.7x vs SB3	9x	docs/benchmark/buffer-ops.md
E2E rollout (256×2048)	3.1x vs SB3	42x	docs/benchmark/e2e-rollout.md
GRPO advantages (64×8)	41x vs NumPy	42x vs PyTorch	docs/benchmark/llm-ops.md
Tokenwise KL (seq=128)	5x vs NumPy	9x vs PyTorch	docs/benchmark/llm-ops.md

Full methodology, raw data, and reproducibility instructions: docs/benchmark/

Performance

Key numbers at a glance (Apple M3 Pro, CPU only, median of 100–200 timed iterations):

Operation	rlox	Baseline	Speedup
GAE (32K steps)	69 µs	9,376 µs (NumPy loop)	135x
Replay buffer push (10K, obs_dim=4)	3.2 ms	14.7 ms (SB3)	4.6x
Replay buffer sample (batch=32)	2.0 µs	19.3 µs (SB3)	9.7x
E2E rollout (256×2048 trans)	669 ms	2,057 ms (SB3)	3.1x
GRPO advantages (64×8 groups)	7.8 µs	318 µs (NumPy)	41x
Tokenwise KL (seq=128)	0.3 µs	3.0 µs (PyTorch)	9x

Convergence (rlox vs SB3)

Same hyperparameters (rl-zoo3 defaults), same evaluation harness, 5 seeds per cell, IQM + bootstrap 95% CI per Agarwal et al. 2021. 10 cells, convergence parity on every one.

Algo	Environment	rlox IQM	rlox CI	SB3 IQM	SB3 CI
PPO	CartPole-v1	450.8	[440.5, 454.2]	438.2	[389.7, 500.0]
PPO	Acrobot-v1	-86.0	[-89.7, -83.0]	-83.7	[-97.0, -77.4]
PPO	Hopper-v4	932.8	[706.0, 2190.4]	1173.1	[719.4, 1578.8]
PPO	HalfCheetah-v4	1854.6	[1381.3, 2598.8]	1568.7	[1516.9, 3094.3]
SAC	Pendulum-v1	-152.1	[-173.9, -129.5]	—	—
SAC	HalfCheetah-v4	10871.9	[10294.9, 11293.1]	10795.5	[10499.7, 11542.2]
TD3	Pendulum-v1	-149.1	[-171.7, -134.2]	—	—
TD3	HalfCheetah-v4	10880.1	[7584.4, 11299.1]	—	—
DQN	CartPole-v1	500.0	[195.8, 500.0]	500.0	[217.6, 500.0]
A2C	CartPole-v1	417.8	[82.5, 500.0]	491.6	[167.5, 500.0]

SAC HalfCheetah: rlox 10872 vs SB3 10796 — statistically identical, both beat the zoo reference (9656) by ~12%. TD3 HalfCheetah: rlox 10880 beats the zoo reference (9709) by 12%. DQN CartPole: both frameworks hit 500 (perfect). A2C CartPole: rlox 418 vs SB3 492, CIs overlap. The PPO MuJoCo "gap" vs zoo references is a protocol-and-version artifact (v4 vs v3, different eval protocol), not a framework deficit — both rlox and SB3 show the same gap when measured in the same harness.

Full convergence results, learning curves, and performance profiles: docs/benchmark/convergence-results.md

Features

• 22 Algorithms: PPO, SAC, DQN, TD3, A2C, VPG, TRPO, MAPPO, DreamerV3, IMPALA, and more (+ GRPO, DPO for LLM)
• Trainers: Each algorithm has a high-level Trainer with train(), save(), from_checkpoint(), predict()
• Environments: Gymnasium-compatible, Rayon-parallel VecEnv, CartPole and Pendulum-v1 built-in
• Visual RL wrappers: FrameStack, ImagePreprocess, AtariWrapper, DMControlWrapper for pixel-based RL
• Language RL wrappers: LanguageWrapper, GoalConditionedWrapper for language-grounded tasks
• Plugin ecosystem: ENV_REGISTRY, BUFFER_REGISTRY, REWARD_REGISTRY, discover_plugins for extensibility
• Model zoo: ModelZoo.register, ModelZoo.load for sharing and reusing pretrained agents
• VecNormalize: Obs/reward normalization at the environment boundary (SB3-compatible)
• Buffers: ring, mmap, priority replay — all in Rust with zero-copy Python access
• Config-driven training: YAML/TOML configs via TrainingConfig and python -m rlox train --config config.yaml
• Diagnostics dashboard: TerminalDashboard, HTMLReport, entropy/KL/gradient monitoring
• LLM post-training: GRPO, DPO, token KL, sequence packing, vLLM/TGI/SGLang backends
• Cloud deploy: Dockerfile generator, Kubernetes manifest generator, SageMaker integration
• Distributed: pipeline parallelism (crossbeam), gRPC workers, multi-GPU (DDP)
• Evaluation: Trainer.evaluate(), Trainer.enjoy(), VideoRecordingCallback, score normalization
• Non-stationary RL: EMA running stats, CUSUM change-point detection, sliding window replay, dynamic regret metrics
• Asymmetric actor-critic: AsymmetricPolicy for privileged critic observations (sim-to-real)
• Production: callbacks, checkpointing, eval toolkit (IQM, bootstrap CI, performance profiles)
• NN backends: Burn (NdArray) and Candle (CPU) for pure-Rust inference, PyTorch for training
• 444 Rust tests, ~1094 Python tests — comprehensive coverage

Tutorials & Documentation

Guide	Description
Getting Started	Installation, first training run, basic API
Custom Rewards & Training Loops	Reward shaping, GRPO reward functions, custom algorithms
Python Guide	Python API reference and patterns
Rust Guide	Rust crate architecture and extending in Rust
Math Reference	GAE, V-trace, GRPO, DPO derivations
Benchmark Details	Full methodology, per-benchmark analysis, reproducibility
DeepWiki	Auto-generated architecture docs and API reference

Running Tests

# Rust tests (444 tests across all crates)
cargo test --workspace

# Python tests (~1094 tests, after maturin develop)
pip install -e ".[all]"
pytest tests/python/ -q

# Quick smoke test (skip slow tests)
pytest tests/python/ -m "not slow" -q

# Single crate
cargo test --package rlox-core

# All tests (Rust + Python)
./scripts/test.sh

# Full benchmark suite (rlox vs TorchRL vs SB3)
.venv/bin/python benchmarks/run_all.py

Project Layout

crates/
  rlox-core/       Pure Rust: envs, buffers (ring, mmap, priority), GAE,
                   V-trace, GRPO, pipeline (crossbeam), sequence packing
  rlox-nn/         RL algorithm traits (ActorCritic, QFunction, etc.)
  rlox-burn/       Burn backend (Autodiff<NdArray>)
  rlox-candle/     Candle backend (CPU)
  rlox-python/     PyO3 bindings
  rlox-bench/      Criterion benchmarks (env stepping, NN backends)
python/rlox/
  algorithms/      PPO, SAC, DQN, TD3, A2C, GRPO, DPO, MAPPO, DreamerV3, IMPALA
  distributed/     Pipeline, vLLM/TGI/SGLang backends, multi-GPU (DDP)
  llm/             LLM environment, reward model serving
  *.py             Collectors, configs, callbacks, policies, trainers,
                   evaluation toolkit, diagnostics, checkpointing
benchmarks/        Three-framework benchmark suite + convergence tests
tests/python/      Python integration & benchmark TDD tests
docs/              Guides, tutorials, benchmark methodology

Citation

If you use rlox in your research, please cite:

@software{kowalinski2026rlox,
  author       = {Kowalinski, Wojciech},
  title        = {rlox: Rust-Accelerated Reinforcement Learning},
  year         = {2026},
  url          = {https://github.com/wojciechkpl/rlox},
  version      = {1.0.0},
  license      = {MIT OR Apache-2.0}
}

License

Dual-licensed under MIT or Apache 2.0, at your option.