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`.
