thrust-rl 0.4.0

High-performance reinforcement learning in Rust with the Burn tensor backend
Documentation

🚀 Thrust

High-performance reinforcement learning in Rust

Crates.io Documentation License

Give your agents some thrust 🚀

🎮 Try the Live Demo

Watch trained RL agents play CartPole, Snake, and Pong in real-time, running entirely in your browser via WebAssembly!

Thrust is a modern reinforcement learning library built from the ground up in Rust, designed for maximum speed, memory safety, and scalability. Inspired by PufferLib, Thrust combines the raw performance of Rust with the Burn tensor framework's multi-backend GPU support (CUDA / ROCm / Metal / Vulkan / WebGPU / NdArray CPU) to deliver fast training without any C++ FFI dependencies.

🎯 Vision

Our goal is to create the fastest, safest, and most ergonomic reinforcement learning library in existence, with:

  • Live browser demos where trained agents play games in real-time via WebAssembly
  • Production-grade performance that scales from research to deployment
  • Best-in-class developer experience with type safety and clear abstractions
  • Full feature parity with leading Python RL libraries, but faster

✨ Features

  • 🚀 Blazing Fast: optimized Rust pipeline with multi-backend GPU support
  • 🦀 Memory Safe: Leverage Rust's ownership system for fearless concurrency
  • 🔥 Burn Powered: Neural networks via Burn — pure-Rust autodiff, no libtorch FFI
  • ⚡ Async Vectorization: High-performance environment parallelization with Tokio / Rayon
  • 🎮 Live Demos: Train agents and deploy them in the browser via WebAssembly
  • 🎯 Production Ready: Built for research and industry use cases

📦 Project Status

v0.3.0 — published on crates.io. The full loop is shipped: train in Rust, export the policy, and run the trained agent in the browser via WebAssembly. That covers five single-agent algorithms (PPO, A2C, DQN, SAC, BC) plus recurrent (LSTM) PPO with V-trace off-policy correction and an async actor-learner, two multi-agent meta-solvers (PSRO, NFSP), a growing environment zoo (including POMDP variants), runnable examples, WASM bindings (src/wasm.rs), and a live demo with CartPole, Snake, Pong, and a contextual-bandit playground. The public API is documented on docs.rs, warning-free, and packaged for crates.io.

Active work is multi-agent communication (Phase 3), distributed training, and mixed-precision support (blocked on upstream Burn/Metal kernels).

See CHANGELOG.md for release notes, docs/RELEASING.md for the publish process, and ROADMAP.md for the detailed development schedule.

🎯 Roadmap

Phase 1: Foundation (Complete ✅)

  • Experience buffer implementation
  • Policy wrapper (Burn)
  • EnvPool vectorization
  • CartPole environment (301.6 avg reward achieved)
  • PPO training loop with GPU support
  • DQN training loop (replay buffer + target network, CartPole)
  • Checkpoint saving/loading
  • Snake environment (multi-agent support)
  • SimpleBandit environment (contextual bandits)

Phase 2: Multi-Agent & WASM (Mostly Complete ✅)

  • Multi-agent training infrastructure (multi_agent::PolicyLearner is experimental — API may change)
  • Population-based training design
  • PSRO with α-rank meta-solver (N-player)
  • NFSP (approximate, N-player, multi-discrete)
  • Bucket Brigade cooperative-MARL integration
  • Pure Rust inference (Burn backend or hand-rolled WASM path)
  • Universal inference system (JSON model format)
  • Complete WASM bindings (src/wasm.rs)
  • Browser-based demos (live: CartPole, Snake, Pong, bandit)
  • Multi-agent communication channels (Phases 1–2: CommunicatingEnvironment + SignalingGame + trainer hook; differentiable Phase 3 deferred)

Phase 3: Features

  • SAC for continuous control (twin critics, auto-entropy, Polyak targets)
  • Continuous-control environments (PendulumSwingUp, ContinuousLqr, MountainCarContinuous)
  • A2C (Advantage Actor-Critic)
  • Behavioral cloning / imitation learning
  • Prioritized experience replay (sum-tree buffer; DqnConfig::prioritized_replay)
  • LSTM policy support (recurrent PPO validated on a FlickeringCartPole POMDP)
  • V-trace importance sampling (IMPALA; corrects staleness in the async actor-learner)
  • Async actor-learner (single-host, 2.1× synchronous throughput)
  • Mixed precision training (blocked on upstream Burn/Metal bf16 kernels)
  • Distributed training (multi-host; design note in docs/DISTRIBUTED_TRAINING_DESIGN.md)

Phase 4: Demo Site (Live ✅)

  • WebAssembly policy compilation
  • Browser inference engine
  • Live training dashboard (recorded CartPole curves)
  • Public demo deployment (rjwalters.github.io/thrust)

🏗️ Architecture

┌─────────────────────────────────────────────────┐
│         Core Library (thrust-rl)                │
│  ┌──────────────┐  ┌──────────────────────┐    │
│  │ Policy       │  │ Vectorization        │    │
│  │ (Burn)       │  │ (Tokio/Rayon)        │    │
│  └──────────────┘  └──────────────────────┘    │
│  ┌──────────────┐  ┌──────────────────────┐    │
│  │ Experience   │  │ Environments         │    │
│  │ Buffers      │  │ (Pure Rust)          │    │
│  └──────────────┘  └──────────────────────┘    │
│  ┌─────────────────────────────────────────┐   │
│  │ PPO / DQN / SAC + PSRO / NFSP trainers   │   │
│  │            (Burn autodiff)               │   │
│  └─────────────────────────────────────────┘   │
└─────────────────────────────────────────────────┘

