Actors and actions — rl_mind.core

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).

The base container, with the tensor-like operations every other container inherits:

class rl_mind.core.TensorStruct

Bases: object

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.

__getitem__(index) → Self

Index or slice all the tensor fields along the batch dimension

Parameters:

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)

__len__() → 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

set_(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.

Parameters:

• index – Where to write, as for __getitem__

• value – A structure of the same type, whose batch size matches what index selects

classmethod cat(items: Sequence[Self]) → Self

Concatenate a sequence of structures along the batch dimension

Parameters:

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)

classmethod stack(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.

Parameters:

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)

What an actor returns:

class rl_mind.core.Action(value: Tensor)

Bases: TensorStruct

What an actor returns: a batch of actions

Parameters:

value – The action tensor, shape [B, action_dim] (continuous actions) or [B] (discrete actions)

class rl_mind.core.StochasticAction(value: Tensor, log_prob: Tensor)

Bases: Action

An action sampled from a distribution

Parameters:

log_prob – The log-probability of the sampled action, shape [B]

rl_mind.core.ActionT = ~ActionT

Type of the action returned by an actor

And the actor itself:

class rl_mind.core.Actor(*args, **kwargs)

Bases: 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 forward(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.

Parameters:

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

act(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.

Parameters:

obs – The observations, shape [B, obs_dim]

Returns:

The action tensor, shape [B, action_dim] (continuous actions) or [B] (discrete actions)

__call__(obs: Tensor) → ActionT

Call forward, i.e. compute the training-time action for obs