LeRobot 項目部署運行邏輯（六）——visualize_dataset_html.py/visualize

可視化腳本包括了兩個方法：遠程下載 huggingface 上的數據集和使用本地數據集

腳本主要使用兩個：

目前來說，ACT 采集訓練用的是統一時間長度的數據集，此外，這兩個腳本最大的問題在于不能裁剪，這也是比較好的升級方向；

1 可視化運行

1.1 遠程 html

1.2 本地數據集

2 代碼詳解 visualize_dataset_html.py

2.1 綜述

2.2 流程概覽

2.3 庫引用

2.4?mian() 函數

2.5?關鍵函數

2.5.1 run_server() —— Flask 應用核心

2.5.2 get_ep_csv_fname(episode_id) ?

2.5.3?get_episode_data()

2.5.4?get_episode_video_paths

2.5.5?get_episode_language_instruction

2.5.6 get_dataset_info(repo_id)

2.5.7 visualize_dataset_html

3 代碼詳解 visualize_dataset.py

3.1 綜述

3.2 流程概覽

3.3 庫引用

3.4?mian() 函數

3.5 關鍵函數

3.5.1 采樣器（EpisodeSampler）

3.5.2 圖像轉換（to_hwc_uint8_numpy）

3.5.3 核心可視化函數（visualize_dataset()） ?

1 可視化運行

1.1 遠程 html

對于開源數據集，只需要在?huggingface 上查看 id，比如 aloha_static_coffee 這個：

點進去選擇 use this dataset，可以看到id

然后運行腳本：

python lerobot/scripts/visualize_dataset_html.py \--repo-id lerobot/aloha_static_coffee

下載數據集后生成 web browser：http://127.0.0.1:9090

可以看到采集的各類信息：

可以看到結果保存地址：

其中，下載的數據集默認存儲在了 /home/yejiangchen/.cache/huggingface/lerobot/lerobot

下次再運行會直接調用無需下載

此外，如果想運行本地數據集，則需要指定 --root：

python lerobot/scripts/visualize_dataset_html.py \--root /home/yejiangchen/.cache/huggingface/lerobot/lerobot/aloha_static_coffee \--repo-id lerobot/aloha_static_coffee

即可正常運行：

1.2 本地數據集

本地的話可以直接使用 visualize_dataset.py 腳本，測試一下之前下載的數據

python lerobot/scripts/visualize_dataset.py \--repo-id lerobot/aloha_static_coffee \--episode-index 0

2 代碼詳解 visualize_dataset_html.py

2.1 綜述

此腳本將?LeRobotDataset?中的視頻+時序傳感數據（動作、狀態等）渲染成交互式網頁，方便快速瀏覽與排查

視頻：在瀏覽器原生 <video> 標簽播放
時序數值：轉成 CSV 字符串，交給前端 Dygraphs JavaScript 庫即刻繪制折線圖
語言任務描述：展示在同一頁面
部署：內置 Flask 服務器（默認 127.0.0.1:9090）即可本地或經 SSH?tunnel 遠程查看

2.2 流程概覽

main() -> visualize_dataset_html() ->（配置、軟鏈接）-> run_server() ->（HTTP 請求）-> get_dataset_info()、get_episode_data()、get_episode_language_instruction() 等

main()└─ 解析 CLI 參數└─ （可選）加載本地/遠程數據集 → LeRobotDataset 或 IterableNamespace└─ visualize_dataset_html()├─ 創建/復用輸出目錄（含模板與靜態文件）├─ （本地數據集）軟鏈接視頻到 static/videos└─ run_server()                       ←– 關鍵：注冊所有 Flask 路由├─ "/"                            : 首頁 / 數據集選擇頁├─ "/<ns>/<name>"                : 自動跳到 episode_0└─ "/<ns>/<name>/episode_<id>"   : 主可視化頁面

2.3 庫引用

