Replay Buffer 使用教程#

本教程聚焦 TrajectoryReplayBuffer 的 实际使用 与 配置建议。 更完整的设计说明与数据流细节见 API 文档：Replay Buffer。

快速开始#

from rlinf.data.replay_buffer import TrajectoryReplayBuffer

buffer = TrajectoryReplayBuffer(
    seed=1234,
    enable_cache=True,
    cache_size=5,
    sample_window_size=100,
    auto_save=True,
    auto_save_path="/path/to/buffer",
    trajectory_format="pt",
)

常用参数#

  • enable_cache / cache_size：启用并控制缓存数量，用于提升采样吞吐。

  • sample_window_size：仅在最近 N 条轨迹内采样；0 表示全量。

  • auto_save：是否自动落盘；为 False 时仅缓存并在保存 checkpoint 时落盘。

  • auto_save_path：开启 auto_save 时的轨迹存储目录。

  • trajectory_format：pt`（默认）或 `pkl。

写入轨迹#

# trajectories 为 List[Trajectory]
buffer.add_trajectories(trajectories)

写入阶段的关键行为：

  • 为每条轨迹生成 uuid 与 trajectory_id

  • 更新 _trajectory_index 与计数器

  • 在后台线程异步保存轨迹文件（若 auto_save=True）

采样训练#

batch = buffer.sample(num_chunks=256)
# batch 形状: [num_chunks, ...]

采样在滑动窗口内随机抽取 transition，并返回与 rollout 对齐的 batch 字典。

保存与加载#

buffer.save_checkpoint("/path/to/ckpt")

buffer.load_checkpoint(
    load_path="/path/to/ckpt",
    is_distributed=True,
    local_rank=0,
    world_size=4,
)

保存 checkpoint 时会把缓存轨迹与 metadata 一并写入 checkpoint 路径。 加载时需要设置 load_path 指向包含 metadata 和轨迹文件的 checkpoint 目录。 轨迹数据保存格式为 trajectory_{trajectory_id}_{model_weights_uuid}_{model_update_count}.{trajectory_format}。

命令行测试#

python rlinf/data/replay_buffer.py \
  --load-path /path/to/buffer \
  --num-chunks 1024 \
  --cache-size 10 \
  --enable-cache

该命令会加载 buffer checkpoint 并进行一次采样，输出 batch 的 key 与 shape。

合并 / 拆分工具#

脚本位置：toolkits/replay_buffer/merge_or_split_replay_buffer.py

# 合并多个 rank（按原 trajectory_id 交错）
python toolkits/replay_buffer/merge_or_split_replay_buffer.py \
  --source-path /path/to/buffer \
  --save-path /path/to/merged \
  --copy

# 拆分单个 buffer，取前 N 条轨迹
python toolkits/replay_buffer/merge_or_split_replay_buffer.py \
  --source-path /path/to/buffer \
  --save-path /path/to/split \
  --split-count 30 \
  --copy

资源释放与重置#

buffer.close()        # 关闭异步保存线程
buffer.clear()        # 清空索引与计数
buffer.clear_cache()  # 清空缓存并关闭线程

实践建议#

  • 吞吐优先：开启 enable_cache，cache_size 设为近期活跃轨迹数。

  • 数据新鲜度：使用 sample_window_size 限制采样窗口。

可视化工具#

RLinf 提供了交互式可视化工具，用于检查 replay buffer 保存的轨迹数据。

功能特性#

  • 延迟加载：使用 TrajectoryReplayBuffer 按需加载轨迹，避免将所有数据加载到内存

  • 自动切换：到达最后一帧时自动前进到下一条轨迹

  • 跳转轨迹：在文本框中输入轨迹 ID 直接跳转

  • 多相机支持：查看主相机、腕部相机或额外视角相机

  • 批次导航：在 B > 1 时可在批次索引间导航

  • SSH/无显示器模式支持：保存图像以便查看

交互模式（本地机器有显示）#

python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0

键盘导航：

  • ← / → (或 p / n)：上一步/下一步（在边界自动切换轨迹）

  • ↑ / ↓：上一条/下一条轨迹

  • b / v：在批次索引间切换（如果 B > 1）

  • s：保存当前视图到图像文件

  • Home / End：跳转到第一步/最后一步

  • q / Esc：退出

  • 在文本框中输入轨迹 ID 直接跳转

SSH/无显示器模式#

使用无显示器交互脚本：

python toolkits/replay_buffer/visualize_headless.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --output viz.png

然后在 VSCode 中：

  1. 在编辑器中打开 viz.png

  2. 使用命令行提示进行导航

  3. 图像自动更新 - VSCode 会显示变化

命令：

  • n / next：下一步（在末尾自动切换到下一条轨迹）

  • p / prev：上一步

  • nt / nexttraj：下一条轨迹

  • pt / prevtraj：上一条轨迹

  • j <id>：跳转到轨迹 ID（例如 j 42）

  • info：显示当前位置

  • q / quit：退出

带 X11 转发的自动保存#

如果启用了 X11 转发：

python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --save_image --output viz.png

使用键盘导航，图像文件会自动更新。在 VSCode 中打开 viz.png 查看当前视图。

静态图像导出#

保存单帧而不进行交互：

python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --save_image --output viz.png --no_display

显示信息#

可视化工具显示：

  • 当前观察 （左面板）

  • 下一个观察 （右面板）

  • 轨迹 ID 和索引位置

  • 步骤 和 批次 索引

  • 每个转换的 动作、 奖励 和 完成 标志

查看不同相机角度#

# 主相机（默认）
python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --camera main_images

# 腕部相机
python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --camera wrist_images

# 额外视角相机
python toolkits/replay_buffer/visualize.py \
    --replay_dir logs/my_run/replay_buffer/rank_0 \
    --camera extra_view_images

注意事项#

  • 工具使用 TrajectoryReplayBuffer.load_checkpoint() 读取元数据和索引文件

  • 轨迹使用 _load_trajectory() 按需延迟加载

  • 缓存大小设置为 5 条轨迹以平衡内存和性能

  • 当在轨迹 i 的最后一帧按 → 时，会自动跳转到轨迹 i+1 的第 0 帧

  • 当在轨迹 i 的第一帧按 ← 时，会自动跳转到轨迹 i-1 的最后一帧

  • 图像文件以 150 DPI 保存，在保持良好质量的同时控制文件大小