memmap-replay-buffer 0.2.0__tar.gz → 0.2.1__tar.gz

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memmap-replay-buffer
-Version: 0.2.0
+Version: 0.2.1
 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.2.0 → memmap_replay_buffer-0.2.1}/memmap_replay_buffer/__init__.py

@@ -10,3 +10,5 @@ try:
     from memmap_replay_buffer.replay_buffer_h5py import ReplayBufferH5PY as ReplayBufferH5PY
 except ImportError:
     pass
+
+from memmap_replay_buffer.episode_collector import EpisodeCollector as EpisodeCollector

memmap_replay_buffer-0.2.1/memmap_replay_buffer/episode_collector.py

@@ -0,0 +1,284 @@
+from __future__ import annotations
+
+from collections import defaultdict, namedtuple
+from typing import Any
+
+import numpy as np
+from beartype import beartype
+from numpy import ndarray
+import torch
+from torch import is_tensor
+from torch.utils.data import DataLoader, Dataset, default_collate
+
+from memmap_replay_buffer.replay_buffer import tree_map_to_device
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def is_group_batched(value, num_groups):
+    if is_tensor(value) or isinstance(value, ndarray):
+        return value.ndim > 0 and value.shape[0] == num_groups
+
+    if isinstance(value, (list, tuple)):
+        return len(value) == num_groups
+
+    return False
+
+def stack(values, dim = 0):
+    if is_tensor(values[0]):
+        return torch.stack(values, dim = dim)
+
+    return np.stack(values, axis = dim)
+
+# dataset
+
+class TensorDictDataset(Dataset):
+    def __init__(self, data: dict[str, torch.Tensor]):
+        self.data = data
+        self.num_samples = len(next(iter(data.values())))
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        return {name: value[index] for name, value in self.data.items()}
+
+# class
+
+class EpisodeCollector:
+    """ accumulates a group-batched rollout step by step, then either stores it or trains on it in memory
+
+    layouts are explicit, so the time and group axes can never be silently swapped
+
+        `append` - one step, group-batched (num_groups, *shape) or (num_groups, 1, *shape)
+        `extend` - whole trajectories, time-major (time, num_groups, *shape), e.g. returns / advantages
+        `env_major` / `time_major` - read back (num_groups, time, ...) / (time, num_groups, ...)
+        `all_data` / `dataloader` - env-major tensors / flattened minibatches, entirely in memory
+        `store` - write one episode per group, then `clear`
+        `clear` - discard the current rollout
+    """
+
+    @beartype
+    def __init__(
+        self,
+        replay_buffer: Any,        # ReplayBuffer | ReplayBufferH5PY, without a hard dependency on h5py
+        num_groups: int,
+        fieldnames: tuple[str, ...] | None = None,
+        meta_fieldnames: tuple[str, ...] | None = None,
+        strict: bool = True
+    ):
+        self.replay_buffer = replay_buffer
+        self.num_groups = num_groups
+
+        self.fieldnames = set(default(fieldnames, tuple(replay_buffer.fieldnames)))
+        self.meta_fieldnames = set(default(meta_fieldnames, tuple(replay_buffer.meta_fieldnames)))
+
+        self.shapes = replay_buffer.shapes
+        self.meta_shapes = replay_buffer.meta_shapes
+        self.strict = strict
+
+        self.clear()
+
+    def __len__(self):
+        return self.step_count
+
+    def clear(self):
+        self.episode_data = [defaultdict(list) for _ in range(self.num_groups)]
+        self.step_count = 0
+
+    # fields
+
+    def field_kind(self, name):
+        if name in self.fieldnames:
+            return 'data'
+
+        if name in self.meta_fieldnames:
+            return 'meta'
+
+        if self.strict:
+            raise KeyError(f'unknown field {name} - valid fields are {sorted(self.fieldnames | self.meta_fieldnames)}')
+
+        return None
+
+    def field_shape(self, name):
+        return self.meta_shapes[name] if name in self.meta_fieldnames else self.shapes[name]
+
+    def per_group(self, name, value):
+        # (num_groups, *shape) / (num_groups, 1, *shape) -> num_groups values of shape (*shape)
+
+        if isinstance(value, (list, tuple)):
+            return list(value)
+
+        shape = tuple(self.field_shape(name))
+        rest = tuple(value.shape[1:])
+
+        assert rest in (shape, (1, *shape)), f'field {name} - per-group shape {rest} should be {shape} or (1, *shape)'
+
+        if rest != shape:
+            value = value[:, 0]
+
+        return list(value)
+
+    # adding data
+
+    @beartype
+    def append(self, **data):
+        # one step, group-batched values
+
+        for name, value in data.items():
+            kind = self.field_kind(name)
+
+            if not exists(kind):
+                continue
+
+            if kind == 'meta':
+                values = [value] * self.num_groups
+            else:
+                shape = tuple(self.field_shape(name))
+
+                assert is_group_batched(value, self.num_groups), (
+                    f'field {name} - `append` expects per-step group-batched values of shape ({self.num_groups}, *{shape}) '
+                    f'or ({self.num_groups}, 1, *{shape}) - for whole trajectories in time-major (time, {self.num_groups}, *{shape}), use `extend`'
+                )
+
+                values = self.per_group(name, value)
+
+            for group_index, group_value in enumerate(values):
+                self.episode_data[group_index][name].append(group_value)
+
+        self.step_count += 1
+
+    @beartype
+    def extend(self, **trajectories):
+        # whole trajectories, time-major (time, num_groups, *shape)
+
+        for name, value in trajectories.items():
+            kind = self.field_kind(name)
+
+            if not exists(kind):
+                continue
+
+            assert is_tensor(value) or isinstance(value, ndarray), f'field {name} - `extend` expects a tensor / ndarray of shape (time, {self.num_groups}, *shape)'
+
+            time = value.shape[0]
+
+            assert value.ndim > 1 and value.shape[1] == self.num_groups, (
+                f'field {name} - `extend` expects time-major (time, {self.num_groups}, *shape), got {tuple(value.shape)}'
+            )
+
+            if len(self) == 0:
+                self.step_count = time
+            else:
+                assert time == len(self), f'field {name} - time dim {time} does not match the {len(self)} collected steps'
+
+            shape = tuple(self.field_shape(name))
+            assert tuple(value.shape[2:]) == shape, f'field {name} - feature shape {tuple(value.shape[2:])} does not match {shape}'
+
+            for group_index in range(self.num_groups):
+                self.episode_data[group_index][name].extend(list(value[:, group_index]))
+
+    # reading back
+
+    def groups(self, name):
+        # one value per group, i.e. a list of num_groups tensors / arrays
+
+        if self.field_kind(name) == 'meta':
+            return [data[name][-1] for data in self.episode_data]
+
+        return [stack(data[name], dim = 0) for data in self.episode_data]
+
+    def time_major(self, name):
+        # (time, num_groups, *shape) for data fields, (num_groups,) for meta fields
+
+        return stack(self.groups(name), dim = 1 if self.field_kind(name) == 'data' else 0)
+
+    def env_major(self, name):
+        # (num_groups, time, *shape), the layout generalized advantage estimation expects
+
+        values = self.time_major(name)
+
+        if self.field_kind(name) == 'meta':
+            return values
+
+        return values.swapaxes(0, 1)
+
+    # in-memory training - no storing required
+
+    def all_data(self, fields = None, device: torch.device | str | None = None):
+        # dict of env-major tensors, (num_groups, time, *shape), for the collected data fields
+
+        if len(self) == 0:
+            raise ValueError('collector is empty')
+
+        collected = tuple(name for name in self.episode_data[0] if name in self.fieldnames)
+
+        data = {name: self.env_major(name) for name in default(fields, collected)}
+        return tree_map_to_device(data, device)
+
+    def dataloader(
+        self,
+        batch_size,
+        fields = None,
+        filter_fields = None,
+        to_named_tuple = None,
+        shuffle = False,
+        device: torch.device | str | None = None,
+        **kwargs
+    ):
+        # minibatches flattened over (num_groups, time) for the current rollout, entirely in memory
+
+        data = {name: torch.as_tensor(value).reshape(-1, *value.shape[2:]) for name, value in self.all_data().items()}
+
+        if exists(filter_fields) and len(filter_fields) > 0:
+            mask = None
+            for name, filter_value in filter_fields.items():
+                field_mask = data[name] == filter_value
+                mask = field_mask if not exists(mask) else (mask & field_mask)
+
+            data = {name: value[mask] for name, value in data.items()}
+
+        fields = default(fields, tuple(data.keys()))
+        data = {name: data[name] for name in fields}
+
+        dataset = TensorDictDataset(data)
+
+        NamedTupleCls = None
+        if exists(to_named_tuple):
+            sanitized_fields = tuple(f.lstrip('_') if f.startswith('_') else f for f in to_named_tuple)
+            NamedTupleCls = namedtuple('Batch', sanitized_fields)
+
+        def collate_fn(data):
+            batch = default_collate(data)
+
+            if exists(NamedTupleCls):
+                for field in to_named_tuple:
+                    if field not in batch:
+                        raise ValueError(f'field `{field}` not found in batch. available fields: {list(batch.keys())}')
+
+                batch = NamedTupleCls(**{san: batch[orig] for orig, san in zip(to_named_tuple, sanitized_fields)})
+
+            return tree_map_to_device(batch, device)
+
+        return DataLoader(dataset, batch_size = batch_size, shuffle = shuffle, collate_fn = collate_fn, **kwargs)
+
+    # storing
+
+    @beartype
+    def store(self, **meta_data):
+        for group_index, group_data in enumerate(self.episode_data):
+            store_kwargs = {name: stack(values, dim = 0) if name in self.fieldnames else values[-1] for name, values in group_data.items()}
+
+            for name, value in meta_data.items():
+                if name not in self.meta_fieldnames:
+                    continue
+
+                store_kwargs[name] = value[group_index] if is_group_batched(value, self.num_groups) else value
+
+            self.replay_buffer.store_episode(**store_kwargs)
+
+        self.clear()

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.1}/memmap_replay_buffer/replay_buffer.py