import argparse  # 用于解析命令行參數
import csv       # 用于生成 CSV 格式字符串
import json      # 用于解析和生成 JSON 數據
import logging   # 用于日志記錄
import re        # 用于正則表達式處理
import shutil    # 用于文件和目錄操作，如復制、刪除
import tempfile  # 用于創建臨時目錄
from io import StringIO  # 用于將字符串當作文件讀寫
from pathlib import Path  # 用于跨平臺路徑操作import numpy as np     # 數值計算庫
import pandas as pd    # 數據處理庫
import requests        # 用于發起 HTTP 請求
from flask import Flask, redirect, render_template, request, url_for  # Flask Web 框架核心組件from lerobot import available_datasets  # 導入可用數據集列表
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset  # LeRobotDataset 類
from lerobot.common.datasets.utils import IterableNamespace  # 簡單 namespace 類型
from lerobot.common.utils.utils import init_logging  # 初始化日志設置

2.4?mian() 函數

作為腳本入口，負責解析所有命令行參數并據此準備數據集實例，最后調用 visualize_dataset_html() 啟動可視化流程

核心流程： ?

用 argparse 定義并讀取參數（如 --repo-id、--root、--episodes、--serve 等）
根據 --load-from-hf-hub 決定是實例化完整的 LeRobotDataset（加載本地/緩存數據與視頻），還是只拉取元信息 (get_dataset_info)
將解析好的 dataset 對象與其它參數傳入 visualize_dataset_html()

參數	作用	典型值
`--repo-id`	HF Hub 上的數據集 `namespace/name`	`lerobot/pusht`
`--root`	本地數據集根目錄	`./data`
`--load-from-hf-hub`	整數；為 1 時只下拉 meta / parquet / mp4，不構造完整 `LeRobotDataset`	`0/1`
`--episodes`	想看的 episode 索引列表	`0 3 5`
`--host`, `--port`	Flask 服務地址	默認 `127.0.0.1:9090`
`--tolerance-s`	時間戳容差，保證?`fps`?一致性	`1e-4`

def main():# 入口：解析命令行并調用可視化函數parser = argparse.ArgumentParser()parser.add_argument("--repo-id",type=str,default=None,help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht`).",)parser.add_argument("--root",type=Path,default=None,help="Root directory for a dataset stored locally (e.g. `--root data`).",)parser.add_argument("--load-from-hf-hub",type=int,default=0,help="Load videos and parquet files from HF Hub rather than local system.",)parser.add_argument("--episodes",type=int,nargs="*",default=None,help="Episode indices to visualize (e.g. `0 1 5 6`).",)parser.add_argument("--output-dir",type=Path,default=None,help="Directory path to write html files and kickoff a web server.",)parser.add_argument("--serve",type=int,default=1,help="Launch web server.",)parser.add_argument("--host",type=str,default="127.0.0.1",help="Web host used by the http server.",)parser.add_argument("--port",type=int,default=9090,help="Web port used by the http server.",)parser.add_argument("--force-override",type=int,default=0,help="Delete the output directory if it exists already.",)parser.add_argument("--tolerance-s",type=float,default=1e-4,help=("Tolerance in seconds used to ensure data timestamps respect the dataset fps value""If not given, defaults to 1e-4."),)args = parser.parse_args()  # 解析命令行參數kwargs = vars(args)repo_id = kwargs.pop("repo-id")  # 獲取 repo-id 并從 kwargs 刪除load_from_hf_hub = kwargs.pop("load_from_hf_hub")root = kwargs.pop("root")tolerance_s = kwargs.pop("tolerance_s")dataset = Noneif repo_id:# 根據 load_from_hf_hub 決定實例化 LeRobotDataset 還是只讀 metadataset = (LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)if not load_from_hf_hubelse get_dataset_info(repo_id))visualize_dataset_html(dataset, **vars(args))  # 調用主可視化入口if __name__ == "__main__":# 腳本直接運行時進入 mainmain()

2.5?關鍵函數

2.5.1 run_server() —— Flask 應用核心

全局配置：app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0 # 每次刷新都拉最新資源

路由配置：

路由	功能
`/`	? 如果腳本在“單數據集模式”（已傳 `dataset`），立刻重定向到 episode?0 ? 否則渲染首頁，列出推薦 (`featured_datasets`) + 全部可用數據集 (`lerobot_datasets`)
`/<ns>/<name>`	純跳轉：把 `<dataset>/episode_0` 作為入口
`/<ns>/<name>/episode_<id>`	主工作函數： 1.若腳本啟動時沒載數據，就動態 `get_dataset_info()` 2.檢查數據集版本 <2 則拒絕（舊格式） 3.調用 `get_episode_data()` → CSV + 列信息；拼裝 Video?URL / Tasks?Text 4.把所有信息喂給 `visualize_dataset_template.html` 渲染

