task5_setting_mousetraps_da3/code/b1k_da3.py

"""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 <cache>/chunk-*_g<G>/ep-<n>.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)