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.
Files changed (25) hide show
  1. {rl_mind-0.2.2 → rl_mind-0.2.4}/PKG-INFO +6 -4
  2. {rl_mind-0.2.2 → rl_mind-0.2.4}/README.md +5 -0
  3. {rl_mind-0.2.2 → rl_mind-0.2.4}/pyproject.toml +5 -3
  4. {rl_mind-0.2.2 → rl_mind-0.2.4}/rl_mind.md +3 -1
  5. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/collectors.py +77 -9
  6. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/core.py +58 -8
  7. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/data.py +53 -6
  8. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/env.py +26 -5
  9. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/envs.py +18 -1
  10. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/evaluation.py +52 -7
  11. rl_mind-0.2.4/src/rl_mind/mazemdp_env.py +317 -0
  12. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/nn.py +10 -1
  13. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/notebook.py +17 -2
  14. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/rendering.py +71 -6
  15. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/PKG-INFO +6 -4
  16. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/requires.txt +1 -1
  17. rl_mind-0.2.2/src/rl_mind/mazemdp_env.py +0 -182
  18. {rl_mind-0.2.2 → rl_mind-0.2.4}/setup.cfg +0 -0
  19. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/__init__.py +0 -0
  20. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/profiles.py +0 -0
  21. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind/py.typed +0 -0
  22. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/SOURCES.txt +0 -0
  23. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/dependency_links.txt +0 -0
  24. {rl_mind-0.2.2 → rl_mind-0.2.4}/src/rl_mind.egg-info/top_level.txt +0 -0
  25. {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.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://pypi.org/project/rl-mind/
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.0
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. Importing the
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.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.0",
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
- Documentation = "https://pypi.org/project/rl-mind/"
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. Importing the
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`s (for episodic
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
- Returns the recorded transitions (at most `n_steps * num_envs`;
60
- steps where an environment was auto-resetting are skipped).
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; returns the
170
- corresponding `[n_steps, num_envs]` rollout"""
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
- `index` can be anything a tensor accepts: an integer, a slice, an
52
- index tensor or a boolean mask.
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 (shape `[B, obs_dim]`)"""
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 (defaults to `forward().value`)"""
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 in the episode"""
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; returns the initial observations
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
- Shape `[B, obs_dim]` for a flat space, or a `TensorStruct` of `[B, ...]`
250
- tensors for a `Dict` space.
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: `[B, action_dim]` (continuous) or `[B]` (discrete)
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 at least `every` steps have passed since the
86
- last evaluation; returns the result (or None if no evaluation was run)
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; returns the
106
- cumulated rewards (shape `[num_envs]`)"""
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
- :returns: the path to the recorded video
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: