### PPO Training Loop Example

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Demonstrates a complete PPO training loop using the ActorCritic model, SyncVecEnv, ppo_collect, and ppo_update. Includes PPO configuration, optimizer setup, and logging. Requires specific imports and backend setup.

```rust
type AB = Autodiff<NdArray>;

fn main() {
    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let mut rng = rand::rngs::SmallRng::seed_from_u64(42);
    let n_envs = 4;

    let envs: Vec<CartPole<_>> = (0..n_envs)
        .map(|i| CartPole::new(rand::rngs::SmallRng::seed_from_u64(i as u64)))
        .collect();
    let mut vec_env = SyncVecEnv::new(envs);
    let mut model: ActorCritic<AB> = ActorCritic::new(&device);
    let mut optim = AdamConfig::new().with_epsilon(1e-5).init();
    let config = PpoConfig::default(); // lr=2.5e-4, gamma=0.99, clip_eps=0.2
    let mut logger = PrintLogger::new(0);
    let mut ep_acc = vec![0.0f32; n_envs];
    let mut current_obs = vec_env.reset();

    for iter in 0..100 {
        // Linear LR annealing
        let frac = 1.0 - iter as f64 / 100.0;
        let lr = config.lr * frac;

        let rollout = ppo_collect::<NdArray, _, _>(
            &model.valid(), &mut vec_env, &config, &device, &mut rng,
            &mut current_obs, &mut ep_acc,
        );
        let (new_model, stats) = ppo_update(
            model, &mut optim, &rollout, &config, lr, &device, &mut rng,
        );
        model = new_model;

        let step = (iter + 1) as u64 * (config.n_steps * n_envs) as u64;
        stats.log(&mut logger, step);
        if !rollout.episode_returns.is_empty() {
            let avg = rollout.episode_returns.iter().sum::<f32>()
                / rollout.episode_returns.len() as f32;
            logger.log_scalar("rollout/avg_return", avg as f64, step);
            // Expect returns to reach 500 (solved) within ~50 iterations
        }
    }
    logger.flush();
}
```

--------------------------------

### GAE Usage Example

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/gae.md

Example demonstrating how to call the GAE function with sample rewards, values, and done flags. Ensure the 'rl4burn::gae' module is imported.

```rust
use rl4burn::gae;

let rewards = vec![1.0, 1.0, 1.0, 0.0];
let values = vec![5.0, 4.0, 3.0, 2.0];
let dones = vec![false, false, false, true];
let (advantages, returns) = gae::gae(&rewards, &values, &dones, 0.0, 0.99, 0.95);
```

--------------------------------

### PPO on CartPole (discrete) Example

Source: https://github.com/rpp1011/rl4burn/blob/master/README.md

This example demonstrates training a PPO agent on the CartPole environment using the rl4burn library. It sets up the environment, defines an ActorCritic model, and iterates through collection and update steps.

```rust
use burn::backend::{Autodiff, NdArray};
use burn::module::{AutodiffModule, Module};
use burn::nn::{Linear, LinearConfig};
use burn::optim::AdamConfig;
use burn::prelude::*;
use rand::SeedableRng;

use rl4burn::envs::CartPole;
use rl4burn::{DiscreteAcOutput, DiscreteActorCritic, ppo_collect, ppo_update, PpoConfig, SyncVecEnv};

#[derive(Module, Debug)]
struct ActorCritic<B: Backend> {
    actor_fc1: Linear<B>,
    actor_fc2: Linear<B>,
    actor_out: Linear<B>,
    critic_fc1: Linear<B>,
    critic_fc2: Linear<B>,
    critic_out: Linear<B>,
}

impl<B: Backend> ActorCritic<B> {
    fn new(device: &B::Device) -> Self {
        Self {
            actor_fc1: LinearConfig::new(4, 64).init(device),
            actor_fc2: LinearConfig::new(64, 64).init(device),
            actor_out: LinearConfig::new(64, 2).init(device),
            critic_fc1: LinearConfig::new(4, 64).init(device),
            critic_fc2: LinearConfig::new(64, 64).init(device),
            critic_out: LinearConfig::new(64, 1).init(device),
        }
    }
}

impl<B: Backend> DiscreteActorCritic<B> for ActorCritic<B> {
    fn forward(&self, obs: Tensor<B, 2>) -> DiscreteAcOutput<B> {
        let a = self.actor_fc1.forward(obs.clone()).tanh();
        let a = self.actor_fc2.forward(a).tanh();
        let logits = self.actor_out.forward(a);

        let c = self.critic_fc1.forward(obs).tanh();
        let c = self.critic_fc2.forward(c).tanh();
        let values = self.critic_out.forward(c).squeeze_dim::<1>(1);

        DiscreteAcOutput { logits, values }
    }
}

type AB = Autodiff<NdArray>;

fn main() {
    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let mut rng = rand::rngs::SmallRng::seed_from_u64(42);

    let n_envs = 4;
    let envs: Vec<CartPole<rand::rngs::SmallRng>> = (0..n_envs)
        .map(|i| CartPole::new(rand::rngs::SmallRng::seed_from_u64(i as u64)))
        .collect();
    let mut vec_env = SyncVecEnv::new(envs);
    let mut model: ActorCritic<AB> = ActorCritic::new(&device);
    let mut optim = AdamConfig::new().with_epsilon(1e-5).init();
    let config = PpoConfig::default();
    let mut current_obs = vec_env.reset();
    let mut ep_acc = vec![0.0f32; n_envs];

    for iter in 0..100 {
        let rollout = ppo_collect::<burn::backend::NdArray, _, _>(
            &model.valid(), &mut vec_env, &config, &device, &mut rng,
            &mut current_obs, &mut ep_acc,
        );
        let (new_model, _stats) = ppo_update(
            model, &mut optim, &rollout, &config, config.lr, &device, &mut rng,
        );
        model = new_model;

        if !rollout.episode_returns.is_empty() {
            let avg: f32 = rollout.episode_returns.iter().sum::<f32>()
                / rollout.episode_returns.len() as f32;
            println!("iter {iter}: avg_return={avg:.0}");
        }
    }
}

```

