wds mixture: unitree+nav converters (gnm/vamos), wds source branch in train_fast, AMP autoselect (V100 fp16), embedding expand 16->32, push/fetch/validate tooling, DATA.md
b152c62 verified | #!/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}") | |