r2l-core 0.0.3

A rust reinforcement learning library
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

use crate::{
    buffers::{Memory, TrajectoryBatch},
    tensor::R2lTensor,
};

/// Owned, structure-of-arrays storage for one environment trajectory.
#[derive(Clone)]
pub struct TrajectoryBuffer<T: R2lTensor> {
    states: Vec<T>,
    next_states: Vec<T>,
    actions: Vec<T>,
    rewards: Vec<f32>,
    terminated: Vec<bool>,
    truncated: Vec<bool>,
}

impl<T: R2lTensor> Default for TrajectoryBuffer<T> {
    fn default() -> Self {
        Self {
            states: Vec::default(),
            next_states: Vec::default(),
            actions: Vec::default(),
            rewards: Vec::default(),
            terminated: Vec::default(),
            truncated: Vec::default(),
        }
    }
}

impl<T: R2lTensor> TrajectoryBuffer<T> {
    /// Removes all stored transitions while retaining allocated capacity.
    pub fn clear(&mut self) {
        self.states.clear();
        self.next_states.clear();
        self.actions.clear();
        self.rewards.clear();
        self.terminated.clear();
        self.truncated.clear();
    }

    /// Appends one transition to the buffer.
    pub fn push(&mut self, memory: Memory<T>) {
        let Memory {
            state,
            next_state,
            action,
            reward,
            terminated,
            truncated,
        } = memory;
        self.states.push(state);
        self.next_states.push(next_state);
        self.actions.push(action);
        self.rewards.push(reward);
        self.terminated.push(terminated);
        self.truncated.push(truncated);
    }

    /// Replaces the most recently stored next state, if one exists.
    pub fn replace_last_next_state(&mut self, next_state: T) {
        if let Some(last_next_state) = self.next_states.last_mut() {
            *last_next_state = next_state;
        }
    }

    /// Returns the number of stored transitions.
    #[must_use]
    pub fn len(&self) -> usize {
        self.states.len()
    }

    /// Returns `true` when the buffer contains no transitions.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.states.is_empty()
    }

    /// Returns terminal-state flags.
    #[must_use]
    pub fn terminated(&self) -> &[bool] {
        &self.terminated
    }

    /// Returns truncation flags.
    #[must_use]
    pub fn truncated(&self) -> &[bool] {
        &self.truncated
    }

    /// Returns stored rewards.
    #[must_use]
    pub fn rewards(&self) -> &[f32] {
        &self.rewards
    }

    /// Returns mutable access to the stored rewards.
    pub fn rewards_mut(&mut self) -> &mut [f32] {
        &mut self.rewards
    }

    /// Borrows the buffer as aligned trajectory slices.
    #[must_use]
    pub fn to_trajectory_view(&self) -> TrajectoryView<'_, T> {
        TrajectoryView {
            states: &self.states,
            next_states: &self.next_states,
            actions: &self.actions,
            rewards: &self.rewards,
            terminated: &self.terminated,
            truncated: &self.truncated,
        }
    }
}

/// Borrowed view over the aligned fields of a [`TrajectoryBuffer`].
pub struct TrajectoryView<'a, T: R2lTensor> {
    /// Observations before each action.
    pub states: &'a [T],
    /// Observations after each action.
    pub next_states: &'a [T],
    /// Actions selected at each step.
    pub actions: &'a [T],
    /// Rewards produced at each step.
    pub rewards: &'a [f32],
    /// Terminal-state flags for each step.
    pub terminated: &'a [bool],
    /// Truncation flags for each step.
    pub truncated: &'a [bool],
}

impl<T: R2lTensor> TrajectoryBatch<T> for TrajectoryView<'_, T> {
    fn len(&self) -> usize {
        self.states.len()
    }

    fn is_empty(&self) -> bool {
        self.states.is_empty()
    }

    fn states(&self) -> &[T] {
        self.states
    }

    fn next_states(&self) -> &[T] {
        self.next_states
    }

    fn actions(&self) -> &[T] {
        self.actions
    }

    fn rewards(&self) -> &[f32] {
        self.rewards
    }

    fn terminated(&self) -> &[bool] {
        self.terminated
    }

    fn truncated(&self) -> &[bool] {
        self.truncated
    }
}

impl<T: R2lTensor> TrajectoryView<'_, T> {
    /// Iterates over combined termination and truncation flags.
    pub fn dones(&self) -> impl Iterator<Item = bool> {
        self.terminated
            .iter()
            .zip(self.truncated.iter())
            .map(|(terminated, truncated)| *terminated || *truncated)
    }

    /// Counts transitions that end an episode.
    #[must_use]
    pub fn episode_terminations(&self) -> usize {
        self.dones().filter(|x| *x).count()
    }
}