rl-mind 0.2.2__tar.gz → 0.2.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {rl_mind-0.2.2 → rl_mind-0.2.4}/PKG-INFO +6 -4
- {rl_mind-0.2.2 → rl_mind-0.2.4}/README.md +5 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/pyproject.toml +5 -3
- {rl_mind-0.2.2 → rl_mind-0.2.4}/rl_mind.md +3 -1
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/collectors.py +77 -9
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/core.py +58 -8
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/data.py +53 -6
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/env.py +26 -5
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/envs.py +18 -1
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/evaluation.py +52 -7
- rl_mind-0.2.4/src/rl_mind/mazemdp_env.py +317 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/nn.py +10 -1
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/notebook.py +17 -2
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/rendering.py +71 -6
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/PKG-INFO +6 -4
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/requires.txt +1 -1
- rl_mind-0.2.2/src/rl_mind/mazemdp_env.py +0 -182
- {rl_mind-0.2.2 → rl_mind-0.2.4}/setup.cfg +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/__init__.py +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/profiles.py +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/py.typed +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/SOURCES.txt +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/dependency_links.txt +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/top_level.txt +0 -0
- {rl_mind-0.2.2 → rl_mind-0.2.4}/tests/test_rl_mind.py +0 -0
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: rl-mind
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.4
|
|
4
4
|
Summary: A small, typed reinforcement-learning toolkit for the Master MIND RL practicals
|
|
5
5
|
Author-email: Benjamin Piwowarski <benjamin@piwowarski.fr>
|
|
6
|
-
Project-URL: Documentation, https://
|
|
6
|
+
Project-URL: Documentation, https://master-dac.isir.upmc.fr/rl/rl-mind/
|
|
7
7
|
Keywords: reinforcement-learning,gymnasium,torch,teaching
|
|
8
8
|
Classifier: Intended Audience :: Education
|
|
9
9
|
Classifier: Programming Language :: Python :: 3
|
|
@@ -16,7 +16,7 @@ Description-Content-Type: text/markdown
|
|
|
16
16
|
Requires-Dist: gymnasium[box2d,classic-control,mujoco]>=1.2.0
|
|
17
17
|
Requires-Dist: ipython>=8
|
|
18
18
|
Requires-Dist: ipywidgets>=8
|
|
19
|
-
Requires-Dist: mazemdp>=1.3.
|
|
19
|
+
Requires-Dist: mazemdp>=1.3.1
|
|
20
20
|
Requires-Dist: moviepy>=1.0
|
|
21
21
|
Requires-Dist: numpy>=1.26
|
|
22
22
|
Requires-Dist: tensorboard>=2.19
|
|
@@ -32,7 +32,9 @@ dictionaries and the editor can autocomplete field names.
|
|
|
32
32
|
|
|
33
33
|
The library is published on PyPI as
|
|
34
34
|
[`rl-mind`](https://pypi.org/project/rl-mind/) (`pip install rl-mind`; the
|
|
35
|
-
practicals install it for you) and the notebooks simply import it.
|
|
35
|
+
practicals install it for you) and the notebooks simply import it. This page is
|
|
36
|
+
the guide; every class, method and attribute is documented in the [API
|
|
37
|
+
reference](https://master-dac.isir.upmc.fr/rl/rl-mind/api.html). Importing the
|
|
36
38
|
package has **no side effect** — in particular the extra gymnasium environments
|
|
37
39
|
are only registered by an explicit `import rl_mind.envs` (see below).
|
|
38
40
|
|
|
@@ -1,5 +1,10 @@
|
|
|
1
1
|
# BBRL Notebooks
|
|
2
2
|
|
|
3
|
+
This file documents the **notebook authoring conventions**. The library the
|
|
4
|
+
notebooks build on is documented elsewhere: `rl_mind.md` is the guide and
|
|
5
|
+
`docs/` the Sphinx sources of the full API reference (`make docs`, deployed by
|
|
6
|
+
`make rsync` to <https://master-dac.isir.upmc.fr/rl/rl-mind/>).
|
|
7
|
+
|
|
3
8
|
## Notebooks
|
|
4
9
|
|
|
5
10
|
The notebooks are the python files in [percent
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "rl-mind"
|
|
3
|
-
version = "0.2.
|
|
3
|
+
version = "0.2.4"
|
|
4
4
|
description = "A small, typed reinforcement-learning toolkit for the Master MIND RL practicals"
|
|
5
5
|
readme = "rl_mind.md"
|
|
6
6
|
authors = [{ name = "Benjamin Piwowarski", email = "benjamin@piwowarski.fr" }]
|
|
@@ -21,7 +21,7 @@ dependencies = [
|
|
|
21
21
|
"gymnasium[mujoco,classic-control,box2d]>=1.2.0",
|
|
22
22
|
"ipython>=8",
|
|
23
23
|
"ipywidgets>=8",
|
|
24
|
-
"mazemdp>=1.3.
|
|
24
|
+
"mazemdp>=1.3.1",
|
|
25
25
|
"moviepy>=1.0",
|
|
26
26
|
"numpy>=1.26",
|
|
27
27
|
"tensorboard>=2.19",
|
|
@@ -29,7 +29,9 @@ dependencies = [
|
|
|
29
29
|
]
|
|
30
30
|
|
|
31
31
|
[project.urls]
|
|
32
|
-
|
|
32
|
+
# Built from docs/ (`make docs`) and deployed by `make rsync`; the API
|
|
33
|
+
# reference is generated from the docstrings of src/rl_mind
|
|
34
|
+
Documentation = "https://master-dac.isir.upmc.fr/rl/rl-mind/"
|
|
33
35
|
|
|
34
36
|
[build-system]
|
|
35
37
|
# swig is not needed to build rl-mind itself: it is listed here so that the
|
|
@@ -8,7 +8,9 @@ dictionaries and the editor can autocomplete field names.
|
|
|
8
8
|
|
|
9
9
|
The library is published on PyPI as
|
|
10
10
|
[`rl-mind`](https://pypi.org/project/rl-mind/) (`pip install rl-mind`; the
|
|
11
|
-
practicals install it for you) and the notebooks simply import it.
|
|
11
|
+
practicals install it for you) and the notebooks simply import it. This page is
|
|
12
|
+
the guide; every class, method and attribute is documented in the [API
|
|
13
|
+
reference](https://master-dac.isir.upmc.fr/rl/rl-mind/api.html). Importing the
|
|
12
14
|
package has **no side effect** — in particular the extra gymnasium environments
|
|
13
15
|
are only registered by an explicit `import rl_mind.envs` (see below).
|
|
14
16
|
|
|
@@ -9,8 +9,8 @@ auto-reset bookkeeping so that the learning code never has to.
|
|
|
9
9
|
`Transitions` (for replay-buffer based algorithms). The environments
|
|
10
10
|
keep their state between calls.
|
|
11
11
|
- `EpisodeCollector.collect(n_episodes)` resets the environments and
|
|
12
|
-
returns at least `n_episodes` complete `Episode`
|
|
13
|
-
algorithms, which need *on-policy* full episodes
|
|
12
|
+
returns at least `n_episodes` complete episodes (`Episode`), for
|
|
13
|
+
episodic algorithms, which need *on-policy* full episodes.
|
|
14
14
|
- `RolloutCollector.collect(n_steps)` returns a `[T, B]` `Rollout` (for
|
|
15
15
|
on-policy algorithms such as A2C and PPO). It requires an environment
|
|
16
16
|
with `same_step_reset=True`, so that every row of the rollout is a
|
|
@@ -43,10 +43,25 @@ def _stack(items):
|
|
|
43
43
|
|
|
44
44
|
|
|
45
45
|
class TransitionCollector(Generic[ActionT]):
|
|
46
|
-
"""Collect batches of transitions by running an actor in a `VecEnv`
|
|
46
|
+
"""Collect batches of transitions by running an actor in a `VecEnv`
|
|
47
|
+
|
|
48
|
+
The collector for off-policy algorithms (DQN, DDPG, TD3, SAC, TQC): the
|
|
49
|
+
transitions it returns are meant to be pushed into a
|
|
50
|
+
`rl_mind.data.ReplayBuffer`, so their order does not matter and the
|
|
51
|
+
(invalid) auto-reset steps are simply dropped.
|
|
52
|
+
"""
|
|
47
53
|
|
|
48
54
|
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
|
|
55
|
+
"""Create the collector
|
|
56
|
+
|
|
57
|
+
:param env: The training environments; the collector keeps their
|
|
58
|
+
state across `collect` calls
|
|
59
|
+
:param actor: The actor to run — its *training* behavior
|
|
60
|
+
(`actor(obs)`, i.e. with exploration) is what gets collected
|
|
61
|
+
"""
|
|
62
|
+
#: The environments being stepped
|
|
49
63
|
self.env = env
|
|
64
|
+
#: The actor being run
|
|
50
65
|
self.actor = actor
|
|
51
66
|
#: Total number of recorded transitions
|
|
52
67
|
self.steps = 0
|
|
@@ -56,8 +71,16 @@ class TransitionCollector(Generic[ActionT]):
|
|
|
56
71
|
def collect(self, n_steps: int) -> Transitions[ActionT]:
|
|
57
72
|
"""Perform `n_steps` steps in every environment
|
|
58
73
|
|
|
59
|
-
|
|
60
|
-
|
|
74
|
+
The environments are reset on the first call only: successive calls
|
|
75
|
+
continue the same episodes, so a training loop can interleave
|
|
76
|
+
`collect(1)` with gradient steps.
|
|
77
|
+
|
|
78
|
+
:param n_steps: The number of steps performed in *each* of the
|
|
79
|
+
`env.num_envs` environments
|
|
80
|
+
:returns: The recorded transitions, a flat batch of at most
|
|
81
|
+
`n_steps * num_envs` of them (the steps where an environment was
|
|
82
|
+
auto-resetting are skipped, and `steps` is increased by the
|
|
83
|
+
number actually recorded)
|
|
61
84
|
"""
|
|
62
85
|
if self._obs is None:
|
|
63
86
|
self._obs = self.env.reset()
|
|
@@ -93,10 +116,23 @@ class TransitionCollector(Generic[ActionT]):
|
|
|
93
116
|
|
|
94
117
|
|
|
95
118
|
class EpisodeCollector(Generic[ActionT]):
|
|
96
|
-
"""Collect full episodes by running an actor in a `VecEnv`
|
|
119
|
+
"""Collect full episodes by running an actor in a `VecEnv`
|
|
120
|
+
|
|
121
|
+
The collector for episodic algorithms (REINFORCE), which need complete
|
|
122
|
+
*on-policy* episodes: the environments are therefore reset at the
|
|
123
|
+
beginning of every `collect`.
|
|
124
|
+
"""
|
|
97
125
|
|
|
98
126
|
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
|
|
127
|
+
"""Create the collector
|
|
128
|
+
|
|
129
|
+
:param env: The training environments (reset at each `collect`)
|
|
130
|
+
:param actor: The actor to run — its *training* behavior
|
|
131
|
+
(`actor(obs)`) is what gets collected
|
|
132
|
+
"""
|
|
133
|
+
#: The environments being stepped
|
|
99
134
|
self.env = env
|
|
135
|
+
#: The actor being run
|
|
100
136
|
self.actor = actor
|
|
101
137
|
#: Total number of recorded steps
|
|
102
138
|
self.steps = 0
|
|
@@ -106,6 +142,11 @@ class EpisodeCollector(Generic[ActionT]):
|
|
|
106
142
|
|
|
107
143
|
The environments are reset at the beginning of the call, so that all
|
|
108
144
|
the returned episodes are collected with the *current* actor.
|
|
145
|
+
|
|
146
|
+
:param n_episodes: The minimum number of episodes to return
|
|
147
|
+
:returns: The completed episodes, of varying lengths — more than
|
|
148
|
+
`n_episodes` when several environments finish during the same
|
|
149
|
+
step (the episodes still running are discarded)
|
|
109
150
|
"""
|
|
110
151
|
num_envs = self.env.num_envs
|
|
111
152
|
obs = self.env.reset()
|
|
@@ -151,23 +192,50 @@ class EpisodeCollector(Generic[ActionT]):
|
|
|
151
192
|
|
|
152
193
|
|
|
153
194
|
class RolloutCollector(Generic[ActionT]):
|
|
154
|
-
"""Collect fixed-length `[T, B]` rollouts by running an actor in a `VecEnv`
|
|
195
|
+
"""Collect fixed-length `[T, B]` rollouts by running an actor in a `VecEnv`
|
|
196
|
+
|
|
197
|
+
The collector for on-policy algorithms (A2C, PPO, GAIL): the time axis is
|
|
198
|
+
preserved, since the advantages are computed by a backward recursion over
|
|
199
|
+
it. This is why the environment must use `same_step_reset=True` — every
|
|
200
|
+
row of the `[T, B]` block is then a valid transition.
|
|
201
|
+
"""
|
|
155
202
|
|
|
156
203
|
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
|
|
204
|
+
"""Create the collector
|
|
205
|
+
|
|
206
|
+
:param env: The training environments, created with
|
|
207
|
+
`same_step_reset=True`; the collector keeps their state across
|
|
208
|
+
`collect` calls
|
|
209
|
+
:param actor: The actor to run — its *training* behavior
|
|
210
|
+
(`actor(obs)`) is what gets collected
|
|
211
|
+
:raises AssertionError: If `env` was not created with
|
|
212
|
+
`same_step_reset=True`
|
|
213
|
+
"""
|
|
157
214
|
assert env.same_step_reset, (
|
|
158
215
|
"RolloutCollector requires a VecEnv created with "
|
|
159
216
|
"same_step_reset=True (so that every rollout row is a valid "
|
|
160
217
|
"transition)"
|
|
161
218
|
)
|
|
219
|
+
#: The environments being stepped
|
|
162
220
|
self.env = env
|
|
221
|
+
#: The actor being run
|
|
163
222
|
self.actor = actor
|
|
164
223
|
#: Total number of recorded steps
|
|
165
224
|
self.steps = 0
|
|
166
225
|
self._obs: torch.Tensor | None = None
|
|
167
226
|
|
|
168
227
|
def collect(self, n_steps: int) -> Rollout[ActionT]:
|
|
169
|
-
"""Perform `n_steps` steps in every environment
|
|
170
|
-
|
|
228
|
+
"""Perform `n_steps` steps in every environment
|
|
229
|
+
|
|
230
|
+
The environments are reset on the first call only: successive calls
|
|
231
|
+
continue where the previous one stopped.
|
|
232
|
+
|
|
233
|
+
:param n_steps: The number of steps performed in *each* environment,
|
|
234
|
+
i.e. the length `T` of the rollout
|
|
235
|
+
:returns: A `[n_steps, num_envs]` rollout (every row is a valid
|
|
236
|
+
transition, episode boundaries being marked by `terminated` /
|
|
237
|
+
`truncated`)
|
|
238
|
+
"""
|
|
171
239
|
if self._obs is None:
|
|
172
240
|
self._obs = self.env.reset()
|
|
173
241
|
|
|
@@ -48,13 +48,20 @@ class TensorStruct:
|
|
|
48
48
|
def __getitem__(self, index) -> Self:
|
|
49
49
|
"""Index or slice all the tensor fields along the batch dimension
|
|
50
50
|
|
|
51
|
-
|
|
52
|
-
|
|
51
|
+
:param index: Anything a tensor accepts: an integer, a slice, an
|
|
52
|
+
index tensor or a boolean mask
|
|
53
|
+
:returns: A new structure holding the selected entries (the fields of
|
|
54
|
+
a `TensorStruct` are never modified in place)
|
|
53
55
|
"""
|
|
54
56
|
return self._map(lambda tensor: tensor[index])
|
|
55
57
|
|
|
56
58
|
def __len__(self) -> int:
|
|
57
|
-
"""The size of the batch dimension
|
|
59
|
+
"""The size of the batch dimension
|
|
60
|
+
|
|
61
|
+
:returns: The size of the leading dimension, read from the first
|
|
62
|
+
tensor (or nested structure) field
|
|
63
|
+
:raises ValueError: If the structure holds no tensor field
|
|
64
|
+
"""
|
|
58
65
|
for field in fields(self):
|
|
59
66
|
value = getattr(self, field.name)
|
|
60
67
|
if isinstance(value, TensorStruct):
|
|
@@ -64,7 +71,16 @@ class TensorStruct:
|
|
|
64
71
|
raise ValueError(f"{type(self).__name__} has no tensor field")
|
|
65
72
|
|
|
66
73
|
def set_(self, index, value: Self) -> None:
|
|
67
|
-
"""Write `value` at position(s) `index` (in-place)
|
|
74
|
+
"""Write `value` at position(s) `index` (in-place)
|
|
75
|
+
|
|
76
|
+
The only mutating operation: instances are frozen, but the *contents*
|
|
77
|
+
of their tensor fields can be overwritten. This is how `ReplayBuffer`
|
|
78
|
+
fills its preallocated storage.
|
|
79
|
+
|
|
80
|
+
:param index: Where to write, as for `__getitem__`
|
|
81
|
+
:param value: A structure of the same type, whose batch size matches
|
|
82
|
+
what `index` selects
|
|
83
|
+
"""
|
|
68
84
|
for field in fields(self):
|
|
69
85
|
target = getattr(self, field.name)
|
|
70
86
|
if isinstance(target, TensorStruct):
|
|
@@ -74,7 +90,13 @@ class TensorStruct:
|
|
|
74
90
|
|
|
75
91
|
@classmethod
|
|
76
92
|
def cat(cls, items: Sequence[Self]) -> Self:
|
|
77
|
-
"""Concatenate a sequence of structures along the batch dimension
|
|
93
|
+
"""Concatenate a sequence of structures along the batch dimension
|
|
94
|
+
|
|
95
|
+
:param items: The structures to concatenate (same type, at least
|
|
96
|
+
one); non-tensor fields are taken from `items[0]`
|
|
97
|
+
:returns: One structure of batch size `sum(len(item) for item in
|
|
98
|
+
items)`
|
|
99
|
+
"""
|
|
78
100
|
values = {}
|
|
79
101
|
for field in fields(items[0]):
|
|
80
102
|
value = getattr(items[0], field.name)
|
|
@@ -90,7 +112,16 @@ class TensorStruct:
|
|
|
90
112
|
|
|
91
113
|
@classmethod
|
|
92
114
|
def stack(cls, items: Sequence[Self]) -> Self:
|
|
93
|
-
"""Stack a sequence of structures along a new (first) dimension
|
|
115
|
+
"""Stack a sequence of structures along a new (first) dimension
|
|
116
|
+
|
|
117
|
+
This is how the collectors turn a list of per-step actions into a
|
|
118
|
+
time-indexed `[T, ...]` structure.
|
|
119
|
+
|
|
120
|
+
:param items: The structures to stack (same type and batch size, at
|
|
121
|
+
least one); non-tensor fields are taken from `items[0]`
|
|
122
|
+
:returns: One structure with a new leading dimension of size
|
|
123
|
+
`len(items)`
|
|
124
|
+
"""
|
|
94
125
|
values = {}
|
|
95
126
|
for field in fields(items[0]):
|
|
96
127
|
value = getattr(items[0], field.name)
|
|
@@ -140,13 +171,32 @@ class Actor(nn.Module, ABC, Generic[ActionT]):
|
|
|
140
171
|
|
|
141
172
|
@abstractmethod
|
|
142
173
|
def forward(self, obs: Tensor) -> ActionT:
|
|
143
|
-
"""Compute the action for a batch of observations (
|
|
174
|
+
"""Compute the action for a batch of observations (training behavior)
|
|
175
|
+
|
|
176
|
+
Called by `actor(obs)`: this is the behavior used to *collect* data,
|
|
177
|
+
so it is where sampling and exploration noise belong.
|
|
178
|
+
|
|
179
|
+
:param obs: The observations, shape `[B, obs_dim]` (or a
|
|
180
|
+
`TensorStruct` for a structured observation space)
|
|
181
|
+
:returns: The actions, an `Action` (or a subclass) whose tensors have
|
|
182
|
+
batch size `B`
|
|
183
|
+
"""
|
|
144
184
|
...
|
|
145
185
|
|
|
146
186
|
def act(self, obs: Tensor) -> Tensor:
|
|
147
|
-
"""Deterministic action used for evaluation
|
|
187
|
+
"""Deterministic action used for evaluation
|
|
188
|
+
|
|
189
|
+
Defaults to `forward(obs).value`; override it when the training-time
|
|
190
|
+
behavior is stochastic, so that evaluation uses e.g. the mode of the
|
|
191
|
+
distribution instead of a sample.
|
|
192
|
+
|
|
193
|
+
:param obs: The observations, shape `[B, obs_dim]`
|
|
194
|
+
:returns: The action tensor, shape `[B, action_dim]` (continuous
|
|
195
|
+
actions) or `[B]` (discrete actions)
|
|
196
|
+
"""
|
|
148
197
|
return self(obs).value
|
|
149
198
|
|
|
150
199
|
# Overriden so that type checkers know the return type of `actor(obs)`
|
|
151
200
|
def __call__(self, obs: Tensor) -> ActionT:
|
|
201
|
+
"""Call `forward`, i.e. compute the training-time action for `obs`"""
|
|
152
202
|
return super().__call__(obs)
|
|
@@ -45,20 +45,42 @@ class Transitions(TensorStruct, Generic[ActionT]):
|
|
|
45
45
|
|
|
46
46
|
class ReplayBuffer(Generic[ActionT]):
|
|
47
47
|
"""A fixed-capacity buffer of transitions (a "ring buffer": when full,
|
|
48
|
-
new transitions overwrite the oldest ones)
|
|
48
|
+
new transitions overwrite the oldest ones)
|
|
49
|
+
|
|
50
|
+
Used by the off-policy notebooks: `add` the transitions returned by a
|
|
51
|
+
`rl_mind.collectors.TransitionCollector`, then `sample` minibatches from
|
|
52
|
+
the whole (shuffled) history.
|
|
53
|
+
"""
|
|
49
54
|
|
|
50
55
|
def __init__(self, capacity: int):
|
|
56
|
+
"""Create an empty buffer
|
|
57
|
+
|
|
58
|
+
The storage itself is allocated on the first `add` (it needs the
|
|
59
|
+
shapes and dtypes of the transitions).
|
|
60
|
+
|
|
61
|
+
:param capacity: The maximum number of transitions kept; beyond it,
|
|
62
|
+
the oldest ones are overwritten
|
|
63
|
+
"""
|
|
64
|
+
#: The maximum number of transitions the buffer holds
|
|
51
65
|
self.capacity = int(capacity)
|
|
52
66
|
self._storage: Transitions[ActionT] | None = None
|
|
53
67
|
self._size = 0
|
|
54
68
|
self._pos = 0
|
|
55
69
|
|
|
56
70
|
def __len__(self) -> int:
|
|
57
|
-
"""The number of transitions currently stored
|
|
71
|
+
"""The number of transitions currently stored
|
|
72
|
+
|
|
73
|
+
:returns: A number between 0 and `capacity`
|
|
74
|
+
"""
|
|
58
75
|
return self._size
|
|
59
76
|
|
|
60
77
|
def add(self, transitions: Transitions[ActionT]) -> None:
|
|
61
|
-
"""Add a batch of transitions to the buffer
|
|
78
|
+
"""Add a batch of transitions to the buffer
|
|
79
|
+
|
|
80
|
+
:param transitions: The transitions to store (a batch of any size;
|
|
81
|
+
the storage is allocated from the first batch, so every later
|
|
82
|
+
batch must have the same fields, shapes and dtypes)
|
|
83
|
+
"""
|
|
62
84
|
if self._storage is None:
|
|
63
85
|
# Allocate the storage on first use, using the shapes/dtypes of
|
|
64
86
|
# the incoming transitions
|
|
@@ -74,7 +96,12 @@ class ReplayBuffer(Generic[ActionT]):
|
|
|
74
96
|
self._size = min(self._size + len(transitions), self.capacity)
|
|
75
97
|
|
|
76
98
|
def sample(self, batch_size: int) -> Transitions[ActionT]:
|
|
77
|
-
"""Sample a batch of transitions uniformly at random (with replacement)
|
|
99
|
+
"""Sample a batch of transitions uniformly at random (with replacement)
|
|
100
|
+
|
|
101
|
+
:param batch_size: The number of transitions to draw
|
|
102
|
+
:returns: A `Transitions` of batch size `batch_size`
|
|
103
|
+
:raises AssertionError: If the buffer is still empty
|
|
104
|
+
"""
|
|
78
105
|
assert self._storage is not None, "The replay buffer is empty"
|
|
79
106
|
indices = torch.randint(0, self._size, (batch_size,))
|
|
80
107
|
return self._storage[indices]
|
|
@@ -102,7 +129,11 @@ class Episode(Generic[ActionT]):
|
|
|
102
129
|
terminated: bool
|
|
103
130
|
|
|
104
131
|
def __len__(self) -> int:
|
|
105
|
-
"""The number of steps
|
|
132
|
+
"""The number of steps `T` of the episode
|
|
133
|
+
|
|
134
|
+
:returns: The number of actions taken (`obs`, `action` and `reward`
|
|
135
|
+
all have `T` entries; `all_obs` has `T+1`)
|
|
136
|
+
"""
|
|
106
137
|
return self.reward.shape[0]
|
|
107
138
|
|
|
108
139
|
@property
|
|
@@ -159,7 +190,15 @@ class Rollout(TensorStruct, Generic[ActionT]):
|
|
|
159
190
|
return self.terminated | self.truncated
|
|
160
191
|
|
|
161
192
|
def flatten(self) -> "Rollout[ActionT]":
|
|
162
|
-
"""Merge the time and environment dimensions: `[T, B, ...] -> [T*B, ...]`
|
|
193
|
+
"""Merge the time and environment dimensions: `[T, B, ...] -> [T*B, ...]`
|
|
194
|
+
|
|
195
|
+
The time structure is lost, so this is done *after* the quantities
|
|
196
|
+
that need it (returns, GAE advantages) have been computed; the flat
|
|
197
|
+
result is what `minibatches` iterates over.
|
|
198
|
+
|
|
199
|
+
:returns: The same rollout with a single leading dimension of size
|
|
200
|
+
`T * B`
|
|
201
|
+
"""
|
|
163
202
|
return self._map(lambda tensor: tensor.flatten(0, 1))
|
|
164
203
|
|
|
165
204
|
|
|
@@ -173,6 +212,14 @@ def minibatches(
|
|
|
173
212
|
|
|
174
213
|
The data is partitioned: each sample belongs to exactly one minibatch
|
|
175
214
|
(the last one may be smaller). With `shuffle`, the partition is random.
|
|
215
|
+
|
|
216
|
+
:param data: Any `TensorStruct` (e.g. `Transitions`, or a `Rollout` that
|
|
217
|
+
was `flatten`ed), of batch size `N`
|
|
218
|
+
:param batch_size: The size of the minibatches
|
|
219
|
+
:param shuffle: Whether to shuffle the samples before partitioning them
|
|
220
|
+
(a different partition at every pass)
|
|
221
|
+
:returns: An iterator over `ceil(N / batch_size)` structures of the same
|
|
222
|
+
type as `data`
|
|
176
223
|
"""
|
|
177
224
|
indices = torch.randperm(len(data)) if shuffle else torch.arange(len(data))
|
|
178
225
|
for start in range(0, len(data), batch_size):
|
|
@@ -130,7 +130,12 @@ class EnvStep:
|
|
|
130
130
|
|
|
131
131
|
|
|
132
132
|
class VecEnv:
|
|
133
|
-
"""`num_envs` copies of a gymnasium environment, with torch tensors in/out
|
|
133
|
+
"""`num_envs` copies of a gymnasium environment, with torch tensors in/out
|
|
134
|
+
|
|
135
|
+
The environments are stepped **synchronously** (in the same process): one
|
|
136
|
+
call to `step` advances every copy by one step, and the returned tensors
|
|
137
|
+
have the number of environments as their leading dimension.
|
|
138
|
+
"""
|
|
134
139
|
|
|
135
140
|
def __init__(
|
|
136
141
|
self,
|
|
@@ -153,14 +158,20 @@ class VecEnv:
|
|
|
153
158
|
used for training, evaluation and video recording
|
|
154
159
|
:param kwargs: Extra arguments passed to `gymnasium.make`
|
|
155
160
|
"""
|
|
161
|
+
#: The gymnasium environment name
|
|
156
162
|
self.env_name = env_name
|
|
163
|
+
#: The number of parallel environments (the batch size `B`)
|
|
157
164
|
self.num_envs = num_envs
|
|
165
|
+
#: Whether ended episodes are reset within the ending step
|
|
158
166
|
self.same_step_reset = same_step_reset
|
|
167
|
+
#: The wrappers applied to each sub-environment
|
|
159
168
|
self.wrappers = tuple(wrappers)
|
|
169
|
+
#: The extra arguments passed to `gymnasium.make`
|
|
160
170
|
self.env_kwargs = kwargs
|
|
161
171
|
vector_kwargs = {}
|
|
162
172
|
if same_step_reset:
|
|
163
173
|
vector_kwargs["autoreset_mode"] = gym.vector.AutoresetMode.SAME_STEP
|
|
174
|
+
#: The underlying gymnasium vector environment
|
|
164
175
|
self.gym_env = gym.make_vec(
|
|
165
176
|
env_name,
|
|
166
177
|
num_envs=num_envs,
|
|
@@ -244,10 +255,15 @@ class VecEnv:
|
|
|
244
255
|
return int(space.n)
|
|
245
256
|
|
|
246
257
|
def reset(self) -> Observation:
|
|
247
|
-
"""Reset all the environments
|
|
258
|
+
"""Reset all the environments, starting a fresh episode in each
|
|
259
|
+
|
|
260
|
+
Only needed once, at the very beginning: afterwards the environments
|
|
261
|
+
reset themselves when an episode ends (see the module documentation).
|
|
262
|
+
The collectors call it for you.
|
|
248
263
|
|
|
249
|
-
|
|
250
|
-
|
|
264
|
+
:returns: The initial observations, shape `[B, obs_dim]` for a flat
|
|
265
|
+
space, `[B]` (state indices) for a tabular one, or a
|
|
266
|
+
`TensorStruct` of `[B, ...]` tensors for a `Dict` space
|
|
251
267
|
"""
|
|
252
268
|
obs, _ = self.gym_env.reset(seed=self._seed)
|
|
253
269
|
self._seed = None # seed only the first reset
|
|
@@ -256,7 +272,12 @@ class VecEnv:
|
|
|
256
272
|
def step(self, actions: Tensor) -> EnvStep:
|
|
257
273
|
"""Perform one step in every environment
|
|
258
274
|
|
|
259
|
-
:param actions:
|
|
275
|
+
:param actions: One action per environment: `[B, action_dim]`
|
|
276
|
+
(continuous) or `[B]` (discrete). The tensor is detached and moved
|
|
277
|
+
to the CPU, so gradients never flow into the environment.
|
|
278
|
+
:returns: The `EnvStep` of this step (all fields `[B, ...]`). Ended
|
|
279
|
+
episodes are reset automatically, in the mode chosen at
|
|
280
|
+
construction time.
|
|
260
281
|
"""
|
|
261
282
|
np_actions = actions.detach().cpu().numpy()
|
|
262
283
|
np_actions = np_actions.astype(self.gym_env.action_space.dtype, copy=False)
|
|
@@ -15,13 +15,30 @@ __all__ = ["ContinuousCartPoleEnv"]
|
|
|
15
15
|
|
|
16
16
|
|
|
17
17
|
class ContinuousCartPoleEnv(CartPoleEnv):
|
|
18
|
-
"""Continuous-action version of the CartPole-v1 environment
|
|
18
|
+
"""Continuous-action version of the CartPole-v1 environment
|
|
19
|
+
|
|
20
|
+
Same observation space, reward (+1 per step) and termination condition as
|
|
21
|
+
`CartPole-v1`, but the action is the *signed fraction* of the maximum
|
|
22
|
+
force applied to the cart, a `Box(-1, 1, shape=(1,))`. Registered as
|
|
23
|
+
`CartPoleContinuous-v1` (500 steps at most) when `rl_mind.envs` is
|
|
24
|
+
imported.
|
|
25
|
+
"""
|
|
19
26
|
|
|
20
27
|
def __init__(self, **kwargs):
|
|
28
|
+
""":param kwargs: Passed to `gymnasium`'s `CartPoleEnv` (e.g.
|
|
29
|
+
`render_mode`)
|
|
30
|
+
"""
|
|
21
31
|
super().__init__(**kwargs)
|
|
22
32
|
self.action_space = spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float64)
|
|
23
33
|
|
|
24
34
|
def step(self, action):
|
|
35
|
+
"""Apply a force to the cart and advance the simulation by one step
|
|
36
|
+
|
|
37
|
+
:param action: The force, as a fraction of `force_mag` in $[-1, 1]$
|
|
38
|
+
(clipped); a scalar or a one-element array
|
|
39
|
+
:returns: The gymnasium 5-tuple `(obs, reward, terminated, truncated,
|
|
40
|
+
info)`, with a reward of 1.0 at every step
|
|
41
|
+
"""
|
|
25
42
|
force_scale = float(np.clip(np.asarray(action), -1.0, 1.0).item())
|
|
26
43
|
assert self.state is not None, "Call reset before using step method."
|
|
27
44
|
|
|
@@ -50,7 +50,13 @@ class EvalResult:
|
|
|
50
50
|
|
|
51
51
|
|
|
52
52
|
class Evaluator(Generic[ActionT]):
|
|
53
|
-
"""Periodically evaluates an actor, keeping track of the best one
|
|
53
|
+
"""Periodically evaluates an actor, keeping track of the best one
|
|
54
|
+
|
|
55
|
+
The evaluation environments are deliberately *separate* from the training
|
|
56
|
+
ones (their own seed, their own episode count, no same-step reset), so
|
|
57
|
+
that what is measured does not depend on the data the agent is training
|
|
58
|
+
on. Evaluation uses `Actor.act`, the deterministic behavior.
|
|
59
|
+
"""
|
|
54
60
|
|
|
55
61
|
def __init__(
|
|
56
62
|
self,
|
|
@@ -68,9 +74,13 @@ class Evaluator(Generic[ActionT]):
|
|
|
68
74
|
:param writer: A tensorboard writer where the evaluation rewards are
|
|
69
75
|
logged (optional)
|
|
70
76
|
"""
|
|
77
|
+
#: The evaluation environments
|
|
71
78
|
self.env = env
|
|
79
|
+
#: The number of training steps between two evaluations
|
|
72
80
|
self.every = every
|
|
81
|
+
#: Where videos are recorded (see `visualize_best`)
|
|
73
82
|
self.run_dir = run_dir
|
|
83
|
+
#: The tensorboard writer the mean reward is logged to, if any
|
|
74
84
|
self.writer = writer
|
|
75
85
|
|
|
76
86
|
#: All the evaluations so far
|
|
@@ -82,8 +92,18 @@ class Evaluator(Generic[ActionT]):
|
|
|
82
92
|
self._last_eval_step = 0
|
|
83
93
|
|
|
84
94
|
def run_if_needed(self, steps: int, actor: Actor[ActionT]) -> EvalResult | None:
|
|
85
|
-
"""Evaluate `actor` if
|
|
86
|
-
|
|
95
|
+
"""Evaluate `actor` if it is time to do so
|
|
96
|
+
|
|
97
|
+
The call to make from the training loop: it does nothing until `every`
|
|
98
|
+
steps have been performed since the last evaluation. When it does
|
|
99
|
+
evaluate, the result is appended to `history`, logged to `writer`, and
|
|
100
|
+
a deep copy of the actor is kept in `best_actor` if it improves on
|
|
101
|
+
`best_reward`.
|
|
102
|
+
|
|
103
|
+
:param steps: The number of environment steps performed so far
|
|
104
|
+
(typically a collector's `steps`)
|
|
105
|
+
:param actor: The actor to evaluate
|
|
106
|
+
:returns: The `EvalResult`, or `None` when no evaluation was run
|
|
87
107
|
"""
|
|
88
108
|
if steps - self._last_eval_step < self.every:
|
|
89
109
|
return None
|
|
@@ -102,8 +122,16 @@ class Evaluator(Generic[ActionT]):
|
|
|
102
122
|
return result
|
|
103
123
|
|
|
104
124
|
def run(self, actor: Actor[ActionT]) -> Tensor:
|
|
105
|
-
"""Run one evaluation episode in each environment
|
|
106
|
-
|
|
125
|
+
"""Run one evaluation episode in each environment
|
|
126
|
+
|
|
127
|
+
Steps every environment until *all* of them are done, ignoring the
|
|
128
|
+
rewards of an environment once its episode has ended. Does not update
|
|
129
|
+
`history` nor `best_actor` (see `run_if_needed`).
|
|
130
|
+
|
|
131
|
+
:param actor: The actor to evaluate (its `act`, deterministic,
|
|
132
|
+
behavior is used)
|
|
133
|
+
:returns: The cumulated reward of each episode, shape `[num_envs]`
|
|
134
|
+
"""
|
|
107
135
|
obs = self.env.reset()
|
|
108
136
|
cumulated = torch.zeros(self.env.num_envs)
|
|
109
137
|
finished = torch.zeros(self.env.num_envs, dtype=torch.bool)
|
|
@@ -117,7 +145,16 @@ class Evaluator(Generic[ActionT]):
|
|
|
117
145
|
return cumulated
|
|
118
146
|
|
|
119
147
|
def visualize_best(self):
|
|
120
|
-
"""Record and display a video of the best actor
|
|
148
|
+
"""Record and display a video of the best actor
|
|
149
|
+
|
|
150
|
+
The video is written under `<run_dir>/videos` (the working directory
|
|
151
|
+
if no `run_dir` was given), with the same wrappers and environment
|
|
152
|
+
arguments as the evaluation environments.
|
|
153
|
+
|
|
154
|
+
:returns: An `IPython.display.Video` in a notebook, `None` in a script
|
|
155
|
+
(where the path is printed instead)
|
|
156
|
+
:raises AssertionError: If no evaluation has been run yet
|
|
157
|
+
"""
|
|
121
158
|
assert self.best_actor is not None, "No actor has been evaluated yet"
|
|
122
159
|
directory = (self.run_dir or Path(".")) / "videos"
|
|
123
160
|
path = record_video(
|
|
@@ -135,9 +172,17 @@ def record_video(
|
|
|
135
172
|
) -> Path | None:
|
|
136
173
|
"""Record a video of one episode played by `actor` (deterministic mode)
|
|
137
174
|
|
|
175
|
+
A fresh single environment is created with `render_mode="rgb_array"` and
|
|
176
|
+
played until the episode ends; only `actor.act` is used, so the video
|
|
177
|
+
shows the behavior the actor has learned, without exploration noise.
|
|
178
|
+
|
|
179
|
+
:param actor: The actor to film
|
|
180
|
+
:param env_name: The gymnasium environment name
|
|
181
|
+
:param directory: Where to write the video file
|
|
138
182
|
:param wrappers: gymnasium wrappers applied to the environment (must match
|
|
139
183
|
the ones used during training)
|
|
140
|
-
:
|
|
184
|
+
:param kwargs: Extra arguments passed to `gymnasium.make`
|
|
185
|
+
:returns: The path to the recorded video, or `None` if no file was written
|
|
141
186
|
"""
|
|
142
187
|
env = gym.make(env_name, render_mode="rgb_array", **kwargs)
|
|
143
188
|
for wrapper in wrappers:
|