"""DA3 spatial inputs for the 2026 v3 pipeline. Extends BehaviorV3Dataset with per-frame DA3 inputs (3 cams at DA3 resolution + robot->cam OpenCV extrinsics + intrinsics + ModernBERT task-language), and provides a data-loader factory that runs the frozen DA3-GIANT extractor once per BATCH (GPU) via the loader's batch hook. Geometry (empirically calibrated against GT depth, see /work/jack/behavior1k/calib): * robot2cam_pose[7] = [pos(3), quat_wxyz(4)] = the CAMERA POSE IN THE ROBOT FRAME, already OpenCV-convention (+Z optical axis). robot->cam = inv(pose_matrix). * intrinsics: fx = fy = W * 17.0/20.995 (OmniGibson VisionSensor defaults), cx=cy=W/2. """ import logging import os import pickle import glob import numpy as np from b1k.training.b1k_2026 import BehaviorV3Dataset, B1kInputs2026 logger = logging.getLogger(__name__) FOCAL_RATIO = 17.0 / 20.995 # OmniGibson VisionSensor default focal/aperture # dst rgb key -> pose parquet column (same camera) POSE_COLS = { "observation.images.rgb.head": "observation.robot2cam_pose.zed_link_camera_0", "observation.images.rgb.left_wrist": "observation.robot2cam_pose.left_realsense_link_camera_0", "observation.images.rgb.right_wrist": "observation.robot2cam_pose.right_realsense_link_camera_0", } # view order MUST match the bank builder: 0=main(head), 1=left, 2=right VIEW_ORDER = ( "observation.images.rgb.head", "observation.images.rgb.left_wrist", "observation.images.rgb.right_wrist", ) DEPTH_SRCS = { "observation.images.rgb.head": "observation.depth_linear.zed_link_camera_0", "observation.images.rgb.left_wrist": "observation.depth_linear.left_realsense_link_camera_0", "observation.images.rgb.right_wrist": "observation.depth_linear.right_realsense_link_camera_0", } def quat_wxyz_to_R(q): w, x, y, z = q / (np.linalg.norm(q) + 1e-12) return np.array([ [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], ]) # OmniGibson cameras use the OpenGL optical convention (-Z forward, +Y up); the DA3 pinhole # projection assumes OpenCV (+Z forward, +Y down). This diag(1,-1,-1) flips the camera Y/Z axes. # WITHOUT it, cross-view GT-depth reprojection is 0.00 (cameras point the wrong way); WITH it, 0.20+ # (best of all 8 conventions), and the head camera lands at its true +1.56 m height. (calib/calibrate_v3.py) _GL2CV = np.diag([1.0, -1.0, -1.0, 1.0]).astype(np.float32) def pose7_to_robot2cam(p7: np.ndarray) -> np.ndarray: """[pos3, quat_wxyz] camera-pose-in-robot-frame -> 4x4 robot->cam (OpenCV). Convention (validated in calib/): quat=wxyz, pose is the CAMERA-IN-ROBOT transform so robot->cam = inv(T), then GL->CV optical flip. """ T = np.eye(4, dtype=np.float32) T[:3, :3] = quat_wxyz_to_R(np.asarray(p7[3:], np.float64)) T[:3, 3] = p7[:3] return (_GL2CV @ np.linalg.inv(T)).astype(np.float32) class BehaviorV3DA3Dataset(BehaviorV3Dataset): """BehaviorV3Dataset + DA3 inputs (frames @ da3_hw, extrinsics, intrinsics, task language).""" def __init__(self, *args, da3_hw=(224, 224), lang_cache: str | None = None, lang_max_len: int = 32, **kwargs): super().__init__(*args, **kwargs) self._da3_hw = tuple(da3_hw) # ONE decode per (clip, frame), shared by the VLM path and the DA3 path. Tiny and # always on: only a few frames are ever in flight for a single sample (3 views), so 8 # entries is slack. Not the base class's _frame_mem_cache -- that one is off by default # (B1K_FRAME_MEM_CACHE_GB=0) and its key omits the resize, so it cannot safely serve two # resolutions from one entry. self._shared_frames: dict = {} self._shared_frames_max = 8 self._lang_max_len = int(lang_max_len) self._lang = None if lang_cache: with open(lang_cache, "rb") as f: self._lang = pickle.load(f) logger.info("DA3 lang cache: %d tasks from %s", len(self._lang), lang_cache) # GT depth grid cache (build with build_depth_cache.py). Checked BEFORE the video # path setup: with a complete cache the depth videos are not needed at all, so a task # can train from the cache alone without the ~7 GiB/chunk of depth mp4s on disk. self._depth_cache = None self._depth_mm: dict = {} if os.environ.get("B1K_USE_GT_DEPTH", "1") == "1": self._init_depth_cache() if os.environ.get("B1K_USE_GT_DEPTH", "1") == "1" and self._depth_cache is None: import pandas as pd depth_meta = pd.concat( [pd.read_parquet(f) for f in sorted(glob.glob( os.path.join(self.root, "meta", "episodes", "**", "*.parquet"), recursive=True ))], ignore_index=True, ).set_index("episode_index") for rec in self.episodes: row = depth_meta.loc[rec["episode_index"]] rec["depth_video"] = {} rec["depth_from_ts"] = {} for dst, src in DEPTH_SRCS.items(): rec["depth_video"][dst] = os.path.join( self.root, "videos", src, f"chunk-{int(row[f'videos/{src}/chunk_index']):03d}", f"file-{int(row[f'videos/{src}/file_index']):03d}.mp4", ) rec["depth_from_ts"][dst] = float(row[f"videos/{src}/from_timestamp"]) if not os.path.exists(rec["depth_video"][dst]): _ch = os.path.basename(os.path.dirname(rec["depth_video"][dst])) raise FileNotFoundError( f"GT depth requested but missing: {rec['depth_video'][dst]}\n" f" Depth is only downloaded for a subset of chunks. Fetch this one with:\n" f" hf download behavior-1k/2026-challenge-demos --repo-type dataset \\\n" f" --local-dir {self.root} --include " f"'videos/observation.depth_linear.*/{_ch}/*'\n" f" (~7 GiB per chunk.) Or set B1K_USE_GT_DEPTH=0 to fall back to the " f"model's predicted depth -- note plain DA3-GIANT has no metric head, " f"so that requires DA3_MODEL_NAME=depth-anything/DA3NESTED-GIANT-LARGE-1.1." ) logger.info("GT metric depth enabled for %d episodes", len(self.episodes)) def _init_depth_cache(self): """Index /chunk-*_g/ep-.npy. Requires EVERY episode of this run to be present -- a partial cache would silently mix cached and freshly-decoded depth.""" import glob as _glob root = os.environ.get("B1K_DEPTH_CACHE", "/work/jack/behavior1k/depth_cache") if not root or not os.path.isdir(root): return gh = self._da3_hw[0] // 14 idx = {} for d in sorted(_glob.glob(os.path.join(root, f"*_g{gh}"))): for f in _glob.glob(os.path.join(d, "ep-*.npy")): try: idx[int(os.path.basename(f)[3:-4])] = f except ValueError: continue if not idx: return want = [int(r["episode_index"]) for r in self.episodes] missing = [e for e in want if e not in idx] if missing: logger.warning("GT depth cache at %s covers %d/%d episodes (missing e.g. %s); " "falling back to depth-video decode", root, len(want) - len(missing), len(want), missing[:3]) return self._depth_cache = idx logger.info("GT depth CACHE: %d episodes at grid %dx%d from %s " "(no depth-video decode in the loader)", len(want), gh, gh, root) def _depth_grid(self, ep: int, t: int): """[V,gh,gw] float32 for episode `ep` frame `t`, from the memory-mapped cache.""" mm = self._depth_mm.get(ep) if mm is None: mm = np.load(self._depth_cache[ep], mmap_mode="r") self._depth_mm[ep] = mm if len(self._depth_mm) > 64: # bound open maps in long-lived workers self._depth_mm.pop(next(iter(self._depth_mm))) return np.asarray(mm[min(t, mm.shape[0] - 1)], dtype=np.float32) def _episode_pose_arrays(self, rec): """Compact per-episode pose arrays {view: [L,7]}, cached. Same fix as BehaviorV3Dataset._episode_arrays: the 8-file LRU below thrashed once training spanned 3800 episodes over 190 parquet files, costing a full parquet read per sample. One read now caches every episode in the file. Values identical to the old filter+sort+iloc path. """ import pandas as pd if not hasattr(self, "_ep_pose_cache"): from collections import OrderedDict self._ep_pose_cache = OrderedDict() self._ep_pose_cache_max = int(os.environ.get("B1K_EPISODE_CACHE", "8000")) key = (rec["data"], int(rec["episode_index"])) hit = self._ep_pose_cache.get(key) if hit is not None: self._ep_pose_cache.move_to_end(key) return hit df = pd.read_parquet(rec["data"], columns=["episode_index", "frame_index", *POSE_COLS.values()]) want = None for ep, grp in df.groupby("episode_index", sort=False): grp = grp.sort_values("frame_index") entry = {dst: np.stack(grp[col].to_numpy()).astype(np.float64) for dst, col in POSE_COLS.items()} k2 = (rec["data"], int(ep)) self._ep_pose_cache[k2] = entry self._ep_pose_cache.move_to_end(k2) if k2 == key: want = entry while len(self._ep_pose_cache) > self._ep_pose_cache_max: self._ep_pose_cache.popitem(last=False) return want if want is not None else self._ep_pose_cache[key] def _episode_poses(self, rec): """Cached per-episode pose table (the base reader's parquet cache omits pose columns).""" import pandas as pd if not hasattr(self, "_pose_cache"): from collections import OrderedDict self._pose_cache = OrderedDict() key = rec["data"] if key not in self._pose_cache: df = pd.read_parquet(key, columns=["episode_index", "frame_index", *POSE_COLS.values()]) self._pose_cache[key] = df if len(self._pose_cache) > 8: self._pose_cache.popitem(last=False) df = self._pose_cache[key] return df[df["episode_index"] == rec["episode_index"]].sort_values("frame_index") def _decode_rgb(self, path: str, ts: float) -> np.ndarray: """Base (VLM) frame, memoised so the DA3 pass can reuse it instead of decoding again. _AttachDA3Fields evaluates the transformed sample (which lands here) BEFORE calling da3_fields (which lands in _decode_da3), so the cache is warm by construction. The cached entry is a COPY: the transform stack may mutate its array in place (openpi's DeltaActions does exactly that to `actions`), and DA3 must not observe the mutation. One 150 KB memcpy against a saved video decode is a trivial trade. """ img = super()._decode_rgb(path, ts) key = (path, int(round(ts * self.fps))) if key not in self._shared_frames: self._shared_frames[key] = img.copy() while len(self._shared_frames) > self._shared_frames_max: self._shared_frames.pop(next(iter(self._shared_frames))) return img def _decode_da3(self, path: str, ts: float) -> np.ndarray: """DA3 frame at da3_hw, taken from the SAME decode the VLM used. With 224x224 source video and da3_hw=(224,224) this is an exact hit: no second decode and no resample at all -- the identical array feeds both the VLM and DA3. For larger source video it still decodes once and only resizes. """ h, w = self._da3_hw key = (path, int(round(ts * self.fps))) img = self._shared_frames.get(key) if img is None: # DA3 asked first (not the _AttachDA3Fields order); decode and seed the cache. img = super()._decode_rgb(path, ts) self._shared_frames[key] = img.copy() while len(self._shared_frames) > self._shared_frames_max: self._shared_frames.pop(next(iter(self._shared_frames))) if img.shape[0] == h and img.shape[1] == w: return img import cv2 interp = cv2.INTER_AREA if img.shape[0] >= h else cv2.INTER_LINEAR return cv2.resize(img, (w, h), interpolation=interp) def _decode_gt_depth(self, path: str, ts: float, grid_hw: tuple[int, int]) -> np.ndarray: """Decode BEHAVIOR-1K gray12le linear depth and align it to the DA3 patch grid. Native gray12le samples are metric millimeters (not gray16-scaled values). The depth and RGB videos share the same square camera raster/FOV; resizing native depth directly to the (H/14,W/14) token grid therefore aligns each depth cell with the corresponding DA3 patch. Area resampling averages metric Z-depth over the same image support represented by a token. """ import cv2 container = self._cached_container(path) vs = container.streams.video[0] container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True) frame = None for fr in container.decode(vs): if fr.time is not None and fr.time >= ts - 1e-3: frame = fr break if frame is None: container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True) for fr in container.decode(vs): frame = fr # Do not request gray16le: swscale left-shifts the native 12-bit samples by four bits. depth_m = frame.to_ndarray().astype(np.float32) / 1000.0 gh, gw = grid_hw depth_grid = cv2.resize(depth_m, (gw, gh), interpolation=cv2.INTER_AREA) return depth_grid[None] # [1,gh,gw] def _lang_entry(self, task_name: str): if self._lang is None: L = self._lang_max_len return np.zeros((L, 1024), np.float32), np.zeros((L,), bool) feat, mask = self._lang[task_name] return np.asarray(feat, np.float32), np.asarray(mask, bool) def da3_fields(self, i): """Compute ONLY the DA3 input fields for sample i (attached AFTER the transform stack, which constructs fresh dicts and would drop unknown keys).""" item = {} ei, t = self.samples[i] rec = self.episodes[ei] poses = self._episode_pose_arrays(rec) # cached; no per-sample parquet h, w = self._da3_hw imgs, extr, intr, gt_depth = [], [], [], [] use_gt_depth = os.environ.get("B1K_USE_GT_DEPTH", "1") == "1" grid_hw = (h // 14, w // 14) # Intrinsics AT da3_hw. Independent of the SOURCE video resolution, because # fx = (FOCAL_RATIO * native_w) * (w / native_w) = FOCAL_RATIO * w # cx = (native_w / 2) * (w / native_w) = w / 2 # native_w cancels exactly. So feeding 224 source instead of 720 does not move the # geometry by one bit, and the DA3 frame no longer has to report its native width -- # which is what lets it come from the shared VLM decode. K = np.array([ [FOCAL_RATIO * w, 0.0, w / 2.0], [0.0, FOCAL_RATIO * h, h / 2.0], [0.0, 0.0, 1.0], ], np.float32) for dst in VIEW_ORDER: frame_ts = rec["from_ts"][dst] + t / self.fps imgs.append(self._decode_da3(rec["video"][dst], frame_ts)) extr.append(pose7_to_robot2cam(poses[dst][t])) intr.append(K) if use_gt_depth and self._depth_cache is None: depth_ts = rec["depth_from_ts"][dst] + t / self.fps gt_depth.append(self._decode_gt_depth( rec["depth_video"][dst], depth_ts, grid_hw )) item["da3_images"] = np.stack(imgs, 0) # [V,H,W,3] uint8 at da3_hw (224 by default) item["camera_extrinsics"] = np.stack(extr, 0) # [V,4,4] robot->cam OpenCV item["camera_intrinsics"] = np.stack(intr, 0) # [V,3,3] @ da3_hw if use_gt_depth: if self._depth_cache is not None: # [V,gh,gw] -> [V,1,gh,gw]; one array slice instead of three video decodes item["gt_metric_depth"] = self._depth_grid( int(rec["episode_index"]), t)[:, None, :, :] else: item["gt_metric_depth"] = np.stack(gt_depth, 0).astype(np.float32) # [V,1,gh,gw] lf, lm = self._lang_entry(rec["task0"]) item["lang_feat"] = lf item["lang_mask"] = lm return item class _AttachDA3Fields: """Wraps the TRANSFORMED dataset; merges the raw dataset's DA3 fields into each sample.""" def __init__(self, transformed, raw: BehaviorV3DA3Dataset): self._transformed = transformed self._raw = raw def __len__(self): return len(self._transformed) def __getitem__(self, i): out = dict(self._transformed[i]) out.update(self._raw.da3_fields(i)) return out def create_v3_behavior_da3_loader(config, root_2026, activities, task_data_json, *, lang_cache, sharding=None, shuffle=True, num_workers=None, seed=0, da3_hw=(224, 224)): """v3 loader with DA3 inputs + a per-batch frozen DA3-GIANT extraction hook (GPU).""" import jax import dataclasses as _dc from b1k.policies import b1k_policy from b1k.training.data_loader import transform_dataset, DataLoaderImpl from openpi.training.data_loader import TorchDataLoader from b1k.training import da3_extractor as _ex data_config = config.data.create(config.assets_dirs, config.model) new_inputs = tuple( B1kInputs2026(model_type=config.model.model_type) if isinstance(x, b1k_policy.B1kInputs) else x for x in data_config.data_transforms.inputs ) data_config = _dc.replace( data_config, data_transforms=_dc.replace(data_config.data_transforms, inputs=new_inputs)) # VGGT-Omega extractor: decode/patchify at process_res (patch 16), not the DA3 252 grid. _use_vggt = os.environ.get("USE_VGGT") == "1" if _use_vggt: da3_hw = (int(os.environ.get("VGGT_PROCESS_RES", "256")),) * 2 # The model declares da3_features as [b, layers, views, C, gh, gw] from config.da3.grid_hw, # while the extractor's grid is da3_hw/patch. A mismatch is a shape error thousands of steps # into a run (or, worse, a silent reinterpretation); check it here where the message is clear. _patch = 16 if _use_vggt else 14 _want = (da3_hw[0] // _patch, da3_hw[1] // _patch) _have = tuple(config.model.da3.grid_hw) if _want != _have: raise ValueError( f"grid mismatch: da3_hw={da3_hw} with patch {_patch} gives grid {_want}, but " f"config.model.da3.grid_hw={_have}. Set DA3_GRID_H/DA3_GRID_W (or grid_hw) to {_want}." ) logger.info("DA3 input %s, patch %d -> grid %s (%d patches)", da3_hw, _patch, _want, _want[0] * _want[1]) ds = BehaviorV3DA3Dataset( root_2026, activities=activities, action_horizon=config.model.action_horizon, task_data_json=task_data_json, seed=seed, da3_hw=da3_hw, lang_cache=lang_cache, lang_max_len=config.model.da3.lang_max_len, ) tds = transform_dataset(ds, data_config) tds = _AttachDA3Fields(tds, ds) logger.info("Building inline DA3-GIANT extractor (da3_hw=%s) ...", da3_hw) # Extraction devices: default single-GPU (cuda:0). Set B1K_EXTRACT_DEVICES to spread the frozen # DA3-GIANT forward across GPUs (one replica per device, batch split, run concurrently) so the # ~2.7s single-GPU extraction shrinks and better overlaps the JAX train step. _dev_env = os.environ.get("B1K_EXTRACT_DEVICES", "").strip() _devices = [d.strip() for d in _dev_env.split(",") if d.strip()] or None _fchunk = int(os.environ.get("B1K_DA3_FWD_CHUNK", "16")) logger.info("DA3 extractor: devices=%s forward_chunk=%d", _devices or ["cuda:0"], _fchunk) if _use_vggt: from b1k.training import vggt_extractor as _vex extractor = _vex.VGGTInlineExtractor(process_res=da3_hw[0], forward_chunk=_fchunk, devices=_devices) logger.info("Using VGGT-Omega extractor (process_res=%d, grid=%d)", da3_hw[0], da3_hw[0] // 16) else: # Plain GIANT, not the NESTED variant: nested bolts a 334 M monocular metric model # (19.8% of params) onto the 1355.7 M any-view GIANT purely to scale depth, and we # discard that depth in favour of GT. The features the bank builder consumes come from # the GIANT either way. Set DA3_MODEL_NAME to go back to nested. model_name = os.environ.get( "DA3_MODEL_NAME", "depth-anything/DA3-GIANT-1.1" ) out_layers = tuple(int(x) for x in os.environ.get( "DA3_OUT_LAYERS", "19,26,33,39" ).split(",")) extractor = _ex.DA3InlineExtractor( model_name=model_name, out_layers=out_layers, da3_hw=da3_hw, forward_chunk=_fchunk, devices=_devices, ) logger.info("Using DA3 extractor model=%s out_layers=%s", model_name, out_layers) # DLPack GPU->GPU handoff: skip the ~2.6s/batch host round-trip by moving extractor features # straight from the extraction GPUs to the training GPUs over NVLink. Requires CUDA extraction # devices whose count matches the training mesh size (contiguous batch split aligns 1:1). # Holds the last few batches' torch source shards alive so the async NVLink copies (device_put) # can never read freed memory — replaces a blocking block_until_ready that serialized the producer. import collections as _collections _keepalive = _collections.deque(maxlen=4) def _dlpack_ok(): try: m = getattr(sharding, "mesh", None) return (os.environ.get("B1K_DLPACK") == "1" and m is not None and len(list(m.devices.flat)) == len(extractor.devices) and all(str(d).startswith("cuda") for d in extractor.devices)) except Exception: return False # Output field order MUST match the extractor's extract()/extract_shards_torch() tuple order. _field_names = (["da3_features", "da3_ray", "da3_depth", "da3_depth_conf", "da3_pose_enc", "da3_cam_tokens"] if _use_vggt else ["da3_features", "da3_ray", "da3_depth"]) def batch_transform(batch): use_gt_depth = os.environ.get("B1K_USE_GT_DEPTH", "1") == "1" if _dlpack_ok(): import torch import jax parts = extractor.extract_shards_torch( batch["da3_images"], batch["camera_extrinsics"], batch["camera_intrinsics"]) tdevs = list(sharding.mesh.devices.flat) # training devices, batch-chunk k -> tdevs[k] if use_gt_depth: gt = batch.pop("gt_metric_depth") bounds = [round(i * int(gt.shape[0]) / len(parts)) for i in range(len(parts) + 1)] replaced = [] for k, part in enumerate(parts): d = gt[bounds[k]:bounds[k + 1]] if not isinstance(d, torch.Tensor): d = torch.as_tensor(d) d = d.to(extractor.devices[k], dtype=torch.float32, non_blocking=True) replaced.append((part[0], part[1], d)) parts = replaced def _asm(fi): # assemble per-shard torch tensors (field fi) into one sharded jax array js = [jax.device_put(jax.dlpack.from_dlpack(parts[k][fi]), tdevs[k]) for k in range(len(parts))] gshape = (sum(int(s.shape[0]) for s in js),) + tuple(int(d) for d in js[0].shape[1:]) return jax.make_array_from_single_device_arrays(gshape, sharding, js) for fi, nm in enumerate(_field_names): batch[nm] = _asm(fi) # Do NOT block here: the device_put queues behind the in-flight train step on the target # GPUs, so blocking would serialize the producer with training (killing the overlap). # Instead keep the torch source shards referenced for a few batches so the async NVLink # copy can't read freed memory. _keepalive.append(parts) else: outs = extractor.extract( batch["da3_images"], batch["camera_extrinsics"], batch["camera_intrinsics"]) if use_gt_depth: gt = batch.pop("gt_metric_depth") if hasattr(gt, "detach"): gt = gt.detach().cpu().numpy() outs = (outs[0], outs[1], np.asarray(gt, dtype=np.float32)) for nm, arr in zip(_field_names, outs): batch[nm] = arr # da3_features = uint16 bf16-bits; rest fp32 batch.pop("da3_images", None) batch.pop("camera_intrinsics", None) return batch loader = TorchDataLoader( tds, local_batch_size=config.batch_size // jax.process_count(), sharding=sharding, shuffle=shuffle, num_workers=config.num_workers if num_workers is None else num_workers, seed=seed, batch_transform=batch_transform, ) return DataLoaderImpl(data_config, loader)