"""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"]
[docs]
@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
[docs]
class ReplayBuffer(Generic[ActionT]):
"""A fixed-capacity buffer of transitions (a "ring buffer": when full,
new transitions overwrite the oldest ones)
Used by the off-policy notebooks: `add` the transitions returned by a
`rl_mind.collectors.TransitionCollector`, then `sample` minibatches from
the whole (shuffled) history.
"""
def __init__(self, capacity: int):
"""Create an empty buffer
The storage itself is allocated on the first `add` (it needs the
shapes and dtypes of the transitions).
:param capacity: The maximum number of transitions kept; beyond it,
the oldest ones are overwritten
"""
#: The maximum number of transitions the buffer holds
self.capacity = int(capacity)
self._storage: Transitions[ActionT] | None = None
self._size = 0
self._pos = 0
[docs]
def __len__(self) -> int:
"""The number of transitions currently stored
:returns: A number between 0 and `capacity`
"""
return self._size
[docs]
def add(self, transitions: Transitions[ActionT]) -> None:
"""Add a batch of transitions to the buffer
:param transitions: The transitions to store (a batch of any size;
the storage is allocated from the first batch, so every later
batch must have the same fields, shapes and dtypes)
"""
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)
[docs]
def sample(self, batch_size: int) -> Transitions[ActionT]:
"""Sample a batch of transitions uniformly at random (with replacement)
:param batch_size: The number of transitions to draw
:returns: A `Transitions` of batch size `batch_size`
:raises AssertionError: If the buffer is still empty
"""
assert self._storage is not None, "The replay buffer is empty"
indices = torch.randint(0, self._size, (batch_size,))
return self._storage[indices]
[docs]
@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
[docs]
def __len__(self) -> int:
"""The number of steps `T` of the episode
:returns: The number of actions taken (`obs`, `action` and `reward`
all have `T` entries; `all_obs` has `T+1`)
"""
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())
[docs]
@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
[docs]
def flatten(self) -> "Rollout[ActionT]":
"""Merge the time and environment dimensions: `[T, B, ...] -> [T*B, ...]`
The time structure is lost, so this is done *after* the quantities
that need it (returns, GAE advantages) have been computed; the flat
result is what `minibatches` iterates over.
:returns: The same rollout with a single leading dimension of size
`T * B`
"""
return self._map(lambda tensor: tensor.flatten(0, 1))
StructT = TypeVar("StructT", bound=TensorStruct)
[docs]
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.
:param data: Any `TensorStruct` (e.g. `Transitions`, or a `Rollout` that
was `flatten`ed), of batch size `N`
:param batch_size: The size of the minibatches
:param shuffle: Whether to shuffle the samples before partitioning them
(a different partition at every pass)
:returns: An iterator over `ceil(N / batch_size)` structures of the same
type as `data`
"""
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]]