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. 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 (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)
• Helpers (rl_mind.nn, rl_mind.notebook)

Actors and actions

rl_mind.core

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

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

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:

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

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

rl_mind.notebook

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.