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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memmap-replay-buffer
-Version: 0.2.1
+Version: 0.2.3
 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,23 @@ 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 episodes 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.]`. Pass a list of episodes to concatenate them in time.
+
+```python
+buffer.save_video('./episode-0.mp4', episode_index = 0, fps = 30)
+buffer.save_video('./episodes-0-3-5.mp4', episode_index = [0, 3, 5], 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. `--episode` and `--field` accept comma-delimited lists - with multiple cameras (image-like fields), one video is written per field (`<name>.<field>.<ext>`) and the frontend plays them side by side, synchronized, with a live frame counter to the side. The frontend has frame-level controls: a timeline scrubber, single-frame stepping (buttons, left/right arrows, shift for 10, or mouse wheel over the videos), play/pause, and slowdown playback (down to 0.1x).
+
+```bash
+$ inspect-replay-buffer ./image_data --episode 0,3,5 --field front,wrist --output replay.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.3}/README.md

@@ -161,6 +161,23 @@ 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 episodes 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.]`. Pass a list of episodes to concatenate them in time.
+
+```python
+buffer.save_video('./episode-0.mp4', episode_index = 0, fps = 30)
+buffer.save_video('./episodes-0-3-5.mp4', episode_index = [0, 3, 5], 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. `--episode` and `--field` accept comma-delimited lists - with multiple cameras (image-like fields), one video is written per field (`<name>.<field>.<ext>`) and the frontend plays them side by side, synchronized, with a live frame counter to the side. The frontend has frame-level controls: a timeline scrubber, single-frame stepping (buttons, left/right arrows, shift for 10, or mouse wheel over the videos), play/pause, and slowdown playback (down to 0.1x).
+
+```bash
+$ inspect-replay-buffer ./image_data --episode 0,3,5 --field front,wrist --output replay.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.3}/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.3/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.3/memmap_replay_buffer/frontend/index.html

