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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memmap-replay-buffer
-Version: 0.2.2
+Version: 0.2.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
@@ -205,26 +205,67 @@ buffer.update(episode_ids, returns = returns_for_ids)   # specific episodes
 buffer.update(0, returns = returns_for_one)             # scalar index
 ```
 
+### Computed fields and selective timestep training
+
+Derived fields must be declared in `fields` when the buffer is created, then written after collection. Use `timestep_start` to write contiguous time chunks without loading a whole episode into memory. A common self-play recipe is to compute per-timestep advantages, threshold them into a boolean field, and then let the timestep-level dataset fetch only the selected timesteps.
+
+```python
+buffer = ReplayBuffer(
+    './replay_data',
+    max_episodes = 1000,
+    max_timesteps = 500,
+    fields = dict(
+        state = ('float', (8,)),
+        reward = 'float',
+        advantage = 'float',
+        high_advantage = ('bool', (), False),   # declares the derived mask
+    )
+)
+
+# compute advantages in episode batches
+buffer.update(episode_ids, advantage = advantages)
+
+# or write contiguous time chunks for all episodes
+for start in range(0, max_len, 128):
+    buffer.update(
+        timestep_start = start,
+        advantage = advantage_chunk
+    )
+
+# set the mask from a comparison of the computed field
+buffer.update(high_advantage = advantage > threshold)
+
+# only selected timesteps are returned
+dataloader = buffer.dataloader(
+    batch_size = 64,
+    timestep_level = True,
+    filter_fields = dict(high_advantage = True)
+)
+```
+
+`demo_advantage_filtering.py` contains a complete self-play loop that computes advantages in batches, marks extreme positive advantages, and trains only on those timesteps.
+
 ### Pulling everything at once
 
 ```python
 all_data = buffer.get_all_data()  # dict of tensors, time-padded to the longest episode
 ```
 
-### Rendering an episode to video
+### 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.]`.
+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 (video centered, field stats and a live frame counter to the side).
+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 --output episode-0.mp4 --serve 8000
+$ inspect-replay-buffer ./image_data --episode 0,3,5 --field front,wrist --output replay.mp4 --serve 8000
 ```
 
 ### Reopening, read-only, and clearing

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

@@ -155,26 +155,67 @@ buffer.update(episode_ids, returns = returns_for_ids)   # specific episodes
 buffer.update(0, returns = returns_for_one)             # scalar index
 ```
 
+### Computed fields and selective timestep training
+
+Derived fields must be declared in `fields` when the buffer is created, then written after collection. Use `timestep_start` to write contiguous time chunks without loading a whole episode into memory. A common self-play recipe is to compute per-timestep advantages, threshold them into a boolean field, and then let the timestep-level dataset fetch only the selected timesteps.
+
+```python
+buffer = ReplayBuffer(
+    './replay_data',
+    max_episodes = 1000,
+    max_timesteps = 500,
+    fields = dict(
+        state = ('float', (8,)),
+        reward = 'float',
+        advantage = 'float',
+        high_advantage = ('bool', (), False),   # declares the derived mask
+    )
+)
+
+# compute advantages in episode batches
+buffer.update(episode_ids, advantage = advantages)
+
+# or write contiguous time chunks for all episodes
+for start in range(0, max_len, 128):
+    buffer.update(
+        timestep_start = start,
+        advantage = advantage_chunk
+    )
+
+# set the mask from a comparison of the computed field
+buffer.update(high_advantage = advantage > threshold)
+
+# only selected timesteps are returned
+dataloader = buffer.dataloader(
+    batch_size = 64,
+    timestep_level = True,
+    filter_fields = dict(high_advantage = True)
+)
+```
+
+`demo_advantage_filtering.py` contains a complete self-play loop that computes advantages in batches, marks extreme positive advantages, and trains only on those timesteps.
+
 ### Pulling everything at once
 
 ```python
 all_data = buffer.get_all_data()  # dict of tensors, time-padded to the longest episode
 ```
 