@@ -1394,12 +1394,16 @@ class ReplayBuffer:
         self,
         num_groups: int,
         fieldnames: tuple[str, ...] | None = None,
-        meta_fieldnames: tuple[str, ...] | None = None
+        meta_fieldnames: tuple[str, ...] | None = None,
+        **collector_kwargs
     ):
         from memmap_replay_buffer.episode_collector import EpisodeCollector
         return EpisodeCollector(
             self,
             num_groups,
             fieldnames = fieldnames,
-            meta_fieldnames = meta_fieldnames
+            meta_fieldnames = meta_fieldnames,
+            **collector_kwargs
         )
+
+    create_rollout_collector = create_collector

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.1}/memmap_replay_buffer/replay_buffer_h5py.py

@@ -917,16 +917,20 @@ class ReplayBufferH5PY:
         self,
         num_groups: int,
         fieldnames: tuple[str, ...] | None = None,
-        meta_fieldnames: tuple[str, ...] | None = None
+        meta_fieldnames: tuple[str, ...] | None = None,
+        **collector_kwargs
     ):
         from memmap_replay_buffer.episode_collector import EpisodeCollector
         return EpisodeCollector(
             self,
             num_groups,
             fieldnames = fieldnames,
-            meta_fieldnames = meta_fieldnames
+            meta_fieldnames = meta_fieldnames,
+            **collector_kwargs
         )
 