def run_server(dataset: LeRobotDataset | IterableNamespace | None,episodes: list[int] | None,host: str,port: str,static_folder: Path,template_folder: Path,
):"""啟動 Flask HTTP 服務，渲染可視化頁面。參數:- dataset: 已加載的數據集實例或 None- episodes: 要展示的 episode 列表或 None- host, port: 服務監聽地址與端口- static_folder: 靜態文件目錄（視頻、JS、CSS）- template_folder: Jinja2 模板目錄"""app = Flask(__name__,static_folder=static_folder.resolve(),  # 靜態資源路徑template_folder=template_folder.resolve()  # 模板文件路徑)app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0  # 禁用瀏覽器緩存，確保每次都拉最新的資源@app.route("/")def hommepage(dataset=dataset):"""應用根路由：根據有無 dataset 參數決定重定向或渲染選擇頁"""if dataset:# 如果在腳本啟動時傳入 dataset，直接跳轉到第 0 集dataset_namespace, dataset_name = dataset.repo_id.split("/")return redirect(url_for("show_episode",dataset_namespace=dataset_namespace,dataset_name=dataset_name,episode_id=0,))# 否則嘗試從 query 參數讀取 dataset & episode 再跳轉dataset_param, episode_param = None, Noneall_params = request.argsif "dataset" in all_params:dataset_param = all_params["dataset"]if "episode" in all_params:episode_param = int(all_params["episode"])if dataset_param:dataset_namespace, dataset_name = dataset_param.split("/")return redirect(url_for("show_episode",dataset_namespace=dataset_namespace,dataset_name=dataset_name,episode_id=episode_param if episode_param is not None else 0,))# 默認渲染首頁，列出 featured + 全部 available datasetsfeatured_datasets = ["lerobot/aloha_static_cups_open","lerobot/columbia_cairlab_pusht_real","lerobot/taco_play",]return render_template("visualize_dataset_homepage.html",featured_datasets=featured_datasets,lerobot_datasets=available_datasets,)@app.route("/<string:dataset_namespace>/<string:dataset_name>")

2.5.2 get_ep_csv_fname(episode_id) ?

簡單工具，按約定返回某集 CSV 文件名 episode_{id}.csv

def get_ep_csv_fname(episode_id: int):# 根據 episode 索引構造 CSV 文件名ep_csv_fname = f"episode_{episode_id}.csv"return ep_csv_fname

2.5.3?get_episode_data()

把單個 episode 的多通道數值數據 ->?二維列表 ->?CSV 字符串（返給前端 JS）

1. 列挑選

selected = [col for col, ft in ds.features.items()if ft["dtype"] in ["float32", "int32"]]
selected.remove("timestamp")

2.?過濾高維張量：shape 維度 > 1 的列記入 ignored_columns，避免動態圖崩潰

3. 列名展開：如果在 meta 里有 names 用定義好的；否則按 col_0 … col_n 生成

4. 取數據：本地 LeRobotDataset 利用 .episode_data_index 截取 parquet；Hub?Only 直接 pd.read_parquet(url)

5. 轉換為 CSV string（StringIO + csv.writer）

def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index):"""獲取 episode 的時序數據，并將其轉換為 CSV 字符串返回。Returns:- csv_string: CSV 格式的整個 episode 數據- columns: [{key: 原始列名, value: 展開后子列名列表}, ...]- ignored_columns: 被忽略的高維列名稱列表"""columns = []  # 存儲展開后列的信息# 選出所有 dtype 為 float32/int32 的數值列selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]selected_columns.remove("timestamp")  # timestamp 先單獨處理ignored_columns = []  # 高維列名稱for column_name in selected_columns:shape = dataset.features[column_name]["shape"]  # 列的原始 shapeshape_dim = len(shape)if shape_dim > 1:# 如果維度 >1，則忽略，不支持 Dygraph 繪多維張量selected_columns.remove(column_name)ignored_columns.append(column_name)# CSV header: timestamp + 各子列名header = ["timestamp"]# 遍歷每個一維列，展開成多列子名稱for column_name in selected_columns:dim_state = (dataset.meta.shapes[column_name][0]if isinstance(dataset, LeRobotDataset)else dataset.features[column_name].shape[0])if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]:# 如果 meta 中定義了 names，則使用自定義子列名column_names = dataset.features[column_name]["names"]while not isinstance(column_names, list):column_names = list(column_names.values())[0]else:# 否則按 col_0...col_n 展開column_names = [f"{column_name}_{i}" for i in range(dim_state)]columns.append({"key": column_name, "value": column_names})header += column_names  # 累加到 CSV header# timestamp 放回最前selected_columns.insert(0, "timestamp")if isinstance(dataset, LeRobotDataset):# 本地模式：根據 index 范圍 select pandas DataFramefrom_idx = dataset.episode_data_index["from"][episode_index]to_idx   = dataset.episode_data_index["to"][episode_index]data = (dataset.hf_dataset.select(range(from_idx, to_idx)).select_columns(selected_columns).with_format("pandas"))else:# 遠程模式：通過 HTTP 拉取 parquet，然后篩列repo_id = dataset.repo_idurl = (f"https://huggingface.co/datasets/{repo_id}/resolve/main/"+ dataset.data_path.format(episode_chunk=int(episode_index) // dataset.chunks_size,episode_index=episode_index))df = pd.read_parquet(url)data = df[selected_columns]# 構造 numpy 二維數組：首列 timestamp，其余為各子列值rows = np.hstack((np.expand_dims(data["timestamp"], axis=1),*[np.vstack(data[col]) for col in selected_columns[1:]],)).tolist()# 寫 CSV 到內存字符串csv_buffer = StringIO()csv_writer = csv.writer(csv_buffer)csv_writer.writerow(header)csv_writer.writerows(rows)csv_string = csv_buffer.getvalue()return csv_string, columns, ignored_columns

