Data containers — rl_mind.data

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.

Off-policy data — a flat batch of transitions, and the buffer they are kept in:

class rl_mind.data.Transitions(obs: Tensor, action: ActionT, reward: Tensor, next_obs: Tensor, terminated: Tensor)

Bases: TensorStruct, Generic[ActionT]

A batch of N transitions \((s, a, r, s', \textrm{terminated})\)

Parameters:

• obs – The observations \(s\), shape [N, obs_dim]

• action – The actions \(a\) (an Action, tensors of shape [N, ...])

• reward – The rewards \(r\), shape [N]

• next_obs – The next observations \(s'\), shape [N, obs_dim]

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

class rl_mind.data.ReplayBuffer(capacity: int)

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

capacity

The maximum number of transitions the buffer holds

__len__() → int

The number of transitions currently stored

Returns:

A number between 0 and capacity

add(transitions: Transitions[ActionT]) → None

Add a batch of transitions to the buffer

Parameters:

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)

sample(batch_size: int) → Transitions[ActionT]

Sample a batch of transitions uniformly at random (with replacement)

Parameters:

batch_size – The number of transitions to draw

Returns:

A Transitions of batch size batch_size

Raises:

AssertionError – If the buffer is still empty

Episodic and on-policy data:

class rl_mind.data.Episode(obs: Tensor, action: ActionT, reward: Tensor, final_obs: Tensor, terminated: bool)

Bases: Generic[ActionT]

One full episode of length T

Parameters:

• obs – The observations \(s_0, \ldots, s_{T-1}\), shape [T, obs_dim]

• action – The actions \(a_0, \ldots, a_{T-1}\) (tensors [T, ...])

• reward – The rewards \(r_1, \ldots, r_T\) where \(r_{t+1}\) is received after performing \(a_t\) in \(s_t\); shape [T]

• final_obs – The last observation \(s_T\), shape [obs_dim] (see also the all_obs property, which appends it to obs)

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

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

property all_obs

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

property cumulated_reward: float

The (undiscounted) sum of rewards of the episode

class rl_mind.data.Rollout(obs: Tensor, action: ActionT, reward: Tensor, next_obs: Tensor, terminated: Tensor, truncated: Tensor)

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

Parameters:

• obs – The observations \(s_t\), shape [T, B, obs_dim]

• action – The actions \(a_t\) (an Action, tensors [T, B, ...])

• reward – The rewards \(r_{t+1}\), shape [T, B]

• next_obs – The successor states \(s_{t+1}\), shape [T, B, obs_dim]

• terminated – Whether \(s_{t+1}\) is a terminal state, shape [T, B]

• truncated – Whether the episode was truncated at this step, shape [T, B]

property done: Tensor

True at episode boundaries (terminated or truncated), shape [T, B]

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

rl_mind.data.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.

Parameters:

• data – Any TensorStruct (e.g. Transitions, or a Rollout that was flatten`ed), of batch size `N

• batch_size – The size of the minibatches

• 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