rl-mind 0.1.0__py3-none-any.whl

rl_mind/__init__.py

@@ -0,0 +1,55 @@
+"""Typed, minimal RL infrastructure for the Master MIND RL practicals.
+
+Published on PyPI as ``rl-mind``: the notebooks import it like any other
+package. See ``rl_mind.md`` (the package README) for the documentation.
+"""
+
+from rl_mind.collectors import (
+    EpisodeCollector,
+    RolloutCollector,
+    TransitionCollector,
+)
+from rl_mind.core import Action, ActionT, Actor, StochasticAction, TensorStruct
+from rl_mind.data import Episode, ReplayBuffer, Rollout, Transitions, minibatches
+
+# NOTE: rl_mind.envs is deliberately NOT imported here: importing it registers
+# the extra gymnasium environments, and importing the package must have no
+# side effect. Use an explicit `import rl_mind.envs` where needed.
+from rl_mind.env import EnvStep, VecEnv
+from rl_mind.evaluation import EvalResult, Evaluator, record_video
+from rl_mind.nn import build_mlp, soft_update
+from rl_mind.notebook import (
+    is_notebook,
+    outputs_directory,
+    run_directory,
+    setup_tensorboard,
+    video_display,
+)
+
+__all__ = [
+    "Action",
+    "ActionT",
+    "Actor",
+    "EnvStep",
+    "EpisodeCollector",
+    "EvalResult",
+    "Evaluator",
+    "Episode",
+    "ReplayBuffer",
+    "Rollout",
+    "RolloutCollector",
+    "StochasticAction",
+    "TensorStruct",
+    "Transitions",
+    "TransitionCollector",
+    "VecEnv",
+    "build_mlp",
+    "is_notebook",
+    "minibatches",
+    "outputs_directory",
+    "record_video",
+    "run_directory",
+    "setup_tensorboard",
+    "soft_update",
+    "video_display",
+]

rl_mind/collectors.py

