diff --git a/egomimic/eval/latent_dataset.py b/egomimic/eval/latent_dataset.py index dcd078fc2a..22ad0393b5 100644 --- a/egomimic/eval/latent_dataset.py +++ b/egomimic/eval/latent_dataset.py @@ -43,7 +43,7 @@ import logging from egomimic.rldb.filters import DatasetFilter -from egomimic.rldb.zarr.zarr_dataset_multi import EvenStrideDataset, MultiDataset +from egomimic.rldb.multi import EvenStrideDataset, MultiDataset logger = logging.getLogger(__name__) diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index dbd5f47d17..03db6396b6 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -5,7 +5,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/aria_keypoints.yaml b/egomimic/hydra_configs/data/aria_keypoints.yaml index 7ed22d6eb7..9b5ea50824 100644 --- a/egomimic/hydra_configs/data/aria_keypoints.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints.yaml @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml index 1a4edb25d3..00855cb0ed 100644 --- a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/aria_pilot_lance.yaml b/egomimic/hydra_configs/data/aria_pilot_lance.yaml index 126932df3b..e9a1e6b651 100644 --- a/egomimic/hydra_configs/data/aria_pilot_lance.yaml +++ b/egomimic/hydra_configs/data/aria_pilot_lance.yaml @@ -5,7 +5,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.lance.dataset.LanceEpisodeResolver lance_root: ${oc.env:EGOVERSE_LANCE_PILOT,/coc/flash7/scratch/lance_pilot}/lance diff --git a/egomimic/hydra_configs/data/aria_pilot_zarr.yaml b/egomimic/hydra_configs/data/aria_pilot_zarr.yaml index c60cabae3a..c679c8f872 100644 --- a/egomimic/hydra_configs/data/aria_pilot_zarr.yaml +++ b/egomimic/hydra_configs/data/aria_pilot_zarr.yaml @@ -4,7 +4,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver folder_path: ${oc.env:EGOVERSE_LANCE_PILOT,/coc/flash7/scratch/lance_pilot}/zarr diff --git a/egomimic/hydra_configs/data/aria_qwen.yaml b/egomimic/hydra_configs/data/aria_qwen.yaml index cd976a643d..2618a5d1db 100644 --- a/egomimic/hydra_configs/data/aria_qwen.yaml +++ b/egomimic/hydra_configs/data/aria_qwen.yaml @@ -6,7 +6,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index 7fb6122886..08429526df 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -18,7 +18,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} @@ -33,7 +33,7 @@ train_datasets: mode: train valid_ratio: 0.05 human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml index c56f51c4e9..6cbf6bebcd 100644 --- a/egomimic/hydra_configs/data/eva.yaml +++ b/egomimic/hydra_configs/data/eva.yaml @@ -2,7 +2,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/human.yaml b/egomimic/hydra_configs/data/human.yaml index 9c4a0b27eb..9b045ba01b 100644 --- a/egomimic/hydra_configs/data/human.yaml +++ b/egomimic/hydra_configs/data/human.yaml @@ -14,7 +14,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper # ('aria_bimanual'). Until that migration lands, this dataset matches zero rows. train_datasets: human_bimanual: &hb - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index b904d4dde0..c0c67f6bde 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -11,7 +11,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/mecka_all_6d.yaml b/egomimic/hydra_configs/data/mecka_all_6d.yaml index 94471a1a3f..8ad3ec105c 100644 --- a/egomimic/hydra_configs/data/mecka_all_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_all_6d.yaml @@ -20,7 +20,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver # Per-cluster zarr mirror (paths/default.yaml; $EGOVERSE_DATASET_DIR diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml index 3fab676fe4..efe9b25638 100644 --- a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -68,7 +68,7 @@ valid_prefix: seen_op_valid train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml index 90fd9c6cac..fb5a6f304a 100644 --- a/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml @@ -36,7 +36,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml index 58f01426ac..1b6b53df8b 100644 --- a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml @@ -84,7 +84,7 @@ held_out_operators: train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 5b97388514..b5de74acb8 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -7,7 +7,7 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + _target_: egomimic.rldb.multi.MultiDataset._from_resolver resolver: _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver folder_path: ${paths.dataset_dir} diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 67caf2c401..e16720972c 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -11,7 +11,7 @@ import egomimic.utils.tensor_utils as TensorUtils from egomimic.pl_utils.pl_data_utils import head_of_loader, video_loader_name -from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.rldb.multi import MultiDataset def _barrier() -> None: diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 369ce445a3..b006986877 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -6,7 +6,7 @@ import numpy as np import torch -from egomimic.rldb.zarr.action_chunk_transforms import Transform +from egomimic.rldb.transforms import Transform from egomimic.utils.type_utils import _to_numpy from egomimic.utils.viz_utils import ( _viz_annotations, diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index b3edfa4f48..21cde701e1 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -10,7 +10,7 @@ _strip_pi_keymap_mode, ) from egomimic.rldb.embodiment.human import ARIA_INTRINSICS -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, CartesianRot6DToYPR, diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index 010942a961..e5ebe3d121 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -11,7 +11,7 @@ _reject_legacy_rotation, _strip_pi_keymap_mode, ) -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, CartesianRot6DToYPR, diff --git a/egomimic/rldb/filters.py b/egomimic/rldb/filters.py index a8dc62260a..66af480330 100644 --- a/egomimic/rldb/filters.py +++ b/egomimic/rldb/filters.py @@ -145,3 +145,19 @@ def matches(self, row: Mapping[str, Any]) -> bool: if row.get("episode_hash") not in self.completed_episode_hashes: return False return super().matches(row) + + +class PinError(ValueError): + """A pinned episode_hash cannot be used: missing, deleted, no processed + path, or its embodiment does not match the dataset it was pinned under.""" + + +def _ensure_dataset_filter(filters: DatasetFilter | None) -> DatasetFilter: + if filters is None: + return DatasetFilter() + if isinstance(filters, DatasetFilter): + return filters + raise TypeError( + "filters must be a DatasetFilter or None; plain dict filters are no longer " + "supported." + ) diff --git a/egomimic/rldb/lance/annotations.py b/egomimic/rldb/lance/annotations.py new file mode 100644 index 0000000000..12c1ed0c99 --- /dev/null +++ b/egomimic/rldb/lance/annotations.py @@ -0,0 +1,143 @@ +"""Review of annotation change files before `commit.apply_annotations` commits them. + +A change file is a Lance or parquet table with an `op` column ("add" or "delete") and the +ANNOTATIONS fields. An add row may leave `annotation_id` null, since it is computed from +the row's content; a delete row needs only `annotation_id`. +""" + +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import lance +import pyarrow as pa +import pyarrow.parquet as pq + +from egomimic.rldb.lance.commit import _exists, annotations_uri +from egomimic.rldb.lance.schema import ANNOTATION_TYPES, ANNOTATIONS, annotation_id + +CONTENT = ("recording_id", "annotation_vendor", "annotation_type", "t_start_ns", "t_end_ns", "text") # fmt: skip + + +@dataclass +class Review: + adds: pa.Table + deletes: list[str] + errors: list[str] = field(default_factory=list) + unchanged: list[str] = field(default_factory=list) + + +def read_changes(path: str | Path) -> pa.Table: + path = str(path).rstrip("/") + return ( + lance.dataset(path).to_table() + if path.endswith(".lance") + else pq.read_table(path) + ) + + +def committed(root: Path) -> pa.Table: + uri = annotations_uri(root) + if not _exists(uri): + return ANNOTATIONS.empty_table() + return lance.dataset(str(uri)).to_table() + + +def _task_bounds(root: Path) -> dict[str, list[tuple[int, int]]]: + cols = ["recording_id", "t_start_ns", "t_end_ns"] + out = defaultdict(list) + for t in ( + lance.dataset(str(Path(root) / "tasks.lance")) + .to_table(columns=cols) + .to_pylist() + ): + out[t["recording_id"]].append((t["t_start_ns"], t["t_end_ns"])) + return out + + +def _add_errors(r: dict, tasks: dict) -> list[str]: + errs = [f"{k} is missing" for k in CONTENT if r.get(k) in (None, "")] + if errs: + return errs + if r["annotation_type"] not in ANNOTATION_TYPES: + errs.append( + f"annotation_type {r['annotation_type']!r} not in {ANNOTATION_TYPES}" + ) + t0, t1 = r["t_start_ns"], r["t_end_ns"] + if t0 >= t1: + errs.append(f"t_start_ns {t0} is not before t_end_ns {t1}") + if r["recording_id"] not in tasks: + errs.append(f"recording {r['recording_id']} has no tasks") + elif not any(s <= t0 and t1 <= e for s, e in tasks[r["recording_id"]]): + errs.append(f"[{t0}, {t1}) is not inside any task of its recording") + q = r.get("quality_score") + if q is not None and not 0 <= q <= 5: + errs.append(f"quality_score {q} is outside 0-5") + return errs + + +def review(root: Path, changes: pa.Table) -> Review: + """Validate `changes` against the committed tables; nothing is written.""" + tasks = _task_bounds(root) + have = set(committed(root).column("annotation_id").to_pylist()) + adds, deletes, errors, unchanged, seen = [], [], [], [], set() + for i, r in enumerate(changes.to_pylist()): + op, given = r.get("op"), r.get("annotation_id") + if op == "delete": + if not given: + errors.append(f"row {i}: delete without annotation_id") + elif given not in have: + errors.append(f"row {i}: delete of {given}, which is not committed") + elif given in seen: + errors.append(f"row {i}: {given} appears twice") + else: + deletes.append(given) + seen.add(given) + continue + if op != "add": + errors.append(f"row {i}: op {op!r} is not add or delete") + continue + errs = _add_errors(r, tasks) + if errs: + errors += [f"row {i}: {e}" for e in errs] + continue + aid = annotation_id(*(r[k] for k in CONTENT)) + if given and given != aid: + errors.append( + f"row {i}: annotation_id {given} does not match its content ({aid})" + ) + elif aid in seen: + errors.append(f"row {i}: {aid} appears twice") + elif aid in have: + unchanged.append(aid) + else: + adds.append( + {**{k: r.get(k) for k in ANNOTATIONS.names}, "annotation_id": aid} + ) + seen.add(aid) + return Review( + pa.Table.from_pylist(adds, schema=ANNOTATIONS), deletes, errors, unchanged + ) + + +def diff(root: Path, rv: Review) -> str: + """Per recording: rows added, deleted and already committed, with their texts.""" + old = {r["annotation_id"]: r for r in committed(root).to_pylist()} + per = defaultdict(lambda: {"+": [], "-": [], "=": 0}) + for r in rv.adds.to_pylist(): + per[r["recording_id"]]["+"].append(r) + for aid in rv.deletes: + per[old[aid]["recording_id"]]["-"].append(old[aid]) + for aid in rv.unchanged: + per[old[aid]["recording_id"]]["="] += 1 + lines = [] + for rid in sorted(per): + p = per[rid] + lines.append(f"{rid}: +{len(p['+'])} -{len(p['-'])} ={p['=']}") + for sign in "-+": + for r in sorted(p[sign], key=lambda r: r["t_start_ns"]): + lines.append( + f" {sign} [{r['t_start_ns']}, {r['t_end_ns']}) " + f"{r['annotation_vendor']}/{r['annotation_type']}: {r['text']}" + ) + return "\n".join(lines) diff --git a/egomimic/rldb/lance/commit.py b/egomimic/rldb/lance/commit.py index 773655b8ef..9931514d82 100644 --- a/egomimic/rldb/lance/commit.py +++ b/egomimic/rldb/lance/commit.py @@ -16,12 +16,14 @@ import lance import pyarrow as pa +import pyarrow.compute as pc from lance import LanceOperation from lance.fragment import write_fragments from egomimic.rldb.embodiment.embodiment import END_EFFECTOR from egomimic.rldb.lance.schema import ( ALIASES, + ANNOTATIONS, RECORDINGS, STEPS_KEY, TASKS, @@ -88,6 +90,16 @@ def steps_uris(root: Path) -> list[Path]: return [Path(root) / p for p in steps_paths(recordings)] +ANNOTATIONS_TABLE = "annotations.lance" +# annotation_id hashes an annotation's content, so `edit` may change only its grade. +ANNOTATION_GRADING = ("quality_score", "graded_by", "rubric_version") +EDITABLE = (*CATALOG, ANNOTATIONS_TABLE) + + +def annotations_uri(root: Path) -> Path: + return Path(root) / ANNOTATIONS_TABLE + + def chunk_uris(root: Path) -> list[Path]: """The `video_chunks/*.lance` tables (chunk_videos.py).""" return sorted( @@ -96,10 +108,11 @@ def chunk_uris(root: Path) -> list[Path]: def table_uris(root: Path) -> list[Path]: - """Every table: steps, catalog in write order, then video chunks.""" + """Every table: steps, catalog in write order, annotations, then video chunks.""" return [ *steps_uris(root), *(Path(root) / n for n in CATALOG), + annotations_uri(root), *chunk_uris(root), ] @@ -131,7 +144,9 @@ def _append(uri: Path, table: pa.Table) -> None: lance.write_dataset(table, uri, mode="append" if _exists(uri) else "create") -def _replace(uri: Path, predicate: str, table: pa.Table) -> None: +def _replace( + uri: Path, predicate: str | None, table: pa.Table, message: str | None = None +) -> None: """Delete rows matching `predicate` and append `table` in one commit. merge_insert would also be one commit, but it writes matched rows before inserted @@ -139,7 +154,7 @@ def _replace(uri: Path, predicate: str, table: pa.Table) -> None: """ ds = lance.dataset(uri) removed, updated = [], [] - for frag in ds.get_fragments(): + for frag in ds.get_fragments() if predicate else []: if frag.count_rows(predicate) == 0: continue meta = frag.delete(predicate) @@ -147,11 +162,11 @@ def _replace(uri: Path, predicate: str, table: pa.Table) -> None: removed.append(frag.fragment_id) else: updated.append(meta) - new = write_fragments(table, uri, schema=ds.schema) + new = write_fragments(table, uri, schema=ds.schema) if table.num_rows else [] op = LanceOperation.Update( removed_fragment_ids=removed, updated_fragments=updated, new_fragments=new ) - lance.LanceDataset.commit(uri, op, read_version=ds.version) + lance.LanceDataset.commit(uri, op, read_version=ds.version, commit_message=message) def _append_new(uri: Path, predicate: str, table: pa.Table) -> None: @@ -219,9 +234,45 @@ def append(root: Path, deliveries: list[Delivery]) -> list[str]: _append_new( Path(root) / name, pred, pa.concat_tables([s[name] for s in staged]) ) + _ensure_annotations(root) return [d.recording_id for d in todo] +def _ensure_annotations(root: Path) -> None: + # Exists from the first commit on, so a release tag always has an annotations version to pin. + if not _exists(annotations_uri(root)): + _append(annotations_uri(root), ANNOTATIONS.empty_table()) + + +def apply_annotations( + root: Path, adds: pa.Table, deletes: list[str], message: str | None = None +) -> dict[str, int]: + """Add rows and delete rows by `annotation_id`, in one commit. Adds whose id is already + committed are skipped, so resubmitting a file is a no-op. `message` (e.g. who + approved it) is stored on the commit: `read_transaction(v).transaction_properties`.""" + _ensure_annotations(root) + uri = annotations_uri(root) + have = lance.dataset(uri).to_table(columns=["annotation_id"]).column(0) + adds = adds.select(ANNOTATIONS.names).cast(ANNOTATIONS) + adds = adds.filter( + pc.invert(pc.is_in(adds["annotation_id"], have.combine_chunks())) + ) + if not adds.num_rows and not deletes: + return {"added": 0, "deleted": 0} + pred = ( + f"annotation_id IN ({', '.join(map(_sql_str, deletes))})" if deletes else None + ) + _replace(uri, pred, adds, message) + return {"added": adds.num_rows, "deleted": len(deletes)} + + +def drop_task_annotations(root: Path) -> None: + """One-off for tables written while annotations were nested in tasks.""" + ds = lance.dataset(Path(root) / "tasks.lance") + if "annotations" in ds.schema.names: + ds.drop_columns(["annotations"]) + + def replace(root: Path, d: Delivery) -> None: """Swap a committed recording for a reprocessed delivery; its steps rows move to the end of the table.""" check_required(d) @@ -245,10 +296,21 @@ def replace(root: Path, d: Delivery) -> None: def edit(root: Path, table: str, rows: pa.Table, on: str) -> None: - """Update columns of existing rows keyed by `on` (label and metadata fixes); every key must exist.""" - if table not in CATALOG: + """Update columns of existing rows keyed by `on` (label and metadata fixes); every key must exist. + On annotations.lance only the grading columns can change; content goes through + `apply_annotations`.""" + if table not in EDITABLE: # merge_insert rewrites matched rows at the end of the table, which would split a recording's steps. - raise ValueError(f"edit only applies to {list(CATALOG)}; use replace for steps") + raise ValueError( + f"edit only applies to {list(EDITABLE)}; use replace for steps" + ) + if table == ANNOTATIONS_TABLE: + other = set(rows.column_names) - {"annotation_id", *ANNOTATION_GRADING} + if on != "annotation_id" or other: + raise ValueError( + f"edit on {table} is keyed by annotation_id and changes only " + f"{list(ANNOTATION_GRADING)}; got on={on!r}, other columns {sorted(other)}" + ) ds = lance.dataset(Path(root) / table) moved = [c for c in STEPS_KEY if c in rows.column_names] if table == "recordings.lance" and moved: @@ -272,6 +334,7 @@ def add_columns(root: Path, table: str, columns: dict[str, str]) -> None: def tag(root: Path, name: str) -> dict[str, int]: """Tag the current version of every table with `name`; returns {table: version}.""" + _ensure_annotations(root) out = {} for uri in table_uris(root): ds = lance.dataset(uri) @@ -288,7 +351,7 @@ def index(root: Path) -> None: def compact(root: Path) -> None: - for uri in table_uris(root): + for uri in filter(_exists, table_uris(root)): ds = lance.dataset(uri) ds.optimize.compact_files() if lance.dataset(uri).describe_indices(): @@ -303,7 +366,7 @@ def cleanup( `older_than` must exceed the longest training run that resolved an untagged version. """ - for uri in table_uris(root): + for uri in filter(_exists, table_uris(root)): lance.dataset(uri).cleanup_old_versions( older_than=older_than, error_if_tagged_old_versions=False ) diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index 3302f250c0..41ff470109 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/egomimic/rldb/lance/dataset.py @@ -14,6 +14,7 @@ import os from collections import OrderedDict from pathlib import Path +from typing import NamedTuple import lance import numpy as np @@ -22,13 +23,9 @@ import pyarrow.compute as pc import torch -from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.filters import DatasetFilter, PinError, _ensure_dataset_filter from egomimic.rldb.lance.schema import STEPS_KEY, steps_path -from egomimic.rldb.zarr.zarr_dataset_multi import ( - PinError, - ZarrDataset, - _ensure_dataset_filter, -) +from egomimic.rldb.leaf import EpisodeLeaf logger = logging.getLogger(__name__) @@ -257,7 +254,7 @@ def _chunks_of(table: dict, recording_id: str, num_frames: int, name: str) -> tu return (table["uri"], table["version"], table["rows"][recording_id], table["k"]) -class LanceRecordingDataset(ZarrDataset): +class LanceRecordingDataset(EpisodeLeaf): """One recording from local Lance tables: steps rows plus its MP4s on disk, read one sample at a time. `LanceStreamingDataset` is the object-storage variant.""" @@ -280,6 +277,8 @@ def __init__( key_map: dict, transform_list: list | None = None, content_digest: str | None = None, + annotations: tuple[np.ndarray, np.ndarray, list[str]] | None = None, + clamp_actions: bool = False, ): self.episode_path = name self.content_digest = content_digest @@ -296,14 +295,13 @@ def __init__( self._intrinsics = intrinsics self.videos = videos self.task_text = task_text - self.key_map = key_map - self.transform = transform_list self._image_keys = {k for k in videos} - self._json_keys = set() - self._annotations = [] self.anchors = np.asarray(anchors, dtype=np.int64).reshape(-1, 2) self._cum = np.r_[0, np.cumsum(self.anchors[:, 1] - self.anchors[:, 0])] - torch.utils.data.Dataset.__init__(self) + # (start frames, end frames, texts) sorted by start; None: the task text stands in. + self._ann = annotations + self._clamp_actions = clamp_actions + super().__init__(key_map=key_map, transform_list=transform_list) def cache_fingerprint(self) -> str: """Keys on the recording's content digest, so commits to other recordings and @@ -321,6 +319,8 @@ def cache_fingerprint(self) -> str: for k in sorted(self.videos): h.update(f"{k}\0{self.videos[k]}\0".encode()) h.update(self.anchors.tobytes()) + if self._clamp_actions and self._ann is not None: + h.update(b"clamp\0" + self._ann[0].tobytes() + self._ann[1].tobytes()) return h.hexdigest() @property @@ -338,8 +338,21 @@ def anchor(self, i: int) -> int: j = int(np.searchsorted(self._cum, i, side="right")) - 1 return int(self.anchors[j, 0] + (i - self._cum[j])) + def _covering(self, frame_idx: int) -> np.ndarray: + s, e, _ = self._ann + return np.flatnonzero((s <= frame_idx) & (frame_idx < e)) + def _annotation_text_for_frame(self, frame_idx: int): - return [self.task_text] + if self._ann is None: + return [self.task_text] + return [self._ann[2][j] for j in self._covering(frame_idx)] + + def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: + end = super()._chunk_end_idx(start_idx, horizon, key_type) + if not self._clamp_actions or key_type != "action_keys" or self._ann is None: + return end + cover = self._covering(start_idx) + return min(end, int(self._ann[1][cover].min())) if len(cover) else end def _read_numeric(self, needs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: """One `take` for every numeric column; returns {zarr key: rows at the requested frames}.""" @@ -380,7 +393,7 @@ def assemble_sample(self, plan: dict, fetched: tuple[dict, dict]): data[k] = self._pad_to_horizon( rows[0] if single else rows, spec.get("horizon") ) - return self._finalize_sample(idx, data) + return self._finish(idx, data) def __getitem__(self, i: int) -> dict[str, torch.Tensor]: plan = self.plan_sample(i) @@ -398,6 +411,8 @@ class LanceStreamingDataset(LanceRecordingDataset): instead of whole MP4s, and `MultiDataset.__getitems__` fetches a DataLoader batch through `fetch_planned`, so a batch costs about two round trips.""" + batched_fetch = True + def __init__( self, *, @@ -432,6 +447,16 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor]: return self.assemble_sample(plan, self.fetch_planned([(self, plan)])[0]) +class StepsRows(NamedTuple): + """Where a recording's frames sit in its steps table.""" + + uri: str + version: int + row_start: int + num_frames: int + valid: dict[str, np.ndarray] + + class LanceEpisodeResolver: """Resolves recordings from local Lance tables into LanceRecordingDataset leaves. @@ -442,8 +467,22 @@ class LanceEpisodeResolver: Each table is read at `versions[name]` (a number or tag), else at `tag`, else at its latest version. Tags are resolved to numbers once, here, so a tag moved later does not change what an existing resolver reads. + + `annotations` switches annotation keys from the task description to the rows of + `annotations.lance` covering the frame, in start order. Its keys (all optional): + `types` (default ["language"]), `vendors`, `min_quality` (ungraded rows fail it), + `clamp_actions` (stop action chunks at the end of the covering annotation) and + `require` (sample only anchors some annotation covers; drops leaves with none). """ + ANNOTATION_DEFAULTS = { + "types": ["language"], + "vendors": None, + "min_quality": None, + "clamp_actions": False, + "require": False, + } + def __init__( self, lance_root: str | Path, @@ -452,6 +491,7 @@ def __init__( cameras: dict[str, str] | None = None, versions: dict[str, int | str] | None = None, tag: str | None = None, + annotations: dict | None = None, ): self.root = str(lance_root).rstrip("/") self.key_map = key_map @@ -459,6 +499,12 @@ def __init__( self.cameras = cameras or {"images.front_1": "front_1"} self.tag = tag self.versions = dict(versions or {}) + self.annotations = None + if annotations is not None: + unknown = set(annotations) - set(self.ANNOTATION_DEFAULTS) + if unknown: + raise ValueError(f"unknown annotations options {sorted(unknown)}") + self.annotations = {**self.ANNOTATION_DEFAULTS, **dict(annotations)} def _open(self, name: str): uri = f"{self.root}/{name}" @@ -480,6 +526,39 @@ def catalog(self) -> pd.DataFrame: df = r.merge(t, on="recording_id").merge(a, on="recording_id") return df.rename(columns={"task_name": "task"}).reset_index(drop=True) + def _annotation_rows(self) -> dict[str, pd.DataFrame]: + """{recording_id: rows passing the `annotations` filters, sorted by start}.""" + cfg = self.annotations + _, _, tbl = self._open("annotations.lance") + df = tbl.to_table().to_pandas() + keep = df.annotation_type.isin(cfg["types"]) + if cfg["vendors"] is not None: + keep &= df.annotation_vendor.isin(cfg["vendors"]) + if cfg["min_quality"] is not None: + keep &= df.quality_score.notna() & (df.quality_score >= cfg["min_quality"]) + df = df[keep].sort_values(["recording_id", "t_start_ns", "t_end_ns"]) + return {rid: g for rid, g in df.groupby("recording_id")} + + @staticmethod + def _frame_spans(rows: pd.DataFrame, ts: np.ndarray): + """Half-open frame ranges of half-open time ranges; empty ones are dropped.""" + a = np.searchsorted(ts, rows.t_start_ns.to_numpy(), "left") + b = np.searchsorted(ts, rows.t_end_ns.to_numpy(), "left") + keep = a < b + return a[keep], b[keep], rows.text[keep].tolist() + + @staticmethod + def _covered(anchors: np.ndarray, starts, ends, n: int) -> np.ndarray: + on = np.zeros(n + 1, np.int64) + np.add.at(on, starts, 1) + np.add.at(on, ends, -1) + cover = np.cumsum(on)[:n] > 0 + mask = np.zeros(n, bool) + for s, e in anchors: + mask[s:e] = True + edges = np.flatnonzero(np.diff(np.r_[0, (mask & cover).astype(np.int8), 0])) + return edges.reshape(-1, 2) + def _horizon(self) -> int: hs = [ int(s["horizon"]) @@ -526,11 +605,17 @@ def resolve( located = self._steps_rows(keep) horizon = self._horizon() tables = self._leaf_tables() + ann_rows, ann_ts = {}, {} + if self.annotations is not None: + ann_rows = self._annotation_rows() + ann_ts = self._timestamps( + located, [r for r in keep.recording_id if r in ann_rows] + ) + no_ann = [] out = {} for row in keep.itertuples(): - uri, version, r0, valid = located[row.recording_id] - n = int(row.num_frames) + uri, version, r0, n, valid = located[row.recording_id] anchors = eligible_intervals( valid, HANDS_OF.get(row.embodiment, ("left", "right")), @@ -540,6 +625,21 @@ def resolve( if not len(anchors): logger.warning("recording %s has no eligible anchors", row.episode_hash) continue + ann_kw = {} + if self.annotations is not None: + spans = (np.empty(0, np.int64), np.empty(0, np.int64), []) + if row.recording_id in ann_ts: + spans = self._frame_spans( + ann_rows[row.recording_id], ann_ts[row.recording_id] + ) + if self.annotations["require"]: + anchors = self._covered(anchors, spans[0], spans[1], n) + if not len(anchors): + no_ann.append(row.episode_hash) + continue + ann_kw = dict( + annotations=spans, clamp_actions=self.annotations["clamp_actions"] + ) intr = dict(row.intrinsics) if row.intrinsics is not None else {} vids = dict(row.video_path) cam = next(iter(self.cameras.values())) @@ -560,11 +660,19 @@ def resolve( key_map=self.key_map, transform_list=self.transform_list, content_digest=getattr(row, "content_digest", None), + **ann_kw, + ) + if self.annotations is not None and self.annotations["require"]: + logger.info( + "kept %d/%d recordings with annotated anchors; dropped e.g. %s", + len(out), + len(out) + len(no_ann), + no_ann[:5], ) return out def _steps_rows(self, recs: pd.DataFrame) -> dict: - """{recording_id: (steps uri, version, first row, {hand: validity per frame})}, + """{recording_id: StepsRows}, each recording read from the steps table of its (lab, embodiment, end_effector).""" out = {} for key, group in recs.groupby(list(STEPS_KEY)): @@ -590,7 +698,31 @@ def _steps_rows(self, recs: pd.DataFrame) -> dict: ) r0 = int(starts[rid]) valid = {h: a[r0 : r0 + n] for h, a in all_valid.items()} - out[rid] = (uri, version, r0, valid) + out[rid] = StepsRows(uri, version, r0, n, valid) + return out + + @staticmethod + def _timestamps(located: dict, rids) -> dict[str, np.ndarray]: + """{recording_id: timestamp_ns per frame}, one take per steps table.""" + by_table: dict = {} + for rid in rids: + loc = located[rid] + by_table.setdefault((loc.uri, loc.version), []).append(rid) + out = {} + for (uri, version), ids in by_table.items(): + rows = np.concatenate( + [ + np.arange( + located[r].row_start, + located[r].row_start + located[r].num_frames, + ) + for r in ids + ] + ) + ts = _table(uri, version).take(rows.tolist(), columns=["timestamp_ns"]) + ts = ts.column(0).to_numpy() + at = np.cumsum([0] + [located[r].num_frames for r in ids]) + out.update({r: ts[at[j] : at[j + 1]] for j, r in enumerate(ids)}) return out def _leaf_tables(self) -> dict: diff --git a/egomimic/rldb/lance/schema.py b/egomimic/rldb/lance/schema.py index 46bbe9af84..e7575692ca 100644 --- a/egomimic/rldb/lance/schema.py +++ b/egomimic/rldb/lance/schema.py @@ -39,6 +39,40 @@ def only_steps_path(recordings) -> str: return paths[0] +def annotation_id( + recording_id: str, + annotation_vendor: str, + annotation_type: str, + t_start_ns: int, + t_end_ns: int, + text: str, +) -> str: + """Same content, same id: vendors can compute the id of a row they sent, to delete it.""" + key = "\0".join( + map( + str, + ( + recording_id, + annotation_vendor, + annotation_type, + t_start_ns, + t_end_ns, + text, + ), + ) + ) + return str(uuid.uuid5(RECORDING_NAMESPACE, f"annotation/{key}")) + + +def task_instance_id(recording_id: str, task_name: str | None, t_start_ns: int) -> str: + """Stable across rebuilds and computable by vendors; old episode_hashes resolve through aliases.""" + return str( + uuid.uuid5( + RECORDING_NAMESPACE, f"task/{recording_id}/{task_name or ''}/{t_start_ns}" + ) + ) + + def _vec(n: int) -> pa.DataType: return pa.list_(pa.float32(), n) @@ -92,15 +126,6 @@ def _vec(n: int) -> pa.DataType: ] ) -ANNOTATION = pa.struct( - [ - ("t_start_ns", pa.int64()), - ("t_end_ns", pa.int64()), - ("text", pa.string()), - ("annotator", pa.string()), - ] -) - TASKS = pa.schema( [ ("task_instance_id", pa.string()), @@ -110,7 +135,6 @@ def _vec(n: int) -> pa.DataType: ("task_name", pa.string()), ("task_description", pa.string()), ("objects", pa.list_(pa.string())), - ("annotations", pa.list_(ANNOTATION)), ] ) @@ -122,3 +146,22 @@ def _vec(n: int) -> pa.DataType: ("t_end_ns", pa.int64()), ] ) + +ANNOTATION_TYPES = ("language",) + +# Intervals are half-open [t_start_ns, t_end_ns) on the recording clock; a point event is one frame wide. +ANNOTATIONS = pa.schema( + [ + ("annotation_id", pa.string()), + ("recording_id", pa.string()), + ("t_start_ns", pa.int64()), + ("t_end_ns", pa.int64()), + ("annotation_type", pa.string()), + ("text", pa.string()), + ("annotation_vendor", pa.string()), + # Rubric items passed (0-5); null until graded. + ("quality_score", pa.int8()), + ("graded_by", pa.string()), + ("rubric_version", pa.string()), + ] +) diff --git a/egomimic/rldb/leaf.py b/egomimic/rldb/leaf.py new file mode 100644 index 0000000000..29bcf1a66a --- /dev/null +++ b/egomimic/rldb/leaf.py @@ -0,0 +1,240 @@ +"""Base class for one episode's samples, whatever the storage: keymap windowing +(chunks, history, lag), padding and masks, and the shared sample finishing.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from egomimic.rldb.embodiment.embodiment import HISTORY_MASK_SUFFIX, get_embodiment_id + + +class EpisodeLeaf(torch.utils.data.Dataset): + """One episode under a ``MultiDataset``. + + Subclasses set ``episode_path``, ``total_frames``, ``embodiment``, ``metadata`` + (with ``fps``) and ``intrinsics``, read the windows ``_read_windows`` names, and + hand the raw arrays to ``_finish``. + + A leaf with ``batched_fetch`` set has ``MultiDataset.__getitems__`` read a whole + DataLoader batch at once: ``plan_sample`` per sample, one ``fetch_planned`` per + leaf class, then ``assemble_sample`` per sample. + """ + + # False when the leaf runs threads that a forked DataLoader worker would + # inherit mid-flight (see multi.worker_context). + fork_safe = True + batched_fetch = False + + def __init__(self, *, key_map: dict, transform_list: list | None = None): + self.key_map = key_map + self.transform = transform_list + super().__init__() + + def cache_fingerprint(self) -> str | None: + """Digest the norm cache keys this episode's samples on; None disables it.""" + from egomimic.rldb.norm.norm_cache import episode_fingerprint + + return episode_fingerprint(getattr(self, "episode_path", None)) + + def plan_sample(self, i: int): + raise NotImplementedError + + @staticmethod + def fetch_planned(items: list) -> list: + raise NotImplementedError + + def assemble_sample(self, plan, fetched) -> dict: + raise NotImplementedError + + def __len__(self) -> int: + return self.total_frames + + def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: + """End index (exclusive) for a windowed read starting at ``start_idx``. + + Subclasses can override to add per-key-type clamping (e.g., annotation EOS). + """ + return min(start_idx + horizon, self.total_frames) + + def _action_horizon(self) -> int | None: + """Horizon of the keymap's ``action_keys`` entries, if it has any.""" + for spec in self.key_map.values(): + if ( + spec.get("key_type") == "action_keys" + and spec.get("horizon") is not None + ): + return int(spec["horizon"]) + return None + + def _action_pad_mask(self, idx: int) -> np.ndarray | None: + """1.0 per real action frame, 0.0 per repeat-last padded one, (horizon,). + + ``_pad_to_horizon`` pads a short action chunk by repeating the last + frame, which in a wrist-relative action space supervises "stop moving" + at every episode (or annotation) tail. This mask says which steps are + real so ``loss_fn`` can drop the rest; ``None`` when the keymap has no + horizoned action key. + """ + horizon = self._action_horizon() + if horizon is None: + return None + n_real = self._chunk_end_idx(idx, horizon, "action_keys") - idx + mask = np.zeros(horizon, dtype=np.float32) + mask[: max(n_real, 0)] = 1.0 + return mask + + def _fps(self) -> float: + """The episode's frame rate (zarr ``fps`` attr; 30 when absent).""" + return float(self.metadata.get("fps") or 30) + + def _seconds_to_frames(self, seconds: float) -> int: + """``seconds`` in this episode's fps, at least one frame.""" + return max(1, round(float(seconds) * self._fps())) + + def _lagged_idx(self, idx: int, lag_s: float) -> int: + """Frame ``lag_s`` seconds before ``idx`` (at least one frame back), + clamped at the episode start -- so the first samples read frame 0, + the same front padding a history window gets.""" + return max(0, idx - self._seconds_to_frames(lag_s)) + + def _history_stride(self, spec: dict) -> int: + """Frames between history steps: ``history_stride_s`` (seconds, so a + window spans the same time at any fps) wins over ``history_stride``.""" + stride_s = spec.get("history_stride_s") + if stride_s is not None: + return self._seconds_to_frames(stride_s) + return int(spec.get("history_stride", 1) or 1) + + def _history_indices(self, idx: int, spec: dict) -> tuple[np.ndarray, np.ndarray]: + """``(K,)`` frame indices of a ``history`` entry, oldest first, and their + ``(K,)`` float32 mask. + + The newest step is ``idx`` or, with ``lag_s``, that many seconds before + it. Steps before the episode start are masked 0 and repeat the oldest + real step (frame 0 when none is real), so the newest step is always + ``[-1]``. + """ + history = int(spec["history"]) + newest = idx + if spec.get("lag_s") is not None: + newest = idx - self._seconds_to_frames(spec["lag_s"]) + steps = newest - self._history_stride(spec) * np.arange(history - 1, -1, -1) + real = steps >= 0 + fill = steps[real][0] if real.any() else 0 + return np.where(real, steps, fill), real.astype(np.float32) + + @staticmethod + def _pad_to_horizon(array, horizon: int | None): + """Repeat-last out to ``horizon`` a chunk the episode (or annotation) + ended before filling.""" + if horizon is None or not isinstance(array, np.ndarray): + return array + missing = horizon - array.shape[0] + if missing <= 0: + return array + return np.concatenate([array, np.repeat(array[-1:], missing, axis=0)], axis=0) + + def _read_windows( + self, idx: int + ) -> dict[str, tuple[str, tuple[int, int | None] | np.ndarray]]: + """``{keymap key: (zarr key, window)}`` for frame *idx*. + + ``window`` is ``(start, end)``, ``end`` None meaning the single frame + at ``start``, or for a ``history`` entry the ``(K,)`` frame indices of + ``_history_indices``. Annotation keys read nothing and are left out. + """ + windows = {} + for k, spec in self.key_map.items(): + if spec.get("key_type") == "annotation_keys": + continue + horizon = spec.get("horizon") + history = spec.get("history") + lag_s = spec.get("lag_s") + if horizon is not None: + window = ( + idx, + self._chunk_end_idx(idx, horizon, spec.get("key_type")), + ) + elif history is not None and ( + int(history) > 1 or spec.get("key_type") == "camera_keys" + ): + # A camera window keeps its (K, ...) axis even at K = 1. Point + # reads, not a contiguous span: on a JPEG key (one frame per + # chunk) a span would fetch every frame between the steps. + # ZarrEpisode.plan_reads still merges them on numeric keys. + window = self._history_indices(idx, spec)[0] + elif lag_s is not None: + # A single PAST frame (image history): the same zarr array as + # the current-frame entry, read ``lag_s`` seconds back. + window = (self._lagged_idx(idx, lag_s), None) + else: + window = (idx, None) + windows[k] = (spec["zarr_key"], window) + return windows + + def _finish(self, idx: int, data: dict) -> dict: + """Masks, transforms and per-episode metadata on the raw arrays read for + frame ``idx``; returns the sample as tensors.""" + pad_mask = self._action_pad_mask(idx) + if pad_mask is not None: + # Travels with the chunk through the interpolators + # (InterpolatePadMask) to the model horizon. + data["action_pad_mask"] = pad_mask + + for k, spec in self.key_map.items(): + if spec.get("key_type") == "camera_keys" and spec.get("history"): + # Unlike proprio, a padded frame has to be masked out of + # attention, so each image window carries its own mask. + data[f"{k}{HISTORY_MASK_SUFFIX}"] = self._history_indices(idx, spec)[1] + + if self.transform: + for transform in self.transform or []: + data = transform.transform(data) + + for k, v in data.items(): + if isinstance(v, np.ndarray): + data[k] = torch.from_numpy(v).to(torch.float32) + + data["embodiment"] = get_embodiment_id(self.embodiment) + # Control frequency of this episode (RDT's ctrl-freq token). + data["fps"] = torch.tensor(self._fps(), dtype=torch.float32) + data["intrinsics"] = torch.from_numpy(self._intrinsics_3x4()) + data["episode_hash"] = self._episode_hash() + return data + + def _intrinsics_3x4(self) -> np.ndarray: + """``intrinsics`` as the (3, 4) K the batch carries.""" + # Per-episode camera intrinsics travel with the batch so a single + # data-driven embodiment can project without a hardcoded class const. + # Single-camera -> (3,4) K matrix; no/multi-camera -> NaN sentinel, + # which _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". + K = self.intrinsics + if isinstance(K, dict): + # Multi-camera rig: project against the front camera (viz/projection + # target the front image). Prefer a "front"-keyed entry, else the + # first available camera. + K = next( + (v for k, v in K.items() if "front" in str(k).lower()), + next(iter(K.values()), None) if K else None, + ) + if K is not None: + K = np.asarray(K, dtype=np.float32) + if K.size == 9: + # Lance stores K flat. + K = K.reshape(3, 3) + if K.shape == (3, 3): + # Normalize a 3x3 K to the canonical 3x4 (zeros last column); + # some contributors store 3x3 (e.g. microagi). + K = np.concatenate([K, np.zeros((3, 1), dtype=np.float32)], axis=1) + if K.shape != (3, 4): # unexpected -> sentinel (viz falls back to const) + K = np.full((3, 4), np.nan, dtype=np.float32) + else: + K = np.full((3, 4), np.nan, dtype=np.float32) + return np.ascontiguousarray(K) + + def _episode_hash(self) -> str: + name = Path(self.episode_path).name + return name[:-5] if name.endswith(".zarr") else name diff --git a/egomimic/rldb/multi.py b/egomimic/rldb/multi.py new file mode 100644 index 0000000000..057149eed3 --- /dev/null +++ b/egomimic/rldb/multi.py @@ -0,0 +1,1272 @@ +"""The training-loop dataset: a dict of episode leaves (zarr or Lance) with the +index map, bounds check, normalization stats and batched-fetch dispatch on top.""" + +from __future__ import annotations + +import concurrent.futures as cf +import copy +import functools +import inspect +import json +import logging +import math +import os +import random +import time +from pathlib import Path +from typing import Mapping + +import numpy as np +import torch +from tqdm import tqdm + +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.filters import DatasetFilter, PinError +from egomimic.utils.action_utils import ( + _apply_norm_one, + _apply_unnorm_one, + degenerate_mask, +) +from egomimic.utils.gpu_orphans import die_with_parent +from egomimic.utils.pose_utils import bounds_check_channels, rot6d_channels + +logger = logging.getLogger(__name__) + + +@functools.cache +def _fetch_pool() -> cf.ThreadPoolExecutor: + """Per-process pool for the concurrent parts of a batched leaf fetch.""" + return cf.ThreadPoolExecutor(max_workers=8, thread_name_prefix="batch-fetch") + + +SEED = 42 + + +def split_dataset_names(dataset_names, valid_ratio=0.2, seed=SEED): + """ + Split a list of dataset names into train/valid sets. + Args: + dataset_names (Iterable[str]) + valid_ratio (float): fraction of datasets to put in valid. + seed (int): for deterministic shuffling. + + + Returns: + train_set (set[str]), valid_set (set[str]) + """ + names = sorted(dataset_names) + if not names: + return set(), set() + + rng = random.Random(seed) + rng.shuffle(names) + + if not (0.0 <= valid_ratio <= 1.0): + raise ValueError(f"valid_ratio must be in [0,1], got {valid_ratio}") + + n_valid = int(len(names) * valid_ratio) + if valid_ratio > 0.0: + n_valid = max(1, n_valid) + + valid = set(names[:n_valid]) + train = set(names[n_valid:]) + return train, valid + + +def _infer_key_type(key_name: str) -> str | None: + """ + Heuristically infer a key_type for a post-transform key that wasn't + declared in the leaf's key_map (typically produced by a transform like + ``ConcatKeys``). Returns one of the recognised key_type strings or None + when the key looks like metadata that NormStats shouldn't normalize. + + Recognition is conservative — only positive matches return a type. + """ + name = key_name.lower() + # Camera image keys + if "image" in name or "/img" in name or name.endswith("_img") or ".img" in name: + return "camera_keys" + # Action keys (model outputs) + if name.startswith("actions") or "/cmd_" in name or ".cmd_" in name: + return "action_keys" + # Proprioception (observed state) + if ( + name.startswith("observations.state") + or "/obs_" in name + or ".obs_" in name + or "joint_positions" in name + or "ee_pose" in name + or "keypoints" in name + ): + return "proprio_keys" + # Language / annotations + if "annotation" in name or "tokenized" in name or "lang" in name: + return "annotation_keys" + return None + + +def worker_context(leaves) -> str | None: + """DataLoader start method for ``leaves``: forkserver if any is not fork-safe + (Lance runs Rust threads; https://lance.org/integrations/pytorch/), else the + default fork.""" + if any( + not getattr(sub, "fork_safe", True) + for leaf in leaves + for sub in MultiDataset._iter_leaves(leaf) + ): + return "forkserver" + return None + + +class _IndexMap: + """Global index -> (dataset name, local index), as offsets rather than one tuple per sample. + + A per-sample list costs ~100 B a sample in every process: tens of GB at 10^7 samples, + copied page by page into forked workers and pickled whole into spawned ones. + """ + + def __init__(self, datasets: dict): + self.names = list(datasets) + self.offsets = np.cumsum( + [0] + [len(d) for d in datasets.values()], dtype=np.int64 + ) + + def __len__(self) -> int: + return int(self.offsets[-1]) + + def __getitem__(self, idx: int) -> tuple[str, int]: + if idx < 0: + idx += len(self) + if not 0 <= idx < len(self): + raise IndexError(idx) + j = int(np.searchsorted(self.offsets, idx, side="right")) - 1 + return self.names[j], int(idx - self.offsets[j]) + + def __iter__(self): + for j, name in enumerate(self.names): + for local in range(int(self.offsets[j + 1] - self.offsets[j])): + yield name, local + + def ranges(self) -> dict[str, range]: + return { + n: range(int(self.offsets[j]), int(self.offsets[j + 1])) + for j, n in enumerate(self.names) + } + + +class MultiDataset(torch.utils.data.Dataset): + """ + Wraps a dict of child datasets (Zarr leaves or other MultiDatasets) and + also owns the normalization-stats descriptor and helpers that used to live + in separate ``NormStats`` / ``NormalizingMultiDataset`` classes. + + Two construction modes: + - **Data mode** (default): pass ``datasets`` to wrap a real dataset graph + (existing behaviour, used during training). Stats fields start empty; + call ``populate_from_datasets()`` and ``infer_norm_from_dataset(...)`` + to fill them in, then ``attach_normalize_transforms()`` to wire + normalize/reject transforms onto each leaf's ``transform_list``. + - **State mode** (``state=...``, ``datasets=None``): construct a + stats-only instance for deploy/eval where the dataset graph isn't + available. ``self.datasets`` is empty; only the stats fields are + populated. Used by checkpoint reconstruction. + """ + + NORMALIZE_KEY_TYPES = ("proprio_keys", "action_keys") + # Saved in the state: a model trained before rot6d passed through + # normalization must keep normalizing it (see _load_state). + rot6d_passthrough = True + + # Bounds-check slack per (timestep, channel) cell; see _check_bounds. + BOUNDS_ABS_TOL = 1e-6 + BOUNDS_REL_SLACK = 0.5 + # Bad-sample retries stay in the failing episode this many times, then + # draw from the whole index space so one bad episode can't exhaust them. + GLOBAL_FALLBACK_ATTEMPTS = 25 + # Consecutive bad samples before raising: a systemic-failure detector. + MAX_FALLBACK_ATTEMPTS = 100 + # Retry warnings: the first RETRY_LOG_FIRST per worker, then every RETRY_LOG_EVERY. + RETRY_LOG_FIRST = 20 + RETRY_LOG_EVERY = 500 + # Episodes tried per key_map/transform group before populate_from_datasets + # falls back to the raw key_map. + POPULATE_PROBES = 8 + _retry_log_count = 0 + + def __init__( + self, + datasets: (dict[str, MultiDataset | torch.utils.data.Dataset] | None) = None, + mode: str = "train", + percent: float = 0.1, + valid_ratio: float = 0.2, + norm_mode: str = "zscore", + state: dict | None = None, + **kwargs, + ): + """ + Args: + datasets: Dict of child datasets. None for state-only construction. + mode: One of "train", "valid", "total", "percent" — which split to keep. + percent: Fraction (when mode="percent"). + valid_ratio: Train/valid split ratio. + norm_mode: One of "zscore", "minmax", "quantile". + state: If provided, populate stats fields from this dict (deploy mode). + """ + super().__init__() + + # ---- Stats fields (always present, may be empty) ---- + self.norm_mode = norm_mode + self.embodiments: set[int] = set() + self.key_types: dict[int, dict[str, str]] = {} + self.zarr_keys: dict[int, dict[str, str]] = {} + self.shapes: dict[int, dict[str, tuple]] = {} + self.norm_stats: dict[int, dict[str, dict[str, np.ndarray]]] = {} + self._norm_run_metadata: dict[str, float | int | None] | None = None + + # ---- Dataset graph fields ---- + self.datasets: dict = {} + self.index_map: list = [] + self._global_indices_by_dataset: dict[str, list[int]] = {} + # Dedup bounds-check warnings: keyed by f"bounds:{episode}:{zarr_key}" + self._warned_violations: set[str] = set() + self.train_collections: set = set() + self.valid_collections: set = set() + + if state is not None: + # Deploy / state-only construction — no dataset graph. + self._load_state(state) + return + + if datasets is None: + raise ValueError("MultiDataset requires either `datasets` or `state`.") + + # ---- Normal data-mode construction ---- + self.train_collections, self.valid_collections = split_dataset_names( + datasets.keys(), valid_ratio=valid_ratio, seed=SEED + ) + + if mode == "train": + chosen = self.train_collections + elif mode == "valid": + chosen = self.valid_collections + elif mode == "total": + chosen = set(datasets.keys()) + elif mode == "percent": + all_names = sorted(datasets.keys()) + rng = random.Random(SEED) + rng.shuffle(all_names) + n_keep = int(len(all_names) * percent) + if percent > 0.0: + n_keep = max(1, n_keep) + chosen = set(all_names[:n_keep]) + else: + raise ValueError(f"Unknown mode: {mode}") + + self.datasets = {rid: ds for rid, ds in datasets.items() if rid in chosen} + assert self.datasets, "No datasets left after applying mode split." + + self.index_map = _IndexMap(self.datasets) + self._global_indices_by_dataset = self.index_map.ranges() + + def __len__(self) -> int: + return len(self.index_map) + + @staticmethod + def _episode_name_for_dataset(dataset, dataset_name: str) -> str: + episode_path = getattr(dataset, "episode_path", None) + if episode_path is None: + return dataset_name + return Path(episode_path).name + + def set_norm_stats_from(self, source: "MultiDataset") -> None: + """Share stats with this dataset (and its nested MultiDatasets) by + reference. After this call ``__getitem__`` will bounds-check + normalize + each sample using ``source``'s ``norm_stats``/``key_types``/``zarr_keys``. + + Use this *instead of* ``attach_normalize_transforms`` — it doesn't mutate + any leaf-level ``transform`` list, so it can't accumulate duplicate + passes when leaves share a transform list reference. + """ + self.norm_stats = source.norm_stats + self.key_types = source.key_types + self.zarr_keys = source.zarr_keys + self.shapes = source.shapes + self.embodiments = source.embodiments + self.norm_mode = source.norm_mode + self.rot6d_passthrough = source.rot6d_passthrough + # Each MultiDataset keeps its own warning-dedup state. + self._warned_violations = set() + for ds in self.datasets.values(): + if isinstance(ds, MultiDataset): + ds.set_norm_stats_from(source) + + def _check_bounds( + self, data: dict, dataset, idx: int, dataset_name: str + ) -> str | None: + """Return a violation message if any tracked key in ``data`` has NaN/Inf + or values outside per-key quantile bounds. ``None`` means the sample + passes. Logs each (episode, key) violation once. + """ + embodiment_id = data.get("embodiment") + if embodiment_id is None: + return None + per_emb_stats = self.norm_stats.get(embodiment_id, {}) + if not per_emb_stats: + return None + + episode_name = self._episode_name_for_dataset(dataset, dataset_name) + + for key_name, stats in per_emb_stats.items(): + zarr_key = self.zarr_keys.get(embodiment_id, {}).get(key_name) + if zarr_key is None or zarr_key not in data: + continue + v = data[zarr_key] + if isinstance(v, torch.Tensor): + arr = v.float() + elif isinstance(v, np.ndarray): + arr = torch.from_numpy(v).float() + else: + continue + + q_low = stats.get( + "quantile_0_01", stats.get("quantile_0_1", stats["quantile_1"]) + ) + q_high = stats.get( + "quantile_99_99", stats.get("quantile_99_9", stats["quantile_99"]) + ) + q_low = torch.as_tensor(q_low, device=arr.device, dtype=torch.float32) + q_high = torch.as_tensor(q_high, device=arr.device, dtype=torch.float32) + try: + q_low = torch.broadcast_to(q_low, arr.shape) + q_high = torch.broadcast_to(q_high, arr.shape) + except RuntimeError: + # Stats were computed for a different layout than this sample + # (e.g. a stale precomputed norm_stats.json). Say so once + # instead of silently disabling the bounds check for the key; + # normalize() raises on the same mismatch anyway. + warn_key = f"bounds-shape:{zarr_key}" + if warn_key not in self._warned_violations: + self._warned_violations.add(warn_key) + logger.warning( + f"[MultiDataset] bounds check skipped for {zarr_key}: " + f"stats shape {tuple(q_low.shape)} does not broadcast to " + f"sample shape {tuple(arr.shape)} (norm stats computed " + "for a different layout?)" + ) + continue + + if torch.any(torch.isnan(arr)) or torch.any(torch.isinf(arr)): + prefix = f"NaN/Inf in {zarr_key} ep={episode_name} frame={idx}" + warn_key = f"naninf:{episode_name}:{zarr_key}" + if warn_key not in self._warned_violations: + self._warned_violations.add(warn_key) + logger.warning(prefix) + return prefix + + # Rotation channels are either Euler ypr (wraps at +-pi) or + # continuous 6D columns; quantile bounds on them are meaningless + # and reject otherwise-valid frames, so only the translation / + # gripper / keypoint channels are bounds-checked. Unrecognized + # widths fall through to a full-vector check; NaN/Inf above still + # covers the full vector. + check_idx = bounds_check_channels(zarr_key, arr.shape[-1]) + if check_idx is not None: + arr_q = arr[..., check_idx] + q_low = q_low[..., check_idx] + q_high = q_high[..., check_idx] + else: + arr_q = arr + + # Per-cell any-cell checks against subsampled stats reject valid + # frames in bulk, so allow BOUNDS_REL_SLACK of each cell's quantile + # range (plus an absolute floor). Corrupt values sit orders of + # magnitude outside and are still caught. + tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low) + below = arr_q < q_low - tol + above = arr_q > q_high + tol + # Cells normalize() treats as constant (wrist-frame t=0, the kp0 / + # kp9 structural zeros) cannot blow up the loss whatever offset + # they carry, so don't reject on them. The exemption follows + # norm_mode's range stat, the same one normalize() uses. + constant = self._degenerate_cells(stats, arr.shape) + if constant is not None: + if check_idx is not None: + constant = constant[..., check_idx] + below = below & ~constant + above = above & ~constant + if torch.any(below) or torch.any(above): + prefix = f"Bounds violation in {zarr_key} ep={episode_name} frame={idx}" + warn_key = f"bounds:{episode_name}:{zarr_key}" + if warn_key not in self._warned_violations: + self._warned_violations.add(warn_key) + n_below = int(below.sum().item()) + n_above = int(above.sum().item()) + logger.warning( + f"{prefix} | n_below={n_below} n_above={n_above} " + f"arr_range=[{arr_q.min().item():.4f}, {arr_q.max().item():.4f}]" + ) + return prefix + return None + + def _degenerate_cells(self, stats: dict, shape) -> torch.Tensor | None: + """Mask of the cells ``normalize()`` treats as constant, or None when + this key's stats do not carry the range ``norm_mode`` reads.""" + try: + mask = degenerate_mask(stats, self.norm_mode) + except (KeyError, ValueError): + return None + return torch.broadcast_to(mask, shape) + + def __getitem__(self, idx, _attempts: int | None = None): + attempts = _attempts + while True: + dataset_name, local_idx = self.index_map[idx] + dataset = self.datasets[dataset_name] + try: + data = dataset[local_idx] + except Exception as e: + reason = ( + f"Sample failed ({type(e).__name__}: {e}) at " + f"{dataset_name}[{local_idx}]" + ) + idx, attempts = self._next_after_failure( + idx, dataset_name, attempts, reason=reason + ) + continue + + # If this leaf is itself a MultiDataset, it already ran bounds + + # normalize for its returned sample. Pass through unchanged. + if isinstance(dataset, MultiDataset): + return data + + violation = self._check_bounds(data, dataset, local_idx, dataset_name) + if violation is not None: + idx, attempts = self._next_after_failure( + idx, dataset_name, attempts, reason=violation + ) + continue + + # Bounds passed: normalize and return. + if self.norm_stats and data.get("embodiment") in self.norm_stats: + data = self.normalize(data, data["embodiment"]) + return data + + def __getitems__(self, idxs: list) -> list: + """A DataLoader batch. Leaves with ``batched_fetch`` set (Lance streaming) do + their reads for the whole batch at once, which is what makes random access from + object storage affordable; every other sample, and any batched one that fails + or trips the bounds check, goes through ``__getitem__`` and its fallback.""" + out = [None] * len(idxs) + groups: dict[type, list] = {} + for pos, idx in enumerate(idxs): + name, local = self.index_map[idx] + leaf = self.datasets[name] + if getattr(leaf, "batched_fetch", False): + groups.setdefault(type(leaf), []).append((pos, name, leaf, local)) + for cls, members in groups.items(): + planned = [] + for pos, name, leaf, local in members: + try: + planned.append((pos, name, leaf, local, leaf.plan_sample(local))) + except Exception: + continue + if not planned: + continue + step = getattr(planned[0][2], "fetch_split", None) or len(planned) + parts = [planned[i : i + step] for i in range(0, len(planned), step)] + + def fetch(part, cls=cls): + return cls.fetch_planned([(m[2], m[4]) for m in part]) + + if len(parts) > 1: + # Parts download concurrently; each is assembled as soon as it lands, + # so CPU work overlaps the remaining reads. + futures = {_fetch_pool().submit(fetch, part): part for part in parts} + done = ((futures[f], f.result) for f in cf.as_completed(futures)) + else: + done = [(parts[0], functools.partial(fetch, parts[0]))] + for part, result in done: + try: + fetched = result() + except Exception: + logger.warning( + "[MultiDataset] batched fetch failed; reading samples one by one", + exc_info=True, + ) + continue + for (pos, name, leaf, local, plan), got in zip(part, fetched): + try: + data = leaf.assemble_sample(plan, got) + except Exception: + continue + if self._check_bounds(data, leaf, local, name) is not None: + continue + if self.norm_stats and data.get("embodiment") in self.norm_stats: + data = self.normalize(data, data["embodiment"]) + out[pos] = data + return [self[idx] if d is None else d for d, idx in zip(out, idxs)] + + def _next_after_failure( + self, idx: int, dataset_name: str, attempts: int | None, *, reason: str + ) -> tuple[int, int]: + """Pick the next index to try after a bad sample at ``idx``. + + Raises once ``MAX_FALLBACK_ATTEMPTS`` consecutive samples failed.""" + attempts = (attempts or 0) + 1 + if attempts >= self.MAX_FALLBACK_ATTEMPTS: + raise RuntimeError( + f"[MultiDataset] {attempts} consecutive bad samples, the last " + f"{attempts - self.GLOBAL_FALLBACK_ATTEMPTS} drawn uniformly " + "from the whole index space: the store, the keymap or the norm " + f"stats are broken, not one episode. Last: {reason}" + ) + local = self._global_indices_by_dataset[dataset_name] + if attempts <= self.GLOBAL_FALLBACK_ATTEMPTS and len(local) > 1: + next_idx = idx + while next_idx == idx: + next_idx = random.choice(local) + else: + next_idx = random.randrange(len(self.index_map)) + next_dataset_name, next_local_idx = self.index_map[next_idx] + self._retry_log_count += 1 + if ( + self._retry_log_count <= self.RETRY_LOG_FIRST + or self._retry_log_count % self.RETRY_LOG_EVERY == 0 + ): + logger.warning( + f"{reason} | attempt {attempts}, " + f"trying {next_dataset_name}[{next_local_idx}] " + f"(retries so far in this worker: {self._retry_log_count})" + ) + return next_idx, attempts + + @classmethod + def _from_resolver(cls, resolver, dataset_name: str | None = None, **kwargs): + """create a MultiDataset from a resolver: anything with + ``resolve(filters=, expected_embodiment=) -> {name: leaf}``. + + dataset_name: the config key this dataset is built under (trainHydra passes + it). When it names an embodiment, pinned episodes must carry that + embodiment (PinError otherwise) and other mismatches are logged. + + Every pinned episode must also end up in the dataset: a pin that + filter_lambdas or the debug limit exclude, or that fails to load or is + skipped by the resolver, is a PinError rather than a silent drop.""" + sync_from_s3 = kwargs.pop("sync_from_s3", False) + filters = kwargs.pop("filters", None) + expected = None + if dataset_name is not None: + try: + get_embodiment_id(dataset_name) + expected = dataset_name + except (KeyError, AttributeError): + expected = None + + # Only resolvers whose resolve() takes the flag get it (LocalEpisodeResolver + # warns and ignores it; S3 resolvers always sync). + takes_sync = "sync_from_s3" in inspect.signature(resolver.resolve).parameters + extra = {"sync_from_s3": True} if sync_from_s3 and takes_sync else {} + resolved = resolver.resolve( + filters=filters, expected_embodiment=expected, **extra + ) + + pins = filters.episode_hashes if isinstance(filters, DatasetFilter) else () + dropped = sorted(set(pins) - set(resolved)) + if dropped: + raise PinError( + f"{len(dropped)} pinned episode(s) not in dataset " + f"'{dataset_name or ''}': {dropped}. They were excluded by " + "filter_lambdas or the resolver's debug limit, or failed to load / " + "were skipped by the resolver (see the log above)." + ) + + return cls(datasets=resolved, **kwargs) + + # ===================================================================== + # Stats / normalization + # ===================================================================== + + @staticmethod + def _iter_leaves(ds): + """Yield non-MultiDataset leaves from possibly nested wrappers.""" + if isinstance(ds, MultiDataset): + for child in ds.datasets.values(): + yield from MultiDataset._iter_leaves(child) + else: + yield ds + + def populate_from_datasets(self, datasets: dict | None = None) -> None: + """ + Populate per-embodiment key inventory by walking leaves and probing + one post-transform sample per (embodiment, key_map, transform list). + ``datasets`` defaults to ``self.datasets`` so the typical call is just + ``mds.populate_from_datasets()``. + """ + graph = datasets if datasets is not None else self.datasets + groups: dict[tuple, list] = {} + for ds in graph.values(): + for leaf in self._iter_leaves(ds): + emb = getattr(leaf, "embodiment", None) + key_map = getattr(leaf, "key_map", None) + if emb is None or key_map is None: + continue + emb_id = emb if isinstance(emb, int) else get_embodiment_id(emb) + # A resolver hands every episode the same key_map and transform + # objects, so one probe covers them all: probing each leaf read a + # full sample per episode (~20 min on 41k mecka episodes). + key = (emb_id, id(key_map), id(getattr(leaf, "transform", None))) + groups.setdefault(key, [emb_id, key_map, []])[2].append(leaf) + + for emb_id, key_map, leaves in groups.values(): + self.embodiments.add(emb_id) + self.key_types.setdefault(emb_id, {}) + self.zarr_keys.setdefault(emb_id, {}) + self.shapes.setdefault(emb_id, {}) + self.norm_stats.setdefault(emb_id, {}) + + sample_keys: set | None = None + for leaf in leaves[: self.POPULATE_PROBES]: + try: + sample = leaf[0] + except Exception as e: + logger.warning( + f"[MultiDataset] Could not probe leaf for post-transform " + f"keys (emb={emb_id}): {e}. Trying the next episode." + ) + continue + if isinstance(sample, dict): + sample_keys = set(sample.keys()) + break + if sample_keys is None: + logger.warning( + f"[MultiDataset] No probe succeeded (emb={emb_id}). " + "Falling back to raw key_map." + ) + for key_name, info in key_map.items(): + self.key_types[emb_id][key_name] = info.get( + "key_type", "metadata_keys" + ) + self.zarr_keys[emb_id][key_name] = info["zarr_key"] + continue + + # Identity zarr_keys map (data_key is the algo-side name). + for data_key in sample_keys: + if data_key in key_map: + info = key_map[data_key] + self.key_types[emb_id][data_key] = info.get( + "key_type", "metadata_keys" + ) + else: + inferred = _infer_key_type(data_key) + if inferred is None: + continue + self.key_types[emb_id][data_key] = inferred + self.zarr_keys[emb_id][data_key] = data_key + + # ---- key lookups ---- + + def keys_of_type(self, key_type: str, embodiment_id: int) -> list[str]: + return [ + k for k, t in self.key_types.get(embodiment_id, {}).items() if t == key_type + ] + + def is_key_with_embodiment(self, key_name: str, embodiment_id: int) -> bool: + return key_name in self.key_types.get(embodiment_id, {}) + + def keyname_to_zarr_key(self, key_name: str, embodiment_id: int) -> str | None: + return self.zarr_keys.get(embodiment_id, {}).get(key_name) + + def zarr_key_to_keyname(self, zarr_key: str, embodiment_id: int) -> str | None: + for k, v in self.zarr_keys.get(embodiment_id, {}).items(): + if v == zarr_key: + return k + return None + + def key_shape(self, key_name: str, embodiment_id: int) -> tuple: + if key_name not in self.shapes.get(embodiment_id, {}): + raise ValueError( + f"Shape for key {key_name!r} on embodiment {embodiment_id} not inferred yet." + ) + return self.shapes[embodiment_id][key_name] + + # ---- shape & norm inference ---- + + def infer_shapes_from_batch(self, batch: dict) -> None: + for emb_id, per_emb in self.zarr_keys.items(): + for key_name, zarr_key in per_emb.items(): + if zarr_key in batch: + val = batch[zarr_key] + if hasattr(val, "shape"): + self.shapes.setdefault(emb_id, {})[key_name] = tuple(val.shape) + elif isinstance(val, int): + self.shapes.setdefault(emb_id, {})[key_name] = (1,) + + def infer_norm_from_dataset( + self, + dataset, + dataset_name, + sample_frac: float = 0.10, + seed: int = 42, + max_samples: int | None = None, + batch_size: int = 512, + num_workers: int = 4, + precomputed_norm_path: str | None = None, + pool_horizon: bool = False, + episode_cache: str | Path | None = None, + episode_fingerprints: Mapping[str, str | None] | None = None, + ): + """Fill ``self.norm_stats[embodiment]`` from a precomputed file or by + sampling ``dataset``. + + ``episode_cache``: a per-episode sample store + (``episode_norm_samples.cache_root``). Rows are read from, and new + episodes written to, it instead of sampling the whole dataset. + ``episode_fingerprints`` saves recomputing the episodes' fingerprints. + + ``pool_horizon``: compute each chunked key's stats over every timestep + of the chunk (one set per channel, like openpi) instead of per + (timestep, channel) cell. The pooled stats are tiled back to the + ``(T, D)`` shape, so normalize / bounds check / checkpoints see the + same layout either way. + """ + embodiment = dataset_name + if isinstance(embodiment, str): + embodiment = get_embodiment_id(embodiment) + + norm_keys = list(self.keys_of_type("proprio_keys", embodiment)) + norm_keys.extend(self.keys_of_type("action_keys", embodiment)) + if not norm_keys: + logger.warning( + f"[MultiDataset] No proprio/action keys for embodiment={embodiment}" + ) + return + + self.norm_stats.setdefault(embodiment, {}) + + if precomputed_norm_path is not None: + if os.path.isdir(precomputed_norm_path): + precomputed_file = os.path.join( + precomputed_norm_path, "norm_stats.json" + ) + elif os.path.isfile(precomputed_norm_path): + precomputed_file = precomputed_norm_path + else: + logger.warning( + f"[MultiDataset] precomputed_norm_path={precomputed_norm_path} is not valid" + ) + return + if os.path.isfile(precomputed_file): + with open(precomputed_file, "r") as f: + payload = json.load(f) + if str(embodiment) not in payload["stats"]: + raise ValueError( + f"norm_stats file {precomputed_file} has no entry for " + f"embodiment id {embodiment} (available: " + f"{sorted(payload['stats'])}). Stats are keyed by numeric " + "EMBODIMENT id, and ids were renumbered by the human/eva " + "embodiment collapse — recompute norm stats instead of " + "reusing a pre-collapse norm_stats.json." + ) + # JSON gives back nested lists; a fresh computation stores + # float32 arrays, so convert to keep both paths identical. + self.norm_stats[embodiment] = { + k: {n: np.asarray(a, dtype=np.float32) for n, a in d.items()} + for k, d in payload["stats"][str(embodiment)].items() + } + self._norm_run_metadata = payload.get("norm_run_metadata", None) + logger.info( + f"[MultiDataset] Loaded precomputed stats for embodiment={embodiment}" + ) + return + + if episode_cache is not None and any( + isinstance(ds, MultiDataset) or getattr(ds, "episode_path", None) is None + for ds in dataset.datasets.values() + ): + logger.warning( + "[MultiDataset] per-episode norm cache needs a flat dataset of " + "zarr episodes; sampling the whole dataset instead" + ) + episode_cache = None + + N = len(dataset) + if N <= 0: + raise ValueError("Dataset is empty") + n_samples = int(math.ceil(sample_frac * N)) + n_samples = max(1, min(n_samples, N)) + if max_samples is not None: + n_samples = min(n_samples, max_samples) + + logger.info(f"[MultiDataset] embodiment={embodiment} norm_keys={norm_keys}") + logger.info( + f"[MultiDataset] sampling {n_samples}/{N} (~{100 * sample_frac:.1f}%)" + ) + + loading_start = time.time() + episode_meta = {} + if episode_cache is not None: + from egomimic.rldb.norm import episode_norm_samples + + zarr_keys = { + k: z + for k in norm_keys + if (z := self.keyname_to_zarr_key(k, embodiment)) is not None + } + collected, episode_meta = episode_norm_samples.collect( + Path(episode_cache), + dataset.datasets, + zarr_keys, + n_samples, + seed=seed, + num_workers=num_workers, + fingerprints=episode_fingerprints, + ) + for k in norm_keys: + collected.setdefault(k, []) + else: + loader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + num_workers=num_workers, + shuffle=True, + generator=torch.Generator().manual_seed(seed), + worker_init_fn=die_with_parent, + multiprocessing_context=( + worker_context([dataset]) if num_workers > 0 else None + ), + ) + collected = self._collect_norm_samples( + loader, norm_keys, embodiment, n_samples, batch_size, num_workers + ) + for k in [k for k, v in collected.items() if not v]: + del collected[k] + norm_keys.remove(k) + loading_time = time.time() - loading_start + + computing_start = time.time() + for k in norm_keys: + collected[k] = self._drop_nonfinite_rows( + np.concatenate(collected[k], axis=0), k + ) + X = collected[k] + if pool_horizon and X.ndim > 2: + stats_np = self._compute_stats_for_array(X.reshape(-1, X.shape[-1])) + stats_np = { + name: np.broadcast_to(arr, X.shape[1:]).copy() + for name, arr in stats_np.items() + } + else: + stats_np = self._compute_stats_for_array(X) + self.norm_stats[embodiment][k] = { + name: np.asarray(arr, dtype=np.float32) + for name, arr in stats_np.items() + } + logger.info( + f"[MultiDataset] key={k} samples={collected[k].shape[0]} stat_shape={stats_np['mean'].shape}" + ) + computing_time = time.time() - computing_start + + self._norm_run_metadata = { + "loading_time": loading_time, + "computing_time": computing_time, + "frames": n_samples, + **episode_meta, + } + logger.info( + f"[MultiDataset] Finished norm inference, loading={loading_time:.2f}s, computing={computing_time:.2f}s" + ) + + def _collect_norm_samples( + self, loader, norm_keys, embodiment, n_samples, batch_size, num_workers + ): + collected = {k: [] for k in norm_keys} + cur = 0 + with tqdm(total=n_samples, unit="sample") as pbar: + for batch in loader: + remaining = n_samples - cur + if remaining <= 0: + break + batch_len = None + for value in batch.values(): + if hasattr(value, "shape") and len(value.shape) > 0: + batch_len = int(value.shape[0]) + break + if batch_len is None: + raise ValueError( + "[MultiDataset] Could not infer batch size from DataLoader batch" + ) + take = min(remaining, batch_len) + for k in norm_keys: + zarr_key = self.keyname_to_zarr_key(k, embodiment) + if zarr_key is None or zarr_key not in batch: + continue + x = batch[zarr_key][:take] + if hasattr(x, "detach"): + x = x.detach().cpu().numpy() + collected[k].append(x) + cur += take + pbar.update(take) + return collected + + @staticmethod + def _drop_nonfinite_rows(X: np.ndarray, key: str) -> np.ndarray: + """Drop norm samples (rows of X, shape (N, ...)) containing NaN/Inf. + + np.percentile/mean propagate NaN, and a NaN quantile then poisons every + normalized sample of that key (the per-sample bounds check cannot flag + NaN bounds). Episodes with a non-finite feature (e.g. an all-NaN gripper + command) are still rejected sample-by-sample at train time; they must + not shape the statistics. + """ + finite = np.isfinite(X).reshape(X.shape[0], -1).all(axis=1) + n_bad = int((~finite).sum()) + if n_bad == 0: + return X + if n_bad == X.shape[0]: + raise ValueError( + f"[MultiDataset] key={key}: every collected norm sample is non-finite" + ) + logger.warning( + f"[MultiDataset] key={key}: dropping {n_bad}/{X.shape[0]} norm samples " + "with non-finite values" + ) + return X[finite] + + @staticmethod + def _compute_stats_for_array(X): + return { + "mean": np.mean(X, axis=0), + "std": np.std(X, axis=0), + "min": np.min(X, axis=0), + "max": np.max(X, axis=0), + "median": np.median(X, axis=0), + "quantile_1": np.percentile(X, 1, axis=0), + "quantile_99": np.percentile(X, 99, axis=0), + # Bounds-check quantiles snap to observed samples ("lower" / + # "higher") instead of interpolating: with fewer than 10k samples + # the default linear percentile lands strictly inside the sample + # range, and the extreme frames of the very episodes the stats + # were computed on would be rejected at train time. + "quantile_0_01": np.percentile(X, 0.01, axis=0, method="lower"), + "quantile_99_99": np.percentile(X, 99.99, axis=0, method="higher"), + } + + def cache_stats(self, save_cache_dir: str): + cache_dir = os.path.join(save_cache_dir, "norm_stats") + os.makedirs(cache_dir, exist_ok=True) + out_path = os.path.join(cache_dir, "norm_stats.json") + + stats_out: dict[str, dict[str, dict[str, list]]] = {} + for emb, keys_dict in self.norm_stats.items(): + stats_out[str(emb)] = { + k: {name: np.asarray(arr).tolist() for name, arr in stat_dict.items()} + for k, stat_dict in keys_dict.items() + } + payload = { + "stats": stats_out, + "loading_time": None, + "computing_time": None, + "frames": None, + } + if self._norm_run_metadata is not None: + for k in ("loading_time", "computing_time", "frames"): + if k in self._norm_run_metadata: + payload[k] = self._norm_run_metadata[k] + with open(out_path, "w") as f: + json.dump(payload, f, indent=4) + logger.info(f"[MultiDataset] Cached stats to {out_path}") + + # ---- normalize / unnormalize ---- + + def _apply_norm_one(self, tensor, stats, zarr_key=None): + return _apply_norm_one( + tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) + ) + + def _apply_unnorm_one(self, tensor, stats, zarr_key=None): + return _apply_unnorm_one( + tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) + ) + + def _identity_channels(self, zarr_key, tensor): + """Channels this key leaves unnormalized: the rot6d columns. + + A 6D rotation is two columns of a rotation matrix, already bounded in + [-1, 1], and per-channel scaling and shifting pulls the pair off the + manifold for no gain in resolution -- which is why the two projects + that anchor an action chunk at near-identity, TRI's LBM and UMI, both + exclude rotation from normalization outright. It also removes the t=0 + identity-rotation degeneracy at its source instead of catching it with + NORM_MIN_RANGE. + """ + if zarr_key is None or not self.rot6d_passthrough: + return None + return rot6d_channels(zarr_key, tensor.shape[-1]) + + def normalize(self, data: dict, embodiment_id: int) -> dict: + if not self.norm_stats.get(embodiment_id): + return data + out = dict(data) + for key_name, key_type in self.key_types.get(embodiment_id, {}).items(): + if key_type not in self.NORMALIZE_KEY_TYPES: + continue + stats = self.norm_stats[embodiment_id].get(key_name) + if stats is None: + continue + zarr_key = self.zarr_keys[embodiment_id][key_name] + if zarr_key not in out: + continue + tensor = out[zarr_key] + if not isinstance(tensor, torch.Tensor): + if isinstance(tensor, np.ndarray): + tensor = torch.from_numpy(tensor).float() + else: + continue + out[zarr_key] = self._apply_norm_one(tensor, stats, zarr_key) + return out + + def unnormalize(self, data: dict, embodiment_id: int) -> dict: + if not self.norm_stats.get(embodiment_id): + return data + out = dict(data) + zk_to_kn = {v: k for k, v in self.zarr_keys.get(embodiment_id, {}).items()} + for data_key, value in list(data.items()): + key_name = ( + data_key + if data_key in self.norm_stats[embodiment_id] + else zk_to_kn.get(data_key) + ) + if key_name is None: + continue + stats = self.norm_stats[embodiment_id].get(key_name) + if stats is None: + continue + if not isinstance(value, torch.Tensor): + if isinstance(value, np.ndarray): + value = torch.from_numpy(value).float() + else: + continue + out[data_key] = self._apply_unnorm_one( + value, stats, self.zarr_keys.get(embodiment_id, {}).get(key_name) + ) + return out + + # ---- transform attachment ---- + + def attach_normalize_transforms( + self, datasets: dict | None = None, reject_outliers: bool = True + ) -> None: + """Deprecated. Use ``set_norm_stats_from`` on each training/valid + MultiDataset instead. Bounds-check + normalize now run at the + MultiDataset level in ``__getitem__``, not as per-leaf transforms. + + Kept as a thin shim that calls ``set_norm_stats_from(self)`` on each + MultiDataset in ``datasets`` so existing callers keep working. The + ``reject_outliers`` flag is no longer honored — bounds checking is + always on when stats are populated. To disable, clear ``norm_stats``. + """ + del reject_outliers # unused + graph = datasets if datasets is not None else self.datasets + for ds in graph.values(): + if isinstance(ds, MultiDataset): + ds.set_norm_stats_from(self) + + # ---- serialization (checkpoint roundtrip) ---- + + @staticmethod + def _clone_norm_stats(norm_stats): + out = {} + for emb, per_emb in (norm_stats or {}).items(): + out[emb] = { + key: { + name: ( + v.detach().cpu().clone() + if torch.is_tensor(v) + else copy.deepcopy(v) + ) + for name, v in stats.items() + } + for key, stats in per_emb.items() + } + return out + + def to_state(self) -> dict: + """Serialize stats only (not the dataset graph). Suitable for checkpoint.""" + return { + "norm_mode": self.norm_mode, + "embodiments": sorted(self.embodiments), + "key_types": copy.deepcopy(self.key_types), + "zarr_keys": copy.deepcopy(self.zarr_keys), + "shapes": copy.deepcopy(self.shapes), + "norm_stats": self._clone_norm_stats(self.norm_stats), + "rot6d_passthrough": self.rot6d_passthrough, + } + + @classmethod + def from_state(cls, state: dict) -> "MultiDataset": + """Reconstruct a stats-only MultiDataset (no dataset graph) from state.""" + if state is None: + raise ValueError("MultiDataset state must be provided for reconstruction.") + return cls(state=state) + + def _load_state(self, state: dict) -> None: + """Populate stats fields from a state dict. Used by state-only construction.""" + self.norm_mode = state.get("norm_mode", self.norm_mode) + self.embodiments = set(state.get("embodiments", [])) + self.key_types = copy.deepcopy(state.get("key_types", {})) + self.zarr_keys = copy.deepcopy(state.get("zarr_keys", {})) + self.shapes = copy.deepcopy(state.get("shapes", {})) + self.norm_stats = self._clone_norm_stats(state.get("norm_stats", {})) + # A checkpoint state saved before the flag existed normalized the rot6d + # channels; an empty state (fresh stats, e.g. trainHydra) keeps the default. + if "rot6d_passthrough" in state: + self.rot6d_passthrough = bool(state["rot6d_passthrough"]) + elif state.get("norm_stats"): + self.rot6d_passthrough = False + for emb in self.embodiments: + self.key_types.setdefault(emb, {}) + self.zarr_keys.setdefault(emb, {}) + self.shapes.setdefault(emb, {}) + self.norm_stats.setdefault(emb, {}) + + +# --------------------------------------------------------------------------- +# Per-episode subsampling wrapper around MultiDataset (K evenly-spaced +# frames OR every Sth frame within each episode). +# --------------------------------------------------------------------------- +def _evenly_spaced_indices(n: int, k: int) -> list[int]: + if n <= 0: + return [] + if k >= n: + return list(range(n)) + if k == 1: + return [n // 2] + return [int(round(i * (n - 1) / (k - 1))) for i in range(k)] + + +def _strided_indices(n: int, stride: int) -> list[int]: + if n <= 0 or stride <= 0: + return [] + return list(range(0, n, stride)) + + +class EvenStrideDataset(MultiDataset): + """Wraps a `MultiDataset` to subsample frames per underlying episode. + Pass exactly one of `frames_per_episode` (K evenly-spaced) or + `stride` (every Sth frame). Subclasses `MultiDataset` so trainHydra's + isinstance check passes; we skip the parent `__init__` and adopt its + class identity only, delegating to `self.base`.""" + + def __init__( + self, base, frames_per_episode: int | None = None, stride: int | None = None + ): + if (frames_per_episode is None) == (stride is None): + raise ValueError( + "EvenStrideDataset requires exactly ONE of " + "`frames_per_episode` or `stride` to be set " + f"(got frames_per_episode={frames_per_episode}, stride={stride})." + ) + gibd = getattr(base, "_global_indices_by_dataset", None) + if gibd is None: + raise ValueError( + f"EvenStrideDataset requires a MultiDataset exposing " + f"_global_indices_by_dataset, got {type(base).__name__}" + ) + torch.utils.data.Dataset.__init__(self) + self.base = base + self.frames_per_episode = frames_per_episode + self.stride = stride + self.datasets = base.datasets + + keep: list[int] = [] + for ep_name, idxs in gibd.items(): + n = len(idxs) + if stride is not None: + picks = _strided_indices(n, stride) + rule = f"stride={stride}" + else: + picks = _evenly_spaced_indices(n, frames_per_episode) + rule = f"k={frames_per_episode}" + chosen = [idxs[i] for i in picks] + actual_stride = n / max(1, len(chosen)) + logger.info( + "EvenStrideDataset: %s -> %d / %d frames (%s, actual stride ~%.1f)", + ep_name, + len(chosen), + n, + rule, + actual_stride, + ) + keep.extend(chosen) + self.indices = sorted(keep) + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx): + return self.base[self.indices[idx]] + + def __getitems__(self, idxs: list) -> list: + return self.base.__getitems__([self.indices[i] for i in idxs]) + + def set_data_schematic(self, data_schematic) -> None: + self.base.set_data_schematic(data_schematic) + self.data_schematic = data_schematic + + def set_norm_stats_from(self, source: "MultiDataset") -> None: + """Delegate to ``self.base`` as well: ``__getitem__`` here forwards to + ``self.base[...]``, so the bounds-check + normalize pass reads the + *base*'s stats. Setting them on the wrapper alone (the inherited + implementation) would leave the subsampled val loader unnormalized.""" + self.base.set_norm_stats_from(source) + MultiDataset.set_norm_stats_from(self, source) + + def __getattr__(self, item): + if item == "base": + raise AttributeError(item) + base = self.__dict__.get("base") + if base is None: + raise AttributeError(item) + return getattr(base, item) + + +def pinned_episode_subset( + base: MultiDataset, episode_hashes, dataset_name: str | None = None +) -> MultiDataset: + """A contiguous, unsubsampled ``MultiDataset`` over just ``episode_hashes`` + of an already-built ``base``. + + Used for the video-only val loaders: the overlay video must be watchable, so + it renders whole episodes in frame order rather than the metric loader's + per-episode subsample. + + The episodes are taken from ``base.datasets`` (the same leaf + objects), NOT re-resolved: a second ``MultiDataset._from_resolver`` with + ``filters.episode_hashes`` set would be a different ``resolve_once`` memo key + and so a second path resolution / zarr open per head. A pin that is not in + ``base`` is a ``PinError``, matching ``_from_resolver``'s pin contract. + + ``episode_hashes`` order is preserved, so the video plays the pinned + episodes in config order. + """ + hashes = [str(h) for h in (episode_hashes or [])] + missing = [h for h in hashes if h not in base.datasets] + if missing: + raise PinError( + f"{len(missing)} pinned video episode(s) not in dataset " + f"'{dataset_name or ''}': {sorted(missing)}. The pin must " + "name an episode of THIS head's split (e.g. a seen-val episode for " + "`valid`, a held-out-operator episode for `unseen_op_valid`)." + ) + if not hashes: + raise PinError( + f"pinned_episode_subset called with no episode hashes for " + f"'{dataset_name or ''}'" + ) + subset = MultiDataset( + datasets={h: base.datasets[h] for h in hashes}, + mode="total", + norm_mode=base.norm_mode, + ) + subset.set_norm_stats_from(base) + return subset diff --git a/egomimic/rldb/norm/__init__.py b/egomimic/rldb/norm/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/rldb/zarr/episode_norm_samples.py b/egomimic/rldb/norm/episode_norm_samples.py similarity index 98% rename from egomimic/rldb/zarr/episode_norm_samples.py rename to egomimic/rldb/norm/episode_norm_samples.py index 75a047da56..f2aae3ee61 100644 --- a/egomimic/rldb/zarr/episode_norm_samples.py +++ b/egomimic/rldb/norm/episode_norm_samples.py @@ -26,7 +26,8 @@ import torch from tqdm import tqdm -from egomimic.rldb.zarr import norm_cache, zarr_dataset_multi +from egomimic.rldb.multi import worker_context +from egomimic.rldb.norm import norm_cache from egomimic.utils.gpu_orphans import die_with_parent from egomimic.utils.type_utils import _to_numpy @@ -214,9 +215,7 @@ def _finish(h, idx, rows): collate_fn=_as_list, worker_init_fn=die_with_parent if num_workers > 0 else None, multiprocessing_context=( - zarr_dataset_multi.worker_context(leaves.values()) - if num_workers > 0 - else None + worker_context(leaves.values()) if num_workers > 0 else None ), ) idx_buf: dict[str, list[int]] = {} diff --git a/egomimic/rldb/zarr/norm_cache.py b/egomimic/rldb/norm/norm_cache.py similarity index 95% rename from egomimic/rldb/zarr/norm_cache.py rename to egomimic/rldb/norm/norm_cache.py index 285623ab02..563e4bef77 100644 --- a/egomimic/rldb/zarr/norm_cache.py +++ b/egomimic/rldb/norm/norm_cache.py @@ -44,12 +44,17 @@ # Store-level metadata of a zarr v3 (zarr.json) or v2 group/array. _V3_METADATA = ("zarr.json",) _V2_METADATA = (".zgroup", ".zarray", ".zattrs") -# Holds the episode reader, resolvers and stats methods; hashed for every recipe. -_READER_MODULE = "egomimic.rldb.zarr.zarr_dataset_multi" +# Leaf windowing, the stats methods, and the zarr reader and resolvers; hashed for +# every recipe. +_READER_MODULES = ( + "egomimic.rldb.leaf", + "egomimic.rldb.multi", + "egomimic.rldb.zarr.zarr_dataset_multi", +) def leaf_fingerprint(ds) -> str | None: - """A leaf's own ``cache_fingerprint()`` if it has one (non-zarr storage), else + """A leaf's own ``cache_fingerprint()`` (every ``EpisodeLeaf`` has one), else ``episode_fingerprint`` of its zarr directory.""" own = getattr(ds, "cache_fingerprint", None) if own is not None: @@ -101,10 +106,10 @@ def _is_egomimic(name) -> bool: def code_modules(recipe: Mapping) -> list[str]: - """The reader module, every egomimic module defining a ``_target_`` in the + """The reader modules, every egomimic module defining a ``_target_`` in the recipe, and the egomimic modules those reach through top-level imports.""" targets = {get_object(t).__module__ for t in _targets(recipe)} - todo = {_READER_MODULE, *filter(_is_egomimic, targets)} + todo = {*_READER_MODULES, *filter(_is_egomimic, targets)} seen: set[str] = set() while todo: name = todo.pop() diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/transforms.py similarity index 100% rename from egomimic/rldb/zarr/action_chunk_transforms.py rename to egomimic/rldb/transforms.py diff --git a/egomimic/rldb/zarr/utils.py b/egomimic/rldb/zarr/utils.py index 6f6bb1ca54..b46f754141 100644 --- a/egomimic/rldb/zarr/utils.py +++ b/egomimic/rldb/zarr/utils.py @@ -1,21 +1,4 @@ -import logging -import os -import random - import numpy as np -import torch - -logger = logging.getLogger(__name__) - - -def set_global_seed(seed: int = 42): - random.seed(seed) # Python RNG - np.random.seed(seed) # NumPy RNG - torch.manual_seed(seed) # PyTorch CPU - torch.cuda.manual_seed(seed) # PyTorch GPU - torch.cuda.manual_seed_all(seed) - - os.environ["PYTHONHASHSEED"] = str(seed) def jpeg_pixels_to_float(decoded: np.ndarray) -> np.ndarray: diff --git a/egomimic/rldb/zarr/zarr_dataset_action_expert.py b/egomimic/rldb/zarr/zarr_dataset_action_expert.py index 05402df78d..59b48ab7a8 100644 --- a/egomimic/rldb/zarr/zarr_dataset_action_expert.py +++ b/egomimic/rldb/zarr/zarr_dataset_action_expert.py @@ -18,10 +18,10 @@ import torch from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.multi import MultiDataset from egomimic.rldb.zarr.utils import jpeg_pixels_to_float from egomimic.rldb.zarr.zarr_dataset_multi import ( LocalEpisodeResolver, - MultiDataset, S3EpisodeResolver, ZarrDataset, ZarrEpisode, diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 32c2a756cc..16b3c9c0a6 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -20,45 +20,38 @@ from __future__ import annotations -import concurrent.futures as cf -import copy -import functools import json import logging -import math import os import random import subprocess import tempfile -import time from pathlib import Path -from typing import TYPE_CHECKING, Any, Iterable, Mapping +from typing import Any, Iterable, Mapping import numpy as np import pandas as pd import simplejpeg import torch import zarr -from tqdm import tqdm from egomimic.rldb.embodiment.embodiment import ( - HISTORY_MASK_SUFFIX, canonical_embodiment_name, - get_embodiment_id, ) - -# from action_chunk_transforms import Transform -from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.filters import DatasetFilter, PinError, _ensure_dataset_filter +from egomimic.rldb.leaf import EpisodeLeaf + +# Re-exported for imports from before they moved; new code imports egomimic.rldb.multi. +from egomimic.rldb.multi import ( # noqa: F401 + EvenStrideDataset, + MultiDataset, + pinned_episode_subset, + split_dataset_names, + worker_context, +) from egomimic.rldb.resolve_memo import memoized from egomimic.rldb.zarr.utils import jpeg_pixels_to_float -from egomimic.utils.action_utils import ( - _apply_norm_one, - _apply_unnorm_one, - degenerate_mask, -) from egomimic.utils.env import load_env -from egomimic.utils.gpu_orphans import die_with_parent -from egomimic.utils.pose_utils import bounds_check_channels, rot6d_channels def create_default_engine(): @@ -73,73 +66,15 @@ def episode_table_to_df(engine): return _f(engine) -if TYPE_CHECKING: - # Annotation-only import — avoids a runtime circular import with - # zarr_dataset_action_expert (which itself imports MultiDataset from here). - from egomimic.rldb.zarr.zarr_dataset_action_expert import ( - ZarrActionExpertDataset, - ) - logger = logging.getLogger(__name__) -@functools.cache -def _fetch_pool() -> cf.ThreadPoolExecutor: - """Per-process pool for the concurrent parts of a batched leaf fetch.""" - return cf.ThreadPoolExecutor(max_workers=8, thread_name_prefix="batch-fetch") - - -SEED = 42 - - def _episode_table() -> pd.DataFrame: return memoized( ("episode_table",), lambda: episode_table_to_df(create_default_engine()) ) -def split_dataset_names(dataset_names, valid_ratio=0.2, seed=SEED): - """ - Split a list of dataset names into train/valid sets. - Args: - dataset_names (Iterable[str]) - valid_ratio (float): fraction of datasets to put in valid. - seed (int): for deterministic shuffling. - - - Returns: - train_set (set[str]), valid_set (set[str]) - """ - names = sorted(dataset_names) - if not names: - return set(), set() - - rng = random.Random(seed) - rng.shuffle(names) - - if not (0.0 <= valid_ratio <= 1.0): - raise ValueError(f"valid_ratio must be in [0,1], got {valid_ratio}") - - n_valid = int(len(names) * valid_ratio) - if valid_ratio > 0.0: - n_valid = max(1, n_valid) - - valid = set(names[:n_valid]) - train = set(names[n_valid:]) - return train, valid - - -def _ensure_dataset_filter(filters: DatasetFilter | None) -> DatasetFilter: - if filters is None: - return DatasetFilter() - if isinstance(filters, DatasetFilter): - return filters - raise TypeError( - "filters must be a DatasetFilter or None in the zarr resolver path. " - "Plain dict filters are no longer supported." - ) - - def _is_missing_filter_value(value: object) -> bool: if value is None: return True @@ -170,38 +105,6 @@ def _normalize_filter_row( return normalized -def _infer_key_type(key_name: str) -> str | None: - """ - Heuristically infer a key_type for a post-transform key that wasn't - declared in the leaf's key_map (typically produced by a transform like - ``ConcatKeys``). Returns one of the recognised key_type strings or None - when the key looks like metadata that NormStats shouldn't normalize. - - Recognition is conservative — only positive matches return a type. - """ - name = key_name.lower() - # Camera image keys - if "image" in name or "/img" in name or name.endswith("_img") or ".img" in name: - return "camera_keys" - # Action keys (model outputs) - if name.startswith("actions") or "/cmd_" in name or ".cmd_" in name: - return "action_keys" - # Proprioception (observed state) - if ( - name.startswith("observations.state") - or "/obs_" in name - or ".obs_" in name - or "joint_positions" in name - or "ee_pose" in name - or "keypoints" in name - ): - return "proprio_keys" - # Language / annotations - if "annotation" in name or "tokenized" in name or "lang" in name: - return "annotation_keys" - return None - - def get_fallback_idx( idx: int, candidates: Iterable[int], @@ -229,11 +132,6 @@ def _embodiment_matches(name: object, expected: str) -> bool: return False -class PinError(ValueError): - """A pinned episode_hash cannot be used: missing, deleted, no processed - path, or its embodiment does not match the dataset it was pinned under.""" - - def _check_pins( df: pd.DataFrame, filters: DatasetFilter, expected_embodiment: str | None ) -> None: @@ -498,7 +396,7 @@ def _check_left_wrist_convention(self, datasets: dict) -> None: business. It warns, naming the knob. An episode that matches neither convention (a handful of 2025 ones do) warns too. """ - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + from egomimic.rldb.transforms import RotateLocalFrame from egomimic.scripts.mecka_process.mecka_to_zarr import left_wrist_convention fix_enabled = any( @@ -987,1176 +885,7 @@ def resolve( return self.load(paths) -def worker_context(leaves) -> str | None: - """DataLoader start method for ``leaves``: forkserver if any is not fork-safe - (Lance runs Rust threads; https://lance.org/integrations/pytorch/), else the - default fork.""" - if any( - not getattr(sub, "fork_safe", True) - for leaf in leaves - for sub in MultiDataset._iter_leaves(leaf) - ): - return "forkserver" - return None - - -class _IndexMap: - """Global index -> (dataset name, local index), as offsets rather than one tuple per sample. - - A per-sample list costs ~100 B a sample in every process: tens of GB at 10^7 samples, - copied page by page into forked workers and pickled whole into spawned ones. - """ - - def __init__(self, datasets: dict): - self.names = list(datasets) - self.offsets = np.cumsum( - [0] + [len(d) for d in datasets.values()], dtype=np.int64 - ) - - def __len__(self) -> int: - return int(self.offsets[-1]) - - def __getitem__(self, idx: int) -> tuple[str, int]: - if idx < 0: - idx += len(self) - if not 0 <= idx < len(self): - raise IndexError(idx) - j = int(np.searchsorted(self.offsets, idx, side="right")) - 1 - return self.names[j], int(idx - self.offsets[j]) - - def __iter__(self): - for j, name in enumerate(self.names): - for local in range(int(self.offsets[j + 1] - self.offsets[j])): - yield name, local - - def ranges(self) -> dict[str, range]: - return { - n: range(int(self.offsets[j]), int(self.offsets[j + 1])) - for j, n in enumerate(self.names) - } - - -class MultiDataset(torch.utils.data.Dataset): - """ - Wraps a dict of child datasets (Zarr leaves or other MultiDatasets) and - also owns the normalization-stats descriptor and helpers that used to live - in separate ``NormStats`` / ``NormalizingMultiDataset`` classes. - - Two construction modes: - - **Data mode** (default): pass ``datasets`` to wrap a real dataset graph - (existing behaviour, used during training). Stats fields start empty; - call ``populate_from_datasets()`` and ``infer_norm_from_dataset(...)`` - to fill them in, then ``attach_normalize_transforms()`` to wire - normalize/reject transforms onto each leaf's ``transform_list``. - - **State mode** (``state=...``, ``datasets=None``): construct a - stats-only instance for deploy/eval where the dataset graph isn't - available. ``self.datasets`` is empty; only the stats fields are - populated. Used by checkpoint reconstruction. - """ - - NORMALIZE_KEY_TYPES = ("proprio_keys", "action_keys") - # Saved in the state: a model trained before rot6d passed through - # normalization must keep normalizing it (see _load_state). - rot6d_passthrough = True - - # Bounds-check slack per (timestep, channel) cell; see _check_bounds. - BOUNDS_ABS_TOL = 1e-6 - BOUNDS_REL_SLACK = 0.5 - # Bad-sample retries stay in the failing episode this many times, then - # draw from the whole index space so one bad episode can't exhaust them. - GLOBAL_FALLBACK_ATTEMPTS = 25 - # Consecutive bad samples before raising: a systemic-failure detector. - MAX_FALLBACK_ATTEMPTS = 100 - # Retry warnings: the first RETRY_LOG_FIRST per worker, then every RETRY_LOG_EVERY. - RETRY_LOG_FIRST = 20 - RETRY_LOG_EVERY = 500 - # Episodes tried per key_map/transform group before populate_from_datasets - # falls back to the raw key_map. - POPULATE_PROBES = 8 - _retry_log_count = 0 - - def __init__( - self, - datasets: ( - dict[str, MultiDataset | ZarrDataset | ZarrActionExpertDataset] | None - ) = None, - mode: str = "train", - percent: float = 0.1, - valid_ratio: float = 0.2, - norm_mode: str = "zscore", - state: dict | None = None, - **kwargs, - ): - """ - Args: - datasets: Dict of child datasets. None for state-only construction. - mode: One of "train", "valid", "total", "percent" — which split to keep. - percent: Fraction (when mode="percent"). - valid_ratio: Train/valid split ratio. - norm_mode: One of "zscore", "minmax", "quantile". - state: If provided, populate stats fields from this dict (deploy mode). - """ - super().__init__() - - # ---- Stats fields (always present, may be empty) ---- - self.norm_mode = norm_mode - self.embodiments: set[int] = set() - self.key_types: dict[int, dict[str, str]] = {} - self.zarr_keys: dict[int, dict[str, str]] = {} - self.shapes: dict[int, dict[str, tuple]] = {} - self.norm_stats: dict[int, dict[str, dict[str, np.ndarray]]] = {} - self._norm_run_metadata: dict[str, float | int | None] | None = None - - # ---- Dataset graph fields ---- - self.datasets: dict = {} - self.index_map: list = [] - self._global_indices_by_dataset: dict[str, list[int]] = {} - # Dedup bounds-check warnings: keyed by f"bounds:{episode}:{zarr_key}" - self._warned_violations: set[str] = set() - self.train_collections: set = set() - self.valid_collections: set = set() - - if state is not None: - # Deploy / state-only construction — no dataset graph. - self._load_state(state) - return - - if datasets is None: - raise ValueError("MultiDataset requires either `datasets` or `state`.") - - # ---- Normal data-mode construction ---- - self.train_collections, self.valid_collections = split_dataset_names( - datasets.keys(), valid_ratio=valid_ratio, seed=SEED - ) - - if mode == "train": - chosen = self.train_collections - elif mode == "valid": - chosen = self.valid_collections - elif mode == "total": - chosen = set(datasets.keys()) - elif mode == "percent": - all_names = sorted(datasets.keys()) - rng = random.Random(SEED) - rng.shuffle(all_names) - n_keep = int(len(all_names) * percent) - if percent > 0.0: - n_keep = max(1, n_keep) - chosen = set(all_names[:n_keep]) - else: - raise ValueError(f"Unknown mode: {mode}") - - self.datasets = {rid: ds for rid, ds in datasets.items() if rid in chosen} - assert self.datasets, "No datasets left after applying mode split." - - self.index_map = _IndexMap(self.datasets) - self._global_indices_by_dataset = self.index_map.ranges() - - def __len__(self) -> int: - return len(self.index_map) - - @staticmethod - def _episode_name_for_dataset(dataset, dataset_name: str) -> str: - episode_path = getattr(dataset, "episode_path", None) - if episode_path is None: - return dataset_name - return Path(episode_path).name - - def set_norm_stats_from(self, source: "MultiDataset") -> None: - """Share stats with this dataset (and its nested MultiDatasets) by - reference. After this call ``__getitem__`` will bounds-check + normalize - each sample using ``source``'s ``norm_stats``/``key_types``/``zarr_keys``. - - Use this *instead of* ``attach_normalize_transforms`` — it doesn't mutate - any leaf-level ``transform`` list, so it can't accumulate duplicate - passes when leaves share a transform list reference. - """ - self.norm_stats = source.norm_stats - self.key_types = source.key_types - self.zarr_keys = source.zarr_keys - self.shapes = source.shapes - self.embodiments = source.embodiments - self.norm_mode = source.norm_mode - self.rot6d_passthrough = source.rot6d_passthrough - # Each MultiDataset keeps its own warning-dedup state. - self._warned_violations = set() - for ds in self.datasets.values(): - if isinstance(ds, MultiDataset): - ds.set_norm_stats_from(source) - - def _check_bounds( - self, data: dict, dataset, idx: int, dataset_name: str - ) -> str | None: - """Return a violation message if any tracked key in ``data`` has NaN/Inf - or values outside per-key quantile bounds. ``None`` means the sample - passes. Logs each (episode, key) violation once. - """ - embodiment_id = data.get("embodiment") - if embodiment_id is None: - return None - per_emb_stats = self.norm_stats.get(embodiment_id, {}) - if not per_emb_stats: - return None - - episode_name = self._episode_name_for_dataset(dataset, dataset_name) - - for key_name, stats in per_emb_stats.items(): - zarr_key = self.zarr_keys.get(embodiment_id, {}).get(key_name) - if zarr_key is None or zarr_key not in data: - continue - v = data[zarr_key] - if isinstance(v, torch.Tensor): - arr = v.float() - elif isinstance(v, np.ndarray): - arr = torch.from_numpy(v).float() - else: - continue - - q_low = stats.get( - "quantile_0_01", stats.get("quantile_0_1", stats["quantile_1"]) - ) - q_high = stats.get( - "quantile_99_99", stats.get("quantile_99_9", stats["quantile_99"]) - ) - q_low = torch.as_tensor(q_low, device=arr.device, dtype=torch.float32) - q_high = torch.as_tensor(q_high, device=arr.device, dtype=torch.float32) - try: - q_low = torch.broadcast_to(q_low, arr.shape) - q_high = torch.broadcast_to(q_high, arr.shape) - except RuntimeError: - # Stats were computed for a different layout than this sample - # (e.g. a stale precomputed norm_stats.json). Say so once - # instead of silently disabling the bounds check for the key; - # normalize() raises on the same mismatch anyway. - warn_key = f"bounds-shape:{zarr_key}" - if warn_key not in self._warned_violations: - self._warned_violations.add(warn_key) - logger.warning( - f"[MultiDataset] bounds check skipped for {zarr_key}: " - f"stats shape {tuple(q_low.shape)} does not broadcast to " - f"sample shape {tuple(arr.shape)} (norm stats computed " - "for a different layout?)" - ) - continue - - if torch.any(torch.isnan(arr)) or torch.any(torch.isinf(arr)): - prefix = f"NaN/Inf in {zarr_key} ep={episode_name} frame={idx}" - warn_key = f"naninf:{episode_name}:{zarr_key}" - if warn_key not in self._warned_violations: - self._warned_violations.add(warn_key) - logger.warning(prefix) - return prefix - - # Rotation channels are either Euler ypr (wraps at +-pi) or - # continuous 6D columns; quantile bounds on them are meaningless - # and reject otherwise-valid frames, so only the translation / - # gripper / keypoint channels are bounds-checked. Unrecognized - # widths fall through to a full-vector check; NaN/Inf above still - # covers the full vector. - check_idx = bounds_check_channels(zarr_key, arr.shape[-1]) - if check_idx is not None: - arr_q = arr[..., check_idx] - q_low = q_low[..., check_idx] - q_high = q_high[..., check_idx] - else: - arr_q = arr - - # Per-cell any-cell checks against subsampled stats reject valid - # frames in bulk, so allow BOUNDS_REL_SLACK of each cell's quantile - # range (plus an absolute floor). Corrupt values sit orders of - # magnitude outside and are still caught. - tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low) - below = arr_q < q_low - tol - above = arr_q > q_high + tol - # Cells normalize() treats as constant (wrist-frame t=0, the kp0 / - # kp9 structural zeros) cannot blow up the loss whatever offset - # they carry, so don't reject on them. The exemption follows - # norm_mode's range stat, the same one normalize() uses. - constant = self._degenerate_cells(stats, arr.shape) - if constant is not None: - if check_idx is not None: - constant = constant[..., check_idx] - below = below & ~constant - above = above & ~constant - if torch.any(below) or torch.any(above): - prefix = f"Bounds violation in {zarr_key} ep={episode_name} frame={idx}" - warn_key = f"bounds:{episode_name}:{zarr_key}" - if warn_key not in self._warned_violations: - self._warned_violations.add(warn_key) - n_below = int(below.sum().item()) - n_above = int(above.sum().item()) - logger.warning( - f"{prefix} | n_below={n_below} n_above={n_above} " - f"arr_range=[{arr_q.min().item():.4f}, {arr_q.max().item():.4f}]" - ) - return prefix - return None - - def _degenerate_cells(self, stats: dict, shape) -> torch.Tensor | None: - """Mask of the cells ``normalize()`` treats as constant, or None when - this key's stats do not carry the range ``norm_mode`` reads.""" - try: - mask = degenerate_mask(stats, self.norm_mode) - except (KeyError, ValueError): - return None - return torch.broadcast_to(mask, shape) - - def __getitem__(self, idx, _attempts: int | None = None): - attempts = _attempts - while True: - dataset_name, local_idx = self.index_map[idx] - dataset = self.datasets[dataset_name] - try: - data = dataset[local_idx] - except Exception as e: - reason = ( - f"Sample failed ({type(e).__name__}: {e}) at " - f"{dataset_name}[{local_idx}]" - ) - idx, attempts = self._next_after_failure( - idx, dataset_name, attempts, reason=reason - ) - continue - - # If this leaf is itself a MultiDataset, it already ran bounds + - # normalize for its returned sample. Pass through unchanged. - if isinstance(dataset, MultiDataset): - return data - - violation = self._check_bounds(data, dataset, local_idx, dataset_name) - if violation is not None: - idx, attempts = self._next_after_failure( - idx, dataset_name, attempts, reason=violation - ) - continue - - # Bounds passed: normalize and return. - if self.norm_stats and data.get("embodiment") in self.norm_stats: - data = self.normalize(data, data["embodiment"]) - return data - - def __getitems__(self, idxs: list) -> list: - """A DataLoader batch. Leaves whose class defines ``fetch_planned`` (Lance) do - their reads for the whole batch at once, which is what makes random access from - object storage affordable; every other sample, and any batched one that fails - or trips the bounds check, goes through ``__getitem__`` and its fallback.""" - out = [None] * len(idxs) - groups: dict[type, list] = {} - for pos, idx in enumerate(idxs): - name, local = self.index_map[idx] - leaf = self.datasets[name] - if not isinstance(leaf, MultiDataset) and hasattr(leaf, "fetch_planned"): - groups.setdefault(type(leaf), []).append((pos, name, leaf, local)) - for cls, members in groups.items(): - planned = [] - for pos, name, leaf, local in members: - try: - planned.append((pos, name, leaf, local, leaf.plan_sample(local))) - except Exception: - continue - if not planned: - continue - step = getattr(planned[0][2], "fetch_split", None) or len(planned) - parts = [planned[i : i + step] for i in range(0, len(planned), step)] - - def fetch(part, cls=cls): - return cls.fetch_planned([(m[2], m[4]) for m in part]) - - if len(parts) > 1: - # Parts download concurrently; each is assembled as soon as it lands, - # so CPU work overlaps the remaining reads. - futures = {_fetch_pool().submit(fetch, part): part for part in parts} - done = ((futures[f], f.result) for f in cf.as_completed(futures)) - else: - done = [(parts[0], functools.partial(fetch, parts[0]))] - for part, result in done: - try: - fetched = result() - except Exception: - logger.warning( - "[MultiDataset] batched fetch failed; reading samples one by one", - exc_info=True, - ) - continue - for (pos, name, leaf, local, plan), got in zip(part, fetched): - try: - data = leaf.assemble_sample(plan, got) - except Exception: - continue - if self._check_bounds(data, leaf, local, name) is not None: - continue - if self.norm_stats and data.get("embodiment") in self.norm_stats: - data = self.normalize(data, data["embodiment"]) - out[pos] = data - return [self[idx] if d is None else d for d, idx in zip(out, idxs)] - - def _next_after_failure( - self, idx: int, dataset_name: str, attempts: int | None, *, reason: str - ) -> tuple[int, int]: - """Pick the next index to try after a bad sample at ``idx``. - - Raises once ``MAX_FALLBACK_ATTEMPTS`` consecutive samples failed.""" - attempts = (attempts or 0) + 1 - if attempts >= self.MAX_FALLBACK_ATTEMPTS: - raise RuntimeError( - f"[MultiDataset] {attempts} consecutive bad samples, the last " - f"{attempts - self.GLOBAL_FALLBACK_ATTEMPTS} drawn uniformly " - "from the whole index space: the store, the keymap or the norm " - f"stats are broken, not one episode. Last: {reason}" - ) - local = self._global_indices_by_dataset[dataset_name] - if attempts <= self.GLOBAL_FALLBACK_ATTEMPTS and len(local) > 1: - next_idx = idx - while next_idx == idx: - next_idx = random.choice(local) - else: - next_idx = random.randrange(len(self.index_map)) - next_dataset_name, next_local_idx = self.index_map[next_idx] - self._retry_log_count += 1 - if ( - self._retry_log_count <= self.RETRY_LOG_FIRST - or self._retry_log_count % self.RETRY_LOG_EVERY == 0 - ): - logger.warning( - f"{reason} | attempt {attempts}, " - f"trying {next_dataset_name}[{next_local_idx}] " - f"(retries so far in this worker: {self._retry_log_count})" - ) - return next_idx, attempts - - @classmethod - def _from_resolver( - cls, resolver: EpisodeResolver, dataset_name: str | None = None, **kwargs - ): - """create a MultiDataset from an EpisodeResolver. - - dataset_name: the config key this dataset is built under (trainHydra passes - it). When it names an embodiment, pinned episodes must carry that - embodiment (PinError otherwise) and other mismatches are logged. - - Every pinned episode must also end up in the dataset: a pin that - filter_lambdas or the debug limit exclude, or that fails to load or is - skipped by the resolver, is a PinError rather than a silent drop.""" - sync_from_s3 = kwargs.pop("sync_from_s3", False) - filters = kwargs.pop("filters", None) - expected = None - if dataset_name is not None: - try: - get_embodiment_id(dataset_name) - expected = dataset_name - except (KeyError, AttributeError): - expected = None - - if isinstance(resolver, LocalEpisodeResolver): - resolved = resolver.resolve( - sync_from_s3=sync_from_s3, filters=filters, expected_embodiment=expected - ) - else: - resolved = resolver.resolve(filters=filters, expected_embodiment=expected) - - pins = filters.episode_hashes if isinstance(filters, DatasetFilter) else () - dropped = sorted(set(pins) - set(resolved)) - if dropped: - raise PinError( - f"{len(dropped)} pinned episode(s) not in dataset " - f"'{dataset_name or ''}': {dropped}. They were excluded by " - "filter_lambdas or the resolver's debug limit, or failed to load / " - "were skipped by the resolver (see the log above)." - ) - - return cls(datasets=resolved, **kwargs) - - # ===================================================================== - # Stats / normalization - # ===================================================================== - - @staticmethod - def _iter_leaves(ds): - """Yield non-MultiDataset leaves from possibly nested wrappers.""" - if isinstance(ds, MultiDataset): - for child in ds.datasets.values(): - yield from MultiDataset._iter_leaves(child) - else: - yield ds - - def populate_from_datasets(self, datasets: dict | None = None) -> None: - """ - Populate per-embodiment key inventory by walking leaves and probing - one post-transform sample per (embodiment, key_map, transform list). - ``datasets`` defaults to ``self.datasets`` so the typical call is just - ``mds.populate_from_datasets()``. - """ - graph = datasets if datasets is not None else self.datasets - groups: dict[tuple, list] = {} - for ds in graph.values(): - for leaf in self._iter_leaves(ds): - emb = getattr(leaf, "embodiment", None) - key_map = getattr(leaf, "key_map", None) - if emb is None or key_map is None: - continue - emb_id = emb if isinstance(emb, int) else get_embodiment_id(emb) - # A resolver hands every episode the same key_map and transform - # objects, so one probe covers them all: probing each leaf read a - # full sample per episode (~20 min on 41k mecka episodes). - key = (emb_id, id(key_map), id(getattr(leaf, "transform", None))) - groups.setdefault(key, [emb_id, key_map, []])[2].append(leaf) - - for emb_id, key_map, leaves in groups.values(): - self.embodiments.add(emb_id) - self.key_types.setdefault(emb_id, {}) - self.zarr_keys.setdefault(emb_id, {}) - self.shapes.setdefault(emb_id, {}) - self.norm_stats.setdefault(emb_id, {}) - - sample_keys: set | None = None - for leaf in leaves[: self.POPULATE_PROBES]: - try: - sample = leaf[0] - except Exception as e: - logger.warning( - f"[MultiDataset] Could not probe leaf for post-transform " - f"keys (emb={emb_id}): {e}. Trying the next episode." - ) - continue - if isinstance(sample, dict): - sample_keys = set(sample.keys()) - break - if sample_keys is None: - logger.warning( - f"[MultiDataset] No probe succeeded (emb={emb_id}). " - "Falling back to raw key_map." - ) - for key_name, info in key_map.items(): - self.key_types[emb_id][key_name] = info.get( - "key_type", "metadata_keys" - ) - self.zarr_keys[emb_id][key_name] = info["zarr_key"] - continue - - # Identity zarr_keys map (data_key is the algo-side name). - for data_key in sample_keys: - if data_key in key_map: - info = key_map[data_key] - self.key_types[emb_id][data_key] = info.get( - "key_type", "metadata_keys" - ) - else: - inferred = _infer_key_type(data_key) - if inferred is None: - continue - self.key_types[emb_id][data_key] = inferred - self.zarr_keys[emb_id][data_key] = data_key - - # ---- key lookups ---- - - def keys_of_type(self, key_type: str, embodiment_id: int) -> list[str]: - return [ - k for k, t in self.key_types.get(embodiment_id, {}).items() if t == key_type - ] - - def is_key_with_embodiment(self, key_name: str, embodiment_id: int) -> bool: - return key_name in self.key_types.get(embodiment_id, {}) - - def keyname_to_zarr_key(self, key_name: str, embodiment_id: int) -> str | None: - return self.zarr_keys.get(embodiment_id, {}).get(key_name) - - def zarr_key_to_keyname(self, zarr_key: str, embodiment_id: int) -> str | None: - for k, v in self.zarr_keys.get(embodiment_id, {}).items(): - if v == zarr_key: - return k - return None - - def key_shape(self, key_name: str, embodiment_id: int) -> tuple: - if key_name not in self.shapes.get(embodiment_id, {}): - raise ValueError( - f"Shape for key {key_name!r} on embodiment {embodiment_id} not inferred yet." - ) - return self.shapes[embodiment_id][key_name] - - # ---- shape & norm inference ---- - - def infer_shapes_from_batch(self, batch: dict) -> None: - for emb_id, per_emb in self.zarr_keys.items(): - for key_name, zarr_key in per_emb.items(): - if zarr_key in batch: - val = batch[zarr_key] - if hasattr(val, "shape"): - self.shapes.setdefault(emb_id, {})[key_name] = tuple(val.shape) - elif isinstance(val, int): - self.shapes.setdefault(emb_id, {})[key_name] = (1,) - - def infer_norm_from_dataset( - self, - dataset, - dataset_name, - sample_frac: float = 0.10, - seed: int = 42, - max_samples: int | None = None, - batch_size: int = 512, - num_workers: int = 4, - precomputed_norm_path: str | None = None, - pool_horizon: bool = False, - episode_cache: str | Path | None = None, - episode_fingerprints: Mapping[str, str | None] | None = None, - ): - """Fill ``self.norm_stats[embodiment]`` from a precomputed file or by - sampling ``dataset``. - - ``episode_cache``: a per-episode sample store - (``episode_norm_samples.cache_root``). Rows are read from, and new - episodes written to, it instead of sampling the whole dataset. - ``episode_fingerprints`` saves recomputing the episodes' fingerprints. - - ``pool_horizon``: compute each chunked key's stats over every timestep - of the chunk (one set per channel, like openpi) instead of per - (timestep, channel) cell. The pooled stats are tiled back to the - ``(T, D)`` shape, so normalize / bounds check / checkpoints see the - same layout either way. - """ - embodiment = dataset_name - if isinstance(embodiment, str): - embodiment = get_embodiment_id(embodiment) - - norm_keys = list(self.keys_of_type("proprio_keys", embodiment)) - norm_keys.extend(self.keys_of_type("action_keys", embodiment)) - if not norm_keys: - logger.warning( - f"[MultiDataset] No proprio/action keys for embodiment={embodiment}" - ) - return - - self.norm_stats.setdefault(embodiment, {}) - - if precomputed_norm_path is not None: - if os.path.isdir(precomputed_norm_path): - precomputed_file = os.path.join( - precomputed_norm_path, "norm_stats.json" - ) - elif os.path.isfile(precomputed_norm_path): - precomputed_file = precomputed_norm_path - else: - logger.warning( - f"[MultiDataset] precomputed_norm_path={precomputed_norm_path} is not valid" - ) - return - if os.path.isfile(precomputed_file): - with open(precomputed_file, "r") as f: - payload = json.load(f) - if str(embodiment) not in payload["stats"]: - raise ValueError( - f"norm_stats file {precomputed_file} has no entry for " - f"embodiment id {embodiment} (available: " - f"{sorted(payload['stats'])}). Stats are keyed by numeric " - "EMBODIMENT id, and ids were renumbered by the human/eva " - "embodiment collapse — recompute norm stats instead of " - "reusing a pre-collapse norm_stats.json." - ) - # JSON gives back nested lists; a fresh computation stores - # float32 arrays, so convert to keep both paths identical. - self.norm_stats[embodiment] = { - k: {n: np.asarray(a, dtype=np.float32) for n, a in d.items()} - for k, d in payload["stats"][str(embodiment)].items() - } - self._norm_run_metadata = payload.get("norm_run_metadata", None) - logger.info( - f"[MultiDataset] Loaded precomputed stats for embodiment={embodiment}" - ) - return - - if episode_cache is not None and any( - isinstance(ds, MultiDataset) or getattr(ds, "episode_path", None) is None - for ds in dataset.datasets.values() - ): - logger.warning( - "[MultiDataset] per-episode norm cache needs a flat dataset of " - "zarr episodes; sampling the whole dataset instead" - ) - episode_cache = None - - N = len(dataset) - if N <= 0: - raise ValueError("Dataset is empty") - n_samples = int(math.ceil(sample_frac * N)) - n_samples = max(1, min(n_samples, N)) - if max_samples is not None: - n_samples = min(n_samples, max_samples) - - logger.info(f"[MultiDataset] embodiment={embodiment} norm_keys={norm_keys}") - logger.info( - f"[MultiDataset] sampling {n_samples}/{N} (~{100 * sample_frac:.1f}%)" - ) - - loading_start = time.time() - episode_meta = {} - if episode_cache is not None: - from egomimic.rldb.zarr import episode_norm_samples - - zarr_keys = { - k: z - for k in norm_keys - if (z := self.keyname_to_zarr_key(k, embodiment)) is not None - } - collected, episode_meta = episode_norm_samples.collect( - Path(episode_cache), - dataset.datasets, - zarr_keys, - n_samples, - seed=seed, - num_workers=num_workers, - fingerprints=episode_fingerprints, - ) - for k in norm_keys: - collected.setdefault(k, []) - else: - loader = torch.utils.data.DataLoader( - dataset, - batch_size=batch_size, - num_workers=num_workers, - shuffle=True, - generator=torch.Generator().manual_seed(seed), - worker_init_fn=die_with_parent, - multiprocessing_context=( - worker_context([dataset]) if num_workers > 0 else None - ), - ) - collected = self._collect_norm_samples( - loader, norm_keys, embodiment, n_samples, batch_size, num_workers - ) - for k in [k for k, v in collected.items() if not v]: - del collected[k] - norm_keys.remove(k) - loading_time = time.time() - loading_start - - computing_start = time.time() - for k in norm_keys: - collected[k] = self._drop_nonfinite_rows( - np.concatenate(collected[k], axis=0), k - ) - X = collected[k] - if pool_horizon and X.ndim > 2: - stats_np = self._compute_stats_for_array(X.reshape(-1, X.shape[-1])) - stats_np = { - name: np.broadcast_to(arr, X.shape[1:]).copy() - for name, arr in stats_np.items() - } - else: - stats_np = self._compute_stats_for_array(X) - self.norm_stats[embodiment][k] = { - name: np.asarray(arr, dtype=np.float32) - for name, arr in stats_np.items() - } - logger.info( - f"[MultiDataset] key={k} samples={collected[k].shape[0]} stat_shape={stats_np['mean'].shape}" - ) - computing_time = time.time() - computing_start - - self._norm_run_metadata = { - "loading_time": loading_time, - "computing_time": computing_time, - "frames": n_samples, - **episode_meta, - } - logger.info( - f"[MultiDataset] Finished norm inference, loading={loading_time:.2f}s, computing={computing_time:.2f}s" - ) - - def _collect_norm_samples( - self, loader, norm_keys, embodiment, n_samples, batch_size, num_workers - ): - collected = {k: [] for k in norm_keys} - cur = 0 - with tqdm(total=n_samples, unit="sample") as pbar: - for batch in loader: - remaining = n_samples - cur - if remaining <= 0: - break - batch_len = None - for value in batch.values(): - if hasattr(value, "shape") and len(value.shape) > 0: - batch_len = int(value.shape[0]) - break - if batch_len is None: - raise ValueError( - "[MultiDataset] Could not infer batch size from DataLoader batch" - ) - take = min(remaining, batch_len) - for k in norm_keys: - zarr_key = self.keyname_to_zarr_key(k, embodiment) - if zarr_key is None or zarr_key not in batch: - continue - x = batch[zarr_key][:take] - if hasattr(x, "detach"): - x = x.detach().cpu().numpy() - collected[k].append(x) - cur += take - pbar.update(take) - return collected - - @staticmethod - def _drop_nonfinite_rows(X: np.ndarray, key: str) -> np.ndarray: - """Drop norm samples (rows of X, shape (N, ...)) containing NaN/Inf. - - np.percentile/mean propagate NaN, and a NaN quantile then poisons every - normalized sample of that key (the per-sample bounds check cannot flag - NaN bounds). Episodes with a non-finite feature (e.g. an all-NaN gripper - command) are still rejected sample-by-sample at train time; they must - not shape the statistics. - """ - finite = np.isfinite(X).reshape(X.shape[0], -1).all(axis=1) - n_bad = int((~finite).sum()) - if n_bad == 0: - return X - if n_bad == X.shape[0]: - raise ValueError( - f"[MultiDataset] key={key}: every collected norm sample is non-finite" - ) - logger.warning( - f"[MultiDataset] key={key}: dropping {n_bad}/{X.shape[0]} norm samples " - "with non-finite values" - ) - return X[finite] - - @staticmethod - def _compute_stats_for_array(X): - return { - "mean": np.mean(X, axis=0), - "std": np.std(X, axis=0), - "min": np.min(X, axis=0), - "max": np.max(X, axis=0), - "median": np.median(X, axis=0), - "quantile_1": np.percentile(X, 1, axis=0), - "quantile_99": np.percentile(X, 99, axis=0), - # Bounds-check quantiles snap to observed samples ("lower" / - # "higher") instead of interpolating: with fewer than 10k samples - # the default linear percentile lands strictly inside the sample - # range, and the extreme frames of the very episodes the stats - # were computed on would be rejected at train time. - "quantile_0_01": np.percentile(X, 0.01, axis=0, method="lower"), - "quantile_99_99": np.percentile(X, 99.99, axis=0, method="higher"), - } - - def cache_stats(self, save_cache_dir: str): - cache_dir = os.path.join(save_cache_dir, "norm_stats") - os.makedirs(cache_dir, exist_ok=True) - out_path = os.path.join(cache_dir, "norm_stats.json") - - stats_out: dict[str, dict[str, dict[str, list]]] = {} - for emb, keys_dict in self.norm_stats.items(): - stats_out[str(emb)] = { - k: {name: np.asarray(arr).tolist() for name, arr in stat_dict.items()} - for k, stat_dict in keys_dict.items() - } - payload = { - "stats": stats_out, - "loading_time": None, - "computing_time": None, - "frames": None, - } - if self._norm_run_metadata is not None: - for k in ("loading_time", "computing_time", "frames"): - if k in self._norm_run_metadata: - payload[k] = self._norm_run_metadata[k] - with open(out_path, "w") as f: - json.dump(payload, f, indent=4) - logger.info(f"[MultiDataset] Cached stats to {out_path}") - - # ---- normalize / unnormalize ---- - - def _apply_norm_one(self, tensor, stats, zarr_key=None): - return _apply_norm_one( - tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) - ) - - def _apply_unnorm_one(self, tensor, stats, zarr_key=None): - return _apply_unnorm_one( - tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) - ) - - def _identity_channels(self, zarr_key, tensor): - """Channels this key leaves unnormalized: the rot6d columns. - - A 6D rotation is two columns of a rotation matrix, already bounded in - [-1, 1], and per-channel scaling and shifting pulls the pair off the - manifold for no gain in resolution -- which is why the two projects - that anchor an action chunk at near-identity, TRI's LBM and UMI, both - exclude rotation from normalization outright. It also removes the t=0 - identity-rotation degeneracy at its source instead of catching it with - NORM_MIN_RANGE. - """ - if zarr_key is None or not self.rot6d_passthrough: - return None - return rot6d_channels(zarr_key, tensor.shape[-1]) - - def normalize(self, data: dict, embodiment_id: int) -> dict: - if not self.norm_stats.get(embodiment_id): - return data - out = dict(data) - for key_name, key_type in self.key_types.get(embodiment_id, {}).items(): - if key_type not in self.NORMALIZE_KEY_TYPES: - continue - stats = self.norm_stats[embodiment_id].get(key_name) - if stats is None: - continue - zarr_key = self.zarr_keys[embodiment_id][key_name] - if zarr_key not in out: - continue - tensor = out[zarr_key] - if not isinstance(tensor, torch.Tensor): - if isinstance(tensor, np.ndarray): - tensor = torch.from_numpy(tensor).float() - else: - continue - out[zarr_key] = self._apply_norm_one(tensor, stats, zarr_key) - return out - - def unnormalize(self, data: dict, embodiment_id: int) -> dict: - if not self.norm_stats.get(embodiment_id): - return data - out = dict(data) - zk_to_kn = {v: k for k, v in self.zarr_keys.get(embodiment_id, {}).items()} - for data_key, value in list(data.items()): - key_name = ( - data_key - if data_key in self.norm_stats[embodiment_id] - else zk_to_kn.get(data_key) - ) - if key_name is None: - continue - stats = self.norm_stats[embodiment_id].get(key_name) - if stats is None: - continue - if not isinstance(value, torch.Tensor): - if isinstance(value, np.ndarray): - value = torch.from_numpy(value).float() - else: - continue - out[data_key] = self._apply_unnorm_one( - value, stats, self.zarr_keys.get(embodiment_id, {}).get(key_name) - ) - return out - - # ---- transform attachment ---- - - def attach_normalize_transforms( - self, datasets: dict | None = None, reject_outliers: bool = True - ) -> None: - """Deprecated. Use ``set_norm_stats_from`` on each training/valid - MultiDataset instead. Bounds-check + normalize now run at the - MultiDataset level in ``__getitem__``, not as per-leaf transforms. - - Kept as a thin shim that calls ``set_norm_stats_from(self)`` on each - MultiDataset in ``datasets`` so existing callers keep working. The - ``reject_outliers`` flag is no longer honored — bounds checking is - always on when stats are populated. To disable, clear ``norm_stats``. - """ - del reject_outliers # unused - graph = datasets if datasets is not None else self.datasets - for ds in graph.values(): - if isinstance(ds, MultiDataset): - ds.set_norm_stats_from(self) - - # ---- serialization (checkpoint roundtrip) ---- - - @staticmethod - def _clone_norm_stats(norm_stats): - out = {} - for emb, per_emb in (norm_stats or {}).items(): - out[emb] = { - key: { - name: ( - v.detach().cpu().clone() - if torch.is_tensor(v) - else copy.deepcopy(v) - ) - for name, v in stats.items() - } - for key, stats in per_emb.items() - } - return out - - def to_state(self) -> dict: - """Serialize stats only (not the dataset graph). Suitable for checkpoint.""" - return { - "norm_mode": self.norm_mode, - "embodiments": sorted(self.embodiments), - "key_types": copy.deepcopy(self.key_types), - "zarr_keys": copy.deepcopy(self.zarr_keys), - "shapes": copy.deepcopy(self.shapes), - "norm_stats": self._clone_norm_stats(self.norm_stats), - "rot6d_passthrough": self.rot6d_passthrough, - } - - @classmethod - def from_state(cls, state: dict) -> "MultiDataset": - """Reconstruct a stats-only MultiDataset (no dataset graph) from state.""" - if state is None: - raise ValueError("MultiDataset state must be provided for reconstruction.") - return cls(state=state) - - def _load_state(self, state: dict) -> None: - """Populate stats fields from a state dict. Used by state-only construction.""" - self.norm_mode = state.get("norm_mode", self.norm_mode) - self.embodiments = set(state.get("embodiments", [])) - self.key_types = copy.deepcopy(state.get("key_types", {})) - self.zarr_keys = copy.deepcopy(state.get("zarr_keys", {})) - self.shapes = copy.deepcopy(state.get("shapes", {})) - self.norm_stats = self._clone_norm_stats(state.get("norm_stats", {})) - # A checkpoint state saved before the flag existed normalized the rot6d - # channels; an empty state (fresh stats, e.g. trainHydra) keeps the default. - if "rot6d_passthrough" in state: - self.rot6d_passthrough = bool(state["rot6d_passthrough"]) - elif state.get("norm_stats"): - self.rot6d_passthrough = False - for emb in self.embodiments: - self.key_types.setdefault(emb, {}) - self.zarr_keys.setdefault(emb, {}) - self.shapes.setdefault(emb, {}) - self.norm_stats.setdefault(emb, {}) - - -# --------------------------------------------------------------------------- -# Per-episode subsampling wrapper around MultiDataset (K evenly-spaced -# frames OR every Sth frame within each episode). -# --------------------------------------------------------------------------- -def _evenly_spaced_indices(n: int, k: int) -> list[int]: - if n <= 0: - return [] - if k >= n: - return list(range(n)) - if k == 1: - return [n // 2] - return [int(round(i * (n - 1) / (k - 1))) for i in range(k)] - - -def _strided_indices(n: int, stride: int) -> list[int]: - if n <= 0 or stride <= 0: - return [] - return list(range(0, n, stride)) - - -class EvenStrideDataset(MultiDataset): - """Wraps a `MultiDataset` to subsample frames per underlying episode. - Pass exactly one of `frames_per_episode` (K evenly-spaced) or - `stride` (every Sth frame). Subclasses `MultiDataset` so trainHydra's - isinstance check passes; we skip the parent `__init__` and adopt its - class identity only, delegating to `self.base`.""" - - def __init__( - self, base, frames_per_episode: int | None = None, stride: int | None = None - ): - if (frames_per_episode is None) == (stride is None): - raise ValueError( - "EvenStrideDataset requires exactly ONE of " - "`frames_per_episode` or `stride` to be set " - f"(got frames_per_episode={frames_per_episode}, stride={stride})." - ) - gibd = getattr(base, "_global_indices_by_dataset", None) - if gibd is None: - raise ValueError( - f"EvenStrideDataset requires a MultiDataset exposing " - f"_global_indices_by_dataset, got {type(base).__name__}" - ) - torch.utils.data.Dataset.__init__(self) - self.base = base - self.frames_per_episode = frames_per_episode - self.stride = stride - self.datasets = base.datasets - - keep: list[int] = [] - for ep_name, idxs in gibd.items(): - n = len(idxs) - if stride is not None: - picks = _strided_indices(n, stride) - rule = f"stride={stride}" - else: - picks = _evenly_spaced_indices(n, frames_per_episode) - rule = f"k={frames_per_episode}" - chosen = [idxs[i] for i in picks] - actual_stride = n / max(1, len(chosen)) - logger.info( - "EvenStrideDataset: %s -> %d / %d frames (%s, actual stride ~%.1f)", - ep_name, - len(chosen), - n, - rule, - actual_stride, - ) - keep.extend(chosen) - self.indices = sorted(keep) - - def __len__(self): - return len(self.indices) - - def __getitem__(self, idx): - return self.base[self.indices[idx]] - - def __getitems__(self, idxs: list) -> list: - return self.base.__getitems__([self.indices[i] for i in idxs]) - - def set_data_schematic(self, data_schematic) -> None: - self.base.set_data_schematic(data_schematic) - self.data_schematic = data_schematic - - def set_norm_stats_from(self, source: "MultiDataset") -> None: - """Delegate to ``self.base`` as well: ``__getitem__`` here forwards to - ``self.base[...]``, so the bounds-check + normalize pass reads the - *base*'s stats. Setting them on the wrapper alone (the inherited - implementation) would leave the subsampled val loader unnormalized.""" - self.base.set_norm_stats_from(source) - MultiDataset.set_norm_stats_from(self, source) - - def __getattr__(self, item): - if item == "base": - raise AttributeError(item) - base = self.__dict__.get("base") - if base is None: - raise AttributeError(item) - return getattr(base, item) - - -def pinned_episode_subset( - base: MultiDataset, episode_hashes, dataset_name: str | None = None -) -> MultiDataset: - """A contiguous, unsubsampled ``MultiDataset`` over just ``episode_hashes`` - of an already-built ``base``. - - Used for the video-only val loaders: the overlay video must be watchable, so - it renders whole episodes in frame order rather than the metric loader's - per-episode subsample. - - The episodes are taken from ``base.datasets`` (the same leaf ``ZarrDataset`` - objects), NOT re-resolved: a second ``MultiDataset._from_resolver`` with - ``filters.episode_hashes`` set would be a different ``resolve_once`` memo key - and so a second path resolution / zarr open per head. A pin that is not in - ``base`` is a ``PinError``, matching ``_from_resolver``'s pin contract. - - ``episode_hashes`` order is preserved, so the video plays the pinned - episodes in config order. - """ - hashes = [str(h) for h in (episode_hashes or [])] - missing = [h for h in hashes if h not in base.datasets] - if missing: - raise PinError( - f"{len(missing)} pinned video episode(s) not in dataset " - f"'{dataset_name or ''}': {sorted(missing)}. The pin must " - "name an episode of THIS head's split (e.g. a seen-val episode for " - "`valid`, a held-out-operator episode for `unseen_op_valid`)." - ) - if not hashes: - raise PinError( - f"pinned_episode_subset called with no episode hashes for " - f"'{dataset_name or ''}'" - ) - subset = MultiDataset( - datasets={h: base.datasets[h] for h in hashes}, - mode="total", - norm_mode=base.norm_mode, - ) - subset.set_norm_stats_from(base) - return subset - - -class ZarrDataset(torch.utils.data.Dataset): +class ZarrDataset(EpisodeLeaf): """ Base Zarr Dataset object, Just intializes as pass through to read from zarr episode """ @@ -2179,10 +908,7 @@ def __init__( self._json_keys = None # Lazy-loaded set of JSON-encoded keys self._annotations = None self.init_episode() - - self.key_map = key_map - self.transform = transform_list - super().__init__() + super().__init__(key_map=key_map, transform_list=transform_list) def init_episode(self): """ @@ -2264,189 +990,6 @@ def _annotation_text_for_frame(self, frame_idx: int) -> str: valid_annotations.append(ann.get("text", "")) return valid_annotations - def __len__(self) -> int: - return self.total_frames - - def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: - """End index (exclusive) for a windowed read starting at ``start_idx``. - - Subclasses can override to add per-key-type clamping (e.g., annotation EOS). - """ - return min(start_idx + horizon, self.total_frames) - - def _action_horizon(self) -> int | None: - """Horizon of the keymap's ``action_keys`` entries, if it has any.""" - for spec in self.key_map.values(): - if ( - spec.get("key_type") == "action_keys" - and spec.get("horizon") is not None - ): - return int(spec["horizon"]) - return None - - def _action_pad_mask(self, idx: int) -> np.ndarray | None: - """1.0 per real action frame, 0.0 per repeat-last padded one, (horizon,). - - ``_pad_to_horizon`` pads a short action chunk by repeating the last - frame, which in a wrist-relative action space supervises "stop moving" - at every episode (or annotation) tail. This mask says which steps are - real so ``loss_fn`` can drop the rest; ``None`` when the keymap has no - horizoned action key. - """ - horizon = self._action_horizon() - if horizon is None: - return None - n_real = self._chunk_end_idx(idx, horizon, "action_keys") - idx - mask = np.zeros(horizon, dtype=np.float32) - mask[: max(n_real, 0)] = 1.0 - return mask - - def _fps(self) -> float: - """The episode's frame rate (zarr ``fps`` attr; 30 when absent).""" - return float(self.metadata.get("fps") or 30) - - def _seconds_to_frames(self, seconds: float) -> int: - """``seconds`` in this episode's fps, at least one frame.""" - return max(1, round(float(seconds) * self._fps())) - - def _lagged_idx(self, idx: int, lag_s: float) -> int: - """Frame ``lag_s`` seconds before ``idx`` (at least one frame back), - clamped at the episode start -- so the first samples read frame 0, - the same front padding a history window gets.""" - return max(0, idx - self._seconds_to_frames(lag_s)) - - def _history_stride(self, spec: dict) -> int: - """Frames between history steps: ``history_stride_s`` (seconds, so a - window spans the same time at any fps) wins over ``history_stride``.""" - stride_s = spec.get("history_stride_s") - if stride_s is not None: - return self._seconds_to_frames(stride_s) - return int(spec.get("history_stride", 1) or 1) - - def _history_indices(self, idx: int, spec: dict) -> tuple[np.ndarray, np.ndarray]: - """``(K,)`` frame indices of a ``history`` entry, oldest first, and their - ``(K,)`` float32 mask. - - The newest step is ``idx`` or, with ``lag_s``, that many seconds before - it. Steps before the episode start are masked 0 and repeat the oldest - real step (frame 0 when none is real), so the newest step is always - ``[-1]``. - """ - history = int(spec["history"]) - newest = idx - if spec.get("lag_s") is not None: - newest = idx - self._seconds_to_frames(spec["lag_s"]) - steps = newest - self._history_stride(spec) * np.arange(history - 1, -1, -1) - real = steps >= 0 - fill = steps[real][0] if real.any() else 0 - return np.where(real, steps, fill), real.astype(np.float32) - - @staticmethod - def _pad_to_horizon(array, horizon: int | None): - """Repeat-last out to ``horizon`` a chunk the episode (or annotation) - ended before filling.""" - if horizon is None or not isinstance(array, np.ndarray): - return array - missing = horizon - array.shape[0] - if missing <= 0: - return array - return np.concatenate([array, np.repeat(array[-1:], missing, axis=0)], axis=0) - - def _read_windows( - self, idx: int - ) -> dict[str, tuple[str, tuple[int, int | None] | np.ndarray]]: - """``{keymap key: (zarr key, window)}`` for frame *idx*. - - ``window`` is ``(start, end)``, ``end`` None meaning the single frame - at ``start``, or for a ``history`` entry the ``(K,)`` frame indices of - ``_history_indices``. Annotation keys read nothing and are left out. - """ - windows = {} - for k, spec in self.key_map.items(): - if spec.get("key_type") == "annotation_keys": - continue - horizon = spec.get("horizon") - history = spec.get("history") - lag_s = spec.get("lag_s") - if horizon is not None: - window = ( - idx, - self._chunk_end_idx(idx, horizon, spec.get("key_type")), - ) - elif history is not None and ( - int(history) > 1 or spec.get("key_type") == "camera_keys" - ): - # A camera window keeps its (K, ...) axis even at K = 1. Point - # reads, not a contiguous span: on a JPEG key (one frame per - # chunk) a span would fetch every frame between the steps. - # ZarrEpisode.plan_reads still merges them on numeric keys. - window = self._history_indices(idx, spec)[0] - elif lag_s is not None: - # A single PAST frame (image history): the same zarr array as - # the current-frame entry, read ``lag_s`` seconds back. - window = (self._lagged_idx(idx, lag_s), None) - else: - window = (idx, None) - windows[k] = (spec["zarr_key"], window) - return windows - - def _finalize_sample(self, idx: int, data: dict) -> dict: - """Masks, transforms, tensor cast and per-episode keys of a read sample; - shared with the Lance leaves.""" - pad_mask = self._action_pad_mask(idx) - if pad_mask is not None: - # Travels with the chunk through the interpolators - # (InterpolatePadMask) to the model horizon. - data["action_pad_mask"] = pad_mask - - for k, spec in self.key_map.items(): - if spec.get("key_type") == "camera_keys" and spec.get("history"): - # Unlike proprio, a padded frame has to be masked out of - # attention, so each image window carries its own mask. - data[f"{k}{HISTORY_MASK_SUFFIX}"] = self._history_indices(idx, spec)[1] - - if self.transform: - for transform in self.transform or []: - data = transform.transform(data) - - for k, v in data.items(): - if isinstance(v, np.ndarray): - data[k] = torch.from_numpy(v).to(torch.float32) - - data["embodiment"] = get_embodiment_id(self.embodiment) - # Control frequency of this episode (RDT's ctrl-freq token). - data["fps"] = torch.tensor(self._fps(), dtype=torch.float32) - # Per-episode camera intrinsics travel with the batch so a single - # data-driven embodiment can project without a hardcoded class const. - # Single-camera -> (3,4) K matrix; no/multi-camera -> NaN sentinel, - # which _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". - K = self.intrinsics - if isinstance(K, dict): - # Multi-camera rig: project against the front camera (viz/projection - # target the front image). Prefer a "front"-keyed entry, else the - # first available camera. - K = next( - (v for k, v in K.items() if "front" in str(k).lower()), - next(iter(K.values()), None) if K else None, - ) - if K is not None: - K = np.asarray(K, dtype=np.float32) - if K.shape == (3, 3): - # Normalize a 3x3 K to the canonical 3x4 (zeros last column); - # some contributors store 3x3 (e.g. microagi). - K = np.concatenate([K, np.zeros((3, 1), dtype=np.float32)], axis=1) - if K.shape != ( - 3, - 4, - ): # unexpected -> sentinel (viz falls back to const) - K = np.full((3, 4), np.nan, dtype=np.float32) - else: - K = np.full((3, 4), np.nan, dtype=np.float32) - data["intrinsics"] = torch.from_numpy(np.ascontiguousarray(K)) - ep_name = Path(self.episode_path).name - data["episode_hash"] = ep_name[:-5] if ep_name.endswith(".zarr") else ep_name - return data - def __getitem__( self, idx: int, @@ -2563,8 +1106,9 @@ def _rows(zarr_key: str, start: int, stop: int): data[k] = self._decode_json_entry(data[k]) if retry: continue + _ = origin # preserved for symmetry with prior API - return self._finalize_sample(idx, data) + return self._finish(idx, data) class ZarrAnnotationCutoffDataset(ZarrDataset): diff --git a/egomimic/scripts/abc_process/viz_eva_episode.py b/egomimic/scripts/abc_process/viz_eva_episode.py index 788638738f..0a81910e48 100644 --- a/egomimic/scripts/abc_process/viz_eva_episode.py +++ b/egomimic/scripts/abc_process/viz_eva_episode.py @@ -116,7 +116,7 @@ class AxisOverlay: """Project EE / world triads into front_1 pixels via the pipeline transform.""" def __init__(self, extrinsics_key: str, intrinsics_key: str = "eva"): - from egomimic.rldb.zarr.action_chunk_transforms import ( + from egomimic.rldb.transforms import ( PoseCoordinateFrameTransform, ) from egomimic.utils.pose_utils import _matrix_to_xyzwxyz, _xyzwxyz_to_matrix diff --git a/egomimic/scripts/aria_process/aria_to_lance_ray.py b/egomimic/scripts/aria_process/aria_to_lance_ray.py new file mode 100644 index 0000000000..98614288c2 --- /dev/null +++ b/egomimic/scripts/aria_process/aria_to_lance_ray.py @@ -0,0 +1,222 @@ +"""Convert raw Aria recordings (VRS + MPS) into the Lance tables, one Ray task per recording. + +Workers run aria_to_lance.convert into `//`; the driver is the only +Lance writer. As conversions finish it validates them (lance_pilot/commit.py), appends the +recordings' video chunk rows when --chunk-k is set, then commits them in batches. Reruns +skip recordings already delivered or committed, so a killed run resumes where it stopped. + + python -m egomimic.scripts.aria_process.aria_to_lance_ray \\ + --raw-dir /coc/flash7/scratch/lance_pilot/raw \\ + --deliveries /coc/flash7/scratch/lance_pilot/deliveries \\ + --lance /coc/flash7/scratch/lance_pilot/lance +""" + +import argparse +import json +import logging +import time +import traceback +from pathlib import Path +from types import SimpleNamespace + +import lance +import ray + +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import recording_id +from egomimic.scripts.lance_pilot.chunk_videos import CHUNK_ROWS_PER_FILE, chunk_table +from egomimic.utils.aws.aws_sql import timestamp_ms_to_episode_hash + +log = logging.getLogger("aria_to_lance_ray") + + +def discover(raw_dir: Path) -> list[str]: + """Vendor ids (VRS stems) under `raw_dir` that have both a VRS and an MPS folder.""" + return sorted( + p.stem + for p in raw_dir.glob("*.vrs") + if (raw_dir / f"mps_{p.stem}_vrs").is_dir() + ) + + +def embodiments(vids: list[str]) -> dict[str, str]: + from egomimic.scripts.lance_pilot.commit import sql_rows + + sql = sql_rows([timestamp_ms_to_episode_hash(v) for v in vids]) + out = {} + for v in vids: + h = timestamp_ms_to_episode_hash(v) + if h in sql.index and sql.loc[h].embodiment: + out[v] = str(sql.loc[h].embodiment).lower() + return out + + +@ray.remote(max_retries=1) +def convert_one(vid: str, embodiment: str, opts: dict) -> tuple[str, str | None]: + """(vendor id, None) on success, else (vendor id, traceback).""" + from egomimic.scripts.aria_process.aria_to_lance import convert + + logging.basicConfig(level=logging.INFO) + try: + convert( + SimpleNamespace( + vendor_recording_id=vid, + recording_id=recording_id("aria", vid), + embodiment=embodiment, + **opts, + ) + ) + return vid, None + except Exception: + return vid, traceback.format_exc() + + +@ray.remote +def chunk_one(mp4: str, k: int) -> list[bytes]: + from egomimic.scripts.lance_pilot.chunk_videos import chunk_mp4 + + return chunk_mp4(Path(mp4), k) + + +def chunked_ids(uri: Path) -> set[str]: + if not commit._exists(uri): + return set() + col = lance.dataset(str(uri)).to_table(columns=["recording_id"]).column(0) + return set(col.unique().to_pylist()) + + +def append_chunks( + uri: Path, deliveries: list, camera: str, k: int, redo: bool = False +) -> None: + """Chunk rows for deliveries not yet in the chunk table, each recording's rows + contiguous. `redo` first deletes the rows a delivery already has, which a crash + between chunking and committing leaves behind for a recording now reconverted.""" + have = chunked_ids(uri) + stale = [d.recording_id for d in deliveries if d.recording_id in have] + if redo and stale: + ids = ", ".join(commit._sql_str(r) for r in stale) + lance.dataset(str(uri)).delete(f"recording_id IN ({ids})") + have -= set(stale) + todo = [d for d in deliveries if d.recording_id not in have] + refs = [chunk_one.remote(str(d.videos[camera]), k) for d in todo] + for d, chunks in zip(todo, ray.get(refs), strict=True): + tbl = chunk_table(d.recording_id, chunks, d.recording["num_frames"], k) + uri.parent.mkdir(parents=True, exist_ok=True) + lance.write_dataset( + tbl, + str(uri), + mode="append" if commit._exists(uri) else "create", + max_rows_per_file=CHUNK_ROWS_PER_FILE, + ) + + +def commit_dirs(args, dirs: list[Path], report: dict) -> None: + if not dirs: + return + # Imports torchcodec, which needs the FFmpeg libraries. + from egomimic.scripts.lance_pilot.commit import load_deliveries + + deliveries, rejected = load_deliveries(dirs) + report["rejected"].update(rejected) + # Chunks go first: chunk rows for an uncommitted recording are ignored by readers, + # but a committed recording without chunk rows fails the streaming resolver. + if args.chunk_k: + uri = args.lance / "video_chunks" / f"{args.camera}_k{args.chunk_k}.lance" + append_chunks(uri, deliveries, args.camera, args.chunk_k, args.reconvert) + added = commit.append(args.lance, deliveries) + report["committed"] += added + log.info("committed %d recordings (%d rejected)", len(added), len(rejected)) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--raw-dir", type=Path, required=True) + ap.add_argument("--deliveries", type=Path, required=True) + ap.add_argument("--lance", type=Path, required=True) + ap.add_argument( + "--ray-address", help="existing cluster; default starts a local one" + ) + ap.add_argument("--cpus-per-task", type=float, default=4) + ap.add_argument("--commit-every", type=int, default=32) + ap.add_argument( + "--chunk-k", + type=int, + default=8, + help="frames per video chunk row for streaming reads; 0 skips chunking", + ) + ap.add_argument("--camera", default="front_1") + ap.add_argument( + "--embodiment", + help="use this for every recording instead of looking it up in SQL", + ) + ap.add_argument("--reconvert", action="store_true", help="redo existing deliveries") + ap.add_argument("--limit", type=int, help="convert at most this many recordings") + ap.add_argument("--mano", action=argparse.BooleanOptionalAction, default=True) + ap.add_argument("--mano-iters", type=int, default=2400) + ap.add_argument("--ffmpeg", default="ffmpeg") + ap.add_argument("--tag", help="tag every table's version after the last commit") + args = ap.parse_args() + logging.basicConfig(level=logging.INFO) + if args.chunk_k % 2: + ap.error("--chunk-k must be even: deliveries are encoded with GOP 2") + + have = commit.committed_ids(args.lance) + vids = [v for v in discover(args.raw_dir) if recording_id("aria", v) not in have] + delivered = ( + set() + if args.reconvert + else {v for v in vids if (args.deliveries / v / "meta.json").exists()} + ) + todo = [v for v in vids if v not in delivered][: args.limit] + emb = dict.fromkeys(todo, args.embodiment) if args.embodiment else embodiments(todo) + report = { + "committed": [], + "failed": {v: "no embodiment in SQL" for v in todo if v not in emb}, + "rejected": {}, + } + log.info( + "%d uncommitted recordings: %d to convert, %d already delivered", + len(vids), len(emb), len(delivered), + ) # fmt: skip + + ray.init(address=args.ray_address) + opts = dict( + raw_dir=args.raw_dir, + out_dir=args.deliveries, + mano=args.mano, + mano_iters=args.mano_iters, + ffmpeg=args.ffmpeg, + ) + task = convert_one.options(num_cpus=args.cpus_per_task) + pending = [task.remote(v, e, opts) for v, e in emb.items()] + ready = [args.deliveries / v for v in sorted(delivered)] + t0, done = time.time(), 0 + while pending: + finished, pending = ray.wait(pending, num_returns=1) + vid, err = ray.get(finished[0]) + done += 1 + if err: + report["failed"][vid] = err + log.warning("%s failed:\n%s", vid, err) + else: + ready.append(args.deliveries / vid) + log.info("%d/%d converted, %.0fs", done, len(emb), time.time() - t0) + if len(ready) >= args.commit_every: + commit_dirs(args, ready, report) + ready = [] + commit_dirs(args, ready, report) + + if args.tag and report["committed"]: + report["tag"] = commit.tag(args.lance, args.tag) + out = args.deliveries / f"ray_report_{time.strftime('%Y%m%d_%H%M%S')}.json" + out.write_text(json.dumps(report, indent=1)) + print( + f"committed {len(report['committed'])}, failed {len(report['failed'])}, " + f"rejected {len(report['rejected'])}; report at {out}" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/benchmark/norm_stats.py b/egomimic/scripts/benchmark/norm_stats.py index 5a7fb3930b..254b9c4aa8 100644 --- a/egomimic/scripts/benchmark/norm_stats.py +++ b/egomimic/scripts/benchmark/norm_stats.py @@ -9,10 +9,10 @@ import pandas as pd from omegaconf import DictConfig, OmegaConf -from egomimic.rldb.zarr.utils import set_global_seed from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.pylogger import RankedLogger +from egomimic.utils.seed import set_global_seed from egomimic.utils.utils import extras log = RankedLogger(__name__, rank_zero_only=True) diff --git a/egomimic/scripts/eva_process/check_zero.py b/egomimic/scripts/eva_process/check_zero.py index 7518630e23..ea93feea18 100644 --- a/egomimic/scripts/eva_process/check_zero.py +++ b/egomimic/scripts/eva_process/check_zero.py @@ -7,9 +7,9 @@ from torch.utils.data import DataLoader, default_collate from tqdm import tqdm -from egomimic.rldb.zarr.utils import set_global_seed from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.utils.pylogger import RankedLogger +from egomimic.utils.seed import set_global_seed OmegaConf.register_new_resolver("eval", eval) log = RankedLogger(__name__, rank_zero_only=True) diff --git a/egomimic/scripts/format_bench/bench_update.py b/egomimic/scripts/format_bench/bench_update.py index 2fb578f048..9caa06490d 100644 --- a/egomimic/scripts/format_bench/bench_update.py +++ b/egomimic/scripts/format_bench/bench_update.py @@ -16,7 +16,7 @@ zarr `annotations` array of JSON bytes per episode (ZarrWriter._write_annotations) lerobot `language_persistent` (the episode's spans, broadcast to every frame) and `language_events` columns in the data parquet (LeRobot 0.6 annotation writer) - lance the `annotations` column of tasks.lance, one row per episode + lance annotations.lance, one row per span What each edit writes: zarr objects built in memory and PUT, nothing read back except the group zarr.json on @@ -26,7 +26,7 @@ stats, stats.json and info.json where the edit changes them lance replace: one new data file per fragment (merge_columns + Merge commit), then drop_columns + alter_columns, both metadata only; append: merge_columns + Merge; - annotations: the same column replace on tasks.lance + annotations: overwrite annotations.lance """ import argparse @@ -753,8 +753,8 @@ def steps_uri(root: str) -> str: return f"{root.rstrip('/')}/steps.lance" -def tasks_uri(root: str) -> str: - return f"{root.rstrip('/')}/tasks.lance" +def annotations_uri(root: str) -> str: + return f"{root.rstrip('/')}/annotations.lance" def rec_id(ep: int) -> str: @@ -787,31 +787,29 @@ def hand(s): ) -def lance_annotations(lay: Layout, eps: np.ndarray, k: int) -> pa.Array: - from egomimic.rldb.lance.schema import ANNOTATION +def lance_annotations(lay: Layout, eps, k: int) -> pa.Table: + from egomimic.rldb.lance.schema import ANNOTATIONS - t0 = lay.starts[eps] * DT_NS - return pa.array( - [ - [ - { - "t_start_ns": int(t + s * DT_NS), - "t_end_ns": int(t + e * DT_NS), - "text": txt, - "annotator": "bench", - } - for txt, s, e in lay.spans(int(ep), k) - ] - for ep, t in zip(eps, t0) - ], - pa.list_(ANNOTATION), - ) + rows = [ + { + "annotation_id": f"{rec_id(ep)}-{i}", + "recording_id": rec_id(ep), + "t_start_ns": int((lay.starts[ep] + s) * DT_NS), + "t_end_ns": int((lay.starts[ep] + e) * DT_NS), + "annotation_type": "language", + "text": txt, + "annotation_vendor": "bench", + } + for ep in eps + for i, (txt, s, e) in enumerate(lay.spans(int(ep), k)) + ] + return pa.Table.from_pylist(rows, schema=ANNOTATIONS) def make_lance(root: str, lay: Layout, rows_per_frag: int) -> None: import lance - from egomimic.rldb.lance.schema import HUMAN_STEPS, TASKS + from egomimic.rldb.lance.schema import HUMAN_STEPS opts = lance_opts(root) batches = (steps_batch(lay.episode(ep), rec_id(ep)) for ep in range(lay.n_eps)) @@ -822,22 +820,11 @@ def make_lance(root: str, lay: Layout, rows_per_frag: int) -> None: max_rows_per_file=rows_per_frag, storage_options=opts, ) - eps = np.arange(lay.n_eps) - t0 = lay.starts * DT_NS - tasks = pa.table( - { - "task_instance_id": [f"task-{ep:07d}" for ep in eps], - "recording_id": [rec_id(ep) for ep in eps], - "t_start_ns": t0, - "t_end_ns": t0 + (lay.lens - 1) * DT_NS, - "task_name": ["fold_clothes"] * lay.n_eps, - "task_description": ["folding clothes"] * lay.n_eps, - "objects": pa.array([["shirt"]] * lay.n_eps, pa.list_(pa.string())), - "annotations": lance_annotations(lay, eps, 0), - }, - schema=TASKS, + lance.write_dataset( + lance_annotations(lay, range(lay.n_eps), 0), + annotations_uri(root), + storage_options=opts, ) - lance.write_dataset(tasks, tasks_uri(root), storage_options=opts) def lance_add(uri: str, opts, name: str, read: str, fn, workers: int): @@ -873,15 +860,15 @@ def lance_replace(uri: str, opts, col: str, read: str, fn, workers: int) -> int: def update_lance(root: str, lay: Layout, op: str, k: int, workers: int) -> dict: opts = lance_opts(root) if op == "annotations": + import lance - def anns(rids): - eps = np.array([int(r[4:]) for r in rids.to_pylist()]) - return lance_annotations(lay, eps, k) - - n = lance_replace( - tasks_uri(root), opts, "annotations", "recording_id", anns, workers + ds = lance.write_dataset( + lance_annotations(lay, range(lay.n_eps), k), + annotations_uri(root), + mode="overwrite", + storage_options=opts, ) - return {"objects_written": n} + return {"objects_written": sum(len(f.data_files()) for f in ds.get_fragments())} def hp(ts): return fsl(lay.new_head_pose(ts.to_numpy() // DT_NS, k)) @@ -900,18 +887,17 @@ def hp(ts): def verify_lance(root: str, lay: Layout, op: str, k: int) -> None: import lance - from egomimic.rldb.lance.schema import HUMAN_STEPS, TASKS + from egomimic.rldb.lance.schema import HUMAN_STEPS opts = lance_opts(root) if op == "annotations": - ds = lance.dataset(tasks_uri(root), storage_options=opts) - assert {f.name: f.type for f in ds.schema} == {f.name: f.type for f in TASKS} eps = sorted({0, lay.n_eps // 2, lay.n_eps - 1}) - t = ds.take(eps, columns=["recording_id", "annotations"]) - got_eps = np.array([int(r[4:]) for r in t.column("recording_id").to_pylist()]) - assert ( - t.column("annotations").to_pylist() - == lance_annotations(lay, got_eps, k).to_pylist() + rids = ", ".join(f"'{rec_id(ep)}'" for ep in eps) + got = lance.dataset(annotations_uri(root), storage_options=opts).to_table( + filter=f"recording_id IN ({rids})" + ) + assert got.sort_by("annotation_id").equals( + lance_annotations(lay, eps, k).sort_by("annotation_id") ) return ds = lance.dataset(steps_uri(root), storage_options=opts) diff --git a/egomimic/scripts/format_bench/zarr_fast.py b/egomimic/scripts/format_bench/zarr_fast.py index 4c24aa4a18..979145d0aa 100644 --- a/egomimic/scripts/format_bench/zarr_fast.py +++ b/egomimic/scripts/format_bench/zarr_fast.py @@ -195,6 +195,7 @@ class FastZarrDataset(ZarrDataset): `assemble_sample` hands the prefetched rows to the stock `__getitem__`, so decoding, padding, masks and the JPEG-failure retry (a different index reads afresh) are unchanged.""" + batched_fetch = True fetch_split = None def init_episode(self): diff --git a/egomimic/scripts/lance_pilot/annotations.py b/egomimic/scripts/lance_pilot/annotations.py new file mode 100644 index 0000000000..46fe4da2df --- /dev/null +++ b/egomimic/scripts/lance_pilot/annotations.py @@ -0,0 +1,179 @@ +"""Review and commit annotation change files, and build one from existing zarr annotations. + + review validate a change file against the tables and print the per-recording diff + commit the same, then apply it as one commit on annotations.lance (optionally tagged) + migrate-zarr write a change file from the zarr episodes' annotations, mapped onto the + recording clock through each episode's obs_rgb_timestamps_ns + +See egomimic/rldb/lance/annotations.py for the change file format. +""" + +import argparse +import datetime +import json +import sys +from pathlib import Path + +import lance +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import zarr + +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.annotations import diff, read_changes, review +from egomimic.rldb.lance.schema import STEPS_KEY, steps_path + +CHANGE_SCHEMA = pa.schema( + [ + ("op", pa.string()), + ("recording_id", pa.string()), + ("t_start_ns", pa.int64()), + ("t_end_ns", pa.int64()), + ("annotation_type", pa.string()), + ("text", pa.string()), + ("annotation_vendor", pa.string()), + ] +) + + +def _review(args): + rv = review(args.lance, read_changes(args.file)) + print(diff(args.lance, rv) or "no changes") + for e in rv.errors: + print(f"error: {e}") + return rv + + +def cmd_review(args): + sys.exit(1 if _review(args).errors else 0) + + +def cmd_commit(args): + rv = _review(args) + if rv.errors: + sys.exit("not committed: fix the errors above") + message = f"{args.file}, approved by {args.approver} on {datetime.date.today()}" + counts = commit.apply_annotations(args.lance, rv.adds, rv.deletes, message) + print(f"committed {counts} from {message}") + if args.tag: + print(f"tagged {args.tag}: {commit.tag(args.lance, args.tag)}") + + +def _zarr_annotations(path: Path) -> tuple[list[dict], np.ndarray]: + z = zarr.open_group(str(path), mode="r") + if "annotations" not in z: + return [], np.empty(0, np.int64) + anns = [] + for raw in z["annotations"][:]: + d = json.loads(raw.decode() if isinstance(raw, bytes) else raw) + if isinstance(d, dict): + anns.append(d) + # Arrays are padded to the chunk length; total_frames is the real count. + n = int(z.attrs["total_frames"]) + return anns, np.asarray(z["obs_rgb_timestamps_ns"][:n], np.int64).reshape(-1) + + +def cmd_migrate_zarr(args): + root = args.lance + aliases = lance.dataset(str(root / "aliases.lance")).to_table().to_pylist() + recs = ( + lance.dataset(str(root / "recordings.lance")) + .to_table(columns=["recording_id", "fps", *STEPS_KEY]) + .to_pylist() + ) + fps = {r["recording_id"]: r["fps"] for r in recs} + steps_of = {r["recording_id"]: steps_path(*(r[c] for c in STEPS_KEY)) for r in recs} + tables = {p: lance.dataset(str(root / p)) for p in set(steps_of.values())} + rows, dropped = [], {} + for a in aliases: + h, rid = a["episode_hash"], a["recording_id"] + path = next( + (p for p in (args.zarr / h, args.zarr / f"{h}.zarr") if p.exists()), None + ) + if path is None: + continue + anns, zts = _zarr_annotations(path) + if not anns: + continue + rts = ( + tables[steps_of[rid]] + .to_table(columns=["timestamp_ns"], filter=f"recording_id = '{rid}'") + .column(0) + ) + rts = np.sort(rts.to_numpy()) + period = int(round(1e9 / fps[rid])) + + def frame_of(t) -> int | None: + i = int(np.searchsorted(rts, t)) + return i if i < len(rts) and rts[i] == t else None + + for d in anns: + s, e = int(d.get("start_idx", -1)), int(d.get("end_idx", -1)) + if not 0 <= s < e <= len(zts): + dropped.setdefault(h, []).append(f"bad range [{s}, {e})") + continue + i0 = frame_of(int(zts[s])) + # An annotation running to the clip's end ends at the recording frame after its last. + last = frame_of(int(zts[e if e < len(zts) else -1])) + if i0 is None or last is None: + dropped.setdefault(h, []).append( + f"[{s}, {e}) is off the recording clock" + ) + continue + t0 = int(rts[i0]) + if e < len(zts): + t1 = int(rts[last]) + else: + t1 = ( + int(rts[last + 1]) if last + 1 < len(rts) else int(rts[-1]) + period + ) + rows.append( + dict( + op="add", + recording_id=rid, + t_start_ns=t0, + t_end_ns=t1, + annotation_type="language", + text=str(d.get("text", "")), + annotation_vendor=args.vendor, + ) + ) + pq.write_table(pa.Table.from_pylist(rows, schema=CHANGE_SCHEMA), args.out) + n_dropped = sum(map(len, dropped.values())) + print(f"wrote {len(rows)} annotations to {args.out}; dropped {n_dropped}") + for h, why in sorted(dropped.items()): + print(f" {h}: {len(why)} dropped, e.g. {why[0]}") + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--lance", type=Path, required=True, help="Lance root") + sub = ap.add_subparsers(required=True) + + p = sub.add_parser("review", help="validate a change file and print its diff") + p.add_argument("file", type=Path) + p.set_defaults(fn=cmd_review) + + p = sub.add_parser("commit", help="review, then commit a change file") + p.add_argument("file", type=Path) + p.add_argument("--approver", required=True) + p.add_argument("--tag", help="tag every table's resulting version") + p.set_defaults(fn=cmd_commit) + + p = sub.add_parser("migrate-zarr", help="change file from zarr annotations") + p.add_argument( + "--zarr", type=Path, required=True, help="dir of zarrs" + ) + p.add_argument("--vendor", required=True, help="annotation_vendor for every row") + p.add_argument("--out", type=Path, required=True, help="parquet change file") + p.set_defaults(fn=cmd_migrate_zarr) + + args = ap.parse_args() + args.fn(args) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/chunk_videos.py b/egomimic/scripts/lance_pilot/chunk_videos.py index d4ccbd6fcc..7c9d792476 100644 --- a/egomimic/scripts/lance_pilot/chunk_videos.py +++ b/egomimic/scripts/lance_pilot/chunk_videos.py @@ -30,6 +30,30 @@ ] ) +CHUNK_ROWS_PER_FILE = 1 << 20 + + +def chunk_table( + recording_id: str, chunks: list[bytes], num_frames: int, k: int +) -> pa.Table: + """One recording's chunk rows; raises unless the chunks cover exactly `num_frames`.""" + starts = list(range(0, len(chunks) * k, k)) + sizes = [min(k, num_frames - s) for s in starts] + if sum(sizes) != num_frames: + raise ValueError( + f"{recording_id}: chunks cover {len(chunks)} x {k}, num_frames {num_frames}" + ) + return pa.table( + { + "recording_id": [recording_id] * len(chunks), + "chunk_index": list(range(len(chunks))), + "frame_start": starts, + "num_frames": sizes, + "data": blob_array(chunks), + }, + schema=SCHEMA, + ) + def chunk_mp4(path: Path, k: int) -> list[bytes]: """Stream-copy `path` into consecutive k-frame MP4s.""" @@ -103,11 +127,7 @@ def main(): path = args.lance_root / dict(recs[rid]["video_path"])[args.camera] chunks = chunk_mp4(path, k) nf = recs[rid]["num_frames"] - starts = list(range(0, len(chunks) * k, k)) - if sum(min(k, nf - s) for s in starts) != nf: - raise ValueError( - f"{rid}: chunks cover {len(chunks)} x {k}, num_frames {nf}" - ) + tbl = chunk_table(rid, chunks, nf, k) want = sorted(rng.sample(range(nf), min(args.check_frames, nf))) full, _ = frames(str(path), want) for j, f in enumerate(want): @@ -116,18 +136,11 @@ def main(): raise ValueError( f"{rid}: frame {f} differs between chunk and full MP4" ) - tbl = pa.table( - { - "recording_id": [rid] * len(chunks), - "chunk_index": list(range(len(chunks))), - "frame_start": starts, - "num_frames": [min(k, nf - s) for s in starts], - "data": blob_array(chunks), - }, - schema=SCHEMA, - ) lance.write_dataset( - tbl, uri, mode="append" if i else "overwrite", max_rows_per_file=1 << 20 + tbl, + uri, + mode="append" if i else "overwrite", + max_rows_per_file=CHUNK_ROWS_PER_FILE, ) mp4_bytes += path.stat().st_size chunk_bytes += sum(map(len, chunks)) diff --git a/egomimic/scripts/lance_pilot/commit.py b/egomimic/scripts/lance_pilot/commit.py index 474aa77e63..eee5b233b5 100644 --- a/egomimic/scripts/lance_pilot/commit.py +++ b/egomimic/scripts/lance_pilot/commit.py @@ -7,7 +7,6 @@ import argparse import json -import uuid from pathlib import Path import numpy as np @@ -18,6 +17,7 @@ from torchcodec.decoders import VideoDecoder from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import task_instance_id from egomimic.scripts.lance_pilot import PILOT from egomimic.utils.aws.aws_sql import ( create_default_engine, @@ -64,23 +64,10 @@ def objects_list(v) -> list[str]: return [str(o).strip('"') for o in items if o and str(o) not in ("None", "NULL")] -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument( - "--deliveries", - type=Path, - default=PILOT / "deliveries", - ) - ap.add_argument("--lance", type=Path, default=PILOT / "lance") - ap.add_argument( - "--replace", - action="store_true", - help="replace committed recordings with these deliveries", - ) - ap.add_argument("--tag", help="tag every table's resulting version") - args = ap.parse_args() - - dirs = sorted(p for p in args.deliveries.iterdir() if (p / "meta.json").exists()) +def load_deliveries( + dirs: list[Path], +) -> tuple[list[commit.Delivery], dict[str, list[str]]]: + """Validated deliveries for `dirs` (each holding a meta.json) and {vendor id: errors} for the rest.""" metas = {d: json.loads((d / "meta.json").read_text()) for d in dirs} hashes = [ timestamp_ms_to_episode_hash(m["vendor_recording_id"]) for m in metas.values() @@ -123,14 +110,15 @@ def main(): license=text_or_none(s.license), ) task = dict( - task_instance_id=str(uuid.uuid4()), + task_instance_id=task_instance_id( + rid, text_or_none(s.task), m["t_start_ns"] + ), recording_id=rid, t_start_ns=m["t_start_ns"], t_end_ns=m["t_end_ns"], task_name=text_or_none(s.task), task_description=text_or_none(s.task_description), objects=objects_list(s.objects), - annotations=[], ) alias = dict( episode_hash=h, @@ -146,6 +134,27 @@ def main(): rejected[m["vendor_recording_id"]] = [str(e)] continue deliveries.append(delivery) + return deliveries, rejected + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--deliveries", + type=Path, + default=PILOT / "deliveries", + ) + ap.add_argument("--lance", type=Path, default=PILOT / "lance") + ap.add_argument( + "--replace", + action="store_true", + help="replace committed recordings with these deliveries", + ) + ap.add_argument("--tag", help="tag every table's resulting version") + args = ap.parse_args() + + dirs = sorted(p for p in args.deliveries.iterdir() if (p / "meta.json").exists()) + deliveries, rejected = load_deliveries(dirs) have = commit.committed_ids(args.lance) added = commit.append(args.lance, deliveries) diff --git a/egomimic/scripts/mecka_process/viz_mecka_episode.py b/egomimic/scripts/mecka_process/viz_mecka_episode.py index 1e8c413d28..a20be3dbf6 100644 --- a/egomimic/scripts/mecka_process/viz_mecka_episode.py +++ b/egomimic/scripts/mecka_process/viz_mecka_episode.py @@ -51,7 +51,7 @@ def __init__(self): # Intrinsics live on the embodiment module; imported lazily so --help # stays cheap. from egomimic.rldb.embodiment.human import MECKA_INTRINSICS - from egomimic.rldb.zarr.action_chunk_transforms import ( + from egomimic.rldb.transforms import ( PoseCoordinateFrameTransform, ) from egomimic.utils.pose_utils import _xyzwxyz_to_matrix diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py index 945a37d02f..256dd68888 100644 --- a/egomimic/scripts/precompute_norm_stats.py +++ b/egomimic/scripts/precompute_norm_stats.py @@ -41,7 +41,7 @@ import egomimic from egomimic.rldb.embodiment.embodiment import get_embodiment_id -from egomimic.rldb.zarr import episode_norm_samples, norm_cache +from egomimic.rldb.norm import episode_norm_samples, norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.trainHydra import _instantiate_dataset from egomimic.utils.env import load_env diff --git a/egomimic/scripts/tutorials/aria_to_mano_convert.py b/egomimic/scripts/tutorials/aria_to_mano_convert.py index 930e73b57a..af8d081f22 100644 --- a/egomimic/scripts/tutorials/aria_to_mano_convert.py +++ b/egomimic/scripts/tutorials/aria_to_mano_convert.py @@ -84,9 +84,7 @@ def fetch_episode_loader(): filters = DatasetFilter( filter_lambdas=[f"lambda row: row['episode_hash'] == {episode_hash!r}"] ) - dataset = MultiDataset._from_resolver( - resolver, filters=filters, sync_from_s3=True, mode="total" - ) + dataset = MultiDataset._from_resolver(resolver, filters=filters, mode="total") return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False) diff --git a/egomimic/scripts/tutorials/mecka_keypoints_viz.py b/egomimic/scripts/tutorials/mecka_keypoints_viz.py index 2364e04efc..58cb6ab1da 100644 --- a/egomimic/scripts/tutorials/mecka_keypoints_viz.py +++ b/egomimic/scripts/tutorials/mecka_keypoints_viz.py @@ -44,9 +44,7 @@ def main() -> None: filters = DatasetFilter( filter_lambdas=[f"lambda row: row['episode_hash'] == {episode_hash!r}"] ) - dataset = MultiDataset._from_resolver( - resolver, filters=filters, sync_from_s3=True, mode="total" - ) + dataset = MultiDataset._from_resolver(resolver, filters=filters, mode="total") loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False) frames = [] diff --git a/egomimic/scripts/tutorials/zarr_data_viz.py b/egomimic/scripts/tutorials/zarr_data_viz.py index 9f55b1a3be..fa4e56f3f3 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.py +++ b/egomimic/scripts/tutorials/zarr_data_viz.py @@ -27,9 +27,7 @@ filters = DatasetFilter( filter_lambdas=["lambda row: row['episode_hash'] in {'2026-03-17-01-32-38-000000'}"] ) -multi_ds = MultiDataset._from_resolver( - resolver, filters=filters, sync_from_s3=True, mode="total" -) +multi_ds = MultiDataset._from_resolver(resolver, filters=filters, mode="total") loader = torch.utils.data.DataLoader(multi_ds, batch_size=1, shuffle=False) diff --git a/egomimic/scripts/tutorials/zarr_viz_annotation.py b/egomimic/scripts/tutorials/zarr_viz_annotation.py index f647dcee00..f2a0eb0b3c 100644 --- a/egomimic/scripts/tutorials/zarr_viz_annotation.py +++ b/egomimic/scripts/tutorials/zarr_viz_annotation.py @@ -27,9 +27,7 @@ filter_lambdas=["lambda row: row['episode_hash'] in {'2026-03-15-23-45-35-137000'}"] ) -cloudflare_ds = MultiDataset._from_resolver( - resolver, filters=filters, sync_from_s3=True, mode="total" -) +cloudflare_ds = MultiDataset._from_resolver(resolver, filters=filters, mode="total") loader = torch.utils.data.DataLoader( cloudflare_ds, batch_size=1, shuffle=False, collate_fn=annotation_collate diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 3dde26a78c..42b9ee96e5 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -22,14 +22,9 @@ from egomimic.eval.eval_train_viz import PrefixedEvalVideo from egomimic.pl_utils.pl_model import ModelWrapper from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.filters import PinError +from egomimic.rldb.multi import EvenStrideDataset, MultiDataset, pinned_episode_subset from egomimic.rldb.resolve_memo import resolve_once -from egomimic.rldb.zarr.utils import set_global_seed -from egomimic.rldb.zarr.zarr_dataset_multi import ( - EvenStrideDataset, - MultiDataset, - PinError, - pinned_episode_subset, -) from egomimic.utils.checkpoint_utils import load_checkpoint_weights from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env @@ -37,6 +32,7 @@ from egomimic.utils.instantiators import instantiate_callbacks, instantiate_loggers from egomimic.utils.logging_utils import log_hyperparameters from egomimic.utils.pylogger import RankedLogger +from egomimic.utils.seed import set_global_seed from egomimic.utils.utils import extras, task_wrapper OmegaConf.register_new_resolver("eval", eval) @@ -649,7 +645,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ) norm_stats.populate_from_datasets(datamodule.train_datasets) - from egomimic.rldb.zarr import episode_norm_samples, norm_cache + from egomimic.rldb.norm import episode_norm_samples, norm_cache sample_frac = OmegaConf.select(cfg, "norm_stats.sample_frac", default=1.0) max_samples = OmegaConf.select(cfg, "norm_stats.max_samples", default=None) diff --git a/egomimic/utils/seed.py b/egomimic/utils/seed.py new file mode 100644 index 0000000000..3c55af5ed8 --- /dev/null +++ b/egomimic/utils/seed.py @@ -0,0 +1,18 @@ +import logging +import os +import random + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + + +def set_global_seed(seed: int = 42): + random.seed(seed) # Python RNG + np.random.seed(seed) # NumPy RNG + torch.manual_seed(seed) # PyTorch CPU + torch.cuda.manual_seed(seed) # PyTorch GPU + torch.cuda.manual_seed_all(seed) + + os.environ["PYTHONHASHSEED"] = str(seed) diff --git a/tests/fixtures/lance_tables.py b/tests/fixtures/lance_tables.py index 35053d9500..d65ad4b29c 100644 --- a/tests/fixtures/lance_tables.py +++ b/tests/fixtures/lance_tables.py @@ -1,9 +1,38 @@ """Synthetic Lance steps rows for unit tests.""" +import lance import numpy as np import pyarrow as pa -from egomimic.rldb.lance.schema import HUMAN_STEPS +from egomimic.rldb.lance.schema import ( + ALIASES, + ANNOTATIONS, + HUMAN_STEPS, + RECORDINGS, + TASKS, + annotation_id, + recording_id, + steps_path, +) + +PERIOD_NS = 33_333_333 + + +def annotation(rid, f0, f1, text, vendor="scale", quality=None): + """An ANNOTATIONS row over frames [f0, f1) of a `steps_table` recording.""" + t0, t1 = f0 * PERIOD_NS, f1 * PERIOD_NS + return dict( + annotation_id=annotation_id(rid, vendor, "language", t0, t1, text), + recording_id=rid, + t_start_ns=t0, + t_end_ns=t1, + annotation_type="language", + text=text, + annotation_vendor=vendor, + quality_score=quality, + graded_by=None, + rubric_version=None, + ) def vec(x, n): @@ -29,7 +58,7 @@ def steps_table(rid, n, offset): return pa.table( { "recording_id": pa.array([rid] * n), - "timestamp_ns": pa.array(np.arange(n, dtype=np.int64) * 33_333_333), + "timestamp_ns": pa.array(np.arange(n, dtype=np.int64) * PERIOD_NS), "head_pose": vec(np.c_[t + offset, np.zeros((n, 6))], 7), "left": hand(0.5), "right": hand(0.25), @@ -40,3 +69,85 @@ def steps_table(rid, n, offset): }, schema=HUMAN_STEPS, ) + + +def write_catalog(tmp_path, annotations=lambda recs: []): + """Steps, recordings, tasks, aliases and annotations tables for two aria recordings + under ``tmp_path``; returns their sorted ``(recording_id, episode_hash, num_frames)``. + ``annotations(recs)`` gives the annotation rows.""" + recs = [ + (recording_id("aria", "1"), "2025-01-01-00-00-00-000000", 12), + (recording_id("aria", "2"), "2025-01-01-00-00-01-000000", 8), + ] + recs.sort() + steps = pa.concat_tables( + [steps_table(rid, n, 100 * i) for i, (rid, _, n) in enumerate(recs)] + ) + lance.write_dataset( + steps, tmp_path / steps_path("rl2", "human_bimanual", "human_hand") + ) + k = [200.0, 0, 240, 0, 200.0, 180, 0, 0, 1] + lance.write_dataset( + pa.Table.from_pylist( + [ + dict( + recording_id=rid, + vendor_recording_id=h, + lab="rl2", + embodiment="human_bimanual", + end_effector="human_hand", + rig_name="aria_gen1", + operator="op", + scene="1", + captured_at=None, + fps=30.0, + num_frames=n, + video_path=[("front_1", f"videos/{rid}/front_1.mp4")], + intrinsics=[("front_1", k)], + extrinsics=[], + license="x", + ) + for rid, h, n in recs + ], + schema=RECORDINGS, + ), + tmp_path / "recordings.lance", + ) + lance.write_dataset( + pa.Table.from_pylist( + [ + dict( + task_instance_id=f"t{i}", + recording_id=rid, + t_start_ns=0, + t_end_ns=n * PERIOD_NS, + task_name="fold_clothes", + task_description="folding clothes", + objects=[], + ) + for i, (rid, _, n) in enumerate(recs) + ], + schema=TASKS, + ), + tmp_path / "tasks.lance", + ) + lance.write_dataset( + pa.Table.from_pylist( + [ + dict( + episode_hash=h, + recording_id=rid, + t_start_ns=0, + t_end_ns=n * PERIOD_NS, + ) + for rid, h, n in recs + ], + schema=ALIASES, + ), + tmp_path / "aliases.lance", + ) + lance.write_dataset( + pa.Table.from_pylist(annotations(recs), schema=ANNOTATIONS), + tmp_path / "annotations.lance", + ) + return recs diff --git a/tests/unit/test_action_chunk_transforms.py b/tests/unit/test_action_chunk_transforms.py index 0f3f0ab4aa..dfa7fc5ba6 100644 --- a/tests/unit/test_action_chunk_transforms.py +++ b/tests/unit/test_action_chunk_transforms.py @@ -8,7 +8,7 @@ from egomimic.rldb.embodiment.human import ( _build_human_cartesian_bimanual_transform_list as build_human_bimanual_transform_list, ) -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( ActionChunkCoordinateFrameTransform, ConcatKeys, InterpolatePose, diff --git a/tests/unit/test_aria_to_lance_ray.py b/tests/unit/test_aria_to_lance_ray.py new file mode 100644 index 0000000000..0056befffa --- /dev/null +++ b/tests/unit/test_aria_to_lance_ray.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace + +import pytest + +lance = pytest.importorskip("lance") +pytest.importorskip("ray") + +from egomimic.scripts.aria_process import aria_to_lance_ray as m # noqa: E402 + + +@pytest.fixture +def chunks(monkeypatch): + """{mp4 path: chunk bytes} that the fake chunk_one returns.""" + out = {} + monkeypatch.setattr(m, "chunk_one", SimpleNamespace(remote=lambda mp4, k: out[mp4])) + monkeypatch.setattr(m.ray, "get", list) + return out + + +def _delivery(rid, n): + return SimpleNamespace( + recording_id=rid, recording={"num_frames": n}, videos={"front_1": rid} + ) + + +def _frames(uri): + t = lance.dataset(str(uri)).to_table(columns=["recording_id", "num_frames"]) + out = {} + for r, n in zip(*(c.to_pylist() for c in t.columns)): + out[r] = out.get(r, 0) + n + return out + + +def test_reconvert_replaces_chunk_rows_left_by_a_crashed_run(tmp_path, chunks): + uri = tmp_path / "video_chunks" / "front_1_k2.lance" + chunks.update(a=[b"a0", b"a1"], b=[b"b0"]) + m.append_chunks(uri, [_delivery("a", 4), _delivery("b", 2)], "front_1", 2) + + chunks["a"] = [b"a0", b"a1", b"a2"] + m.append_chunks(uri, [_delivery("a", 6)], "front_1", 2) + assert _frames(uri) == {"a": 4, "b": 2} + + m.append_chunks(uri, [_delivery("a", 6)], "front_1", 2, redo=True) + assert _frames(uri) == {"a": 6, "b": 2} + ids = lance.dataset(str(uri)).to_table(columns=["recording_id"]).column(0) + assert ids.to_pylist() == ["b", "a", "a", "a"] + + +def test_chunks_must_cover_the_recording(tmp_path, chunks): + chunks["a"] = [b"a0"] + with pytest.raises(ValueError, match="num_frames 4"): + m.append_chunks(tmp_path / "c.lance", [_delivery("a", 4)], "front_1", 2) diff --git a/tests/unit/test_episode_norm_samples.py b/tests/unit/test_episode_norm_samples.py index 32a40203d2..c3899f8acd 100644 --- a/tests/unit/test_episode_norm_samples.py +++ b/tests/unit/test_episode_norm_samples.py @@ -6,7 +6,7 @@ from omegaconf import OmegaConf from egomimic.rldb.embodiment.eva import Eva -from egomimic.rldb.zarr import episode_norm_samples as ens +from egomimic.rldb.norm import episode_norm_samples as ens from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset T = 48 diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index 2d805d76f6..850ebc9717 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -8,7 +8,7 @@ from hydra import compose, initialize_config_module from egomimic.rldb.embodiment.eva import Eva -from egomimic.rldb.zarr import norm_cache +from egomimic.rldb.norm import norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset EVA_WRIST_REVERT = "egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list" diff --git a/tests/unit/test_format_bench_update.py b/tests/unit/test_format_bench_update.py index 005a295d06..f00810e419 100644 --- a/tests/unit/test_format_bench_update.py +++ b/tests/unit/test_format_bench_update.py @@ -138,22 +138,20 @@ def test_lance(tmp_path, lay, op): steps = lance.dataset(bu.steps_uri(root)) assert len(steps.get_fragments()) > 1 steps_before = steps.to_table() - tasks_before = lance.dataset(bu.tasks_uri(root)).to_table() + anns_before = lance.dataset(bu.annotations_uri(root)).to_table() files_before = {d.path for f in steps.get_fragments() for d in f.data_files()} bu.update_lance(root, lay, op, k=3, workers=4) bu.verify_lance(root, lay, op, k=3) steps = lance.dataset(bu.steps_uri(root)) steps_after = steps.to_table() - tasks_after = lance.dataset(bu.tasks_uri(root)).to_table() + anns_after = lance.dataset(bu.annotations_uri(root)).to_table() files_after = {d.path for f in steps.get_fragments() for d in f.data_files()} if op == "annotations": assert steps_after.equals(steps_before) assert files_after == files_before - assert tasks_after.drop_columns(["annotations"]).equals( - tasks_before.drop_columns(["annotations"]) - ) + assert anns_after.num_rows == anns_before.num_rows return - assert tasks_after.equals(tasks_before) + assert anns_after.equals(anns_before) # One new file per fragment, and every existing data file is left in place. assert files_before < files_after assert len(files_after) == 2 * len(files_before) diff --git a/tests/unit/test_index_map.py b/tests/unit/test_index_map.py index 013ef5bb48..306e04da1c 100644 --- a/tests/unit/test_index_map.py +++ b/tests/unit/test_index_map.py @@ -1,6 +1,6 @@ import pytest -from egomimic.rldb.zarr.zarr_dataset_multi import _IndexMap +from egomimic.rldb.multi import _IndexMap def test_index_map_matches_per_sample_list(): diff --git a/tests/unit/test_keypoints_6d.py b/tests/unit/test_keypoints_6d.py index c4a523006f..103792683a 100644 --- a/tests/unit/test_keypoints_6d.py +++ b/tests/unit/test_keypoints_6d.py @@ -24,7 +24,7 @@ _build_human_keypoints_revert_6d_transform_list, _build_human_keypoints_revert_6d_wristframe_transform_list, ) -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( KeypointsRot6DToYPR, KeypointsYPRToRot6D, ) diff --git a/tests/unit/test_lance_annotations.py b/tests/unit/test_lance_annotations.py new file mode 100644 index 0000000000..d967aa0659 --- /dev/null +++ b/tests/unit/test_lance_annotations.py @@ -0,0 +1,292 @@ +import pyarrow as pa +import pytest + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import PERIOD_NS, annotation, write_catalog # noqa: E402 + +from egomimic.rldb.lance import commit # noqa: E402 +from egomimic.rldb.lance.annotations import diff, review # noqa: E402 +from egomimic.rldb.lance.dataset import ( # noqa: E402 + LanceEpisodeResolver, + LanceStreamingResolver, +) +from egomimic.rldb.lance.schema import ANNOTATIONS, annotation_id # noqa: E402 + +KEY_MAP = { + "left.action_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.obs_ee_pose", + "horizon": 3, + }, + "obs_head_pose": {"key_type": "proprio_keys", "zarr_key": "obs_head_pose"}, + "annotations": {"key_type": "annotation_keys", "zarr_key": "annotations"}, +} + + +def _rows(recs): + rid = recs[0][0] + return [ + annotation(rid, 0, 6, "reach"), + annotation(rid, 5, 9, "grasp", quality=4), + annotation(rid, 1, 3, "look", vendor="rl2"), + ] + + +@pytest.fixture +def annotated(tmp_path): + return tmp_path, write_catalog(tmp_path, annotations=_rows) + + +def _leaves(root, resolver=LanceEpisodeResolver, **ann): + kw = {"video_chunks": {}} if resolver is LanceStreamingResolver else {} + return resolver( + root, key_map=KEY_MAP, cameras={}, annotations=ann or None, **kw + ).resolve() + + +def _at_frame(leaf, frame): + i = next(i for i in range(len(leaf)) if leaf.anchor(i) == frame) + return leaf[i] + + +def _change(op, **row): + return {"op": op, **row} + + +def _changes(rows): + keys = list(dict.fromkeys(k for r in rows for k in r)) + return pa.Table.from_pylist([{k: r.get(k) for k in keys} for r in rows]) + + +def test_annotation_id_is_content_addressed(): + a = annotation_id("r", "scale", "language", 0, 10, "reach") + assert a == annotation_id("r", "scale", "language", 0, 10, "reach") + assert a != annotation_id("r", "scale", "language", 0, 10, "Reach") + assert a != annotation_id("r", "rl2", "language", 0, 10, "reach") + + +def test_without_annotations_config_the_task_text_stands_in(annotated): + root, recs = annotated + leaf = _leaves(root)[recs[0][1]] + assert leaf[0]["annotations"] == ["folding clothes"] + + +def test_covering_annotations_in_start_order(annotated): + root, recs = annotated + leaf = _leaves(root, types=["language"])[recs[0][1]] + assert _at_frame(leaf, 1)["annotations"] == ["reach", "look"] + assert _at_frame(leaf, 5)["annotations"] == ["reach", "grasp"] + assert _at_frame(leaf, 9)["annotations"] == [] + other = _leaves(root, types=["language"])[recs[1][1]] + assert other[0]["annotations"] == [] + + +def test_vendor_and_quality_filters(annotated): + root, recs = annotated + h = recs[0][1] + assert _at_frame(_leaves(root, vendors=["scale"])[h], 1)["annotations"] == ["reach"] + graded = _leaves(root, min_quality=4)[h] + assert _at_frame(graded, 5)["annotations"] == ["grasp"] + assert _at_frame(graded, 0)["annotations"] == [] + + +def test_clamp_actions_stops_at_the_covering_annotation_end(annotated): + root, recs = annotated + h = recs[0][1] + plain = _leaves(root)[h] + clamped = _leaves(root, clamp_actions=True)[h] + # Frame 5 is covered by "reach" [0, 6) and "grasp" [5, 9): the earlier end wins. + assert _at_frame(plain, 5)["action_pad_mask"].tolist() == [1, 1, 1] + assert _at_frame(clamped, 5)["action_pad_mask"].tolist() == [1, 0, 0] + assert _at_frame(clamped, 0)["action_pad_mask"].tolist() == [1, 1, 1] + assert plain.cache_fingerprint() != clamped.cache_fingerprint() + + +def test_require_keeps_only_annotated_anchors(annotated): + root, recs = annotated + leaves = _leaves(root, require=True) + assert list(leaves) == [recs[0][1]] + # Eligible anchors are 0, 1 and 5.. (frame 4 is invalid); annotations cover [0, 9). + assert [leaves[recs[0][1]].anchor(i) for i in range(6)] == [0, 1, 5, 6, 7, 8] + assert len(leaves[recs[0][1]]) == 6 + + +def test_streaming_batch_gives_the_same_annotations(annotated): + from egomimic.rldb.multi import MultiDataset + + root, _ = annotated + local = MultiDataset(datasets=_leaves(root, clamp_actions=True), mode="total") + stream = MultiDataset( + datasets=_leaves(root, LanceStreamingResolver, clamp_actions=True), + mode="total", + ) + idxs = list(range(len(local))) + for a, b in zip(stream.__getitems__(idxs), [local[i] for i in idxs], strict=True): + assert a["annotations"] == b["annotations"] + assert a["action_pad_mask"].tolist() == b["action_pad_mask"].tolist() + + +def test_unknown_annotation_option_is_rejected(annotated): + root, _ = annotated + with pytest.raises(ValueError, match="unknown annotations options"): + _leaves(root, typo=True) + + +def _add(rid, t0, t1, text="x", **kw): + row = dict( + recording_id=rid, + annotation_vendor="scale", + annotation_type="language", + t_start_ns=t0, + t_end_ns=t1, + text=text, + ) + return _change("add", **{**row, **kw}) + + +def test_review_rejects_each_bad_row(annotated): + root, recs = annotated + rid, n = recs[0][0], recs[0][2] + ok = _add(rid, 0, PERIOD_NS) + bad = [ + _add(rid, 0, PERIOD_NS, annotation_type="subtask"), + _add(rid, PERIOD_NS, PERIOD_NS), + _add(rid, 0, (n + 1) * PERIOD_NS), + _add("no-such-recording", 0, PERIOD_NS), + _add(rid, 0, PERIOD_NS, text="y", quality_score=9), + _add(rid, 0, PERIOD_NS, text="z", annotation_id="wrong"), + ok, + _change("delete", annotation_id="never-committed"), + _change("upsert", annotation_id="x"), + ] + rv = review(root, _changes([ok, *bad])) + assert rv.adds.num_rows == 1 + joined = "\n".join(rv.errors) + for needle in ( + "row 1: annotation_type 'subtask'", + "row 2: t_start_ns", + "row 3: [0, ", + "row 4: recording no-such-recording has no tasks", + "row 5: quality_score 9", + "row 6: annotation_id wrong does not match", + "row 7: ", + "row 8: delete of never-committed", + "row 9: op 'upsert'", + ): + assert needle in joined, needle + assert len(rv.errors) == 9 + + +def test_apply_adds_and_deletes_in_one_version_and_resubmit_is_a_noop(annotated): + root, recs = annotated + rid = recs[0][0] + uri = str(root / "annotations.lance") + old = _rows(recs)[0] + edited = _add(rid, 0, 6 * PERIOD_NS, text="reach for the shirt") + changes = _changes([_change("delete", annotation_id=old["annotation_id"]), edited]) + rv = review(root, changes) + assert not rv.errors + assert "-1" in diff(root, rv) and "reach for the shirt" in diff(root, rv) + v = lance.dataset(uri).version + assert commit.apply_annotations(root, rv.adds, rv.deletes, "approved by me") == { + "added": 1, + "deleted": 1, + } + ds = lance.dataset(uri) + assert ds.version == v + 1 + props = ds.read_transaction(ds.version).transaction_properties + assert "approved by me" in props.values() + texts = set(ds.to_table(columns=["text"]).column(0).to_pylist()) + assert texts == {"reach for the shirt", "grasp", "look"} + + again = review(root, _changes([edited])) + assert again.adds.num_rows == 0 and len(again.unchanged) == 1 + assert commit.apply_annotations(root, again.adds, again.deletes) == { + "added": 0, + "deleted": 0, + } + assert lance.dataset(uri).version == v + 1 + + +def test_grading_through_edit_keeps_the_id(annotated): + root, recs = annotated + aid = _rows(recs)[0]["annotation_id"] + commit.edit( + root, + "annotations.lance", + pa.table( + { + "annotation_id": [aid], + "quality_score": pa.array([3], pa.int8()), + "graded_by": ["grader"], + "rubric_version": ["v1"], + } + ), + on="annotation_id", + ) + rows = lance.dataset(str(root / "annotations.lance")).to_table().to_pylist() + row = next(r for r in rows if r["annotation_id"] == aid) + assert (row["quality_score"], row["graded_by"], row["text"]) == ( + 3, + "grader", + "reach", + ) + assert lance.dataset(str(root / "annotations.lance")).schema == ANNOTATIONS + + +def test_edit_refuses_annotation_content_columns(annotated): + root, recs = annotated + aid = _rows(recs)[0]["annotation_id"] + with pytest.raises(ValueError, match="quality_score"): + commit.edit( + root, + "annotations.lance", + pa.table({"annotation_id": [aid], "text": ["grasp"]}), + on="annotation_id", + ) + + +def test_migrate_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): + import argparse + + import numpy as np + import pyarrow.parquet as pq + + from egomimic.rldb.zarr.zarr_writer import ZarrWriter + from egomimic.scripts.lance_pilot.annotations import cmd_migrate_zarr + + root = tmp_path / "lance" + recs = write_catalog(root) + rid, h, _ = recs[0] + # A clip of recording frames 2..9 with frame 5 dropped, as zarr cut and filtered it. + frames = np.array([2, 3, 4, 6, 7, 8, 9]) + ZarrWriter.create_and_write( + tmp_path / "zarr" / f"{h}.zarr", + numeric_data={ + "obs_rgb_timestamps_ns": (frames * PERIOD_NS).astype(np.int64), + "obs_head_pose": np.zeros((len(frames), 7), np.float32), + }, + embodiment="human_bimanual", + annotations=[("reach", 0, 3), ("fold", 4, 7), ("bad", 3, 2)], + intrinsics={"front_1": np.eye(3, 4)}, + task_name="synthetic", + operator="op_00", + lab="synthetic", + scene="test_scene", + objects=["cube"], + vendor_recording_id="synthetic_session", + ) + out = tmp_path / "changes.parquet" + cmd_migrate_zarr( + argparse.Namespace(lance=root, zarr=tmp_path / "zarr", vendor="scale", out=out) + ) + rows = pq.read_table(out).to_pylist() + spans = {r["text"]: (r["t_start_ns"], r["t_end_ns"]) for r in rows} + # [0, 3) ends at the clip's next kept frame (recording frame 6), so frame 5 is covered. + assert spans["reach"] == (2 * PERIOD_NS, 6 * PERIOD_NS) + # Running to the clip's end, it ends at the recording frame after the clip's last. + assert spans["fold"] == (7 * PERIOD_NS, 10 * PERIOD_NS) + assert "bad" not in spans + rv = review(root, pq.read_table(out)) + assert not rv.errors and rv.adds.num_rows == 2 diff --git a/tests/unit/test_lance_commit.py b/tests/unit/test_lance_commit.py index 21e9de13a7..179d1d44d1 100644 --- a/tests/unit/test_lance_commit.py +++ b/tests/unit/test_lance_commit.py @@ -6,11 +6,16 @@ lance = pytest.importorskip("lance") -from fixtures.lance_tables import steps_table # noqa: E402 +from fixtures.lance_tables import PERIOD_NS, steps_table # noqa: E402 from egomimic.rldb.lance import commit # noqa: E402 from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 -from egomimic.rldb.lance.schema import recording_id, steps_path # noqa: E402 +from egomimic.rldb.lance.schema import ( # noqa: E402 + TASKS, + recording_id, + steps_path, + task_instance_id, +) STEPS = steps_path("rl2", "human_bimanual", "human_hand") @@ -49,13 +54,14 @@ def _delivery(tmp, i, n, offset=None, video=b"mp4", **rec): task_instance_id=f"t{i}", recording_id=rid, t_start_ns=0, - t_end_ns=n, + t_end_ns=n * PERIOD_NS, task_name="fold", task_description=f"task {i}", objects=[], - annotations=[], ), - alias=dict(episode_hash=h, recording_id=rid, t_start_ns=0, t_end_ns=n), + alias=dict( + episode_hash=h, recording_id=rid, t_start_ns=0, t_end_ns=n * PERIOD_NS + ), steps=steps_table(rid, n, 100 * i if offset is None else offset), videos={"front_1": src / "front_1.mp4"}, ) @@ -246,3 +252,42 @@ def crash_on_recordings(uri, table): assert lance.dataset(root / "aliases.lance").count_rows() == 3 leaves = _resolve(root) assert sorted(len(leaf) for leaf in leaves.values()) == [8 - 3, 9 - 3, 10 - 3] + + +def test_task_instance_id_is_stable_and_keyed_on_name_and_start(): + rid = recording_id("aria", "1") + a = task_instance_id(rid, "fold_clothes", 100) + assert a == task_instance_id(rid, "fold_clothes", 100) + assert a != task_instance_id(rid, "fold_clothes", 101) + assert a != task_instance_id(rid, "iron_clothes", 100) + assert a != task_instance_id(recording_id("aria", "2"), "fold_clothes", 100) + + +def test_release_tag_covers_annotations(tmp_path): + root = tmp_path / "lance" + commit.append(root, [_delivery(tmp_path, 0, 10)]) + tagged = commit.tag(root, "release-1") + assert "annotations.lance" in tagged + resolver = LanceEpisodeResolver( + root, key_map=KEY_MAP, cameras={}, tag="release-1", annotations={} + ) + (leaf,) = resolver.resolve().values() + assert leaf[0]["annotations"] == [] + + +def test_drop_task_annotations_brings_old_tables_to_the_schema(tmp_path): + root = tmp_path / "lance" + first, second = _delivery(tmp_path, 0, 10), _delivery(tmp_path, 1, 8) + commit.append(root, [first]) + tasks = root / "tasks.lance" + old = lance.dataset(tasks).to_table() + nested = pa.list_(pa.struct([("text", pa.string())])) + lance.write_dataset( + old.append_column("annotations", pa.array([[]] * old.num_rows, nested)), + tasks, + mode="overwrite", + ) + commit.drop_task_annotations(root) + assert lance.dataset(tasks).schema.names == TASKS.names + assert commit.append(root, [second]) == [second.recording_id] + commit.drop_task_annotations(root) diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py index 19bdbc019b..08554fd8a7 100644 --- a/tests/unit/test_lance_dataset.py +++ b/tests/unit/test_lance_dataset.py @@ -4,7 +4,7 @@ lance = pytest.importorskip("lance") -from fixtures.lance_tables import steps_table # noqa: E402 +from fixtures.lance_tables import write_catalog # noqa: E402 from egomimic.rldb.filters import DatasetFilter # noqa: E402 from egomimic.rldb.lance.dataset import ( # noqa: E402 @@ -15,13 +15,6 @@ _runs, eligible_intervals, ) -from egomimic.rldb.lance.schema import ( # noqa: E402 - ALIASES, - RECORDINGS, - TASKS, - recording_id, - steps_path, -) def test_eligible_intervals_requires_whole_chunk_valid(): @@ -90,74 +83,7 @@ def test_runs_merge_across_chunks_and_keep_repeats(): @pytest.fixture def catalog(tmp_path): - recs = [ - (recording_id("aria", "1"), "2025-01-01-00-00-00-000000", 12), - (recording_id("aria", "2"), "2025-01-01-00-00-01-000000", 8), - ] - recs.sort() - steps = pa.concat_tables( - [steps_table(rid, n, 100 * i) for i, (rid, _, n) in enumerate(recs)] - ) - lance.write_dataset( - steps, tmp_path / steps_path("rl2", "human_bimanual", "human_hand") - ) - k = [200.0, 0, 240, 0, 200.0, 180, 0, 0, 1] - lance.write_dataset( - pa.Table.from_pylist( - [ - dict( - recording_id=rid, - vendor_recording_id=h, - lab="rl2", - embodiment="human_bimanual", - end_effector="human_hand", - rig_name="aria_gen1", - operator="op", - scene="1", - captured_at=None, - fps=30.0, - num_frames=n, - video_path=[("front_1", f"videos/{rid}/front_1.mp4")], - intrinsics=[("front_1", k)], - extrinsics=[], - license="x", - ) - for rid, h, n in recs - ], - schema=RECORDINGS, - ), - tmp_path / "recordings.lance", - ) - lance.write_dataset( - pa.Table.from_pylist( - [ - dict( - task_instance_id=f"t{i}", - recording_id=rid, - t_start_ns=0, - t_end_ns=n, - task_name="fold_clothes", - task_description="folding clothes", - objects=[], - annotations=[], - ) - for i, (rid, _, n) in enumerate(recs) - ], - schema=TASKS, - ), - tmp_path / "tasks.lance", - ) - lance.write_dataset( - pa.Table.from_pylist( - [ - dict(episode_hash=h, recording_id=rid, t_start_ns=0, t_end_ns=n) - for rid, h, n in recs - ], - schema=ALIASES, - ), - tmp_path / "aliases.lance", - ) - return tmp_path, recs + return tmp_path, write_catalog(tmp_path) KEY_MAP = { diff --git a/tests/unit/test_leaf_equivalence.py b/tests/unit/test_leaf_equivalence.py new file mode 100644 index 0000000000..e4279f7c73 --- /dev/null +++ b/tests/unit/test_leaf_equivalence.py @@ -0,0 +1,91 @@ +"""The same recording read as a zarr leaf and as a Lance leaf gives the same samples.""" + +import numpy as np +import pytest +import torch + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import write_catalog # noqa: E402 + +from egomimic.rldb.lance.dataset import ( # noqa: E402 + ZARR_TO_LANCE, + LanceEpisodeResolver, + _column, +) +from egomimic.rldb.leaf import EpisodeLeaf # noqa: E402 +from egomimic.rldb.zarr.zarr_dataset_multi import ZarrDataset # noqa: E402 +from egomimic.rldb.zarr.zarr_writer import ZarrWriter # noqa: E402 + +KEY_MAP = { + "left.action_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.obs_ee_pose", + "horizon": 3, + }, + "right.obs_ee_pose": { + "key_type": "proprio_keys", + "zarr_key": "right.obs_ee_pose", + "history": 3, + "history_stride": 2, + }, + "obs_head_pose": {"key_type": "proprio_keys", "zarr_key": "obs_head_pose"}, + "annotations": {"key_type": "annotation_keys", "zarr_key": "annotations"}, +} +TASK = "folding clothes" + + +def _zarr_twin(root, leaf) -> ZarrDataset: + """A zarr episode holding the Lance leaf's steps rows, intrinsics and task.""" + steps = lance.dataset(leaf.steps_uri, version=leaf.steps_version).to_table( + offset=leaf.row_start, limit=leaf.total_frames + ) + keys = {s["zarr_key"] for s in KEY_MAP.values()} - {"annotations"} + numeric = {zk: _column(steps, ZARR_TO_LANCE[zk]) for zk in keys} + path = ZarrWriter.create_and_write( + root / f"{leaf.episode_path}.zarr", + numeric_data=numeric, + embodiment=leaf.embodiment, + fps=int(leaf.metadata["fps"]), + annotations=[(TASK, 0, leaf.total_frames)], + intrinsics={ + "front_1": np.c_[np.asarray(leaf.intrinsics).reshape(3, 3), np.zeros(3)] + }, + task_name="synthetic", + operator="op_00", + lab="synthetic", + scene="test_scene", + objects=["cube"], + vendor_recording_id="synthetic_session", + ) + return ZarrDataset(path, key_map=KEY_MAP) + + +def test_zarr_and_lance_leaves_agree(tmp_path): + lance_root = tmp_path / "lance" + write_catalog(lance_root) + leaves = LanceEpisodeResolver(lance_root, key_map=KEY_MAP, cameras={}).resolve() + assert leaves + for leaf in leaves.values(): + twin = _zarr_twin(tmp_path, leaf) + assert isinstance(leaf, EpisodeLeaf) and isinstance(twin, EpisodeLeaf) + for i in range(len(leaf)): + a, b = leaf[i], twin[leaf.anchor(i)] + assert a.keys() == b.keys() + for k in a: + if isinstance(a[k], torch.Tensor): + torch.testing.assert_close(a[k], b[k], equal_nan=True, msg=k) + else: + assert a[k] == b[k], k + + +def test_only_streaming_leaves_batch_fetch(): + from egomimic.rldb.lance.dataset import ( + LanceRecordingDataset, + LanceStreamingDataset, + ) + + assert not EpisodeLeaf.batched_fetch + assert not ZarrDataset.batched_fetch + assert not LanceRecordingDataset.batched_fetch + assert LanceStreamingDataset.batched_fetch diff --git a/tests/unit/test_loss_masking.py b/tests/unit/test_loss_masking.py index 055471b16c..e98c202fd6 100644 --- a/tests/unit/test_loss_masking.py +++ b/tests/unit/test_loss_masking.py @@ -21,7 +21,7 @@ from egomimic.models.denoising_policy import DenoisingPolicy from egomimic.rldb.embodiment.embodiment import get_embodiment_id from egomimic.rldb.embodiment.human import Human -from egomimic.rldb.zarr.action_chunk_transforms import InterpolatePadMask +from egomimic.rldb.transforms import InterpolatePadMask from egomimic.rldb.zarr.zarr_dataset_multi import ( LocalAnnotationCutoffEpisodeResolver, LocalEpisodeResolver, diff --git a/tests/unit/test_norm_cache.py b/tests/unit/test_norm_cache.py index d8050fbd3c..43815d0f9d 100644 --- a/tests/unit/test_norm_cache.py +++ b/tests/unit/test_norm_cache.py @@ -17,7 +17,7 @@ from omegaconf import OmegaConf import egomimic.trainHydra as train_hydra -from egomimic.rldb.zarr import episode_norm_samples, norm_cache +from egomimic.rldb.norm import episode_norm_samples, norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.rldb.zarr.zarr_writer import write_annotations @@ -102,11 +102,12 @@ def test_key_changes_with_code(monkeypatch) -> None: def test_code_modules_cover_reader_targets_and_their_imports() -> None: mods = norm_cache.code_modules(OmegaConf.to_container(CFG)) for m in ( - "egomimic.rldb.zarr.zarr_dataset_multi", # reader, resolvers, stats methods + "egomimic.rldb.multi", # stats methods + "egomimic.rldb.zarr.zarr_dataset_multi", # reader, resolvers "egomimic.rldb.embodiment.human", # keymap + transform factory (targets) "egomimic.rldb.filters", - "egomimic.rldb.zarr.action_chunk_transforms", # imported by human - "egomimic.utils.pose_utils", # imported by action_chunk_transforms + "egomimic.rldb.transforms", # imported by human + "egomimic.utils.pose_utils", # imported by transforms ): assert m in mods assert all(m.startswith("egomimic.") for m in mods) diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 92e044e5b9..8bd28c66ff 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -14,7 +14,7 @@ import pytest import torch -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( CartesianRot6DToYPR, CartesianYPRToRot6D, ) @@ -223,7 +223,7 @@ def test_base_converter_rejects_norm_6d_encoding(): def test_unpad_gripper_zeros_inverts_pad_and_noops_unpadded(): - from egomimic.rldb.zarr.action_chunk_transforms import ( + from egomimic.rldb.transforms import ( PadGripperZeros, UnpadGripperZeros, ) @@ -245,7 +245,7 @@ def test_human_6d_reverts_unpad_proprio(): _build_human_cartesian_revert_6d_transform_list, _build_human_cartesian_revert_6d_wristframe_transform_list, ) - from egomimic.rldb.zarr.action_chunk_transforms import UnpadGripperZeros + from egomimic.rldb.transforms import UnpadGripperZeros for build in ( _build_human_cartesian_revert_6d_transform_list, @@ -260,7 +260,7 @@ def test_rotate_local_frame_flips_left_wrist_convention(): # handle both (7,) poses and (T, 7) chunks. from scipy.spatial.transform import Rotation as R - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + from egomimic.rldb.transforms import RotateLocalFrame rng = np.random.default_rng(4) q = R.random(3, random_state=5) @@ -308,7 +308,7 @@ def test_keypoint_converter_is_identity_and_pads_to_model_width(): def test_fix_left_wrist_convention_flag_prepends_correction(): from egomimic.rldb.embodiment.human import Human - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + from egomimic.rldb.transforms import RotateLocalFrame tl = Human.get_transform_list( "cartesian_wristframe_6d", stride=1, fix_left_wrist_convention=True diff --git a/tests/unit/test_pins.py b/tests/unit/test_pins.py index 2e65cf4ec3..7a83db7c51 100644 --- a/tests/unit/test_pins.py +++ b/tests/unit/test_pins.py @@ -179,6 +179,22 @@ def test_from_resolver_skips_check_for_non_embodiment_names(tmp_path) -> None: assert set(ds.datasets) == {"eva_00"} +def test_from_resolver_passes_sync_from_s3_only_where_resolve_takes_it( + tmp_path, +) -> None: + write_episode(tmp_path, "eva", seed=0) + + class NoSyncFlag(LocalEpisodeResolver): + def resolve(self, filters=None, expected_embodiment=None): + return super().resolve( + filters=filters, expected_embodiment=expected_embodiment + ) + + for resolver in (NoSyncFlag(tmp_path), LocalEpisodeResolver(tmp_path)): + ds = MultiDataset._from_resolver(resolver, sync_from_s3=True, mode="total") + assert set(ds.datasets) == {"eva_00"} + + def test_train_fails_fast_on_bad_pin(tmp_path, monkeypatch) -> None: hermetic_env(monkeypatch) recipe = RECIPES[("eva", "hpt")] diff --git a/tests/unit/test_proprio_history.py b/tests/unit/test_proprio_history.py index 80ff41dacc..3bad6bc907 100644 --- a/tests/unit/test_proprio_history.py +++ b/tests/unit/test_proprio_history.py @@ -705,7 +705,7 @@ def _run(proprio): def test_select_current_step_leaves_a_flat_value_alone_and_is_idempotent(): - from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + from egomimic.rldb.transforms import SelectCurrentStep t = SelectCurrentStep(keys=["observations.state.ee_pose", "missing"]) flat = np.arange(20, dtype=np.float32) @@ -721,7 +721,7 @@ def test_select_current_step_leaves_a_flat_value_alone_and_is_idempotent(): def test_every_human_revert_list_selects_the_current_proprio_step(): from egomimic.rldb.embodiment import human as human_mod - from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + from egomimic.rldb.transforms import SelectCurrentStep builders = { "_build_human_cartesian_revert_6d_transform_list": "observations.state.ee_pose", diff --git a/tests/unit/test_wrist6d_roundtrip.py b/tests/unit/test_wrist6d_roundtrip.py index 882e86c2c1..749977caac 100644 --- a/tests/unit/test_wrist6d_roundtrip.py +++ b/tests/unit/test_wrist6d_roundtrip.py @@ -22,7 +22,7 @@ from scipy.spatial.transform import Rotation as R from egomimic.rldb.embodiment.embodiment import Embodiment -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( CartesianYPRToRot6D, SplitKeys, )