burn-rl 0.22.0-pre.3

RL crate for the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

use burn_core::prelude::*;

/// Trait for types that support tensor-like slice operations,
/// enabling storage in a [`TransitionBuffer`](super::TransitionBuffer).
///
/// Implement this trait for any type that wraps tensors and can be stored
/// in a replay buffer. The buffer uses these operations for:
/// - Pre-allocating storage (`zeros_like`)
/// - Writing transitions (`slice_assign_inplace`)
/// - Sampling batches (`select`)
pub trait SliceAccess: Clone + Sized {
    /// Create zeroed storage matching the shape of `sample` but with `capacity` rows
    /// along the first dimension.
    fn zeros_like(sample: &Self, capacity: usize, device: &Device) -> Self;

    /// Select rows at the given indices along the specified dimension.
    fn select(self, dim: usize, indices: Tensor<1, Int>) -> Self;

    /// Assign `value` at row `index` along the first dimension, in place.
    fn slice_assign_inplace(&mut self, index: usize, value: Self);
}

impl SliceAccess for Tensor<2> {
    fn zeros_like(sample: &Self, capacity: usize, device: &Device) -> Self {
        let feature_dim = sample.dims()[1];
        Tensor::zeros([capacity, feature_dim], device)
    }

    fn select(self, dim: usize, indices: Tensor<1, Int>) -> Self {
        Tensor::select(self, dim, indices)
    }

    fn slice_assign_inplace(&mut self, index: usize, value: Self) {
        self.inplace(|t| t.slice_assign(index..index + 1, value));
    }
}