+    create_rollout_collector = create_collector
+
     @can_write
     def _store_episodes_batch(self, data: dict[str, np.ndarray]):
         batch_size = next(iter(data.values())).shape[0]

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.1}/pyproject.toml

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

memmap_replay_buffer-0.2.1/tests/test_rollout_collector.py

@@ -0,0 +1,187 @@
+import pytest
+import torch
+
+from memmap_replay_buffer import ReplayBuffer
+
+FIELDS = dict(
+    obs = ('float', (3,)),
+    action = 'int',
+    reward = 'float',
+    next_obs = ('float', (3,)),
+    value = 'float',
+    returns = 'float',
+    advantage = 'float',
+)
+
+def make_buffer(tmp_path):
+    return ReplayBuffer(
+        tmp_path / 'buffer',
+        max_episodes = 4,
+        max_timesteps = 16,
+        fields = FIELDS,
+        circular = True,
+        overwrite = True
+    )
+
+def test_rollout_collector_alignment_roundtrip(tmp_path):
+    T, N = 5, 2
+
+    buffer = make_buffer(tmp_path)
+    collector = buffer.create_collector(num_groups = N)
+
+    obs_all = torch.randn(T + 1, N, 3)
+    obs, next_obs = obs_all[:-1], obs_all[1:]
+    action = torch.randint(0, 4, (T, N))
+    reward = torch.randn(T, N)
+    value = torch.randn(T, N)
+
+    for t in range(T):
+        collector.append(
+            obs = obs[t],
+            action = action[t],
+            reward = reward[t],
+            next_obs = next_obs[t],
+            value = value[t]
+        )
+
+    # whole-trajectory fields, explicitly time-major
+
+    returns = torch.randn(T, N)
+    advantage = torch.randn(T, N)
+    collector.extend(returns = returns, advantage = advantage)
+
+    assert len(collector) == T
+
+    # read-back views must agree with the inputs
+
+    assert torch.equal(collector.time_major('obs'), obs)
+    assert torch.equal(collector.env_major('obs'), obs.transpose(0, 1))
+    assert torch.equal(collector.time_major('returns'), returns)
+    assert torch.equal(collector.env_major('advantage'), advantage.transpose(0, 1))
+
+    collector.store()
+
+    # stored rows are per group and time-ordered, never scrambled
+
+    for i in range(N):
+        assert torch.allclose(torch.tensor(buffer.data['obs'][i, :T]), obs[:, i])
+        assert torch.allclose(torch.tensor(buffer.data['next_obs'][i, :T]), next_obs[:, i])
+        assert torch.allclose(torch.tensor(buffer.data['returns'][i, :T]), returns[:, i])
+        assert torch.allclose(torch.tensor(buffer.data['advantage'][i, :T]), advantage[:, i])
+        assert torch.allclose(torch.tensor(buffer.data['value'][i, :T]), value[:, i])
+
+        # the transition continuity invariant that catches any permutation
+
+        assert torch.allclose(torch.tensor(buffer.data['next_obs'][i, :T - 1]), torch.tensor(buffer.data['obs'][i, 1:T]))
+
+def test_collector_in_memory_training(tmp_path):
+    T, N = 5, 3
+
+    buffer = make_buffer(tmp_path)
+    collector = buffer.create_collector(num_groups = N)
+
+    obs = torch.randn(T, N, 3)
+    action = torch.randint(0, 4, (T, N))
+    value = torch.randn(T, N)
+    returns = torch.randn(T, N)
+    advantage = torch.randn(T, N)
+
+    for t in range(T):
+        collector.append(obs = obs[t], action = action[t], value = value[t])
+
+    collector.extend(returns = returns, advantage = advantage)
+
+    # env-major snapshot, no disk access
+
+    all_data = collector.all_data()
+    assert torch.equal(all_data['obs'], obs.transpose(0, 1))
+
+    # flattened minibatches, optionally filtered
+
+    dataloader = collector.dataloader(
+        batch_size = 4,
+        filter_fields = dict(action = 1),
+        to_named_tuple = ('obs', 'action', 'returns', 'advantage')
+    )
+
+    batches = list(dataloader)
+    assert sum(len(batch.obs) for batch in batches) == (action == 1).sum().item()
+
+    expected = obs.transpose(0, 1).reshape(-1, 3)[(action.transpose(0, 1) == 1).reshape(-1)]
+    assert torch.allclose(torch.cat([batch.obs for batch in batches]), expected)
+
+    # nothing was stored
+
+    assert buffer.num_episodes == 0
+
+    # clearing resets the rollout
+
+    collector.clear()
+    assert len(collector) == 0
+
+    with pytest.raises(ValueError, match = 'empty'):
+        collector.all_data()
+
+def test_rollout_collector_rejects_time_major_append(tmp_path):
+    T, N = 5, 2
+
+    buffer = make_buffer(tmp_path)
+    collector = buffer.create_collector(num_groups = N)
+
+    with pytest.raises(AssertionError, match = 'extend'):
+        collector.append(obs = torch.randn(T, N, 3))
+
+def test_rollout_collector_rejects_wrong_time_dim_in_extend(tmp_path):
+    T, N = 4, 3
+
+    buffer = make_buffer(tmp_path)
+    collector = buffer.create_collector(num_groups = N)
+
+    for t in range(T):
+        collector.append(
+            obs = torch.randn(N, 3),
+            action = torch.zeros(N, dtype = torch.long),
+            reward = torch.zeros(N),
+            next_obs = torch.randn(N, 3),
+            value = torch.zeros(N)
+        )
+
+    with pytest.raises(AssertionError, match = 'time dim'):
+        collector.extend(returns = torch.randn(T + 1, N))
+
+def test_rollout_collector_unit_time_dim_and_strict(tmp_path):
+    N = 2
+
+    buffer = make_buffer(tmp_path)
+    collector = buffer.create_collector(num_groups = N)
+
+    # the (N, 1, *shape) form is accepted and stored identically to (N, *shape)
+
+    for t in range(3):
+        collector.append(
+            obs = torch.full((N, 1, 3), float(t)),
+            action = torch.full((N, 1), t),
+            reward = torch.full((N, 1), float(t)),
+            next_obs = torch.zeros(N, 1, 3),
+            value = torch.zeros(N, 1)
+        )
+
+    collector.store()
+
+    data = buffer.get_all_data()
+    assert data['obs'].shape[:2] == (N, 3)
+    assert data['action'][0].tolist() == [0, 1, 2]
+    assert data['action'][1].tolist() == [0, 1, 2]
+
+    with pytest.raises(KeyError, match = 'unknown field'):
+        collector.append(typo_field = torch.zeros(N))
+
+    collector_lax = buffer.create_collector(num_groups = N, strict = False)
+    collector_lax.append(
+        obs = torch.zeros(N, 3),
+        action = torch.zeros(N, dtype = torch.long),
+        reward = torch.zeros(N),
+        next_obs = torch.zeros(N, 3),
+        value = torch.zeros(N),
+        typo_field = torch.zeros(N)
+    )

