memmap-replay-buffer 0.0.6

Simple Replay Buffer for RL

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

import torch
from memmap_replay_buffer import ReplayBuffer

# initialize buffer

buffer = ReplayBuffer(
    './replay_data',
    max_episodes = 1000,
    max_timesteps = 500,
    fields = dict(
        state = ('float', (8,), 0.5),   # 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(8).numpy(),
                action = torch.randint(0, 4, (2,)),
                reward = 1.0
            )

# rehydrate from disk

buffer_rehydrated = ReplayBuffer.from_config('./replay_data')
assert buffer_rehydrated.num_episodes == 4

# learn 2 episodes at a time

dataloader = buffer.dataloader(batch_size = 2)

for batch in dataloader:
    state = batch['state']    # (2, 100, 8)
    action = batch['action']  # (2, 100)
    reward = batch['reward']  # (2, 100)
    lens = batch['_lens']     # (2,)

    assert state.shape  == (2, 100, 8)
    assert action.shape == (2, 100, 2)
    assert reward.shape == (2, 100)
    assert lens.shape   == (2,)