imitation.data.buffer#

Buffers to store NumPy arrays and transitions in.

Functions

num_samples(data)

Computes the number of samples contained in data.

Classes

Buffer(capacity, sample_shapes, dtypes)

A FIFO ring buffer for NumPy arrays of a fixed shape and dtype.

ReplayBuffer(capacity[, venv, obs_shape, ...])

Buffer for Transitions.
class imitation.data.buffer.Buffer(capacity, sample_shapes, dtypes)[source]#

Bases: object

A FIFO ring buffer for NumPy arrays of a fixed shape and dtype.

Supports random sampling with replacement.

__init__(capacity, sample_shapes, dtypes)[source]#

Constructs a Buffer.

Parameters

  • capacity (int) – The number of samples that can be stored.

  • sample_shapes (Mapping[str, Tuple[int, ...]]) – A dictionary mapping string keys to the shape of samples associated with that key.

  • dtypes (np.dtype-like) – A dictionary mapping string keys to the dtype of samples associated with that key.

Raises

KeyError – sample_shapes and dtypes have different keys.

capacity: int#

The number of data samples that can be stored in this buffer.

classmethod from_data(data, capacity=None, truncate_ok=False)[source]#

Constructs and return a Buffer containing the provided data.

Shapes and dtypes are automatically inferred.

Parameters

  • data (Mapping[str, ndarray]) – A dictionary mapping keys to data arrays. The arrays may differ in their shape, but should agree in the first axis.

  • capacity (Optional[int]) – The Buffer capacity. If not provided, then this is automatically set to the size of the data, so that the returned Buffer is at full capacity.

  • truncate_ok (bool) – Whether to error if capacity < the number of samples in data. If False, then only store the last capacity samples from data when overcapacity.

Examples

In the follow examples, suppose the arrays in data are length-1000.

Buffer with same capacity as arrays in data:

Buffer.from_data(data)

Buffer with larger capacity than arrays in data:

Buffer.from_data(data, 10000)

Buffer with smaller capacity than arrays in `data. Without truncate_ok=True, from_data will error:

Buffer.from_data(data, 5, truncate_ok=True)
Return type

Buffer

Returns

Buffer of specified capacity containing provided data.

Raises

  • ValueError – data is empty.

  • ValueError – data has items mapping to arrays differing in the length of their first axis.

sample(n_samples)[source]#

Uniformly sample n_samples samples from the buffer with replacement.

Parameters

n_samples (int) – The number of samples to randomly sample.

Returns

An array with shape

(n_samples) + self.sample_shape.

Return type

samples (np.ndarray)

Raises

ValueError – The buffer is empty.

sample_shapes: Mapping[str, Tuple[int, ...]]#

The shapes of each data sample stored in this buffer.

size()[source]#

Returns the number of samples stored in the buffer.

Return type

int

store(data, truncate_ok=False)[source]#

Stores new data samples, replacing old samples with FIFO priority.

Parameters

  • data (Mapping[str, ndarray]) – A dictionary mapping keys k to arrays with shape (n_samples,) + self.sample_shapes[k], where n_samples is less than or equal to self.capacity.

  • truncate_ok (bool) – If False, then error if the length of transitions is greater than self.capacity. Otherwise, store only the final self.capacity transitions.

Raises

  • ValueError – data is empty.

  • ValueError – If n_samples is greater than self.capacity.

  • ValueError – data is the wrong shape.

Return type

None

class imitation.data.buffer.ReplayBuffer(capacity, venv=None, *, obs_shape=None, act_shape=None, obs_dtype=None, act_dtype=None)[source]#

Bases: object

Buffer for Transitions.

__init__(capacity, venv=None, *, obs_shape=None, act_shape=None, obs_dtype=None, act_dtype=None)[source]#

Constructs a ReplayBuffer.

Parameters

  • capacity (int) – The number of samples that can be stored.

  • venv (Optional[VecEnv]) – The environment whose action and observation spaces can be used to determine the data shapes of the underlying buffers. Mutually exclusive with shape and dtype arguments.

  • obs_shape (Optional[Tuple[int, ...]]) – The shape of the observation space.

  • act_shape (Optional[Tuple[int, ...]]) – The shape of the action space.

  • obs_dtype (Optional[dtype]) – The dtype of the observation space.

  • act_dtype (Optional[dtype]) – The dtype of the action space.

Raises

  • ValueError – Couldn’t infer the observation and action shapes and dtypes from the arguments.

  • ValueError – Specified both venv and shapes/dtypes.

capacity: int#

The number of data samples that can be stored in this buffer.

classmethod from_data(transitions, capacity=None, truncate_ok=False)[source]#

Construct and return a ReplayBuffer containing the provided data.

Shapes and dtypes are automatically inferred, and the returned ReplayBuffer is ready for sampling.

Parameters

  • transitions (Transitions) – Transitions to store.

  • capacity (Optional[int]) – The ReplayBuffer capacity. If not provided, then this is automatically set to the size of the data, so that the returned Buffer is at full capacity.

  • truncate_ok (bool) – Whether to error if capacity < the number of samples in data. If False, then only store the last capacity samples from data when overcapacity.

Examples

ReplayBuffer with same capacity as arrays in data:

ReplayBuffer.from_data(data)

ReplayBuffer with larger capacity than arrays in data:

ReplayBuffer.from_data(data, 10000)

ReplayBuffer with smaller capacity than arrays in `data. Without truncate_ok=True, from_data will error:

ReplayBuffer.from_data(data, 5, truncate_ok=True)
Return type

ReplayBuffer

Returns

A new ReplayBuffer.

sample(n_samples)[source]#

Sample obs-act-obs triples.

Parameters

n_samples (int) – The number of samples.

Return type

Transitions

Returns

A Transitions named tuple containing n_samples transitions.

size()[source]#

Returns the number of samples stored in the buffer.

Return type

Optional[int]

store(transitions, truncate_ok=True)[source]#

Store obs-act-obs triples.

Parameters

  • transitions (Transitions) – Transitions to store.

  • truncate_ok (bool) – If False, then error if the length of transitions is greater than self.capacity. Otherwise, store only the final self.capacity transitions.

Raises

ValueError – The arguments didn’t have the same length.

Return type

None

imitation.data.buffer.num_samples(data)[source]#

Computes the number of samples contained in data.

Parameters

data (Mapping[Any, ndarray]) – A Mapping from keys to NumPy arrays.

Return type

int

Returns

The unique length of the first dimension of arrays contained in data.

Raises

ValueError – The length is not unique.