--------------------------------

### Provider-Agnostic Code Example

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Demonstrates how to write provider-agnostic code by working with a slice of `&dyn CloudProvider` to find the cheapest offer across multiple providers.

```APIDOC
## Comparing Providers

You can write provider-agnostic code by working with `&dyn CloudProvider`:

```rust,ignore
fn cheapest_offer(
    providers: &[&dyn CloudProvider],
    reqs: &InstanceRequirements,
) -> Option<(GpuOffer, &'static str)> {
    let mut best: Option<(GpuOffer, &str)> = None;
    for provider in providers {
        if let Ok(offers) = provider.search_offers(reqs) {
            for offer in offers {
                if best.as_ref().map_or(true, |(b, _)| offer.price_per_hour < b.price_per_hour) {
                    best = Some((offer, provider.name()));
                }
            }
        }
    }
    best
}
```
```

--------------------------------

### View TensorBoard Logs Locally

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/saving-sharing.md

Launches a local TensorBoard server to view training logs. Ensure TensorBoard is installed and the log directory is correctly specified.

```bash
tensorboard --logdir runs/
```

--------------------------------

### RunPod Provider Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Example of how to use the RunPod provider to search for GPU offers, launch an instance, and stop it.

```APIDOC
## RunPod

```rust,ignore
use rl4burn::cloud::{RunPodProvider, CloudProvider};

let provider = RunPodProvider::new(std::env::var("RUNPOD_API_KEY").unwrap());

let offers = provider.search_offers(&reqs)?;
let instance = provider.launch(&offers[0])?;
println!("Pod {} status: {}", instance.instance_id, instance.status);

provider.stop(&instance.instance_id)?;
```
```

--------------------------------

### PPO LR Annealing Example

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/algorithms/ppo.md

Demonstrates linear learning rate annealing for PPO updates. For constant LR, simply pass config.lr.

```rust
let frac = 1.0 - iter as f64 / n_iterations as f64;
let current_lr = config.lr * frac;
```

--------------------------------

### Example QNet Implementation

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/algorithms/dqn.md

An example implementation of the `QNetwork` trait using Burn's `Linear` layers and ReLU activation. Input shape is `[batch, obs_dim]` and output shape is `[batch, n_actions]`.

```rust
use burn::tensor::activation::relu;

#[derive(Module, Debug)]
struct QNet<B: Backend> {
    fc1: Linear<B>,
    fc2: Linear<B>,
    q_head: Linear<B>,
}

impl<B: Backend> QNetwork<B> for QNet<B> {
    fn q_values(&self, obs: Tensor<B, 2>) -> Tensor<B, 2> {
        let h = relu(self.fc1.forward(obs));
        let h = relu(self.fc2.forward(h));
        self.q_head.forward(h)
    }
}
```

--------------------------------

### Verify rl4burn Installation with a Rust Program

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/getting-started/installation.md

Create a main.rs file to test the installation by initializing a CartPole environment and performing a step. This verifies that rl4burn and its dependencies are correctly set up.

```rust
use rl4burn::envs::CartPole;
use rl4burn::env::Env;
use rand::SeedableRng;

fn main() {
    let mut env = CartPole::new(rand::rngs::SmallRng::seed_from_u64(42));
    let obs = env.reset();
    println!("CartPole observation: {:?}", obs);

    let step = env.step(1); // push right
    println!("Reward: {}, Done: {}", step.reward, step.done());
}
```

--------------------------------

### Multi-discrete ActionDist Initialization

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/algorithms/ppo.md

Example of initializing ActionDist for a multi-discrete action space, specifying the number of actions for each discrete dimension.

```rust
use rl4burn::{ActionDist, MaskedActorCritic, masked_ppo_collect, masked_ppo_update};

// Action space: [action_type(5), target(10)]
let action_dist = ActionDist::MultiDiscrete(vec![5, 10]);
```

--------------------------------

### PPO Action Masking Setup

Source: https://github.com/rpp1011/rl4burn/blob/master/README.md

Configures ActionDist for multi-discrete actions and demonstrates how environments can provide action masks via the Env trait.

```rust
let action_dist = ActionDist::MultiDiscrete(vec![3, 3]);

// Environments provide masks via the Env trait
impl Env for MyEnv {
    fn action_mask(&self) -> Option<Vec<f32>> {
        let mut mask = vec![1.0; 6]; // 3+3 logits
        if self.at_left_wall()  { mask[0] = 0.0; } // can't go left
        if self.at_right_wall() { mask[2] = 0.0; } // can't go right
        Some(mask)
    }
    // ...
}
```