@@ -0,0 +1,285 @@
+<!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;
+  }
+  #stage {
+    display: flex;
+    flex-direction: column;
+    gap: 12px;
+    width: min(calc(100vw - 340px), 1200px);
+  }
+  main {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(320px, 1fr));
+    gap: 12px;
+  }
+  figure { margin: 0; min-width: 0; }
+  figure video { width: 100%; max-height: 74vh; display: block; border-radius: 6px; background: #000; }
+  figcaption { margin-top: 6px; color: #8a8f98; font-size: 11px; text-transform: uppercase; letter-spacing: .08em; }
+  #controls {
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    padding: 8px 12px;
+    background: #17181b;
+    border: 1px solid #26282c;
+    border-radius: 6px;
+  }
+  #controls button {
+    width: 40px;
+    padding: 4px 0;
+    background: #1f2126;
+    color: #e8e8e8;
+    border: 1px solid #2e3136;
+    border-radius: 4px;
+    font: inherit;
+    cursor: pointer;
+  }
+  #controls button:hover { background: #26282c; }
+  #timeline { flex: 1; accent-color: #7fb4e6; min-width: 0; }
+  #controls .time { color: #8a8f98; white-space: nowrap; }
+  #speed {
+    background: #1f2126;
+    color: #e8e8e8;
+    border: 1px solid #2e3136;
+    border-radius: 4px;
+    padding: 4px;
+    font: inherit;
+  }
+  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; }
+    #stage { width: 100%; max-width: 640px; }
+    aside { position: static; transform: none; width: 100%; max-width: 480px; }
+  }
+</style>
+</head>
+<body>
+  <div id="stage">
+    <main id="replays"></main>
+    <footer id="controls">
+      <button id="step-back" title="previous frame (left arrow)">&#10216;</button>
+      <button id="play-pause" title="play / pause (space)">&#9654;</button>
+      <button id="step-forward" title="next frame (right arrow)">&#10217;</button>
+      <input id="timeline" type="range" min="0" value="0" step="1" />
+      <span id="time" class="time">0 / 0</span>
+      <select id="speed">
+        <option value="0.1">0.1x</option>
+        <option value="0.25">0.25x</option>
+        <option value="0.5">0.5x</option>
+        <option value="1" selected>1x</option>
+        <option value="2">2x</option>
+        <option value="4">4x</option>
+      </select>
+    </footer>
+  </div>
+  <aside>
+    <h1>replay buffer</h1>
+    <dl id="stats"></dl>
+    <h2>fields</h2>
+    <dl id="field-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 fields = stats.fields ?? [{ name: 'video' }];
+  const extras = document.getElementById('extras');
+
+  const statsList = document.getElementById('stats');
+  for (const [key, value] of Object.entries(stats)) {
+    if (key === 'fields') continue;
+
+    const dt = document.createElement('dt');
+    dt.textContent = key;
+    const dd = document.createElement('dd');
+    dd.textContent = Array.isArray(value) ? value.join(', ') : value;
+    statsList.append(dt, dd);
+  }
+
+  const fieldList = document.getElementById('field-stats');
+  for (const field of fields) {
+    const dt = document.createElement('dt');
+    dt.textContent = field.name;
+    const dd = document.createElement('dd');
+    dd.textContent = [field.shape, field.dtype].filter(Boolean).join(' ') || 'video';
+    fieldList.append(dt, dd);
+  }
+
+  // one video per field, every video follows the first
+
+  const replays = document.getElementById('replays');
+  const single = fields.length === 1;
+
+  const videos = fields.map((field, index) => {
+    const figure = document.createElement('figure');
+    const video = document.createElement('video');
+
+    video.src = single ? '/video' : `/video/${encodeURIComponent(field.name)}`;
+    video.controls = false;
+    video.autoplay = true;
+    video.loop = true;
+    video.muted = true;
+    video.playsInline = true;
+
+    const caption = document.createElement('figcaption');
+    caption.textContent = field.name;
+
+    figure.append(video, caption);
+    replays.append(figure);
+    return video;
+  });
+
+  const [master, ...slaves] = videos;
+
+  const sync = () => {
+    for (const slave of slaves) {
+      if (Math.abs(slave.currentTime - master.currentTime) > 0.08) {
+        slave.currentTime = master.currentTime;
+      }
+    }
+  };
+
+  slaves.forEach(slave => slave.addEventListener('loadedmetadata', sync));
+
+  master.addEventListener('play', () => {
+    playPause.textContent = '⏸';
+    slaves.forEach(video => video.play());
+  });
+
+  master.addEventListener('pause', () => {
+    playPause.textContent = '▶';
+    slaves.forEach(video => video.pause());
+  });
+
+  master.addEventListener('ratechange', () => {
+    slaves.forEach(video => video.playbackRate = master.playbackRate);
+    speed.value = String(master.playbackRate);
+  });
+
+  master.addEventListener('seeking', sync);
+  master.addEventListener('seeked', sync);
+  master.addEventListener('timeupdate', sync);
+
+  const tick = () => {
+    if (!master.paused) sync();
+    requestAnimationFrame(tick);
+  };
+  requestAnimationFrame(tick);
+
+  // frame-level controls
+
+  const playPause = document.getElementById('play-pause');
+  const timeline = document.getElementById('timeline');
+  const timeLabel = document.getElementById('time');
+  const speed = document.getElementById('speed');
+  const frameLabel = document.getElementById('frame');
+
+  const totalFrames = () => stats.frames ?? (Number.isFinite(master.duration) ? Math.round(master.duration * stats.fps) : 0);
+  const frameFromTime = (time) => Math.max(0, Math.min(Math.round(time * stats.fps), totalFrames() - 1));
+
+  const updateTime = () => {
+    const frame = frameFromTime(master.currentTime);
+    const total = totalFrames();
+
+    frameLabel.textContent = `${frame} / ${total}`;
+    timeLabel.textContent = `${frame} / ${total}`;
+    timeline.max = String(Math.max(total - 1, 0));
+
+    if (!scrubbing) timeline.value = String(frame);
+  };
+
+  let scrubbing = false;
+  timeline.addEventListener('pointerdown', () => scrubbing = true);
+  window.addEventListener('pointerup', () => { if (scrubbing) { scrubbing = false; updateTime(); } });
+
+  const seekToFrame = (frame) => {
+    master.currentTime = Math.max(0, Math.min(frame, totalFrames() - 1)) / stats.fps;
+    sync();
+    updateTime();
+  };
+
+  timeline.addEventListener('input', () => {
+    master.pause();
+    seekToFrame(Number(timeline.value));
+  });
+
+  playPause.addEventListener('click', () => master.paused ? master.play() : master.pause());
+  document.getElementById('step-back').addEventListener('click', () => { master.pause(); seekToFrame(frameFromTime(master.currentTime) - 1); });
+  document.getElementById('step-forward').addEventListener('click', () => { master.pause(); seekToFrame(frameFromTime(master.currentTime) + 1); });
+
+  speed.addEventListener('change', () => master.playbackRate = Number(speed.value));
+
+  // mouse wheel over the videos steps through frames
+
+  replays.addEventListener('wheel', (event) => {
+    event.preventDefault();
+    master.pause();
+    seekToFrame(frameFromTime(master.currentTime) + Math.sign(event.deltaY));
+  }, { passive: false });
+
+  window.addEventListener('keydown', (event) => {
+    if (event.target.matches('input, select, button, textarea')) return;
+
+    const step = event.shiftKey ? 10 : 1;
+
+    if (event.key === ' ') {
+      event.preventDefault();
+      master.paused ? master.play() : master.pause();
+    } else if (event.key === 'ArrowLeft') {
+      event.preventDefault();
+      master.pause();
+      seekToFrame(frameFromTime(master.currentTime) - step);
+    } else if (event.key === 'ArrowRight') {
+      event.preventDefault();
+      master.pause();
+      seekToFrame(frameFromTime(master.currentTime) + step);
+    }
+  });
+
+  master.addEventListener('timeupdate', updateTime);
+  master.addEventListener('loadedmetadata', updateTime);
+  master.addEventListener('play', updateTime);
+  master.addEventListener('pause', updateTime);
+  updateTime();
+
+  // future: sync value predictions / actions / other per-timestep elements to `master.currentTime`
+  window.replay = { stats, fields, videos, master, timeline, seekToFrame, extras };
+</script>
+</body>
+</html>