2.5.4?get_episode_video_paths

僅在本地 LeRobotDataset 場景下，獲取指定 episode 在底層 HF 數據集中的視頻文件路徑列表（內部沒用到，備用）

找到該集第一幀在整表中的行索引
針對每個 dataset.meta.video_keys，在對應列讀取 ["path"] 字段

def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:# hack: 取該 episode 第一幀索引以定位 video pathfirst_frame_idx = dataset.episode_data_index["from"][ep_index].item()return [dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"]for key in dataset.meta.video_keys]

2.5.5?get_episode_language_instruction

僅在數據集包含 language_instruction 特征時調用，從對應行抽取并清洗掉 Tensor 的包裝字符串，返回指令文本

判斷 dataset.features 是否存在 language_instruction
取該集第一幀索引，讀取字段，去掉前后綴冗余信息

def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:# 如果數據集含 language_instruction 特征，則提取并清洗字符串if "language_instruction" not in dataset.features:return Nonefirst_frame_idx = dataset.episode_data_index["from"][ep_index].item()language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]# 去除 Tensor 格式冗余包裝return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")

2.5.6 get_dataset_info(repo_id)

遠程數據輔助：拉 meta/info.json 并包成 IterableNamespace

額外用 episodes.jsonl 找每一集的 tasks 列表

def get_dataset_info(repo_id: str) -> IterableNamespace:# 遠程拉取 meta/info.json 并轉為 IterableNamespaceresponse = requests.get(f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json",timeout=5)response.raise_for_status()dataset_info = response.json()dataset_info["repo_id"] = repo_idreturn IterableNamespace(dataset_info)

2.5.7 visualize_dataset_html

搭建靜態目錄結構（HTML 模板 + 靜態資源），并根據是否已有數據集對象決定是否創建視頻軟鏈接，最后根據 serve 標志調用 run_server()

調用 init_logging() 初始化日志設置
計算模板目錄 templates，創建或清空（若 force_override）輸出目錄以及 static 子目錄
若傳入本地 LeRobotDataset，在 static/videos 下打軟鏈接指向數據集的 videos 文件夾
若 serve 為真，調用 run_server() 啟動 Flask 服務