--------------------------------

### Vast.ai Provider Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Example of how to use the Vast.ai provider to search for GPU offers, launch an instance, check its status, and stop it.

```APIDOC
## Vast.ai

```rust,ignore
use rl4burn::cloud::{VastAiProvider, CloudProvider};

let provider = VastAiProvider::new(std::env::var("VASTAI_API_KEY").unwrap());

// Search for offers
let offers = provider.search_offers(&reqs)?;
println!("Found {} offers, cheapest: ${}/hr",
    offers.len(), offers[0].price_per_hour);

// Launch the cheapest offer
let instance = provider.launch(&offers[0])?;
println!("Instance {} status: {}", instance.instance_id, instance.status);

// Check status
let inst = provider.status(&instance.instance_id)?;
if let Some(ssh) = &inst.ssh_connection {
    println!("Connect: {}", ssh);
}

// Cleanup
provider.stop(&instance.instance_id)?;
```
```

--------------------------------

### Custom GradientSync Implementation Example

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/distributed.md

Illustrates how to implement the `GradientSync` trait for custom cluster communication libraries like MPI. This example shows a placeholder for calling MPI_Allreduce.

```rust
struct MpiSync { /* ... */ }

impl GradientSync for MpiSync {
    fn all_reduce_f32(&self, values: &[f32], strategy: ReduceStrategy) -> Vec<f32> {
        // Call MPI_Allreduce
    }
    // ...
}
```

--------------------------------

### Initialize CartPole Environment with SyncVecEnv

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/getting-started/first-agent.md

Wrap multiple CartPole environments in SyncVecEnv for parallel execution. This setup automatically handles episode resets.

```rust
use rl4burn::envs::CartPole;
use rl4burn::SyncVecEnv;
use rand::SeedableRng;

let n_envs = 4;
let envs: Vec<CartPole<rand::rngs::SmallRng>> = (0..n_envs)
    .map(|i| CartPole::new(rand::rngs::SmallRng::seed_from_u64(i as u64)))
    .collect();
let mut vec_env = SyncVecEnv::new(envs);
```

--------------------------------

### DQN QNetwork Implementation

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Implement the QNetwork trait for your Q-network architecture. This example shows a simple feedforward network for Q-value estimation.

```rust
use burn::tensor::activation::relu;
use rl4burn::{
    dqn_update,
    epsilon_greedy,
    epsilon_schedule,
    DqnConfig,
    ReplayBuffer,
    Transition,
    polyak_update,
    QNetwork,
};

#[derive(Module, Debug)]
struct QNet<B: Backend> { fc1: Linear<B>, fc2: Linear<B>, q_head: Linear<B> }

impl<B: Backend> QNetwork<B> for QNet<B> {
    fn q_values(&self, obs: Tensor<B, 2>) -> Tensor<B, 2> {
        let h = relu(self.fc2.forward(relu(self.fc1.forward(obs))));
        self.q_head.forward(h) // [batch, n_actions]
    }
}
```

--------------------------------

### Compose Multiple Environment Wrappers

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/concepts/wrappers.md

Wrappers can be composed naturally to apply multiple transformations sequentially. This example shows composing EpisodeStats, RewardClip, and NormalizeObservation.

```rust
let env = EpisodeStats::new(
    RewardClip::new(
        NormalizeObservation::new(my_env, 10.0).unwrap(),
        1.0
    )
);
```

--------------------------------

### Custom Environment Implementation

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/concepts/environments.md

An example of implementing the `Env` trait for a custom environment named `MyEnv`. It includes state management, reset, step logic, and space definitions.

```rust
use rl4burn::env::{Env, Step};
use rl4burn::space::Space;

struct MyEnv {
    state: f32,
    step_count: usize,
}

impl Env for MyEnv {
    type Observation = Vec<f32>;
    type Action = usize;

    fn reset(&mut self) -> Vec<f32> {
        self.state = 0.0;
        self.step_count = 0;
        vec![self.state]
    }

    fn step(&mut self, action: usize) -> Step<Vec<f32>> {
        self.state += if action == 0 { -0.1 } else { 0.1 };
        self.step_count += 1;
        Step {
            observation: vec![self.state],
            reward: -self.state.abs(), // reward for staying near 0
            terminated: self.state.abs() > 1.0,
            truncated: self.step_count >= 200,
        }
    }

    fn observation_space(&self) -> Space {
        Space::Box { low: vec![-2.0], high: vec![2.0] }
    }

    fn action_space(&self) -> Space {
        Space::Discrete(2)
    }
}
```

--------------------------------

### Compose Environment Wrappers for Observation and Reward Transformation

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Chain environment wrappers to modify observations and rewards. This example composes `EpisodeStats`, `RewardClip`, and `NormalizeObservation` to track episode statistics, clip rewards, and normalize observations.

```rust
use rl4burn::wrapper::{EpisodeStats, RewardClip, NormalizeObservation};

// Compose: normalize observations, clip rewards, track episode stats
let env = EpisodeStats::new(
    RewardClip::new(
        NormalizeObservation::new(my_env, 10.0).unwrap(), // clip to [-10, 10]
        1.0,  // clip rewards to [-1, 1]
    )
);

// After each episode ends:
// env.last_episode_reward => Some(f32)
// env.last_episode_length => Some(usize)
```

