RL4Burn

https://docs.rs/rl4burn/0.1.1
Admin
RL4Burn provides generic, backend-agnostic reinforcement learning algorithms and building blocks for...
Tokens:288,507
Snippets:2,837
Trust Score:9.1
Update:5 months ago
Show doc for...
Tokens:
Raw
### Example QNetwork Implementation

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/dqn/trait.QNetwork.html

Provides an example implementation of the QNetwork trait using a simple feedforward neural network. Requires the `Module` derive macro.

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

impl<B: Backend> QNetwork<B> for MyQNet<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)
    }
}
```

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

### CartPole Environment Setup

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/envs/cartpole.rs.html

Demonstrates how to create and initialize the CartPole environment. Ensure the 'cartpole' feature is enabled.

```rust
use rl4burn::envs::cartpole::CartPole;

fn main() {
    let env = CartPole::new();
    // Environment is now ready for use
}
```

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

### PfspMatchmaking struct and new method

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/pfsp.rs.html

Initializes the PFSP matchmaking system with a given configuration. The matchmaking pool starts empty.

```rust
pub struct PfspMatchmaking {
    config: PfspConfig,
    /// Win rates of the current agent against each opponent.
    records: Vec<PlayerRecord>,
}

impl PfspMatchmaking {
    pub fn new(config: PfspConfig) -> Self {
        Self {
            config,
            records: Vec::new(),
        }
    }

    /// Add a new opponent to the pool.
    pub fn add_opponent(&mut self, id: u64) {
        self.records.push(PlayerRecord::new(id));
    }

    /// Record game result (from perspective of current agent).
    pub fn record_result(&mut self, opponent_id: u64, win: bool, draw: bool) {
        if let Some(record) = self.records.iter_mut().find(|r| r.id == opponent_id) {
            if draw {
                record.draws += 1;
            } else if win {
                record.wins += 1;
            } else {
                record.losses += 1;
            }
        }
    }

    /// Sample an opponent using PFSP weighting.
    /// Harder opponents (lower win rate) are sampled more frequently.
    pub fn sample_opponent(&self, rng: &mut impl Rng) -> Option<u64> {
        if self.records.is_empty() {
            return None;
        }

        let weights: Vec<f64> = self
            .records
            .iter()
            .map(|r| {
                (1.0 - r.win_rate())
                    .powf(self.config.power)
                    .max(self.config.min_prob)
            })
            .collect();

        let total: f64 = weights.iter().sum();
        let probs: Vec<f64> = weights.iter().map(|w| w / total).collect();

        let u: f64 = rng.random::<f64>();
        let mut cum = 0.0;
        for (i, &p) in probs.iter().enumerate() {
            cum += p;
            if u < cum {
                return Some(self.records[i].id);
            }
        }
        Some(self.records.last().unwrap().id)
    }

    /// Get all records.
    pub fn records(&self) -> &[PlayerRecord] {
        &self.records
    }