memmap_replay_buffer-0.2.0/memmap_replay_buffer/episode_collector.py

@@ -1,105 +0,0 @@
-from __future__ import annotations
-
-from collections import defaultdict
-
-import numpy as np
-from beartype import beartype
-from numpy import ndarray
-from torch import cat, is_tensor, stack
-
-from memmap_replay_buffer.replay_buffer import ReplayBuffer
-from memmap_replay_buffer.replay_buffer_h5py import ReplayBufferH5PY
-
-# helpers
-
-def exists(v):
-    return v is not None
-
-def default(v, d):
-    return v if exists(v) else d
-
-def is_group_batched(value, num_groups):
-    if is_tensor(value) or isinstance(value, ndarray):
-        return value.shape[0] == num_groups
-
-    if isinstance(value, (list, tuple)):
-        return len(value) == num_groups
-
-    return False
-
-def combine_trajectory(values):
-    first = values[0]
-    if is_tensor(first):
-        return stack(values) if first.ndim == 0 else cat(values, dim = 0)
-
-    if isinstance(first, ndarray):
-        return np.stack(values) if first.ndim == 0 else np.concatenate(values, axis = 0)
-
-    return np.array(values)
-
-class EpisodeCollector:
-    @beartype
-    def __init__(
-        self,
-        replay_buffer: ReplayBuffer | ReplayBufferH5PY,
-        num_groups: int,
-        fieldnames: tuple[str, ...] | None = None,
-        meta_fieldnames: tuple[str, ...] | None = None
-    ):
-        self.replay_buffer = replay_buffer
-        self.num_groups = num_groups
-
-        self.fieldnames = default(fieldnames, tuple(replay_buffer.fieldnames))
-        self.meta_fieldnames = default(meta_fieldnames, tuple(replay_buffer.meta_fieldnames))
-
-        self.reset()
-
-    def reset(self):
-        self.episode_data = [defaultdict(list) for _ in range(self.num_groups)]
-
-    @beartype
-    def append(
-        self,
-        **data
-    ):
-        for name, value in data.items():
-            if name not in self.fieldnames and name not in self.meta_fieldnames:
-                continue
-
-            if is_group_batched(value, self.num_groups):
-                for i in range(self.num_groups):
-                    self.episode_data[i][name].append(value[i])
-            else:
-                for i in range(self.num_groups):
-                    self.episode_data[i][name].append(value)
-
-    @beartype
-    def store(self, **meta_data):
-        for i in range(self.num_groups):
-            group_data = self.episode_data[i]
-
-            store_kwargs = dict()
-
-            # Handle collected data
-            for name, values in group_data.items():
-                if name in self.fieldnames:
-                    combined = combine_trajectory(values)
-                else:
-                    # Meta fields: take the last appended value
-                    combined = values[-1]
-
-                store_kwargs[name] = combined
-
-            # Merge with extra meta data provided to store()
-            for name, value in meta_data.items():
-                if name not in self.meta_fieldnames:
-                    continue
-
-                if is_group_batched(value, self.num_groups):
-                    store_kwargs[name] = value[i]
-                else:
-                    store_kwargs[name] = value
-
-            self.replay_buffer.store_episode(**store_kwargs)
-
-        self.reset()