def visualize_dataset_html(dataset: LeRobotDataset | None,episodes: list[int] | None = None,output_dir: Path | None = None,serve: bool = True,host: str = "127.0.0.1",port: int = 9090,force_override: bool = False,
) -> Path | None:# 主函數：準備靜態目錄 & 啟動服務器init_logging()  # 配置根日志級別等template_dir = Path(__file__).resolve().parent.parent / "templates"if output_dir is None:# 未指定輸出目錄時，創建臨時目錄output_dir = tempfile.mkdtemp(prefix="lerobot_visualize_dataset_")output_dir = Path(output_dir)if output_dir.exists():if force_override:shutil.rmtree(output_dir)  # 強制覆蓋時先刪掉else:logging.info(f"Output directory already exists. Loading from it: '{output_dir}'")output_dir.mkdir(parents=True, exist_ok=True)static_dir = output_dir / "static"static_dir.mkdir(parents=True, exist_ok=True)if dataset is None:# 僅在無本地 dataset 且 serve=True 時進入 run_serverif serve:run_server(dataset=None,episodes=None,host=host,port=port,static_folder=static_dir,template_folder=template_dir,)else:# 本地數據集：在 static/videos 創建軟鏈接到 dataset.root/videosif isinstance(dataset, LeRobotDataset):ln_videos_dir = static_dir / "videos"if not ln_videos_dir.exists():ln_videos_dir.symlink_to((dataset.root / "videos").resolve())# 啟動服務器if serve:run_server(dataset, episodes, host, port, static_dir, template_dir)

3 代碼詳解 visualize_dataset.py

3.1 綜述

此腳本基于 Rerun SDK，實現對 LeRobotDataset 中單個 episode 進行可視化或記錄，主要有三種模式：

本地交互模式（mode="local"）直接在當前機器彈出可視化窗口，用于快速調試與觀測
遠端服務模式（mode="distant"）在數據存放的遠端機器上啟動 WebSocket+HTTP 服務，本地通過 rerun ws://… 連接瀏覽
離線保存模式（--save 1）將整次會話記錄到一個 .rrd 文件，后續可通過 rerun path/to/file.rrd 離線回放

其中，腳本既能實時顯示視頻幀，也能同步繪制動作、狀態、獎勵等時序數值

3.2 流程概覽

main()├─ 解析 CLI 參數├─ LeRobotDataset(repo_id, root, tolerance_s)└─ visualize_dataset(...)├─ EpisodeSampler(dataset, episode_index)├─ DataLoader(dataset, sampler, batch_size, num_workers)├─ rr.init(namespace, spawn=local_viewer?)├─ gc.collect()  # 避免多 worker 卡死├─ (mode=="distant")? rr.serve(web_port, ws_port)├─ for batch in DataLoader:│    ├─ for each frame in batch:│    │    ├─ rr.set_time_sequence/frame_index│    │    ├─ rr.set_time_seconds/timestamp│    │    ├─ rr.log(Image) for each camera│    │    └─ rr.log(Scalar) for each numeric field└─ 會話結束├─ local+save → rr.save(.rrd)└─ distant     → 阻塞等待 Ctrl–C

3.3 庫引用

import argparse               # 解析命令行參數模塊
import gc                     # 垃圾回收模塊，用于手動觸發回收
import logging                # 日志記錄模塊
import time                   # 時間相關函數模塊
from pathlib import Path      # 跨平臺路徑操作
from typing import Iterator   # 類型提示：迭代器import numpy as np            # 數值計算庫
import rerun as rr            # Rerun SDK，用于實時可視化
import torch                  # PyTorch 深度學習庫
import torch.utils.data       # PyTorch 數據加載工具
import tqdm                   # 進度條庫from lerobot.common.datasets.lerobot_dataset import LeRobotDataset  # 自定義 LeRobotDataset 數據集類

3.4?mian() 函數

強制要求：--repo-id（數據集標識）和 --episode-index（要可視化的集號）
可選：數據集根目錄、DataLoader 配置（--batch-size、--num-workers）、模式切換（--mode、--save、--output-dir、--web-port、--ws-port）等
最終實例化 LeRobotDataset 并調用 visualize_dataset()

