`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.Tensors, so there are no untyped string dictionaries and the editor can autocomplete field names.

The library is published on PyPI as rl-mind (pip install rl-mind; the practicals install it for you) and the notebooks simply import it. This page is the guide; every class, method and attribute is documented in the API reference. 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) — except one: it sets SDL_NO_SIGNAL_HANDLERS=1 (unless already set), so that a kernel that rendered with pygame still stops on SIGTERM.

Actors and actions (rl_mind.core)
Environments (rl_mind.env, rl_mind.envs)
Data containers (rl_mind.data)
Collectors (rl_mind.collectors) — with the comparison table
Evaluation (rl_mind.evaluation)
Rendering (rl_mind.rendering)
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.

@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:

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:

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:

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:

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.

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
`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 TensorStructs, 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()` (plays an episode of the best actor with `rollout_frames`/`animate_frames` in a notebook; `visualize_best(video=True)`, or any call from a script, records a video instead — embedded, it stays in the saved notebook).
`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.
`actor_policy(actor, action_space)`	Turn an actor into a single-environment `obs -> action` policy (deterministic `act`), e.g. for `rollout_frames`.

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.

Rendering¶

rl_mind.rendering

Class / symbol	Role
`RenderWidget(env, k=0)`	Wraps an env (`render_mode="rgb_array"`) so every `reset`/`step` renders a frame and pushes it to a live `ipywidgets.Image`, with a slider/`Play` button to scrub back through history. Accepts either a single `gym.Env`, or a `VecEnv` (which steps every copy together) — `k` then picks which copy to render, via `VecEnv.gym_unwrapped`. Attribute access is forwarded to the wrapped env.
`FrameViewer`	The widget `RenderWidget` is built on, usable on its own for a frame source with neither API: call `update(frame)` each time a new frame is available and `reset()` at the start of an episode.
`Trimmer`	Base class bounding how many `(step, frame)` entries a history keeps. Subclass and override `trim` for custom strategies.
`WindowTrimmer(size)`	Sliding window: keeps only the last `size` entries.
`DecimateTrimmer(capacity, keep_ratio=0.5)`	Bounded memory, but spans the whole history: once `capacity` is hit, subsamples with a stride so the buffer settles around `capacity * keep_ratio` entries — older parts get sparser rather than being dropped. `keep_ratio` bounds the fraction kept per trim event, not of the whole episode (unbounded episode length in O(1) memory always means older segments thin out).
`rollout_frames(env_name, policy, seed=0, trimmer=None, wrappers=(), **kwargs)`	Play back a fixed episode instead of a live one: step `policy` (`obs -> action`) through one episode and collect `(step, encoded frame)` pairs.
`animate_frames(history, fps=30, width="400px")`	Display a `(step, encoded frame)` history (e.g. from `rollout_frames`) with play/pause, step and rewind controls.
`encode_frame(frame)` / `decode_frame(data)`	Histories keep frames encoded as lossy WebP (~10 KB for a MuJoCo frame, ~700 KB raw), so a whole episode fits in memory without a trimmer; `decode_frame` gives an `[H, W, 3]` array back.

RenderWidget/FrameViewer are for interactively watching an env react to actions chosen live (as in the A01-environments practical); rollout_frames/animate_frames replay a fixed episode, which is what Evaluator.visualize_best (above) uses to show a trained actor — record_video saves one as an actual video file instead.

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()`	In Jupyter, a chooser between external (the default: prints the command to launch tensorboard from a shell) and in the notebook (the dashboard inline); always inline on Colab, always the command from a script. 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¶

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 — a small, typed RL toolkit¶