@@ -0,0 +1,195 @@
+"""Collecting interaction data.
+
+The collectors run an actor in a vectorized environment and package the
+result into the containers of `rl_mind.data`. They deal with the
+auto-reset bookkeeping so that the learning code never has to.
+
+- `TransitionCollector.collect(n_steps)` performs `n_steps` steps in
+  *every* environment and returns the corresponding batch of
+  `Transitions` (for replay-buffer based algorithms). The environments
+  keep their state between calls.
+- `EpisodeCollector.collect(n_episodes)` resets the environments and
+  returns at least `n_episodes` complete `Episode`s (for episodic
+  algorithms, which need *on-policy* full episodes).
+- `RolloutCollector.collect(n_steps)` returns a `[T, B]` `Rollout` (for
+  on-policy algorithms such as A2C and PPO). It requires an environment
+  with `same_step_reset=True`, so that every row of the rollout is a
+  valid transition. The environments keep their state between calls.
+
+All count the total number of recorded environment steps in `steps`.
+Actions are computed with `torch.no_grad()`: if the learning algorithm
+needs gradients (e.g. of action log-probabilities), it must recompute them
+from the stored observations and actions.
+"""
+
+from rl_mind.core import ActionT, Actor, TensorStruct
+from rl_mind.data import Episode, Rollout, Transitions
+from rl_mind.env import EnvStep, VecEnv
+
+from typing import Generic
+
+import torch
+
+__all__ = ["TransitionCollector", "EpisodeCollector", "RolloutCollector"]
+
+
+def _stack(items):
+    """Stack a list of observations (plain tensors, or `TensorStruct`s for
+    `Dict` observation spaces) along a new leading dimension."""
+    first = items[0]
+    if isinstance(first, TensorStruct):
+        return type(first).stack(items)
+    return torch.stack(items)
+
+
+class TransitionCollector(Generic[ActionT]):
+    """Collect batches of transitions by running an actor in a `VecEnv`"""
+
+    def __init__(self, env: VecEnv, actor: Actor[ActionT]):
+        self.env = env
+        self.actor = actor
+        #: Total number of recorded transitions
+        self.steps = 0
+        self._obs: torch.Tensor | None = None
+        self._just_reset = torch.zeros(env.num_envs, dtype=torch.bool)
+
+    def collect(self, n_steps: int) -> Transitions[ActionT]:
+        """Perform `n_steps` steps in every environment
+
+        Returns the recorded transitions (at most `n_steps * num_envs`;
+        steps where an environment was auto-resetting are skipped).
+        """
+        if self._obs is None:
+            self._obs = self.env.reset()
+
+        chunks: list[Transitions[ActionT]] = []
+        with torch.no_grad():
+            for _ in range(n_steps):
+                action = self.actor(self._obs)
+                step = self.env.step(action.value)
+
+                # Do not record the transitions of environments that were
+                # auto-resetting during this step (their action was ignored)
+                valid = ~self._just_reset
+                chunks.append(
+                    Transitions(
+                        obs=self._obs,
+                        action=action,
+                        reward=step.reward,
+                        next_obs=step.next_obs,
+                        terminated=step.terminated,
+                    )[valid]
+                )
+
+                self._obs = step.obs
+                if not self.env.same_step_reset:
+                    # With next-step reset, the step following `done` is the
+                    # (invalid) auto-reset step
+                    self._just_reset = step.done
+
+        transitions = Transitions.cat(chunks)
+        self.steps += len(transitions)
+        return transitions
+
+
+class EpisodeCollector(Generic[ActionT]):
+    """Collect full episodes by running an actor in a `VecEnv`"""
+
+    def __init__(self, env: VecEnv, actor: Actor[ActionT]):
+        self.env = env
+        self.actor = actor
+        #: Total number of recorded steps
+        self.steps = 0
+
+    def collect(self, n_episodes: int) -> list[Episode[ActionT]]:
+        """Collect at least `n_episodes` complete episodes
+
+        The environments are reset at the beginning of the call, so that all
+        the returned episodes are collected with the *current* actor.
+        """
+        num_envs = self.env.num_envs
+        obs = self.env.reset()
+        just_reset = torch.zeros(num_envs, dtype=torch.bool)
+
+        # Per-environment lists of observations, actions and rewards
+        obs_lists: list[list[torch.Tensor]] = [[] for _ in range(num_envs)]
+        action_lists: list[list[ActionT]] = [[] for _ in range(num_envs)]
+        reward_lists: list[list[torch.Tensor]] = [[] for _ in range(num_envs)]
+
+        episodes: list[Episode[ActionT]] = []
+        with torch.no_grad():
+            while len(episodes) < n_episodes:
+                action = self.actor(obs)
+                step = self.env.step(action.value)
+
+                for i in range(num_envs):
+                    if just_reset[i]:
+                        # This environment was auto-resetting: nothing to record
+                        continue
+                    obs_lists[i].append(obs[i])
+                    action_lists[i].append(action[i])
+                    reward_lists[i].append(step.reward[i])
+                    self.steps += 1
+
+                    if step.done[i]:
+                        episodes.append(
+                            Episode(
+                                obs=_stack(obs_lists[i]),
+                                action=type(action).stack(action_lists[i]),
+                                reward=torch.stack(reward_lists[i]),
+                                final_obs=step.next_obs[i],
+                                terminated=bool(step.terminated[i]),
+                            )
+                        )
+                        obs_lists[i], action_lists[i], reward_lists[i] = [], [], []
+
+                obs = step.obs
+                if not self.env.same_step_reset:
+                    just_reset = step.done
+
+        return episodes
+
+
+class RolloutCollector(Generic[ActionT]):
+    """Collect fixed-length `[T, B]` rollouts by running an actor in a `VecEnv`"""
+
+    def __init__(self, env: VecEnv, actor: Actor[ActionT]):
+        assert env.same_step_reset, (
+            "RolloutCollector requires a VecEnv created with "
+            "same_step_reset=True (so that every rollout row is a valid "
+            "transition)"
+        )
+        self.env = env
+        self.actor = actor
+        #: Total number of recorded steps
+        self.steps = 0
+        self._obs: torch.Tensor | None = None
+
+    def collect(self, n_steps: int) -> Rollout[ActionT]:
+        """Perform `n_steps` steps in every environment; returns the
+        corresponding `[n_steps, num_envs]` rollout"""
+        if self._obs is None:
+            self._obs = self.env.reset()
+
+        obs_list: list[torch.Tensor] = []
+        action_list: list[ActionT] = []
+        steps: list[EnvStep] = []
+        with torch.no_grad():
+            for _ in range(n_steps):
+                action = self.actor(self._obs)
+                step = self.env.step(action.value)
+
+                obs_list.append(self._obs)
+                action_list.append(action)
+                steps.append(step)
+                self._obs = step.obs
+
+        self.steps += n_steps * self.env.num_envs
+        return Rollout(
+            obs=_stack(obs_list),
+            action=type(action_list[0]).stack(action_list),
+            reward=torch.stack([step.reward for step in steps]),
+            next_obs=_stack([step.next_obs for step in steps]),
+            terminated=torch.stack([step.terminated for step in steps]),
+            truncated=torch.stack([step.truncated for step in steps]),
+        )

