rlx_rl/
buffer.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15// RLX — replay buffer.
16
17/// One environment transition.
18#[derive(Debug, Clone)]
19pub struct Transition {
20    pub state: Vec<f32>,
21    pub action: Vec<f32>,
22    pub reward: f32,
23    pub next_state: Vec<f32>,
24    pub done: bool,
25}
26
27/// FIFO replay buffer with uniform sampling.
28#[derive(Debug, Clone, Default)]
29pub struct ReplayBuffer {
30    data: Vec<Transition>,
31    capacity: usize,
32}
33
34impl ReplayBuffer {
35    pub fn with_capacity(capacity: usize) -> Self {
36        Self {
37            data: Vec::new(),
38            capacity,
39        }
40    }
41
42    pub fn len(&self) -> usize {
43        self.data.len()
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.data.is_empty()
48    }
49
50    pub fn push(&mut self, t: Transition) {
51        if self.data.len() >= self.capacity {
52            self.data.remove(0);
53        }
54        self.data.push(t);
55    }
56
57    pub fn extend<I: IntoIterator<Item = Transition>>(&mut self, iter: I) {
58        for t in iter {
59            self.push(t);
60        }
61    }
62
63    /// Sample `n` transition indices (with replacement if `n` > len).
64    pub fn sample_indices(&self, n: usize, seed: &mut u64) -> Vec<usize> {
65        if self.data.is_empty() {
66            return Vec::new();
67        }
68        (0..n)
69            .map(|_| {
70                *seed = rand_like(*seed);
71                (*seed as usize) % self.data.len()
72            })
73            .collect()
74    }
75
76    pub fn get(&self, idx: usize) -> &Transition {
77        &self.data[idx]
78    }
79
80    pub fn iter(&self) -> impl Iterator<Item = &Transition> {
81        self.data.iter()
82    }
83}
84
85/// Minimal LCG for sampling without pulling in `rand`.
86pub fn rand_like(seed: u64) -> u64 {
87    seed.wrapping_mul(6364136223846793005).wrapping_add(1)
88}
rlx_rl/buffer.rs

rlx_rl/
buffer.rs
