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

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

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

@@ -12,3 +12,5 @@ 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/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