memmap-replay-buffer
An easy-to-use numpy memmap replay buffer for RL and other sequence-based learning tasks.
Install
$ pip install memmap-replay-buffer
Usage
Supports trajectory-level, timestep-level, and n-step transition dataloading from a single stored buffer.
import torch
from memmap_replay_buffer import ReplayBuffer
# initialize buffer
buffer = ReplayBuffer(
'./replay_data',
max_episodes = 1000,
max_timesteps = 500,
fields = dict(
state = ('float', (3, 16, 16), 0.), # type, shape, and optional default value
action = ('int', 2),
reward = 'float' # default shape is ()
),
meta_fields = dict(
task_id = 'int'
),
circular = True,
overwrite = True
)
# store 4 episodes
for _ in range(4):
with buffer.one_episode(task_id = 1):
for _ in range(100):
buffer.store(
state = torch.randn(3, 16, 16),
action = torch.randint(0, 4, (2,)).numpy(),
reward = 1.0
)
# rehydrate from disk
buffer_rehydrated = ReplayBuffer.from_folder('./replay_data')
assert buffer_rehydrated.num_episodes == 4
Trajectory-level
Variable-length trajectories, automatically padded with mask and lengths.
dataloader = buffer.dataloader(
batch_size = 2,
return_mask = True,
to_named_tuple = ('state', 'action', 'reward', 'task_id', '_mask', '_lens')
)
for state, action, reward, task_id, mask, lens in dataloader:
assert state.shape == (2, 100, 3, 16, 16)
assert action.shape == (2, 100, 2)
assert reward.shape == (2, 100)
assert task_id.shape == (2,)
assert lens.shape == (2,)
assert mask.shape == (2, 100)
Timestep-level
Individual timesteps across episodes, with optional filter_meta for conditioning.
dataloader = buffer.dataloader(
batch_size = 8,
filter_meta = dict(
task_id = 1
),
to_named_tuple = ('state', 'action', 'task_id'),
timestep_level = True,
drop_last = True
)
for state, action, task_id in dataloader:
assert state.shape == (8, 3, 16, 16)
assert action.shape == (8, 2)
assert task_id.shape == (8,)
N-step transitions
Fetches current_fields at $t$, next_fields at $t + n$ (prefixed next_), and sequence_fields from $t$ to $t + n$ (prefixed seq_, zero-padded at episode boundaries). Use fieldname_map to remap to your model's kwargs.
dataloader = buffer.dataloader(
batch_size = 4,
n_steps = 5,
current_fields = ('state',),
next_fields = ('state',),
sequence_fields = ('action', 'reward'),
to_named_tuple = ('state', 'next_state', 'action_chunk', 'rewards', 'n_step_lens'),
fieldname_map = {
'seq_action': 'action_chunk',
'seq_reward': 'rewards'
}
)
for state, next_state, action_chunk, rewards, n_step_lens in dataloader:
assert state.shape == (4, 3, 16, 16)
assert next_state.shape == (4, 3, 16, 16)
assert action_chunk.shape == (4, 5, 2)
assert rewards.shape == (4, 5)
assert n_step_lens.shape == (4,)
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file memmap_replay_buffer-0.1.11.tar.gz.
File metadata
- Download URL: memmap_replay_buffer-0.1.11.tar.gz
- Upload date:
- Size: 25.6 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.8.17
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
09bfd82e0f3ca9c1627976bb38008e35d07f8c70b863bed68d1a1e204883abb6
|
|
| MD5 |
fcf2e74996d6c004b1c6c13212f4104c
|
|
| BLAKE2b-256 |
6568e4ecc4a46cb7047d92d9c0c8f123e72f13ee7fbd47d565f2bef033632e29
|
File details
Details for the file memmap_replay_buffer-0.1.11-py3-none-any.whl.
File metadata
- Download URL: memmap_replay_buffer-0.1.11-py3-none-any.whl
- Upload date:
- Size: 35.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.8.17
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
32268d89df549342e163ce6274eba72debd53e2bd79fe7ac4367018bc5d31c10
|
|
| MD5 |
15ab8fbed0113d883f701140bd0afa4a
|
|
| BLAKE2b-256 |
6c320e4b50c0744943860c11b3639d6956a0aeb61d91b0987bb49d2698005e01
|
