diff --git a/egomimic/eval/latent_dataset.py b/egomimic/eval/latent_dataset.py index 441be8341..dcd078fc2 100644 --- a/egomimic/eval/latent_dataset.py +++ b/egomimic/eval/latent_dataset.py @@ -71,6 +71,7 @@ def build_dataset( task: str, embodiment: str, resolver, + dataset_name: str | None = None, hashes: list[str] | None = None, frames_per_episode: int | None = 128, stride: int | None = None, @@ -94,6 +95,7 @@ def build_dataset( logger.info("[build_dataset] %s | random mode | filter=%s", embodiment, lam) return MultiDataset._from_resolver( resolver, + dataset_name=dataset_name, filters=filters, mode="total", valid_ratio=valid_ratio, @@ -123,6 +125,7 @@ def build_dataset( ) base = MultiDataset._from_resolver( resolver, + dataset_name=dataset_name, filters=filters, mode="total", valid_ratio=valid_ratio, diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index a1b7ea2f1..f96c9750c 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -15,8 +15,7 @@ train_datasets: stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: row['episode_hash'] == '2025-09-20-17-47-54-000000'" + episode_hashes: ["2025-09-20-17-47-54-000000"] mode: total valid_datasets: diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml index d7cfe85ac..c1b366a3c 100644 --- a/egomimic/hydra_configs/data/eva.yaml +++ b/egomimic/hydra_configs/data/eva.yaml @@ -14,8 +14,7 @@ train_datasets: mode: cartesian filters: _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: row['episode_hash'] == '2025-12-26-18-07-46-296000'" + episode_hashes: ["2025-12-26-18-07-46-296000"] mode: total valid_datasets: diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 69ac9d756..e038820b8 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -15,8 +15,7 @@ train_datasets: stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: row['episode_hash'] == '69199812208123403bbdb24f'" + episode_hashes: ["6972164b391867c01c9d3265"] mode: total valid_datasets: diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index aa045cc22..593685ac7 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -15,8 +15,7 @@ train_datasets: stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: row['episode_hash'] in {'2026-03-16-01-22-26-448000'}" + episode_hashes: ["2026-04-30-09-41-51-255837"] mode: total valid_datasets: diff --git a/egomimic/rldb/filters.py b/egomimic/rldb/filters.py index 395cd5f75..373e43fa3 100644 --- a/egomimic/rldb/filters.py +++ b/egomimic/rldb/filters.py @@ -9,8 +9,19 @@ class DatasetFilter: - def __init__(self, filter_lambdas: Sequence[str] | None = None) -> None: + def __init__( + self, + filter_lambdas: Sequence[str] | None = None, + episode_hashes: Sequence[str] | None = None, + ) -> None: self.filter_lambdas = list(filter_lambdas or []) + # Pinned episode hashes. Empty = no pin. Validated at resolve time by the + # resolvers (missing / deleted / wrong embodiment is an error there). + if isinstance(episode_hashes, str): # `filters.episode_hashes=abc` override + episode_hashes = [episode_hashes] + self.episode_hashes: frozenset[str] = frozenset( + str(h) for h in (episode_hashes or []) + ) self.filters = [] for expr in self.filter_lambdas: try: @@ -24,12 +35,17 @@ def __init__(self, filter_lambdas: Sequence[str] | None = None) -> None: self.filters.append(predicate) def __repr__(self) -> str: - return f"DatasetFilter(filter_lambdas={self.filter_lambdas!r})" + return ( + f"DatasetFilter(filter_lambdas={self.filter_lambdas!r}, " + f"episode_hashes={sorted(self.episode_hashes)!r})" + ) def matches(self, row: Mapping[str, Any]) -> bool: row = dict(row) if row.get("is_deleted", False): return False + if self.episode_hashes and row.get("episode_hash") not in self.episode_hashes: + return False for expr, predicate in zip(self.filter_lambdas, self.filters, strict=True): result = predicate(row) if not isinstance(result, bool): @@ -41,14 +57,17 @@ def matches(self, row: Mapping[str, Any]) -> bool: class ScaleAnnotationDatasetFilter(DatasetFilter): def __init__( - self, project_name: str, filter_lambdas: Sequence[str] | None = None + self, + project_name: str, + filter_lambdas: Sequence[str] | None = None, + episode_hashes: Sequence[str] | None = None, ) -> None: self.project_name = project_name self.api_key = os.environ["SCALE_API_KEY"] self.tasks = get_completed_tasks(self.project_name, self.api_key) self.df = build_df_from_tasks(self.tasks) self.completed_episode_hashes = set(self.df["SEQUENCE_ID"].unique().tolist()) - super().__init__(filter_lambdas) + super().__init__(filter_lambdas, episode_hashes) def matches(self, row: Mapping[str, Any]) -> bool: if row.get("episode_hash") not in self.completed_episode_hashes: diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 739e89eb7..8b570626e 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -181,6 +181,75 @@ def get_fallback_idx( return random.choice(valid_candidates), attempts +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: + pins = filters.episode_hashes + if not pins: + return + # Duplicate rows for one hash would make `row[...]` a Series instead of a + # scalar, so collapse to one row per hash before indexing. + by_hash = ( + df.drop_duplicates("episode_hash").set_index("episode_hash", drop=False) + if "episode_hash" in df.columns + else None + ) + problems: list[str] = [] + for h in sorted(pins): + if by_hash is None or h not in by_hash.index: + problems.append(f"{h}: not in app.episodes") + continue + row = by_hash.loc[h] + if not _is_missing_filter_value(row.get("is_deleted")) and bool( + row["is_deleted"] + ): + problems.append(f"{h}: is_deleted") + if ( + _is_missing_filter_value(row.get("zarr_processed_path")) + or not str(row["zarr_processed_path"]).strip() + ): + problems.append(f"{h}: empty zarr_processed_path") + if expected_embodiment is not None: + emb = str(row.get("embodiment", "")) + if emb != expected_embodiment: + problems.append( + f"{h}: embodiment '{emb}' != dataset '{expected_embodiment}'" + ) + if problems: + name = expected_embodiment or "" + raise PinError( + f"{len(problems)} pinned episode(s) invalid for dataset '{name}':\n " + + "\n ".join(problems) + ) + + +def _warn_embodiment_mismatch( + rows: pd.DataFrame, expected_embodiment: str | None, source: str +) -> None: + if expected_embodiment is None or "embodiment" not in rows: + return + mismatched = rows.loc[rows["embodiment"].astype(str) != expected_embodiment] + if len(mismatched): + details = sorted( + f"{h} ({e})" + for h, e in zip( + mismatched["episode_hash"], mismatched["embodiment"].astype(str) + ) + ) + logger.warning( + "%s: %d resolved episode(s) have embodiment != dataset '%s': %s", + source, + len(mismatched), + expected_embodiment, + details, + ) + + class EpisodeResolver: """ Base class for episode resolution utilities. @@ -272,6 +341,7 @@ def __init__( def resolve( self, filters: DatasetFilter | None = None, + expected_embodiment: str | None = None, ) -> dict[str, "ZarrDataset"]: """ Outputs a dict of ZarrDatasets with relevant filters. @@ -291,6 +361,7 @@ def resolve( filters=filters, local_dir=self.folder_path, debug=self.debug, + expected_embodiment=expected_embodiment, ) valid_hashes = {hashes for _, hashes in filtered_paths} @@ -309,7 +380,9 @@ def resolve( @staticmethod def _get_filtered_paths( - filters: DatasetFilter | None = None, debug: int | bool | None = None + filters: DatasetFilter | None = None, + debug: int | bool | None = None, + expected_embodiment: str | None = None, ) -> list[tuple[str, str]]: """ Filters episodes from the SQL episode table according to the criteria specified in `filters` @@ -327,6 +400,9 @@ def _get_filtered_paths( filters = _ensure_dataset_filter(filters) engine = create_default_engine() df = episode_table_to_df(engine) + # Before the empty-table early return: pins against an empty table are + # still invalid and must raise rather than silently resolve to nothing. + _check_pins(df, filters, expected_embodiment) if df.empty: logger.info("Episode table is empty.") return [] @@ -335,19 +411,23 @@ def _get_filtered_paths( lambda row: filters.matches(_normalize_filter_row(row.to_dict())), axis=1, ) - output = df.loc[mask, ["zarr_processed_path", "episode_hash"]] - n_matched_sql = len(output) + matched = df.loc[mask] + n_matched_sql = len(matched) - output = output[ - output["zarr_processed_path"].fillna("").astype(str).str.strip() != "" + matched = matched[ + matched["zarr_processed_path"].fillna("").astype(str).str.strip() != "" ] - n_skipped_null = n_matched_sql - len(output) + n_skipped_null = n_matched_sql - len(matched) if n_skipped_null: logger.info( "Skipped %d episodes with null/empty zarr_processed_path.", n_skipped_null, ) + # Warn only about episodes that actually make it into the dataset. + _warn_embodiment_mismatch(matched, expected_embodiment, "S3EpisodeResolver") + output = matched[["zarr_processed_path", "episode_hash"]] + if debug is not None and debug is not False: k = min(10 if debug is True else int(debug), len(output)) if k < len(output): @@ -453,6 +533,7 @@ def sync_from_filters( local_dir: Path, numworkers: int = 10, debug: int | bool | None = None, + expected_embodiment: str | None = None, ): """ Public API: @@ -469,7 +550,9 @@ def sync_from_filters( filters = _ensure_dataset_filter(filters) # 1) Resolve episodes from DB - filtered_paths = cls._get_filtered_paths(filters, debug=debug) + filtered_paths = cls._get_filtered_paths( + filters, debug=debug, expected_embodiment=expected_embodiment + ) if not filtered_paths: logger.warning("No episodes matched filters.") return [] @@ -554,8 +637,12 @@ def __init__( self.require_annotations = require_annotations self.annotation_key = annotation_key - def resolve(self, filters=None) -> dict[str, "ZarrDataset"]: - datasets = super().resolve(filters=filters) + def resolve( + self, filters=None, expected_embodiment=None + ) -> dict[str, "ZarrDataset"]: + datasets = super().resolve( + filters=filters, expected_embodiment=expected_embodiment + ) if self.key_map is None: return datasets @@ -622,44 +709,74 @@ def __init__( super().__init__(folder_path, key_map, transform_list) self.debug = debug - @staticmethod - def _local_filters_match( - metadata: dict, - episode_hash: str, - filters: DatasetFilter, - ) -> bool: - return filters.matches( - _normalize_filter_row(metadata, episode_hash=episode_hash) - ) - @classmethod def _get_local_filtered_paths( cls, search_path: Path, filters: DatasetFilter | None = None, debug: int | bool | None = None, + expected_embodiment: str | None = None, ): filters = _ensure_dataset_filter(filters) if not search_path.is_dir(): logger.warning("Local path does not exist: %s", search_path) return [] - filtered = [] + rows: list[dict] = [] # one normalized attrs row per episode dir + paths_by_hash: dict[str, str] = {} for p in sorted(search_path.iterdir()): if not p.is_dir(): continue - episode_hash = p.name[:-5] if p.name.endswith(".zarr") else p.name - try: store = zarr.open_group(str(p), mode="r") metadata = dict(store.attrs) except Exception as e: logger.warning("Failed to read metadata for %s: %s", p, e) continue + rows.append(_normalize_filter_row(metadata, episode_hash=episode_hash)) + paths_by_hash[episode_hash] = str(p) + + if filters.episode_hashes: + present = {r["episode_hash"]: r for r in rows} + problems = [] + for h in sorted(filters.episode_hashes): + if h not in present: + problems.append(f"{h}: not in local directory {search_path}") + elif expected_embodiment is not None: + emb = str(present[h].get("embodiment", "")) + if emb != expected_embodiment: + problems.append( + f"{h}: embodiment '{emb}' != dataset '{expected_embodiment}'" + ) + if problems: + raise PinError( + f"{len(problems)} pinned episode(s) invalid for dataset " + f"'{expected_embodiment or ''}':\n " + + "\n ".join(problems) + ) - if cls._local_filters_match(metadata, episode_hash, filters): - filtered.append((str(p), episode_hash)) + # Evaluate the filters once; both the path list and the mismatch warning + # below are derived from the same matched rows. + matched = [r for r in rows if filters.matches(r)] + filtered = [ + (paths_by_hash[r["episode_hash"]], r["episode_hash"]) for r in matched + ] + if expected_embodiment is not None: + bad = sorted( + (r["episode_hash"], str(r.get("embodiment", ""))) + for r in matched + if str(r.get("embodiment", "")) != expected_embodiment + ) + if bad: + details = [f"{h} ({e})" for h, e in bad] + logger.warning( + "LocalEpisodeResolver: %d resolved episode(s) have embodiment != " + "dataset '%s': %s", + len(bad), + expected_embodiment, + details, + ) if debug is not None and debug is not False: k = min(10 if debug is True else int(debug), len(filtered)) @@ -674,6 +791,7 @@ def resolve( self, sync_from_s3=False, filters: DatasetFilter | None = None, + expected_embodiment: str | None = None, ) -> dict[str, "ZarrDataset"]: """ Outputs a dict of ZarrDatasets with relevant filters from local data. @@ -686,7 +804,10 @@ def resolve( filters = _ensure_dataset_filter(filters) filtered_paths = self._get_local_filtered_paths( - self.folder_path, filters, debug=self.debug + self.folder_path, + filters, + debug=self.debug, + expected_embodiment=expected_embodiment, ) valid_folder_names = {folder_name for _, folder_name in filtered_paths} @@ -962,15 +1083,44 @@ def _next_after_failure( return next_idx, attempts @classmethod - def _from_resolver(cls, resolver: EpisodeResolver, **kwargs): - """create a MultiDataset from an EpisodeResolver.""" + 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) + resolved = resolver.resolve( + sync_from_s3=sync_from_s3, filters=filters, expected_embodiment=expected + ) else: - resolved = resolver.resolve(filters=filters) + 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) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 53c0f012a..a4ca7a166 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -20,7 +20,7 @@ from egomimic.eval.eval import Eval from egomimic.pl_utils.pl_model import ModelWrapper from egomimic.rldb.zarr.utils import set_global_seed -from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, PinError from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.instantiators import instantiate_callbacks, instantiate_loggers from egomimic.utils.logging_utils import log_hyperparameters @@ -66,6 +66,22 @@ def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> Non log.info("Dataset frame counts:\n" + table) +def _instantiate_dataset(*args, **kwargs): + """``hydra.utils.instantiate`` wraps any exception raised by the target in an + ``InstantiationException``. A bad pin must surface at launch as a bare + ``PinError`` (not a generic Hydra error the caller has to unwrap), so this + intercepts only that case and re-raises the original ``PinError`` -- keeping + the ``InstantiationException`` as its context (``from e``, not ``from None``) + so the Hydra target/config info isn't lost. Every other target exception is + re-raised exactly as Hydra raised it.""" + try: + return hydra.utils.instantiate(*args, **kwargs) + except hydra.errors.InstantiationException as e: + if isinstance(e.__cause__, PinError): + raise e.__cause__ from e + raise + + class MmapCheckpointIO(TorchCheckpointIO): """``TorchCheckpointIO`` that memory-maps tensor storages instead of reading them. @@ -126,14 +142,14 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: train_datasets = {} for dataset_name in cfg.data.train_datasets: - train_datasets[dataset_name] = hydra.utils.instantiate( - cfg.data.train_datasets[dataset_name] + train_datasets[dataset_name] = _instantiate_dataset( + cfg.data.train_datasets[dataset_name], dataset_name=dataset_name ) valid_datasets = {} for dataset_name in cfg.data.valid_datasets: - valid_datasets[dataset_name] = hydra.utils.instantiate( - cfg.data.valid_datasets[dataset_name] + valid_datasets[dataset_name] = _instantiate_dataset( + cfg.data.valid_datasets[dataset_name], dataset_name=dataset_name ) log.info(f"Instantiating datamodule <{cfg.data._target_}>") @@ -163,7 +179,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: km["norm_mode"] = True instantiate_copy.resolver.key_map = km - norm_dataset = hydra.utils.instantiate(instantiate_copy) + norm_dataset = _instantiate_dataset(instantiate_copy, dataset_name=dataset_name) # infer_norm_from_dataset: load from precomputed JSON/dir if set, else compute (no disk write). norm_stats.infer_norm_from_dataset( norm_dataset, diff --git a/tests/fixtures/recipes.py b/tests/fixtures/recipes.py index 7864f0799..05d57476b 100644 --- a/tests/fixtures/recipes.py +++ b/tests/fixtures/recipes.py @@ -57,13 +57,29 @@ def _pi(data: str, model: str, emb: str) -> Recipe: def common_overrides( - emb: str, data_dir, out_dir, *, batch_size: int, num_workers: int + emb: str, + data_dir, + out_dir, + *, + batch_size: int, + num_workers: int, + episode_hashes: tuple[str, ...] | None = None, ) -> list[str]: - return [ + overrides = [ f"paths.dataset_dir={data_dir}", f"paths.output_dir={out_dir}", f"data.train_datasets.{emb}.resolver._target_={LOCAL_RESOLVER}", + # Reset to null first: some real configs' filters dict (e.g. eva_pi's + # embodiment-only lambda) is struct-locked to its own keys, so a plain + # merge that introduces `episode_hashes` would fail. Force-adding onto + # a null value has no such restriction. f"data.train_datasets.{emb}.filters=null", + ] + if episode_hashes: + pins = ", ".join(episode_hashes) + filters = f"{{_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [{pins}]}}" + overrides.append(f"++data.train_datasets.{emb}.filters={filters}") + overrides += [ f"data.train_dataloader_params.{emb}.batch_size={batch_size}", f"data.train_dataloader_params.{emb}.num_workers={num_workers}", f"data.valid_dataloader_params.{emb}.batch_size={batch_size}", @@ -75,6 +91,7 @@ def common_overrides( "norm_stats.save_cache_dir=null", f"norm_stats.num_workers={num_workers}", ] + return overrides def cpu_trainer_overrides(steps: int = 2) -> list[str]: diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index e65dd72a3..9d7e3828a 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -72,13 +72,15 @@ def _hf_offline(monkeypatch) -> None: monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) -def write_fixtures(tmp_path: Path, vendor: str, n: int = 3) -> tuple[Path, Path]: +def write_fixtures( + tmp_path: Path, vendor: str, n: int = 3 +) -> tuple[Path, Path, tuple[str, ...]]: data, out = tmp_path / "data", tmp_path / "out" data.mkdir() out.mkdir() for i in range(n): write_episode(data, vendor, seed=i) - return data, out + return data, out, tuple(f"{vendor}_{i:02d}" for i in range(n)) class BatchKeySpy: diff --git a/tests/integration/test_train_step_gpu.py b/tests/integration/test_train_step_gpu.py index f4aa112fd..4b6babada 100644 --- a/tests/integration/test_train_step_gpu.py +++ b/tests/integration/test_train_step_gpu.py @@ -33,22 +33,18 @@ VENDOR_NAMES = ["eva", "aria", "mecka", "scale"] # Known-good episodes: the ones the shipped per-vendor HPT configs pin -# (data/eva.yaml, data/aria.yaml). Deterministic, and a regression here means -# the recipe broke, not that the S3 resolver handed us a different episode. -# Two shipped pins are unusable (checked against app.episodes 2026-09-12): -# data/scale.yaml pins 2026-03-16-01-22-26-448000, recorded as eva_bimanual -# (stale after the 2026-05 scale reprocess); data/mecka.yaml pins -# 69199812208123403bbdb24f, which is not in the table at all. -# Those vendors fall back to lab + embodiment with resolver.debug=2. The S3 -# resolver's order is not deterministic, and one rl2 eva episode -# (2026-03-28-19-49-52-786000) has an all-NaN gripper command, so pin -# whenever a pin exists. +# (data/eva.yaml, data/aria.yaml, data/scale.yaml, data/mecka.yaml). +# Deterministic, and a regression here means the recipe broke, not that the +# S3 resolver handed us a different episode. All four vendors are pinned now; +# data/scale.yaml and data/mecka.yaml carry the same pins recorded below +# (checked against app.episodes 2026-09-12, after the 2026-05 scale reprocess +# and the 2026-07-27 mecka reload). PINNED_EPISODES = { "eva": ["2025-12-26-18-07-46-296000"], "aria": ["2025-09-20-17-47-54-000000"], + "scale": ["2026-04-30-09-41-51-255837"], + "mecka": ["6972164b391867c01c9d3265"], } -# `lab` column of app.episodes for the fallback vendors (queried 2026-09-12). -VENDOR_SQL_LAB = {"mecka": "mecka", "scale": "scale"} STEPS = 5 @@ -59,12 +55,7 @@ def _pi_ckpt() -> Path: def _gpu_overrides(emb: str, out: Path, vendor: str) -> list[str]: - debug = ( - [] - if vendor in PINNED_EPISODES - else [f"+data.train_datasets.{emb}.resolver.debug=2"] - ) - return debug + [ + return [ f"paths.output_dir={out}", f"data.train_dataloader_params.{emb}.batch_size=4", f"data.train_dataloader_params.{emb}.num_workers=4", @@ -91,19 +82,11 @@ def _gpu_overrides(emb: str, out: Path, vendor: str) -> list[str]: def _vendor_filter(cfg, emb: str, vendor: str) -> None: - """Real configs pick episodes by SQL; pin this vendor's known-good episodes - (or its lab, limited by debug=2) so the run is deterministic.""" - if vendor in PINNED_EPISODES: - hashes = ", ".join(f"'{h}'" for h in PINNED_EPISODES[vendor]) - lam = f"lambda row: row['episode_hash'] in {{{hashes}}}" - else: - lam = ( - f"lambda row: (row['lab'] == '{VENDOR_SQL_LAB[vendor]}') " - f"& (row['embodiment'] == '{emb}') & (row['is_deleted'] == False)" - ) + """Pin this vendor's known-good episodes through DatasetFilter.episode_hashes + so the run is deterministic and a stale pin fails at resolve (PinError).""" filters = { "_target_": "egomimic.rldb.filters.DatasetFilter", - "filter_lambdas": [lam], + "episode_hashes": list(PINNED_EPISODES[vendor]), } with open_dict(cfg): cfg.data.train_datasets[emb].filters = filters diff --git a/tests/unit/test_dataset_filter.py b/tests/unit/test_dataset_filter.py index 7afdf4c6f..bf98aadde 100644 --- a/tests/unit/test_dataset_filter.py +++ b/tests/unit/test_dataset_filter.py @@ -133,3 +133,33 @@ def test_sync_s3_parser_accepts_named_filter_key() -> None: def test_sync_s3_parser_rejects_unknown_filter_key() -> None: with pytest.raises(ValueError, match="Available filter keys"): parse_dataset_filter_key("does-not-exist") + + +def test_dataset_filter_episode_hashes_pin_rows() -> None: + filters = DatasetFilter(episode_hashes=["a", "b"]) + + assert filters.episode_hashes == frozenset({"a", "b"}) + assert filters.matches({"episode_hash": "a"}) + assert not filters.matches({"episode_hash": "c"}) + assert not filters.matches({"episode_hash": "a", "is_deleted": True}) + + +def test_dataset_filter_episode_hashes_combine_with_lambdas() -> None: + filters = DatasetFilter( + filter_lambdas=["lambda row: row['task'] == 'fold'"], episode_hashes=["a"] + ) + + assert filters.matches({"episode_hash": "a", "task": "fold"}) + assert not filters.matches({"episode_hash": "a", "task": "stack"}) + assert not filters.matches({"episode_hash": "b", "task": "fold"}) + + +def test_dataset_filter_empty_episode_hashes_means_no_pin() -> None: + assert DatasetFilter(episode_hashes=[]).matches({"episode_hash": "anything"}) + assert "episode_hashes" in repr(DatasetFilter(episode_hashes=["a"])) + + +def test_dataset_filter_single_string_pin_is_one_hash() -> None: + filters = DatasetFilter(episode_hashes="2026-04-30-09-41-51-255837") + + assert filters.episode_hashes == {"2026-04-30-09-41-51-255837"} diff --git a/tests/unit/test_latent_dataset_kwargs.py b/tests/unit/test_latent_dataset_kwargs.py new file mode 100644 index 000000000..237f9715a --- /dev/null +++ b/tests/unit/test_latent_dataset_kwargs.py @@ -0,0 +1,98 @@ +"""`trainHydra` passes `dataset_name=` to every dataset target, so the latent +data configs' factory (`egomimic.eval.latent_dataset.build_dataset`) must accept +it and forward it to `MultiDataset._from_resolver` -- otherwise +`data=cotrain_pi_latent` dies at launch with a TypeError.""" + +from __future__ import annotations + +import inspect + +import hydra +import pytest +from fixtures.synthetic_episodes import write_episode +from hydra import compose, initialize_config_module +from omegaconf import OmegaConf, open_dict + +from egomimic.eval import latent_dataset +from egomimic.eval.latent_dataset import build_dataset + +LOCAL_RESOLVER = "egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" + + +def test_build_dataset_accepts_dataset_name() -> None: + assert "dataset_name" in inspect.signature(build_dataset).parameters + + +@pytest.mark.parametrize("mode", ["random", "custom"]) +def test_build_dataset_forwards_dataset_name(monkeypatch, mode) -> None: + seen: dict = {} + + def fake_from_resolver(resolver, **kwargs): + seen.update(kwargs) + return object() + + monkeypatch.setattr( + latent_dataset.MultiDataset, "_from_resolver", fake_from_resolver + ) + monkeypatch.setattr(latent_dataset, "EvenStrideDataset", lambda base, **kw: base) + + build_dataset( + mode, + task="pick_place", + embodiment="eva_bimanual", + resolver=object(), + dataset_name="eva_bimanual", + hashes=["eva_00"], + ) + assert seen["dataset_name"] == "eva_bimanual" + + +def _latent_nodes(): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose( + config_name="train_zarr_cartesian_pi", + overrides=["data=cotrain_pi_latent", "model=pi0.5_bc_eva"], + ) + return cfg + + +def test_latent_config_nodes_instantiate_without_typeerror(tmp_path) -> None: + """Every node of `data=cotrain_pi_latent` must survive + `instantiate(node, dataset_name=)`. The synthetic episodes carry + task_name='synthetic', so the 'random' mode filter legitimately matches + nothing -- a ValueError about episodes is fine; a TypeError is the bug.""" + data = tmp_path / "data" + data.mkdir() + write_episode(data, "eva", seed=0) + write_episode(data, "aria", seed=0) + + cfg = _latent_nodes() + with open_dict(cfg): + cfg.paths.dataset_dir = str(data) + + checked = 0 + for group in ("train_datasets", "valid_datasets"): + for key in cfg.data[group]: + node = OmegaConf.create( + OmegaConf.to_container(cfg.data[group][key], resolve=True) + ) + with open_dict(node): + node.resolver._target_ = LOCAL_RESOLVER + # LocalEpisodeResolver has no `require_annotations` kwarg. + del node.resolver["require_annotations"] + try: + hydra.utils.instantiate(node, dataset_name=key) + except Exception as e: # noqa: BLE001 - we only care about the type + chain = [] + cur: BaseException | None = e + while cur is not None: + chain.append(cur) + cur = cur.__cause__ + assert not any(isinstance(c, TypeError) for c in chain), ( + f"{group}.{key} raised TypeError: " + + "; ".join(f"{type(c).__name__}: {c}" for c in chain) + ) + checked += 1 + assert checked == 4 diff --git a/tests/unit/test_pins.py b/tests/unit/test_pins.py new file mode 100644 index 000000000..db0f4090d --- /dev/null +++ b/tests/unit/test_pins.py @@ -0,0 +1,219 @@ +import logging + +import pandas as pd +import pytest +import zarr +from fixtures.recipes import ( + RECIPES, + common_overrides, + cpu_trainer_overrides, + hpt_small_overrides, +) +from fixtures.synthetic_episodes import write_episode +from fixtures.train_harness import compose_recipe, hermetic_env, write_fixtures + +import egomimic.trainHydra as train_hydra +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr import zarr_dataset_multi +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, + PinError, + S3EpisodeResolver, +) + +TABLE = pd.DataFrame( + [ + { + "episode_hash": "h-ok", + "zarr_processed_path": "s3://rldb/p/h-ok/", + "embodiment": "human_bimanual", + "is_deleted": False, + }, + { + "episode_hash": "h-deleted", + "zarr_processed_path": "s3://rldb/p/h-deleted/", + "embodiment": "human_bimanual", + "is_deleted": True, + }, + { + "episode_hash": "h-nopath", + "zarr_processed_path": "", + "embodiment": "human_bimanual", + "is_deleted": False, + }, + { + "episode_hash": "e-ok", + "zarr_processed_path": "s3://rldb/p/e-ok/", + "embodiment": "eva_bimanual", + "is_deleted": False, + }, + ] +) + + +@pytest.fixture +def fake_table(monkeypatch): + monkeypatch.setattr(zarr_dataset_multi, "create_default_engine", lambda: object()) + monkeypatch.setattr(zarr_dataset_multi, "episode_table_to_df", lambda engine: TABLE) + + +def test_valid_pin_resolves_only_that_episode(fake_table) -> None: + paths = S3EpisodeResolver._get_filtered_paths( + DatasetFilter(episode_hashes=["h-ok"]), expected_embodiment="human_bimanual" + ) + assert paths == [("s3://rldb/p/h-ok/", "h-ok")] + + +def test_bad_pins_are_reported_together(fake_table) -> None: + with pytest.raises(PinError) as exc: + S3EpisodeResolver._get_filtered_paths( + DatasetFilter( + episode_hashes=["h-ok", "h-deleted", "h-nopath", "h-missing", "e-ok"] + ), + expected_embodiment="human_bimanual", + ) + msg = str(exc.value) + assert "4 pinned episode(s) invalid for dataset 'human_bimanual'" in msg + assert "h-missing: not in app.episodes" in msg + assert "h-deleted: is_deleted" in msg + assert "h-nopath: empty zarr_processed_path" in msg + assert "e-ok: embodiment 'eva_bimanual' != dataset 'human_bimanual'" in msg + assert "h-ok" not in msg.split("invalid")[1] + + +def test_embodiment_check_skipped_without_dataset_name(fake_table) -> None: + paths = S3EpisodeResolver._get_filtered_paths( + DatasetFilter(episode_hashes=["e-ok"]) + ) + assert [h for _, h in paths] == ["e-ok"] + + +def test_unpinned_embodiment_mismatch_only_warns(fake_table, caplog) -> None: + with caplog.at_level(logging.WARNING): + paths = S3EpisodeResolver._get_filtered_paths( + DatasetFilter(), expected_embodiment="human_bimanual" + ) + assert {h for _, h in paths} == {"h-ok", "e-ok"} + assert any( + "e-ok" in r.message and "eva_bimanual" in r.message for r in caplog.records + ) + + +def _attrs_only_episode(root, name, **attrs): + zarr.open_group(str(root / f"{name}.zarr"), mode="w").attrs.update(attrs) + + +def test_local_pin_missing_dir_is_error(tmp_path) -> None: + _attrs_only_episode(tmp_path, "l-ok", embodiment="human_bimanual") + with pytest.raises(PinError, match="l-missing: not in local directory"): + LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, + DatasetFilter(episode_hashes=["l-ok", "l-missing"]), + expected_embodiment="human_bimanual", + ) + + +def test_local_pin_embodiment_mismatch_is_error(tmp_path) -> None: + _attrs_only_episode(tmp_path, "l-eva", embodiment="eva_bimanual") + with pytest.raises( + PinError, match="l-eva: embodiment 'eva_bimanual' != dataset 'human_bimanual'" + ): + LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, + DatasetFilter(episode_hashes=["l-eva"]), + expected_embodiment="human_bimanual", + ) + + +def test_local_unpinned_mismatch_warns(tmp_path, caplog) -> None: + _attrs_only_episode(tmp_path, "l-eva", embodiment="eva_bimanual") + _attrs_only_episode(tmp_path, "l-hum", embodiment="human_bimanual") + with caplog.at_level(logging.WARNING): + paths = LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, DatasetFilter(), expected_embodiment="human_bimanual" + ) + assert {h for _, h in paths} == {"l-eva", "l-hum"} + assert any("l-eva" in r.message for r in caplog.records) + + +def test_from_resolver_forwards_dataset_name_as_expected_embodiment(tmp_path) -> None: + write_episode(tmp_path, "eva", seed=0) # hash eva_00, embodiment eva_bimanual + resolver = LocalEpisodeResolver(tmp_path) + ok = MultiDataset._from_resolver( + resolver, + filters=DatasetFilter(episode_hashes=["eva_00"]), + dataset_name="eva_bimanual", + mode="total", + ) + assert set(ok.datasets) == {"eva_00"} + with pytest.raises( + PinError, match="eva_00: embodiment 'eva_bimanual' != dataset 'human_bimanual'" + ): + MultiDataset._from_resolver( + resolver, + filters=DatasetFilter(episode_hashes=["eva_00"]), + dataset_name="human_bimanual", + mode="total", + ) + + +def test_from_resolver_skips_check_for_non_embodiment_names(tmp_path) -> None: + write_episode(tmp_path, "eva", seed=0) + ds = MultiDataset._from_resolver( + LocalEpisodeResolver(tmp_path), + filters=DatasetFilter(episode_hashes=["eva_00"]), + dataset_name="my_mixture", + 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")] + data, out, _ = write_fixtures(tmp_path, "eva") + overrides = common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=("eva_00", "nope"), + ) + cfg = compose_recipe( + recipe, + overrides + cpu_trainer_overrides(2) + hpt_small_overrides(recipe.embodiment), + out, + ) + with pytest.raises(PinError, match="nope: not in local directory"): + train_hydra.train(cfg) + + +def test_from_resolver_rejects_pin_excluded_by_filter_lambdas(tmp_path) -> None: + write_episode(tmp_path, "eva", seed=0) + write_episode(tmp_path, "eva", seed=1) + with pytest.raises(PinError, match=r"1 pinned episode\(s\) not in dataset.*eva_01"): + MultiDataset._from_resolver( + LocalEpisodeResolver(tmp_path), + filters=DatasetFilter( + filter_lambdas=["lambda row: row['episode_hash'] != 'eva_01'"], + episode_hashes=["eva_00", "eva_01"], + ), + dataset_name="eva_bimanual", + mode="total", + ) + + +def test_from_resolver_rejects_pin_that_fails_to_load(tmp_path) -> None: + write_episode(tmp_path, "eva", seed=0) + # Readable attrs (so path resolution accepts the pin) but no arrays or + # features, so constructing its ZarrDataset raises and the loader skips it. + _attrs_only_episode(tmp_path, "eva_broken", embodiment="eva_bimanual") + with pytest.raises(PinError, match="eva_broken"): + MultiDataset._from_resolver( + LocalEpisodeResolver(tmp_path), + filters=DatasetFilter(episode_hashes=["eva_00", "eva_broken"]), + dataset_name="eva_bimanual", + mode="total", + ) diff --git a/tests/unit/test_train_step.py b/tests/unit/test_train_step.py index cafb65eb8..e5fbe9114 100644 --- a/tests/unit/test_train_step.py +++ b/tests/unit/test_train_step.py @@ -47,10 +47,17 @@ def _run(cfg) -> float: def test_hpt_train_step(tmp_path, monkeypatch, vendor): hermetic_env(monkeypatch) recipe = RECIPES[(vendor, "hpt")] - data, out = write_fixtures(tmp_path, vendor) + data, out, hashes = write_fixtures(tmp_path, vendor) cfg = compose_recipe( recipe, - common_overrides(recipe.embodiment, data, out, batch_size=2, num_workers=0) + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + cpu_trainer_overrides(STEPS) + hpt_small_overrides(recipe.embodiment), out, @@ -91,10 +98,17 @@ def test_pi_train_step(tmp_path, monkeypatch, vendor): "expect", {"state_dim": STATE_DIM[recipe.embodiment], "wrist_present": is_eva}, ) - data, out = write_fixtures(tmp_path, vendor) + data, out, hashes = write_fixtures(tmp_path, vendor) cfg = compose_recipe( recipe, - common_overrides(recipe.embodiment, data, out, batch_size=2, num_workers=0) + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + cpu_trainer_overrides(STEPS) + pi_cpu_overrides(), out,