--------------------------------

### Composite Logger Setup and Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/logging.md

Initialize a CompositeLogger with PrintLogger and TensorBoardLogger. Log scalar values and flush at the end of training iterations. Ensure TensorBoardLogger is enabled via feature flags.

```rust
use rl4burn::{CompositeLogger, Loggable, Logger, PrintLogger, TensorBoardLogger};

let mut logger = CompositeLogger::new(vec![
    Box::new(PrintLogger::new(5000)),
    Box::new(TensorBoardLogger::new("runs/ppo").unwrap()),
]);

for iter in 0..n_iterations {
    let rollout = ppo_collect::<NdArray, _, _>(&model.valid(), &mut vec_env, &config, &device, &mut rng, &mut current_obs, &mut ep_acc);

    let step = (iter + 1) as u64 * steps_per_iter as u64;
    if !rollout.episode_returns.is_empty() {
        let avg = rollout.episode_returns.iter().sum::<f32>() / rollout.episode_returns.len() as f32;
        logger.log_scalar("rollout/avg_return", avg as f64, step);
    }

    let (new_model, stats) = ppo_update(model, &mut optim, &rollout, &config, lr, &device, &mut rng);
    model = new_model;
    stats.log(&mut logger, step);
}
logger.flush();
```

--------------------------------

### PfspMatchmaking Initialization and Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/pfsp.md

Demonstrates how to initialize the PfspMatchmaking system, add opponents, record match results, and sample an opponent for the next match.

```APIDOC
## PfspMatchmaking Initialization and Usage

### Description
This snippet shows how to set up and use the `PfspMatchmaking` struct. It covers initializing the matchmaking configuration, adding opponents to the system, recording the outcomes of matches against specific opponents, and sampling the next opponent to play against.

### Initialization
Initialize `PfspMatchmaking` with a `PfspConfig` that defines the `power` for weighting opponent selection and `min_prob` to ensure a minimum selection probability for all opponents.

```rust
use rl4burn::{PfspMatchmaking, PfspConfig};

let mut mm = PfspMatchmaking::new(PfspConfig {
    power: 1.0,    // higher = more focus on hard opponents
    min_prob: 0.01, // every opponent has at least 1% chance
});
```

### Adding Opponents
Add opponents to the matchmaking system using their unique identifiers.

```rust
mm.add_opponent(0);
mm.add_opponent(1);
mm.add_opponent(2);
```

### Recording Results
Record the outcome of a match. The `record_result` function takes the opponent ID, a boolean indicating if the match was won, and a boolean indicating if it was a draw.

```rust
// Record results
mm.record_result(0, true, false);   // beat opponent 0
mm.record_result(1, false, false);  // lost to opponent 1
```

### Sampling Opponent
Sample an opponent for the next match. The sampling logic prioritizes opponents based on the `(1 - win_rate) ^ power` formula.

```rust
// Sample: opponent 1 (harder) is sampled more often
let opponent = mm.sample_opponent(&mut rng);
```

### Weighting Formula

The selection probability for an opponent is calculated using the formula: `(1 - win_rate) ^ power`.

- Win rate 90% → weight 0.1
- Win rate 50% → weight 0.5
- Win rate 10% → weight 0.9

A higher `power` value results in a more extreme distribution, emphasizing harder opponents more strongly.
```

--------------------------------

### Determine Network Layer Sizes Using Space Dimensions

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Use `flat_dim()` from the `Space` trait to dynamically size network layers based on environment observation and action spaces. This example shows how to get dimensions for discrete action spaces.

```rust
use rl4burn::space::Space;

let obs_dim = env.observation_space().flat_dim(); // e.g. 4 for CartPole
let n_actions = match env.action_space() {
    Space::Discrete(n) => n,
    Space::MultiDiscrete(dims) => dims.iter().sum(),
    _ => panic!("expected discrete"),
};

// Size linear layers from the environment directly
let fc1 = LinearConfig::new(obs_dim, 64).init(&device);
let out = LinearConfig::new(64, n_actions).init(&device);
```

--------------------------------

### Initialize and Use ReplayBuffer

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/replay-buffer.md

Demonstrates initializing a ReplayBuffer with a specific capacity and a seeded RNG for reproducibility. Shows methods for extending the buffer with transitions, sampling random batches (by reference and by clone), and grouping samples by a key.

```rust
use rand::SeedableRng;

let mut buffer = ReplayBuffer::new(10_000, rand::rngs::SmallRng::seed_from_u64(42));

buffer.extend(transitions);
let batch = buffer.sample(64);
let batch = buffer.sample_cloned(64);
let groups = buffer.group_by(|t| t.episode_id);
```

--------------------------------

### Environment Action Masking Example

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/algorithms/ppo.md

Example implementation of the action_mask method for an environment, returning a vector indicating valid actions.

```rust
fn action_mask(&self) -> Option<Vec<f32>> {
    let mut mask = vec![0.0; 15]; // 5 + 10
    for valid_type in &self.valid_action_types { mask[*valid_type] = 1.0; }
    for valid_target in &self.valid_targets { mask[5 + *valid_target] = 1.0; }
    Some(mask)
}
```

--------------------------------

### BetaVae Initialization and Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/beta-vae.md

Demonstrates how to initialize a BetaVae model with custom configurations and use it for training and inference.

