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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memmap-replay-buffer
-Version: 0.2.0
+Version: 0.2.2
 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
@@ -40,7 +40,12 @@ Requires-Dist: numpy
 Requires-Dist: torch>=2.4
 Provides-Extra: test
 Requires-Dist: h5py; extra == 'test'
+Requires-Dist: imageio; extra == 'test'
+Requires-Dist: imageio-ffmpeg; extra == 'test'
 Requires-Dist: pytest; extra == 'test'
+Provides-Extra: video
+Requires-Dist: imageio; extra == 'video'
+Requires-Dist: imageio-ffmpeg; extra == 'video'
 Description-Content-Type: text/markdown
 
 ## memmap-replay-buffer
@@ -206,6 +211,22 @@ buffer.update(0, returns = returns_for_one)             # scalar index
 all_data = buffer.get_all_data()  # dict of tensors, time-padded to the longest episode
 ```
 
+### Rendering an episode to video
+
+When a field holds one frame per timestep, `save_video` renders an episode to `.mp4` or `.gif` (needs `pip install memmap-replay-buffer[video]`). The image field is auto-detected, and frames may be `uint8` or float in `[0., 1.]`.
+
+```python
+buffer.save_video('./episode-0.mp4', episode_index = 0, fps = 30)
+```
+
+`trajectory_to_video(frames, path, fps)` is exported for plain `(time, height, width, channels)` arrays.
+
+The same is available from the command line, with `--sheet` for a contact sheet and `--serve` for a minimal local frontend (video centered, field stats and a live frame counter to the side).
+
+```bash
+$ inspect-replay-buffer ./image_data --episode 0 --output episode-0.mp4 --serve 8000
+```
+
 ### Reopening, read-only, and clearing
 
 `overwrite = False` (or `from_folder`) rehydrates an existing buffer from disk; the stored config is validated against the requested one, so mismatched `fields`/`max_episodes` raise a clear error. `read_only = True` guarantees no files are created or written. `clear()` wipes all episodes.

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.2}/README.md

@@ -161,6 +161,22 @@ buffer.update(0, returns = returns_for_one)             # scalar index
 all_data = buffer.get_all_data()  # dict of tensors, time-padded to the longest episode
 ```
 
+### Rendering an episode to video
+
+When a field holds one frame per timestep, `save_video` renders an episode to `.mp4` or `.gif` (needs `pip install memmap-replay-buffer[video]`). The image field is auto-detected, and frames may be `uint8` or float in `[0., 1.]`.
+
+```python
+buffer.save_video('./episode-0.mp4', episode_index = 0, fps = 30)
+```
+
+`trajectory_to_video(frames, path, fps)` is exported for plain `(time, height, width, channels)` arrays.
+
+The same is available from the command line, with `--sheet` for a contact sheet and `--serve` for a minimal local frontend (video centered, field stats and a live frame counter to the side).
+
+```bash
+$ inspect-replay-buffer ./image_data --episode 0 --output episode-0.mp4 --serve 8000
+```
+
 ### Reopening, read-only, and clearing
 
 `overwrite = False` (or `from_folder`) rehydrates an existing buffer from disk; the stored config is validated against the requested one, so mismatched `fields`/`max_episodes` raise a clear error. `read_only = True` guarantees no files are created or written. `clear()` wipes all episodes.

{memmap_replay_buffer-0.2.0 → memmap_replay_buffer-0.2.2}/memmap_replay_buffer/__init__.py

@@ -10,3 +10,7 @@ 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
+from memmap_replay_buffer.video import frames_to_sheet as frames_to_sheet
+from memmap_replay_buffer.video import trajectory_to_video as trajectory_to_video