def main():# 腳本入口：解析參數并調用可視化函數parser = argparse.ArgumentParser()parser.add_argument("--repo-id",type=str,required=True,help="HF Hub 上數據集標識，例如 `lerobot/pusht`。",)parser.add_argument("--episode-index",type=int,required=True,help="要可視化的 episode 索引。",)parser.add_argument("--root",type=Path,default=None,help="本地數據集根目錄，例如 `--root data`。默認使用 HuggingFace 緩存。",)parser.add_argument("--output-dir",type=Path,default=None,help="保存 .rrd 文件的目錄，當 `--save 1` 時生效。",)parser.add_argument("--batch-size",type=int,default=32,help="DataLoader 的 batch 大小。",)parser.add_argument("--num-workers",type=int,default=4,help="DataLoader 的并行工作進程數。",)parser.add_argument("--mode",type=str,default="local",help=("可視化模式：'local' 或 'distant'。`"local` 會本地彈出 viewer；`"distant` 則啟動服務供遠程瀏覽。"),)parser.add_argument("--web-port",type=int,default=9090,help="`--mode distant` 時的 HTTP 服務端口。",)parser.add_argument("--ws-port",type=int,default=9087,help="`--mode distant` 時的 WebSocket 服務端口。",)parser.add_argument("--save",type=int,default=0,help=("是否保存為 .rrd 文件，啟用后會禁用彈窗。""使用 `--output-dir path` 指定目錄。"),)parser.add_argument("--tolerance-s",type=float,default=1e-4,help=("時間戳容差，保證與 fps 一致。""傳入 LeRobotDataset 構造參數。"),)args = parser.parse_args()         # 解析命令行kwargs = vars(args)repo_id = kwargs.pop("repo_id")   # 提取 repo_idroot = kwargs.pop("root")         # 提取 roottolerance_s = kwargs.pop("tolerance_s")  # 提取容差參數logging.info("Loading dataset")  # 日志：開始加載數據集dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)  # 構造數據集visualize_dataset(dataset, **vars(args))  # 調用可視化主函數if __name__ == "__main__":# 如果腳本被直接執行，則運行 main()main()

3.5 關鍵函數

3.5.1 采樣器（EpisodeSampler）

只遍歷指定 episode 在底層數據表（Parquet）中的幀索引范圍，供 PyTorch DataLoader 使用

class EpisodeSampler(torch.utils.data.Sampler):# 自定義數據采樣器，僅遍歷指定 episode 的幀索引def __init__(self, dataset: LeRobotDataset, episode_index: int):# 根據 episode_index 從 dataset 中獲取起始和結束的全局幀索引from_idx = dataset.episode_data_index["from"][episode_index].item()to_idx = dataset.episode_data_index["to"][episode_index].item()# 保存幀索引范圍，用于 DataLoader 的 samplerself.frame_ids = range(from_idx, to_idx)def __iter__(self) -> Iterator:# 返回一個針對幀索引的迭代器return iter(self.frame_ids)def __len__(self) -> int:# 返回此 sampler 的總采樣數量（即幀數）return len(self.frame_ids)

3.5.2 圖像轉換（to_hwc_uint8_numpy）

把 PyTorch 的 C×H×W 浮點圖像張量（float32, 值域 [0,1]）轉換為 NumPy 的 H×W×C uint8 數組（值域 [0,255]），以便 Rerun 顯示

def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:# 將 C×H×W 的 float32 Torch 張量轉換為 H×W×C 的 uint8 NumPy 數組assert chw_float32_torch.dtype == torch.float32  # 確保數據類型為 float32assert chw_float32_torch.ndim == 3              # 確保是 3 維c, h, w = chw_float32_torch.shape               # 解包通道、高度、寬度assert c < h and c < w, f"expect channel first images, but instead {chw_float32_torch.shape}"# 先乘 255，再轉 uint8，然后 permute 到 HWC，最后轉換為 NumPyhwc_uint8_numpy = (chw_float32_torch * 255).type(torch.uint8).permute(1, 2, 0).numpy()return hwc_uint8_numpy                            # 返回處理后的圖像數組

3.5.3 核心可視化函數（visualize_dataset()） ?

1. 初始化

構造 DataLoader(dataset, sampler=EpisodeSampler, batch_size, num_workers)
調用 rr.init() 啟動 Rerun 會話
在遠端模式下額外執行 rr.serve() 開啟 WebSocket+HTTP 服務

2. 數據記錄 ?