-### Rendering an episode to video
+### 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.]`.
+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 (video centered, field stats and a live frame counter to the side).
+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 --output episode-0.mp4 --serve 8000
+$ inspect-replay-buffer ./image_data --episode 0,3,5 --field front,wrist --output replay.mp4 --serve 8000
 ```
 
 ### Reopening, read-only, and clearing

memmap_replay_buffer-0.2.4/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.2 → memmap_replay_buffer-0.2.4}/memmap_replay_buffer/frontend.py

@@ -1,17 +1,40 @@
 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 urlparse
+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(stats))
+    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='):
@@ -28,7 +51,7 @@ def parse_range(range_header: str | None, size: int):
 # server
 
 class ReplayBufferHandler(BaseHTTPRequestHandler):
-    video_path: Path
+    video_paths: dict[str, Path]
     page: str
 
     protocol_version = 'HTTP/1.1'
@@ -40,11 +63,25 @@ class ReplayBufferHandler(BaseHTTPRequestHandler):
         path = urlparse(self.path).path
 
         if path == '/':
-            self.send_page()
-        elif path == '/video':
-            self.send_video()
-        else:
-            self.send_error(404)
+            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()
@@ -55,15 +92,15 @@ class ReplayBufferHandler(BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(body)
 
-    def send_video(self):
-        size = self.video_path.stat().st_size
+    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', 'video/mp4')
+        self.send_header('Content-Type', content_type(video_path))
         self.send_header('Content-Length', str(end - start + 1))
         self.send_header('Accept-Ranges', 'bytes')
 
@@ -74,7 +111,7 @@ class ReplayBufferHandler(BaseHTTPRequestHandler):
 
         # stream the selected slice, seeking works via repeated range requests
 
-        with open(self.video_path, 'rb') as f:
+        with open(video_path, 'rb') as f:
             f.seek(start)
             remaining = end - start + 1
 
@@ -87,16 +124,16 @@ class ReplayBufferHandler(BaseHTTPRequestHandler):
                 self.wfile.write(chunk)
                 remaining -= len(chunk)
 
-def make_server(video_path: str | Path, stats: dict, host = '127.0.0.1', port = 8000):
+def make_server(videos: str | Path | dict[str, Path], stats: dict, host = '127.0.0.1', port = 8000):
     handler = type('Handler', (ReplayBufferHandler,), dict(
-        video_path = Path(video_path),
+        video_paths = normalize_videos(videos),
         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)
+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)')

{memmap_replay_buffer-0.2.2 → memmap_replay_buffer-0.2.4}/memmap_replay_buffer/replay_buffer.py

@@ -11,7 +11,7 @@ import einx
 import numpy as np
 import torch
 from beartype import beartype
-from beartype.typing import Any, Literal
+from beartype.typing import Any, Literal, Sequence
 from einops import rearrange
 from numpy import ndarray
 from numpy.lib.format import open_memmap
@@ -987,10 +987,18 @@ class ReplayBufferBase:
 
     @can_write
     @beartype
-    def update(self, indices: int | list | ndarray | slice | None = None, **data):
+    def update(
+        self,
+        indices: int | list | ndarray | slice | None = None,
+        timestep_start: int = 0,
+        **data
+    ):
         if len(data) == 0:
             raise ValueError('No data provided to `update`')
 
+        if timestep_start < 0:
+            raise ValueError(f'timestep_start must be non-negative, got {timestep_start}')
+
         indices, scalar_index = expand_indices(indices, self.max_episodes)
 
         if not exists(indices):
@@ -1028,11 +1036,11 @@ class ReplayBufferBase:
 
                 time_dim = value.shape[1]
 
-                if time_dim > self.max_timesteps:
+                if timestep_start + time_dim > self.max_timesteps:
                     raise ValueError(f'You exceeded the `max_timesteps` ({self.max_timesteps}) set on the replay buffer. Please increase it on init.')
 
                 value = np.broadcast_to(value, (len(indices), time_dim, *self.shapes[name]))
-                self.data[name][indices, :time_dim] = value[order]
+                self.data[name][indices, timestep_start:timestep_start + time_dim] = value[order]
 
             else:
                 value = cast_to_target_shape(value, self.meta_shapes[name], num_leading_dims = 1)
@@ -1235,8 +1243,12 @@ class ReplayBufferBase:
 
     # video
 
+    def image_fields(self):
+        """ all fields that look like images, in field declaration order """
+        return [name for name in self.shapes if len(self.shapes[name]) >= 2]
+
     def default_image_field(self):
-        candidates = sorted(name for name in self.fieldnames if len(self.shapes[name]) >= 2)
+        candidates = self.image_fields()
 
         if len(candidates) == 1:
             return candidates[0]
@@ -1246,17 +1258,28 @@ class ReplayBufferBase:
 
         raise ValueError(f'multiple image-like fields {candidates} - pass `field` explicitly')
 
-    def episode_frames(self, episode_index: int = 0, field: str | None = None):
-        field = default(field, self.default_image_field())
+    def episode_frames(self, episode_index: int | Sequence[int] = 0, field: str | None = None):
+        if not exists(field):
+            field = self.default_image_field()
 
         if field not in self.fieldnames:
             raise ValueError(f'invalid field name `{field}` - must be one of {self.fieldnames}')
 
-        if not (0 <= episode_index < self.num_episodes):
-            raise ValueError(f'episode_index {episode_index} out of range - buffer has {self.num_episodes} episodes')
+        episode_indices = [episode_index] if isinstance(episode_index, (int, np.integer)) else list(episode_index)
+
+        if len(episode_indices) == 0:
+            raise ValueError('at least one episode index is required')
+
+        for index in episode_indices:
+            if not (0 <= index < self.num_episodes):
+                raise ValueError(f'episode_index {index} out of range - buffer has {self.num_episodes} episodes')
+
+        frames = [self.data[field][index, :int(self.episode_lens[index])] for index in episode_indices]
+
+        if len(frames) == 1:
+            return frames[0].copy()
 
-        episode_len = int(self.episode_lens[episode_index])
-        return self.data[field][episode_index, :episode_len].copy()
+        return np.concatenate(frames, axis = 0)
 
     def save_video(self, path: str | Path, episode_index: int = 0, field: str | None = None, fps: int = 30, **kwargs):
         from memmap_replay_buffer.video import trajectory_to_video