memmap_replay_buffer-0.2.2/memmap_replay_buffer/concat_replay_buffer.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import Tensor, cat
+from torch.utils.data import ConcatDataset, Dataset
+
+from memmap_replay_buffer.replay_buffer import ReplayBuffer, default, exists, pad_right_at_dim
+
+class ConcatReplayBuffer:
+    def __init__(self, folders: list[str | Path]):
+        if len(folders) == 0:
+            raise ValueError('ConcatReplayBuffer requires at least one folder')
+
+        self.buffers = [ReplayBuffer.from_folder(folder, read_only = True) for folder in folders]
+
+        first_buffer = self.buffers[0]
+        self.fieldnames = first_buffer.fieldnames
+        self.meta_fieldnames = first_buffer.meta_fieldnames
+
+        for buffer in self.buffers[1:]:
+            if buffer.fieldnames != self.fieldnames:
+                raise ValueError(f'All buffers must have the same fieldnames - got {buffer.fieldnames} and {self.fieldnames}')
+
+            if buffer.meta_fieldnames != self.meta_fieldnames:
+                raise ValueError(f'All buffers must have the same meta_fieldnames - got {buffer.meta_fieldnames} and {self.meta_fieldnames}')
+
+        self.read_only = True
+
+    @property
+    def num_episodes(self):
+        return sum(buffer.num_episodes for buffer in self.buffers)
+
+    @property
+    def max_episodes(self):
+        return sum(buffer.max_episodes for buffer in self.buffers)
+
+    @property
+    def max_timesteps(self):
+        return max(buffer.max_timesteps for buffer in self.buffers)
+
+    def __len__(self):
+        return sum(len(buffer) for buffer in self.buffers)
+
+    def dataset(self, **kwargs) -> Dataset:
+        if len(self) == 0:
+            raise ValueError('replay buffer is empty')
+
+        datasets = [buffer.dataset(**kwargs) for buffer in self.buffers if len(buffer) > 0]
+        return ConcatDataset(datasets)
+
+    def dataloader(self, *args, **kwargs):
+        return ReplayBuffer.dataloader(self, *args, **kwargs)
+
+    def get_all_data(
+        self,
+        fields: tuple[str, ...] | None = None,
+        meta_fields: tuple[str, ...] | None = None,
+        device: torch.device | str | None = None
+    ) -> dict[str, Tensor]:
+
+        all_data = [buffer.get_all_data(fields = fields, meta_fields = meta_fields, device = device) for buffer in self.buffers]
+        all_data = [data for data in all_data if len(data) > 0]
+
+        if not all_data:
+            return dict()
+
+        data_fields = self.fieldnames if not exists(fields) and not exists(meta_fields) else default(fields, ())
+
+        out = dict()
+
+        for key in all_data[0].keys():
+            tensors = [data[key] for data in all_data if key in data]
+
+            if key in data_fields and tensors[0].ndim > 1:
+                # pad trajectories to the longest across buffers
+                max_time = max(t.shape[1] for t in tensors)
+                tensors = [pad_right_at_dim(t, max_time - t.shape[1], dim = 1) for t in tensors]
+
+            out[key] = cat(tensors, dim = 0)
+
+        return out
+
+    def flush(self):
+        pass
+
+    def _read_only(self, *args, **kwargs):
+        raise NotImplementedError('ConcatReplayBuffer is read-only')
+
+    clear = reset_ = advance_episode = store_batch = store_meta_batch = \
+    one_episode = batched_episode = store_datapoint = store_meta_datapoint = \
+    store_batch_datapoint = store_batch_meta_datapoint = store = store_episode = update = _read_only

memmap_replay_buffer-0.2.2/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.2/memmap_replay_buffer/frontend/index.html