memmap_replay_buffer-0.2.3/memmap_replay_buffer/frontend.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import json
+import mimetypes
+import threading
+import webbrowser
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from urllib.parse import unquote, urlparse
+
+# helpers
+
+def normalize_stats(stats: dict):
+    """ fold the legacy single-field stats keys into a `fields` list """
+    stats = dict(stats)
+
+    if 'fields' not in stats and 'field' in stats:
+        stats['fields'] = [dict(
+            name = stats.pop('field'),
+            shape = stats.pop('shape', ''),
+            dtype = stats.pop('dtype', '')
+        )]
+
+    return stats
+
+def normalize_videos(videos):
+    if isinstance(videos, dict):
+        return {str(name): Path(path) for name, path in videos.items()}
+
+    return dict(video = Path(videos))
+
+def content_type(path: Path):
+    return mimetypes.guess_type(path.name)[0] or 'video/mp4'
+
+def render_page(stats: dict):
+    template_path = Path(__file__).parent / 'frontend' / 'index.html'
+    return template_path.read_text().replace('__REPLAY_STATS__', json.dumps(normalize_stats(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_paths: dict[str, 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 == '/':
+            return self.send_page()
+
+        video_path = self.video_for_path(path)
+
+        if video_path is not None:
+            return self.send_video(video_path)
+
+        self.send_error(404)
+
+    def video_for_path(self, path):
+        # `/video` serves the first video for backward compatibility
+
+        if path == '/video':
+            return next(iter(self.video_paths.values()), None)
+
+        if path.startswith('/video/'):
+            return self.video_paths.get(unquote(path[len('/video/'):]))
+
+        return None
+
+    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, video_path):
+        size = 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', content_type(video_path))
+        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(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(videos: str | Path | dict[str, Path], stats: dict, host = '127.0.0.1', port = 8000):
+    handler = type('Handler', (ReplayBufferHandler,), dict(
+        video_paths = normalize_videos(videos),
+        page = render_page(stats)
+    ))
+
+    return ThreadingHTTPServer((host, port), handler)
+
+def serve(videos: str | Path | dict[str, Path], stats: dict, host = '127.0.0.1', port = 8000, open_browser = True, block = True):
+    httpd = make_server(videos, 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