Crate thrust_rl 

Thrust

High-performance reinforcement learning in Rust.

Thrust is a modern RL library built on top of the Burn tensor framework. After phase 5 of the Burn migration (#65), Burn is the only tensor backend in the workspace; the previous tch/libtorch path has been removed in favour of Burn’s multi-backend stack (CPU NdArray, WebGPU, CUDA, ROCm, Metal, Vulkan).

Quick Start

The prelude re-exports the types you need to stand up a training run: the Environment trait and a couple of built-in environments, the MlpBurnPolicy actor-critic network, and the trainer configs/trainers (A2C, PPO, DQN, SAC, BC). The example below builds a CartPole env, an A2C config, and a policy on Burn’s CPU NdArray backend — the same pieces the train_cartpole_a2c example wires into a full rollout/update loop.

use burn::backend::{Autodiff, NdArray};
use thrust_rl::prelude::*;

type Backend = Autodiff<NdArray<f32>>;

// 1. Build an environment and read its observation / action dims.
let env = CartPole::new();
let obs_dim = env.observation_space().shape[0];
let action_dim = match env.action_space().space_type {
    SpaceType::Discrete(n) => n,
    SpaceType::Box => panic!("CartPole is discrete"),
};

// 2. Construct the actor-critic policy on the CPU backend.
let device = Default::default();
let policy = MlpBurnPolicy::<Backend>::with_config(
    obs_dim,
    action_dim,
    MlpBurnConfig { hidden_dim: 64, ..Default::default() },
    &device,
);
let _ = policy;

// 3. Pick a trainer config. See the `train_cartpole_a2c` example for
//    the full rollout-collect + GAE + `A2cTrainer::train_step` loop.
let config = A2cConfig { num_envs: 16, n_steps: 5, ..Default::default() };
assert_eq!(config.num_envs, 16);

Modules

buffer

Experience buffers and replay management (requires training feature).

env

Environment traits and implementations Environment traits and implementations

inference

Pure Rust inference for WASM compilation Pure Rust inference module for WASM compatibility

multi_agent

Multi-agent training infrastructure (Burn backend).

policy

Policy and neural network implementations inference submodule available for WASM, training modules require training feature Policy and neural network wrappers.

prelude

Prelude module for convenient imports.

train

Training algorithms (PPO, DQN). Training algorithms (Burn backend).

utils

Utility functions and helpers Utility functions and helpers

Constants

VERSION

Current version of thrust-rl