🎮 Environments

  • CartPole ✅ - Classic control benchmark (solved: 301.6 avg reward)
  • FlickeringCartPole ✅ - POMDP variant (whole-observation dropout, Hausknecht & Stone); the recurrent-PPO memory benchmark
  • MaskedCartPole ✅ - Velocity-masked variant (kept as a documented negative result: not actually a POMDP)
  • GridWorld ✅ - Discrete navigation benchmark (DQN)
  • PendulumSwingUp ✅ - Continuous-control benchmark (Gym Pendulum-v1; SAC's reference env)
  • MountainCarContinuous ✅ - Deceptive-reward continuous-control benchmark (SAC)
  • ContinuousLqr ✅ - Linear-quadratic regulator (continuous-action trait existence proof)
  • Snake ✅ - Multi-agent grid world with torus wrapping
  • Pong ✅ - Two-player competitive self-play
  • SimpleBandit ✅ - Contextual multi-armed bandits
  • Matching Pennies ✅ - Two-player and N-player zero-sum (PSRO/NFSP smoke tests)
  • SignalingGame ✅ - Sender/receiver communication reference env (CommunicatingEnvironment)
  • Bucket Brigade ✅ - Cooperative multi-agent coordination (Slepian-Wolf MARL adapter)
  • More coming soon!

🧠 Algorithms

  • PPO ✅ - Proximal Policy Optimization (on-policy, actor-critic; single- and multi-agent), plus a recurrent (LSTM) variant for POMDPs and an async actor-learner with V-trace off-policy correction
  • A2C ✅ - Advantage Actor-Critic (synchronous, on-policy; un-clipped policy gradient + MSE value, one update per rollout)
  • DQN ✅ - Deep Q-Network (off-policy, replay buffer + target network), including Double-DQN target computation and optional Polyak (soft) target updates
  • SAC ✅ - Soft Actor-Critic (off-policy, continuous control; twin critics, automatic entropy tuning, Polyak target updates)
  • BC ✅ - Behavioral Cloning (supervised imitation learning from expert demonstrations)
  • PSRO ✅ - Policy-Space Response Oracles with α-rank meta-solver (multi-agent, N-player)
  • NFSP ✅ - Neural Fictitious Self-Play (approximate, N-player, multi-discrete actions)

See the example gallery for a runnable trainer per algorithm.

📚 Inspiration

Thrust is inspired by:

  • PufferLib - Python RL library achieving 1M+ SPS
  • Burn - Pure-Rust deep learning framework with multi-backend GPU support
  • Border - Rust RL library

🚀 Quick Start

Local Development (CPU)

# Build and run the bandit training example (CPU NdArray backend)
cargo run --release --example train_simple_bandit

The default training feature builds against Burn's pure-Rust NdArray backend; no external libraries (libtorch, CUDA, etc.) are required.

📚 Tutorials — a guided learning path

New to Thrust? Start with the Tutorial Series: a dependency-ordered path from your first agent to a trained-and-deployed policy. Every code snippet is CI-doc-tested, so nothing rots. Begin with Your first agent and Solving CartPole with PPO.

GPU Training

Burn ships several GPU backends behind feature flags. Compose them with the default training feature:

# wgpu (cross-platform: Vulkan / Metal / DX12 / WebGPU)
cargo run --release --features "training,wgpu" --example train_simple_bandit

# CUDA (Linux + NVIDIA)
cargo run --release --features "training,cuda" --example train_simple_bandit

More Examples

The bandit trainer is just one of eighteen runnable examples. See the Example Gallery for the full set — a trainer per algorithm (PPO, recurrent PPO, A2C, DQN, SAC, BC, PSRO, NFSP) across all environments, plus the async actor-learner and research harnesses, with copy-paste run commands and the env vars each one honors.

Library Usage

use thrust_rl::prelude::*;

// See examples/games/bandit/train_simple_bandit.rs for the end-to-end
// rollout/loss/update loop using the Burn backend.

📊 Performance

A criterion throughput harness (benches/trainer_throughput.rs) measures per-update and full-loop steps/sec for PPO, A2C, DQN, and SAC. Measured CPU-vs-GPU numbers (Burn NdArray vs. wgpu on an RTX 4090) are in docs/BURN_BACKENDS.md.

🤝 Contributing

We welcome contributions! Thrust is published and actively developed.

Ways to contribute:

  • 🐛 Report bugs and issues
  • 💡 Suggest features or improvements
  • 📝 Improve documentation
  • 🔧 Implement environments or algorithms
  • ⚡ Optimize performance
  • 🎨 Design the demo website

See CONTRIBUTING.md for guidelines and ROADMAP.md for areas where we need help.

📄 License

Licensed under either of:

at your option.

🌟 Star the project!

If you find Thrust interesting, give it a star to help others discover it!

Built with 🦀 Rust and ❤️ for reinforcement learning