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

use itertools::izip;

use crate::tensor::R2lTensor;

/// Owned and borrowed trajectory buffer types.
pub mod buffer;

/// One transition collected from an environment.
#[derive(Debug)]
pub struct Memory<T> {
    /// Observation before the action.
    pub state: T,
    /// Observation after the action.
    pub next_state: T,
    /// Action selected by the actor.
    pub action: T,
    /// Reward emitted by the environment.
    pub reward: f32,
    /// Whether the transition ended in a terminal state.
    pub terminated: bool,
    /// Whether the transition ended because of a time limit or external cutoff.
    pub truncated: bool,
}

impl<T> Memory<T> {
    /// Returns `true` when the transition ends the episode for any reason.
    pub fn is_done(&self) -> bool {
        self.terminated || self.truncated
    }
}

#[derive(Debug)]
/// A set of transitions collected from multiple environments in one step.
pub struct MultiMemory<T: R2lTensor> {
    last_states: Vec<T>,
    next_states: Vec<T>,
    actions: Vec<T>,
    rewards: Vec<f32>,
    terminateds: Vec<bool>,
    truncateds: Vec<bool>,
}

impl<T: R2lTensor> MultiMemory<T> {
    /// Creates empty transition storage for up to `capacity` environments.
    #[must_use]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            last_states: Vec::with_capacity(capacity),
            next_states: Vec::with_capacity(capacity),
            actions: Vec::with_capacity(capacity),
            rewards: Vec::with_capacity(capacity),
            terminateds: Vec::with_capacity(capacity),
            truncateds: Vec::with_capacity(capacity),
        }
    }

    /// Adds one environment transition.
    pub fn push_memory(&mut self, memory: Memory<T>) {
        let Memory {
            state,
            next_state,
            action,
            reward,
            terminated,
            truncated,
        } = memory;
        self.last_states.push(state);
        self.next_states.push(next_state);
        self.actions.push(action);
        self.rewards.push(reward);
        self.terminateds.push(terminated);
        self.truncateds.push(truncated);
    }

    /// Returns mutable access to the observations produced by environment steps.
    pub fn next_states_mut(&mut self) -> &mut [T] {
        &mut self.next_states
    }

    /// Converts the stored transitions into individual memories.
    #[must_use]
    pub fn into_stored_memories(self) -> Vec<Memory<T>> {
        let mut memories = Vec::with_capacity(self.last_states.len());
        let Self {
            last_states: states,
            next_states,
            actions,
            rewards,
            terminateds,
            truncateds,
        } = self;
        for (state, next_state, action, reward, terminated, truncated) in izip!(
            states,
            next_states,
            actions,
            rewards,
            terminateds,
            truncateds
        ) {
            memories.push(Memory {
                state,
                next_state,
                action,
                reward,
                terminated,
                truncated,
            });
        }
        memories
    }

    /// Completes the stored transitions with their corresponding next states.
    ///
    /// Extra values on either side are ignored.
    pub fn into_memories(self, next_states: &[T]) -> Vec<Memory<T>> {
        let mut memories = Vec::with_capacity(self.last_states.len());
        let Self {
            last_states: states,
            next_states: _,
            actions,
            rewards,
            terminateds,
            truncateds,
        } = self;
        for (state, next_state, action, reward, terminated, truncated) in izip!(
            states,
            next_states,
            actions,
            rewards,
            terminateds,
            truncateds
        ) {
            memories.push(Memory {
                state,
                next_state: next_state.clone(),
                action,
                reward,
                terminated,
                truncated,
            });
        }
        memories
    }
}

/// Read-only access to a batch of aligned trajectory values.
pub trait TrajectoryBatch<T: R2lTensor> {
    /// Returns the number of transitions in the batch.
    fn len(&self) -> usize;

    /// Returns `true` when the batch contains no transitions.
    fn is_empty(&self) -> bool;

    /// Returns observations before each action.
    fn states(&self) -> &[T];

    /// Returns observations after each action.
    fn next_states(&self) -> &[T];

    /// Returns actions selected at each step.
    fn actions(&self) -> &[T];

    /// Returns rewards produced at each step.
    fn rewards(&self) -> &[f32];

    /// Returns terminal-state flags for each step.
    fn terminated(&self) -> &[bool];

    /// Returns truncation flags for each step.
    fn truncated(&self) -> &[bool];
}