```APIDOC
## BetaVae Initialization and Usage

### Description
Initializes a Beta-VAE model with specified dimensions and beta value, then shows how to perform forward passes for training and extract strategy embeddings for inference.

### Methods

- `BetaVaeConfig::new(obs_dim)`: Constructor for BetaVaeConfig.
- `with_latent_dim(latent_dim)`: Sets the dimension of the latent space.
- `with_beta(beta)`: Sets the beta value for the VAE.
- `init(device)`: Initializes the BetaVae model on the specified device.
- `forward(opponent_features)`: Performs a forward pass to get model output.
- `loss(opponent_features, &output)`: Calculates the VAE loss.
- `strategy_embedding(opponent_features)`: Extracts the strategy embedding from the input features.

### Parameters

#### Initialization Parameters
- **obs_dim** (usize): The dimensionality of the observation space.
- **latent_dim** (usize): The desired dimensionality of the latent space (default: 32).
- **beta** (f32): The beta value for the VAE, controlling disentanglement (default: 4.0).
- **device** (&Device): The computational device to initialize the model on.

#### Input Parameters
- **opponent_features** (Tensor): Input features representing opponent behavior.

### Request Example
```rust
use rl4burn::nn::vae::{BetaVae, BetaVaeConfig};

let vae = BetaVaeConfig::new(obs_dim)
    .with_latent_dim(32)
    .with_beta(4.0)
    .init(&device);

// Training
let output = vae.forward(opponent_features);
let loss = vae.loss(opponent_features, &output);

// Inference: extract strategy embedding
let z = vae.strategy_embedding(opponent_features);
// z: [batch, 32] — feed this as extra context to the policy
```

### Response

#### Success Response
- **output** (Tensor): The output of the VAE model during a forward pass.
- **loss** (Tensor): The calculated loss value.
- **z** (Tensor): The strategy embedding vector.
```

--------------------------------

### League Initialization and Agent Addition

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/league.md

Demonstrates how to initialize a new League, set an initial model, and add agents with specific roles and configurations.

```APIDOC
## League Initialization and Agent Addition

### Description
Initializes a new League instance, sets the starting model for all agents, and adds agents with specified roles and training configurations.

### Method
```rust
use rl4burn::{League, AgentRole, LeagueAgentConfig};

let mut league = League::new();
league.set_initial_model(initial_model.clone());

// Add agents
league.add_agent(model.clone(), LeagueAgentConfig {
    role: AgentRole::MainAgent,
    checkpoint_interval: 1000,
    reset_threshold: 0,
});
league.add_agent(model.clone(), LeagueAgentConfig {
    role: AgentRole::MainExploiter,
    checkpoint_interval: 2000,
    reset_threshold: 50000,
});
```

### Parameters
- `initial_model`: The initial model weights to set for the league.
- `model`: The model weights for the agent being added.
- `LeagueAgentConfig`:
  - `role`: The role of the agent (e.g., `AgentRole::MainAgent`, `AgentRole::MainExploiter`).
  - `checkpoint_interval`: The number of steps after which the agent's weights are checkpointed.
  - `reset_threshold`: The threshold for exploiter improvement before a reset is triggered.
```

--------------------------------

### MctsTree Initialization and Search

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/mcts.md

Demonstrates how to initialize an MCTS tree with specific configurations and perform a search to evaluate action sequences.

```APIDOC
## MctsTree Initialization and Search

### Description
Initializes an MCTS tree with a given configuration and performs a search using a provided evaluation function.

### Method
Rust (SDK)

### Initialization
```rust
use rl4burn::algo::planning::mcts::{MctsTree, MctsConfig};

let mut tree = MctsTree::new(MctsConfig {
    n_simulations: 800,
    exploration_constant: 1.41,
    n_actions: 30,  // number of possible picks
});
```

### Search
```rust
let visit_counts = tree.search(|action_path| {
    // Evaluate this sequence of picks.
    // Return estimated win rate (0.0 to 1.0).
    evaluate_composition(action_path)
}, &mut rng);
```

### Retrieving Best Action and Probabilities
```rust
let best_pick = tree.best_action();
let pick_probs = tree.action_probs();
```

### Parameters
#### MctsConfig
- **n_simulations** (usize) - The total number of simulations to run.
- **exploration_constant** (f64) - The constant used in the UCT formula for balancing exploration and exploitation.
- **n_actions** (usize) - The number of possible actions (e.g., picks) available at each step.

#### Search Closure
- **action_path** (Vec<Action>) - A sequence of actions taken so far.
- **Returns** (f64) - The estimated win rate (0.0 to 1.0) for the given action sequence.

### Return Value
- **visit_counts** (HashMap<Action, usize>) - A map of actions to their visit counts after the search.
- **best_pick** (Action) - The action with the highest visit count.
- **pick_probs** (HashMap<Action, f64>) - A map of actions to their probabilities based on visit counts.
```

--------------------------------

### SequenceReplayBuffer Initialization and Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/sequence-replay.md

Demonstrates how to initialize a SequenceReplayBuffer, add steps to it, and sample batches of sequences. The buffer is initialized with a capacity and a sequence length. Steps are added using the `push` method, and sequences are sampled using the `sample` method.

