`evorl.replay_buffers.replay_buffer`¶

Module Contents¶

Classes¶

`AbstractReplayBuffer`	A ReplyBuffer Interface.
`ReplayBuffer`	ReplayBuffer with uniform sampling.
`ReplayBufferState`	Contains data related to a replay buffer.

API¶

class evorl.replay_buffers.replay_buffer.AbstractReplayBuffer[source]¶

Bases: evorl.types.PyTreeNode

A ReplyBuffer Interface.

abstract add(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState, xs: chex.ArrayTree) → evorl.replay_buffers.replay_buffer.ReplayBufferState[source]¶

Add data to the replay buffer.

Parameters:

buffer_state – The current state of the replay buffer.
xs – The data to add to the replay buffer.

Returns:

Updated state of the replay buffer.

abstract can_sample(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState) → bool[source]¶

Check if the current replay buffer state can be used to sample.

Parameters:: buffer_state – The current state of the replay buffer.
Returns:: Whether the replay buffer is ready to call sample().

abstract init(sample_spec: chex.ArrayTree) → evorl.replay_buffers.replay_buffer.ReplayBufferState[source]¶

Initialize the state of the replay buffer.

Parameters:: sample_spec – A single sample or sample spec that contains the pytree structure and their dtype and shape.
Returns:: The initial state of the replay buffer.

abstract is_full(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState) → bool[source]¶

abstract sample(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState, key: chex.PRNGKey) → chex.ArrayTree[source]¶

Sample a batch of data from the replay buffer.

Parameters:

buffer_state – The current state of the replay buffer.
key – JAX PRNGKey.

Returns:

A batch of data sampled from the replay buffer.

class evorl.replay_buffers.replay_buffer.ReplayBuffer[source]¶

Bases: evorl.replay_buffers.replay_buffer.AbstractReplayBuffer

ReplayBuffer with uniform sampling.

Data are added and sampled in 1d-like structure.

Variables:

capacity – the maximum capacity of the replay buffer.
sample_batch_size – the batch size for sample().
min_sample_timesteps – the minimum number of timesteps before the replay buffer can sample.

add(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState, xs: chex.ArrayTree, mask: chex.Array | None = None) → evorl.replay_buffers.replay_buffer.ReplayBufferState[source]¶

can_sample(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState) → bool[source]¶

capacity: int¶: None

init(spec: chex.ArrayTree) → evorl.replay_buffers.replay_buffer.ReplayBufferState[source]¶

is_full(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState) → bool[source]¶

min_sample_timesteps: int¶: 0

sample(buffer_state: evorl.replay_buffers.replay_buffer.ReplayBufferState, key: chex.ArrayTree) → chex.ArrayTree[source]¶

sample_batch_size: int¶: None

class evorl.replay_buffers.replay_buffer.ReplayBufferState[source]¶

Bases: evorl.types.PyTreeData

Contains data related to a replay buffer.

Variables:

data – the stored replay buffer data.
current_index – the pointer used for adding data.
buffer_size – the current size of the replay buffer.

buffer_size: chex.Array¶: ‘zeros(…)’

current_index: chex.Array¶: ‘zeros(…)’

data: chex.ArrayTree¶: None