r2l_core/buffers/
buffer.rs

1use crate::{
2    buffers::{Memory, TrajectoryBatch},
3    tensor::R2lTensor,
4};
5
6/// Owned, structure-of-arrays storage for one environment trajectory.
7#[derive(Clone)]
8pub struct TrajectoryBuffer<T: R2lTensor> {
9    states: Vec<T>,
10    next_states: Vec<T>,
11    actions: Vec<T>,
12    rewards: Vec<f32>,
13    terminated: Vec<bool>,
14    truncated: Vec<bool>,
15}
16
17impl<T: R2lTensor> Default for TrajectoryBuffer<T> {
18    fn default() -> Self {
19        Self {
20            states: Vec::default(),
21            next_states: Vec::default(),
22            actions: Vec::default(),
23            rewards: Vec::default(),
24            terminated: Vec::default(),
25            truncated: Vec::default(),
26        }
27    }
28}
29
30impl<T: R2lTensor> TrajectoryBuffer<T> {
31    /// Removes all stored transitions while retaining allocated capacity.
32    pub fn clear(&mut self) {
33        self.states.clear();
34        self.next_states.clear();
35        self.actions.clear();
36        self.rewards.clear();
37        self.terminated.clear();
38        self.truncated.clear();
39    }
40
41    /// Appends one transition to the buffer.
42    pub fn push(&mut self, memory: Memory<T>) {
43        let Memory {
44            state,
45            next_state,
46            action,
47            reward,
48            terminated,
49            truncated,
50        } = memory;
51        self.states.push(state);
52        self.next_states.push(next_state);
53        self.actions.push(action);
54        self.rewards.push(reward);
55        self.terminated.push(terminated);
56        self.truncated.push(truncated);
57    }
58
59    /// Replaces the most recently stored next state, if one exists.
60    pub fn replace_last_next_state(&mut self, next_state: T) {
61        if let Some(last_next_state) = self.next_states.last_mut() {
62            *last_next_state = next_state;
63        }
64    }
65
66    /// Returns the number of stored transitions.
67    #[must_use]
68    pub fn len(&self) -> usize {
69        self.states.len()
70    }
71
72    /// Returns `true` when the buffer contains no transitions.
73    #[must_use]
74    pub fn is_empty(&self) -> bool {
75        self.states.is_empty()
76    }
77
78    /// Returns terminal-state flags.
79    #[must_use]
80    pub fn terminated(&self) -> &[bool] {
81        &self.terminated
82    }
83
84    /// Returns truncation flags.
85    #[must_use]
86    pub fn truncated(&self) -> &[bool] {
87        &self.truncated
88    }
89
90    /// Returns stored rewards.
91    #[must_use]
92    pub fn rewards(&self) -> &[f32] {
93        &self.rewards
94    }
95
96    /// Returns mutable access to the stored rewards.
97    pub fn rewards_mut(&mut self) -> &mut [f32] {
98        &mut self.rewards
99    }
100
101    /// Borrows the buffer as aligned trajectory slices.
102    #[must_use]
103    pub fn to_trajectory_view(&self) -> TrajectoryView<'_, T> {
104        TrajectoryView {
105            states: &self.states,
106            next_states: &self.next_states,
107            actions: &self.actions,
108            rewards: &self.rewards,
109            terminated: &self.terminated,
110            truncated: &self.truncated,
111        }
112    }
113}
114
115/// Borrowed view over the aligned fields of a [`TrajectoryBuffer`].
116pub struct TrajectoryView<'a, T: R2lTensor> {
117    /// Observations before each action.
118    pub states: &'a [T],
119    /// Observations after each action.
120    pub next_states: &'a [T],
121    /// Actions selected at each step.
122    pub actions: &'a [T],
123    /// Rewards produced at each step.
124    pub rewards: &'a [f32],
125    /// Terminal-state flags for each step.
126    pub terminated: &'a [bool],
127    /// Truncation flags for each step.
128    pub truncated: &'a [bool],
129}
130
131impl<T: R2lTensor> TrajectoryBatch<T> for TrajectoryView<'_, T> {
132    fn len(&self) -> usize {
133        self.states.len()
134    }
135
136    fn is_empty(&self) -> bool {
137        self.states.is_empty()
138    }
139
140    fn states(&self) -> &[T] {
141        self.states
142    }
143
144    fn next_states(&self) -> &[T] {
145        self.next_states
146    }
147
148    fn actions(&self) -> &[T] {
149        self.actions
150    }
151
152    fn rewards(&self) -> &[f32] {
153        self.rewards
154    }
155
156    fn terminated(&self) -> &[bool] {
157        self.terminated
158    }
159
160    fn truncated(&self) -> &[bool] {
161        self.truncated
162    }
163}
164
165impl<T: R2lTensor> TrajectoryView<'_, T> {
166    /// Iterates over combined termination and truncation flags.
167    pub fn dones(&self) -> impl Iterator<Item = bool> {
168        self.terminated
169            .iter()
170            .zip(self.truncated.iter())
171            .map(|(terminated, truncated)| *terminated || *truncated)
172    }
173
174    /// Counts transitions that end an episode.
175    #[must_use]
176    pub fn episode_terminations(&self) -> usize {
177        self.dones().filter(|x| *x).count()
178    }
179}
r2l_core/buffers/buffer.rs

r2l_core/buffers/
buffer.rs