```APIDOC
## SequenceReplayBuffer

### Description
A FIFO buffer that samples contiguous sequences of a fixed length, respecting episode boundaries. Used by DreamerV3 for world model training.

### Initialization
```rust
use rl4burn::SequenceReplayBuffer;

let mut buffer = SequenceReplayBuffer::new(1_000_000, 64);
// capacity: 1M steps, sequence_length: 64
```

### Adding Transitions
```rust
use rl4burn::SequenceStep;

buffer.push(SequenceStep {
    observation: obs.clone(),
    action: vec![1.0, 0.0],
    reward: 1.0,
    done: false,
});
```

### Sampling Sequences
```rust
use rand::Rng;

let mut rng = rand::thread_rng();
let sequences = buffer.sample(16, &mut rng);
// sequences: Vec<Vec<SequenceStep<O>>>, each of length 64
```
```

--------------------------------

### Initialize and Use SequenceReplayBuffer

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/sequence-replay.md

Demonstrates how to initialize a SequenceReplayBuffer with a given capacity and sequence length, add transitions, and sample batches of sequences. Ensure to provide a random number generator for sampling.

```rust
use rl4burn::{SequenceReplayBuffer, SequenceStep};

let mut buffer = SequenceReplayBuffer::new(1_000_000, 64);
// capacity: 1M steps, sequence_length: 64

// Add transitions
buffer.push(SequenceStep {
    observation: obs.clone(),
    action: vec![1.0, 0.0],
    reward: 1.0,
    done: false,
});

// Sample batch of sequences
let sequences = buffer.sample(16, &mut rng);
// sequences: Vec<Vec<SequenceStep<O>>>, each of length 64
```

--------------------------------

### Initialize and Use PFSP Matchmaking

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/pfsp.md

Initializes PFSP matchmaking with a configuration, adds opponents, records match results, and samples an opponent based on the PFSP algorithm. Ensure 'rng' is a mutable random number generator.

```rust
use rl4burn::{PfspMatchmaking, PfspConfig};

let mut mm = PfspMatchmaking::new(PfspConfig {
    power: 1.0,    // higher = more focus on hard opponents
    min_prob: 0.01, // every opponent has at least 1% chance
});

mm.add_opponent(0);
mm.add_opponent(1);
mm.add_opponent(2);

// Record results
mm.record_result(0, true, false);   // beat opponent 0
mm.record_result(1, false, false);  // lost to opponent 1

// Sample: opponent 1 (harder) is sampled more often
let opponent = mm.sample_opponent(&mut rng);
```

--------------------------------

### CSPL Pipeline Initialization and Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cspl.md

Demonstrates how to initialize and use the CSPL pipeline for training AI agents. The pipeline progresses through different phases, and its state can be monitored.

```APIDOC
## CSPL Pipeline Usage

### Description
Initializes and runs the CSPL training pipeline.

### Method
Instantiation and method calls on `CsplPipeline`.

### Endpoint
N/A (Rust SDK)

### Parameters
#### `CsplConfig` struct:
- **phase1_steps** (u64) - Required - Number of steps for Phase 1 (Specialist Training).
- **phase2_steps** (u64) - Required - Number of steps for Phase 2 (Distillation).
- **phase3_steps** (u64) - Required - Number of steps for Phase 3 (Generalization).
- **n_specialists** (usize) - Required - Number of specialist models to train in Phase 1.

### Request Example
```rust
use rl4burn::{CsplPipeline, CsplConfig, CsplPhase};

let mut pipeline = CsplPipeline::new(CsplConfig {
    phase1_steps: 100_000,
    phase2_steps: 50_000,
    phase3_steps: 200_000,
    n_specialists: 10,
});

loop {
    let phase_changed = pipeline.step();

    match pipeline.current_phase() {
        CsplPhase::SpecialistTraining => { /* train specialists via self-play */ }
        CsplPhase::Distillation => { /* distill into student */ }
        CsplPhase::Generalization => { /* continue RL with random compositions */ }
    }

    if pipeline.is_complete() { break; }
}
```

### Response
#### `step()` method:
- Returns `bool`: `true` if the phase changed during this step, `false` otherwise.

#### `current_phase()` method:
- Returns `CsplPhase`: The current phase of the pipeline (`SpecialistTraining`, `Distillation`, or `Generalization`).

#### `is_complete()` method:
- Returns `bool`: `true` if the pipeline has completed all phases, `false` otherwise.

#### Example Usage:
See the Request Example for a full usage pattern.
```

--------------------------------

### Initialize and Use SyncVecEnv

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/concepts/vec-env.md

Demonstrates how to create and use SyncVecEnv with multiple instances of a CartPole environment. It shows resetting all environments and taking a step with actions.

```rust
use rl4burn::vec_env::SyncVecEnv;
use rl4burn::envs::CartPole;
use rand::SeedableRng;

let envs: Vec<CartPole<_>> = (0..8)
    .map(|i| CartPole::new(rand::rngs::SmallRng::seed_from_u64(i as u64)))
    .collect();
let mut vec_env = SyncVecEnv::new(envs);

// Reset all environments
let observations = vec_env.reset(); // Vec of 8 observations

// Step all environments with one action each
let actions = vec![0, 1, 0, 1, 1, 0, 1, 0];
let steps = vec_env.step(actions); // Vec of 8 Step results
```

--------------------------------

### PPO Continuous Actions Setup

