在本教程中，我们围绕 NVIDIA Cosmos3-DROID 数据集设计了一个端到端的流式机器人学习流水线，而无需在本地下载其 707 GB 的仓库。我们首先审视 LeRobotDataset v3.0 的结构，并从 info.json、任务元数据、episode 表和数据集统计信息构建元数据图，然后使用 HTTP 字节范围访问配合 PyArrow 来有选择地读取 Parquet 行组和列。我们将单个 episode 转换为状态-动作轨迹，并分析关节运动、夹爪事件、笛卡尔末端执行器路径和动作频率谱，之后仅通过基于 seek 的 PyAV/FFmpeg 访问解码所需的 AV1 视频窗口。接着，我们使用数据集统计信息对观测和动作进行归一化，构建一个带有可选视觉条件的 ACT 风格分块 PyTorch 数据集，并训练一个多模态行为克隆策略。最后，我们通过带有时间集成动作块的开环 rollout 评估学习到的策略，报告相对于平均动作基线的逐关节 MSE 和 R^2，可视化预测动作与真实动作的对比，并保存完整的策略检查点以供下游使用。
复制代码已复制使用其他浏览器
import subprocess, sys, os, json, math, time, warnings, random, tempfile
warnings.filterwarnings("ignore")
subprocess.run([sys.executable, "-m", "pip", "install", "-q",
"huggingface_hub>=0.34.0", "pyarrow>=15.0", "av>=12.0",
"pandas", "matplotlib", "tqdm"], check=False)
import numpy as np, pandas as pd, pyarrow as pa, pyarrow.parquet as pq
import matplotlib.pyplot as plt
from huggingface_hub import HfApi, HfFileSystem, hf_hub_download, hf_hub_url
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
REPO_ID = "nvidia/Cosmos3-DROID"
ROOT = "success"
VIDEO_KEY = "observation.image.wrist_image_left"
FPS = 15
N_EPISODES = 48
HORIZON = 8
OBS_HISTORY = 2
USE_VISION = True
N_VIS_EPS = 6
VIS_SIZE = 96
EPOCHS = 12
BATCH = 256
SEED = 0
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[env] torch={torch.__version__} device={DEV}")
if os.environ.get("HF_TOKEN"):
from huggingface_hub import login; login(os.environ["HF_TOKEN"])
api = HfApi()
fs = HfFileSystem()
HFS = lambda rel: f"datasets/{REPO_ID}/{rel}"
URL = lambda rel: hf_hub_url(REPO_ID, rel, repo_type="dataset")
print("\n" + "="*78 + "\n1. REPO INTROSPECTION\n" + "="*78)
all_files = api.list_repo_files(REPO_ID, repo_type="dataset")
print(f"total files in repo : {len(all_files):,}")
for prefix in ("success/data", "success/videos", "success/meta",
"failure/data", "failure/videos", "failure/meta"):
print(f" {prefix:<18} {sum(f.startswith(prefix) for f in all_files):>6,} files")
data_shards = sorted(f for f in all_files if f.startswith(f"{ROOT}/data/") and f.endswith(".parquet"))
vid_shards = sorted(f for f in all_files if f.startswith(f"{ROOT}/videos/{VIDEO_KEY}/"))
meta_files = sorted(f for f in all_files if f.startswith(f"{ROOT}/meta/"))
print(f"\n[{ROOT}] data shards={len(data_shards)} video shards({VIDEO_KEY})={len(vid_shards)}")
print("first data shard :", data_shards[0])
print("first video shard:", vid_shards[0])
print("\n" + "="*78 + "\n2. METADATA\n" + "="*78)
info = json.load(open(hf_hub_download(REPO_ID, f"{ROOT}/meta/info.json", repo_type="dataset")))
print(f"episodes={info.get('total_episodes'):,} frames={info.get('total_frames'):,} "
f"tasks={info.get('total_tasks'):,} fps={info.get('fps')}")
print("data_path template :", info.get("data_path"))
print("video_path template:", info.get("video_path"))
FEATURES = info["features"]
state_keys = sorted(k for k in FEATURES if k.startswith("observation.state"))
action_keys = sorted(k for k in FEATURES if k.startswith("action."))
video_keys = sorted(k for k in FEATURES if FEATURES[k]["dtype"] == "video")
print("\nstate :", [f"{k.split('.')[-1]}{tuple(FEATURES[k]['shape'])}" for k in state_keys])
print("action :", [f"{k.split('.')[-1]}{tuple(FEATURES[k]['shape'])}" for k in action_keys])
print("video :", video_keys)
tdf = pd.read_parquet(hf_hub_download(REPO_ID, f"{ROOT}/meta/tasks.parquet", repo_type="dataset"))
tdf = tdf.reset_index()
tcol = "task" if "task" in tdf.columns else tdf.columns[0]
TASKS = dict(zip(tdf["task_index"].astype(int), tdf[tcol].astype(str))) if "task_index" in tdf \
else {i: str(v) for i, v in enumerate(tdf[tcol])}
print(f"\n{len(TASKS):,} task strings. Random sample:")
for t in random.sample(list(TASKS.values()), min(8, len(TASKS))): print(" ·", t[:90])
ep_files = [f for f in meta_files if "/episodes/" in f and f.endswith(".parquet")]
eps = pd.concat([pd.read_parquet(hf_hub_download(REPO_ID, f, repo_type="dataset"))
for f in ep_files[:4]], ignore_index=True)
print(f"\nepisodes table: {len(eps):,} rows")
print("columns:", [c for c in eps.columns if not c.startswith("stats")][:14], "...")
print(eps[[c for c in ("episode_index", "length", "data/chunk_index", "data/file_index")
if c in eps.columns]].head())
我们初始化 Colab 环境，安装所需的库，并配置 Cosmos3-DROID 数据集、episode、视频和训练参数。我们检查仓库结构，识别可用的数据、视频和元数据分片，而无需下载完整数据集。随后我们加载核心元数据和任务描述，以了解数据集 schema、可用的 state/action 特征以及 episode 组织方式。
Copy CodeCopiedUse a different Browser
print("\n" + "="*78 + "\n3. BYTE-RANGE PARQUET READER\n" + "="*78)
def open_pf(rel_path):
return pq.ParquetFile(fs.open(HFS(rel_path), "rb"))
def rowgroup_span(pf):
md, starts, c = pf.metadata, [], 0
for i in range(md.num_row_groups):
starts.append(c); c += md.row_group(i).num_rows
return np.array(starts), c
def read_rows(pf, lo, hi, columns):
starts, total = rowgroup_span(pf)
ends = np.append(starts[1:], total)
rgs = [i for i in range(len(starts)) if starts[i] < hi and ends[i] > lo]
tbl = pf.read_row_groups(rgs, columns=columns)
return tbl.slice(lo - starts[rgs[0]], hi - lo)
def col2np(tbl, name):
ca = tbl.column(name).combine_chunks()
if pa.types.is_list(ca.type) or pa.types.is_large_list(ca.type) or pa.types.is_fixed_size_list(ca.type):
flat = np.asarray(ca.flatten().to_numpy(zero_copy_only=False))
return flat.reshape(len(ca), -1).astype(np.float32)
return np.asarray(ca.to_numpy(zero_copy_only=False)).reshape(-1, 1).astype(np.float32)
SHARD = data_shards[0]
pf = open_pf(SHARD)
md = pf.metadata
print(f"shard : {SHARD}")
print(f"rows : {md.num_rows:,} row_groups: {md.num_row_groups} "
f"compressed: {md.serialized_size/1e6:.1f} MB footer")
print(f"columns : {len(pf.schema_arrow.names)}")
t0 = time.time()
ep_idx_all = pf.read(columns=["episode_index"]).column("episode_index").to_numpy()
print(f"pulled episode_index column ({len(ep_idx_all):,} rows) in {time.time()-t0:.1f}s")
uniq, first_pos = np.unique(ep_idx_all, return_index=True)
order = np.argsort(first_pos)
uniq = uniq[order]; first_pos = first_pos[order]
last_pos = np.append(first_pos[1:], len(ep_idx_all))
EP_BOUNDS = {int(e): (int(a), int(b)) for e, a, b in zip(uniq, first_pos, last_pos)}
print(f"{len(EP_BOUNDS)} episodes live in this shard "
f"(ids {uniq.min()}..{uniq.max()}, mean len {np.mean(last_pos-first_pos):.0f} frames)")
STATE_USE = ["observation.state.joint_positions", "observation.state.gripper_position",
"observation.state.cartesian_position"]
ACTION_USE = ["action.joint_velocity", "action.gripper_position"]
READ_COLS = STATE_USE + ACTION_USE + ["timestamp", "frame_index", "task_index", "episode_index"]
def load_episode(ep):
lo, hi = EP_BOUNDS[ep]
tbl = read_rows(pf, lo, hi, READ_COLS)
out = {k: col2np(tbl, k) for k in STATE_USE + ACTION_USE}
out["timestamp"] = col2np(tbl, "timestamp").ravel()
out["task_index"] = int(col2np(tbl, "task_index").ravel()[0])
out["task"] = TASKS.get(out["task_index"], "<unknown>")
out["state"] = np.concatenate([out[k] for k in STATE_USE], axis=1)
out["action"] = np.concatenate([out[k] for k in ACTION_USE], axis=1)
return out
EP0 = int(uniq[0]); traj = load_episode(EP0)
print(f"\nepisode {EP0}: T={len(traj['state'])} state_dim={traj['state'].shape[1]} "
f"action_dim={traj['action'].shape[1]}")
print(f"task: {traj['task']!r}")
print("\n" + "="*78 + "\n5. TRAJECTORY ANALYTICS\n" + "="*78)
q = traj["observation.state.joint_positions"]
grip = traj["observation.state.gripper_position"].ravel()
cart = traj["observation.state.cartesian_position"]
dq = traj["action.joint_velocity"]
t = traj["timestamp"]
fig = plt.figure(figsize=(15, 9))
ax = fig.add_subplot(2, 3, 1)
for j in range(q.shape[1]): ax.plot(t, q[:, j], lw=1.1, label=f"j{j+1}")
ax.set_title("joint positions [rad]"); ax.set_xlabel("s"); ax.legend(fontsize=6, ncol=2)
ax = fig.add_subplot(2, 3, 2)
ax.plot(t, grip, color="crimson", lw=1.4)
opens = np.where(np.abs(np.diff(grip)) > 0.05)[0]
for k in opens[:40]: ax.axvline(t[k], color="k", alpha=.15, lw=.8)
ax.set_title(f"gripper (|Δ|>0.05 events: {len(opens)})"); ax.set_xlabel("s")
ax = fig.add_subplot(2, 3, 3, projection="3d")
ax.plot(cart[:, 0], cart[:, 1], cart[:, 2], lw=1.2)
ax.scatter(*cart[0, :3], c="g", s=45, label="start"); ax.scatter(*cart[-1, :3], c="r", s=45, label="end")
ax.set_title("EE cartesian path [m]"); ax.legend(fontsize=7)
ax = fig.add_subplot(2, 3, 4)
im = ax.imshow(dq.T, aspect="auto", cmap="RdBu_r", vmin=-np.abs(dq).max(), vmax=np.abs(dq).max())
ax.set_title("action.joint_velocity (7 x T)"); ax.set_ylabel("joint"); plt.colorbar(im, ax=ax)
ax = fig.add_subplot(2, 3, 5)
freqs = np.fft.rfftfreq(len(dq), d=1/FPS)
for j in range(dq.shape[1]):
ax.semilogy(freqs, np.abs(np.fft.rfft(dq[:, j] - dq[:, j].mean())) + 1e-9, lw=.9)
ax.set_title("action spectra (Nyquist=7.5 Hz)"); ax.set_xlabel("Hz")
ax = fig.add_subplot(2, 3, 6)
lens = [EP_BOUNDS[e][1] - EP_BOUNDS[e][0] for e in list(EP_BOUNDS)[:2000]]
ax.hist(np.array(lens)/FPS, bins=40, color="steelblue")
ax.set_title(f"episode duration [s] (n={len(lens)})"); ax.set_xlabel("s")
plt.suptitle(f"{REPO_ID} · {ROOT} · ep {EP0} · {traj['task'][:70]}", y=1.0)
plt.tight_layout(); plt.show()
我们实现了一个字节范围 Parquet 读取器，仅通过 Hugging Face 文件系统访问所需的行组和列。我们在数据分片中识别回合边界，并将选定的状态和动作字段转换为 NumPy 轨迹。随后我们可视化关节位置、夹爪活动、笛卡尔运动、动作分布、频谱以及回合时长统计。
复制代码已复制使用其他浏览器
print("\n" + "="*78 + "\n6. VIDEO: SEEK-BASED AV1 DECODE (no full download)\n" + "="*78)
def video_window(ep):
row = eps.loc[eps["episode_index"] == ep]
if len(row) == 0: return None
row = row.iloc[0]
ci = int(row.get(f"videos/{VIDEO_KEY}/chunk_index", row.get("data/chunk_index", 0)))
fi = int(row.get(f"videos/{VIDEO_KEY}/file_index", row.get("data/file_index", 0)))
f0 = float(row.get(f"videos/{VIDEO_KEY}/from_timestamp", 0.0))
f1 = float(row.get(f"videos/{VIDEO_KEY}/to_timestamp",
f0 + int(row.get("length", 100))/FPS))
return f"{ROOT}/videos/{VIDEO_KEY}/chunk-{ci:03d}/file-{fi:03d}.mp4", f0, f1
def decode_pyav(url, t0, t1, max_frames, stride, size):
import av
c = av.open(url, options={"rw_timeout": "30000000"})
s = c.streams.video[0]; s.thread_type = "AUTO"
if t0 > 0: c.seek(int(t0 / s.time_base), stream=s)
out, k = [], 0
for fr in c.decode(s):
ts = float(fr.pts * s.time_base)
if ts < t0 - 1e-3: continue
if ts > t1 + 1e-3 or len(out) >= max_frames: break
if k % stride == 0:
out.append(fr.reformat(width=size, height=size, format="rgb24").to_ndarray())
k += 1
c.close()
return np.stack(out) if out else None
def decode_ffmpeg(url, t0, t1, max_frames, stride, size):
cmd = ["ffmpeg", "-v", "error", "-ss", f"{t0:.3f}", "-i", url,
"-t", f"{max(t1-t0, 0.5):.3f}",
"-vf", f"select=not(mod(n\\,{stride})),scale={size}:{size}",
"-vsync", "0", "-frames:v", str(max_frames),
"-f", "rawvideo", "-pix_fmt", "rgb24", "-"]
buf = subprocess.run(cmd, capture_output=True).stdout
n = len(buf) // (size*size*3)
return np.frombuffer(buf[:n*size*size*3], np.uint8).reshape(n, size, size, 3) if n else None
def get_frames(ep, max_frames=64, stride=2, size=VIS_SIZE):
w = video_window(ep)
if w is None: return None
rel, t0, t1 = w; url = URL(rel)
for fn in (decode_pyav, decode_ffmpeg):
try:
f = fn(url, t0, t1, max_frames, stride, size)
if f is not None and len(f): return f
except Exception as e:
print(f" {fn.__name__} failed: {type(e).__name__}: {str(e)[:80]}")
return None
frames = get_frames(EP0, max_frames=12, stride=max(1, len(q)//12), size=160)
if frames is not None:
print(f"decoded {frames.shape} from {video_window(EP0)[0]}")
fig, axs = plt.subplots(2, 6, figsize=(15, 5.2))
for i, ax in enumerate(axs.ravel()):
ax.axis("off")
if i < len(frames):
ax.imshow(frames[i]); ax.set_title(f"t≈{i*(len(q)//12)/FPS:.1f}s", fontsize=8)
plt.suptitle(f"{VIDEO_KEY} · ep {EP0} · {traj['task'][:60]}"); plt.tight_layout(); plt.show()
else:
print("video decode unavailable (AV1 codec missing) — continuing state-only.")
USE_VISION = False
print("\n" + "="*78 + "\n7. NORMALIZATION\n" + "="*78)
try:
stats = json.load(open(hf_hub_download(REPO_ID, f"{ROOT}/meta/stats.json", repo_type="dataset")))
def cat_stat(keys, field):
return np.concatenate([np.atleast_1d(np.asarray(stats[k][field], dtype=np.float32).ravel())
for k in keys])
S_MEAN, S_STD = cat_stat(STATE_USE, "mean"), cat_stat(STATE_USE, "std")
A_MEAN, A_STD = cat_stat(ACTION_USE, "mean"), cat_stat(ACTION_USE, "std")
print("using dataset-level stats from meta/stats.json")
except Exception as e:
print("stats.json unusable, will compute empirically:", type(e).__name__)
S_MEAN = S_STD = A_MEAN = A_STD = None
我们构建了一个基于 seek 的视频流水线，它只从某个 episode 中检索所需的时间窗口，而不是下载整个视频分片。我们同时支持 PyAV 和 FFmpeg 解码路径，以高效处理 AV1 视频，并对选中的帧进行缩放以便轻量处理。我们还从 stats.json 加载数据集级别的归一化统计量，并在这些统计量不可用时使用经验回退方案。
Copy CodeCopiedUse a different Browser
print("\n" + "="*78 + "\n8. BUILDING TRAINING SET\n" + "="*78)
ep_ids = [e for e in list(EP_BOUNDS) if EP_BOUNDS[e][1]-EP_BOUNDS[e][0] > HORIZON+OBS_HISTORY+4][:N_EPISODES]
EPISODES = {}
for i, e in enumerate(ep_ids):
EPISODES[e] = load_episode(e)
if (i+1) % 8 == 0: print(f" loaded {i+1}/{len(ep_ids)} episodes")
print(f"loaded {len(EPISODES)} episodes, {sum(len(v['state']) for v in EPISODES.values()):,} frames")
VIS_CACHE = {}
if USE_VISION:
for e in ep_ids[:N_VIS_EPS]:
T = len(EPISODES[e]["state"])
f = get_frames(e, max_frames=min(T, 200), stride=1, size=VIS_SIZE)
if f is not None:
VIS_CACHE[e] = f
print(f" video ep {e}: {f.shape}")
USE_VISION = len(VIS_CACHE) >= 2
print(f"vision enabled: {USE_VISION} ({len(VIS_CACHE)} episodes cached)")
if S_MEAN is None:
allS = np.concatenate([v["state"] for v in EPISODES.values()])
allA = np.concatenate([v["action"] for v in EPISODES.values()])
S_MEAN, S_STD = allS.mean(0), allS.std(0) + 1e-6
A_MEAN, A_STD = allA.mean(0), allA.std(0) + 1e-6
S_STD = np.maximum(S_STD, 1e-4); A_STD = np.maximum(A_STD, 1e-4)
class DroidChunks(Dataset):
def __init__(self, episodes, ep_list, vision):
self.eps, self.vision, self.items = episodes, vision, []
for e in ep_list:
if vision and e not in VIS_CACHE: continue
T = len(episodes[e]["state"])
if vision: T = min(T, len(VIS_CACHE[e]))
for i in range(OBS_HISTORY-1, T-HORIZON): self.items.append((e, i))
def __len__(self): return len(self.items)
def __getitem__(self, k):
e, i = self.items[k]; d = self.eps[e]
s = (d["state"][i-OBS_HISTORY+1:i+1] - S_MEAN) / S_STD
a = (d["action"][i:i+HORIZON] - A_MEAN) / A_STD
out = [torch.from_numpy(s.ravel().astype(np.float32)),
torch.from_numpy(a.astype(np.float32))]
if self.vision:
img = VIS_CACHE[e][i].astype(np.float32) / 255.0
out.insert(1, torch.from_numpy(img.transpose(2, 0, 1)))
return tuple(out)
pool = list(VIS_CACHE) if USE_VISION else ep_ids
tr_eps, te_eps = pool[:-2], pool[-2:]
tr, te = DroidChunks(EPISODES, tr_eps, USE_VISION), DroidChunks(EPISODES, te_eps, USE_VISION)
tl = DataLoader(tr, batch_size=BATCH, shuffle=True, num_workers=2, drop_last=True)
vl = DataLoader(te, batch_size=BATCH, shuffle=False, num_workers=2)
print(f"train windows={len(tr):,} ({len(tr_eps)} eps) val windows={len(te):,} ({len(te_eps)} eps)")
S_DIM, A_DIM = EPISODES[ep_ids[0]]["state"].shape[1], EPISODES[ep_ids[0]]["action"].shape[1]
class ChunkPolicy(nn.Module):
def __init__(self, s_dim, a_dim, horizon, vision, h=512):
super().__init__()
self.vision, self.horizon, self.a_dim = vision, horizon, a_dim
feat = h
if vision:
self.cnn = nn.Sequential(
nn.Conv2d(3, 32, 5, 2, 2), nn.GroupNorm(8, 32), nn.SiLU(),
nn.Conv2d(32, 64, 3, 2, 1), nn.GroupNorm(8, 64), nn.SiLU(),
nn.Conv2d(64,128, 3, 2, 1), nn.GroupNorm(8, 128), nn.SiLU(),
nn.Conv2d(128,256,3, 2, 1), nn.GroupNorm(8, 256), nn.SiLU(),
nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(256, 256))
feat += 256
self.smlp = nn.Sequential(nn.Linear(s_dim*OBS_HISTORY, h), nn.SiLU(), nn.Linear(h, h))
self.trunk = nn.Sequential(nn.Linear(feat, h), nn.SiLU(), nn.LayerNorm(h),
nn.Linear(h, h), nn.SiLU(), nn.LayerNorm(h))
self.head = nn.Linear(h, horizon*a_dim)
def forward(self, s, img=None):
z = self.smlp(s)
if self.vision: z = torch.cat([z, self.cnn(img)], -1)
return self.head(self.trunk(z)).view(-1, self.horizon, self.a_dim)
我们加载一个可配置的 episode 集合，并可选地为一个小型子集缓存同步的视觉观测，以保持训练在计算上可管理。我们构建一个 ACT 风格的 PyTorch 数据集，将观测历史和可选的图像与归一化的未来动作块相结合。然后我们定义一个分块策略架构，将 MLP 状态编码器与可选的 CNN 视觉编码器相结合，并预测一系列未来动作。
复制代码已复制使用其他浏览器
model = ChunkPolicy(S_DIM, A_DIM, HORIZON, USE_VISION).to(DEV)
print(f"\nmodel params: {sum(p.numel() for p in model.parameters())/1e6:.2f} M (vision={USE_VISION})")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, 3e-4, total_steps=EPOCHS*max(len(tl), 1), pct_start=.15)
scaler = torch.amp.GradScaler(DEV, enabled=(DEV == "cuda"))
hist = {"train": [], "val": []}
def run(loader, train):
model.train(train); tot = n = 0
for batch in loader:
batch = [b.to(DEV, non_blocking=True) for b in batch]
s, img, a = (batch[0], batch[1], batch[2]) if USE_VISION else (batch[0], None, batch[1])
with torch.set_grad_enabled(train), torch.amp.autocast(DEV, enabled=(DEV == "cuda")):
loss = F.smooth_l1_loss(model(s, img), a, beta=0.1)
if train:
opt.zero_grad(set_to_none=True); scaler.scale(loss).backward()
scaler.unscale_(opt); nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt); scaler.update(); sched.step()
tot += loss.item()*len(s); n += len(s)
return tot/max(n, 1)
print("\n" + "="*78 + "\n9. TRAINING\n" + "="*78)
for ep in range(EPOCHS):
t0 = time.time(); trl = run(tl, True); vll = run(vl, False)
hist["train"].append(trl); hist["val"].append(vll)
print(f"epoch {ep+1:>2}/{EPOCHS} train={trl:.5f} val={vll:.5f} ({time.time()-t0:.1f}s)")
我们初始化分块策略，并使用 AdamW、OneCycle 学习率调度、混合精度执行、梯度缩放和梯度裁剪对其进行优化。我们使用 Smooth L1 损失，使行为克隆对噪声较大或变化不定的遥操作动作更加鲁棒。我们按配置的 epoch 数训练策略，同时跟踪训练损失和验证损失以监控学习行为。
Copy CodeCopiedUse a different Browser
print("\n" + "="*78 + "\n10. OPEN-LOOP ROLLOUT (temporal ensembling)\n" + "="*78)
@torch.no_grad()
def rollout(ep, m=0.1):
d = EPISODES[ep]; T = len(d["state"])
if USE_VISION: T = min(T, len(VIS_CACHE[ep]))
acc = np.zeros((T, HORIZON, A_DIM), np.float32); cnt = np.zeros((T, HORIZON), np.float32)
model.eval()
for i in range(OBS_HISTORY-1, T-HORIZON):
s = torch.from_numpy(((d["state"][i-OBS_HISTORY+1:i+1]-S_MEAN)/S_STD).ravel()
.astype(np.float32))[None].to(DEV)
img = None
if USE_VISION:
img = torch.from_numpy((VIS_CACHE[ep][i].astype(np.float32)/255.)
.transpose(2, 0, 1))[None].to(DEV)
p = model(s, img)[0].float().cpu().numpy()*A_STD + A_MEAN
for k in range(HORIZON):
if i+k < T: acc[i+k, k] = p[k]; cnt[i+k, k] = math.exp(-m*k)
w = cnt[..., None]; pred = (acc*w).sum(1) / np.maximum(w.sum(1), 1e-8)
valid = cnt.sum(1) > 0
return pred, d["action"][:T], valid
ep_eval = te_eps[0]
pred, gt, valid = rollout(ep_eval)
mse = ((pred[valid]-gt[valid])**2).mean(0)
base = ((gt[valid].mean(0)-gt[valid])**2).mean(0)
names = [f"jvel_{i+1}" for i in range(7)] + ["gripper"]
print(f"episode {ep_eval} · task: {EPISODES[ep_eval]['task'][:70]}")
print(f"{'dim':<10}{'MSE':>12}{'mean-baseline':>16}{'R²':>10}")
for i, nm in enumerate(names[:A_DIM]):
print(f"{nm:<10}{mse[i]:>12.5f}{base[i]:>16.5f}{1-mse[i]/max(base[i],1e-9):>10.3f}")
print(f"{'OVERALL':<10}{mse.mean():>12.5f}{base.mean():>16.5f}{1-mse.mean()/base.mean():>10.3f}")
fig, axs = plt.subplots(3, 3, figsize=(15, 8), sharex=True)
for i, ax in enumerate(axs.ravel()):
if i >= A_DIM: ax.axis("off"); continue
ax.plot(gt[:, i], "k", lw=1.3, label="ground truth")
ax.plot(np.where(valid, pred[:, i], np.nan), "r", lw=1.1, alpha=.85, label="policy")
ax.set_title(names[i], fontsize=9)
if i == 0: ax.legend(fontsize=7)
axs.ravel()[-1].axis("off")
inset = fig.add_axes([0.71, 0.08, 0.24, 0.2])
inset.plot(hist["train"], label="train"); inset.plot(hist["val"], label="val")
inset.set_yscale("log"); inset.set_title("loss", fontsize=8); inset.legend(fontsize=6)
plt.suptitle(f"Open-loop chunked BC · {ROOT} ep {ep_eval} · vision={USE_VISION}")
plt.tight_layout(); plt.show()
torch.save({"model": model.state_dict(), "s_mean": S_MEAN, "s_std": S_STD,
"a_mean": A_MEAN, "a_std": A_STD, "cfg": dict(
state_keys=STATE_USE, action_keys=ACTION_USE, horizon=HORIZON,
obs_history=OBS_HISTORY, vision=USE_VISION, root=ROOT)},
"droid_chunk_policy.pt")
print("\nsaved -> droid_chunk_policy.pt")
print(f"""
{'='*78}
DONE. Everything above streamed from a 707 GB repo; peak disk use ≈ a few hundred MB.
Scale-up levers
· N_EPISODES / more shards -> data_shards[1:], rebuild EP_BOUNDS per shard
· ROOT="failure" -> 14,268 negative episodes for success classifiers
· VIDEO_KEY -> exterior_image_1_left / exterior_image_2_left
(3 synced views: multi-view or view-randomization)
· language -> 53,086 task strings; add a text encoder for VLA-style
conditioning instead of the state-only trunk
· targets -> swap ACTION_USE to action.cartesian_velocity for
end-effector control, or predict deltas
· real training -> pip install lerobot; LeRobotDataset("<local>/{ROOT}")
once you have local disk (v3.0 native loader)
{'='*78}""")
我们通过开环 rollout 评估训练好的策略，并使用指数加权时间集成来合并重叠的动作预测。我们计算逐关节 MSE、基线误差和 R^2，并将预测动作与真值轨迹一起绘制，同时绘制训练/验证损失曲线。最后，我们保存训练好的模型及其归一化统计量和配置元数据，以便在后续实验中复用该策略。
In conclusion, we showed how to turn a massive real-world robot dataset into a learning pipeline while keeping storage and data-transfer requirements extremely low. We used metadata-driven episode discovery, column- and row-group-level Parquet projection, and seek-based video decoding to access only the information required for analysis and training rather than materializing the full dataset. We combined proprioceptive state history with optional visual observations to train a chunked behavior-cloning policy and used temporal ensembling to obtain smoother action predictions during open-loop evaluation. The resulting workflow gives us a compact but extensible foundation that we can scale across additional shards, failure demonstrations, camera views, language instructions, or alternative action representations for more sophisticated robotics and vision-language-action experiments.
Check out the FULL CODES here. All credit goes to the researcher of this project. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
The post Building a Streaming Robotics Learning Pipeline Using NVIDIA Cosmos3-DROID appeared first on MarkTechPost.