"""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 episodes (`Episode`), 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)
[docs]
class TransitionCollector(Generic[ActionT]):
"""Collect batches of transitions by running an actor in a `VecEnv`
The collector for off-policy algorithms (DQN, DDPG, TD3, SAC, TQC): the
transitions it returns are meant to be pushed into a
`rl_mind.data.ReplayBuffer`, so their order does not matter and the
(invalid) auto-reset steps are simply dropped.
"""
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
"""Create the collector
:param env: The training environments; the collector keeps their
state across `collect` calls
:param actor: The actor to run — its *training* behavior
(`actor(obs)`, i.e. with exploration) is what gets collected
"""
#: The environments being stepped
self.env = env
#: The actor being run
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)
[docs]
def collect(self, n_steps: int) -> Transitions[ActionT]:
"""Perform `n_steps` steps in every environment
The environments are reset on the first call only: successive calls
continue the same episodes, so a training loop can interleave
`collect(1)` with gradient steps.
:param n_steps: The number of steps performed in *each* of the
`env.num_envs` environments
:returns: The recorded transitions, a flat batch of at most
`n_steps * num_envs` of them (the steps where an environment was
auto-resetting are skipped, and `steps` is increased by the
number actually recorded)
"""
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
[docs]
class EpisodeCollector(Generic[ActionT]):
"""Collect full episodes by running an actor in a `VecEnv`
The collector for episodic algorithms (REINFORCE), which need complete
*on-policy* episodes: the environments are therefore reset at the
beginning of every `collect`.
"""
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
"""Create the collector
:param env: The training environments (reset at each `collect`)
:param actor: The actor to run — its *training* behavior
(`actor(obs)`) is what gets collected
"""
#: The environments being stepped
self.env = env
#: The actor being run
self.actor = actor
#: Total number of recorded steps
self.steps = 0
[docs]
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.
:param n_episodes: The minimum number of episodes to return
:returns: The completed episodes, of varying lengths — more than
`n_episodes` when several environments finish during the same
step (the episodes still running are discarded)
"""
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
[docs]
class RolloutCollector(Generic[ActionT]):
"""Collect fixed-length `[T, B]` rollouts by running an actor in a `VecEnv`
The collector for on-policy algorithms (A2C, PPO, GAIL): the time axis is
preserved, since the advantages are computed by a backward recursion over
it. This is why the environment must use `same_step_reset=True` — every
row of the `[T, B]` block is then a valid transition.
"""
def __init__(self, env: VecEnv, actor: Actor[ActionT]):
"""Create the collector
:param env: The training environments, created with
`same_step_reset=True`; the collector keeps their state across
`collect` calls
:param actor: The actor to run — its *training* behavior
(`actor(obs)`) is what gets collected
:raises AssertionError: If `env` was not created with
`same_step_reset=True`
"""
assert env.same_step_reset, (
"RolloutCollector requires a VecEnv created with "
"same_step_reset=True (so that every rollout row is a valid "
"transition)"
)
#: The environments being stepped
self.env = env
#: The actor being run
self.actor = actor
#: Total number of recorded steps
self.steps = 0
self._obs: torch.Tensor | None = None
[docs]
def collect(self, n_steps: int) -> Rollout[ActionT]:
"""Perform `n_steps` steps in every environment
The environments are reset on the first call only: successive calls
continue where the previous one stopped.
:param n_steps: The number of steps performed in *each* environment,
i.e. the length `T` of the rollout
:returns: A `[n_steps, num_envs]` rollout (every row is a valid
transition, episode boundaries being marked by `terminated` /
`truncated`)
"""
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]),
)