rl-mind 0.1.0__tar.gz

rl_mind-0.1.0/PKG-INFO

@@ -0,0 +1,313 @@
+Metadata-Version: 2.4
+Name: rl-mind
+Version: 0.1.0
+Summary: A small, typed reinforcement-learning toolkit for the Master MIND RL practicals
+Author-email: Benjamin Piwowarski <benjamin@piwowarski.fr>
+Project-URL: Documentation, https://pypi.org/project/rl-mind/
+Keywords: reinforcement-learning,gymnasium,torch,teaching
+Classifier: Intended Audience :: Education
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Education
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: gymnasium[box2d,classic-control,mujoco]>=1.2.0
+Requires-Dist: ipython>=8
+Requires-Dist: mazemdp>=1.3.0
+Requires-Dist: moviepy>=1.0
+Requires-Dist: numpy>=1.26
+Requires-Dist: tensorboard>=2.19
+Requires-Dist: torch>=2.4
+
+# `rl_mind` — a small, typed RL toolkit
+
+`rl_mind` is the minimal reinforcement-learning library used by the RL
+practicals. It replaces the heavier BBRL machinery with a handful of **typed**
+building blocks: every piece of data an agent exchanges with its environment is
+a plain (frozen) dataclass of `torch.Tensor`s, so there are no untyped string
+dictionaries and the editor can autocomplete field names.
+
+The library is published on PyPI as
+[`rl-mind`](https://pypi.org/project/rl-mind/) (`pip install rl-mind`; the
+practicals install it for you) and the notebooks simply import it. Importing the
+package has **no side effect** — in particular the extra gymnasium environments
+are only registered by an explicit `import rl_mind.envs` (see below).
+
+- [Actors and actions](#actors-and-actions) (`rl_mind.core`)
+- [Environments](#environments) (`rl_mind.env`, `rl_mind.envs`)
+- [Data containers](#data-containers) (`rl_mind.data`)
+- [Collectors](#collectors) (`rl_mind.collectors`) — with the comparison table
+- [Evaluation](#evaluation) (`rl_mind.evaluation`)
+- [Helpers](#helpers) (`rl_mind.nn`, `rl_mind.notebook`)
+
+---
+
+## Actors and actions
+
+`rl_mind.core`
+
+| Class | Role |
+|---|---|
+| `TensorStruct` | Base class: a frozen dataclass of tensors. Supports tensor-like operations applied field by field — `struct[idx]` (index/slice/mask), `TensorStruct.cat([...])`, `TensorStruct.stack([...])`, `struct.set_(idx, value)`. Recurses into nested `TensorStruct` fields. |
+| `Action` | What an actor returns for a batch of observations. One field: `value` — the action tensor (`[B, action_dim]` for continuous actions, `[B]` for discrete). |
+| `StochasticAction` | An `Action` that also stores `log_prob` (`[B]`), the log-probability of the sampled action. Used by REINFORCE / PPO / SAC. |
+| `Actor[A]` | A `torch.nn.Module` mapping a batch of observations to an action of type `A`. `actor(obs)` is the **training-time** behavior (sampling, exploration noise); `actor.act(obs)` is the **deterministic evaluation** behavior (defaults to `forward().value`). |
+| `ActionT` | The generic type variable (`TypeVar` bound to `Action`) that parameterizes `Actor`, `Transitions`, collectors, ... so the action type flows through the API. |
+
+### `TensorStruct` in practice
+
+`TensorStruct` is the foundation of every data container in the library
+(`Action`, `Transitions`, `Episode`, `Rollout`, ...). The idea: you write a
+plain frozen dataclass whose fields are tensors sharing a common leading (batch)
+dimension, and you get tensor-like operations that apply to **all fields at
+once**, while each field keeps its name and type.
+
+```python
+@dataclass(frozen=True)
+class Transitions(TensorStruct):
+    obs: Tensor          # [N, obs_dim]
+    action: Action       # a nested TensorStruct
+    reward: Tensor       # [N]
+    next_obs: Tensor     # [N, obs_dim]
+    terminated: Tensor   # [N] (bool)
+```
+
+**Indexing / slicing / masking** — `struct[index]` applies `index` to every
+tensor field along the batch dimension and returns a new struct. `index` can be
+anything a tensor accepts:
+
+```python
+batch = buffer.sample(64)          # a Transitions with len(batch) == 64
+batch[0]                           # int      -> a single transition
+batch[:32]                         # slice    -> first 32 transitions
+batch[torch.tensor([0, 5, 9])]     # fancy    -> transitions 0, 5, 9
+batch[~batch.terminated]           # bool mask-> only the non-terminal ones
+len(batch)                         # 64 (size of the leading dimension)
+```
+
+**Concatenating and stacking** — the two class methods build a big struct from
+small ones:
+
+```python
+Transitions.cat([chunk_a, chunk_b])   # concatenate along the batch dim: N_a + N_b
+Action.stack([a0, a1, a2])            # add a NEW leading dim: 3 actions -> [3, ...]
+```
+
+`stack` is exactly how the collectors turn a list of per-step actions into a
+time-indexed `[T, ...]` tensor; `cat` is how `TransitionCollector` merges the
+per-step chunks it records into one flat batch.
+
+**Nesting recurses automatically** — a field that is itself a `TensorStruct`
+(here `action`) is sliced/stacked along with the rest, so the alignment between
+observations and actions can never drift:
+
+```python
+sub = batch[mask]        # slices batch.obs AND batch.action.value together
+sub.action.log_prob      # still lined up with sub.obs, sub.reward, ...
+```
+
+**In-place writes** — instances are frozen (immutable), so `[]`, `cat` and
+`stack` all return *new* structs. The one mutating operation is `set_(index,
+value)`, used by `ReplayBuffer` to overwrite slots of its preallocated storage:
+
+```python
+storage.set_(indices, transitions)   # write a batch of transitions at `indices`
+```
+
+Because every container shares this behaviour, the learning code reads the same
+way whether the batch came from a replay buffer, an episode or a rollout —
+`batch.reward`, `batch.action.value`, `batch.terminated` are always named,
+typed, and mutually aligned.
+
+The generic parameter is what makes the typing pay off: an `Actor[StochasticAction]`
+guarantees that the actions it produces carry a `.log_prob`, and the type
+checker will flag a `TransitionCollector[StochasticAction]` whose batches you try
+to use as if they had none.
+
+```python
+class DiscretePolicy(Actor[StochasticAction]):
+    def dist(self, obs): return torch.distributions.Categorical(logits=self.model(obs))
+    def forward(self, obs):
+        d = self.dist(obs); a = d.sample()
+        return StochasticAction(value=a, log_prob=d.log_prob(a))
+    def act(self, obs): return self.model(obs).argmax(-1)   # deterministic
+```
+
+---
+
+## Environments
+
+`rl_mind.env`, `rl_mind.envs`
+
+| Class / symbol | Role |
+|---|---|
+| `VecEnv` | Runs `num_envs` copies of a gymnasium environment in parallel, talking **torch tensors**. `reset()` → `[B, obs_dim]`; `step(actions)` → `EnvStep`. Exposes `observation_dim` (flat `Box` spaces) or `n_states` (tabular `Discrete` spaces), `action_dim` / `n_actions`, `is_continuous`, `num_envs`, `env_name`, `same_step_reset`. Extra keyword arguments (and `wrappers=`) are forwarded to each sub-environment. |
+| `EnvStep` | Result of one step (all fields `[B, ...]`): `obs` (what to act on next), `next_obs` (the true successor $s_{t+1}$), `reward`, `terminated`, `truncated`, and the derived `done = terminated | truncated`. |
+| `ContinuousCartPoleEnv` | `CartPole-v1` with a continuous force action in $[-1, 1]$. **Importing `rl_mind.envs` registers `CartPoleContinuous-v1`** in gymnasium (the opt-in side effect). |
+
+### Observation spaces
+
+`VecEnv` adapts what it returns to the gymnasium observation space:
+
+- **`Box`** (the usual case) → a `[B, obs_dim]` float tensor; `observation_dim`
+  gives `obs_dim`.
+- **`Discrete`** (tabular environments) → a `[B]` tensor of **state indices**
+  (`torch.long`), which index a Q-table directly (`q_table[obs, actions]`);
+  `n_states` gives the number of states.
+- **`Dict`** (structured observations) → a `TensorStruct` with one named tensor
+  field per key; inspect `observation_space` to write the encoder.
+
+### `terminated` vs `truncated`
+
+Both flags end an episode, but they mean different things for learning:
+
+- **`terminated`** — a real terminal state (the pole fell). The future is worth
+  0, so you **do not bootstrap**.
+- **`truncated`** — the episode was cut short, e.g. a time limit. The agent
+  could have continued, so you **do bootstrap** with the value of `next_obs`.
+
+### Auto-reset modes
+
+When an episode ends, `VecEnv` resets it automatically, in one of two modes:
+
+- **next-step reset** (default): the ending step returns the episode's final
+  observation; the *following* `step` ignores its action and returns the first
+  observation of a fresh episode.
+- **same-step reset** (`same_step_reset=True`): the ending step already returns
+  the fresh episode's first observation in `obs`, while `next_obs` holds the
+  final observation of the episode that just ended. **Every step is then a valid
+  transition** — this is what `RolloutCollector` needs.
+
+---
+
+## Data containers
+
+`rl_mind.data`
+
+| Class / symbol | Shape | Role |
+|---|---|---|
+| `Transitions[A]` | flat `[N, ...]` | A batch of independent transitions $(s, a, r, s', \text{terminated})$: `obs`, `action`, `reward`, `next_obs`, `terminated`. Off-policy data. |
+| `ReplayBuffer[A]` | — | A fixed-capacity ring buffer of `Transitions`. `add(transitions)`, `sample(batch_size)` (uniform, with replacement), `len(buffer)`. |
+| `Episode[A]` | time `[T, ...]` | One full episode: `obs` (`[T, obs_dim]`), `action` (`[T, ...]`), `reward` (`[T]`), `final_obs` (`[obs_dim]`), `terminated` (bool). Plus `len(ep)`, `ep.cumulated_reward` and `ep.all_obs` (`[T+1, obs_dim]`): the observations $s_0, \ldots, s_T$, i.e. `obs` with `final_obs` appended — `all_obs[1:]` are the successors of `obs`, and `critic(all_obs)` gives $V(s_0), \ldots, V(s_T)$. |
+| `Rollout[A]` | time×env `[T, B, ...]` | A fixed-length on-policy segment: `obs`, `action`, `reward`, `next_obs`, `terminated`, `truncated`, the derived `done`, and `flatten()` → `[T*B, ...]`. |
+| `minibatches(data, size)` | — | Iterate over random minibatches of any `TensorStruct`, using each sample **exactly once** per pass (the last batch may be smaller). |
+
+`Transitions` and `Rollout` are `TensorStruct`s, so indexing, slicing and
+stacking work uniformly: `batch.action.log_prob`, `rollout[t]`,
+`transitions[mask]`, ... `ReplayBuffer` wraps one; `Episode` is a plain
+dataclass, because its fields do not share a batch dimension (`final_obs` has
+no time axis and `terminated` is a Python `bool`).
+
+---
+
+## Collectors
+
+`rl_mind.collectors`
+
+The three collectors all run an actor in a `VecEnv` and keep the environment
+state across successive `.collect()` calls, counting total steps in `.steps`.
+They differ in the *shape* of data they produce, matching the three families of
+algorithms:
+
+| | `TransitionCollector` | `EpisodeCollector` | `RolloutCollector` |
+|---|---|---|---|
+| Returns | `Transitions` — flat `[N, …]` | `list[Episode]` | `Rollout` — time×env `[T, B, …]` |
+| `collect(...)` arg | `n_steps` | `n_episodes` | `n_steps` |
+| For | off-policy (DQN, DDPG, TD3, SAC) | episodic (REINFORCE) | on-policy (A2C, PPO) |
+| Time structure | none (goes to a shuffled buffer) | whole episodes | preserved (needed for GAE) |
+| `VecEnv` reset mode | next-step (default) | next-step (default) | requires `same_step_reset=True` |
+| Length | `≤ n_steps × num_envs` (reset steps dropped) | `≥ n_episodes` whole episodes | exactly `[n_steps, num_envs]` |
+| Ending flags kept | `terminated` | `terminated` (per episode) | `terminated` **and** `truncated` |
+| Env state across calls | kept | reset at each call (on-policy) | kept |
+
+### Why two step-based collectors (the subtle part)
+
+`TransitionCollector` and `RolloutCollector` both walk a fixed number of steps,
+but they treat episode boundaries differently — which is exactly why both exist:
+
+- **`TransitionCollector`** produces an *unordered bag* of transitions that will
+  be shuffled in a replay buffer. When an episode ends under next-step reset, the
+  following "reset" step is invalid, so the collector simply **drops** those rows
+  (hence `≤ n_steps × num_envs`). Order doesn't matter, so holes are fine.
+
+- **`RolloutCollector`** produces a *rectangular `[T, B]` block* whose time axis
+  must stay intact: on-policy algorithms compute GAE as a **backward recursion
+  over time**, and need to know at each step whether the episode ended. Dropping
+  rows would punch holes in the grid, so it instead requires
+  `same_step_reset=True`: the environment resets *within* the ending step, so
+  every row is a valid transition (`obs` = fresh state, `next_obs` = true final
+  state) and the block stays dense. This is why `Rollout` also carries
+  `truncated` — GAE must stop propagating across boundaries.
+
+`EpisodeCollector` is the odd one out: it returns *variable-length whole
+episodes* and resets the environments at the start of each `collect()`, so the
+episodes are strictly on-policy (all collected with the current actor).
+
+---
+
+## Evaluation
+
+`rl_mind.evaluation`
+
+| Class / symbol | Role |
+|---|---|
+| `Evaluator[A]` | Periodically evaluates the current actor on a separate `VecEnv` (using `actor.act`, the deterministic behavior) and keeps a **copy of the best actor so far**. Call `run_if_needed(steps, actor)` in the training loop — it no-ops until `every` steps have passed. Exposes `best_actor`, `best_reward`, `history`, an optional tensorboard `writer`, and `visualize_best()`. |
+| `EvalResult` | One evaluation: `step`, `rewards` (`[n_eval_envs]`), `is_best`, and the derived `.mean`. `Evaluator.history` is a list of these — handy for learning-curve plots and Welch t-tests. |
+| `record_video(actor, env_name, directory)` | Record a video of one deterministic episode and return the video path. |
+
+The evaluation environment is intentionally *separate* from the training env
+(different seed, its own episode count, no `same_step_reset`) so that evaluation
+is independent of the data the agent is training on.
+
+---
+
+## Helpers
+
+`rl_mind.nn`
+
+| Symbol | Role |
+|---|---|
+| `build_mlp(sizes, activation=ReLU(), output_activation=None)` | Build a `nn.Sequential` MLP from a list of layer sizes. |
+| `soft_update(source, target, tau)` | Polyak update of a target network: $\theta' \leftarrow \tau\theta + (1-\tau)\theta'$. |
+
+`rl_mind.notebook`
+
+| Symbol | Role |
+|---|---|
+| `run_directory(name)` | Create and return a fresh timestamped output directory for a run (under `outputs/`, or `outputs-testing/` in test mode). |
+| `outputs_directory()` | The base output directory (test-mode aware). |
+| `setup_tensorboard()` | Show the tensorboard dashboard inline (Jupyter, Colab), or print the command to launch it from a shell. Warns if the `tensorboard` package is missing, and always prints the **absolute** log directory (it is `outputs/` relative to the kernel's working directory). |
+| `silence_known_warnings()` | Silence the `pkg_resources` deprecation warning emitted by `pygame` and `tensorboard`. |
+| `video_display(path)` | Display a video in the notebook, or print its path when run as a script. |
+| `is_notebook()` | True when running inside Jupyter / Colab. |
+
+---
+
+## A minimal off-policy loop
+
+```python
+import rl_mind.envs  # register CartPoleContinuous-v1 (explicit opt-in)
+from rl_mind.env import VecEnv
+from rl_mind.data import ReplayBuffer
+from rl_mind.collectors import TransitionCollector
+from rl_mind.evaluation import Evaluator
+
+env = VecEnv("CartPoleContinuous-v1", num_envs=1, seed=1)
+collector = TransitionCollector(env, GaussianNoise(actor, sigma=0.1))
+buffer = ReplayBuffer(200_000)
+evaluator = Evaluator(VecEnv("CartPoleContinuous-v1", 10, seed=101), every=2_000)
+
+while collector.steps < 30_000:
+    buffer.add(collector.collect(1))
+    if len(buffer) < 1_000:
+        continue
+    batch = buffer.sample(64)              # Transitions[Action]
+    target = batch.reward + gamma * next_q * (~batch.terminated).float()
+    ...                                    # critic / actor updates
+    evaluator.run_if_needed(collector.steps, actor)
+```
+
+The on-policy notebooks (A2C, PPO) swap the replay buffer for a
+`RolloutCollector` + `minibatches`; REINFORCE uses an `EpisodeCollector`.

rl_mind-0.1.0/README.md

@@ -0,0 +1,85 @@
+# BBRL Notebooks
+
+## Notebooks
+
+The notebooks are the python files in [percent
+format](https://jupytext.readthedocs.io/en/latest/formats.html) which are
+transformed into student (with a Colab version) and teacher versions using
+special markers (as described below). Images are inlined.
+
+### Imports
+
+To regroup all local imports, use 
+
+```py3
+# %% tags=["imports"]
+import numpy as np
+import torch.nn as nn
+```
+
+And in the place where all the imports should be placed, put this instruction
+```py3
+# [[imports]]
+```
+
+### Collapsing a cell
+
+To collapse a cell in the notebook, use the `hide-input` tag
+
+```md
+# %% tags=["hide-input"]
+
+some_not_interesting_code()
+```
+
+### Including a file
+
+To include all the cells from another `py:percent` file, use 
+a `copy` tag with a `from ... import ...`
+
+```py3
+# %% tags=["copy"]
+
+from common.env import *
+```
+
+### Filtering cells
+
+Tags can be used to filter cells depending on the output:
+
+- `teacher` removes the cell when generating a student version
+- `colab` removes the cell when not generating a google Colab version
+- `not-colab` removes the cell when generating a google Colab version
+
+
+### Filtering within cells
+
+Within a cell, special markers can be used to transform the text.
+
+To remove altogether the content, use `[[remove]]`:
+
+```py3
+# [[remove]]
+test_mode = True
+# [[/remove]]
+```
+
+To remove the content only for student, use `[[student]]`:
+```py3
+# [[student]] Your code to obtain an optimal policy here
+def value_iteration_q(mdp: MazeMDPEnv, render: bool = True) -> Tuple[np.ndarray, List[float]]:
+    q = np.zeros((mdp.nb_states, mdp.action_space.n))  # initial action values are set to 0
+    q_list = []
+    # ># Set stop to False
+    # >stop = ...
+    stop = False
+    ...
+# [[/student]]
+```
+is transformed into
+```py3
+# Your code to obtain an optimal policy here
+assert False, "Your code to obtain an optimal policy here"
+# Set stop to False
+stop =
+```

rl_mind-0.1.0/pyproject.toml

@@ -0,0 +1,74 @@
+[project]
+name = "rl-mind"
+version = "0.1.0"
+description = "A small, typed reinforcement-learning toolkit for the Master MIND RL practicals"
+readme = "rl_mind.md"
+authors = [{ name = "Benjamin Piwowarski", email = "benjamin@piwowarski.fr" }]
+keywords = ["reinforcement-learning", "gymnasium", "torch", "teaching"]
+classifiers = [
+    "Intended Audience :: Education",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Education",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+requires-python = ">=3.11"
+# Runtime dependencies of the library only (what `import rl_mind` needs). The
+# gymnasium extras cover every environment used by the practicals, so a bare
+# `pip install rl-mind` is enough to run the notebooks.
+dependencies = [
+    "gymnasium[mujoco,classic-control,box2d]>=1.2.0",
+    "ipython>=8",
+    "mazemdp>=1.3.0",
+    "moviepy>=1.0",
+    "numpy>=1.26",
+    "tensorboard>=2.19",
+    "torch>=2.4",
+]
+
+[project.urls]
+Documentation = "https://pypi.org/project/rl-mind/"
+
+[build-system]
+# swig is not needed to build rl-mind itself: it is listed here so that the
+# generated Colab install cell installs it first (box2d-py builds from source).
+requires = ["swig", "setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["rl_mind*"]
+
+[tool.setuptools.package-data]
+rl_mind = ["py.typed"]
+
+[tool.pytest.ini_options]
+# Only our own tests: `stk-competition/repositories/` holds student projects
+# (unrelated dependencies, not collectable)
+testpaths = ["tests"]
+
+[dependency-groups]
+# What the notebooks need on top of the library (jupyter, plotting, the
+# instructor build helper, ...). Installed by `uv sync` (default groups) but
+# not by `pip install rl-mind`.
+notebooks = [
+    "hydra-core>=1.3.2",
+    "jupyter>=1.1.1",
+    "jupyterlab>=4.4.4",
+    "jupytext-notebook-helper>=0.4.1",
+    "matplotlib>=3.8",
+    "optuna>=4.9.0",
+    "pip>=25.2",
+    "scipy>=1.11",
+    "tqdm>=4.66",
+]
+dev = [
+    "imgcat>=0.5.0",
+    "jupytext>=1.17.2",
+    "pre-commit>=4.2.0",
+    "pytest>=8",
+]
+
+[tool.uv]
+default-groups = ["notebooks", "dev"]