Source: https://github.com/rpp1011/rl4burn/blob/master/README.md

Sets up the ActionDist for continuous actions. The MaskedActorCritic trait is implemented to return policy logits and value estimates.

```rust
use rl4burn::{ActionDist, LogStdMode, MaskedActorCritic, masked_ppo_collect, masked_ppo_update};

let action_dist = ActionDist::Continuous {
    action_dim: 1,
    log_std_mode: LogStdMode::ModelOutput,
};

// MaskedActorCritic returns (logits, values) — ActionDist handles the rest
impl<B: Backend> MaskedActorCritic<B> for ContinuousAgent<B> {
    fn forward(&self, obs: Tensor<B, 2>) -> (Tensor<B, 2>, Tensor<B, 1>) {
        let h = self.encoder.forward(obs);
        let logits = self.policy_head.forward(h.clone());   // [batch, 2] = [mean, log_std]
        let values = self.value_head.forward(h).squeeze(1);
        (logits, values)
    }
}
```

--------------------------------

### Initialize and Use PFSP Matchmaking

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Initializes PFSP matchmaking with configurable parameters for opponent selection bias and minimum probability. Add opponents and record results to influence future sampling.

```rust
use rl4burn::{PfspMatchmaking, PfspConfig};

let mut mm = PfspMatchmaking::new(PfspConfig {
    power: 1.0,     // higher = more skewed toward hard opponents
    min_prob: 0.01, // every opponent has ≥1% chance (prevents starvation)
});

mm.add_opponent(0);
mm.add_opponent(1);
mm.add_opponent(2);

mm.record_result(0, /*win=*/true,  /*draw=*/false); // beat opponent 0
mm.record_result(1, /*win=*/false, /*draw=*/false); // lost to opponent 1

// Opponent 1 (win_rate≈0%) gets sampled far more than opponent 0 (win_rate≈100%)
let opponent_id = mm.sample_opponent(&mut rng);
```

--------------------------------

### JSON Log Line Format

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/logging.md

Example of a single log entry in JSON Lines format, including type, key, value, step, and wall time.

```json
{"type":"scalar","key":"train/loss","value":0.42,"step":1000,"wall_time":1234567890.123}
```

--------------------------------

### Initialize and Use RSSM API

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/world-models/rssm.md

Demonstrates initializing the RSSM with a configuration and then using its methods for observation-based updates and imagination-based predictions. It also shows how to predict rewards and episode continuation.

```rust
use rl4burn::{Rssm, RssmConfig, RssmState};

let config = RssmConfig::new(obs_dim, action_dim);
let rssm = config.init(&device);

// Initial state (all zeros)
let state = rssm.initial_state(batch_size, &device);

// Training: observe → update
let (new_state, posterior_logits, prior_logits) = rssm.obs_step(&state, action, obs);

// Imagination: predict without observing
let new_state = rssm.imagine_step(&state, action);

// Predictions
let reward_logits = rssm.predict_reward(state.h, state.z);   // [batch, 255]
let cont_logits = rssm.predict_continue(state.h, state.z);   // [batch, 1]
```

--------------------------------

### Vast.ai Provider Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Demonstrates how to use the VastAiProvider to search for GPU offers, launch an instance, check its status, and stop it. Requires a VASTAI_API_KEY environment variable.

```rust
use rl4burn::cloud::{VastAiProvider, CloudProvider};

let provider = VastAiProvider::new(std::env::var("VASTAI_API_KEY").unwrap());

// Search for offers
let offers = provider.search_offers(&reqs)?;
println!("Found {} offers, cheapest: ${}/hr",
    offers.len(), offers[0].price_per_hour);

// Launch the cheapest offer
let instance = provider.launch(&offers[0])?;
println!("Instance {} status: {}", instance.instance_id, instance.status);

// Check status
let inst = provider.status(&instance.instance_id)?;
if let Some(ssh) = &inst.ssh_connection {
    println!("Connect: {}", ssh);
}

// Cleanup
provider.stop(&instance.instance_id)?;
```

--------------------------------

### LSTM, GRU, and Block GRU Cells

Source: https://context7.com/rpp1011/rl4burn/llms.txt

Provides examples for initializing and using LSTM, GRU, and Block GRU recurrent cells for processing sequences with hidden states.

```APIDOC
## LSTM, GRU, and Block GRU — Recurrent Cells

Recurrent cells for policies under partial observability. Block GRU (DreamerV3) reduces parameters by splitting the recurrent matrix into independent blocks.

```rust,ignore
use rl4burn::{ LstmCell, LstmCellConfig, LstmState, GruCell, GruCellConfig,
              BlockGruCell, BlockGruCellConfig};

// LSTM: 3-gate cell with cell and hidden state
let lstm = LstmCellConfig::new(/*input=*/64, /*hidden=*/256).init(&device);
let state = LstmState::zeros(batch_size, 256, &device);
let new_state = lstm.forward(input, &state);                     // single step
let (outputs, final_state) = lstm.forward_seq(inputs, &state);  // full sequence
// outputs: [batch, seq_len, hidden_size]

// GRU: 2-gate cell, simpler than LSTM
let gru = GruCellConfig::new(64, 256).init(&device);
let h = Tensor::zeros([batch_size, 256], &device);
let new_h = gru.forward(input, h);

// Block GRU (DreamerV3): n_blocks=8 reduces recurrent params by 8x
// DreamerV3 uses hidden=4096, n_blocks=8 → 16M params → 2M params
let block_gru = BlockGruCellConfig::new(64, 4096)
    .with_n_blocks(8)
    .init(&device);
// API identical to GruCell; n_blocks=1 is standard GRU
```
```

--------------------------------

### Initialize and Use Beta-VAE Model

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/beta-vae.md

Configure and initialize a Beta-VAE model for opponent behavior prediction. Use the `forward` method for training and `strategy_embedding` for inference to extract latent strategy embeddings.

```rust
use rl4burn::nn::vae::{BetaVae, BetaVaeConfig};

let vae = BetaVaeConfig::new(obs_dim)
    .with_latent_dim(32)
    .with_beta(4.0)
    .init(&device);

// Training
let output = vae.forward(opponent_features);
let loss = vae.loss(opponent_features, &output);

// Inference: extract strategy embedding
let z = vae.strategy_embedding(opponent_features);
// z: [batch, 32] — feed this as extra context to the policy
```

--------------------------------

### Calling clip_grad_norm between backward and step

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/grad-clip.md

This example shows how to use `clip_grad_norm` in the training loop, specifically after calling `backward()` and before `optim.step()`. Set `max_grad_norm` to 0.0 to disable clipping.

```rust
let grads = loss.backward();
let mut grads = GradientsParams::from_grads(grads, &model);
grads = clip_grad_norm(&model, grads, 0.5);  // max_norm = 0.5
model = optim.step(lr, model, grads);
```

--------------------------------

### Get Opponent and Update Agent in Training Loop

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/league.md

Retrieves an opponent for a given agent during the training loop and updates the agent's status, which handles checkpointing. Requires a random number generator `rng`.

```rust
// Training loop
let opponent = league.get_opponent(agent_idx, &mut rng);
// ... play game, update model ...
league.update_agent(agent_idx); // handles checkpointing
```

--------------------------------

### SelfPlayPool Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/self-play.md

Demonstrates how to initialize a SelfPlayPool, add model snapshots, sample opponents, and retain a specific number of recent snapshots.

```APIDOC
## SelfPlayPool

### Description
Manages a collection of past model versions (snapshots) for self-play training.

### Methods

#### `new()`
- **Description**: Creates a new, empty `SelfPlayPool`.
- **Usage**: `let mut pool = SelfPlayPool::new();`

#### `add_snapshot(&self, model: &Model, training_step: usize)`
- **Description**: Adds a deep copy (clone) of the current model as a snapshot to the pool.
- **Parameters**:
  - `model` (Model): The model to snapshot.
  - `training_step` (usize): The training step associated with this snapshot.
- **Usage**: `pool.add_snapshot(&model, training_step);`

#### `sample<R: Rng>(&self, rng: &mut R) -> Option<Model>`
- **Description**: Samples a random past opponent model from the pool.
- **Parameters**:
  - `rng` (&mut R): A mutable reference to a random number generator.
- **Returns**: An `Option<Model>` containing a cloned opponent model if the pool is not empty, otherwise `None`.
- **Usage**: `if let Some(opponent) = pool.sample(&mut rng) { ... }`

#### `retain_recent(&mut self, count: usize)`
- **Description**: Retains only the `count` most recent snapshots in the pool, discarding older ones.
- **Parameters**:
  - `count` (usize): The maximum number of recent snapshots to keep.
- **Usage**: `pool.retain_recent(50);`

### Important Note
Snapshots are deep copies (`.clone()`'d). Mutating the original model after adding a snapshot does not affect the stored snapshot.
```

--------------------------------

### RunPod Provider Usage

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Demonstrates how to use the RunPodProvider to search for GPU offers, launch an instance, check its status, and stop it. Requires a RUNPOD_API_KEY environment variable.

```rust
use rl4burn::cloud::{RunPodProvider, CloudProvider};

let provider = RunPodProvider::new(std::env::var("RUNPOD_API_KEY").unwrap());

let offers = provider.search_offers(&reqs)?;
let instance = provider.launch(&offers[0])?;
println!("Pod {} status: {}", instance.instance_id, instance.status);

provider.stop(&instance.instance_id)?;
```

--------------------------------

### Polyak Update Usage Examples

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/building-blocks/polyak.md

Demonstrates how to use the `polyak_update` function for both hard and soft target updates. A hard update replaces the target network entirely, while a soft update gradually tracks the source network.

```rust
use rl4burn::polyak::polyak_update;

// Hard target update (DQN-style, every N steps)
if step % 250 == 0 {
    target = polyak_update(target, &online, 1.0);
}

// Soft target update (SAC/TD3-style, every step)
target = polyak_update(target, &online, 0.005);
```

--------------------------------

### Custom HTTP Function Integration

Source: https://github.com/rpp1011/rl4burn/blob/master/book/src/game-ai/cloud-deploy.md

Example of integrating a custom HTTP client function with a cloud provider. The `with_http` method allows you to provide your own function to handle HTTP requests, enabling the use of clients like reqwest or ureq.

```rust
fn my_http(req: &rl4burn::cloud::vastai::HttpRequest) -> Result<String, String> {
    // Use reqwest, ureq, curl, etc.
    todo!()
}

let provider = VastAiProvider::new("key").with_http(my_http);
```