Replay buffer module

class mbrl.util.replay_buffer.BootstrapIterator(transitions: mbrl.types.TransitionBatch, batch_size: int, ensemble_size: int, shuffle_each_epoch: bool = False, permute_indices: bool = True, rng: Optional[numpy.random._generator.Generator] = None)

Bases: mbrl.util.replay_buffer.TransitionIterator

A transition iterator that can be used to train ensemble of bootstrapped models.

When iterating, this iterator samples from a different set of indices for each model in the ensemble, essentially assigning a different dataset to each model. Each batch is of shape (ensemble_size x batch_size x obs_size).

Parameters

  • batch_size (int) – the batch size to use when iterating over the stored data.

  • ensemble_size (int) – the number of models in the ensemble.

  • shuffle_each_epoch (bool) – if True the iteration order is shuffled everytime a loop over the data is completed. Defaults to False.

  • permute_indices (boot) – if True the bootstrap datasets are just permutations of the original data. If False they are sampled with replacement. Defaults to True.

  • rng (np.random.Generator, optional) – a random number generator when sampling batches. If None (default value), a new default generator will be used.

toggle_bootstrap()

Toggles whether the iterator returns a batch per model or a single batch.

class mbrl.util.replay_buffer.ReplayBuffer(capacity: int, obs_shape: Tuple[int], action_shape: Tuple[int], obs_type=<class 'numpy.float32'>, action_type=<class 'numpy.float32'>, rng: Optional[numpy.random._generator.Generator] = None, max_trajectory_length: Optional[int] = None)

Bases: object

A replay buffer with support for training/validation iterators and ensembles.

This buffer can be pushed to and sampled from as a typical replay buffer. Additionally, it provides a method called get_iterators(), which returns training/validation iterators for the data stored in the buffer.

Parameters

  • capacity (int) – the maximum number of transitions that the buffer can store. When the capacity is reached, the contents are overwritten in FIFO fashion.

  • obs_shape (tuple of ints) – the shape of the observations to store.

  • action_shape (tuple of ints) – the shape of the actions to store.

  • obs_type (type) – the data type of the observations (defaults to np.float32).

  • action_type (type) – the data type of the actions (defaults to np.float32).

  • rng (np.random.Generator, optional) – a random number generator when sampling batches. If None (default value), a new default generator will be used.

  • max_trajectory_length (int, optional) – if given, indicates that trajectory information should be stored and that trajectories will be at most this number of steps. Defaults to None in which case no trajectory information will be kept. The buffer will keep trajectory information automatically using the done value when calling add().

Warning

When using max_trajectory_length it is the user’s responsibility to ensure that trajectories are stored continuously in the replay buffer.

add(obs: numpy.ndarray, action: numpy.ndarray, next_obs: numpy.ndarray, reward: float, done: bool)

Adds a transition (s, a, s’, r, done) to the replay buffer.

Parameters

  • obs (np.ndarray) – the observation at time t.

  • action (np.ndarray) – the action at time t.

  • next_obs (np.ndarray) – the observation at time t + 1.

  • reward (float) – the reward at time t + 1.

  • done (bool) – a boolean indicating whether the episode ended or not.

get_all() → mbrl.types.TransitionBatch

Returns all data stored in the replay buffer.

get_iterators(batch_size: int, val_ratio: float, train_ensemble: bool = False, ensemble_size: Optional[int] = None, shuffle_each_epoch: bool = True, bootstrap_permutes: bool = False) → Tuple[mbrl.util.replay_buffer.TransitionIterator, Optional[mbrl.util.replay_buffer.TransitionIterator]]

Returns training/validation iterators for the data in the replay buffer.

Parameters

  • batch_size (int) – the batch size for the iterators.

  • val_ratio (float) – the proportion of data to use for validation. If 0., the validation buffer will be set to None.

  • train_ensemble (bool) – if True, the training iterator will be and instance of BootstrapIterator. Defaults to False.

  • ensemble_size (int) – the size of the ensemble being trained. Must be provided if train_ensemble == True.

  • shuffle_each_epoch (bool) – if True, the iterator will shuffle the order each time a loop starts. Otherwise the iteration order will be the same. Defaults to True.

  • bootstrap_permutes (bool) – if True, the bootstrap iterator will create the bootstrap data using permutations of the original data. Otherwise it will use sampling with replacement. Defaults to False.

load(load_dir: Union[pathlib.Path, str])

Loads transition data from a given directory.

Parameters

load_dir (str) – the directory where the buffer is stored.

sample(batch_size: int) → Sized

Samples a batch of transitions from the replay buffer.

Parameters

batch_size (int) – the number of samples required.

Returns

the sampled values of observations, actions, next observations, rewards and done indicators, as numpy arrays, respectively. The i-th transition corresponds to (obs[i], act[i], next_obs[i], rewards[i], dones[i]).

Return type

(tuple)

sample_trajectory() → Optional[mbrl.types.TransitionBatch]

Samples a full trajectory and returns it as a batch.

Returns

A tuple with observations, actions, next observations, rewards and done indicators, as numpy arrays, respectively; these will correspond to a full trajectory. The i-th transition corresponds to (obs[i], act[i], next_obs[i], rewards[i], dones[i]).

Return type

(tuple)

save(save_dir: Union[pathlib.Path, str])

Saves the data in the replay buffer to a given directory.

Parameters

save_dir (str) – the directory to save the data to. File name will be replay_buffer.npz.

class mbrl.util.replay_buffer.TransitionIterator(transitions: mbrl.types.TransitionBatch, batch_size: int, shuffle_each_epoch: bool = False, rng: Optional[numpy.random._generator.Generator] = None)

Bases: object

An iterator for batches of transitions.

The iterator can be used doing:

for batch in batch_iterator:
    do_something_with_batch()

Rather than be constructed directly, the preferred way to use objects of this class is for the user to obtain them from ReplayBuffer.

Parameters

  • batch_size (int) – the batch size to use when iterating over the stored data.

  • shuffle_each_epoch (bool) – if True the iteration order is shuffled everytime a loop over the data is completed. Defaults to False.

  • rng (np.random.Generator, optional) – a random number generator when sampling batches. If None (default value), a new default generator will be used.