rl_mind/core.py

@@ -0,0 +1,152 @@
+"""Actions, actors and tensor containers.
+
+All the data exchanged between the agent and the environment is stored in
+**typed containers**: frozen dataclasses whose fields are ``torch.Tensor``s
+(`TensorStruct`). They can be indexed (``batch[idx]``), concatenated and
+stacked like tensors, but each field keeps its name and its type.
+
+- An `Action` is what an `Actor` returns: at minimum the action tensor itself
+  (``value``). Subclasses can carry more information, e.g. a
+  `StochasticAction` also stores the log-probability of the sampled action.
+- An `Actor` is a ``torch.nn.Module`` mapping a batch of observations to an
+  `Action`. Calling ``actor(obs)`` gives the *training-time* behavior
+  (sampling, exploration noise); ``actor.act(obs)`` gives the *evaluation*
+  action (deterministic).
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields, replace
+from typing import Callable, Generic, Self, Sequence, TypeVar
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["TensorStruct", "Action", "StochasticAction", "Actor", "ActionT"]
+
+
+@dataclass(frozen=True)
+class TensorStruct:
+    """A frozen dataclass of tensors (possibly nested).
+
+    Provides tensor-like operations applied field by field: indexing,
+    concatenation, stacking. All tensor fields must share the same first
+    (batch) dimension.
+    """
+
+    def _map(self, fn: Callable[[Tensor], Tensor]) -> Self:
+        """Return a copy where `fn` has been applied to every tensor field"""
+        values = {}
+        for field in fields(self):
+            value = getattr(self, field.name)
+            if isinstance(value, TensorStruct):
+                values[field.name] = value._map(fn)
+            elif isinstance(value, Tensor):
+                values[field.name] = fn(value)
+        return replace(self, **values)
+
+    def __getitem__(self, index) -> Self:
+        """Index or slice all the tensor fields along the batch dimension
+
+        `index` can be anything a tensor accepts: an integer, a slice, an
+        index tensor or a boolean mask.
+        """
+        return self._map(lambda tensor: tensor[index])
+
+    def __len__(self) -> int:
+        """The size of the batch dimension"""
+        for field in fields(self):
+            value = getattr(self, field.name)
+            if isinstance(value, TensorStruct):
+                return len(value)
+            if isinstance(value, Tensor):
+                return value.shape[0]
+        raise ValueError(f"{type(self).__name__} has no tensor field")
+
+    def set_(self, index, value: Self) -> None:
+        """Write `value` at position(s) `index` (in-place)"""
+        for field in fields(self):
+            target = getattr(self, field.name)
+            if isinstance(target, TensorStruct):
+                target.set_(index, getattr(value, field.name))
+            elif isinstance(target, Tensor):
+                target[index] = getattr(value, field.name)
+
+    @classmethod
+    def cat(cls, items: Sequence[Self]) -> Self:
+        """Concatenate a sequence of structures along the batch dimension"""
+        values = {}
+        for field in fields(items[0]):
+            value = getattr(items[0], field.name)
+            if isinstance(value, TensorStruct):
+                values[field.name] = type(value).cat(
+                    [getattr(item, field.name) for item in items]
+                )
+            elif isinstance(value, Tensor):
+                values[field.name] = torch.cat(
+                    [getattr(item, field.name) for item in items]
+                )
+        return replace(items[0], **values)
+
+    @classmethod
+    def stack(cls, items: Sequence[Self]) -> Self:
+        """Stack a sequence of structures along a new (first) dimension"""
+        values = {}
+        for field in fields(items[0]):
+            value = getattr(items[0], field.name)
+            if isinstance(value, TensorStruct):
+                values[field.name] = type(value).stack(
+                    [getattr(item, field.name) for item in items]
+                )
+            elif isinstance(value, Tensor):
+                values[field.name] = torch.stack(
+                    [getattr(item, field.name) for item in items]
+                )
+        return replace(items[0], **values)
+
+
+@dataclass(frozen=True)
+class Action(TensorStruct):
+    """What an actor returns: a batch of actions
+
+    :param value: The action tensor, shape `[B, action_dim]` (continuous
+        actions) or `[B]` (discrete actions)
+    """
+
+    value: Tensor
+
+
+@dataclass(frozen=True)
+class StochasticAction(Action):
+    """An action sampled from a distribution
+
+    :param log_prob: The log-probability of the sampled action, shape `[B]`
+    """
+
+    log_prob: Tensor
+
+
+#: Type of the action returned by an actor
+ActionT = TypeVar("ActionT", bound=Action)
+
+
+class Actor(nn.Module, ABC, Generic[ActionT]):
+    """Maps a batch of observations to actions
+
+    An actor is a `torch.nn.Module`; subclasses implement `forward` (the
+    training-time behavior, e.g. sampling) and can override `act` (the
+    evaluation-time behavior, deterministic by default).
+    """
+
+    @abstractmethod
+    def forward(self, obs: Tensor) -> ActionT:
+        """Compute the action for a batch of observations (shape `[B, obs_dim]`)"""
+        ...
+
+    def act(self, obs: Tensor) -> Tensor:
+        """Deterministic action used for evaluation (defaults to `forward().value`)"""
+        return self(obs).value
+
+    # Overriden so that type checkers know the return type of `actor(obs)`
+    def __call__(self, obs: Tensor) -> ActionT:
+        return super().__call__(obs)

rl_mind/data.py

@@ -0,0 +1,179 @@
+"""Transitions, episodes, rollouts and the replay buffer.
+
+Two ways of storing interaction data, depending on the algorithm:
+
+- `Transitions` is a flat batch of $(s, a, r, s', \textrm{terminated})$
+  tuples, used by *off-policy* algorithms (DQN, DDPG, TD3, SAC) through a
+  `ReplayBuffer`;
+- `Episode` is one full episode with time-indexed tensors, used by
+  *episodic* algorithms (REINFORCE);
+- `Rollout` is a fixed-length `[T, B]` segment of interaction, used by
+  *on-policy* algorithms (A2C, PPO), with `minibatches` to iterate over
+  its (flattened) steps.
+"""
+
+from rl_mind.core import ActionT, TensorStruct
+
+from dataclasses import dataclass
+from typing import Generic, Iterator, TypeVar
+
+import torch
+from torch import Tensor
+
+__all__ = ["Transitions", "ReplayBuffer", "Episode", "Rollout", "minibatches"]
+
+
+@dataclass(frozen=True)
+class Transitions(TensorStruct, Generic[ActionT]):
+    """A batch of `N` transitions $(s, a, r, s', \\textrm{terminated})$
+
+    :param obs: The observations $s$, shape `[N, obs_dim]`
+    :param action: The actions $a$ (an `Action`, tensors of shape `[N, ...]`)
+    :param reward: The rewards $r$, shape `[N]`
+    :param next_obs: The next observations $s'$, shape `[N, obs_dim]`
+    :param terminated: Whether $s'$ is a terminal state, shape `[N]` (bool).
+        When the episode was *truncated* (time limit), this is `False`: the
+        value of $s'$ must be bootstrapped.
+    """
+
+    obs: Tensor
+    action: ActionT
+    reward: Tensor
+    next_obs: Tensor
+    terminated: Tensor
+
+
+class ReplayBuffer(Generic[ActionT]):
+    """A fixed-capacity buffer of transitions (a "ring buffer": when full,
+    new transitions overwrite the oldest ones)"""
+
+    def __init__(self, capacity: int):
+        self.capacity = int(capacity)
+        self._storage: Transitions[ActionT] | None = None
+        self._size = 0
+        self._pos = 0
+
+    def __len__(self) -> int:
+        """The number of transitions currently stored"""
+        return self._size
+
+    def add(self, transitions: Transitions[ActionT]) -> None:
+        """Add a batch of transitions to the buffer"""
+        if self._storage is None:
+            # Allocate the storage on first use, using the shapes/dtypes of
+            # the incoming transitions
+            self._storage = transitions._map(
+                lambda tensor: torch.empty(
+                    (self.capacity, *tensor.shape[1:]), dtype=tensor.dtype
+                )
+            )
+
+        indices = (self._pos + torch.arange(len(transitions))) % self.capacity
+        self._storage.set_(indices, transitions)
+        self._pos = int((self._pos + len(transitions)) % self.capacity)
+        self._size = min(self._size + len(transitions), self.capacity)
+
+    def sample(self, batch_size: int) -> Transitions[ActionT]:
+        """Sample a batch of transitions uniformly at random (with replacement)"""
+        assert self._storage is not None, "The replay buffer is empty"
+        indices = torch.randint(0, self._size, (batch_size,))
+        return self._storage[indices]
+
+
+@dataclass(frozen=True)
+class Episode(Generic[ActionT]):
+    """One full episode of length `T`
+
+    :param obs: The observations $s_0, \\ldots, s_{T-1}$, shape `[T, obs_dim]`
+    :param action: The actions $a_0, \\ldots, a_{T-1}$ (tensors `[T, ...]`)
+    :param reward: The rewards $r_1, \\ldots, r_T$ where $r_{t+1}$ is
+        received after performing $a_t$ in $s_t$; shape `[T]`
+    :param final_obs: The last observation $s_T$, shape `[obs_dim]` (see also
+        the `all_obs` property, which appends it to `obs`)
+    :param terminated: True if $s_T$ is a terminal state; False if the
+        episode was truncated (e.g. time limit), in which case the value of
+        $s_T$ must be bootstrapped.
+    """
+
+    obs: Tensor
+    action: ActionT
+    reward: Tensor
+    final_obs: Tensor
+    terminated: bool
+
+    def __len__(self) -> int:
+        """The number of steps in the episode"""
+        return self.reward.shape[0]
+
+    @property
+    def all_obs(self):
+        """The whole observation sequence $s_0, \\ldots, s_T$, shape `[T+1, ...]`
+
+        `obs` only holds the `T` observations on which an action was taken;
+        `all_obs` appends the final one. Hence `all_obs[t+1]` is the successor
+        of `all_obs[t]`: `all_obs[1:]` are the successors of `obs`, and
+        `critic(all_obs)` gives the `T+1` values $V(s_0), \\ldots, V(s_T)$.
+        """
+        if isinstance(self.obs, TensorStruct):
+            # Structured (`Dict`) observations: `final_obs` has no time
+            # dimension, add it before concatenating
+            final_obs = self.final_obs._map(lambda tensor: tensor.unsqueeze(0))
+            return type(self.obs).cat([self.obs, final_obs])
+        return torch.cat([self.obs, self.final_obs.unsqueeze(0)])
+
+    @property
+    def cumulated_reward(self) -> float:
+        """The (undiscounted) sum of rewards of the episode"""
+        return float(self.reward.sum())
+
+
+@dataclass(frozen=True)
+class Rollout(TensorStruct, Generic[ActionT]):
+    """A fixed-length segment of interaction with `B` parallel environments
+
+    All the tensors are time-indexed: shape `[T, B, ...]`. A rollout may
+    span several episodes: episode boundaries are marked by `terminated` /
+    `truncated`, and `next_obs[t]` is always the true successor state of
+    `obs[t]` (the *final* observation when the episode ends at step `t`,
+    while `obs[t+1]` is then the first observation of the next episode).
+
+    :param obs: The observations $s_t$, shape `[T, B, obs_dim]`
+    :param action: The actions $a_t$ (an `Action`, tensors `[T, B, ...]`)
+    :param reward: The rewards $r_{t+1}$, shape `[T, B]`
+    :param next_obs: The successor states $s_{t+1}$, shape `[T, B, obs_dim]`
+    :param terminated: Whether $s_{t+1}$ is a terminal state, shape `[T, B]`
+    :param truncated: Whether the episode was truncated at this step, shape
+        `[T, B]`
+    """
+
+    obs: Tensor
+    action: ActionT
+    reward: Tensor
+    next_obs: Tensor
+    terminated: Tensor
+    truncated: Tensor
+
+    @property
+    def done(self) -> Tensor:
+        """True at episode boundaries (terminated or truncated), shape `[T, B]`"""
+        return self.terminated | self.truncated
+
+    def flatten(self) -> "Rollout[ActionT]":
+        """Merge the time and environment dimensions: `[T, B, ...] -> [T*B, ...]`"""
+        return self._map(lambda tensor: tensor.flatten(0, 1))
+
+
+StructT = TypeVar("StructT", bound=TensorStruct)
+
+
+def minibatches(
+    data: StructT, batch_size: int, shuffle: bool = True
+) -> Iterator[StructT]:
+    """Iterate over minibatches of a batch of data
+
+    The data is partitioned: each sample belongs to exactly one minibatch
+    (the last one may be smaller). With `shuffle`, the partition is random.
+    """
+    indices = torch.randperm(len(data)) if shuffle else torch.arange(len(data))
+    for start in range(0, len(data), batch_size):
+        yield data[indices[start : start + batch_size]]