r2l_core/buffers/
mod.rs

1use itertools::izip;
2
3use crate::tensor::R2lTensor;
4
5/// Owned and borrowed trajectory buffer types.
6pub mod buffer;
7
8/// One transition collected from an environment.
9#[derive(Debug)]
10pub struct Memory<T> {
11    /// Observation before the action.
12    pub state: T,
13    /// Observation after the action.
14    pub next_state: T,
15    /// Action selected by the actor.
16    pub action: T,
17    /// Reward emitted by the environment.
18    pub reward: f32,
19    /// Whether the transition ended in a terminal state.
20    pub terminated: bool,
21    /// Whether the transition ended because of a time limit or external cutoff.
22    pub truncated: bool,
23}
24
25impl<T> Memory<T> {
26    /// Returns `true` when the transition ends the episode for any reason.
27    pub fn is_done(&self) -> bool {
28        self.terminated || self.truncated
29    }
30}
31
32#[derive(Debug)]
33/// A set of transitions collected from multiple environments in one step.
34pub struct MultiMemory<T: R2lTensor> {
35    last_states: Vec<T>,
36    next_states: Vec<T>,
37    actions: Vec<T>,
38    rewards: Vec<f32>,
39    terminateds: Vec<bool>,
40    truncateds: Vec<bool>,
41}
42
43impl<T: R2lTensor> MultiMemory<T> {
44    /// Creates empty transition storage for up to `capacity` environments.
45    #[must_use]
46    pub fn with_capacity(capacity: usize) -> Self {
47        Self {
48            last_states: Vec::with_capacity(capacity),
49            next_states: Vec::with_capacity(capacity),
50            actions: Vec::with_capacity(capacity),
51            rewards: Vec::with_capacity(capacity),
52            terminateds: Vec::with_capacity(capacity),
53            truncateds: Vec::with_capacity(capacity),
54        }
55    }
56
57    /// Adds one environment transition.
58    pub fn push_memory(&mut self, memory: Memory<T>) {
59        let Memory {
60            state,
61            next_state,
62            action,
63            reward,
64            terminated,
65            truncated,
66        } = memory;
67        self.last_states.push(state);
68        self.next_states.push(next_state);
69        self.actions.push(action);
70        self.rewards.push(reward);
71        self.terminateds.push(terminated);
72        self.truncateds.push(truncated);
73    }
74
75    /// Returns mutable access to the observations produced by environment steps.
76    pub fn next_states_mut(&mut self) -> &mut [T] {
77        &mut self.next_states
78    }
79
80    /// Converts the stored transitions into individual memories.
81    #[must_use]
82    pub fn into_stored_memories(self) -> Vec<Memory<T>> {
83        let mut memories = Vec::with_capacity(self.last_states.len());
84        let Self {
85            last_states: states,
86            next_states,
87            actions,
88            rewards,
89            terminateds,
90            truncateds,
91        } = self;
92        for (state, next_state, action, reward, terminated, truncated) in izip!(
93            states,
94            next_states,
95            actions,
96            rewards,
97            terminateds,
98            truncateds
99        ) {
100            memories.push(Memory {
101                state,
102                next_state,
103                action,
104                reward,
105                terminated,
106                truncated,
107            });
108        }
109        memories
110    }
111
112    /// Completes the stored transitions with their corresponding next states.
113    ///
114    /// Extra values on either side are ignored.
115    pub fn into_memories(self, next_states: &[T]) -> Vec<Memory<T>> {
116        let mut memories = Vec::with_capacity(self.last_states.len());
117        let Self {
118            last_states: states,
119            next_states: _,
120            actions,
121            rewards,
122            terminateds,
123            truncateds,
124        } = self;
125        for (state, next_state, action, reward, terminated, truncated) in izip!(
126            states,
127            next_states,
128            actions,
129            rewards,
130            terminateds,
131            truncateds
132        ) {
133            memories.push(Memory {
134                state,
135                next_state: next_state.clone(),
136                action,
137                reward,
138                terminated,
139                truncated,
140            });
141        }
142        memories
143    }
144}
145
146/// Read-only access to a batch of aligned trajectory values.
147pub trait TrajectoryBatch<T: R2lTensor> {
148    /// Returns the number of transitions in the batch.
149    fn len(&self) -> usize;
150
151    /// Returns `true` when the batch contains no transitions.
152    fn is_empty(&self) -> bool;
153
154    /// Returns observations before each action.
155    fn states(&self) -> &[T];
156
157    /// Returns observations after each action.
158    fn next_states(&self) -> &[T];
159
160    /// Returns actions selected at each step.
161    fn actions(&self) -> &[T];
162
163    /// Returns rewards produced at each step.
164    fn rewards(&self) -> &[f32];
165
166    /// Returns terminal-state flags for each step.
167    fn terminated(&self) -> &[bool];
168
169    /// Returns truncation flags for each step.
170    fn truncated(&self) -> &[bool];
171}
r2l_core/buffers/mod.rs

r2l_core/buffers/
mod.rs
