"""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"]
[docs]
@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)
[docs]
def __getitem__(self, index) -> Self:
"""Index or slice all the tensor fields along the batch dimension
:param index: Anything a tensor accepts: an integer, a slice, an
index tensor or a boolean mask
:returns: A new structure holding the selected entries (the fields of
a `TensorStruct` are never modified in place)
"""
return self._map(lambda tensor: tensor[index])
[docs]
def __len__(self) -> int:
"""The size of the batch dimension
:returns: The size of the leading dimension, read from the first
tensor (or nested structure) field
:raises ValueError: If the structure holds no tensor field
"""
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")
[docs]
def set_(self, index, value: Self) -> None:
"""Write `value` at position(s) `index` (in-place)
The only mutating operation: instances are frozen, but the *contents*
of their tensor fields can be overwritten. This is how `ReplayBuffer`
fills its preallocated storage.
:param index: Where to write, as for `__getitem__`
:param value: A structure of the same type, whose batch size matches
what `index` selects
"""
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)
[docs]
@classmethod
def cat(cls, items: Sequence[Self]) -> Self:
"""Concatenate a sequence of structures along the batch dimension
:param items: The structures to concatenate (same type, at least
one); non-tensor fields are taken from `items[0]`
:returns: One structure of batch size `sum(len(item) for item in
items)`
"""
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)
[docs]
@classmethod
def stack(cls, items: Sequence[Self]) -> Self:
"""Stack a sequence of structures along a new (first) dimension
This is how the collectors turn a list of per-step actions into a
time-indexed `[T, ...]` structure.
:param items: The structures to stack (same type and batch size, at
least one); non-tensor fields are taken from `items[0]`
:returns: One structure with a new leading dimension of size
`len(items)`
"""
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)
[docs]
@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
[docs]
@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)
[docs]
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).
"""
[docs]
@abstractmethod
def forward(self, obs: Tensor) -> ActionT:
"""Compute the action for a batch of observations (training behavior)
Called by `actor(obs)`: this is the behavior used to *collect* data,
so it is where sampling and exploration noise belong.
:param obs: The observations, shape `[B, obs_dim]` (or a
`TensorStruct` for a structured observation space)
:returns: The actions, an `Action` (or a subclass) whose tensors have
batch size `B`
"""
...
[docs]
def act(self, obs: Tensor) -> Tensor:
"""Deterministic action used for evaluation
Defaults to `forward(obs).value`; override it when the training-time
behavior is stochastic, so that evaluation uses e.g. the mode of the
distribution instead of a sample.
:param obs: The observations, shape `[B, obs_dim]`
:returns: The action tensor, shape `[B, action_dim]` (continuous
actions) or `[B]` (discrete actions)
"""
return self(obs).value
# Overriden so that type checkers know the return type of `actor(obs)`
[docs]
def __call__(self, obs: Tensor) -> ActionT:
"""Call `forward`, i.e. compute the training-time action for `obs`"""
return super().__call__(obs)