#!/usr/bin/env python """Shared direct-to-webdataset writer for converters whose sources are NOT LeRobot repos (rosbags, hdf5, processed nav parquet). Skips the intermediate shard-builder layout and emits exactly what pack_wds.py emits: shard-%05d.tar __key__ / cam0.jpg / cam1.jpg / meta.npz manifest.json {name, samples, chunk, embodiment_id, norm_stats, shards, ...} tasks.json {task_index(str) -> instruction} Normalization contract matches pack_wds.py: state/action are normalized HERE with per-dataset mean/std computed in the converter's cheap first pass (no image decoding), and the applied stats land in manifest["norm_stats"]. Converters use: stats = StatsAccum(); stats.add(states, actions) per episode; st = stats.finalize() w = WdsWriter(out, name, st, chunk=50, embodiment_id=...) w.add_episode(ep_idx, states, actions, task_index, cam0_jpegs, cam1_jpegs=None) w.close(tasks={idx: text}, manifest_extra={...}) """ from __future__ import annotations import io import json from pathlib import Path import numpy as np def encode_jpeg(arr: np.ndarray, size: int = 256, quality: int = 92) -> bytes: """Same convention as build_shards_v2.encode_jpeg: plain squash resize.""" from PIL import Image img = Image.fromarray(arr) if img.size != (size, size): img = img.resize((size, size), Image.BILINEAR) buf = io.BytesIO() img.save(buf, format="JPEG", quality=quality, subsampling=0) return buf.getvalue() class StatsAccum: """Pooled mean/std over frames, same math as shards.py stats pooling. Separate counters: callers may feed different numbers of state rows and action rows (e.g. VAMOS feeds T states but T*chunk trajectory points).""" def __init__(self): self.s_n = self.a_n = 0 self.s_sum = self.s_sq = self.a_sum = self.a_sq = None def add(self, states: np.ndarray, actions: np.ndarray): states = np.asarray(states, np.float64) actions = np.asarray(actions, np.float64) if self.s_sum is None: self.s_sum = states.sum(0); self.s_sq = (states ** 2).sum(0) self.a_sum = actions.sum(0); self.a_sq = (actions ** 2).sum(0) else: self.s_sum += states.sum(0); self.s_sq += (states ** 2).sum(0) self.a_sum += actions.sum(0); self.a_sq += (actions ** 2).sum(0) self.s_n += len(states) self.a_n += len(actions) def finalize(self) -> dict: sn, an = max(1, self.s_n), max(1, self.a_n) sm = self.s_sum / sn am = self.a_sum / an sv = np.clip(self.s_sq / sn - sm ** 2, 0, None) av = np.clip(self.a_sq / an - am ** 2, 0, None) return { "observation.state": {"mean": sm.tolist(), "std": np.sqrt(sv).tolist()}, "action": {"mean": am.tolist(), "std": np.sqrt(av).tolist()}, "count": int(self.s_n), } class WdsWriter: def __init__(self, out: Path, name: str, stats: dict, chunk: int = 50, embodiment_id: int = 0, maxcount: int = 3000): import webdataset as wds self.out = Path(out) self.out.mkdir(parents=True, exist_ok=True) self.name, self.chunk, self.embodiment_id = name, chunk, embodiment_id self.stats = stats self.s_mean = np.asarray(stats["observation.state"]["mean"], np.float32) self.s_std = np.clip(np.asarray(stats["observation.state"]["std"], np.float32), 1e-6, None) self.a_mean = np.asarray(stats["action"]["mean"], np.float32) self.a_std = np.clip(np.asarray(stats["action"]["std"], np.float32), 1e-6, None) self.sink = wds.ShardWriter(str(self.out / "shard-%05d.tar"), maxcount=maxcount, verbose=0) self.n = 0 self.episodes = 0 def add_episode(self, ep_idx: int, states, actions, task_index, cam0_jpegs: list, cam1_jpegs: list | None = None, action_chunks: np.ndarray | None = None): """states f32[T,S], actions f32[T,A], cam0_jpegs list[bytes] of length T. task_index: one int for the whole episode, or an int array [T] when the source labels every frame (VAMOS rows). action_chunks f32[T, chunk, A] overrides the default consecutive-future chunking (used when each frame carries its own precomputed trajectory). """ states = (np.asarray(states, np.float32) - self.s_mean) / self.s_std actions = (np.asarray(actions, np.float32) - self.a_mean) / self.a_std T = len(states) assert len(cam0_jpegs) == T, f"ep {ep_idx}: {len(cam0_jpegs)} jpegs vs {T} frames" task_per_frame = (np.full(T, task_index, np.int32) if np.isscalar(task_index) else np.asarray(task_index, np.int32)) S, A, CH = states.shape[1], actions.shape[1], self.chunk for f in range(T): if action_chunks is not None: chunk = ((np.asarray(action_chunks[f], np.float32) - self.a_mean) / self.a_std) is_pad = np.zeros(CH, bool) else: avail = min(CH, T - f) chunk = np.empty((CH, A), np.float32) chunk[:avail] = actions[f : f + avail] if avail < CH: chunk[avail:] = actions[f + avail - 1] is_pad = np.zeros(CH, bool) is_pad[avail:] = True buf = io.BytesIO() np.savez(buf, state=states[f], action_chunk=chunk, action_is_pad=is_pad, action_dim=np.int32(A), state_dim=np.int32(S), task_index=task_per_frame[f], embodiment_id=np.int32(self.embodiment_id)) cam1 = cam1_jpegs[f] if cam1_jpegs else b"" self.sink.write({"__key__": f"{self.name}_{ep_idx:06d}_{f:05d}", "cam0.jpg": cam0_jpegs[f], "cam1.jpg": cam1, "meta.npz": buf.getvalue()}) self.n += 1 self.episodes += 1 def close(self, tasks: dict, manifest_extra: dict | None = None): self.sink.close() (self.out / "tasks.json").write_text( json.dumps({str(k): v for k, v in tasks.items()}, ensure_ascii=False)) (self.out / "manifest.json").write_text(json.dumps({ **(manifest_extra or {}), "name": self.name, "samples": self.n, "chunk": self.chunk, "episodes": self.episodes, "embodiment_id": self.embodiment_id, "norm_stats": self.stats, "shards": sorted(p.name for p in self.out.glob("shard-*.tar")), }, indent=2)) print(f"{self.name}: {self.n:,} семплов, {self.episodes} эпизодов -> " f"{len(list(self.out.glob('shard-*.tar')))} тарболов в {self.out}")