memmap-replay-buffer 0.1.2__tar.gz → 0.1.4__tar.gz

{memmap_replay_buffer-0.1.2 → memmap_replay_buffer-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: memmap-replay-buffer
-Version: 0.1.2
+Version: 0.1.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

{memmap_replay_buffer-0.1.2 → memmap_replay_buffer-0.1.4}/memmap_replay_buffer/replay_buffer.py

@@ -988,7 +988,8 @@ class ReplayBuffer:
     def get_all_data(
         self,
         fields: tuple[str, ...] | None = None,
-        meta_fields: tuple[str, ...] | None = None
+        meta_fields: tuple[str, ...] | None = None,
+        device: torch.device | str | None = None
     ):
         self.flush()
 
@@ -1018,7 +1019,7 @@ class ReplayBuffer:
             memmap = self.meta_data[name]
             all_data[name] = from_numpy(memmap[:n].copy())
 
-        return all_data
+        return tree_map_to_device(all_data, device)
 
     @beartype
     def dataset(

{memmap_replay_buffer-0.1.2 → memmap_replay_buffer-0.1.4}/memmap_replay_buffer/replay_buffer_h5py.py

@@ -582,7 +582,12 @@ class ReplayBufferH5PY:
         self.timestep_index = time_dim
         self.advance_episode()
 
-    def get_all_data(self, fields = None, meta_fields = None):
+    def get_all_data(
+        self,
+        fields = None,
+        meta_fields = None,
+        device: torch.device | str | None = None
+    ):
         self.flush()
 
         n = self.num_episodes
@@ -606,7 +611,7 @@ class ReplayBufferH5PY:
         for name in meta_data_fields:
             all_data[name] = from_numpy(self.meta_data[name][:n])
 
-        return all_data
+        return tree_map_to_device(all_data, device)
 
     @beartype
     def dataset(

{memmap_replay_buffer-0.1.2 → memmap_replay_buffer-0.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "memmap-replay-buffer"
-version = "0.1.2"
+version = "0.1.4"
 description = "Simple Replay Buffer for RL"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -47,6 +47,9 @@ test = [
 pythonpath = [
     "."
 ]
+testpaths = [
+    "tests"
+]
 
 [build-system]
 requires = ["hatchling"]

memmap_replay_buffer-0.1.4/tests/test_mixed_storage.py

@@ -0,0 +1,81 @@
+import pytest
+import shutil
+import torch
+from pathlib import Path
+
+from memmap_replay_buffer import ReplayBuffer
+
+h5py = pytest.importorskip("h5py")
+from memmap_replay_buffer import ReplayBufferH5PY
+
+
+@pytest.fixture(params = [ReplayBuffer, ReplayBufferH5PY], ids = ['memmap', 'h5py'])
+def buffer(request, tmp_path):
+    buffer_class = request.param
+    folder = tmp_path / f'mixed_storage_{buffer_class.__name__}'
+
+    buf = buffer_class(
+        folder,
+        max_episodes = 5,
+        max_timesteps = 10,
+        fields = dict(
+            state = 'float',
+            action = 'int'
+        ),
+        meta_fields = dict(
+            text_embedding = ('float', 4),
+            goal_id = 'int'
+        )
+    )
+
+    yield buf
+
+    if folder.exists():
+        shutil.rmtree(folder)
+
+
+def test_store_episode_mixed(buffer):
+    buffer.store_episode(
+        state = torch.randn(5),
+        action = torch.ones(5),
+        text_embedding = torch.randn(4),
+        goal_id = 1
+    )
+
+    buffer.store_episode(
+        state = torch.randn(3),
+        action = torch.zeros(3),
+        text_embedding = torch.randn(4),
+        goal_id = 2
+    )
+
+    assert buffer.num_episodes == 2
+
+
+def test_dataloader_mixed(buffer):
+    buffer.store_episode(
+        state = torch.randn(5),
+        action = torch.ones(5),
+        text_embedding = torch.randn(4),
+        goal_id = 1
+    )
+
+    buffer.store_episode(
+        state = torch.randn(3),
+        action = torch.zeros(3),
+        text_embedding = torch.randn(4),
+        goal_id = 2
+    )
+
+    loader = buffer.dataloader(
+        batch_size = 2,
+        timestep_level = True,
+        shuffle = False
+    )
+
+    batch = next(iter(loader))
+
+    assert 'state' in batch
+    assert 'text_embedding' in batch
+    assert batch['text_embedding'].shape == (2, 4)
+    assert batch['goal_id'].shape == (2,)