@@ -0,0 +1,87 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8" />
+<meta name="viewport" content="width=device-width, initial-scale=1" />
+<title>inspect-replay-buffer</title>
+<style>
+  :root { color-scheme: dark; }
+  * { box-sizing: border-box; }
+  body {
+    margin: 0;
+    min-height: 100vh;
+    display: grid;
+    place-items: center;
+    background: #0f1012;
+    color: #e8e8e8;
+    font: 13px/1.5 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
+  }
+  video {
+    max-width: min(75vw, 1200px);
+    max-height: 88vh;
+    border-radius: 6px;
+    background: #000;
+  }
+  aside {
+    position: fixed;
+    top: 50%;
+    right: 20px;
+    transform: translateY(-50%);
+    width: 240px;
+    max-height: 88vh;
+    overflow: auto;
+    padding: 16px;
+    background: #17181b;
+    border: 1px solid #26282c;
+    border-radius: 6px;
+  }
+  aside h1 { margin: 0 0 12px; font-size: 13px; font-weight: 600; color: #7fb4e6; }
+  aside h2 { margin: 16px 0 8px; font-size: 11px; font-weight: 600; text-transform: uppercase; letter-spacing: .08em; color: #8a8f98; }
+  dl { display: grid; grid-template-columns: auto minmax(0, 1fr); gap: 4px 12px; margin: 0; }
+  dt { color: #8a8f98; }
+  dd { margin: 0; overflow-wrap: anywhere; }
+  #extras:empty { display: none; }
+  @media (max-width: 1000px) {
+    body { display: flex; flex-direction: column; align-items: center; gap: 16px; padding: 16px; }
+    video { max-width: 100%; }
+    aside { position: static; transform: none; width: 100%; max-width: 480px; }
+  }
+</style>
+</head>
+<body>
+  <video id="replay" src="/video" controls autoplay loop muted></video>
+  <aside>
+    <h1>replay buffer</h1>
+    <dl id="stats"></dl>
+    <h2>time</h2>
+    <dl>
+      <dt>frame</dt><dd id="frame">0</dd>
+    </dl>
+    <div id="extras"></div>
+  </aside>
+<script>
+  const stats = __REPLAY_STATS__;
+  const video = document.getElementById('replay');
+  const extras = document.getElementById('extras');
+
+  const statsList = document.getElementById('stats');
+  for (const [key, value] of Object.entries(stats)) {
+    const dt = document.createElement('dt');
+    dt.textContent = key;
+    const dd = document.createElement('dd');
+    dd.textContent = value;
+    statsList.append(dt, dd);
+  }
+
+  const frameLabel = document.getElementById('frame');
+  const updateTime = () => {
+    frameLabel.textContent = `${Math.round(video.currentTime * stats.fps)} / ${stats.frames}`;
+  };
+  video.addEventListener('timeupdate', updateTime);
+  updateTime();
+
+  // future: sync value predictions / actions / other per-timestep elements to `video.currentTime`
+  window.replay = { stats, video, extras };
+</script>
+</body>
+</html>

memmap_replay_buffer-0.2.2/memmap_replay_buffer/frontend.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import json
+import threading
+import webbrowser
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from urllib.parse import urlparse
+
+# helpers
+
+def render_page(stats: dict):
+    template_path = Path(__file__).parent / 'frontend' / 'index.html'
+    return template_path.read_text().replace('__REPLAY_STATS__', json.dumps(stats))
+
+def parse_range(range_header: str | None, size: int):
+    if not range_header or not range_header.startswith('bytes='):
+        return None
+
+    start_str, _, end_str = range_header[len('bytes='):].split(',')[0].partition('-')
+
+    if not start_str:
+        # suffix range, e.g. bytes=-500
+        return max(size - int(end_str), 0), size - 1
+
+    return int(start_str), min(int(end_str), size - 1) if end_str else size - 1
+
+# server
+
+class ReplayBufferHandler(BaseHTTPRequestHandler):
+    video_path: Path
+    page: str
+
+    protocol_version = 'HTTP/1.1'
+
+    def log_message(self, *args):
+        pass
+
+    def do_GET(self):
+        path = urlparse(self.path).path
+
+        if path == '/':
+            self.send_page()
+        elif path == '/video':
+            self.send_video()
+        else:
+            self.send_error(404)
+
+    def send_page(self):
+        body = self.page.encode()
+
+        self.send_response(200)
+        self.send_header('Content-Type', 'text/html; charset=utf-8')
+        self.send_header('Content-Length', str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def send_video(self):
+        size = self.video_path.stat().st_size
+        byte_range = parse_range(self.headers.get('Range'), size)
+
+        start, end = byte_range if byte_range else (0, size - 1)
+        is_partial = byte_range is not None
+
+        self.send_response(206 if is_partial else 200)
+        self.send_header('Content-Type', 'video/mp4')
+        self.send_header('Content-Length', str(end - start + 1))
+        self.send_header('Accept-Ranges', 'bytes')
+
+        if is_partial:
+            self.send_header('Content-Range', f'bytes {start}-{end}/{size}')
+
+        self.end_headers()
+
+        # stream the selected slice, seeking works via repeated range requests
+
+        with open(self.video_path, 'rb') as f:
+            f.seek(start)
+            remaining = end - start + 1
+
+            while remaining > 0:
+                chunk = f.read(min(remaining, 1 << 20))
+
+                if not chunk:
+                    break
+
+                self.wfile.write(chunk)
+                remaining -= len(chunk)
+
+def make_server(video_path: str | Path, stats: dict, host = '127.0.0.1', port = 8000):
+    handler = type('Handler', (ReplayBufferHandler,), dict(
+        video_path = Path(video_path),
+        page = render_page(stats)
+    ))
+
+    return ThreadingHTTPServer((host, port), handler)
+
+def serve(video_path: str | Path, stats: dict, host = '127.0.0.1', port = 8000, open_browser = True, block = True):
+    httpd = make_server(video_path, stats, host, port)
+    url = f'http://{host}:{httpd.server_address[1]}/'
+
+    print(f'serving replay at {url} (ctrl-c to stop)')
+
+    if open_browser:
+        threading.Thread(target = webbrowser.open, args = (url,), daemon = True).start()
+
+    if not block:
+        threading.Thread(target = httpd.serve_forever, daemon = True).start()
+        return httpd, url
+
+    try:
+        httpd.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        httpd.server_close()
+
+    return httpd, url