遍歷每個 batch、每幀
用 rr.set_time_sequence／rr.set_time_seconds 標注時間信息
對所有攝像頭鍵（camera_keys）逐幀 rr.log(Image)
逐維 rr.log(Scalar) 記錄 action、observation.state、next.reward、next.done、next.success 等數值

3. 會話收尾

本地保存模式：rr.save() 寫出 .rrd 文件并返回路徑
遠端服務模式：進入阻塞循環以保持 WebSocket 連接，直至 Ctrl–C 退出

def visualize_dataset(dataset: LeRobotDataset,episode_index: int,batch_size: int = 32,num_workers: int = 0,mode: str = "local",web_port: int = 9090,ws_port: int = 9087,save: bool = False,output_dir: Path | None = None,
) -> Path | None:# 主可視化函數，根據模式(Local/Distant)實時或離線記錄并展示數據if save:# 如果要保存為 .rrd 文件，必須傳入 output_dirassert output_dir is not None, ("Set an output directory where to write .rrd files with `--output-dir path/to/directory`.")repo_id = dataset.repo_id                          # 獲取數據集唯一標識logging.info("Loading dataloader")               # 日志：開始加載 DataLoaderepisode_sampler = EpisodeSampler(dataset, episode_index)  # 創建只遍歷指定 episode 的 samplerdataloader = torch.utils.data.DataLoader(dataset,                                     # 數據集num_workers=num_workers,                     # 并行加載進程數batch_size=batch_size,                       # 每個 batch 的幀數sampler=episode_sampler,                     # 自定義 sampler)logging.info("Starting Rerun")                   # 日志：啟動 Rerun 會話if mode not in ["local", "distant"]:# 不支持其它模式時拋錯raise ValueError(mode)# 本地模式且不保存時，自動 spawn viewer；否則不彈出spawn_local_viewer = mode == "local" and not saverr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer)# Rerun v0.16 前的 workaround：觸發垃圾回收，避免多進程 DataLoader 卡住gc.collect()if mode == "distant":# 遠端模式：啟動 WebSocket + HTTP 服務，不自動打開瀏覽器rr.serve(open_browser=False, web_port=web_port, ws_port=ws_port)logging.info("Logging to Rerun")                  # 日志：開始寫入 Rerun 數據for batch in tqdm.tqdm(dataloader, total=len(dataloader)):# 遍歷每個 batch，顯示進度條for i in range(len(batch["index"])):# 記錄時間序列：幀索引與時間戳rr.set_time_sequence("frame_index", batch["frame_index"][i].item())rr.set_time_seconds("timestamp", batch["timestamp"][i].item())# 遍歷所有 camera key，記錄圖像for key in dataset.meta.camera_keys:rr.log(key, rr.Image(to_hwc_uint8_numpy(batch[key][i])))# 如果存在 action 字段，則按維度記錄每個動作值if "action" in batch:for dim_idx, val in enumerate(batch["action"][i]):rr.log(f"action/{dim_idx}", rr.Scalar(val.item()))# 如果存在 observation.state，則按維度記錄狀態值if "observation.state" in batch:for dim_idx, val in enumerate(batch["observation.state"][i]):rr.log(f"state/{dim_idx}", rr.Scalar(val.item()))# 可選字段：next.done, next.reward, next.successif "next.done" in batch:rr.log("next.done", rr.Scalar(batch["next.done"][i].item()))if "next.reward" in batch:rr.log("next.reward", rr.Scalar(batch["next.reward"][i].item()))if "next.success" in batch:rr.log("next.success", rr.Scalar(batch["next.success"][i].item()))if mode == "local" and save:# 本地保存模式：寫入 .rrd 文件并返回路徑output_dir = Path(output_dir)output_dir.mkdir(parents=True, exist_ok=True)repo_id_str = repo_id.replace("/", "_")rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd"rr.save(rrd_path)return rrd_pathelif mode == "distant":# 遠端模式：阻塞當前進程，直到手動按 Ctrl-Ctry:while True:time.sleep(1)except KeyboardInterrupt:print("Ctrl-C received. Exiting.")

4 本地數據集效果

python lerobot/scripts/visualize_dataset.py  --repo-id loacalhost/square_into_box --root=./collections/square_into_box/   --episode-index 0