    /// Get selection probabilities.
    pub fn selection_probs(&self) -> Vec<f64> {
        if self.records.is_empty() {
            return vec![];
        }
        let weights: Vec<f64> = self
            .records
            .iter()
            .map(|r| {
                (1.0 - r.win_rate())
                    .powf(self.config.power)

```

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

### SequenceReplayBuffer::new

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/collect/sequence_replay.rs.html

Creates a new SequenceReplayBuffer with specified capacity and sequence length. Initializes episode starts.

```rust
/// Create a new buffer with the given capacity and sequence length.
pub fn new(capacity: usize, sequence_length: usize) -> Self {
    Self {
        steps: Vec::with_capacity(capacity.min(1024)),
        episode_starts: vec![0], // first episode starts at 0
        capacity,
        sequence_length,
    }
}

```

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

### Local Sync Sum Strategy Example

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/distributed.rs.html

Demonstrates the LocalSync all_reduce_f32 method with the Sum reduction strategy. Verifies that summing the provided values results in the original values.

```rust
#[test]
    fn local_sync_sum_strategy() {
        let sync = LocalSync;
        let values = vec![10.0, 20.0];
        let result = sync.all_reduce_f32(&values, ReduceStrategy::Sum);
        assert_eq!(result, values);
    }
```

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

### Create a SelfPlayPool

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/self_play.rs.html

Initializes an empty pool for storing policy snapshots. Use this when starting a new self-play training session.

```rust
pub fn new() -> Self {
        Self {
            snapshots: Vec::new(),
            next_id: 0,
        }
    }
```

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

### Create a New League

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Instantiates a new League object. Use this to start the league training orchestrator.

```rust
pub fn new() -> Self
```

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

### Example Implementation of DiscreteActorCritic

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/policy/trait.DiscreteActorCritic.html

An example demonstrating how to implement the `DiscreteActorCritic` trait for a custom actor-critic model using linear layers.

```APIDOC
## Example Implementation
```rust
#[derive(Module)]
struct MyActorCritic<B: Backend> {
    shared: Linear<B>,
    policy_head: Linear<B>,
    value_head: Linear<B>,
}

impl<B: Backend> DiscreteActorCritic<B> for MyActorCritic<B> {
    fn forward(&self, obs: Tensor<B, 2>) -> DiscreteAcOutput<B> {
        let h = self.shared.forward(obs);
        DiscreteAcOutput {
            logits: self.policy_head.forward(h.clone()),
            values: self.value_head.forward(h).squeeze(1),
        }
    }
}
```
```

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

### PfspMatchmaking::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/pfsp/struct.PfspMatchmaking.html

Initializes a new PfspMatchmaking instance with the provided configuration.

```APIDOC
## PfspMatchmaking::new

### Description
Initializes a new PfspMatchmaking instance with the provided configuration.

### Signature
```rust
pub fn new(config: PfspConfig) -> Self
```

### Parameters
* `config` (PfspConfig) - The configuration for the matchmaking.

### Returns
* `Self` - A new instance of PfspMatchmaking.
```

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

### Initialize BetaVae with Configuration

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/vae/struct.BetaVaeConfig.html

Initializes a Beta-VAE model on a specified device using the current BetaVaeConfig.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> BetaVae<B>
```

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

### SequenceReplayBuffer::valid_sequence_starts

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/collect/sequence_replay.rs.html

Helper function to find all valid starting indices for sequences. A valid start index ensures the sequence of `sequence_length` stays within a single episode.

```rust
/// Find all valid starting indices for sequences.
///
/// A valid start index `i` means `steps[i..i+sequence_length]` are all
/// within the same episode.
fn valid_sequence_starts(&self) -> Vec<usize> {
    if self.steps.len() < self.sequence_length {
        return vec![];
    }

    // Build episode ranges
    let mut sorted_starts = self.episode_starts.clone();
    sorted_starts.sort();
    sorted_starts.dedup();

    let mut valid = Vec::new();
    for i in 0..sorted_starts.len() {
        let start = sorted_starts[i];
        let end = if i + 1 < sorted_starts.len() {
            sorted_starts[i + 1]
        } else {
            self.steps.len()
        };
        if end - start >= self.sequence_length {
            for s in start..=end - self.sequence_length {

```

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

### Example Implementation of DiscreteActorCritic

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/policy/trait.DiscreteActorCritic.html

An example struct implementing the DiscreteActorCritic trait. It uses linear layers for shared feature extraction, policy head, and value head, demonstrating a typical actor-critic architecture.

```rust
#[derive(Module)]
struct MyActorCritic<B: Backend> {
    shared: Linear<B>,
    policy_head: Linear<B>,
    value_head: Linear<B>,
}

impl<B: Backend> DiscreteActorCritic<B> for MyActorCritic<B> {
    fn forward(&self, obs: Tensor<B, 2>) -> DiscreteAcOutput<B> {
        let h = self.shared.forward(obs);
        DiscreteAcOutput {
            logits: self.policy_head.forward(h.clone()),
            values: self.value_head.forward(h).squeeze(1),
        }
    }
}
```

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

### CountBasedReward Implementation

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/intrinsic/trait.IntrinsicReward.html

Example implementation of the IntrinsicReward trait for CountBasedReward.

```APIDOC
## impl IntrinsicReward for CountBasedReward

### Associated Types

#### type Observation = Vec<f32>

Observation type for CountBasedReward.
```

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

### BetaVaeConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/vae/struct.BetaVaeConfig.html

Initializes a Beta-VAE model on a specified device using the current configuration.

```APIDOC
## `init` Function

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> BetaVae<B>
```

### Description
Initialize the Beta-VAE on the given device.

### Parameters
* `device: &B::Device` - The device on which to initialize the Beta-VAE.
```

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

### type_id

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.AttentionPool.html

Gets the `TypeId` of `self`. This method is part of the `Any` trait implementation.

```APIDOC
## fn type_id(&self) -> TypeId

### Description
Gets the `TypeId` of `self`.

### Method
`type_id`

### Parameters
None

### Response
- **TypeId**: The type identifier of the object.
```

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

### Get Number of Agents

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/league.rs.html

Returns the total number of agents currently in the league.

```APIDOC
## n_agents

### Description
Returns the total count of agents participating in the league.

### Method Signature
`pub fn n_agents(&self) -> usize`

### Returns
The number of agents in the league.
```

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

### GridWorld Environment Initialization

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/envs/gridworld.rs.html

Initializes a new GridWorld environment with a given random number generator. The goal is fixed, and the agent's starting position is randomized upon creation.

```rust
pub fn new(rng: R) -> Self {
        let mut env = Self {
            agent: (0, 0),
            goal: (GRID_SIZE - 1, GRID_SIZE - 1),
            step_count: 0,
            rng,
        };
        env.reset();
        env
    }
```

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

### Get Agent Reference

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/league.rs.html

Provides read-only access to a specific agent in the league.

```APIDOC
## agent

### Description
Returns an immutable reference to the agent at the specified index.

### Method Signature
`pub fn agent(&self, idx: usize) -> &LeagueAgent<M>`

### Parameters
* **idx** (usize) - The index of the agent to retrieve.

### Returns
An immutable reference to the `LeagueAgent` at the given index.
```

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

### League::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Creates a new, empty League instance.

```APIDOC
## League::new

### Description
Create a new league.

### Method
`new()`

### Returns
- `Self`: A new instance of the League.
```

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

### Get Current CSPL Phase

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/cspl.rs.html

Returns the current phase of the CSPL pipeline.

```rust
    /// Current phase of the pipeline.
    pub fn current_phase(&self) -> CsplPhase {
        self.current_phase
    }
```

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

### BetaVaeConfig Initialization

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/vae.rs.html

Initializes the Beta-VAE configuration with specified dimensions and beta value. Requires a backend device.

```rust
impl BetaVaeConfig {
    /// Initialize the Beta-VAE on the given device.
    pub fn init<B: Backend>(&self, device: &B::Device) -> BetaVae<B> {
        let encoder = VaeEncoder {
            fc1: LinearConfig::new(self.input_dim, self.hidden_dim).init(device),
            fc_mu: LinearConfig::new(self.hidden_dim, self.latent_dim).init(device),
            fc_logvar: LinearConfig::new(self.hidden_dim, self.latent_dim).init(device),
        };
        let decoder = VaeDecoder {
            fc1: LinearConfig::new(self.latent_dim, self.hidden_dim).init(device),
            fc_out: LinearConfig::new(self.hidden_dim, self.input_dim).init(device),
        };
        BetaVae {
            encoder,
            decoder,
            beta: self.beta,
            latent_dim: self.latent_dim,
        }
    }
}
```

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

### Test: Pool Starts Empty

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/self_play.rs.html

Verifies that a newly created SelfPlayPool is initially empty.

```rust
let pool: SelfPlayPool<i32> = SelfPlayPool::new();
        assert!(pool.is_empty());
```

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

### TargetAttentionConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TargetAttentionConfig.html

Initializes a TargetAttention module on the given device using the configuration.

```APIDOC
### impl TargetAttentionConfig

#### pub fn init<B: Backend>(&self, device: &B::Device) -> TargetAttention<B>
Initialize a target attention module on the given device.
The projection dimension is `min(query_dim, key_dim)`.
```

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

### Get Number of Frozen Snapshots

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/league.rs.html

Returns the count of models stored in the frozen pool.

```APIDOC
## n_frozen

### Description
Returns the number of model snapshots currently stored in the frozen pool. These are typically checkpoints of agents' models.

### Method Signature
`pub fn n_frozen(&self) -> usize`

### Returns
The count of frozen model snapshots.
```

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

### RewardClip::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/env/wrapper/struct.RewardClip.html

Constructs a new RewardClip instance.

```APIDOC
## pub fn new(inner: E, limit: f32) -> Self

### Description
Constructs a new RewardClip instance.

### Parameters
#### Path Parameters
- **inner** (E) - The inner environment.
- **limit** (f32) - The limit for clipping rewards.
```

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

### TargetAttentionConfig::init

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/attention.rs.html

Initializes a TargetAttention module with specified dimensions and backend.

```APIDOC
## TargetAttentionConfig::init<B: Backend>

### Description
Initialize a target attention module on the given device.

The projection dimension is `min(query_dim, key_dim)`.

### Parameters
* `self`: `&TargetAttentionConfig` - The configuration for the target attention.
* `device`: `&B::Device` - The device on which to initialize the module.

### Returns
`TargetAttention<B>` - An initialized TargetAttention module.
```

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

### Get Size of SelfPlayPool

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/self_play.rs.html

Returns the number of policy snapshots currently stored in the pool.

```rust
pub fn len(&self) -> usize {
        self.snapshots.len()
    }
```

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

### Run Greedy Action and Step Environment

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/policy/fn.greedy_action.html

Use this function to get the greedy action from a model for a single observation and then step the environment with that action. It's suitable for inference and evaluation loops.

```rust
let action = greedy_action(&model, &obs, &device);
let step = env.step(action);
```

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

### EntropyReductionReward::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/intrinsic/struct.EntropyReductionReward.html

Creates a new instance of the EntropyReductionReward module. This is the entry point for initializing the reward calculation system.

```APIDOC
## EntropyReductionReward::new

### Description
Create a new entropy-reduction reward module.

### Method
`new()`

### Returns
- `Self`: A new instance of `EntropyReductionReward`.
```

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

### Get Number of Agents

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Returns the total number of agents currently participating in the league.

```rust
pub fn n_agents(&self) -> usize
```

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

### Initialize CSPL Pipeline

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/cspl.rs.html

Creates a new CSPL pipeline instance, starting at the SpecialistTraining phase.

```rust
impl CsplPipeline {
    /// Create a new pipeline starting at Phase 1 (SpecialistTraining).
    pub fn new(config: CsplConfig) -> Self {
        Self {
            config,
            current_phase: CsplPhase::SpecialistTraining,
            phase_step: 0,
        }
    }
// ... other methods
```

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

### Initialize ZConditioning Module

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/z_conditioning/struct.ZConditioningConfig.html

Initializes a z-conditioning module on a specified device using the current configuration. Requires a Backend and its associated Device.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> ZConditioning<B>
```

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

### Initialize EpisodeStats Wrapper

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/env/wrapper.rs.html

Creates a new `EpisodeStats` wrapper around an existing environment to track cumulative episode reward and length. Initialize with `EpisodeStats::new(inner_env)`.

```rust
pub fn new(inner: E) -> Self {
        Self {
            inner,
            episode_reward: 0.0,
            episode_length: 0,
            last_episode_reward: None,
            last_episode_length: None,
        }
    }
```

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

### Get Agent Indices by Role

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/league.rs.html

Retrieves a list of agent indices that match a specified role.

```APIDOC
## agents_with_role

### Description
Filters and returns the indices of all agents that are assigned a specific role within the league.

### Method Signature
`pub fn agents_with_role(&self, role: AgentRole) -> Vec<usize>`

### Parameters
* **role** (AgentRole) - The role to filter agents by.

### Returns
A vector containing the indices of agents matching the specified role.
```

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

### new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/envs/pendulum/struct.Pendulum.html

Creates a new Pendulum environment instance.

```APIDOC
## fn new<R: Rng>(rng: R) -> Self

### Description
Create a new Pendulum environment with the given RNG.

### Parameters
- **rng** (R): A random number generator implementing the `Rng` trait.
```

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

### Get Mutable Agent Model Reference

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/league.rs.html

Provides mutable access to an agent's model.

```APIDOC
## agent_model_mut

### Description
Returns a mutable reference to the model of the agent at the specified index, allowing for direct modification.

### Method Signature
`pub fn agent_model_mut(&mut self, idx: usize) -> &mut M`

### Parameters
* **idx** (usize) - The index of the agent whose model should be accessed.

### Returns
A mutable reference to the agent's model.
```

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

### PointerNetConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.PointerNetConfig.html

Initializes a PointerNet module on a given device using the configuration.

```APIDOC
## PointerNetConfig::init

### Description
Initialize a pointer network module on the given device.

### Signature
```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> PointerNet<B>
```

### Parameters
* `device` (&B::Device) - The device to initialize the PointerNet on.
```

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

### Get Owned Version

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.AttentionPoolConfig.html

Returns the owned version of the value. This is part of the ToOwned trait implementation.

```rust
type Owned = T
```

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

### Initialize RewardClip

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/env/wrapper/struct.RewardClip.html

Creates a new RewardClip instance with a given environment and a reward limit.

```rust
pub fn new(inner: E, limit: f32) -> Self>
```

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

### Get ZConditioning Devices

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/z_conditioning/struct.ZConditioning.html

Returns a vector of all unique devices on which the ZConditioning module and its sub-modules are located.

```rust
fn devices(&self) -> Vec<<B as Backend>::Device>
```

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

### PPO Training Loop with Gradient Clipping and KL Divergence Early Stopping

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/ppo.rs.html

Implements the PPO training loop, including backward pass, gradient clipping by global norm, optimizer step, and target KL divergence early stopping.

```rust
            total_entropy += entropy.clone().into_data().to_vec::<f32>().unwrap()[0];
            total_kl += approx_kl_val;
            n_updates += 1;

            // Backward pass with global gradient norm clipping
            let grads = loss.backward();
            let mut grads = GradientsParams::from_grads(grads, &model);
            if config.max_grad_norm > 0.0 {
                grads = clip_grad_norm(&model, grads, config.max_grad_norm);
            }
            model = optim.step(current_lr, model, grads);
        }

        // Target KL early stopping (matches CleanRL)
        if let Some(target) = config.target_kl {
            if n_updates > 0 && (total_kl / n_updates as f32) > target {
                break;
            }
        }
    }

    let n = n_updates.max(1) as f32;
    let stats = PpoStats {
        policy_loss: total_policy_loss / n,
        value_loss: total_value_loss / n,
        entropy: total_entropy / n,
        approx_kl: total_kl / n,
    };

    (model, stats)
}
```

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

### PercentileNormalizer::scale

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/percentile_normalize/struct.PercentileNormalizer.html

Gets the current normalization scale, calculated as max(1.0, ema_high - ema_low).

```APIDOC
## PercentileNormalizer::scale

### Description
Current normalization scale: `max(1.0, ema_high - ema_low)`. Returns `1.0` before the first `update` call.

### Returns
- `f32`: The current normalization scale.
```

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

### AttentionPoolConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.AttentionPoolConfig.html

Initializes an AttentionPool module on the specified device using the current configuration.

```APIDOC
## pub fn init<B: Backend>(&self, device: &B::Device) -> AttentionPool<B>

Initialize an attention pool module on the given device.
```

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

### Get Type ID

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/intrinsic/struct.CountBasedReward.html

Returns the `TypeId` of the current object. This is a blanket implementation for any type `T` that is `'static`.

```rust
fn type_id(&self) -> TypeId>
```

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

### CsplPipeline::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/cspl/struct.CsplPipeline.html

Creates a new CsplPipeline instance, initializing it to Phase 1 (SpecialistTraining).

```APIDOC
## CsplPipeline::new

### Description
Create a new pipeline starting at Phase 1 (SpecialistTraining).

### Signature
```rust
pub fn new(config: CsplConfig) -> Self
```
```

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

### Initialize SequenceReplayBuffer and Check Capacity

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/collect/sequence_replay.rs.html

Demonstrates how to create a new SequenceReplayBuffer with a specified capacity and number of elements, and then asserts its capacity.

```rust
let mut buf: SequenceReplayBuffer<i32> = SequenceReplayBuffer::new(42, 5);
assert_eq!(buf.capacity(), 42);
```

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

### Get Number of Nodes in MCTS Tree

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/mcts.rs.html

Returns the total number of nodes currently in the MCTS tree.

```rust
pub fn n_nodes(&self) -> usize {
    self.nodes.len()
}

```

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

### TargetAttentionConfig::load_binary

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TargetAttentionConfig.html

Loads the configuration from a binary buffer.

```APIDOC
### impl Config for TargetAttentionConfig

#### fn load_binary(data: &[u8]) -> Result<Self, ConfigError>
Loads the configuration from a binary buffer.
```

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

### Get number of action heads

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/autoregressive/struct.CompositeDistribution.html

Returns the number of action dimensions, which corresponds to the number of heads in the distribution.

```rust
pub fn n_heads(&self) -> usize
```

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

### FilmConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/film/struct.FilmConfig.html

Initializes a Film layer on a given device using the current configuration.

```APIDOC
### impl FilmConfig

#### pub fn init<B: Backend>(&self, device: &B::Device) -> Film<B>

Initialize a FiLM layer on the given device.
```

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

### Get total number of logits

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/autoregressive/struct.CompositeDistribution.html

Calculates the total number of logits required across all action heads.

```rust
pub fn total_logits(&self) -> usize
```

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

### Get the maximum capacity of SequenceReplayBuffer

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/sequence_replay/struct.SequenceReplayBuffer.html

Returns the maximum number of steps the buffer can hold. This is determined at initialization.

```rust
pub fn capacity(&self) -> usize
```

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

### SequenceReplayBuffer Initialization and Basic Operations

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/collect/sequence_replay.rs.html

Demonstrates the creation of a SequenceReplayBuffer and its basic operations like push, is_empty, and len. The buffer has a maximum capacity and a sequence length.

```rust
impl<T: Clone> SequenceReplayBuffer<T> {
    /// Creates a new SequenceReplayBuffer.
    ///
    /// # Arguments
    ///
    /// * `capacity` - Maximum number of steps to store.
    /// * `seq_len` - The length of sequences to sample.
    pub fn new(capacity: usize, seq_len: usize) -> Self {
        SequenceReplayBuffer {
            steps: Vec::with_capacity(capacity),
            capacity,
            seq_len,
        }
    }

    /// Adds a step to the buffer. If the buffer is full, the oldest step is evicted.
    pub fn push(&mut self, step: SequenceStep<T>) {
        if self.steps.len() == self.capacity {
            self.steps.remove(0);
        }
        self.steps.push(step);
    }

    /// Extends the buffer with a collection of steps.
    pub fn extend<I: IntoIterator<Item = SequenceStep<T>>>(&mut self, iter: I) {
        for step in iter {
            self.push(step);
        }
    }

    /// Samples a batch of sequences from the buffer.
    ///
    /// # Arguments
    ///
    /// * `batch_size` - The number of sequences to sample.
    /// * `rng` - A random number generator.
    ///
    /// # Returns
    ///
    /// A vector of sampled sequences. Returns an empty vector if the buffer does not contain enough steps to form a sequence of `seq_len`.
    pub fn sample<R: Rng + ?Sized>(&self, batch_size: usize, rng: &mut R) -> Vec<Vec<SequenceStep<T>>> {
        if self.steps.len() < self.seq_len {
            return Vec::new();
        }

        let mut sequences = Vec::with_capacity(batch_size);
        let mut valid_starts = Vec::new();

        // Collect all possible start indices for sequences, respecting episode boundaries.
        for i in 0..=(self.steps.len() - self.seq_len) {
            // Check if the sequence starting at `i` crosses an episode boundary.
            let mut crosses_boundary = false;
            for j in 0..self.seq_len {
                if self.steps[i + j].done {
                    // If a step within the sequence is marked as done, and it's not the last step of the sequence,
                    // then this sequence crosses an episode boundary.
                    if j < self.seq_len - 1 {
                        crosses_boundary = true;
                        break;
                    }
                }
            }

            if !crosses_boundary {
                valid_starts.push(i);
            }
        }

        if valid_starts.is_empty() {
            return Vec::new();
        }

        // Sample batch_size start indices with replacement.
        for _ in 0..batch_size {
            let start_index = valid_starts[rng.gen_range(0..valid_starts.len())];
            let sequence = self.steps[start_index..start_index + self.seq_len].to_vec();
            sequences.push(sequence);
        }

        sequences
    }

    /// Number of steps currently stored.
    pub fn len(&self) -> usize {
        self.steps.len()
    }

    /// Whether the buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.steps.is_empty()
    }

    /// Maximum capacity of the buffer.
    pub fn capacity(&self) -> usize {
        self.capacity
    }
}

```

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

### Get Number of Visited States

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/collect/intrinsic/struct.CountBasedReward.html

Returns the total number of distinct states that have been visited by the CountBasedReward module.

```rust
pub fn num_visited(&self) -> usize>
```

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

### Create BetaVaeConfig Instance

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/vae/struct.BetaVaeConfig.html

Initializes a new BetaVaeConfig with the specified input dimension. Other parameters can be chained.

```rust
pub fn new(input_dim: usize) -> Self
```

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

### Get Number of Snapshots in SelfPlayPool

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/self_play/struct.SelfPlayPool.html

Returns the total number of policy snapshots currently stored in the pool.

```rust
pub fn len(&self) -> usize
```

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

### Sample and Calculate Log Probability for Continuous Actions

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/dist.rs.html

Demonstrates sampling an action from a continuous distribution and then calculating the log probability of that sampled action. This is crucial for policy gradient methods.

```rust
let dist = ActionDist::Continuous {
            action_dim: 2,
            log_std_mode: LogStdMode::ModelOutput,
        };
        let device = device();
        let mut rng = rand::rngs::SmallRng::seed_from_u64(42);

        // means=[1.0, -1.0], log_stds=[0.0, 0.0] → stds=[1.0, 1.0]
        let logits: Tensor<B, 2> = Tensor::from_data(
            TensorData::new(vec![1.0f32, -1.0, 0.0, 0.0], [1, 4]),
            &device,
        );

        let actions = dist.sample(&logits, None, None, &mut rng);
        assert_eq!(actions.len(), 1);
        assert_eq!(actions[0].len(), 2);

        // Log prob of the sampled action
        let lp: f32 = dist
            .log_prob(
                Tensor::<B, 2>::from_data(
                    TensorData::new(vec![1.0f32, -1.0, 0.0, 0.0], [1, 4]),
                    &device,
                ),
                &actions,
                None,
                None,
                &device,
            )
            .into_data()
            .to_vec()
            .unwrap()[0];

        // Should be finite and negative
        assert!(lp.is_finite(), "log_prob should be finite");
        assert!(lp <= 0.0, "log_prob should be <= 0");
```

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

### Initialize TargetAttention Module

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TargetAttentionConfig.html

Initializes a TargetAttention module on a given device using the configuration. The projection dimension is determined by the minimum of query_dim and key_dim.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> TargetAttention<B>
```

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

### Get Number of Frozen Snapshots

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Returns the count of policy snapshots that have been frozen and are available for opponents to play against.

```rust
pub fn n_frozen(&self) -> usize
```

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

### Get Number of Nodes in MctsTree

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/mcts/struct.MctsTree.html

Returns the total number of nodes currently present in the MCTS tree.

```rust
pub fn n_nodes(&self) -> usize
```

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

### Initialize AttentionPool Module

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.AttentionPoolConfig.html

Initializes an AttentionPool module on a specified device using the current configuration. Requires a Backend and a Device.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> AttentionPool<B>
```

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

### TargetAttentionConfig::load

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TargetAttentionConfig.html

Loads the configuration from a file.

```APIDOC
### impl Config for TargetAttentionConfig

#### fn load<P>(file: P) -> Result<Self, ConfigError>
where P: AsRef<Path>,
Loads the configuration from a file.
```

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

### Test CSPL Pipeline Initialization

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/cspl.rs.html

Verifies that a new CSPL pipeline starts in the 'SpecialistTraining' phase with step 0.

```rust
#[test]
    fn starts_in_specialist_training() {
        let pipeline = CsplPipeline::new(default_config());
        assert_eq!(pipeline.current_phase(), CsplPhase::SpecialistTraining);
        assert_eq!(pipeline.phase_step(), 0);
    }
```

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

### MultiHeadAttentionConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.MultiHeadAttentionConfig.html

Initializes a multi-head attention module on the given device.

```APIDOC
## pub fn init<B: Backend>(&self, device: &B::Device) -> MultiHeadAttention<B>

### Description
Initialize a multi-head attention module on the given device.

### Parameters
- **device** (*B::Device*) - Required - The device to initialize the module on.

### Panics
Panics if `d_model` is not divisible by `n_heads`.

### Returns
- MultiHeadAttention<B> - An initialized multi-head attention module.
```

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

### Initialize RewardClip Wrapper

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/env/wrapper.rs.html

Creates a `RewardClip` wrapper that limits rewards to the range `[-limit, limit]`. The `limit` must be a positive value.

```rust
pub fn new(inner: E, limit: f32) -> Self {
        assert!(limit > 0.0, "limit must be positive");
        Self { inner, limit }
    }
```

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

### Get Bin Centers

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/symlog/struct.TwohotEncoder.html

Calculates and returns the centers of the bins as a 1-D float tensor. The shape of the tensor is `[n_bins]`.

```rust
pub fn bin_centers<B: Backend>(&self, device: &B::Device) -> Tensor<B, 1>
```

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

### Get RewardClip Action Space

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/env/wrapper/struct.RewardClip.html

Returns the description of the action space for the RewardClip environment. This delegates to the inner environment.

```rust
fn action_space(&self) -> Space>
```

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

### Load BetaVaeConfig from File or Binary

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/vae/struct.BetaVaeConfig.html

Loads BetaVaeConfig from a file path or a binary buffer. Returns a Result with the loaded configuration or an error.

```rust
fn load<P>(file: P) -> Result<Self, ConfigError> where P: AsRef<Path>
```

```rust
fn load_binary(data: &[u8]) -> Result<Self, ConfigError>
```

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

### FilmConfig::init

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/film.rs.html

Initializes a new FiLM layer configuration on the specified device.

```APIDOC
## FilmConfig::init

### Description
Initializes a new FiLM layer configuration on the specified device.

### Method
`init`

### Arguments
- `self` - The `FilmConfig` instance.
- `device` (&B::Device) - The device on which to initialize the layer.

### Returns
- `Film<B>` - An initialized FiLM layer.

### Request Example
```rust
let config = FilmConfig::new(8, 16);
let device = Default::default(); // Or specify a concrete device
let film_layer = config.init::<NdArray>(&device);
```

### Response
#### Success Response (Film<B>)
- An instance of the `Film` layer initialized with the specified configuration and device.
```

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

### Get RewardClip Observation Space

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/env/wrapper/struct.RewardClip.html

Returns the description of the observation space for the RewardClip environment. This delegates to the inner environment.

```rust
fn observation_space(&self) -> Space>
```

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

### AttentionPoolConfig::init

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/attention.rs.html

Initializes an AttentionPool module with specified dimensions and backend.

```APIDOC
## AttentionPoolConfig::init<B: Backend>

### Description
Initialize an attention pool module on the given device.

### Parameters
* `self`: `&AttentionPoolConfig` - The configuration for the attention pool.
* `device`: `&B::Device` - The device on which to initialize the module.

### Returns
`AttentionPool<B>` - An initialized AttentionPool module.
```

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

### Get Output Dimension

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/z_conditioning/struct.ZConditioning.html

Returns the output dimension of the ZConditioning module after conditioning. This dimension is calculated as `obs_dim + hidden_dim`.

```rust
pub fn output_dim(&self) -> usize
```

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

### Initialize TargetAttentionConfig

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/attention.rs.html

Initializes a TargetAttention module, calculating the projection dimension as the minimum of query and key dimensions.

```rust
impl TargetAttentionConfig {
    /// Initialize a target attention module on the given device.
    ///
    /// The projection dimension is `min(query_dim, key_dim)`.
    pub fn init<B: Backend>(&self, device: &B::Device) -> TargetAttention<B> {
        let d_k = self.query_dim.min(self.key_dim);
        TargetAttention {
            w_q: LinearConfig::new(self.query_dim, d_k).init(device),
            w_k: LinearConfig::new(self.key_dim, d_k).init(device),
            d_k,
        }
    }
}
```

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

### Get League Agent Reference

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Provides an immutable reference to a specific agent within the league using its index.

```rust
pub fn agent(&self, idx: usize) -> &LeagueAgent<M>
```

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

### Default MctsConfig Implementation

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/mcts/struct.MctsConfig.html

Provides a default configuration for MctsConfig. This is useful for quick setup when specific parameters are not critical.

```rust
impl Default for MctsConfig {
    fn default() -> Self {
        // Default values are typically defined here
        // For example:
        MctsConfig {
            n_simulations: 100,
            exploration_constant: 2.0f64.sqrt(),
            n_actions: 0, // This might need to be set explicitly based on the environment
        }
    }
}
```

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

### Load PointerNetConfig from File

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.PointerNetConfig.html

Loads the configuration from a specified file path.

```rust
fn load<P>(file: P) -> Result<Self, ConfigError>
where P: AsRef<Path>,
```

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

### PointerNet Initialization

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/attention.rs.html

Initializes a PointerNet module with specified dimensions for query, key, and hidden layers. Requires a Backend and Device.

```rust
impl PointerNetConfig {
    /// Initialize a pointer network module on the given device.
    pub fn init<B: Backend>(&self, device: &B::Device) -> PointerNet<B> {
        PointerNet {
            w_query: LinearConfig::new(self.query_dim, self.hidden_dim).init(device),
            w_key: LinearConfig::new(self.key_dim, self.hidden_dim).init(device),
            v: LinearConfig::new(self.hidden_dim, 1).init(device),
        }
    }
}
```

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

### Initialize Film Layer

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/film/struct.FilmConfig.html

Initializes a FiLM layer on a given device using the current FilmConfig. Requires a Backend and a device.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> Film<B>
```

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

### Default DistillationConfig

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/distillation/struct.DistillationConfig.html

Provides the default configuration for DistillationConfig. This is useful when you need a starting point or don't want to specify all parameters.

```rust
impl Default for DistillationConfig
Source§
#### fn default() -> Self
Returns the “default value” for a type. Read more
```

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

### Get VAE Latent Dimension

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/nn/vae.rs.html

Returns the configured latent dimension size of the VAE. This indicates the dimensionality of the latent space.

```rust
/// Latent dimension accessor.
    pub fn latent_dim(&self) -> usize {
        self.latent_dim
    }
```

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

### Create New TargetAttentionConfig

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TargetAttentionConfig.html

Instantiates a new TargetAttentionConfig with specified query and key dimensions. Use this to set up the configuration before initializing the attention module.

```rust
pub fn new(query_dim: usize, key_dim: usize) -> Self
```

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

### Initialize PointerNet Module

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.PointerNetConfig.html

Initializes a PointerNet module on a given device using the current configuration.

```rust
pub fn init<B: Backend>(&self, device: &B::Device) -> PointerNet<B>
```

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

### Get Number of Specialists for Phase 1

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/cspl.rs.html

Retrieves the configured number of specialists to be trained in Phase 1 of the CSPL pipeline.

```rust
    /// Get the number of specialists for Phase 1.
    pub fn n_specialists(&self) -> usize {
        self.config.n_specialists
    }
```

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

### TransformerBlockConfig::init

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/attention/struct.TransformerBlockConfig.html

Initializes a TransformerBlock on a specified device using the current configuration.

```APIDOC
### impl TransformerBlockConfig

#### pub fn init<B: Backend>(&self, device: &B::Device) -> TransformerBlock<B>

Initialize a transformer block on the given device.
```

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

### Get Current Phase Step Count

Source: https://docs.rs/rl4burn/0.1.1/src/rl4burn/algo/cspl.rs.html

Returns the number of steps taken within the current phase of the CSPL pipeline.

```rust
    /// Number of steps taken in the current phase.
    pub fn phase_step(&self) -> u64 {
        self.phase_step
    }
```

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

### Get Number of Parameters

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/z_conditioning/struct.ZConditioning.html

Returns the total number of parameters within the ZConditioning module, including those of any sub-modules it might contain.

```rust
fn num_params(&self) -> usize
```

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

### BetaVaeConfig::new

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/nn/vae/struct.BetaVaeConfig.html

Creates a new BetaVaeConfig instance with the specified input dimension.

```APIDOC
## `new` Function

```rust
pub fn new(input_dim: usize) -> Self
```

### Description
Create a new instance of the config.

### Parameters
* `input_dim: usize` - The input dimension for the VAE.
```

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

### Get Latest Snapshot from SelfPlayPool

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/self_play/struct.SelfPlayPool.html

Retrieves the most recently added policy snapshot from the pool. Returns None if the pool is empty.

```rust
pub fn latest(&self) -> Option<&M>
```

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

### Initialize NormalizeReward

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/env/wrapper/struct.NormalizeReward.html

Creates a new NormalizeReward wrapper with a given environment, discount factor (gamma), and clipping value.

```rust
pub fn new(inner: E, gamma: f32, clip: f32) -> Self
```

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

### Get Mutable Agent Model

Source: https://docs.rs/rl4burn/0.1.1/rl4burn/algo/league/struct.League.html

Provides a mutable reference to an agent's model, allowing direct modification of its weights.

```rust
pub fn agent_model_mut(&mut self, idx: usize) -> &mut M
```
Tokens:
Raw