memmap-replay-buffer 0.2.3__tar.gz → 0.2.4__tar.gz

{memmap_replay_buffer-0.2.3 → memmap_replay_buffer-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memmap-replay-buffer
-Version: 0.2.3
+Version: 0.2.4
 Summary: Simple Replay Buffer for RL
 Project-URL: Homepage, https://pypi.org/project/memmap-replay-buffer/
 Project-URL: Repository, https://codeberg.org/lucidrains/memmap-replay-buffer
@@ -205,6 +205,46 @@ buffer.update(episode_ids, returns = returns_for_ids)   # specific episodes
 buffer.update(0, returns = returns_for_one)             # scalar index
 ```
 
+### Computed fields and selective timestep training
+
+Derived fields must be declared in `fields` when the buffer is created, then written after collection. Use `timestep_start` to write contiguous time chunks without loading a whole episode into memory. A common self-play recipe is to compute per-timestep advantages, threshold them into a boolean field, and then let the timestep-level dataset fetch only the selected timesteps.
+
+```python
+buffer = ReplayBuffer(
+    './replay_data',
+    max_episodes = 1000,
+    max_timesteps = 500,
+    fields = dict(
+        state = ('float', (8,)),
+        reward = 'float',
+        advantage = 'float',
+        high_advantage = ('bool', (), False),   # declares the derived mask
+    )
+)
+
+# compute advantages in episode batches
+buffer.update(episode_ids, advantage = advantages)
+
+# or write contiguous time chunks for all episodes
+for start in range(0, max_len, 128):
+    buffer.update(
+        timestep_start = start,
+        advantage = advantage_chunk
+    )
+
+# set the mask from a comparison of the computed field
+buffer.update(high_advantage = advantage > threshold)
+
+# only selected timesteps are returned
+dataloader = buffer.dataloader(
+    batch_size = 64,
+    timestep_level = True,
+    filter_fields = dict(high_advantage = True)
+)
+```
+
+`demo_advantage_filtering.py` contains a complete self-play loop that computes advantages in batches, marks extreme positive advantages, and trains only on those timesteps.
+
 ### Pulling everything at once
 
 ```python

{memmap_replay_buffer-0.2.3 → memmap_replay_buffer-0.2.4}/README.md

@@ -155,6 +155,46 @@ buffer.update(episode_ids, returns = returns_for_ids)   # specific episodes
 buffer.update(0, returns = returns_for_one)             # scalar index
 ```
 
+### Computed fields and selective timestep training
+
+Derived fields must be declared in `fields` when the buffer is created, then written after collection. Use `timestep_start` to write contiguous time chunks without loading a whole episode into memory. A common self-play recipe is to compute per-timestep advantages, threshold them into a boolean field, and then let the timestep-level dataset fetch only the selected timesteps.
+
+```python
+buffer = ReplayBuffer(
+    './replay_data',
+    max_episodes = 1000,
+    max_timesteps = 500,
+    fields = dict(
+        state = ('float', (8,)),
+        reward = 'float',
+        advantage = 'float',
+        high_advantage = ('bool', (), False),   # declares the derived mask
+    )
+)
+
+# compute advantages in episode batches
+buffer.update(episode_ids, advantage = advantages)
+
+# or write contiguous time chunks for all episodes
+for start in range(0, max_len, 128):
+    buffer.update(
+        timestep_start = start,
+        advantage = advantage_chunk
+    )
+
+# set the mask from a comparison of the computed field
+buffer.update(high_advantage = advantage > threshold)
+
+# only selected timesteps are returned
+dataloader = buffer.dataloader(
+    batch_size = 64,
+    timestep_level = True,
+    filter_fields = dict(high_advantage = True)
+)
+```
+
+`demo_advantage_filtering.py` contains a complete self-play loop that computes advantages in batches, marks extreme positive advantages, and trains only on those timesteps.
+
 ### Pulling everything at once
 
 ```python

{memmap_replay_buffer-0.2.3 → memmap_replay_buffer-0.2.4}/memmap_replay_buffer/replay_buffer.py

@@ -987,10 +987,18 @@ class ReplayBufferBase:
 
     @can_write
     @beartype
-    def update(self, indices: int | list | ndarray | slice | None = None, **data):
+    def update(
+        self,
+        indices: int | list | ndarray | slice | None = None,
+        timestep_start: int = 0,
+        **data
+    ):
         if len(data) == 0:
             raise ValueError('No data provided to `update`')
 
+        if timestep_start < 0:
+            raise ValueError(f'timestep_start must be non-negative, got {timestep_start}')
+
         indices, scalar_index = expand_indices(indices, self.max_episodes)
 
         if not exists(indices):
@@ -1028,11 +1036,11 @@ class ReplayBufferBase:
 
                 time_dim = value.shape[1]
 
-                if time_dim > self.max_timesteps:
+                if timestep_start + time_dim > self.max_timesteps:
                     raise ValueError(f'You exceeded the `max_timesteps` ({self.max_timesteps}) set on the replay buffer. Please increase it on init.')
 
                 value = np.broadcast_to(value, (len(indices), time_dim, *self.shapes[name]))
-                self.data[name][indices, :time_dim] = value[order]
+                self.data[name][indices, timestep_start:timestep_start + time_dim] = value[order]
 
             else:
                 value = cast_to_target_shape(value, self.meta_shapes[name], num_leading_dims = 1)

{memmap_replay_buffer-0.2.3 → memmap_replay_buffer-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "memmap-replay-buffer"
-version = "0.2.3"
+version = "0.2.4"
 description = "Simple Replay Buffer for RL"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

memmap_replay_buffer-0.2.4/tests/test_advantage_filtering.py

@@ -0,0 +1,143 @@
+import numpy as np
+import pytest
+import torch
+
+from memmap_replay_buffer import ReplayBuffer
+from memmap_replay_buffer.replay_buffer_h5py import ReplayBufferH5PY
+
+
+@pytest.mark.parametrize('buffer_class', [ReplayBuffer, ReplayBufferH5PY])
+def test_update_timestep_start(buffer_class, tmp_path):
+    buffer = buffer_class(
+        tmp_path,
+        max_episodes = 2,
+        max_timesteps = 6,
+        fields = dict(advantage = 'float'),
+    )
+
+    buffer.store_episode(advantage = np.zeros(6, dtype = np.float32))
+    buffer.store_episode(advantage = np.ones(6, dtype = np.float32))
+
+    # unsorted indices -> values must follow the passed order
+    buffer.update(
+        indices = [1, 0],
+        timestep_start = 2,
+        advantage = torch.tensor([[20., 21.], [10., 11.]])
+    )
+
+    data = buffer.get_all_data(fields = ('advantage',))['advantage']
+
+    assert torch.equal(data[0], torch.tensor([0., 0., 10., 11., 0., 0.]))
+    assert torch.equal(data[1], torch.tensor([1., 1., 20., 21., 1., 1.]))
+
+    # scalar episode index writes (time, *shape) at the offset
+    buffer.update(0, timestep_start = 4, advantage = torch.tensor([40., 41.]))
+    data = buffer.get_all_data(fields = ('advantage',))['advantage']
+    assert torch.equal(data[0], torch.tensor([0., 0., 10., 11., 40., 41.]))
+
+    with pytest.raises(ValueError, match = 'non-negative'):
+        buffer.update(timestep_start = -1, advantage = torch.zeros(2, 1))
+
+    with pytest.raises(ValueError, match = 'max_timesteps'):
+        buffer.update(timestep_start = 5, advantage = torch.zeros(2, 2))
+
+
+@pytest.mark.parametrize('buffer_class', [ReplayBuffer, ReplayBufferH5PY])
+def test_advantage_threshold_filtering_end_to_end(buffer_class, tmp_path):
+    buffer = buffer_class(
+        tmp_path / 'advantage',
+        max_episodes = 6,
+        max_timesteps = 8,
+        fields = dict(
+            reward = 'float',
+            advantage = 'float',
+            high_advantage = ('bool', (), False),
+        ),
+        circular = False,
+        overwrite = True,
+    )
+
+    rewards = [
+        [1., 1., 1., -1., 1., 1.],
+        [-1., -1., -1., -1.],
+        [1., -1., 1., -1., 1., -1., 1.],
+        [-1., 1., -1.],
+        [1., 1., 1., 1., 1.],
+        [-1., -1.],
+    ]
+
+    for episode_rewards in rewards:
+        buffer.store_episode(reward = np.asarray(episode_rewards, dtype = np.float32))
+
+    num_episodes = buffer.num_episodes
+    episode_lens = torch.from_numpy(buffer.episode_lens[:num_episodes].copy())
+    max_len = int(episode_lens.max())
+
+    # 1. compute advantage in episode batches, writing each batch back
+
+    gamma = 0.9
+
+    for start in range(0, num_episodes, 2):
+        episode_indices = np.arange(start, min(start + 2, num_episodes))
+        batch_len = int(episode_lens[episode_indices].max())
+
+        rewards_batch = torch.from_numpy(buffer.data['reward'][episode_indices, :batch_len].copy())
+        valid = torch.arange(batch_len)[None] < episode_lens[episode_indices][:, None]
+
+        returns = torch.zeros_like(rewards_batch)
+        running = torch.zeros(len(episode_indices))
+
+        for t in reversed(range(batch_len)):
+            running = rewards_batch[:, t] + gamma * running
+            returns[:, t] = running
+
+        buffer.update(episode_indices, advantage = returns * valid)
+
+    # 2. set a mask from a threshold comparison, in time chunks
+
+    threshold = 1.
+
+    for start in range(0, max_len, 4):
+        end = min(start + 4, max_len)
+
+        advantages = torch.from_numpy(buffer.data['advantage'][:num_episodes, start:end].copy())
+        valid = torch.arange(start, end)[None] < episode_lens[:, None]
+        high_advantage = (advantages > threshold) & valid
+
+        buffer.update(timestep_start = start, high_advantage = high_advantage)
+
+    # 3. verify the mask matches the comparison exactly
+
+    advantages = torch.from_numpy(buffer.data['advantage'][:num_episodes].copy())
+    high_advantage = torch.from_numpy(buffer.data['high_advantage'][:num_episodes].copy()).bool()
+
+    expected = {
+        (episode, timestep)
+        for episode in range(num_episodes)
+        for timestep in range(int(episode_lens[episode]))
+        if advantages[episode, timestep].item() > threshold
+    }
+
+    assert len(expected) > 0
+    assert len(expected) < int(episode_lens.sum())
+
+    for episode in range(num_episodes):
+        for timestep in range(int(episode_lens[episode])):
+            assert high_advantage[episode, timestep].item() == ((episode, timestep) in expected)
+
+    # 4. timestep-level learning only receives the selected timesteps
+
+    dataset = buffer.dataset(
+        timestep_level = True,
+        fields = ('advantage',),
+        filter_fields = dict(high_advantage = True),
+        return_indices = True,
+    )
+
+    fetched = {
+        tuple(dataset[index]['_indices'].tolist())
+        for index in range(len(dataset))
+    }
+
+    assert fetched == expected
+    assert all(dataset[index]['advantage'].item() > threshold for index in range(len(dataset)))