diff --git a/.gitignore b/.gitignore index 8f3d361847..e9a7bc1e58 100644 --- a/.gitignore +++ b/.gitignore @@ -45,12 +45,26 @@ annotations_test/** log_conversion/** debug_model_inputs/** temp_dir/** +/datasets/ +apikey.txt +slurm-*.out +slurmoutputs/ +*.log +.inductor_cache/ + scratch/ sample_data/ +/fb/ external_ckpts/ external/MANO/ +ann_backup/ +aria_ann_backup/ +sort_annotations/ # setuptools build output build/ dist/ docs/ + +# per-worktree venv for the Pi benchmarks (the shared emimic/ venv is off-pin) +.venv-pi/ diff --git a/AGENTS.md b/AGENTS.md index c081b5e36f..0d7644225f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,9 +20,55 @@ Tests live under `tests/`, never next to source in `egomimic/`. `pytest --integration -m gpu tests/integration/test_train_step_gpu.py` on a GPU node. `tests/unit/test_data_configs_compose.py` checks every `data/*.yaml` composes. +## Lance +When implementing anything that reads or writes Lance tables (`egomimic/rldb/lance/`, converters, +committers, loaders), check the Lance format docs first: https://lance.org/format/table/ (table +format, versioning, fragments, schema evolution). Follow them over memory of older Lance releases. +Lance is not fork-safe (internal Rust threads): DataLoader workers that touch Lance use +`multiprocessing_context="forkserver"` (or `spawn`), never the default fork, and reopen the dataset +from its URI inside the worker; see https://lance.org/integrations/pytorch/ and +https://docs.lancedb.com/training/torch. + ## Model settings use plan mode for anything except extremely simple tasks ## Slurm rules If you're on a slurm cluster, request a GPU before running or testing training. -On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G \ No newline at end of file +On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G +Use the `rl2-lab` or `hoffman-lab` partition, with `-A` matching `-p`. If rl2-lab pends on +`QOSGrpGRES` (group GPU quota), switch to `-p hoffman-lab -A hoffman-lab`. + +## Cancelled jobs and leaked GPU memory +A cancelled DDP job can leave whole GPUs allocated, with `nvidia-smi` blaming PIDs that are not in +`/proc`. Nothing is wrong with the driver: the memory is held by **live** `pt_data_worker` processes +reparented to init. + +Dataloader workers are forked after CUDA init, so each inherits the rank's `/dev/nvidia*` fds — and +ranks open *every* visible GPU, so one survivor pins the whole node. The worker cannot exit on its +own: it inherited both ends of its result pipe, so a dead rank never gives it an `EPIPE`, it blocks +in `pipe_write`, and PyTorch's parent check (which only runs between `index_queue` timeouts) is +never reached. `ProctrackType=proctrack/linuxproc` then walks the PPID tree and never sees it. + +Whole ranks leak the same way. A rank wedged in `torch.cuda.synchronize()` inside inductor's +`cudagraph_trees` does not answer SIGTERM, Slurm hits `UnkillableStepTimeout` (60s) and gives up +(`srun: error: Timed out waiting for job step to complete`), and the job leaves the queue with its +ranks still spinning on every GPU. Enabling `torch.compile` makes this more likely, not less. + +`egomimic/utils/gpu_orphans.py` owns the fix: + +- `orphan_guarded(params)` / `die_with_parent` set `PR_SET_PDEATHSIG` on every worker, so the kernel + kills it with its rank regardless of what it is blocked on. **Any new `DataLoader` on a training + path needs one of them** — they are already on the train/val loaders and the norm-stats pass. +- `arm_rank_pdeathsig()` does the same for the rank process itself, so a launcher that gives up + takes its ranks with it. Slurm-gated, and skipped when SIGHUP is ignored, so a `nohup` run in + a `salloc` shell is not shot when that shell exits. +- `reap_orphans()` runs from `trainHydra.main` before CUDA init and clears what an earlier job left + on the node, so a leak cannot outlive the job that made it. It only ever signals processes that + are yours, reparented to init and holding an nvidia fd; a live job's workers have a live parent, + so they are never touched. Off outside Slurm and on non-zero `SLURM_LOCALID`; disable with + `EGOMIMIC_REAP_GPU_ORPHANS=0`. +- `./reap_gpu_orphans.sh [--kill] ` is the manual version, for nodes you are not about to + submit to. It attaches to your own job on a busy node rather than queueing behind it. + +The cluster-level fix is `ProctrackType=proctrack/cgroup`, which sweeps by cgroup instead of by PPID +tree and would make all of this unnecessary. That needs admin; `/etc/slurm` is read-only to us. diff --git a/README.md b/README.md index 6573e5419f..4d77d988be 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ Every cluster-specific default in the configs can be overridden without editing | `EGOVERSE_CACHE_DIR` | `paths.cache_dir`; norm stats cached by content under `/norm_stats/` | `$EGOVERSE_LOG_DIR/cache` | | `EGOVERSE_ROOT` | `paths.root_dir` | the directory you launch from | | `EGOVERSE_PI05_WEIGHTS` | pi0.5 base weights directory (`model.robomimic_model.config.pytorch_weight_path`) | lab PACE path | +| `EGOVERSE_PALIGEMMA_WEIGHTS` | PaliGemma checkpoint for the from-scratch pi arm (`...config.paligemma_weight_path`), a local dir or an HF repo id | `google/paligemma-3b-pt-224` | | `WANDB_ENTITY`, `WANDB_PROJECT` | `logger.wandb.entity` / `.project` | `rl2-group` / `zarr_test` | ### Other Settings diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 0000000000..aa009203cf --- /dev/null +++ b/bench/README.md @@ -0,0 +1,260 @@ +# Training-speed experiments (H100, `ryanco/torch-compile-train`) + +Where the HPT and Pi train steps actually spend their time on the `loaner` +H100 nodes, what was changed, and what each change bought. Every number below +is one H100 (gpu11), bf16 autocast, measured with the harnesses in this folder. + +## The harnesses + +| file | what it does | +| --- | --- | +| `bench_hpt.py` | HPT train step (`forward_training` + `compute_losses` + backward + optimizer) on a synthetic batch, real model config. `--variants` A/Bs one change at a time; `--profile` dumps a torch-profiler table. | +| `bench_pi.py` | The same for Pi (openpi pi0.5), random init -- step time does not depend on the base weights. | +| `check_compile_parity.py` | Eager vs compiled, same weights, same batch: is `torch.compile` changing the numbers? | +| `check_attn_parity.py` | `need_weights=True` (the old path) vs `False`, same weights, same batch. | +| `check_aug_distribution.py` | Two-sample KS between the per-image and batched jitter, over per-sample image statistics. | +| `ab_augs.sh`, `ab_compare.py` | Accuracy A/B for the aug rewrite: both paths, two seeds, held-out loss on disjoint episodes. | +| `e2e_hpt.sh`, `e2e_sweep.sh` | The real `trainHydra` path on locally cached mecka episodes, so the dataloader is in the measurement. | + +```bash +salloc -p loaner -A loaner --gres=gpu:h100:1 -c 28 --mem=200G +source ./wt-env.sh +python bench/bench_hpt.py --variants base,compile_hook --steps 25 --warmup 10 + +# accuracy A/B: ~12 min an arm on one H100 +for s in 42 43; do for arm in loop batched; do + AB_SEED=$s AB_EPOCHS=20 bench/ab_augs.sh $arm /${arm}_s${s}.json +done; done +python bench/ab_compare.py +``` + +`bench_pi.py` needs a venv on the pinned transformers (see "Loose ends", below). + +## What the baseline was doing + +The 300M HPT step at batch 64 was **not** limited by its own arithmetic. The +profiler put ~14,600 CUDA kernel launches in one step, 230 ms of CPU against +111 ms of GPU: the step was bound by the rate the host could issue work. The +model's matmuls accounted for ~26 ms of that. + +Two things dominated the launch count: + +1. **Per-sample image augmentation.** `PerSampleAugs` called torchvision's + `ColorJitter` once per image so each sample drew its own factors. At batch + 64 that is 64 Python round trips, each launching a few dozen kernels and + copying its factors from the host: ~470 `cudaMemcpyAsync` per step, and + **89 ms of a 196 ms step**. +2. **Single-tensor AdamW.** Not a bug in the repo -- `torch.optim.AdamW` with + neither `foreach` nor `fused` set picks `foreach` automatically, which the + shipped configs get. It is called out because passing `fused=False` + *explicitly* silently turns the foreach default off too, which is how the + first version of this benchmark ended up 17 ms/step slower than the real + training loop. + +## Changes + +| change | where | +| --- | --- | +| Vectorized per-sample `ColorJitter`: per-sample factors drawn on-device, each adjustment run once over the batch | `egomimic/models/image_augs.py` | +| `torch.compile` knob: `model.compile.enabled=true` -> `.compile_for_training()` | `pl_utils/pl_model.py`, `utils/compile_utils.py`, `algo/{algo,hpt,pi}.py` | +| `need_weights=False` on the flow head's attentions, so they dispatch to `scaled_dot_product_attention` instead of materialising per-head scores | `egomimic/models/denoising_nets.py` | +| Decode frames to float32 instead of numpy's float64 default, halving what the collate, the pin and the H2D copy move (338 -> 169 MB per camera per batch at 360x640, batch 64; bit-identical) | `egomimic/rldb/zarr/zarr_dataset_multi.py` | +| `non_blocking=True` on the batch's H2D copy (a no-op today, needed by anything that later pins) | `egomimic/algo/{hpt,pi}.py` | +| Drop a per-camera, per-step device->host sync (`if torch.all(img == 0)`) | `egomimic/algo/hpt.py` | +| Skip the per-step deep copy of the batch unless the OT loss is on | `egomimic/algo/hpt.py` | + +`compile_for_training` uses `nn.Module.compile`, not `torch.compile(module)`: +the former patches the call in place, the latter returns an `OptimizedModule` +whose parameters are all renamed `_orig_mod.*`, which would break every +checkpoint written before the flag was turned on. The flip side is that +`Module.compile` only routes `__call__`, so a submodule the algo invokes +through a named method (`stem.compute_latent`, `head.compute_loss`) stays +uncompiled no matter what -- that is why HPT compiles the trunk, the flow +head's denoiser and the image encoders and nothing else, and why `PI`'s call +site moved from `.forward(...)` to `(...)`. + +## Results + +### HPT model step, isolated (`bench_hpt.py`) + +One H100, `hpt_bc_mecka_6d_300M` (298 M trainable), batch 64, bf16 autocast, +synthetic batch, `AdamW` as the configs have it (no explicit `foreach`/`fused`), +including two grad-norm passes (`--grad-norm current`). That models a wrapper +older than this branch: `ModelWrapper` has run one pass since the MAD detector +went log-only, so every absolute time below carries one extra ~3 ms sweep; the +ratios are unaffected. + +360x640 -- the resolution mecka frames actually decode to: + +| | ms/step | vs before | peak GB | +| --- | ---: | ---: | ---: | +| before (per-image `ColorJitter`) | 245.2 | 1.00x | 16.8 | +| vectorized augs | 173.1 | 1.42x | 16.8 | +| + `model.compile.enabled=true` | 109.4 | **2.24x** | 16.3 | + +224x224, for comparison with anything already measured at that size: + +| | ms/step | vs before | peak GB | +| --- | ---: | ---: | ---: | +| before | 205.3 | 1.00x | 14.0 | +| vectorized augs | 120.2 | 1.71x | 13.9 | +| + compile | 82.1 | 2.50x | 13.8 | +| + compile, `mode=reduce-overhead` | 77.7 | 2.64x | 13.8 | + +Compile costs ~15-40 s at the first batch (it grows with image size) and +`reduce-overhead` roughly doubles that for another 5%. Inductor's cache is +keyed per Slurm job (`utils/compile_cache.py`), so every job -- including every +requeue -- pays it once. + +### Pi model step, isolated (`bench_pi.py`) + +One H100, `pi0.5_bc_mecka_6d` (3.6 B), batch 16, random init: + +| | ms/step | vs base | +| --- | ---: | ---: | +| base | 826.8 | 1.00x | +| `fused` AdamW | 794.8 | 1.04x | +| `model.compile.enabled=true` | 588.1 | 1.41x | +| compile + `fused` AdamW | 560.1 | **1.48x** | + +Compile costs ~4 min at the first batch. Forcing the gemma tower back to `sdpa` +attention did nothing for training -- openpi only pins it to `eager` inside +`sample_actions`, which the training pass never enters. + +### Is compile changing the numbers? (`check_compile_parity.py`) + +With every source of randomness pinned (dropout, drop-path, colour jitter, the +flow-matching noise and time) and both copies starting from identical weights: + +* **fp32: the loss is bit-identical.** The compiled graph is the same function. +* The remaining fp32 gradient differences are confined to the ResNet's conv + weights (cosine >= 0.9989), which is cudnn's non-deterministic + backward-filter, not compile. +* Under bf16 autocast the loss drifts ~0.3% on a random-init network, which is + reduction-order noise amplified through 23 trunk blocks. + +### Does the aug rewrite change what the model learns? (`ab_augs.sh`, `ab_compare.py`) + +The vectorized `ColorJitter` is the only change here that is not exactly +equivalent to what it replaced. Everything else is either bit-identical +(float32 decode, the OT clone, the `torch.where` camera select, `non_blocking`) +or floating-point reassociation (compile, `need_weights`). The jitter is +different in two ways, both deliberate: + +* the **op order** is drawn once per call, not once per image, so samples in a + batch share a permutation. A single sample's marginal is unaffected -- the + order is still uniform over the 4! permutations and the factors are still + per-sample -- but within-batch independence is gone. +* the RNG stream differs, so no run reproduces an old seed. + +`check_aug_distribution.py` covers the marginal: 6400 augmented samples through +each path, two-sample KS on eight per-sample statistics (per-channel mean and +std, min, max). Worst p = 0.61. The marginals are indistinguishable. + +The within-batch correlation only a training run can see. `ab_augs.sh` runs both +arms at two seeds: same config, same 32 train episodes, compile off, 2000 steps, +held-out loss on 15 **disjoint** episodes every 100 steps. The held-out pass runs +in eval mode -- deterministic eval augs -- and reseeds, so it measures what the +model learned, not what it was shown. The `loop` arm stubs out `_vectorize` so +`PerSampleAugs` takes its own per-image fallback: the old code path, not a +reimplementation of it. + +| arm | seed | train (last 100) | held-out (mean of last 5) | wall s | +| --- | ---: | ---: | ---: | ---: | +| loop | 42 | 0.617 | 0.752 | 740 | +| loop | 43 | 0.647 | 0.822 | 734 | +| batched | 42 | 0.585 | 0.685 | 589 | +| batched | 43 | 0.598 | 0.631 | 566 | + +**No regression: the batched arm is better at both seeds** (-9.0%, -23.2%), and +the 1.26x wall-clock gap independently confirms the `loop` arm really ran the +per-image path. + +Do NOT read that as the rewrite improving accuracy. Two seeds cannot support +that claim, the two arms consume different amounts of RNG so each arm/seed pair +is effectively its own seed, and the held-out loss still swings 0.6-1.2 between +adjacent evals at step 2000. `ab_compare.py` prints `ARM EFFECT EXCEEDS SEED +NOISE` because the mean gap (0.129) is larger than the within-arm seed spread +(0.070); with n=2 that is weak evidence of a real difference and no evidence at +all of its direction. What it does rule out is the thing worth ruling out -- +sharing the op order across a batch is not costing accuracy. Confirming the +apparent gain would need ~6 seeds an arm. + +### Does `need_weights=False` change the numbers? (`check_attn_parity.py`) + +The flow head's attentions used to take `nn.MultiheadAttention`'s default +`need_weights=True`, which forces the materialized-softmax path; they now +dispatch to `scaled_dot_product_attention`. Forcing the old path back on and +comparing a full fp32 train step: + +* loss 197.009293 vs 197.009033, a relative delta of 1.3e-6 +* 581 gradient tensors, min cosine similarity 0.999999, worst relative max-abs + delta 1.1e-3 (all on ResNet conv weights, i.e. the same cudnn + backward-filter nondeterminism the compile parity check sees) + +Same function, different reduction order. + +### End to end on real data (`e2e_hpt.sh`) + +The real `trainHydra` path, 20 locally cached mecka episodes, 5 x 40 batches, +median of the steps after the first epoch (`bench/step_timer.py` -- Lightning's +`simple` profiler averages the compile and the worker spin-up into every step, +which at 240 steps is a 70 ms/step artefact): + +| | step ms (median) | step ms (p10) | batch wait ms | +| --- | ---: | ---: | ---: | +| before | 298.6 | 266.2 | 41.1 | +| vectorized augs + float32 decode | 274.4 | 196.7 | 41.6 | +| + compile | 245.8 | 157.9 | 40.0 | + +Read the p10 column. The isolated model step drops 64 ms when compiled +(173 -> 109) and the real run's *fast* steps drop 59 ms (217 -> 158), which +matches; the median only drops 33 ms because the real step time is bimodal, +with roughly +/-70 ms of input-side jitter on top of every step. Compile is +doing what the microbenchmark says it does -- the run just spends a chunk of +each step somewhere else. + +That "somewhere else" is not worker count. 8 / 24 / 48 workers give 270.9 / +279.4 / 279.2 ms with a 40 ms batch wait throughout, so the pool is not the +limit and raising `num_workers` off the configs' 8 buys nothing here. The +constant ~40 ms wait is the main process receiving a 169 MB batch out of the +worker queue, which more workers cannot parallelise. + +## Things that were tried and did not pay + +Worth recording so nobody re-runs them: + +| | result | +| --- | --- | +| `persistent_workers=True` on the train loader | No effect. The epoch-start wait is ~1.6 s either way (measured per epoch), and the steady-state step is identical (272.9 vs 274.4 ms). | +| `pin_memory=True` (+ `non_blocking` H2D) | A wash. It moves the 41 ms batch wait into the step (wait 0.9 ms, step 351 vs 312) -- the pinning thread does the copy the main thread used to. It would only win with a stream prefetcher that issues the next batch's H2D during the current step. | +| `prefetch_factor=4` | No measurable effect. | +| `torch.set_float32_matmul_precision("high")`, `cudnn.benchmark` | Under bf16 autocast there is almost no fp32 matmul left to speed up: <1%. | +| `channels_last` for the ResNet | 2-3% in the isolated benchmark, inside the run-to-run noise end to end. | +| `mode="max-autotune-no-cudagraphs"` | 697 s of autotuning for 0.3%. Not worth it. | +| `mode="reduce-overhead"` (CUDA graphs) | A further 5% for roughly double the compile time. Left off by default. | +| Caching the sinusoid position tables | ~0.2%. The tables are rebuilt on the CPU and copied to the device every forward, which is ugly, but it is not where the time goes. | +| `fused=True` AdamW | 5% on HPT, 4% on Pi. Not wired into the configs: it is a numerics change to the optimizer and belongs in its own PR, not this one. Note that passing `fused=False` **explicitly** also disables the `foreach` default and costs 17 ms/step -- the configs correctly pass neither. | + +## Loose ends + +* **`enable_grad_norm` was never passed to `ModelWrapper`**, so + `model/pi0.5_base.yaml`'s `false` did nothing. It is wired now; Pi drops the + `false` and keeps logging its grad norm, as openpi does, and clips at openpi's + 1.0 (`gradient_clip_val` in the model config). HPT pays one pass, a + `clip_grad_norm_(inf)` sweep plus a host sync, about half the 6.5 ms/step the + two-pass bench measured; Pi pays that plus Lightning's clipping sweep, as + openpi does. +* **Which transformers you get depends on the venv.** `pyproject.toml` pins + 4.53.2 plus openpi's `transformers_replace` patch. The venv these numbers were + taken in had 5.17.0, where that patch breaks Pi (`PaliGemmaForConditionalGeneration` + has no `language_model`) and `tests/fixtures/train_harness.py` + (`transformers.utils.hub._is_offline_mode` was removed in 5.x); the Pi + benchmarks used a pinned, git-ignored `.venv-pi`, in which the whole unit suite + passes (1127 passed, 8 skipped). +* **A stream prefetcher** is the obvious next input-side move: issue the next + batch's H2D on a side stream during the current step, which is what would + make `pin_memory` pay. +* **uint8 images** would cut the batch another 4x (169 MB -> 42 MB) over the + float32 change, moving the `/255` and the transpose onto the GPU. Bigger + blast radius -- the transform list and norm stats would need checking. diff --git a/bench/__init__.py b/bench/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/bench/ab_augs.sh b/bench/ab_augs.sh new file mode 100755 index 0000000000..43305aae1a --- /dev/null +++ b/bench/ab_augs.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Accuracy A/B for the vectorized augmentations: same seed, same data, same +# config, only the augmentation path differs. +# bench/ab_augs.sh [extra hydra overrides...] +# Train and held-out episodes are disjoint, so the held-out loss the ABMetrics +# callback records is a generalization measure, not a fit measure. +set -euo pipefail +cd "$(dirname "$0")/.." +source ./wt-env.sh +export EGOVERSE_DATASET_DIR="${EGOVERSE_DATASET_DIR:-/workspace/users/copierceryan/datasets}" +export TOKENIZERS_PARALLELISM=false +ALL=($(ls "$EGOVERSE_DATASET_DIR")) +TRAIN=$(printf '%s\n' "${ALL[@]:0:32}" | paste -sd, -) +VALID=$(printf '%s\n' "${ALL[@]:32}" | paste -sd, -) + +ARM="$1"; OUT="$2"; shift 2 +case "$ARM" in + batched) ENTRY=egomimic/trainHydra.py ;; + loop) ENTRY=bench/train_loop_augs.py ;; + *) echo "arm must be batched or loop" >&2; exit 2 ;; +esac + +EPOCHS="${AB_EPOCHS:-20}" +python "$ENTRY" --config-name=train_zarr_cartesian \ + data=mecka_fold_freeform_opscale1_6d model=hpt_bc_mecka_6d_300M \ + "seed=${AB_SEED:-42}" \ + "data.train_datasets.human_bimanual.resolver._target_=egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" \ + "data.train_datasets.human_bimanual.filters=null" \ + "++data.train_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${TRAIN}]}" \ + "data.valid_datasets.human_bimanual.filters=null" \ + "++data.valid_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${VALID}]}" \ + "~data.unseen_op_valid_datasets" "~data.unseen_op_valid_dataloader_params" \ + "~data.video_episodes" \ + trainer=default trainer.accelerator=gpu trainer.devices=1 \ + "trainer.max_epochs=${EPOCHS}" "trainer.min_epochs=${EPOCHS}" \ + trainer.limit_train_batches=100 +trainer.enable_checkpointing=false \ + trainer.limit_val_batches=0 trainer.check_val_every_n_epoch=100000 \ + +trainer.enable_progress_bar=false logger=debug '~evaluator' '~callbacks' train_viz=false \ + "+callbacks={ab: {_target_: bench.ab_metrics.ABMetrics, out_path: '${OUT}', eval_batches: 12}}" \ + norm_stats.sample_frac=0.02 name=ab_augs "description=$ARM" "$@" diff --git a/bench/ab_compare.py b/bench/ab_compare.py new file mode 100644 index 0000000000..bd52945f2d --- /dev/null +++ b/bench/ab_compare.py @@ -0,0 +1,108 @@ +"""Compare ABMetrics runs: is the arm gap bigger than the seed gap? + +A single pair of runs cannot answer whether an augmentation change moved +accuracy, because two seeds of the SAME arm already differ. This reports the +between-arm difference next to the between-seed difference; the change is only +suspicious if the first is the larger of the two. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + + +def load(path): + with open(path) as f: + return json.load(f) + + +def tail_mean(series, k): + return statistics.fmean(v for _, v in series[-k:]) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("results_dir") + ap.add_argument("--arms", nargs=2, default=["loop", "batched"]) + ap.add_argument("--seeds", nargs="+", type=int, default=[42, 43]) + ap.add_argument("--tail", type=int, default=5, help="eval points to average") + args = ap.parse_args() + + d = Path(args.results_dir) + runs = {} + for arm in args.arms: + for seed in args.seeds: + p = d / f"{arm}_s{seed}.json" + if p.exists(): + runs[(arm, seed)] = load(p) + else: + print(f"missing {p}") + + print( + f"{'arm':>9} {'seed':>5} {'steps':>6} {'train(last100)':>15} " + f"{'heldout(final)':>15} {f'heldout(last{args.tail})':>17} {'wall s':>8}" + ) + for (arm, seed), r in sorted(runs.items()): + tr = r["train_losses"] + ho = r["heldout"] + print( + f"{arm:>9} {seed:>5} {len(tr):>6} " + f"{statistics.fmean(v for _, v in tr[-100:]):>15.3f} " + f"{ho[-1][1]:>15.3f} {tail_mean(ho, args.tail):>17.3f} " + f"{r.get('wall_s') or 0:>8.0f}" + ) + + a0, a1 = args.arms + print(f"\nheld-out loss, mean of last {args.tail} eval points") + arm_gaps = [] + for seed in args.seeds: + if (a0, seed) in runs and (a1, seed) in runs: + x = tail_mean(runs[(a0, seed)]["heldout"], args.tail) + y = tail_mean(runs[(a1, seed)]["heldout"], args.tail) + arm_gaps.append(y - x) + print( + f" seed {seed}: {a0} {x:.3f} -> {a1} {y:.3f} " + f"delta {y - x:+.3f} ({100 * (y - x) / x:+.2f}%)" + ) + + seed_gaps = [] + for arm in args.arms: + vals = [ + tail_mean(runs[(arm, s)]["heldout"], args.tail) + for s in args.seeds + if (arm, s) in runs + ] + if len(vals) > 1: + seed_gaps.append(max(vals) - min(vals)) + print( + f" {arm}: seed spread {max(vals) - min(vals):.3f} " + f"over seeds {args.seeds}" + ) + + if arm_gaps and seed_gaps: + mean_arm = statistics.fmean(arm_gaps) + max_seed = max(seed_gaps) + print(f"\nmean arm effect {mean_arm:+.3f}, largest seed spread {max_seed:.3f}") + if abs(mean_arm) <= max_seed: + print("ARM EFFECT WITHIN SEED NOISE") + else: + print("ARM EFFECT EXCEEDS SEED NOISE -- investigate") + + print( + "\nheld-out curve (step: " + + ", ".join(f"{a}/{s}" for a in args.arms for s in args.seeds) + + ")" + ) + keys = [(a, s) for a in args.arms for s in args.seeds if (a, s) in runs] + n = min(len(runs[k]["heldout"]) for k in keys) + for i in range(n): + step = runs[keys[0]]["heldout"][i][0] + cells = " ".join(f"{runs[k]['heldout'][i][1]:9.3f}" for k in keys) + print(f" {step:>6} {cells}") + + +if __name__ == "__main__": + main() diff --git a/bench/ab_metrics.py b/bench/ab_metrics.py new file mode 100644 index 0000000000..50ee198ba1 --- /dev/null +++ b/bench/ab_metrics.py @@ -0,0 +1,84 @@ +"""Per-step train loss and a fixed held-out loss, written to JSONL. + +For comparing two training runs that differ only in the augmentation path. The +held-out pass runs in eval mode, so it uses the deterministic eval augs and is +unaffected by the change under test -- it measures what the model learned, not +what it was shown. The RNG is reseeded at the start of every pass so the +flow-matching noise and timesteps are identical across passes and across runs. +""" + +from __future__ import annotations + +import json +import statistics +import time + +import torch +from lightning import Callback + + +class ABMetrics(Callback): + def __init__(self, out_path: str, eval_batches: int = 12, eval_seed: int = 4242): + self.out_path = out_path + self.eval_batches = int(eval_batches) + self.eval_seed = int(eval_seed) + self.train_losses: list[tuple[int, float]] = [] + self.heldout: list[tuple[int, float]] = [] + self._loader = None + self._t0 = None + + def on_train_start(self, trainer, pl_module): + self._t0 = time.perf_counter() + dm = trainer.datamodule + self._loader = dm._build_val_style_loader( + dm.valid_datasets, dm.valid_dataloader_params, "valid" + ) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + loss = outputs["loss"] if isinstance(outputs, dict) else outputs + self.train_losses.append((int(trainer.global_step), float(loss))) + + @torch.no_grad() + def on_train_epoch_end(self, trainer, pl_module): + if self._loader is None: + return + model = pl_module.model + was_training = model.nets.training + model.nets.eval() + torch.manual_seed(self.eval_seed) + losses = [] + for i, batch in enumerate(self._loader): + if i >= self.eval_batches: + break + if ( + isinstance(batch, tuple) + and len(batch) == 3 + and isinstance(batch[0], dict) + ): + batch = batch[0] + processed = model.process_batch_for_training(batch) + with torch.autocast("cuda", dtype=torch.bfloat16): + preds = model.forward_training(processed) + losses.append( + float(model.compute_losses(preds, processed)["action_loss"]) + ) + if was_training: + model.nets.train() + if losses: + value = statistics.fmean(losses) + self.heldout.append((int(trainer.global_step), value)) + print( + f"[ab] step {trainer.global_step} heldout_loss {value:.6f} " + f"({len(losses)} batches, {time.perf_counter() - self._t0:.0f}s)", + flush=True, + ) + + def on_fit_end(self, trainer, pl_module): + payload = { + "train_losses": self.train_losses, + "heldout": self.heldout, + "wall_s": time.perf_counter() - self._t0 if self._t0 else None, + } + with open(self.out_path, "w") as f: + json.dump(payload, f) + print(f"[ab] wrote {self.out_path}", flush=True) diff --git a/bench/bench_hpt.py b/bench/bench_hpt.py new file mode 100644 index 0000000000..3d520e5dd8 --- /dev/null +++ b/bench/bench_hpt.py @@ -0,0 +1,544 @@ +"""Isolated HPT train-step benchmark: real model config, synthetic batches. + +No dataloader, no Lightning -- just forward_training + compute_losses + +backward + optimizer step, so a change to the model path shows up undiluted. + + python bench/bench_hpt.py --variants base,tf32,sdpa,compile_trunk + +Variants are monkeypatches applied before the model is built, so an A/B needs +no source edit; the winners get ported into egomimic/ afterwards. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +EMB_HUMAN = 3 +EMB_EVA = 6 + + +class FakeNormStats: + """The slice of MultiDataset that HPT.__init__ reads: key discovery.""" + + norm_mode = "quantile" + + def __init__(self, keys: dict, embodiment_id: int): + self._keys = keys + self._emb = embodiment_id + + def keys_of_type(self, kind, embodiment_id): + return list(self._keys.get(kind, [])) + + def is_key_with_embodiment(self, key, embodiment_id): + return True + + def zarr_key_to_keyname(self, key, embodiment_id): + return key + + def unnormalize(self, preds, embodiment_id): + return preds + + +# ---------------------------------------------------------------- variants + + +def v_tf32(cfg): + torch.set_float32_matmul_precision("high") + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + + +def v_sdpa(cfg): + """need_weights=False on the flow head's attentions. + + nn.MultiheadAttention defaults to need_weights=True, which forces the slow + path: it materialises B x H x L x L scores and averages them over heads, + instead of dispatching to scaled_dot_product_attention. + """ + from egomimic.models import denoising_nets as dn + + def forward_cross(self, x, cond): + res = x + x = self.ln1(x) + x, _ = self.mha(x, x, x, need_weights=False) + x = x + res + res = x + x = self.ln2(x) + x, _ = self.cmha(x, cond, cond, need_weights=False) + x = x + res + res = x + x = self.ln3(x) + x = self.mlp(x) + return x + res + + dn.CrossBlock.forward_cross = forward_cross + + +def v_poscache(cfg): + """Cache the sinusoid position tables instead of rebuilding them on the CPU + (and H2D-copying them) inside every forward.""" + import egomimic.algo.hpt as hpt_mod + import egomimic.utils.tensor_utils as tu + + real = tu.get_sinusoid_encoding_table + cache: dict = {} + + def cached(position_start, position_end, d_hid): + key = (position_start, position_end, d_hid) + if key not in cache: + cache[key] = real(position_start, position_end, d_hid) + return cache[key] + + tu.get_sinusoid_encoding_table = cached + hpt_mod.get_sinusoid_encoding_table = cached + + +def v_noclone(cfg): + """_clone_batch deep-copies every image tensor each step; only the optional + OT loss reads the copy.""" + import egomimic.algo.hpt as hpt_mod + + hpt_mod.HPT._clone_batch = lambda self, batch: batch + + +def v_nozerosync(cfg): + """`if not torch.all(img == 0)` in _robomimic_to_hpt_data is a blocking + device->host sync per camera per step. Always augment instead.""" + import egomimic.algo.hpt as hpt_mod + + def patched(self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[]): + data = {} + for key in proprio_keys: + if key in batch: + short = key.rsplit(".", 1)[-1] + modality = f"state_{short}" + value = batch[key] + data[modality] = value.unsqueeze(1) if value.ndim == 2 else value + for key in cam_keys: + if key in batch: + short = key.rsplit(".", 1)[-1] + _data = self._apply_image_augs(batch[key], short) + data[short] = _data.unsqueeze(1).unsqueeze(1) + for key in lang_keys: + if key in batch: + data[key] = batch[key] + if "sampled_prompt" in batch: + data[self.annotation_modality] = batch["sampled_prompt"] + data["is_6dof"] = self.is_6dof + data["pad_mask"] = batch["pad_mask"] + data["embodiment"] = batch["embodiment"] + for aux in aux_ac_keys: + data[aux] = batch[aux] + data["action"] = batch[self.shared_ac_key or ac_key] + return data + + hpt_mod.HPT._robomimic_to_hpt_data = patched + + +# Every attribute a variant may replace. Without restoring these between +# variants the second row of a sweep measures the first row plus itself, and +# the per-variant column becomes a running total. +def _patch_targets(): + import egomimic.algo.hpt as hpt_mod + import egomimic.utils.tensor_utils as tu + from egomimic.models import denoising_nets, image_augs + + return [ + (image_augs.PerSampleAugs, "forward"), + (hpt_mod.HPT, "_apply_image_augs"), + (hpt_mod.HPT, "_robomimic_to_hpt_data"), + (hpt_mod.HPT, "_clone_batch"), + (denoising_nets.CrossBlock, "forward_cross"), + (tu, "get_sinusoid_encoding_table"), + (hpt_mod, "get_sinusoid_encoding_table"), + ] + + +_PRISTINE = None +_PRISTINE_TF32 = None + + +def _restore_patched_globals(): + global _PRISTINE, _PRISTINE_TF32 + targets = _patch_targets() + if _PRISTINE is None: + _PRISTINE = [getattr(obj, name) for obj, name in targets] + _PRISTINE_TF32 = ( + torch.get_float32_matmul_precision(), + torch.backends.cudnn.benchmark, + torch.backends.cudnn.allow_tf32, + ) + return + for (obj, name), original in zip(targets, _PRISTINE): + setattr(obj, name, original) + precision, benchmark, cudnn_tf32 = _PRISTINE_TF32 + torch.set_float32_matmul_precision(precision) + torch.backends.cudnn.benchmark = benchmark + torch.backends.cudnn.allow_tf32 = cudnn_tf32 + + +def v_loopaugs(cfg): + """The pre-2026-09-19 image-aug path: torchvision's ColorJitter called once + per image. Kept so the table can show what the vectorized one replaced.""" + import torch + + from egomimic.models import image_augs + + def forward(self, images): + return torch.stack([self.augs(image) for image in images]) + + image_augs.PerSampleAugs.forward = forward + + +def v_noaugs(cfg): + """Upper bound for the image-aug path: PerSampleAugs loops ColorJitter over + the batch one image at a time.""" + import egomimic.algo.hpt as hpt_mod + + hpt_mod.HPT._apply_image_augs = lambda self, images, short: images + + +def v_channels_last(cfg): + os.environ["BENCH_CHANNELS_LAST"] = "1" + + +def v_compile_trunk(cfg): + os.environ["BENCH_COMPILE"] = "trunk" + + +def v_compile_head(cfg): + os.environ["BENCH_COMPILE"] = "head" + + +def v_compile_both(cfg): + os.environ["BENCH_COMPILE"] = "both" + + +def v_compile_block(cfg): + os.environ["BENCH_COMPILE"] = "block" + + +def v_compile_all(cfg): + os.environ["BENCH_COMPILE"] = "all" + + +def v_compile_hook(cfg): + """What the shipped knob does: HPT.compile_for_training via nn.Module.compile.""" + os.environ["BENCH_COMPILE"] = "hook" + + +def v_maxautotune(cfg): + os.environ["BENCH_COMPILE_MODE"] = "max-autotune-no-cudagraphs" + + +def v_cudagraphs(cfg): + os.environ["BENCH_COMPILE_MODE"] = "reduce-overhead" + + +VARIANTS = { + "base": lambda cfg: None, + "tf32": v_tf32, + "sdpa": v_sdpa, + "poscache": v_poscache, + "noclone": v_noclone, + "nozerosync": v_nozerosync, + "channels_last": v_channels_last, + "noaugs": v_noaugs, + "loopaugs": v_loopaugs, + "compile_trunk": v_compile_trunk, + "compile_head": v_compile_head, + "compile_both": v_compile_both, + "compile_block": v_compile_block, + "compile_all": v_compile_all, + "compile_hook": v_compile_hook, + "maxautotune": v_maxautotune, + "cudagraphs": v_cudagraphs, +} + + +# ---------------------------------------------------------------- build + + +def build_cfg(model_name: str, overrides: list[str]): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + return compose( + config_name="train_zarr_cartesian", + overrides=[f"model={model_name}", *overrides], + ) + + +def build_algo(cfg, embodiment: str, emb_id: int, keys: dict, device): + node = cfg.model.robomimic_model + algo = hydra.utils.instantiate(node, norm_stats=FakeNormStats(keys, emb_id)) + algo.device = device + algo.nets = algo.nets.to(device) + return algo + + +def synth_batch( + keys, + emb_id, + batch_size, + device, + action_dim, + proprio_dim, + horizon, + img_hw, + history_len, + prompt, +): + b = {} + for k in keys["camera_keys"]: + b[k] = torch.rand(batch_size, 3, *img_hw, device=device) + for k in keys["proprio_keys"]: + shape = ( + (batch_size, proprio_dim) + if history_len == 1 + else (batch_size, history_len, proprio_dim) + ) + b[k] = torch.randn(*shape, device=device) + for k in keys["action_keys"]: + b[k] = torch.randn(batch_size, horizon, action_dim, device=device) + b["pad_mask"] = torch.ones(batch_size, horizon, 1, device=device) + b["embodiment"] = torch.tensor([emb_id], device=device, dtype=torch.int64) + if prompt is not None: + b["sampled_prompt"] = [prompt] * batch_size + return {emb_id: b} + + +def maybe_compile(algo): + what = os.environ.get("BENCH_COMPILE") + if not what: + return + mode = os.environ.get("BENCH_COMPILE_MODE") or None + kw = {"dynamic": False} + if mode: + kw["mode"] = mode + policy = algo.nets["policy"] + if what == "hook": + algo.compile_for_training(mode=mode, dynamic=False) + return + if what == "all": + policy.trunk["trunk"] = torch.compile(policy.trunk["trunk"], **kw) + for name, head in policy.heads.items(): + head.model = torch.compile(head.model, **kw) + for name, enc in policy.encoders.items(): + policy.encoders[name] = torch.compile(enc, **kw) + for name, st in policy.stems.items(): + if name.endswith("_annotation"): + continue # frozen + text-feature cached; input is list[str] + policy.stems[name] = torch.compile(st, **kw) + return + if what in ("trunk", "both"): + policy.trunk["trunk"] = torch.compile(policy.trunk["trunk"], **kw) + if what in ("head", "both"): + for name, head in policy.heads.items(): + head.model = torch.compile(head.model, **kw) + if what == "block": + # per-block compile: one small graph reused 23x, so compile time stays + # flat in depth and the whole trunk still runs compiled kernels + blocks = policy.trunk["trunk"].blocks + for i, blk in enumerate(blocks): + blocks[i] = torch.compile(blk, **kw) + for name, head in policy.heads.items(): + for i, layer in enumerate(head.model.layers): + head.model.layers[i] = torch.compile(layer, **kw) + + +def maybe_channels_last(algo): + if os.environ.get("BENCH_CHANNELS_LAST"): + algo.nets["policy"].encoders.to(memory_format=torch.channels_last) + + +# ---------------------------------------------------------------- run + + +def run(args, variant: str): + for key in ("BENCH_COMPILE", "BENCH_COMPILE_MODE", "BENCH_CHANNELS_LAST"): + os.environ.pop(key, None) + _restore_patched_globals() + torch.manual_seed(0) + device = torch.device("cuda") + + emb_id = EMB_EVA if args.embodiment == "eva_bimanual" else EMB_HUMAN + keys = { + "camera_keys": [f"observations.images.{c}" for c in args.cameras.split(",")], + "proprio_keys": [f"observations.state.{args.proprio}"], + "action_keys": [args.action_key], + "lang_keys": [], + } + + for name in variant.split("+"): + VARIANTS[name](None) + + cfg = build_cfg(args.model, args.override or []) + algo = build_algo(cfg, args.embodiment, emb_id, keys, device) + maybe_channels_last(algo) + maybe_compile(algo) + + params = [p for p in algo.nets.parameters() if p.requires_grad] + # An explicit fused=False would also switch OFF the foreach auto-default, + # i.e. benchmark the single-tensor path against the real config's foreach. + adam_kw = {"foreach": None, "fused": None} + if args.adam == "fused": + adam_kw["fused"] = True + elif args.adam == "foreach": + adam_kw["foreach"] = True + elif args.adam == "single": + adam_kw["foreach"] = False + opt = torch.optim.AdamW(params, lr=1e-5, **adam_kw) + + stem = algo.nets["policy"].stems + hist = 1 + for n, s in stem.items(): + if n.endswith(f"_state_{args.proprio}"): + hist = int(getattr(s, "history_len", 1)) + prompt = "fold the towel" if algo.annotation_key else None + batch = synth_batch( + keys, + emb_id, + args.batch_size, + device, + args.action_dim, + args.proprio_dim, + args.horizon, + tuple(int(x) for x in args.img.split("x")), + hist, + prompt, + ) + + algo.nets.train() + torch.cuda.reset_peak_memory_stats() + amp = torch.autocast("cuda", dtype=torch.bfloat16, enabled=not args.fp32) + + def grad_norm_hooks(): + """What ModelWrapper does between backward and the optimizer step. + + ``once`` is today's wrapper: on_after_backward runs a full + clip_grad_norm_(inf) over every parameter and reads it on the host. + ``current`` adds the second pass on_before_optimizer_step used to run, + which is what the README's numbers were taken with. + """ + if args.grad_norm == "none": + return + params = list(algo.nets.parameters()) + n1 = torch.nn.utils.clip_grad_norm_(params, max_norm=float("inf")) + float(n1) + if args.grad_norm == "current": + n2 = torch.nn.utils.clip_grad_norm_(params, max_norm=float("inf")) + float(n2) + + def step(): + opt.zero_grad(set_to_none=True) + with amp: + preds = algo.forward_training({k: dict(v) for k, v in batch.items()}) + losses = algo.compute_losses(preds, batch) + losses["action_loss"].backward() + grad_norm_hooks() + opt.step() + return losses["action_loss"] + + t_compile = time.perf_counter() + for _ in range(args.warmup): + step() + torch.cuda.synchronize() + t_compile = time.perf_counter() - t_compile + + if args.profile: + from torch.profiler import ProfilerActivity, profile + + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=False, + ) as prof: + for _ in range(3): + step() + torch.cuda.synchronize() + print( + prof.key_averages().table( + sort_by="self_cuda_time_total", row_limit=args.profile_rows + ) + ) + print( + prof.key_averages().table( + sort_by="self_cpu_time_total", row_limit=args.profile_rows + ) + ) + + times = [] + for _ in range(args.steps): + torch.cuda.synchronize() + t0 = time.perf_counter() + step() + torch.cuda.synchronize() + times.append((time.perf_counter() - t0) * 1e3) + + return { + "variant": variant, + "median_ms": round(statistics.median(times), 2), + "p10_ms": round(sorted(times)[len(times) // 10], 2), + "warmup_s": round(t_compile, 1), + "peak_mem_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), + "params_M": round(sum(p.numel() for p in params) / 1e6, 1), + } + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="hpt_bc_mecka_6d_300M") + p.add_argument("--embodiment", default="human_bimanual") + p.add_argument("--cameras", default="front_img_1") + p.add_argument("--proprio", default="ee_pose") + p.add_argument("--action-key", default="actions_cartesian") + p.add_argument("--action-dim", type=int, default=18) + p.add_argument("--proprio-dim", type=int, default=20) + p.add_argument("--horizon", type=int, default=100) + p.add_argument("--img", default="224x224") + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--steps", type=int, default=30) + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--fp32", action="store_true") + p.add_argument( + "--adam", default="default", choices=["default", "foreach", "fused", "single"] + ) + p.add_argument("--grad-norm", default="none", choices=["none", "once", "current"]) + p.add_argument("--override", action="append") + p.add_argument("--variants", default="base") + p.add_argument("--json", default=None) + p.add_argument("--profile", action="store_true") + p.add_argument("--profile-rows", type=int, default=25) + args = p.parse_args() + + rows = [] + for variant in args.variants.split(","): + rows.append(run(args, variant)) + print(json.dumps(rows[-1]), flush=True) + + base = rows[0]["median_ms"] + print(f"\n{'variant':<34}{'ms':>9}{'speedup':>9}{'mem GB':>9}{'warmup s':>10}") + for r in rows: + print( + f"{r['variant']:<34}{r['median_ms']:>9.2f}{base / r['median_ms']:>8.2f}x" + f"{r['peak_mem_gb']:>9.2f}{r['warmup_s']:>10.1f}" + ) + if args.json: + with open(args.json, "w") as f: + json.dump(rows, f, indent=2) + + +if __name__ == "__main__": + main() diff --git a/bench/bench_pi.py b/bench/bench_pi.py new file mode 100644 index 0000000000..51fd761e0c --- /dev/null +++ b/bench/bench_pi.py @@ -0,0 +1,220 @@ +"""Isolated Pi (openpi pi0.5) train-step benchmark: real model config, synthetic +batches, random init (step time does not depend on the base weights). + + python bench/bench_pi.py --variants base,compile_expert +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from bench.bench_hpt import FakeNormStats # noqa: E402 + + +def _patch_transformers_5_image_features() -> bool: + """The venv ships transformers 5.x while openpi pins 4.53.2, where + `get_image_features` returned a bare tensor. Unwrap it so the benchmark can + run; the numbers are indicative of the pinned stack, not a substitute.""" + from openpi.models_pytorch import gemma_pytorch as gp + + real = gp.PaliGemmaWithExpertModel.embed_image + + def embed_image(self, image): + out = real(self, image) + return getattr(out, "last_hidden_state", out) + + gp.PaliGemmaWithExpertModel.embed_image = embed_image + return True + + +EMB_HUMAN = 3 + + +def v_sdpa(cfg): + """openpi pins the gemma tower's attention to `eager` inside sample_actions + and never puts it back, so a train step that follows a val pass runs + materialised-score attention. Force sdpa before every forward.""" + os.environ["BENCH_PI_SDPA"] = "1" + + +def v_compile_expert(cfg): + os.environ["BENCH_PI_COMPILE"] = "expert" + + +def v_compile_forward(cfg): + os.environ["BENCH_PI_COMPILE"] = "forward" + + +def v_compile_hook(cfg): + """What the shipped knob does: PI.compile_for_training via Module.compile.""" + os.environ["BENCH_PI_COMPILE"] = "hook" + + +def v_grad_ckpt(cfg): + os.environ["BENCH_PI_GRAD_CKPT"] = "1" + + +def v_fused_adam(cfg): + os.environ["BENCH_PI_FUSED"] = "1" + + +VARIANTS = { + "base": lambda cfg: None, + "sdpa": v_sdpa, + "compile_expert": v_compile_expert, + "compile_forward": v_compile_forward, + "compile_hook": v_compile_hook, + "grad_ckpt": v_grad_ckpt, + "fused_adam": v_fused_adam, +} + + +def run(args, variant): + for key in ( + "BENCH_PI_COMPILE", + "BENCH_PI_GRAD_CKPT", + "BENCH_PI_FUSED", + "BENCH_PI_SDPA", + ): + os.environ.pop(key, None) + for name in variant.split("+"): + VARIANTS[name](None) + + if args.compat_transformers5: + _patch_transformers_5_image_features() + torch.manual_seed(0) + device = torch.device("cuda") + keys = { + "camera_keys": [f"observations.images.{c}" for c in args.cameras.split(",")], + "proprio_keys": [f"observations.state.{args.proprio}"], + "action_keys": [args.action_key], + "lang_keys": [], + } + + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose( + config_name="train_zarr_cartesian_pi", + overrides=[ + f"model={args.model}", + "model.robomimic_model.config.pytorch_weight_path=null", + *(args.override or []), + ], + ) + algo = hydra.utils.instantiate( + cfg.model.robomimic_model, norm_stats=FakeNormStats(keys, EMB_HUMAN) + ) + algo.device = device + algo.nets = algo.nets.to(device) + + if os.environ.get("BENCH_PI_GRAD_CKPT"): + algo.nets["policy"].gradient_checkpointing_enable() + if os.environ.get("BENCH_PI_COMPILE") == "expert": + pwe = algo.nets["policy"].paligemma_with_expert + pwe.forward = torch.compile(pwe.forward, dynamic=False) + if os.environ.get("BENCH_PI_COMPILE") == "hook": + algo.compile_for_training() + if os.environ.get("BENCH_PI_COMPILE") == "forward": + policy = algo.nets["policy"] + policy.forward = torch.compile(policy.forward, dynamic=False) + + params = [p for p in algo.nets.parameters() if p.requires_grad] + opt = torch.optim.AdamW( + params, lr=1e-5, fused=True if os.environ.get("BENCH_PI_FUSED") else None + ) + + b, hz = args.batch_size, args.horizon + hw = tuple(int(x) for x in args.img.split("x")) + raw = { + keys["camera_keys"][0]: torch.rand(b, 3, *hw, device=device), + keys["proprio_keys"][0]: torch.randn(b, args.proprio_dim, device=device), + args.action_key: torch.randn(b, hz, args.action_dim, device=device), + "annotations": [["fold the towel"] for _ in range(b)], + } + batch = algo.process_batch_for_training({"human_bimanual": raw}) + + algo.nets.train() + torch.cuda.reset_peak_memory_stats() + + def force_sdpa(): + if not os.environ.get("BENCH_PI_SDPA"): + return + pwe = algo.nets["policy"].paligemma_with_expert + pwe.paligemma.language_model.config._attn_implementation = "sdpa" + pwe.gemma_expert.model.config._attn_implementation = "sdpa" + + def step(): + force_sdpa() + opt.zero_grad(set_to_none=True) + preds = algo.forward_training(batch) + loss = algo.compute_losses(preds, batch)["action_loss"] + loss.backward() + opt.step() + + t0 = time.perf_counter() + for _ in range(args.warmup): + step() + torch.cuda.synchronize() + warm = time.perf_counter() - t0 + + times = [] + for _ in range(args.steps): + torch.cuda.synchronize() + t = time.perf_counter() + step() + torch.cuda.synchronize() + times.append((time.perf_counter() - t) * 1e3) + + return { + "variant": variant, + "median_ms": round(statistics.median(times), 2), + "warmup_s": round(warm, 1), + "peak_mem_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), + "params_M": round(sum(p.numel() for p in params) / 1e6, 1), + } + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="pi0.5_bc_mecka_6d") + p.add_argument("--cameras", default="front_img_1") + p.add_argument("--proprio", default="ee_pose") + p.add_argument("--action-key", default="actions_cartesian") + p.add_argument("--action-dim", type=int, default=18) + p.add_argument("--proprio-dim", type=int, default=20) + p.add_argument("--horizon", type=int, default=100) + p.add_argument("--img", default="224x224") + p.add_argument("--batch-size", type=int, default=16) + p.add_argument("--steps", type=int, default=15) + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--override", action="append") + p.add_argument("--variants", default="base") + p.add_argument("--compat-transformers5", action="store_true") + args = p.parse_args() + + rows = [run(args, v) for v in args.variants.split(",")] + for r in rows: + print(json.dumps(r), flush=True) + base = rows[0]["median_ms"] + print(f"\n{'variant':<28}{'ms':>9}{'speedup':>9}{'mem GB':>9}{'warmup s':>10}") + for r in rows: + print( + f"{r['variant']:<28}{r['median_ms']:>9.2f}{base / r['median_ms']:>8.2f}x" + f"{r['peak_mem_gb']:>9.2f}{r['warmup_s']:>10.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/bench/check_attn_parity.py b/bench/check_attn_parity.py new file mode 100644 index 0000000000..9ea765758d --- /dev/null +++ b/bench/check_attn_parity.py @@ -0,0 +1,122 @@ +"""Does need_weights=False change the numbers? + +The flow head's attentions used to call nn.MultiheadAttention with the default +need_weights=True, which forces the materialized-softmax path; they now pass +False and dispatch to scaled_dot_product_attention. That is the same function +with a different reduction order, so the loss and gradients should agree to +floating-point tolerance and nothing more. This forces the old path back on by +wrapping MultiheadAttention.forward and compares a full train step. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import torch +import torch.nn as nn + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from bench.bench_hpt import EMB_HUMAN, synth_batch # noqa: E402 +from bench.check_compile_parity import build, make_deterministic, run # noqa: E402 + +_MHA_FORWARD = nn.MultiheadAttention.forward + + +def force_need_weights(enabled: bool): + """Restore the pre-change call: need_weights defaulted to True.""" + if not enabled: + nn.MultiheadAttention.forward = _MHA_FORWARD + return + + def forward(self, *a, **kw): + kw["need_weights"] = True + kw.pop("average_attn_weights", None) + return _MHA_FORWARD(self, *a, **kw) + + nn.MultiheadAttention.forward = forward + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--rtol", type=float, default=2e-2) + ap.add_argument("--atol", type=float, default=2e-3) + ap.add_argument("--min-cos", type=float, default=0.999) + ap.add_argument("--fp32", action="store_true") + args = ap.parse_args() + + keys = { + "camera_keys": ["observations.images.front_img_1"], + "proprio_keys": ["observations.state.ee_pose"], + "action_keys": ["actions_cartesian"], + "lang_keys": [], + } + torch.manual_seed(1234) + batch = synth_batch( + keys, + EMB_HUMAN, + args.batch_size, + torch.device("cuda"), + 18, + 20, + 100, + (224, 224), + 1, + "fold the towel", + ) + noise = torch.randn(args.batch_size, 100, 18, device="cuda") + time = torch.rand(args.batch_size, device="cuda") * 0.9 + 0.05 + + results = {} + for label, old_path in ( + ("sdpa (need_weights=False)", False), + ("materialized (need_weights=True)", True), + ): + force_need_weights(old_path) + algo = build(keys, compile_it=False, mode=None) + make_deterministic(algo, noise, time) + results[label] = run(algo, batch, seed=7, bf16=not args.fp32) + del algo + torch.cuda.empty_cache() + force_need_weights(False) + + (new_label, old_label) = list(results) + new_loss, new_grads = results[new_label] + old_loss, old_grads = results[old_label] + print( + f"loss {old_label}={old_loss:.6f} {new_label}={new_loss:.6f} " + f"delta={abs(new_loss - old_loss):.3e}" + ) + assert abs(new_loss - old_loss) <= args.atol + args.rtol * abs( + old_loss + ), "loss differs beyond floating-point tolerance" + + assert set(new_grads) == set(old_grads), "gradient key sets differ" + rows, dead = [], 0 + for name, g in old_grads.items(): + h = new_grads[name] + if g.count_nonzero() == 0 and h.count_nonzero() == 0: + dead += 1 + continue + scale = g.abs().max().clamp(min=1e-8) + rel = ((g - h).abs().max() / scale).item() + cos = torch.nn.functional.cosine_similarity( + g.flatten().float(), h.flatten().float(), dim=0 + ).item() + rows.append((rel, cos, name)) + rows.sort(reverse=True) + print(f"{len(rows)} gradient tensors compared ({dead} zero in both)") + print("worst by relative max-abs delta:") + for rel, cos, name in rows[:5]: + print(f" rel={rel:.2e} cos={cos:.6f} {name}") + min_cos = min(c for _, c, _ in rows) + print(f"min cosine similarity: {min_cos:.6f}") + assert min_cos >= args.min_cos, "a gradient points somewhere else" + print("ATTENTION PARITY OK") + + +if __name__ == "__main__": + main() diff --git a/bench/check_aug_distribution.py b/bench/check_aug_distribution.py new file mode 100644 index 0000000000..4d71286ba0 --- /dev/null +++ b/bench/check_aug_distribution.py @@ -0,0 +1,123 @@ +"""Does BatchedColorJitter draw from the same distribution as the per-image loop? + +The two cannot be compared pointwise: they consume different RNG streams, so +every sample differs by construction. What must match is the DISTRIBUTION of +augmented images. This draws the same images through both paths many times and +compares the empirical distributions of a set of per-sample statistics with a +two-sample Kolmogorov-Smirnov test. + +The one known deviation is op ORDER: the loop drew a permutation per image, the +batched version draws one per call. A single sample's marginal is unaffected +(the order is still uniform), so the KS test should pass; what it cannot see is +the within-batch correlation that introduces. --order-check measures that +separately. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import torch +from torchvision import transforms as T + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from egomimic.models.image_augs import PerSampleAugs # noqa: E402 + +JITTER = dict(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05) + + +def ks_2samp(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]: + """Two-sample KS statistic and its asymptotic p-value.""" + a, b = torch.sort(a)[0], torch.sort(b)[0] + n, m = a.numel(), b.numel() + allv = torch.cat([a, b]) + cdf_a = torch.searchsorted(a, allv, right=True).double() / n + cdf_b = torch.searchsorted(b, allv, right=True).double() / m + d = float((cdf_a - cdf_b).abs().max()) + en = (n * m / (n + m)) ** 0.5 + lam = (en + 0.12 + 0.11 / en) * d + p = 2.0 * sum( + (-1) ** (j - 1) * torch.exp(torch.tensor(-2.0 * j * j * lam * lam)).item() + for j in range(1, 101) + ) + return d, min(max(p, 0.0), 1.0) + + +def stats_of(img: torch.Tensor) -> torch.Tensor: + """Per-sample summary: per-channel mean and std, then overall min and max.""" + per_ch_mean = img.mean(dim=(-2, -1)) + per_ch_std = img.std(dim=(-2, -1)) + lo = img.amin(dim=(-3, -2, -1), keepdim=False).unsqueeze(-1) + hi = img.amax(dim=(-3, -2, -1), keepdim=False).unsqueeze(-1) + return torch.cat([per_ch_mean, per_ch_std, lo, hi], dim=-1) + + +STAT_NAMES = ["mean_R", "mean_G", "mean_B", "std_R", "std_G", "std_B", "min", "max"] + + +def collect(module, images, draws, device, seed): + torch.manual_seed(seed) + out = [] + for _ in range(draws): + out.append(stats_of(module(images.clone())).cpu()) + return torch.cat(out, dim=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--draws", type=int, default=200) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--size", type=int, default=64) + ap.add_argument("--alpha", type=float, default=1e-3) + ap.add_argument("--device", default="cuda") + args = ap.parse_args() + + device = torch.device(args.device) + augs = T.Compose([T.ColorJitter(**JITTER)]) + + # A fixed, non-degenerate image set: flat colour would make several of the + # adjustments indistinguishable from each other. + torch.manual_seed(0) + images = torch.rand(args.batch_size, 3, args.size, args.size, device=device) + + batched = PerSampleAugs(augs).to(device) + assert batched.vectorized is not None, "the shipped aug list did not vectorize" + + loop = PerSampleAugs(augs).to(device) + loop.vectorized = None # the pre-change path + + n = args.draws * args.batch_size + print(f"{n} augmented samples per path ({args.draws} draws of {args.batch_size})") + a = collect(loop, images, args.draws, device, seed=1) + b = collect(batched, images, args.draws, device, seed=2) + + print( + f"{'statistic':>10} {'loop mean':>12} {'batched mean':>13} " + f"{'loop sd':>10} {'batched sd':>11} {'KS D':>8} {'p':>8}" + ) + worst_p = 1.0 + for i, name in enumerate(STAT_NAMES): + x, y = a[:, i].double(), b[:, i].double() + d, p = ks_2samp(x, y) + worst_p = min(worst_p, p) + print( + f"{name:>10} {x.mean():12.6f} {y.mean():13.6f} " + f"{x.std():10.6f} {y.std():11.6f} {d:8.4f} {p:8.4f}" + ) + + print( + f"\nworst p = {worst_p:.4f} over {len(STAT_NAMES)} statistics " + f"(alpha = {args.alpha})" + ) + if worst_p < args.alpha: + print("DISTRIBUTIONS DIFFER") + return 1 + print("DISTRIBUTION OK: no detectable difference in the per-sample marginal") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/check_compile_parity.py b/bench/check_compile_parity.py new file mode 100644 index 0000000000..efb97c1878 --- /dev/null +++ b/bench/check_compile_parity.py @@ -0,0 +1,176 @@ +"""Does compile_for_training change the numbers? + +Same seed, same synthetic batch, eager vs compiled: the loss and every gradient +must agree to bf16-autocast tolerance. A mismatch bigger than that means the +compiled graph is not the model we think it is. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import torch.nn as nn # noqa: E402 + +from bench.bench_hpt import EMB_HUMAN, FakeNormStats, synth_batch # noqa: E402 +from egomimic.models.fm_policy import FMPolicy # noqa: E402 + + +def make_deterministic(algo, noise, time): + """Strip every source of randomness from the train step. + + Inductor does not reproduce eager's RNG stream, so a stochastic model gives + two different losses for reasons that have nothing to do with whether the + compiled graph is right. Dropout, drop-path, the colour jitter and the + flow-matching noise/time all get pinned, and what is left must match. + """ + for module in algo.nets.modules(): + if isinstance(module, nn.Dropout): + module.p = 0.0 + if isinstance(module, nn.MultiheadAttention): + module.dropout = 0.0 + if type(module).__name__ == "DropPath": + module.drop_prob = 0.0 + algo.train_image_augs = algo.eval_image_augs + + def predict(self, actions, global_cond): + t = time.to(actions.device).view(-1, 1, 1) + x_t = t * noise + (1 - t) * actions + return self.model(x_t, time.to(actions.device), global_cond), noise - actions + + FMPolicy.predict = predict + + +def build(keys, compile_it, mode): + torch.manual_seed(0) + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose( + config_name="train_zarr_cartesian", overrides=["model=hpt_bc_mecka_6d_300M"] + ) + algo = hydra.utils.instantiate( + cfg.model.robomimic_model, norm_stats=FakeNormStats(keys, EMB_HUMAN) + ) + algo.device = torch.device("cuda") + algo.nets = algo.nets.to("cuda") + if compile_it: + algo.compile_for_training(mode=mode) + return algo + + +def run(algo, batch, seed, bf16=True): + algo.nets.train() + algo.nets.zero_grad(set_to_none=True) + torch.manual_seed(seed) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=bf16): + preds = algo.forward_training({k: dict(v) for k, v in batch.items()}) + loss = algo.compute_losses(preds, batch)["action_loss"] + loss.backward() + grads = { + n: p.grad.detach().clone() + for n, p in algo.nets.named_parameters() + if p.grad is not None + } + return float(loss), grads + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", default=None) + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--rtol", type=float, default=2e-2) + ap.add_argument("--atol", type=float, default=2e-3) + ap.add_argument("--min-cos", type=float, default=0.999) + ap.add_argument("--fp32", action="store_true") + args = ap.parse_args() + + keys = { + "camera_keys": ["observations.images.front_img_1"], + "proprio_keys": ["observations.state.ee_pose"], + "action_keys": ["actions_cartesian"], + "lang_keys": [], + } + torch.manual_seed(1234) + batch = synth_batch( + keys, + EMB_HUMAN, + args.batch_size, + torch.device("cuda"), + 18, + 20, + 100, + (224, 224), + 1, + "fold the towel", + ) + + noise = torch.randn(args.batch_size, 100, 18, device="cuda") + time = torch.rand(args.batch_size, device="cuda") * 0.9 + 0.05 + + eager = build(keys, False, None) + make_deterministic(eager, noise, time) + compiled = build(keys, True, args.mode) + make_deterministic(compiled, noise, time) + + # The comparison is meaningless unless both copies start from the same + # weights, and instantiate() draws from the global RNG. + a = dict(eager.nets.named_parameters()) + b = dict(compiled.nets.named_parameters()) + assert set(a) == set(b) + bad = [n for n in a if not torch.equal(a[n], b[n])] + assert not bad, f"{len(bad)} weights differ before the step, e.g. {bad[:3]}" + + eager_loss, eager_grads = run(eager, batch, seed=7, bf16=not args.fp32) + comp_loss, comp_grads = run(compiled, batch, seed=7, bf16=not args.fp32) + + print(f"loss eager={eager_loss:.6f} compiled={comp_loss:.6f}") + assert abs(eager_loss - comp_loss) <= args.atol + args.rtol * abs( + eager_loss + ), "loss differs beyond tolerance" + + assert set(eager_grads) == set(comp_grads), "gradient key sets differ" + + # Per tensor: max |delta| against that tensor's own scale, and the cosine of + # the two gradient vectors. The elementwise ratio is noisy under bf16 + # autocast (inductor fuses and reorders reductions); the direction is what + # the optimizer actually follows, so it is the honest check. + rows, dead = [], [] + for name, g in eager_grads.items(): + h = comp_grads[name] + if g.count_nonzero() == 0 and h.count_nonzero() == 0: + dead.append(name) # unused parameter in both; no direction to compare + continue + scale = g.abs().max().clamp(min=1e-8) + rel = ((g - h).abs().max() / scale).item() + cos = torch.nn.functional.cosine_similarity( + g.flatten().float(), h.flatten().float(), dim=0 + ).item() + rows.append((rel, cos, name)) + if dead: + print(f"{len(dead)} parameters have a zero gradient in both, e.g. {dead[:2]}") + rows.sort(reverse=True) + print(f"{len(rows)} gradient tensors") + print("worst by relative max-abs delta:") + for rel, cos, name in rows[:5]: + print(f" rel={rel:.2e} cos={cos:.6f} {name}") + min_cos = min(c for _, c, _ in rows) + worst_name = min(rows, key=lambda r: r[1])[2] + print( + f"min cosine similarity: {min_cos:.6f} ({worst_name}) " + f"|eager|={eager_grads[worst_name].norm():.3e} " + f"|compiled|={comp_grads[worst_name].norm():.3e}" + ) + assert min_cos >= args.min_cos, "a gradient points somewhere else" + print("PARITY OK") + + +if __name__ == "__main__": + main() diff --git a/bench/e2e_hpt.sh b/bench/e2e_hpt.sh new file mode 100755 index 0000000000..5ff44515dd --- /dev/null +++ b/bench/e2e_hpt.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# End-to-end HPT throughput: the real trainHydra path on locally cached mecka +# episodes, so the dataloader is in the measurement too. +# bench/e2e_hpt.sh [extra hydra overrides...] +set -euo pipefail +cd "$(dirname "$0")/.." +source ./wt-env.sh +export EGOVERSE_DATASET_DIR="${EGOVERSE_DATASET_DIR:-/workspace/users/copierceryan/datasets}" +export TOKENIZERS_PARALLELISM=false +HASHES=$(ls "$EGOVERSE_DATASET_DIR" | head -20 | paste -sd, -) +RUN="$1"; shift +python egomimic/trainHydra.py --config-name=train_zarr_cartesian \ + data=mecka_fold_freeform_opscale1_6d model=hpt_bc_mecka_6d_300M \ + "data.train_datasets.human_bimanual.resolver._target_=egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" \ + "data.train_datasets.human_bimanual.filters=null" \ + "++data.train_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${HASHES}]}" \ + "data.valid_datasets.human_bimanual.filters=null" \ + "++data.valid_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${HASHES}]}" \ + "~data.unseen_op_valid_datasets" "~data.unseen_op_valid_dataloader_params" \ + "~data.video_episodes" \ + trainer=default trainer.accelerator=gpu trainer.devices=1 trainer.max_epochs=5 trainer.min_epochs=5 \ + trainer.limit_train_batches=40 +trainer.enable_checkpointing=false trainer.limit_val_batches=0 trainer.check_val_every_n_epoch=1000 \ + +trainer.enable_progress_bar=false logger=debug '~evaluator' '~callbacks' train_viz=false \ + '+callbacks={step_timer: {_target_: bench.step_timer.StepTimer, warmup_steps: 40}}' \ + norm_stats.sample_frac=0.01 +trainer.profiler=simple name=perf_e2e "description=$RUN" "$@" diff --git a/bench/e2e_loader.sh b/bench/e2e_loader.sh new file mode 100755 index 0000000000..faaf7d8d99 --- /dev/null +++ b/bench/e2e_loader.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# persistent_workers, isolated: everything else identical. +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="$HERE/../logs/loader" +P="+data.train_dataloader_params.human_bimanual" +for v in false true; do + rm -f "$OUT.$v.log" + bash "$HERE/e2e_hpt.sh" "persist_$v" "$P.persistent_workers=$v" \ + "$P.pin_memory=false" "$P.prefetch_factor=2" > "$OUT.$v.log" 2>&1 +done +for v in false true; do + echo "--- persistent_workers=$v" + grep -h "\[step\]" "$OUT.$v.log" || echo FAILED +done diff --git a/bench/e2e_sweep.sh b/bench/e2e_sweep.sh new file mode 100755 index 0000000000..833ba21db8 --- /dev/null +++ b/bench/e2e_sweep.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# The DataLoader knobs, end to end. Each row is a full trainHydra run on the +# locally cached mecka episodes; read `run_training_batch` (per-step cost) and +# `run_training_epoch` minus 40 x that (per-epoch worker overhead). +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="${1:-$HERE/../logs/e2e}" +P="+data.train_dataloader_params.human_bimanual" + +declare -A ROWS=( + [off]="$P.persistent_workers=false $P.pin_memory=false $P.prefetch_factor=2" + [persist]="$P.persistent_workers=true $P.pin_memory=false $P.prefetch_factor=4" + [persist_pin]="$P.persistent_workers=true $P.pin_memory=true $P.prefetch_factor=4" + [persist_compile]="$P.persistent_workers=true $P.pin_memory=false $P.prefetch_factor=4 model.compile.enabled=true" +) +for row in off persist persist_pin persist_compile; do + rm -f "$OUT.$row.log" + # shellcheck disable=SC2086 + bash "$HERE/e2e_hpt.sh" "$row" ${ROWS[$row]} > "$OUT.$row.log" 2>&1 +done + +for row in off persist persist_pin persist_compile; do + printf '%-18s ' "$row" + grep -h "steady state over" "$OUT.$row.log" || echo FAILED +done diff --git a/bench/e2e_workers.sh b/bench/e2e_workers.sh new file mode 100755 index 0000000000..51d6eefea6 --- /dev/null +++ b/bench/e2e_workers.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# num_workers, end to end. The data configs ship 8; a loaner node has 224 cores. +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="$HERE/../logs/workers" +for n in 8 24 48; do + rm -f "$OUT.$n.log" + bash "$HERE/e2e_hpt.sh" "w$n" "data.train_dataloader_params.human_bimanual.num_workers=$n" \ + > "$OUT.$n.log" 2>&1 +done +rm -f "$OUT.24c.log" +bash "$HERE/e2e_hpt.sh" w24c "data.train_dataloader_params.human_bimanual.num_workers=24" \ + model.compile.enabled=true > "$OUT.24c.log" 2>&1 +for n in 8 24 48 24c; do + printf '%-6s ' "$n" + grep -h "steady state over" "$OUT.$n.log" || echo FAILED +done diff --git a/bench/step_timer.py b/bench/step_timer.py new file mode 100644 index 0000000000..6c90754bbe --- /dev/null +++ b/bench/step_timer.py @@ -0,0 +1,86 @@ +"""Steady-state step timing for a real trainHydra run. + +Lightning's `simple` profiler reports the MEAN over every training batch, which +buries what we want to compare: the first batch carries worker spin-up, cudnn +autotuning and (with `model.compile.enabled`) a whole inductor compile, and at +240 steps a 17 s compile alone shifts the mean by 70 ms/step. This callback +reports the median of the steps after a warmup instead, split into the wait for +the batch and the step itself. +""" + +from __future__ import annotations + +import statistics +import time + +import torch +from lightning import Callback + + +class StepTimer(Callback): + def __init__(self, warmup_steps: int = 40, label: str = "step"): + self.warmup_steps = int(warmup_steps) + self.label = label + self._waits: list[float] = [] + self._steps: list[float] = [] + self._t_end = None + self._t_start = None + self._epoch_start = None + self._epoch_first_waits: list[float] = [] + + def _sync(self): + if torch.cuda.is_available(): + torch.cuda.synchronize() + + def on_train_epoch_start(self, trainer, pl_module): + self._epoch_start = time.perf_counter() + + def on_train_batch_start(self, trainer, pl_module, batch, batch_idx): + now = time.perf_counter() + if batch_idx == 0 and self._epoch_start is not None: + # Wait for the FIRST batch of an epoch. Without persistent_workers + # Lightning rebuilds the iterator here, which forks the pool and + # refills every prefetch queue before this returns. + self._epoch_first_waits.append(now - self._epoch_start) + elif self._t_end is not None: + self._waits.append(now - self._t_end) + self._t_start = now + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + self._sync() + self._t_end = time.perf_counter() + self._steps.append(self._t_end - self._t_start) + + def on_fit_end(self, trainer, pl_module): + n = self.warmup_steps + steps, waits = self._steps[n:], self._waits[n:] + if not steps: + print(f"[{self.label}] only {len(self._steps)} steps, all warmup") + return + + def stats(xs): + xs = sorted(xs) + return ( + statistics.median(xs) * 1e3, + xs[len(xs) // 10] * 1e3, + xs[-max(1, len(xs) // 10)] * 1e3, + ) + + s_med, s_p10, s_p90 = stats(steps) + w_med, w_p10, w_p90 = stats(waits) if waits else (0.0, 0.0, 0.0) + print( + f"[{self.label}] steady state over {len(steps)} steps " + f"(first {n} dropped): " + f"step {s_med:.1f} ms (p10 {s_p10:.1f} / p90 {s_p90:.1f}), " + f"batch wait {w_med:.1f} ms (p10 {w_p10:.1f} / p90 {w_p90:.1f}), " + f"first step {self._steps[0] * 1e3:.0f} ms", + flush=True, + ) + epoch_waits = self._epoch_first_waits + if epoch_waits: + per_epoch = ", ".join(f"{w * 1e3:.0f}" for w in epoch_waits) + print( + f"[{self.label}] epoch-start batch wait (ms), one per epoch: " + f"{per_epoch} (mid-epoch median {w_med:.1f})", + flush=True, + ) diff --git a/bench/train_loop_augs.py b/bench/train_loop_augs.py new file mode 100644 index 0000000000..c992b0afcf --- /dev/null +++ b/bench/train_loop_augs.py @@ -0,0 +1,26 @@ +"""trainHydra with the pre-change, per-image augmentation path. + +PerSampleAugs falls back to the per-image loop whenever _vectorize declines the +aug list, so stubbing it out reproduces the old behaviour exactly -- including +drawing the op order once per image -- without a second checkout. + +trainHydra is executed rather than imported: @hydra.main resolves its relative +config_path against the __file__ of the module that defines the task function, +so importing main from here makes Hydra look for the configs under bench/. +""" + +import os +import runpy + +import egomimic.models.image_augs as image_augs + +image_augs._vectorize = lambda augs: None + +from torchvision import transforms as _T # noqa: E402 + +_probe = image_augs.PerSampleAugs(_T.Compose([_T.ColorJitter(brightness=0.1)])) +assert _probe.vectorized is None, "loop-augs patch did not take" +print("[ab] per-image augmentation path active", flush=True) + +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +runpy.run_path(os.path.join(_ROOT, "egomimic", "trainHydra.py"), run_name="__main__") diff --git a/egomimic/algo/algo.py b/egomimic/algo/algo.py index a01e4cc861..d5040ca26b 100644 --- a/egomimic/algo/algo.py +++ b/egomimic/algo/algo.py @@ -72,3 +72,14 @@ def log_info(self, info): loss_log (dict): name -> summary statistic """ raise NotImplementedError("Must implement log_info in subclass") + + def compile_for_training(self, mode=None, dynamic=False): + """Hand the hot submodules to ``torch.compile``. + + Called once by ``ModelWrapper`` when ``model.compile.enabled`` is set, + before DDP wraps anything -- inductor then sees one module at a time and + DDP's bucketing is unaffected. Algos that have nothing worth compiling + leave this alone; the default no-ops so a run that turns the flag on + against such an algo is slower to start, not broken. + """ + return [] diff --git a/egomimic/algo/eval_determinism.py b/egomimic/algo/eval_determinism.py new file mode 100644 index 0000000000..d25459d248 --- /dev/null +++ b/egomimic/algo/eval_determinism.py @@ -0,0 +1,140 @@ +"""Deterministic evaluation. Mixed into HPT; nothing else uses it yet. + +Every RNG used while an algo is in eval mode is seeded from +``self.eval_base_seed + rank * EVAL_RANK_STRIDE + pass_counter``, where the +pass counter counts ``forward_eval`` calls since the last +``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). +Two validation passes over the same batches therefore replay the same prompts +and the same sampling noise, while different ranks and different batches +within a pass still see different draws. + +PI does NOT mix this in and has no eval seeding of its own, so pi0.5 +validation still replays different prompts and different sampling noise every +pass. Seven of these eight methods are algo-agnostic and PI has everything +they read, so wiring it up is a small follow-up; only ``_build_prompts`` is +HPT-shaped (PI's takes an embodiment and reads ``self.sampling_mode``). + +The mixin reads ``self.nets`` (an ``nn.ModuleDict``, for ``.training``), +``self.device``, ``self.eval_base_seed``, ``self.annotation_key``, +``self.annotation_sampling_mode`` and ``self.default_prompt`` from the algo it +is mixed into; the algo is responsible for setting them. +""" + +from __future__ import annotations + +import logging +import os +import random + +import torch + +logger = logging.getLogger(__name__) + +EVAL_BASE_SEED = 0 +EVAL_RANK_STRIDE = 1_000_003 +# Offset of the val-loss RNG from the sampling generator's seed. It must not be +# a small number: the pass counter advances by 1 per forward_eval call, and a +# torch.Generator seeded with S draws the same stream as a global manual_seed +# of S, so an offset of 1 gave batch i's val loss the noise batch i+1 sampled. +EVAL_LOSS_STRIDE = 7_919_369 + + +class DeterministicEvalMixin: + # Annotated, not defaulted: an algo that forgets to set one of these should + # raise rather than silently take an invented default. + eval_base_seed: int + annotation_key: str | None + annotation_sampling_mode: str + default_prompt: str + + def reset_eval_pass_counter(self) -> None: + """Start a new validation pass: the next ``forward_eval`` call is pass 0. + + Called from ``ModelWrapper.on_validation_start``. Two passes over the + same loader (same order, same number of batches) then draw the same + prompts and the same sampling noise. + """ + self._eval_pass_counter = 0 + + @staticmethod + def _eval_rank() -> int: + """Distributed rank, or 0 outside a process group.""" + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(var) + if value is not None and value.lstrip("-").isdigit(): + return int(value) + return 0 + + def _eval_pass_seed(self) -> int: + """Seed for the current eval pass: base + rank stride + pass counter.""" + return ( + self.eval_base_seed + + self._eval_rank() * EVAL_RANK_STRIDE + + int(getattr(self, "_eval_pass_counter", 0)) + ) + + def _eval_device(self) -> torch.device: + return ( + torch.device(self.device) + if self.device is not None + else torch.device("cpu") + ) + + def _eval_generator(self) -> torch.Generator: + """Seeded generator on the device the sampling noise is drawn on.""" + device = self._eval_device() + generator = torch.Generator(device=device) + generator.manual_seed(self._eval_pass_seed()) + return generator + + def _eval_prompt_rng(self) -> random.Random: + """Seeded Python RNG for annotation choice during eval.""" + return random.Random(self._eval_pass_seed()) + + def _eval_fork_devices(self) -> list: + """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" + device = self._eval_device() + if device.type != "cuda": + return [] + return [ + device.index if device.index is not None else torch.cuda.current_device() + ] + + def _build_prompts(self, _batch, batch_size: int) -> list[str]: + """Sample one annotation per batch item, falling back to default_prompt + on empty / missing annotations. Mirrors the Pi algo flow. + + In eval the choice comes from a per-pass seeded RNG regardless of + ``annotation_sampling_mode``: that keeps prompt variety (unlike always + taking ``sample[0]``) while making two validation passes over the same + batches produce identical prompt lists. Training is unchanged. + """ + if self.annotation_key is None or self.annotation_key not in _batch: + return [self.default_prompt] * batch_size + eval_rng = None if self.nets.training else self._eval_prompt_rng() + prompts = [] + for sample in _batch[self.annotation_key]: + if not sample: + if not self.default_prompt: + self._warn_empty_prompt_once() + prompts.append(self.default_prompt) + elif eval_rng is not None: + prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) + elif self.annotation_sampling_mode == "random": + prompts.append(sample[random.randint(0, len(sample) - 1)]) + else: # "first" + prompts.append(sample[0]) + return prompts + + def _warn_empty_prompt_once(self) -> None: + if getattr(self, "_empty_prompt_warned", False): + return + self._empty_prompt_warned = True + logger.warning( + "prompt fallback: a sample has no %r annotation and default_prompt " + "is empty, so the text stem sees ''. Filter unannotated episodes or " + "set model.robomimic_model.default_prompt.", + self.annotation_key, + ) diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index ad567b0fb6..d79b2847a1 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -1,8 +1,7 @@ import os -import random from collections import OrderedDict from functools import partial -from typing import Literal +from typing import Literal, Union import einops import numpy as np @@ -14,8 +13,20 @@ from tslearn.metrics import SoftDTWLossPyTorch from egomimic.algo.algo import Algo -from egomimic.models.hpt_nets import MultiheadAttention, SimpleTransformer +from egomimic.algo.eval_determinism import ( + EVAL_BASE_SEED, + EVAL_LOSS_STRIDE, + DeterministicEvalMixin, +) +from egomimic.models.hpt_nets import ( + MultiheadAttention, + SimpleTransformer, + apply_skipping_pretrained, + verify_pretrained_weights, +) +from egomimic.models.image_augs import PerSampleAugs from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id +from egomimic.utils.compile_utils import compile_modules from egomimic.utils.hf_utils import download_from_huggingface from egomimic.utils.tensor_utils import EinOpsRearrange, get_sinusoid_encoding_table @@ -111,7 +122,7 @@ def __init__( self.auxiliary_ac_keys = None self.shared_action = False - self.device = None + self.device: Union[torch.device, str, None] = None self.ot_6dof = False self.use_dtw = False @@ -145,6 +156,9 @@ def init_domain_stem(self, domain_name, stem_spec): A specification containing configurations for each modality's stem. """ + # A ``null`` entry lets a child config drop a stem its base declares + # (hydra merges cannot delete keys). + stem_spec = {k: v for k, v in stem_spec.items() if v is not None} self.stem_spec[domain_name] = stem_spec self.modalities[domain_name] = list(stem_spec.keys()) @@ -179,7 +193,12 @@ def finalize_modules(self): """ self.stems = nn.ModuleDict(self.stems) self.heads = nn.ModuleDict(self.heads) - self.apply(self._init_weights) + # Only initialize the modules HPT builds itself. ``self.apply`` would + # walk into pretrained submodules (the HF text encoders, the ImageNet + # ResNet backbone) and overwrite their weights with xavier noise; the + # traversal below skips anything a module declares via + # ``PretrainedWeights``, so future pretrained stems are covered too. + apply_skipping_pretrained(self, self._init_weights) # Shared action tokens if self.token_postprocessing == "action_token": @@ -388,16 +407,24 @@ def stem_process(self, domain, data): horizon = np.random.randint(1, data_horizon + 1) data[modality] = data[modality][:, data_horizon - horizon :] - positional_embedding = get_sinusoid_encoding_table( - 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] - ).to(data[modality]) - positional_embedding = einops.repeat( - positional_embedding, "b h w -> (repeat b) h w", repeat=data_shape[0] - ) + # An observation-history stem owns its own LEARNED time-step + # embedding (added after the projection); the input-side sinusoid + # here is a constant [0, 1, 0, 1, ...] added to the NORMALIZED + # values, which would be amplitude-1 noise on [-1, 1] data once + # axis 1 carries more than the current step. + if getattr(stem, "history_len", 1) == 1: + positional_embedding = get_sinusoid_encoding_table( + 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] + ).to(data[modality]) + positional_embedding = einops.repeat( + positional_embedding, + "b h w -> (repeat b) h w", + repeat=data_shape[0], + ) - data[modality] = data[modality] + positional_embedding.view( - data[modality].shape - ) + data[modality] = data[modality] + positional_embedding.view( + data[modality].shape + ) stem_token = stem.compute_latent(data[modality]) feats.append(stem_token) feat_dict[modality] = stem_token @@ -656,7 +683,7 @@ def compute_loss(self, batch): total_loss = action_loss + shared_action_loss + auxiliary_action_loss return total_loss - def forward(self, domain, data): + def forward(self, domain, data, generator=None): """ Forward pass of the HPTModel to compute actions. @@ -666,6 +693,10 @@ def forward(self, domain, data): The domain corresponding to the input data. data : dict Dictionary containing input data for various modalities. + generator : torch.Generator, optional + RNG handed to the denoising heads' sampler so that evaluation is + reproducible (see ``HPT.forward_eval``). ``None`` (the default) + keeps the global RNG, so every other caller is unchanged. Returns ------- @@ -675,19 +706,23 @@ def forward(self, domain, data): features, block_outputs = self.forward_features(domain, data) action = {} + # Only the denoising (diffusion / flow matching) heads sample noise and + # accept a generator; the plain MLP heads take the features alone. + head_kwargs = {"generator": generator} if self.diffusion else {} + if self.diffusion: features = (features, domain) if domain in self.heads: - action[domain] = self.heads[domain](features) + action[domain] = self.heads[domain](features, **head_kwargs) if self.shared_action: - action["shared"] = self.heads["shared"](features) + action["shared"] = self.heads["shared"](features, **head_kwargs) if domain in self.auxiliary_ac_keys: for key in self.auxiliary_ac_keys[domain]: if f"{domain}_{key}" in self.heads: - action[key] = self.heads[f"{domain}_{key}"](features) + action[key] = self.heads[f"{domain}_{key}"](features, **head_kwargs) return action @@ -781,9 +816,28 @@ def load_pretrained(self, checkpoint_path): self.load_trunk(os.path.join(checkpoint_path, "trunk.pth")) -class HPT(Algo): +class HPT(DeterministicEvalMixin, Algo): """ """ + @property + def device(self): + return self._device + + @device.setter + def device(self, value): + """Set the algo's device, and the policy's with it. + + The policy is built on CPU and moved by Lightning, which writes the + real device here (``ModelWrapper.on_fit_start`` / + ``on_validation_start``). ``HPTModel`` reads it to create the loss + accumulators in ``compute_loss``, so a write that stopped at the algo + would leave those on a different device than the parameters. + """ + self._device = value + nets = getattr(self, "nets", None) + if nets is not None and "policy" in nets: + nets["policy"].device = value + def __init__( self, norm_stats, @@ -820,6 +874,12 @@ def __init__( annotation_modality: str = "annotation", default_prompt: str = "", # --------------------------- + # Loss masking: supervise only the real action steps of a chunk + # (the dataset's ``action_pad_mask``) instead of the repeat-last + # padding at every episode / annotation tail. + # --------------------------- + use_pad_mask: bool = False, + # --------------------------- # Catch-all kwargs # --------------------------- **kwargs, @@ -830,8 +890,13 @@ def __init__( self.annotation_sampling_mode = annotation_sampling_mode self.annotation_modality = annotation_modality self.default_prompt = default_prompt + self.use_pad_mask = use_pad_mask - self.train_image_augs = train_image_augs + # Per-sample so every image draws its own jitter parameters; torchvision + # samples once per call, i.e. once for the whole stacked batch. + self.train_image_augs = ( + PerSampleAugs(train_image_augs) if train_image_augs is not None else None + ) self.eval_image_augs = eval_image_augs self.stem_specs = stem_specs self.head_specs = head_specs @@ -849,13 +914,12 @@ def __init__( self.is_6dof = kwargs.get("6dof", False) self.kinematics_solver = kwargs.get("kinematics_solver", None) - model = HPTModel(**trunk) + model = self._build_policy(trunk) model.auxiliary_ac_keys = self.auxiliary_ac_keys + self.eval_base_seed = int(kwargs.get("eval_base_seed", EVAL_BASE_SEED)) self.multitask = kwargs.get("multitask", False) - self.device = kwargs.get( - "device", torch.device("cuda" if torch.cuda.is_available() else "cpu") - ) + self.device = None model.device = self.device self.diffusion = kwargs.get("diffusion", False) @@ -896,8 +960,6 @@ def __init__( for modality, encoder_cfg in self.encoders.items(): model.init_encoder(modality, encoder_cfg) - model.finalize_modules() - self.ac_keys = {} self.camera_keys = {} self.proprio_keys = {} @@ -909,8 +971,6 @@ def __init__( self.freeze_depth = kwargs.get("freeze_depth", 8) model.depth = self.depth - self.rkl_samples = kwargs.get("reverse_kl_samples", 4) - if self.ot: self.ot_warm_start_steps = kwargs.get("ot_warm_start_steps", 0) self.ot_6dof = kwargs.get("ot_6dof", False) @@ -951,25 +1011,42 @@ def __init__( model.finalize_modules() self.nets["policy"] = model - self.nets = self.nets.float().to(self.device) + self.nets = self.nets.float() + + # Every pretrained submodule must still hold its checkpoint's weights + # once the model is fully built, moved and upcast. Raises on mismatch. + # It costs a second full copy of the encoder (~1.1 GB fp16 for 0.6B) + # and ~3 s of hashing per rank, so it is skippable -- but on by + # default: a silently re-initialised encoder trains to a plausible, + # wrong result. + self.verify_pretrained = bool(kwargs.get("verify_pretrained", True)) + if self.verify_pretrained: + verify_pretrained_weights(self.nets["policy"]) self.training_step = 0 - def _build_prompts(self, _batch, batch_size: int) -> list[str]: - """Sample one annotation per batch item, falling back to default_prompt - on empty / missing annotations. Mirrors the Pi algo flow. - """ - if self.annotation_key is None or self.annotation_key not in _batch: - return [self.default_prompt] * batch_size - prompts = [] - for sample in _batch[self.annotation_key]: - if not sample: - prompts.append(self.default_prompt) - elif self.annotation_sampling_mode == "random": - prompts.append(sample[random.randint(0, len(sample) - 1)]) - else: # "first" - prompts.append(sample[0]) - return prompts + def _build_policy(self, trunk: dict) -> HPTModel: + """The policy network; subclasses (``RDT``) swap in their own.""" + return HPTModel(**trunk) + + @override + def compile_for_training(self, mode=None, dynamic=False): + """Compile the trunk, the flow head's denoiser and the image encoders. + + Only these three are invoked through ``__call__``; the stems run via + ``compute_latent`` and the heads via ``compute_loss``, which + ``Module.compile`` does not route (and which measured no gain anyway). + """ + policy = self.nets["policy"] + # no_trunk: HPT's forward does not call it. RDT's heads do, so its DiT + # is traced into each head's graph without an entry of its own. + targets = [("trunk", policy.trunk["trunk"] if not policy.no_trunk else None)] + targets += [ + (f"head[{name}].model", getattr(head, "model", None)) + for name, head in policy.heads.items() + ] + targets += [(f"encoder[{name}]", enc) for name, enc in policy.encoders.items()] + return compile_modules(targets, mode=mode, dynamic=dynamic) @override def process_batch_for_training(self, batch): @@ -1006,9 +1083,29 @@ def process_batch_for_training(self, batch): B, S, _ = processed_batch[embodiment_id][ac_key].shape device = processed_batch[embodiment_id][ac_key].device - processed_batch[embodiment_id]["pad_mask"] = torch.ones( - B, S, 1, device=device + # Dropped either way: downstream sees the same keys as before. + action_pad_mask = processed_batch[embodiment_id].pop( + "action_pad_mask", None ) + # Front-padded history steps are copies of the first real frame -- + # exactly a history-dropout sample -- so the stem needs nothing + # extra and downstream keys stay as they were. + processed_batch[embodiment_id].pop("proprio_history_mask", None) + # MultiDataset already zeroed an untracked hand's proprio. + processed_batch[embodiment_id].pop("proprio_hand_valid", None) + if self.use_pad_mask and action_pad_mask is not None: + if action_pad_mask.shape[-1] != S: + raise ValueError( + f"action_pad_mask length {action_pad_mask.shape[-1]} does " + f"not match the action horizon {S}" + ) + processed_batch[embodiment_id]["pad_mask"] = action_pad_mask[ + ..., None + ].to(device) + else: + processed_batch[embodiment_id]["pad_mask"] = torch.ones( + B, S, 1, device=device + ) # Sample one annotation per item (random/first, default fallback for # empty). Stays as list[str]; the Qwen stem owns tokenization. @@ -1024,7 +1121,11 @@ def process_batch_for_training(self, batch): # TODO make this work with any fp type for key, value in processed_batch[embodiment_id].items(): if isinstance(value, torch.Tensor): - value = value.to(self.device) + # non_blocking: a no-op from pageable memory, and from a + # pin_memory loader it lets the copy overlap the previous + # step. Everything downstream is queued on the same stream, + # so ordering holds without an explicit wait. + value = value.to(self.device, non_blocking=True) if value.is_floating_point(): value = value.float() processed_batch[embodiment_id][key] = value @@ -1056,13 +1157,22 @@ def forward_training(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side "data": data, } - hpt_batches[embodiment_id] = self._clone_batch(hpt_batch) + # Only the OT loss reads this second copy; cloning every image + # tensor each step for a loss that is off is pure allocation. + if self.ot: + hpt_batches[embodiment_id] = self._clone_batch(hpt_batch) if self.freeze_repr: loss = self.nets["policy"].compute_loss_depth( @@ -1100,6 +1210,12 @@ def forward_eval(self, batch): } """ unnorm_preds = {} + # One seeded generator per pass/batch: the flow-matching noise below and + # the prompts picked in ``process_batch_for_training`` share this pass's + # seed, so replaying the pass replays every draw. + generator = self._eval_generator() + pass_seed = self._eval_pass_seed() + fork_devices = self._eval_fork_devices() for embodiment_id, _batch in batch.items(): embodiment_name = get_embodiment(embodiment_id).lower() cam_keys = self.camera_keys[embodiment_id] @@ -1108,7 +1224,13 @@ def forward_eval(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side @@ -1118,16 +1240,24 @@ def forward_eval(self, batch): # so keep a fresh copy for the forward() call below. forward_data = self._clone_batch(hpt_batch["data"]) - # BC val loss — same call as forward_training. - if self.freeze_repr: - val_loss = self.nets["policy"].compute_loss_depth( - hpt_batch, depth=self.freeze_depth - ) - else: - val_loss = self.nets["policy"].compute_loss(hpt_batch) + # BC val loss — same call as forward_training. The training loss + # draws its own noise / timestep from the global RNG deep inside the + # head (``DenoisingPolicy.predict``), which no generator kwarg + # reaches; fork and seed the global RNG around the call instead so + # the reported val loss is reproducible too. + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(pass_seed + EVAL_LOSS_STRIDE) + if self.freeze_repr: + val_loss = self.nets["policy"].compute_loss_depth( + hpt_batch, depth=self.freeze_depth + ) + else: + val_loss = self.nets["policy"].compute_loss(hpt_batch) unnorm_preds[f"{embodiment_name}_loss"] = val_loss - actions = self.nets["policy"].forward(hpt_batch["domain"], forward_data) + actions = self.nets["policy"].forward( + hpt_batch["domain"], forward_data, generator=generator + ) predictions = OrderedDict() for key in actions: @@ -1152,6 +1282,10 @@ def forward_eval(self, batch): for key in unnorm_actions: unnorm_preds[f"{embodiment_name}_{key}"] = unnorm_actions[key] + # Advance the pass counter once per forward_eval call (after the batch, + # so the prompts built for this batch and the noise drawn for it share a + # seed). ``reset_eval_pass_counter`` rewinds it at on_validation_start. + self._eval_pass_counter = int(getattr(self, "_eval_pass_counter", 0)) + 1 return unnorm_preds @override @@ -1218,8 +1352,53 @@ def _forward_ot(self, batch, embodiment1_id, embodiment2_id): supervised=self.supervised, ) + def _apply_image_augs(self, images, short, frames: int = 1): + """ + helper method that augments one camera's (B, 3, H, W) batch; train augs + are per-sample, the deterministic eval augs stay batched. ``frames`` > 1 + is (B * frames, 3, H, W), a sample's frames adjacent and sharing one + draw (see ``PerSampleAugs``). + + An absent camera arrives as an all-zero image and must stay zero + (normalizing it would make it -mean/std). The test is per sample, so a + batch mixing present and absent samples is covered, and it is a + torch.where rather than an `if`, which would read the predicate on the + host and serialise the step against the next batch's copy. + """ + if self.nets.training and self.train_image_augs and short in self.encoders: + augs = self.train_image_augs + elif self.eval_image_augs and short in self.encoders: + augs = self.eval_image_augs + else: + return images + absent = (images == 0).flatten(1).all(1) + out = ( + augs(images, frames=frames) + if isinstance(augs, PerSampleAugs) + else augs(images) + ) + return torch.where(absent.view(-1, *[1] * (out.dim() - 1)), 0.0, out) + + def _stem_history_len(self, domain: str, modality: str) -> int: + """``history_len`` of ``domain``'s ``modality`` stem (1 when the stem + declares none). Stems are registered as ``f"{domain}_{modality}"``; a + cotrain run has one per domain, so the domain has to be named -- a + suffix match would return whichever came first in the dict. + """ + # an nn.ModuleDict once built: `in` and [] work on it, .get does not + stems = getattr(self.nets["policy"], "stems", None) or {} + name = next((n for n in (f"{domain}_{modality}", modality) if n in stems), None) + return int(getattr(stems[name], "history_len", 1)) if name else 1 + def _robomimic_to_hpt_data( - self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[] + self, + batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys=[], + domain="", ): """ helper method that returns data in the format required for the HPT model @@ -1232,19 +1411,29 @@ def _robomimic_to_hpt_data( for key in proprio_keys: if key in batch: short = key.rsplit(".", 1)[-1] - data[f"state_{short}"] = batch[key].unsqueeze(1) + modality = f"state_{short}" + value = batch[key] + if value.ndim == 2: + data[modality] = value.unsqueeze(1) # (B, D) -> (B, 1, D) + else: + # Observation history: already (B, K, D), one token per step. + history_len = self._stem_history_len(domain, modality) + if value.shape[1] != history_len: + raise ValueError( + f"proprio '{key}' carries {value.shape[1]} history " + f"steps but the '{modality}' stem has " + f"history_len={history_len}; set " + "model...state_ee_pose.history_len and the keymap's " + "proprio_history to the same K" + ) + data[modality] = value for key in cam_keys: if key in batch: short = key.rsplit(".", 1)[-1] - _data = batch[key] - if not torch.all(_data == 0): - if self.nets.training and short in self.encoders: - _data = self.train_image_augs(_data) - elif self.eval_image_augs and short in self.encoders: - _data = self.eval_image_augs(_data) - - data[short] = _data.unsqueeze(1).unsqueeze(1) + data[short] = ( + self._apply_image_augs(batch[key], short).unsqueeze(1).unsqueeze(1) + ) for key in lang_keys: if key in batch: diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index d055eab546..0ed4b1a3c8 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -16,6 +16,8 @@ from transformers import AutoTokenizer from egomimic.algo.algo import Algo +from egomimic.models import openpi_compat +from egomimic.models.paligemma_init import load_paligemma_weights, select_init_source from egomimic.models.preprocess_pi_obs import ( PI_CAMERA_SLOTS, _concat_proprio, @@ -25,7 +27,9 @@ gather_pi_images, ) from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id -from egomimic.utils.action_utils import ConverterRegistry +from egomimic.utils.action_utils import ConverterRegistry, pad_to_width +from egomimic.utils.compile_utils import compile_modules +from egomimic.utils.pose_utils import hand_channel_mask logger = logging.getLogger(__name__) # Ensure logger propagates to root logger and has appropriate level @@ -33,6 +37,10 @@ logger.setLevel(logging.INFO) logger.propagate = True # Explicitly enable propagation (default, but ensures it works) +# openpi targets transformers 4.53 (this repo's pin); these shims let it run on a +# 5.x venv too. See egomimic/models/openpi_compat.py. +openpi_compat.apply() + class PI(Algo): """ """ @@ -142,6 +150,7 @@ def __init__( self.action_converters = action_converters self.action_registry = ConverterRegistry() + self._packed_widths: dict[tuple, int] = {} arcfg = self.action_converters default_ac_key = getattr(arcfg, "ac_key", "actions_cartesian") @@ -171,7 +180,20 @@ def __init__( self.model = openpi.models_pytorch.pi0_pytorch.PI0Pytorch(model_cfg) - if self.config.pytorch_weight_path is not None: + # Three mutually exclusive starting points, in decreasing order of what + # is pretrained: the full pi0.5 base checkpoint, PaliGemma only (the + # action expert and the action/time projections stay at their init), or + # nothing at all. + paligemma_weight_path = getattr(self.config, "paligemma_weight_path", None) + init_source = select_init_source( + self.config.pytorch_weight_path, paligemma_weight_path + ) + target = ( + self.model.module + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) + else self.model + ) + if init_source == "pi05_base": model_path = os.path.join( self.config.pytorch_weight_path, "model.safetensors" ) @@ -179,25 +201,63 @@ def __init__( raise FileNotFoundError( f"Pretrained weight file not found: {model_path}" ) - target = ( - self.model.module - if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) - else self.model - ) safetensors.torch.load_model(target, model_path) logger.info( "Loaded pretrained weights from %s (%d parameters)", model_path, sum(p.numel() for p in target.parameters()), ) + elif init_source == "paligemma": + load_paligemma_weights( + target, + paligemma_weight_path, + allow_mirror=getattr(self.config, "paligemma_allow_mirror", False), + ) else: logger.warning( - "No pytorch_weight_path: no base weights loaded. The weights must " - "come from a checkpoint; otherwise this trains from scratch." + "Neither pytorch_weight_path nor paligemma_weight_path is set: no " + "base weights loaded. The weights must come from a checkpoint; " + "otherwise this trains from scratch, VLM included." ) + # openpi's PyTorch port hard-codes 32-wide action projections + # regardless of ``Pi0Config.action_dim``. Wider action spaces (the + # 144-D hand-keypoint action, or a cotrain that pads eva's 20-D into + # the same vector) need ``action_dim``-wide ones, swapped in AFTER the + # (strict) base-weight load. Only these two layers start fresh; + # ``sample_actions`` already draws its noise at ``config.action_dim``. + self.action_dim = int(self.config.model.action_dim) + if self.action_dim != 32: + self._resize_action_projections(self.action_dim) self.nets = nn.ModuleDict() self.nets["policy"] = self.model + def _resize_action_projections(self, action_dim: int) -> None: + target = ( + self.model.module + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) + else self.model + ) + old_in, old_out = target.action_in_proj, target.action_out_proj + width = old_in.out_features + new_in = nn.Linear(action_dim, width).to( + dtype=old_in.weight.dtype, device=old_in.weight.device + ) + new_out = nn.Linear(width, action_dim).to( + dtype=old_out.weight.dtype, device=old_out.weight.device + ) + target.action_in_proj = new_in + target.action_out_proj = new_out + logger.warning( + "action_dim=%d != 32: re-initialized action_in_proj (%d->%d) and " + "action_out_proj (%d->%d); the pretrained 32-wide projections are " + "discarded, everything else comes from the base checkpoint.", + action_dim, + action_dim, + width, + width, + action_dim, + ) + def _control_mode_for(self, emb_name: str | None) -> str: if self.control_mode and emb_name is not None: for key, val in self.control_mode.items(): @@ -350,10 +410,11 @@ def process_batch_for_training(self, batch): if key_name is not None: processed_batch[embodiment_id][key_name] = value - # Carry forward episode_hash from the data dict (other keys are - # built inline below via _tokenize_prompts). - if "episode_hash" in _batch: - processed_batch[embodiment_id]["episode_hash"] = _batch["episode_hash"] + # Not keymap keys, so the filter above drops them; without intrinsics + # the overlays fall back to the Aria camera. + for key in ("episode_hash", "intrinsics", "action_hand_mask"): + if key in _batch: + processed_batch[embodiment_id][key] = _batch[key] ac_key = self.ac_keys[embodiment_id] if ac_key not in processed_batch[embodiment_id]: @@ -382,7 +443,9 @@ def process_batch_for_training(self, batch): for key, value in processed_batch[embodiment_id].items(): if isinstance(value, torch.Tensor): - value = value.to(self.device) + # See HPT.process_batch_for_training: no-op from pageable + # memory, overlapped copy from a pin_memory loader. + value = value.to(self.device, non_blocking=True) if value.is_floating_point(): value = value.float() processed_batch[embodiment_id][key] = value @@ -395,6 +458,20 @@ def process_batch_for_training(self, batch): return processed_batch + @override + def compile_for_training(self, mode=None, dynamic=False): + """Compile openpi's ``PI0Pytorch.forward`` (the flow-matching training + pass: SigLIP tower, PaliGemma prefix and the action expert). + + ``sample_actions`` is already compiled by openpi itself, so only the + training pass is left. The observation preprocessing at the top of that + forward is Python-level dict work, so expect dynamo to graph-break there + and compile the transformer stack behind it. + """ + return compile_modules( + [("pi0", self.nets["policy"])], mode=mode, dynamic=dynamic + ) + @override def forward_training(self, batch): """ @@ -420,8 +497,16 @@ def forward_training(self, batch): embodiment_name, ) - losses = self.nets["policy"].forward(processed_obs, action) - loss = self._reduce_loss(losses, _batch[ac_key], embodiment_id, ac_key) + # __call__, not .forward: nn.Module.compile only routes __call__, so + # a direct .forward() would silently run uncompiled. + losses = self.nets["policy"](processed_obs, action) + loss = self._reduce_loss( + losses, + _batch[ac_key], + embodiment_id, + ac_key, + _batch.get("action_hand_mask"), + ) predictions[f"{embodiment_name}_{ac_key}"] = _batch[ac_key] predictions[f"{embodiment_name}_loss"] = loss @@ -458,9 +543,13 @@ def forward_eval(self, batch): ) # Flow-matching val loss — same call as forward_training. - losses = self.nets["policy"].forward(processed_obs, action) + losses = self.nets["policy"](processed_obs, action) unnorm_preds[f"{embodiment_name}_loss"] = self._reduce_loss( - losses, _batch[ac_key], embodiment_id, ac_key + losses, + _batch[ac_key], + embodiment_id, + ac_key, + _batch.get("action_hand_mask"), ) pred_actions = self.nets["policy"].sample_actions( @@ -478,20 +567,75 @@ def forward_eval(self, batch): return unnorm_preds - def _reduce_loss(self, losses, action, embodiment_id, ac_key) -> torch.Tensor: - """Mean flow-matching loss for the embodiment.""" + def _reduce_loss( + self, losses, action, embodiment_id, ac_key, hand_mask=None + ) -> torch.Tensor: + """Mean flow-matching loss over the embodiment's packed action width. + + openpi returns the per-element loss ``(B, H, action_dim)``. With + ``action_dim`` widened for the 144-D keypoint action, a narrower + embodiment (eva's 32-slot cartesian layout) would otherwise average + its loss over 112 zero-padded "predict the noise" dims and contribute + a fraction of the gradient signal it had at width 32; restrict the + mean to the slots the converter actually packs. + + ``hand_mask`` (B, 2) [left, right] leaves an untracked hand's packed slots + out of the mean. + """ if isinstance(losses, (list, tuple)): losses = torch.stack(losses) elif not isinstance(losses, torch.Tensor): losses = torch.tensor(losses, device=action.device, dtype=torch.float32) - return losses.mean() + if losses.ndim != 3: + return losses.mean() + losses = losses[..., : self._packed_width(action, embodiment_id, ac_key)] + if hand_mask is None: + return losses.mean() + mask = self._packed_hand_mask(hand_mask, action, embodiment_id, ac_key) + if mask is None: + return losses.mean() + mask = mask[:, None, : losses.shape[-1]].expand_as(losses) + return (losses * mask).sum() / mask.sum().clamp(min=1.0) + + def _packed_hand_mask(self, hand_mask, action, embodiment_id, ac_key): + """``hand_mask`` spread over the converter's packed slots, (B, packed width); + slots the converter inserts (gripper zeros, padding) stay 1.""" + channels = hand_channel_mask( + hand_mask.to(device=action.device, dtype=torch.float32), + action.shape[-1], + ) + if channels is None: + if bool((hand_mask == 1).all()): + return None + raise ValueError( + f"action width {action.shape[-1]} has no bimanual layout to mask a hand in" + ) + converter = self.action_registry.get(embodiment_id, ac_key) + packed = 1.0 - converter.to32_norm_6d((1.0 - channels)[:, None, :])[:, 0] + if not bool(((packed == 0) | (packed == 1)).all()): + raise ValueError( + f"the {type(converter).__name__} converter mixes channels; " + "a per-hand loss mask cannot follow it" + ) + return packed + + def _packed_width(self, action, embodiment_id, ac_key) -> int: + """Width the embodiment's converter packs into, cached per pairing: it + is a property of the converter, not of the batch.""" + key = (embodiment_id, ac_key) + if key not in self._packed_widths: + converter = self.action_registry.get(embodiment_id, ac_key) + self._packed_widths[key] = int( + converter.to32_norm_6d(action[:1, :1]).shape[-1] + ) + return self._packed_widths[key] def _postprocess_sampled_actions(self, pred_actions, _batch, embodiment_id, ac_key): """Raw ``sample_actions`` output -> unnormalized native action dict. - Unpack the normalized native action (xyz+6D(+gripper)) from the - model's action vector, then unnormalize via the standard pipeline - (stats were computed on the native layout).""" + Unpack the normalized native action (xyz+6D(+gripper) or the 144-D + keypoint vector) from the model's action vector, then unnormalize via + the standard pipeline (stats were computed on the native layout).""" ref = _batch[ac_key] _, T, D = ref.shape converter = self.action_registry.get(embodiment_id, ac_key) @@ -549,7 +693,7 @@ def log_info(self, info): return log def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodiment): - """Dataset batch -> (openpi observation, 32-D action chunk).""" + """Dataset batch -> (openpi observation, ``action_dim``-wide action chunk).""" if ac_key not in batch: raise KeyError(f"Missing action key '{ac_key}' in batch") @@ -561,8 +705,8 @@ def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodime converter = self.action_registry.get(emb_id, ac_key) # The action is already normalized and in its native layout (the # ypr->6D conversion happened in the data transforms). Pack it into - # the canonical 32-slot block layout. - action32 = converter.to32_norm_6d(action) + # the canonical block layout and zero-pad up to the model width. + action32 = pad_to_width(converter.to32_norm_6d(action), self.action_dim) # OpenPI expects a fixed camera tuple under its own names. Human data # only has the front camera; the missing wrist slots get a copy of it diff --git a/egomimic/algo/rdt.py b/egomimic/algo/rdt.py new file mode 100644 index 0000000000..74043f0deb --- /dev/null +++ b/egomimic/algo/rdt.py @@ -0,0 +1,286 @@ +"""RDT (Robotics Diffusion Transformer) on the HPT chassis. + +``RDT`` is ``HPT`` with a different policy network: the same batch handling, +stems, encoders, denoising heads, eval and config layout, but no Perceiver +pooling. Every stem hands its full token sequence to the head, whose denoiser +(``rdt_nets.RDTDenoiser``) runs the noisy action chunk through the trunk -- the +DiT blocks (``rdt_nets.RDTBackbone``), shared across embodiments, which +cross-attend to the language and image tokens in alternating blocks. +""" + +import torch +import torch.nn as nn + +from egomimic.algo.hpt import HPT, HPTModel +from egomimic.models.rdt_nets import ( + RDTBackbone, + RDTConditions, + RDTMemory, + mlp_gelu, + sincos_2d, +) +from egomimic.rldb.embodiment.embodiment import ( + HISTORY_MASK_SUFFIX, + IMAGE_HISTORY_SUFFIX, + IMAGE_MEMORY_SUFFIX, +) +from egomimic.utils.tensor_utils import get_sinusoid_encoding_table + +INIT_STD = 0.02 + + +class RDTModel(HPTModel): + """Stems -> ``RDTConditions`` -> head -> shared ``RDTBackbone``. Streams: a + modality with an encoder is an image, a stem with ``forward_with_mask`` is + language, the rest is state. + + ``cond_drop`` (RDT's ``cond_mask_prob``, per stream, training only): the + state is replaced by a learned null token; the prompt by ``""``, i.e. the + no-prompt deployment condition; a camera by a learned null token, never + every camera of a sample, so ``img`` is inert with a single camera. + + ``image_history`` is the number of frames per camera (axis 1 of the image + input, current frame last); ``image_history_dropout`` replaces the past + frames by the current one, which is also what an episode start looks like. + + ``memory`` (``RDTMemory`` kwargs) adds long-range memory tokens to the + DiT's self-attention prefix from the batch's ``memory`` window; padded + steps are masked out. ``short_dropout`` replaces a sample's whole short + image stream by a learned null token (training only), so the policy cannot + lean on the latest frames alone. + """ + + def __init__( + self, + embed_dim: int = 1024, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 1.0, + cond_drop: dict | None = None, + image_history: int = 1, + image_history_dropout: float = 0.0, + memory: dict | None = None, + short_dropout: float = 0.0, + **kwargs, + ): + super().__init__( + embed_dim=embed_dim, no_trunk=True, token_postprocessing="no-op" + ) + self.trunk = nn.ModuleDict( + {"trunk": RDTBackbone(embed_dim, depth, num_heads, mlp_ratio)} + ) + self.cond_drop = {"state": 0.0, "lang": 0.0, "img": 0.0, **(cond_drop or {})} + self.image_history = int(image_history) + self.image_history_dropout = float(image_history_dropout) + self.stem_modality = {} + self.short_dropout = float(short_dropout) + self.long_memory = ( + None if memory is None else RDTMemory(hidden_dim=embed_dim, **memory) + ) + + def _create_policy_trunk(self, *args, **kwargs): + return nn.ModuleDict() + + def init_domain_stem(self, domain_name, stem_spec): + # no init_cross_attn: the stems' Perceiver latents would be dead weights + stem_spec = {k: v for k, v in stem_spec.items() if v is not None} + self.stem_spec[domain_name] = stem_spec + self.modalities[domain_name] = list(stem_spec) + for modality, stem in stem_spec.items(): + self.stems[f"{domain_name}_{modality}"] = stem + self.stem_modality[f"{domain_name}_{modality}"] = modality + + def _stream(self, modality, stem) -> str: + if modality in self.encoders: + return "img" + return "lang" if hasattr(stem, "forward_with_mask") else "state" + + def finalize_modules(self): + D = self.embed_dim + streams = { + name: self._stream(self.stem_modality[name], stem) + for name, stem in self.stems.items() + } + cameras = [name for name, s in streams.items() if s == "img"] + # RDT's mlp2x_gelu language adaptor; image / state adaptors are the stems + self.lang_adaptors = nn.ModuleDict( + {name: mlp_gelu(D, D, 2) for name, s in streams.items() if s == "lang"} + ) + super().finalize_modules() + self.camera_embed = nn.ParameterDict( + {name: torch.zeros(1, 1, 1, D) for name in cameras} + ) + self.frame_embed = nn.Parameter(torch.zeros(1, self.image_history, 1, D)) + # Only what can be used: an unused parameter crashes plain DDP. + if self.cond_drop["state"] > 0 and "state" in streams.values(): + self.null_state = nn.Parameter(torch.randn(1, 1, D) * INIT_STD) + if self.cond_drop["img"] > 0 and len(cameras) > 1: + self.null_img = nn.Parameter(torch.randn(1, 1, D) * INIT_STD) + if self.short_dropout > 0: + self.null_short = nn.Parameter(torch.randn(1, 1, D) * INIT_STD) + self.trunk["trunk"].initialize_weights() + for head in self.heads.values(): + init = getattr(getattr(head, "model", None), "initialize_weights", None) + if callable(init): + init() + + def _drop(self, prob: float, batch_size: int, device): + """(B,) bool: which samples lose a condition; None outside training.""" + if not (self.training and prob > 0): + return None + return torch.rand(batch_size, device=device) < prob + + def _image_tokens(self, name, modality, stem, x): + """(B, T, N, 3, H, W) frames, current last -> (B, T * N * patches, D).""" + B, T = x.shape[:2] + if T != self.image_history: + raise ValueError( + f"'{modality}' carries {T} frames but trunk.image_history=" + f"{self.image_history}; set it to match the data config's " + "image_history_gap_s (2 frames when set, else 1)" + ) + drop = self._drop(self.image_history_dropout, B, x.device) if T > 1 else None + if drop is not None: + x = torch.where(drop.view(B, 1, 1, 1, 1, 1), x[:, -1:].expand_as(x), x) + encoder = self.encoders[modality] + tokens = stem(encoder(x)) + tokens = tokens.reshape(B, T, -1, tokens.shape[-1]) + # DINOv3 is RoPE-only: its tokens carry no absolute position. + grid = getattr(encoder, "grid_size", None) + if grid is not None and grid[0] * grid[1] == tokens.shape[2]: + pos = sincos_2d(tokens.shape[-1], grid) + else: + pos = get_sinusoid_encoding_table(0, tokens.shape[2], tokens.shape[-1])[0] + tokens = tokens + pos.to(tokens) + self.frame_embed + self.camera_embed[name] + return tokens.reshape(B, -1, tokens.shape[-1]) + + def _drop_cameras(self, img: list) -> list: + """Per-camera dropout that never blinds a sample completely.""" + if len(img) < 2 or not (self.training and self.cond_drop["img"] > 0): + return img + B, device = len(img[0]), img[0].device + drop = torch.rand(B, len(img), device=device) < self.cond_drop["img"] + keep = torch.randint(len(img), (B,), device=device) + blind = drop.all(dim=1) + drop[blind, keep[blind]] = False + return [ + torch.where(drop[:, i, None, None], self.null_img.to(tokens), tokens) + for i, tokens in enumerate(img) + ] + + def forward_features(self, domain, data): + if "fps" not in data: + raise ValueError("RDT needs the per-sample 'fps' (ZarrDataset emits it)") + img, state, lang, lang_mask = [], [], [], [] + for modality in self.modalities.get(domain, []) + self.shared_keys: + if modality not in data: + continue + owner = "shared" if modality in self.shared_keys else domain + name = f"{owner}_{modality}" + stem = self.stems[name] + stream = self._stream(modality, stem) + x = data[modality] + + if stream == "lang": + drop = self._drop(self.cond_drop["lang"], len(x), self.device) + if drop is not None: + x = ["" if d else p for p, d in zip(x, drop.tolist())] + tokens, mask = stem.forward_with_mask(x) + if not mask.any(dim=1).all(): + raise ValueError(f"'{modality}': a prompt encoded to zero tokens") + lang.append(self.lang_adaptors[name](tokens)) + lang_mask.append(mask) + elif stream == "img": + img.append(self._image_tokens(name, modality, stem, x)) + else: + tokens = stem(x) + state.append(tokens.reshape(len(tokens), -1, tokens.shape[-1])) + + if not img: + raise ValueError(f"RDT needs at least one camera for domain '{domain}'") + if (self.long_memory is None) != ("memory" not in data): + raise ValueError( + "trunk.memory and the data config's key_map.image_memory must be " + "set together (model memory: " + f"{self.long_memory is not None}, batch memory: {'memory' in data})" + ) + img = torch.cat(self._drop_cameras(img), dim=1) + drop = self._drop(self.short_dropout, len(img), img.device) + if drop is not None: + img = torch.where(drop[:, None, None], self.null_short.to(img), img) + cond = RDTConditions( + img=img, + freq=data["fps"].reshape(-1).expand(len(img)), + backbone=self.trunk["trunk"], + ) + if state: + cond.state = torch.cat(state, dim=1) + drop = self._drop(self.cond_drop["state"], len(img), img.device) + if drop is not None: + null = self.null_state.to(cond.state) + cond.state = torch.where(drop[:, None, None], null, cond.state) + if lang: + cond.lang = torch.cat(lang, dim=1) + cond.lang_mask = torch.cat(lang_mask, dim=1) + if self.long_memory is not None: + # after state dropout, so a fused memory loses its proprio with it + cond.memory, cond.memory_mask = self.long_memory( + data["memory"], data["memory_mask"], cond.state + ) + return cond, None + + +class RDT(HPT): + def __init__(self, *args, **kwargs): + for flag in ("ot", "freeze_repr"): + if kwargs.get(flag): + raise ValueError(f"RDT has no HPT trunk; '{flag}' is HPT-only") + super().__init__(*args, **kwargs) + + def _build_policy(self, trunk: dict) -> RDTModel: + return RDTModel(**trunk) + + def _robomimic_to_hpt_data( + self, + batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys=[], + domain="", + ): + """HPT's layout plus ``fps``; per camera with a ``*_hist`` twin, the + (past, current) frame pair on the image input's time axis; and a + ``*_mem`` window as ``memory`` / ``memory_mask``, un-augmented (its + tower is frozen and the stem normalizes).""" + past = { + key[: -len(IMAGE_HISTORY_SUFFIX)]: key + for key in cam_keys + if key.endswith(IMAGE_HISTORY_SUFFIX) and key in batch + } + memory = [k for k in cam_keys if k.endswith(IMAGE_MEMORY_SUFFIX) and k in batch] + if len(memory) > 1: + raise ValueError(f"RDT takes one memory camera, got {memory}") + single = [ + k + for k in cam_keys + if k not in past and k not in past.values() and k not in memory + ] + data = super()._robomimic_to_hpt_data( + batch, single, proprio_keys, lang_keys, ac_key, aux_ac_keys, domain=domain + ) + for key, past_key in past.items(): + short = key.rsplit(".", 1)[-1] + # (B * 2, 3, H, W), each sample's (past, current) adjacent: one + # jitter draw per sample covers both frames, and every op -- a + # contrast mean, a crop, a resize -- still sees one frame. + pair = torch.stack([batch[past_key], batch[key]], dim=1) + n = pair.shape[0] + pair = self._apply_image_augs(pair.flatten(0, 1), short, frames=2) + data[short] = pair.reshape(n, 2, *pair.shape[1:]).unsqueeze(2) + for key in memory: + data["memory"] = batch[key] + data["memory_mask"] = batch[f"{key}{HISTORY_MASK_SUFFIX}"] > 0.5 + data["fps"] = batch["fps"] + return data diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py new file mode 100644 index 0000000000..55f444e3f2 --- /dev/null +++ b/egomimic/eval/action_metrics.py @@ -0,0 +1,228 @@ +"""Layout-aware action metrics shared by the PI and HPT evaluators. + +Every metric here is computed on the model's native (unnormalized) output and +is invariant to the rigid transform the evaluator's cam-frame revert applies +to prediction and ground truth alike, so it never needs the revert. + +Two layouts are recognised by width: + +- bimanual cartesian (12/14 ypr, 18/20 rot6d; see + ``pose_utils.bimanual_cartesian_layout``): position MSE, geodesic rotation + error, raw rotation-channel MSE, DTW + Frechet on xyz; +- wrist-first bimanual hand keypoints (138 ypr / 144 rot6d wrist; see + ``pose_utils.bimanual_keypoint_layout``): mean keypoint L2 (m), wrist + position MSE, wrist geodesic rotation error, DTW on the keypoint block. + +Unknown widths yield no metrics. +""" + +from __future__ import annotations + +import math + +import torch + +from egomimic.utils.action_utils import _reconstruct_R_from_cols, _ypr_to_matrix +from egomimic.utils.metrics import dtw_distance, frechet_gaussian_over_time +from egomimic.utils.pose_utils import ( + bimanual_cartesian_layout, + bimanual_keypoint_layout, + hand_channel_mask, +) + +_RAD2DEG = 180.0 / math.pi + + +def fill_untracked_from_pred(batch: dict, ac_key: str, pred) -> dict: + """``batch`` with an untracked hand's ``ac_key`` channels (``action_hand_mask``) + taken from ``pred``: they hold fill values, not motion, so they add no error to + the metrics and the viz transforms never see the degenerate fill pose.""" + hands = batch.get("action_hand_mask") + if hands is None or pred is None: + return batch + gt = batch[ac_key] + keep = hand_channel_mask(hands.to(gt.device), gt.shape[-1]) + if keep is None: + return batch + return { + **batch, + ac_key: torch.where(keep[:, None, :] > 0.5, gt, pred.to(gt.device, gt.dtype)), + } + + +def _paired_mse(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: + """Plain elementwise MSE, stateless (no torchmetrics accumulator).""" + return (pred.float() - gt.float()).pow(2).mean() + + +# Horizon segments for the per-segment L2 metrics, as fractions of the chunk: +# early = the first 10 % of timesteps, mid = 10-50 %, late = the second half. +# They separate near-term precision from long-range drift, which one +# chunk-mean error hides (and which per-timestep vs horizon-pooled norm stats +# trade against each other). +_HORIZON_SEGMENTS = (("early", 0.0, 0.1), ("mid", 0.1, 0.5), ("late", 0.5, 1.0)) + + +def _segment_l2(pred: torch.Tensor, gt: torch.Tensor, idx: list, prefix: str) -> dict: + """Mean euclidean error (m) of the xyz triples at ``idx`` per horizon + segment of a ``(B, T, D)`` chunk.""" + T = pred.shape[1] + diff = pred[..., idx].float() - gt[..., idx].float() + l2 = diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1) # (B, T, P) + out = {} + for name, lo, hi in _HORIZON_SEGMENTS: + start = min(int(round(lo * T)), T - 1) + end = max(int(round(hi * T)), start + 1) + out[f"{prefix}_{name}"] = l2[:, start:end].mean() + return out + + +def _split_mse(pred_t: torch.Tensor, gt_t: torch.Tensor): + """(translation MSE, rotation-channel MSE) over a bimanual cartesian + vector, so a translation problem reads apart from a rotation one. The + rotation channels are whatever the layout holds: the continuous 6D + rotation-matrix columns for the 18D (human) / 20D (robot) widths, + yaw/pitch/roll for the 12/14D widths. Either way this is a plain + per-channel MSE (pre-Gram-Schmidt, axis-dependent, unitless for 6D; +-pi + wrap and gimbal lock for ypr), i.e. a training-tracking diagnostic, not a + rotation error; see :func:`_rot_geodesic_error` for that. Returns + (None, None) for an unknown width. + """ + layout = bimanual_cartesian_layout(pred_t.shape[-1]) + if layout is None: + return None, None + xyz_idx = list(layout["xyz"]) + rot_idx = list(layout["rot"]) + xyz = _paired_mse(pred_t[..., xyz_idx], gt_t[..., xyz_idx]) + rot = _paired_mse(pred_t[..., rot_idx], gt_t[..., rot_idx]) + return xyz, rot + + +def _geodesic_from_rot_channels(pred: torch.Tensor, gt: torch.Tensor, rot: list): + """Mean geodesic angle (rad) over both arms given the rotation channel + indices of a per-arm-concatenated vector (3 ypr or 6 rot6d per arm).""" + per_arm = len(rot) // 2 + # float64: arccos near 1 loses ~sqrt(eps), which in float32 puts a + # ~1e-4 rad floor under identical rotations; double makes it ~1e-8. + pred = pred.double() + gt = gt.double() + errs = [] + for arm in (rot[:per_arm], rot[per_arm:]): + p, g = pred[..., arm], gt[..., arm] + if per_arm == 3: + Rp, Rg = _ypr_to_matrix(p), _ypr_to_matrix(g) + else: + Rp = _reconstruct_R_from_cols(p[..., 0:3], p[..., 3:6]) + Rg = _reconstruct_R_from_cols(g[..., 0:3], g[..., 3:6]) + tr = (Rp.transpose(-1, -2) @ Rg).diagonal(dim1=-2, dim2=-1).sum(-1) + errs.append(torch.arccos(((tr - 1.0) / 2.0).clamp(-1.0, 1.0))) + return torch.stack(errs, dim=-1).mean().float() + + +def _rot_geodesic_error(pred: torch.Tensor, gt: torch.Tensor): + """Mean geodesic rotation error in radians over batch / time / both arms + of a bimanual cartesian vector. + + Euler-free: per arm, builds proper rotation matrices (``_ypr_to_matrix`` + for the 12/14-dim ypr widths; Gram-Schmidt on the two 6D columns for the + 18/20-dim widths, the same reconstruction the model decode uses) and + takes ``arccos((tr(R_pred^T R_gt) - 1) / 2)``. Unlike a ypr MSE this is + immune to the +-pi wrap AND to the yaw/roll degeneracy at pitch ~ +-pi/2. + Returns None for an unknown width. + """ + layout = bimanual_cartesian_layout(pred.shape[-1]) + if layout is None: + return None + return _geodesic_from_rot_channels(pred, gt, list(layout["rot"])) + + +def cartesian_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a bimanual cartesian chunk ``(B, T, D)``; ``{}`` if ``D`` + is not a cartesian width. ``pred``/``gt`` may live on any device.""" + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_cartesian_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + metrics = {} + xyz_idx = list(layout["xyz"]) + xyz_p, rot_p = _split_mse(pred_cpu, gt_cpu) + metrics[f"{prefix}_xyz_paired_mse_avg"] = xyz_p + metrics[f"{prefix}_rot6d_paired_mse_avg"] = rot_p + metrics[f"{prefix}_xyz_final_mse_avg"] = _paired_mse( + pred_cpu[:, -1, xyz_idx], gt_cpu[:, -1, xyz_idx] + ) + metrics.update(_segment_l2(pred_cpu, gt_cpu, xyz_idx, f"{prefix}_xyz_l2")) + # Distributional / alignment metrics on the position channels only + # (metres), so a unitless 6D column never trades off against a metre. + # DTW forgives temporal misalignment (a correct motion executed early / + # late scores near zero where paired MSE penalizes it); Frechet compares + # the time-distribution shape of the motion. + pred_xyz = pred[..., xyz_idx] + gt_xyz = gt[..., xyz_idx].to(pred_xyz.device) + metrics[f"{prefix}_xyz_dtw_avg"] = dtw_distance(pred_xyz, gt_xyz).mean().item() + metrics[f"{prefix}_xyz_frechet_gauss_avg"] = ( + frechet_gaussian_over_time(pred_xyz, gt_xyz).mean().item() + ) + # Rotation error on the manifold, in degrees: what the decode actually + # produces after Gram-Schmidt, free of wrap / gimbal artefacts. + metrics[f"{prefix}_rot_err_deg_avg"] = ( + _rot_geodesic_error(pred_cpu, gt_cpu) * _RAD2DEG + ) + metrics[f"{prefix}_rot_err_deg_final"] = ( + _rot_geodesic_error(pred_cpu[:, -1], gt_cpu[:, -1]) * _RAD2DEG + ) + return metrics + + +def keypoint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a wrist-first bimanual keypoint chunk ``(B, T, D)`` + (D = 138 ypr wrist or 144 rot6d wrist); ``{}`` for other widths. + + ``kp_l2`` is the mean euclidean keypoint error in metres, in whatever + frame the layout holds (the wrist frame for the ``*_wristframe_*`` modes, + head frame otherwise); ``wrist_xyz`` / ``wrist_rot_err_deg`` score the + wrist pose channels; ``kp_dtw`` is DTW over the flattened keypoint block. + """ + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_keypoint_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + kp_idx = list(layout["keypoints"]) + wrist_idx = list(layout["wrist_xyz"]) + rot_idx = list(layout["rot"]) + + def _kp_l2(p, g): + diff = p[..., kp_idx].float() - g[..., kp_idx].float() + return diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1).mean() + + metrics = { + f"{prefix}_kp_l2_avg": _kp_l2(pred_cpu, gt_cpu), + f"{prefix}_kp_l2_final": _kp_l2(pred_cpu[:, -1], gt_cpu[:, -1]), + f"{prefix}_wrist_xyz_paired_mse_avg": _paired_mse( + pred_cpu[..., wrist_idx], gt_cpu[..., wrist_idx] + ), + f"{prefix}_wrist_xyz_final_mse_avg": _paired_mse( + pred_cpu[:, -1, wrist_idx], gt_cpu[:, -1, wrist_idx] + ), + f"{prefix}_wrist_rot_err_deg_avg": ( + _geodesic_from_rot_channels(pred_cpu, gt_cpu, rot_idx) * _RAD2DEG + ), + f"{prefix}_wrist_rot_err_deg_final": ( + _geodesic_from_rot_channels(pred_cpu[:, -1], gt_cpu[:, -1], rot_idx) + * _RAD2DEG + ), + } + metrics.update(_segment_l2(pred_cpu, gt_cpu, kp_idx, f"{prefix}_kp_l2")) + metrics.update(_segment_l2(pred_cpu, gt_cpu, wrist_idx, f"{prefix}_wrist_xyz_l2")) + pred_kp = pred[..., kp_idx] + gt_kp = gt[..., kp_idx].to(pred_kp.device) + metrics[f"{prefix}_kp_dtw_avg"] = dtw_distance(pred_kp, gt_kp).mean().item() + return metrics + + +def layout_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Dispatch on the last-dim width: cartesian, keypoint, or ``{}``.""" + metrics = cartesian_metrics(pred, gt, prefix) + if metrics: + return metrics + return keypoint_metrics(pred, gt, prefix) diff --git a/egomimic/eval/eval.py b/egomimic/eval/eval.py index aff032cf33..174d8947fd 100644 --- a/egomimic/eval/eval.py +++ b/egomimic/eval/eval.py @@ -18,5 +18,8 @@ def on_validation_end(self): pass @abstractmethod - def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + """``mode``: "metrics" / "video" / "both". An evaluator that renders + nothing may ignore it, but it must accept it -- ModelWrapper passes it + for every loader of a config that declares ``video_episodes``.""" pass diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index ee7af2a0e2..715ec99aee 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -3,9 +3,10 @@ import torch from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import fill_untracked_from_pred, layout_metrics from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment -from egomimic.utils.metrics import frechet_gaussian_over_time, reverse_kl_from_samples +from egomimic.utils.metrics import frechet_gaussian_over_time class HPTEvalVideo(EvalVideo): @@ -13,14 +14,15 @@ class HPTEvalVideo(EvalVideo): Eval class for HPT models. Per embodiment, computes: - val loss (BC loss, same as training; also aggregated as ``Valid/action_loss``) - paired/final MSE + Frechet over time for the main / shared / auxiliary heads + - the layout metrics of :mod:`egomimic.eval.action_metrics` on the main + ``ac_key`` (position / rotation / keypoint errors, frame-invariant) - paired/final MSE in cam frame on the main ``ac_key``, when a - ``transform_lists`` entry is configured - - optional Reverse KL from samples + ``transform_lists`` entry is configured (viz batches only) The revert transform is applied once and reused by both the cam-frame MSE and the viz video. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) @@ -33,6 +35,9 @@ def compute_metrics_and_viz(self, batch): _batch = algo.norm_stats.unnormalize(_batch, embodiment_id) embodiment_name = get_embodiment(embodiment_id).lower() ac_key = algo.ac_keys[embodiment_id] + _batch = fill_untracked_from_pred( + _batch, ac_key, preds.get(f"{embodiment_name}_{ac_key}") + ) loss_key = f"{embodiment_name}_loss" if loss_key in preds: @@ -44,24 +49,20 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if f"{embodiment_name}_{ac_key}" in preds and ac_key != algo.shared_ac_key: - metrics[f"Valid/{embodiment_name}_{ac_key}_paired_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"].cpu(), _batch[ac_key].cpu() - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_final_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"][:, -1].cpu(), - _batch[ac_key][:, -1].cpu(), - ) - fd = frechet_gaussian_over_time( - preds[f"{embodiment_name}_{ac_key}"], _batch[ac_key] - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_avg"] = ( - fd.mean().item() + pred_key = f"{embodiment_name}_{ac_key}" + metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( + preds[pred_key].cpu(), _batch[ac_key].cpu() ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_min"] = ( - fd.min().item() + metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( + preds[pred_key][:, -1].cpu().contiguous(), + _batch[ac_key][:, -1].cpu().contiguous(), ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_max"] = ( - fd.max().item() + fd = frechet_gaussian_over_time(preds[pred_key], _batch[ac_key]) + metrics[f"Valid/{pred_key}_frechet_gauss_avg"] = fd.mean().item() + metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() + metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() + metrics.update( + layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") ) if embodiment_name in algo.auxiliary_ac_keys: @@ -72,7 +73,8 @@ def compute_metrics_and_viz(self, batch): preds[pred_key].cpu(), _batch[aux_key].cpu() ) metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - preds[pred_key][:, -1].cpu(), _batch[aux_key][:, -1].cpu() + preds[pred_key][:, -1].cpu().contiguous(), + _batch[aux_key][:, -1].cpu().contiguous(), ) fd = frechet_gaussian_over_time( preds[pred_key], _batch[aux_key] @@ -92,8 +94,8 @@ def compute_metrics_and_viz(self, batch): preds[pred_key].cpu(), _batch[algo.shared_ac_key].cpu() ) metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - preds[pred_key][:, -1].cpu(), - _batch[algo.shared_ac_key][:, -1].cpu(), + preds[pred_key][:, -1].cpu().contiguous(), + _batch[algo.shared_ac_key][:, -1].cpu().contiguous(), ) fd = frechet_gaussian_over_time( preds[pred_key], _batch[algo.shared_ac_key] @@ -102,62 +104,8 @@ def compute_metrics_and_viz(self, batch): metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() - if algo.rkl_samples and algo.rkl_samples > 1: - hpt_batch = { - "domain": embodiment_name, - "data": algo._robomimic_to_hpt_data( - batch[embodiment_id], - algo.camera_keys[embodiment_id], - algo.proprio_keys[embodiment_id], - algo.lang_keys[embodiment_id], - ac_key, - algo.auxiliary_ac_keys.get(embodiment_name, []), - ), - } - rkl_targets = [] - - if ( - f"{embodiment_name}_{ac_key}" in preds - and ac_key != algo.shared_ac_key - ): - rkl_targets.append( - ( - f"{embodiment_name}_{ac_key}", - _batch[ac_key].to(algo.device), - embodiment_name, - ) - ) - - if embodiment_name in algo.auxiliary_ac_keys: - for aux_key in algo.auxiliary_ac_keys[embodiment_name]: - aux_pred_key = f"{embodiment_name}_{aux_key}" - if aux_pred_key in preds: - rkl_targets.append( - ( - aux_pred_key, - _batch[aux_key].to(algo.device), - aux_key, - ) - ) - - if algo.shared_ac_key: - shared_pred_key = f"{embodiment_name}_{algo.shared_ac_key}" - if shared_pred_key in preds: - rkl_targets.append( - ( - shared_pred_key, - _batch[algo.shared_ac_key].to(algo.device), - "shared", - ) - ) - - M = int(algo.rkl_samples) - for pred_key_name, gt_tensor, head_key in rkl_targets: - samples = self._collect_policy_samples( - hpt_batch, ref=gt_tensor, key_name=head_key, M=M - ) - rkl = reverse_kl_from_samples(samples, gt_tensor) - metrics[f"Valid/{pred_key_name}_reverse_kl_M{M}"] = rkl.item() + if not do_viz: + continue transform_list = self.transform_lists.get(embodiment_name) main_pred_key = f"{embodiment_name}_{ac_key}" @@ -195,33 +143,3 @@ def compute_metrics_and_viz(self, batch): metrics["Valid/action_loss"] = total_loss / n_loss_embodiments return metrics, images_dict - - def _visualize_preds(self, predictions, batch): - if self.viz_func is None: - raise ValueError("viz_func is not set") - embodiment_id = batch["embodiment"][0].item() - embodiment_name = get_embodiment(embodiment_id).lower() - return self.viz_func[embodiment_name](predictions, batch) - - @torch.no_grad() - def _collect_policy_samples(self, hpt_batch, ref, key_name, M): - """Collect policy samples for Reverse KL.""" - algo = self.model - B, T, D = ref.shape - samples = [] - was_training = algo.nets.training - algo.nets.eval() - for _ in range(M): - out = algo.nets["policy"].forward( - hpt_batch["domain"], algo._clone_batch(hpt_batch["data"]) - ) - if key_name in out: - pred = out[key_name] - else: - pred = out[hpt_batch["domain"]] - - pred = pred[:, :T, :D] - samples.append(pred.unsqueeze(0)) - if was_training: - algo.nets.train() - return torch.cat(samples, dim=0) diff --git a/egomimic/eval/eval_latent.py b/egomimic/eval/eval_latent.py index 604900ce63..f25cb65977 100644 --- a/egomimic/eval/eval_latent.py +++ b/egomimic/eval/eval_latent.py @@ -61,7 +61,8 @@ def __init__( emit_combined: bool = True, color_by: str = "embodiment", # "embodiment" or "hash" ): - super().__init__(limit_val_batches=limit_val_batches) + # Its own _visualize_preds renders overlay frames only. + super().__init__(limit_val_batches=limit_val_batches, viz_mode="overlay") self.compute_umap = compute_umap self.compute_tsne_2d = compute_tsne_2d self.compute_tsne_3d = compute_tsne_3d @@ -276,8 +277,8 @@ def compute_metrics_and_viz(self, batch, do_viz=True): unnorm_preds[pred_key].cpu(), unnorm_batch[ac_key].cpu() ) metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - unnorm_preds[pred_key][:, -1].cpu(), - unnorm_batch[ac_key][:, -1].cpu(), + unnorm_preds[pred_key][:, -1].cpu().contiguous(), + unnorm_batch[ac_key][:, -1].cpu().contiguous(), ) if do_viz and self.viz_func is not None: diff --git a/egomimic/eval/eval_pi.py b/egomimic/eval/eval_pi.py index 7206197701..4742a5eafb 100644 --- a/egomimic/eval/eval_pi.py +++ b/egomimic/eval/eval_pi.py @@ -1,29 +1,44 @@ import copy import torch -from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import ( # noqa: F401 (re-exported for tests) + _paired_mse, + _rot_geodesic_error, + _split_mse, + fill_untracked_from_pred, + layout_metrics, +) from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment class PIEvalVideo(EvalVideo): """ - Eval class for PI models. Per embodiment, computes: - - val loss (flow-matching loss, same as training; also aggregated as ``Valid/action_loss``) - - paired/final MSE in the model's native wrist frame - - paired/final MSE in cam frame, when a ``transform_lists`` entry is configured - The revert transform is applied once and reused for both the cam-frame MSE - and the viz video. + Eval class for PI models. Per embodiment, on the model's native + (unnormalized) output, computes: + - ``action_loss``: flow-matching val loss, same as training (also logged + per embodiment as ``_loss``) + - the layout metrics of :mod:`egomimic.eval.action_metrics`: + cartesian widths get ``xyz_paired_mse_avg`` / ``xyz_final_mse_avg`` + (m^2), ``rot_err_deg_avg`` / ``rot_err_deg_final`` (geodesic, after + Gram-Schmidt), ``rot6d_paired_mse_avg`` (raw rotation channels, a + training-tracking diagnostic), ``xyz_dtw_avg`` and + ``xyz_frechet_gauss_avg``; keypoint widths get ``kp_l2_avg`` / + ``kp_l2_final`` (m), ``wrist_xyz_*_mse``, ``wrist_rot_err_deg_*`` and + ``kp_dtw_avg``. + + Metrics are frame-invariant (the cam-frame revert is a rigid transform + shared by prediction and gt), so the ``transform_lists`` revert runs only + when rendering the viz video (``do_viz``), never for metrics. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) metrics = {} images_dict = {} - mse = MeanSquaredError() total_loss = None n_loss_embodiments = 0 @@ -33,6 +48,7 @@ def compute_metrics_and_viz(self, batch): ac_key = algo.ac_keys[embodiment_id] pred_key = f"{embodiment_name}_{ac_key}" loss_key = f"{embodiment_name}_loss" + _batch = fill_untracked_from_pred(_batch, ac_key, preds.get(pred_key)) if loss_key in preds: loss_val = preds[loss_key] @@ -43,52 +59,31 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if pred_key in preds: - metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( - preds[pred_key].cpu(), _batch[ac_key].cpu() + metrics.update( + layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") ) - metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - preds[pred_key][:, -1].cpu(), _batch[ac_key][:, -1].cpu() - ) - - transform_list = self.transform_lists.get(embodiment_name) - gt_batch_viz = _batch - preds_for_viz = preds - if transform_list is not None and pred_key in preds: - pred_batch = copy.deepcopy(_batch) - pred_batch[ac_key] = preds[pred_key] - gt_t = Embodiment.apply_transform(_batch, transform_list) - pred_t = Embodiment.apply_transform(pred_batch, transform_list) - # apply_transform drops keys whose shape[0] != batch_size - # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. - gt_batch_viz = {**_batch, **gt_t} - pred_batch_viz = {**_batch, **pred_t} - # ``.contiguous()`` because ``apply_transform`` returns CPU tensors, - # so ``.cpu()`` here is a no-op and ``[:, -1]`` leaves a non-contiguous - # view that torchmetrics' MSE doesn't accept. - metrics[f"Valid/{pred_key}_cam_paired_mse_avg"] = mse( - pred_batch_viz[ac_key].cpu().contiguous(), - gt_batch_viz[ac_key].cpu().contiguous(), - ) - metrics[f"Valid/{pred_key}_cam_final_mse_avg"] = mse( - pred_batch_viz[ac_key][:, -1].cpu().contiguous(), - gt_batch_viz[ac_key][:, -1].cpu().contiguous(), - ) - - preds_for_viz = dict(preds) - preds_for_viz[pred_key] = pred_batch_viz[ac_key] - - ims = self._visualize_preds(preds_for_viz, gt_batch_viz) - images_dict[embodiment_id] = ims + if do_viz: + gt_batch_viz = _batch + preds_for_viz = preds + transform_list = self.transform_lists.get(embodiment_name) + if transform_list is not None and pred_key in preds: + # Revert to cam (head) frame + the ypr layout for the + # overlay only; the metrics above are frame-invariant. + pred_batch = copy.deepcopy(_batch) + pred_batch[ac_key] = preds[pred_key] + gt_t = Embodiment.apply_transform(_batch, transform_list) + pred_t = Embodiment.apply_transform(pred_batch, transform_list) + # apply_transform drops keys whose shape[0] != batch_size + # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. + gt_batch_viz = {**_batch, **gt_t} + pred_batch_viz = {**_batch, **pred_t} + preds_for_viz = dict(preds) + preds_for_viz[pred_key] = pred_batch_viz[ac_key] + ims = self._visualize_preds(preds_for_viz, gt_batch_viz) + images_dict[embodiment_id] = ims if total_loss is not None and n_loss_embodiments > 0: metrics["Valid/action_loss"] = total_loss / n_loss_embodiments return metrics, images_dict - - def _visualize_preds(self, predictions, batch): - if self.viz_func is None: - raise ValueError("viz_func is not set") - embodiment_id = batch["embodiment"][0].item() - embodiment_name = get_embodiment(embodiment_id).lower() - return self.viz_func[embodiment_name](predictions, batch) diff --git a/egomimic/eval/eval_train_viz.py b/egomimic/eval/eval_train_viz.py new file mode 100644 index 0000000000..6a803a453b --- /dev/null +++ b/egomimic/eval/eval_train_viz.py @@ -0,0 +1,74 @@ +"""Train-set visualization evaluator. + +Wraps a concrete EvalVideo (PIEvalVideo, HPTEvalVideo, ...) so the same +forward/metric/viz logic can run a second time against a separate +``train_viz`` dataloader. Videos go to ``/videos_train_viz/`` and +metric keys are prefixed with ``train_viz/`` so they don't collide with the +canonical ``Valid/...`` keys. + +Instantiated via Hydra from a config like +``hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml`` and enabled with +``+evaluator@train_viz_evaluator=`` on a data config that +defines ``train_viz_datasets``. +""" + +from __future__ import annotations + +import os + +from egomimic.eval.eval_video import EvalVideo + + +class TrainVizEvalVideo(EvalVideo): + def __init__( + self, base: EvalVideo, limit_val_batches: int = 50, prefix: str = "train_viz" + ): + """``prefix`` names the metric-key prefix and the ``videos_/`` + subdirectory; trainHydra reuses this wrapper for the opsplit heads + (``unseen_op_valid``, and ``seen_op_valid`` via ``data.valid_prefix``).""" + # `base` must be set before super().__init__: the trainer/model + # property setters fire on the base attribute during construction. + self.base = base + self.prefix = prefix + super().__init__( + limit_val_batches=limit_val_batches, + viz_func=base.viz_func, + transform_lists=base.transform_lists, + viz_every_n_epochs=base.viz_every_n_epochs, + viz_max_batches=base.viz_max_batches, + viz_mode=base.viz_mode, + replay_chunks=base.replay_chunks, + replay_trail=base.replay_trail, + action_stride=base.action_stride, + ) + + def _set_replay_now(self, value: bool) -> None: + # The base renders the batch, so it must know which mode this step is. + super()._set_replay_now(value) + self.base._set_replay_now(value) + + @property + def trainer(self): + return self._trainer + + @trainer.setter + def trainer(self, value): + self._trainer = value + self.base.trainer = value + + @property + def model(self): + return self._model + + @model.setter + def model(self, value): + self._model = value + self.base.model = value + + def video_dir(self): + return os.path.join(self.root_dir(), f"videos_{self.prefix}") + + def compute_metrics_and_viz(self, batch, do_viz=True): + metrics, images_dict = self.base.compute_metrics_and_viz(batch, do_viz=do_viz) + metrics = {f"{self.prefix}/{k}": v for k, v in metrics.items()} + return metrics, images_dict diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 9cf44d121c..d524c0203c 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -3,10 +3,22 @@ import torch import torchvision.io as tvio +from lightning.pytorch.loggers import WandbLogger from egomimic.eval.eval import Eval +from egomimic.eval.replay_viz import ReplayClip, infer_stride, render_replay from egomimic.rldb.embodiment.embodiment import get_embodiment +# viz_gt_preds kwargs that ReplayClip / render_replay consume themselves. +_REPLAY_OWN_KWARGS = { + "image_key", + "action_key", + "annotation_key", + "mode", + "gt_alpha", + "pred_alpha", +} + class EvalVideo(Eval): """ @@ -20,14 +32,47 @@ def __init__( limit_val_batches: int = 400, viz_func: dict = None, transform_lists: dict | None = None, + viz_every_n_epochs: int = 1, + viz_max_batches: int | None = None, + viz_mode: str = "replay", + replay_chunks: int = 5, + replay_trail: int = 5, + action_stride: float | None = None, ): super().__init__() self.trainer = None self.model = None self.viz_func = viz_func + # Render the overlay video only on validation passes where + # (current_epoch + 1) is a multiple of this; metrics log every pass. + self.viz_every_n_epochs = viz_every_n_epochs + # On viz epochs, render overlay frames for only the first N val + # batches (None = all). Rendering is CPU-bound -- ~0.1 s/frame, 4864 + # frames in 8.4 min with four ranks rendering at once -- and dominates + # viz-epoch wall time; metrics are still computed on EVERY batch + # regardless. Required on heads with pinned video loaders (trainHydra). + self.viz_max_batches = viz_max_batches + # "replay" (egomimic/eval/replay_viz.py) plays each chunk as GT then + # prediction; it needs consecutive frames, so it runs only on a pinned + # video loader and caps it at ``replay_chunks`` chunks instead of + # ``viz_max_batches``. "overlay" draws each frame's whole GT and + # predicted chunk, and is what a head without a video loader gets. + # ``action_stride`` is frames per chunk step, fractional when the chunk + # is resampled (None: inferred from the GT keypoints, see + # replay_viz.infer_stride). + if viz_mode not in ("overlay", "replay"): + raise ValueError(f"unknown viz_mode {viz_mode!r}") + self.viz_mode = viz_mode + self.replay_chunks = replay_chunks + self._replay_now = False + self._replay_seen = {} + self._replay_stride = {} + self._replay_horizon = {} + self.replay_trail = replay_trail + self.action_stride = action_stride # Per-embodiment list[Transform] applied once during eval to project - # the model's wrist-frame actions back into cam (head) frame. Reused for - # both cam-frame MSE and the viz video so we don't transform twice. + # the model's wrist-frame actions back into cam (head) frame for the + # viz video. self.transform_lists = transform_lists or {} self.val_image_buffer = {} self.val_counter = {} @@ -44,52 +89,191 @@ def __init__( def video_dir(self): return os.path.join(self.root_dir(), "videos") + def _should_viz(self) -> bool: + if not self.viz_every_n_epochs or self.viz_every_n_epochs <= 0: + return False + # Lightning runs validation when (current_epoch + 1) is a multiple of + # check_val_every_n_epoch, with current_epoch still the pre-increment + # value during the val hooks (19, 39, ...). Gate on the same (+1) + # convention: a plain `current_epoch % n` never aligns with a + # validation epoch unless check_val_every_n_epoch == 1. + epoch = self.trainer.current_epoch + 1 + # The last epoch always renders, so short runs (trainer=debug, eval + # mode's single validate) get a video without retuning the interval. + max_epochs = getattr(self.trainer, "max_epochs", None) + if max_epochs is not None and max_epochs > 0 and epoch == max_epochs: + return True + return epoch % self.viz_every_n_epochs == 0 + @abstractmethod - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): """ Run the model's eval forward and compute metrics and visualization frames. Args: batch (dict): processed batch produced by the algo's `process_batch_for_training`. + do_viz (bool): render overlay frames for this batch. When False the + returned ``images_dict`` must be empty. Returns: metrics (dict[str, torch.Tensor | float]) images_dict (dict[embodiment_id, np.ndarray (B, H, W, 3)]) """ raise NotImplementedError + def _video_fps(self, source_fps: int = 30) -> int: + """Playback fps of the overlay video: the source rate. + + The val heads come back from ``val_dataloader()`` inside + ``CombinedLoader``s, which Lightning does NOT wrap in a + ``DistributedSampler`` (every rank runs every val batch; measured + 2026-09-16: 1 rank = 223 val steps, 8 ranks = 8 x 223), so rank 0 + renders every source frame and the video plays in real time at the + source rate. An earlier version divided by ``world_size`` to undo a + sampler stride that never applied, which made every multi-GPU video a + ``world_size``-times slow-motion (416 frames at 4 fps instead of 30). + """ + return int(source_fps) + + def _visualize_preds(self, predictions, batch): + if self.viz_func is None: + raise ValueError("viz_func is not set") + embodiment_name = get_embodiment(batch["embodiment"][0].item()).lower() + viz_fn = self.viz_func[embodiment_name] + if self._replay_now: + return ReplayClip.from_batch(viz_fn, predictions, batch) + return viz_fn(predictions, batch) + + def _render_replay(self, key, final: bool): + """Render the buffered clip's complete chunks; keep the rest unless final.""" + clips = self.val_image_buffer[key] + if not clips: + return None + clip = ReplayClip.concat(clips) + viz_fn = self.viz_func[get_embodiment(key).lower()] + frames, used = render_replay( + clip, + embodiment_cls=viz_fn.func.__self__, + mode=viz_fn.keywords["mode"], + stride=self._replay_stride.get(key, self.action_stride), + trail=self.replay_trail, + viz_kwargs={ + k: v for k, v in viz_fn.keywords.items() if k not in _REPLAY_OWN_KWARGS + }, + ) + self.val_image_buffer[key] = [] if final else [clip.tail(used)] + return torch.from_numpy(frames) if len(frames) else None + + def _set_replay_now(self, value: bool) -> None: + self._replay_now = value + + def _is_replay_buffer(self, key) -> bool: + buffer = self.val_image_buffer[key] + return bool(buffer) and isinstance(buffer[0], ReplayClip) + + def _replay_done(self) -> bool: + return bool(self._replay_seen) and all( + k in self._replay_stride + and seen + >= self.replay_chunks * self._replay_horizon[k] * self._replay_stride[k] + for k, seen in self._replay_seen.items() + ) + + def _buffered_frames(self, key) -> int: + # Overlay buffers hold one (H, W, 3) tensor per frame, not clips. + if self._is_replay_buffer(key): + return sum(len(c) for c in self.val_image_buffer[key]) + return len(self.val_image_buffer[key]) + + def _flush(self, key, final: bool) -> None: + if self._is_replay_buffer(key): + frames = self._render_replay(key, final) + else: + buffer = self.val_image_buffer[key] + frames = torch.stack(buffer) if buffer else None + self.val_image_buffer[key] = [] + if frames is not None: + self._write_video(key, frames) + self.val_counter[key] += 1 + + def _write_video(self, key, frames) -> None: + # Every rank that gets here renders the same frames to the same path + # (see _video_fps on why no DistributedSampler splits them), so without + # this the file is whatever the last rank to finish wrote. + if not self.trainer.is_global_zero: + return + path = os.path.join( + self.video_dir(), + f"epoch_{self.trainer.current_epoch}", + str(get_embodiment(key)), + f"validation_video_{self.val_counter[key]}.mp4", + ) + os.makedirs(os.path.dirname(path), exist_ok=True) + tvio.write_video(path, frames, fps=self._video_fps(), video_codec="h264") + self._upload_video(key, path) + + def _upload_video(self, key, path) -> None: + name = f"{os.path.basename(self.video_dir())}/{get_embodiment(key)}" + if self.val_counter[key]: + name += f"_{self.val_counter[key]}" + for logger in self.trainer.loggers: + if isinstance(logger, WandbLogger): + logger.log_video( + name, + [path], + step=self.trainer.global_step, + caption=[f"epoch {self.trainer.current_epoch}"], + format=["mp4"], + ) + def on_validation_start(self): - if self.trainer.is_global_zero: + if self.trainer.is_global_zero and self._should_viz(): os.makedirs( os.path.join(self.video_dir(), f"epoch_{self.trainer.current_epoch}"), exist_ok=True, ) def on_validation_end(self): - for key, buffer in self.val_image_buffer.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) - if len(buffer) != 0: - frames = torch.stack(buffer) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - + if not self._should_viz(): + return + for key in list(self.val_image_buffer): + self._flush(key, final=True) self.val_counter[key] = 0 self.val_image_buffer[key] = [] + self._replay_seen, self._replay_stride = {}, {} - def on_validation_step(self, batch, batch_idx, dataloader_idx=0): - metrics, images_dict = self.compute_metrics_and_viz(batch) + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + """``mode`` splits metrics from video so one head can serve two loaders: + + * ``"metrics"`` -- the per-episode subsampled metric loader. Never + renders or buffers a frame (rendering is ~1 s/frame of CPU, and a + subsampled video would be an incoherent time-lapse anyway). + * ``"video"`` -- the contiguous pinned-episode loader. Renders and + buffers, and logs NOTHING, so no ``Valid/...`` key is averaged over + it. On a non-viz epoch, or past its cap (``viz_max_batches``, or + ``replay_chunks`` chunks in replay mode), it returns before the + forward pass: there is nothing such a call could produce. + * ``"both"`` (default) -- today's single-loader behaviour, which is what + a data config without ``video_episodes`` still gets. + """ + if mode not in ("both", "metrics", "video"): + raise ValueError(f"unknown validation mode {mode!r}") + replay = mode == "video" and self.viz_mode == "replay" + self._set_replay_now(replay) + if replay: + # viz_max_batches still bounds a loader that yields no clips. + nothing_yet = not self._replay_seen and ( + self.viz_max_batches is not None and batch_idx >= self.viz_max_batches + ) + past_viz_cap = self._replay_done() or nothing_yet + else: + past_viz_cap = ( + self.viz_max_batches is not None and batch_idx >= self.viz_max_batches + ) + if mode == "video" and (past_viz_cap or not self._should_viz()): + return + do_viz = mode != "metrics" and self._should_viz() and not past_viz_cap + metrics, images_dict = self.compute_metrics_and_viz(batch, do_viz=do_viz) device = self.trainer.lightning_module.device metrics = { @@ -97,30 +281,39 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0): for k, v in metrics.items() } - ## images is now a dict - for key, images in images_dict.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) - if key not in self.val_image_buffer or self.val_image_buffer[key] is None: - self.val_image_buffer[key] = [] - self.val_counter[key] = 0 - self.val_image_buffer[key].extend(torch.from_numpy(images)) - if len(self.val_image_buffer[key]) >= 1000: - frames = torch.stack(self.val_image_buffer[key]) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - self.val_image_buffer[key].clear() - self.val_counter[key] += 1 + if do_viz: + for key, images in images_dict.items(): + if ( + key not in self.val_image_buffer + or self.val_image_buffer[key] is None + ): + self.val_image_buffer[key] = [] + self.val_counter[key] = 0 + if isinstance(images, ReplayClip): + self.val_image_buffer[key].append(images) + seen = self._replay_seen.get(key, 0) + len(images) + self._replay_seen[key] = seen + self._replay_horizon[key] = images.gt.shape[1] + # infer_stride compares frames up to 32 on, from up to 5 anchors. + if key not in self._replay_stride and seen >= 40: + buffered = ReplayClip.concat(self.val_image_buffer[key]) + self._replay_stride[key] = ( + self.action_stride or infer_stride(buffered) or 1 + ) + else: + self.val_image_buffer[key].extend(torch.from_numpy(images)) + if self._buffered_frames(key) >= 1000: + self._flush(key, final=False) + + if mode == "video": + # Video-only loader: the forward pass was for the overlay. Logging + # here would mix pinned-episode numbers into the head's metric and + # (with add_dataloader_idx=False) collide key-for-key with it. + return - self.trainer.lightning_module.log_dict(metrics, sync_dist=True) + # add_dataloader_idx=False: with the train_viz loader Lightning would + # otherwise suffix every key with "/dataloader_idx_N"; the train-viz + # wrapper prefixes its keys itself. + self.trainer.lightning_module.log_dict( + metrics, sync_dist=True, add_dataloader_idx=False + ) 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/eval/replay_viz.py b/egomimic/eval/replay_viz.py new file mode 100644 index 0000000000..c3e70b5834 --- /dev/null +++ b/egomimic/eval/replay_viz.py @@ -0,0 +1,257 @@ +"""GT / prediction replay video: play each action chunk twice over the frames it +spans, first drawing where the ground-truth hand is at each step, then where +the prediction puts it. + +``stride`` is video frames per chunk step, and need not be an integer: the +keypoint pipeline resamples 30 raw frames to a 100-step chunk, so a mecka +(action stride 1) step is 29/99 of a frame. Chunks are anchored every +``ceil(H * stride)`` frames and do not overlap; video frame ``a + j`` of the +chunk anchored at ``a`` shows the latest step at or before it, drawn on that +step's nearest frame ``a + round(k * stride)``. + +A chunk lives in its anchor frame's camera, and the head moves while it plays. +Each step is mapped into its frame's camera with a rigid fit of the GT +keypoints (``gt_a[k]`` in cam ``a`` against ``gt_f[0]`` in cam ``f``, the same +hand at the same instant), so it needs no head pose. That fit exists only for +the 126-D head-frame keypoint layout; any other layout is drawn uncompensated, +which is exact for Eva's fixed camera. +""" + +from dataclasses import dataclass + +import cv2 +import numpy as np + +from egomimic.rldb.embodiment.embodiment import _intrinsics_from_batch, get_embodiment +from egomimic.utils.pose_utils import cam_frame_to_cam_pixels +from egomimic.utils.type_utils import _to_numpy +from egomimic.utils.viz_utils import _prepare_viz_image + +_KP_WIDTH = 126 + + +@dataclass +class ReplayClip: + """Consecutive video frames with the GT and predicted chunk of each.""" + + images: np.ndarray # (N, H, W, 3) uint8 + gt: np.ndarray # (N, T, D) + pred: np.ndarray # (N, T, D) + intrinsics: list # N x (K | None) + + def __len__(self): + return len(self.images) + + @classmethod + def from_batch(cls, viz_fn, predictions, batch): + """Pull what ``viz_fn`` (a ``viz_gt_preds`` partial) would draw.""" + kw = viz_fn.keywords + name = get_embodiment(batch["embodiment"][0].item()).lower() + images = _to_numpy(batch[kw["image_key"]]) + return cls( + images=np.stack([_prepare_viz_image(im) for im in images]), + gt=_to_numpy(batch[kw["action_key"]]).astype(np.float32), + pred=_to_numpy(predictions[f"{name}_{kw['action_key']}"]).astype( + np.float32 + ), + intrinsics=[_intrinsics_from_batch(batch, i) for i in range(len(images))], + ) + + @classmethod + def concat(cls, clips): + return cls( + images=np.concatenate([c.images for c in clips]), + gt=np.concatenate([c.gt for c in clips]), + pred=np.concatenate([c.pred for c in clips]), + intrinsics=[k for c in clips for k in c.intrinsics], + ) + + def tail(self, start): + return ReplayClip( + self.images[start:], + self.gt[start:], + self.pred[start:], + self.intrinsics[start:], + ) + + +def _valid_points(x): + return np.isfinite(x).all(-1) & (np.abs(x) < 1e8).all(-1) + + +def rigid_fit(src, dst): + """Least-squares ``R, t`` with ``dst ~= src @ R.T + t`` over valid pairs.""" + ok = _valid_points(src) & _valid_points(dst) + if ok.sum() < 3: + return np.eye(3), np.zeros(3) + src, dst = src[ok].astype(np.float64), dst[ok].astype(np.float64) + mu_s, mu_d = src.mean(0), dst.mean(0) + U, _, Vt = np.linalg.svd((src - mu_s).T @ (dst - mu_d)) + d = np.sign(np.linalg.det(Vt.T @ U.T)) + R = Vt.T @ np.diag([1.0, 1.0, d]) @ U.T + return R, mu_d - mu_s @ R.T + + +def _interp_step(chunk, p): + lo = min(int(np.floor(p)), len(chunk) - 2) + w = p - lo + return (1 - w) * chunk[lo] + w * chunk[lo + 1] + + +def infer_stride(clip: ReplayClip, max_stride: int = 4, offsets: int = 8): + """Frames per chunk step, from the GT alone: the chunk anchored at ``a``, + read at step ``j / stride``, is the same hand as frame ``a + j``'s own + step 0, up to a rigid camera motion, so the right stride has the smallest + fit residual. Candidates are ``D / (H - 1)`` for a whole number ``D`` of + frames spanned by the chunk, up to ``max_stride`` frames per step. + + None when the layout has no keypoints to fit or the clip is too short. + """ + if clip.gt.shape[-1] != _KP_WIDTH or len(clip) < 2: + return None + horizon = clip.gt.shape[1] + anchors = range(0, max(1, len(clip) - 32), 4)[:5] + errors = {} + for span in range(1, max_stride * (horizon - 1) + 1): + stride = span / (horizon - 1) + res = [] + for a in anchors: + js = np.unique( + np.linspace(1, min(span, 32, len(clip) - 1 - a), offsets).astype(int) + ) + for j in js[js >= 1]: + src = _interp_step(clip.gt[a], j / stride).reshape(-1, 3) + dst = clip.gt[a + j, 0].reshape(-1, 3) + ok = _valid_points(src) & _valid_points(dst) + if ok.sum() < 3: + continue + R, t = rigid_fit(src, dst) + res.append(np.linalg.norm(src[ok] @ R.T + t - dst[ok], axis=1).mean()) + if res: + errors[stride] = float(np.mean(res)) + return min(errors, key=errors.get) if errors else None + + +def _apply_rigid(chunk, R, t): + pts = chunk.reshape(*chunk.shape[:-1], -1, 3) + return (pts @ R.T + t).reshape(chunk.shape).astype(np.float32) + + +def _label(image, text, color): + h = image.shape[0] + scale = max(0.5, h / 700) + thick = max(1, int(h / 350)) + org = (int(h * 0.03), int(h * 0.03) + int(24 * scale)) + cv2.putText( + image, + text, + org, + cv2.FONT_HERSHEY_SIMPLEX, + scale, + (0, 0, 0), + thick + 2, + cv2.LINE_AA, + ) + cv2.putText( + image, text, org, cv2.FONT_HERSHEY_SIMPLEX, scale, color, thick, cv2.LINE_AA + ) + return image + + +def _draw_skeleton(image, keypoints, intrinsics, edges, color): + """Both hands' 21 MANO keypoints (``(126,)``, camera frame) with bones.""" + vis = np.ascontiguousarray(_prepare_viz_image(image)).copy() + h, w = vis.shape[:2] + pts = keypoints.reshape(2, -1, 3) + ok = _valid_points(pts) & (pts[..., 2] > 0.01) + safe = np.where(ok[..., None], pts, [0.0, 0.0, 1.0]) + px = cam_frame_to_cam_pixels(safe.reshape(-1, 3), intrinsics)[:, :2] + px = np.round(px).astype(np.int32).reshape(2, -1, 2) + ok &= (px[..., 0] >= 0) & (px[..., 0] < w) & (px[..., 1] >= 0) & (px[..., 1] < h) + bone = tuple(int(0.6 * c + 0.4 * 255) for c in color) + for hand in range(2): + for i, j in edges: + if ok[hand, i] and ok[hand, j]: + cv2.line( + vis, tuple(px[hand, i]), tuple(px[hand, j]), bone, 2, cv2.LINE_AA + ) + for i in np.flatnonzero(ok[hand]): + cv2.circle(vis, tuple(px[hand, i]), 3, color, -1, cv2.LINE_AA) + return vis + + +def render_replay( + clip: ReplayClip, + embodiment_cls, + mode: str, + stride: int | None = None, + trail: int = 5, + viz_kwargs: dict | None = None, +): + """Render every complete chunk of ``clip``. + + The 126-D keypoint layout draws each step as both hands' skeletons; + any other layout draws the last ``trail`` steps with ``embodiment_cls.viz``. + + ``stride`` (frames per chunk step, possibly fractional) is inferred when + None; see infer_stride. + + Returns ``(frames (M, H, W, 3) uint8, consumed)``; ``clip.tail(consumed)`` + holds the frames a later clip needs to complete the next chunk. + """ + viz_kwargs = viz_kwargs or {} + if stride is None: + stride = infer_stride(clip) or 1 + n, horizon = len(clip), clip.gt.shape[1] + span = int(np.ceil(horizon * stride - 1e-6)) + keypoints = clip.gt.shape[-1] == _KP_WIDTH + # (j + 0.5) keeps whole strides exact and lands a fractional stride's + # step on frame j itself. + steps = [min(horizon - 1, int((j + 0.5) / stride)) for j in range(span)] + + out = [] + a = 0 + while a + span <= n: + frame = {k: min(a + round(k * stride), a + span - 1) for k in set(steps)} + fits = {} + if keypoints: + for k, f in frame.items(): + fits[k] = rigid_fit( + clip.gt[a, k].reshape(-1, 3), clip.gt[f, 0].reshape(-1, 3) + ) + + for role, chunk, color, rgb in ( + ("GT", clip.gt[a], "Greens", (80, 220, 80)), + ("PRED", clip.pred[a], "Reds", (240, 80, 80)), + ): + drawn = {} + for k in steps: + if k not in drawn: + f = frame[k] + K = clip.intrinsics[f] + if keypoints: + im = _draw_skeleton( + clip.images[f], + _apply_rigid(chunk[k], *fits[k]), + K if K is not None else embodiment_cls.INTRINSICS, + embodiment_cls.FINGER_EDGES, + rgb, + ) + else: + im = embodiment_cls.viz( + clip.images[f], + chunk[max(0, k - trail) : k + 1], + mode=mode, + color=color, + intrinsics=K, + **viz_kwargs, + ) + drawn[k] = _label( + np.ascontiguousarray(im), f"{role} {k + 1}/{horizon}", rgb + ) + out.append(drawn[k]) + a += span + + if not out: + return np.zeros((0, *clip.images.shape[1:]), np.uint8), 0 + return np.stack(out), a diff --git a/egomimic/hydra_configs/callbacks/checkpoints.yaml b/egomimic/hydra_configs/callbacks/checkpoints.yaml index 86d21a792f..6bb0f6f561 100644 --- a/egomimic/hydra_configs/callbacks/checkpoints.yaml +++ b/egomimic/hydra_configs/callbacks/checkpoints.yaml @@ -6,6 +6,14 @@ model_checkpoint: _target_: lightning.pytorch.callbacks.ModelCheckpoint dirpath: ${paths.output_dir}/checkpoints filename: "epoch_{epoch}" - save_last: true + # A symlink to the epoch checkpoint, not a second full copy of it (~21 GB on + # Pi 0.5). Requires save_top_k: -1 below, which is what keeps the target + # alive; relinks every epoch, not only checkpoint epochs. A requeue refuses to + # start over when last.ckpt dangles or is missing next to other checkpoints. + save_last: link save_top_k: -1 - every_n_epochs: 100 \ No newline at end of file + every_n_epochs: 50 + # Save at the end of the train epoch, not after validation: with + # check_val_every_n_epoch > 1 the validation-end hook only fires on val + # epochs and every_n_epochs would silently miss the others. + save_on_train_epoch_end: true diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index ef1550dd1f..dbd5f47d17 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -1,3 +1,6 @@ +# One pinned aria episode (stored as human_bimanual) on the 144-D wrist-frame +# hand-keypoint action; see data/human.yaml for the layout and the cartesian +# override. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,10 +11,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/aria_keypoints.yaml b/egomimic/hydra_configs/data/aria_keypoints.yaml index 59d0a3ab62..7ed22d6eb7 100644 --- a/egomimic/hydra_configs/data/aria_keypoints.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints.yaml @@ -11,7 +11,7 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_headframe_ypr + mode: keypoints_headframe_6d stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml index 4c22356e61..1a4edb25d3 100644 --- a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml @@ -11,7 +11,7 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_wristframe_ypr + mode: keypoints_wristframe_6d stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_pilot_lance.yaml b/egomimic/hydra_configs/data/aria_pilot_lance.yaml new file mode 100644 index 0000000000..cd9a4cd062 --- /dev/null +++ b/egomimic/hydra_configs/data/aria_pilot_lance.yaml @@ -0,0 +1,41 @@ +# Lance pilot: the same bimanual Aria recordings as aria_pilot_zarr.yaml, read +# through the Lance catalog. Lance is not fork-safe, so workers use forkserver +# (https://lance.org/integrations/pytorch/). +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.lance.dataset.LanceEpisodeResolver + lance_root: /coc/flash7/scratch/lance_pilot/lance + cameras: + images.front_1: front_1 + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: keypoints + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: keypoints_headframe_6d + stride: 3 + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: row['embodiment'] == 'human_bimanual'" + mode: total + +valid_datasets: + human_bimanual: ${data.train_datasets.human_bimanual} + +train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 + multiprocessing_context: forkserver + persistent_workers: true +valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 4 + multiprocessing_context: forkserver + persistent_workers: true diff --git a/egomimic/hydra_configs/data/aria_pilot_lance_stream.yaml b/egomimic/hydra_configs/data/aria_pilot_lance_stream.yaml new file mode 100644 index 0000000000..0b8eced686 --- /dev/null +++ b/egomimic/hydra_configs/data/aria_pilot_lance_stream.yaml @@ -0,0 +1,41 @@ +# Lance pilot, streamed: the aria_pilot_lance.yaml recordings read straight from R2, +# frames from byte ranges of the faststart MP4s, one batched fetch per DataLoader +# batch (MultiDataset.__getitems__). Needs R2_* credentials in the environment. +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.lance.dataset.LanceStreamingResolver + lance_root: s3://rldb/scratch/rco3/lance_pilot/lance_faststart + cameras: + images.front_1: front_1 + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: keypoints + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: keypoints_headframe_6d + stride: 3 + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: row['embodiment'] == 'human_bimanual'" + mode: total + +valid_datasets: + human_bimanual: ${data.train_datasets.human_bimanual} + +train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 16 + multiprocessing_context: forkserver + persistent_workers: true +valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 + multiprocessing_context: forkserver + persistent_workers: true diff --git a/egomimic/hydra_configs/data/aria_pilot_zarr.yaml b/egomimic/hydra_configs/data/aria_pilot_zarr.yaml new file mode 100644 index 0000000000..3fb5431758 --- /dev/null +++ b/egomimic/hydra_configs/data/aria_pilot_zarr.yaml @@ -0,0 +1,35 @@ +# Lance pilot baseline: the pilot's bimanual Aria zarrs, read locally. Pair with +# aria_pilot_lance.yaml (same recordings, keymap and transforms); see +# docs/lance_aria_pilot_plan.md. +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver + folder_path: /coc/flash7/scratch/lance_pilot/zarr + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: keypoints + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: keypoints_headframe_6d + stride: 3 + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: row['embodiment'] == 'human_bimanual'" + mode: total + +valid_datasets: + human_bimanual: ${data.train_datasets.human_bimanual} + +train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 +valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 4 diff --git a/egomimic/hydra_configs/data/aria_qwen.yaml b/egomimic/hydra_configs/data/aria_qwen.yaml index 9612b56990..cd976a643d 100644 --- a/egomimic/hydra_configs/data/aria_qwen.yaml +++ b/egomimic/hydra_configs/data/aria_qwen.yaml @@ -1,7 +1,7 @@ # Aria data with raw text annotations in the batch, so the HPT Qwen stem can -# sample one annotation per item. Mirror of `aria.yaml` plus -# `annotation_key: annotations` in the keymap (which adds a per-frame raw-text -# annotation key — see Embodiment.get_keymap). +# sample one annotation per item. Mirror of `aria.yaml` (144-D wrist-frame +# hand-keypoint action) plus `annotation_key: annotations` in the keymap +# (which adds a per-frame raw-text annotation key; see Embodiment.get_keymap). _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -12,11 +12,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index bfa45b5bba..7fb6122886 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -1,3 +1,19 @@ +# Base for every pi0.5 data config. +# +# eva: wrist-relative cartesian actions (`cartesian_wristframe_6d`: each arm's +# chunk in its current EEF frame, xyz+6D+gripper per arm, 20D) with the +# cam-frame proprio ee_pose in the same 6D layout. Cam-frame runs override +# the mode to `cartesian_6d` on the command line. +# human: the 144-D wrist-first hand-keypoint action (`keypoints_wristframe_6d`: +# per hand [wrist xyz | wrist rot6d | 21 MANO keypoints] in that wrist's +# frame) with the head-frame `observations.state.keypoints` proprio, plus +# (`include_ee_pose`) the cartesian head-frame `observations.state.ee_pose` +# that the pi0.5 prompt discretizes into its `State:` block, so a keypoint +# run shares its prompt state with the cartesian runs. The cartesian human +# modes stay available by overriding +# `resolver.key_map.keymap_mode=cartesian` + `resolver.transform_list.mode=cartesian_6d` +# (and dropping include_ee_pose) together with a cartesian PI model config. +# Norm stats must be computed on these layouts. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -23,11 +39,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/eva_pilot_lance.yaml b/egomimic/hydra_configs/data/eva_pilot_lance.yaml new file mode 100644 index 0000000000..2298ca4573 --- /dev/null +++ b/egomimic/hydra_configs/data/eva_pilot_lance.yaml @@ -0,0 +1,44 @@ +# Lance pilot: bimanual EVA cup_on_saucer demos converted from raw HDF5 by +# eva_to_lance, read through the Lance catalog. Same keymap and transforms as +# eva.yaml. Lance is not fork-safe, so workers use forkserver +# (https://lance.org/integrations/pytorch/). +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.lance.dataset.LanceEpisodeResolver + lance_root: /coc/flash7/scratch/lance_pilot/lance_eva + lab: eva + cameras: + images.front_1: front_1 + images.left_wrist: left_wrist + images.right_wrist: right_wrist + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + keymap_mode: cartesian + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + mode: cartesian_wristframe_6d + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: row['embodiment'] == 'eva_bimanual'" + mode: total + +valid_datasets: + eva_bimanual: ${data.train_datasets.eva_bimanual} + +train_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 8 + multiprocessing_context: forkserver + persistent_workers: true +valid_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 4 + multiprocessing_context: forkserver + persistent_workers: true diff --git a/egomimic/hydra_configs/data/human.yaml b/egomimic/hydra_configs/data/human.yaml index 59a472a346..9c4a0b27eb 100644 --- a/egomimic/hydra_configs/data/human.yaml +++ b/egomimic/hydra_configs/data/human.yaml @@ -1,8 +1,14 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper -# Canonical human (ex-aria) cartesian data config, keyed by the collapsed -# embodiment human_bimanual (episodes reprocessed to embodiment="human_bimanual" -# + intrinsics in zarr.json). +# Canonical human cartesian data config, keyed by the collapsed embodiment +# human_bimanual (episodes reprocessed to embodiment="human_bimanual" + +# intrinsics in zarr.json). +# Action: the 144-D wrist-frame hand-keypoint action (`keypoints_wristframe_6d`) +# with the head-frame `observations.state.keypoints` proprio and, via +# include_ee_pose, the cartesian head-frame `observations.state.ee_pose` the +# pi0.5 prompt discretizes (HPT ignores it: no stem consumes it). Override +# `resolver.key_map.keymap_mode=cartesian resolver.transform_list.mode=cartesian_6d` +# (and include_ee_pose=false on both) for the cartesian action. # NOTE: the filter selects on the SQL value 'human_bimanual'. This ASSUMES the SQL # embodiment/robot_name column has been migrated from the old vendor values # ('aria_bimanual'). Until that migration lands, this dataset matches zero rows. @@ -14,10 +20,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 3fe7f6c6ca..b904d4dde0 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -1,3 +1,12 @@ +# One pinned mecka episode (stored as human_bimanual): the smoke / overfit +# config for HPT and Pi on the 144-D wrist-frame hand-keypoint action. +# fix_left_wrist_convention: this episode predates the mecka_to_zarr rot_left +# fix (LEFT wrist frame double-mirrored); Rz(180) corrects the raw left +# pose before any frame math. With include_ee_pose it corrects both the +# wrist_pose that the 144-D keypoint action is built from and the cartesian +# ee_pose proprio. +# pad_proprio_gripper: 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) +# so the pi0.5 State: prompt bins align with the robot 20-D layout. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,13 +17,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true fix_left_wrist_convention: true - # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 - # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true stride: 1 filters: diff --git a/egomimic/hydra_configs/data/mecka_20pct_6d.yaml b/egomimic/hydra_configs/data/mecka_20pct_6d.yaml new file mode 100644 index 0000000000..de5d99447b --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_20pct_6d.yaml @@ -0,0 +1,4532 @@ +# data=mecka_all_6d pinned to fixed episode lists: train = a 20% subset of the +# lab=mecka episodes (4471 hashes, mode total), valid = 3 seen-operator +# episodes, unseen_op_valid = 3 held-out-operator episodes. Every hash is a +# pin: one that is missing, deleted or filtered out fails the resolve -- except +# pins in exclude_hashes (mecka_bad_egomotion, inherited), which are dropped +# with a warning: 393 of the 4471 train pins, none of the val / video pins. +defaults: + - mecka_all_6d + - _self_ + +valid_datasets: + human_bimanual: + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: ${data.train_datasets.human_bimanual.filters.filter_lambdas} + exclude_hashes: ${data.train_datasets.human_bimanual.filters.exclude_hashes} + episode_hashes: + - "69b5691e4e73acdd60a51040" + - "69b0c52becd6553472c0cb2d" + - "69b9abc3aa6e21f31eb6ba25" + mode: total + +unseen_op_valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: ${data.train_datasets.human_bimanual.filters.filter_lambdas} + exclude_hashes: ${data.train_datasets.human_bimanual.filters.exclude_hashes} + episode_hashes: + - "69b525ff91d6491aa901a64b" + - "69b9b260c004b8e369a3d517" + - "69b9e5e27301ad7dc90077dc" + mode: total + +unseen_op_valid_dataloader_params: + human_bimanual: ${data.valid_dataloader_params.human_bimanual} + +train_viz_datasets: + human_bimanual: + mode: total + +# 80 val batches x 64 = 5120 frames per head; train_viz: floor(5120/4471) = 1. +# valid / unseen_op_valid have 3 episodes each and run in full. +metric_frames_per_episode: + train_viz: 1 + +video_episodes: + valid: + - "69b5691e4e73acdd60a51040" + train_viz: + - "696a1464715ed751e5e0f674" + unseen_op_valid: + - "69b525ff91d6491aa901a64b" + +train_datasets: + human_bimanual: + mode: total + filters: + episode_hashes: + - "696a1464715ed751e5e0f674" + - "696a17de03b3bab4cd3d4f45" + - "696a17f656b90f0047212055" + - "696a1a388e5cc0de46136daa" + - "696a1a5456b90f0047212090" + - "696a20f84e85aa4ea7da5e5f" + - "696a250203b3bab4cd3d509b" + - "696a25297a2fd07226c097c1" + - "696a2534715ed751e5e0f7c6" + - "696a2835ac451788451c3287" + - "696a28494419259d05e51044" + - "696a292c7a2fd07226c09820" + - "696a2bc2ac451788451c32d3" + - "696a2f6bac451788451c331e" + - "696a304f76b41c2326097948" + - "696a3066e54fd810ece69cd1" + - "696a30707a2fd07226c098c1" + - "696a31c18e5cc0de46136f9a" + - "696a328fac451788451c3360" + - "696a32f7e54fd810ece69d0d" + - "696a33644419259d05e5113e" + - "696a34314e85aa4ea7da601a" + - "696a37f04419259d05e51168" + - "696a64714e85aa4ea7da6046" + - "696ac4bdaafcb3130f5ef6c2" + - "696ac57b4e85aa4ea7da6089" + - "696ac6758e5cc0de46137019" + - "696ac67803b3bab4cd3d5269" + - "696ac6e57a2fd07226c099a5" + - "696ac7667a2fd07226c099b1" + - "696ac7bd4e85aa4ea7da60cd" + - "696ac81703b3bab4cd3d529b" + - "696ac81c7d96ed6fad190af5" + - "696ac83784947585c0cb6b2a" + - "696ac83b7a2fd07226c099bc" + - "696ac84e240b1e4a2d82c161" + - "696ac857e58b362e378a581f" + - "696ac8920c1d9ccd5892fc6f" + - "696ac897d87f83facc75f326" + - "696ac8a751175c2cf41f408a" + - "696ac94e7a2fd07226c099c6" + - "696ac9830c1d9ccd5892fc7f" + - "696ac9a5aafcb3130f5ef71f" + - "696aca377ae3f760695d88c1" + - "696aca378e5cc0de46137054" + - "696aca39b476146af0d5ac60" + - "696acb4456b90f0047212366" + - "696acc8fdaa011f7172d1b31" + - "696ace5951175c2cf41f40b8" + - "696ad1bea04610abd4bb2c76" + - "696ad1c6263039fb3592d973" + - "696afd21b17075c2c32fa2fc" + - "696afd784349fe21180212d8" + - "696afd79c19cc80a4d8afb48" + - "696b0435429b192f08775449" + - "696b046d1da78387ae7e54bf" + - "696b08a7900932d25e86937d" + - "696b0a8a3c1c243026d5bf5f" + - "696b0a8f55db5be0d7706cd5" + - "696b0d6e2a9a68766920d94a" + - "696b0da32d6f5bba1f851353" + - "696b0db634614160c6097825" + - "696b0dc0fda9075ccb0c15b3" + - "696b0de037193673e4f8959b" + - "696b0e26ba086f973109ca52" + - "696b0ff2b8537b10f406afe7" + - "696b119c9a6d79db2cf88293" + - "696b14de0c73035b14a75cd8" + - "696b16a93ab4e38fa0319539" + - "696b16bacca65ef14371fae6" + - "696b1703039d54a181f8e6bd" + - "696b170c39dc2438efc84a7e" + - "696b171d0c73035b14a75d20" + - "696b173bcca65ef14371fafc" + - "696b17574292c4c779911fc9" + - "696b176114764da1c7e17e88" + - "696b19b279e3576eb2d4ecf3" + - "696b1f24daa011f7172d1bf5" + - "696b20257aa8f13c7e4a70f8" + - "696b20bcb476146af0d5ada4" + - "696b21209a084c90623c46e2" + - "696b21c2800cd1498944fb06" + - "696b21ed1e187f32c5eb54da" + - "696b220e0d98b1b1a38b1026" + - "696b2634c7e4e8ee05bbdc6a" + - "696b27778d4f708447df8435" + - "696b2c0e90e2545845783bcc" + - "696b2c18abc451024b8269a6" + - "696b2c2df20f919514c6f02d" + - "696b2c37c426c607c88b2793" + - "696b2c531e187f32c5eb5539" + - "696b2c58394c1e86a2c42802" + - "696b2d9d8e6d0bc988967276" + - "696b2de6d3cc9dc76c561164" + - "696b31418149bf4393894bad" + - "696b3875e1df4679fecfd2c7" + - "696b3899d2d7cb0bd150e85e" + - "696b38a2a12eec59c7475065" + - "696b38b10ac58bb141246a1f" + - "696b38cdd87f83facc75f4b9" + - "696b39af43d5aa74f8b8d08f" + - "696b39af80d64f1587ce363a" + - "696b3a4184947585c0cb6d72" + - "696b3a78d2d7cb0bd150e88b" + - "696b3a7d02ce4c977624f538" + - "696b45d08d4f708447df84b5" + - "696b47170c52a10600ae7ff0" + - "696b496421d8e8ab8ddc7d0a" + - "696b4966d9a2a4e1bf83686a" + - "696b4a33d3cc9dc76c561237" + - "696b4a3470e0519696315a2f" + - "696b4a78f7b3293476cab3f5" + - "696b4aa679bc0402722c1fa0" + - "696b4ab9c2352699ffedeb7a" + - "696b4acb8149bf4393894c5b" + - "696b4b023e2d6a68d5268264" + - "696b4ba77aa8f13c7e4a7207" + - "696b4bceb717e87cf71f73c9" + - "696b4bfc9b8e7b7d20758e0b" + - "696b4c3cd2d7cb0bd150e8f6" + - "696b4c9221d8e8ab8ddc7d58" + - "696b4d0676b41c2326097c93" + - "696b4e9a8d4f708447df84de" + - "696b56a2c0aa1c772e76ef98" + - "696b56e402ce4c977624f5a7" + - "696b56f470d628f8b0741d92" + - "696b58a2017ec1a72078eed2" + - "696b5bed894a8c8702a75895" + - "696b6464c426c607c88b2993" + - "696b65bab7080af9dcc93f36" + - "696b65bcdaa011f7172d1e0f" + - "696b67d1101054681ed9734c" + - "696b67d4e13d37dd4293f1c7" + - "696b6993790efbe0d372bf42" + - "696b6a471e187f32c5eb57b5" + - "696b6a4e07b6c7a3d7946347" + - "696b6aae14764da1c7e18164" + - "696b6bdb0ac3ef927c48e145" + - "696b6bea9c8ea4b6e4a1ee90" + - "696b6ca70ac3ef927c48e14d" + - "696b70dd28ab12fee3734fb5" + - "696b749870e0519696315b7c" + - "696b785a2efe1549985cf711" + - "696b78907a53efd0cd3a289e" + - "696b78cb14764da1c7e181c7" + - "696b793584947585c0cb6fce" + - "696b7950ef23530d3c5a32d0" + - "696b79589a6d79db2cf884e5" + - "696b79708a0fa4aac9be4525" + - "696b799f84947585c0cb6fd9" + - "696b79a6778afed01d13c698" + - "696b79b84292c4c7799121e7" + - "696b79bfcac8603dbfca7bf6" + - "696b79c22a9a68766920dbab" + - "696b7a162a9a68766920dbb3" + - "696b7b5f09bf9f435619ea72" + - "696b7bf80ac58bb141246c48" + - "696b7c1e9b8e7b7d20758e64" + - "696b7eeaca59649f142e1758" + - "696b80331e187f32c5eb5895" + - "696b806d4e85aa4ea7da6532" + - "696b80757e6f94bf34143618" + - "696b80cf556c93e45170dbef" + - "696b8159e300630c41633cf6" + - "696b81824efa2e122be5ffc3" + - "696b81c480d0c67f978be41f" + - "696b8206a5b1694665c194d1" + - "696b820d883568f6441c7f88" + - "696b8211101054681ed97467" + - "696b821b632cfa55ccdfa68e" + - "696b821c736a3883c3d67a69" + - "696b825aa8cba78d3b5a4e4d" + - "696b827ccfb39c7ef2d488b6" + - "696b831a93818b7b258ebed1" + - "696b831aca28cac2415e5ba9" + - "696b831b3724208fc971e045" + - "696b831b86833e290580c114" + - "696b84784efa2e122be5fff3" + - "696b895737193673e4f897f3" + - "696b8a7f017ec1a72078f082" + - "696b8b04f2042e02291298cb" + - "696b8dd8c6cb022a9da400ca" + - "696b8de8e4d0988228a6a549" + - "696b8fc172239e0ab4634933" + - "696b900836602141d44d6ac5" + - "696b93c0c3881a5a8cbaa0cd" + - "696b93c2101054681ed974f6" + - "696b93c35bfe694453c6138b" + - "696b9420e4d0988228a6a58d" + - "696b947aa491d5abbd7951f7" + - "696b94d067d3fa0fb372397a" + - "696b94eb894c2e015aaa0cb9" + - "696b9572620da154d5b7a7d2" + - "696b97b0bd0145c530d36d6c" + - "696b993a6fc0c09c00a97b64" + - "696b9c0f4d5f3676bb0aae72" + - "696b9c7cc74b24422e33d902" + - "696b9ce243edc8aed5480230" + - "696b9d2b0ac3ef927c48e22a" + - "696b9d85894c2e015aaa0d27" + - "696b9d873724208fc971e0cc" + - "696b9d8cd2d7cb0bd150eaff" + - "696b9da9675f6396322f9111" + - "696b9dcc9a643f01c844d254" + - "696b9dd6101054681ed97553" + - "696b9e1587b9ac9746b298a4" + - "696b9e18b717e87cf71f7695" + - "696b9e24c25b828c83107b91" + - "696b9e484177e6c8a83880df" + - "696b9f0267f2110cf0fbdf24" + - "696b9f09abc451024b826ddf" + - "696b9f1407b6c7a3d7946469" + - "696b9f5af38155a6224418b9" + - "696b9f75915ce3edd0967850" + - "696b9f96894c2e015aaa0d3b" + - "696b9fa04a936a4d1328a717" + - "696b9fb2ae7ec71f8657150f" + - "696b9fe9c3881a5a8cbaa145" + - "696ba00986f71ea8bc69d099" + - "696ba00fb261b0a7823a1b3a" + - "696ba023b17075c2c32fa674" + - "696ba0496377d0871b6a8ec5" + - "696ba0e379188f9419302880" + - "696ba111afe603760fe1ac04" + - "696ba1138d4f708447df86b8" + - "696ba1391c7532ffed13fe1d" + - "696ba147f249daed52be58fa" + - "696ba14a715ed751e5e0fe57" + - "696ba16e60ebbd99811cbb75" + - "696ba18339424cbffd650c86" + - "696ba19976b41c2326097eeb" + - "696ba19f76b41c2326097eec" + - "696ba1a08e6d0bc9889676bc" + - "696ba1a51e187f32c5eb59f7" + - "696ba1a8c676295b4ae749f6" + - "696ba1c5632cfa55ccdfa6fd" + - "696ba1dd883568f6441c80a4" + - "696ba1e14d5f3676bb0aaef5" + - "696ba1e909bf9f435619eb81" + - "696ba202cfb39c7ef2d48995" + - "696ba219634a1e1f46927714" + - "696ba22d7ed68fc306df6fb9" + - "696ba24379188f941930289f" + - "696ba25786f71ea8bc69d0cc" + - "696ba2b7f53e1ae2cbd86d51" + - "696ba2c970e0519696315d1f" + - "696ba2d00c52a10600ae837e" + - "696ba2e38a0fa4aac9be4655" + - "696ba30793818b7b258ec023" + - "696ba31ce34991e438ab3692" + - "696ba325ac515a4d340b0d6c" + - "696ba3371ee2a915e1934946" + - "696ba356894a8c8702a75b9a" + - "696ba356bd0145c530d36df4" + - "696ba38eb0900198ccf5bffa" + - "696ba3922cdacabcc9e6f224" + - "696ba39660ebbd99811cbb98" + - "696ba3b87a2fd07226c09e8c" + - "696ba3dbed8ed158122ec488" + - "696ba3f4e34991e438ab36a2" + - "696ba3ff39dc2438efc84f3c" + - "696ba402d99e01dc631f3ce4" + - "696ba40493818b7b258ec03b" + - "696ba4091376c8db5e315240" + - "696ba4ed8d4f708447df8706" + - "696ba51fc35216997ae3eeb5" + - "696ba547c74b24422e33d98b" + - "696ba59e8910191c50ed9d6c" + - "696ba5ab7e6f94bf3414382a" + - "696ba5c409b82287df69f517" + - "696ba5c8cca65ef143720007" + - "696ba5c8d9a2a4e1bf836bba" + - "696ba5cee34991e438ab36ca" + - "696ba5f66377d0871b6a8f2e" + - "696ba5f66fc0c09c00a97c44" + - "696ba5f75bfe694453c614bc" + - "696ba607fadd8a32efe99df8" + - "696ba611ed8ed158122ec4b5" + - "696ba6140ac58bb141246de2" + - "696ba622fda9075ccb0c1b33" + - "696ba627ef23530d3c5a345d" + - "696ba62a36602141d44d6bc4" + - "696ba6498d4f708447df8725" + - "696ba6632a9a68766920dd10" + - "696ba691e1b545a3ce074fbf" + - "696ba6e976b41c2326097f7d" + - "696ba6f81894fc36a07d8d47" + - "696ba70295be6b2159258362" + - "696ba70387b9ac9746b2995e" + - "696ba73160e10e8f87b583d0" + - "696ba73dac515a4d340b0dcd" + - "696ba75dc74b24422e33d9b6" + - "696ba762620da154d5b7a8a0" + - "696ba762e4d0988228a6a696" + - "696ba76709b82287df69f52c" + - "696ba767bd0145c530d36e51" + - "696ba76c1894fc36a07d8d50" + - "696ba7754efa2e122be60139" + - "696ba783e43294cde98aa33a" + - "696ba78b8e5cc0de4613741d" + - "696ba7b6ac515a4d340b0dd6" + - "696ba7c1b0c23fa5919d0f7f" + - "696ba7c21d3959e2e052af47" + - "696ba7c2ed8ed158122ec4dc" + - "696ba7c8c676295b4ae74a2c" + - "696ba7d0800cd1498944ffb5" + - "696ba7e770d628f8b074217d" + - "696ba80ccd29fb97fb107ca7" + - "696ba80f5bfe694453c614ee" + - "696ba84639424cbffd650d22" + - "696ba86564c87f3a4d8c5b27" + - "696ba87de13d37dd4293f33c" + - "696ba87ee34991e438ab3706" + - "696ba8908d4f708447df875a" + - "696ba891f98fc50273d5d2ee" + - "696ba8928e6d0bc98896775e" + - "696ba8c5eea8e4fd8069634d" + - "696ba8d3ed8ed158122ec4f6" + - "696ba8e702ce4c977624f78c" + - "696ba9272cdacabcc9e6f294" + - "696ba92f96b4b190357b0bcd" + - "696ba942778afed01d13c87c" + - "696ba94f915ce3edd0967917" + - "696ba9511225df5278db7db2" + - "696ba96db0900198ccf5c07a" + - "696ba9716fc0c09c00a97c89" + - "696ba980e34991e438ab371e" + - "696ba9834d5f3676bb0aaf9b" + - "696ba99537193673e4f898e0" + - "696ba9e879bc0402722c22d1" + - "696baa5d01fc6bd69561e0d0" + - "696baa6970e0519696315d94" + - "696baa6a316c07a505ec65e2" + - "696baa770ac58bb141246e41" + - "696baa901894fc36a07d8d8f" + - "696baadae4d0988228a6a6df" + - "696baae267d3fa0fb3723a31" + - "696baaef14764da1c7e183e0" + - "696bab01d3cc9dc76c56164c" + - "696bab3072ef8d13aa004a5c" + - "696bab3a67f2110cf0fbdff7" + - "696bab4f5e774190a1c976c3" + - "696bab6d07485cd91adbcffc" + - "696bab6e8a74d25e9af38cc5" + - "696bab789ac5bea01f980121" + - "696babb75fedf8533255dda9" + - "696babb802ce4c977624f7c1" + - "696babba6fc0c09c00a97cb8" + - "696babc6cd29fb97fb107cfe" + - "696babd8590d785c13a75e50" + - "696babfb263039fb3592dc7a" + - "696bac126fc0c09c00a97cbf" + - "696bac2db717e87cf71f77c3" + - "696bac4aabc451024b826ee4" + - "696bac4c42eca86da25f1de7" + - "696bac552a9a68766920dd91" + - "696bac60d2d7cb0bd150ec58" + - "696bac6395be6b21592583e7" + - "696bac658149bf439389506f" + - "696bac706fc0c09c00a97cc8" + - "696bac86da6069828e20cec6" + - "696bac8bc35216997ae3ef21" + - "696bac91316c07a505ec6610" + - "696bac921894fc36a07d8db9" + - "696bac9a5309532923d3f214" + - "696bacb77ed68fc306df7066" + - "696bace4017ec1a72078f293" + - "696bacf67bb51948f9de6129" + - "696bacfa101054681ed97677" + - "696bad144c43650cbf6f5bef" + - "696bad15b7080af9dcc941ff" + - "696bad4d1cafc9f413b5b660" + - "696bad8ab0900198ccf5c0d7" + - "696bad8e017ec1a72078f2a2" + - "696bad91476e9a6e777ffce5" + - "696bad96f05ef518f87fff85" + - "696badc04292c4c779912361" + - "696badc5778afed01d13c8df" + - "696badc8715ed751e5e0ff46" + - "696badc99839477cebe21bea" + - "696badd272239e0ab4634ab3" + - "696badebeea8e4fd8069638d" + - "696badf71e187f32c5eb5aff" + - "696bae0e476e9a6e777ffcf1" + - "696bae5dc7e4e8ee05bbdff8" + - "696bae70c75cf6e4320e2846" + - "696bae715fedf8533255dde5" + - "696baefad9a2a4e1bf836c88" + - "696baf041e187f32c5eb5b15" + - "696bafe5b261b0a7823a1c61" + - "696bafe98e6d0bc988967809" + - "696bb00ef05ef518f87fffb9" + - "696bb010778afed01d13c90d" + - "696bb04e4c43650cbf6f5c39" + - "696bb05272239e0ab4634aea" + - "696bb0dd476e9a6e777ffd33" + - "696bb0f2736a3883c3d67c97" + - "696bb10bc426c607c88b2d1f" + - "696bb15a1c34188da33fea18" + - "696bb1a670d628f8b0742238" + - "696bb1b7c35216997ae3ef9b" + - "696bb1d6316c07a505ec667e" + - "696bb2178e5cc0de461374f3" + - "696bb2a9a491d5abbd7953a9" + - "696bb2aa0ac58bb141246ece" + - "696bb2b51cafc9f413b5b6bc" + - "696bb2d007b6c7a3d7946570" + - "696bb2d9a907c8962e22a82c" + - "696bb33f25368af005ca7440" + - "696bb35e67d3fa0fb3723ae9" + - "696bb366c3881a5a8cbaa27e" + - "696bb36bc7e4e8ee05bbe050" + - "696bb405eea8e4fd806963d2" + - "696bb45767d3fa0fb3723af9" + - "696bb459cfb39c7ef2d48b3d" + - "696bb498a5b1694665c1977b" + - "696bb4a01e187f32c5eb5b90" + - "696bb4c7eab4449d0aabe7be" + - "696bb4e7c25b828c83107d4a" + - "696bb4ec800cd149894500e1" + - "696bb4ef67d3fa0fb3723b06" + - "696bb4f31c34188da33fea6a" + - "696bb52c141d888fd2ec5009" + - "696bb5499a140887948c5ddf" + - "696bb598a491d5abbd7953e3" + - "696bb59ca8cba78d3b5a50d3" + - "696bb5bc800cd149894500f1" + - "696bb5beeab4449d0aabe7d2" + - "696bb5c3c75cf6e4320e28e0" + - "696bb62707b6c7a3d794659f" + - "696bb68aae7ec71f865716c8" + - "696bb6eb9a140887948c5e18" + - "696bb75f800cd14989450111" + - "696bb7af07b6c7a3d79465a7" + - "696bb7efc6cb022a9da402ec" + - "696bb7fff38155a622441a3e" + - "696bb8390c52a10600ae853e" + - "696bb843bd0145c530d36fb1" + - "696bb85570e0519696315e4c" + - "696bb8782c7243cd72c85036" + - "696bb91dcca65ef14372012d" + - "696bb9367a2fd07226c0a050" + - "696bb9c98318b8ceb8224842" + - "696bb9cb8d4f708447df8842" + - "696bb9dc1c34188da33feacc" + - "696bba0a0c52a10600ae8550" + - "696bba30c3881a5a8cbaa2f5" + - "696bba3aca28cac2415e5e89" + - "696bba449c8ea4b6e4a1f203" + - "696bba57f64364413f890e3f" + - "696bba6714764da1c7e184bf" + - "696bba78736a3883c3d67d0e" + - "696bba7e4c43650cbf6f5cea" + - "696bba9f9c8ea4b6e4a1f20b" + - "696bbaa5894c2e015aaa0f4a" + - "696bbaaa42eca86da25f1e3e" + - "696bbabb1c34188da33feae1" + - "696bbadd95be6b21592584de" + - "696bbae17e6f94bf3414399e" + - "696bbae9b17075c2c32fa81a" + - "696bbb0a67f2110cf0fbe111" + - "696bbb0b715ed751e5e10073" + - "696bbb15ae7ec71f8657172d" + - "696bbb18f53e1ae2cbd86f00" + - "696bbb47b0c23fa5919d1073" + - "696bbb5f0ac3ef927c48e37a" + - "696bbb7e3c35512701de64e4" + - "696bbc03cef1149d4f7a1b5e" + - "696bbc841376c8db5e3153c7" + - "696bbca30ac58bb141246f4d" + - "696bbca493818b7b258ec21f" + - "696bbcf9620da154d5b7aa0d" + - "696bbd069a140887948c5e76" + - "696bbd5575020a567fd376e9" + - "696bbd74f53e1ae2cbd86f33" + - "696bbdffc676295b4ae74b14" + - "696bbe228688d0ceb96d3ca3" + - "696bbe617ed68fc306df7143" + - "696bbe828149bf43938951e1" + - "696bbe85f53e1ae2cbd86f49" + - "696bbe87ef23530d3c5a361a" + - "696bbe92c676295b4ae74b1a" + - "696bbea3f38155a622441abe" + - "696bbeb595be6b215925852c" + - "696bbf6f21d8e8ab8ddc82e1" + - "696bbf7d95be6b2159258540" + - "696bbfaf5e774190a1c97812" + - "696bc0311ee2a915e1934af4" + - "696bc041e4eb4a667bf0de72" + - "696bc04434614160c6097cab" + - "696bc04d883568f6441c82a5" + - "696bc0761894fc36a07d8ef7" + - "696bc0de07485cd91adbd0e1" + - "696bc0f270d628f8b0742333" + - "696bc13b39424cbffd650e58" + - "696bc14dc75cf6e4320e295b" + - "696bc14e37193673e4f899f0" + - "696bc1a8d3cc9dc76c5617d6" + - "696bc1b0b717e87cf71f7975" + - "696bc1ba02ce4c977624f8cb" + - "696bc1d00c791affadd83768" + - "696bc1df1c34188da33feb46" + - "696bc1eff53e1ae2cbd86f96" + - "696bc222eea8e4fd80696465" + - "696bc233fda9075ccb0c1d68" + - "696bc258ef11876d0045f765" + - "696bc26ac8af0268a6bace1c" + - "696bc26bc2544a8f5862527e" + - "696bc2db37193673e4f89a06" + - "696bc308d87f83facc75f987" + - "696bc30de4eb4a667bf0deab" + - "696bc34aa5b1694665c198aa" + - "696bc3588a74d25e9af38e6e" + - "696bc35e5fedf8533255def0" + - "696bc375c75cf6e4320e2988" + - "696bc3a0e43294cde98aa4f3" + - "696bc3a44c43650cbf6f5d90" + - "696bc422eea8e4fd80696479" + - "696bc4510ac3ef927c48e439" + - "696bc45fc676295b4ae74b97" + - "696bc4641ae707d8c0922b6f" + - "696bc46e9c8ea4b6e4a1f2be" + - "696bc46ec25b828c83107e81" + - "696bc47064c87f3a4d8c5ccf" + - "696bc4b08910191c50ed9e39" + - "696bc4b34d5f3676bb0ab0d3" + - "696bc4fb75020a567fd37772" + - "696bc51a620da154d5b7aab2" + - "696bc54ec6cb022a9da403bc" + - "696bc55028ab12fee3735429" + - "696bc57d4177e6c8a83883ed" + - "696bc588f79f22003e620407" + - "696bc58a70e0519696315f2e" + - "696bc595cef1149d4f7a1bf5" + - "696bc5a52c7243cd72c85121" + - "696bc5c2263039fb3592dd34" + - "696bc62acca65ef143720226" + - "696bc65bfda9075ccb0c1dcd" + - "696bc65fb0c23fa5919d111f" + - "696bc6679c8ea4b6e4a1f2d5" + - "696bc69942eca86da25f1ec1" + - "696bc6d270d628f8b07423b6" + - "696bc73493818b7b258ec2d4" + - "696bc73d43edc8aed5480520" + - "696bc7657bb51948f9de62e9" + - "696bc78d93818b7b258ec2db" + - "696bc7de4c43650cbf6f5deb" + - "696bc831cd29fb97fb107ecf" + - "696bc868a907c8962e22a956" + - "696bc86f1225df5278db7fc7" + - "696bc8974d5f3676bb0ab11e" + - "696bc8a90d98b1b1a38b1472" + - "696bc8b2a8793724a1ef215e" + - "696bc8b7246cb36ca3758045" + - "696bc8e9883568f6441c8341" + - "696bc8fa60e10e8f87b5851c" + - "696bc95be2eca952034ed8f7" + - "696bc95c8910191c50ed9e6e" + - "696bc9b0800cd1498945025e" + - "696bc9cfc74b24422e33db6f" + - "696bc9df522a3f9be47906f9" + - "696bca369c8ea4b6e4a1f301" + - "696bca64e4eb4a667bf0df47" + - "696bca6572239e0ab4634c38" + - "696bca67ca28cac2415e5f3d" + - "696bca9042eca86da25f1eed" + - "696bcab8da6069828e20d03b" + - "696bcaccd2d7cb0bd150eee9" + - "696bcaead87f83facc75f9f8" + - "696bcaf2c8af0268a6bacec7" + - "696bcb2df64364413f890f49" + - "696bcb465309532923d3f377" + - "696bcb76316c07a505ec6829" + - "696bcb7a60e10e8f87b58532" + - "696bcb99b0b2b22cc0d1f7b4" + - "696bcb9efadd8a32efe9a09b" + - "696bcba9c3881a5a8cbaa3cb" + - "696bcbccc75cf6e4320e2a3a" + - "696bcbce93818b7b258ec339" + - "696bcbcf86833e290580c366" + - "696bcbf301fc6bd69561e2b1" + - "696bcc17c19cc80a4d8b0140" + - "696bcc1b778afed01d13c9c0" + - "696bcc1c1ee2a915e1934bdb" + - "696bcc1f4292c4c7799124f8" + - "696bcc24c25b828c83107f18" + - "696bcc265fedf8533255df96" + - "696bcc26c3881a5a8cbaa3d4" + - "696bcc341894fc36a07d8fe0" + - "696bcc3428ab12fee37354bd" + - "696bcc35c75cf6e4320e2a43" + - "696bcc3ce43294cde98aa58c" + - "696bcc44f249daed52be5b4b" + - "696bcc456fc0c09c00a97e57" + - "696bcc52ae7ec71f86571838" + - "696bcc62dae7f926421a6229" + - "696bcc697a53efd0cd3a2ce6" + - "696bcc6a3c35512701de65f7" + - "696bcc6bd2d7cb0bd150ef0d" + - "696bcc711d3959e2e052b1ea" + - "696bcc82afe603760fe1af38" + - "696bcc8eb7080af9dcc943a2" + - "696bcc92f38155a622441b8d" + - "696bcc951ae707d8c0922bfc" + - "696bccb77bb51948f9de635a" + - "696bccd4fe5c03f9da3888db" + - "696bccd521d8e8ab8ddc83d3" + - "696bcd104c43650cbf6f5e60" + - "696bcd12a6f1d19c0183453b" + - "696bcd1f4292c4c77991250d" + - "696bcd2dd3cc9dc76c56187e" + - "696bcd3c55db5be0d770718c" + - "696bcd679a6d79db2cf88810" + - "696bcd69da6069828e20d075" + - "696bcda1c74b24422e33dbc1" + - "696bcdcc316c07a505ec685b" + - "696bcdf14177e6c8a83884aa" + - "696bcdfecfb39c7ef2d48d3c" + - "696bce8ee34991e438ab396a" + - "696bcea59839477cebe21d74" + - "696bcea879bc0402722c253e" + - "696bceac5fedf8533255dfb2" + - "696bcebfd3cc9dc76c561899" + - "696bcee3ce21067f584460e5" + - "696bceecd3cc9dc76c56189d" + - "696bcef60ac58bb14124709d" + - "696bcf114292c4c779912538" + - "696bcf155309532923d3f3b8" + - "696bcf15c2544a8f58625370" + - "696bcf25d2d7cb0bd150ef49" + - "696bcf2da2c887c4e90e36e3" + - "696bcf36c75cf6e4320e2a80" + - "696bcf503897ce0bf0acac4b" + - "696bcf897ed68fc306df7214" + - "696bcf8a883568f6441c83ca" + - "696bcf8b1d3959e2e052b22a" + - "696bcf9ad0bfcc86fe4a6a20" + - "696bcf9fa12eec59c74754ee" + - "696bcfaa70e0519696315fbf" + - "696bcfb41cafc9f413b5b8b5" + - "696bcfd936602141d44d6e3c" + - "696bd00a9a6d79db2cf88848" + - "696bd0198149bf4393895327" + - "696bd01a5309532923d3f3c4" + - "696bd01b8149bf4393895328" + - "696bd036cca8f940bc789dc6" + - "696bd0421894fc36a07d9035" + - "696bd066f53e1ae2cbd8707f" + - "696bd07507485cd91adbd1d7" + - "696bd0b9d99e01dc631f3f07" + - "696bd0d7a12eec59c747550c" + - "696bd0ddb261b0a7823a1dec" + - "696bd0f6d9a2a4e1bf836ea7" + - "696bd108c25b828c83107f7d" + - "696bd14928ab12fee373550c" + - "696bd169c25b828c83107f85" + - "696bd16acd29fb97fb107f7b" + - "696bd16cef23530d3c5a378a" + - "696bd16eda6069828e20d0cb" + - "696bd1763897ce0bf0acac6e" + - "696bd1a260ebbd99811cbe53" + - "696bd1b507b6c7a3d7946696" + - "696bd1bdfe5c03f9da388948" + - "696bd1cd7a2fd07226c0a1c9" + - "696bd1ed43edc8aed54805f0" + - "696bd1efd99e01dc631f3f13" + - "696bd1fca2c887c4e90e3703" + - "696bd22921d8e8ab8ddc8438" + - "696bd22cc2544a8f586253b5" + - "696bd244675f6396322f94bd" + - "696bd2531ae707d8c0922c70" + - "696bd25986833e290580c3a1" + - "696bd261522a3f9be479078c" + - "696bd278b17075c2c32fa942" + - "696bd291cd29fb97fb107f8d" + - "696bd29693818b7b258ec3d0" + - "696bd2a3d1a6773989721fc6" + - "696bd2a8a12eec59c7475536" + - "696bd2bc9c8ea4b6e4a1f376" + - "696bd2db900932d25e86995b" + - "696bd2fd8910191c50ed9eb1" + - "696bd30a4efa2e122be603d5" + - "696bd346620da154d5b7ab9a" + - "696bd34fc25b828c83107fab" + - "696bd34fcca8f940bc789e0c" + - "696bd35cf79f22003e6204b3" + - "696bd37121d8e8ab8ddc8455" + - "696bd379c7e4e8ee05bbe285" + - "696bd39ceea8e4fd80696529" + - "696bd3a41e187f32c5eb5d7f" + - "696bd3ac07485cd91adbd218" + - "696bd3ac715ed751e5e10250" + - "696bd3baba086f973109d117" + - "696bd3beb717e87cf71f7aa3" + - "696bd3bf4d5f3676bb0ab1d4" + - "696bd3cf6fc0c09c00a97ec4" + - "696bd3d18688d0ceb96d3e4d" + - "696bd3d30ac58bb141247106" + - "696bd3d4f38155a622441bd1" + - "696bd3d560e10e8f87b585b4" + - "696bd3d8f7b3293476cab9a1" + - "696bd3f10d98b1b1a38b14eb" + - "696bd3fac426c607c88b2fa0" + - "696bd3fda12eec59c7475556" + - "696bd412f79f22003e6204c2" + - "696bd414316c07a505ec68e0" + - "696bd421590d785c13a7611e" + - "696bd427715ed751e5e1025b" + - "696bd453c75cf6e4320e2aee" + - "696bd4582a9a68766920df10" + - "696bd4695309532923d3f413" + - "696bd4772cdacabcc9e6f56e" + - "696bd4dcd1a6773989721ff4" + - "696bd510a8793724a1ef2255" + - "696bd52fa04610abd4bb3174" + - "696bd5654efa2e122be603eb" + - "696bd5672cdacabcc9e6f583" + - "696bd57ea907c8962e22aa3e" + - "696bd5849839477cebe21e0f" + - "696bd5a0675f6396322f9506" + - "696bd5a3e1b545a3ce075308" + - "696bd5c95e774190a1c97900" + - "696bd5ef7a53efd0cd3a2d5f" + - "696bd5f26fc0c09c00a97eef" + - "696bd606c426c607c88b2fd3" + - "696bd62e2a9a68766920df35" + - "696bd6339a140887948c6079" + - "696bd65d37193673e4f89b38" + - "696bd677b0b2b22cc0d1f850" + - "696bd6afe4eb4a667bf0e01e" + - "696bd6be883568f6441c8461" + - "696bd6c6a491d5abbd79559a" + - "696bd6dd1ae707d8c0922cc0" + - "696bd6fa9839477cebe21e30" + - "696bd700d3cc9dc76c56192b" + - "696bd7172cdacabcc9e6f5a9" + - "696bd7a7590d785c13a76167" + - "696bd7d2f7b3293476cab9eb" + - "696bd7e6da6069828e20d155" + - "696bd812c426c607c88b3008" + - "696bd8255fedf8533255e09a" + - "696bd83572239e0ab4634d12" + - "696bd8541376c8db5e315576" + - "696bd858634a1e1f469279f3" + - "696bd86f4d5f3676bb0ab233" + - "696bd8c493818b7b258ec458" + - "696bd8e8a6f1d19c01834630" + - "696bd8e933191fac7d918067" + - "696bd8f8ae7ec71f86571893" + - "696bd90a4292c4c779912607" + - "696bd91770d628f8b074254c" + - "696bd9a167f2110cf0fbe30e" + - "696bda14c426c607c88b303a" + - "696bda158149bf43938953cd" + - "696bda6d1e187f32c5eb5e12" + - "696bda9e55db5be0d7707242" + - "696bdaf8778afed01d13caea" + - "696bdb0b60ebbd99811cbf19" + - "696bdb0d894c2e015aaa10e2" + - "696bdb58c426c607c88b305a" + - "696bdb6f70d628f8b0742580" + - "696bdb87f79f22003e620560" + - "696bdba6e34991e438ab3a29" + - "696bdbc54a936a4d1328ab54" + - "696bdbd7476e9a6e777ffff1" + - "696bdbd9015eaae35d1dd6af" + - "696bdc031894fc36a07d90fb" + - "696bdc1f1894fc36a07d90fe" + - "696bdc281ee2a915e1934d33" + - "696bdc29b17075c2c32fa9e1" + - "696bdc5128ab12fee3735606" + - "696bdca2675f6396322f959c" + - "696bdcb855db5be0d770726e" + - "696bdcc7d9a2a4e1bf836fad" + - "696bdd44556c93e45170dfdf" + - "696bdd4d8e6d0bc988967b67" + - "696bddb15309532923d3f4de" + - "696bddb9b7080af9dcc944e7" + - "696bde0934614160c6097e3f" + - "696bde13e7200945ac1fd151" + - "696bde2de2eca952034eda08" + - "696bde327a53efd0cd3a2dfa" + - "696bde8c86833e290580c492" + - "696bdeb6fe5c03f9da388a78" + - "696bdedfa12eec59c747563a" + - "696bdef85fedf8533255e0df" + - "696bdf2d0f689873a4d2b470" + - "696bdf3c473aede2ca7effcd" + - "696bdf42900932d25e869a66" + - "696bdfd5c0aa1c772e76f3d8" + - "696bdfe1c426c607c88b30cd" + - "696be0307e6f94bf34143cc5" + - "696be054590d785c13a761f5" + - "696be0622c7243cd72c85217" + - "696be07595be6b21592587dd" + - "696be078d99e01dc631f402a" + - "696be080f38155a622441cc2" + - "696be08660e10e8f87b58647" + - "696be0928e6d0bc988967bb2" + - "696be0bd8a0fa4aac9be4952" + - "696be0c8c19cc80a4d8b0285" + - "696be0d079bc0402722c268b" + - "696be0d29839477cebe21eed" + - "696be0e2c676295b4ae74d93" + - "696be0e4a12eec59c747564f" + - "696be0ee0c52a10600ae878d" + - "696be0f3473aede2ca7efff1" + - "696be110eab4449d0aabe9db" + - "696be12cef11876d0045f9db" + - "696be1353897ce0bf0acad48" + - "696be13afadd8a32efe9a208" + - "696be13e1cafc9f413b5ba3b" + - "696be14d6377d0871b6a93ee" + - "696be154634a1e1f46927a8a" + - "696be167915ce3edd0967ad4" + - "696be16f6fc0c09c00a97f89" + - "696be19aa8793724a1ef2355" + - "696be1f855db5be0d77072d8" + - "696be1fd522a3f9be47908be" + - "696be230675f6396322f95e7" + - "696be252316c07a505ec69d0" + - "696be26779bc0402722c269b" + - "696be2d4894c2e015aaa1159" + - "696be2da9a6d79db2cf88953" + - "696be307f249daed52be5cfb" + - "696be32828ab12fee37356b2" + - "696be3324c43650cbf6f6046" + - "696be335590d785c13a76232" + - "696be371ef23530d3c5a38a0" + - "696be3a372239e0ab4634dcd" + - "696be3aecfb39c7ef2d48f05" + - "696be3eb67f2110cf0fbe3a7" + - "696be3f2a8793724a1ef2383" + - "696be3f9c0aa1c772e76f415" + - "696be3fa3724208fc971e458" + - "696be449a491d5abbd7956a2" + - "696be491800cd1498945048a" + - "696be4a1473aede2ca7f003a" + - "696be54f8a0fa4aac9be49b2" + - "696be585590d785c13a76253" + - "696be5a1715ed751e5e103ef" + - "696be61901fc6bd69561e3fb" + - "696be6331ee2a915e1934d9e" + - "696be64e632cfa55ccdfa954" + - "696be79e28ab12fee37356f1" + - "696be81c1d3959e2e052b369" + - "696be96d7bb51948f9de6504" + - "696bea400f689873a4d2b4a1" + - "696beaf1a046ed65e2904230" + - "696beb98cca8f940bc789ef4" + - "696bee33778afed01d13cb5b" + - "696bef0aca28cac2415e6102" + - "696befb2a8793724a1ef23fb" + - "696befc575020a567fd379c4" + - "696bf105a8793724a1ef2401" + - "696bf5ea522a3f9be4790962" + - "696bf8158d4f708447df89db" + - "696bf8174efa2e122be60510" + - "696bf858c6cb022a9da40587" + - "696bff6eb0c23fa5919d1355" + - "696bff803897ce0bf0acadb0" + - "696c5b784f526320afafb0a5" + - "696c5c22804859de93b596a2" + - "696c5c78199912be320e26d1" + - "696c5c80fb3b46f4c1e49f66" + - "696c5cf94b4572ebf92598f2" + - "696c5d35f469ea2c3cd715db" + - "696c5d6dc0c2d9149de6554c" + - "696c5d73bcc428fa4d1b2ae6" + - "696c5e4f024933da9b5c97cb" + - "696c5e72ead9a484382f2f83" + - "696c601818824db10c3a76a1" + - "696c604d6ed3eb90778f0276" + - "696c60521fe17643f95f9b3d" + - "696c60c6fbe7ab28abed2101" + - "696c6122296b5bace896d3c3" + - "696c6e0ea6fbe50a32562b4c" + - "696c6e39b0a5cdf2b1a73b06" + - "696c6e3d3cdf7d2e94e4acb4" + - "696c765a18824db10c3a7703" + - "696c79ed8cac7ab9693b0c96" + - "696c79ef4b2b73eea2f1c85b" + - "696c92e2804859de93b59864" + - "696c92e3199912be320e287e" + - "696c9376c7470f85da9343c1" + - "696c94c8a534b5c91758882e" + - "696c95efde00c290bacefb68" + - "696c964ec14d3f6041dd68c5" + - "696c96f7c0c2d9149de65720" + - "696c97a9a534b5c91758886a" + - "696c97f746f3fd69cfbea122" + - "696c99dc38f4b4346c5317ef" + - "696c99f446f3fd69cfbea14c" + - "696c9a2fccc462258d648e2b" + - "696c9a4dc7470f85da934457" + - "696c9a611244464a4aac1921" + - "696c9a6ffb3b46f4c1e4a144" + - "696c9bd10576d930a343fc95" + - "696c9c3afb3b46f4c1e4a16e" + - "696c9c42a199224e8d662189" + - "696c9c4638f4b4346c531824" + - "696c9c834b2b73eea2f1c9d1" + - "696c9cff531c58652b43fffa" + - "696c9d06eaaf0411bfe58983" + - "696c9ed42bd5bc335fb2a71c" + - "696ca025d95827836b14b4be" + - "696ca0d642dcbdaae6ac0dff" + - "696ca19a960611157f73e389" + - "696ca253fad368fca385f3b4" + - "696ca322ead9a484382f31f5" + - "696ca3e738f4b4346c531864" + - "696ca5851244464a4aac198f" + - "696ca73738f4b4346c531874" + - "696ca9c5a534b5c91758892e" + - "696ca9c638adf4f372b451aa" + - "696cacfebcc428fa4d1b2d63" + - "696cb0eefb3b46f4c1e4a1fb" + - "696cc896d95827836b14b5b9" + - "696cca3b5a4e3449377212b2" + - "696ccc7218824db10c3a7a09" + - "696ccd09eaaf0411bfe58ad3" + - "696ccd9d296b5bace896d75b" + - "696ccec7de00c290bacefd81" + - "696cd132de00c290bacefd8e" + - "696cd4eb6ed3eb90778f063b" + - "696cd9250576d930a343fe54" + - "696cda6f5a4e34493772131e" + - "696ce23ba6fbe50a32562e40" + - "696ce3d8de00c290bacefdf9" + - "696ce4a7a534b5c917588a38" + - "696ce57da199224e8d662386" + - "696ce73e960611157f73e534" + - "696ce8eda199224e8d662396" + - "696ce9cb38adf4f372b4531c" + - "696ceefa88af7cc0806368bb" + - "696cef00960611157f73e559" + - "696cefe246f3fd69cfbea398" + - "696cefe2de00c290bacefe31" + - "696cfa1f344d3cc10178aa0e" + - "696cfc408f609c3c2d3704b8" + - "696cfd2f0576d930a343fee2" + - "696cfd3fbcc428fa4d1b2f18" + - "696cffd34f526320afafb572" + - "696d0169804859de93b59b51" + - "696d02376902ce4cbbd7e333" + - "696d023ba534b5c917588ab1" + - "696d04fca6fbe50a32562ed5" + - "696d0c3efad368fca385f596" + - "696d0c4338f4b4346c531ab3" + - "696d166c6ed3eb90778f0759" + - "696d2b8dd9f2dbab1eafb813" + - "696d2bb4467a068ac4aa2027" + - "696d2bcb1fe17643f95fa0d5" + - "696d2bde18824db10c3a7c27" + - "696d2c042b50484b18da1d7a" + - "696d2d02de00c290bacf0015" + - "696d2d274b2b73eea2f1cd88" + - "696d2da3804859de93b59cb1" + - "696d2db31b5841990138a156" + - "696d2e319fed0d14156074ae" + - "696d2e7d85dbcb9294c28415" + - "696d2ecfc466796aa2ebcf08" + - "696d2f38bcc428fa4d1b306c" + - "696d2f3d8795da3e1289ad9d" + - "696d2fe2d95827836b14b81d" + - "696d2feb39c7a2604b9e618c" + - "696d2ff439c7a2604b9e618d" + - "696d3003911ed385aedae28e" + - "696d302446f3fd69cfbea577" + - "696d3024961f9960d100021c" + - "696d30287a3764d501282705" + - "696d30352bd5bc335fb2aaae" + - "696d30587a3764d501282709" + - "696d306a00086fc00b80fefe" + - "696d3070822d75373bb0b138" + - "696d30c1619debe3a900b108" + - "696d30db11f27a89bd9d1321" + - "696d30db18824db10c3a7c84" + - "696d30dff576e326d5950493" + - "696d30e948d5271e70d128b5" + - "696d30f06902ce4cbbd7e4e7" + - "696d30f6911ed385aedae29f" + - "696d311d38f4b4346c531c05" + - "696d313a804859de93b59d01" + - "696d317f4f526320afafb71c" + - "696d31874b2b73eea2f1cde6" + - "696d31bb0df30cb8cb92a743" + - "696d31e3804859de93b59d11" + - "696d31e90576d930a3440085" + - "696d31f1031535a033d8c64c" + - "696d32370c056d556963317d" + - "696d323abcc428fa4d1b30a8" + - "696d323b586e0a17ad22d739" + - "696d326d822d75373bb0b160" + - "696d3293d95827836b14b858" + - "696d33068cac7ab9693b11e5" + - "696d330711f27a89bd9d1346" + - "696d334100086fc00b80ff3b" + - "696d3367bf11c8700b4d8101" + - "696d337e0576d930a34400a8" + - "696d3bd77f2429a1932f4c05" + - "696d5ce6ccc462258d6491dd" + - "696d6504476577c0483a4446" + - "696d667021fdc01782029a64" + - "696d6895f63b90865845f905" + - "696d6a83961f9960d100038a" + - "696d6ad5c26c224df3462190" + - "696d6c5200086fc00b810036" + - "696d6da211f27a89bd9d149d" + - "696d6da9344d3cc10178ac3a" + - "696d6f02960611157f73e938" + - "696d6f61089658b1b2efb0b5" + - "696d6f6721fdc01782029aa9" + - "696d6f72b6b3247e544e1d9c" + - "696d6f77e4c876d0ca24450e" + - "696d6fce016bfb234ba80a46" + - "696d7031a1b45cea8a9c4e7e" + - "696d707eb24c0e1530b37d4b" + - "696d70ed52fdf733526bfb8b" + - "696d70fa4d68bde8f98d1a56" + - "696d713b1fe17643f95fa35b" + - "696d719cc97c5f71220eb78c" + - "696d7691158f7b32a240749e" + - "696d7696d1fbdec9580b4382" + - "696d776f887b21b4b1c624a1" + - "696d8d9568a79d2553bad6be" + - "696d8dd68efe8d0866aaf4a9" + - "696d8de53d5391e3c577b003" + - "696d8e308c9f2180b46cf1ac" + - "696d8e3b4f4c8d4f355cd2d5" + - "696d8e40f099d0b0111f1280" + - "696d8e57a888e5c0e1cb42ea" + - "696d8e985327712a4da21a47" + - "696d8ea09c6845efc42cb929" + - "696d8f2ffd57c67d83d7c194" + - "696d8f3f306d82b3e34c850c" + - "696d8f4bd07d4cdba345b775" + - "696d8f6eab7f9706a844f443" + - "696d8fa2bb9cd59b6d5333fe" + - "696d8faba7185124990ed4dc" + - "696d8fba0b245a8a6d1b61a5" + - "696d8ff177bf62ada9a019ed" + - "696d902fa998602bd03d13a2" + - "696d90e1cde72530e232be9b" + - "696d90f87f9c99fea65ee7d7" + - "696d9114a7185124990ed4fc" + - "696d918e85537287d8383bef" + - "696d91ae95f1143f0bed3ed2" + - "696d91dbfda614253850647f" + - "696d9200b392baf61cbedc54" + - "696d930c85537287d8383c0c" + - "696d9336436c4497755a06ec" + - "696d933b5dd80f591821e91b" + - "696d939584bd47b19f096ce0" + - "696d9402fd57c67d83d7c1f2" + - "696d9408f8199fcb171dac84" + - "696d9422436c4497755a06f7" + - "696d943382e1ac2de620ac9d" + - "696d9433f3fb8768f0f682aa" + - "696d944f4c3d4673e391741e" + - "696d9450b392baf61cbedc88" + - "696d94523b7bea24c781535b" + - "696d9458f4dec6e6d6ff7fa6" + - "696d9460557bad554d70809d" + - "696d9490955e9ffeec13b4eb" + - "696d9490ef17af78f93b6785" + - "696d9493306d82b3e34c854f" + - "696d94d3b1bd0ca2b0048821" + - "696d9543cb3d08a3aa82b805" + - "696d95454d91ecc286655fec" + - "696d95941a736a5afd8751e9" + - "696d95b88e3e3494c788246e" + - "696d95fa306d82b3e34c8571" + - "696d962c85537287d8383c4a" + - "696d96a7bb9cd59b6d533472" + - "696d96ac91191c601bacc476" + - "696d96b6d07d4cdba345b7d0" + - "696d96c57c11eaba88f2de64" + - "696d973384a44b99552ff505" + - "696d97398450426c4f113dba" + - "696d97e7c8145b08d49e81d8" + - "696d9809d9ff44e21f26f33c" + - "696d981ba998602bd03d140f" + - "696d9837f1b842325f9e7fc0" + - "696d98621c142afc32b0a6a7" + - "696d987664124db03d5204de" + - "696d989ecb3d08a3aa82b824" + - "696d98e9ed032ce6cfdcd3c0" + - "696d993f76ff4b529c5066f8" + - "696d9947557bad554d7080c9" + - "696d9948188e5ff22744e53d" + - "696d998137dbf74800ad874c" + - "696d99b29c6845efc42cba13" + - "696d99ec4f3681a9d821c166" + - "696d9a05fdc8cef2b8bd651a" + - "696d9a065d7a7b9a57b7d832" + - "696d9a1d1dee433b04ff67be" + - "696d9a3a8061c372104e6a06" + - "696d9a7ebbfffe73f206b83c" + - "696d9b1a7964535dda78b05c" + - "696d9b238efe8d0866aaf54e" + - "696d9b2ed724a55c2ff9a5aa" + - "696d9b30ef17af78f93b6829" + - "696d9b5b2f22175286ab2676" + - "696d9b743761ff9bfa02deaa" + - "696d9b8c1a736a5afd87524b" + - "696d9b99a888e5c0e1cb43c3" + - "696d9bb0d1d395fa4960afac" + - "696d9bb83bf8e39ce3da3231" + - "696d9bdac8b3144c937f8bee" + - "696d9be2c354e2fafcf2026f" + - "696d9c398a3d37c559a148f1" + - "696d9c3c0ca72a6c983442e8" + - "696d9c543943b9f942d7e8da" + - "696d9c685327712a4da21ad4" + - "696d9c7377bf62ada9a01a83" + - "696d9c7b2f22175286ab268e" + - "696d9d05c93f59bdc6e20d81" + - "696d9d3282f4c5e6dfcb4406" + - "696d9d3dcde72530e232bf33" + - "696d9d836a6524c5b9dbd030" + - "696d9da23943b9f942d7e905" + - "696d9e60cdaff5670111cd9f" + - "696d9e7d45d110a892dd8ff9" + - "696d9ebe92f396e3ef80848b" + - "696d9ef23fb0f5175cce0ae0" + - "696d9ef6d7f80a26e0463de9" + - "696d9f01b2ed922c6673bcb5" + - "696d9f10f1b842325f9e8042" + - "696d9f1104c9be98d76d6102" + - "696d9f1936733bd98d9d5c91" + - "696d9f1d76a542b835f9cb0b" + - "696d9f23b762e60f1f34fce1" + - "696d9f453bf8e39ce3da326a" + - "696d9f4a613c6035217a4280" + - "696d9fc948f6e4a7f7b86868" + - "696d9fd20b0e617c76c599e8" + - "696da01feaf21ce266daa5b8" + - "696da02882e1ac2de620ad4c" + - "696da03abad2864f69af2c21" + - "696da08c80f51286a174d2f3" + - "696da0d800b24b8aa54c2371" + - "696da0da1b41dc512e994a58" + - "696da0dcf1b842325f9e8061" + - "696da0f1e8344b3ee3226e0a" + - "696da11c82ce9ecc67d33a97" + - "696da1640e468cf1b7075aff" + - "696da16cc0a5bbe7897d4bba" + - "696da1742366ab3c47e2f2ee" + - "696da1769d0093f9cc8cdaf7" + - "696da17b4e52be266e5c8f5f" + - "696da17e38760cc5e991491a" + - "696da1838d4e41e87521a543" + - "696da185ed032ce6cfdcd44e" + - "696da1af2366ab3c47e2f2f0" + - "696da1e651d1a68eb767cacd" + - "696da1f969fe7e32c96f1b0a" + - "696da22ffcf19b7cd66460da" + - "696da2382ddbc19f3c9b08a4" + - "696da2387d42972c2ce7e077" + - "696da23b84bd47b19f096e25" + - "696da23e5c22ac3f11f64586" + - "696da2414dcd2eaf6d5fb99d" + - "696da243c8145b08d49e82a7" + - "696da24876a542b835f9cb54" + - "696da24ecfe06ac65aaf30a5" + - "696da253e0ac2bc918d4680a" + - "696da31cab7f9706a844f547" + - "696da35151d1a68eb767caed" + - "696da359103547e961005ef0" + - "696da361925cce50c37ecae3" + - "696da37374465a4afde00d11" + - "696da387c93f59bdc6e20dee" + - "696da38d2db0183ebc4d99e4" + - "696da38f7964535dda78b11e" + - "696da39e51d1a68eb767caf4" + - "696da3aeb1bd0ca2b0048905" + - "696da3b369fe7e32c96f1b2d" + - "696da3dafd57c67d83d7c2e0" + - "696da3ec34845fc3cbe39c68" + - "696da3f1265e9641b943c0d7" + - "696da3f2960ada878296efa7" + - "696da3f5f4dec6e6d6ff806a" + - "696da418154a7aa736439d3d" + - "696da41f265e9641b943c0db" + - "696da422997cb67cb492f08d" + - "696da422c1956556950adebf" + - "696da4374d91ecc2866560a1" + - "696da4470fdf525db6859b7a" + - "696da4857c11eaba88f2df4e" + - "696da4867c11eaba88f2df4f" + - "696da48c38760cc5e9914956" + - "696da48e69fe7e32c96f1b3f" + - "696da49076a542b835f9cb83" + - "696da4eefe3639ea55b00a12" + - "696da4fc4d2f71b8d51ea358" + - "696da510ebb60123bf5f69d3" + - "696da51148f6e4a7f7b868cd" + - "696da52ede1537c60354da17" + - "696da53e854c86e90cade7a8" + - "696da558ddefc568d7372620" + - "696da559557bad554d708179" + - "696da57c64a75ad017733b7e" + - "696da57d69b03d8bffc3f0ea" + - "696da58fdb77d8fed97f897e" + - "696da591913cbec0b8bf10ab" + - "696da59382e1ac2de620ad9b" + - "696da5a504c9be98d76d6190" + - "696da5bc2db0183ebc4d9a0f" + - "696da5bdbcbf0c32b4b22b55" + - "696da5bf2a009ecf67991057" + - "696da5c1fcf19b7cd6646100" + - "696da5e97f9c99fea65ee8cb" + - "696da61716e44e8c77d6ab2d" + - "696da61a64a75ad017733b84" + - "696da634819d3d63e28a2cd8" + - "696da647d07d4cdba345b8cc" + - "696da64fc5cd8b02c240f866" + - "696da6860838d9b98d822df0" + - "696da6e5b2ed922c6673bd40" + - "696da6f5188e5ff22744e624" + - "696da71af099d0b0111f13d0" + - "696da7360fdf525db6859b99" + - "696da74033f2abbafee7e633" + - "696da740fd6a4da2c4f2739c" + - "696da76aaac016e2daaa5905" + - "696da7a449652a41264cf4a7" + - "696da7ab261c825cae4b2bdb" + - "696da7dc62c78a1f9f3dbcff" + - "696da830bbfffe73f206b96b" + - "696da837d1d395fa4960b098" + - "696da83cf099d0b0111f13e9" + - "696da90ad355c133d5ff9a32" + - "696da96f2366ab3c47e2f37f" + - "696da97462c78a1f9f3dbd10" + - "696da9f6ab6134e6dd7ccf6a" + - "696da9f8960ada878296effd" + - "696dab0eebb60123bf5f6a16" + - "696dab10d04b4c84035b4363" + - "696dab12683792d4938d4255" + - "696dab55cc2f1314e6821e0d" + - "696dab8f2366ab3c47e2f38c" + - "696dabb20b1d6d58cb8b214b" + - "696dadead2a8ebe062803928" + - "696db384e8344b3ee3226ec4" + - "696db5f833f2abbafee7e677" + - "696db7828802b4780aaa6431" + - "696db7837cd183bae0552acd" + - "696db80053c4c590ef2ade90" + - "696db99c3bf8e39ce3da3343" + - "696dbf0fbbfffe73f206b98c" + - "696dc302e21603c880525749" + - "696dc489f5e5b97b2c96f97e" + - "696dc48b8802b4780aaa643d" + - "696dc493dc787523fe33b682" + - "696dc55477bf62ada9a01b83" + - "696dc555997cb67cb492f10e" + - "696dc555d2a8ebe062803942" + - "696dc5576256b37b6b389aa8" + - "696dc6250b15943f95116162" + - "696dc62c650ccadd25513257" + - "696dc6f2ab7f9706a844f5d5" + - "696dc7a8c5cd8b02c240f8bf" + - "696dc7ac0e468cf1b7075b79" + - "696dc7bc683792d4938d4276" + - "696dc833297aa9d32bbb0c80" + - "696dc83473a971419ff9fdd0" + - "696dc86a3e5d20b5baf4cb42" + - "696dc9821a736a5afd87539f" + - "696dc98d3ae226ec1ef35480" + - "696dca1fab6134e6dd7ccfa6" + - "696dca25ebb60123bf5f6a35" + - "696dccbba7185124990ed657" + - "696dccf9d84f13d3d33d4dd1" + - "696dcd313eec4106b2aa4cba" + - "696dcda530d6b6f26a006f45" + - "696dcdc85c34c44bef01e394" + - "696dcfb8d84f13d3d33d4de3" + - "696dcfd59d65f7c2b62e813c" + - "696dcfde9d0093f9cc8cdb9e" + - "696dcff730d6b6f26a006f50" + - "696dcffd4f3681a9d821c2b5" + - "696dd0110b15943f95116186" + - "696dd04884bd47b19f096ee4" + - "696dd09b2a009ecf67991102" + - "696dd0c6e21603c88052575a" + - "696dd0c8eaf21ce266daa6c1" + - "696dd0ed51d1a68eb767cbcd" + - "696dd1161b41dc512e994b68" + - "696dd117f4dec6e6d6ff811b" + - "696dd18b1b41dc512e994b6c" + - "696dd18e592a67f25a7b8a4a" + - "696dd1fadb77d8fed97f89ff" + - "696dd1fd73a971419ff9fe01" + - "696dd292db77d8fed97f8a05" + - "696dd30cd724a55c2ff9a6fd" + - "696dd46bdb77d8fed97f8a0d" + - "696dd52033f2abbafee7e6d5" + - "696dd524819d3d63e28a2d63" + - "696dd7fcfa3c85a00c23b974" + - "696dd808f091d1b91dfec304" + - "696dd9bc2a009ecf6799111f" + - "696ddc1774465a4afde00e04" + - "696ddc62a7185124990ed6ad" + - "696debf049cc828bcf021a36" + - "696debf0bffa1ca055842140" + - "696dee5328c744a0fcac8a9d" + - "696df3d0ec7b7b20f30a6d7d" + - "696dfb67b057bb83a072af40" + - "696dfbae482496d71b0f533e" + - "696dfbaf1aa1ee40a4aa8d7f" + - "696dfbb0cc5f7c104dea8618" + - "696dfbb37046f49217da33d0" + - "696dfcb62040d53792cdee8d" + - "696dfcb766866953159e99cc" + - "696dfcb8cc1784359b00964d" + - "696e116352cb0443238510e5" + - "696e120b9a00fe678e7399c6" + - "696e15ddc5e0221414f1b677" + - "696e1e0b3e7dae3f86d26cd0" + - "696e2c574ec64da83885be07" + - "696e2d7d9f6952d8c9d855c4" + - "696e2da1cc1784359b0096a2" + - "696e2deb66866953159e9a2d" + - "696e2e3517377c6d0e55793b" + - "696e2e3699548def4412ca86" + - "696e2e3cb1d3ede5703dfa54" + - "696e2e4d37446407d01eae04" + - "696e2e5828c744a0fcac8b20" + - "696e2e7d153584e0fd42c219" + - "696e2ea91944b7bc7a2c6fab" + - "696e2f19946d3894eba261a8" + - "696e2f7c7d60b0a13a324472" + - "696e2fcc7046f49217da346d" + - "696e2fdb6baf731a9a6e689f" + - "696e30051a287a92a578655d" + - "696e30794b3bc2a1bad5769e" + - "696e307e32c6adbc0a678303" + - "696e30bb49cc828bcf021ae1" + - "696e30c2bb1a0af92b1419d5" + - "696e30c9b588c5a56a13cd45" + - "696e30ce57757f9d4e3a19e8" + - "696e312e39a036db70dcac57" + - "696e31b0dd0cb8d4e5206248" + - "696e31cdb5a66458ce86c546" + - "696e31ceeb3c1de699cccc34" + - "696e31d17ce03c4a8965ec87" + - "696e31ee946d3894eba261d1" + - "696e32089189b378a622ab2b" + - "696e3215589599cd079088f7" + - "696e32173644d45b8cb8a8f7" + - "696e321ea61150e395603649" + - "696e323e69f74feb51ca4ded" + - "696e32b055c6d472ce728642" + - "696e32c6378702184538f6d9" + - "696e32ebb588c5a56a13cd79" + - "696e33070299a6b3c6c692d4" + - "696e333357757f9d4e3a1a15" + - "696e3335f72d126a840ed455" + - "696e337121cd3dcb7f05c27f" + - "696e337780a86712b2f3dd8c" + - "696e338f43f23b8635c5a17e" + - "696e3419b79e9d7860a3b613" + - "696e3447946d3894eba261fb" + - "696e3448f0f65980ee0bb86a" + - "696e3991dd0cb8d4e5206289" + - "696e3cd8946d3894eba26227" + - "696e3e0d843b98de1f7e8879" + - "696e3e587ce03c4a8965ecda" + - "696e405b39a036db70dcace9" + - "696e40c64b013c8444628df0" + - "696e4134ac9f0dbc3ed6bc2e" + - "696e43695e4e2b7186cbc6af" + - "696e4373843b98de1f7e88b1" + - "696e43d52040d53792cdefd1" + - "696e4d547d60b0a13a32459c" + - "696e4dc4589599cd07908a13" + - "696e4f900299a6b3c6c693e8" + - "696e51956edf6acdd058baf1" + - "696e51f77e219e09d76217ae" + - "696e5215620da3c71e5ed1ad" + - "696e521d1944b7bc7a2c70b4" + - "696e535b16141fc0a6fa28a5" + - "696e53e69189b378a622ac78" + - "696e54f08587a69d44bf7b79" + - "696e5c683e7dae3f86d26d8b" + - "696e61ba9a00fe678e739b25" + - "696e632d52cb044323851277" + - "696e6652f78a5e583f4e61b1" + - "696e66596599d861b72712ae" + - "696e6663b1a65193a3f4f707" + - "696e67027e74165050bc5938" + - "696e6813cec1370af5381686" + - "696e6a36cc1784359b00982e" + - "696e6a5c3488965375fbfdf8" + - "696e6bd7d5450f3c814f569b" + - "696e6bddc27c731ddfad65fb" + - "696e6c5955c35f540282662d" + - "696e6d3270007c739ba09d13" + - "696e7111d58bf710bcdcd7e2" + - "696e71961ce919914eec71e0" + - "696e7279fb96a896b1b154b7" + - "696e730012a7f5ee4b96045b" + - "696e7303e2287068cc486beb" + - "696e742396d7c013fcbad002" + - "696e749c43c9ca3b047729d0" + - "696e75005adb904750879a48" + - "696e7501a479c0684c886fa2" + - "696e750f4d1c362ce583b634" + - "696e751f1b679d5d037a5463" + - "696e752ab1a65193a3f4f76d" + - "696e752ac72152e151130f2f" + - "696e75460f903881d07fece9" + - "696e757d9ec04eb46fe06643" + - "696e75c22934fcc9b4fff6d1" + - "696e75d5cf86b2bb0757a714" + - "696e75e6235ddc6771e37690" + - "696e7a25d5b9e89f3df82a38" + - "696e7a4649cc828bcf021ccc" + - "696e7a628587a69d44bf7c19" + - "696e7aac16141fc0a6fa29f4" + - "696e7bdce85937dc7bae4f0a" + - "696e7c025f0b34357517d878" + - "696e7c848db2ad069baadfae" + - "696e7c8a8db2ad069baadfaf" + - "696e7cbe7650e29d1f028df6" + - "696e80e2fcee3b967216195c" + - "696e8134fea209e4672a3668" + - "696e81951ce919914eec7276" + - "696e81fe4d222dfe61ac9bb9" + - "696e82a9f30588a493c07df9" + - "696e832c28c744a0fcac8d7b" + - "696e83580ebd516bc2eb4d24" + - "696e836ed2510d367bc0b3a1" + - "696e83929ba986538aa44e02" + - "696e83aa0ebd516bc2eb4d29" + - "696e83afb588c5a56a13d079" + - "696e83c5e2287068cc486c6f" + - "696e83dd7d60b0a13a324710" + - "696e83e14b013c8444629071" + - "696e841d4b3bc2a1bad57884" + - "696e8454196b035bf4865044" + - "696e8470f72d126a840ed773" + - "696e84753fd4b5d896cd24fc" + - "696e8476a95acfe764b582dc" + - "696e8477f72d126a840ed774" + - "696e8481620f4235f99e5087" + - "696e84896f3b64df08ff0ffe" + - "696e8489887a0e735b83d27c" + - "696e84fb4ec64da83885c0e3" + - "696e84fc391564e464168065" + - "696e85232754ca7ba879e3b3" + - "696e85468395ff83f6754752" + - "696e854ce85937dc7bae4f80" + - "696e857e4b3bc2a1bad578a1" + - "696e85883fd4b5d896cd2517" + - "696e859065584499a5713be8" + - "696e861c6d44bc7cd096aaea" + - "696e861cb1cc409b12d985e3" + - "696e86280299a6b3c6c695ab" + - "696e8634283b3b0e6a54c962" + - "696e8639cc1784359b009958" + - "696e86afd5450f3c814f57b9" + - "696e86c6cec1370af53817c8" + - "696e86e32040d53792cdf1ce" + - "696e86ea30d80e9a7b76618d" + - "696e871496d7c013fcbad0c1" + - "696e8772589599cd07908c31" + - "696e87aa2565a5d3a3a29cfa" + - "696e88017e74165050bc5a4a" + - "696e883b6baf731a9a6e6bdd" + - "696e8859cc5f7c104dea894e" + - "696e8900ee9b736af1a796bc" + - "696e893857ad93e55dd089bb" + - "696e89512565a5d3a3a29d21" + - "696e89bc0ebd516bc2eb4d7b" + - "696e8adfa61150e395603722" + - "696e8b331ce919914eec72e2" + - "696e8cbdf04cd7fd4f334597" + - "696e8d1dd80c7b2e84481a1b" + - "696e8d4cdb9542d237e0a88f" + - "696e8e2ffe004d797507e0bd" + - "696e8e7fec6580da1aea7c21" + - "696e8ea84b013c844462911e" + - "696e956f5f0b34357517d9d6" + - "696e96fbf6607f22e6adf1ec" + - "696e97419ec04eb46fe06823" + - "696e974943c9ca3b04772b98" + - "69b0697a6ced4ef08c5edc6b" + - "69b084a725f7612803c0322e" + - "69b0861ce884d9f18d495b9d" + - "69b0862041d7ebf1301a1127" + - "69b095ccdcb67ea86c828550" + - "69b09d8ea65b5efd3e3d397e" + - "69b0cf3ebe31a1b6e41bdaa9" + - "69b1083a05af61de49d9077c" + - "69b130386589f8d682a783ad" + - "69b1c6b6e280adc034ba7e47" + - "69b1cdd6fb0169196bd689fb" + - "69b1ce11e63780415fa33b32" + - "69b1ce3ee4de0a726f6782a3" + - "69b1ce7cbcbda546478e5064" + - "69b1ce850013d80ee905b84f" + - "69b1d4f42ad31f34e7cccd51" + - "69b1dcbae8792f2c51c86f15" + - "69b1dd6db5e7ce82471bede6" + - "69b1ddaf81e27d6f5fba8b16" + - "69b1df34f569725b8b532848" + - "69b1e164206612d7447a399a" + - "69b1e32ec85f511b1136a79f" + - "69b1eed98d14b3a3def4c502" + - "69b1ef4924f33b6e4f1c95c0" + - "69b1f02718993dd73de56131" + - "69b1f08681e27d6f5fba8b80" + - "69b1f122d984740505d53c54" + - "69b1f136f129f1da443bd361" + - "69b1f15c80702dbf9d8a9b76" + - "69b1f17a5a0fb3802c89c853" + - "69b1f1d04c9ec6fa2593d56d" + - "69b1f1ec334b7ef00d6e2e5f" + - "69b1f1f1f8851644abf4781d" + - "69b1f23be63780415fa33c15" + - "69b1f2542befaed4ff388a0b" + - "69b1f26a5118e04b0df5ba73" + - "69b1f2d19aad8d6b722d2d24" + - "69b1f57ca19ef1757a4c3f27" + - "69b1f62ec1eb0780d35cd950" + - "69b1f88fe3dc955cc0c9a048" + - "69b201ffd8ba604bd6793ce7" + - "69b20252dc49c994c7b73087" + - "69b204319cbfbb1b1e4e350e" + - "69b2051dedbd6e568a928f1f" + - "69b20666f37097c863b46706" + - "69b20b02e4de0a726f6783d4" + - "69b20e85f9bf9019ec27e238" + - "69b20ea5f9bf9019ec27e23a" + - "69b20ecb89befb4ecdbbec0c" + - "69b20f01177c2d1ec9b86ffa" + - "69b211af245baa91f8d627ba" + - "69b213e57731343a25ad673c" + - "69b217614bd0612877d5cf3d" + - "69b218739b2ddb196ef8c3ee" + - "69b21d35f3ac3ef1ee1a8036" + - "69b21d664d66b0ca3d76bd5a" + - "69b21df746c3ab7a2fbe205b" + - "69b21e212b591a028953a119" + - "69b22299ba6f75409fc1af7d" + - "69b222bef91939a9214216b4" + - "69b222ebe6c38b58c797950f" + - "69b226f690800ed928cc8900" + - "69b22b0d85ef15ac9d7b76f3" + - "69b22b11f6a3a400d9a7e60f" + - "69b22b55ae3b93109bbfcb39" + - "69b22eb8676b613268fedf19" + - "69b231098f86e88a3c385584" + - "69b231a34f6316a90f6fe0b4" + - "69b23221efbcd2c76460bcf5" + - "69b23521143884432d64a306" + - "69b237166ee0a65b26af56ed" + - "69b2386171a6e3af58486b80" + - "69b238e9bedffb1710d87922" + - "69b23b99e1aeaab880149e3b" + - "69b2428db73374c42d2a4539" + - "69b244e095c254fe3b23fa15" + - "69b244f00761ad56d71ecf50" + - "69b2451e3f4e9d4e5132ce5c" + - "69b245293a695858ae5978be" + - "69b2455d4408ecef199fd978" + - "69b245746f9a74d13ca35798" + - "69b2497e7bb14b001f7a4f64" + - "69b24a2fdde762835c6bd52e" + - "69b254ab7085b84aa6eee761" + - "69b255d8cb73b811533d9ced" + - "69b259c61582e5851578e6d7" + - "69b25d28210c7be6a2e28203" + - "69b25e5cd5ce800627bbd833" + - "69b26119bd4707e07f7fc2e7" + - "69b2611e88f72232a8dc46ed" + - "69b261a4f1c6ddd84cf41e33" + - "69b261af17b4703132d42c94" + - "69b261c773e0c0cb5c6473b8" + - "69b26254c6262d0cfc57e3f0" + - "69b262a9066d6a18e5924ac7" + - "69b2632fbd56caa41b097ed5" + - "69b2633a73169f5dddc9f293" + - "69b26348d19c7f83729c050a" + - "69b26388e9682d2fd58bedea" + - "69b263f3b23752df3d401f3a" + - "69b2648f2b7f513c95c2fa9f" + - "69b265139e3cc49f405fa328" + - "69b2651a6aa2bd4f824942fb" + - "69b265cdc709e412837eff35" + - "69b265e58e72a9ba85629c10" + - "69b26651ef30c1d95173ad19" + - "69b266faf6e55dd6eb5cdce1" + - "69b2675e0d7a870aae92f6f3" + - "69b267823fc70fb71500be68" + - "69b267a079b499a91398b9a2" + - "69b267abee3473fbe74ba351" + - "69b267b04d9f0b36cdc5ee7c" + - "69b26850a434cb5f7ad46d32" + - "69b2689f2a8194a3895e2dde" + - "69b268f69d438d979a2b346d" + - "69b2697f4dba5f9e2800c997" + - "69b269b3676b613268fee10e" + - "69b26a10ffaa869ba4e9ab39" + - "69b26a1a0aa38b10e4ea234e" + - "69b26a66fa70563d77f1453f" + - "69b26b4c3427fd8268efc3ec" + - "69b26e06e04a0a6b138d0d89" + - "69b26e78cc66963ad213ed13" + - "69b26ed91834da7ed8aa7f9c" + - "69b26f90d620fefc09520973" + - "69b26f99080ab0c125ed933d" + - "69b2702dc42185339cf3a126" + - "69b270b64762f62b97d204f6" + - "69b270b86888d75cca00be06" + - "69b270d223ac5094d7d569d9" + - "69b2710e5311b2f9a097f07c" + - "69b2711a7f8ad153dc9a6a33" + - "69b27158b23752df3d401fd0" + - "69b27190a2b513111a6ab711" + - "69b272735e74a765654520af" + - "69b272754e52e56ddb20e7ab" + - "69b2727da38e9e0d25df488e" + - "69b27291475ffb30a4a9e4d7" + - "69b272bb1c93ca972c984f0c" + - "69b272c5822e681009798257" + - "69b272dafffdb363bee89c74" + - "69b272e58de2a41ff64c0158" + - "69b27307ceb18b0a966718ae" + - "69b2731338baeee842aa5673" + - "69b273249793237729c3f3cc" + - "69b273374cdb10d53045af21" + - "69b27338d1408fb781ac0e18" + - "69b2736d71cd4373f30318ba" + - "69b273a00969c1529552921e" + - "69b273a2f90cf8a5373be522" + - "69b273d9fc265cf0ead4c100" + - "69b2e7b4316e4adbf7aee180" + - "69b31b5af025fa51b6483c3a" + - "69b31cf50507ef63cf806ccb" + - "69b322eaac608b469f1c971f" + - "69b3234507c73e5009bf1217" + - "69b3237b1eefcd823bba7fae" + - "69b323b2b3f11250e53e46d5" + - "69b323df90205ccb6a3a6a06" + - "69b323e2a6c190cb1207a618" + - "69b323f6d1c2f2aeebe9fe28" + - "69b323fc34f57e26be193c40" + - "69b3242eef9810a16792d067" + - "69b32476c28f98623d5a2caa" + - "69b32488c6e3f29ca249cf65" + - "69b3249636b59a71b2c5f588" + - "69b3249a0549b628de4a3db3" + - "69b324a929335659c0495a71" + - "69b324b18d7bdd17e47a1164" + - "69b324b23f676d1a02c7983f" + - "69b324bb353d87177cf7c689" + - "69b324d558a3c20f2e3369df" + - "69b324e7bb1bd09a3f78a348" + - "69b324fff0175c96fc1daa21" + - "69b3251089fcae309f49de1d" + - "69b3251e9d48ed0864532ac7" + - "69b325259cff451d5ad614fe" + - "69b3253f5f6bc8050b611626" + - "69b325426ada34d40e952fce" + - "69b32551c362cb6e75562dff" + - "69b3255f2969b8e96b022794" + - "69b3256d37113df794e259e5" + - "69b325755eb75d08482a4a09" + - "69b3257920d7109745fbcadc" + - "69b3257c4a8627413d50cf9f" + - "69b3257c9801d1a8fe97a3a8" + - "69b3257da58d9206c09fa7e1" + - "69b325951de4535f511ac9f7" + - "69b325cfc8c606b507cca680" + - "69b325d5a4fe1768fd5bb108" + - "69b326568696fd815823ce4d" + - "69b3266ad187c83f040eb6cb" + - "69b326ab8a768ab7570a67ac" + - "69b326b0c2d332cd2928bda6" + - "69b327510cb082e8dc6508e6" + - "69b327af3b84bf64c9af927a" + - "69b327b2d60700dd0e9c9526" + - "69b327d3241a6c270e2e2588" + - "69b327db291d11cb9f27feb3" + - "69b327ff837a0ecabff6150e" + - "69b32831b49f0369244a19d2" + - "69b328871b5a39aabdeb94fb" + - "69b328a0c60ae08dcc457b7c" + - "69b328a0f81f3e5bfc92733a" + - "69b3290f4f0d9c46954d46a2" + - "69b3293c3141c306f0648d6a" + - "69b3293ed953dd4a340d4a43" + - "69b32949432fc80fa7450f90" + - "69b329dc27db95f536bba514" + - "69b329e1532205dcd9dd9d04" + - "69b329eea88d24fd9605ea5e" + - "69b32a18bfab667e606e053b" + - "69b32be0a8a94add5376fba7" + - "69b32c80593bf525d3ce0df0" + - "69b32c9f593bf525d3ce0df2" + - "69b32cb5e78d84c7132c76e6" + - "69b32ccb8d7bdd17e47a11b9" + - "69b32cec3a8d3735c57f1f19" + - "69b32d3d837a0ecabff6152a" + - "69b32d7ef985ffb3e825c277" + - "69b32d84c4b58ad48f4fb6d6" + - "69b32e04fd5c75d1e4fd0f48" + - "69b32e2ab074d10ce5378997" + - "69b32e4cad2306b8f71eff0f" + - "69b32e774d00f2be43082dd9" + - "69b32ea6ea276be3f1869737" + - "69b32ea83ea5430cc96c53ac" + - "69b32ea871afa844e920a071" + - "69b32ee19c7cdaa31d3425d5" + - "69b32eff2f4b63d8f5071c61" + - "69b32f01255bebd995db4907" + - "69b32f154e09c0204d2d9c6a" + - "69b32f19eb3ba1dc65f231a9" + - "69b32f24aa247eb135f9d1c8" + - "69b32f24dcfbd323e893ae4c" + - "69b32f2510072caf7763aa04" + - "69b32f2fe538718b22123db6" + - "69b32f34f76fef9434c1f25e" + - "69b32f3f11a3052a5f6b20da" + - "69b32f451fea872861dae3fd" + - "69b32f46a8619846c0a1b414" + - "69b32f55b36470cd3c859fd6" + - "69b32f55d25cd91bfdee9699" + - "69b32f5a03501cccc085eb1e" + - "69b32f6633f47b3ac2cc9968" + - "69b32f6d4f7d5f0251c706de" + - "69b32f6f20ad0da684570a67" + - "69b32f76b35b907423cd01b6" + - "69b32f7dc93043812094d92a" + - "69b32f7f169575c34826457f" + - "69b32f85c40dfdcea6ba8f69" + - "69b32f893f676d1a02c79898" + - "69b32f9689fcae309f49dea5" + - "69b32f9805adab99ca3fb86b" + - "69b32fc2cf7a852768645d3e" + - "69b32fc533aae35229f6ab95" + - "69b32fd529f07b4abf05c695" + - "69b32fd584df4cc812089987" + - "69b32fd9893ed9307bb32633" + - "69b32fd9f59e011f372d0056" + - "69b32fda1091ba7c96161ea7" + - "69b32fde0b59f1a5f0d3c7e9" + - "69b32fe4d5e06a2d34f3bf2b" + - "69b32ff2ed19bd8136cc7ef6" + - "69b33009668c32fee52d4778" + - "69b3300c191697a20305e224" + - "69b3300d8fbf9f1396de7135" + - "69b330193716c8868efa180b" + - "69b33025a43225ffb32c393f" + - "69b33029691dc4a98f9a64db" + - "69b3302c9db030109c25166e" + - "69b3302dc01e2a69c530d297" + - "69b330304a83a977b34c822d" + - "69b33039147ccf34ec3eaf48" + - "69b33044c971c19f7392fc96" + - "69b330515df413b90aaa2d96" + - "69b3305b84df4cc81208998b" + - "69b3306db284d7d2c8bfb945" + - "69b330714d0b14f438e38c4f" + - "69b3307385df5e52dc3f7110" + - "69b33075644ff12da3d093ca" + - "69b3307ad5a9b3562d25f968" + - "69b330842b540726f0f0838a" + - "69b3308589fcae309f49deb8" + - "69b33088a88d24fd9605eaa4" + - "69b3308e2ab5c44e3b76e1ee" + - "69b3308e8bd25d700d7b0dbd" + - "69b3309487ca82b1335a2815" + - "69b330c7fec4f50e27148868" + - "69b330ce668c32fee52d477e" + - "69b330e1c9bbf6efb8566ebc" + - "69b330e978fb5b97c0e044f5" + - "69b330ea5ba13c3f3ef8a5b4" + - "69b330fb4e8da91e001117b0" + - "69b3310b5eb75d08482a4a62" + - "69b331122f4b63d8f5071c83" + - "69b3312098b985b4a8a45665" + - "69b33128a81e94c587138a53" + - "69b331424e8da91e001117b4" + - "69b33142847f334fade0dc53" + - "69b3318e61128ddef947efc6" + - "69b331b534d85de27d024c1d" + - "69b331c28998295a8f0295c5" + - "69b33243366031dc8c47dcf6" + - "69b332717be3eeb7d0d1b439" + - "69b332a48090086eef6d4973" + - "69b332ccfa6fb36d6634356b" + - "69b332ed3d3cd723898aed9e" + - "69b332f6536b3f17198194ae" + - "69b333420ab9031de4d88396" + - "69b333553e01b15c196d20f9" + - "69b3336e3ea5430cc96c53f4" + - "69b3339d25354b73c1be7262" + - "69b333a9f51fe27231b2e18c" + - "69b333c04ce9bb48ef5314fd" + - "69b333c153fce8af8163104e" + - "69b333c834f57e26be193cef" + - "69b333cb676c0cb59674fb46" + - "69b333df37646f9877133919" + - "69b333dfa94d0a5333fc414e" + - "69b333dff52d5934b64fb826" + - "69b333e74ac1cf94381e26e5" + - "69b333ea1b37513a0415b620" + - "69b334154d7476bca74df427" + - "69b33435d4a938348efecd5f" + - "69b334cfbc043d380839e666" + - "69b335093af62e50846a6414" + - "69b335199d48ed0864532b5b" + - "69b3352d4e780ed4635cbf24" + - "69b33546044619b12e1b415d" + - "69b33569c61e266c110d2906" + - "69b335fde59d54ba94056a94" + - "69b3364359cad34ac9d5d215" + - "69b33661dd416cf7ce9ce6bc" + - "69b33669f27356e96596aaed" + - "69b3369a65a3823ad1af856a" + - "69b336a784df4cc8120899c2" + - "69b336b18fd1a5755971f9f8" + - "69b336e492986a67bf306901" + - "69b336e940680cc48732e19b" + - "69b337050755b02085db1e80" + - "69b337707538c5a928643d6b" + - "69b33773547fcab36dc9f5fe" + - "69b3388a82fbb57ea61938d1" + - "69b33a2d045de00b157d98ae" + - "69b33a8856f8acacb6b7b11a" + - "69b33ae39f081fedddbaac56" + - "69b33c17337f3f7fe0d27f65" + - "69b33c7d0de6a32da02917a5" + - "69b33d3afe69842c5bdcd8e7" + - "69b33d3f53fce8af8163109c" + - "69b33e1a0de6a32da02917a7" + - "69b33e6d53fce8af816310a0" + - "69b33e7347da40a1fff3f655" + - "69b33ea7a0b9e58bff02cadf" + - "69b33ee0ed19bd8136cc7f4b" + - "69b33f2a2c02e998c162a9dd" + - "69b33f4f76b9f97b84ac06c1" + - "69b33f5cc012a885e4d5cdc1" + - "69b33f6ac31128ef5b906581" + - "69b34008bb1bd09a3f78a423" + - "69b34066d6de27d232702dad" + - "69b340965c46953805d1c898" + - "69b340e5ded09c8b367e1d34" + - "69b340ef78f24b0a5ce41a3b" + - "69b341021fea872861dae46a" + - "69b34108ecd1b2f90df5483c" + - "69b3419bc7eaf60b737fc4e4" + - "69b342048796b61661c69f26" + - "69b342b598d0431040a77d21" + - "69b342f28460e3ee4fdc5a39" + - "69b343c8ff13a367887b60de" + - "69b3449cdc165ff6bc8ef048" + - "69b344eb8fca50e0c5d4cea1" + - "69b345187f9dcce2c96e4e2f" + - "69b3453de3ebaf6326d3ab9a" + - "69b34812f27356e96596ab0c" + - "69b34837122e58b0d29ad7ae" + - "69b3484bafe7377840af8822" + - "69b349ce815e27488413cb0c" + - "69b349f725354b73c1be7293" + - "69b34a0734d85de27d024c85" + - "69b34a78dcec0551c6fd2075" + - "69b34ac584df4cc8120899fb" + - "69b34af98648e4230a699553" + - "69b34b2dcdd1d821dd00e4f4" + - "69b34b56740e46344116c019" + - "69b34c882c567018655d156e" + - "69b34cba4d00f2be43082e80" + - "69b34debb2387924d4345ba0" + - "69b35078482b50fa8312ada6" + - "69b350a948431c26200c364e" + - "69b350d63dcbaa000ca451d5" + - "69b350f28796b61661c69f72" + - "69b3518e8648e4230a699575" + - "69b351c77f9dcce2c96e4e68" + - "69b3605340aa1566f5e657d8" + - "69b3610ad9db1651a943ff7c" + - "69b3610f977aa45ae4d31ae7" + - "69b36125e78ebe6fc83d66b4" + - "69b36146de1724371ba7e869" + - "69b361fd069e4261884f5e73" + - "69b362214e9ee3fb43ca13e1" + - "69b36305b3ecae14b98c5a22" + - "69b363071e56f6bdf5a3738d" + - "69b363203cc0a128bfeedab2" + - "69b3644fb01f7ef56f93a3f1" + - "69b364a441cb7c6c960e80d1" + - "69b364dd96668ab3a811c4a7" + - "69b365376ef4b52990dfd87a" + - "69b3656486974057bbbebd13" + - "69b3659f586f11a266067968" + - "69b365b77bb6744c92804507" + - "69b365b8277560815f87b30b" + - "69b366439cf9b3f64ecae9cf" + - "69b366670ed50a663d2e793e" + - "69b366af12f60d6933bc30ad" + - "69b366b0b8e7cfef07cf7161" + - "69b36734f9212b8a6cd4a942" + - "69b36788846a8cd39d79fa5f" + - "69b368731e95afea0201344b" + - "69b368c02d6187f3f280a7ab" + - "69b369024af76c8acce9cb81" + - "69b3698f5475daac769f4500" + - "69b369ce9abb9f4fea72287e" + - "69b369d22ffeee6dd3707fdc" + - "69b36a403605be46bd68e20c" + - "69b36afada671b05d56ff305" + - "69b36b94de7d29e8b4260279" + - "69b36bc962bc507ba59bbb95" + - "69b36bd04af76c8acce9cbb9" + - "69b36c440d82a3f1acf9008c" + - "69b36c5f62e496bd343087de" + - "69b36c9ccc00f170ca94142e" + - "69b36d3c4af76c8acce9cbd5" + - "69b36d6cf186f27a705a6ab7" + - "69b36d7e5475daac769f4542" + - "69b36d8b254f96a47b997c5d" + - "69b36db130df4858f6ffacb6" + - "69b36df595cf361e5f36f489" + - "69b36e449afaa0a8c9b49678" + - "69b36e53a241870232dd5f8d" + - "69b3703aac1fb583171b095b" + - "69b3703e158f8ac239e0263b" + - "69b370e86588e26e947e21ee" + - "69b370f18b893e4661ffac03" + - "69b371659dcd140fca123cdf" + - "69b371af3605be46bd68e274" + - "69b371e9a241870232dd5fb1" + - "69b37313430aa61ff5df32dc" + - "69b373195298b5a76ce806ea" + - "69b37681c160c3e85bc87cf9" + - "69b37944210c48e1246ce9d4" + - "69b379ba7efe6d5409e208ce" + - "69b37b03011d1b02416456ee" + - "69b37b3d0142670cde191d71" + - "69b37b69aced996198f0a00f" + - "69b37c695298b5a76ce8073d" + - "69b37cb83c979f1b0ce8d01e" + - "69b37d2140085d3467341e5f" + - "69b37d29cf8a434cad8828c0" + - "69b37d77f76fef9434c1f4d5" + - "69b37d7a91c5174ed83d52ac" + - "69b37dde210c48e1246cea31" + - "69b37e9c68aadc8d4ab2b90a" + - "69b37ecd9cf9b3f64ecaeb22" + - "69b37ee408ac639d193f3074" + - "69b3801a9f541723fb1f5681" + - "69b380387a4da4dd70c077ff" + - "69b381774af76c8acce9cce7" + - "69b38ea5ac1fb583171b0a61" + - "69b390d87e3e8c6729eb003e" + - "69b3914bf391c63596cc9d98" + - "69b391d2f18d6f8a3726cfb9" + - "69b3921ca0f0fd0fb6b03f0b" + - "69b3953a91b55dc4ccf049a7" + - "69b3955aa6aafcdd15af5f9e" + - "69b39576765d6f4746693371" + - "69b395772a66956143e0e441" + - "69b395bc32031c42384b2ead" + - "69b3992083b1788551dd25ff" + - "69b3992bb1416016a1f2474c" + - "69b39b1e55f90884ebfc461d" + - "69b39b802876c0698ba5f2af" + - "69b39b96ced618554bc09128" + - "69b39c0426866068d9207a74" + - "69b39c15f58c31c981bb0fda" + - "69b39c3330df4858f6ffaf40" + - "69b39d38178586606fe1a5a6" + - "69b39d5e707c7f3e935bfd61" + - "69b39da9b064ee0065ec1504" + - "69b39e02bf6326e4fff7092e" + - "69b39e56e8f50389c2a6f962" + - "69b39f6c032adc6054d2aff0" + - "69b3a13154ebc24d61f8c8b3" + - "69b3a312e78ebe6fc83d6a98" + - "69b3a385a69da9855a769612" + - "69b3a39c3dc459c6590ba001" + - "69b3a3ab525e868ba0a33c48" + - "69b3a3b2d00bc76f562760c5" + - "69b3a3ca65743dd8f9bbf253" + - "69b3a3f377858a878e0c34d7" + - "69b3a554fa845d5efc1d0a20" + - "69b3a55c893833adfabf7316" + - "69b3a5a5b9775929f6ce8e53" + - "69b3a5ac3e30f5215d040e61" + - "69b3a5b819a9498410d236f0" + - "69b3a71dc654804dcaabb438" + - "69b3a72e41caa5c1cd45d80e" + - "69b3a7c1d1fbc6dda05240c5" + - "69b3a8bcd8deec15de4522b9" + - "69b3a8c45679ae4949597413" + - "69b3a8df4befb83ac3ee1656" + - "69b3a8f473b37da347a3ef0b" + - "69b3a94edc3f8c0a4b5cee7f" + - "69b3a95d002c594b0bec4478" + - "69b3a982b94de7b724dbf763" + - "69b3a99a92bbbbe19077e093" + - "69b3a9a27c5f742a1e695349" + - "69b3a9f3a34abee57efd09f1" + - "69b3a9f5139de72c2b7d484a" + - "69b3aa073f9e13cfc94f3881" + - "69b3aa1d44419a61906358c9" + - "69b3aa290b50b1a025350113" + - "69b3aa335f688271a8610e2c" + - "69b3aa48dfb2598257c5374a" + - "69b3aa51cd115d5dccbbb195" + - "69b3aa96cacb1c27687de9dd" + - "69b3ab24c3c5c8e9ace1fc12" + - "69b3ab9ace93a3f1bd6e75a7" + - "69b3ac1f1c50ce32f952f0e1" + - "69b3ac2420ee33ba712b3ad5" + - "69b3ac246e13dcd2761105aa" + - "69b3ac5814a8f34b70dd6976" + - "69b3ac956623282edcf5a426" + - "69b3add5c3521073b6e44f70" + - "69b3aed8cf0a894d8a45a338" + - "69b3af02c1685c006bc7a318" + - "69b3afaaaa410d4bfe17b346" + - "69b3b018c8794e8dc3014911" + - "69b3b02398e12c7a0721cafc" + - "69b3b03e5b48470e94ad461a" + - "69b3b0aef9f1a3291c5d6498" + - "69b3b150b221ba854f220d62" + - "69b3b449ef59307d55ca078c" + - "69b3b53ba18f33b785d3c76e" + - "69b3b58e6144f027ad9416a4" + - "69b3bbd29e5d189b4aae79fc" + - "69b3bcd3bd3343ef280e7020" + - "69b3bcf85fbb70bd57cdc739" + - "69b3c07dcc3081da47be2a37" + - "69b3c093a617e88c195369be" + - "69b3c246fef79614fbc1a346" + - "69b3c29f4698a25dfc2f4a68" + - "69b3c2fc77f9aa81573aae99" + - "69b3c455f629796beb98320f" + - "69b3c4ce7b17290cea3ad09f" + - "69b3c61c153b74b9c3da886a" + - "69b3c7532a77cf818fe8bbeb" + - "69b3c7945dbc1cde5c4e6926" + - "69b3c796880ee28a0531867c" + - "69b3d9014f4b006a4b4a2c0b" + - "69b3d9c9039109330b23ac7b" + - "69b3e2752b2d83dfbe1e382c" + - "69b3e2c8f334e0f5a453b833" + - "69b3e586a424e9e7431c241f" + - "69b3f16e138428bc6f7db2c9" + - "69b3f5c189a1a25e88a5c89f" + - "69b4015bcb724fc0989b633d" + - "69b401ac44509a97d3126396" + - "69b4044b15d40a449a6d7a1f" + - "69b40a77165ee7193a2bc1e4" + - "69b40ed3025d49131a128a4f" + - "69b415eceec4d2ca74ded48d" + - "69b41bc9a1699bb8f1b99950" + - "69b41fd602bb51b30b40a20f" + - "69b4264044509a97d31265f6" + - "69b42a0d8042e94e0c0c9b56" + - "69b42b068c925d17b9076439" + - "69b42c4648093a6af8221a1b" + - "69b42c8ce8b23514640551bd" + - "69b42f189891381d9171386b" + - "69b42f4424c3294450e1ce14" + - "69b4364d99857306aa684caf" + - "69b43cd432cefa428f10353b" + - "69b43e314b2afacec3803545" + - "69b43f14cf93b545e6088a37" + - "69b43f37bba20ceb05d7896c" + - "69b44061b716f96bbddc0607" + - "69b444eb9c2948594da14e7b" + - "69b44604a58e309b6da87273" + - "69b4466af53c7574915fe709" + - "69b447ece6ef302e7aafc31f" + - "69b44970151050b0854f8a6d" + - "69b4523417ec01169beb3a49" + - "69b454a07f1404bd9489530b" + - "69b45594cf93b545e6088afc" + - "69b455ef8c925d17b90765db" + - "69b45832a10eeb6132170926" + - "69b45980680bedd399fcb763" + - "69b45a0313e3c10f8bd28e49" + - "69b45bc9a7585840264c6367" + - "69b45f2932643e552e8851b3" + - "69b460b5ffea0a68ed9d00bd" + - "69b46639868748401c8c0cd5" + - "69b46bb1bb4aaf9805d8b7e1" + - "69b4766b7dbf454e8599221d" + - "69b478486371c861a1709fb9" + - "69b4787d5da194a21ed4091a" + - "69b4792e1b02c068cd2ad021" + - "69b479b07ab5cb8d5871894f" + - "69b4852bc9baeb064babf553" + - "69b485bb90295b261f21be5c" + - "69b485e6e9de50b7d0f6eea7" + - "69b485efbbfb18a9cc6ca6dd" + - "69b487348faeb1c44a3eba9b" + - "69b489ba626f7e5120173a41" + - "69b489d7edb7298bd2610e27" + - "69b48b30bceee9fb4ea30c9d" + - "69b48b4be267367ac3f251e0" + - "69b48c4688ffae5aad001c5d" + - "69b48cc9c09466979bc2934a" + - "69b48d5077f4dc69622ce819" + - "69b48f0c868748401c8c0e03" + - "69b48fbe8a313c3e6edaaf98" + - "69b4910bdef7cddf961d3ee6" + - "69b49187bd343104cdb6dc32" + - "69b4931db534b03c31ab8d32" + - "69b49338c9baeb064babf61d" + - "69b49360f19ed0467e1715be" + - "69b493699a0f35ae6bd284d2" + - "69b493722463c9433cf2f321" + - "69b493767211457ab4709ea0" + - "69b49385ba67778804afbd2a" + - "69b493b2bc1711c9daf2def8" + - "69b493b4a83a7b34affb4525" + - "69b493cdb3512f3d7ede1a29" + - "69b493d6c59480ef5c9304e6" + - "69b494303a7542779b574fbb" + - "69b494316ac521ce17db28d1" + - "69b4943223b8cfce653bc8b4" + - "69b494356b5ffecb54b2f39c" + - "69b494394009cbbd3f6c2061" + - "69b4943be9f454e6cf0d8cd4" + - "69b4944ddb4d0ad1b2f7ffa9" + - "69b4944f969b18fb10eb6f2e" + - "69b494533f21e8be2d394124" + - "69b4945786e9e7f64031e598" + - "69b4948bd62080354d81cbc8" + - "69b49493206ac02073e9ef7c" + - "69b4949b7acc5d78c8107ce9" + - "69b4949f8f419c69e88fad39" + - "69b494a27255ccd7c00a8227" + - "69b494b15d6b48c59a932005" + - "69b494c598e0f543f389f755" + - "69b494e5e4c79c23abfd3c07" + - "69b494ecc4d1de4eaf0d2a32" + - "69b494ed40272c2234cd2c39" + - "69b494f0c7834bedf1971fbd" + - "69b494f6c060218e378a2711" + - "69b494f9690ce5eb0225ca23" + - "69b49505bd343104cdb6dc61" + - "69b49506c9baeb064babf647" + - "69b49507f7197703d5fbd13e" + - "69b4951662d831a509f7c9bb" + - "69b4951b32580579bce8e535" + - "69b4951f09ee8318e185a1b7" + - "69b495278ae23973608291c6" + - "69b4954e252dbaf13297f278" + - "69b49556403cae238d5bebe7" + - "69b4955f6ac521ce17db28e3" + - "69b4956b1a5f291b1f8c9c66" + - "69b49576827f8909e8826704" + - "69b4957be97dc27c11192a5f" + - "69b4957e749e4dd4042ce646" + - "69b4957e913307677d706c55" + - "69b4958919cc1d0464e1557c" + - "69b4958ac3b5f2075c096374" + - "69b4958ccf5923d784013172" + - "69b4958ea506f1c6355d0473" + - "69b4958fd393f848d2e87498" + - "69b495a28574928ac9a8de30" + - "69b495c430fc0dc05fe97577" + - "69b495cc578cf4bf69341a64" + - "69b495d79bc2ce0e6641c7bd" + - "69b495e0389b1557c35068f0" + - "69b495e24db6a7b70da80591" + - "69b495e416633723bd0b2804" + - "69b495e5bba5adb44c3fc02f" + - "69b495ef2cea4472917452df" + - "69b495f19b6d7e4a378e4e79" + - "69b4960b86c1932505c2ed06" + - "69b4960e2e074a12e9ef472d" + - "69b49612cc419d655dacddb5" + - "69b496163dd44d5c018db6a0" + - "69b49616ba3960335c1914c9" + - "69b49618c060218e378a2721" + - "69b4961c6ff5d3c8f6ed4d2c" + - "69b4962bbbfb18a9cc6ca830" + - "69b4962d999caf0c82b98a54" + - "69b49631edfc29666460354d" + - "69b496384f139f508842886a" + - "69b49647938a328b92f667fa" + - "69b49654b7150166535cbe27" + - "69b49663955261bbbd495b09" + - "69b4966917e3a9ccf08b22ff" + - "69b4966a3271aaa5f51ec369" + - "69b4966f8bf4ffe2743b001b" + - "69b4967cf0f8027b635c759e" + - "69b4967fb344fe224b57c728" + - "69b49683d20be0630bf035c6" + - "69b4969b018cc5373adce91c" + - "69b496add51e511a082342b0" + - "69b496c49f8e3016601d9bd1" + - "69b496ca699f525ceebed0fb" + - "69b496e848a0f860e5987fa2" + - "69b496f484179943aec08118" + - "69b496f66b262d6ad3db4444" + - "69b497212e3fada4b3275364" + - "69b49723389b1557c350690a" + - "69b49790b33d5dd889db0a74" + - "69b4979ad19e8afa1f0c4a5a" + - "69b4988c99f7c80d1fa6fdb5" + - "69b49a1772d022b8d44697d0" + - "69b49a4ac2a8c7a3ea5b0b1d" + - "69b49ad2690f91fba74c0dd3" + - "69b49b23d61065cc8a576ec7" + - "69b49bca0e55f47ba88c2de3" + - "69b49bcceb89585320c07443" + - "69b49bccf83a2bb9efdf32cd" + - "69b49c045c9a724fbfbc1c2e" + - "69b49c0dccbdf1e86f172679" + - "69b49c149162b9797784c3e8" + - "69b49c2754e8c8b2eab17562" + - "69b49c2b35a8ae849c5f0686" + - "69b49c7c45a65ae24232ad36" + - "69b49c8dbb2597d3ebd2365d" + - "69b49ca2a4783e6eb9a67267" + - "69b49cae72c0bd4b17ff1e9d" + - "69b49cb8ac8bbc3ef42bb63d" + - "69b49cc68a4c8aae9bc764d2" + - "69b49ccc51d1604985dd2e55" + - "69b49cdb5d4941345a07519c" + - "69b49ce329b01cb732732326" + - "69b49ce54e58e40d052102e2" + - "69b49d0af83a2bb9efdf32e6" + - "69b49d2854de5fc3aac8bb57" + - "69b49d2e8781fdde8bbaa9ca" + - "69b49d44f0f63eb5f69b83e2" + - "69b49d458bf4ffe2743b0086" + - "69b49d4bf3e8b9dc1caf9609" + - "69b49d536325cbcd4232ac72" + - "69b49d5e04a889619e4bdd1f" + - "69b49d64c3f9ed0a1125a7d4" + - "69b49d66138428bc6f7db85d" + - "69b49d6f2363db0e6da1e95e" + - "69b49d7723ef41cf5c6064c8" + - "69b49d8025116d80e5210dc3" + - "69b49d8320714062f4b0b556" + - "69b49d84763fc35351a59864" + - "69b49da1ab37121e9e93bfb8" + - "69b49da3a2e0f23556e6a5d0" + - "69b49dac965ee7245be20a2d" + - "69b49dc32af7475cd15b25ae" + - "69b49e10e664f7d0eeafe654" + - "69b49f3e758deefb5ec00fbc" + - "69b49f9ae27d2b8e1ab6afe9" + - "69b49fb9b534b03c31ab8e44" + - "69b49fd0948f2c5bdccbb944" + - "69b49ff1c8f3797a5f1f667f" + - "69b4a0036462aeb78fd68806" + - "69b4a004671bfb1f5254dfa4" + - "69b4a00a9b2ae0fb15ca35c8" + - "69b4a01d5fbe65b3def7e203" + - "69b4a02416ad6cbe9700c3ce" + - "69b4a04161c85d2534e8f5e1" + - "69b4a05066535f9abb184a9b" + - "69b4a058ca9fe7e1f7add277" + - "69b4a05cd51e511a08234320" + - "69b4a0647c8402f77c793cd6" + - "69b4a07ce8279eda4266ffda" + - "69b4a0c49bc2ce0e6641c85a" + - "69b4a0c4f327a9232dd1b2c6" + - "69b4a0c8c042529718bf8d22" + - "69b4a0cdbc18efd78ffecdd8" + - "69b4a0e2d49f14c6bda8be80" + - "69b4a0f504a889619e4bdd61" + - "69b4a0faab575f0d17062793" + - "69b4a0fe2503a081f0cd3379" + - "69b4a10d02f955d3294489a7" + - "69b4a11e3326dd7f8185c820" + - "69b4a11e85a10d02e9a92ae6" + - "69b4a1302c8d1a50237642cb" + - "69b4a133acea937039fea39f" + - "69b4a13bd61065cc8a576f49" + - "69b4a13de8279eda4266ffea" + - "69b4a13ee27d2b8e1ab6b001" + - "69b4a158e21a4f95872c4b77" + - "69b4a15a955261bbbd495bae" + - "69b4a15f2bcba0117f9f0261" + - "69b4a168487fa1638cd2839f" + - "69b4a16d5804688e3b558e9f" + - "69b4a175017628a28b4a463a" + - "69b4a18df7197703d5fbd1d3" + - "69b4a1912a514e82dec0d5a2" + - "69b4a191711ca2642f803a6d" + - "69b4a19d41432f34e358862c" + - "69b4a1a59338b259898a1f9a" + - "69b4a1c8c3b5f2075c09647e" + - "69b4a1d33f40381200be1517" + - "69b4a1d43825ffd4f25135c3" + - "69b4a1d69bc2ce0e6641c864" + - "69b4a1eab19c9858cf19d262" + - "69b4a1ed8574928ac9a8ded9" + - "69b4a1f1969b18fb10eb6fa0" + - "69b4a21bb3c5edb04398e813" + - "69b4a22f8ca5423dab74c32e" + - "69b4a23f3282e47c873535c3" + - "69b4a240c445c3accc571a94" + - "69b4a2459b2ae0fb15ca35f4" + - "69b4a24cd393f848d2e87544" + - "69b4a25f4593e8cbb107a026" + - "69b4a26ed58343cbfa9adf94" + - "69b4a27da42b8c74f2348ddc" + - "69b4a2d08505130f8b0e6d34" + - "69b4a2f623d820df9de53247" + - "69b4a2f841dea3220fc9b2a8" + - "69b4a304b3b74805360993bd" + - "69b4a316857c353e098d81ff" + - "69b4a359486a23130a8381b5" + - "69b4a359506144ea4e04aea9" + - "69b4a36de27d2b8e1ab6b01f" + - "69b4a3a5578cf4bf69341b9d" + - "69b4a407e7694858dd9574a0" + - "69b4a4990c2d4a47811c682b" + - "69b4a4a7d292e88e08f7f9d3" + - "69b4a4f93593ece06d37408f" + - "69b4a512c3f9ed0a1125a883" + - "69b4a59c479220aa237ec4a7" + - "69b4a5cf2530ca2e1b04b318" + - "69b4a6dae27d2b8e1ab6b059" + - "69b4a6ee8faeb1c44a3ebbe4" + - "69b4a70a214150b696113763" + - "69b4a72bf9274b0a852545fc" + - "69b4a73784ac85f46f77563b" + - "69b4a8bb00ce0b618c19ca89" + - "69b4a8f92f4f45a63823ff9d" + - "69b4a923831d1b8c4d630744" + - "69b4a931f83765e0c840cd5b" + - "69b4a941fd8b3ba702656425" + - "69b4a984a03afaf54dbcba93" + - "69b4a9aa8f4228ae45998443" + - "69b4a9c12fca00d0b66add26" + - "69b4a9ea388ccecd105c5b92" + - "69b4a9ea75c68475da2a245a" + - "69b4aa195c9a724fbfbc1c9b" + - "69b4aa30f629796beb983f64" + - "69b4aa4f01650da7131e2079" + - "69b4aacaa68ffd2fbd98a8bf" + - "69b4aacab3a779ee9e027c26" + - "69b4aafe99e58efaf1232aaa" + - "69b4ab24e2a75fae64879a4a" + - "69b4ab4e4e5ac9677e3a2445" + - "69b4ab675fadf00f5a3f1a01" + - "69b4ab728415a3e84cfa8a4b" + - "69b4ab8bfe501abf8bdb4057" + - "69b4ab8f41432f34e35886d0" + - "69b4abad96ac497b0841a96c" + - "69b4abcc138b755ea6edc792" + - "69b4abe29705148536964d66" + - "69b4abe82176c8614d96df49" + - "69b4abf51ce35576286e7879" + - "69b4ac00230639e424fd1c38" + - "69b4ac02d61065cc8a576fda" + - "69b4ac077ae6fe5f6c5a3b3e" + - "69b4ac3c89c93d14505f40f1" + - "69b4ac425d5efd59f80d46fb" + - "69b4ac4ae328b8f0f2640e49" + - "69b4ac6691b851da8626ac9c" + - "69b4ac76f4719f087672d087" + - "69b4ac99e9de50b7d0f6efd0" + - "69b4aca3fc5f3e6e2e68bfa5" + - "69b4aca9857c353e098d82c2" + - "69b4acc8cf026361c7c654a7" + - "69b4acd49f8e3016601d9cc4" + - "69b4acd545a65ae24232ae3d" + - "69b4acd9b8745dc1f264d4d8" + - "69b4acdb29e4b6d94a2f9f13" + - "69b4ace2b2b6f95180884a03" + - "69b4aceb9944ebcdcaaf6976" + - "69b4acf03271aaa5f51ec465" + - "69b4ad061460ea5ea1cc3fb3" + - "69b4ad06ea0f93b1d3bf9f9f" + - "69b4ad15d62080354d81cd65" + - "69b4ad1772d022b8d44698d6" + - "69b4ad18d131be6d95f77101" + - "69b4ad190dd40e72ddf4ed3b" + - "69b4ad1b455ff604c8193839" + - "69b4ad2684ac85f46f7756a6" + - "69b4ad299162b9797784c552" + - "69b4ad441279cedd289bc30d" + - "69b4ad4aba6c35ba5800efff" + - "69b4ad4e8a313c3e6edab137" + - "69b4ad5be5be71666d494451" + - "69b4ad5bf05424f9a3a2e147" + - "69b4ad7542b321609bc429d6" + - "69b4ad763ea50ce70f2d3f4f" + - "69b4ad9938586bc8a41a5dd6" + - "69b4ada75de136447199ebea" + - "69b4adae6a8d780c96378bc6" + - "69b4adb172afcd77da62feff" + - "69b4adeea03afaf54dbcbaf6" + - "69b4adff37cb8c569b487789" + - "69b4ae067eb149d6a67098a1" + - "69b4ae083282e47c87353676" + - "69b4ae08a13067e31cf8af9c" + - "69b4ae0d90be9c6b2a9d9950" + - "69b4ae0f0cd58dc901b88fac" + - "69b4ae2838720da38b3e50f3" + - "69b4ae2c51fb8a949dbb7f1d" + - "69b4ae2fb7150166535cbf62" + - "69b4ae58fe3e51b5a878e3ac" + - "69b4ae5feb89585320c0754b" + - "69b4ae603e8130e694cfc0e1" + - "69b4ae6735cef09fa62e2a6a" + - "69b4ae6d1279cedd289bc325" + - "69b4ae790c01a69644a21e0b" + - "69b4ae7a6f7d4ee860c5e955" + - "69b4ae7f3758288247afca56" + - "69b4ae89a3213aef75c96019" + - "69b4ae8a2a05fa27685f089a" + - "69b4ae8a9c40481a89baf7e3" + - "69b4ae8e214150b696113811" + - "69b4ae8f656475939c9bda52" + - "69b4ae93711ca2642f803b81" + - "69b4ae9486c1932505c2ede3" + - "69b4ae99c7b864206a9a6f56" + - "69b4ae9a84179943aec08336" + - "69b4ae9c90be9c6b2a9d9959" + - "69b4ae9cd292e88e08f7fb40" + - "69b4aeaafd18e37694a9434a" + - "69b4aeb1230db4d63a859b28" + - "69b4aebadb402eaf0fc5a867" + - "69b4aebc3638c8d1a91a7cf2" + - "69b4aed2adbe5164b37f9b05" + - "69b4aedb557734daed3d7d3b" + - "69b4aee026ecc3ea6063cd5f" + - "69b4aee770be8493faa0a243" + - "69b4aeef51fb8a949dbb7f32" + - "69b4aef2698982b08fa376f1" + - "69b4aef8a506f1c6355d0622" + - "69b4aef8cf6532061788557b" + - "69b4af10c7b864206a9a6f62" + - "69b4af2b749e4dd4042ce71b" + - "69b4af3274ac87c55e05a96c" + - "69b4af3bdef7cddf961d4093" + - "69b4af5024715e745fc74f13" + - "69b4af54b3c5edb04398e8a8" + - "69b4af74aa4961e178763687" + - "69b4af86bc1711c9daf2e08f" + - "69b4afb6a6f94f418b4d56b5" + - "69b4afcf98e0f543f389fa03" + - "69b4aff516633723bd0b2952" + - "69b4aff5aa5c1d155aec3139" + - "69b4aff6bc56af1cde4fbc19" + - "69b4b0061d8d8d4ad85ca0d9" + - "69b4b01a1d57a6d358e479ee" + - "69b4b01bbb2330fd69279bc2" + - "69b4b020da81c162c79f382c" + - "69b4b04466de3768833bddc3" + - "69b4b04a8ae239736082938d" + - "69b4b05ab84aad9597edc45b" + - "69b4b06006857a01e6a7391e" + - "69b4b08a4e1b15ae18625ce8" + - "69b4b0927d741d012b44f8c5" + - "69b4b0a8a14f5f8fb96d7240" + - "69b4b0b98f9cabb9ff5c26e6" + - "69b4b0c0319e61c1e723a8d3" + - "69b4b0c5472f820b6876e600" + - "69b4b0d5fd24ff75565f6ace" + - "69b4b0f6c445c3accc571b52" + - "69b4b104c420444b72115951" + - "69b4b10e90be9c6b2a9d998d" + - "69b4b139050842bc0e0e9eb1" + - "69b4b13a9b2ecfe4a474bb74" + - "69b4b142fd24ff75565f6ad8" + - "69b4b148bc18efd78ffeceef" + - "69b4b176fdce20eb7f3241a6" + - "69b4b17b25fd3eda026f20cb" + - "69b4b185b8745dc1f264d51c" + - "69b4b18a9a77d2d77af87c35" + - "69b4b1a7c042529718bf8e51" + - "69b4b1ae4db6a7b70da806fb" + - "69b4b1c414a4d59b08be0031" + - "69b4b1c96ff5d3c8f6ed4ef5" + - "69b4b1d21e1240fd6bcfe9d2" + - "69b4b1e4138b755ea6edc7ef" + - "69b4b1f9aea2c4fce0c3ea16" + - "69b4b20dbc56af1cde4fbc2b" + - "69b4b2111813007fddbcbc75" + - "69b4b2112fca00d0b66addb4" + - "69b4b22d0dd5025bc8b2a2bd" + - "69b4b23219dfa98a17d19580" + - "69b4b23f83db69b280d297a4" + - "69b4b2530cd58dc901b8900b" + - "69b4b25889a1a25e88a5ceb5" + - "69b4b25eee737c3c16391c7e" + - "69b4b2696462aeb78fd68935" + - "69b4b293fd8b3ba70265647d" + - "69b4b2aa472f820b6876e62c" + - "69b4b2bbb30844b29a8327ad" + - "69b4b2c2ca9fe7e1f7add3d0" + - "69b4b2e1a3dc51170950653a" + - "69b4b2f0146676112722811a" + - "69b4b2ff9b2185b61d193cae" + - "69b4b308cfe53ea1af65fb2a" + - "69b4b31d38fbfdf7b20f11ae" + - "69b4b3223ae16cd77636e1a0" + - "69b4b329cfe53ea1af65fb2c" + - "69b4b335d781f5e3eea1e391" + - "69b4b339c3b5f2075c0965fc" + - "69b4b33c73ca03704a5029ad" + - "69b4b33d0dd5025bc8b2a2cf" + - "69b4b33e04a889619e4bdeb5" + - "69b4b34286c1932505c2ee22" + - "69b4b345394e490d4592db76" + - "69b4b3497c84b98b9831e4cc" + - "69b4b34d66bbb8ba908da0bb" + - "69b4b360e985757268c73578" + - "69b4b3677c687258b13ab8dd" + - "69b4b3696868654392eb8511" + - "69b4b375936e114f0730b2fa" + - "69b4b37ba6f94f418b4d5702" + - "69b4b37e56077778e4b51509" + - "69b4b38129e555555c6488b1" + - "69b4b3965bb6442829fb418a" + - "69b4b3966b9429ed3f34daf9" + - "69b4b398c9baeb064babf840" + - "69b4b39e6260efb9125c1bb2" + - "69b4b3a3e248420994a82e58" + - "69b4b3a8e5be71666d494489" + - "69b4b3b35fbb70bd57cdcd9d" + - "69b4b3b554775729157ed36b" + - "69b4b3baba9b1b2ff17fd083" + - "69b4b3c3cde059c2090359b6" + - "69b4b3d182377b6dbb0155c0" + - "69b4b3df680bedd399fcba28" + - "69b4b3e426611d16819bb590" + - "69b4b3fb5bceecbc49ebecff" + - "69b4b40362d831a509f7cae2" + - "69b4b40402219e539949137a" + - "69b4b4166641c34ca8db1550" + - "69b4b43686c1932505c2ee2f" + - "69b4b4371ce35576286e78ed" + - "69b4b45a460cd47c547c8743" + - "69b4b45cba9b1b2ff17fd090" + - "69b4b479f29408487e3faa0c" + - "69b4b492b3b74805360994e8" + - "69b4b495da06e853603f5a29" + - "69b4b4bf65bc47b3b31eac67" + - "69b4b4c460299e2df8f99cb1" + - "69b4b4c9c04c33841ec2198d" + - "69b4b4cd394e490d4592db86" + - "69b4b4e00d48c3f6ede856e4" + - "69b4b50766de3768833bde16" + - "69b4b50da83a7b34affb485d" + - "69b4b51e65bc47b3b31eac6f" + - "69b4b524389b1557c3506b5d" + - "69b4b5246a8d780c96378c24" + - "69b4b52d10fc8b8f166659b3" + - "69b4b52d1a5f291b1f8c9d65" + - "69b4b540ddcebec248dbb1a7" + - "69b4b5442bcba0117f9f03f5" + - "69b4b5542b9d7dec7f449339" + - "69b4b57a53c39a7b7561c5f4" + - "69b4b57a9f66a2746698f81d" + - "69b4b58be024ce298b096dea" + - "69b4b59d017628a28b4a47e2" + - "69b4b5a0b504f66b2436a748" + - "69b4b5ade62059bdb6db7e33" + - "69b4b5b389a1a25e88a5cee0" + - "69b4b5bb7dbf454e8599247c" + - "69b4b5c572c0bd4b17ff1fcd" + - "69b4b5ca8bf4ffe2743b0216" + - "69b4b5ea48a0f860e5988099" + - "69b4b5f841c2dcb85b2e2d7e" + - "69b4b6066868654392eb8550" + - "69b4b60dd0103fb8e4fb4bfa" + - "69b4b6158fea875fbba16b66" + - "69b4b61c9944ebcdcaaf6a4f" + - "69b4b6245dc2bb787e7c0bba" + - "69b4b63996658ef2ddbd75f0" + - "69b4b649ca9fe7e1f7add405" + - "69b4b66e699f525ceebed34b" + - "69b4b676138428bc6f7db902" + - "69b4b69d31cc112bbe13df5b" + - "69b4b6af8f419c69e88faf99" + - "69b4b6c34dd9410ce4c3c52b" + - "69b4b6f31d57a6d358e47a88" + - "69b4b70283793f4df344532c" + - "69b4b72cad9ebb37e8776a42" + - "69b4b7611667be999ff443d5" + - "69b4b768153b74b9c3da90a4" + - "69b4b77055b20ccbe5761da6" + - "69b4b7a4dabcf1361edc3bb3" + - "69b4b7b9b564ee80404a3967" + - "69b4b7ba0cd58dc901b8907a" + - "69b4b80783793f4df344533c" + - "69b4b858ac58e3fe36d85b86" + - "69b4b88a515d4294ccc79262" + - "69b4b88cd95f0cb5ffa2ba06" + - "69b4b8989b3c11bddef4aee1" + - "69b4b8ac146676112722819b" + - "69b4b8c153168422516b2547" + - "69b4b8c62fb9ffa8355ede46" + - "69b4b8de23ef41cf5c6065ea" + - "69b4b8ea10fc8b8f166659f6" + - "69b4b92e4f712d2ccd55b59f" + - "69b4b937eec4d2ca74ded91a" + - "69b4b94cb70fb5b833849b2c" + - "69b4b963f29408487e3faad3" + - "69b4b966611e81861dad4859" + - "69b4b96c6bb53d53abcf63b2" + - "69b4b96de7e57ee3fa956e08" + - "69b4b9842e0e1cd2280ed2ef" + - "69b4b99fb70fb5b833849b30" + - "69b4b9ef1667be999ff443f4" + - "69b4b9fbedba17448480b788" + - "69b4ba049b6d7e4a378e50ca" + - "69b4ba17654dd32a33b3b96f" + - "69b4ba2df6333e2fa8a2c5f1" + - "69b4ba2e4f6aaaa9cd1bdb96" + - "69b4ba4c9338b259898a2198" + - "69b4ba7c02f50e8a066ea46a" + - "69b4baa0bf0ab5f0fcd34abd" + - "69b4baacbd343104cdb6df8b" + - "69b4bab69adf23002bb00987" + - "69b4bab88a1f883a68a40359" + - "69b4badb99f7c80d1fa6ff05" + - "69b4bae8af3cc1eb766df7f1" + - "69b4baf36dab335bc90fe4d9" + - "69b4bb0714d880d0812831fd" + - "69b4bb08ccdcb17165f9b7a5" + - "69b4bb0cba3960335c191768" + - "69b4bb21c042529718bf8f37" + - "69b4bb2ec5fa25330433f460" + - "69b4bb49daf18bc1bcd1522f" + - "69b4bb4d7927e426d4bdd88b" + - "69b4bb67fb51c10a1b8f800f" + - "69b4bb6bc4c5496f9afd47b9" + - "69b4bbbee89d73f4ec383efa" + - "69b4bbbfaea2c4fce0c3ea6f" + - "69b4bbc2f76dedb2b40db3bf" + - "69b4bbc673ca03704a502a56" + - "69b4bbc7e7bad87e1f52729b" + - "69b4bbcc72d022b8d44699c1" + - "69b4bbcf3282e47c8735374c" + - "69b4bbd260b9979b587b0f04" + - "69b4bbd3f0c88edf01a6323e" + - "69b4bbd62917034a1c53ad1e" + - "69b4bbd6bb2330fd69279c78" + - "69b4bbd7a3c0ce02c5a55a46" + - "69b4bbe58993ecc25d2718ef" + - "69b4bbee2990d2ad579a2dd0" + - "69b4bbf0b7c8fd45d1462fd8" + - "69b4bbf2970f3d65c454a178" + - "69b4bc0719ebdadc8764a6df" + - "69b4bc1fd158ec4a43dcf819" + - "69b4bc390c02cfba7bd9d004" + - "69b4bc3c3aabc6ea561ff8c3" + - "69b4bc4160172888461a82c1" + - "69b4bc438c1c37de293e74ff" + - "69b4bc4e8dfc6b6c4f7d2e87" + - "69b4bc4e9a0f35ae6bd28772" + - "69b4bc7e8993ecc25d2718f7" + - "69b4bcac2e40b83afba4a1cd" + - "69b4bcb2d480c010eccb2139" + - "69b4bcb374320179c62b46d8" + - "69b4bcc3b4fe7bfb5208d32d" + - "69b4bcdf9a0f35ae6bd28778" + - "69b4bd464c035879577bfda0" + - "69b4bd4a9cc44bc26f6a1bd1" + - "69b4bd54948f2c5bdccbbb7f" + - "69b4bd6b19cc1d0464e15886" + - "69b4bd7c230db4d63a859bb3" + - "69b4bdabb49ed0379faeb92c" + - "69b4be3e05a258967f4ab34e" + - "69b4be5aee25266722a758f0" + - "69b4bee26e1591125967e8aa" + - "69b4bef4e8279eda426702a1" + - "69b4bf04f109813904030e47" + - "69b4bf431813007fddbcbd76" + - "69b4bf6913a3425f7ea68252" + - "69b4bfcc74ac87c55e05aa9e" + - "69b4bfed8415a3e84cfa8c53" + - "69b4c03a16d52405b4a2c260" + - "69b4c0e0d6e27d500501f805" + - "69b4c245c3b5f2075c096744" + - "69b4c271bb3fe1af3a8734c5" + - "69b4c278bbfb18a9cc6cab5c" + - "69b4c28644de405906e81098" + - "69b4c2c3edb7298bd26111ab" + - "69b4c2cb0ea33062b1c27745" + - "69b4c2d5f268396b4bc3445c" + - "69b4c2e717e3a9ccf08b25bf" + - "69b4c30660b9979b587b0f83" + - "69b4c30c3725ced819eb2447" + - "69b4c3722e0701a08ff7f411" + - "69b4c3a61431a94b0cf6deea" + - "69b4c3be728fdfa816cbf707" + - "69b4c4e0c042529718bf8fdf" + - "69b4c5312990d2ad579a2e04" + - "69b4c55cf0f63eb5f69b87dc" + - "69b4c57ad902e000df8b9623" + - "69b4c58e7d741d012b44f983" + - "69b4c6774e1b15ae18625e56" + - "69b4c8918579b939089f489d" + - "69b4c938a0c2188e1d8f48b9" + - "69b4ca4dab4b984623699693" + - "69b4caccef58e30ec42f89ff" + - "69b4cb0ee8211b84899cf012" + - "69b4cb230b666ac17c526536" + - "69b4cb88ea0f93b1d3bfa257" + - "69b4cc34c0682c81b27b2344" + - "69b4cca4f01bb0c8128714e0" + - "69b4cd5225116d80e5211173" + - "69b4cd7d7440b684d16cfbf2" + - "69b4cd812c83b2c34dfe5562" + - "69b4ce9e72c0bd4b17ff20c3" + - "69b4ceb01e1240fd6bcfebf7" + - "69b4cfc470be8493faa0a497" + - "69b4cffabba5adb44c3fc596" + - "69b4d022fb51c10a1b8f80ce" + - "69b4d09e06e01260c1ac9415" + - "69b4d1089cc44bc26f6a1d1b" + - "69b4d18465bc47b3b31eadec" + - "69b4d2a863a359f1f0847a79" + - "69b4d2aac792eda587dd9e15" + - "69b4d30d1ca41b2857617021" + - "69b4d355bd343104cdb6e189" + - "69b4d39f66de3768833be0b0" + - "69b4d4145964d373bc945c08" + - "69b4d464dc0c3de83c8eb5b8" + - "69b4d4cd6ff5d3c8f6ed5122" + - "69b4d4d0485757cc553ece92" + - "69b4d53d4184c502d4080baf" + - "69b4d54dce95af8cea9a3635" + - "69b4d5820c2d4a47811c6b6a" + - "69b4d5a720714062f4b0b8f1" + - "69b4d5aee41b2381f0f55eba" + - "69b4d5b5f4719f087672d2d2" + - "69b4d5ebfe501abf8bdb41ae" + - "69b4d5f136c56722006292a7" + - "69b4d5f569d10ac4ac82aa94" + - "69b4d63af52b66f1a3088a30" + - "69b4d6709e5d189b4aae82a0" + - "69b4d739980f80025fc5df33" + - "69b4d7515e012a4936eb6e21" + - "69b4d753fd24ff75565f6d52" + - "69b4d7b258a78ed4c167d383" + - "69b4d7d19c40481a89bafa49" + - "69b4d7f829b01cb73273264c" + - "69b4d83a16ad6cbe9700c505" + - "69b4d885398ef280fe770c48" + - "69b4d88f24e02ea6eaf54cdc" + - "69b4d89229b01cb732732652" + - "69b4d8ab7bc9e8b3eac05cd3" + - "69b4d8f31c33cb7c3a42b266" + - "69b4d916a098085fcfe41bcf" + - "69b4d9193db5e101c3b651f5" + - "69b4d95183a8c659e0d05e2a" + - "69b4d98e44de405906e81239" + - "69b4d9a0def7cddf961d42ee" + - "69b4d9ab52783435e5edeb67" + - "69b4d9c33ae16cd77636e427" + - "69b4d9db52e5c59d347ad84d" + - "69b4d9e0fa004527e64458dc" + - "69b4da40f327a9232dd1b57f" + - "69b4da57060a6199caac1226" + - "69b4da642c8d1a50237646e0" + - "69b4daa4c3a8dbc608ff8e9c" + - "69b4daba98e0f543f389fe21" + - "69b4dadb2994f223142b83b4" + - "69b4dae9590494387f207eb7" + - "69b4daf4de3dc7ed68294f4a" + - "69b4db1e4f6aaaa9cd1bdca9" + - "69b4db3289deb7a84c2d572b" + - "69b4db5a7acc5d78c81082cc" + - "69b4db949605a9793b6b909a" + - "69b4dba6c02d505205328d7e" + - "69b4dc0a699f525ceebed5bd" + - "69b4dc0abfce163f13246b84" + - "69b4dcbba8ea953037b9fca4" + - "69b4dd709b3c11bddef4b1da" + - "69b4dd70b8816c7c6f3c92b3" + - "69b4dd9e711ca2642f803f48" + - "69b4ddef84179943aec0869e" + - "69b4de167bc9e8b3eac05d15" + - "69b4de5428737cee656ce8f0" + - "69b4de7b00ce0b618c19cd17" + - "69b4df2ce21a4f95872c4ea4" + - "69b4df473825ffd4f2513925" + - "69b4dfab182fde7bbe2252e9" + - "69b4e0174b8693d11c3b9a7d" + - "69b4e0fa8c0e050035cc5f6f" + - "69b4e1268d89f01acb82fc00" + - "69b4e15bad9ebb37e8776c02" + - "69b4e16ea0835130796403f5" + - "69b4e1914b8693d11c3b9a84" + - "69b4e207d781f5e3eea1e62c" + - "69b4e386e6e3d585253ca7ba" + - "69b4e3eeca3a06a9f7b0d857" + - "69b4e41202f955d329448e3e" + - "69b4e4342fca00d0b66ae076" + - "69b4e4973b674a7b354e6b5c" + - "69b4e498d20be0630bf03840" + - "69b4e4b85f4789f4f7292b02" + - "69b4e4d0c04c33841ec21af7" + - "69b4e594c5fa25330433f627" + - "69b4e6b3cfe023119bddd05a" + - "69b4e73abf16bae4bd7ced2b" + - "69b4e79789c93d14505f44b1" + - "69b4e7ffbb2597d3ebd23831" + - "69b4e80117ec01169beb4249" + - "69b4e86b1425999d967049a1" + - "69b4e89a50ea53288aeda28d" + - "69b4e8a63ea50ce70f2d421f" + - "69b4e8f68c9ba98495e6e2e6" + - "69b4ea4063b26e9384dd3b54" + - "69b4eb0996658ef2ddbd7850" + - "69b4ec068559e49d8af4e57d" + - "69b4ed3de025fc2bec3b0599" + - "69b4ed3e018cc5373adcedf5" + - "69b4ed4fd480c010eccb23cb" + - "69b4ed7ff1166dec72cfcfe2" + - "69b4ed964ff0a19d651b18b8" + - "69b4edf6130037c5b4ba5e55" + - "69b4edf6270553bea390ae42" + - "69b4ee1f1e9a5acce0dd0e31" + - "69b4eea32917034a1c53ae67" + - "69b4eeb28f4228ae459986a3" + - "69b4eedd698982b08fa37b54" + - "69b4eee58559e49d8af4e588" + - "69b4eee6270553bea390ae46" + - "69b4efa5f1c3826f1de62361" + - "69b4efb590295b261f21c1a2" + - "69b4f01e880ee28a05318e23" + - "69b4f125c69c46f0a2021584" + - "69b4f14040d43b4590183193" + - "69b4f196a97937a3e2fb3a62" + - "69b4f1d38a313c3e6edab466" + - "69b4f258def7cddf961d43ed" + - "69b4f2c45c17e11ba0da268b" + - "69b4f37c388ccecd105c5fe1" + - "69b4f472590494387f207f90" + - "69b4f5687ae6fe5f6c5a3cfe" + - "69b4f57e56742533296f951c" + - "69b4f5e6eb2eeb714c5ec126" + - "69b4f5fd3b674a7b354e6baa" + - "69b4f61f880ee28a05318e41" + - "69b4f684394e490d4592dd05" + - "69b4f6b15fa639c4e24e76d8" + - "69b4f6ddf268396b4bc345ae" + - "69b4f716a55386e5af547817" + - "69b4f729f327a9232dd1b6ac" + - "69b4f7c3b2b6f95180884e2b" + - "69b4f7cbb90f582cfc1bfbbd" + - "69b4f7cc6b262d6ad3db473e" + - "69b4f7d4a13ae541201c75b7" + - "69b4f7f141a4401e4c3f549d" + - "69b4f854060a6199caac129e" + - "69b4f8773593ece06d374571" + - "69b4f88a8c9ba98495e6e335" + - "69b4f895817936a8e39f4346" + - "69b4f89752b439fcc613b58a" + - "69b4f8cdeaf185a9827a39df" + - "69b4f92141a4401e4c3f54b4" + - "69b4f9505d6b48c59a932758" + - "69b4f9d0a424e9e7431c2bfe" + - "69b4fa9d4e58e40d05210719" + - "69b4fad6c314b4adc2d818ee" + - "69b4fad99cd671b2e770b251" + - "69b4faf54963ef0ac28d59ff" + - "69b4fbd57c84b98b9831e66b" + - "69b4fc18de3dc7ed68295098" + - "69b4fc3bc02d505205328ea7" + - "69b4fcd31bbb62ca41aef27b" + - "69b4fcf06b262d6ad3db4756" + - "69b4fd248faeb1c44a3ebe5c" + - "69b4fddace95af8cea9a37be" + - "69b4fe9309a939795bdb7d06" + - "69b4fe98b8745dc1f264d729" + - "69b4ff2c3f21e8be2d3943de" + - "69b4ff5a515d4294ccc7942e" + - "69b4ff78f14d5ff1cd3b027a" + - "69b4ffdb17b1b3bbefd3fe67" + - "69b5007641a4401e4c3f54ee" + - "69b500ff24715e745fc75271" + - "69b50118de3dc7ed682950b2" + - "69b5015a58f5afe39b1de817" + - "69b5022fba5771467aa5c78a" + - "69b502636b262d6ad3db476a" + - "69b5026485847273dc74d475" + - "69b502985bcb582c92e292f6" + - "69b502cc85e36f35da0bb203" + - "69b50375179ecfd3b33aa17c" + - "69b5037a96658ef2ddbd78fd" + - "69b5037d8faeb1c44a3ebe77" + - "69b503892f2fa993841ebcb2" + - "69b5038a857c353e098d87fa" + - "69b503cd8559e49d8af4e5cf" + - "69b503d50c01a69644a221cc" + - "69b504267565ed366761942e" + - "69b5044591d6491aa901a591" + - "69b504a547e0e2e54a0ed6f4" + - "69b505516371c861a170a388" + - "69b505f5a55386e5af547875" + - "69b505fbda5e069c8ac8f54f" + - "69b50605fdefe9ea069f4a5a" + - "69b50616969b18fb10eb7265" + - "69b506f7c02d505205328ef5" + - "69b5070046452e48802c2711" + - "69b5070dd5bde809a4eec104" + - "69b507829f8e3016601da0e5" + - "69b508e544509a97d31272c8" + - "69b508eb980f80025fc5e15a" + - "69b50946a13067e31cf8b31a" + - "69b509c61b02c068cd2ad238" + - "69b50a4bb30844b29a832d57" + - "69b50a6d16f9cc224f6655ad" + - "69b50ac2bb3fe1af3a87371a" + - "69b50b1f9cd671b2e770b292" + - "69b50bce70be8493faa0a6c4" + - "69b50bfd8f74bcceeacdadf8" + - "69b50c7a460cd47c547c8aaf" + - "69b50cdc980f80025fc5e17d" + - "69b50d3dea50b19996d7e87e" + - "69b50d454a206179cbd004bf" + - "69b50dc79a082448a776e9af" + - "69b50eedc7b864206a9a746b" + - "69b50f235d5efd59f80d4ab0" + - "69b50fdd519b84b08472f7b7" + - "69b51087cf7413b53d1929cf" + - "69b510cd35a8ae849c5f0b4a" + - "69b51120c4d1de4eaf0d2db4" + - "69b513525fa639c4e24e77ca" + - "69b513e7f83765e0c840d2ff" + - "69b5143b7c962028610d50aa" + - "69b51556ea0478dac47671a8" + - "69b5155ec5efd21514309118" + - "69b5158aa2bb2f3fc33fd4be" + - "69b5166ae38bde62855cf92d" + - "69b51690a55386e5af547926" + - "69b516b2557b5c79a6d7a80b" + - "69b5179eea50b19996d7e8f0" + - "69b51838da5e069c8ac8f5a3" + - "69b51860e9de50b7d0f6f218" + - "69b518dc6641c34ca8db19a4" + - "69b51924e8b235146405581f" + - "69b5194beaf185a9827a3aef" + - "69b51952f3077050a50791d5" + - "69b51985d5bde809a4eec16c" + - "69b51a64654dd32a33b3bbc6" + - "69b51aa19f66a2746698fb7a" + - "69b51ab9827f8909e8826d69" + - "69b51ada8f419c69e88fb37a" + - "69b51b2f6c13d04069dccd8e" + - "69b51bad52b439fcc613b688" + - "69b51bc2eb6fc92ec9812bbc" + - "69b51c11ab575f0d17062a4a" + - "69b51c162e9983f31a0b5b54" + - "69b51c18f1d531c517f55b3f" + - "69b51c303e8130e694cfc360" + - "69b51c50d90644203dac6b31" + - "69b51c836c13d04069dccd9c" + - "69b51ca13ac25998bd7be354" + - "69b51cc0980f80025fc5e20b" + - "69b51d00aa4f5d834935856c" + - "69b51d0c7565ed366761953a" + - "69b51d16ba5771467aa5c816" + - "69b51d4b3758288247afcce8" + - "69b51d4e65776c1b020fdcc7" + - "69b51dcd23ef41cf5c606a15" + - "69b51df9ad8cb8d985c9a61a" + - "69b51dfa8579b939089f4a6b" + - "69b51e0400d0e98b2aab033f" + - "69b51e1ee8b235146405583d" + - "69b51f6441dea3220fc9b7a0" + - "69b51f7cd19166322bb8979a" + - "69b51f80242a79d532ab7627" + - "69b52023c792eda587dd9f86" + - "69b52053f14d5ff1cd3b03f9" + - "69b520cb7c84b98b9831e712" + - "69b520e01bbb62ca41aef3d6" + - "69b520eb2994f223142b8674" + - "69b5213ed197af22134e7125" + - "69b5216acf7413b53d192a6c" + - "69b521b4d9ae83850b4beb91" + - "69b521eaf52b66f1a3088baa" + - "69b52315c420444b72115be2" + - "69b5232ccf653206178859a7" + - "69b52348817936a8e39f442e" + - "69b5238a038aad71cceef737" + - "69b523ab50e213ed974b1a42" + - "69b523f8287b1e524e4af0c1" + - "69b5244ab1370a7b23cf17ac" + - "69b524fe665ebcb51d1e0ff4" + - "69b5253fc02d505205329025" + - "69b5256f179ecfd3b33aa218" + - "69b525a86a8d780c96378e70" + - "69b525b926871cfad9ed17c7" + - "69b525dba7585840264c68fe" + - "69b52643a55386e5af5479c1" + - "69b5280016f9cc224f6656d1" + - "69b5280d53e765a22cddd811" + - "69b528d4618df5d7d36be836" + - "69b528dc79a6a023093deb6d" + - "69b528e688f4b4a6e37d9a55" + - "69b528f3590494387f20819f" + - "69b5295d48a0f860e5988311" + - "69b5296d4e73acdd60a50e75" + - "69b529a9164aaac0a2a6f085" + - "69b529bd60172888461a8501" + - "69b529c21cdb3cb1bcb168e5" + - "69b529c6f83765e0c840d3d8" + - "69b52a5b980f80025fc5e298" + - "69b52a7e870fd7d45c1f3411" + - "69b52acec995b04b13458402" + - "69b52b03d9bd4ed2c8cda811" + - "69b52bda025d49131a129110" + - "69b52bdad2ab5ebd52f6e590" + - "69b52bdd16f9cc224f6656fb" + - "69b52c05806149738b4e5aed" + - "69b52c2cab11feb553052f01" + - "69b52c501bbb62ca41aef466" + - "69b52c92e025fc2bec3b0808" + - "69b52c98a424e9e7431c2cf3" + - "69b52ca99f8e3016601da28a" + - "69b52d35e024ce298b097257" + - "69b52d7eb7823508aa10c04a" + - "69b52d959a0f35ae6bd28c5e" + - "69b52e01daf18bc1bcd15684" + - "69b52e0470be8493faa0a838" + - "69b52e81bc415a1f0b593af0" + - "69b52eb44f4b006a4b4a3d73" + - "69b52ee1a1b0f0badb996814" + - "69b52f13096a5342c3dc7944" + - "69b52f1e9c7daebcace61f7a" + - "69b52f48d158ec4a43dcfa7a" + - "69b52f63de62960a07de8989" + - "69b52f6f79a6a023093deb92" + - "69b52f7ed3ef4c69cdc63108" + - "69b52fc00403c14bba1f60e7" + - "69b52fd735a8ae849c5f0bd1" + - "69b53039ee2ba8e55c296a41" + - "69b530e4d309464df7115a78" + - "69b5311d83db69b280d29ad3" + - "69b53157a10eeb613217100f" + - "69b5315a8061ae794120aea6" + - "69b53177d3d86d024f123ca7" + - "69b5318016f9cc224f665747" + - "69b531d599790d8da7837d20" + - "69b531f4dd4d1616e5387598" + - "69b531f9c69c46f0a2021711" + - "69b532392814c4be9c919b4c" + - "69b5325cc4d1de4eaf0d2e7e" + - "69b5327db70fb5b833849d9c" + - "69b5331b230db4d63a859e04" + - "69b53343967351f71e33ec56" + - "69b53398060a6199caac13dd" + - "69b533d417e3a9ccf08b298d" + - "69b5340d83db69b280d29ae7" + - "69b5348088f4b4a6e37d9ada" + - "69b534c2d19166322bb8981d" + - "69b534f0d9bd4ed2c8cda85a" + - "69b5351bad9ebb37e8776ef7" + - "69b5353548a0f860e5988356" + - "69b53538bd3343ef280e83d1" + - "69b5357941dea3220fc9b84e" + - "69b535c52843c2f113eb1593" + - "69b535ca85847273dc74d6a0" + - "69b535d3bba20ceb05d79014" + - "69b535d4969b18fb10eb7373" + - "69b535fbeaf185a9827a3c3d" + - "69b536063dd17eb53d7ddbcc" + - "69b53616f70d1949244a6350" + - "69b5361f9a082448a776ea8c" + - "69b536265c9a724fbfbc1fa1" + - "69b53649c7e02bc8d412b02c" + - "69b5366f153b74b9c3da93b9" + - "69b536aa4009cbbd3f6c26f5" + - "69b536c66b9429ed3f34df87" + - "69b536d9e1d17e38b58e07a1" + - "69b53760da5e069c8ac8f65e" + - "69b53813c6100a03dc6cf802" + - "69b53917c995b04b1345846f" + - "69b539268559e49d8af4e71d" + - "69b539308a4c8aae9bc76b46" + - "69b539331a32c53ddbd16dba" + - "69b53967423da797309e328e" + - "69b53a70cecf38bf0cf6dfc1" + - "69b53af9f429dfde897d9448" + - "69b53afdf629796beb9844cd" + - "69b53b135bbfec0f5b716689" + - "69b53b15967351f71e33ec8e" + - "69b53b2f7565ed36676196b3" + - "69b53b6b3326dd7f8185cbdd" + - "69b53b72969b18fb10eb739c" + - "69b53be72cea447291745729" + - "69b53c05b4e21c90299dd92a" + - "69b53c366b4827eb68cf134a" + - "69b53c73bb3647e88d660351" + - "69b53cdfc6100a03dc6cf85c" + - "69b53cff230639e424fd1fa5" + - "69b53d085a2141ade9a3b0b5" + - "69b53d306dab335bc90fe78f" + - "69b53d74a865ba37a7ea9a6d" + - "69b53ea6423da797309e32b8" + - "69b53ebd7e64d72afd3bf44c" + - "69b53efb5fbb70bd57cdd074" + - "69b53f1da9f25eecea6b58e9" + - "69b53f3809a939795bdb7e7f" + - "69b53f5fcd67f708f6e64f2d" + - "69b53f7d6c13d04069dccf86" + - "69b53ff0c5efd21514309218" + - "69b54004bb4aaf9805d8bf48" + - "69b5409bf52b66f1a3088c79" + - "69b540a069d10ac4ac82ae6d" + - "69b540a3618df5d7d36be984" + - "69b540e883db69b280d29b40" + - "69b540f19a082448a776eadf" + - "69b540ffb8816c7c6f3c954a" + - "69b5412d99790d8da7837e11" + - "69b5418c62ac9501398f3ecf" + - "69b541d5a2bb2f3fc33fd64a" + - "69b541e44e58e40d05210921" + - "69b54225cecf38bf0cf6e002" + - "69b54236b8816c7c6f3c9556" + - "69b54243cb97446b898f425f" + - "69b542739cd671b2e770b420" + - "69b542828c61b5e3f72f9c93" + - "69b5429bd9ae83850b4bec52" + - "69b5431bf4110359f68f7bc8" + - "69b543252cea447291745765" + - "69b5433f8061ae794120af3a" + - "69b543df9a082448a776eaf9" + - "69b5441ab1370a7b23cf1886" + - "69b544354e6b0afb54824163" + - "69b544ed5a10c4a2d8fa0371" + - "69b545290660b04f45ae5a20" + - "69b5456ac0248818ec109dfe" + - "69b545ecbb4aaf9805d8bf7a" + - "69b545ed52b439fcc613b7de" + - "69b54612cf026361c7c65864" + - "69b5465f618df5d7d36be9ea" + - "69b546c462fdc2f575078c37" + - "69b546d8654dd32a33b3bcde" + - "69b5472abd3343ef280e84a1" + - "69b5477003e5822473e2d7d5" + - "69b5485b164aaac0a2a6f1a3" + - "69b54893f55514359677d994" + - "69b548cab84aad9597edc795" + - "69b548f7124502aa36325381" + - "69b5490116f9cc224f6658d3" + - "69b54912aea2c4fce0c3ed7b" + - "69b549e46ac521ce17db2daf" + - "69b54a5eea0478dac47672ff" + - "69b54a5fea50b19996d7eb80" + - "69b54a8349b632b7cc703071" + - "69b54ae8df63a3e00aac2817" + - "69b54b3be38bde62855cfa7a" + - "69b54b621bbb62ca41aef696" + - "69b54b8c35a8ae849c5f0c75" + - "69b54b9f99790d8da7837f69" + - "69b54bb4bb3647e88d6603cd" + - "69b54be75656a2300b64f08d" + - "69b54c075bbfec0f5b71670b" + - "69b54c0ee8b235146405597e" + - "69b54c16b1370a7b23cf18cc" + - "69b54d202e9983f31a0b5ce8" + - "69b54da5f3077050a50792d6" + - "69b54db88579b939089f4bc4" + - "69b54dc0e9de50b7d0f6f388" + - "69b54dde00b486e6cf3754cd" + - "69b54de1db5484827307bedd" + - "69b54e2abba20ceb05d790d1" + - "69b54e7300b486e6cf3754d1" + - "69b54ea8b8816c7c6f3c95d0" + - "69b54ecca321d4f716b75222" + - "69b54ecfd0103fb8e4fb4f78" + - "69b54ee27e46044c0318cf0d" + - "69b54f09c8f3797a5f1f6d1e" + - "69b54f139f66a2746698fd45" + - "69b54f192aa2afaba2bd5ad5" + - "69b55009ced344cb0df766e2" + - "69b5501d32cefa428f103d69" + - "69b5505f9b2ae0fb15ca3ce1" + - "69b550fc5bcb582c92e294cc" + - "69b55106bd3343ef280e850e" + - "69b5524660172888461a8612" + - "69b5529c16f9cc224f66597b" + - "69b552b061c85d2534e8fa79" + - "69b552cb37cb8c569b487bb3" + - "69b552ce86e9e7f64031e9fc" + - "69b552e032643e552e8857aa" + - "69b553aa5656a2300b64f0cf" + - "69b553c19a3e364072f0061d" + - "69b553d9817936a8e39f457d" + - "69b55407124502aa363253cc" + - "69b5542c95b405a6f71061e3" + - "69b5543632643e552e8857b4" + - "69b554703d0a537fcffa3d2e" + - "69b5548f5b4047fbbfc4d86d" + - "69b5549b89a1a25e88a5d2cd" + - "69b554d4ae7c8a432518f9d6" + - "69b5553a88f4b4a6e37d9ca0" + - "69b55541d5bde809a4eec3be" + - "69b556043fbc85f1fd117512" + - "69b556df9d6df5f24703fee1" + - "69b556ee3fbc85f1fd117520" + - "69b55740cee7dbe8803e4112" + - "69b5574da9f25eecea6b59a0" + - "69b55753151c8cc559daadac" + - "69b5577a8505130f8b0e73f8" + - "69b5577b09a939795bdb7f4b" + - "69b557bd78673e5de4228922" + - "69b558341ca41b285761733d" + - "69b558b866687323993e8f13" + - "69b558e5f55514359677daa6" + - "69b558fadaf18bc1bcd15934" + - "69b55912c995b04b134585a2" + - "69b5592049b632b7cc703170" + - "69b55965d3d86d024f123e33" + - "69b559df6325cbcd4232b177" + - "69b55a37210430db9a877f35" + - "69b55a60124502aa3632540f" + - "69b55ab117b1b3bbefd40092" + - "69b55b6a6371c861a170a698" + - "69b55b7c626f7e5120173e16" + - "69b55b7cfe3e51b5a878e7f5" + - "69b55ba61bbb62ca41aef7e1" + - "69b55bcf9d6df5f24703ff0d" + - "69b55be443041b5f0fab06d5" + - "69b55c79287b1e524e4af24b" + - "69b55ca3d422ae5db7e71de3" + - "69b55d7717ec01169beb45ed" + - "69b55d8f48093a6af822225f" + - "69b55e544963ef0ac28d5eb8" + - "69b55e7b557b5c79a6d7ab15" + - "69b55eb27dbf454e859928ff" + - "69b55f2d8f4228ae45998972" + - "69b55f4678673e5de4228964" + - "69b55f6872c7691edd73ba76" + - "69b55f8fb416a351f4c4292c" + - "69b55ffa85847273dc74d95a" + - "69b55ffd99790d8da78380be" + - "69b560496b9429ed3f34e172" + - "69b56054ab11feb553053232" + - "69b56059e025fc2bec3b0ba3" + - "69b5613c210430db9a877f6b" + - "69b5615899790d8da78380d3" + - "69b5620a3246ab5abc88b128" + - "69b5626ba424e9e7431c2ef1" + - "69b5627d4f4b006a4b4a4032" + - "69b5628616f9cc224f665a88" + - "69b562a27ab5cb8d58718fa4" + - "69b562e3e267367ac3f25614" + - "69b5630d46452e48802c2970" + - "69b5632ee025fc2bec3b0bcb" + - "69b56343c7e02bc8d412b17b" + - "69b5639170c4570003841cc2" + - "69b563a0d3d86d024f123e9a" + - "69b563ccba6c35ba5800f46f" + - "69b563f6eac64bf65ebcd96d" + - "69b56404519b84b08472fa23" + - "69b56449758deefb5ec01510" + - "69b5650affea0a68ed9d068f" + - "69b565223ac25998bd7be599" + - "69b56526eb6fc92ec9812e2c" + - "69b56537df63a3e00aac2a9c" + - "69b565658cd7957ebe4793f8" + - "69b5658390295b261f21c4af" + - "69b5659de6ecb1abc97eda04" + - "69b565c5cdf4d3962d10b17e" + - "69b565d47f0b85103365790e" + - "69b565f43ac25998bd7be5a1" + - "69b56617f629796beb9846a0" + - "69b5663535a8ae849c5f0d1c" + - "69b5664eea50b19996d7ee5a" + - "69b5667e153b74b9c3da952b" + - "69b566c761c85d2534e8fb22" + - "69b567192f2fa993841ebf4a" + - "69b5671e3fbc85f1fd1175ff" + - "69b56730519b84b08472fa41" + - "69b56776e8b2351464055ae9" + - "69b567875301ac33f39cba19" + - "69b567b941dea3220fc9bbb4" + - "69b567c188f4b4a6e37d9e83" + - "69b5685bc4d1de4eaf0d3035" + - "69b568ceea50b19996d7ee80" + - "69b56918c155de971bf1bf40" + - "69b56951c375af84c0577131" + - "69b569d217e3a9ccf08b2d10" + - "69b56a095301ac33f39cba2d" + - "69b56ace1bbb62ca41aef8d5" + - "69b56b80cb724fc0989b6c99" + - "69b56bbb0c3122b34d6e8d03" + - "69b56c0072c7691edd73bafa" + - "69b56c8d8061ae794120b090" + - "69b56cc119ae335e59eb6882" + - "69b56d56cee7c43263a475cb" + - "69b56d93460cf9a11ded49c7" + - "69b56db8a2bb2f3fc33fd84d" + - "69b56dc28c9ba98495e6e6f0" + - "69b56de9310c703fd24107b0" + - "69b56eca6371c861a170a739" + - "69b56f0e0403c14bba1f6486" + - "69b56f8a0afa99ac56e21eac" + - "69b56fbdc995b04b13458672" + - "69b5702c827f8909e8827181" + - "69b570a76abfa602a5495724" + - "69b570bab6fd585b6fa4d20a" + - "69b570ca6462aeb78fd690dc" + - "69b5710a9a3e364072f00839" + - "69b57123cd67f708f6e650d2" + - "69b5712d2a68f4514f1d66b4" + - "69b57132df63a3e00aac2b6f" + - "69b571479cd671b2e770b75a" + - "69b5715b3fbc85f1fd117696" + - "69b5716f1666a9ba0092ea72" + - "69b5719ed7dc8311ecd92ea1" + - "69b5728cab11feb553053378" + - "69b5728e039109330b23be65" + - "69b572aeafe4badf8009e7f7" + - "69b572cd250b26ab0bf1c6f7" + - "69b572f72a77cf818fe8ce57" + - "69b572fddd3c32e8914daa26" + - "69b57326c0248818ec109ffb" + - "69b5734fd673f8a5bdac0fa0" + - "69b573673242599fceea3076" + - "69b573676c44b53e645bf9e6" + - "69b57372e0be395bb48be461" + - "69b5737c210430db9a877fe6" + - "69b5737c76fea7858f6baa96" + - "69b573a8eaa2f14deed26a5c" + - "69b573a9da06e853603f5ed7" + - "69b57416135e07fa47414d12" + - "69b57465ea9519dc9ae2c1c0" + - "69b574ca9c27f8b9d1b5a69c" + - "69b5755eea0478dac47674fe" + - "69b57590fc7b36de840aeec5" + - "69b575cf50e213ed974b1d05" + - "69b575efe817baf56a934842" + - "69b5766185847273dc74dad4" + - "69b576a4571b01a65fb6acbb" + - "69b576ab460cf9a11ded4a05" + - "69b576ccc04c33841ec21f28" + - "69b577a5098f0b40252c8a44" + - "69b577db2aa2afaba2bd5c9b" + - "69b57889d7a5dd9bef79705a" + - "69b5788ac4c5496f9afd4bf3" + - "69b578d1d43107729f703725" + - "69b578def182bd6c8b6798b6" + - "69b57921baf2957a279cfb3a" + - "69b5798a8cd7957ebe4794a9" + - "69b579cb250b26ab0bf1c73a" + - "69b57bad1666a9ba0092eac6" + - "69b57c3919fd2396f8df5f10" + - "69b57c5644509a97d31278bf" + - "69b57ca826871cfad9ed1ad9" + - "69b57ce29a3e364072f008bf" + - "69b57cec98fbfeb784ab5391" + - "69b57d434e58e40d05210ba2" + - "69b57d4ce1d17e38b58e0a94" + - "69b57da51a315c79b1b51751" + - "69b57da7a6262a0eb585abf9" + - "69b57daabe537f8fabc6bb10" + - "69b57dbb17e3a9ccf08b2dfb" + - "69b57dbb5bcb582c92e2962c" + - "69b57dfb43041b5f0fab07ba" + - "69b57dfc096a5342c3dc7bf9" + - "69b57e17da5e069c8ac8f940" + - "69b57e4322471f0894f4cd80" + - "69b57e4499790d8da78382bf" + - "69b57e536c44b53e645bfa50" + - "69b57e550dd5025bc8b2aa56" + - "69b57fc68a313c3e6edab93a" + - "69b57fcff3077050a5079410" + - "69b57fffde5165de31d69066" + - "69b5804552b439fcc613ba4a" + - "69b580831666a9ba0092eae8" + - "69b5808f60172888461a8775" + - "69b580a0793c697e1df2dde6" + - "69b580d4b7bba0cd0bb6583e" + - "69b580fd6ac521ce17db2f6a" + - "69b581042c1553f6d605e2ae" + - "69b581baaf0cd74d1a856ed3" + - "69b581c9e5aa43f09bb6c04d" + - "69b581cf287b1e524e4af361" + - "69b58227b48aa02ad85674f7" + - "69b5829ce664f7d0eeafec04" + - "69b582e3e664f7d0eeafec06" + - "69b5833800b486e6cf3756dd" + - "69b58345d422ae5db7e71eed" + - "69b583627d58b76fe17ef197" + - "69b583e0c8f3797a5f1f710c" + - "69b583fd99790d8da783830a" + - "69b584168c3e4c65ea3397e5" + - "69b58428db5484827307c1e2" + - "69b5842ab8745dc1f264dbe1" + - "69b58436f55514359677dd35" + - "69b584550c3122b34d6e8e18" + - "69b584932cea4472917459bd" + - "69b584b679a6a023093deded" + - "69b584b99cd671b2e770b814" + - "69b584c9e7f9e51220ce641d" + - "69b584dee817baf56a9348a5" + - "69b5853d1b02c068cd2ad68d" + - "69b5855ea2bb2f3fc33fd946" + - "69b58624185e4fd6a7484217" + - "69b587b9279e4b209d33526f" + - "69b587f29ee0b627f1fcfd71" + - "69b587fd394e490d4592e0a6" + - "69b588066a8d780c96379158" + - "69b5883486477fa003cbf583" + - "69b588bfe2d927d8bec08303" + - "69b588fa6a8d780c9637915e" + - "69b588faa25cc84d5c20796e" + - "69b5892b0273df551d93c380" + - "69b5895cb8816c7c6f3c98f4" + - "69b5896f17ec01169beb47ac" + - "69b589771e8afb702ffa97e8" + - "69b589b57ae6fe5f6c5a40e1" + - "69b589ba4963ef0ac28d60f4" + - "69b589c145a65ae24232b770" + - "69b589d0230639e424fd22e4" + - "69b589d2025d49131a129439" + - "69b58a43164aaac0a2a6f441" + - "69b58b35c6100a03dc6cfd49" + - "69b58b4ce6e3d585253cac01" + - "69b58b6a7d741d012b44fe46" + - "69b58b6e86477fa003cbf599" + - "69b58bfa3e4093f9b85c7b9d" + - "69b58c67f3077050a507945a" + - "69b58cf3793c697e1df2de34" + - "69b58d2c8b61ffe2823cbd75" + - "69b58d376b4827eb68cf18f6" + - "69b58d4b8781fdde8bbab280" + - "69b58d541a59a1957c5886de" + - "69b58db0ee2ba8e55c296df0" + - "69b58db5f334e0f5a453c372" + - "69b58ef3b1370a7b23cf1a9b" + - "69b58efdb49ed0379faebe4d" + - "69b58f91f3f27a8574b3c1be" + - "69b59007f70c81ef9281d8a3" + - "69b59011f30a828e1ab2f8b6" + - "69b590456b5ffecb54b2fee0" + - "69b59048c6100a03dc6cfd8f" + - "69b590573e7a3f6666bbdaa7" + - "69b5912b1b02c068cd2ad6dd" + - "69b5917b164aaac0a2a6f477" + - "69b5919ec4c5496f9afd4c8b" + - "69b591b2151050b0854f9405" + - "69b5920285847273dc74dc3c" + - "69b5921d2e9983f31a0b5fd7" + - "69b683edaf933f49003771a3" + - "69b683f7390a74c74ee0ac3f" + - "69b684d888c67edd437b4d5a" + - "69b684dccd225473d6956765" + - "69b685cf6fd488b5dec5a3ee" + - "69b685d35f1d2dd786b35e16" + - "69b68621e2cc1b1f8eff101d" + - "69b68675b498972c0d5a1530" + - "69b686af7c71d96b40a21c86" + - "69b686d7539983322c06ef0d" + - "69b686e46f44e1498bdf94f7" + - "69b687001f451cbdb2c2284a" + - "69b68700346a130f980d34af" + - "69b6879d4da1fe51454a07a2" + - "69b687b9d88f1dc9ad09524a" + - "69b687bd6b67e6fb30537306" + - "69b6882bc90ad0858c925341" + - "69b688c49e1d6ce3c40e022e" + - "69b6896ca3d57c71da0e64d8" + - "69b68983d3c9fd08cb4d2158" + - "69b689d7d2cce6bcc4f795eb" + - "69b68a4b42cfcbbbdd230c8c" + - "69b68a7a8663d16d970b6f2d" + - "69b68a9696e604419ceb31ab" + - "69b68adabb7ac89497113506" + - "69b68adc67c4fd5814ce969a" + - "69b68adc6b67e6fb30537332" + - "69b68b831cd819c3c4e43a10" + - "69b68b99aa96b63b2bddb9e9" + - "69b68bc4792dc732a00dd65e" + - "69b68c22f1f7e85beb93d10a" + - "69b68c9f1ccf2f13f516980c" + - "69b68cc189a1870e36b45fe5" + - "69b68d0d33702a0c61de0b6e" + - "69b68ddebd2c26f28ba3eed3" + - "69b68de21aac15c691cf2928" + - "69b68e27e226858f5f982155" + - "69b68e6897994f5347e03472" + - "69b68ec19fd8af79803a6fbd" + - "69b68f2e7c2e4161f098b7b3" + - "69b68f3c686ec2d47340ae9a" + - "69b68f477c2e4161f098b7b5" + - "69b6903d0248d6f194ddb858" + - "69b6906e33702a0c61de0ba8" + - "69b690c3fb47cf874d7289e8" + - "69b690d16fd488b5dec5a4a2" + - "69b6912056bb38e9cf352577" + - "69b691ffb52dd3401b147729" + - "69b6921a62f6b9a70b374ed5" + - "69b692404da1fe51454a0806" + - "69b6925219af0367402916ed" + - "69b6927261533cd7ffd1c9a1" + - "69b69293e4547fd5917ae31e" + - "69b692b5beaf9aa0229cd479" + - "69b692e8792dc732a00dd6ba" + - "69b69337dfde53b5ccd0ecd4" + - "69b69383b498972c0d5a15ae" + - "69b693f5aa2c6d4ae2f424c1" + - "69b6944e90ff96d53410f3c3" + - "69b6945bbda1ce3e8a6e982a" + - "69b694724d8c5f73d519add8" + - "69b694d6eae2ffd4a87212ae" + - "69b695de346a130f980d356b" + - "69b695e961533cd7ffd1c9c5" + - "69b6960ee9a7a1d4e945b5d4" + - "69b6961cf69d549b769dda1d" + - "69b6966c6af251b30b21ff17" + - "69b696852e915c8fc8b73004" + - "69b696fdfe41e6a150240379" + - "69b697a76af251b30b21ff27" + - "69b697d34d8e0ab375337d14" + - "69b6986c0869e51f5860efe9" + - "69b69878fa71ea99e3eb88f1" + - "69b698e1a967df5ed1873702" + - "69b6990902c15af4600ee850" + - "69b6990be226858f5f9821ac" + - "69b699af15c15130fd7d5860" + - "69b699f8cb8a8960356ecebc" + - "69b69a03f2710b05b13a6e00" + - "69b69a46dc720e64f95ebac1" + - "69b69a52f2710b05b13a6e04" + - "69b69b0aa967df5ed187371e" + - "69b69b5ce9530a92e03235b8" + - "69b69bed4f8988c123b3377f" + - "69b69c38bd3b9c472951f27b" + - "69b69c4f4f8988c123b33783" + - "69b69c6cf8a03559f744e8ca" + - "69b69c815f1d2dd786b35f38" + - "69b69d0ec90ad0858c925431" + - "69b69d283615389300a1da02" + - "69b69d57346a130f980d35bf" + - "69b69d75c764c30eb512053a" + - "69b69e1748733b2339f0ce1f" + - "69b69ee734f0d55328d45f4b" + - "69b69f37ace85918a0bf18e8" + - "69b69fdef85248dd0877e4ae" + - "69b69fdf1ccf2f13f51698e8" + - "69b6a0163166ab375e94886d" + - "69b6a01d1ccf2f13f51698ec" + - "69b6a31bd0ec56f9fd959b04" + - "69b6ab1db8c28a7597e26488" + - "69b6b455082198b0a5621de1" + - "69b6b6e48d468328c426b598" + - "69b6bacbab43941e8ba87df2" + - "69b6bb92e30380ba2b80f358" + - "69b6bd0ce1b8ba5278f6c1c9" + - "69b6bd33a7b2977fd6c83b44" + - "69b6be34c3671d8c2b7156f6" + - "69b6c0084f8988c123b33886" + - "69b6c417d540ab4fa8bbea33" + - "69b6c475385e5c8f59ae4f60" + - "69b6c5e6d0b8510d6762be8a" + - "69b6c6f8dcf7032ff6b94cdd" + - "69b6c82951caaccc5aa8cb8c" + - "69b6c853630f2c9338a93758" + - "69b6cf5898e255bbb5dbc2bc" + - "69b6d0dbff263e4b37caf58d" + - "69b6d3ca686dad175412e066" + - "69b6d4db1dd637de2ce9a14b" + - "69b6d5169a639e1ee4efead7" + - "69b6d54816dd4d92eabd0521" + - "69b6d7b732d3c1e1ff597211" + - "69b6df49c5cad2c02e779dfd" + - "69b843ccd45ee7cfab24a03d" + - "69b84450be6261da4d580a44" + - "69b845a1cdbc7aede6943858" + - "69b848ddfb1a61ea23a43d60" + - "69b84b6ddf08149cccad9d89" + - "69b84efeefd45f0907856e21" + - "69b84f1de1a23f007e48100c" + - "69b8513122a04882a7cc8397" + - "69b85276a964762cef314a08" + - "69b8535084cde035d66462f2" + - "69b8536bc5d0f737ab74b690" + - "69b8547bfb1a61ea23a43da0" + - "69b8561ca7abcaf7ad7cf556" + - "69b8565da964762cef314a32" + - "69b857d48464a9f85cba6c88" + - "69b858bed498a456d3063f76" + - "69b859c42f17571d7e345dc7" + - "69b859e8044006a7f056ed0e" + - "69b85a3dc5ab7070deca48c5" + - "69b85bf0044006a7f056ed27" + - "69b85c3cc33b3ff3d53759ab" + - "69b85d46391d0dae38fd04bb" + - "69b85e90e8b3dd6c427b2aab" + - "69b868d4fba39064bbb002f4" + - "69b868e20cacc0e708815be3" + - "69b86f25e5eb700352941bf3" + - "69b86fad2cc8314fbad4f335" + - "69b870cb235a79b35e08d730" + - "69b870cdc5a1357004e8f8fb" + - "69b8716dfb9cb514d86f2821" + - "69b871e14268a1fb1bdcab8a" + - "69b87262be2eb07817677dbe" + - "69b87269c43b2ff00fc793cc" + - "69b87309942d4a851e0510da" + - "69b8731dd363bdf0ddeabf36" + - "69b8736b1a97e051d5434c55" + - "69b874e42a89e0bab91ac961" + - "69b875672629a9d06dec3354" + - "69b8763cb5d55fd4bab77027" + - "69b876754763c6a9d3add5f3" + - "69b8768c2f17571d7e345ed4" + - "69b8787c3e8588946b17ed4d" + - "69b8793f10740bcffd3fd554" + - "69b87ad269b798b0f6217167" + - "69b87b58dc9fd6c40b127f53" + - "69b87ba3044006a7f056eea5" + - "69b87dbec580d0f1025c606a" + - "69b87e450be6b02b2a3cb3e4" + - "69b87e69724098ae21210301" + - "69b87f23c43b2ff00fc7942a" + - "69b8804319aecd94aedc36c0" + - "69b88119c580d0f1025c608d" + - "69b881a62c7dcf0ebc7bc69a" + - "69b88219a64179ee4e4cd36d" + - "69b8830473edb68839f24c8f" + - "69b8835865810dd91c215b99" + - "69b883ee690f748de4d813f4" + - "69b884e8789976e408acf7ab" + - "69b885a5e1a23f007e4810bd" + - "69b886224b754305a80bf16a" + - "69b88658fd53261c0abe1408" + - "69b889f6d92d792f69edd841" + - "69b88a1f48502cefbff211bf" + - "69b88a55446a7ee1a670d557" + - "69b88b111d448e9ea86d3633" + - "69b88bf5b1ffd7875e7ccccf" + - "69b88c5a86b9f227f633a562" + - "69b88cd3ece5a565a8f2dc65" + - "69b88ce03b814dc5ab0c6117" + - "69b88d52329155a5c8427469" + - "69b88da8329155a5c842746b" + - "69b88e6f8db6c1dd3e4f20c2" + - "69b88e708bc0613700293e2a" + - "69b88ea6489516b28573476c" + - "69b88f8eadb6e6a77d23fa26" + - "69b890f3adb6e6a77d23fa2e" + - "69b8918f0b1ec7e2efa00d8b" + - "69b892975782dbb059cc79ed" + - "69b894846cc7dde328d4f1ed" + - "69b894c3f8984dcc5fcf28c2" + - "69b8952ba3ef2775e0e988b2" + - "69b8956afb9cb514d86f28c8" + - "69b8982884e004c157308759" + - "69b89831e25d009d215f1363" + - "69b89856de59539246182670" + - "69b89a3e119f6fd572d61bb8" + - "69b89ba63b814dc5ab0c615b" + - "69b89c1b392f510eb0b1f9ba" + - "69b89c46d93479be080560e8" + - "69b89ef01cf7f6f00d43647e" + - "69b8a107e03168d093e511e9" + - "69b8a1d4cb90c111e3107ee5" + - "69b8a26df8984dcc5fcf2901" + - "69b8a2b260258e584a7dd00d" + - "69b8a4715d93f21ab156bb86" + - "69b8a5660cacc0e708815cf2" + - "69b8a5f7a0cd91c7e41039dd" + - "69b8a600de595392461826ba" + - "69b8a972708396dd182c1d67" + - "69b8a975186d8785b78d9808" + - "69b8a99cc5133d089fd7c49a" + - "69b8acd27a7328ccf6ae474e" + - "69b8ad1550e8ec10497acbcf" + - "69b8b00bb6710a234aafc5e4" + - "69b8b09bf95728a8dcc1ef83" + - "69b8b4308271c173426e33aa" + - "69b8b57d0f92f74da61ce439" + - "69b8b6e2cc1860339d1627e7" + - "69b8b74ddba5beeb5e1fb8f1" + - "69b8b7abc5133d089fd7c4db" + - "69b8b7b068c639144c96f60b" + - "69b8b82f44a2e80ec56560da" + - "69b8b898f95728a8dcc1efa7" + - "69b8b9632553f1058c598180" + - "69b8b9c41cf7f6f00d436504" + - "69b8b9dae32488e3141b7c02" + - "69b8bb85f95728a8dcc1efb8" + - "69b8bb9f86b9f227f633a63d" + - "69b8bc0bc1c83acb1695f886" + - "69b8bce6617de250ae32d58e" + - "69b8bdd692a8add8f9ed97c9" + - "69b8beac1d448e9ea86d3728" + - "69b8c123c9e8f1de0533dfeb" + - "69b8c13b87a8695556dc0e92" + - "69b8c1782030ce071fb48f4e" + - "69b8c17c6be1c6d054a42dd8" + - "69b8c1cf17abc716e5e013ff" + - "69b8c2f8b90d624e8e254402" + - "69b8c35fa0469793c24becf0" + - "69b8c3b417abc716e5e0140d" + - "69b8c3c4446a7ee1a670d675" + - "69b8c40ea8f5c8c8ef9bacae" + - "69b8c41a0e99f9971a7f7eb5" + - "69b8c57786b9f227f633a686" + - "69b8c5be52c2c675c36ca19c" + - "69b8c6dd1eba61b69ddad9a9" + - "69b8c773c9e8f1de0533e017" + - "69b8c78f119f6fd572d61cb6" + - "69b8cbfd1eba61b69ddad9c7" + - "69b8ce77f50ffaf8aea15b6c" + - "69b8d08a2030ce071fb48fc7" + - "69b8d0c0ad4a5ebbe874b1b9" + - "69b8d0f7c8858e187272b7bb" + - "69b8d0fde094c1fa8268f736" + - "69b8d1092f5695d1d2f5316d" + - "69b8d13c976bda2d64157752" + - "69b8d3bca52e1a2126f460eb" + - "69b8d41dadb6e6a77d23fbeb" + - "69b8d631d5481a3e9ebca328" + - "69b8d729e094c1fa8268f76c" + - "69b8d7b087a8695556dc0f56" + - "69b8d98906cdf6f5e21c77f6" + - "69b8da0d70551c776df516e6" + - "69b8db45004576ba1550a6eb" + - "69b8dd1cad4a5ebbe874b220" + - "69b8ddc9a8f5c8c8ef9bad62" + - "69b8de4a06964923c7124afa" + - "69b8e17808daef11eac8a132" + - "69b8e2ae2565bc2e9320aa6b" + - "69b8e319f95728a8dcc1f0fe" + - "69b8e35de094c1fa8268f7db" + - "69b8e3c587d69b550a82a905" + - "69b8e40239078534320789a6" + - "69b8e42ab94ccdc21b873015" + - "69b8e5ab43b675efc7201237" + - "69b8e6040f92f74da61ce5d0" + - "69b8e66e68c639144c96f784" + - "69b8e6d448502cefbff2145f" + - "69b8e73d87a8695556dc0fe0" + - "69b8e88ab443baf0e08286e0" + - "69b8e8db004576ba1550a763" + - "69b8e9275bd3642672ea0745" + - "69b8e9bb119f6fd572d61de4" + - "69b8eac3f787e78f3caf7e93" + - "69b8ec060e9f2891ea5e73f0" + - "69b8ec68de595392461828c4" + - "69b8ec86a8f416005104fd75" + - "69b8ef0187a8695556dc1028" + - "69b8ef0b1740c71b0874fa52" + - "69b8f01e6be1c6d054a42f98" + - "69b8f0f1c3d83cdf8646ba12" + - "69b8f17c06964923c7124bb4" + - "69b8f1a887a8695556dc1042" + - "69b8f24ae97bec8c72c5e23c" + - "69b8f39670551c776df517d3" + - "69b8f43d7a7328ccf6ae496a" + - "69b8f46fe4dd0d3cf59fff48" + - "69b8f49970551c776df517de" + - "69b8f5f646c7cc5009e4bcaf" + - "69b8f66ec147ee9b771dcaa0" + - "69b8f726a531a4418172fdda" + - "69b8f762438c3d18fbb76ffc" + - "69b8f851e094c1fa8268f8a9" + - "69b8f8c4bb8856ba14a76f7d" + - "69b8f950e97bec8c72c5e27e" + - "69b8fa22e97bec8c72c5e287" + - "69b8fa44b3cc90fa8d9acbae" + - "69b8faa446c7cc5009e4bce2" + - "69b8facda8f416005104fdfc" + - "69b8fb0ab13a790889c052d8" + - "69b8fb0c0b0ebaecc7d82599" + - "69b8fb2fc5133d089fd7c6d9" + - "69b8fc6186b9f227f633a889" + - "69b8fecde7f0e45fa5e486fe" + - "69b8ffae6cc7dde328d4f49e" + - "69b9010b73c73837a4222996" + - "69b9010d4e1bd59d5fded4f2" + - "69b9017f63e615dae6ebc9d8" + - "69b9019ce094c1fa8268f907" + - "69b901dd77b94f867cf5688a" + - "69b9027ef25718a726e343c0" + - "69b90379a89e9b47645a9db3" + - "69b903c145f54b69a3dc86cc" + - "69b90405345333ff05a8cf76" + - "69b904cf004576ba1550a86c" + - "69b9066d7ee0ff33631f083f" + - "69b90710b3cc90fa8d9acc16" + - "69b90712e25d009d215f16a9" + - "69b907c410490e396dd979ab" + - "69b907d14458d25a26dc68e1" + - "69b9086b55294f107bfbc7f6" + - "69b90a64b90d624e8e254633" + - "69b90ae4aacf0a241873b5c9" + - "69b90fd6b443baf0e0828851" + - "69b910044e1bd59d5fded55a" + - "69b911fe4a18f0c8d9bf7a31" + - "69b913492ab1219ca569be3e" + - "69b913b43907853432078b55" + - "69b914043c35a3c417d1802d" + - "69b91522345333ff05a8cfea" + - "69b91694fe5dea81e805457f" + - "69b917892ab1219ca569be5e" + - "69b917f217abc716e5e01694" + - "69b91897c3d83cdf8646bb66" + - "69b918e74e1bd59d5fded596" + - "69b9197c1740c71b0874fbaf" + - "69b91b63b3cc90fa8d9acca4" + - "69b91b7473c73837a4222a67" + - "69b91c85446a7ee1a670d8dc" + - "69b91cc62030ce071fb4928d" + - "69b91d14976bda2d64157a0a" + - "69b91df8a531a4418172ff21" + - "69b91e27a89e9b47645a9e71" + - "69b9200677b94f867cf56968" + - "69b9202186b9f227f633a9ad" + - "69b920f31264ecdcc598b290" + - "69b9217a2555407a34bad442" + - "69b921b6b3fe52e7b3246aa4" + - "69b922d6ee7ae1002b6a72ae" + - "69b922f0984ca3f76894ce74" + - "69b923513907853432078bd7" + - "69b923d31743775241d36070" + - "69b9240fd5481a3e9ebca5b8" + - "69b92452c147ee9b771dcbf8" + - "69b92490e0d38179fe8302a6" + - "69b924df9bcf11a7dc66d67f" + - "69b925704458d25a26dc69b5" + - "69b9258bb1ffd7875e7cd17b" + - "69b925a4cf3be3afac7d751e" + - "69b9263f22b409ba79d215e7" + - "69b9266c6be1c6d054a43173" + - "69b927aa6a84c22b595578c8" + - "69b928b17ee0ff33631f0910" + - "69b929045088ac62355d6c47" + - "69b929b963e615dae6ebcb11" + - "69b92a527ee0ff33631f0920" + - "69b92a5a2030ce071fb49307" + - "69b92b54bb8856ba14a7713d" + - "69b92daab94ccdc21b873237" + - "69b92e26976bda2d64157aac" + - "69b92eadfb73b7a0dcf4ca08" + - "69b92f277ee0ff33631f094a" + - "69b9301717abc716e5e01750" + - "69b9305027838502f33a93a0" + - "69b930a3b1ffd7875e7cd1e3" + - "69b9310c3907853432078c63" + - "69b932266be7e1c4139f8195" + - "69b932ef2e71fa6c931cbf3a" + - "69b9358687a8695556dc1249" + - "69b9358fcf3be3afac7d75a2" + - "69b9376f9bcf11a7dc66d717" + - "69b93810b3fe52e7b3246b74" + - "69b9387bde59539246182b6f" + - "69b938a1de59539246182b71" + - "69b938b1cf3be3afac7d75bc" + - "69b9391ca89e9b47645a9f5d" + - "69b93b90507132953a1af05d" + - "69b93ba25088ac62355d6cfa" + - "69b93c698973f21812765cc1" + - "69b93caf2030ce071fb493bf" + - "69b93cf197d25b237a7cd96c" + - "69b93d3f9c31f4f8be85d99d" + - "69b93de510490e396dd97b6f" + - "69b93e06507132953a1af077" + - "69b93f68605829a110d6cac6" + - "69b93fb32ab1219ca569bfc2" + - "69b9401097d25b237a7cd983" + - "69b9407f507132953a1af091" + - "69b9427317abc716e5e017e8" + - "69b942d5d3f6bebb1b51ffb1" + - "69b9438aadb6e6a77d240027" + - "69b94403bb8856ba14a77249" + - "69b9448987a8695556dc12e3" + - "69b944f6bb8856ba14a77253" + - "69b945dca7178c0e45b56dc3" + - "69b94661c5133d089fd7c922" + - "69b946b0d98341b3bb349a15" + - "69b946ec1743775241d361c2" + - "69b9484c2f5695d1d2f5353f" + - "69b9494ba0f9b58bf114a63e" + - "69b949cd3d7564e60197049e" + - "69b94a17984ca3f76894cfe0" + - "69b94a19b3fe52e7b3246c38" + - "69b94acdfe5dea81e8054745" + - "69b94c30bd0a1226e8720c66" + - "69b94c588bc0613700294408" + - "69b94c7cf50ffaf8aea15fd4" + - "69b94cabea8b4cce23afead0" + - "69b94d4ed3f6bebb1b52001b" + - "69b94fa2e094c1fa8268fbf1" + - "69b94fdbe8575b54b34add3a" + - "69b95092507132953a1af135" + - "69b951c8e094c1fa8268fc0d" + - "69b951fcd184f934b8b27494" + - "69b952b60e9f2891ea5e77a9" + - "69b953442030ce071fb494c7" + - "69b95346cf21eaeeccc6b2d7" + - "69b95359b3fe52e7b3246cb4" + - "69b9555663e615dae6ebcce5" + - "69b955689c31f4f8be85daa5" + - "69b956d9fa0ca5629329ecce" + - "69b957d0507132953a1af187" + - "69b957f9249d9593ce92e2a9" + - "69b958497ee0ff33631f0ac8" + - "69b958c4b3a349ac0a3d8d48" + - "69b959bbfe5dea81e80547fb" + - "69b95a7dace9d1c4923cff39" + - "69b95af0bd0a1226e8720d06" + - "69b95b4eaacf0a241873b89b" + - "69b95b7e73c73837a4222c13" + - "69b95b8d36bea964e3c3b00f" + - "69b95ba697d25b237a7cdab4" + - "69b95bceadb6e6a77d240169" + - "69b95bd3e7f0e45fa5e4894e" + - "69b95d23f50ffaf8aea16094" + - "69b95d6c85ed8de7009108d0" + - "69b95e3fa89e9b47645aa10f" + - "69b95e598bc06137002944f8" + - "69b95e97d056890bb2ca2233" + - "69b95f0b0b1ec7e2efa013e6" + - "69b95f279f40b5816ba082a2" + - "69b95f27e02b54ac36b08430" + - "69b95f5887a8695556dc140b" + - "69b95f74de59539246182d03" + - "69b95fd76a6fcc387e6fcc8f" + - "69b9606ee97bec8c72c5e685" + - "69b960ff6a84c22b59557b34" + - "69b961386a84c22b59557b38" + - "69b9618acf21eaeeccc6b380" + - "69b962063907853432078e77" + - "69b962702030ce071fb49584" + - "69b96282bd0a1226e8720d60" + - "69b962a34a18f0c8d9bf7d21" + - "69b9633af50ffaf8aea160f8" + - "69b9636d345333ff05a8d2f4" + - "69b963e9d1b1254fcf5880ea" + - "69b965116e470e7c633d7835" + - "69b966dddaa78337a3cf86c4" + - "69b966e27ee0ff33631f0b8a" + - "69b9674c8973f21812765e8d" + - "69b96862a89e9b47645aa19f" + - "69b96922446a7ee1a670dbf2" + - "69b96934b13a790889c05656" + - "69b96958b3cc90fa8d9acf77" + - "69b9698af25718a726e3477a" + - "69b96a3a73c73837a4222cd1" + - "69b96b27c3d83cdf8646bdb2" + - "69b96b4d9bcf11a7dc66d973" + - "69b96c59f50ffaf8aea16184" + - "69b96c78249d9593ce92e3a7" + - "69b96cf87ee0ff33631f0bd4" + - "69b96cfafe7b4065d8ee639c" + - "69b96d28b3a349ac0a3d8e5a" + - "69b96d493d7564e601970660" + - "69b96de7d3f6bebb1b5201b9" + - "69b96edfbd0a1226e8720df6" + - "69b96f559f40b5816ba08370" + - "69b96ff3e749b83b1a3332db" + - "69b97028020e9efa2add56e2" + - "69b971226a84c22b59557c1e" + - "69b971a9e094c1fa8268fdb1" + - "69b972438bc061370029461c" + - "69b97260fb73b7a0dcf4cd62" + - "69b972888973f21812765f2d" + - "69b972cffe5dea81e805492b" + - "69b97354de59539246182e2d" + - "69b9737d6e470e7c633d78f5" + - "69b973ba0e9f2891ea5e7967" + - "69b97440a89e9b47645aa245" + - "69b97449507132953a1af30b" + - "69b974d36a6fcc387e6fcd93" + - "69b97527b2ef41e07b9ad54c" + - "69b97593f25718a726e34818" + - "69b975abadb6e6a77d24031d" + - "69b975eca3ef2775e0e98fe9" + - "69b97601b3a349ac0a3d8ed6" + - "69b976672023be3344a8d8b6" + - "69b9776a4a7adca60e60cc8b" + - "69b9776c4458d25a26dc6d25" + - "69b9779fcf3be3afac7d78b3" + - "69b9786abb8856ba14a77581" + - "69b978bca0469793c24bf37f" + - "69b978e12e71fa6c931cc288" + - "69b9791db3cc90fa8d9ad039" + - "69b9794c6be1c6d054a435e1" + - "69b979854e1bd59d5fded938" + - "69b979b508daef11eac8a7a2" + - "69b979e5d056890bb2ca23be" + - "69b979fc4a18f0c8d9bf7e6b" + - "69b97a15e3e5accb88abd54d" + - "69b97b6d08daef11eac8a7be" + - "69b97b6d976bda2d64157eca" + - "69b97bd1c147ee9b771dcfac" + - "69b97c810f92f74da61ceb19" + - "69b97c92571aa3769df569ad" + - "69b97ce2b743dd78f6a7a132" + - "69b98142d056890bb2ca2427" + - "69b981bc118acc9f17ae3194" + - "69b981f3ec7e30bfdad0ec04" + - "69b982e4c5d47355c732ec84" + - "69b98498a6c46b5fcef51753" + - "69b984a23627655cc9c899ac" + - "69b98529e29a845eca678185" + - "69b986ef827d40654acf55bc" + - "69b9881cdefc1ea2052c568a" + - "69b98823d28517af7a845e92" + - "69b98ab013f22dc5a5ec5acd" + - "69b98b609d8f2a52dc2e03a6" + - "69b98bbe3d0324014543c330" + - "69b98bd7fe7b4065d8ee64d1" + - "69b98cf3a6c46b5fcef517bc" + - "69b98d08a5fa7432c21c9645" + - "69b98d15332eddd8ed1a352a" + - "69b98d3b45605c1e0a2b1100" + - "69b98d48d28517af7a845edb" + - "69b98da255ee96e0d7396778" + - "69b98ea72f58f0ebb008753a" + - "69b98f0af25718a726e348f6" + - "69b9900274e1f11b66f84223" + - "69b9901614b5b0e24e4fdde4" + - "69b990230404ba0716d1bc1f" + - "69b9904502fc9c70406fa58b" + - "69b990fab90b68a041f912af" + - "69b9911f0a20f57947d95a32" + - "69b991e75532d119a1e752f6" + - "69b9921b868d7a82dcebacc8" + - "69b9924502ef422eb23d8ae1" + - "69b992693f50bd74b8f9b37b" + - "69b99354a85730cbc3e51f18" + - "69b9942e7afa66d9c67317b6" + - "69b99470ed5649dcd05e73a5" + - "69b9947fd5316084ad4c4247" + - "69b9976d3627655cc9c89a6b" + - "69b99772e006dacdfca0af25" + - "69b997c774d8ca57221c45c7" + - "69b997f53673ca659e070aad" + - "69b9980711531b4a4388f9ff" + - "69b9980c345333ff05a8d50a" + - "69b9982e0708f95ac3fce45c" + - "69b998647460200d7164617b" + - "69b998c374d8ca57221c45d3" + - "69b9996a9d8f2a52dc2e0406" + - "69b99980c04056253aea576d" + - "69b999b046911c9d0fdcbb4f" + - "69b99acbbc74703b35644e77" + - "69b99ad128a3d930ce23de46" + - "69b99b4cfdfee6bde9ea42c0" + - "69b99ec420c9cd9395eb7ac4" + - "69b99ec7a85730cbc3e51fb4" + - "69b99f488433ca9c37c44e51" + - "69b99fba20c9cd9395eb7aca" + - "69b99fe25888159a2960d68e" + - "69b9a059c69a19757c9ec3e7" + - "69b9a06718c008a21ec40045" + - "69b9a0ec54bdf791a9510dd1" + - "69b9a13b5ffb742f9e71165c" + - "69b9a1737d726cc6462425de" + - "69b9a184fdfee6bde9ea42e7" + - "69b9a24814b5b0e24e4fded7" + - "69b9a2c0eddc130379669f2d" + - "69b9a2c4e29a845eca678243" + - "69b9a30a46911c9d0fdcbbab" + - "69b9a35702795913b379aac5" + - "69b9a371f25718a726e34982" + - "69b9a37621f3c78fd919b7f4" + - "69b9a3f182216e8e92e3c60c" + - "69b9a42c00d85d9c3fe9c12d" + - "69b9a46f740847c0e44f34eb" + - "69b9a535d647ee40234c80b0" + - "69b9a5d574d8ca57221c4671" + - "69b9a60b8ae31a8e66653a4d" + - "69b9a6a202795913b379aadf" + - "69b9a724615fa2ddcffeff85" + - "69b9a759426603a5b4624d4f" + - "69b9a785a34871e32cd951ca" + - "69b9a819d9bfbe60f0573653" + - "69b9a85800d85d9c3fe9c15f" + - "69b9a94f03c5e6f8a5e6d27f" + - "69b9a988eddc130379669f76" + - "69b9a9ff5a52e001e0cbdf47" + - "69b9aa13d9bfbe60f0573663" + - "69b9aa31f9568308a55add66" + - "69b9aa37f626d61b0f19cf28" + - "69b9aa817c1f9b2d4c5f8be3" + - "69b9aa8ba6c46b5fcef51939" + - "69b9aa950a4d21fe795c1d53" + - "69b9aaa44a24dede6ab84f51" + - "69b9aae431c65d1c8cb6c55a" + - "69b9aafb2ac3284b987726b2" + - "69b9abdaf5c97437d7e74986" + - "69b9ac17d5c8a084192357e1" + - "69b9ac26d3b355e380a94bdb" + - "69b9ac2ecc5ab1c922041b60" + - "69b9ac3630e465706c590ffc" + - "69b9ac4e46911c9d0fdcbc00" + - "69b9ac79eccb0b08b29a520b" + - "69b9acb86b5e7703a13b3c4b" + - "69b9ad1a332eddd8ed1a3604" + - "69b9ad9b6434c7a796518dce" + - "69b9ae2652dad08d449b01bd" + - "69b9ae2adb7fde1db3e0deda" + - "69b9ae9d6bbdec4efc1c48e0" + - "69b9aeaa4673c55fc915ba45" + - "69b9b059c5d47355c732ee40" + - "69b9b09941bad320ee133b65" + - "69b9b0b588b0b7e363d2fcdf" + - "69b9b1c455ee96e0d7396949" + - "69b9b1e3c69a19757c9ec498" + - "69b9b207f99fb11c7b078412" + - "69b9b2169fd2f95542baefa7" + - "69b9b2924673c55fc915ba66" + - "69b9b2a88e04b030b1532a8f" + - "69b9b2f9e5c216ef54bdabc0" + - "69b9b310c4c32649b8a8946b" + - "69b9b3261b71fa5cbda7fc41" + - "69b9bd847ae2957b84651aeb" + - "69b9beca253521f049c31ee2" + - "69b9bfc2d28517af7a8460fb" + - "69b9c0cd5532d119a1e7543c" + - "69b9c21ba6f41015b52033ab" + - "69b9c23d740847c0e44f360a" + - "69b9c253c1c0794c8dcd496f" + - "69b9c282565746ee03cbb3dd" + - "69b9c32607283983c22dd9ac" + - "69b9c3c4ef8c2b6d9612f24f" + - "69b9c4c796eb780c06ea5ddd" + - "69b9c5a94c322fa28c9a956a" + - "69b9c5b42e91c763d8d888fd" + - "69b9c6b10a168c5d1f814632" + - "69b9c6bb9563c2ec63a04af9" + - "69b9c72adea38cbf40e8242f" + - "69b9c94b30e465706c591106" + - "69b9c998dea38cbf40e8244f" + - "69b9c9dc5e8bf02d9dd9d6f6" + - "69b9ca16192d2ab7670dc0c6" + - "69b9ca2c97ebd3eddbcf2397" + - "69b9ca33529ed4055306149f" + - "69b9ca7ec004b8e369a3d5d3" + - "69b9cb12df7d76fec0802291" + - "69b9cb5c5edbea42cde66f65" + - "69b9cb8f07283983c22dda1e" + - "69b9cb9330e465706c59112a" + - "69b9cc0a5f7ef497d251f185" + - "69b9cc1daf16ec3f4defb6f7" + - "69b9cc3c5d426ff433fda099" + - "69b9cc40c004b8e369a3d5e5" + - "69b9cc6c97ebd3eddbcf23b3" + - "69b9cc700a20f57947d95cf7" + - "69b9cd2f8b764038e339aca9" + - "69b9cd44db7ef5ce54cd7f15" + - "69b9cd67dea38cbf40e8247d" + - "69b9ceb1192d2ab7670dc106" + - "69b9cf0b6522a44325c950c2" + - "69b9cf78e7f8ff8bcf4ee784" + - "69b9cfffd28517af7a8461f8" + - "69b9d106bb20394db7cb8439" + - "69b9d1177460200d71646424" + - "69b9d135913f43e764571a38" + - "69b9d2206522a44325c950f2" + - "69b9d2795532d119a1e7550a" + - "69b9d27e7460200d7164643c" + - "69b9d31fd963aad3aa6c37d7" + - "69b9d3c7cc5ab1c922041cde" + - "69b9d41ab64dff59fd1895b2" + - "69b9d479136f1bb5f505f42c" + - "69b9d589114b0bb34640372a" + - "69b9d5997ae2957b84651bf3" + - "69b9d63988b0b7e363d2fe4e" + - "69b9d99f30e465706c5911fb" + - "69b9da39f8228b781f97b13f" + - "69b9da3ddb7ef5ce54cd7fa2" + - "69b9daeda6e02486fe080c4d" + - "69b9db020c7a6ded83fe14a7" + - "69b9db477a815d66acd21665" + - "69b9db867c2f56356e9e9c8c" + - "69b9db98ab3326fefa02a813" + - "69b9dd84b9a2a0ab7ccd37c3" + - "69b9e01a3dc51caf38b77aeb" + - "69b9e17edcf3f4ec2e48a0bb" + - "69b9e2186522a44325c951df" + - "69b9e22b1622d2fd0809afe4" + - "69b9e27573b4ee215ccb1c1c" + - "69b9e309f4432734dc1b3a5d" + - "69b9e3184a24dede6ab85117" + - "69b9e338b1f61839c2d3d2e3" + - "69b9e3b2259f8dad0074f0a7" + - "69b9e407bf6faf79a0b9f6b9" + - "69b9e45c136f1bb5f505f4be" + - "69b9e4627c2f56356e9e9d28" + - "69b9e473700b828eaaaa7eda" + - "69b9e48bc2ba99f1e9bbada8" + - "69b9e498af16ec3f4defb843" + - "69b9e4aceaca3559cb5a9903" + - "69b9e4af7e993aad6a355d10" + - "69b9e4da7c7691048de8d1bf" + - "69b9e52e6e28a6e6253238e7" + - "69b9e545bbb970524d3cf45d" + - "69b9e547fccae70c7a7f12f1" + - "69b9e58a5532d119a1e755aa" + - "69b9e58ab64dff59fd189681" + - "69b9e59585af398716d00143" + - "69b9e5a57bb83b592af1b814" + - "69b9e5aae2f3e19f53b175a1" + - "69b9e5b099820bddbe86ae45" + - "69b9e5c08d17d8a50177c868" + - "69b9e60811e27738047cbba9" + - "69b9e60bfa8594e2e697cfba" + - "69b9e636c744d9c59679801f" + - "69b9e76d9c6c9ff2c936c7ab" + - "69b9e8790a5db606ab760cd1" + - "69b9e8c6a4be4f88ef786438" + - "69b9e906765c34e61fcc5923" + - "69b9eb8fbbb970524d3cf496" + - "69b9ec75868c6ba99f13ad89" + - "69b9ec7fa6e02486fe080d6f" + - "69b9f00304ac9216d9b54336" + - "69b9f08e27ee5763b67cb054" + - "69b9f176e6e790699c51eea4" + - "69b9f1843404c306f8caa8fa" + - "69b9f1d70ed8c646a6770f89" + - "69b9f1f11468d5fad08b6e10" + - "69b9f1f9208bd760e93f5191" + - "69b9f23ceb0d19098758c911" + - "69b9f27436c4cbf22e304eb0" + - "69b9f2a9630c5e7b37c10a7d" + - "69b9f2b00f404e75c2bb55ad" + - "69b9f319cdb35089f2fd4d15" + - "69b9f33a789f73cd03e95c7d" + - "69b9f43741e4e9d70d465295" + - "69b9f52134dc4e1f3b8faaab" + - "69b9f5709f1276ff46fe5b65" + - "69b9f6ffcc5ab1c922041dee" + - "69b9f851eea7173ad053f4c9" + - "69b9f8785a7b6b085811c1a7" + - "69b9f93fb390cfa5f2469f70" + - "69b9f94f75f11ecfeb793349" + - "69b9fa95820274d0e8c21706" + - "69b9faad5ce896083ab0d6ca" + - "69b9faddb4874252fd2ade55" + - "69b9fb285532d119a1e755fd" + - "69b9fbf6750de787958219b8" + - "69b9fc5597a8dc055e571284" + - "69b9fcf51b71fa5cbda7ff0e" + - "69b9fdca31631a97c54106f4" + - "69b9fddea897194fbc62dabd" + - "69b9fe48136f1bb5f505f530" + - "69b9fee711acf66ee75a1bbc" + - "69b9ff08a4dc83d0e3e50b80" + - "69b9ff1cb3eba4062eae11b0" + - "69b9ff350db321af2af1cffb" + - "69b9ff9235bd438e92af49c8" + - "69ba00719bc4a64f7def938c" + - "69ba00a2a862bc58898f4778" + - "69ba00b5afb568e39e92f340" + - "69ba00c3c1ff235ac7befce6" + - "69ba010100acc44fea9f9a8e" + - "69ba0176c31dd0773e99cad2" + - "69ba01e303c8a4aff336258d" + - "69ba027d549143988cf4d879" + - "69ba0409e8a238408e8c935e" + - "69ba049c25e5abf0d324f865" + - "69ba05ea70bed64513114630" + - "69ba065b114b0bb346403873" + - "69ba066a3bf97c04957feabc" + - "69ba06f4862b272d33cfec6f" + - "69ba077a711f472f28e00f57" + - "69ba07b023ed03a377c59f4f" + - "69ba07e5e2f3e19f53b176cc" + - "69ba0828fccae70c7a7f138d" + - "69ba0833fc9975b47da66c2f" + - "69ba094ca6f41015b5203692" + - "69ba09c503aa8c39fb213e6a" + - "69ba09c745fcbca6914670fb" + - "69ba0a3388073a6c9a161ee0" + - "69ba0a8202ef422eb23d8e66" + - "69ba0ad7a862bc58898f47b6" + - "69ba0b2c5d41223954a48138" + - "69ba11ac03c8a4aff33625d0" + - "69ba1387136f1bb5f505f585" + - "69ba157cc004b8e369a3d825" + - "69ba1b4dd81a7c6b0ba59fdc" + - "69ba1bb2565746ee03cbb608" + - "69ba1c0f0199ca05374311d3" + - "69ba1c3745605c1e0a2b157c" + - "69ba1c5141b2115214eb13f8" + - "69ba1ce2e729db2d75ed3293" + - "69ba1e5302ef422eb23d8ea9" + - "69ba1f44cae877cbaf8a1251" + - "69bac7cb40ec44cc3c14daef" + - "69bac7f5133b888b6cabec9d" + - "69bac8edde337ef3bb9f27fc" + - "69bac93236f6f84621c27f27" + - "69bac9ef39a7d1afdf43744b" + - "69bac9f03123a00be5e037cc" + - "69baca448309b566dfb6ff31" + - "69baca662e74bd8072a07637" + - "69bacac24a1ed3f60f69f569" + - "69bacb125ef2ce13c70f1e45" + - "69bacbcb0d127eea9c48b5e3" + - "69bacc10f1c3256ecb58f4fd" + - "69bacc9545e41a2ae14c50bb" + - "69baccbc51f370e3808cb172" + - "69bacd0ab7af83de25a38027" + - "69bacd110d127eea9c48b5f3" + - "69bacdd1fa5e6eb75daa47f1" + - "69bace9591aa8f4b5b84fe72" + - "69bacf3023fe22b05ad5f8e7" + - "69bacf350e0205088620670a" + - "69bacf373a510372caab01f8" + - "69bad0054fea4b718b4a447f" + - "69bad052ba4ac30debb7f974" + - "69bad107fbfb477d1e6941e2" + - "69bad119e3feab0e9044bbf5" + - "69bad135452c948b1ed973ee" + - "69bad1af864173fe45257daa" + - "69bad1cecefe6f3b52bc1a35" + - "69bad1e9fb4c94a31bd85af1" + - "69bad1f586d23e7d740753b0" + - "69bad31993a2e9a8268a7dd0" + - "69bad3c4b1e3407323aa71cd" + - "69bad3dd82d693a700cc3fbd" + - "69bad4911d29cdc2ac1d2dde" + - "69bad510538bfccd32b608bd" + - "69bad52381fa3d8a22c8c0e8" + - "69bad53e8fe2be8eaa451406" + - "69bad56d7f0654debc0e3392" + - "69bad57d935a9aa010abb111" + - "69bad5ce248c31cadc9604fe" + - "69bad65bcefe6f3b52bc1a71" + - "69bad7c868a98b44c75c7d17" + - "69bad7e5b1e3407323aa71f5" + - "69bad7fe12343a4a72551d07" + - "69bad82410ece1f993b3efe0" + - "69bad8e9ecb066bea1c404ed" + - "69bad9b2ed46fe5d7036ab16" + - "69bada1e02f0142ffeba0e6e" + - "69bada5b790ccc5cb89c1083" + - "69bada9d828bb1912f3ef86f" + - "69badabe62fb9126dcbeebfc" + - "69badaec1f65cac1b9f76846" + - "69badb4886da1211813755db" + - "69badb57a5df1ac4a38af33e" + - "69badb663abe312e772beb4f" + - "69badbbef1c3256ecb58f623" + - "69badbc3012a47486fe5a684" + - "69badc26012a47486fe5a68a" + - "69badc6a93a2e9a8268a7e52" + - "69badc8653dcecdfbfad8595" + - "69badcc5fbff68b7b12b8cc3" + - "69badcecc27f2c10349cce90" + - "69badd62579f289c85b84ed6" + - "69badd7ada5440bb6aa6697b" + - "69badda6df5b9967d32bc318" + - "69baddb60e02050886206792" + - "69badddac0ba9e5e22a17f7c" + - "69bade09828bb1912f3ef89b" + - "69bade5b3c08def5bd53bcdb" + - "69bade7324937b609e379019" + - "69bade76fa5e6eb75daa48b0" + - "69bade8d4a1d1416e5738005" + - "69bade8d6a34fc83ba539999" + - "69bade9cf7dd56e062995b38" + - "69badea43c312d27a6078c1b" + - "69badf14ce45e4d85b430f10" + - "69badf536c7eaf28bc3b6aae" + - "69bae02e4cfec299236c85e5" + - "69bae03cd738810497993a47" + - "69bae05927b7b28b5b27ed05" + - "69bae0840826d6aa1fc63c7d" + - "69bae0de8c02130ff9f0c641" + - "69bae15aa61c0d91b91506d6" + - "69bae17897c3826a73c085ce" + - "69bae18c8c02130ff9f0c64b" + - "69bae1d7d10b1c996b2a37a7" + - "69bae1eea9cd03e5237ff0cd" + - "69bae27f1bfde911ab6bfeb2" + - "69bae2a9c27f2c10349cced2" + - "69bae2c6acebd405a2f98c2f" + - "69bae2ec7defa4aeece5f682" + - "69bae2ed51ca711667bb8066" + - "69bae34f1e0eb5c4ba7078b0" + - "69bae36d99f35e9578425f54" + - "69bae37a45e41a2ae14c51ad" + - "69bae3c84012b22f2ea5f449" + - "69bae3d9fb4c94a31bd85bc7" + - "69bae41a02f0142ffeba0ef6" + - "69bae4352eba6823ab4afc00" + - "69bae46d7f9d2a9ee33a6de3" + - "69bae48851f370e3808cb2cf" + - "69bae510363717331941bbc0" + - "69bae57012343a4a72551d9d" + - "69bae733b7af83de25a3812a" + - "69bae7683bc7f73174aaecdb" + - "69bae7f7c40add85137b6f3c" + - "69bae7febae66100347d2988" + - "69bae9002c3d669e7fabae2b" + - "69baeb290d0367fe463d8c01" + - "69baeb77a2eb4da46d6a85b2" + - "69baec3afb4c94a31bd85c06" + - "69baec83a9cd03e5237ff11d" + - "69baecd2a61c0d91b915072d" + - "69baed3a01b99fbbf664aebe" + - "69baeea1a9cd03e5237ff12d" + - "69baeee0ab6213a943bb8811" + - "69baeeeb65807ea20da0f505" + - "69baef0fc9d1064915f12829" + - "69baef27d10b1c996b2a3813" + - "69baef48c9d1064915f1282b" + - "69baef992bbc3c64164dbe08" + - "69baf076e0792196122e23e1" + - "69baf08c8c02130ff9f0c6cd" + - "69baf0e87e311f13fc6df155" + - "69baf10efb4c94a31bd85c2c" + - "69baf149daa28f952f17b7a7" + - "69baf1907236446f4f672422" + - "69baf1c4e0792196122e23ed" + - "69baf1ddf2ab75b66f8bb41a" + - "69baf218f82295aec4769096" + - "69baf2c1fc73c01f8801608b" + - "69baf2f8db46c0031ed32795" + - "69baf3561e737760229bc509" + - "69baf439a2eb4da46d6a85f2" + - "69baf5d3ea5dc1c74861cc01" + - "69baf6a05d1a1f3632b7fd41" + - "69baf6ddf205b66bf504f4fa" + - "69baf75be14594caf653b016" + - "69baf910c45e3578443baca1" + - "69baf94252a0f72bd574a423" + - "69baf988026035f7322e5976" + - "69baf9b2da5440bb6aa66a6d" + - "69baf9e7d0b23f532acb4f05" + - "69baf9f6f742d6da3389512b" + - "69bafa8d071c4967be381d1c" + - "69bafd01e423ad04bb17091c" + - "69bafe20363717331941bc70" + - "69bafe26466ebc139e53c718" + - "69bafef4e75c835551cd08d1" + - "69baff4929ed959db4a3f2d8" + - "69baff5ec233bd773a54a9dd" + - "69bb00a60bb085884dbc50b7" + - "69bb00c55d1a1f3632b7fd8b" + - "69bb0122b89ba1de8852d4d5" + - "69bb017767738c4ecdc63b2e" + - "69bb02624529d34e79e7d8ab" + - "69bb02e5d50175523d0489a5" + - "69bb030d8e236ab0291cd024" + - "69bb04234a46b3886b8c12f5" + - "69bb05f996c7ea4843e12bf5" + - "69bb06af0826d6aa1fc63da5" + - "69bb07117ac42379425acbb2" + - "69bb078cd1413a767f2ec14d" + - "69bb08363abe312e772becaf" + - "69bb0986d738810497993b87" + - "69bb0995e0792196122e24c5" + - "69bb0a998606a40f2af2a772" + - "69bb0ac701b99fbbf664af9e" + - "69bb0b14daa28f952f17b872" + - "69bb0b6d3abe312e772beccb" + - "69bb0bf302f0142ffeba103c" + - "69bb0c09c27f2c10349cd00a" + - "69bb0c7e411dd3347c32cacf" + - "69bb0cef4a1d1416e5738177" + - "69bb0f9d7c0ffa66dece83ee" + - "69bb105ffc73c01f8801617a" + - "69bb1284d385c444f34222a7" diff --git a/egomimic/hydra_configs/data/mecka_all_6d.yaml b/egomimic/hydra_configs/data/mecka_all_6d.yaml new file mode 100644 index 0000000000..9312e609c7 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_all_6d.yaml @@ -0,0 +1,94 @@ +# MECKA_ALL_PI_6D: pi0.5 on ALL mecka episodes in the SQL table (filters +# lab=mecka with no task filter), stored under the human_bimanual embodiment. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation representation +# (18D = xyz+6D per arm, no gripper): each arm's action chunk is expressed +# relative to that wrist's current pose, so both the arbitrary world frame AND +# the moving head frame cancel out of the action targets. The proprio ee_pose +# stays in the head/cam frame (it defines where the wrists are for the model +# and for the eval revert) and is 6D-encoded. stride=1: mecka actions are +# recorded at the target rate (Aria uses 3). Language prompts are sampled +# from the `annotations` track; prompt assembly (proprio-in-prompt on, +# embodiment label off) lives in model=pi0.5_bc_mecka_6d, which this config +# pairs with, alongside evaluator=eval_pi_wristframe_6d (reverts wrist -> +# head frame via the proprio, then 6D -> ypr, for the viz). +# +# This is the explicit cartesian recipe; the repo default for human data is +# the hand-keypoint action (data/mecka.yaml, data/cotrain_pi_base.yaml). + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + # Per-cluster zarr mirror (paths/default.yaml; $EGOVERSE_DATASET_DIR + # overrides). The resolver s5cmd-syncs only the SQL-matched episodes + # missing here. + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + # processed_v3 mecka zarrs predate the mecka_to_zarr rot_left fix: + # the LEFT wrist frame was double-mirrored onto the right hand's + # spatial convention. Rz(180)-corrects the raw left pose keys before + # any frame math (exactly equivalent to reconverting). + fix_left_wrist_convention: true + # 18 -> 20-dim proprio (zero grip slots at 9/19): State: prompt bins + # align positionally with the robot 20-dim layout. + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + # Failed vendor egomotion (constant or jumping head pose) corrupts the + # head-frame action chunks; see scan_egomotion.py. + exclude_hashes: mecka_bad_egomotion + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# Train-set spot check: a second val-style loader over the TRAIN split, +# consumed by the TrainVizEvalVideo wrapper on every validation epoch +# (metrics prefixed `train_viz/`, videos under `videos_train_viz/`). No +# shuffle -> the same leading slice every val, so the curves are comparable +# across time. Declared here only to set its loader params; without it the +# default train_viz head reuses the train split with the valid loader params +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: train + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false 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 new file mode 100644 index 0000000000..36d76ec617 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -0,0 +1,226 @@ +# MECKA_FOLD_FLAGSHIP_OPSPLIT_HPT_6D: HPT BC on the mecka FLAGSHIP fold-clothes +# episodes (SQL lab='mecka', task='fold_clothes', all 1607 eps / 316k frames +# under processed_v3/mecka/flagship/, 20 operators; a DIFFERENT set from the +# 462-ep FREEFORM 'folding_clothes' episodes of mecka_fold_freeform_opsplit_6d, +# zero operator overlap), with a SEEN- and an UNSEEN-operator validation set. +# Standalone (no defaults inheritance). +# +# THREE-WAY SPLIT (chosen 2026-09-11 from SQL app.episodes): +# unseen val = the 2 operators in `held_out_operators`, ALL their episodes +# (191 eps / 30,956 frames, 9.8 %). The two mid-sized operators +# with >= 70 episodes each; the three largest (65 % of frames) +# stay in train. +# seen val = the other 18 operators, seeded 5 % episode split +# (`mode: valid`, ~70 eps / 13.8k frames): known operators, +# unseen episodes. +# train = the other 18 operators, the remaining 95 % (~1346 eps / +# 271k frames). +# +# VAL LOADERS (three heads, one validation pass): +# valid_datasets -> seen_op_valid/Valid/human_bimanual_*, +# videos_seen_op_valid/: SEEN operators, unseen +# episodes (the canonical loader; `valid_prefix` +# renames its metrics and video dir) +# train_viz (default) -> train_viz/Valid/human_bimanual_*, +# videos_train_viz/: the TRAIN split itself (no +# train_viz_datasets here, so trainHydra reuses +# train_datasets, unshuffled); `train_viz=false` +# drops it +# unseen_op_valid_datasets -> unseen_op_valid/Valid/human_bimanual_*, +# videos_unseen_op_valid/: UNSEEN operators +# The train and seen-val entries share `filters` and `valid_ratio`, so the +# seeded split (split_dataset_names, SEED 42) makes them exact complements. +# +# METRIC vs VIDEO LOADERS (2026-09-15, lane A). Each head above is served by +# TWO loaders: a per-episode subsampled METRIC loader (no rendering) and a +# contiguous VIDEO-only loader over pinned episodes (no metrics). See +# `metric_frames_per_episode` / `video_episodes` at the bottom of this file. +# +# PIPELINE: WRIST-frame CARTESIAN actions with the continuous 6D rotation +# (18D = xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose; the head-frame proprio ee_pose is +# 6D-encoded too and padded 18 -> 20 (zero grip slots at 9/19). +# fix_left_wrist_convention: processed_v3 mecka zarrs have the LEFT wrist frame +# double-mirrored; Rz(180)-corrected before any frame math. stride 1 (mecka +# is recorded at the target rate). Prompts from the `annotations` track. +# +# NORMALIZATION: quantile-normalize proprio + action keys from stats +# precomputed on THIS config (train side = the 18 operators): +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_opsplit_hpt_6d --model hpt_bc_mecka_6d_300M --out +# The freeform stats share the key set but not the distribution; do not reuse. +# +# Pairs with model=hpt_bc_mecka_6d_300M and evaluator=eval_hpt_wristframe_6d +# (the train_viz head wraps it automatically); see +# train_zarr_mecka_flagship_6d_hpt.yaml for the composed recipe. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 2 held-out (UNSEEN) operator ids (SQL app.episodes.operator). Referenced +# by the filter lambdas below via the mounted path `data.held_out_operators`. +# MultiDataModuleWrapper accepts this key (config-only). +held_out_operators: + - 6871d82a1aaf88556f0c6041 + - 6903a36c466be157ec64562d + +# Metric prefix / video dir for the canonical valid loader (config-only; read by +# trainHydra), so the seen-operator val is not logged under a bare Valid/. +valid_prefix: seen_op_valid + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + # Proprio observation history: the ee_pose keys read K steps spaced + # `history_stride` frames apart, (K, 20) with the current frame last. + # At 30 fps the median wrist translation is 6.8 mm over 1 frame but + # 42.9 mm over 5, so an unstrided window is mostly sensor noise while + # the action chunk it conditions spans a full second; K = 3, s = 5 is + # 0.33 s of past. Must match the stem's `history_len` (the run fails + # loudly if not), and the K = 1 baseline needs + # ...key_map.proprio_history=1 ...key_map.history_stride=1 + # model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len=1 + # The two val loaders interpolate this resolver, so they follow. + proprio_history: 3 + history_stride: 5 + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: train = every operator NOT in the held-out list. + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + # Seeded 95/5 episode split within the train operators; the 5 % is the + # seen-operator val (valid_datasets below). + mode: train + valid_ratio: 0.05 + +# SEEN operators, held-out episodes (the complement of train's 95 %) -> +# seen_op_valid/ metrics via the canonical evaluator. +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# UNSEEN operators: all episodes of the held-out operators -> unseen_op_valid/ +# metrics via a second copy of the canonical evaluator. +unseen_op_valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: unseen val = the held-out operators only. + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +# Unshuffled (both val loaders): the same leading slice every val, comparable +# across epochs with a coherent overlay video. Capped per pass by the +# trainer's limit_val_batches. The train_viz loader reuses these params. +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +unseen_op_valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +# --------------------------------------------------------------------------- +# METRIC LOADERS: K evenly spaced frames per EPISODE (lane A, 2026-09-15). +# --------------------------------------------------------------------------- +# Config-only (trainHydra reads them; MultiDataModuleWrapper declares them +# because hydra.instantiate(cfg.data) forwards every root key). +# +# WHY: the val loaders are unshuffled so the overlay video is coherent, so the +# metric window was just the first limit_val_batches x batch_size x world_size +# = 80 x 64 x W frames in episode-hash order (DistributedSampler strides, so +# the ranks between them walk one contiguous head of the split). Measured +# 2026-09-15 at W = 1 (lane-0 diagnostic, docs .../reports/lane0-coverage.md): +# the seen-op head scored 25/70 episodes and 6/11 operators, and the unseen +# head 33/191 episodes with an 81/19 operator mix on a split that is 52/48. +# trainHydra now wraps each val head in EvenStrideDataset(base, +# frames_per_episode=K), which takes K evenly spaced frames from EVERY episode +# (all of them when the episode is shorter than K), so every episode and +# operator is scored at the SAME 80-batch cost. +# +# K below is the SINGLE-RANK value, floor(5120 / episodes-in-that-split); +# trainHydra multiplies it by trainer.devices x num_nodes, so a 4-GPU run +# scores 4x these frames in the same 80 batches per rank rather than a quarter +# of what it scored before subsampling existed. Episode counts measured from +# the real splits, 2026-09-15: +# valid 70 eps / 13,744 frames -> K = floor(5120/70) = 73 +# -> 70 x 73 = 5,110 frames = 79.8 batches +# train_viz 1,346 eps / 271,247 frames -> K = floor(5120/1346) = 3 +# -> 1346 x 3 = 4,038 frames = 63.1 batches +# unseen_op_valid 191 eps / 30,937 frames -> K = floor(5120/191) = 26 +# -> 191 x 26 = 4,966 frames = 77.6 batches +# NOTE: per-batch Frechet _min/_max now span many episodes instead of ~2 s of +# one, so those curves are NOT comparable with pre-2026-09-15 runs. +metric_frames_per_episode: + valid: 73 + train_viz: 3 + unseen_op_valid: 26 + +# --------------------------------------------------------------------------- +# VIDEO LOADERS: pinned, contiguous episodes (`_video` loaders). +# --------------------------------------------------------------------------- +# One episode_hash per operator per head. Pinning them here (rather than +# letting the loader order decide) means every run renders the SAME episodes, +# so videos are comparable across runs. The pinned episodes are taken out of +# that head's already-resolved split -- no second SQL pull -- and played in +# frame order; `evaluator.viz_max_batches` (10 batches = 640 frames) caps it. +# These loaders log NO metrics, and the metric loaders above render nothing. +# +# Chosen 2026-09-15 from the real splits; format is (operator, frames): +# valid 692e771c39719ab57395b408 (690364170e94ce070afd1da5, 498f) +# the longest seen-val episode: 498f = 7.8 of the 10 video +# batches from a single contiguous episode (the split mean +# is ~196f, which would only fill ~3). +# train_viz 692e9bda21fc595fbe6f941e (6903686e0e94ce070afd1f24, 496f) +# longest train episode of the largest operator; it is also +# in the topop train split, so that child inherits this pin. +# unseen_op_valid 692eadbaaec602a46af10686 (6871d82a1aaf88556f0c6041, 204f) +# 692ea49a727c13b350cb7c48 (6903a36c466be157ec64562d, 212f) +# one per held-out operator, both of typical length +# (416f = 6.5 of the 10 batches). Operator 6903a36c... does +# have a 661f episode (692fe8d889e6f7f45d175b78) that would +# fill the budget alone, but one-per-operator keeps both +# unseen operators on screen. Shared by every flagship +# config, since the unseen head is identical in all of them. +video_episodes: + valid: + - 692e771c39719ab57395b408 + train_viz: + - 692e9bda21fc595fbe6f941e + unseen_op_valid: + - 692eadbaaec602a46af10686 + - 692ea49a727c13b350cb7c48 diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml new file mode 100644 index 0000000000..11b8820795 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml @@ -0,0 +1,32 @@ +# mecka_fold_flagship_opsplit_hpt_6d with the wrist-frame KEYPOINT pipeline +# (the keymap / transform of mecka_fold_freeform_opsplit_hpt_keypoints): +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action, per hand [wrist xyz+rot6d +# (9) | 21 keypoints (63)] +# observations.state.keypoints (144,) per hand [wrist pose in HEAD +# frame | keypoints in wrist frame] +# No `observations.state.ee_pose` (include_ee_pose off), so +# pad_proprio_gripper is off. The operator split, filters, val loaders and +# video episodes are inherited unchanged. Pairs with +# model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. Norm stats are data-specific: +# precompute them for this config (egomimic/scripts/precompute_norm_stats.py). +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + # The keypoint proprio is in the WRIST frame, where a history step + # would carry the wrist's own displacement into the articulation + # channels; stay at the parent's pre-history setting until the + # per-step framing lands. + proprio_history: 1 + history_stride: 1 + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml new file mode 100644 index 0000000000..ba728fd656 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml @@ -0,0 +1,93 @@ +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +# MECKA_FOLD_FLAGSHIP_TOP3_HPT_6D: THREE-operator twin of +# mecka_fold_flagship_opsplit_hpt_6d (and the wider sibling of the +# single-operator mecka_fold_flagship_topop_hpt_6d). TRAIN on the 3 largest +# flagship fold-clothes operators only; everything else (pipeline, +# unseen-operator val, loader params) is inherited from the parent. +# +# SPLIT (operator totals from SQL app.episodes, 2026-09-15; the train/val +# numbers measured from the real seeded split on a compute node the same day): +# train ops = the 3 largest of the 20 flagship operators, by frames: +# 6903686e0e94ce070afd1f24 405 eps / 76,013 frames +# 690366b20e94ce070afd1e8a 354 eps / 67,085 frames +# 683785ac01ca734152093448 320 eps / 63,178 frames +# total 1,079 eps / 206,276 frames +# = 1.91 h at 30 fps, 65.3 % of the flagship set. +# None of the three is in the parent's `held_out_operators`. +# train = 95 % of those episodes: 1,026 eps / 196,450 frames +# (6903686e 382 eps/71,722 f, 690366b2 334/63,397, +# 683785ac 310/61,331). +# seen val = the other 5 % (53 eps / 9,826 frames) -> seen_op_valid/. The +# parent's valid_datasets interpolates the train filters and +# valid_ratio, so it follows this override automatically. +# train_viz = the train split above (trainHydra's default) -> train_viz/. +# unseen val = UNCHANGED from the parent: all episodes of the 2 +# `held_out_operators` (191 eps / 30,937 frames) -> +# unseen_op_valid/, directly comparable with the 18-operator +# and single-operator runs. +# The other 15 operators are unused. +# +# Hydra replaces lists wholesale, so the train filter_lambdas are restated in +# full; only the last (operator) lambda differs from the parent. The operator +# ids are inlined rather than put in a config key: MultiDataModuleWrapper +# rejects root keys it does not declare. +# +# NORMALIZATION: the train distribution differs from the parent's and from +# topop's, so precompute stats on THIS config: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_top3_hpt_6d --model hpt_bc_mecka_6d_300M --out +# then train with +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# data=mecka_fold_flagship_top3_hpt_6d norm_stats.precomputed_norm_path=/norm_stats + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # TOP-3 OPERATOR TRAIN: the three largest flagship operators only. + - "lambda row: str(row.get('operator', '')) in ['6903686e0e94ce070afd1f24', '690366b20e94ce070afd1e8a', '683785ac01ca734152093448']" + +# --------------------------------------------------------------------------- +# METRIC / VIDEO LOADERS: only the SEEN-op and train_viz entries differ from +# the parent (this config trains on 3 of the 18 operators, so those two splits +# are smaller); the unseen head is identical, so its K and pins are inherited. +# --------------------------------------------------------------------------- +# K is the SINGLE-RANK value, floor(5120 / episodes-in-that-split), 5120 = +# limit_val_batches (80) x batch_size (64); trainHydra scales it by the world +# size. Episode counts measured from the real splits 2026-09-15: +# valid 53 eps / 9,826 frames -> K = floor(5120/53) = 96. The +# shortest seen-val episode is 127f > 96, so every episode +# contributes exactly 96: 53 x 96 = 5,088 f = 79.5 batches. +# train_viz 1,026 eps / 196,450 frames -> K = floor(5120/1026) = 4 +# -> 1026 x 4 = 4,104 f = 64.1 batches. +# `unseen_op_valid: 26` is inherited unchanged (191 eps, same split as the +# parent), and hydra merges dicts, so only the two changed keys are restated. +metric_frames_per_episode: + valid: 96 + train_viz: 4 + +# The parent's seen-val pin belongs to operator 690364170e94ce070afd1da5, +# which is not in this config's 3-operator split, so `valid` is replaced +# (hydra replaces lists wholesale). Seen-val now spans 3 operators, so ONE +# episode is pinned per operator and the video shows all three. The train_viz +# pin 692e9bda21fc595fbe6f941e IS an episode of operator 6903686e..., hence in +# this train split too (confirmed on the measured split), and the unseen pins +# are the same split as the parent's -- both inherited. +# valid 692ea44dc621d7f4aac3aae6 (683785ac01ca734152093448, 191f) +# 692eb08997ffd8ec90340290 (690366b20e94ce070afd1e8a, 223f) +# 692fe79f34a99e18e25289ad (6903686e0e94ce070afd1f24, 226f) +# All three are at or above the split's 177f median length, and +# 191 + 223 + 226 = 640 f = EXACTLY the 10-batch +# (evaluator.viz_max_batches x 64) video budget, so each operator's +# episode is rendered whole whatever order the loader uses. +video_episodes: + valid: + - 692ea44dc621d7f4aac3aae6 + - 692eb08997ffd8ec90340290 + - 692fe79f34a99e18e25289ad diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml new file mode 100644 index 0000000000..226b603e2e --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml @@ -0,0 +1,23 @@ +# mecka_fold_flagship_top3_hpt_6d (the 3-operator flagship train split) with +# the wrist-frame KEYPOINT pipeline; see +# mecka_fold_flagship_opsplit_hpt_keypoints for the key layout. Norm stats: +# logs/hpt_architecture/norm_stats/mecka_fold_flagship_top3_hpt_keypoints/norm_stats +defaults: + - mecka_fold_flagship_top3_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + # The keypoint proprio is in the WRIST frame, where a history step + # would carry the wrist's own displacement into the articulation + # channels; stay at the parent's pre-history setting until the + # per-step framing lands. + proprio_history: 1 + history_stride: 1 + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml new file mode 100644 index 0000000000..063c9b5285 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml @@ -0,0 +1,74 @@ +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +# MECKA_FOLD_FLAGSHIP_TOPOP_HPT_6D: single-operator twin of +# mecka_fold_flagship_opsplit_hpt_6d. TRAIN on the ONE largest flagship +# fold-clothes operator only; everything else (pipeline, unseen-operator val, +# loader params) is inherited from the parent. +# +# SPLIT (from SQL app.episodes, 2026-09-15): +# train = operator 6903686e0e94ce070afd1f24, the largest of the 20 +# (405 eps / 76,013 frames / ~0.70 h at 30 fps, 24 % of +# flagship), 95 % of its episodes (~385 eps / ~72k frames). +# seen val = the other 5 % of that operator's episodes (20 eps / 3,731 +# frames) -> seen_op_valid/. The parent's valid_datasets interpolates +# the train filters and valid_ratio, so it follows this +# override automatically. +# train_viz = the train split above (trainHydra's default) -> train_viz/. +# unseen val = UNCHANGED from the parent: all episodes of the 2 +# `held_out_operators` -> unseen_op_valid/, directly comparable +# with the 18-operator run. +# The other 17 operators are unused. +# +# Hydra replaces lists wholesale, so the train filter_lambdas are restated in +# full; only the last (operator) lambda differs from the parent. The operator id +# is inlined rather than put in a config key: MultiDataModuleWrapper rejects +# root keys it does not declare. +# +# NORMALIZATION: the train distribution differs from the parent's, so +# precompute stats on THIS config: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_topop_hpt_6d --model hpt_bc_mecka_6d_300M --out +# then train with +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# data=mecka_fold_flagship_topop_hpt_6d norm_stats.precomputed_norm_path=/norm_stats + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # SINGLE-OPERATOR TRAIN: the largest flagship operator only. + - "lambda row: str(row.get('operator', '')) == '6903686e0e94ce070afd1f24'" + +# --------------------------------------------------------------------------- +# METRIC / VIDEO LOADERS: only the SEEN-op and train_viz entries differ from the +# parent (this config trains on one operator, so those two splits are smaller). +# --------------------------------------------------------------------------- +# K is the SINGLE-RANK value, floor(5120 / episodes-in-that-split); trainHydra +# scales it by the world size. Episode counts measured 2026-09-15: +# valid 20 eps / 3,731 frames -> K = floor(5120/20) = 256. Every +# episode here is shorter than 256 (longest 239f), so this keeps +# the WHOLE seen-val split: 3,731 frames = 58.3 of the 80 batches. +# train_viz 385 eps / 72,282 frames -> K = floor(5120/385) = 13 +# -> 385 x 13 = 5,005 frames = 78.2 batches. +# `unseen_op_valid: 26` is inherited unchanged (same split as the parent), and +# hydra merges dicts, so only the two changed keys are restated. +metric_frames_per_episode: + valid: 256 + train_viz: 13 + +# The parent's seen-val pin belongs to operator 690364170e94ce070afd1da5, which +# is not in this config's single-operator split, so it is replaced (hydra +# replaces lists wholesale). The train_viz pin 692e9bda21fc595fbe6f941e IS an +# episode of operator 6903686e..., hence in this train split too, and the +# unseen pins are the same split as the parent's -- both inherited. +# valid 692e83ad567c97525f4011a5 (6903686e0e94ce070afd1f24, 239f) +# the longest episode in this 20-episode seen-val split: 3.7 of the 10 +# video batches. No longer contiguous episode exists here. +video_episodes: + valid: + - 692e83ad567c97525f4011a5 diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml new file mode 100644 index 0000000000..90fd9c6cac --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml @@ -0,0 +1,96 @@ +# MECKA_FOLD_FREEFORM_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes only, stored under the human_bimanual embodiment: SQL +# task='folding_clothes', 462 eps / 1.13M frames under +# processed_v3/mecka/freeform/. "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts +# either fold-clothes task spelling (the flagship 'fold_clothes' rows fail the +# path check, so this resolves to the 462). +# +# RANDOM-EPISODE split (valid_ratio 0.05). For the held-out-OPERATOR split use +# mecka_fold_freeform_opsplit_6d instead. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose, so the world frame and the moving +# head frame cancel out of the targets; the head-frame proprio ee_pose is +# 6D-encoded too (it is the only head-frame signal the model sees). +# +# Same pipeline knobs as mecka_all_6d: fix_left_wrist_convention (processed_v3 +# mecka zarrs have the LEFT wrist frame double-mirrored; Rz(180)-corrected +# before any frame math), pad_proprio_gripper (18 -> 20-dim proprio so the +# State: prompt bins align with the robot layout), stride 1 (mecka is +# recorded at the target rate), prompts from the `annotations` track. +# Pairs with model=pi0.5_bc_mecka_6d and evaluator=eval_pi_wristframe_6d (the +# train-set spot-check head wraps it automatically). For CAM (head)-frame +# actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# Norm stats: precompute on a CPU node with +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_6d \ +# --model pi0.5_bc_mecka_6d --out +# and point training at norm_stats.precomputed_norm_path=/norm_stats +# (the script also fills the content-keyed cache, so a plain run hits it). + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# Train-set spot check for the TrainVizEvalVideo head (train_viz/ metrics, +# videos_train_viz/). No shuffle -> same leading slice every val +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: train + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml new file mode 100644 index 0000000000..e3b1396be3 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml @@ -0,0 +1,44 @@ +# MECKA_FOLD_FREEFORM_OPSCALE1_6D: the k=1-operator rung of the freeform +# fold-clothes OPERATOR-SCALING ladder. Everything -- pipeline, the annotation +# guard, the shared seen-op val, the 31-operator unseen val, the video pins, +# loader params -- is inherited from mecka_fold_freeform_opscale4_6d; only the +# train OPERATOR SET changes, so each rung only ADDS an operator. Edit the ladder +# only in the parent. +# +# TRAIN: ALL annotated episodes of the top 1 operator by annotated hours, +# minus the 3 shared seen-op val episodes -> +# 9 eps / 25633 f / 14.24 min = 0.237 h +# (1.00x the k=1 arm): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# Counts measured 2026-09-17; the filter is live, so they grow if SQL gains more +# episodes for these operators. +# +# NORM STATS: this arm's train distribution differs from every other rung -- +# precompute on THIS config, do not reuse the parent's: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale1_6d \ +# --model pi0.5_bc_mecka_6d --out + +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +# Hydra replaces lists wholesale, so the train lambdas are restated IN FULL; only +# the operator tuple differs from the parent. The resolver and `mode: total` are +# inherited untouched. +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see mecka_fold_freeform_opscale4_6d. + - "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 1 operator by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792',)" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml new file mode 100644 index 0000000000..094efc9dc5 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE1_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale1_6d -- the SAME k=1-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 9 eps / 25633 f / 14.24 min = 0.237 h across 1 operator +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale1_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale1_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml new file mode 100644 index 0000000000..ab7943eb4b --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml @@ -0,0 +1,45 @@ +# MECKA_FOLD_FREEFORM_OPSCALE2_6D: the k=2-operator rung of the freeform +# fold-clothes OPERATOR-SCALING ladder. Everything -- pipeline, the annotation +# guard, the shared seen-op val, the 31-operator unseen val, the video pins, +# loader params -- is inherited from mecka_fold_freeform_opscale4_6d; only the +# train OPERATOR SET changes, so each rung only ADDS an operator. Edit the ladder +# only in the parent. +# +# TRAIN: ALL annotated episodes of the top 2 operators by annotated hours, +# minus the 3 shared seen-op val episodes -> +# 24 eps / 60374 f / 33.54 min = 0.559 h +# (2.36x the k=1 arm): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# #2 68e0b875c33a1abcb8fc55b9 15 eps 34741 f = 19.30 min +# Counts measured 2026-09-17; the filter is live, so they grow if SQL gains more +# episodes for these operators. +# +# NORM STATS: this arm's train distribution differs from every other rung -- +# precompute on THIS config, do not reuse the parent's: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale2_6d \ +# --model pi0.5_bc_mecka_6d --out + +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +# Hydra replaces lists wholesale, so the train lambdas are restated IN FULL; only +# the operator tuple differs from the parent. The resolver and `mode: total` are +# inherited untouched. +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see mecka_fold_freeform_opscale4_6d. + - "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 2 operators by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792', '68e0b875c33a1abcb8fc55b9')" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml new file mode 100644 index 0000000000..c6fd633ed4 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE2_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale2_6d -- the SAME k=2-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 24 eps / 60374 f / 33.54 min = 0.559 h across 2 operators +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale2_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale2_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml new file mode 100644 index 0000000000..b6d6580247 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml @@ -0,0 +1,288 @@ +# MECKA_FOLD_FREEFORM_OPSCALE4_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes (SQL lab='mecka', task='folding_clothes', 462 eps / 1.13 M frames / +# 178 operators under processed_v3/mecka/freeform/; "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts either +# fold-clothes task spelling). A DIFFERENT set from the FLAGSHIP 'fold_clothes' +# episodes -- zero operator overlap between the two. +# Standalone (no defaults inheritance); the k=1/2 arms inherit THIS file. +# +# OPERATOR-SCALING LADDER (chosen 2026-09-17 from SQL app.episodes). The question +# is how many DIFFERENT OPERATORS the training data should come from. Each arm +# trains on ALL the annotated data its operators have, minus the held-out +# validation episodes -- the per-operator amount is NOT capped, so the total +# grows both with the operator count and with how much those operators recorded: +# +# k=1 9 eps 25633 f 14.24 min = 0.237 h ( 1.00x the k=1 arm) +# k=2 24 eps 60374 f 33.54 min = 0.559 h ( 2.36x the k=1 arm) +# k=4 48 eps 118388 f 65.77 min = 1.096 h ( 4.62x the k=1 arm) +# +# Per-operator contributions (k=4; the smaller arms are prefixes of this list): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# #2 68e0b875c33a1abcb8fc55b9 15 eps 34741 f = 19.30 min +# #3 695d09fc83a9fdf2d84d9c11 13 eps 30549 f = 16.97 min +# #4 69439ad5a2a8b5aee76602f5 11 eps 27465 f = 15.26 min +# +# So the arms are NOT a clean 1:2:4 in hours (1.00x, 2.36x, 4.62x) -- +# operator #1 is the smallest of the four after funding the val set. Operator +# count and data volume therefore move TOGETHER across this ladder; it measures +# "does adding another operator's data help", not operator diversity at a fixed +# budget. For the fixed-budget question, cap each operator instead. +# +# The three train operator sets are NESTED (k=1 subset k=2 subset k=4), so +# operator identity never changes underneath the ladder -- each rung only ADDS an +# operator. Because the split is an operator filter rather than an episode pin, +# that nesting cannot drift. +# +# LIVE FILTERS, NOT A SNAPSHOT: the train split is "every annotated episode of +# these operators", so it will GROW if more episodes for them land in SQL. The +# episode counts above were measured 2026-09-17 by hand and are documentation +# only -- nothing asserts them, because there is nothing to assert them against +# without the DB. tests/unit/test_opscale_ladder.py checks the parts that are +# in the config: the operator nesting, the exclusions and the metric budgets. +# +# ANNOTATED EPISODES ONLY. 112 of the 462 freeform episodes have segments = NULL, +# which becomes an empty `annotations` track and makes _build_prompts fall back to +# `default_prompt` ("" by default) for every one of their frames. Restricting to +# the 350 annotated episodes (7.91 h / 156 operators) makes the text stem see a +# real instruction on ~100 % of frames. The operator RANKING is by ANNOTATED +# hours, which is NOT the same order as raw hours -- do not re-derive it from +# unfiltered totals. +# +# THREE-WAY SPLIT, and both val sets are IDENTICAL in all three arms: +# train = all annotated episodes of the arm's operators, minus the +# 3 seen-op val episodes below. +# seen val = 3 episodes of operator #1 (6964d38183a9fdf2d8663792), +# 9745 f = 5.41 min -> seen_op_valid/. Operator #1 is in EVERY +# arm (the sets are nested), so this is the one seen-op val +# that is directly comparable across the whole ladder. These +# are the ONLY episodes withheld from training. +# unseen val = the ANNOTATED episodes of the 31 `held_out_operators` +# (35 of their 50 eps, 79850 f = 0.739 h, 26 of the 31 operators +# contribute) -> unseen_op_valid/. The operator LIST is reused +# verbatim from mecka_fold_freeform_opsplit_6d, so no unseen +# operator is ever trained on and the split is comparable with +# those checkpoints at the OPERATOR level; the annotation guard +# makes it a subset of that config's 50-episode val side, so it +# is NOT episode-identical. +# train_viz = the train split itself (no train_viz_datasets here, so +# trainHydra reuses train_datasets unshuffled). +# +# VAL LOADERS (three heads, one validation pass), as in the flagship opsplit +# configs: valid_datasets -> seen_op_valid/ (renamed by `valid_prefix`), +# unseen_op_valid_datasets -> unseen_op_valid/, train_viz -> train_viz/. Each is +# served by a subsampled METRIC loader and a pinned contiguous VIDEO loader; see +# `metric_frames_per_episode` / `video_episodes` at the bottom. +# +# PIPELINE: WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed relative to +# that wrist's current pose, so the world frame and the moving head frame cancel +# out of the targets; the head-frame proprio ee_pose is 6D-encoded too and padded +# 18 -> 20 (zero grip slots at 9/19) so the State: prompt bins align with the +# robot layout. fix_left_wrist_convention: processed_v3 mecka zarrs have the LEFT wrist +# frame double-mirrored; Rz(180)-corrected before any frame math. stride 1 (mecka +# is recorded at the target rate). Prompts come from the `annotations` track -- +# one segment label per frame, sampled by _build_prompts. +# +# Pairs with model=pi0.5_bc_mecka_6d and evaluator=eval_pi_wristframe_6d. For CAM +# (head)-frame actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# NORMALIZATION: every arm has a DIFFERENT train distribution, so precompute +# stats per arm -- do NOT share them across the ladder or reuse the opsplit / +# flagship stats: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale4_6d \ +# --model pi0.5_bc_mecka_6d --out +# then train with norm_stats.precomputed_norm_path=/norm_stats. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 31 UNSEEN-operator ids, reused verbatim from mecka_fold_freeform_opsplit_6d +# so no unseen operator is ever trained on and the split stays operator-level +# comparable with the existing freeform opsplit checkpoints. Verified disjoint +# from the 4 train operators. Referenced by the unseen filter via +# `data.held_out_operators`; MultiDataModuleWrapper declares this key. +held_out_operators: + - 68a8e70d1ffcf0417e08f0e5 + - 68e2fc680b158ea94ab4f5bc + - 6930bf8a6e496b45b1e1122e + - 6931784c6e496b45b1e1b1cd + - 693a399ec67c9e4814b0c09f + - 693b9530c67c9e4814b20ab2 + - 693c78165cb6c2828a4c2087 + - 694604945c0479a0fb7f48a4 + - 6948d23b5c0479a0fb8302b3 + - 69529a85c0d9448524bc7880 + - 6954f0b720b100982d812ed8 + - 6957464720b100982d83b76c + - 695d07d783a9fdf2d84d9371 + - 695d084183a9fdf2d84d955a + - 6961c50f83a9fdf2d85ec95d + - 69624a8683a9fdf2d86013c0 + - 6973c3ae89fa4bc311d1bbe5 + - 69842643f8ca4464d9920799 + - 698937a5f8ca4464d99dce94 + - 6992b19ab9d55c9c2d48432c + - 699429e18ef0252970048624 + - 6996d5eded1708063b2fe419 + - 6999d17c8f173e14271f9890 + - 69a24256c4216c3e3ce33834 + - 69a39ec9c4216c3e3ce49734 + - 69a57a78c4216c3e3ce6323c + - 69a725a347e12c9b3390ba33 + - 69a7943547e12c9b33911c02 + - 69a84eed47e12c9b339208fd + - 69ab87207bf9cbce29a3cb36 + - 69ae2463d22d99655253f287 + +# Metric prefix / video dir for the canonical valid loader (config-only, read by +# trainHydra), so the SEEN-operator val is not logged under a bare Valid/. +valid_prefix: seen_op_valid + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + # No proprio_history: the keymap's K and the model's stem history_len are + # one setting in two files and HPTModel raises when they disagree, so a + # rung that asks for K > 1 only composes with a model configured for the + # same K. Raise both together (and see PoseCoordinateFrameTransform's + # per_step_target, which the keypoint twins rely on above K=1). + # Language: the per-frame segment label track. Every episode in every + # split of this config is annotation-guarded, so this is never empty. + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # LANGUAGE ANNOTATIONS REQUIRED. 112 of the 462 freeform episodes have a + # NULL `segments` column, which becomes an EMPTY `annotations` track in the + # zarr; _annotation_text_for_frame then returns [] for every frame of them + # and _build_prompts silently substitutes `default_prompt` (an EMPTY string + # unless overridden), so they would train the text stem on nothing. This + # guard keeps the 350 annotated episodes, whose segments tile 99.7-100 % of + # their frames. `segments` is a list (or None) on the SQL row, so isinstance + # is the safe test -- do NOT use truthiness (a numpy round-trip makes + # `bool(...)` raise "truth value of an array is ambiguous"). + # Anchored so valid/unseen reuse this exact lambda. + - &annotated "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 4 operators by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792', '68e0b875c33a1abcb8fc55b9', '695d09fc83a9fdf2d84d9c11', '69439ad5a2a8b5aee76602f5')" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" + # The filters ARE the split; no random episode split on top of them. + mode: total + +# SEEN operators, held-out episodes -> seen_op_valid/ metrics. 3 episodes of +# operator #1, which is in every arm of the ladder, so this val set is shared and +# comparable across all three rungs. Same resolver as train (so it follows any +# pipeline override); pinned by hash, and `mode: total` because the pin is the +# split. The train filter excludes exactly these 3 hashes. +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see train_datasets. + - *annotated + # op #1 6964d38183a9fdf2d8663792: 3 eps, 9745 f = 5.41 min + episode_hashes: + - 696bcbe0e13d37dd4293f58b + - 696bcbe86377d0871b6a9253 + - 696bcc71ed8ed158122ec794 + mode: total + +# UNSEEN operators: the ANNOTATED episodes of the 31 held-out operators -> +# unseen_op_valid/ metrics via a second copy of the canonical evaluator. +# Identical in all three arms. +unseen_op_valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see train_datasets. + - *annotated + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +# Unshuffled (both val loaders): the same leading slice every val, comparable +# across epochs with a coherent overlay video. The train_viz loader reuses these. +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +unseen_op_valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +# --------------------------------------------------------------------------- +# METRIC LOADERS: K evenly spaced frames per EPISODE, so every episode and +# operator is scored inside the same limit_val_batches window (80 x 64 = 5120 +# frames per rank on the default trainer). `auto` sets +# K = floor(limit_val_batches * batch_size / episodes-in-split) from the RESOLVED +# split at startup, so K follows the live filters as SQL grows and needs no +# per-rung override. At the 2026-09-17 counts that is: +# valid 3 eps / 9745 f -> K = floor(5120/3) = 1706 +# unseen_op_valid 35 eps / 79850 f -> K = floor(5120/35) = 146 +# train_viz 48 eps / 118388 f -> K = floor(5120/48) = 106 (k=4) +# 24 eps / 60374 f -> K = 213 (k=2) +# 9 eps / 25633 f -> K = 568 (k=1) +metric_frames_per_episode: + valid: auto + unseen_op_valid: auto + train_viz: auto + +# --------------------------------------------------------------------------- +# VIDEO LOADERS: pinned contiguous episodes, so every run in the ladder renders +# the SAME episodes. Taken from that head's already-resolved split (no second SQL +# pull); `evaluator.viz_max_batches` (10 batches = 640 frames) caps them. Each +# pin is the LONGEST episode of its split, so one contiguous episode fills the +# whole video budget. Format (operator, frames, segments): +# valid 696bcbe86377d0871b6a9253 (6964d38183a9fdf2d8663792, 3688 f, 8 seg) +# unseen_op_valid 696bc7e021d8e8ab8ddc8380 (6957464720b100982d83b76c, 3358 f, 10 seg) +# train_viz 696dd090a56f03a27e59f69c (6964d38183a9fdf2d8663792, 3298 f, 8 seg) +# -- an operator #1 episode, so it is in EVERY arm's train +# split and the children inherit this pin unchanged. +video_episodes: + valid: + - 696bcbe86377d0871b6a9253 + unseen_op_valid: + - 696bc7e021d8e8ab8ddc8380 + train_viz: + - 696dd090a56f03a27e59f69c diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml new file mode 100644 index 0000000000..cc67215aa5 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE4_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale4_6d -- the SAME k=4-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 48 eps / 118388 f / 65.77 min = 1.096 h across 4 operators +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale4_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml new file mode 100644 index 0000000000..0eacf0e873 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml @@ -0,0 +1,151 @@ +# MECKA_FOLD_FREEFORM_OPSPLIT_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes (SQL task='folding_clothes', 462 eps / 1.13M frames under +# processed_v3/mecka/freeform/; "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts +# either fold-clothes task spelling), stored under the human_bimanual +# embodiment, with the held-out-OPERATOR train/val split. +# +# OPERATOR SPLIT: validation = 31 operators held out ENTIRELY, so no val +# episode shares an operator with training. The list lives ONCE, in +# `held_out_operators` below; the train and valid filter lambdas both +# interpolate it, so the two sides cannot drift apart. The HPT keypoint twin +# (mecka_fold_freeform_opsplit_hpt_keypoints) inherit this file. +# +# Chosen deterministically on 2026-09-01 from SQL app.episodes (462 freeform +# fold eps / 178 operators): skip the 10 largest operators (their data stays in +# train), seeded shuffle (random.Random(42)) of the remaining operator ids, +# accumulate until >= 10 % of frames -> 31 operators, 50 episodes, +# 115,681 frames (10.2 %) held out; train = 147 operators, 412 episodes, +# 1,013,220 frames. Both sides use mode: total (no random episode split); the +# list IS the split. Do not edit it: checkpoints were trained against it. +# +# NO ANNOTATION GUARD: 112 of the 462 freeform episodes have segments = NULL, so +# their `annotations` track is empty and the prompt falls back to +# `default_prompt` ('' by default) on every frame: Pi sees "Task: , State: ...", +# HPT feeds '' to its text stem. Left as is because checkpoints were trained on +# this split; the opscale ladder (mecka_fold_freeform_opscale4_6d) reuses only +# its held-out OPERATOR list and filters to annotated episodes. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose; the head-frame proprio ee_pose is +# 6D-encoded too. Pairs with model=pi0.5_bc_mecka_6d. +# +# Shared pipeline knobs (same as mecka_all_6d): fix_left_wrist_convention +# (processed_v3 mecka zarrs have the LEFT wrist frame double-mirrored; +# Rz(180)-corrected before any frame math), pad_proprio_gripper (18 -> 20-dim +# proprio so the State: prompt bins align with the robot layout), stride 1, +# prompts from the `annotations` track. Pairs with evaluator=eval_pi_wristframe_6d +# (the train_viz head wraps it automatically). For CAM (head)-frame actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# Norm stats: precompute on a CPU node +# (egomimic/scripts/precompute_norm_stats.py) and point training at +# norm_stats.precomputed_norm_path=/norm_stats. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 31 held-out operator ids (SQL app.episodes.operator), sorted. See the +# OPERATOR SPLIT note above. Referenced by the filter lambdas via the mounted +# path `data.held_out_operators`. MultiDataModuleWrapper accepts this key +# (config-only). +held_out_operators: + - 68a8e70d1ffcf0417e08f0e5 + - 68e2fc680b158ea94ab4f5bc + - 6930bf8a6e496b45b1e1122e + - 6931784c6e496b45b1e1b1cd + - 693a399ec67c9e4814b0c09f + - 693b9530c67c9e4814b20ab2 + - 693c78165cb6c2828a4c2087 + - 694604945c0479a0fb7f48a4 + - 6948d23b5c0479a0fb8302b3 + - 69529a85c0d9448524bc7880 + - 6954f0b720b100982d812ed8 + - 6957464720b100982d83b76c + - 695d07d783a9fdf2d84d9371 + - 695d084183a9fdf2d84d955a + - 6961c50f83a9fdf2d85ec95d + - 69624a8683a9fdf2d86013c0 + - 6973c3ae89fa4bc311d1bbe5 + - 69842643f8ca4464d9920799 + - 698937a5f8ca4464d99dce94 + - 6992b19ab9d55c9c2d48432c + - 699429e18ef0252970048624 + - 6996d5eded1708063b2fe419 + - 6999d17c8f173e14271f9890 + - 69a24256c4216c3e3ce33834 + - 69a39ec9c4216c3e3ce49734 + - 69a57a78c4216c3e3ce6323c + - 69a725a347e12c9b3390ba33 + - 69a7943547e12c9b33911c02 + - 69a84eed47e12c9b339208fd + - 69ab87207bf9cbce29a3cb36 + - 69ae2463d22d99655253f287 + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: train = every operator NOT in the held-out list. + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + mode: total + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: valid = the held-out operators only. + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +# Train-set spot check for the TrainVizEvalVideo head (train_viz/ metrics, +# videos_train_viz/): same operators as train, no shuffle -> same slice every val +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: total + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml new file mode 100644 index 0000000000..66ecee19bb --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml @@ -0,0 +1,51 @@ +defaults: + - mecka_fold_freeform_opsplit_6d + - _self_ + +# MECKA_FOLD_FREEFORM_OPSPLIT_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opsplit_6d: the SAME mecka FREEFORM fold-clothes +# episodes and the SAME held-out-OPERATOR split (`held_out_operators`, the +# filters, the train_viz slice and the loader params are all inherited; edit +# the operator list only in the parent). Only the keymap and transform differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations task-description track -> the +# `sampled_prompt` for a Qwen text stem +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no +# stem for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: +# in keypoint modes it corrects the left `obs/action_wrist_pose` keys, which +# carry the same double-mirrored rotation as the ee_pose on processed_v3. +# +# NORMALIZATION: proprio + action keys are quantile-normalized; precompute on +# a CPU node with +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opsplit_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist (the inherited train_viz slice gets the +# default head wrapping it); see train_zarr_mecka_kp_wrist_hpt.yaml for the +# composed recipe. + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 81a486e8cc..5b97388514 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -1,3 +1,8 @@ +# One pinned scale episode (stored as human_bimanual) on the 144-D wrist-frame +# hand-keypoint action; see data/human.yaml for the layout and the cartesian +# override. fix_left_wrist_convention is on: scale runs its own converter, +# but it writes the same pre-fix left-wrist convention as mecka (measured on +# all 20 production episodes). _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,10 +13,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # Scale exports were written by the pre-fix hand-pose converter too: # measured on all 20 production episodes, the left wrist matches # rot_left to <= 2.06 deg with the right-hand control clean. diff --git a/egomimic/hydra_configs/evaluator/eval_hpt.yaml b/egomimic/hydra_configs/evaluator/eval_hpt.yaml index 901be30749..55e0acf4ec 100644 --- a/egomimic/hydra_configs/evaluator/eval_hpt.yaml +++ b/egomimic/hydra_configs/evaluator/eval_hpt.yaml @@ -1,5 +1,21 @@ +# Default HPT evaluator, matching the per-vendor data defaults: eva on the +# wrist-frame cartesian 6D action (reverted to cam-frame ypr for the overlay), +# human on the 144-D wrist-frame hand-keypoint action (reverted to head-frame +# keypoints for the overlay). +# For cartesian human data use eval_hpt_cartesian / eval_hpt_wrist. defaults: - - viz@viz_func: cartesian + - viz@viz_func: keypoints_wrist - _self_ _target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation run. +viz_every_n_epochs: 200 + +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml new file mode 100644 index 0000000000..17a3ef6d3d --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml @@ -0,0 +1,11 @@ +# HPT evaluator for cam/head-frame CARTESIAN ypr data on every embodiment +# (the pre-keypoint default; e.g. model=hpt_bc_flow_human_cartesian with a +# cartesian data override). No revert: the actions are already in the frame +# the overlay draws in. +defaults: + - viz@viz_func: cartesian + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 200 diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml new file mode 100644 index 0000000000..abb99f2f8f --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml @@ -0,0 +1,17 @@ +# HPT head-frame keypoint evaluator: use with a data config that applies the +# `keypoints_headframe_6d` transform mode (e.g. data=aria_keypoints). The +# keypoints are already in head frame; the revert only converts the wrist +# rotation rot6d -> ypr (144 -> 138) so the overlay's keypoint splitter sees +# the ypr layout. +defaults: + - viz@viz_func: keypoints + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..1549a3d080 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml @@ -0,0 +1,27 @@ +# HPT wrist-frame keypoint evaluator: use with a data config that applies the +# `keypoints_wristframe_6d` transform mode (e.g. +# mecka_fold_freeform_opsplit_hpt_keypoints, or the per-vendor defaults). The +# model predicts the 144-D wrist-first action [wrist xyz+rot6d (9) | 21 +# keypoints (63)] x 2 in each hand's wrist frame; the revert converts the +# wrist rotation back to ypr and projects the two 63-D keypoint chunks back to +# head (camera) frame via the head-frame wrist pose in +# `observations.state.keypoints`, giving a 126-D head-frame +# `actions_keypoints` for the cam-frame keypoint MSE +# (Valid/human_bimanual_actions_keypoints_cam_{paired,final}_mse_avg) and the +# overlay video. Native-frame metrics (paired/final MSE, Frechet, the +# keypoint layout metrics) run on the full 144-D. +defaults: + - viz@viz_func: keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..ad998e6498 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml @@ -0,0 +1,24 @@ +# HPT wrist-frame 6D CARTESIAN evaluator: HPT twin of eval_pi_wristframe_6d. +# Use with a data config that applies the `cartesian_wristframe_6d` transform +# mode with `keymap_mode: cartesian` (e.g. mecka_fold_flagship_opsplit_hpt_6d). +# The model predicts the 18-D [xyz | 6D rot] x 2 action in each wrist's frame; +# the revert converts 6D -> ypr (Gram-Schmidt), unpads the 20-D proprio, and +# projects predictions + gt back to head (cam) frame, giving the cam-frame MSE +# (Valid/human_bimanual_actions_cartesian_cam_{paired,final}_mse_avg) and the +# overlay video. Native-frame metrics (paired/final MSE, Frechet, xyz / +# rotation errors) run on the 18-D output. +defaults: + - viz@viz_func: cartesian_wrist_lang + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi.yaml b/egomimic/hydra_configs/evaluator/eval_pi.yaml index 97b9c1a1f3..62fdd2c39b 100644 --- a/egomimic/hydra_configs/evaluator/eval_pi.yaml +++ b/egomimic/hydra_configs/evaluator/eval_pi.yaml @@ -1,18 +1,30 @@ +# Default Pi evaluator, matching the cotrain_pi_base data defaults: eva on the +# wrist-frame cartesian 6D action, human on the 144-D wrist-frame hand-keypoint +# action. Metrics are computed on the native output (frame-invariant; see +# egomimic/eval/action_metrics.py) and never go through the reverts below, +# which only feed the viz overlay: +# eva: xyz+6D -> xyz+ypr, then each arm's chunk wrist -> cam frame via +# the cam-frame proprio ee_pose +# human: wrist rot6d -> ypr, then each hand's keypoints wrist -> head frame +# via the head-frame wrist pose in observations.state.keypoints +# For the all-cartesian recipes use eval_pi_camframe_6d / eval_pi_wristframe_6d. defaults: - - viz@viz_func: pi_cartesian_lang + - viz@viz_func: pi_keypoints_lang_wrist - _self_ _target_: egomimic.eval.eval_pi.PIEvalVideo -# Per-embodiment revert transform. Applied once during validation to project -# the model's wrist-frame action chunks back to cam (head) frame, then reused -# for both the cam-frame MSE and the viz video. Each value resolves to a -# list[Transform] via its ``_target_``. Must match whichever viz config is -# mounted above. +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. Without it the EvalVideo default of 1 renders video on every pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame CPU +# draw); metrics still run on every batch. +viz_max_batches: 10 + transform_lists: eva_bimanual: - _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list human_bimanual: - _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml new file mode 100644 index 0000000000..eceaf649de --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +# All-cartesian cam-frame evaluator: use when every embodiment expresses its +# cartesian action in head/camera frame with the continuous 6D rotation +# (`transform_list: mode: cartesian_6d`, e.g. data=mecka_fold_freeform_6d +# with data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d). +# Actions are already in cam (head) frame, so +# no frame change is needed; the revert only converts the rotation back from +# xyz+6D (9/arm) to xyz+ypr (6/arm) for the viz overlay, which draws in ypr. +# Metrics are computed on the native 6D output and never go through this +# revert. Each value resolves to a list[Transform] via its ``_target_``. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml new file mode 100644 index 0000000000..62d092ddc1 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml @@ -0,0 +1,17 @@ +# Wrist-frame twin of eval_pi (data=cotrain_pi_lang_wrist): eva on the +# wrist-frame cartesian 6D action (reverted 6D -> ypr, then wrist -> cam frame +# via the proprio ee_pose), human on the wrist-frame keypoint action. +defaults: + - viz@viz_func: pi_keypoints_lang_wrist + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..581d454fc0 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this (see eval_pi_camframe_6d.yaml); without it the EvalVideo default of 1 +# renders video on every validation pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame +# CPU draw; unbounded rendering added ~2h per viz epoch on the combined +# pretrain). Metrics still run on every batch. The train_viz wrapper +# inherits this cap. +viz_max_batches: 10 + +# All-cartesian wrist-frame evaluator: use with a data config that applies the +# `cartesian_wristframe_6d` transform mode (e.g. mecka_all_6d, +# mecka_fold_freeform*_6d). The model predicts in each wrist's local frame +# using the continuous 6D rotation; these revert transforms first convert the +# rotation xyz+6D -> xyz+ypr and then project predictions + gt back to cam +# (head) frame for the viz video only (metrics are frame-invariant and +# computed on the native output). Each value resolves to a list[Transform] +# via its ``_target_``. Must match the viz config mounted above. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..d9631fefd5 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml @@ -0,0 +1,23 @@ +# Keypoint HPT twin of train_viz_hpt_wristframe_6d: wraps +# eval_hpt_keypoints_wrist so the same forward/metric/viz logic runs a second +# time against the `train_viz` dataloader (metrics `train_viz/`, videos +# `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_hpt_keypoints_wrist` + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_hpt_keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..3cc9e4b13b --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml @@ -0,0 +1,24 @@ +# Wraps eval_hpt_wristframe_6d so the same forward/metric/viz logic runs a +# second time against the `train_viz` dataloader (metrics `train_viz/`, videos +# `videos_train_viz/`). HPT twin of train_viz_pi_wristframe_6d. +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_hpt_wristframe_6d` +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_hpt_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml new file mode 100644 index 0000000000..9f1bdbd797 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml @@ -0,0 +1,25 @@ +# Wraps the default Pi evaluator (eval_pi) so the same forward/metric/viz +# logic runs a second time against the `train_viz` dataloader. The wrapped +# `base` evaluator owns compute_metrics_and_viz; this wrapper writes videos to +# `videos_train_viz/` and prefixes metric keys with `train_viz/`. +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi` +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml new file mode 100644 index 0000000000..5eef381d55 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml @@ -0,0 +1,26 @@ +# Cam-frame all-cartesian twin of train_viz_pi: wraps eval_pi_camframe_6d so +# the same forward/metric/viz logic runs a second time against the +# `train_viz` dataloader (metrics `train_viz/`, videos `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi_camframe_6d` +# Needs the human transform mode overridden to `cartesian_6d`. + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi_camframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..15863c7cee --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml @@ -0,0 +1,25 @@ +# Wraps the all-cartesian wrist-frame Pi evaluator (eval_pi_wristframe_6d) so +# the same forward/metric/viz logic runs a second time against the +# `train_viz` dataloader (metrics `train_viz/`, videos `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi_wristframe_6d` + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml index e4aae52a72..9863646876 100644 --- a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml @@ -1,12 +1,12 @@ eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml index cf5d1d218c..8f000e5ca7 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml @@ -1,12 +1,15 @@ +# Head-frame keypoint overlay: eva draws its cartesian trajectory, human its +# MANO keypoints (the 138-D ypr layout after the evaluator's 6d -> ypr revert, +# or 126-D head-frame keypoints after the wrist-frame revert). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml index cf5d1d218c..2f92baf889 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml @@ -1,12 +1,14 @@ +# Wrist-frame keypoint overlay (same drawing as keypoints.yaml; the evaluator's +# transform_lists revert the wrist-frame prediction to head frame first). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml new file mode 100644 index 0000000000..65d79d73a4 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml @@ -0,0 +1,10 @@ +# Pi twin of keypoints_wrist: eva cartesian trajectory + human MANO keypoint +# overlay with the sampled prompt drawn on the frame. +defaults: + - keypoints_wrist + - _self_ + +eva_bimanual: + annotation_key: sampled_prompt +human_bimanual: + annotation_key: sampled_prompt diff --git a/egomimic/hydra_configs/hydra/launcher/submitit.yaml b/egomimic/hydra_configs/hydra/launcher/submitit.yaml index c56f2cd508..f154af112e 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit.yaml @@ -10,7 +10,7 @@ account: "rl2-lab" # Slurm account (e.g., 'my_accou cpus_per_task: 12 # Number of CPUs per task nodes: ${launch_params.nodes} # Number of nodes tasks_per_node: ${launch_params.gpus_per_node} # Use variable for tasks per node -gres: "gpu:a40:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count +gres: "gpu:a40:${launch_params.gpus_per_node}" # GPU type and count per node qos: "short" # Slurm QoS timeout_min: 2880 # Timeout in minutes (48 hours) exclude: "protocol, puma" # Nodes to exclude diff --git a/egomimic/hydra_configs/hydra/launcher/submitit_loaner.yaml b/egomimic/hydra_configs/hydra/launcher/submitit_loaner.yaml new file mode 100644 index 0000000000..9d443e4986 --- /dev/null +++ b/egomimic/hydra_configs/hydra/launcher/submitit_loaner.yaml @@ -0,0 +1,28 @@ +defaults: + - submitit_slurm + +_target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher + +# Loaner H100 cluster: 5 bare-metal nodes, 8x H100 80GB, 224 logical CPUs, +# ~2 TiB RAM per node. Nodes are allocated exclusively to one job. +name: ${hydra.job.name} +partition: "loaner" +account: "loaner" + +# 224 logical CPUs / 8 tasks-per-node (one task per GPU) = 28 +cpus_per_task: 28 + +nodes: ${launch_params.nodes} +tasks_per_node: ${launch_params.gpus_per_node} + +# NOTE: Slurm --gres is PER NODE, so this is deliberately NOT multiplied by +# ${launch_params.nodes} the way submitit.yaml / submitit_pace.yaml do. Those +# over-request whenever nodes > 1; here that would exceed the 8 GPUs a node has. +gres: "gpu:h100:${launch_params.gpus_per_node}" + +mem_per_gpu: 200G +timeout_min: 2880 # 48h + +# This partition has no QoS and no nodes worth excluding. +additional_parameters: + requeue: true diff --git a/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml b/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml index ca724575bf..e61539521e 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml @@ -10,7 +10,7 @@ account: "gts-dxu345-rl2" # Slurm account cpus_per_task: 8 # Number of CPUs per task (max 4:1 CPU:GPU ratio) nodes: ${launch_params.nodes} # Number of nodes tasks_per_node: ${launch_params.gpus_per_node} # Use variable for tasks per node -gres: "gpu:h200:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count (h100 for H100 GPUs) +gres: "gpu:h200:${launch_params.gpus_per_node}" # GPU type and count per node (h100 for H100 GPUs) qos: "inferno" # Slurm QoS mem_per_gpu: 250G timeout_min: 2880 # Timeout in minutes (48 hours) diff --git a/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml b/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml index ca83317d8a..35bcdaae03 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml @@ -11,7 +11,7 @@ cpus_per_task: 12 nodes: ${launch_params.nodes} tasks_per_node: ${launch_params.gpus_per_node} -gres: "gpu:l40s:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" +gres: "gpu:l40s:${launch_params.gpus_per_node}" qos: "short" timeout_min: 2880 # 80 hours diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml index eb345dd1e9..3389976760 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 16 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (aria, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/aria.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml index 196dc7a716..6440e4dc08 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml @@ -32,7 +32,6 @@ robomimic_model: multitask: false pretrained: false pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 domains: ["eva_bimanual"] shared_obs_keys: ["front_img_1"] diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml index 21c3f7fb01..227cc40e1d 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml @@ -1,23 +1,25 @@ +# HPT flow-matching BC on human data with the 144-D wrist-frame hand-keypoint +# action (`keypoints_wristframe_6d`): per hand [wrist xyz | wrist rot6d | 21 +# MANO keypoints] in that wrist's frame. Proprio = the 144-D head-frame +# `observations.state.keypoints` through the `state_keypoints` stem (the +# ee_pose a data config may also emit for the pi0.5 prompt has no stem here +# and is ignored). There is no vendor notion at the embodiment level, so this +# is the canonical human keypoint config; hpt_bc_flow_{mecka,aria,scale} are +# aliases of it, one per data config. For the cartesian action pair a +# cartesian data override with model=hpt_bc_flow_human_cartesian. _target_: egomimic.pl_utils.pl_model.ModelWrapper robomimic_model: _target_: egomimic.algo.hpt.HPT - # camera_transforms is a dead HPT.__init__ param after refactor 80e5523a - # ("move intrinsics/extrinsics into embodiment classes" — CameraTransforms was - # removed; HPT never stored/used it). Per-episode intrinsics now travel in - # batch["intrinsics"] (zarr.json) and viz reads them via the embodiment classes. - # HPT still takes the positional arg, so pass an empty mapping (it is ignored). camera_transforms: {} dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} + human_bimanual: {proprio: 144, action: 144} - diffusion: true + diffusion: true 6dof: true ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" trunk: embed_dim: 256 @@ -30,11 +32,10 @@ robomimic_model: use_domain_embedding: true drop_path: 0.1 weight_init_style: "pytorch" - + multitask: false pretrained: false - pretrained_checkpoint: "" - reverse_kl_samples: 8 + pretrained_checkpoint: "" # TODO domains: ["human_bimanual"] shared_obs_keys: ["front_img_1"] @@ -53,10 +54,10 @@ robomimic_model: crossattn_dim_head: 64 crossattn_modality_dropout: 0.1 modality_embed_dim: 256 - + stem_specs: human_bimanual: - state_ee_pose: + state_keypoints: _target_: egomimic.models.hpt_nets.MLPPolicyStem input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 256 @@ -71,7 +72,7 @@ robomimic_model: modality_embed_dim: 256 head_specs: - human_bimanual: + human_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml new file mode 100644 index 0000000000..bdb2775b90 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml @@ -0,0 +1,134 @@ +# HPT flow-matching BC on human data with the CARTESIAN head-frame action +# (18-D xyz+rot6d per arm pair) and the 20-D gripper-padded ee_pose proprio. +# The repo default for human data is the 144-D hand-keypoint action +# (hpt_bc_flow_human); use this with a cartesian data override, e.g. +# data=mecka data.train_datasets.human_bimanual.resolver.key_map.keymap_mode=cartesian \ +# data.train_datasets.human_bimanual.resolver.key_map.include_ee_pose=false \ +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d \ +# data.train_datasets.human_bimanual.resolver.transform_list.include_ee_pose=false +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + # camera_transforms is a dead HPT.__init__ param after refactor 80e5523a + # ("move intrinsics/extrinsics into embodiment classes" — CameraTransforms was + # removed; HPT never stored/used it). Per-episode intrinsics now travel in + # batch["intrinsics"] (zarr.json) and viz reads them via the embodiment classes. + # HPT still takes the positional arg, so pass an empty mapping (it is ignored). + camera_transforms: {} + + # Per-embodiment widths, declared once. Stems and heads below interpolate + # these; override e.g. model.robomimic_model.dims..proprio=16. + dims: + human_bimanual: {proprio: 20, action: 18} + + diffusion: true + 6dof: true + + ac_keys: + human_bimanual: "actions_cartesian" + + trunk: + embed_dim: 256 + num_blocks: 16 + num_heads: 8 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: "" + + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1"] + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 256 + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 + + stem_specs: + human_bimanual: + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 256 + hidden_dim: 128 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 4 + dropout: 0.1 + mlp_layers: 4 + mlp_ratio: 4 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 256 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 3e-4 + weight_decay: 0.0001 + +scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + _partial_: true + T_max: 1400 + eta_min: 1e-5 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml index eb345dd1e9..30c9af5165 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 16 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (mecka, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/mecka.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml index 6ae2b6d1d4..68235751ab 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 64 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (scale, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/scale.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml index 51dfa890bc..9ed1ae6a03 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml @@ -3,14 +3,16 @@ robomimic_model: _target_: egomimic.algo.hpt.HPT camera_transforms: {} + # 144 = per hand [wrist xyz (3) | wrist rot6d (6) | 21 MANO keypoints (63)] + # x 2, the `keypoints_*_6d` transform modes (the 138-D ypr layout of the + # `keypoints_*_ypr` modes is no longer the default). dims: - human_bimanual: {proprio: 138, action: 138} + human_bimanual: {proprio: 144, action: 144} ac_keys: human_bimanual: "actions_keypoints" diffusion: true - reverse_kl_samples: 8 trunk: embed_dim: 256 diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml new file mode 100644 index 0000000000..c992e14b64 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -0,0 +1,177 @@ +# HPT_BC_KEYPOINTS_WRIST_300M: ~300M-parameter HPT flow-matching BC on the +# 144-D wrist-frame MANO keypoint action, with a frozen Qwen3-Embedding text +# stem for the task prompt ("300M hpt qwen"). +# +# Data: mecka_fold_freeform_opsplit_hpt_keypoints; action `actions_keypoints`, +# per hand [wrist xyz+rot6d in its own wrist frame (9) | 21 keypoints in that +# wrist frame (63)] x {left, right}, quantile-normalized from precomputed +# stats. Inputs: front image (ResNet), the 144-D +# `observations.state.keypoints` (wrist pose in head frame + keypoints in +# wrist frame, per hand), and the sampled annotation through the Qwen +# per-token stem. Drop the text stem with +# ~model.robomimic_model.shared_stem_specs.annotation model.robomimic_model.shared_obs_keys=[front_img_1] +# +# Widths follow the lab's 300M reference recipe (hpt_bc_flow_mecka_300M.yaml on +# aniketh/hpt-run-configs; also the basis of the arc-E1 300M cotrain configs): +# embed_dim 840, 10 trunk heads, cross-attn 18 latents x 10 heads x 140 +# dim_head, head hidden_dim 320 / 5 heads / mlp 5x5. embed_dim is the shared +# bus width: every stem in/out, modality_embed_dim, encoder output and the +# head cond_dim move with it. num_blocks is set so the TRAINABLE count +# (frozen Qwen excluded) lands at ~300M (23 blocks ~= 298M measured on the +# 138-D layout; the 144-D stem/head add < 0.1M). +# +# Constant LR 5e-5, no scheduler (the setting of the mecka_fold_*_300M +# constlr_5e5 ablation whose keypoint arm fit best). +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + camera_transforms: {} + dims: + human_bimanual: {proprio: 144, action: 144} + ac_keys: + human_bimanual: "actions_keypoints" + + diffusion: true + + # Prompt for the Qwen stem: one annotation sampled per batch item from the + # `annotations` track (the data config's keymap exposes it). + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + # Mask the repeat-last padding at every episode / annotation tail out of + # the action loss (the dataset's `action_pad_mask`). + use_pad_mask: false + + trunk: + embed_dim: 840 + num_blocks: 23 + num_heads: 10 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 840 + hidden_dim: 320 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 5 + dropout: 0.1 + mlp_layers: 5 + mlp_ratio: 5 + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 840 + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + # The frozen encoder's per-token features are cached per unique prompt + # string (LRU, in memory), so the 600M-parameter forward runs once per + # annotation instead of once per training step. Ignored when + # `freeze: false` - a trainable encoder must stay online. + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: 840 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + stem_specs: + human_bimanual: + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 840 + widths: [840] + # Proprio observation history (as hpt_bc_mecka_6d_300M's state_ee_pose + # stem): K tokens, learned zero-init time-step embedding, whole-past + # dropout. Must equal the data config's `proprio_history`; inert at 1. + history_len: 1 + history_dropout: 0.2 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 840 + num_of_copy: 1 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-5 + weight_decay: 0.0001 + +scheduler: null diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml new file mode 100644 index 0000000000..894b747092 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -0,0 +1,202 @@ +# HPT_BC_MECKA_6D_300M: ~300M-parameter HPT flow-matching BC on the 18-D +# WRIST-frame cartesian action with the continuous 6D rotation, plus a frozen +# Qwen3-Embedding text stem for the task prompt. HPT twin of pi0.5_bc_mecka_6d. +# +# Data: a `cartesian_wristframe_6d` data config with `keymap_mode: cartesian`, +# e.g. mecka_fold_flagship_opsplit_hpt_6d. Action `actions_cartesian` = per arm +# [xyz | 6D rotation] (18) relative to that wrist's current pose; proprio +# `observations.state.ee_pose` = head-frame wrist poses, 6D-encoded and padded +# 18 -> 20 (pad_proprio_gripper; zero grip slots at 9/19). Both are +# quantile-normalized from precomputed stats. Inputs: front image (ResNet), the +# 20-D ee_pose through the `state_ee_pose` stem, and one annotation sampled per +# batch item through the Qwen per-token stem. Drop the text stem with +# ~model.robomimic_model.shared_stem_specs.annotation model.robomimic_model.shared_obs_keys=[front_img_1] +# +# Widths follow the lab's 300M reference (hpt_bc_flow_mecka_300M on +# aniketh/hpt-run-configs): embed_dim 840, 10 trunk heads, cross-attn 18 +# latents x 10 heads x 140 dim_head, head hidden_dim 320 / 5 heads / mlp 5x5. +# num_blocks is set so the TRAINABLE count (frozen Qwen excluded) lands at +# ~300M: 23 blocks = 298.1M trainable (trunk 195.0M = 8.48M/block, head 74.2M, +# ResNet 11.6M, image stem 6.1M, Qwen stem projection 5.6M, ee_pose stem 5.4M; +# Qwen 595.8M frozen). 24 blocks would be 306.5M. +# +# Constant LR 5e-5, no scheduler (the setting of the mecka_fold_*_300M +# constlr_5e5 ablations; the CosineAnnealingLR of hpt_bc_flow_mecka sawtooths +# over a long run). +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + camera_transforms: {} + dims: + human_bimanual: {proprio: 20, action: 18} + ac_keys: + human_bimanual: "actions_cartesian" + + diffusion: true + 6dof: true + + # Prompt for the Qwen stem: one annotation sampled per batch item from the + # `annotations` track (the data config's keymap exposes it). + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + # Mask the repeat-last padding at every episode / annotation tail out of + # the action loss (the dataset's `action_pad_mask`). + use_pad_mask: false + + trunk: + embed_dim: 840 + num_blocks: 23 + num_heads: 10 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 840 + hidden_dim: 320 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 5 + dropout: 0.1 + mlp_layers: 5 + mlp_ratio: 5 + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 840 + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + # The frozen encoder's per-token features are cached per unique prompt + # string (LRU, in memory), so the 600M-parameter forward runs once per + # annotation instead of once per training step. Ignored when + # `freeze: false` - a trainable encoder must stay online. + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: 840 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + stem_specs: + human_bimanual: + # 20-D: 6D-encoded head-frame wrist poses (18) + zero grip pads (2). + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 840 + widths: [840] + # Proprio observation history: K tokens (one per step, current step + # last) plus a learned zero-init time-step embedding, instead of the + # single current-step token. Must equal the data config's + # `proprio_history` (which also sets how far apart the steps are). + # `history_dropout` replaces the whole past by the current step for a + # sample (train only) so the policy still works when history is + # missing at episode start -- it is live at history_len 3, and it is + # what keeps the policy from reading its own past as the answer. + history_len: 3 + history_dropout: 0.2 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 840 + num_of_copy: 1 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-5 + weight_decay: 0.0001 + +scheduler: null + +# torch.compile for the training step: ModelWrapper._maybe_compile hands the +# trunk, the flow head's denoiser and the image encoders to torch.compile +# (HPT.compile_for_training) before DDP wraps anything. Measured on one H100 at +# batch 64 and the 360x640 frames mecka decodes to (2026-09-19): 173 ms/step -> +# 109 ms (1.58x), for ~40 s of compile at the first batch. mode="reduce-overhead" +# is worth a further ~5% for roughly double that compile, but CUDA-graph modes +# (reduce-overhead, max-autotune) are where a rank can wedge in inductor's +# cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel. +# +# Off by default -- a run that never reaches steady state pays the compile for +# nothing. Per run: model.compile.enabled=true. `mode` goes straight to +# torch.compile; null is the default inductor mode the numbers were taken with. +compile: + enabled: false + mode: null + dynamic: false diff --git a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml index e515212854..7edd9e4904 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml @@ -4,14 +4,15 @@ robomimic_model: camera_transforms: {} dims: + # eva proprio is joint positions (7/arm), not a pose. eva_bimanual: {proprio: 14, action: 14} - human_bimanual: {proprio: 12, action: 12} + # human cartesian_6d: 18-D action, 20-D gripper-padded ee_pose proprio. + human_bimanual: {proprio: 20, action: 18} ac_keys: human_bimanual: "actions_cartesian" eva_bimanual: "actions_joints" - reverse_kl_samples: 8 6dof: true trunk: diff --git a/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml index 499472930d..a3d73c5697 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml @@ -1,12 +1,41 @@ defaults: - hpt_cotrain_enc_dec_base +# Separate-head cotrain: eva on the base's 14-D joint-position action +# (`actions_joints` ac_key), human on the 144-D wrist-frame hand-keypoint +# action through a `state_keypoints` stem; the base's `state_ee_pose` stem is +# dropped (null) so a keypoint data config's extra ee_pose proprio is ignored. +# No shipped HPT cotrain data config emits this pairing yet; build one from a +# per-vendor keypoint config plus data/eva.yaml (whose eva action is the 20-D +# wrist-frame 6D layout, so override dims.eva_bimanual to match). The +# shared-head cotrain configs (hpt_cotrain_*_shared_head) need one action +# space and stay cartesian on both sides. robomimic_model: domains: ["eva_bimanual", "human_bimanual"] - diffusion: true + diffusion: true + dims: + human_bimanual: {proprio: 144, action: 144} + ac_keys: + human_bimanual: "actions_keypoints" + stem_specs: + human_bimanual: + state_ee_pose: null + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 head_specs: shared: null - human_bimanual: + human_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 @@ -26,7 +55,7 @@ robomimic_model: mlp_layers: 4 mlp_ratio: 4 - eva_bimanual: + eva_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 @@ -43,9 +72,9 @@ robomimic_model: act_seq: 100 n_heads: 4 dropout: 0.1 - mlp_layers: 4 + mlp_layers: 4 mlp_ratio: 4 - + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml index 56fde2ad4e..fad676946b 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml @@ -13,7 +13,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml index 0598b67e83..b2811a2000 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml @@ -32,7 +32,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml index 42dc191dff..f927f22cd2 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml @@ -29,7 +29,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index e94ff5c095..0658feca5d 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -5,8 +5,24 @@ robomimic_model: pytorch_training_precision: bfloat16 # directory holding model.safetensors of the pi0.5 base ($EGOVERSE_PI05_WEIGHTS overrides) pytorch_weight_path: ${oc.env:EGOVERSE_PI05_WEIGHTS,/storage/home/hcoda1/4/paphiwetsa3/r-dxu345-0/projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch} + # The alternative to pytorch_weight_path: a local directory or an HF repo id + # holding the paligemma-3b-*-224 safetensors. Only the VLM prefix (SigLIP + # tower, multimodal projector, Gemma-2B) is loaded from it; everything pi0.5 + # adds -- the action expert, the action in/out projections, the adaRMS time + # MLPs -- trains from scratch. Setting both paths is an error. Flip it on + # with the model/pi0.5_paligemma_init.yaml overlay. + paligemma_weight_path: null + # google/paligemma-3b-pt-224 is license-gated. true = without an HF token + # that has accepted the Gemma terms, fall back to a third-party mirror whose + # shards sha256-match the canonical repo; false (default) = raise. + paligemma_allow_mirror: false model: pi05: true + # Width of the model's action vector. 32 is openpi's pretrained width + # (the cartesian block layout fits in it). Configs with a human domain + # raise it to 144 for the wrist-first hand-keypoint action; PI then + # re-initializes only openpi's two action projections at that width + # and zero-pads narrower embodiments (eva's 20-D) into the same vector. action_dim: 32 action_horizon: 100 max_token_len: 180 @@ -53,7 +69,11 @@ robomimic_model: size: 224 interpolation: 3 -enable_grad_norm: false +# openpi's trainers clip the global grad norm at 1.0 (AdamW +# clip_gradient_norm) and log the pre-clip norm every step. trainHydra hands this +# to the Trainer unless trainer.gradient_clip_val is set; ModelWrapper logs the +# norm as Train/policy_grad_norms_raw (enable_grad_norm: false turns that off). +gradient_clip_val: 1.0 optimizer: _target_: torch.optim.AdamW @@ -63,11 +83,31 @@ optimizer: eps: 1e-8 weight_decay: 0.0 +# Constant LR after warmup (matches the remote pi0.5 training stack). The +# previous cosine schedule decayed to 0 over num_training_steps=60000, so runs +# past ~60k steps trained at ~0 lr and shorter runs never saw the configured lr +# for most of training. scheduler: - _target_: transformers.get_cosine_schedule_with_warmup + _target_: transformers.optimization.get_constant_schedule_with_warmup _partial_: true num_warmup_steps: 2000 - num_training_steps: 60000 - num_cycles: 0.5 -scheduler_interval: step \ No newline at end of file +scheduler_interval: step + +# torch.compile for the training step: ModelWrapper._maybe_compile compiles +# openpi's PI0Pytorch.forward (PI.compile_for_training). Measured on one H100 at +# batch 16 (2026-09-19): 827 ms/step -> 588 ms (1.41x), for ~4 min of compile at +# the first batch. `sample_actions` is compiled by openpi itself and is not +# affected by this flag. +# +# CUDA-graph modes (reduce-overhead, max-autotune) are where a rank can wedge in +# inductor's cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel; +# openpi already compiles sample_actions with max-autotune whatever this says. +# +# Off by default -- a run that never reaches steady state pays the compile for +# nothing. Per run: model.compile.enabled=true. `mode` goes straight to +# torch.compile; null is the default inductor mode the numbers were taken with. +compile: + enabled: false + mode: null + dynamic: false diff --git a/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml index 8782b49304..7dcfe5dc02 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml @@ -2,17 +2,21 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (aria, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/aria.yaml); see pi0.5_bc_mecka.yaml. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter - diff --git a/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml index c08fe331ce..f69fe9810a 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml @@ -11,7 +11,7 @@ robomimic_model: action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml index 8782b49304..fa4d5a0deb 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml @@ -2,17 +2,24 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (mecka, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/mecka.yaml, `keypoints_wristframe_6d`). +# The action vector is widened to 144: PI re-initializes openpi's two action +# projections at that width, everything else loads from pi05_base. The prompt +# state is the cartesian head-frame ee_pose (proprio_keys_for_prompt default). robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter - diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml new file mode 100644 index 0000000000..3a557c0b5b --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml @@ -0,0 +1,30 @@ +defaults: + - pi0.5_base + - _self_ + +# pi0.5 BC on mecka data alone (stored as human_bimanual) with the CARTESIAN +# action: actions arrive already in xyz+6D layout (18D per arm pair, no +# gripper, wrist-relative via the `cartesian_wristframe_6d` transform mode) +# and the head-frame proprio ee_pose arrives 6D-encoded too; the forward pass +# only packs them into openpi's 32D vector (pretrained projections kept). +# Pairs with data=mecka_all_6d / mecka_fold_freeform*_6d and +# evaluator=eval_pi_wristframe_6d (or, with the human transform mode +# overridden to `cartesian_6d`, eval_pi_camframe_6d). +robomimic_model: + camera_transforms: {} + ac_keys: + human_bimanual: "actions_cartesian" + domains: ["human_bimanual"] + + # Keep the discretized-proprio "State: " prompt block (pi0.5_base + # defaults proprio_in_prompt on) but do not splice "Embodiment: " into + # the prompt. + embodiment_label: false + + action_converters: + rules: + HUMAN_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml new file mode 100644 index 0000000000..91abc72867 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml @@ -0,0 +1,15 @@ +defaults: + - pi0.5_bc_mecka_6d + - pi0.5_paligemma_init + - _self_ + +# pi0.5_bc_mecka_6d with the pi0.5 base checkpoint swapped for a PaliGemma-only +# init -- the ablation that asks how much of the mecka 6D result comes from +# pi0.5's robot pretraining rather than from the VLM. Everything else (the +# cartesian 6D action layout, the action converter, the prompt blocks, the +# optimizer) is the parent's; see pi0.5_paligemma_init.yaml for the two keys +# this arm changes and for the lr knobs worth sweeping. +# +# Pairs with the same data configs as the parent (data=mecka_all_6d, +# mecka_fold_freeform*_6d) and evaluator=eval_pi_wristframe_6d. Norm stats are +# data-config-specific, not model-specific, so this arm reuses the parent's. diff --git a/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml index 79ed54e8f2..dae132236f 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml @@ -2,16 +2,21 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (scale, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/scale.yaml); see pi0.5_bc_mecka.yaml. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml index d9c2695b68..a11a9a24a0 100644 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml @@ -2,19 +2,27 @@ defaults: - pi0.5_base - _self_ +# pi0.5 cotrain: eva on the cam-frame cartesian 6D action (20-D, packed into +# the canonical 32-slot block layout and zero-padded), human on the 144-D +# wrist-frame hand-keypoint action (data/cotrain_pi_base.yaml). The action +# vector is 144 wide; PI re-initializes openpi's two action projections at +# that width, everything else loads from pi05_base. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: eva_bimanual: "actions_cartesian" - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["eva_bimanual", "human_bimanual"] action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: - _target_: egomimic.utils.action_utils.BaseActionConverter \ No newline at end of file + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml index 301122433b..57159a8908 100644 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml @@ -25,9 +25,9 @@ robomimic_model: action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml new file mode 100644 index 0000000000..4e56103988 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml @@ -0,0 +1,30 @@ +# OVERLAY, not a runnable model config. Compose it AFTER any pi0.5_* config to +# swap the pi0.5 base checkpoint for a PaliGemma-only init: the VLM prefix comes +# from the public paligemma-3b-pt-224 release and the action expert, the action +# in/out projections and the adaRMS time MLPs start from their PyTorch init. +# +# defaults: +# - pi0.5_bc_mecka_6d +# - pi0.5_paligemma_init +# - _self_ +# +# It sets only the two init keys, so it never clobbers the config it is layered +# on (action_dim, ac_keys, domains, the prompt flags all survive). Per-run, +# without a new file: +# model.robomimic_model.config.pytorch_weight_path=null \ +# model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +# +# The optimizer is left at pi0.5_base's fine-tuning recipe (constant 3e-5 after +# 2000 warmup steps). A randomly initialized action expert may well want a +# higher peak and a longer warmup; that is a per-run decision, e.g. +# model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +# +# Multi-rank runs should pre-fetch the checkpoint once, or every rank races on +# the same Hugging Face download lock: +# python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +robomimic_model: + config: + pytorch_weight_path: null + # $EGOVERSE_PALIGEMMA_WEIGHTS overrides: point it at a local snapshot dir to + # skip the Hub entirely. + paligemma_weight_path: ${oc.env:EGOVERSE_PALIGEMMA_WEIGHTS,google/paligemma-3b-pt-224} diff --git a/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml new file mode 100644 index 0000000000..d60059a7b1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml @@ -0,0 +1,11 @@ +# RDT_1B_BC_MECKA_6D_DINOV3_VITB: rdt_bc_mecka_6d_dinov3_vitb at the paper's +# RDT-1B shape (width 2048 x 28 blocks x 32 heads, ~1.2B). From scratch on one +# task this is a scaling-curve point, not the default. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + width: 2048 + trunk: + num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_bc_aria_keypoints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_bc_aria_keypoints_dinov3_vitb.yaml new file mode 100644 index 0000000000..16118fd8c2 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_aria_keypoints_dinov3_vitb.yaml @@ -0,0 +1,120 @@ +# RDT ~300M (rdt_bc_mecka_6d_dinov3_vitb) on Aria wrist-first keypoints: 144-D +# head-frame action and proprio (keypoints_headframe_6d), one front frame, 4:3 +# Aria images. Aria samples carry no annotation key, so language is the fixed +# default_prompt. +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.rdt.RDT + camera_transforms: {} + dims: + human_bimanual: {proprio: 144, action: 144} + width: 1024 + ac_keys: + human_bimanual: "actions_keypoints" + + diffusion: true + 6dof: true + + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + use_pad_mask: false + + trunk: + embed_dim: ${model.robomimic_model.width} + depth: 28 + num_heads: 16 + mlp_ratio: 1.0 + cond_drop: {state: 0.1, lang: 0.1, img: 0.0} + image_history: 1 + image_history_dropout: 0.0 + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + n_state_tokens: 1 + time_scale: 1000.0 + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.width} + output_dim: ${model.robomimic_model.width} + widths: ["${model.robomimic_model.width}"] + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: ${model.robomimic_model.width} + + stem_specs: + human_bimanual: + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.0 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.DINOv3Stem + model_name: "vit_base_patch16_dinov3.lvd1689m" + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + # Aria front_1 frames are 4:3 (360x480); 15x20 patches. + image_size: [240, 320] + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + eval_image_augs: null + +backbone_lr_scale: 0.1 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + _convert_: all + lr: 1e-4 + weight_decay: 0.01 + betas: [0.9, 0.95] + +scheduler: + _target_: torch.optim.lr_scheduler.LinearLR + _partial_: true + start_factor: 1.0e-3 + end_factor: 1.0 + total_iters: 1000 +scheduler_interval: step diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml new file mode 100644 index 0000000000..f0539898f9 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml @@ -0,0 +1,173 @@ +# RDT_BC_MECKA_6D_DINOV3_VITB: Robotics Diffusion Transformer +# (thu-ml/RoboticsDiffusionTransformer) on the HPT chassis, ~300M trainable +# DiT (param-matched to hpt_bc_mecka_6d_300M), same data / action space: 18-D +# wrist-frame 6D-rotation cartesian action, 20-D ee_pose proprio, front camera, +# one sampled annotation. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_rdt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# egomimic.algo.rdt.RDT is HPT with the policy network swapped: no Perceiver +# pooling. The stems are RDT's condition adaptors and hand their full token +# sequences to the head. `trunk` is the DiT, shared across embodiments: +# self-attention over [timestep; ctrl freq; state; noisy action chunk], +# cross-attention alternating between the language tokens (even blocks) and the +# image tokens (odd blocks). A head holds one embodiment's action adaptor, +# position table and output layer, so co-training shares the blocks at native +# action widths. +# +# Against the paper (decided with ABC-DiT, abc.bot, as the second reference): +# - flow matching (the repo's FMPolicy, 10 Euler steps like pi0.5 / ABC) for +# DDPM sample-prediction + DPM-Solver; +# - DINOv3 ViT-B/16 for SigLIP-so400m, fine-tuned at 0.1x LR (ABC's setting; +# encoder_specs.front_img_1.freeze_backbone=true is the frozen control), at +# the native 16:9 aspect, 224x400 = 350 patch tokens per frame, all of them +# cross-attended. The stem resizes and normalizes; the augs must NOT; +# - Qwen3-Embedding-0.6B per-token states for T5-XXL (frozen, cached); +# - native action width + per-embodiment adaptors for the 128-D unified vector; +# - 2 frames per camera, the past one 0.1 s back (RDT: the previous step); +# needs data...key_map.image_history_gap_s (the train config sets it); +# - DINOv3 is RoPE-only, so patches get a fixed normalized 2D sin-cos code on +# top of RDT's (learned, zero-init) frame / camera tags; +# - condition dropout per stream: state and language 0.1, no image dropout +# with a single camera. +# +# DiT: width 1024 x 28 blocks x 16 heads (head_dim 64, RDT-1B's), RDT's FFN +# ratio 1 -> 10.5M / block. +# Trainable 392.7M = DiT 296.6M + DINOv3 85.6M + adaptors; Qwen 596M frozen. +# H100, batch 32, bf16: 212 ms / step, 30 GiB (ViT-L: 365 ms, 49 GiB; the 1B +# DiT: 359 ms, 57 GiB); sampling 0.3 s per batch. +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.rdt.RDT + camera_transforms: {} + dims: + human_bimanual: {proprio: 20, action: 18} + # DiT width; every stem / encoder output_dim and the head's hidden_dim. + width: 1024 + ac_keys: + human_bimanual: "actions_cartesian" + + diffusion: true + 6dof: true + + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + use_pad_mask: false + + trunk: + embed_dim: ${model.robomimic_model.width} + depth: 28 + num_heads: 16 + mlp_ratio: 1.0 + # RDT's cond_mask_prob, per stream (train only). A dropped prompt becomes + # "", the no-prompt deployment condition; `img` is per camera and never + # blinds a sample, so it is inert with one camera. + cond_drop: {state: 0.1, lang: 0.1, img: 0.0} + # frames per camera (current last); 2 <-> key_map.image_history_gap_s set + image_history: 2 + # past frame := current frame, per sample (what an episode start looks like) + image_history_dropout: 0.2 + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + # proprio stems x history_len + n_state_tokens: 1 + # FMPolicy's t in [0, 1] -> RDT's [0, 1000) timestep embedding + time_scale: 1000.0 + + shared_stem_specs: + # image adaptor (DINOv3Stem.proj + this MLP ~ RDT's mlp2x_gelu) + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.width} + output_dim: ${model.robomimic_model.width} + widths: ["${model.robomimic_model.width}"] + # RDTModel adds the mlp2x_gelu language adaptor on top of these tokens. + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: ${model.robomimic_model.width} + + stem_specs: + human_bimanual: + # state adaptor (RDT's mlp3x_gelu); one token per proprio step + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.2 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.DINOv3Stem + # timm pulls model_name from the HF hub; prefetch it (the ViT-L config's + # too) on a login node before an offline or multi-rank run: + # python -c "import timm; timm.create_model('', pretrained=True)" + model_name: "vit_base_patch16_dinov3.lvd1689m" + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + # [H, W]; 360x640 frames at their native aspect + image_size: [224, 400] + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + eval_image_augs: null + +# DINOv3 tower at 1e-5 against 1e-4 everywhere else (ABC-DiT's 0.1). +backbone_lr_scale: 0.1 + +# RDT and ABC-DiT agree: AdamW 1e-4, wd 1e-2, constant LR. ABC's 1k-step +# warmup and (0.9, 0.95) betas; gradient clipping is in the train config. +optimizer: + _target_: torch.optim.AdamW + _partial_: true + _convert_: all + lr: 1e-4 + weight_decay: 0.01 + betas: [0.9, 0.95] + +# LinearLR holds end_factor after total_iters: warmup, then constant. +scheduler: + _target_: torch.optim.lr_scheduler.LinearLR + _partial_: true + start_factor: 1.0e-3 + end_factor: 1.0 + total_iters: 1000 +scheduler_interval: step diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml new file mode 100644 index 0000000000..47096dd4d1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml @@ -0,0 +1,11 @@ +# RDT_BC_MECKA_6D_DINOV3_VITL: rdt_bc_mecka_6d_dinov3_vitb with the DINOv3 +# ViT-L/16 tower (300M, 1024-wide tokens) -- the two differ only in the image +# encoder. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + encoder_specs: + front_img_1: + model_name: "vit_large_patch16_dinov3.lvd1689m" diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 333a2372de..7ccfe6bc95 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -28,13 +28,36 @@ launch_params: seed: 42 +# Optional pre-fit baseline validation (metrics at epoch 0; the overlay video +# still follows evaluator.viz_every_n_epochs). Off by default; enable with +# `val_at_start=true`. Skipped automatically on requeues/checkpoint resumes. +val_at_start: false + +# Second eval head on TRAIN data at every validation (metrics `train_viz/`, +# videos under `videos_train_viz/`). Wraps the canonical evaluator and reuses +# the train split unless the data config declares `train_viz_datasets` / +# the top-level config a `train_viz_evaluator`. `train_viz=false` turns it off. +train_viz: true + # Normalization and norm-stat cache settings (override per-run as needed). norm_stats: norm_mode: quantile + # true: one set of stats per channel over the whole action chunk (openpi, + # RDT, GR00T, Diffusion Policy, lerobot); false: per (timestep, channel) + # cell, which keeps resolution on the near-future steps of a relative-frame + # chunk (TRI's LBM makes the same choice for the same reason). Turn it on + # for checkpoint parity with an openpi-shaped recipe; the two have not been + # compared on a real run here. + pool_horizon: false sample_frac: 1.0 + # Hard cap on sampled frames. Memory is ~2 x N x T x D x 4 bytes per key, so + # set it on large splits (200000 on full mecka keypoints is ~23 GB). + max_samples: null num_workers: 4 save_cache_dir: ${hydra:runtime.output_dir} precomputed_norm_path: null - cache_dir: ${paths.cache_dir}/norm_stats # content-keyed cache shared across runs + # Content-keyed cache shared across runs: one entry per split, plus per-episode + # samples under episodes/ that any other split of the same recipe reuses. + cache_dir: ${paths.cache_dir}/norm_stats use_cache: true # false: always recompute, never write reject_outliers: true diff --git a/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml b/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml index c625cbd71c..09ab5c5a56 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml @@ -5,3 +5,4 @@ defaults: - override trainer: ddp_pi - override evaluator: eval_pi - override hydra/launcher: submitit_pace + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml index e45cfeea0a..b1571252c1 100644 --- a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints_wrist - - override evaluator/viz@evaluator.viz_func: keypoints_wrist + - override evaluator: eval_hpt_keypoints_wrist - _self_ diff --git a/egomimic/hydra_configs/train_zarr_keypoints.yaml b/egomimic/hydra_configs/train_zarr_keypoints.yaml index 170bffd287..3cbe07a386 100644 --- a/egomimic/hydra_configs/train_zarr_keypoints.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoints.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints - - override evaluator/viz@evaluator.viz_func: keypoints + - override evaluator: eval_hpt_keypoints_head - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_base.yaml b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_base.yaml new file mode 100644 index 0000000000..07d987ec23 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_base.yaml @@ -0,0 +1,41 @@ +# RDT memory-representation ablation, BASELINE arm: RDT-300M (model/ +# rdt_bc_mecka_6d_dinov3_vitb) on the pinned 20% mecka subset (data/ +# mecka_20pct_6d), with the short image pair (current + 0.1 s back) and the +# current proprio only. The other arms change one thing each (2x2): +# _proprio + proprio history: 11 steps 1 s apart (t-10 s .. t) +# _vision + long-range visual memory: 10 frames 1 s apart (t-9 s .. t), +# one CLS token each from a second fine-tuned DINOv3, in the +# DiT's self-attention prefix; the short pair is dropped for 30 % +# of training samples +# _both both, each memory token fused with that step's proprio +# +# One 8xH100 node per arm, the same budget everywhere: global batch 256 +# (32 / GPU x 8), 1000 epochs x 100 steps, RDT's AdamW 1e-4 / 1k-step warmup, +# DINOv3 at 0.1x, grad clip 1.0, val every 25 epochs. Every arm shares one set +# of norm stats (proprio stats are current-step, cameras are not normalized): +# python egomimic/trainHydra.py -m --config-name=train_zarr_mecka_20pct_6d_rdt_mem_base \ +# name=rdt_mem description=base norm_stats.precomputed_norm_path=/norm_stats +defaults: + - train_zarr_mecka_flagship_6d_rdt + - override data: mecka_20pct_6d + - override hydra/launcher: submitit_loaner + - _self_ + +launch_params: + gpus_per_node: 8 + nodes: 1 + +data: + train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 12 + valid_dataloader_params: + human_bimanual: + batch_size: 32 + +trainer: + max_epochs: 1000 + min_epochs: 1000 + check_val_every_n_epoch: 25 + limit_train_batches: 100 diff --git a/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_both.yaml b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_both.yaml new file mode 100644 index 0000000000..ef38651d69 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_both.yaml @@ -0,0 +1,30 @@ +# RDT memory ablation, BOTH arm: the vision arm's memory plus the proprio arm's +# history on the same 1 s grid (11 steps, t-10 s .. t). Each memory token also +# gets a projection of its step's proprio token (the newest 10); the 11 state +# tokens stay in the prefix as in the proprio arm. +# See train_zarr_mecka_20pct_6d_rdt_mem_base.yaml. +defaults: + - train_zarr_mecka_20pct_6d_rdt_mem_vision + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + proprio_history: 11 + history_stride_s: 1.0 + +model: + robomimic_model: + trunk: + memory: + fuse_proprio: true + stem_specs: + human_bimanual: + state_ee_pose: + history_len: 11 + head_specs: + human_bimanual: + model: + n_state_tokens: 11 diff --git a/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_proprio.yaml b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_proprio.yaml new file mode 100644 index 0000000000..64a9b0d310 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_proprio.yaml @@ -0,0 +1,25 @@ +# RDT memory ablation, PROPRIO arm: the baseline plus proprio history, 11 +# ee_pose steps 1 s apart (t-10 s .. t, current last), one state token each. +# See train_zarr_mecka_20pct_6d_rdt_mem_base.yaml for the arms and recipe. +defaults: + - train_zarr_mecka_20pct_6d_rdt_mem_base + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + proprio_history: 11 + history_stride_s: 1.0 + +model: + robomimic_model: + stem_specs: + human_bimanual: + state_ee_pose: + history_len: 11 + head_specs: + human_bimanual: + model: + n_state_tokens: 11 diff --git a/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_vision.yaml b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_vision.yaml new file mode 100644 index 0000000000..ac1a5fb9db --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_20pct_6d_rdt_mem_vision.yaml @@ -0,0 +1,48 @@ +# RDT memory ablation, VISION arm: the baseline plus a long-range visual +# memory, 10 front frames 1 s apart (t-9 s .. t). A second DINOv3 tower, +# fine-tuned at the same 0.1x LR as the short one, pools each frame to its CLS +# token, tagged by the frame's age (sin-cos, seconds); the 10 tokens join the +# DiT's self-attention prefix, with steps before the episode start masked. +# 30 % of training samples lose the short pair (a learned null token), so the +# policy cannot lean on the latest frames alone. +# +# Cost, 8xH100, 32 / GPU, bf16 (2026-09-23): both arm 1.19 s / step (base 0.40), +# ~33 h per 100k steps. +# See train_zarr_mecka_20pct_6d_rdt_mem_base.yaml. +defaults: + - train_zarr_mecka_20pct_6d_rdt_mem_base + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + image_memory: 10 + image_memory_stride_s: 1.0 + +model: + robomimic_model: + trunk: + short_dropout: 0.3 + memory: + encoder: + _target_: egomimic.models.hpt_nets.DINOv3Stem + model_name: ${model.robomimic_model.encoder_specs.front_img_1.model_name} + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + image_size: ${model.robomimic_model.encoder_specs.front_img_1.image_size} + pool: cls + # 320 frames / GPU: their activations OOM an H100 without it + grad_checkpointing: true + # Must equal key_map.image_memory / image_memory_stride_s above: the + # model is built from the model config alone, so it cannot interpolate + # them. A frame-count mismatch raises at the first batch. + frames: 10 + stride_s: 1.0 + fuse_proprio: false + head_specs: + human_bimanual: + model: + # one token per memory frame + n_memory_tokens: 10 diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml new file mode 100644 index 0000000000..8a93235d67 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml @@ -0,0 +1,20 @@ +# Composed recipe: ~300M HPT (+ Qwen text stem) BC on the 18-D wrist-frame +# 6D-rotation cartesian action, mecka FLAGSHIP fold-clothes, seen/unseen +# operator split (see data/mecka_fold_flagship_opsplit_hpt_6d.yaml). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The norm stats must come from data=mecka_fold_flagship_opsplit_hpt_6d (see +# that file for the precompute command). +defaults: + - train_zarr_cartesian + # Second eval head: metrics + videos on the TRAIN split each val + # (train_viz/). The data config renames the canonical head seen_op_valid/ + # and adds a third, unseen_op_valid/, for the held-out operators. (A plain + # entry, so it must precede the overrides.) + - evaluator@train_viz_evaluator: train_viz_hpt_wristframe_6d + - override model: hpt_bc_mecka_6d_300M + - override data: mecka_fold_flagship_opsplit_hpt_6d + - override evaluator: eval_hpt_wristframe_6d + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml new file mode 100644 index 0000000000..6ca6da36e4 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml @@ -0,0 +1,32 @@ +# Composed recipe: train_zarr_mecka_flagship_6d_hpt with the RDT policy +# (model/rdt_bc_mecka_6d_dinov3_vitb.yaml; model=rdt_bc_mecka_6d_dinov3_vitl for +# the ViT-L tower, model=rdt_1b_bc_mecka_6d_dinov3_vitb for the 1B DiT). Same +# data, action space, evaluators and norm stats as the HPT recipe. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_rdt \ +# norm_stats.precomputed_norm_path=/norm_stats +defaults: + - train_zarr_mecka_flagship_6d_hpt + - override model: rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + # RDT's 2-frame image history: also load the front frame this many + # seconds back (3 frames at 30 fps). The val heads interpolate this + # resolver, so they follow. Pairs with model...trunk.image_history=2. + image_history_gap_s: 0.1 + # PROPRIO history stays off for this arm: the HPT recipe ships K = 3 + # / stride 5, but RDT's state adaptor is declared at history_len 1 + # and the two must agree or process_batch_for_training raises. Raise + # both together if the RDT arm should carry history too. + proprio_history: 1 + history_stride: 1 + +trainer: + # RDT clips at 1.0 (ABC at 10): protects the fine-tuned tower from the + # from-scratch DiT's early gradients. + gradient_clip_val: 1.0 diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml new file mode 100644 index 0000000000..791868c538 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml @@ -0,0 +1,23 @@ +# Composed recipe: HPT 300M flow-matching BC on the 144-D wrist-frame MANO +# keypoint action, mecka FLAGSHIP fold-clothes, seen/unseen operator split. +# The keypoint counterpart of train_zarr_mecka_flagship_6d_hpt, and the +# reference arm any other architecture on this split is compared against +# (the same recipe with `override model:` swapped). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_kp_hpt \ +# data=mecka_fold_flagship_top3_hpt_keypoints \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The flagship annotations are empty, so every prompt is the default one +# ("fold clothes"), so a comparison arm has to set the same string. +defaults: + - train_zarr_cartesian + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: hpt_bc_keypoints_wrist_300M + - override data: mecka_fold_flagship_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +model: + robomimic_model: + default_prompt: "fold clothes" diff --git a/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml new file mode 100644 index 0000000000..59479cd4ed --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml @@ -0,0 +1,17 @@ +# Composed recipe: ~300M HPT (+ Qwen text stem) BC on the 144-D wrist-frame +# MANO keypoint action, mecka freeform fold-clothes, held-out-operator split. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_kp_wrist_hpt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The norm stats must come from data=mecka_fold_freeform_opsplit_hpt_keypoints +# (see that file for the precompute command). +defaults: + - train_zarr_cartesian + # Second eval head: metrics + videos on TRAIN data each val (train_viz/). + # (A plain entry, so it must precede the overrides.) + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: hpt_bc_keypoints_wrist_300M + - override data: mecka_fold_freeform_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ diff --git a/egomimic/hydra_configs/trainer/ddp.yaml b/egomimic/hydra_configs/trainer/ddp.yaml index d3d90aca03..8b8c98f072 100644 --- a/egomimic/hydra_configs/trainer/ddp.yaml +++ b/egomimic/hydra_configs/trainer/ddp.yaml @@ -4,7 +4,8 @@ defaults: strategy: ddp accelerator: gpu -devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} +# Per node: Lightning's world size is devices * num_nodes. +devices: ${launch_params.gpus_per_node} num_nodes: ${launch_params.nodes} sync_batchnorm: True check_val_every_n_epoch: 200 diff --git a/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml b/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml new file mode 100644 index 0000000000..c266e66452 --- /dev/null +++ b/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml @@ -0,0 +1,17 @@ +# Plain DDP plus a bf16 gradient all-reduce, for a model big enough that the +# gradient reduction, not the math, sets the step time. On the lab's 8 x A40 +# nodes -- no NVLink, two NUMA islands joined by cross-socket links -- a ~1.2 B +# -parameter policy measured 6.4 s/step at 8 x batch 16 against 1.4 s/step on a +# single GPU, and 3.6 s/step once the reduction was halved (2026-09-16). Only +# the reduction is bf16; master weights and the optimizer stay fp32. +# +# This is plain DDP: no unused-parameter detection. A model that leaves any +# trainable tensor without a gradient needs `ddp_pi`-style +# `ddp_find_unused_parameters_true` instead, and will crash under this one. +defaults: + - ddp + +strategy: + _target_: lightning.pytorch.strategies.DDPStrategy + ddp_comm_hook: + _target_: egomimic.utils.ddp_hooks.make_bf16_compress_hook diff --git a/egomimic/hydra_configs/trainer/ddp_pi.yaml b/egomimic/hydra_configs/trainer/ddp_pi.yaml index 731316953a..48ab3b76f5 100644 --- a/egomimic/hydra_configs/trainer/ddp_pi.yaml +++ b/egomimic/hydra_configs/trainer/ddp_pi.yaml @@ -3,9 +3,17 @@ defaults: strategy: ddp_find_unused_parameters_true +# Validate under torch.no_grad instead of torch.inference_mode: the compiled +# pi sampler (max-autotune -> cudagraphs) does in-place setup on its static +# buffers during graph capture, which is illegal on inference tensors +# ("Inplace update to inference tensor outside InferenceMode"). no_grad is +# functionally equivalent for eval and creates no inference tensors. +inference_mode: false + accelerator: gpu -devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} +# Per node: Lightning's world size is devices * num_nodes. +devices: ${launch_params.gpus_per_node} num_nodes: ${launch_params.nodes} sync_batchnorm: True -check_val_every_n_epoch: 200 -num_sanity_val_steps: 0 \ No newline at end of file +check_val_every_n_epoch: 20 +num_sanity_val_steps: 0 diff --git a/egomimic/hydra_configs/trainer/default.yaml b/egomimic/hydra_configs/trainer/default.yaml index 62f147c27d..e75492a53f 100644 --- a/egomimic/hydra_configs/trainer/default.yaml +++ b/egomimic/hydra_configs/trainer/default.yaml @@ -17,4 +17,11 @@ check_val_every_n_epoch: 100 # set True to to ensure deterministic results # makes training slower but gives more reproducibility than just setting seeds -deterministic: False \ No newline at end of file +deterministic: False + +# Fixed gradient clipping, applied by Lightning in configure_gradient_clipping. +# The MAD spike detector in ModelWrapper.on_after_backward is log-only and never +# touches gradients; this is the only thing that clips. null = no clipping (the +# value is picked from the grad-norm curve of a log-only baseline run). +gradient_clip_val: null +gradient_clip_algorithm: norm diff --git a/egomimic/models/denoising_nets.py b/egomimic/models/denoising_nets.py index 18471c59f5..6407340423 100644 --- a/egomimic/models/denoising_nets.py +++ b/egomimic/models/denoising_nets.py @@ -575,7 +575,7 @@ def forward(self, x, t, cond, timestep_table): res = x x = self.ln1(x) x = self.aln1(x, t, cond, timestep_table) - (x, _) = self.mha(x, x, x) + (x, _) = self.mha(x, x, x, need_weights=False) x = self.alns1(x, t, cond, timestep_table) x = x + res res = x @@ -617,13 +617,17 @@ def __init__( self.cond_proj = nn.Linear(cond_dim, hidden_dim) def forward_cross(self, x, cond): + # need_weights=False: the default True makes nn.MultiheadAttention take + # the slow path -- it materialises the B x H x L x L scores and averages + # them over heads to return weights nobody here reads -- instead of + # dispatching to scaled_dot_product_attention. The outputs are the same. res = x x = self.ln1(x) - x, _ = self.mha(x, x, x) + x, _ = self.mha(x, x, x, need_weights=False) x = x + res res = x x = self.ln2(x) - x, _ = self.cmha(x, cond, cond) + x, _ = self.cmha(x, cond, cond, need_weights=False) x = x + res res = x x = self.ln3(x) diff --git a/egomimic/models/denoising_policy.py b/egomimic/models/denoising_policy.py index 645a8c44e4..2ebef651d4 100644 --- a/egomimic/models/denoising_policy.py +++ b/egomimic/models/denoising_policy.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from egomimic.models.denoising_nets import ConditionalUnet1D +from egomimic.utils.pose_utils import hand_channel_mask class DenoisingPolicy(nn.Module): @@ -80,9 +81,16 @@ def sample_action(self, global_cond, embodiment_name, generator=None): ) return self.inference(noise, global_cond, generator) - def forward(self, global_cond): + def forward(self, global_cond, generator=None): + """Sample actions from the conditioning. + + ``generator`` is threaded down to ``sample_action`` so the eval path can + make sampling reproducible (see ``HPT.forward_eval``). It defaults to + ``None``, i.e. the global RNG, so training and every other caller are + unchanged. + """ cond, embodiment = global_cond - return self.sample_action(cond, embodiment) + return self.sample_action(cond, embodiment, generator=generator) def predict(self, actions, global_cond) -> Tuple[torch.Tensor, torch.Tensor]: """ @@ -90,11 +98,24 @@ def predict(self, actions, global_cond) -> Tuple[torch.Tensor, torch.Tensor]: """ raise NotImplementedError - def loss_fn(self, pred, target): + def loss_fn(self, pred, target, pad_mask=None): """ Computes loss, function to override for stuff like adaptive loss weighting + + ``pad_mask`` (B, S, 1) or (B, S) marks the real action steps; padded + tail steps (0.0) drop out of the mean so the model is not supervised on + the dataset's repeat-last padding. ``None`` (or an all-ones mask) is + plain ``F.mse_loss``. """ - return F.mse_loss(pred, target) + if pad_mask is None: + return F.mse_loss(pred, target) + mask = pad_mask.to(dtype=pred.dtype, device=pred.device) + if mask.dim() == pred.dim() - 1: + mask = mask.unsqueeze(-1) + mask = mask.expand_as(pred) + # At least one element: an all-padding chunk contributes 0, not NaN. + denom = mask.sum().clamp(min=1.0) + return (((pred - target) ** 2) * mask).sum() / denom def preprocess_compute_loss(self, global_cond, data): if self.pooling == "mean": @@ -124,4 +145,28 @@ def preprocess_compute_loss(self, global_cond, data): def compute_loss(self, global_cond, data): actions, global_cond = self.preprocess_compute_loss(global_cond, data) pred, target = self.predict(actions, global_cond) - return self.loss_fn(pred, target) + # Unlike the pad mask, always applied: an untracked hand's action channels + # hold fill values, not motion. + mask = _with_hand_mask( + data.get("pad_mask"), data.get("action_hand_mask"), actions + ) + return self.loss_fn(pred, target, pad_mask=mask) + + +def _with_hand_mask(pad_mask, hand_mask, actions): + """``pad_mask`` (B, S, 1) times the per-channel mask of ``hand_mask`` (B, 2) + [left, right] over the bimanual layout of ``actions`` (B, S, D), so an untracked + hand's channels drop out of the loss.""" + if hand_mask is None: + return pad_mask + channels = hand_channel_mask( + hand_mask.to(actions.device, actions.dtype), actions.shape[-1] + ) + if channels is None: + if bool((hand_mask == 1).all()): + return pad_mask + raise ValueError( + f"action width {actions.shape[-1]} has no bimanual layout to mask a hand in" + ) + channels = channels[:, None, :] + return channels if pad_mask is None else pad_mask.to(actions.dtype) * channels diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 3e76bcee43..ac351effb9 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -1,3 +1,7 @@ +import hashlib +import json +import os +from collections import OrderedDict from functools import partial from typing import Callable, List, Optional, Union @@ -7,6 +11,7 @@ import torch.utils.checkpoint as checkpoint import torchvision from einops import rearrange, repeat +from termcolor import cprint from timm.models.layers import DropPath, trunc_normal_ from timm.models.vision_transformer import VisionTransformer from torch import einsum @@ -15,6 +20,186 @@ from egomimic.utils.tensor_utils import get_sinusoid_encoding_table +# -------------------------------------------------------------------------- +# Pretrained-weight protection +# +# ``HPTModel.finalize_modules`` runs a xavier pass (``_init_weights``) over the +# whole model once the stems / heads are ModuleDicts. That pass used to visit +# every ``nn.Linear`` in every stem, including the HF text encoders loaded from +# a pretrained snapshot, silently randomising them. Modules that own pretrained +# weights declare them here; ``apply_skipping_pretrained`` then walks the tree +# without ever entering such a subtree, so any future pretrained stem (e.g. a +# Qwen VLM stem) is protected without touching ``finalize_modules`` again. +# -------------------------------------------------------------------------- + + +class PretrainedWeights: + """Mixin marking a module whose weights come from a pretrained checkpoint. + + Subclasses list the attribute names holding the pretrained submodules in + ``_hpt_pretrained_attrs``. Two things follow: + + * ``finalize_modules``'s init pass skips those subtrees entirely. + * if the subclass can produce a reference state dict for the checkpoint + (``pretrained_reference_state_dict``), the weights are hashed and + compared with that reference once the model is fully built. + """ + + _hpt_pretrained = True + _hpt_pretrained_attrs: tuple = () + + def pretrained_submodules(self) -> List[nn.Module]: + """The submodules whose weights must never be re-initialised.""" + mods = [] + for name in self._hpt_pretrained_attrs: + mod = getattr(self, name, None) + if isinstance(mod, nn.Module): + mods.append(mod) + return mods + + def pretrained_state_dict(self) -> dict: + """Flat ``{".": tensor}`` view of the pretrained weights.""" + state = {} + for name in self._hpt_pretrained_attrs: + mod = getattr(self, name, None) + if isinstance(mod, nn.Module): + for key, value in mod.state_dict().items(): + state[f"{name}.{key}"] = value + return state + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The checkpoint's own weights, read back independently of this model. + + ``None`` (the default) means "cannot be verified offline"; the checker + then skips this module instead of failing, so only the stems that own + an HF snapshot are hash-checked. ``ResNet`` is protected from the init + pass but not verified: its torchvision weights are not a file this can + re-read. + """ + return None + + def pretrained_hash_dtype(self) -> Optional[torch.dtype]: + """dtype both sides are cast to before hashing (``None`` = as stored).""" + return None + + +def pretrained_module_ids(root: nn.Module) -> set: + """ids of every module living under a declared pretrained submodule.""" + skip = set() + for module in root.modules(): + get = getattr(module, "pretrained_submodules", None) + if not callable(get): + continue + for sub in get(): + skip.update(id(m) for m in sub.modules()) + return skip + + +def apply_skipping_pretrained(root: nn.Module, fn: Callable) -> None: + """``root.apply(fn)`` that never enters a pretrained submodule. + + Same post-order traversal (and therefore the same RNG consumption) as + ``nn.Module.apply`` for a tree with no pretrained submodules. + """ + skip = pretrained_module_ids(root) + + def _walk(module: nn.Module) -> None: + if id(module) in skip: + return + for child in module.children(): + _walk(child) + fn(module) + + _walk(root) + + +def hash_state_dict(state: dict, dtype: Optional[torch.dtype] = None) -> str: + """Deterministic sha256 over parameter names + raw bytes, in name order.""" + digest = hashlib.sha256() + for name in sorted(state): + tensor = state[name].detach().to("cpu") + if dtype is not None: + tensor = tensor.to(dtype) + digest.update(name.encode("utf-8")) + digest.update(str(tuple(tensor.shape)).encode("utf-8")) + digest.update(str(tensor.dtype).encode("utf-8")) + digest.update(tensor.flatten().contiguous().view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def verify_pretrained_weights(root: nn.Module, verbose: bool = True) -> dict: + """Hash every declared pretrained submodule and compare with its checkpoint. + + Raises ``RuntimeError`` naming the module and both hashes on a mismatch. + Returns ``{module_name: sha256}`` for the modules that could be checked. + """ + checked = {} + for name, module in root.named_modules(): + get_ref = getattr(module, "pretrained_reference_state_dict", None) + if not callable(get_ref): + continue + reference = get_ref() + if reference is None: + continue + dtype = module.pretrained_hash_dtype() + current = module.pretrained_state_dict() + if set(current) != set(reference): + raise RuntimeError( + f"pretrained weight check failed for '{name or root.__class__.__name__}'" + f" ({module.__class__.__name__}): the checkpoint and the live module" + f" disagree on which tensors exist; only in model:" + f" {sorted(set(current) - set(reference))[:5]}, only in checkpoint:" + f" {sorted(set(reference) - set(current))[:5]}" + ) + live_hash = hash_state_dict(current, dtype) + ref_hash = hash_state_dict(reference, dtype) + label = name or root.__class__.__name__ + if live_hash != ref_hash: + raise RuntimeError( + f"pretrained weights of '{label}' ({module.__class__.__name__}) do not" + f" match their checkpoint: model sha256={live_hash}," + f" checkpoint sha256={ref_hash}. Something modified them between" + " the load and here -- an init pass that descended into the" + " subtree, or an in-place edit. This runs during HPT.__init__," + " before any checkpoint restore, so a resumed checkpoint is not" + " the cause." + ) + checked[label] = live_hash + if verbose: + cprint( + f"[hpt] pretrained weights of '{label}' verified: sha256={live_hash}", + color="green", + ) + return checked + + +def _snapshot_state_dict( + snapshot_dir: str, prefix: str = "", dtype: Optional[torch.dtype] = None +) -> Optional[dict]: + """Read a local HF snapshot's safetensors shards into a flat state dict. + + ``None`` if the snapshot has no safetensors (the caller then skips the + check rather than hitting the network). + """ + from safetensors.torch import load_file + + index = os.path.join(snapshot_dir, "model.safetensors.index.json") + if os.path.isfile(index): + with open(index, encoding="utf-8") as handle: + shards = sorted(set(json.load(handle)["weight_map"].values())) + files = [os.path.join(snapshot_dir, shard) for shard in shards] + else: + single = os.path.join(snapshot_dir, "model.safetensors") + if not os.path.isfile(single): + return None + files = [single] + + state = {} + for path in files: + for key, tensor in load_file(path, device="cpu").items(): + state[f"{prefix}{key}"] = tensor if dtype is None else tensor.to(dtype) + return state + ## Taken directly from hpt/models/transformer with no modifications class CrossAttention(nn.Module): @@ -540,10 +725,35 @@ def __init__( tanh_end: bool = False, ln: bool = True, num_of_copy: int = 1, + history_len: int = 1, + history_dropout: float = 0.0, **kwargs, ) -> None: - """vanilla MLP class""" + """vanilla MLP class + + ``history_len`` (K) > 1 makes this an observation-history stem: the + input carries the last K steps on axis 1 (current step last), one token + per step, and a LEARNED per-step embedding (zero-init, so at init the K + tokens are exactly today's projection of each step) is added AFTER the + projection. ``HPT.stem_process`` skips its input-side sinusoid for such + a stem: that sinusoid is a constant [0, 1, 0, 1, ...] added to the + NORMALIZED values, i.e. amplitude-1 noise on [-1, 1] data. + + ``history_dropout`` is the per-sample probability, in training only, of + replacing every past step by the current one. That is exactly what an + episode-start sample looks like after the dataset's front padding, so + the policy trains on "history missing" the way it will see it at + deployment (the mitigation for copycat / causal confusion). Inert at + K = 1 and in eval. + """ super().__init__(**kwargs) + if history_len > 1 and num_of_copy > 1: + raise ValueError( + f"MLPPolicyStem: history_len={history_len} and " + f"num_of_copy={num_of_copy} both consume axis 1; use one." + ) + self.history_len = int(history_len) + self.history_dropout = float(history_dropout) modules = [nn.Linear(input_dim, widths[0]), nn.SiLU()] for i in range(len(widths) - 1): @@ -561,6 +771,15 @@ def __init__( self.net = nn.ModuleList( [nn.Sequential(*modules) for _ in range(num_of_copy)] ) + if self.history_len > 1: + self.time_embed = nn.Parameter(torch.zeros(self.history_len, output_dim)) + + def _drop_history(self, x: torch.Tensor) -> torch.Tensor: + """Per sample, with probability ``history_dropout``, replace every past + step by the current one (the last slot of axis 1).""" + drop = torch.rand(x.shape[0], device=x.device) < self.history_dropout + current = x[:, -1:].expand_as(x) + return torch.where(drop.view(-1, *([1] * (x.ndim - 1))), current, x) def forward(self, x: torch.Tensor) -> torch.Tensor: """ @@ -571,6 +790,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: Flatten tensor with shape [B, M, 512] """ + if self.history_len > 1: + if x.shape[1] != self.history_len: + raise ValueError( + f"MLPPolicyStem got {x.shape[1]} history steps on axis 1 but " + f"history_len={self.history_len} (input shape {tuple(x.shape)})" + ) + if self.training and self.history_dropout > 0: + x = self._drop_history(x) if self.num_of_copy > 1: out = [] iter_num = min(self.num_of_copy, x.shape[1]) @@ -581,10 +808,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = torch.stack(out, dim=1) else: y = self.net(x) + if self.history_len > 1: + # Learned time-step embedding on axis 1, last slot = current step. + y = y + self.time_embed.view( + 1, self.history_len, *([1] * (y.ndim - 3)), y.shape[-1] + ) return y -class ResNet(PolicyStem): +class ResNet(PretrainedWeights, PolicyStem): + # ``self.net`` is the ImageNet-pretrained torchvision backbone; ``self.proj`` + # is built by HPT and must keep its xavier init. + _hpt_pretrained_attrs = ("net",) + def __init__( self, output_dim: int = 10, @@ -616,6 +852,7 @@ def __init__( self.avgpool = nn.AvgPool2d(7, stride=1) # Freeze the backbone if specified + self.freeze_backbone = freeze_backbone if freeze_backbone: self._freeze_backbone() @@ -629,6 +866,27 @@ def _freeze_backbone(self): for param in self.net.parameters(): param.requires_grad = False + def train(self, mode: bool = True): + """Keep a frozen backbone in eval mode regardless of the outer flag. + + ``requires_grad = False`` alone leaves the BatchNorm layers in train + mode: they would keep normalizing with batch statistics and keep + updating their running stats, so a "frozen" backbone would not be + frozen at all. + """ + super().train(mode) + if self.freeze_backbone: + if isinstance(self.net, nn.ModuleList): + for net in self.net: + net.eval() + else: + self.net.eval() + return self + + def backbone_parameters(self): + """Pretrained-backbone parameters (``proj`` stays in the main group).""" + return list(self.net.parameters()) + def forward(self, x: torch.Tensor) -> torch.Tensor: """ Performs a forward pass of the model. @@ -676,7 +934,159 @@ def _qwen_last_token_pool( return last_hidden_states[batch_idx, seq_lens] -class _Qwen3BaseEncoder(PolicyStem): +def _local_snapshot_if_offline(model_name: str) -> str: + """Resolve a HF repo id to its cached snapshot directory when the process + runs offline (``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE``); a local path + or an online process is returned unchanged.""" + if os.path.isdir(model_name): + return model_name + # huggingface_hub parses these with _is_true, i.e. 1 / ON / YES / TRUE in + # any case, and reads them once at import; match the values but read the + # live environment, since a launcher or a test may set them after import. + from huggingface_hub.constants import ENV_VARS_TRUE_VALUES + + offline = any( + (os.environ.get(var) or "").upper() in ENV_VARS_TRUE_VALUES + for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ) + if not offline: + return model_name + from huggingface_hub import snapshot_download + + return snapshot_download(model_name, local_files_only=True) + + +class DINOv3Stem(PretrainedWeights, PolicyStem): + """Pretrained DINOv3 ViT (timm), a drop-in replacement for ``ResNet``: + ``(B, T, N, 3, H, W)`` images in [0, 1] in, ``(B, patches, output_dim)`` out. + The stem owns the resize and the checkpoint's mean / std, so the config + must NOT normalize these images in its augs. + + Loaded through timm rather than transformers: the ``timm/*_dinov3`` hub + repos are not gated (``facebook/dinov3-*`` is), and DINOv3 in transformers + needs >= 4.56, above our pin. + + Args: + model_name: timm model id, e.g. ``vit_base_patch16_dinov3.lvd1689m`` + (86M) or ``vit_large_patch16_dinov3.lvd1689m`` (300M). + output_dim: width of the projected tokens. + freeze_backbone: freeze the tower, keep it in eval mode and run it + under ``no_grad``. + image_size: ``S`` or ``[H, W]``, multiples of the patch size. DINOv3 is + RoPE-based, so non-square inputs need no position interpolation. + pretrained: ``False`` builds a randomly initialised tower (tests). + tower_kwargs: timm architecture overrides (tests). + pool: ``"patch"`` (every patch token) or ``"cls"`` (one token per image). + grad_checkpointing: recompute the tower's blocks in backward, for a + fine-tuned tower over many frames. + """ + + DEFAULT_MODEL = "vit_base_patch16_dinov3.lvd1689m" + _hpt_pretrained_attrs = ("tower",) + + def __init__( + self, + model_name: str = DEFAULT_MODEL, + output_dim: int = 10, + freeze_backbone: bool = True, + image_size=256, + pretrained: bool = True, + tower_kwargs: Optional[dict] = None, + pool: str = "patch", + grad_checkpointing: bool = False, + **kwargs, + ) -> None: + super().__init__(**kwargs) + import timm + + if pool not in ("patch", "cls"): + raise ValueError(f"DINOv3Stem pool must be 'patch' or 'cls', got {pool!r}") + self.pool = pool + + self.model_name = model_name + self.pretrained = pretrained + self.tower = timm.create_model( + model_name, pretrained=pretrained, num_classes=0, **dict(tower_kwargs or {}) + ) + self.tower.set_grad_checkpointing(grad_checkpointing) + self.hidden_size = int(self.tower.embed_dim) + self.num_prefix_tokens = int(self.tower.num_prefix_tokens) + size = [image_size] * 2 if isinstance(image_size, int) else list(image_size) + patch = self.tower.patch_embed.patch_size[0] + if any(s % patch for s in size): + raise ValueError( + f"DINOv3Stem image_size {size} is not a multiple of {patch}" + ) + self.image_size = tuple(int(s) for s in size) + self.grid_size = tuple(s // patch for s in self.image_size) + cfg = self.tower.pretrained_cfg + self.register_buffer( + "image_mean", torch.tensor(cfg["mean"]).view(1, 3, 1, 1), persistent=False + ) + self.register_buffer( + "image_std", torch.tensor(cfg["std"]).view(1, 3, 1, 1), persistent=False + ) + + self.out_dim = output_dim + self.proj = nn.Linear(self.hidden_size, output_dim) + + self.freeze_backbone = freeze_backbone + if freeze_backbone: + for param in self.tower.parameters(): + param.requires_grad = False + + def train(self, mode: bool = True): + """Keep a frozen tower in eval mode regardless of the outer flag.""" + super().train(mode) + if self.freeze_backbone: + self.tower.eval() + return self + + def backbone_parameters(self): + """Pretrained-tower parameters (``proj`` stays in the main group).""" + return list(self.tower.parameters()) + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The cached timm checkpoint, keyed ``tower.*``. ``None`` for a random + tower, or when the file is not in the local HF cache.""" + if not self.pretrained: + return None + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + try: + path = hf_hub_download( + self.tower.pretrained_cfg["hf_hub_id"], + "model.safetensors", + local_files_only=True, + ) + except Exception: # not cached, or not an HF-hosted timm model + return None + return {f"tower.{k}": v for k, v in load_file(path).items()} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, *_, H, W = x.shape + x = x.reshape(-1, 3, H, W) + x = F.interpolate( + x, + size=self.image_size, + mode="bilinear", + align_corners=False, + antialias=True, + ) + x = (x - self.image_mean) / self.image_std + with torch.set_grad_enabled( + not self.freeze_backbone and torch.is_grad_enabled() + ): + feat = self.tower.forward_features(x) + if self.pool == "cls": + feat = feat[:, :1] + else: + feat = feat[:, self.num_prefix_tokens :] + return self.proj(feat.reshape(B, -1, self.hidden_size)) + + +class _Qwen3BaseEncoder(PretrainedWeights, PolicyStem): """Shared base for Qwen3-Embedding stems used by HPT. Owns the tokenizer + transformer encoder so HPT.process_batch_for_training @@ -693,9 +1103,19 @@ class _Qwen3BaseEncoder(PolicyStem): cost down when frozen). normalize_pooled: only used by the pooled subclass; L2-normalizes the sentence embedding (Qwen3 official recipe). + cache_text_features: cache the frozen per-token hidden states per + unique prompt string (default True). Ignored when ``freeze`` is + False - a trainable encoder must stay online. + text_cache_max_entries: cache capacity in ENTRIES; LRU eviction beyond + it. An entry is the prompt's un-padded (L, D) hidden states, so + its size follows the prompt: 4-25 KB for a 2-12 token prompt at + 1024-dim fp16, up to max_length * D * 2 bytes (262 KB at + max_length 128). The 8192 default is therefore ~100-200 MB of + device memory on real annotations and 2 GB in the worst case. """ DEFAULT_MODEL = "Qwen/Qwen3-Embedding-0.6B" + _hpt_pretrained_attrs = ("encoder",) def __init__( self, @@ -704,6 +1124,8 @@ def __init__( freeze: bool = True, dtype: str = "float16", output_dim: int | None = None, + cache_text_features: bool = True, + text_cache_max_entries: int = 8192, **kwargs, ) -> None: super().__init__(**kwargs) @@ -713,8 +1135,18 @@ def __init__( self.max_length = max_length self.freeze_encoder = freeze torch_dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype - self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") - self.encoder = AutoModel.from_pretrained(model_name, torch_dtype=torch_dtype) + # transformers >= 4.57 probes the Hub while loading a tokenizer by repo + # id even under HF_HUB_OFFLINE=1 (tokenization_utils_base + # ._patch_mistral_regex -> model_info), which raises on compute nodes + # with no internet. Loading from the cached snapshot DIRECTORY skips + # that probe, so resolve the repo id to its local snapshot when offline. + load_path = _local_snapshot_if_offline(model_name) + self.tokenizer = AutoTokenizer.from_pretrained(load_path, padding_side="left") + self.encoder = AutoModel.from_pretrained(load_path, torch_dtype=torch_dtype) + # Remembered so the weight-hash check can read the same snapshot back + # and compare in the dtype the encoder was loaded in. + self._snapshot_dir = load_path + self._load_dtype = torch_dtype if freeze: for p in self.encoder.parameters(): p.requires_grad = False @@ -730,8 +1162,47 @@ def __init__( if self.output_dim != self.hidden_size else nn.Identity() ) + self._init_text_feature_cache(cache_text_features, text_cache_max_entries) - def _encode(self, prompts): + # ---------------------------------------------------------------- + # frozen text-feature cache + # ---------------------------------------------------------------- + + def _init_text_feature_cache( + self, cache_text_features: bool = True, text_cache_max_entries: int = 8192 + ) -> None: + """Set up the (empty) prompt -> per-token-features cache.""" + self.cache_text_features = bool(cache_text_features) + self.text_cache_max_entries = int(text_cache_max_entries) + self._text_feature_cache = OrderedDict() + self._text_cache_hits = 0 + self._text_cache_misses = 0 + + def clear_text_feature_cache(self) -> None: + """Drop every cached prompt (the hit/miss counters are cumulative).""" + self._text_feature_cache.clear() + + def text_feature_cache_stats(self) -> dict: + """``{"entries": n, "hits": h, "misses": m}``; hits/misses count prompt + occurrences, duplicates inside one batch included.""" + return { + "entries": len(self._text_feature_cache), + "hits": self._text_cache_hits, + "misses": self._text_cache_misses, + } + + def _text_cache_enabled(self) -> bool: + # A trainable encoder must stay online: its outputs change every step + # and the loss needs the graph. Only a frozen encoder is cacheable. + return bool( + self.cache_text_features + and self.freeze_encoder + and self.text_cache_max_entries > 0 + ) + + def _encode_online(self, prompts): + """Tokenize + run the encoder, returning ``(last_hidden_state, + attention_mask)`` left-padded (the tokenizer is ``padding_side="left"``).""" tokens = self.tokenizer( prompts, padding=True, @@ -746,6 +1217,95 @@ def _encode(self, prompts): out = self.encoder(**tokens) return out.last_hidden_state, tokens["attention_mask"] + def _encode(self, prompts): + """``(last_hidden_state, attention_mask)`` for ``prompts``, served from + the per-prompt cache whenever the encoder is frozen. + + The annotation vocabulary is tiny next to the number of training steps, + so a frozen encoder recomputes the same 600M-parameter forward over and + over. Misses go through ``_encode_online`` in ONE sub-batch of unique + strings; each prompt's un-padded ``(L_i, D)`` rows are stored (the mask + is implied by ``L_i``) and the batch is re-assembled with the same LEFT + padding the tokenizer produces, so callers cannot tell the paths apart. + + One deliberate difference: padded positions come back as ZEROS here, + whereas the online path leaves the encoder's output for the pad tokens + there. Every consumer masks those positions out + (``forward_with_mask``; last-token pooling reads the final, always-real + position under left padding), so this only removes junk. + """ + if not self._text_cache_enabled(): + return self._encode_online(prompts) + + prompts = list(prompts) + device = self.device + misses = [] + for prompt in prompts: + entry = self._text_feature_cache.get(prompt) + # Entries computed on another device are stale (the module moved): + # drop them rather than silently copy across devices. + if entry is not None and entry.device != device: + del self._text_feature_cache[prompt] + entry = None + if entry is None: + self._text_cache_misses += 1 + if prompt not in misses: + misses.append(prompt) + else: + self._text_cache_hits += 1 + self._text_feature_cache.move_to_end(prompt) + + fresh = {} + if misses: + hidden, mask = self._encode_online(misses) + mask = mask.bool() + for i, prompt in enumerate(misses): + # un-padded per-token states; detached, no graph is kept + fresh[prompt] = hidden[i][mask[i]].detach() + self._text_feature_cache[prompt] = fresh[prompt] + self._text_feature_cache.move_to_end(prompt) + while len(self._text_feature_cache) > self.text_cache_max_entries: + self._text_feature_cache.popitem(last=False) # LRU victim + + # Read this batch's rows from what was just encoded first: a batch with + # more unique prompts than text_cache_max_entries evicts its own + # earliest entries in the trim above, and a KeyError here would be a + # crash rather than a slow path. + rows = [ + fresh[p] if p in fresh else self._text_feature_cache[p] for p in prompts + ] + longest = max(row.shape[0] for row in rows) + hidden = rows[0].new_zeros((len(rows), longest, rows[0].shape[-1])) + mask = torch.zeros((len(rows), longest), dtype=torch.long, device=device) + for i, row in enumerate(rows): # LEFT padding, as padding_side="left" + length = row.shape[0] + hidden[i, longest - length :] = row + mask[i, longest - length :] = 1 + return hidden, mask + + def _apply(self, *args, **kwargs): + """``.to()`` / ``.cuda()`` / ``.float()`` invalidate the cache: the + stored tensors belong to the old device / dtype.""" + out = super()._apply(*args, **kwargs) + if getattr(self, "_text_feature_cache", None): + self.clear_text_feature_cache() + return out + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The snapshot's own safetensors, cast exactly the way + ``from_pretrained(torch_dtype=...)`` casts them.""" + if not os.path.isdir(self._snapshot_dir): + return None + return _snapshot_state_dict( + self._snapshot_dir, prefix="encoder.", dtype=self._load_dtype + ) + + def pretrained_hash_dtype(self) -> Optional[torch.dtype]: + """Hash in the load dtype. ``HPT.__init__`` upcasts the whole model with + ``nets.float()``; fp16 -> fp32 -> fp16 is exact, so casting the live + weights back down compares like for like without a second model load.""" + return self._load_dtype + def train(self, mode: bool = True): """Keep the frozen encoder in eval mode regardless of outer train flag.""" super().train(mode) @@ -777,18 +1337,26 @@ class QwenPerTokenEncoder(_Qwen3BaseEncoder): information from the learnable stem tokens. """ - def forward(self, prompts): + def forward_with_mask(self, prompts): + """Returns ``(features, attention_mask)``; the mask is what padded + positions must be excluded by, since zeroing them before the biased + ``proj`` turns them into the learned constant ``proj.bias``.""" hidden, mask = self._encode(prompts) feat = hidden.float() * mask.unsqueeze(-1).float() - return self.proj(feat) # (B, L, output_dim) + return self.proj(feat), mask.bool() # (B, L, output_dim), (B, L) + + def forward(self, prompts): + return self.forward_with_mask(prompts)[0] def compute_latent(self, prompts): - feat = self(prompts) # (B, L, hidden_size) + feat, mask = self.forward_with_mask(prompts) # (B, L, output_dim), (B, L) stem_tokens = self.tokens.repeat(feat.shape[0], 1, 1) - return self.cross_attention(stem_tokens, feat) + return self.cross_attention(stem_tokens, feat, mask=mask) + +class T5Encoder(PretrainedWeights, PolicyStem): + _hpt_pretrained_attrs = ("encoder",) -class T5Encoder(PolicyStem): def __init__(self, per_token=True, **kwargs) -> None: """T5 Encoder that expects pre-tokenized inputs diff --git a/egomimic/models/image_augs.py b/egomimic/models/image_augs.py new file mode 100644 index 0000000000..e95e12b39e --- /dev/null +++ b/egomimic/models/image_augs.py @@ -0,0 +1,166 @@ +"""Per-sample image augmentation. + +torchvision's random transforms draw one parameter set per call, so applying a +``ColorJitter`` to a stacked (B, 3, H, W) batch jitters every image in the batch +by the same factors (and in the same op order). ``PerSampleAugs`` gives every +sample its own parameters. + +The obvious way to do that -- call the transform once per image -- costs one +Python round trip and a few dozen tiny CUDA launches *per image*: measured at +89 ms of a 196 ms HPT train step at batch 64 (H100, 2026-09-19), i.e. 46% of +the step, because every factor is a CPU scalar that has to be copied to the +device. ``BatchedColorJitter`` draws the factors as one on-device (B,) tensor +and runs each adjustment once over the whole batch instead. +""" + +from typing import List, Optional, Tuple + +import torch +from torch import nn +from torchvision import transforms as T +from torchvision.transforms import _functional_tensor as _ft + +_GRAY_W = (0.2989, 0.587, 0.114) + + +def _rgb_to_gray(img: torch.Tensor) -> torch.Tensor: + w = torch.tensor(_GRAY_W, dtype=img.dtype, device=img.device).view(3, 1, 1) + return (img * w).sum(dim=-3, keepdim=True) + + +def _blend(img1: torch.Tensor, img2: torch.Tensor, ratio: torch.Tensor) -> torch.Tensor: + return (ratio * img1 + (1.0 - ratio) * img2).clamp(0.0, 1.0) + + +class BatchedColorJitter(nn.Module): + """``torchvision.transforms.ColorJitter`` with per-sample factors, batched. + + Same adjustments and same factor distributions as torchvision, drawn once + per image instead of once per call. The op ORDER is drawn once per call + (torchvision's own behaviour), not once per image as the per-image loop + did; each adjustment clamps to [0, 1], so the order is not quite inert, but + it is a permutation of the same four factors either way. + + Input is (B, 3, H, W) float in [0, 1]. With ``frames`` > 1 it is + (B * frames, 3, H, W) with a sample's frames adjacent: they share one draw, + and every adjustment (contrast's mean included) still runs per frame. + """ + + def __init__( + self, + brightness: Optional[Tuple[float, float]] = None, + contrast: Optional[Tuple[float, float]] = None, + saturation: Optional[Tuple[float, float]] = None, + hue: Optional[Tuple[float, float]] = None, + ): + super().__init__() + self.brightness = brightness + self.contrast = contrast + self.saturation = saturation + self.hue = hue + + @classmethod + def from_color_jitter(cls, jitter: T.ColorJitter) -> "BatchedColorJitter": + return cls(jitter.brightness, jitter.contrast, jitter.saturation, jitter.hue) + + @staticmethod + def _factors(bounds, n, img) -> torch.Tensor: + # Drawn in float32: uniform_ into bfloat16 yields ~40 distinct values, + # and into uint8 it raises. + lo, hi = bounds + f = torch.empty(n, 1, 1, 1, dtype=torch.float32, device=img.device) + return f.uniform_(lo, hi).to(img.dtype) + + def _apply_brightness(self, img, f): + return (img * f).clamp(0.0, 1.0) + + def _apply_contrast(self, img, f): + mean = _rgb_to_gray(img).mean(dim=(-3, -2, -1), keepdim=True) + return _blend(img, mean, f) + + def _apply_saturation(self, img, f): + return _blend(img, _rgb_to_gray(img), f) + + def _apply_hue(self, img, f): + hsv = _ft._rgb2hsv(img) + h, s, v = hsv.unbind(dim=-3) + h = (h + f.squeeze(-3)) % 1.0 + return _ft._hsv2rgb(torch.stack((h, s, v), dim=-3)).clamp(0.0, 1.0) + + def forward(self, img: torch.Tensor, frames: int = 1) -> torch.Tensor: + n = img.shape[0] // frames + ops = [ + (self.brightness, self._apply_brightness), + (self.contrast, self._apply_contrast), + (self.saturation, self._apply_saturation), + (self.hue, self._apply_hue), + ] + for idx in torch.randperm(len(ops)).tolist(): + bounds, fn = ops[idx] + if bounds is None: + continue + img = fn(img, self._factors(bounds, n, img).repeat_interleave(frames, 0)) + return img + + def extra_repr(self) -> str: + return ( + f"brightness={self.brightness}, contrast={self.contrast}, " + f"saturation={self.saturation}, hue={self.hue}" + ) + + +# Transforms that are already batch-safe AND sample no parameters, so running +# them on the stacked batch is identical to running them per image. +_DETERMINISTIC = (T.Normalize, T.Resize, T.CenterCrop, T.ConvertImageDtype) + + +def _vectorize(augs) -> Optional[List[nn.Module]]: + """The batched equivalent of ``augs``, or None if any member has no batched + form (then the caller keeps the per-image loop).""" + members = augs.transforms if isinstance(augs, T.Compose) else [augs] + out: List[nn.Module] = [] + for t in members: + if isinstance(t, T.ColorJitter): + out.append(BatchedColorJitter.from_color_jitter(t)) + elif isinstance(t, _DETERMINISTIC): + out.append(t) + else: + return None + return out + + +class PerSampleAugs(nn.Module): + """Apply ``augs`` to each image of a (B, C, H, W) batch independently. + + ``augs`` is the configured callable (typically a + ``torchvision.transforms.Compose``); a deterministic one is unaffected by + the wrapping, so this is safe for any existing aug list. + + ``frames`` > 1 takes (B * frames, C, H, W), each sample's frames adjacent: + the frames of a sample share one parameter draw, and every transform still + sees one frame at a time, so a crop, resize or contrast mean never spans + two frames. + """ + + def __init__(self, augs): + super().__init__() + self.augs = augs + vectorized = _vectorize(augs) + self.vectorized = nn.ModuleList(vectorized) if vectorized is not None else None + + def forward(self, images, frames: int = 1): + if self.vectorized is None: + if frames == 1: + return torch.stack([self.augs(image) for image in images]) + # torchvision draws once per call and works per image of a stack + return torch.cat([self.augs(clip) for clip in images.split(frames)]) + for t in self.vectorized: + images = ( + t(images, frames=frames) + if isinstance(t, BatchedColorJitter) + else t(images) + ) + return images + + def extra_repr(self): + return repr(self.augs) diff --git a/egomimic/models/openpi_compat.py b/egomimic/models/openpi_compat.py new file mode 100644 index 0000000000..ce6f5e15a3 --- /dev/null +++ b/egomimic/models/openpi_compat.py @@ -0,0 +1,129 @@ +"""Adapters that let openpi's PyTorch pi0 run against transformers 5. + +`external/openpi` pins `transformers==4.53.2` and ships a `transformers_replace` +patch written against it (its own `check.py` asserts that version), which is also +this repo's pin. A venv on transformers 5 works too, but there the stock modules +have moved on, so a few of openpi's call sites no longer line up. Nothing here +runs below transformers 5, and each shim is a no-op when it finds openpi's own +spelling already in place -- so porting the fork's patch forward makes this +module inert. Delete it then. + +Applied from `egomimic.algo.pi`, which is the only importer of openpi. +""" + +import logging + +import torch + +logger = logging.getLogger(__name__) + +_APPLIED = False + + +def _patch_embed_image() -> bool: + """`get_image_features` returns a `BaseModelOutputWithPooling`, not a tensor. + + transformers 5's `PaliGemmaModel.get_image_features` returns the vision + output object with the projected features stashed on `pooler_output`, while + openpi's `embed_prefix` reads `.shape` off the result. Unwrap it. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.embed_image + + def embed_image(self, image: torch.Tensor) -> torch.Tensor: + features = original(self, image) + if isinstance(features, torch.Tensor): + return features + return features.pooler_output + + PaliGemmaWithExpertModel.embed_image = embed_image + return True + + +def _patch_paligemma_bc_properties() -> bool: + """`paligemma.language_model` / `.vision_tower` / `.multi_modal_projector`. + + transformers 5 moved those submodules under `PaliGemmaModel` and dropped the + back-compat properties that used to forward to them from + `PaliGemmaForConditionalGeneration`. openpi still reads all three off the + outer model. Put the forwarders back. + """ + from transformers import PaliGemmaForConditionalGeneration + + added = [] + for name in ("language_model", "vision_tower", "multi_modal_projector"): + if hasattr(PaliGemmaForConditionalGeneration, name): + continue + setattr( + PaliGemmaForConditionalGeneration, + name, + property(lambda self, _name=name: getattr(self.model, _name)), + ) + added.append(name) + return bool(added) + + +# openpi keeps SigLIP's patch and position embeddings in float32 by substring on +# `vision_tower.vision_model.embeddings.*`; the flattened 5.x spelling of the same +# three tensors. +_FP32_KEEP_FLAT = ( + "vision_tower.embeddings.patch_embedding.", + "vision_tower.embeddings.position_embedding.", +) + + +def _patch_fp32_keep_list() -> bool: + """openpi's float32 keep-list misses the SigLIP embeddings. + + `to_bfloat16_for_selected_params` casts everything to bf16 and then restores + float32 on names containing `vision_tower.vision_model.embeddings.*`. + transformers 5 flattened `SiglipVisionModel`, so nothing matches and the + patch embedding and position embedding stay bf16 -- the tensors openpi + protects because small updates vanish in bf16. Re-apply the keep under the + flattened names; with the nested names present they match nothing. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + + def to_bfloat16_for_selected_params(self, precision="bfloat16"): + original(self, precision) + if precision != "bfloat16": + return + for name, param in self.named_parameters(): + if any(selector in name for selector in _FP32_KEEP_FLAT): + param.data = param.data.to(dtype=torch.float32) + + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = ( + to_bfloat16_for_selected_params + ) + return True + + +def apply() -> None: + """Install every shim once. Safe to call repeatedly.""" + global _APPLIED + if _APPLIED: + return + _APPLIED = True + + import transformers + + major = int(transformers.__version__.split(".")[0]) + if major < 5: + return + + shims = ( + ("embed_image", _patch_embed_image), + ("paligemma_bc_properties", _patch_paligemma_bc_properties), + ("fp32_keep_list", _patch_fp32_keep_list), + ) + applied = [name for name, fn in shims if fn()] + if applied: + logger.info( + "openpi targets transformers 4.53 and this environment has %s; " + "applied compatibility shims: %s", + transformers.__version__, + ", ".join(applied), + ) diff --git a/egomimic/models/paligemma_init.py b/egomimic/models/paligemma_init.py new file mode 100644 index 0000000000..55799580ea --- /dev/null +++ b/egomimic/models/paligemma_init.py @@ -0,0 +1,296 @@ +"""PaliGemma-only initialization for the pi0.5 chassis. + +``egomimic.algo.pi.PI`` normally starts from the full pi0.5 base checkpoint. +This module is the other entry point: the VLM half (SigLIP tower, multimodal +projector, Gemma-2B prefix) comes from the public PaliGemma release, and +everything pi0.5 adds on top -- the action expert, the action in/out +projections, the adaRMS time MLPs -- keeps its PyTorch init. + +openpi builds the prefix as a stock ``PaliGemmaForConditionalGeneration`` whose +config is the ``gemma_2b`` variant (width 2048, depth 18, head_dim 256) next to +a SigLIP-so400m/14 tower, which is exactly the ``paligemma-3b-*-224`` +architecture, so the released checkpoint drops straight in. +""" + +import hashlib +import logging +import os +import re +from pathlib import Path + +import torch +from safetensors import safe_open + +logger = logging.getLogger(__name__) + +CANONICAL_REPO = "google/paligemma-3b-pt-224" + +# `CANONICAL_REPO` is license-gated: a machine without an HF token that has +# accepted the Gemma terms gets a 401 on the weight shards (the repo *metadata* +# stays anonymously readable, which is what makes the check below possible). +# These repos carry the same three shards. Every one is sha256-checked against +# the canonical repo's file metadata before it is used, so a mirror cannot +# substitute different weights -- it is only a transport. Opt-in +# (`paligemma_allow_mirror`): integrity is checked, but the weights of record +# should come from the canonical repo unless a run says otherwise. +MIRROR_REPOS = ( + "leo009/paligemma-3b-pt-224", + "hehe156/paligemma-3b-pt-224", +) + +# Only the shards: the loader reads nothing else and openpi builds the config in +# code, so an unverified json would only sit in the HF cache. +_ALLOW_PATTERNS = ["*.safetensors"] + +# Released PaliGemma checkpoint keys -> the module tree. These are what +# transformers applies via +# ``PaliGemmaForConditionalGeneration._checkpoint_conversion_mapping``; spelled +# out here because we stream the shards ourselves rather than going through +# ``from_pretrained`` (which would build a second full-size model alongside the +# one openpi already made). The vision rule is the transformers >= 5 spelling, +# which flattened ``SiglipVisionModel``; below 5 the ``vision_model.`` hop is +# still there, and ``_module_name`` picks whichever the live tree has. +# +# ``load_paligemma_weights`` fails loudly if these leave any parameter +# uncovered, so a future rename surfaces as an error, not a silent partial load. +_KEY_REMAP = ( + (re.compile(r"^language_model\.model\."), "model.language_model."), + (re.compile(r"^language_model\.lm_head\."), "lm_head."), + (re.compile(r"^vision_tower\.vision_model\."), "model.vision_tower."), + (re.compile(r"^vision_tower\."), "model.vision_tower."), + (re.compile(r"^multi_modal_projector\."), "model.multi_modal_projector."), +) + + +# The released checkpoint pads the token embedding from PaliGemma's 257152-token +# vocabulary up to 257216 rows. Row 257152 is -- which openpi never looks +# up, because ``embed_prefix`` concatenates SigLIP features directly instead of +# scattering them into an placeholder -- and the 63 rows after it are +# unused padding. openpi builds the table at exactly 257152, so these two +# tensors are truncated on their vocabulary axis; every real text token keeps its +# row, since the padding is at the end. +_VOCAB_TRUNCATABLE = frozenset( + {"model.language_model.embed_tokens.weight", "lm_head.weight"} +) + + +def _remap_key(key: str) -> str: + for pattern, replacement in _KEY_REMAP: + new, n = pattern.subn(replacement, key, count=1) + if n: + return new + return key + + +def _module_name(key: str, own: dict) -> str: + """Where checkpoint ``key`` lives in ``own``, the model's named tensors.""" + name = _remap_key(key) + if name not in own and key.startswith("vision_tower.vision_model."): + nested = "model." + key + if nested in own: + return nested + return name + + +def _fit(name: str, tensor: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + if tensor.shape == param.shape: + return tensor + if ( + name in _VOCAB_TRUNCATABLE + and tensor.ndim == param.ndim == 2 + and tensor.shape[0] > param.shape[0] + and tensor.shape[1] == param.shape[1] + ): + logger.info( + "Truncating %s from %d to %d vocabulary rows (padding tokens).", + name, + tensor.shape[0], + param.shape[0], + ) + return tensor[: param.shape[0]] + raise ValueError( + f"Shape mismatch for {name}: checkpoint {tuple(tensor.shape)} " + f"vs model {tuple(param.shape)}" + ) + + +def select_init_source( + pytorch_weight_path: str | None, paligemma_weight_path: str | None +) -> str: + """Which set of starting weights a PI run uses: the full pi0.5 base + checkpoint (``"pi05_base"``), PaliGemma with a fresh action expert + (``"paligemma"``), or nothing at all (``"none"``).""" + if pytorch_weight_path is not None and paligemma_weight_path is not None: + raise ValueError( + "pytorch_weight_path and paligemma_weight_path are both set. The " + "pi0.5 base checkpoint already contains a trained PaliGemma; set " + "pytorch_weight_path=null to start from PaliGemma with a fresh " + "action expert." + ) + if pytorch_weight_path is not None: + return "pi05_base" + if paligemma_weight_path is not None: + return "paligemma" + return "none" + + +def _shard_digests(repo: str) -> dict[str, str]: + """``{shard filename: sha256}`` from a repo's file metadata.""" + from huggingface_hub import HfApi + + info = HfApi().model_info(repo, files_metadata=True) + digests = {} + for sibling in info.siblings: + if not sibling.rfilename.endswith(".safetensors"): + continue + lfs = getattr(sibling, "lfs", None) + sha = ( + lfs.get("sha256") if isinstance(lfs, dict) else getattr(lfs, "sha256", None) + ) + if sha: + digests[sibling.rfilename] = sha + return digests + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(16 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def _verify_shards(directory: str, expected: dict[str, str]) -> None: + root = Path(directory) + for name, sha in expected.items(): + shard = root / name + if not shard.is_file(): + raise FileNotFoundError(f"{shard} is missing from the mirror snapshot") + actual = _sha256(shard) + if actual != sha: + raise ValueError( + f"{shard} does not match {CANONICAL_REPO}: expected sha256 {sha}, " + f"got {actual}. Refusing to initialize from it." + ) + + +def resolve_paligemma_dir(source: str, *, allow_mirror: bool = False) -> str: + """Local directory holding the PaliGemma checkpoint named by ``source``. + + ``source`` is a local path (returned as-is) or a Hugging Face repo id. When + it is the gated canonical repo, this machine is not authorized for it (401 + or 403, not a rate limit or an outage) and ``allow_mirror`` is set, fall + back to a sha256-verified mirror. + """ + if os.path.isdir(source): + return source + + from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError, HfHubHTTPError + + try: + return snapshot_download(source, allow_patterns=_ALLOW_PATTERNS) + except HfHubHTTPError as gated: + status = getattr(getattr(gated, "response", None), "status_code", None) + unauthorized = isinstance(gated, GatedRepoError) or status in (401, 403) + if not (allow_mirror and unauthorized and source == CANONICAL_REPO): + raise + logger.warning( + "%s is gated for this machine (%s) and paligemma_allow_mirror is set. " + "Falling back to a mirror, sha256-verified against %s. Set HF_TOKEN " + "to a token that has accepted the Gemma terms to pull the canonical " + "repo instead.", + CANONICAL_REPO, + type(gated).__name__, + CANONICAL_REPO, + ) + + expected = _shard_digests(CANONICAL_REPO) + if not expected: + raise RuntimeError( + f"Cannot read shard checksums for {CANONICAL_REPO}; refusing to use " + "an unverified mirror." + ) + + failures = [] + for mirror in MIRROR_REPOS: + try: + if _shard_digests(mirror) != expected: + raise ValueError("shard checksums differ from the canonical repo") + directory = snapshot_download(mirror, allow_patterns=_ALLOW_PATTERNS) + _verify_shards(directory, expected) + except Exception as exc: # try the next mirror, report them all if none works + failures.append(f"{mirror}: {type(exc).__name__}: {exc}") + continue + logger.warning( + "Initializing PaliGemma from mirror %s (all %d shards sha256-match %s).", + mirror, + len(expected), + CANONICAL_REPO, + ) + return directory + + raise RuntimeError( + f"{CANONICAL_REPO} is gated and no verified mirror worked:\n " + + "\n ".join(failures) + ) + + +def load_paligemma_weights(model, source: str, *, allow_mirror: bool = False) -> str: + """Copy a PaliGemma checkpoint into ``model``'s VLM prefix, in place. + + ``model`` is an ``openpi.models_pytorch.pi0_pytorch.PI0Pytorch``. Only + ``paligemma_with_expert.paligemma`` is touched; the action expert and the + action/time projections are left at their init. Returns the directory the + weights came from. + """ + directory = resolve_paligemma_dir(source, allow_mirror=allow_mirror) + target = model.paligemma_with_expert.paligemma + + # remove_duplicate=False so a tied lm_head.weight is reachable under both + # names; the checkpoint only carries one of them. + params = dict(target.named_parameters(remove_duplicate=False)) + # Buffers are accepted if the checkpoint ships them, but are not required: + # rotary inv_freq, Gemma's embed_scale and the SigLIP position_ids are all + # derived from the config at construction time and are absent from the + # release. + own = dict(target.named_buffers(remove_duplicate=False)) + own.update(params) + + shards = sorted(Path(directory).glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"No .safetensors under {directory}") + + loaded, unexpected = set(), [] + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + name = _module_name(key, own) + param = own.get(name) + if param is None: + unexpected.append(key) + continue + tensor = _fit(name, handle.get_tensor(key), param) + with torch.no_grad(): + param.copy_(tensor.to(param.dtype)) + loaded.add(name) + + # lm_head is tied to embed_tokens and, either way, dead weight here: no code + # path in openpi's PyTorch pi0 reads it. + tied = set(getattr(type(target), "_tied_weights_keys", None) or ()) + missing = {n for n in params if n not in loaded} - tied + if missing or unexpected: + raise RuntimeError( + f"PaliGemma load from {directory} did not cover the prefix exactly: " + f"{len(missing)} missing (e.g. {sorted(missing)[:5]}), " + f"{len(unexpected)} unexpected (e.g. {unexpected[:5]})." + ) + + logger.info( + "Initialized the VLM prefix from PaliGemma at %s (%d tensors, %d parameters). " + "The action expert and the action/time projections stay randomly initialized.", + directory, + len(loaded), + sum(p.numel() for p in target.parameters()), + ) + return directory diff --git a/egomimic/models/rdt_nets.py b/egomimic/models/rdt_nets.py new file mode 100644 index 0000000000..cacc31f9b2 --- /dev/null +++ b/egomimic/models/rdt_nets.py @@ -0,0 +1,378 @@ +"""Robotics Diffusion Transformer (RDT, https://github.com/thu-ml/RoboticsDiffusionTransformer) +building blocks. ``RDTBackbone`` is the DiT shared by every embodiment (the +policy's trunk); ``RDTDenoiser`` is one embodiment's adaptor / output layer +around it, used as the denoiser of a ``DenoisingPolicy`` head.""" + +import math +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def sincos_1d(embed_dim: int, positions) -> np.ndarray: + """(M,) positions -> (M, embed_dim) sin-cos table.""" + assert embed_dim % 2 == 0 + omega = 1.0 / 10000 ** ( + np.arange(embed_dim // 2, dtype=np.float64) / (embed_dim / 2) + ) + out = np.einsum("m,d->md", np.asarray(positions, dtype=np.float64), omega) + return np.concatenate([np.sin(out), np.cos(out)], axis=1) + + +def multimodal_sincos(embed_dim: int, lens: list) -> torch.Tensor: + """RDT's input position table: the first half of each row encodes which + segment the token belongs to, the second half its index inside it.""" + half = embed_dim // 2 + segment = sincos_1d(half, np.arange(len(lens))) + rows = [] + for idx, length in enumerate(lens): + row = np.zeros((length, embed_dim)) + row[:, :half] = segment[idx] + row[:, half:] = sincos_1d(embed_dim - half, np.arange(length)) + rows.append(row) + return torch.from_numpy(np.concatenate(rows, axis=0)).float() + + +def sincos_2d(embed_dim: int, grid_hw, ref_grid: float = 16.0) -> torch.Tensor: + """(h * w, embed_dim) raster-order table, half the channels per axis. + + Coordinates are normalized to [0, ref_grid] on both axes, so an edge or the + centre gets the same code at any resolution or aspect ratio. + """ + h, w = grid_hw + half = embed_dim // 2 + rows = np.linspace(0.0, ref_grid, h) if h > 1 else np.zeros(1) + cols = np.linspace(0.0, ref_grid, w) if w > 1 else np.zeros(1) + table = np.concatenate( + [ + np.repeat(sincos_1d(half, rows), w, axis=0), + np.tile(sincos_1d(embed_dim - half, cols), (h, 1)), + ], + axis=1, + ) + return torch.from_numpy(table).float() + + +def mlp_gelu(in_dim: int, out_dim: int, depth: int) -> nn.Sequential: + """RDT's ``mlp{depth}x_gelu`` condition adaptor.""" + layers = [nn.Linear(in_dim, out_dim)] + for _ in range(1, depth): + layers += [nn.GELU(approximate="tanh"), nn.Linear(out_dim, out_dim)] + return nn.Sequential(*layers) + + +@dataclass +class RDTConditions: + """What the denoiser is conditioned on, already at the DiT width. + + Quacks like the single conditioning tensor ``DenoisingPolicy`` expects + (``len`` / ``dtype`` / ``device``), so the FM / diffusion heads take it + unchanged. + """ + + img: torch.Tensor # (B, N_img, D) + freq: torch.Tensor # (B,) control frequency, Hz + backbone: "RDTBackbone" # the shared trunk the denoiser runs through + state: Optional[torch.Tensor] = None # (B, K, D) + lang: Optional[torch.Tensor] = None # (B, L, D) + lang_mask: Optional[torch.Tensor] = None # (B, L) bool, True = attend + memory: Optional[torch.Tensor] = None # (B, M, D), self-attended + memory_mask: Optional[torch.Tensor] = None # (B, M) bool, True = attend + + def __len__(self): + return self.img.shape[0] + + @property + def dtype(self): + return self.img.dtype + + @property + def device(self): + return self.img.device + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size), + ) + self.frequency_embedding_size = frequency_embedding_size + + def forward(self, t: torch.Tensor) -> torch.Tensor: + half = self.frequency_embedding_size // 2 + freqs = torch.exp( + -math.log(10000) + * torch.arange(half, dtype=torch.float32, device=t.device) + / half + ) + args = t[:, None].float() * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.mlp(emb.to(self.mlp[0].weight.dtype)) + + +class _Attention(nn.Module): + """Self- (``context=None``) or cross-attention with QK RMSNorm.""" + + def __init__(self, dim: int, num_heads: int): + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.q = nn.Linear(dim, dim) + self.kv = nn.Linear(dim, dim * 2) + self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6) + self.proj = nn.Linear(dim, dim) + + def forward(self, x, context=None, mask=None): + B, N, C = x.shape + context = x if context is None else context + L = context.shape[1] + q = self.q(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2) + k, v = ( + self.kv(context) + .view(B, L, 2, self.num_heads, self.head_dim) + .permute(2, 0, 3, 1, 4) + ) + q, k = self.q_norm(q), self.k_norm(k) + if mask is not None: + mask = mask.view(B, 1, 1, L) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + return self.proj(out.transpose(1, 2).reshape(B, N, C)) + + +class _Mlp(nn.Module): + def __init__(self, dim: int, hidden: int, out_dim: Optional[int] = None): + super().__init__() + self.fc1 = nn.Linear(dim, hidden) + self.act = nn.GELU(approximate="tanh") + self.fc2 = nn.Linear(hidden, out_dim or dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class RDTBlock(nn.Module): + """Pre-RMSNorm self-attention, cross-attention to one condition, FFN.""" + + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 1.0): + super().__init__() + self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6) + self.attn = _Attention(hidden_size, num_heads) + self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6) + self.cross_attn = _Attention(hidden_size, num_heads) + self.norm3 = nn.RMSNorm(hidden_size, eps=1e-6) + self.ffn = _Mlp(hidden_size, int(hidden_size * mlp_ratio)) + + def forward(self, x, c, mask=None, self_mask=None): + x = x + self.attn(self.norm1(x), mask=self_mask) + x = x + self.cross_attn(self.norm2(x), c, mask) + return x + self.ffn(self.norm3(x)) + + +def _xavier(module): + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class RDTBackbone(nn.Module): + """The embodiment-agnostic part of RDT: timestep / control-frequency + embedders and the blocks, which alternate the condition they cross-attend + to (language, image, language, ...; image only when there is no language). + """ + + def __init__( + self, + hidden_dim: int = 1024, + depth: int = 28, + n_heads: int = 16, + mlp_ratio: float = 1.0, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.t_embedder = TimestepEmbedder(hidden_dim) + self.freq_embedder = TimestepEmbedder(hidden_dim) + self.blocks = nn.ModuleList( + [RDTBlock(hidden_dim, n_heads, mlp_ratio) for _ in range(depth)] + ) + self.initialize_weights() + + def initialize_weights(self): + self.apply(_xavier) + for embedder in (self.t_embedder, self.freq_embedder): + nn.init.normal_(embedder.mlp[0].weight, std=0.02) + nn.init.normal_(embedder.mlp[2].weight, std=0.02) + + def prefix(self, t: torch.Tensor, freq: torch.Tensor, batch_size: int): + """(B, 2, D): the timestep and control-frequency tokens.""" + t = self.t_embedder(t.reshape(-1)).unsqueeze(1).expand(batch_size, -1, -1) + return torch.cat([t, self.freq_embedder(freq.reshape(-1)).unsqueeze(1)], dim=1) + + def forward( + self, + h: torch.Tensor, + cond: "RDTConditions", + self_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``self_mask`` (B, N) bool: which of ``h``'s tokens may be attended + to (padded memory steps may not).""" + conds = [(cond.img, None)] + if cond.lang is not None: + conds.insert(0, (cond.lang, cond.lang_mask)) + for i, block in enumerate(self.blocks): + c, mask = conds[i % len(conds)] + h = block(h, c, mask, self_mask) + return h + + +class RDTDenoiser(nn.Module): + """One embodiment's view of RDT: the action adaptor, the position table of + ``[timestep; ctrl freq; state tokens; memory tokens; noisy action chunk]`` + and the output layer. The blocks in between are ``cond.backbone``, shared across + embodiments. ``forward`` has the ``model(x_t, t, cond)`` signature of the + other denoising nets, with ``cond`` an ``RDTConditions``. + + ``time_scale`` maps the head's time onto RDT's [0, 1000) embedding range: + 1000 for flow matching (t in [0, 1]), 1 for a DDPM head (integer steps). + """ + + def __init__( + self, + act_dim: int, + act_seq: int, + hidden_dim: int = 1024, + n_state_tokens: int = 1, + time_scale: float = 1000.0, + n_memory_tokens: int = 0, + ): + super().__init__() + self.act_seq = act_seq + self.n_state_tokens = n_state_tokens + self.n_memory_tokens = n_memory_tokens + self.time_scale = time_scale + self.action_adaptor = mlp_gelu(act_dim, hidden_dim, depth=3) + self.x_pos_embed = nn.Parameter( + torch.zeros(1, 2 + n_state_tokens + n_memory_tokens + act_seq, hidden_dim) + ) + self.norm_final = nn.RMSNorm(hidden_dim, eps=1e-6) + self.ffn_final = _Mlp(hidden_dim, hidden_dim, act_dim) + self.initialize_weights() + + def initialize_weights(self): + """RDT's init. ``RDTModel.finalize_modules`` calls it again, after HPT's + xavier pass over every Linear.""" + self.apply(_xavier) + lens = [1, 1, self.n_state_tokens, self.n_memory_tokens, self.act_seq] + self.x_pos_embed.data.copy_( + multimodal_sincos(self.x_pos_embed.shape[-1], [n for n in lens if n > 0]) + ) + nn.init.constant_(self.ffn_final.fc2.weight, 0) + nn.init.constant_(self.ffn_final.fc2.bias, 0) + + def forward(self, x, timesteps, cond: RDTConditions, *args, **kwargs): + n_state = 0 if cond.state is None else cond.state.shape[1] + if n_state != self.n_state_tokens: + raise ValueError( + f"RDTDenoiser got {n_state} state tokens but n_state_tokens=" + f"{self.n_state_tokens}; set it to (proprio stems x history_len)" + ) + n_memory = 0 if cond.memory is None else cond.memory.shape[1] + if n_memory != self.n_memory_tokens: + raise ValueError( + f"RDTDenoiser got {n_memory} memory tokens but n_memory_tokens=" + f"{self.n_memory_tokens}; set it to the memory's frames" + ) + tokens = [cond.backbone.prefix(timesteps * self.time_scale, cond.freq, len(x))] + if n_state: + tokens.append(cond.state) + self_mask = None + if n_memory: + tokens.append(cond.memory) + B = len(x) + keep = torch.ones(B, 2 + n_state, dtype=torch.bool, device=x.device) + self_mask = torch.cat( + [keep, cond.memory_mask, keep.new_ones(B, self.act_seq)], dim=1 + ) + tokens.append(self.action_adaptor(x)) + h = torch.cat(tokens, dim=1) + self.x_pos_embed + h = cond.backbone(h, cond, self_mask) + return self.ffn_final(self.norm_final(h))[:, -self.act_seq :] + + +class RDTMemory(nn.Module): + """Long-range visual memory: ``frames`` frames of one camera ``stride_s`` + apart, the newest the current one, one token each for the DiT's + self-attention prefix. + + ``encoder`` must pool each frame to a single token (``DINOv3Stem`` with + ``pool="cls"``); a fixed sin-cos code of the frame's age in seconds tags + it. ``fuse_proprio`` adds a projection of that step's proprio token, which + needs the proprio history on the same grid, current step last. + """ + + def __init__( + self, + encoder: nn.Module, + hidden_dim: int, + frames: int = 10, + stride_s: float = 1.0, + fuse_proprio: bool = False, + ): + super().__init__() + self.encoder = encoder + self.frames = int(frames) + self.fuse_proprio = bool(fuse_proprio) + if self.fuse_proprio: + self.proprio_proj = nn.Linear(hidden_dim, hidden_dim) + # ages in 0.1 s units, oldest first, so 0.1 s apart is one sin-cos step + age_s = float(stride_s) * np.arange(self.frames - 1, -1, -1) + self.register_buffer( + "age_embed", + torch.from_numpy(sincos_1d(hidden_dim, age_s * 10.0)).float(), + persistent=False, + ) + + @property + def n_tokens(self) -> int: + return self.frames + + def forward( + self, + images: torch.Tensor, + mask: torch.Tensor, + state: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(B, N, 3, H, W)`` frames, oldest first, and their ``(B, N)`` + validity -> ``(B, N, D)`` tokens and their ``(B, N)`` bool mask.""" + B, N = images.shape[:2] + if N != self.frames: + raise ValueError( + f"memory window has {N} frames but trunk.memory.frames=" + f"{self.frames}; set it to the data config's key_map.image_memory" + ) + tokens = self.encoder(images) + if tokens.shape[1] != N: + raise ValueError( + f"the memory encoder gave {tokens.shape[1] // N} tokens per frame, " + "not 1; use a DINOv3Stem with pool: cls" + ) + tokens = tokens + self.age_embed.to(tokens) + if self.fuse_proprio: + if state is None or state.shape[1] < N: + got = None if state is None else state.shape[1] + raise ValueError( + f"fuse_proprio needs >= {N} proprio steps on the memory's grid " + f"(got {got}); set key_map.proprio_history and " + "history_stride_s to match the memory" + ) + tokens = tokens + self.proprio_proj(state[:, -N:]) + return tokens, mask.bool() diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index e16a1a1b13..d7fd4ae114 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -1,11 +1,68 @@ import logging from lightning import LightningDataModule -from lightning.pytorch.utilities.combined_loader import CombinedLoader +from lightning.pytorch.utilities.combined_loader import CombinedLoader, _MaxSizeCycle from torch.utils.data import DataLoader, default_collate +from egomimic.utils.gpu_orphans import orphan_guarded + logger = logging.getLogger(__name__) +# Val heads that can carry a metric loader, in val_dataloader() order. +VAL_HEADS = ("valid", "train_viz", "unseen_op_valid") +# Suffix of the video-only companion loader of a head (``valid`` -> +# ``valid_video``). ModelWrapper.validation_step splits on it. +VIDEO_SUFFIX = "_video" + + +def video_loader_name(head: str) -> str: + return f"{head}{VIDEO_SUFFIX}" + + +def head_of_loader(name: str) -> tuple[str, bool]: + """``("valid", False)`` for a metric loader, ``("valid", True)`` for its + video-only companion.""" + if name.endswith(VIDEO_SUFFIX): + return name[: -len(VIDEO_SUFFIX)], True + return name, False + + +class _VizGatedLoader(CombinedLoader): + """A val loader that yields nothing when its gate is closed. + + ``ModelWrapper.validation_step`` already returns before the forward for a + ``_video`` loader it does not want, but Lightning has built the batch by + then, so the pinned frames are decoded and discarded anyway. + + The val loaders are built once (``_EvaluationLoop.setup_data`` returns early + once ``_combined_loader`` is set), so the gate cannot be a length; it is an + empty iterator. ``_Sequential`` already reads an early ``StopIteration`` from + one dataloader as "move to the next". + """ + + def __init__(self, iterables, mode, gate): + if mode != "max_size_cycle": + raise ValueError(f"only max_size_cycle is gated, not {mode!r}") + super().__init__(iterables, mode) + self._gate = gate + + def __iter__(self): + if not self._gate(): + return iter(()) + return super().__iter__() + + def __len__(self) -> int: + if self._iterator is not None: + return super().__len__() + # `CombinedLoader.__len__` reads a live iterator, and `__iter__` is what + # forks this loader's workers and starts them prefetching. Lightning asks + # for the length once per rank in `setup_data`, where the gate is already + # closed on every rank but 0, so answering WITHOUT an iterator is what + # keeps those ranks off the pinned episode. The length does not depend on + # the gate. `_MaxSizeCycle` (the mode is pinned in __init__) only stores + # its iterables until __iter__, so asking it forks nothing. + return len(_MaxSizeCycle(self.flattened, self.limits)) + class MultiDataModuleWrapper(LightningDataModule): """ @@ -20,6 +77,15 @@ def __init__( valid_datasets: dict, train_dataloader_params: dict, valid_dataloader_params: dict, + train_viz_datasets: dict | None = None, + train_viz_dataloader_params: dict | None = None, + unseen_op_valid_datasets: dict | None = None, + unseen_op_valid_dataloader_params: dict | None = None, + valid_prefix: str | None = None, + held_out_operators: list | None = None, + video_datasets: dict | None = None, + metric_frames_per_episode: dict | None = None, + video_episodes: dict | None = None, ): """ Args: @@ -27,10 +93,48 @@ def __init__( valid_datasets: dictionary of valid datasets train_dataloader_params: dictionary of train dataloader parameters valid_dataloader_params: dictionary of valid dataloader parameters + train_viz_datasets: optional dict of datasets iterated like a + second val loader. Used by TrainVizEvalVideo to visualize the + policy on training data alongside the canonical validation. + train_viz_dataloader_params: dict of per-dataset DataLoader kwargs + for the train_viz loader. + unseen_op_valid_datasets: optional dict of datasets iterated as a + third val loader, evaluated by a prefixed copy of the canonical + evaluator (metrics ``unseen_op_valid/``). The held-out-operator split configs put the + unseen operators here, so ``valid_datasets`` can hold the seen + operators' held-out episodes and ``train_viz`` the train split. + unseen_op_valid_dataloader_params: dict of per-dataset DataLoader + kwargs for the unseen_op_valid loader. + valid_prefix: config-only. When set, trainHydra wraps the canonical + evaluator so the valid loader logs ``/...`` and writes + ``videos_/`` (``seen_op_valid`` in the opsplit configs). + held_out_operators: config-only. The held-out-operator split + configs (data/mecka_fold_*_opsplit_*.yaml) keep the operator + id list once at the data-config root and interpolate it from + the train/valid filter lambdas; hydra.instantiate(cfg.data) + forwards every root key here, so it must be accepted. Kept as + an attribute for provenance only. + video_datasets: optional ``{head: {dataset_name: dataset}}`` for the + video-ONLY companion loaders (``_video``). Built by + trainHydra from ``data.video_episodes``: the head's own split + restricted to the pinned episode hashes, contiguous and + unsubsampled so the overlay video stays watchable while the + metric loader is per-episode subsampled. A head with no pins + gets no video loader (today's behaviour). The video loader + reuses that head's ``*_dataloader_params`` (unshuffled). + metric_frames_per_episode: config-only ``{head: K}``. trainHydra + wraps each val head's dataset in ``EvenStrideDataset(base, + frames_per_episode=K)`` so the ``limit_val_batches`` window + covers every episode of the split instead of the leading + hash-sorted slice. Forwarded here by hydra.instantiate (it + rejects undeclared root keys); kept for provenance only. + video_episodes: config-only ``{head: [episode_hash, ...]}`` -- the + pinned video episodes, one per operator per head. Provenance + only; trainHydra is what reads it. Tokenization (sampling a prompt from per-sample annotation lists, splicing in embodiment / control-mode / proprio blocks, and running - the HF tokenizer) lives on the algo side now — see + the HF tokenizer) lives on the algo side now; see ``PI.process_batch_for_training``. The collate here only stacks tensors and preserves variable-length list-valued keys (e.g. raw ``annotations``) so the algo can consume them downstream. @@ -43,6 +147,25 @@ def __init__( self.valid_datasets = {k: v for k, v in valid_datasets.items() if v is not None} self.train_dataloader_params = train_dataloader_params self.valid_dataloader_params = valid_dataloader_params + self.train_viz_datasets = { + k: v for k, v in (train_viz_datasets or {}).items() if v is not None + } + self.train_viz_dataloader_params = train_viz_dataloader_params or {} + self.unseen_op_valid_datasets = { + k: v for k, v in (unseen_op_valid_datasets or {}).items() if v is not None + } + self.unseen_op_valid_dataloader_params = unseen_op_valid_dataloader_params or {} + self.valid_prefix = valid_prefix + self.held_out_operators = list(held_out_operators or []) + # {head: {dataset_name: dataset}}; heads with no pinned episodes are + # dropped so `val_loader_names()` only grows for configs that ask for it. + self.video_datasets = { + head: {k: v for k, v in (datasets or {}).items() if v is not None} + for head, datasets in (video_datasets or {}).items() + } + self.video_datasets = {h: d for h, d in self.video_datasets.items() if d} + self.metric_frames_per_episode = dict(metric_frames_per_episode or {}) + self.video_episodes = dict(video_episodes or {}) self.collate_fn = annotation_collate def train_dataloader(self): @@ -57,20 +180,22 @@ def train_dataloader(self): dataset, shuffle=True, collate_fn=self.collate_fn, - **dataset_params, + **orphan_guarded(dataset_params), ) return CombinedLoader(iterables, "max_size_cycle") - def val_dataloader(self): + def _build_val_style_loader( + self, datasets: dict, params: dict, kind: str, gate=None + ): iterables = dict() - for dataset_name, dataset in self.valid_datasets.items(): - dataset_params = self.valid_dataloader_params.get(dataset_name) + for dataset_name, dataset in datasets.items(): + dataset_params = params.get(dataset_name) if dataset_params is None or len(dataset_params) == 0: raise ValueError( - f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config valid_dataloader_params." + f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) - dataset_params = dict(dataset_params) + dataset_params = orphan_guarded(dataset_params) shuffle = dataset_params.pop("shuffle", False) iterables[dataset_name] = DataLoader( dataset, @@ -78,8 +203,93 @@ def val_dataloader(self): collate_fn=self.collate_fn, **dataset_params, ) + if gate is None: + return CombinedLoader(iterables, "max_size_cycle") + return _VizGatedLoader(iterables, "max_size_cycle", gate) - return CombinedLoader(iterables, "max_size_cycle") + def _video_gate(self, head: str) -> bool: + """Does this validation pass render ``head``'s overlay video? + + Open only on rank 0 and only on a viz pass. Elsewhere the loader's whole + output -- decode, eval forward, render -- is discarded: the val loaders + carry no ``DistributedSampler`` (see ``EvalVideo._video_fps``), so every + rank was producing the same frames and only rank 0 ever wrote them. + + Ranks therefore iterate different numbers of val batches. That is safe + because what they log is not: the metric loaders run everywhere and log + the same keys, the video loader logs nothing, so the ``sync_dist`` + reduction in ``_evaluation_epoch_end`` stays symmetric. It does mean the + other ranks sit in that reduction for as long as rank 0 renders. Measured + 8.4 min for 76 batches back when all four ranks rendered at once, so rank + 0 alone is inside that and well under the 30-minute process-group timeout + -- but that timeout is now the ceiling on ``evaluator.viz_max_batches``, + which trainHydra therefore requires on every head with a video loader. + + Reads the same ``_should_viz()`` ``ModelWrapper.validation_step`` does, + back through the trainer so the datamodule needs no extra wiring, and + falls open whenever a link is missing: eval-only mode, tests, an + evaluator that predates the hook. + """ + trainer = getattr(self, "trainer", None) + if trainer is not None and not getattr(trainer, "is_global_zero", True): + return False + module = getattr(trainer, "lightning_module", None) + val_heads = getattr(module, "_val_heads", None) + evaluator = val_heads().get(head) if callable(val_heads) else None + should_viz = getattr(evaluator, "_should_viz", None) + return should_viz() if callable(should_viz) else True + + def _metric_sources(self) -> dict: + return { + "valid": (self.valid_datasets, self.valid_dataloader_params), + "train_viz": (self.train_viz_datasets, self.train_viz_dataloader_params), + "unseen_op_valid": ( + self.unseen_op_valid_datasets, + self.unseen_op_valid_dataloader_params, + ), + } + + def val_loader_names(self) -> list[str]: + """Names of the val loaders in ``val_dataloader()`` order; position i is + Lightning's ``dataloader_idx`` i. ModelWrapper dispatches on these. + + Metric loaders first, in the historical order, so a config with no + pinned video episodes gets exactly today's list and today's indices; + the video-only loaders are appended after them in the same head order. + """ + metric = [ + name + for name in VAL_HEADS + if name == "valid" or self._metric_sources()[name][0] + ] + return metric + [ + video_loader_name(name) for name in metric if name in self.video_datasets + ] + + def val_dataloader(self): + sources = self._metric_sources() + for head, datasets in self.video_datasets.items(): + # The video loader reuses its head's params (batch size / workers / + # unshuffled); only the dataset differs. + sources[video_loader_name(head)] = (datasets, sources[head][1]) + names = self.val_loader_names() + loaders = [] + for name in names: + head, video_only = head_of_loader(name) + loaders.append( + # kind names the *_dataloader_params block in the error message, + # and a video loader borrows its head's block. + self._build_val_style_loader( + *sources[name], + kind=head, + gate=(lambda h=head: self._video_gate(h)) if video_only else None, + ) + ) + if len(loaders) == 1: + return loaders[0] + # Several heads: return a list so Lightning populates dataloader_idx + # (the position in val_loader_names()) and ModelWrapper can dispatch. + return loaders def _extract_list_keys(batch): diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index b5657baf7c..a77dc9cce2 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -10,7 +10,8 @@ from omegaconf import DictConfig, OmegaConf import egomimic.utils.tensor_utils as TensorUtils -from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.pl_utils.pl_data_utils import head_of_loader, video_loader_name +from egomimic.rldb.multi import MultiDataset def _barrier() -> None: @@ -56,6 +57,7 @@ def __init__( if config_tree is not None: self.model = self._instantiate_model(config_tree, norm_stats_state) + self._maybe_compile(config_tree) elif robomimic_model is not None: # legacy support self.model = robomimic_model else: @@ -75,6 +77,16 @@ def __init__( self.epoch_memory_stats = [] # Store memory stats per epoch self.evaluator = evaluator + # Optional second evaluator that runs against a train-sampled + # dataloader during the same validation pass (dataloader_idx=1). + # See egomimic/eval/eval_train_viz.py. + self.train_viz_evaluator = None + # Optional third evaluator for the unseen_op_valid loader (held-out + # operators in the opsplit data configs). + self.unseen_op_valid_evaluator = None + # Val loader names in dataloader_idx order, from + # MultiDataModuleWrapper.val_loader_names(); None = [valid, train_viz]. + self.val_loader_names = None @staticmethod def _as_config(cfg): @@ -92,6 +104,29 @@ def _instantiate_model(self, config_tree, norm_stats_state): norm_stats=norm_stats, ) + def _maybe_compile(self, config_tree) -> None: + """Hand the algo's hot submodules to ``torch.compile`` when + ``model.compile.enabled`` is set. + + Here, not in the algo's ``__init__``: this runs on every DDP rank right + after the rank builds its own copy on CPU and before Lightning wraps it, + so each rank compiles for the GPU it is about to land on. Compilation + itself is lazy -- nothing is traced until the first batch. + """ + cfg = self._as_config(config_tree) + opts = cfg.model.get("compile") + if not opts or not opts.get("enabled", False): + return + compiled = self.model.compile_for_training( + mode=opts.get("mode"), dynamic=bool(opts.get("dynamic", False)) + ) + if not compiled: + print( + f"[compile] model.compile.enabled is set but " + f"{type(self.model).__name__} compiled nothing", + flush=True, + ) + # batch is now a dict, handle on model side def training_step(self, batch, batch_idx): self.train() @@ -173,62 +208,108 @@ def on_after_backward(self): info["policy_grad_norms_mad_threshold"] = threshold grad_norm_flagged = grad_norm_val > threshold info["policy_grad_norms_mad_flag"] = float(grad_norm_flagged) - if grad_norm_flagged: - torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=median) - if self.trainer.is_global_zero: - print( - "[GRAD_NORM_SPIKE] " - f"step={self.global_step} " - f"grad_norm={grad_norm_val:.4f} " - f"median={median:.4f} " - f"mad={mad:.4f} " - f"threshold={threshold:.4f}", - flush=True, - ) - - if not grad_norm_flagged: - self.grad_norm_history.append(grad_norm_val) + if grad_norm_flagged and self.trainer.is_global_zero: + print( + "[GRAD_NORM_SPIKE] " + f"step={self.global_step} " + f"grad_norm={grad_norm_val:.4f} " + f"median={median:.4f} " + f"mad={mad:.4f} " + f"threshold={threshold:.4f}", + flush=True, + ) + + # The MAD detector is log-only: it never rescales gradients (actual + # clipping is Lightning's fixed ``trainer.gradient_clip_val``), so every + # step goes into the rolling history. The old code excluded flagged + # steps only to stop the clip threshold from drifting upwards; as a pure + # diagnostic the plain rolling median/MAD is the honest statistic. + self.grad_norm_history.append(grad_norm_val) for k, v in info.items(): self.log("Train/" + k, v, on_step=False, on_epoch=True, sync_dist=True) - def on_before_optimizer_step(self, optimizer): - if not self.enable_grad_norm: - return - grad_norm = torch.nn.utils.clip_grad_norm_( - self.parameters(), max_norm=float("inf") - ) - self.log( - "Train/policy_grad_norms_clipped", - float(grad_norm), - on_step=False, - on_epoch=True, - sync_dist=True, - ) + def _val_heads(self) -> dict: + return { + "valid": self.evaluator, + "train_viz": self.train_viz_evaluator, + "unseen_op_valid": getattr(self, "unseen_op_valid_evaluator", None), + } def on_validation_start(self): - if self.evaluator is None: + # Deterministic eval: rewind the algo's per-pass RNG counter so this + # pass replays the same prompts and sampling noise as any other pass + # over the same batches (egomimic/algo/hpt.py). Algos without the hook + # (e.g. PI) are unaffected. + reset_eval_rng = getattr(self.model, "reset_eval_pass_counter", None) + if callable(reset_eval_rng): + reset_eval_rng() + + heads = [h for h in self._val_heads().values() if h is not None] + if not heads: return self.model.device = self.device - self.evaluator.on_validation_start() + for head in heads: + head.on_validation_start() def validation_step(self, batch, batch_idx, dataloader_idx=0): """ Run a validation step on the batch, and save that batch of images into the val_image_buffer. Once the buffer hits 1000 images, save that as a 30fps video using torchvision.io.write_video. + + ``dataloader_idx`` is populated by Lightning when ``val_dataloader()`` + returns a list; ``val_loader_names[dataloader_idx]`` picks the head + (valid -> evaluator, train_viz -> train_viz_evaluator, unseen_op_valid -> + unseen_op_valid_evaluator). + + A head with pinned video episodes also gets a ``_video`` loader, + which routes to the SAME evaluator in video-only mode while the metric + loader runs in metrics-only mode. Heads without one keep the single + combined call, so evaluators (and stub evaluators in tests) that take no + ``mode`` still work. """ - if self.evaluator is None: + names = getattr(self, "val_loader_names", None) or ["valid", "train_viz"] + name = names[dataloader_idx] if dataloader_idx < len(names) else None + if name is None: + return + head, video_only = head_of_loader(name) + active = self._val_heads().get(head) + if active is None: return + if video_only: + mode = "video" + # Non-viz epoch: skip before the forward pass, it would be wasted. + should_viz = getattr(active, "_should_viz", None) + if callable(should_viz) and not should_viz(): + return + elif video_loader_name(head) in names: + mode = "metrics" + else: + mode = "both" + # When val_dataloader returns a list of CombinedLoaders (valid + + # train_viz), Lightning wraps it in an outer sequential CombinedLoader. + # The outer iterator calls next() on each inner CombinedLoader, which + # itself yields a (batch_dict, batch_idx, dataloader_idx) triple; that + # triple lands here as `batch`. Unwrap to recover the dict. + if isinstance(batch, tuple) and len(batch) == 3 and isinstance(batch[0], dict): + batch = batch[0] batch = self.model.process_batch_for_training(batch) print( - f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}", + f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}, " + f"dataloader_idx={dataloader_idx}, loader={name}, mode={mode}", flush=True, ) - self.evaluator.on_validation_step(batch, batch_idx, dataloader_idx) + if mode == "both": + # The historical call: no `mode` kwarg, so evaluators that predate + # the split keep working unchanged. + active.on_validation_step(batch, batch_idx, dataloader_idx) + else: + active.on_validation_step(batch, batch_idx, dataloader_idx, mode=mode) def on_validation_end(self): print(f"[ON_VALIDATION_END] rank={self.global_rank}", flush=True) - if self.evaluator is not None: - self.evaluator.on_validation_end() + for head in self._val_heads().values(): + if head is not None: + head.on_validation_end() print( f"Rank {self.global_rank} on validation end, waiting for all ranks to synchronize", @@ -240,6 +321,32 @@ def on_validation_end(self): flush=True, ) + def _backbone_param_groups(self, lr: float, scale: float): + """Two param groups: the encoders' pretrained backbones at ``lr * + scale``, every other trainable parameter at ``lr``. + + A pretrained image backbone and a randomly initialised trunk do not + want the same learning rate, and a frozen-vs-tuned comparison under one + LR is confounded. Modules that do not declare ``backbone_parameters`` + (and frozen parameters) are left to the main group. + """ + backbone, seen = [], set() + for module in self.model.nets["policy"].modules(): + if not hasattr(module, "backbone_parameters"): + continue + for param in module.backbone_parameters(): + if param.requires_grad and id(param) not in seen: + seen.add(id(param)) + backbone.append(param) + rest = [ + param + for param in self.trainer.model.parameters() + if param.requires_grad and id(param) not in seen + ] + groups = [{"params": backbone, "lr": lr * scale}, {"params": rest}] + # a fully frozen (or absent) backbone leaves the main group alone + return [group for group in groups if group["params"]] + def configure_optimizers(self) -> Dict[str, Any]: """Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you'd need one. But in the case of GANs or similar you might have multiple. @@ -252,12 +359,27 @@ def configure_optimizers(self) -> Dict[str, Any]: config_tree = getattr(self.hparams, "config_tree", None) if config_tree is not None: cfg = self._as_config(config_tree) - optimizer = hydra.utils.instantiate( - cfg.model.optimizer, - params=self.trainer.model.parameters(), - ) - if callable(optimizer): - optimizer = optimizer() + scale = float(cfg.model.get("backbone_lr_scale", 1.0)) + groups = None + if scale != 1.0: + lr = cfg.model.optimizer.get("lr") + if lr is None: + raise ValueError( + "model.backbone_lr_scale needs an explicit model.optimizer.lr" + ) + groups = self._backbone_param_groups(float(lr), scale) + if groups is None: + optimizer = hydra.utils.instantiate( + cfg.model.optimizer, + params=self.trainer.model.parameters(), + ) + if callable(optimizer): + optimizer = optimizer() + else: + # ``instantiate`` would turn the param-group dicts into config + # nodes, so bind them outside it. + factory = hydra.utils.instantiate(cfg.model.optimizer, _partial_=True) + optimizer = factory(params=groups) scheduler_cfg = cfg.model.get("scheduler") if scheduler_cfg is not None: scheduler = hydra.utils.instantiate( diff --git a/egomimic/resources/blocklists/mecka_bad_egomotion.txt b/egomimic/resources/blocklists/mecka_bad_egomotion.txt new file mode 100644 index 0000000000..473e35c070 --- /dev/null +++ b/egomimic/resources/blocklists/mecka_bad_egomotion.txt @@ -0,0 +1,3013 @@ +# Episodes with failed head tracking (obs_head_pose / vendor egomotion), +# from egomimic/scripts/mecka_process/scan_egomotion.py over the +# 41607 episodes in the local dataset dir on 2026-09-23. +# constant: one pose for the whole episode; jump: a single-frame head +# move over 0.3 m. 3008 episodes. +692bb8c4a589f55a386e3ff2 # jump +692de91dbe226dd37c284a29 # jump +693a2eb976fbe3d2eea3effc # constant +693a2ec56d31efd205d000a0 # constant +693a2ec9899a3f42531cad57 # constant +693a31344949f69269a874d9 # constant +693a314d1772bb2c1fefefd5 # constant +693a328dbc6817e55c26b61f # constant +693a32ff9370a5841f6c85e0 # constant +693a3376db9d711abe5b0a2e # constant +693a34770d06f344ee34707b # constant +693a3ac4e6f40c3f4259a6ba # constant +693a3cfab07e0d1ae7613c25 # jump +693a5b2aafeb3b82b7d56cd3 # jump +693a5cef6f71cc238727ab97 # constant +693a5da9c03d30428c3f1a9c # constant +693a5fcd074493bb20f07806 # jump +693bd02a9b1e2efb2972c757 # constant +693bd2035adf90b6a80ffa31 # constant +693bd204b3fb840e8f81e768 # constant +693bd403dffd115628611fb6 # constant +693bd418eae7bb2135331eaa # constant +693bd444caac47f5b0f8696f # constant +693bd44705d6e83890c2bd2b # constant +693bd574dffd115628611ff6 # constant +693bd5c19ea209cc23b7023b # constant +693bd64eb0e6a9d2fa0bb576 # constant +693bd65e80eb971466214796 # constant +693bd6762168dd435d371876 # constant +693bd7286f80c42defd06b04 # constant +693bd96cf4a26a8bd47efc84 # constant +693cce31da7b160436a1054d # jump +693ccf367be631a588f81aab # constant +693cd5863b560ae2a5174457 # constant +693cd9fccfe68b40e919aaa2 # jump +693cde268dd3d9f190574ce5 # constant +693cdf1b0d50f31fa56c0495 # constant +694a0e6834f59ef8535c5e2c # constant +696a15a756b90f004721201b # constant +696a2f4e4419259d05e510df # jump +696a2f6bac451788451c331e # jump +696a304f76b41c2326097948 # constant +696ac68c76b41c23260979f9 # constant +696ac6ba4e85aa4ea7da60ae # constant +696ac6e57a2fd07226c099a5 # constant +696ac73876b41c2326097a0f # constant +696aca92ac451788451c343b # constant +696b37b682fb9f9fc0a7e974 # constant +696b37c070d628f8b0741c57 # constant +696b37c0f8e8293f615f769d # constant +696b37d703b3bab4cd3d549e # constant +696b3891e9d8f7b5d6be6dd6 # constant +696b38a2a12eec59c7475065 # constant +696b38a321d8e8ab8ddc7c71 # constant +696b38b10ac58bb141246a1f # constant +696b38dc7e6f94bf341433cf # jump +696b38e170d628f8b0741c63 # constant +696b45d08d4f708447df84b5 # constant +696b48a984947585c0cb6db4 # constant +696b48d6a907c8962e22a53b # constant +696b49cb240b1e4a2d82c376 # constant +696b4a71cb4744d00cf71ddf # constant +696b689e4efa2e122be5ff15 # constant +696b68b386833e290580c063 # constant +696b699133191fac7d917cb1 # jump +696b6ab31c7532ffed13fc03 # constant +696b6bea9c8ea4b6e4a1ee90 # constant +696b77db7a53efd0cd3a2890 # constant +696b78471d3959e2e052ae4e # constant +696b7847ea14729f866e8ab0 # constant +696b785a2efe1549985cf711 # constant +696b791b476e9a6e777ffa97 # constant +696b7950ef23530d3c5a32d0 # constant +696b795f39424cbffd650b52 # constant +696b79b7c2544a8f586250a5 # constant +696b79bc1e187f32c5eb5860 # constant +696b79bc75020a567fd3746a # constant +696b79be86f71ea8bc69cf91 # constant +696b79c22a9a68766920dbab # constant +696b79e034614160c6097a42 # constant +696b79e21d3959e2e052ae66 # constant +696b7a0cf20f919514c6f34b # constant +696b7a15d87f83facc75f6c6 # constant +696b7a16634a1e1f4692759b # constant +696b806d4e85aa4ea7da6532 # constant +696b8ac9b717e87cf71f75f4 # jump +696b8af186f71ea8bc69cfee # jump +696b9cd3b0b2b22cc0d1f4b1 # constant +696b9ce243edc8aed5480230 # constant +696b9cebeea8e4fd806962cd # constant +696b9ced5617babd55e6589d # constant +696b9d1cd2d7cb0bd150eaf5 # constant +696b9d1f894c2e015aaa0d23 # constant +696b9d298318b8ceb8224679 # constant +696b9d2b522a3f9be479040a # constant +696b9d349ac5bea01f980022 # constant +696b9da9675f6396322f9111 # constant +696b9da9ac515a4d340b0cec # constant +696b9daf21d8e8ab8ddc8053 # constant +696b9db6900932d25e869644 # constant +696b9dbc07485cd91adbcedc # constant +696b9dcc9a643f01c844d254 # constant +696b9dd6101054681ed97553 # constant +696b9ddbcd29fb97fb107bef # constant +696b9eff715ed751e5e0fe25 # constant +696b9f01e43294cde98aa274 # constant +696b9f0267f2110cf0fbdf24 # constant +696b9f09abc451024b826ddf # constant +696b9f0b25368af005ca72fd # constant +696b9f0ea8793724a1ef1eea # constant +696b9f10316c07a505ec650f # constant +696b9f1407b6c7a3d7946469 # constant +696b9f53f98fc50273d5d24b # constant +696b9f596377d0871b6a8eb7 # constant +696b9f64a5b1694665c1958a # constant +696b9f70246cb36ca3757d13 # constant +696b9f705fedf8533255dcae # constant +696b9f9428ab12fee373515d # constant +696b9f98c7e4e8ee05bbded7 # constant +696b9fa04a936a4d1328a717 # constant +696b9fa737193673e4f8987f # constant +696ba19aca28cac2415e5cad # jump +696ba24179188f941930289e # constant +696ba249c70dc8ed89019e99 # constant +696ba25786f71ea8bc69d0cc # constant +696ba2b486f71ea8bc69d0d4 # constant +696ba325ac515a4d340b0d6c # constant +696ba5465bfe694453c614ac # constant +696ba55c4a936a4d1328a789 # constant +696ba5739ac5bea01f980095 # constant +696ba5857a2fd07226c09eb8 # constant +696ba59353fe50fed71373c5 # constant +696ba5a207485cd91adbcf81 # constant +696ba5bb9a6d79db2cf885e0 # constant +696ba5c409b82287df69f517 # constant +696ba5c8d9a2a4e1bf836bba # constant +696ba5cb915ce3edd09678ce # constant +696ba5cdc19cc80a4d8afef6 # constant +696ba5ecc426c607c88b2c2c # constant +696ba5eee2eca952034ed6dd # constant +696ba5ef900932d25e8696f4 # constant +696ba5f29a140887948c5cbd # constant +696ba5f3bd0145c530d36e2f # constant +696ba5f66377d0871b6a8f2e # constant +696ba5f66fc0c09c00a97c44 # constant +696ba68f1c7532ffed13fe9a # jump +696ba6a48e6d0bc988967731 # jump +696ba72475020a567fd375f4 # constant +696ba76c1894fc36a07d8d50 # constant +696ba76c79e3576eb2d4f1e9 # constant +696ba7754efa2e122be60139 # constant +696ba7c8c676295b4ae74a2c # jump +696baa451ee2a915e19349e1 # constant +696bab7693818b7b258ec0e4 # constant +696babaf4a936a4d1328a80a # constant +696babb1f7b3293476cab775 # constant +696babd0eea8e4fd80696370 # constant +696bac2537193673e4f898f4 # jump +696bac706fc0c09c00a97cc8 # constant +696bac7636602141d44d6c43 # constant +696bacf5632cfa55ccdfa79a # constant +696bacfdf53e1ae2cbd86e13 # constant +696bad8d017ec1a72078f2a1 # jump +696baddcb0900198ccf5c0dc # constant +696bade10c52a10600ae8472 # constant +696badec9c8ea4b6e4a1f137 # constant +696badefc6cb022a9da40224 # constant +696baff5316c07a505ec665b # constant +696bb012a5b1694665c19712 # constant +696bb01939424cbffd650dd7 # constant +696bb043590d785c13a75eab # constant +696bb054ac515a4d340b0e9a # constant +696bb06f632cfa55ccdfa7eb # constant +696bb321cef1149d4f7a1af5 # constant +696bb36bc7e4e8ee05bbe050 # constant +696bb7706377d0871b6a90b0 # constant +696bbc3fd3cc9dc76c561769 # constant +696bbc988e6d0bc9889678fc # constant +696bc1a41894fc36a07d8f0f # constant +696bc26c14764da1c7e1852f # constant +696bc27c2c7243cd72c85102 # constant +696bc3982a9a68766920de00 # constant +696bc3ad0c52a10600ae85ee # constant +696bc64c95be6b2159258593 # constant +696bc7bd8688d0ceb96d3d65 # constant +696bc8e44177e6c8a8388438 # constant +696bc9032c7243cd72c85139 # constant +696bca8bafe603760fe1af21 # constant +696bcb22e34991e438ab3922 # constant +696bcb3bc7e4e8ee05bbe1e1 # constant +696bcb46ae7ec71f86571821 # constant +696bcc83d2d7cb0bd150ef10 # constant +696bcc8eb7080af9dcc943a2 # constant +696bccd521d8e8ab8ddc83d3 # constant +696bcd3c55db5be0d770718c # constant +696bcd93915ce3edd09679d4 # constant +696bcedb96b4b190357b0e62 # constant +696bcfd936602141d44d6e3c # constant +696bd06901fc6bd69561e2fe # constant +696bd158473aede2ca7efeaf # constant +696bd16767f2110cf0fbe26d # constant +696bd21fd0bfcc86fe4a6a36 # constant +696bd23360ebbd99811cbe5f # constant +696bd25121d8e8ab8ddc843c # constant +696bd25707b6c7a3d79466a2 # constant +696bd261522a3f9be479078c # constant +696bd273263039fb3592de3a # constant +696bd278b17075c2c32fa942 # constant +696bd299fe5c03f9da38895b # constant +696bd2bdcd29fb97fb107f91 # constant +696bd2cc07b6c7a3d79466aa # constant +696bd2fd8910191c50ed9eb1 # constant +696bd34539424cbffd650f81 # constant +696bd34a8a0fa4aac9be483a # constant +696bd34fc25b828c83107fab # constant +696bd379c7e4e8ee05bbe285 # constant +696bd3a97a53efd0cd3a2d4d # constant +696bd3beb717e87cf71f7aa3 # constant +696bd3bf4d5f3676bb0ab1d4 # constant +696bd3cfa907c8962e22aa1f # constant +696bd3d3cfb39c7ef2d48dbd # constant +696bd3d59a140887948c6044 # constant +696bd421590d785c13a7611e # constant +696bd4261e187f32c5eb5d9e # constant +696bd4262a9a68766920df0c # constant +696bd444620da154d5b7abaf # constant +696bd44a7a2fd07226c0a1ff # constant +696bd44a93818b7b258ec3f8 # constant +696bd44aca28cac2415e5feb # constant +696bd454715ed751e5e1025f # constant +696bd45507b6c7a3d79466c6 # constant +696bd45a4c43650cbf6f5f00 # constant +696bd463da6069828e20d10c # constant +696bd4a528ab12fee3735567 # constant +696bd4c71d3959e2e052b29c # constant +696bd7a7590d785c13a76167 # constant +696bd7a85e774190a1c9791b # constant +696bd8e933191fac7d918067 # constant +696bd8f8ae7ec71f86571893 # constant +696bd90337193673e4f89b68 # constant +696bd98e95be6b2159258746 # constant +696bddedfe5c03f9da388a65 # constant +696bdf2d0f689873a4d2b470 # constant +696bdf5be4eb4a667bf0e0ec # constant +696bdf81b7080af9dcc944f6 # constant +696bdfb92a9a68766920dff5 # jump +696be1fd522a3f9be47908be # jump +696be230675f6396322f95e7 # constant +696be37567d3fa0fb3723cfb # jump +696be37e4177e6c8a838855d # jump +696be557246cb36ca3758254 # jump +696be5a1715ed751e5e103ef # jump +696be67f43edc8aed548076f # constant +696be6edc3881a5a8cbaa54e # constant +696c5a952bd5bc335fb2a4d5 # constant +696c9634089658b1b2efa9e5 # constant +696ca9c746f3fd69cfbea1e0 # constant +696cf19a344d3cc10178a9f3 # constant +696d017788af7cc080636904 # constant +696d2d7fde00c290bacf0021 # constant +696d2da1ead9a484382f35b6 # constant +696d2e73961f9960d10001f7 # constant +696d2ecaf66e78c232fa59a9 # constant +696d2ed10576d930a3440041 # constant +696d2faa42dcbdaae6ac10a4 # constant +696d2fac995fc04852021b2b # constant +696d3024961f9960d100021c # constant +696d302c619debe3a900b0fb # constant +696d30312ed0529037c4e95e # constant +696d30352bd5bc335fb2aaae # constant +696d3067e4c876d0ca2443dd # constant +696d338f5ce51a36f6f90721 # constant +696d33926ed3eb90778f08de # constant +696d339d38f4b4346c531c36 # constant +696d3414296b5bace896da6d # constant +696d6fecb777874eaf05982e # jump +696d791da199224e8d662855 # constant +696d8e3816e13cf003fab385 # constant +696d8eb3d07d4cdba345b767 # constant +696d940329ce27f261bcf00d # constant +696d943a29ce27f261bcf00f # constant +696d94583ae226ec1ef352a7 # constant +696d9724b813e38d6a1a1122 # constant +696d97264c3d4673e3917442 # constant +696d9727d724a55c2ff9a56e # constant +696d9728d920e42eac457cae # constant +696d972936733bd98d9d5c29 # constant +696d973a37dbf74800ad871c # constant +696d973bcde72530e232bee7 # constant +696d9770cb3d08a3aa82b81a # constant +696d97ab3d5391e3c577b080 # constant +696d981385537287d8383c6a # constant +696d98185c2b2b8722b2a077 # constant +696d9a287b9f7e17cc5772a9 # constant +696d9a4a34e6f4bcaf23b053 # constant +696d9b12306d82b3e34c85e2 # constant +696d9bcb0838d9b98d822d16 # constant +696d9bf05c2b2b8722b2a08f # constant +696d9ebb3761ff9bfa02defd # constant +696d9ee6c354e2fafcf202ac # constant +696d9eec92f396e3ef808490 # constant +696d9eed04c9be98d76d60fe # constant +696d9eede25d15f86621b567 # constant +696d9eee9d65f7c2b62e802b # constant +696d9ef23fb0f5175cce0ae0 # constant +696d9ef3d19968f832c03456 # constant +696d9ef4557bad554d708100 # constant +696d9ef63d5391e3c577b104 # constant +696d9ef6d7f80a26e0463de9 # constant +696d9ef7592a67f25a7b8955 # constant +696d9ef938760cc5e99148dc # constant +696d9ef94a0fdf8d0ad096a8 # constant +696d9efb8061c372104e6a4e # constant +696d9efdb762e60f1f34fcde # constant +696d9f0083d3fc2d8b5b956b # constant +696d9f1677bf62ada9a01ac4 # constant +696d9f1936733bd98d9d5c91 # constant +696d9f1a00b24b8aa54c2345 # constant +696d9f1a16e44e8c77d6aa91 # constant +696d9f1c69fe7e32c96f1acd # constant +696d9f1d76a542b835f9cb0b # constant +696d9f2b154a7aa736439ce0 # constant +696d9f4fc93f59bdc6e20da7 # constant +696d9fdacb3d08a3aa82b897 # constant +696da359103547e961005ef0 # constant +696da35c71e72197ed90c810 # constant +696da361925cce50c37ecae3 # constant +696da3803943b9f942d7e96c # constant +696da3b5e21603c8805256e2 # constant +696da3ec7c11eaba88f2df43 # constant +696da420f091d1b91dfec21b # constant +696da46a9984cbf630b0cd46 # constant +696da47a63999329b2a60288 # constant +696da71a76a542b835f9cbb8 # constant +696da743613c6035217a430e # constant +696da75abd7eacba279019d7 # constant +696da75cf1b842325f9e80d4 # constant +696da764102bc82df8e9a8d8 # constant +696da76aaac016e2daaa5905 # constant +696da76f4dcd2eaf6d5fba06 # constant +696da79a34e6f4bcaf23b145 # constant +696da79f4d91ecc2866560e0 # constant +696da7a016e44e8c77d6ab51 # constant +696da7a064124db03d5205db # constant +696da7a0e75aabd3cbfcedf3 # constant +696da7a1840d7d67309c3375 # constant +696da7a449652a41264cf4a7 # constant +696da7aa69b03d8bffc3f0fe # constant +696daa27f1ecd06ff7255038 # constant +696daa2bcfe06ac65aaf312e # constant +696daa2d433fcf31b5691da8 # constant +696daa31a998602bd03d1532 # constant +696db12328a92a09a6839e01 # constant +696db1284f4c8d4f355cd463 # constant +696dc3ad5c2e8fc96acef92d # constant +696dc4873cf2a230e1c65961 # jump +696dc488261c825cae4b2c0c # jump +696dc48963999329b2a602fc # jump +696dc48a04c9be98d76d6218 # jump +696dc48caabe1eee429be102 # jump +696dc493dc787523fe33b682 # jump +696dc6250b15943f95116162 # constant +696dc62c650ccadd25513257 # constant +696dcac93d5391e3c577b1d1 # constant +696dce79960ada878296f047 # constant +696dce7a5b7992e1b3c3e9f3 # constant +696dce7ab2ed922c6673bdb9 # constant +696dce7afe33c2160dc9e7ca # constant +696dce7ba56f03a27e59f690 # constant +696dce7eebb60123bf5f6a47 # constant +696dcfabc354e2fafcf203e0 # constant +696dd11f4d2f71b8d51ea3d9 # constant +696dd13582ce9ecc67d33b3e # constant +696dd37deaf21ce266daa6d0 # constant +696dd448913cbec0b8bf114a # constant +696de4b6b72e3284f3b3768e # constant +696de9ea1d3a134969609728 # constant +696dfbafee9b736af1a79331 # constant +696dfbb02934fcc9b4fff444 # constant +696dfbb07992c9b130b1b63b # constant +696e1f0f39143339d5e0a5ca # constant +696e314217377c6d0e557974 # constant +696e31c843f23b8635c5a168 # constant +696e31cdb5a66458ce86c546 # constant +696e31d66edf6acdd058ba2f # constant +696e31d8e1edeabea6298ac7 # constant +696e31f77cc56e261b1067cf # constant +696e320edd0cb8d4e520624b # constant +696e32173644d45b8cb8a8f7 # constant +696e321ea61150e395603649 # constant +696e32717046f49217da348f # constant +696e338f43f23b8635c5a17e # constant +696e33a7620da3c71e5ed0fd # constant +696e358812a7f5ee4b9602e3 # constant +696e3bfbb79e9d7860a3b636 # constant +696e3c48e23dfd1ea54cecdd # constant +696e3e2f3644d45b8cb8a93d # constant +696e40727cc56e261b106846 # constant +696e412cb79e9d7860a3b6a1 # constant +696e4f97235ddc6771e3752b # constant +696e53638587a69d44bf7b54 # constant +696e54f08587a69d44bf7b79 # constant +696e5613f04cd7fd4f334423 # constant +696e57f312a7f5ee4b9603e2 # constant +696e673fcc5f7c104dea8815 # constant +696e6744c4bd89d6782639d4 # constant +696e7aac16141fc0a6fa29f4 # constant +696e7c112a18d10d394ba61c # constant +696e82bcb79e9d7860a3b80e # jump +696e8336658193f417ff80ba # constant +696e861557ad93e55dd08977 # constant +696e861cb1cc409b12d985e3 # constant +696e861de2287068cc486c9b # constant +696e861ef563b4f8900f2fb3 # constant +696e861f7046f49217da36fb # constant +696e8621103f142a3d3fac36 # constant +696e8621d5450f3c814f57ad # constant +696e8625d5450f3c814f57ae # constant +696e86280299a6b3c6c695ab # constant +696e86294b3bc2a1bad578af # constant +696e86299c0503f8c2b71576 # constant +696e862bd3afac558b18ba37 # constant +696e862c887a0e735b83d28e # constant +696e862cf78a5e583f4e629a # constant +696e862ef30588a493c07e36 # constant +696e863457757f9d4e3a1d3e # constant +696e86afd5450f3c814f57b9 # constant +696e86b0e85937dc7bae4fa1 # constant +696e86bb5e4e2b7186cbc81d # constant +696e87df0ebd516bc2eb4d5a # constant +696e87e68c2d9dcaa805a11b # constant +696e87ec57ad93e55dd0899c # constant +696e87f88a176d6397a7a16b # constant +696e87ff6d44bc7cd096aafd # constant +696e8829f0f65980ee0bba56 # constant +696e883b6baf731a9a6e6bdd # constant +696e8854e2287068cc486cce # constant +696e886e9b200fdeaf2e5fd4 # constant +696e88c761ed9732c88666a0 # jump +696e8e2b70007c739ba09e84 # constant +696e8e2ffe004d797507e0bd # constant +696e8e30c696b4a1ab3ca108 # constant +696e8e31d80c7b2e84481a32 # constant +696e8e7fec6580da1aea7c21 # constant +696e8e928587a69d44bf7c98 # constant +696e8ea84b013c844462911e # constant +696e8eac276bf7d7724c5a9c # constant +696e903b7d60b0a13a32478b # constant +696e955843c9ca3b04772b6b # constant +696e97c69b200fdeaf2e6053 # constant +696e9a1f11a1ac2628c606dc # constant +696e9a9fec7b7b20f30a70ba # constant +696e9c3f149e0f309e2caa63 # jump +696e9c61d5bdf2afafe2cd18 # constant +696e9c8ee4766e44b4d2aa78 # constant +696e9df17e219e09d7621b3d # constant +696e9df7cec1370af53818c9 # constant +696e9e804d3fd8430e0713c3 # constant +696ea955946d3894eba265ab # constant +696eb1a0fea209e4672a38da # constant +696eb9ff6f3b64df08ff12e0 # constant +696ebaef7d60b0a13a324938 # constant +696ebb68b057bb83a072b251 # constant +696ec03006dd95887ea9d12e # constant +696ec1362aa298d03d7cd874 # jump +696ec3828e3319820c7fac37 # constant +696ec3838587a69d44bf7e52 # constant +696ec3d9ec6580da1aea7e47 # constant +696ec4a29ba986538aa450a8 # constant +696ec61e8ea2ad3c2dad7543 # constant +696ec6203e7dae3f86d26ff4 # constant +696ec6ba3e7dae3f86d27002 # constant +696ec76cfe004d797507e436 # constant +696ec77f6edf6acdd058be8a # constant +696ec80c9ed17030548e7b54 # constant +696ec88780a86712b2f3e358 # constant +696ec90f4a5fcb629b3380da # constant +696ec932be978d4f1be335a8 # constant +696eca21f17d4757b6649beb # constant +696ecb89d5f713915440c2d4 # constant +696ecbd7128c825c68738824 # constant +696ed41352cb044323851711 # constant +696eda0fd1aafd24047c3840 # constant +696f331ae6cfc143ed371f4c # constant +696f3821ea3e7b471881463e # constant +696f4e7c5421d1827a1141fc # constant +696f50fb30ccd2968d7c51b7 # constant +696f53773fa392f88766e43b # constant +696f54280fe6e86763b43f0b # constant +696f5749f409f8e80b8a3f06 # constant +696f57b60e6eee9dc9f31ad8 # constant +696f57b77c254e888f239111 # constant +696f5934597ed929a691fa44 # constant +696f59483fa166bffbc464b7 # constant +696f5949337cb3e824164997 # jump +696f5cc4776534b4b8885cc5 # constant +696f5d20d8c3dea1dddc523d # constant +696f5d28b547042a74cec991 # constant +696f5d2a30ccd2968d7c52cd # constant +69706faba7e84318b3fba12a # constant +6970d1ac43399201b8e2315b # constant +6970da497a7337376e2bc4af # constant +6970dedff8ce5fae410d3d9e # constant +6970df0e2e2410cce12eed8a # constant +6970e0efcb8250eb53de1700 # constant +6970e1a61c853339d601e238 # constant +6970e9099148c925c45980ed # constant +6970ec9b8fb9776b4e1b72ad # jump +6970eeb5feabbb0a255d3b1e # jump +6970efa64df645940fddc4ba # constant +6970f0bb700778442cfce7d9 # constant +69716b116ae0fc2841fae537 # constant +697181d50db12b20e3530129 # constant +697186fd7eadbc4841b0ec89 # constant +69718a3d3c7d6692908d86b3 # constant +6971c1f12c357dcf66a6253f # constant +6971ca988187fb0568c2eed9 # constant +6971ce6fa489dfa963954cb1 # constant +6971dce4843fdbc1f74d95f6 # constant +6971f1e1bc8fb2284fc25a79 # constant +697201fb70159bb81c71a1e3 # constant +6972053e86eb2bfcfb7dc961 # constant +697205710864eb2c44f6d2bb # jump +697205cea85a5f59270a7d64 # constant +69721035892d3b1348cb38af # constant +697211379a033e87ff634924 # jump +6972138711add03562e20745 # constant +697219639d94825b9b1d5926 # constant +697219668f2b41a4f98ef3ce # constant +69725a9392854961b34654e0 # constant +69725bafe968e630d03ce34b # constant +6972700975cb2d5034d15223 # constant +697277f77f5862fa2b0cb597 # constant +697278ef90d5a26a3655c663 # constant +69727c63fff3e0b0b1b4ac26 # constant +6972839d266b6be87a581bbd # constant +6972906d66df7c68504aac1e # constant +69729993fcc5ad222d1c20ef # constant +69729a08793ae530061c9e33 # constant +69729ab0949555784c6d7a41 # constant +6972b02b026a2ba2357ed79e # jump +6972b208fbec37ec305b5cd6 # constant +6972b2711ca403a8d6e7a8ab # constant +6972b28262c4aae9294fd0ea # constant +6972be6663204a0c55d781df # constant +6972bf0870159bb81c71a7d5 # constant +6972c1a754f3c620cb3fa797 # constant +6972c4c4f983410986a59849 # constant +6972c5bf54fa95b5ff445f89 # constant +6972c5d6129577bdfff83c91 # constant +6972ca39b7b1624095ab5c24 # constant +6973008c892d3b1348cb3e36 # jump +697305c9441a5d3fca0fc8cd # constant +6973076cbf0b5944b7b573db # constant +69730d8420c41a167913d92f # constant +69731348379486fda752386f # constant +697316880ca9a458aa2fdb35 # constant +69732ceabe2416d834ed15a8 # constant +697346d40cf3c36a559e4214 # constant +697347ce9dce020024863ec4 # constant +697349fcf1d81564aa967810 # constant +69735d942e4df8aa85b112ca # constant +697376c6fd8f95b606f1902e # constant +697379c921d5b32c7f63bae9 # constant +6973942bb78f473064f5c57e # constant +697395aa9a02a5fbd06bcd6c # constant +69739e5192854961b346613a # jump +6973a094471ae3ad2b94bed5 # constant +697542c6e9621a11417acdac # constant +69754f6898d5102275b1480c # constant +69755aba4cf7dda51b6bc9ce # constant +69755c5fc03e10f394f2ed0e # constant +6976b54330bfef5fcc991a9e # constant +6976b5977039753b382700bd # constant +6976b5cb39510854b7622940 # constant +6976bca7b6233b96c3d48bd0 # constant +69a663a37f3184ccf64b0ab1 # constant +69b06c63a41e7fb22dd4a8a6 # constant +69b06f38752636ccc06b75ec # constant +69b06f7e444927c0dcab6ba9 # jump +69b085970563bd22b57167b1 # constant +69b085c1abad9395c0362a13 # constant +69b085dab1c888082e6d34a6 # constant +69b08e76b72bf9dca5afe23b # constant +69b0980be1cfd215d63f63cc # constant +69b0a11d0e7411f4bbebb515 # constant +69b0a14302be46f851bd6dc0 # constant +69b0b67bf44fb212ce962e71 # constant +69b0beecd984d4e5174e34a4 # constant +69b0cf3ebe31a1b6e41bdaa9 # jump +69b0d332e6981bb134c2df07 # constant +69b0ee7a8ad37d8884217dd8 # constant +69b0fd8653d600d0039c40c6 # constant +69b10666a56bff1fe3d65c1f # constant +69b10c18b18a85b173612725 # constant +69b1249b30f0fba8c63a1ced # constant +69b12c3eedbf7d0be142db86 # constant +69b12e0bb2e806badeceb9a0 # constant +69b12ee5efea4a48af846644 # constant +69b12f04cf06f50982a6234e # constant +69b12f18a89ab4b4bcb49bdb # constant +69b12f256787a1a655f767ce # constant +69b12f737e87f4440cdcccfe # constant +69b1309cbf177a6092a31a83 # constant +69b130a523cbfc2fe7dcc276 # constant +69b130e55848ee3f7b14ecb9 # constant +69b13111dba7336f4d4c5a46 # constant +69b1314a5006a5defae2261a # constant +69b1331b73430cd59a783c06 # constant +69b13403016ec66053d1443c # constant +69b136037cf8aa3b41635c70 # constant +69b136c7655fbc57678b852d # constant +69b137d3ec88eec3eaf21538 # constant +69b138a7b8e04b716d90b8b7 # constant +69b1403196805ff364b3665d # constant +69b15384a9843b0463763a16 # constant +69b18c793a82e884664ad0d4 # constant +69b1984bc70989363d216b59 # constant +69b19b8398a5e78c16ea81c8 # constant +69b19d8b8298a64c6992ebc3 # constant +69b1a46271a6e3af584865e8 # constant +69b1accea227efc31f66aee3 # constant +69b1b54b874d2c898fa49e99 # constant +69b1b80a25b08fb36c09ec75 # constant +69b1ba4e07fc78cb00c5cd1e # constant +69b1bf720b8af5e626fffa04 # constant +69b1c108d58394ebe352f7e8 # constant +69b1c1497c6435c837f4fd58 # constant +69b1c16b080ab0c125ed8cf8 # constant +69b1c16d0b4d6349ecf538d3 # constant +69b1c1b2ff452209ed443884 # constant +69b1c2a93ccdee9ce5eb162f # constant +69b1c2cc8d7e52f146936d09 # constant +69b1c3236787a1a655f76be1 # constant +69b1c364f9c3924f061a748c # constant +69b1c3a02e618a52f6d6fde6 # constant +69b1c3b0e0776aaea582c432 # constant +69b1c3b41c9c171f3f389b16 # constant +69b1c3c7ca9ec6cee9fb84da # constant +69b1c3d03780875f4907ba3c # constant +69b1c3dc000ba22d23231eb2 # constant +69b1c3fff247191f7a56467d # constant +69b1c4049446808800d0a2c1 # constant +69b1c426d0d92ee83be9f3cd # constant +69b1c42b7020a6d58a10138c # constant +69b1c4cbcc4b9ea5d8b035ca # constant +69b1c50c905ce054bacb6a60 # constant +69b1c51929d110add752dcd5 # constant +69b1c56ae249e0239ee298fd # constant +69b1c56c677ad5b8f785f57f # constant +69b1c570618f06c7dd18c13c # constant +69b1c5776b7aa68410ac8f5b # constant +69b1c59e6639bd32f1fe8a02 # constant +69b1c61542c6fa43e13cf78d # constant +69b1c628fc265cf0ead4b9c3 # constant +69b1c63de1ef6a6c397722d1 # constant +69b1c64db0ecbcc7408132d6 # constant +69b1c65a9f17df4a975d50ee # constant +69b1c661f247191f7a56468f # constant +69b1c67e323cfca1c6d7b49a # constant +69b1c682824c8a0015e4528d # constant +69b1c690522ab97a962ab5f1 # constant +69b1c69c61f41e36320e760c # constant +69b1c6b46a2b9c841d425cde # constant +69b1c6cbcda50c2db8875034 # constant +69b1c6d30e7e8f406a7ae6a4 # constant +69b1c6f10437e50778c76325 # constant +69b1c7099c77a19156c3a6fb # constant +69b1c71b64685105db818251 # constant +69b1c72684ed4cb5e9ebc2af # constant +69b1c731ec387ac11eb95d1e # constant +69b1c75afa244109e4970571 # constant +69b1c76c670578de9a958581 # constant +69b1c7745ccd991072cb5150 # constant +69b1c77d5c4bfe254a44a7b0 # constant +69b1c78bf717ad4138f084ed # constant +69b1c7996e5ad94e03ef978d # constant +69b1c7a455512547cecff36e # constant +69b1c7b50437e50778c76331 # constant +69b1c7bf74727173470579b3 # constant +69b1c7e06e5ad94e03ef9791 # constant +69b1c803ab59e885a8a11154 # constant +69b1c809743a2514ca64e5c6 # constant +69b1c80b9f0ca7b3e36a6278 # constant +69b1c815ce1ac0c1accf5c32 # constant +69b1c816d61b3a5bbd82334f # constant +69b1c82f850f4b031baebdb3 # jump +69b1c83b3a82e884664ad2bf # constant +69b1c85853ee403b9baa1312 # constant +69b1c86a05829050a08f7582 # constant +69b1c8822bd092414266dd11 # constant +69b1c8857b14ac89f528d75b # constant +69b1c88ee63780415fa33b01 # constant +69b1c892810bde0dbd313c19 # constant +69b1c8a9ef7e64f0dd5a68ec # constant +69b1c8c7895623595997b526 # constant +69b1c8c97efbe2a3a0f8c3c0 # constant +69b1c8d10b1e5afc00dafaff # constant +69b1c8d4e39500b459a40627 # constant +69b1c8da7ba5833e2c72b8e9 # constant +69b1c8fce2257ecc216ba128 # constant +69b1c907811e707a59e4ce49 # constant +69b1c91366e2cf554104fb38 # constant +69b1c9144008f46a714af92a # constant +69b1c9396392801584cbaeb0 # constant +69b1c9ae49b6552c0a558dd8 # constant +69b1c9bfe2226b5f968e9d0c # constant +69b1c9daa9ab356d6af9f8bc # constant +69b1ca11d5bea5bb5487b5a0 # constant +69b1ca160f2b085b6ff3b244 # constant +69b1ca1cf7aaf7b4ebc131b1 # constant +69b1ca25b75f730c09419e51 # constant +69b1ca4961f5ac8ce166f2aa # constant +69b1ca4e88b492ddbcdb3ab2 # constant +69b1ca5a836a1d9136668fd1 # constant +69b1ca5c4a3ae69e7048ed99 # constant +69b1ca5ed58394ebe352f83a # constant +69b1ca7d179c113f82bd4616 # constant +69b1ca7dc650c171c1b99a20 # constant +69b1ca8bfa8b87257c184123 # constant +69b1cab1bdc252e66ed98b48 # constant +69b1cb06e8ff04c40b394b5d # constant +69b1cb32a46af22d7d4e54ae # constant +69b1cb3968839fc7df2d3df0 # constant +69b1cb3a2e76ad18608f78e5 # constant +69b1cb412cd447ba578d82a7 # constant +69b1cb6087aef0a66e63031f # constant +69b1cb71908da4b28922dd30 # constant +69b1cb899677a13a457962e6 # constant +69b1cb8ceaf22e325027e3f0 # constant +69b1cb98a9843b0463763e0e # constant +69b1cb9c73430cd59a78418a # constant +69b1cb9d71a8c14db79ee6f1 # constant +69b1cbf0697828881c839d6b # constant +69b1cc178de9c1647cf4e539 # constant +69b1cc179e32c6db737d14c9 # constant +69b1cc2c9c77a19156c3a74f # constant +69b1cc2df6a3a400d9a7e3cf # constant +69b1cc3d3f4e9d4e5132cb2d # constant +69b1cc3f307325b3fb1f1d12 # constant +69b1cc45d681620c86ae2e04 # constant +69b1cc497601e53f94e14e08 # constant +69b1cc4c39f69bb8177a02a6 # constant +69b1cc50f3410a910fd55087 # constant +69b1cc5595fa8a42f817b31b # constant +69b1cc6100b8839f89342320 # constant +69b1cc656a96791183f0c841 # constant +69b1cc6d607ab2dbbc249bf7 # constant +69b1cc8d6f2b79851798a644 # constant +69b1cc9be7fcca548f22fd66 # constant +69b1ccaa1fa85f581929eccc # constant +69b1ccaa6a2b9c841d425d04 # constant +69b1ccb97366e3614a0143c5 # constant +69b1ccc0e6c38b58c797934a # constant +69b1ccf23584aca46e3a969c # constant +69b1ccfa050263a59047980c # constant +69b1ccffef97f6b74d1f52df # constant +69b1cd0e2cba1134ec023dce # constant +69b1cd25c70989363d216dce # constant +69b1cd3f92d07aa47a442b8b # constant +69b1cd5f00b8839f89342334 # constant +69b1cdb250e3da38b96383d0 # constant +69b1cdb4206612d7447a38f0 # constant +69b1cdbdcc4b9ea5d8b0362a # constant +69b1cddf402b694194fc98e2 # constant +69b1cddff1a94d09b2657210 # constant +69b1cde5f19ad1ac0cbefd5b # constant +69b1cde8e249e0239ee2996f # constant +69b1cdf8fb0169196bd689fd # constant +69b1cdfa6c46a88e72d2d674 # constant +69b1cdfe666aa62abc6104af # constant +69b1ce02fc265cf0ead4ba46 # constant +69b1ce06e8792f2c51c86ece # constant +69b1ce07d7bb0bc5d6c8c7d1 # constant +69b1ce1f3d4451ef1e9fe1c7 # constant +69b1ce263e73b5b8e5b6ba5e # constant +69b1ce277f8ad153dc9a638c # constant +69b1ce37f3616e385647177e # constant +69b1ce4039f899f2207b2ab6 # constant +69b1ce40fd5324306d5432e4 # constant +69b1ce623dfd90ffca02c3d3 # constant +69b1ce832515ef21966d861b # constant +69b1d116f4b374d812576b1c # constant +69b1d16e6639bd32f1fe8aa6 # constant +69b1d1dfdc49c994c7b72f6e # constant +69b1d1e73561cbf13ab92169 # constant +69b1d2ef9e6318b19bb247b6 # constant +69b1d3436f2b79851798a677 # constant +69b1d3cddc996e5d631e20bf # constant +69b1d4222f3ba15cc222bcd6 # constant +69b1d44750e3da38b9638409 # constant +69b1d55aa7531977d1bfe332 # constant +69b1d7bcb0ecbcc74081336e # constant +69b1d98c7c6435c837f4fea9 # constant +69b1dc66e589ef12931771f2 # constant +69b1dca50b4d6349ecf53998 # constant +69b1de7660260e3bfbdaf314 # constant +69b1dee4467fe92c51e27f89 # constant +69b1df28fb0169196bd68a56 # constant +69b1e018174fcb08ae0a2755 # constant +69b1e13760260e3bfbdaf320 # constant +69b1e716f13f28ee829ae80a # constant +69b1e9524e52e56ddb20e2fa # constant +69b1eb8fdf7ecaecf2f05cd2 # constant +69b1ebae824c8a0015e4536f # constant +69b1ede02ed48da728c4ff44 # constant +69b1ee2f40b88adf88543450 # constant +69b1ee344a4c77ec4448167b # constant +69b1ef37b71aabadada9b9e2 # constant +69b1ef4e6498922f13b26f92 # constant +69b1f00bf02d40438e1d062e # constant +69b1f04228ded605468f269b # constant +69b1f051ccea5bafb0687858 # constant +69b1f0bde8792f2c51c86f91 # constant +69b1f0c2452da522b8ec4eca # constant +69b1f0c5970eeaed1cfd4a70 # constant +69b1f13aa7e483802622d827 # constant +69b1f1457f7f6bba16122c78 # constant +69b1f1f1f8851644abf4781d # constant +69b1f214fa1f51b13782d3cc # jump +69b1f23802b32345b42f1667 # constant +69b1f25168839fc7df2d3efc # constant +69b1f26b0b9284f30cbe1de1 # constant +69b1f278040ec369e9c8096b # constant +69b1f44dc86658441027caf1 # constant +69b1f61f95fa8a42f817b4c3 # constant +69b1f65d1fa85f581929ee11 # constant +69b1f6bcb6eb46262d3e5146 # constant +69b1f88fe3dc955cc0c9a048 # constant +69b1f8f5323cfca1c6d7b5ec # constant +69b1f91581246395cb00b5d5 # constant +69b1f9516a2eb7a29c01c209 # constant +69b1fadd2a8194a3895e2a26 # constant +69b1fb93848e5ebe10d00a70 # constant +69b1fcffe35b059aa3d43fc5 # constant +69b1ff794d95e9221b498eb4 # constant +69b1ffe03fc70fb71500b975 # constant +69b20008f234944494879397 # constant +69b20025b06f8568161285db # constant +69b200d103402af350ad5aa5 # constant +69b200f89d42834d4dd310e7 # constant +69b2013e826e447348ddf911 # constant +69b201694bd0612877d5cea2 # constant +69b201d90f039f41c73469e3 # constant +69b201ffd8ba604bd6793ce7 # constant +69b202afd4ab394764b01ffd # constant +69b203a0d58394ebe352f966 # constant +69b203c798a5e78c16ea84b3 # constant +69b2041b698edefc939fdf81 # jump +69b20447a9ab356d6af9f9e2 # constant +69b204a00aa38b10e4ea1daf # constant +69b204a400828a0a0c869277 # constant +69b204c229d110add752de6a # constant +69b205af73e0c0cb5c646f30 # constant +69b206bf6406f5d89aa9d289 # constant +69b206e7d0d92ee83be9f59c # constant +69b20a18b5a0f7e5d58317e3 # constant +69b20ab1956030fc59477200 # constant +69b20ad89793237729c3ee1c # constant +69b20b2b4901899e19adb707 # constant +69b20b3fe082925304333bf1 # constant +69b20daa2797193a853e433a # constant +69b20e557c22ae1bdf84ab0c # constant +69b20e85f9bf9019ec27e238 # constant +69b20ea5f9bf9019ec27e23a # constant +69b20eafd3a1ac6c37cd72e9 # constant +69b20ef4cda50c2db88751c2 # jump +69b20f38482867e1ebdf0a5a # constant +69b20f49ac171dba7f9cc7d7 # constant +69b20f68da5dabe4210a96a1 # constant +69b20ff0d229d308a8caefba # constant +69b210018f238f3a9d9adb4e # constant +69b21257ba9b7c369f8299b3 # constant +69b2127226ac9f82b256de5d # constant +69b212bc17522960acb2f5c0 # constant +69b212c0e6c38b58c7979491 # constant +69b2142c4d9f0b36cdc5ea24 # constant +69b21438b3ee845714c80660 # constant +69b215f92513c12b3eb4878d # constant +69b2181dcda50c2db88751ef # constant +69b2182db75f730c0941a0e6 # constant +69b21898b088f10190415992 # constant +69b218bf4d41a5486b5c13b2 # constant +69b218f7a9b32465e6c6209b # constant +69b21acacda50c2db887520e # constant +69b21c164fb7611461b271b7 # constant +69b21d296406f5d89aa9d36e # constant +69b21d4671a6e3af58486a3a # constant +69b21d6f44dd6b98769ff9d0 # constant +69b21d7682b2dc4a62c1b1ad # constant +69b21d867366e3614a014668 # constant +69b21df746c3ab7a2fbe205b # constant +69b22163dde762835c6bd38a # constant +69b22195d9d6734431014bfc # constant +69b22219d66bc09df43956c1 # constant +69b222b8308d2b1fb1b01bb4 # constant +69b222dc5b7210edda1fa41c # constant +69b22695302f13a707920b52 # constant +69b226f690800ed928cc8900 # constant +69b22a266db35fe0b3e1b51c # constant +69b22a7f4901899e19adb833 # constant +69b22a9691d453bbb325fad1 # constant +69b22aea5c6a0b0fff5ab727 # constant +69b22b3f1609bd58c95d858e # constant +69b22dabd0d92ee83be9f732 # constant +69b22f0be8d31439937b1ae7 # constant +69b22fa66ca89c92ccabbf1b # constant +69b22fcfba6f75409fc1b02b # constant +69b22fd00e7e8f406a7ae9cb # constant +69b22fff0aa38b10e4ea1fdd # constant +69b231098f86e88a3c385584 # constant +69b2313d341ab72595923043 # constant +69b231a34f6316a90f6fe0b4 # constant +69b231eb6a9477cb194f1123 # constant +69b2327a38afe4964db02cfe # constant +69b2361e62589a2fc271bbbb # constant +69b236725f15f83550ba235a # constant +69b23762c95c551058c60950 # constant +69b238b35326db55a599bc03 # constant +69b23965b08736f01277e9b4 # constant +69b239a7440c90d4b1b4a905 # constant +69b23c5b88f72232a8dc45c4 # constant +69b23e370f039f41c7346b68 # constant +69b2418471cea83d020face6 # constant +69b241a7f9bf9019ec27e41f # constant +69b2428204ba116fc2db1288 # constant +69b2429699d2ae09a233062a # constant +69b244e095c254fe3b23fa15 # constant +69b2451e3f4e9d4e5132ce5c # constant +69b2455d4408ecef199fd978 # jump +69b2457661fd9020e382ee35 # constant +69b2458f7788dece58cf7aee # constant +69b246ccd8ba604bd6793eb2 # constant +69b249246a9477cb194f1230 # constant +69b2497e7bb14b001f7a4f64 # constant +69b2499c7472717347057e98 # constant +69b249e3575202eb1d6de08e # constant +69b2515a85295271d4619057 # constant +69b2553f1a47a5a6d0f15a9f # constant +69b2555dc93d14ed45184738 # constant +69b25563417edec0382bc764 # constant +69b25564d8bad40c1e0685a4 # constant +69b258b6ce9b0a6db91dfc17 # constant +69b259c61582e5851578e6d7 # constant +69b25a750ee6ced376d8fcf1 # constant +69b25b91dde762835c6bd5d5 # constant +69b25c13361580c6a2c0dadd # constant +69b25cf5b46029d141170c8b # constant +69b25cfc1583afad3c9fb1db # constant +69b25d28210c7be6a2e28203 # constant +69b25d32143884432d64a43c # constant +69b25d4091c2de78a452aeb6 # constant +69b25d596db35fe0b3e1b7b3 # constant +69b25e3285ef15ac9d7b7888 # constant +69b25e439aa7a40fb462f5a2 # constant +69b25eaa979179114e07a2c4 # constant +69b25eca35be5cce01bb0aff # constant +69b26102c8d14cdde1ab2668 # constant +69b261300c961ecc5eaf65db # constant +69b26175afa0d5b528c2720a # constant +69b2618d9b9c9569420841a4 # jump +69b261a4f1c6ddd84cf41e33 # constant +69b261daf8a7c5c267f441b1 # constant +69b2621a0caa1714d4f329fb # constant +69b2621b181804f2f1eb1d3b # constant +69b2621beb6efb20f117fb7b # constant +69b2624f0969c15295529103 # constant +69b2625fe4de0a726f678644 # constant +69b26270f6a3a400d9a7e7d7 # constant +69b26277cef1587932e7bf3d # constant +69b2628406e4a575b18cf40a # constant +69b262961834da7ed8aa7ee6 # constant +69b262b3f134f5466995afd4 # constant +69b26302884e6331148898d2 # constant +69b2637436c0e04e86f86d12 # constant +69b265139e3cc49f405fa328 # constant +69b2664f079c97598e6d8af7 # constant +69b2665c6324e399950c1adc # constant +69b2665e48a9e563789c9e4c # constant +69b266fb4dba5f9e2800c97c # constant +69b26727b73762718e5eb6b8 # constant +69b26751361580c6a2c0db6a # constant +69b2681fb92d0e5c40fdd550 # constant +69b268537753d19e1b09fee9 # constant +69b2688825b08fb36c09f0c5 # constant +69b2689fbf2aef4e86d364cb # constant +69b268e4ac23d40ea6190831 # constant +69b268eaf37095194f9814e2 # constant +69b26901b74d46a83eca3a98 # constant +69b269355d1ea87dccb4f0eb # constant +69b26938cbc9621aee962256 # constant +69b2693ac6262d0cfc57e43b # constant +69b2696be27a9bcf3e1ac9bc # jump +69b269ce7e1f7dacd01b557b # constant +69b26a1a0aa38b10e4ea234e # constant +69b26ad97b2a6be8d4bd1011 # constant +69b26b33f170a7b4162b0f8b # constant +69b26c3695fa8a42f817ba6c # constant +69b26c563ade47bf2df05917 # constant +69b26d302f8f84af845b69d4 # jump +69b26d5b46c3ab7a2fbe24d4 # constant +69b26e1065a5b9bc7044d1cf # constant +69b27107a563276bb7763cac # constant +69b2711a7f8ad153dc9a6a33 # constant +69b271dbe634f4e959c8687b # constant +69b271e8b01da9ee1bb516d6 # constant +69b272300b8af5e626ffffec # constant +69b2727684f5d3ba208fc3c1 # constant +69b2727b323cfca1c6d7b90d # constant +69b272843f4e9d4e5132d02b # constant +69b27287e17911b66940efe2 # constant +69b272c97ba3e5cddc08eff7 # constant +69b272dba3a7fae5d0619bea # constant +69b272e135be5cce01bb0c29 # constant +69b29b6ef8effecfb6aeb926 # constant +69b29f330cb082e8dc650627 # constant +69b2bfd398ea9432b2ee0dbd # constant +69b2fefe82df0fdbdca7a017 # constant +69b31b2b7165e989304283dc # constant +69b31fb3122e58b0d29ad6b1 # constant +69b321d1e18f5ea87263039d # constant +69b3222f35bb39eb424af565 # constant +69b3233f8fca50e0c5d4ce14 # constant +69b3237b1eefcd823bba7fae # constant +69b323941f4d2bf0b20ad91b # constant +69b32397b283a20a88e17955 # constant +69b323c2010e5ad85c7da73e # constant +69b323d34d02cc7e23a495ec # constant +69b32401b58b7005e0b3a783 # constant +69b324062433dcedafd51b22 # constant +69b3240eb42d367f838d3d97 # constant +69b3243c01f7f67c1291d374 # constant +69b32450b30324b70ff4492d # constant +69b32460f8cf6ed5ab28932f # constant +69b324616cc9ec3d14a7dd72 # constant +69b3246c3bf215e8fa28700d # constant +69b32483a7689ba6853a5439 # constant +69b3249ebcfac7d916abb9c4 # constant +69b324b23f676d1a02c7983f # jump +69b324b482fbb57ea61937c5 # constant +69b324bd5906bea44a5ddac4 # constant +69b324da363ed901ce3a8d62 # constant +69b324fff0175c96fc1daa21 # constant +69b3250b098c263fb8c36782 # constant +69b3250bffd3937850c8a51e # constant +69b325286e38ef150f637b0a # constant +69b3252b83ecdb45165b4bfd # constant +69b3252c0d1262ccb624d8c9 # constant +69b3253f5f6bc8050b611626 # constant +69b32546c40dfdcea6ba8ee6 # constant +69b325503414f205cdc12850 # constant +69b325717efe6d5409e20582 # constant +69b3257a98ea9432b2ee1084 # constant +69b3257c4a8627413d50cf9f # constant +69b3258727848bc798a731ee # constant +69b3258bc8ec51003a702c15 # constant +69b3258ed0112a77c5d943ab # constant +69b325998998295a8f029555 # constant +69b325a56f67bf880a5e49f9 # constant +69b325ab89601c3154d39ec8 # constant +69b325bb137d10ccc88930d4 # constant +69b325c79eaf89a9a89c2a52 # jump +69b325cdbe19e847fb1a4014 # constant +69b325f67f7fe03c9b69d416 # constant +69b3261590d5a69f2472db72 # constant +69b3268731a915888d2eb0ef # constant +69b326b84de4e40cf051c0dc # constant +69b3275057cff8bc5523a36f # constant +69b327510cb082e8dc6508e6 # constant +69b327b2d56b172436b627da # constant +69b327b92d28bcf464a9b57b # constant +69b327d062d1e5bbf32137db # constant +69b327d3241a6c270e2e2588 # constant +69b327db291d11cb9f27feb3 # constant +69b327e87da6f1e5d5e886ed # constant +69b327f2b3b5402f8573ed83 # jump +69b327f478c4d2bf52285f3c # constant +69b3280f34a638848d0b6064 # constant +69b328288fbf9f1396de70c6 # constant +69b3284534a638848d0b6066 # constant +69b328552c7a598763e72f7a # constant +69b3286ca5696e8fa1144046 # constant +69b3289577e7802c0820d6e9 # constant +69b328a0c60ae08dcc457b7c # constant +69b328cf77e7802c0820d6eb # constant +69b328e3f51fe27231b2e11e # constant +69b3293c3141c306f0648d6a # constant +69b329830dc441a980acab38 # constant +69b329ec104ace7ca9d2b94b # constant +69b32a25e18fcf36cab8e7d6 # constant +69b32b2856a1ff167b595636 # constant +69b32bae6963c69ccf254fb2 # constant +69b32c036ac3fc90331f7882 # constant +69b32c20a01f387ed0a0d110 # constant +69b32c6f6f3c496b94e5534b # constant +69b32c812a265566fbd025ab # constant +69b32cec3a8d3735c57f1f19 # constant +69b32d0da47704554ca335f0 # constant +69b32ddbaceaf3fb3c0dc222 # constant +69b32def931340ed5fe82b83 # constant +69b32e2188155012c3c44bbe # constant +69b32e3ce2c63e32fb9fa556 # constant +69b32e95515dbd55d8c391d8 # constant +69b32e9a8e6de7b3274332ca # constant +69b32eb4e6f429f43f570ec3 # constant +69b32eb8787568937fb40927 # constant +69b32edb27848bc798a73266 # constant +69b32edffdd826d94553cd04 # constant +69b32ee07a1c6f3ad1bf1f2e # constant +69b32ee0f555c89bc668df46 # constant +69b32efce35ae6f7c9cc7d24 # constant +69b32f154e09c0204d2d9c6a # constant +69b32f2395e9f2d8310b5beb # jump +69b32f24dcfbd323e893ae4c # constant +69b32f2510072caf7763aa04 # constant +69b32f32b6743420cb324f5d # constant +69b32f339c82aa4347978f8e # constant +69b32f5a03501cccc085eb1e # constant +69b32f6633f47b3ac2cc9968 # constant +69b32f75fecf5d3702dfc006 # constant +69b32f76b35b907423cd01b6 # constant +69b32f89614469a95a00b746 # constant +69b32f8a87b3db5f42a368ba # constant +69b32f8cdfe7e66f1b54ced2 # constant +69b32f9805adab99ca3fb86b # constant +69b32f9a6d548cf3375ad80f # constant +69b32fb1a046a1dfae210ed1 # constant +69b32fb300a8b752c25014b3 # constant +69b32fbc6bcc95c86de27fa4 # constant +69b32fc512ad73152982ce24 # constant +69b32fc533aae35229f6ab95 # constant +69b32fc6c2e9ac0400a04c48 # constant +69b32fd3d355097cd9aff922 # constant +69b32fe613fccda4be4ce9c3 # constant +69b32feec4b840ddacc8633f # constant +69b32ffac7f30273637b9963 # constant +69b32ffe704b9f2c093d8ec6 # constant +69b3300090ee52de87c986d6 # constant +69b33006ae26e0bf2e8d32a6 # constant +69b33017d355097cd9aff926 # constant +69b3301e3e05c25efa26ccc3 # constant +69b33020d841d4e6bee39166 # constant +69b33023d5a9b3562d25f962 # constant +69b3302877b22b669fa7de03 # constant +69b33041290f064f72218f61 # constant +69b33044c971c19f7392fc96 # constant +69b3304cc22cbfcd46725ada # constant +69b3304f665cf1dc45af8b65 # constant +69b33050740e46344116bfd4 # constant +69b3306db284d7d2c8bfb945 # constant +69b3307108808ecd6ad0c4e0 # constant +69b33073931b75760c6ee820 # constant +69b33075644ff12da3d093ca # constant +69b33080a4fe1768fd5bb181 # constant +69b330842b540726f0f0838a # constant +69b3309c655c0e891fd8ee07 # constant +69b330abbe550c1ec1cd1f3e # constant +69b330b1bdc42bf69a583cb2 # constant +69b330bf13e59a723bbc6fec # constant +69b330c7fec4f50e27148868 # constant +69b330deca7f88e8cd81cadd # constant +69b330e1c9bbf6efb8566ebc # constant +69b330e4b35b907423cd01d3 # constant +69b330ea42c0b84f1a22282d # constant +69b330f521d402c4c8723bbd # constant +69b33100314218860b190ae2 # constant +69b3310514547e661bd507c9 # constant +69b3311bdcd0233fa55bc545 # constant +69b3311c34e4864b3dec0d8b # constant +69b3313966b25b9d418d5fa2 # constant +69b33167bc5201c2b1793d88 # constant +69b3317bb13c7d0ff2ec4037 # constant +69b33243366031dc8c47dcf6 # constant +69b3324e328a28db44f5d0b1 # constant +69b3327e2ab5c44e3b76e209 # constant +69b3329aab764680c23531ca # constant +69b332a2b0cb758673ba56e1 # constant +69b332a48090086eef6d4973 # constant +69b332a9f6fd3ef55b0bba47 # constant +69b332d29e6808b3b5e498c9 # constant +69b332e262b44d4b598096c3 # constant +69b3333b8f71eb09f42050ae # constant +69b3334adcd2f45bbef47929 # constant +69b333567aaac69b6ecbeec3 # jump +69b3336112a1e98061c326fa # constant +69b33372c971c19f7392fcbf # constant +69b333776e38ef150f637bca # constant +69b33382e6f175f3ad6f6e0e # constant +69b333920c5fc9c39cb6cc0c # constant +69b3339d13c67cb9ab2753bc # constant +69b3339e8d537e229a25aff0 # constant +69b333a9f51fe27231b2e18c # constant +69b333aaff90c39ab620338e # constant +69b333ab543751ab8d897122 # constant +69b333ba2a265566fbd02617 # constant +69b333c834f57e26be193cef # constant +69b333d81572f433b1c4ca26 # constant +69b333d99c9ee25fe2398998 # constant +69b333e186ea2fa9708cb300 # constant +69b333ef8d138d8ced227ba5 # constant +69b3344e95e8c2980925b257 # constant +69b33478c368ce0fd209e914 # constant +69b334d7fe6b0d3b57f71136 # constant +69b335c88be9bb70fdde77f4 # constant +69b336195162be9bfc0c9f89 # constant +69b33661f0175c96fc1dab00 # constant +69b3368e3aa694fbd6c719ed # constant +69b336a66717bcea4d6fc7f4 # constant +69b336b319afedabed6d0cd9 # constant +69b336d1a8a94add5376fc36 # constant +69b337707538c5a928643d6b # constant +69b337a09a1f969837d7d9ec # constant +69b33af6e7a1f9d2fa7283fb # constant +69b33b57a9c177462f0277e7 # constant +69b33c08fc135f84a7fc69ad # constant +69b33c72c61e266c110d2918 # constant +69b33c7e246f555f18d96d30 # constant +69b33e27536b3f17198194cc # constant +69b33ea7a0b9e58bff02cadf # constant +69b33f14a4fe1768fd5bb1cd # constant +69b33f2a5031f15daf422269 # constant +69b33f6d6202365215d2e5f4 # constant +69b33ffc4403d08cc7c945f7 # constant +69b340338780567a95e8f5be # constant +69b340aab9be0be7a2f466eb # constant +69b3413745d67ce5881fb493 # constant +69b341ad366031dc8c47dd44 # constant +69b341ca862a7bcaaab3e0e0 # constant +69b342048796b61661c69f26 # constant +69b34219f59e011f372d00b4 # constant +69b3423d5103cf7092357fae # constant +69b3426f7ca9894dff005116 # constant +69b3430d911a770c69a48f3b # constant +69b343419de26e769dc5fca9 # constant +69b3434598f474c17d2d7119 # constant +69b3438f39afc1f2ee252215 # constant +69b343bc3bf215e8fa2870fe # constant +69b344966656f58acf8a0e17 # constant +69b344b739152d04fb349b3f # constant +69b344b9d9c8e72b11700919 # constant +69b3456be6f429f43f570f06 # constant +69b345b6d9c8e72b1170091d # constant +69b347059ad4198e1f47a608 # constant +69b3475f7538c5a928643d95 # constant +69b3478c80ef0cbdf4f3e8fe # constant +69b347cc0e4b935c6f9f4e9d # constant +69b348b84888be6ca72f7b59 # constant +69b349a8af2033fbacbee6fa # constant +69b34ac56aa346452d1ef78d # constant +69b34af98648e4230a699553 # constant +69b34b13a6a7f2318ff02ae7 # constant +69b34b56740e46344116c019 # constant +69b34cba4d00f2be43082e80 # constant +69b34cd62b037d820b4965a8 # constant +69b34d29482b50fa8312ad9e # constant +69b34e6a59bc7c30d2471cbc # constant +69b34fb94e780ed4635cbf6f # constant +69b35078482b50fa8312ada6 # constant +69b3507d5f038b19a31cff59 # constant +69b350a2b7fca73bbb147055 # constant +69b350fba46d1a7a1498bc18 # constant +69b3511a95aafca64473e092 # constant +69b3518a4ed12a926a0ecb28 # constant +69b352d2bae757f659fdf346 # constant +69b360e109abffe5b35da395 # constant +69b3610f977aa45ae4d31ae7 # constant +69b361891edd75c6bfdefcac # constant +69b363a51e56f6bdf5a37397 # constant +69b363e6db737c35559c53f4 # constant +69b3641bebb927a3ce3dbb7a # constant +69b36455f86ba06c04272d6c # constant +69b364f99afaa0a8c9b4963b # constant +69b3656486974057bbbebd13 # constant +69b3656f77947a6c4ce3a3e4 # constant +69b36587448cacd57cdbef10 # constant +69b367987f8035cd758e6a32 # constant +69b368d9c68645217fa07e02 # constant +69b369057efe6d5409e20823 # constant +69b369c903e21a10856d461b # constant +69b369f07b083c90587bb9b1 # constant +69b36a7fb8e7cfef07cf7191 # constant +69b36aeaac604c29b4a33284 # constant +69b36b4331acde346400f5e2 # constant +69b36b922ffeee6dd3707fe7 # constant +69b36ba33a81791a4497ac81 # constant +69b36be741ba7893dfe001e3 # constant +69b36beb550b55accef36a30 # constant +69b36c04c00a906d239edf8e # constant +69b36c198b7e31850a2c3ef5 # constant +69b36c23d03e216772694110 # constant +69b36c2d1137f38a3a2dd9f4 # constant +69b36cac069e4261884f5eb3 # constant +69b36cbb7efe6d5409e2085c # constant +69b36db130df4858f6ffacb6 # constant +69b36e82dfb2598257c534d2 # constant +69b36f5448431c26200c37b8 # constant +69b36f6414a9f4b4e4fc57a1 # constant +69b371845475daac769f456c # constant +69b3718fbfec88fa6221bf47 # constant +69b37308ebb927a3ce3dbc3d # constant +69b373b9bb87ec6694cf2094 # constant +69b37863fcdb29e7ff12c807 # constant +69b37971806f8dde7843c998 # constant +69b37bd08fad355b49bb1aab # constant +69b37bfe56f8acacb6b7b378 # constant +69b37c695298b5a76ce8073d # constant +69b37de014a9f4b4e4fc5826 # constant +69b37e2803323db35c6d48f8 # constant +69b37eadda671b05d56ff3cf # constant +69b37f54f186f27a705a6b5d # jump +69b37f77254f96a47b997d08 # constant +69b37fd530df4858f6ffad73 # constant +69b38e6d8c0921f26676aa53 # constant +69b38ea5ac1fb583171b0a61 # constant +69b391223e30f5215d040d9e # constant +69b39145b93646b92ef2d877 # constant +69b391e3338f88c169971895 # constant +69b3922b210c48e1246ceb50 # constant +69b395313fb2fc1e7790afb7 # constant +69b3955e663d4dc7043eb4f9 # constant +69b395dcc054f66bb940cb25 # constant +69b398b7daef557c66e76a3d # constant +69b3992369f908ca24f60e0c # constant +69b39aa6b4a3e4257c829439 # constant +69b39b363c979f1b0ce8d126 # constant +69b39b93cd115d5dccbbb0c8 # jump +69b39c0426866068d9207a74 # constant +69b39caf30df4858f6ffaf48 # constant +69b39cf0f7ff0aaf85af89c5 # constant +69b39d002225c757449566e2 # constant +69b39df24641838925653c7e # constant +69b39e153e7466349feaf73f # constant +69b39e56e8f50389c2a6f962 # constant +69b3a0fa81da7a597ce95881 # constant +69b3a1950f672f8186f91ccd # constant +69b3a1ab7be90cb966537bb1 # constant +69b3a30f662586b764e59fc3 # constant +69b3a42b17c5294b8621c552 # constant +69b3a42f575e091651bb433f # constant +69b3a73f3a3f89b0cbc05f89 # constant +69b3a752ff53bfe0d77b1229 # constant +69b3a76709211fe5a7944933 # constant +69b3a7bb90cbf3aae2e30c12 # constant +69b3a7cd6ba06e41288994f5 # constant +69b3a7f65c1b524afe3b2d48 # constant +69b3a8c45679ae4949597413 # constant +69b3a8df4befb83ac3ee1656 # constant +69b3a9b189e4a30fd19da14b # constant +69b3a9b19faa41a631ebc3cf # constant +69b3a9f5139de72c2b7d484a # constant +69b3aa059623b99a976ea7bd # constant +69b3aa331ebebadcc7a4a262 # constant +69b3aa39f6b72542f6910b4d # constant +69b3aa3ec8794e8dc30148b7 # constant +69b3aa8eb66719dba38e32c3 # constant +69b3aae4acadccdf28af12d4 # constant +69b3ab275546cf946e250803 # constant +69b3ab4a434b2fe40e146f38 # constant +69b3abcf4a1e3cf594e0044c # constant +69b3abd34ae110d0d2931af0 # constant +69b3ac246e13dcd2761105aa # constant +69b3ac7467f866e3883fa3c0 # constant +69b3ac956623282edcf5a426 # constant +69b3aea8ce45115aef988e34 # constant +69b3af32a2d29f6909d85216 # constant +69b3af3be5cc6a9c801d520e # constant +69b3af6993237f6691b1d8e7 # constant +69b3af89cd1838f0f0844ed7 # constant +69b3b00f25b5da7c1c93fbe1 # constant +69b3b038dea844ed80048ff8 # constant +69b3b03dd4d65c71b3e458fb # constant +69b3b03e5b48470e94ad461a # constant +69b3b0576c27226f07d12536 # constant +69b3b35215919ee04e7c74a7 # constant +69b3b99a479e26e892851da8 # constant +69b3baccda9a3efeac01bcb2 # constant +69b3bbacac9b11a4d5640676 # constant +69b3bc80de082ecce36d7c87 # constant +69b3bcf85fbb70bd57cdc739 # constant +69b3c1e59e9cd7379cfb9705 # constant +69b3c211beeef798925f72ed # constant +69b3c2a08c925d17b9076024 # constant +69b3c35df9d3e87d7b77d8df # constant +69b3c360d1665a2c6f45255d # constant +69b3c441f498361f83ea2bae # constant +69b3c4ede9de50b7d0f6e9ca # constant +69b3c796880ee28a0531867c # constant +69b3c81091b8aab9d333b1e1 # constant +69b3c9a772d022b8d4468c26 # jump +69b3d9014f4b006a4b4a2c0b # constant +69b3dbdead93f729269e551f # constant +69b3e5a200e6d908f86d7ec0 # constant +69b3e8565c01f1f10da71161 # constant +69b3e9ec46dbfcf8d0484d75 # constant +69b3eb2263b4d10580cf794a # constant +69b3ec5ebb2330fd692791d9 # constant +69b3f3253e8130e694cfbb81 # constant +69b3fda2d43107729f702ae0 # constant +69b4018da13ae541201c6f55 # constant +69b40701f3077050a50788d0 # constant +69b409cd72d022b8d4468fb1 # constant +69b40a77165ee7193a2bc1e4 # constant +69b4291c7ae6fe5f6c5a382c # constant +69b42f3a5b4047fbbfc4cf4e # constant +69b42fa785e36f35da0ba68b # constant +69b43048665ebcb51d1e0799 # constant +69b4316360172888461a7ed1 # constant +69b436f1ce1481ad9038690e # constant +69b43bc08579b939089f4489 # constant +69b43ccb6462aeb78fd681e6 # constant +69b43e1c9b16ea798a6fae4e # constant +69b43e314b2afacec3803545 # constant +69b43f198a313c3e6edaab44 # constant +69b43f5b7f0b851033657119 # constant +69b441ecc375af84c0576a0c # constant +69b4436a78673e5de422822b # constant +69b445f8a2bb2f3fc33fca82 # constant +69b446f8515d4294ccc78de7 # constant +69b44781a9f25eecea6b5396 # constant +69b44970151050b0854f8a6d # constant +69b449c2035c3af18343724d # constant +69b44fb69a3e364072eff9b8 # constant +69b451c3f109813904030ae1 # constant +69b4584bb8816c7c6f3c8dac # constant +69b459299ee0b627f1fcf40f # constant +69b45b101e6f0569de9c5d9c # constant +69b45c071843f1c2dd2bb763 # constant +69b45f2932643e552e8851b3 # constant +69b4603fc8f3797a5f1f6313 # constant +69b460b5ffea0a68ed9d00bd # constant +69b4666923d820df9de53074 # constant +69b4719e9f66a2746698f42b # constant +69b475015676713aae693420 # constant +69b475d31ce35576286e75dc # constant +69b477b9cfe53ea1af65f902 # constant +69b4792e1b02c068cd2ad021 # constant +69b47ebb969b18fb10eb6e71 # constant +69b47ece2cea44729174520f # constant +69b485856d78ad6736570afd # constant +69b485b5b344fe224b57c69c # constant +69b487286144f027ad941eb8 # constant +69b4883ab504f66b2436a39a # constant +69b48ae59f8e3016601d9b16 # constant +69b48cbacd6846904ec6cc09 # constant +69b48dfed7a5dd9bef796a58 # constant +69b48f7c4ee473f8af6c3b98 # constant +69b48f8c5b4a7686bbe200bb # constant +69b48fba71e1c59c5b4c3d1f # constant +69b4913966535f9abb1849b1 # constant +69b49220e2c65ab505ae0aa2 # constant +69b49263738fa090902cb75f # constant +69b492a709a939795bdb7985 # constant +69b492de4e5ac9677e3a22e1 # constant +69b4932af62243c65718e02a # constant +69b4932f87abff9f5759793e # constant +69b49335cf653206178853da # constant +69b493541ee47c072a367c5d # constant +69b4937ea3213aef75c95eba # constant +69b493ad4f4b006a4b4a3717 # constant +69b493af7255ccd7c00a821b # constant +69b493d8e985757268c73336 # constant +69b493e7e2e62141d97cc879 # constant +69b4940ebb067226df72d0ea # constant +69b494122f2fa993841eb9a4 # constant +69b49431e27d2b8e1ab6af18 # constant +69b49440868748401c8c0e29 # constant +69b49445ca9fe7e1f7add1b5 # constant +69b49459a321d4f716b74d55 # constant +69b4945a53e765a22cddd159 # constant +69b4947c95bb5d149ad3dcdd # constant +69b494875058b9f00b9fa38f # constant +69b49491bb3fe1af3a8731f2 # constant +69b4949b7acc5d78c8107ce9 # constant +69b4949f0c02cfba7bd9cda8 # constant +69b494a44e58e40d05210269 # constant +69b494a988ffae5aad001ce1 # constant +69b494aad93db6d47a9fa32d # constant +69b494b15d6b48c59a932005 # constant +69b494caffea0a68ed9d01d4 # constant +69b494ddb8e17c5baf01db63 # constant +69b494e17fe01837a6e53fc9 # constant +69b4951290be9c6b2a9d97c9 # constant +69b4951662d831a509f7c9bb # constant +69b495248f03e03188e8a819 # constant +69b4952574320179c62b44ed # constant +69b495278ae23973608291c6 # constant +69b4952973ca03704a502741 # constant +69b495369b3c11bddef4abed # constant +69b495447f92b7bf22b566db # constant +69b49552c7e02bc8d412ac5e # constant +69b49554bf16bae4bd7cea85 # constant +69b4955ba13067e31cf8ae78 # constant +69b495659c40481a89baf665 # constant +69b49566c9e23487ffa60fb0 # constant +69b49572af135b8270eb1f6a # constant +69b49573eac64bf65ebcd462 # constant +69b495844f6aaaa9cd1bd923 # constant +69b4958df62243c65718e072 # constant +69b4959d09ee8318e185a1c5 # constant +69b495a28574928ac9a8de30 # constant +69b495aa6b3066ada527b6be # constant +69b495adcfe023119bddcd9e # constant +69b495da7230230e5fb80c02 # constant +69b495e416633723bd0b2804 # constant +69b495e558a78ed4c167d063 # constant +69b495e96bb53d53abcf6175 # constant +69b495fce2259218030e569e # constant +69b495ff496f6f8e424401e2 # constant +69b49601a83a7b34affb4573 # constant +69b4960619cc1d0464e1558d # jump +69b4960b86c1932505c2ed06 # constant +69b4960e8505130f8b0e6c7f # constant +69b4960f6b3066ada527b6c8 # constant +69b49618ad5941fd367947f7 # constant +69b4962b58f5afe39b1de493 # constant +69b4962b5c1d84f9a8fe83d2 # constant +69b4962cbaf2957a279cf59a # constant +69b4963020714062f4b0b4c5 # constant +69b4964a17f7000b991af0de # constant +69b4966917e3a9ccf08b22ff # constant +69b4966e0273df551d93bd4d # constant +69b49676817936a8e39f406b # constant +69b4968d2463c9433cf2f33f # constant +69b496a6389b1557c3506902 # constant +69b496add51e511a082342b0 # constant +69b496ba1e229be1408463e4 # constant +69b496c930ac9bf923d7f170 # constant +69b496d2fea01fd1b81d68a9 # jump +69b496e991b851da8626ab48 # constant +69b496f66b262d6ad3db4444 # constant +69b49719d61065cc8a576e8a # constant +69b49730843d1189f0e964a3 # constant +69b4975b9275ad7c73fdb59d # constant +69b4975cad8cb8d985c9a12d # constant +69b49763b30844b29a8324d6 # constant +69b497a33ae16cd77636dfe3 # constant +69b497f71aa03dd5179621c7 # constant +69b4988c99f7c80d1fa6fdb5 # constant +69b49a4ac2a8c7a3ea5b0b1d # constant +69b49bcceb89585320c07443 # constant +69b49bf4827f8909e882675a # constant +69b49c1a200b571aab71442d # constant +69b49c7d09a939795bdb79fe # constant +69b49c96d5bde809a4eebb79 # constant +69b49ca7c5590bdee651851d # constant +69b49ca9738fa090902cb7e6 # constant +69b49cc68a4c8aae9bc764d2 # constant +69b49cc8dcdac3bf3149070b # constant +69b49cd14f712d2ccd55b363 # constant +69b49cd3ab37121e9e93bfa8 # constant +69b49cf7a3c0ce02c5a55881 # constant +69b49cfdc042529718bf8cce # constant +69b49d04335d8a8a3c8df6f0 # constant +69b49d2419cc1d0464e155fd # constant +69b49d2e8781fdde8bbaa9ca # constant +69b49d2fcee7dbe8803e396f # constant +69b49d3307347b484b3fd740 # constant +69b49d4a6641c34ca8db13f0 # constant +69b49d4d124502aa36324f17 # constant +69b49d536325cbcd4232ac72 # constant +69b49d5f7927e426d4bdd65f # constant +69b49d7723ef41cf5c6064c8 # constant +69b49d796e8f13cd1310c6c2 # constant +69b49d7a51d1604985dd2e62 # constant +69b49de09e3c339e3d3e891d # constant +69b49fa71aae85955a5916d9 # constant +69b49fb3f429dfde897d9093 # constant +69b49fc7b9c83720b195b67f # constant +69b49fd14593e8cbb1079ff6 # constant +69b4a022763fc35351a5989e # constant +69b4a043d292e88e08f7f96d # jump +69b4a0a03b17fa6e06ad16ac # constant +69b4a0ac142b1d1af817d8ba # constant +69b4a0b4d9bd4ed2c8cda3ab # constant +69b4a0bd61b1e5e3f04929c0 # constant +69b4a0c1cd67f708f6e64aa6 # constant +69b4a0c4f327a9232dd1b2c6 # constant +69b4a0f504a889619e4bdd61 # constant +69b4a0f8d6e27d500501f5b6 # constant +69b4a0f95d4941345a0751d5 # constant +69b4a0fcce3174949ada5597 # constant +69b4a12e9c7daebcace61d82 # constant +69b4a16a970f3d65c4549ff4 # constant +69b4a16bd8e164bfe686509c # constant +69b4a16bfea01fd1b81d6958 # constant +69b4a17b02e594024304cc6e # constant +69b4a18e87a763f35834ab2f # constant +69b4a1d6c7834bedf19720cc # constant +69b4a1e1bc1711c9daf2dfb1 # constant +69b4a1e85907c64c2ffa31c5 # constant +69b4a1f3ba3960335c191583 # constant +69b4a21653102801f5122735 # constant +69b4a21b6462aeb78fd68830 # constant +69b4a21d85786b94b9725d73 # constant +69b4a235bf16bae4bd7ceb27 # constant +69b4a239ee25266722a7562c # constant +69b4a23a1aa03dd51796224a # constant +69b4a2ddca9fe7e1f7add2b5 # constant +69b4a2e230fc0dc05fe97616 # constant +69b4a33c39b41b04e92756e4 # constant +69b4a359486a23130a8381b5 # constant +69b4a38a4c035879577bfc55 # constant +69b4a3954e5ac9677e3a23ba # constant +69b4a395d49f14c6bda8beba # constant +69b4a3984c035879577bfc57 # constant +69b4a39dfdce20eb7f3240b2 # constant +69b4a3aabc4c7ea1d5d299f1 # constant +69b4a3abbb2597d3ebd236a2 # constant +69b4a3d9230639e424fd1c03 # constant +69b4a530c3f9ed0a1125a886 # constant +69b4a550af135b8270eb207e # constant +69b4a599ea0478dac4766f5a # constant +69b4a59c479220aa237ec4a7 # constant +69b4a5a91460ea5ea1cc3f5c # constant +69b4a6cf84179943aec08278 # constant +69b4a6e645a65ae24232add6 # constant +69b4a72b9bc2ce0e6641c8d6 # constant +69b4a77d53c39a7b7561c4c5 # constant +69b4a7f6506144ea4e04aee6 # constant +69b4a80cce95af8cea9a33d5 # constant +69b4a868acea937039fea414 # constant +69b4a87c06e01260c1ac90c6 # constant +69b4a8b0618df5d7d36be42f # constant +69b4a8b3a571e5b46d954610 # constant +69b4a8b7433b19a99c73126b # constant +69b4a8db80026d5c84e9cad4 # constant +69b4a9127eb149d6a6709832 # constant +69b4a93040272c2234cd2dd0 # constant +69b4a981938a328b92f66949 # constant +69b4a9d541c2dcb85b2e2c36 # constant +69b4a9ea388ccecd105c5b92 # constant +69b4a9f273ca03704a5028d9 # constant +69b4aa4d2127166166f18bc1 # constant +69b4aa74a14f5f8fb96d7201 # constant +69b4aa9b3db5e101c3b64d77 # constant +69b4aab66a7faf595b5f2413 # constant +69b4aab9c61a380b28232572 # constant +69b4aabc728fdfa816cbf558 # constant +69b4aadaf9274b0a85254637 # constant +69b4ab19af135b8270eb20d4 # constant +69b4ab333076173dad5af052 # constant +69b4ab4e4e5ac9677e3a2445 # constant +69b4ab70e328b8f0f2640e3d # constant +69b4ab710b666ac17c526211 # constant +69b4ab73e2e62141d97cc9ed # constant +69b4aba7a97937a3e2fb3736 # constant +69b4abd6f629796beb983f82 # constant +69b4abde04f0e48996918d31 # constant +69b4abe2f19ed0467e17172f # constant +69b4abfc9f8e3016601d9cb0 # constant +69b4ac05c61a380b28232587 # jump +69b4ac164184c502d40809eb # constant +69b4ac216e9a953c6e6c97f6 # constant +69b4ac346f257391be694bdf # constant +69b4ac399f0b814536fde190 # constant +69b4ac3df688821d294a1af3 # constant +69b4ac49569652660a4f4ae1 # constant +69b4ac5899e58efaf1232ad1 # constant +69b4ac6f2f3ab8fec51cfa89 # constant +69b4ac71969b18fb10eb6ffb # constant +69b4ac71b3512f3d7ede1c18 # constant +69b4ac7a53e0e51b6dd346a7 # constant +69b4ac7c8ca5423dab74c3d5 # constant +69b4ac7f77f4dc69622ceaff # constant +69b4ac9b83db69b280d2975e # constant +69b4aca0f0c88edf01a631a0 # constant +69b4aca3fc5f3e6e2e68bfa5 # constant +69b4acb26a0a108e820b110e # constant +69b4acc45b4a7686bbe20277 # constant +69b4acd9b8745dc1f264d4d8 # constant +69b4acea281cb4ce74d81ebc # constant +69b4aced699f525ceebed277 # constant +69b4acf23282e47c8735365c # constant +69b4ad0a335d8a8a3c8df7d6 # constant +69b4ad37f70c81ef9281d2cb # constant +69b4ad3c9730ecddb446bc37 # constant +69b4ad4faa4961e178763677 # constant +69b4ad51d95f0cb5ffa2b961 # constant +69b4ad5be5be71666d494451 # constant +69b4ad7f8c0e050035cc5b50 # constant +69b4ad88b564ee80404a389e # constant +69b4ad96d5bde809a4eebd26 # constant +69b4ad9c4a206179cbd000b9 # constant +69b4ada8e38bde62855cf6a8 # jump +69b4adb172afcd77da62feff # constant +69b4adb42634ac70cd74ec56 # constant +69b4adc289deb7a84c2d5601 # constant +69b4adc4f691e55c03ff810b # constant +69b4adda654dd32a33b3b8e4 # constant +69b4addd360d4a79220912d0 # constant +69b4adea98e0f543f389f9c7 # constant +69b4ae005907c64c2ffa3214 # constant +69b4ae067eb149d6a67098a1 # constant +69b4ae126bac906263f2491a # constant +69b4ae265a10c4a2d8f9fe84 # constant +69b4ae3453e0e51b6dd346d7 # constant +69b4ae63831d1b8c4d6307c4 # constant +69b4ae6735cef09fa62e2a6a # constant +69b4ae6f0b666ac17c52626b # constant +69b4ae7914667611272280a1 # jump +69b4ae7f3758288247afca56 # constant +69b4ae85d7a5dd9bef796b51 # constant +69b4ae8f656475939c9bda52 # constant +69b4ae916efa9e5f0edc79b4 # constant +69b4ae9b17ec01169beb3f88 # constant +69b4ae9cd292e88e08f7fb40 # constant +69b4aeadb8816c7c6f3c90f3 # constant +69b4aebb1d8d8d4ad85ca0c4 # constant +69b4aede2e3fada4b32754d6 # constant +69b4aee88415a3e84cfa8aa1 # constant +69b4aeee31cf18f433028965 # constant +69b4af0038720da38b3e510a # constant +69b4af0eb83417cc7109b92f # constant +69b4af10c7b864206a9a6f62 # constant +69b4af1c9b2ae0fb15ca36c2 # constant +69b4af1db8e17c5baf01dd40 # constant +69b4af29039109330b23b7a2 # constant +69b4af2cedfc296664603752 # constant +69b4af3b38fbfdf7b20f1174 # constant +69b4af4c4efa38975a2e1561 # constant +69b4af5a2634ac70cd74ec72 # constant +69b4af5b5f4789f4f729289e # constant +69b4af5d590494387f207c16 # constant +69b4af6c8c61b5e3f72f9917 # constant +69b4af75a8ea953037b9fa03 # constant +69b4af8389267616a20eac50 # constant +69b4af8451fb8a949dbb7f40 # constant +69b4afcf53e0e51b6dd34705 # constant +69b4afea6e93f1bed9b4d6b8 # constant +69b4b00b2fd6713c1d3cb814 # constant +69b4b010f6333e2fa8a2c522 # constant +69b4b039bc4c7ea1d5d29a9d # constant +69b4b043d5a535d5c7d1737d # constant +69b4b05ab84aad9597edc45b # constant +69b4b05f36c567220062903d # constant +69b4b06823b8cfce653bcb1e # constant +69b4b0852bcba0117f9f0379 # constant +69b4b0927d741d012b44f8c5 # constant +69b4b095ba67778804afbf9d # constant +69b4b09ef14d5ff1cd3afd8f # constant +69b4b0b2e21a4f95872c4c7d # constant +69b4b11deac64bf65ebcd537 # constant +69b4b120388ccecd105c5c35 # constant +69b4b124a2e0f23556e6a69a # constant +69b4b1364e5ac9677e3a24ab # constant +69b4b138d43107729f7031ef # constant +69b4b143c7834bedf19721f2 # constant +69b4b14cf01bb0c81287136e # constant +69b4b154f0f8027b635c7745 # constant +69b4b159396463451d97c689 # constant +69b4b15ead910fda4c7fc4bb # constant +69b4b16cf9c34971b9f6d333 # constant +69b4b174ab39f27198e315cb # constant +69b4b179a7cd7cb244b9eb30 # constant +69b4b17f9a0f35ae6bd286b0 # constant +69b4b1c414a4d59b08be0031 # constant +69b4b1c9a68ffd2fbd98a96d # constant +69b4b1d1857c353e098d8345 # constant +69b4b1de025d49131a128ea5 # constant +69b4b1f15a10c4a2d8f9fec2 # constant +69b4b1f90e55f47ba88c2fd8 # constant +69b4b2112fca00d0b66addb4 # constant +69b4b213ccbdf1e86f1727ff # constant +69b4b2336bac906263f24963 # constant +69b4b234b504f66b2436a6f4 # constant +69b4b23afd18e37694a9439d # constant +69b4b2617fe01837a6e54183 # constant +69b4b26a30fc0dc05fe9771e # constant +69b4b26b0d48c3f6ede856b1 # constant +69b4b26e02c72bff69b7ecf9 # constant +69b4b27e3f77b99a0a79c214 # constant +69b4b29b0c01a69644a21e61 # jump +69b4b29ba85a9038236c3e71 # constant +69b4b2b4ea0f93b1d3bfa01c # constant +69b4b2b59375cfad34293d14 # constant +69b4b2c81aae85955a59196c # constant +69b4b2c923b8cfce653bcb5e # constant +69b4b2e8af135b8270eb2178 # constant +69b4b2f63f77b99a0a79c221 # constant +69b4b3173282e47c873536c0 # constant +69b4b31bc3a8dbc608ff8c58 # constant +69b4b31eaf085e525cb19381 # constant +69b4b320435f9e0bc40152d8 # constant +69b4b3211ee47c072a367e08 # constant +69b4b334515d4294ccc79228 # constant +69b4b366d40b3f62385738ac # constant +69b4b369b8993611e542417e # constant +69b4b372cde059c2090359ae # constant +69b4b38197224d24a5b38317 # constant +69b4b3843ac24440ed54629c # constant +69b4b38628737cee656ce740 # constant +69b4b39451fb8a949dbb7fa3 # constant +69b4b395025d49131a128eb7 # constant +69b4b3966b9429ed3f34daf9 # constant +69b4b3982e9983f31a0b57b7 # constant +69b4b3a5c3f9ed0a1125a9eb # constant +69b4b3afd902e000df8b958e # constant +69b4b3c16366bd7d9755e02e # constant +69b4b3d6d99f32c1d5332245 # constant +69b4b403d99f32c1d5332249 # constant +69b4b40b5fa639c4e24e7473 # constant +69b4b4149730ecddb446bcb7 # constant +69b4b4166641c34ca8db1550 # constant +69b4b4327e46044c0318cb62 # constant +69b4b440060a6199caac111c # constant +69b4b45a460cd47c547c8743 # constant +69b4b4670b666ac17c526312 # constant +69b4b471335d8a8a3c8df881 # constant +69b4b477cd6846904ec6cfa7 # constant +69b4b479f29408487e3faa0c # constant +69b4b47b2a05fa27685f091c # constant +69b4b495e2e62141d97ccac6 # constant +69b4b4ae99f7c80d1fa6fecc # constant +69b4b4decd6846904ec6cfad # constant +69b4b4f31466761127228152 # constant +69b4b4fc32580579bce8e750 # constant +69b4b5079bc2ce0e6641c9af # constant +69b4b5530b666ac17c526323 # constant +69b4b555b564ee80404a393b # constant +69b4b561a7585840264c6679 # constant +69b4b579f0f8027b635c7798 # constant +69b4b580b564ee80404a393f # constant +69b4b588260e0757c5fecd1b # constant +69b4b5988b0d76bdb828823d # constant +69b4b5c3f05424f9a3a2e1dc # constant +69b4b5d6bb2b154f8ddd6749 # constant +69b4b60dd0103fb8e4fb4bfa # constant +69b4b634135e07fa474147c7 # constant +69b4b63e45a65ae24232aed5 # constant +69b4b677ad910fda4c7fc4fb # constant +69b4b68323b8cfce653bcbc4 # constant +69b4b6896144f027ad9420b0 # constant +69b4b6a36b9429ed3f34db26 # constant +69b4b6ec16d52405b4a2c1e1 # constant +69b4b72cad9ebb37e8776a42 # constant +69b4b73665f80a5330cb3a04 # constant +69b4b7548579b939089f4810 # constant +69b4b754f83765e0c840cf03 # constant +69b4b76f018cc5373adcebd9 # constant +69b4b78d5bb6442829fb41be # constant +69b4b8197440b684d16cfa99 # constant +69b4b838f0c88edf01a6322a # constant +69b4b86452783435e5ede910 # constant +69b4b8670ae018e8fdd3a087 # constant +69b4b86c557734daed3d7df7 # constant +69b4b8752634ac70cd74ed24 # constant +69b4b88505a258967f4ab2e1 # constant +69b4b89029b01cb7327324ca # constant +69b4b89bcde059c209035a2b # constant +69b4b8e732580579bce8e79c # constant +69b4b8f5f0d863f7b395423b # constant +69b4b8f8fc41f34c724be519 # constant +69b4b93cc04c33841ec219af # constant +69b4b949cc5ffef12a7697bb # constant +69b4b9782b9d7dec7f4493a1 # constant +69b4b97b47952ada75929047 # constant +69b4b99944663caa553ff3e6 # constant +69b4b9b405b203a8a8cb853f # constant +69b4b9c3d236eadaea23d2d1 # constant +69b4b9c80ea33062b1c276da # constant +69b4b9fbf76dedb2b40db3a1 # constant +69b4ba004f4b006a4b4a3948 # constant +69b4ba081af1a96290db7a19 # constant +69b4ba22b8993611e542422b # constant +69b4ba2bdef7cddf961d4151 # constant +69b4ba2df6333e2fa8a2c5f1 # constant +69b4ba5e3133ed309198da97 # constant +69b4ba9022bc85670f1d8392 # jump +69b4baa0bf0ab5f0fcd34abd # constant +69b4baacbd343104cdb6df8b # constant +69b4bab88a1f883a68a40359 # constant +69b4bac0230639e424fd1cdc # constant +69b4bae3adb32e579d38f553 # constant +69b4bafd6e9a953c6e6c988c # constant +69b4bb329944ebcdcaaf6ab1 # constant +69b4bb3eea15fe5d53dd07fc # constant +69b4bb41610bc1e395ebfe3a # constant +69b4bb68728fdfa816cbf6c2 # constant +69b4bb6ec230d56ec85a761f # jump +69b4bb8471e1c59c5b4c3fe4 # constant +69b4bba6e5809912f66ab263 # constant +69b4bba7569652660a4f4c73 # constant +69b4bbbc4dd9410ce4c3c58e # jump +69b4bbc673ca03704a502a56 # constant +69b4bbe0dae1cdda260d8960 # constant +69b4bbe4cd6846904ec6d01e # constant +69b4bbf0b7c8fd45d1462fd8 # constant +69b4bc03db4d0ad1b2f80302 # constant +69b4bc0719ebdadc8764a6df # constant +69b4bc15e3eb896b22fc2e1c # constant +69b4bc32b8816c7c6f3c9190 # constant +69b4bc4160172888461a82c1 # constant +69b4bc4e8dfc6b6c4f7d2e87 # constant +69b4bc7cc8b94b2da296ad24 # constant +69b4bc938ae239736082947d # constant +69b4bc976bb2557b3d868031 # constant +69b4bcac2e40b83afba4a1cd # constant +69b4bcad9a0f35ae6bd28776 # constant +69b4bcbe2c8d1a50237644c9 # constant +69b4bcc9ca3a06a9f7b0d6cd # constant +69b4bcdc4e323e0ee801a84d # constant +69b4bce58934299fa1119498 # constant +69b4bd11153b74b9c3da90ea # constant +69b4bd3282377b6dbb015687 # constant +69b4bd4a9cc44bc26f6a1bd1 # constant +69b4bd65e8279eda4267027c # constant +69b4bd6ebc4c7ea1d5d29b7c # constant +69b4be228b0d76bdb82882da # constant +69b4be42f9274b0a852547ca # constant +69b4be6aa0c2188e1d8f47cd # constant +69b4be7c5da194a21ed40b80 # constant +69b4be9086e9e7f64031e6c4 # constant +69b4bec3360d4a7922091398 # constant +69b4bf19e8211b84899cef50 # constant +69b4bf2b8415a3e84cfa8c44 # constant +69b4bf6b4f6aaaa9cd1bdbd7 # constant +69b4bfa58b5380aac2061e46 # constant +69b4c06a039109330b23b884 # constant +69b4c07d2917034a1c53ad4c # constant +69b4c29d8a1f883a68a403cb # constant +69b4c2cae2a75fae64879bc2 # constant +69b4c2db1ee47c072a367ed4 # constant +69b4c2de2530ca2e1b04b545 # constant +69b4c2fb73ca03704a502adb # jump +69b4c2fe487fa1638cd28557 # constant +69b4c303f2dba405a5777bc0 # jump +69b4c30c3725ced819eb2447 # constant +69b4c3e7d3d86d024f123905 # constant +69b4c42c9b2ae0fb15ca3810 # constant +69b4c57ad902e000df8b9623 # constant +69b4c58985dcc0aae0d2814a # constant +69b4c595142b1d1af817db6b # jump +69b4c598e2e62141d97ccc0a # constant +69b4c5b232580579bce8e856 # constant +69b4c5f57fe01837a6e54278 # constant +69b4c69cf9274b0a85254836 # constant +69b4c7e802f50e8a066ea50c # constant +69b4c7ece7f9e51220ce5e3d # constant +69b4c8c6f2dba405a5777bf8 # constant +69b4c8d39cc44bc26f6a1c9a # constant +69b4c8d8c289a137df248ec9 # constant +69b4c96c0c02cfba7bd9d049 # constant +69b4ca4dab4b984623699693 # constant +69b4cc18ca778a960281456f # constant +69b4cc34c0682c81b27b2344 # constant +69b4cc9b5afe1d75d15ab846 # constant +69b4cca4f01bb0c8128714e0 # constant +69b4cce4999caf0c82b98c50 # constant +69b4cd5225116d80e5211173 # constant +69b4cdfd398ef280fe770b46 # constant +69b4ceb01e1240fd6bcfebf7 # constant +69b4ceebd40b3f6238573a61 # constant +69b4cf3029e555555c648b1f # jump +69b4cfb0cbd2b424438be0aa # constant +69b4cfb50ce632fcb0bdd9ba # jump +69b4d0404a9710254bb97991 # constant +69b4d0739a082448a776e879 # constant +69b4d0c7d5bde809a4eebf3e # constant +69b4d101ca3a06a9f7b0d7b6 # constant +69b4d15d473f617df3e42f1c # constant +69b4d19fcf65320617885786 # constant +69b4d1a02bcba0117f9f0631 # constant +69b4d26b00941d1749b0fe68 # jump +69b4d2a03ea50ce70f2d411d # constant +69b4d3617b8718172195a2e5 # constant +69b4d3b17e64d72afd3bf0f4 # constant +69b4d3fd6bb2557b3d8681ab # constant +69b4d43b5fadf00f5a3f1d39 # constant +69b4d468acea937039fea6b1 # constant +69b4d472f66794101aa6a706 # constant +69b4d5820c2d4a47811c6b6a # constant +69b4d59acaa6ed5689ec2b34 # constant +69b4d5ebfe501abf8bdb41ae # constant +69b4d6176c8f3e4f2e164a62 # constant +69b4d69055b20ccbe5761f21 # constant +69b4d69f36c56722006292b5 # constant +69b4d6b044c32e2c9984ce96 # constant +69b4d7d19c40481a89bafa49 # constant +69b4d7e86366bd7d9755e142 # constant +69b4d887fc520bb21dcd54ff # constant +69b4d89229b01cb732732652 # constant +69b4d905d6a047d8467b653c # constant +69b4d93930fc0dc05fe97925 # constant +69b4d939e2d3e9207de45b0d # constant +69b4d97be3eb896b22fc2fcb # constant +69b4d98e44de405906e81239 # jump +69b4d995a7cd7cb244b9ed9d # constant +69b4d9992cea447291745529 # constant +69b4d9a0def7cddf961d42ee # constant +69b4d9afa55386e5af5476fd # constant +69b4da0f52e5c59d347ad84f # constant +69b4da14f0a360216793659f # constant +69b4da2ec330d24667261b57 # jump +69b4da354593e8cbb107a4f6 # constant +69b4da495f4789f4f7292aa8 # constant +69b4da517702a3b3a7aa6fff # constant +69b4da9f2bf7176cf0892297 # jump +69b4dabec7834bedf197254b # constant +69b4dabf19ebdadc8764a92e # constant +69b4dadb2994f223142b83b4 # constant +69b4daed0b666ac17c5266a0 # constant +69b4db82cee7dbe8803e3d5f # constant +69b4dbb1f93256bcb41f201e # constant +69b4dbb9bc1711c9daf2e280 # constant +69b4dbd91bd7394361b51da9 # constant +69b4dbff4f13e26864ae1480 # constant +69b4dd709b3c11bddef4b1da # constant +69b4dd959e5df516fced68dc # constant +69b4de167bc9e8b3eac05d15 # constant +69b4de83fd2b3f7e99af4eb1 # constant +69b4e0174b8693d11c3b9a7d # constant +69b4e04b6cb5a1c1a2eb4f74 # constant +69b4e10800ce0b618c19cd41 # constant +69b4e1247702a3b3a7aa7044 # constant +69b4e1b12af7475cd15b2946 # constant +69b4e207d781f5e3eea1e62c # constant +69b4e2379a0f35ae6bd28950 # constant +69b4e27953e0e51b6dd34ad7 # constant +69b4e358758deefb5ec011f2 # constant +69b4e3dd60b4c94b607be392 # constant +69b4e3eeca3a06a9f7b0d857 # constant +69b4e47c8f37b66ca9f7b2d8 # constant +69b4e4973133ed309198dd4e # constant +69b4e557a2e0f23556e6a860 # constant +69b4e696acea937039fea768 # constant +69b4e808db402eaf0fc5ab62 # constant +69b4e85aab39f27198e3195b # jump +69b4e8db827f8909e8826af6 # constant +69b4e8f68c9ba98495e6e2e6 # constant +69b4e8f9f76dedb2b40db687 # constant +69b4e94341a4401e4c3f5421 # constant +69b4ea2015d40a449a6d8129 # constant +69b4eb26230db4d63a859cba # constant +69b4ec24ad35495e55532ef8 # constant +69b4ec29610bc1e395ec0179 # constant +69b4ec2be9f454e6cf0d8ee2 # jump +69b4ec441af1a96290db7bab # constant +69b4ed964ff0a19d651b18b8 # constant +69b4edb86c8f3e4f2e164b60 # constant +69b4ee71bb2b154f8ddd6b39 # constant +69b4eed847e0e2e54a0ed63d # constant +69b4eedd698982b08fa37b54 # constant +69b4ef55db5484827307bc41 # constant +69b4efa5f1c3826f1de62361 # constant +69b4efaf570289759dbd7cc2 # constant +69b4f01ec314b4adc2d81827 # constant +69b4f03f25fd3eda026f237c # constant +69b4f051515d4294ccc793f1 # constant +69b4f1d38a313c3e6edab466 # constant +69b4f1f82c83b2c34dfe57d1 # constant +69b4f2a13e423be2d663a84d # constant +69b4f2ce60172888461a83f7 # constant +69b4f31d95b405a6f7105dd8 # constant +69b4f346a506f1c6355d092f # constant +69b4f37c388ccecd105c5fe1 # constant +69b4f4188f419c69e88fb220 # constant +69b4f438a3c0ce02c5a55c2b # constant +69b4f472590494387f207f90 # constant +69b4f4885d5efd59f80d4a17 # constant +69b4f4bf728fdfa816cbf8b5 # jump +69b4f561c69c46f0a2021592 # constant +69b4f56bab4b9846236997a8 # constant +69b4f5969891381d917142d3 # constant +69b4f5b93e423be2d663a85e # constant +69b4f62df01bb0c8128716af # constant +69b4f6b15fa639c4e24e76d8 # constant +69b4f6b4ac8bbc3ef42bbb7d # constant +69b4f6d67d68d7e49aa56623 # constant +69b4f759fc41f34c724be975 # constant +69b4f768b2b6f95180884e29 # constant +69b4f7fbe5be71666d49464c # constant +69b4f895817936a8e39f4346 # constant +69b4fa0fcf93b545e608900d # constant +69b4fa10bdfca6e9a5a1e5e8 # constant +69b4fa8553e765a22cddd639 # constant +69b4fad99cd671b2e770b251 # constant +69b4fae21b02c068cd2ad204 # constant +69b4faee31cf18f433028d38 # constant +69b4fb02f334e0f5a453c030 # constant +69b4fb052843c2f113eb1409 # constant +69b4fc258579b939089f49c5 # constant +69b4fc49bdfca6e9a5a1e5f1 # constant +69b4fc619b6d7e4a378e52dd # constant +69b4fc82bdb3ad2519469084 # constant +69b4fd27b1d07bc9b4781a7c # constant +69b4fdf990295b261f21c1e7 # constant +69b4fe2532643e552e8855b0 # constant +69b4ff458a4c8aae9bc769aa # constant +69b4ff90e51948e7b02d9fdc # constant +69b50044671bfb1f5254e50e # constant +69b500ff24715e745fc75271 # constant +69b50117a83a7b34affb4d1f # constant +69b501caf916ea27abf01a59 # constant +69b502cc85e36f35da0bb203 # constant +69b502edaa4f5d834935847c # constant +69b5032114a4d59b08be0396 # constant +69b50351c7b864206a9a73f6 # constant +69b504205fa639c4e24e772f # constant +69b50445aa4f5d8349358486 # constant +69b5046285e36f35da0bb20b # constant +69b5049f3f21e8be2d3943f9 # constant +69b504f361c85d2534e8f88a # constant +69b50518460cd47c547c8a65 # constant +69b505293d0a537fcffa394b # constant +69b5053515d40a449a6d81a9 # constant +69b5055c7c84b98b9831e689 # constant +69b505f060b4c94b607be46f # constant +69b505f3edba17448480bc74 # constant +69b50617db5484827307bcaa # constant +69b5063c47e0e2e54a0ed6fe # constant +69b50673bbfb18a9cc6cae09 # constant +69b5070bd20be0630bf038da # constant +69b5071f0c01a69644a221dc # constant +69b50739df9603316e69d63d # constant +69b5075a8ae23973608297cf # constant +69b50799f435d2bd800ba868 # constant +69b50800fe3e51b5a878e5e0 # constant +69b50803d2ab5ebd52f6e4d1 # constant +69b5087aab91168cc457b93d # constant +69b508f5c5fa25330433f70d # constant +69b50a18515d4294ccc7945e # constant +69b50a3aa9f25eecea6b57cd # constant +69b50a89c7e02bc8d412af3d # constant +69b50ab866687323993e8d8f # constant +69b50ac708a0268c1c38beec # constant +69b50ba3335d8a8a3c8dfd9e # constant +69b50bae16f9cc224f6655c2 # constant +69b50bce70be8493faa0a6c4 # constant +69b50be2124502aa363251f2 # constant +69b50cb29fb2a70d14310292 # constant +69b50cc94a206179cbd004bb # constant +69b50d3dea50b19996d7e87e # constant +69b50d454a206179cbd004bf # constant +69b50e1a91d6491aa901a5c2 # constant +69b50e2fe8b23514640557e1 # constant +69b50e35da06e853603f5c10 # constant +69b50e9cde62960a07de88c5 # constant +69b50f235d5efd59f80d4ab0 # constant +69b50f273065bf627cde2d4c # constant +69b50f4bab11feb553052dbb # constant +69b510139a3e364072f003c1 # constant +69b5105abb2330fd69279f44 # constant +69b5106d394e490d4592dd82 # constant +69b51078ea0478dac4767193 # constant +69b51087cf7413b53d1929cf # constant +69b510f9cd67f708f6e64df0 # constant +69b5113e6b4f61a1748f5026 # constant +69b5123ffc1abe8f43fd56d8 # constant +69b512d7a865ba37a7ea986f # constant +69b512fde025fc2bec3b06e1 # constant +69b51333c4c5496f9afd496c # constant +69b513c1b416a351f4c42736 # constant +69b51556ea0478dac47671a8 # constant +69b516535c9a724fbfbc1eec # constant +69b5166ae38bde62855cf92d # constant +69b5169e8ca5423dab74c835 # constant +69b516a53ac25998bd7be336 # constant +69b516ff53168422516b2674 # constant +69b5173085bdd8b99b29802b # constant +69b51795e5be71666d4946d9 # constant +69b5179eea50b19996d7e8f0 # constant +69b517f4e5be71666d4946db # constant +69b518391b5d00f6e2fca70b # constant +69b51860e9de50b7d0f6f218 # constant +69b518dc6641c34ca8db19a4 # constant +69b5196362ac9501398f3dd9 # constant +69b51963d197af22134e70f8 # constant +69b51985d5bde809a4eec16c # constant +69b519c8eb7f329ec6720f99 # constant +69b519f872c7691edd73b6da # constant +69b51a4782377b6dbb015a84 # constant +69b51aa691d6491aa901a604 # constant +69b51ac7a25cc84d5c2076a2 # constant +69b51add5d5efd59f80d4ade # constant +69b51b863326dd7f8185cb14 # constant +69b51b8e98e0f543f38a008c # constant +69b51bad52b439fcc613b688 # constant +69b51c17e248420994a830ad # constant +69b51c2f16f9cc224f665653 # jump +69b51c37b30844b29a832e01 # constant +69b51c62f1d531c517f55b41 # constant +69b51ca2d93db6d47a9fa6e7 # constant +69b51caa17219d378d74a988 # constant +69b51cc0980f80025fc5e20b # constant +69b51cecf629796beb9843e1 # constant +69b51d00aa4f5d834935856c # constant +69b51d21ee737c3c16391ec4 # constant +69b51d29ab91168cc457b9fc # constant +69b51d4e65776c1b020fdcc7 # constant +69b51d5aea50b19996d7e92e # constant +69b51d7adf63a3e00aac2586 # constant +69b51e1ee8b235146405583d # constant +69b51e9d5de136447199ee60 # jump +69b51fece7f9e51220ce60c8 # constant +69b5204416ad6cbe9700c6df # constant +69b520898c9ba98495e6e43a # constant +69b520a4b7bba0cd0bb65566 # constant +69b520d9a13067e31cf8b3e7 # constant +69b52107e025fc2bec3b0781 # constant +69b521162994f223142b8676 # constant +69b5217cffbcef2361fd446a # constant +69b5218b3b674a7b354e6c79 # constant +69b521be78673e5de422879f # constant +69b521eaf52b66f1a3088baa # constant +69b5220d130037c5b4ba5f57 # constant +69b5225e6b4827eb68cf11f2 # constant +69b5229dd673f8a5bdac0d46 # constant +69b522d4edba17448480bdd6 # constant +69b522e0803d8a9edd67ed66 # constant +69b522f7cf5923d78401374f # constant +69b52300d309464df71159bf # constant +69b52303aea2c4fce0c3ec7b # constant +69b523208f419c69e88fb3e3 # constant +69b5232ccf653206178859a7 # constant +69b523af6b4827eb68cf1201 # constant +69b523cae2e62141d97ccfc7 # constant +69b523d3d197af22134e7133 # constant +69b524b500d0e98b2aab0383 # constant +69b524b7c375af84c0576f41 # constant +69b524cc6641c34ca8db1a30 # constant +69b5257da1b0f0badb9967eb # constant +69b5258608a0268c1c38bf9b # constant +69b525b926871cfad9ed17c7 # constant +69b525fd5d5efd59f80d4b2e # constant +69b526e321e3db52d1460e56 # constant +69b526e48ae2397360829919 # constant +69b52738a9f25eecea6b5860 # constant +69b5276d24715e745fc753a5 # constant +69b527a9ab37121e9e93c3e8 # constant +69b5281f335d8a8a3c8dfed5 # constant +69b5287b1b5d00f6e2fca7b3 # constant +69b5290bab575f0d17062a96 # constant +69b529589f66a2746698fbe1 # constant +69b52a46870fd7d45c1f340f # constant +69b52a59135e07fa47414a16 # constant +69b52a6fc420444b72115c10 # constant +69b52ab5a97937a3e2fb3bbf # constant +69b52b03d9bd4ed2c8cda811 # constant +69b52b04bb2330fd69279ff6 # constant +69b52b62d158ec4a43dcfa62 # constant +69b52b6df637eb3c0a12e39e # constant +69b52bdd16f9cc224f6656fb # constant +69b52c05806149738b4e5aed # constant +69b52c55473f617df3e430dd # constant +69b52c594009cbbd3f6c26ad # constant +69b52cabe7f9e51220ce6114 # constant +69b52cd16c13d04069dcce5d # constant +69b52d38f4110359f68f7b1f # constant +69b52d6c5c9a724fbfbc1f68 # constant +69b52de57ff8637604061e4c # constant +69b52e018c9ba98495e6e496 # constant +69b52e0a5656a2300b64efaf # constant +69b52e208c1c37de293e7763 # constant +69b52e880dd5025bc8b2a713 # constant +69b52e8db70fb5b833849d85 # constant +69b52eb0bd3343ef280e838f # constant +69b52ee1a1b0f0badb996814 # constant +69b52f155656a2300b64efb5 # constant +69b52f2745a65ae24232b2ee # constant +69b52f470403c14bba1f60e1 # constant +69b52f48d158ec4a43dcfa7a # constant +69b52f92025d49131a129125 # constant +69b52fb2cee7dbe8803e3f7c # constant +69b52fbbeec4d2ca74dedb83 # constant +69b52ff9ad9ebb37e8776eba # constant +69b530608c61b5e3f72f9c0c # constant +69b5307463d17484ef03a2e5 # constant +69b531021ca41b2857617211 # constant +69b5315a8061ae794120aea6 # constant +69b531f414a4d59b08be05a7 # constant +69b53205ab11feb553052f40 # constant +69b53243060a6199caac13d6 # constant +69b5324f60b4c94b607be66e # constant +69b5327db70fb5b833849d9c # constant +69b532921b02c068cd2ad337 # constant +69b5329926871cfad9ed1864 # constant +69b532a9c8f3797a5f1f6bf8 # constant +69b532dad309464df7115a8e # constant +69b533026b5ffecb54b2fa96 # constant +69b5334d7ff8637604061e8b # constant +69b5336c9e5d189b4aae8437 # jump +69b533c07c962028610d514d # constant +69b533d417e3a9ccf08b298d # constant +69b534169a3e364072f004c2 # constant +69b53441ab37121e9e93c47c # constant +69b5346b28176f06fd328512 # constant +69b534ca8c1c37de293e778b # constant +69b534ecb7823508aa10c0bc # constant +69b534f215f14cb9c4fb64ad # constant +69b53501f70d1949244a6348 # constant +69b53509e9f454e6cf0d905f # constant +69b53538bd3343ef280e83d1 # constant +69b5355fb70fb5b833849daf # constant +69b535656d90f36faea35151 # constant +69b5358b9275ad7c73fdb973 # constant +69b535ca85847273dc74d6a0 # constant +69b53616f70d1949244a6350 # constant +69b536298faeb1c44a3ebf83 # constant +69b53630aea2c4fce0c3ecee # constant +69b53637ad9ebb37e8776f08 # constant +69b536a98505130f8b0e72c2 # constant +69b5374648093a6af8222133 # constant +69b537608505130f8b0e72ca # constant +69b537bfa424e9e7431c2d3d # constant +69b537c73e423be2d663a9e4 # constant +69b537ed557b5c79a6d7a93b # constant +69b5386246dbfcf8d04858ee # constant +69b538b1a2bb2f3fc33fd5d6 # constant +69b538f41666a9ba0092e873 # constant +69b5393c9cd671b2e770b3ca # constant +69b53975eec4d2ca74dedbc4 # constant +69b539a151dbfa3ebcfcd605 # constant +69b539ce3d0a537fcffa3bb1 # constant +69b539df70be8493faa0a8ee # constant +69b539f61a315c79b1b51418 # constant +69b53a09124502aa36325304 # constant +69b53a0b153b74b9c3da93dd # constant +69b53a70cecf38bf0cf6dfc1 # constant +69b53abd7c962028610d517f # constant +69b53aff671bfb1f5254e6b8 # constant +69b53b02096a5342c3dc799a # constant +69b53b36e2d927d8bec080aa # constant +69b53b5b17b1b3bbefd3ffb4 # constant +69b53b7aae7c8a432518f90e # constant +69b53d48b70fb5b833849df3 # constant +69b53df0b70fb5b833849df9 # constant +69b53ec1d5bde809a4eec2b6 # constant +69b53f38ea50b19996d7eace # constant +69b53f6976b3ece51f04532e # constant +69b53f6f3a7542779b575391 # constant +69b53faca1b0f0badb996891 # constant +69b53fd453168422516b27a2 # constant +69b53fdff70d1949244a639c # constant +69b5400316f9cc224f665838 # constant +69b54012d0103fb8e4fb4eff # constant +69b54012e248420994a83197 # constant +69b540f19a082448a776eadf # constant +69b54106a25cc84d5c20778f # constant +69b54185c5efd21514309224 # constant +69b5419d06857a01e6a73c6b # constant +69b5419f08a0268c1c38c095 # constant +69b541e5230639e424fd1fd1 # constant +69b542225bcb582c92e29453 # constant +69b54245ab34715a79255963 # constant +69b5425689deb7a84c2d5943 # constant +69b5428bf3077050a5079288 # constant +69b5431bf4110359f68f7bc8 # constant +69b543e75a2141ade9a3b0f3 # constant +69b544354e6b0afb54824163 # constant +69b5445a8781fdde8bbaaffb # constant +69b5445b7dbf454e859927f6 # constant +69b544782cea44729174576f # constant +69b544eb44509a97d31274ce # constant +69b54502a424e9e7431c2d99 # constant +69b545064009cbbd3f6c2781 # constant +69b546345d5efd59f80d4c34 # constant +69b5464ed422ae5db7e71d39 # constant +69b54657164aaac0a2a6f18d # constant +69b54666e025fc2bec3b097a # constant +69b546a4ffbcef2361fd4555 # constant +69b547d6a865ba37a7ea9b22 # constant +69b547fe5b4047fbbfc4d7f6 # constant +69b548f7124502aa36325381 # constant +69b54902fd7b6f90a9fb576c # constant +69b549087ae6fe5f6c5a3f0f # constant +69b549953a7542779b5753e6 # constant +69b54af432cefa428f103d3f # constant +69b54b46a119235ef859a82a # constant +69b54cdba7585840264c6a19 # constant +69b54d14e38bde62855cfa8a # constant +69b54d36d0103fb8e4fb4f6a # constant +69b54da5f3077050a50792d6 # constant +69b54dabc19d1c22701b78e3 # constant +69b54db88579b939089f4bc4 # constant +69b54e81f52b66f1a3088ce6 # constant +69b54e9ce7f9e51220ce6240 # constant +69b54ea06a8d780c96378fab # constant +69b54f30d673f8a5bdac0e6c # jump +69b54f36da06e853603f5dad # constant +69b54f5e1b02c068cd2ad468 # constant +69b54fa5c4d1de4eaf0d2f6c # constant +69b55009ced344cb0df766e2 # constant +69b5500afd7b6f90a9fb57a2 # constant +69b5505f9b2ae0fb15ca3ce1 # constant +69b55106bd3343ef280e850e # constant +69b551ccf4110359f68f7c3f # jump +69b551eacfe53ea1af65ff2b # constant +69b55292bb4aaf9805d8bfe9 # constant +69b5529c16f9cc224f66597b # constant +69b552cb37cb8c569b487bb3 # constant +69b552ed742c9a8a3dfa9b4c # constant +69b5532bc0248818ec109e8d # constant +69b55387f268396b4bc347f3 # constant +69b5539ecf93b545e608920d # constant +69b553fcc4c5496f9afd4aef # constant +69b554008c1c37de293e787b # constant +69b5542c95b405a6f71061e3 # constant +69b5544d15f14cb9c4fb65be # constant +69b554bcb1d07bc9b4781cb4 # constant +69b5551c752d20a393bd09c6 # constant +69b55541d5bde809a4eec3be # constant +69b55579d197af22134e7297 # constant +69b5568d060a6199caac14d4 # constant +69b556f3c19d1c22701b792b # constant +69b5578ed0cc9cadf7a25857 # constant +69b55799230db4d63a859f1a # constant +69b557a46325cbcd4232b161 # constant +69b558fee2a75fae6487a0be # constant +69b559361b02c068cd2ad4d0 # constant +69b559438a4c8aae9bc76cbe # constant +69b5596013cfede97b789d2c # constant +69b559754e73acdd60a50fc8 # constant +69b55a1cbb3647e88d660444 # constant +69b55a836325cbcd4232b17d # constant +69b55a90bd3343ef280e856e # constant +69b55aee26871cfad9ed1a14 # constant +69b55af85c9a724fbfbc20af # constant +69b55ca3d422ae5db7e71de3 # constant +69b55d5c1a32c53ddbd16ed8 # constant +69b55df6752d20a393bd0a0a # constant +69b55e44817936a8e39f45df # constant +69b55eb27dbf454e859928ff # constant +69b55f3b3fbc85f1fd117591 # constant +69b55f8fb416a351f4c4292c # constant +69b55f9088f4b4a6e37d9e13 # constant +69b55faea2bb2f3fc33fd7a5 # constant +69b560379605a9793b6b95b6 # constant +69b560a26a8d780c9637903e # constant +69b560d343041b5f0fab06fe # constant +69b561736325cbcd4232b1b2 # constant +69b561d9151050b0854f92a9 # constant +69b561f753168422516b2950 # constant +69b562575bbfec0f5b7167b2 # constant +69b56282124502aa3632544e # constant +69b5628616f9cc224f665a88 # constant +69b5629ce9f454e6cf0d921d # constant +69b562bbced344cb0df767f5 # constant +69b562ccda06e853603f5e48 # constant +69b562d9ea50b19996d7ee1a # constant +69b562e8b7c7ad88170d4148 # constant +69b5630d46452e48802c2970 # constant +69b56320d9bd4ed2c8cdab06 # constant +69b56329ea0478dac4767464 # constant +69b56343c7e02bc8d412b17b # constant +69b563798ca5423dab74cae9 # constant +69b563be9aeb80834f1fad28 # constant +69b563f03acb25c0d76dbbc2 # constant +69b564047f0b851033657900 # constant +69b5649a21cd70aecfce4def # constant +69b565aebd3343ef280e85d5 # constant +69b566b97d68d7e49aa568fa # constant +69b566dc423da797309e33f9 # constant +69b5672e90295b261f21c4bd # constant +69b56761a7585840264c6afc # constant +69b56762df9603316e69da4d # constant +69b567c4c155de971bf1bf34 # constant +69b567e9bb4aaf9805d8c0a9 # constant +69b5686083db69b280d29c7d # constant +69b5686c5a2141ade9a3b221 # jump +69b568ba41dea3220fc9bbbe # constant +69b568e7e267367ac3f2564c # constant +69b56a523acb25c0d76dbc11 # constant +69b56a66ee2ba8e55c296cf1 # constant +69b56a6e32643e552e8858e8 # constant +69b56a9f6b4f61a1748f52e2 # constant +69b56ad89e3c339e3d3e8ed2 # constant +69b56bbb0c3122b34d6e8d03 # constant +69b56be70403c14bba1f6450 # constant +69b56c135bcb582c92e295b0 # constant +69b56c15075e5c4012d20845 # constant +69b56cb35656a2300b64f1a1 # constant +69b56cd617219d378d74ac52 # constant +69b56d0824715e745fc75696 # constant +69b56d248c1c37de293e796f # constant +69b56d7bd673f8a5bdac0f69 # constant +69b56d8d758deefb5ec0155d # constant +69b56dbb7e46044c0318d0f5 # constant +69b56e31827f8909e8827161 # constant +69b56e518c9ba98495e6e6fa # constant +69b56e952843c2f113eb179a # constant +69b56ed717b1b3bbefd4013d # constant +69b56f16c04c33841ec21ee9 # constant +69b570a76abfa602a5495724 # constant +69b5710a9a3e364072f00839 # constant +69b571266144f027ad942626 # constant +69b57132df63a3e00aac2b6f # constant +69b5715b3fbc85f1fd117696 # constant +69b5716d31d39e2036432894 # constant +69b57180153b74b9c3da9591 # constant +69b571adc8e2220dbb9e4868 # constant +69b57275a2bb2f3fc33fd892 # constant +69b572ee287b1e524e4af2f8 # constant +69b5732f132acce764a592a8 # constant +69b573673242599fceea3076 # constant +69b5737c9c7daebcace62308 # constant +69b573ecf85fc301e7132041 # constant +69b5742a88f4b4a6e37d9f23 # constant +69b5742f4d610f6c3568ee02 # constant +69b574ca9c27f8b9d1b5a69c # constant +69b5757506857a01e6a73eaf # constant +69b5764588f4b4a6e37d9f54 # constant +69b5765615f14cb9c4fb6770 # constant +69b5766185847273dc74dad4 # constant +69b576ab460cf9a11ded4a05 # constant +69b576d5ab575f0d17062c8c # constant +69b57703b416a351f4c429df # jump +69b5773854db284caf2f805f # constant +69b57767151050b0854f9361 # constant +69b577f47c962028610d5338 # constant +69b577f6d90644203dac6d9f # constant +69b578216bceb806c655f5eb # constant +69b578bbc155de971bf1c001 # constant +69b578cf165ee7193a2bcee6 # constant +69b57932c773307d2dd79413 # constant +69b5798a8cd7957ebe4794a9 # constant +69b579deffea0a68ed9d072a # constant +69b579e5de5165de31d6902d # constant +69b57ca826871cfad9ed1ad9 # constant +69b57ce3c8e2220dbb9e48b8 # constant +69b57d434e58e40d05210ba2 # constant +69b57d4a3dd17eb53d7de101 # constant +69b57d8244509a97d31278c9 # constant +69b57da51a315c79b1b51751 # constant +69b57db233d4fbca9435a095 # constant +69b57e38c7e771d4de7ec384 # constant +69b57e550dd5025bc8b2aa56 # constant +69b57ed48c9ba98495e6e7b3 # constant +69b57f834009cbbd3f6c2975 # constant +69b57fa7bde18229db6d05a2 # constant +69b57fafcf93b545e608935a # constant +69b57fc48ca5423dab74cc2a # constant +69b5801976b3ece51f0454fa # constant +69b58029c5efd215143093c2 # constant +69b5807c526f02661a0c1d5d # constant +69b5807ee2a75fae6487a277 # constant +69b581771a32c53ddbd16fd1 # constant +69b5823a096a21940282e7bb # constant +69b5825d5d5efd59f80d4e42 # constant +69b583bf37cb8c569b487d1c # constant +69b584406142729908de7d63 # constant +69b584932cea4472917459bd # constant +69b5855ea2bb2f3fc33fd946 # constant +69b58624185e4fd6a7484217 # constant +69b5862b5a10c4a2d8fa05a8 # constant +69b586a385e36f35da0bb7ea # constant +69b587159d6df5f247040033 # constant +69b5874148a0f860e59885a2 # constant +69b587885da194a21ed411aa # constant +69b587a7a7fff82d19902465 # constant +69b587f29ee0b627f1fcfd71 # constant +69b588066a8d780c96379158 # constant +69b588191dea4182847f781c # constant +69b5884025ade459f7bbbb0e # constant +69b588c0c375af84c05771e9 # constant +69b588db6f92ac8c2497d1ef # constant +69b589336b4827eb68cf18c2 # constant +69b5894dacb24cf31692e9ca # constant +69b589b65de136447199f141 # constant +69b589c145a65ae24232b770 # constant +69b589d57f0b8510336579f5 # constant +69b589de0a007568736e92d9 # constant +69b58a16098f0b40252c8aa5 # constant +69b58a219d6df5f247040047 # constant +69b58a222814c4be9c919dd4 # constant +69b58a6c250b26ab0bf1c7c5 # constant +69b58ab251dbfa3ebcfcd943 # constant +69b58b0585847273dc74dbe3 # constant +69b58b4ac7e4a63bf5d7305b # constant +69b58b8e1c8cd0adc71287fd # constant +69b58b938c1c37de293e7a72 # jump +69b58b9eb8745dc1f264dc11 # constant +69b58c53ac89a92803af52cc # constant +69b58c67f3077050a507945a # constant +69b58cde287b1e524e4af3a7 # constant +69b58d158c3e4c65ea33981f # constant +69b58d6f61ba083d8cc3b026 # constant +69b58d7f5907c64c2ffa3737 # constant +69b58dac1a315c79b1b517bf # constant +69b58e29afd18a8a6a3efc69 # constant +69b58e2d9aeb80834f1fafd4 # constant +69b58eee2bf7176cf08927f6 # constant +69b58f2d060a6199caac164a # constant +69b58fc6e248420994a83398 # constant +69b59011f30a828e1ab2f8b6 # constant +69b590573e7a3f6666bbdaa7 # constant +69b5907f4f85a34e0a4ad7dd # constant +69b5909aacb24cf31692ea07 # constant +69b590d552fb78d78806ed1e # constant +69b59106422d3b866caf843a # constant +69b5c30c3ec321e7831ca9f0 # constant +69b5cc0c32cd5d0d72d53ec7 # constant +69b67ec04bd4a214b726bdb0 # constant +69b680cbb3c6644588cd1e0c # constant +69b682e6bf9d2da1280256db # constant +69b682ec256c89916b55db90 # constant +69b682fc6f44e1498bdf94b7 # constant +69b68415cb7101c1e4d945d1 # constant +69b686386dc1fae16521f154 # constant +69b68741d7831248001c5a74 # constant +69b687676dc1fae16521f16a # constant +69b687b9d88f1dc9ad09524a # constant +69b687f36394554b8c0f5bf0 # constant +69b688c8ad89c76ef7a57e35 # constant +69b689b8a044ae4efb6a92d5 # constant +69b68afe346a130f980d34e9 # constant +69b68b0da4ecc6ce792cc99f # constant +69b68b0e1f451cbdb2c2287a # constant +69b68b831cd819c3c4e43a10 # constant +69b68b99aa96b63b2bddb9e9 # constant +69b68caf76299d1a1459270f # constant +69b68d20fe41e6a1502402d7 # constant +69b68d6533702a0c61de0b74 # constant +69b68ddab3c6644588cd1ef4 # constant +69b68f3c686ec2d47340ae9a # constant +69b68ff285b0e68bd9a831b2 # constant +69b6900f62f6b9a70b374eb9 # constant +69b69039f8528d928cf6eda1 # constant +69b690c4863376b642ddc3d1 # constant +69b690d16fd488b5dec5a4a2 # constant +69b690ee28ce90cf8b274827 # constant +69b691750fa6b58f60de7e9a # constant +69b691e933f0a195b7e7283c # constant +69b6921bbeaf9aa0229cd46f # constant +69b69239f617163cea7d391a # constant +69b6927261533cd7ffd1c9a1 # constant +69b692cd3455db50a4d7dae8 # constant +69b6933a2434e5fdd31b08dd # constant +69b693a508a239398230892e # constant +69b693ddc10ecf2093fda98b # constant +69b693ebb22a75796cb10eee # constant +69b694d6eae2ffd4a87212ae # constant +69b694eaace85918a0bf18a3 # constant +69b696d47dd881e790ced65e # constant +69b696fdfe41e6a150240379 # constant +69b69a20b926e6af70bbc643 # constant +69b69acf792dc732a00dd708 # constant +69b69ade179ff685576ec894 # constant +69b69bed4f8988c123b3377f # constant +69b69c0bb52dd3401b14777e # constant +69b69c815f1d2dd786b35f38 # constant +69b69de1f69d549b769dda62 # constant +69b69ffb12e974fed9084f21 # constant +69b6a021c90ad0858c925459 # constant +69b6b6a5e374dd52788b1f06 # constant +69b6bbba1b42c6b7c1b800ae # constant +69b6bd69c7d146edfbde10df # constant +69b6bdfe3f5db8bc2edec76b # constant +69b6c82951caaccc5aa8cb8c # constant +69b6c8460fdc71adf8a0058f # jump +69b6cfe3b84861ef35e3975c # constant +69b6d0dbff263e4b37caf58d # constant +69b6d54816dd4d92eabd0521 # jump +69b8421ccdbc7aede694381a # constant +69b843291b24140def911b28 # constant +69b84448e5eb7003529419b2 # constant +69b845131f8050f3429449d7 # constant +69b8451aeb6da1bebb644dda # constant +69b8491de1a23f007e480fce # constant +69b84ad1ac8b66a014689f36 # jump +69b84e9d098cd47aa604ca11 # constant +69b84f1de1a23f007e48100c # constant +69b8534cdf08149cccad9dde # constant +69b855a3c5ab7070deca488b # constant +69b855b085467b6bbdf2874f # constant +69b8569604dc9d7077aa25fa # constant +69b859377e21dd23ed5a5c1b # constant +69b85994b5fe113100c26712 # constant +69b859f4c33b3ff3d5375992 # constant +69b85a3dc5ab7070deca48c5 # constant +69b85df660e13456873a46ca # constant +69b85fc6d9fab0624b95183c # constant +69b8688bc64b14a0c104d71a # constant +69b869662c2daf3cb076c974 # constant +69b869e6b370f3d1f8c93305 # constant +69b86b6ae14e49e6e3f2eb2f # constant +69b86c72ef2b840df479ed3c # constant +69b86eec850bd49ca23e582b # constant +69b870cdc5a1357004e8f8fb # constant +69b8710f0a20ea4e2d821e50 # constant +69b8727fd498a456d3063fa7 # jump +69b8731dd363bdf0ddeabf36 # constant +69b874a385467b6bbdf2893c # constant +69b8784a85467b6bbdf28964 # constant +69b87a6a0f529cac3ae53fbc # constant +69b8804319aecd94aedc36c0 # constant +69b8810171158352535117ee # constant +69b882a9a16ad76e6d97f1f3 # constant +69b88368c43b2ff00fc7944b # constant +69b884512a89e0bab91ac9fc # constant +69b8858da228f21b38777efe # constant +69b885a5e1a23f007e4810bd # constant +69b887b1e97bec8c72c5df62 # constant +69b889f6d92d792f69edd841 # constant +69b88b6843b675efc7200fda # constant +69b88bd122b409ba79d211d7 # constant +69b88c5a86b9f227f633a562 # constant +69b88c7b50e8ec10497acb46 # constant +69b88caccdbc7aede6943c89 # constant +69b88ddb2565bc2e9320a882 # constant +69b88e0c17abc716e5e01308 # constant +69b88eb4789976e408acf7ee # constant +69b8956afb9cb514d86f28c8 # jump +69b89807186d8785b78d97b4 # constant +69b89edbc5133d089fd7c46a # jump +69b89f9e8c04b58b6d952ae9 # constant +69b8a1221a97e051d5434d17 # constant +69b8a72cf29b404cd4040750 # constant +69b8a7871a97e051d5434d31 # constant +69b8a972708396dd182c1d67 # constant +69b8ac92d363bdf0ddeac08c # constant +69b8b623c8858e187272b705 # constant +69b8b7b068c639144c96f60b # constant +69b8b805a52e1a2126f46019 # jump +69b8b84a63d67d4ab9032c6e # constant +69b8b84cf62f97e1ce06ab88 # constant +69b8b986f787e78f3caf7cd0 # constant +69b8b9e2a3ef2775e0e98954 # constant +69b8bb33f62f97e1ce06ab9a # constant +69b8bb822555407a34bad0af # constant +69b8bc26ffcdacfe68dc34bb # constant +69b8bced3907853432078864 # constant +69b8bd8ee25d009d215f1415 # constant +69b8be47c147ee9b771dc8b0 # constant +69b8c7a19fe903ae72105944 # jump +69b8c893adb6e6a77d23fb79 # constant +69b8cd7fe03168d093e51304 # constant +69b8cf32d184f934b8b26fa6 # jump +69b8d273b1ffd7875e7cce8d # constant +69b8d5e6617de250ae32d662 # jump +69b8da8d390785343207895b # jump +69b8dfca6e470e7c633d732a # constant +69b8e2795f6aeb7be1dd3c6d # constant +69b8e2ae2565bc2e9320aa6b # constant +69b8e35de094c1fa8268f7db # constant +69b8e6040f92f74da61ce5d0 # constant +69b8e6206a6fcc387e6fc7f0 # constant +69b8e626a8f416005104fd3c # constant +69b8e66e68c639144c96f784 # constant +69b8e88ab443baf0e08286e0 # constant +69b8e9225088ac62355d6a02 # constant +69b8eb5bc9e8f1de0533e153 # jump +69b8edeadaa78337a3cf8330 # jump +69b8f8c7daa78337a3cf839c # constant +69b8f91fbb8856ba14a76f81 # constant +69b8f95ac9e8f1de0533e1dc # constant +69b8f9ca4b754305a80bf48a # constant +69b8fb0c0b0ebaecc7d82599 # constant +69b9022e55294f107bfbc7c7 # jump +69b9027e2f5695d1d2f53348 # constant +69b9027ef25718a726e343c0 # constant +69b903521740c71b0874fb10 # constant +69b904cf004576ba1550a86c # jump +69b904f7002659f12899f641 # constant +69b907f4b1ffd7875e7cd087 # constant +69b90f54e749b83b1a332f19 # constant +69b910a55088ac62355d6b77 # constant +69b911fe4a18f0c8d9bf7a31 # constant +69b915d7345333ff05a8cff0 # constant +69b917f217abc716e5e01694 # constant +69b91dbc87a8695556dc1177 # constant +69b91efc9bcf11a7dc66d64d # jump +69b9239f87a8695556dc11a9 # constant +69b92511345333ff05a8d066 # constant +69b926adee7ae1002b6a72ca # constant +69b92749e3e5accb88abd1c8 # jump +69b927fce094c1fa8268fa49 # constant +69b928101743775241d36094 # constant +69b92819a3ef2775e0e98ca4 # constant +69b92837a0469793c24bf05b # jump +69b928b17ee0ff33631f0910 # constant +69b92b04c147ee9b771dcc2e # constant +69b92e266be7e1c4139f8173 # jump +69b932266be7e1c4139f8195 # constant +69b93386e4dd0d3cf5a00134 # constant +69b93466adb6e6a77d23ff79 # constant +69b9356b4458d25a26dc6a2f # constant +69b935bfa52e1a2126f46472 # constant +69b9386263e615dae6ebcb97 # constant +69b938a1de59539246182b71 # constant +69b93ab155294f107bfbc968 # constant +69b93b8e0b1ec7e2efa01286 # constant +69b93c50bd0a1226e8720bce # constant +69b93c8b6e470e7c633d7665 # constant +69b9407f507132953a1af091 # constant +69b941ca2030ce071fb493fb # constant +69b947bee8575b54b34adcdc # constant +69b9491abb8856ba14a77283 # constant +69b9491b3d7564e601970496 # constant +69b94fdbe8575b54b34add3a # jump +69b951c8e094c1fa8268fc0d # constant +69b955c7345333ff05a8d23c # constant +69b956ecea8b4cce23afeb4a # jump +69b956f84a7adca60e60cac3 # constant +69b95750a3ef2775e0e98e55 # constant +69b958aa87a8695556dc13b1 # constant +69b95987e02b54ac36b083dc # constant +69b959b2605829a110d6cbd1 # constant +69b95a8394c5e2cd4f0fd69d # constant +69b95b5ee7f0e45fa5e4894a # constant +69b95bd3e7f0e45fa5e4894e # constant +69b95be0e3e5accb88abd3c8 # constant +69b95e1da0f9b58bf114a724 # constant +69b95e3c507132953a1af1e3 # constant +69b95e8d1d448e9ea86d3ce6 # constant +69b95f3b8973f21812765e25 # constant +69b960b2b3fe52e7b3246d94 # constant +69b961190dc8a1a65089c305 # jump +69b96357c147ee9b771dce6c # constant +69b96498adb6e6a77d2401f5 # jump +69b9669336bea964e3c3b06b # jump +69b966e27ee0ff33631f0b8a # constant +69b9675b4e1bd59d5fded86a # constant +69b96958b3cc90fa8d9acf77 # constant +69b96af04e1bd59d5fded898 # jump +69b96d28b3a349ac0a3d8e5a # constant +69b96d89e094c1fa8268fd6f # constant +69b96fbc1d448e9ea86d3dcc # constant +69b96fe5e8575b54b34adede # constant +69b97094c147ee9b771dcf10 # constant +69b972ec984ca3f76894d1da # constant +69b973a1119f6fd572d62368 # constant +69b973ea10490e396dd97e19 # constant +69b97415b2ef41e07b9ad53e # constant +69b974e7adb6e6a77d24030f # constant +69b97593f25718a726e34818 # constant +69b9785e0f92f74da61ceadf # constant +69b978bca0469793c24bf37f # constant +69b9792eaacf0a241873ba49 # constant +69b979854e1bd59d5fded938 # constant +69b979fc4a18f0c8d9bf7e6b # constant +69b97e5f106a23d254069a71 # constant +69b97fc74dfc9b16e5786bc7 # constant +69b98048a8a5d3d57bad06d8 # constant +69b982d28ae31a8e6665394b # jump +69b988e154b561b4d199a12f # constant +69b98b739fd2f95542baed84 # constant +69b98d48d28517af7a845edb # constant +69b98dd19d10fb75da97f3c5 # constant +69b98e097d726cc6462424c3 # constant +69b990fab90b68a041f912af # constant +69b991273673ca659e070a6f # constant +69b99173529ed405530612af # constant +69b993b7a34871e32cd950ca # constant +69b997f53673ca659e070aad # constant +69b998acb7cbeeb258d5db9c # constant +69b9999d84c9baa55d4a2033 # constant +69b999bceddc130379669eae # jump +69b99a44104929c3da276d12 # constant +69b99a5eb1518c9453935d58 # constant +69b99a9e8433ca9c37c44e0f # jump +69b99b4cfdfee6bde9ea42c0 # constant +69b99c7f3f50bd74b8f9b411 # jump +69b99f66bc74703b35644e99 # constant +69b9a11356a215184096c352 # constant +69b9a21b31c65d1c8cb6c514 # constant +69b9a2b41f7d7f2fff171d6e # constant +69b9a42ec004b8e369a3d47e # constant +69b9a4cc106a23d254069b95 # constant +69b9a4ddb64dff59fd1893fc # constant +69b9a50aa85730cbc3e52005 # constant +69b9a58ce199cca1f42f0f9a # constant +69b9a5dfd647ee40234c80b9 # constant +69b9a6b1114b0bb346403597 # constant +69b9a7355a52e001e0cbdf33 # jump +69b9a819d9bfbe60f0573653 # constant +69b9a84211531b4a4388fa98 # constant +69b9a98d5ffb742f9e711698 # constant +69b9aa8ba6c46b5fcef51939 # constant +69b9aae431c65d1c8cb6c55a # constant +69b9ab4784c9baa55d4a20a6 # constant +69b9abf24e450ef9374d6212 # constant +69b9ac2e332eddd8ed1a35fc # constant +69b9ad0e5f7ef497d251f086 # constant +69b9ad29527eab922ae76a3d # constant +69b9b0b60a4d21fe795c1dab # constant +69b9b1b1121953ab4566f429 # constant +69b9b30502795913b379ab56 # constant +69b9c36811531b4a4388fb63 # jump +69b9c5d7529ed40553061466 # constant +69b9c6e5ed661313f0b40a4a # constant +69b9c72adea38cbf40e8242f # constant +69b9c805bb20394db7cb83d4 # constant +69b9c88ac004b8e369a3d5b9 # constant +69b9c8da862b272d33cfeabd # constant +69b9c8f8a6f41015b5203409 # constant +69b9c94e96eb780c06ea5e0d # constant +69b9c982d9ad7e205674eccb # constant +69b9c998dea38cbf40e8244f # constant +69b9c9c37460200d7164639e # constant +69b9ca0de7f8ff8bcf4ee745 # constant +69b9ca2c97ebd3eddbcf2397 # constant +69b9ca33529ed4055306149f # constant +69b9caaa4a24dede6ab8501f # constant +69b9cb12df7d76fec0802291 # constant +69b9cb4c565746ee03cbb443 # constant +69b9cb7cd3c9da58970a20b6 # constant +69b9cc30253521f049c31f86 # constant +69b9cc3c5d426ff433fda099 # constant +69b9ccf007283983c22dda36 # constant +69b9cd3e165676f7ee537967 # constant +69b9cd56c3a1b340d7445f1c # constant +69b9cd602e91c763d8d88977 # constant +69b9cd790043f7da4b9d7c6b # constant +69b9ce4007283983c22dda48 # constant +69b9cf3fd862e56ee1312a6f # constant +69b9d08a9563c2ec63a04b66 # constant +69b9d0ebe07f3be422c3cc99 # constant +69b9d0f70199ca0537430f6b # jump +69b9d1a1277e8f89504e69a5 # constant +69b9d279799879b98b39c13d # constant +69b9d29b14b5b0e24e4fe0c7 # constant +69b9d39303c8a4aff3362441 # constant +69b9d41ab64dff59fd1895b2 # constant +69b9d479136f1bb5f505f42c # constant +69b9d4c1cc5ab1c922041cec # constant +69b9d696a85730cbc3e52265 # constant +69b9d7160199ca0537430fc7 # constant +69b9d7ffc2ba99f1e9bbad35 # constant +69b9d91321ace8909aabceb0 # constant +69b9d91aef8c2b6d9612f353 # constant +69b9da2197ebd3eddbcf2434 # constant +69b9da39f8228b781f97b13f # constant +69b9db01f57489e9b3730acc # constant +69b9db24f21025ebcfd87328 # constant +69b9db51d8e76c936c822e6f # constant +69b9db53f95f50d65fb62405 # constant +69b9db6d8bdfa372d2c33bca # constant +69b9db6fed7939980848e149 # constant +69b9db867c2f56356e9e9c8c # constant +69b9db99bd54a5a03ecaeb40 # constant +69b9dbc73d69f59ff4b241b0 # constant +69b9e22b1622d2fd0809afe4 # constant +69b9e460b1f61839c2d3d301 # constant +69b9e571ad973eae519634ad # constant +69b9e573a897194fbc62d9b1 # jump +69b9e636c744d9c59679801f # constant +69b9e64fa9b6357abf470de9 # constant +69b9e7113e562d65f6438621 # constant +69b9e8790a5db606ab760cd1 # constant +69b9e95cc5a5ba0a39ec3105 # jump +69b9ea1964f8c8adcee87b87 # constant +69b9ea8863d76eda11ee7c8e # jump +69b9ea8dd35a5120e2a470ad # constant +69b9eb80bcc1b6028e39a4ba # constant +69b9ec7fa6e02486fe080d6f # jump +69b9ecc479ee15de09ec56d6 # constant +69b9f6ab4c180db3d3ac4ff1 # constant +69b9f70d50927f2015f832ac # constant +69b9f82e2e43faed7e596853 # constant +69b9f93d0199ca0537431159 # constant +69b9f96927ee5763b67cb08b # constant +69b9fc98c6f89ce2f65b5754 # constant +69b9fcb9055065a1832b575a # constant +69b9fcfd4fa8dc79b0a5e3c9 # constant +69b9fefc529ed405530615c1 # jump +69b9ff9235bd438e92af49c8 # constant +69b9ffed4b70ca9a566cbadd # constant +69ba00995ce896083ab0d6fe # constant +69ba00c3c1ff235ac7befce6 # constant +69ba01132c8c79039814c481 # constant +69ba0af27c203dd91dc168fc # constant +69ba1741192d2ab7670dc2d1 # constant +69ba1c326f992b7e0bad1319 # jump +69baca662e74bd8072a07637 # constant +69bad38636f6f84621c27f8e # constant +69bad4168e236ab0291cce64 # constant +69bad4de87716f787d182391 # constant +69bad675e14594caf653aee2 # constant +69bad6968fc471cc504cbb84 # constant +69badaec1f65cac1b9f76846 # jump +69badb9308c4df18fef54b1a # constant +69badbac27b843cc211ddf15 # jump +69badc72c25038ef965c7100 # jump +69bade0cda5440bb6aa66983 # jump +69bade428967797cb7430c0c # jump +69bade857edeefbfe14a4653 # constant +69bade9cf7dd56e062995b38 # constant +69bae062da5440bb6aa6699f # constant +69bae09736a13e0af52e6dea # constant +69bae0e840ec44cc3c14dbe3 # constant +69bae12da2eb4da46d6a855e # constant +69bae15aa61c0d91b91506d6 # constant +69bae16c3866120d05d2618a # jump +69bae1eea9cd03e5237ff0cd # jump +69bae36fda5440bb6aa669c5 # jump +69bae7f7c40add85137b6f3c # jump +69bae7febae66100347d2988 # constant +69bae9002c3d669e7fabae2b # constant +69baec21071c4967be381cb2 # jump +69baed12f492c4f8a835a99a # constant +69baed27f205b66bf504f4b6 # constant +69baef3f10ece1f993b3f0aa # constant +69baef7ab7de678fb06af1d1 # constant +69baf0849e7c65d18be11133 # constant +69baf08c8c02130ff9f0c6cd # constant +69baf3550826d6aa1fc63d15 # constant +69baf376466ebc139e53c6ce # constant +69baf3d500a72d9655373ab6 # constant +69baf579f7dd56e062995c03 # constant +69baf77aed46fe5d7036ac14 # constant +69baff10466ebc139e53c720 # constant +69baff92071c4967be381d48 # constant +69bb012be14594caf653b066 # constant +69bb02e5d50175523d0489a5 # constant +69bb0a998606a40f2af2a772 # jump +69bb0c26527daa899b4c7ceb # constant +69bb0e991428e3c5a32c047c # constant +69bb0f12f3f248f510ffc2b6 # jump diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index c6fb9d1870..05eeb10eff 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, @@ -15,6 +15,14 @@ _viz_traj, ) +# A camera keymap entry "_hist" is the same camera, `lag_s` seconds earlier. +IMAGE_HISTORY_SUFFIX = "_hist" +# "_mem" is the same camera's long-range window: `history` frames, +# `history_stride_s` apart, the newest `lag_s` seconds back. +IMAGE_MEMORY_SUFFIX = "_mem" +# A camera window "" comes with "_mask", (K,) 1.0 per real frame. +HISTORY_MASK_SUFFIX = "_mask" + class EMBODIMENT(Enum): # All human demonstration data is one embodiment (HUMAN_*); the robot Eva is @@ -106,6 +114,10 @@ def _strip_pi_keymap_mode(cls, keymap_mode: str) -> str: "cartesian_padded": "cartesian_6d", "cartesian_wristframe_ypr": "cartesian_wristframe_6d", "cartesian_wristframe_quat": "cartesian_wristframe_6d", + "keypoints_headframe_ypr": "keypoints_headframe_6d", + "keypoints_headframe_quat": "keypoints_headframe_6d", + "keypoints_wristframe_ypr": "keypoints_wristframe_6d", + "keypoints_wristframe_quat": "keypoints_wristframe_6d", } diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 6067ad16fc..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, @@ -18,6 +18,7 @@ ConcatKeys, DeleteKeys, InterpolateLinear, + InterpolatePadMask, InterpolatePose, NumpyToTensor, PoseCoordinateFrameTransform, @@ -396,6 +397,7 @@ def _build_eva_bimanual_eef_frame_transform_list( output_action_key=right_cmd_gripper, stride=stride, ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), # Step 2: transform camera-frame actions into EEF-relative (wrist) frame ActionChunkCoordinateFrameTransform( target_world=left_obs_camframe, @@ -559,6 +561,7 @@ def _build_eva_bimanual_transform_list( output_action_key=right_cmd_gripper, stride=stride, ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if is_quat: diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index cadef181c6..387716c2df 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -5,23 +5,29 @@ import numpy as np from egomimic.rldb.embodiment.embodiment import ( + IMAGE_HISTORY_SUFFIX, + IMAGE_MEMORY_SUFFIX, Embodiment, _reject_legacy_rotation, _strip_pi_keymap_mode, ) -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, CartesianRot6DToYPR, CartesianYPRToRot6D, ConcatKeys, DeleteKeys, + InterpolatePadMask, InterpolatePose, + KeypointsRot6DToYPR, + KeypointsYPRToRot6D, PadGripperZeros, PoseCoordinateFrameTransform, QuaternionPoseToYPR, Reshape, RotateLocalFrame, + SelectCurrentStep, SplitKeys, Transform, UnpadGripperZeros, @@ -30,7 +36,7 @@ from egomimic.utils.viz_utils import ( ColorPalette, _viz_gaze, - _viz_keypoints, + _viz_keypoints_horizon_trace, ) ARIA_INTRINSICS = np.array( @@ -83,38 +89,6 @@ ) -# Aria's raw 21-keypoint layout (0-4 fingertips, 5 palm root) — NOT MANO. Used -# only for the opt-in raw-Aria-keypoint viz; the canonical keypoints are MANO. -ARIA_FINGER_EDGES = [ - (5, 6), - (6, 7), - (7, 0), # thumb - (5, 8), - (8, 9), - (9, 10), - (10, 1), # index - (5, 11), - (11, 12), - (12, 13), - (13, 2), # middle - (5, 14), - (14, 15), - (15, 16), - (16, 3), # ring - (5, 17), - (17, 18), - (18, 19), - (19, 4), # pinky -] -ARIA_FINGER_EDGE_RANGES = [ - ("thumb", 0, 3), - ("index", 3, 7), - ("middle", 7, 11), - ("ring", 11, 15), - ("pinky", 15, 19), -] - - class Human(Embodiment): """Single data-driven human embodiment shared by every vendor's human data. @@ -179,8 +153,6 @@ def viz( "traj", "traj+rotation", "axes", "annotations", "keypoints", "gaze" ], intrinsics=None, - finger_edges=None, - finger_edge_ranges=None, **kwargs, ): K = intrinsics if intrinsics is not None else cls.INTRINSICS @@ -200,22 +172,24 @@ def viz( finger: ColorPalette.to_rgb(color, value=(i + 1) / (n + 1)) for i, finger in enumerate(cls.FINGER_COLORS) } - dot_color = ColorPalette.to_rgb(color, value=0.7) else: colors = cls.FINGER_COLORS - dot_color = cls.DOT_COLOR - return _viz_keypoints( + # The horizon trace is the ONLY keypoint renderer. A chunk draws + # as its trajectory; a lone timestep degenerates to its start + # markers. There is deliberately no static-skeleton path to fall + # back to: that fallback is what silently rendered motionless + # keypoint overlays for a whole 50k-step campaign, and a renderer + # you can reach by accident is one you will reach. + return _viz_keypoints_horizon_trace( image=image, - actions=viz_data, + actions=np.asarray(viz_data), intrinsics=K, - edges=finger_edges if finger_edges is not None else cls.FINGER_EDGES, - edge_ranges=( - finger_edge_ranges - if finger_edge_ranges is not None - else cls.FINGER_EDGE_RANGES - ), colors=colors, - dot_color=dot_color, + wrist_color=( + ColorPalette.to_rgb(color, value=0.85) + if color is not None and ColorPalette.is_valid(color) + else cls.DOT_COLOR + ), **kwargs, ) return super().viz(image, viz_data, mode=mode, intrinsics=intrinsics, **kwargs) @@ -226,19 +200,74 @@ def get_keymap( keymap_mode: str, has_head_pose: bool = True, include_aria_keypoints: bool = False, + include_ee_pose: bool = False, norm_mode: bool = False, annotation_key: str = None, + proprio_history: int = 1, + history_stride: int = 1, + history_stride_s: float | None = None, + image_history_gap_s: float | None = None, + image_memory: int = 0, + image_memory_stride_s: float = 1.0, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (False only for contributors that truly omit - ``obs_head_pose``; incompatible with cartesian modes, which pivot on it) and ``include_aria_keypoints`` - (Aria=True). ``norm_mode``/``annotation_key`` behave as in the base. + ``obs_head_pose``; incompatible with cartesian modes, which pivot on + it), ``include_aria_keypoints`` (Aria=True) and ``include_ee_pose`` + (keypoints mode only: also read the palm-origin ``obs_ee_pose`` keys so + the keypoint transform modes can build the cartesian + ``observations.state.ee_pose`` proprio, the pi0.5 prompt state). + ``norm_mode``/``annotation_key`` behave as in the base. + + ``proprio_history`` (K) makes the proprio keys the model consumes read + the last K frames instead of one: the dataset emits ``(K, D)`` with the + current frame last, plus a ``proprio_history_mask``. K = 1 (the + default) leaves the keymap exactly as it was. ``history_stride`` (s) + spaces that window out -- the steps are ``idx - s*(K-1) ... idx`` -- + because at 30 fps consecutive frames are nearly a duplicate of the + current pose (median wrist translation 6.8 mm at lag 1 against 42.9 mm + at lag 5), so s = 1 spends tokens on sensor noise. + + ``image_history_gap_s`` adds a second front-camera entry, + ``_hist``: the frame that many seconds before the current + one (frames = gap x the episode's fps; frame 0 at an episode start). + The algo must consume it explicitly: HPT has no stem for it and skips + augs and normalization on cameras outside its encoders, so under plain + HPT it is decoded and dropped. + + ``history_stride_s`` replaces ``history_stride`` by a spacing in + seconds, read in each episode's fps. ``image_memory`` (N) adds + ``_mem``: N front frames ``image_memory_stride_s`` apart, + the newest the current frame, plus a ``_mem_mask`` of the + steps that fall inside the episode. Same caveat as ``_hist``: only an algo + that consumes it uses it. """ key_map = cls._get_keymap( keymap_mode, has_head_pose=has_head_pose, include_aria_keypoints=include_aria_keypoints, + include_ee_pose=include_ee_pose, + proprio_history=proprio_history, + history_stride=history_stride, + history_stride_s=history_stride_s, ) + if image_history_gap_s is not None: + front = key_map[cls.VIZ_IMAGE_KEY] + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_HISTORY_SUFFIX}"] = { + **front, + "lag_s": float(image_history_gap_s), + } + if image_memory > 0: + if image_memory_stride_s <= 0: + raise ValueError( + f"image_memory_stride_s must be > 0, got {image_memory_stride_s}" + ) + front = key_map[cls.VIZ_IMAGE_KEY] + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_MEMORY_SUFFIX}"] = { + **front, + "history": int(image_memory), + "history_stride_s": float(image_memory_stride_s), + } if annotation_key is not None and not norm_mode: key_map[annotation_key] = { "key_type": "annotation_keys", @@ -260,10 +289,17 @@ def _get_keymap( keymap_mode: str, has_head_pose: bool = True, include_aria_keypoints: bool = False, + include_ee_pose: bool = False, + proprio_history: int = 1, + history_stride: int = 1, + history_stride_s: float | None = None, ): """Canonical MANO keymap. ``include_aria_keypoints`` additionally exposes - the raw Aria-layout proprio keypoints alongside the MANO ones. The front - image is always ``VIZ_IMAGE_KEY``; Pi renames it onto its own slot. + the raw Aria-layout proprio keypoints alongside the MANO ones; + ``include_ee_pose`` adds the palm-origin ee_pose keys to the keypoints + keymap. The front image is always ``VIZ_IMAGE_KEY``; Pi renames it + onto its own slot. ``proprio_history`` > 1 marks the proprio entries + the model consumes with ``"history": K`` (see ``get_keymap``). """ base_mode = _strip_pi_keymap_mode(cls, keymap_mode) front_key = cls.VIZ_IMAGE_KEY @@ -338,6 +374,23 @@ def _get_keymap( "zarr_key": "right.obs_wrist_pose", }, } + if include_ee_pose: + # Palm-origin ee_pose alongside the keypoints: lets the + # keypoint transform modes build the SAME head-frame + # ``observations.state.ee_pose`` proprio the cartesian pi runs + # put in the prompt (left-wrist fix, 6D, grip-padded), so a + # keypoint-action run differs from them only in the action + # space. Pair with ``get_transform_list(include_ee_pose=True)``. + for side in ("left", "right"): + key_map[f"{side}.action_ee_pose"] = { + "key_type": "action_keys", + "zarr_key": f"{side}.obs_ee_pose", + "horizon": horizon, + } + key_map[f"{side}.obs_ee_pose"] = { + "key_type": "proprio_keys", + "zarr_key": f"{side}.obs_ee_pose", + } if include_aria_keypoints: # Raw Aria-layout keypoints exposed alongside the canonical MANO # ones (proprio, no horizon: no transform consumes them). @@ -357,6 +410,30 @@ def _get_keymap( "key_type": "proprio_keys", "zarr_key": "obs_head_pose", } + + if proprio_history > 1: + if history_stride < 1: + raise ValueError(f"history_stride must be >= 1, got {history_stride}") + # Backward window on the proprio the MODEL consumes. Not on + # ``obs_head_pose``: every history step is expressed in the CURRENT + # step's head frame (a per-step head frame would make the history + # unreadable), and not on the horizoned ``action_*`` entries, which + # are forward chunks. + if base_mode == "cartesian": + history_keys = ["left.obs_ee_pose", "right.obs_ee_pose"] + else: + history_keys = [ + f"{side}.{name}" + for side in ("left", "right") + for name in ("obs_keypoints", "obs_wrist_pose") + ] + if include_ee_pose: + history_keys += ["left.obs_ee_pose", "right.obs_ee_pose"] + for key in history_keys: + key_map[key]["history"] = int(proprio_history) + key_map[key]["history_stride"] = int(history_stride) + if history_stride_s is not None: + key_map[key]["history_stride_s"] = float(history_stride_s) return key_map @classmethod @@ -370,12 +447,15 @@ def get_transform_list( "cartesian_wristframe_6d", "keypoints_headframe_ypr", "keypoints_headframe_quat", + "keypoints_headframe_6d", "keypoints_wristframe_ypr", "keypoints_wristframe_quat", + "keypoints_wristframe_6d", ], stride: int = 3, fix_left_wrist_convention: bool = False, pad_proprio_gripper: bool = False, + include_ee_pose: bool = False, allow_legacy_rotation: bool = False, ) -> list[Transform]: """Transform pipeline. ``stride`` is the per-vendor action stride @@ -389,15 +469,15 @@ def get_transform_list( wrist-frame convention of zarrs written before the ``rot_left`` fix in ``mecka_to_zarr.compute_hand_pose_xyzquat`` (which double-mirrored the left hand onto the right hand's spatial convention): the raw left - pose keys (ee_pose in cartesian modes, wrist_pose in keypoints modes; - the converter writes the same rotation into both) are right-multiplied - by Rz(180 deg) before any frame math, exactly equivalent to - reconverting. It is a property of the DATA, not of the vendor: enable - it for any zarr written by the pre-fix hand-pose converter (measured: - every mecka and scale episode in the production set), and disable it - for data reconverted after the fix, which would double-flip. Aria and - LightWheel are not covered by that measurement -- no such episodes are - in the production set -- so leave them unset. + pose keys (ee_pose in cartesian modes, wrist_pose in keypoints modes, + both when ``include_ee_pose``; the converter writes the same rotation + into both) are right-multiplied by Rz(180 deg) before any frame math, + exactly equivalent to reconverting. It is a property of the DATA, not + of the vendor: enable it for any zarr written by the pre-fix hand-pose + converter (measured: every mecka and scale episode in the production + set), and disable it for data reconverted after the fix, which would + double-flip. Aria and LightWheel are not covered by that measurement -- + no such episodes are in the production set -- so leave them unset. ``left_wrist_convention`` re-derives the verdict from a zarr's own keypoints when a setting has to be checked. @@ -405,21 +485,31 @@ def get_transform_list( ``observations.state.ee_pose`` 18 -> 20 (zero grip slots at 9/19) so the pi0.5 ``State:`` prompt bins align positionally with the robot 20D layout. 6D modes only. + + ``include_ee_pose`` (keypoints modes only; pair with + ``get_keymap(include_ee_pose=True)``) additionally builds the + cartesian pipeline's head-frame ``observations.state.ee_pose`` proprio + (left-wrist fix, 6D-encoded, optionally grip-padded) next to the + keypoint action, so a keypoint-action pi0.5 run shares its prompt + state with the cartesian runs and differs only in the action space. """ _reject_legacy_rotation(cls, mode, allow_legacy_rotation) is_keypoints = mode.startswith("keypoints") prefix: list[Transform] = [] if fix_left_wrist_convention: - fix_keys = ( - ["left.action_wrist_pose", "left.obs_wrist_pose"] - if is_keypoints - else ["left.action_ee_pose", "left.obs_ee_pose"] - ) + fix_keys = [] + if is_keypoints: + fix_keys += ["left.action_wrist_pose", "left.obs_wrist_pose"] + if not is_keypoints or include_ee_pose: + fix_keys += ["left.action_ee_pose", "left.obs_ee_pose"] prefix = [RotateLocalFrame(keys=fix_keys)] - if pad_proprio_gripper and not mode.endswith("_6d"): + if include_ee_pose and not is_keypoints: + raise ValueError("include_ee_pose only applies to the keypoints modes") + has_6d_ee_pose = (not is_keypoints and mode.endswith("_6d")) or include_ee_pose + if pad_proprio_gripper and not has_6d_ee_pose: raise ValueError( "pad_proprio_gripper needs a 6D-encoded ee_pose proprio: a " - "cartesian *_6d mode" + "cartesian *_6d mode, or a keypoints mode with include_ee_pose" ) proprio_6d: list[Transform] = [ @@ -434,10 +524,8 @@ def get_transform_list( ) if mode == "cartesian_6d": # Head/camera-frame cartesian (12D xyz+ypr per arm pair) with the - # rotation re-expressed as the continuous 6D representation (18D) - # for pi0.5. The proprio ee_pose is 6D-encoded too: normalized YPR - # saturates yaw/roll at +-pi, so per-dim normalization needs the - # continuous representation. + # rotation as continuous 6D (18D) for pi0.5, on the proprio ee_pose + # as well as the actions -- see CartesianYPRToRot6D. return ( prefix + _build_human_cartesian_bimanual_transform_list(stride=stride) @@ -464,22 +552,47 @@ def get_transform_list( + [CartesianYPRToRot6D(action_key="actions_cartesian")] + proprio_6d ) - if mode == "keypoints_headframe_ypr": - return prefix + _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=False - ) - if mode == "keypoints_headframe_quat": - return prefix + _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=True - ) - if mode == "keypoints_wristframe_ypr": - return prefix + _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=False - ) - if mode == "keypoints_wristframe_quat": - return prefix + _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=True + + if is_keypoints: + frame, _, rot = mode[len("keypoints_") :].partition("_") + if frame not in ("headframe", "wristframe") or rot not in ( + "ypr", + "quat", + "6d", + ): + raise ValueError( + f"Unsupported transform_list mode '{mode}' for {cls.__name__}" + ) + ee_pose: list[Transform] = [] + if include_ee_pose: + # The cartesian builder keeps ``obs_head_pose`` alive for the + # keypoint builder (delete_target_world=False); its + # ``actions_cartesian`` by-product is dropped. + ee_pose = ( + _build_human_cartesian_bimanual_transform_list( + stride=stride, delete_target_world=False + ) + + proprio_6d + + [DeleteKeys(keys_to_delete=["actions_cartesian"])] + ) + build = ( + _build_human_keypoints_eef_frame_transform_list + if frame == "wristframe" + else _build_human_keypoints_bimanual_transform_list ) + if rot == "quat": + return prefix + ee_pose + build(stride=stride, is_quat=True) + keypoints = build(stride=stride, is_quat=False) + if rot == "6d": + # 138-D ypr layout -> 144-D: per hand [wrist xyz (3) | wrist + # rot6d (6) | 21 keypoints (63)], for both the action and the + # proprio ``observations.state.keypoints`` (head-frame wrist + # pose + keypoints), continuous like the cartesian 6d modes. + keypoints = keypoints + [ + KeypointsYPRToRot6D(action_key="actions_keypoints"), + KeypointsYPRToRot6D(action_key="observations.state.keypoints"), + ] + return prefix + ee_pose + keypoints raise ValueError(f"Unsupported transform_list mode '{mode}' for {cls.__name__}") @@ -507,6 +620,9 @@ def _build_human_keypoints_revert_eef_frame_transform_list( else: pose_shape = 6 transform_list = [ + # proprio_history: the proprio defines the frame, so keep the CURRENT + # step (no-op at K = 1, idempotent after the 6D revert). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -645,17 +761,22 @@ def _build_human_keypoints_eef_frame_transform_list( output_key=right_keypoints_action_headframe, shape=(chunk_length, 21, 3), ), + # target_history: the proprio wrist pose may carry a leading time + # axis (proprio_history); everything is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_wrist_obs_headframe, chunk_world=left_keypoints_action_headframe, transformed_key_name=left_keypoints_action_wristframe, mode="xyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_keypoints_action_headframe, transformed_key_name=right_keypoints_action_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_action_wristframe, @@ -677,17 +798,30 @@ def _build_human_keypoints_eef_frame_transform_list( output_key=right_keypoints_obs_headframe, shape=(21, 3), ), + # per_step_target: with proprio_history K, step k's keypoints go + # into step k's OWN wrist frame, not the current step's. Otherwise + # a past hand is measured from today's wrist, which smears wrist + # motion into the articulation channels -- the palm-rigid knuckles + # (MANO 1/5/13/17) have a very tight legitimate range and land far + # outside it, worsening with age (measured 2026-09-17: 10.3/6.9/3.3 + # out-of-range cells per sample at steps 0/1/2, 0.0 at the current + # step). Per-step framing keeps every step a genuine state, so + # keypoint 0 is 0 at EVERY step and the norm stats / bounds check + # describe all of them. The ACTION chunk keeps the current-step + # frame: its leading axis is the forward horizon, not history. PoseCoordinateFrameTransform( target_world=left_wrist_obs_headframe, pose_world=left_keypoints_obs_headframe, transformed_key_name=left_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), PoseCoordinateFrameTransform( target_world=right_wrist_obs_headframe, pose_world=right_keypoints_obs_headframe, transformed_key_name=right_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_obs_wristframe, @@ -704,12 +838,14 @@ def _build_human_keypoints_eef_frame_transform_list( chunk_world=left_wrist_action_headframe, transformed_key_name=left_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_wrist_action_headframe, transformed_key_name=right_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ] ) @@ -930,6 +1066,7 @@ def _build_human_keypoints_bimanual_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if not is_quat: transform_list.extend( @@ -1003,6 +1140,10 @@ def _build_human_cartesian_revert_eef_frame_transform_list( pose_shape = 7 if is_quat else 6 mode = "xyzwxyz" if is_quat else "xyzypr" transform_list = [ + # The proprio defines the frame the action chunk is projected out of; + # with proprio_history it arrives (K, D), so take the CURRENT step + # (no-op at K = 1, and idempotent when a 6D revert already did it). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -1055,6 +1196,7 @@ def _build_human_cartesian_revert_6d_transform_list( grip-padded 18 -> 20) is reverted the same way. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1076,6 +1218,7 @@ def _build_human_cartesian_revert_6d_wristframe_transform_list( actions back into head frame using that ypr proprio to define the frame. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1083,6 +1226,45 @@ def _build_human_cartesian_revert_6d_wristframe_transform_list( ] +def _build_human_keypoints_revert_6d_transform_list( + *, + action_key: str = "actions_keypoints", + obs_key: str = "observations.state.keypoints", +) -> list[Transform]: + """Revert head-frame 144-D keypoint actions (``keypoints_headframe_6d``) + to the 138-D ypr layout the keypoint overlay draws (the keypoint blocks + are already in head frame; only the wrist rotation representation + changes). The proprio is reverted the same way. + """ + return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step + KeypointsRot6DToYPR(action_key=action_key), + KeypointsRot6DToYPR(action_key=obs_key), + ] + + +def _build_human_keypoints_revert_6d_wristframe_transform_list( + *, + action_key: str = "actions_keypoints", + obs_key: str = "observations.state.keypoints", +) -> list[Transform]: + """Revert wrist-frame 144-D keypoint actions (``keypoints_wristframe_6d``) + to head-frame keypoints for the overlay: (1) wrist rot6d -> ypr on the + action and the proprio (144 -> 138 each); (2) the standard keypoint eef + revert projects each hand's 63-D keypoint chunk back to head frame via the + head-frame wrist pose in the proprio, giving a 126-D head-frame + ``actions_keypoints`` (the wrist-pose slices are consumed by the revert). + """ + return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step + KeypointsRot6DToYPR(action_key=action_key), + KeypointsRot6DToYPR(action_key=obs_key), + *_build_human_keypoints_revert_eef_frame_transform_list( + action_key=action_key, obs_key=obs_key, is_quat=False + ), + ] + + def _build_human_cartesian_eef_frame_transform_list( *, target_world: str = "obs_head_pose", @@ -1166,17 +1348,23 @@ def _build_human_cartesian_eef_frame_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), + # target_history: the proprio obs pose may carry a leading time axis + # (proprio_history); the action chunk is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_obs_headframe, chunk_world=left_action_headframe, transformed_key_name=left_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_obs_headframe, chunk_world=right_action_headframe, transformed_key_name=right_action_wristframe, mode="xyzwxyz", + target_history=True, ), XYZWXYZ_to_XYZYPR( keys=[ @@ -1279,6 +1467,7 @@ def _build_human_cartesian_bimanual_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if target_world_is_quat: diff --git a/egomimic/rldb/filters.py b/egomimic/rldb/filters.py index 2f06da6f14..15ca71d940 100644 --- a/egomimic/rldb/filters.py +++ b/egomimic/rldb/filters.py @@ -1,27 +1,72 @@ from __future__ import annotations +import logging import os import sys from collections.abc import Mapping, Sequence +from pathlib import Path from typing import Any from egomimic.rldb.resolve_memo import memoized +logger = logging.getLogger(__name__) + +BLOCKLIST_DIR = Path(__file__).resolve().parents[1] / "resources" / "blocklists" + + +def load_hash_list(source: str | Sequence[str] | None) -> frozenset[str]: + """Episode hashes from a list, a text file, or a blocklist name. + + A string is a path if one exists there, else ``BLOCKLIST_DIR/.txt``. + Files hold one hash per line; ``#`` starts a comment. + """ + if source is None: + return frozenset() + if not isinstance(source, str): + return frozenset(str(h) for h in source) + path = Path(source) + if not path.is_file(): + path = BLOCKLIST_DIR / f"{source}.txt" + if not path.is_file(): + raise FileNotFoundError( + f"exclude_hashes {source!r}: no such file, and no blocklist " + f"{path.name} in {BLOCKLIST_DIR}" + ) + hashes = (line.split("#", 1)[0].strip() for line in path.read_text().splitlines()) + return frozenset(h for h in hashes if h) + class DatasetFilter: def __init__( self, filter_lambdas: Sequence[str] | None = None, episode_hashes: Sequence[str] | None = None, + exclude_hashes: str | 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 []) - ) + pins = frozenset(str(h) for h in (episode_hashes or [])) + # Episodes that never match, e.g. known-bad data (see load_hash_list). + self.exclude_hashes = load_hash_list(exclude_hashes) + # An excluded pin is dropped rather than raised on, so one blocklist can + # apply under hand-pinned splits; it is logged, and a pin list excluded + # entirely is an error because an empty pin set means "no pin". + dropped = pins & self.exclude_hashes + if dropped: + if dropped == pins: + raise ValueError( + f"all {len(pins)} pinned episode(s) are in exclude_hashes" + ) + logger.warning( + "Dropping %d of %d pinned episode(s) listed in exclude_hashes: %s", + len(dropped), + len(pins), + sorted(dropped)[:10] + (["..."] if len(dropped) > 10 else []), + ) + # Pinned episode hashes. Empty = no pin. Validated at resolve time by the + # resolvers (missing / deleted / wrong embodiment is an error there). + self.episode_hashes: frozenset[str] = pins - dropped self.filters = [] for expr in self.filter_lambdas: try: @@ -37,7 +82,8 @@ def __init__( def __repr__(self) -> str: return ( f"DatasetFilter(filter_lambdas={self.filter_lambdas!r}, " - f"episode_hashes={sorted(self.episode_hashes)!r})" + f"episode_hashes={sorted(self.episode_hashes)!r}, " + f"exclude_hashes=<{len(self.exclude_hashes)} hashes>)" ) def cache_key(self) -> tuple | None: @@ -49,14 +95,26 @@ def cache_key(self) -> tuple | None: """ if "cache_key" not in type(self).__dict__: return None - return (type(self), tuple(self.filter_lambdas), self.episode_hashes) + return ( + type(self), + tuple(self.filter_lambdas), + self.episode_hashes, + self.exclude_hashes, + ) def matches(self, row: Mapping[str, Any]) -> bool: + if self.episode_hashes and row.get("episode_hash") not in self.episode_hashes: + return False + if row.get("episode_hash") in self.exclude_hashes: + return False + return self.passes_lambdas(row) + + def passes_lambdas(self, row: Mapping[str, Any]) -> bool: + """`matches` without the pins and excludes, for resolvers that apply those to + time ranges themselves (Lance).""" 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): @@ -72,6 +130,7 @@ def __init__( project_name: str, filter_lambdas: Sequence[str] | None = None, episode_hashes: Sequence[str] | None = None, + exclude_hashes: str | Sequence[str] | None = None, ) -> None: from egomimic.utils.scale_utils import build_df_from_tasks, get_completed_tasks @@ -87,7 +146,7 @@ def __init__( self.completed_episode_hashes = frozenset( self.df["SEQUENCE_ID"].unique().tolist() ) - super().__init__(filter_lambdas, episode_hashes) + super().__init__(filter_lambdas, episode_hashes, exclude_hashes) def cache_key(self) -> tuple | None: base = super().cache_key() @@ -99,3 +158,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 in the zarr resolver path. " + "Plain dict filters are no longer supported." + ) diff --git a/egomimic/rldb/lance/__init__.py b/egomimic/rldb/lance/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/rldb/lance/annotations.py b/egomimic/rldb/lance/annotations.py new file mode 100644 index 0000000000..81b73c25f2 --- /dev/null +++ b/egomimic/rldb/lance/annotations.py @@ -0,0 +1,139 @@ +"""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, _open, annotations_uri, join +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) -> pa.Table: + uri = annotations_uri(root) + if not _exists(uri): + return ANNOTATIONS.empty_table() + return _open(uri).to_table() + + +def _task_bounds(root) -> dict[str, list[tuple[int, int]]]: + cols = ["recording_id", "t_start_ns", "t_end_ns"] + out = defaultdict(list) + for t in _open(join(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, 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, 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 new file mode 100644 index 0000000000..ac63810a85 --- /dev/null +++ b/egomimic/rldb/lance/commit.py @@ -0,0 +1,566 @@ +"""Committer operations on the Lance tables (https://lance.org/format/table/). + +Each operation is one Lance commit per table it touches. Lance has no multi-table +transaction, so writes go steps -> aliases -> tasks -> recordings: a recording is +published by its `recordings` row, and readers ignore steps rows nobody references. +Videos live outside Lance under `videos///` and are never +overwritten, so older table versions keep pointing at the files they were written with. + +`root` is a local directory or an s3:// prefix; credentials come from the R2_* +environment (dataset.storage_options). +""" + +import datetime +import getpass +import hashlib +import json +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path + +import lance +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq +from lance import LanceOperation +from lance.fragment import write_fragments + +from egomimic.rldb.lance.dataset import _is_remote, _s3_client, storage_options +from egomimic.rldb.lance.schema import ALIASES, ANNOTATIONS, CHANGES, RECORDINGS, TASKS + +CATALOG = { + "aliases.lance": ALIASES, + "tasks.lance": TASKS, + "recordings.lance": RECORDINGS, +} + + +@dataclass +class Delivery: + """One validated recording: catalog rows without `video_path`/`content_digest`, its steps and local MP4s.""" + + recording: dict + tasks: list[dict] # a recording can hold several tasks (e.g. iron, fold, pack) + alias: dict + steps: pa.Table + videos: dict[str, Path] + + @property + def recording_id(self) -> str: + return self.recording["recording_id"] + + +def join(root, *parts: str) -> str: + return "/".join([str(root).rstrip("/"), *parts]) + + +def _split(uri: str) -> tuple[str, str]: + bucket, key = uri[len("s3://") :].split("/", 1) + return bucket, key + + +def list_files(prefix) -> dict[str, tuple[int, str]]: + """{path relative to `prefix`: (size, etag or mtime)} of every file under it.""" + prefix = str(prefix).rstrip("/") + if not _is_remote(prefix): + root = Path(prefix) + return { + str(p.relative_to(root)): (p.stat().st_size, str(p.stat().st_mtime_ns)) + for p in (root.rglob("*") if root.exists() else ()) + if p.is_file() + } + bucket, key = _split(prefix) + out = {} + for page in ( + _s3_client() + .get_paginator("list_objects_v2") + .paginate(Bucket=bucket, Prefix=key + "/") + ): + for o in page.get("Contents", []): + out[o["Key"][len(key) + 1 :]] = (o["Size"], o["ETag"]) + return out + + +def list_dirs(prefix) -> list[str]: + """Names of the directories (common prefixes) directly under `prefix`.""" + prefix = str(prefix).rstrip("/") + if not _is_remote(prefix): + p = Path(prefix) + return sorted(d.name for d in p.iterdir() if d.is_dir()) if p.exists() else [] + bucket, key = _split(prefix) + out = [] + for page in ( + _s3_client() + .get_paginator("list_objects_v2") + .paginate(Bucket=bucket, Prefix=key + "/", Delimiter="/") + ): + out += [ + c["Prefix"][len(key) + 1 :].rstrip("/") + for c in page.get("CommonPrefixes", []) + ] + return sorted(out) + + +def _nonempty(prefix: str) -> bool: + if not _is_remote(prefix): + p = Path(prefix) + return p.is_dir() and any(p.iterdir()) + bucket, key = _split(prefix.rstrip("/")) + resp = _s3_client().list_objects_v2(Bucket=bucket, Prefix=key + "/", MaxKeys=1) + return resp.get("KeyCount", 0) > 0 + + +def put_file(src: Path, dst: str) -> None: + if _is_remote(dst): + _s3_client().upload_file(str(src), *_split(dst)) + else: + Path(dst).parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) + + +def put_bytes(data: bytes, dst: str) -> None: + if _is_remote(dst): + bucket, key = _split(dst) + _s3_client().put_object(Bucket=bucket, Key=key, Body=data) + else: + Path(dst).parent.mkdir(parents=True, exist_ok=True) + Path(dst).write_bytes(data) + + +def get_bytes(src: str) -> bytes: + if _is_remote(src): + bucket, key = _split(src) + return _s3_client().get_object(Bucket=bucket, Key=key)["Body"].read() + return Path(src).read_bytes() + + +def get_file(src: str, dst: Path) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if _is_remote(src): + _s3_client().download_file(*_split(src), str(dst)) + else: + shutil.copy2(src, dst) + + +def remove(uris: list[str]) -> None: + remote = [u for u in uris if _is_remote(u)] + for u in uris: + if not _is_remote(u): + Path(u).unlink(missing_ok=True) + for i in range(0, len(remote), 1000): + by_bucket: dict[str, list] = {} + for u in remote[i : i + 1000]: + bucket, key = _split(u) + by_bucket.setdefault(bucket, []).append({"Key": key}) + for bucket, objs in by_bucket.items(): + _s3_client().delete_objects(Bucket=bucket, Delete={"Objects": objs}) + + +def stage(delivery: Path, dst) -> None: + """Upload a delivery directory to the staging prefix `dst`, meta.json last: the + committer only picks up deliveries whose meta.json exists, so it never sees a + half-uploaded one.""" + files = sorted(p for p in Path(delivery).iterdir() if p.is_file()) + for p in sorted(files, key=lambda p: p.name == "meta.json"): + put_file(p, join(dst, p.name)) + + +# A delivery whose catalog rows travel with it, so the committer needs no lookup (SQL for +# Aria, meta.json fields for EVA): meta.json has "format": "delivery". +DESCRIBED = "delivery" +_CATALOG_FILES = { + "recordings.parquet": RECORDINGS, + "tasks.parquet": TASKS, + "aliases.parquet": ALIASES, +} +_UNSET = ("video_path", "content_digest") # filled by the commit + + +def stage_delivery( + d: Delivery, dst, labels: pa.Table | None = None, meta: dict | None = None +) -> None: + """Stage `d` with its catalog rows, and optionally its labels (ANNOTATIONS rows), at + `dst`; meta.json last, as `stage`.""" + rows = { + "recordings.parquet": [d.recording], + "tasks.parquet": d.tasks, + "aliases.parquet": [d.alias], + } + with tempfile.TemporaryDirectory() as t: + t = Path(t) + pq.write_table(d.steps, t / "steps.parquet") + for name, schema in _CATALOG_FILES.items(): + pq.write_table(pa.Table.from_pylist(rows[name], schema=schema), t / name) + if labels is not None: + pq.write_table( + labels.select(ANNOTATIONS.names).cast(ANNOTATIONS), t / "labels.parquet" + ) + for f in sorted(t.iterdir()): + put_file(f, join(dst, f.name)) + for cam, src in sorted(d.videos.items()): + put_file(Path(src), join(dst, f"{cam}.mp4")) + meta = { + **(meta or {}), + "format": DESCRIBED, + "recording_id": d.recording_id, + "vendor_recording_id": d.recording["vendor_recording_id"], + "num_frames": d.recording["num_frames"], + "videos": sorted(d.videos), + } + put_bytes(json.dumps(meta, indent=1, default=str).encode(), join(dst, "meta.json")) + + +def load_staged_delivery(d: Path) -> tuple[Delivery, pa.Table | None]: + """The Delivery and labels `stage_delivery` wrote, read from a local copy `d`.""" + meta = json.loads((d / "meta.json").read_text()) + rec, tasks, alias = (pq.read_table(d / name).to_pylist() for name in _CATALOG_FILES) + recording = {k: v for k, v in rec[0].items() if k not in _UNSET} + labels = ( + pq.read_table(d / "labels.parquet") if (d / "labels.parquet").exists() else None + ) + videos = {cam: d / f"{cam}.mp4" for cam in meta["videos"]} + steps = pq.read_table(d / "steps.parquet") + return Delivery(recording, tasks, alias[0], steps, videos), labels + + +def _size(uri: str) -> int | None: + if not _is_remote(uri): + return Path(uri).stat().st_size if Path(uri).exists() else None + from botocore.exceptions import ClientError + + bucket, key = _split(uri) + try: + return _s3_client().head_object(Bucket=bucket, Key=key)["ContentLength"] + except ClientError: + return None + + +def _open(uri: str, version=None) -> lance.LanceDataset: + return lance.dataset(uri, version=version, storage_options=storage_options(uri)) + + +def steps_uri(root, lab: str) -> str: + return join(root, "steps", f"{lab}.lance") + + +def annotations_uri(root) -> str: + return join(root, "annotations.lance") + + +def changes_uri(root) -> str: + return join(root, "changes.lance") + + +def table_uris(root) -> list[str]: + """Every table, in write order: each lab's steps, the catalog, changes, annotations.""" + return [ + *( + join(root, "steps", d) + for d in list_dirs(join(root, "steps")) + if d.endswith(".lance") + ), + *(join(root, n) for n in CATALOG), + changes_uri(root), + annotations_uri(root), + ] + + +def content_digest(steps: pa.Table, videos: dict[str, Path]) -> str: + h = hashlib.sha256() + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, steps.schema) as w: + w.write_table(steps.combine_chunks()) + h.update(sink.getvalue()) + for cam in sorted(videos): + h.update(cam.encode() + b"\0") + with open(videos[cam], "rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + return h.hexdigest() + + +def _exists(uri: str) -> bool: + return _nonempty(join(uri, "_versions")) + + +def _sql_str(s: str) -> str: + return "'" + s.replace("'", "''") + "'" + + +def _append(uri: str, table: pa.Table) -> None: + lance.write_dataset( + table, + uri, + mode="append" if _exists(uri) else "create", + storage_options=storage_options(uri), + ) + + +def _replace(uri: str, predicate: str | None, table: pa.Table) -> 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 + ones, so a recording's steps would no longer be in timestamp order. + """ + ds = _open(uri) + removed, updated = [], [] + for frag in ds.get_fragments() if predicate else []: + if frag.count_rows(predicate) == 0: + continue + meta = frag.delete(predicate) + (removed.append(frag.fragment_id) if meta is None else updated.append(meta)) + opts = storage_options(uri) + new = ( + write_fragments(table, uri, schema=ds.schema, storage_options=opts) + 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, storage_options=opts) + + +def _stage(root, d: Delivery) -> dict[str, pa.Table]: + """Copy the videos into place and build each table's rows for one delivery.""" + digest = content_digest(d.steps, d.videos) + paths = {} + for cam, src in d.videos.items(): + rel = f"videos/{d.recording_id}/{digest[:16]}/{cam}.mp4" + dst = join(root, rel) + if _size(dst) != Path(src).stat().st_size: + put_file(src, dst) + paths[cam] = rel + rec = {**d.recording, "video_path": sorted(paths.items()), "content_digest": digest} + return { + "steps": d.steps, + "aliases.lance": pa.Table.from_pylist([d.alias], schema=ALIASES), + "tasks.lance": pa.Table.from_pylist(d.tasks, schema=TASKS), + "recordings.lance": pa.Table.from_pylist([rec], schema=RECORDINGS), + } + + +def committed_ids(root) -> set[str]: + uri = join(root, "recordings.lance") + if not _exists(uri): + return set() + return set(_open(uri).to_table(columns=["recording_id"]).column(0).to_pylist()) + + +def append(root, deliveries: list[Delivery], lab: str = "aria") -> list[str]: + """Commit new recordings; ones already in `recordings` are skipped, so re-running is a no-op.""" + have = committed_ids(root) + todo = [d for d in deliveries if d.recording_id not in have] + if not todo: + return [] + staged = [_stage(root, d) for d in todo] + _append(steps_uri(root, lab), pa.concat_tables([s["steps"] for s in staged])) + for name in CATALOG: + _append(join(root, name), pa.concat_tables([s[name] for s in staged])) + _ensure_tables(root) + return [d.recording_id for d in todo] + + +def _ensure_tables(root) -> None: + # Exist from the first commit on, so a tag always has a version of them to pin. + for uri, schema in ( + (annotations_uri(root), ANNOTATIONS), + (changes_uri(root), CHANGES), + ): + if not _exists(uri): + _append(uri, schema.empty_table()) + changes = _open(changes_uri(root)) + if "provenance" not in changes.schema.names: + changes.add_columns({"provenance": "CAST(NULL AS STRING)"}) + + +def apply_annotations(root, adds: pa.Table, deletes: list[str]) -> 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.""" + _ensure_tables(root) + uri = annotations_uri(root) + have = _open(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) + return {"added": adds.num_rows, "deleted": len(deletes)} + + +def commit_annotations( + root, + adds: pa.Table, + deletes: list[str], + operator: str | None = None, + provenance: dict | None = None, +) -> dict | None: + """`apply_annotations`, recorded in `changes` under lab "annotations" (added: recordings + gaining rows, replaced: recordings losing rows) and tagged. None when nothing changed.""" + _ensure_tables(root) + gone = [] + if deletes: + pred = f"annotation_id IN ({', '.join(map(_sql_str, deletes))})" + rids = _open(annotations_uri(root)).to_table( + columns=["recording_id"], filter=pred + ) + gone = sorted(set(rids.column(0).to_pylist())) + counts = apply_annotations(root, adds, deletes) + if not any(counts.values()): + return None + added = sorted(set(adds.column("recording_id").to_pylist())) + return { + "rows": counts, + **_record(root, "annotations", added, gone, operator, provenance), + } + + +def drop_task_annotations(root) -> None: + """One-off for tables written while annotations were nested in tasks.""" + ds = _open(join(root, "tasks.lance")) + if "annotations" in ds.schema.names: + ds.drop_columns(["annotations"]) + + +def replace(root, d: Delivery, lab: str = "aria") -> None: + """Swap a committed recording for a reprocessed delivery; its steps rows move to the end of the table.""" + if d.recording_id not in committed_ids(root): + raise KeyError(f"recording {d.recording_id} is not committed") + staged = _stage(root, d) + pred = f"recording_id = {_sql_str(d.recording_id)}" + _replace(steps_uri(root, lab), pred, staged["steps"]) + for name in CATALOG: + _replace(join(root, name), pred, staged[name]) + + +def commit_batch( + root, + deliveries: list[Delivery], + lab: str, + replace_committed: bool = False, + operator: str | None = None, + provenance: dict | None = None, +) -> dict | None: + """Append new recordings (and replace committed ones when asked), record the batch + in `changes` and tag every table `-`. None when nothing changed.""" + added = append(root, deliveries, lab) + replaced = [] + if replace_committed: + for d in deliveries: + if d.recording_id not in added: + replace(root, d, lab) + replaced.append(d.recording_id) + if not added and not replaced: + return None + return _record(root, lab, added, replaced, operator, provenance) + + +def _record( + root, lab: str, added: list, replaced: list, operator, provenance=None +) -> dict: + now = datetime.datetime.now(datetime.timezone.utc) + name = f"{lab}-{now:%Y%m%dT%H%M%S.%f}"[:-3] + row = dict( + tag=name, + lab=lab, + committed_at=now, + operator=operator or getpass.getuser(), + added=added, + replaced=replaced, + provenance=None + if provenance is None + else json.dumps(provenance, sort_keys=True), + ) + _ensure_tables(root) + _append(changes_uri(root), pa.Table.from_pylist([row], schema=CHANGES)) + return { + "tag": name, + "added": added, + "replaced": replaced, + "versions": tag(root, name), + } + + +def edit(root, 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 and table != "annotations.lance": + # 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") + ds = _open(join(root, table)) + stats = ds.merge_insert(on).when_matched_update_all().execute(rows) + if stats["num_updated_rows"] != rows.num_rows: + raise KeyError( + f"{table}: {rows.num_rows - stats['num_updated_rows']} of {rows.num_rows} keys not found" + ) + + +def add_columns(root, table: str, columns: dict[str, str]) -> None: + """New columns from SQL expressions over existing ones; written as new data files, nothing is rewritten.""" + _open(join(root, table)).add_columns(columns) + + +def tag(root, name: str) -> dict[str, int]: + """Tag the current version of every table with `name`; returns {table: version}.""" + _ensure_tables(root) + out = {} + for uri in table_uris(root): + ds = _open(uri) + if name in ds.tags.list(): + raise ValueError(f"tag {name} already exists on {uri}") + ds.tags.create(name, ds.version) + out[uri[len(str(root).rstrip("/")) + 1 :]] = ds.version + return out + + +def index(root, lab: str = "aria") -> None: + _open(steps_uri(root, lab)).create_scalar_index("recording_id", "BTREE") + + +def compact(root) -> None: + for uri in filter(_exists, table_uris(root)): + _open(uri).optimize.compact_files() + if _open(uri).describe_indices(): + _open(uri).optimize.optimize_indices() + + +def cleanup( + root, + older_than: datetime.timedelta = datetime.timedelta(days=14), +) -> None: + """Drop untagged versions older than `older_than`, then videos no remaining version references. + + `older_than` must exceed the longest training run that resolved an untagged version. + """ + for uri in filter(_exists, table_uris(root)): + _open(uri).cleanup_old_versions( + older_than=older_than, error_if_tagged_old_versions=False + ) + gc_videos(root) + + +def gc_videos(root) -> list[str]: + """Delete MP4s under videos/ that no version of `recordings` references.""" + uri = join(root, "recordings.lance") + keep = set() + for v in _open(uri).versions(): + tbl = _open(uri, v["version"]).to_table(columns=["video_path"]) + for m in tbl.column(0).to_pylist(): + keep.update(p for _, p in m) + removed = [ + join(root, "videos", p) + for p in list_files(join(root, "videos")) + if p.endswith(".mp4") and f"videos/{p}" not in keep + ] + remove(removed) + if not _is_remote(str(root)): + vids = Path(root) / "videos" + for d in sorted(vids.rglob("*"), key=lambda p: -len(p.parts)): + if d.is_dir() and not any(d.iterdir()): + d.rmdir() + return removed diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py new file mode 100644 index 0000000000..8215ebb458 --- /dev/null +++ b/egomimic/rldb/lance/dataset.py @@ -0,0 +1,1068 @@ +"""Lance-backed episode datasets: one leaf per recording window, drop-in for ZarrDataset. + +Leaves reuse ZarrDataset's keymap semantics (horizons, history, lag, pad masks, +annotations, transforms); only storage access differs. A leaf's index runs over +eligible anchors, not every frame. Lance and torchcodec handles are opened lazily per +process: Lance is not fork-safe (https://lance.org/integrations/pytorch/), so loaders +over these datasets must use a forkserver or spawn worker context. +""" + +import concurrent.futures as cf +import copy +import functools +import hashlib +import logging +import os +from collections import OrderedDict +from pathlib import Path + +import lance +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.compute as pc +import torch + +from egomimic.rldb.filters import DatasetFilter, PinError, _ensure_dataset_filter +from egomimic.rldb.lance import mp4_ranges +from egomimic.rldb.lance.schema import ARM, task_instance_id +from egomimic.rldb.leaf import EpisodeLeaf + +logger = logging.getLogger(__name__) + +# zarr key -> Lance steps column (nested fields as dotted paths). +ZARR_TO_LANCE = { + "obs_head_pose": "head_pose", + "obs_eye_gaze": "eye_gaze", + **{ + f"{s}.{z}": f"{s}.{c}" + for s in ("left", "right") + for z, c in ( + ("obs_wrist_pose", "wrist_pose"), + ("obs_ee_pose", "ee_pose"), + ("obs_keypoints", "keypoints"), + ("obs_aria_keypoints", "aria_keypoints"), + ) + }, +} +# Robot steps columns already carry the zarr key names. +COLUMNS_OF_LAB = { + "eva": {f"{s}.{f}": f"{s}.{f}" for s in ("left", "right") for f in ARM.names} +} + + +def zarr_to_lance(lab: str) -> dict[str, str]: + return COLUMNS_OF_LAB.get(lab, ZARR_TO_LANCE) + + +# A flag missing from a steps table (robots have no head) is True everywhere. +VALIDITY = {"left": "left_valid", "right": "right_valid", "head": "head_pose_valid"} +HANDS_OF = { + "human_bimanual": ("left", "right"), + "human_left_arm": ("left",), + "human_right_arm": ("right",), + "eva_bimanual": ("left", "right"), + "eva_left_arm": ("left",), + "eva_right_arm": ("right",), +} + +_TABLES: dict[tuple[str, int], "lance.LanceDataset"] = {} +_DECODERS: "OrderedDict[str, object]" = OrderedDict() +_MAX_DECODERS = 16 + + +def _is_remote(uri: str) -> bool: + return uri.startswith("s3://") + + +def storage_options(uri: str) -> dict | None: + """R2 credentials for an s3:// URI, read from the environment in each process so + they never travel in pickled datasets, configs or cache keys.""" + if not _is_remote(uri): + return None + return { + "access_key_id": os.environ["R2_ACCESS_KEY_ID"], + "secret_access_key": os.environ["R2_SECRET_ACCESS_KEY"], + "endpoint": os.environ["R2_ENDPOINT_URL"], + "region": "auto", + } + + +def open_table(uri: str, version: int | str | None = None): + return lance.dataset(uri, version=version, storage_options=storage_options(uri)) + + +def _table(uri: str, version: int): + key = (uri, version) + if key not in _TABLES: + _TABLES[key] = open_table(uri, version) + return _TABLES[key] + + +@functools.cache +def _s3_client(): + import boto3 + + return boto3.client( + "s3", + endpoint_url=os.environ["R2_ENDPOINT_URL"], + aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"], + aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"], + region_name="auto", + ) + + +def _video_source(path: str) -> str: + """Local path as is; s3:// as a presigned HTTPS URL that FFmpeg reads with range requests.""" + if not _is_remote(path): + return path + bucket, key = path[len("s3://") :].split("/", 1) + return _s3_client().generate_presigned_url( + "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=7 * 86400 + ) + + +def _decoder(path: str): + from torchcodec.decoders import VideoDecoder + + dec = _DECODERS.pop(path, None) + if dec is None: + # Our MP4s are constant-rate encodes, so header metadata indexes frames exactly. + dec = VideoDecoder( + _video_source(path), seek_mode="approximate", dimension_order="NHWC" + ) + _DECODERS[path] = dec + while len(_DECODERS) > _MAX_DECODERS: + _DECODERS.popitem(last=False) + return dec + + +def _runs(ids: pa.ChunkedArray) -> tuple[list, np.ndarray]: + """(value, start row) of each run of equal values; a value can recur in later runs.""" + # A numpy object array of every row's id costs seconds at 50M rows; run-end encoding stays in Arrow. + values, starts, offset = [], [], 0 + for chunk in pc.run_end_encode(ids).chunks: + ends = chunk.run_ends.to_numpy() + for s, v in zip(np.r_[0, ends[:-1]] + offset, chunk.values.to_pylist()): + if not values or v != values[-1]: + values.append(v) + starts.append(int(s)) + offset += len(chunk) + return values, np.asarray(starts, dtype=np.int64) + + +REQUIRE_HANDS = ("none", "either", "both", "left", "right") + + +def _chunk_valid(v: np.ndarray, horizon: int) -> np.ndarray: + """Per anchor t: `v` holds over the whole chunk [t, min(t + horizon, n)).""" + n = len(v) + c = np.r_[0, np.cumsum(v.astype(np.int64))] + end = np.minimum(np.arange(n) + horizon, n) + return (c[end] - c[:n]) == (end - np.arange(n)) + + +def eligible_intervals( + valid: dict[str, np.ndarray], hands, horizon: int, require: str = "both" +) -> np.ndarray: + """(K, 2) half-open frame intervals of the anchors `require` admits. + + `hands` are the embodiment's hands. `require` names which of them must be valid + over the action chunk [t, min(t + horizon, n)): all ("both"), any ("either"), one + ("left" / "right"), or none ("none", which also skips the head check). Otherwise + the head pose must be valid at t, because actions are expressed in the anchor's + head frame. + """ + n = len(valid["head"]) + if require == "none": + return np.array([[0, n]] if n else [], np.int64).reshape(-1, 2) + if require in ("left", "right"): + if require not in hands: + raise ValueError( + f"require_hands_in_frame {require!r}: the embodiment has hands {tuple(hands)}" + ) + hands = (require,) + oks = [_chunk_valid(valid[h], horizon) for h in hands] + ok = valid["head"].copy() + if oks: + ok &= ( + np.logical_or.reduce(oks) + if require == "either" + else np.logical_and.reduce(oks) + ) + return _intervals(ok) + + +def _mask(spans, n: int) -> np.ndarray: + """Frames [0, n) inside any half-open (start, end) span.""" + on = np.zeros(n + 1, np.int64) + for a, b in spans: + if a < b: + on[a] += 1 + on[b] -= 1 + return np.cumsum(on)[:n] > 0 + + +def _intervals(mask: np.ndarray) -> np.ndarray: + """(K, 2) half-open runs of True in `mask`.""" + return np.flatnonzero(np.diff(np.r_[0, mask.astype(np.int8), 0])).reshape(-1, 2) + + +def _spans_of(df: pd.DataFrame | None) -> list[tuple[int, int]]: + if df is None: + return [] + return list(zip(df.t_start_ns.tolist(), df.t_end_ns.tolist())) + + +def _spans_array(spans, texts) -> tuple[np.ndarray, np.ndarray, list[str]]: + """(start frames, end frames, texts) of (start, end) frame `spans`, empty ones dropped.""" + se = np.asarray(spans, np.int64).reshape(-1, 2) + s, e = se[:, 0], se[:, 1] + keep = s < e + return s[keep], e[keep], [t for t, k in zip(texts, keep) if k] + + +class _Clock: + """Recording time to frames: frame f covers [ts[f], ts[f + 1]). `ts` is None when + every range asked of it covers the whole recording.""" + + def __init__(self, n: int, first: int, last: int, ts: np.ndarray | None): + self.n, self.first, self.last, self.ts = n, int(first), int(last), ts + + def frames(self, span: tuple[int, int] | None) -> tuple[int, int]: + """Half-open frames whose timestamps fall in the half-open `span`; None is all.""" + if span is None or (span[0] <= self.first and span[1] > self.last): + return 0, self.n + a, b = np.searchsorted(self.ts, span, "left") + return int(a), int(b) + + def time(self, frame: int) -> int: + if frame >= self.n: + return self.last + 1 + return self.first if frame == 0 else int(self.ts[frame]) + + +def _frames_of(window) -> np.ndarray: + if isinstance(window, np.ndarray): + return window.astype(np.int64) + start, end = window + return np.arange(start, start + 1 if end is None else end) + + +def _column(tbl: pa.Table, dotted: str) -> np.ndarray: + path = dotted.split(".") + col = tbl.column(path[0]).combine_chunks() + for p in path[1:]: + col = col.field(p) + arr = np.asarray(col.flatten() if pa.types.is_fixed_size_list(col.type) else col) + return arr.reshape(len(tbl), -1) + + +def _fetch_numeric(items: list) -> list[dict]: + """One `take` per steps table over every item's rows and columns.""" + out = [{} for _ in items] + groups: dict[tuple, list[int]] = {} + for j, (leaf, plan) in enumerate(items): + if plan["needs"]: + groups.setdefault((leaf.steps_uri, leaf.steps_version), []).append(j) + for (uri, version), js in groups.items(): + rows = np.unique( + np.concatenate( + [ + f + items[j][0].row_start + for j in js + for f in items[j][1]["needs"].values() + ] + ) + ) + cols = sorted( + {items[j][0].columns[zk] for j in js for zk in items[j][1]["needs"]} + ) + tbl = _table(uri, version).take(rows.tolist(), columns=cols) + arrays = {c: _column(tbl, c) for c in cols} + for j in js: + leaf, plan = items[j] + for zk, frames in plan["needs"].items(): + at = np.searchsorted(rows, frames + leaf.row_start) + out[j][zk] = arrays[leaf.columns[zk]][at] + return out + + +def read_range(url: str, start: int, end: int) -> bytes: + """Bytes [start, end) of a local file or an s3:// object.""" + if not _is_remote(url): + with open(url, "rb") as f: + f.seek(start) + return f.read(end - start) + bucket, key = url[len("s3://") :].split("/", 1) + resp = _s3_client().get_object( + Bucket=bucket, Key=key, Range=f"bytes={start}-{end - 1}" + ) + return resp["Body"].read() + + +_MP4S: "OrderedDict[str, mp4_ranges.Decoder]" = OrderedDict() +_MAX_MP4S = 1024 + + +def _mp4(url: str) -> "mp4_ranges.Decoder": + dec = _MP4S.pop(url, None) + if dec is None: + dec = mp4_ranges.Decoder( + mp4_ranges.read_index(lambda a, b: read_range(url, a, b)) + ) + _MP4S[url] = dec + while len(_MP4S) > _MAX_MP4S: + _MP4S.popitem(last=False) + return dec + + +def _fetch_range_frames(items: list) -> list[dict]: + """Frames per image key straight from the MP4s: each MP4's sample table is read once + per process, then one range GET per merged run of needed packets, all concurrent.""" + out = [{} for _ in items] + want: dict[str, list] = {} + for j, (leaf, plan) in enumerate(items): + for k, (zk, frames) in plan["images"].items(): + want.setdefault(leaf.videos[zk], []).append((j, k, frames, leaf)) + if not want: + return out + if any(map(_is_remote, want)): + # Built before the pool's threads race to build it; concurrent first calls stall for seconds. + _s3_client() + pool = _get_pool() + urls = list(want) + decs = dict(zip(urls, pool.map(_mp4, urls))) + for url, uses in want.items(): + n = uses[0][3].total_frames + if decs[url].index.num_frames != n: + raise ValueError( + f"{url}: {decs[url].index.num_frames} frames, num_frames {n}" + ) + jobs = [ + (url, rg) + for url, uses in want.items() + for rg in mp4_ranges.ranges( + decs[url].index, np.concatenate([u[2] for u in uses]) + ) + ] + fetched: dict[str, dict] = {url: {} for url in urls} + for (url, rg), data in zip(jobs, pool.map(lambda j: read_range(j[0], *j[1]), jobs)): + fetched[url][rg] = data + for url, uses in want.items(): + for j, k, frames, _ in uses: + out[j][k] = decs[url].frames(frames, fetched[url]) + return out + + +@functools.cache +def _read_pool() -> cf.ThreadPoolExecutor: + """Per-process pool for the frame fetch that overlaps a batch's steps `take`.""" + return cf.ThreadPoolExecutor(max_workers=8, thread_name_prefix="lance-frames") + + +@functools.cache +def _get_pool() -> cf.ThreadPoolExecutor: + """Per-process pool for range GETs; separate from `_read_pool`, which waits on it.""" + return cf.ThreadPoolExecutor(max_workers=16, thread_name_prefix="lance-gets") + + +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.""" + + fork_safe = False + + def __init__( + self, + *, + name: str, + steps_uri: str, + steps_version: int, + row_start: int, + num_frames: int, + anchors: np.ndarray, + embodiment: str, + fps: float, + intrinsics: np.ndarray | None, + videos: dict[str, str], + task_text: str, + key_map: dict, + clips: dict[str, tuple[int, int]] | None = None, + transform_list: list | None = None, + content_digest: str | None = None, + annotations: tuple[np.ndarray, np.ndarray, list[str]] | None = None, + clamp_actions: bool = False, + segments: np.ndarray | None = None, + action_clamp: str = "task", + history_floor: str = "recording", + hand_valid: np.ndarray | None = None, + columns: dict[str, str] | None = None, + ): + self.episode_path = name + # zarr key -> steps column; the steps table's lab decides it. + self.columns = columns or ZARR_TO_LANCE + self.content_digest = content_digest + self.steps_uri = steps_uri + self.steps_version = steps_version + self.row_start = row_start + self.total_frames = num_frames + self.embodiment = embodiment + self.metadata = { + "fps": fps, + "total_frames": num_frames, + "embodiment": embodiment, + } + self._intrinsics = intrinsics + self.videos = videos + self.task_text = task_text + self._image_keys = {k for k in videos} + self.anchors = np.asarray(anchors, dtype=np.int64).reshape(-1, 2) + self._cum = np.r_[0, np.cumsum(self.anchors[:, 1] - self.anchors[:, 0])] + # Old clip episode_hashes aliased to this recording: {hash: (start, end frame)}. + self.clips = dict(clips or {}) + # (start frames, end frames, texts) sorted by start: the tasks rows, or the + # annotations rows under the resolver's `annotations` config. + self._ann = annotations + self._clamp_actions = clamp_actions + # (S, 2) frame spans of the segments sampled, sorted by start. + self.segments = ( + np.asarray(segments, np.int64).reshape(-1, 2) + if segments is not None + else np.array([[0, num_frames]], np.int64) + ) + self.action_clamp = action_clamp + self.history_floor = history_floor + # (num_frames, 2) [left, right] tracked per frame; a hand the embodiment + # lacks is always True so it never masks anything. + self.hand_valid = ( + np.asarray(hand_valid, bool) + if hand_valid is not None + else np.ones((num_frames, 2), bool) + ) + 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 + compaction keep it; without a digest, falls back to the table version and row + position. Never on a tag name.""" + h = hashlib.sha256() + if self.content_digest: + h.update(f"digest\0{self.content_digest}\0{self.total_frames}\0".encode()) + for k in sorted(self.videos): + h.update(f"{k}\0{Path(self.videos[k]).name}\0".encode()) + else: + h.update( + f"{self.steps_uri}\0{self.steps_version}\0{self.row_start}\0{self.total_frames}\0".encode() + ) + for k in sorted(self.videos): + h.update(f"{k}\0{self.videos[k]}\0".encode()) + h.update(self.anchors.tobytes()) + h.update(f"{self.action_clamp}\0{self.history_floor}\0".encode()) + h.update(self.segments.tobytes()) + h.update(b"hands\0" + np.packbits(self.hand_valid).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 + def intrinsics(self): + return self._intrinsics + + def __len__(self) -> int: + return int(self._cum[-1]) + + 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 clip(self, episode_hash: str) -> "LanceRecordingDataset": + """This leaf sampling only the frames of one aliased clip.""" + a, b = self.clips[episode_hash] + leaf = copy.copy(self) + leaf.anchors = _intervals( + _mask(self.anchors, self.total_frames) & _mask([(a, b)], self.total_frames) + ) + leaf._cum = np.r_[0, np.cumsum(leaf.anchors[:, 1] - leaf.anchors[:, 0])] + return leaf + + def _annotation_text_for_frame(self, frame_idx: int): + if self._ann is None: + return [] + return [self._ann[2][j] for j in self._covering(frame_idx)] + + def _segment(self, frame: int) -> tuple[int, int]: + """The innermost sampled segment holding `frame` (the recording when none does).""" + s, e = self.segments[:, 0], self.segments[:, 1] + on = (s <= frame) & (frame < e) + if not on.any(): + return 0, self.total_frames + return int(s[on].max()), int(e[on].min()) + + def _history_floor(self, idx: int) -> int: + return self._segment(idx)[0] if self.history_floor == "task" else 0 + + def _hand_masks(self, idx: int) -> dict: + out = {} + horizon = self._action_horizon() + if horizon is not None: + end = self._chunk_end_idx(idx, horizon, "action_keys") + out["action_hand_mask"] = self.hand_valid[idx:end].all(0).astype(np.float32) + spec = self._proprio_history_spec() + frames = self._history_indices(idx, spec)[0] if spec is not None else idx + out["proprio_hand_valid"] = self.hand_valid[frames].astype(np.float32) + return out + + 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 key_type != "action_keys": + return end + if self.action_clamp == "task": + end = min(end, self._segment(start_idx)[1]) + if not self._clamp_actions 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}.""" + return _fetch_numeric([(self, {"needs": needs})])[0] + + def plan_sample(self, i: int) -> dict: + """What sample `i` reads: numeric frames per zarr key and frames per image key.""" + idx = self.anchor(i) + windows = self._read_windows(idx) + needs, images = {}, {} + for k, (zk, window) in windows.items(): + frames = _frames_of(window) + if zk in self._image_keys: + images[k] = (zk, frames) + continue + needs[zk] = np.unique(np.r_[needs.get(zk, np.empty(0, np.int64)), frames]) + return {"idx": idx, "windows": windows, "needs": needs, "images": images} + + def assemble_sample(self, plan: dict, fetched: tuple[dict, dict]): + idx, windows, needs = plan["idx"], plan["windows"], plan["needs"] + numeric, frames_by_key = fetched + data = {} + for k, spec in self.key_map.items(): + zk = spec["zarr_key"] + if spec.get("key_type") == "annotation_keys": + data[k] = self._annotation_text_for_frame(idx) + continue + window = windows[k][1] + frames = _frames_of(window) + single = not isinstance(window, np.ndarray) and window[1] is None + if zk in self._image_keys: + img = np.moveaxis(frames_by_key[k], -1, -3).astype( + np.float32 + ) / np.float32(255.0) + data[k] = img[0] if single else img + continue + rows = numeric[zk][np.searchsorted(needs[zk], frames)].astype(np.float32) + data[k] = self._pad_to_horizon( + rows[0] if single else rows, spec.get("horizon") + ) + return self._finish(idx, data) + + def __getitem__(self, i: int) -> dict[str, torch.Tensor]: + plan = self.plan_sample(i) + numeric = _fetch_numeric([(self, plan)])[0] + frames = { + k: _decoder(self.videos[zk]).get_frames_at(indices=f.tolist()).data.numpy() + for k, (zk, f) in plan["images"].items() + } + return self.assemble_sample(plan, (numeric, frames)) + + +class LanceStreamingDataset(LanceRecordingDataset): + """One recording read straight from object storage (or any Lance URI): frames come + from byte ranges of the faststart MP4s (`mp4_ranges`), not whole files, and + `MultiDataset.__getitems__` fetches a DataLoader batch through `fetch_planned`, so a + batch costs one `take` plus one concurrent round of range GETs.""" + + batched_fetch = True + + def __init__( + self, + *, + fetch_split: int | None = None, + overlap_reads: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + # MultiDataset.__getitems__ fetches a batch in parts of this many samples, + # concurrently, and assembles each part as it lands; None = one fetch. + self.fetch_split = fetch_split + self.overlap_reads = overlap_reads + + @staticmethod + def fetch_planned(items: list) -> list[tuple[dict, dict]]: + """Reads for many planned samples at once: one `take` per steps table and the + MP4 range GETs, issued concurrently. `items` is [(leaf, plan)]; + returns [(numeric {zarr key: rows}, frames {key: uint8 NHWC})] in order.""" + if items and items[0][0].overlap_reads: + frames = _read_pool().submit(_fetch_range_frames, items) + numeric = _fetch_numeric(items) + frames = frames.result() + else: + numeric, frames = _fetch_numeric(items), _fetch_range_frames(items) + return list(zip(numeric, frames, strict=True)) + + def __getitem__(self, i: int) -> dict[str, torch.Tensor]: + plan = self.plan_sample(i) + return self.assemble_sample(plan, self.fetch_planned([(self, plan)])[0]) + + +class LanceEpisodeResolver: + """Resolves recordings from local Lance tables into LanceRecordingDataset leaves. + + Leaves are keyed by recording_id. Filter lambdas run once per segment of a + recording: each tasks row, plus each untasked gap (`task` None) when + `include_taskless_frames`. They see the recording joined with the segment, with + SQL-style field names (`episode_hash` is the recording_id, `lab` is the site, + `task`, `task_instance_id`, ...), so existing filter configs apply. Pins and + excludes name a recording_id or an old clip episode_hash, meaning that clip's time + range. + + `action_clamp` "task" ends action chunks at the end of the anchor's segment + (repeat-last padding, flagged in `action_pad_mask`); "recording" runs them on. + `history_floor` "recording" lets history reach into earlier segments; "task" pads + and masks it at the segment start. + + `require_hands_in_frame` (none | either | both | left | right, default both) picks + the anchors: the named hands of the embodiment valid over the whole action chunk, + and the head pose valid at the anchor unless "none". It only gates sampling. Each + sample carries `action_hand_mask` ([left, right] valid over the chunk) and + `proprio_hand_valid` ([left, right] per proprio frame), which the losses, norm + stats and bounds check use to leave out an untracked hand. + + 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, + lab: str = "aria", + key_map: dict | None = None, + transform_list: list | None = None, + cameras: dict[str, str] | None = None, + versions: dict[str, int | str] | None = None, + tag: str | None = None, + annotations: dict | None = None, + include_taskless_frames: bool = False, + action_clamp: str = "task", + history_floor: str = "recording", + require_hands_in_frame: str = "both", + ): + if require_hands_in_frame not in REQUIRE_HANDS: + raise ValueError( + f"require_hands_in_frame {require_hands_in_frame!r}: one of {REQUIRE_HANDS}" + ) + if action_clamp not in ("task", "recording"): + raise ValueError(f"action_clamp {action_clamp!r}: 'task' or 'recording'") + if history_floor not in ("task", "recording"): + raise ValueError(f"history_floor {history_floor!r}: 'task' or 'recording'") + self.root = str(lance_root).rstrip("/") + self.lab = lab + self.key_map = key_map + self.transform_list = transform_list + self.cameras = cameras or {"images.front_1": "front_1"} + self.tag = tag + self.versions = dict(versions or {}) + self.include_taskless_frames = include_taskless_frames + self.action_clamp = action_clamp + self.history_floor = history_floor + self.require_hands_in_frame = require_hands_in_frame + 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}" + want = self.versions.get(name, self.tag) + if not isinstance(want, int): + want = open_table(uri, want).version + self.versions[name] = want + return uri, want, _table(uri, want) + + def manifest(self) -> dict: + """What this resolver read: {"root", "versions": {table: version}, "tags": + {table: [tags on that version]}}.""" + tags = { + name: sorted( + t + for t, info in _table(f"{self.root}/{name}", v).tags.list().items() + if info["version"] == v + ) + for name, v in self.versions.items() + } + return {"root": self.root, "versions": dict(self.versions), "tags": tags} + + def recordings(self) -> pd.DataFrame: + """One row per recording of `lab`, with SQL-style names (`lab` is the site).""" + _, _, rec = self._open("recordings.lance") + r = rec.to_table().to_pandas() + r = r[r.lab == self.lab].rename(columns={"lab": "vendor", "site": "lab"}) + return r.reset_index(drop=True) + + def _table_df(self, name: str, columns: list[str]) -> pd.DataFrame: + _, _, tbl = self._open(name) + return tbl.to_table(columns=columns).to_pandas() + + def dataset_keys(self, hashes) -> dict[str, str]: + """{pin: dataset key} for pins naming a recording_id or an aliased clip hash.""" + hashes = set(hashes) + if not hashes: + return {} + rids = set(self._table_df("recordings.lance", ["recording_id"]).recording_id) + a = self._table_df("aliases.lance", ["episode_hash", "recording_id"]) + out = dict(zip(a.episode_hash, a.recording_id)) + out.update({r: r for r in rids}) + return {h: out[h] for h in hashes if h in out} + + 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")} + + def _horizon(self) -> int: + hs = [ + int(s["horizon"]) + for s in self.key_map.values() + if s.get("key_type") == "action_keys" and s.get("horizon") + ] + return max(hs) if hs else 1 + + def _ranges(self, hashes, recs: pd.DataFrame, what: str) -> dict[str, list]: + """{recording_id: [(t_start_ns, t_end_ns) | None]} for pins or excludes; None + is the whole recording. A hash is a recording_id or an aliased clip hash.""" + out: dict[str, list] = {} + if not hashes: + return out + rids = set(recs.recording_id) + a = self._table_df( + "aliases.lance", ["episode_hash", "recording_id", "t_start_ns", "t_end_ns"] + ) + a = a[a.episode_hash.isin(hashes) & a.recording_id.isin(rids)] + for row in a.itertuples(): + out.setdefault(row.recording_id, []).append( + (int(row.t_start_ns), int(row.t_end_ns)) + ) + for h in hashes & rids: + out.setdefault(h, []).append(None) + if what == "pins": + missing = sorted(hashes - rids - set(a.episode_hash)) + if missing: + raise PinError(f"pinned episodes not in the Lance catalog: {missing}") + return out + + def resolve( + self, + filters: DatasetFilter | None = None, + expected_embodiment: str | None = None, + ): + """{recording_id: leaf}. A recording's samples are the eligible anchors inside + its segments (tasks rows, plus the untasked gaps with `include_taskless_frames`) + that pass the filter lambdas, narrowed to pinned and minus excluded ranges.""" + filters = _ensure_dataset_filter(filters) + recs = self.recordings() + if expected_embodiment is not None: + recs = recs[recs.embodiment.str.lower() == expected_embodiment.lower()] + pins = self._ranges(set(filters.episode_hashes), recs, "pins") + excludes = self._ranges(set(filters.exclude_hashes), recs, "excludes") + if filters.episode_hashes: + recs = recs[recs.recording_id.isin(set(pins))] + tasks = { + rid: g.sort_values("t_start_ns") + for rid, g in self._table_df( + "tasks.lance", + [ + "task_instance_id", + "recording_id", + "t_start_ns", + "t_end_ns", + "task_name", + "task_description", + "objects", + ], + ).groupby("recording_id") + } + clips = { + rid: g + for rid, g in self._table_df( + "aliases.lance", + ["episode_hash", "recording_id", "t_start_ns", "t_end_ns"], + ).groupby("recording_id") + } + + uri, version, steps = self._open(f"steps/{self.lab}.lance") + flags = {h: c for h, c in VALIDITY.items() if c in steps.schema.names} + cols = steps.to_table(columns=["recording_id", *flags.values()]) + n_rows = cols.num_rows + first, change = _runs(cols.column("recording_id")) + starts = dict(zip(first, change)) + if len(starts) != len(change): + raise ValueError( + f"{uri}@{version}: a recording's steps rows are not contiguous" + ) + counts = dict(zip(first, np.diff(np.r_[change, n_rows]))) + for row in recs.itertuples(): + if counts.get(row.recording_id) != row.num_frames: + raise ValueError( + f"{row.recording_id}: {counts.get(row.recording_id, 0)} steps rows at {uri}@{version}, num_frames {row.num_frames}" + ) + all_valid = { + h: ( + cols.column(flags[h]).to_numpy(zero_copy_only=False) + if h in flags + else np.ones(n_rows, bool) + ) + for h in VALIDITY + } + ann_rows = self._annotation_rows() if self.annotations is not None else {} + clocks = self._clocks( + steps, + {rid: (int(starts[rid]), int(counts[rid])) for rid in recs.recording_id}, + { + rid: [ + *_spans_of(tasks.get(rid)), + *_spans_of(clips.get(rid)), + *_spans_of(ann_rows.get(rid)), + *(r for r in pins.get(rid, []) + excludes.get(rid, []) if r), + ] + for rid in recs.recording_id + }, + ) + horizon = self._horizon() + + out, no_ann, no_anchor = {}, [], [] + for row in recs.itertuples(): + rid, clock = row.recording_id, clocks[row.recording_id] + r0, n = int(starts[rid]), int(row.num_frames) + segs = self._segments(row, tasks.get(rid), clock) + passing = [ + (a, b) + for seg, (a, b) in segs + if a < b and filters.passes_lambdas({**row._asdict(), **seg}) + ] + mask = _mask(passing, n) + if filters.episode_hashes: + mask &= _mask([clock.frames(r) for r in pins[rid]], n) + mask &= ~_mask([clock.frames(r) for r in excludes.get(rid, [])], n) + valid = {h: a[r0 : r0 + n] for h, a in all_valid.items()} + hands = HANDS_OF.get(row.embodiment, ("left", "right")) + eligible = _mask( + eligible_intervals(valid, hands, horizon, self.require_hands_in_frame), + n, + ) + anchors = _intervals(mask & eligible) + task_text = [ + t.task_description or t.task_name + for t in (tasks[rid].itertuples() if rid in tasks else ()) + ] + ann_kw = dict( + annotations=_spans_array( + [clock.frames(r) for r in _spans_of(tasks.get(rid))], task_text + ) + ) + if self.annotations is not None: + g = ann_rows.get(rid) + spans = _spans_array( + [clock.frames(r) for r in _spans_of(g)], + g.text.tolist() if g is not None else [], + ) + if self.annotations["require"]: + anchors = _intervals(_mask(anchors, n) & _mask(zip(*spans[:2]), n)) + if not len(anchors): + no_ann.append(rid) + continue + ann_kw = dict( + annotations=spans, clamp_actions=self.annotations["clamp_actions"] + ) + if not len(anchors): + no_anchor.append(rid) + continue + intr = dict(row.intrinsics) if row.intrinsics is not None else {} + vids = dict(row.video_path) + cam = next(iter(self.cameras.values())) + out[rid] = self._leaf( + row, + name=rid, + steps_uri=uri, + steps_version=version, + row_start=r0, + num_frames=n, + anchors=anchors, + embodiment=row.embodiment, + fps=float(row.fps), + intrinsics=np.asarray(intr[cam]) if cam in intr else None, + videos={zk: f"{self.root}/{vids[c]}" for zk, c in self.cameras.items()}, + task_text=task_text[0] if task_text else "", + clips={ + c.episode_hash: clock.frames((c.t_start_ns, c.t_end_ns)) + for c in (clips[rid].itertuples() if rid in clips else ()) + }, + key_map=self.key_map, + transform_list=self.transform_list, + content_digest=getattr(row, "content_digest", None), + segments=sorted(passing), + action_clamp=self.action_clamp, + history_floor=self.history_floor, + columns=zarr_to_lance(self.lab), + hand_valid=np.stack( + [ + valid[h] if h in hands else np.ones(n, bool) + for h in ("left", "right") + ], + axis=1, + ), + **ann_kw, + ) + untagged = sorted(n for n, t in self.manifest()["tags"].items() if not t) + if untagged: + logger.warning( + "read untagged versions of %s under %s (%s): this run can be " + "reproduced only until cleanup_old_versions drops them; tag them " + "(commit.tag) to keep them", + untagged, + self.root, + {n: self.versions[n] for n in untagged}, + ) + if no_anchor: + logger.warning( + "%d recordings have no sampleable anchors, e.g. %s", + len(no_anchor), + no_anchor[:5], + ) + 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 _segments(self, row, tasks: pd.DataFrame | None, clock: "_Clock") -> list: + """[(filter fields, (frame start, frame end))]: one per tasks row, plus one per + untasked gap when `include_taskless_frames`.""" + segs, covered = [], np.zeros(clock.n, bool) + for t in tasks.itertuples() if tasks is not None else (): + a, b = clock.frames((t.t_start_ns, t.t_end_ns)) + covered[a:b] = True + seg = dict( + task_instance_id=t.task_instance_id, + t_start_ns=t.t_start_ns, + t_end_ns=t.t_end_ns, + task=t.task_name, + task_description=t.task_description, + objects=t.objects, + ) + segs.append((seg, (a, b))) + if self.include_taskless_frames: + for a, b in _intervals(~covered): + t0 = clock.time(a) + seg = dict( + task_instance_id=task_instance_id(row.recording_id, None, t0), + t_start_ns=t0, + t_end_ns=clock.time(b), + task=None, + task_description=None, + objects=[], + ) + segs.append((seg, (int(a), int(b)))) + for seg, _ in segs: + seg["episode_hash"] = row.recording_id + return segs + + def _clocks(self, steps, rows: dict, ranges: dict) -> dict[str, "_Clock"]: + """{recording_id: _Clock}. Reads every timestamp only of recordings with a range + that starts or ends inside them; the rest map time to frames from their ends.""" + rids = list(rows) + ends = ( + steps.take( + [ + r + for rid in rids + for r in (rows[rid][0], rows[rid][0] + rows[rid][1] - 1) + ], + columns=["timestamp_ns"], + ) + .column(0) + .to_numpy() + .reshape(-1, 2) + ) + bounds = dict(zip(rids, ends)) + need = [ + rid + for rid in rids + if any( + not (t0 <= bounds[rid][0] and t1 > bounds[rid][1]) + for t0, t1 in ranges[rid] + ) + ] + ts = {} + if need: + take = np.concatenate([np.arange(rows[r][0], sum(rows[r])) for r in need]) + col = ( + steps.take(take.tolist(), columns=["timestamp_ns"]).column(0).to_numpy() + ) + at = np.cumsum([0] + [rows[r][1] for r in need]) + ts = {r: col[at[j] : at[j + 1]] for j, r in enumerate(need)} + return {rid: _Clock(rows[rid][1], *bounds[rid], ts.get(rid)) for rid in rids} + + def _leaf(self, row, **kwargs) -> LanceRecordingDataset: + return LanceRecordingDataset(**kwargs) + + +class LanceStreamingResolver(LanceEpisodeResolver): + """Resolves recordings into LanceStreamingDataset leaves, for training straight from + object storage (`lance_root` may be s3://).""" + + def __init__( + self, + lance_root: str | Path, + fetch_split: int | None = None, + overlap_reads: bool = True, + **kwargs, + ): + super().__init__(lance_root, **kwargs) + self.fetch_split = fetch_split + self.overlap_reads = overlap_reads + + def _leaf(self, row, **kwargs) -> LanceStreamingDataset: + return LanceStreamingDataset( + fetch_split=self.fetch_split, overlap_reads=self.overlap_reads, **kwargs + ) diff --git a/egomimic/rldb/lance/mp4_ranges.py b/egomimic/rldb/lance/mp4_ranges.py new file mode 100644 index 0000000000..7b0e8ee20b --- /dev/null +++ b/egomimic/rldb/lance/mp4_ranges.py @@ -0,0 +1,200 @@ +"""Frames from an MP4 by byte range: read its sample table once, then fetch only the +packets a frame needs (its keyframe through itself) and decode them in memory. + +Built for the delivery-spec encodes (one AV1 video track, GOP 2, `+faststart`), where a +frame and its keyframe are adjacent packets, so one range read serves a frame. +""" + +import dataclasses +import io +import struct +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np + +_CONTAINERS = {b"moov", b"trak", b"mdia", b"minf", b"stbl", b"edts"} + + +def _boxes(buf: bytes, start: int = 0, end: int | None = None): + """(type, payload start, box end) of each box in buf[start:end].""" + end = len(buf) if end is None else end + pos = start + while pos + 8 <= end: + size, typ = struct.unpack_from(">I4s", buf, pos) + head = 8 + if size == 1: + size = struct.unpack_from(">Q", buf, pos + 8)[0] + head = 16 + elif size == 0: + size = end - pos + yield typ, pos + head, pos + size + pos += size + + +def _find(buf: bytes, start: int, end: int, path: list[bytes]): + for typ, a, b in _boxes(buf, start, end): + if typ == path[0]: + return (a, b) if len(path) == 1 else _find(buf, a, b, path[1:]) + return None + + +def _full(buf: bytes, a: int, fmt: str): + """Fields after a full box's version/flags.""" + return struct.unpack_from(fmt, buf, a + 4) + + +@dataclass +class Mp4Index: + offsets: np.ndarray # (n,) byte offset of each frame's packet + sizes: np.ndarray # (n,) packet size + keyframe_of: np.ndarray # (n,) the sync sample each frame decodes from + header: bytes # the file up to and including moov, enough to build a decoder + + @property + def num_frames(self) -> int: + return len(self.offsets) + + +def _video_stbl(moov: bytes) -> tuple[int, int]: + for typ, a, b in _boxes(moov): + if typ != b"trak": + continue + hdlr = _find(moov, a, b, [b"mdia", b"hdlr"]) + if hdlr and moov[hdlr[0] + 8 : hdlr[0] + 12] == b"vide": + return _find(moov, a, b, [b"mdia", b"minf", b"stbl"]) + raise ValueError("no video track") + + +def parse_index(header: bytes, moov_at: tuple[int, int]) -> Mp4Index: + """`moov_at` is the moov payload's half-open byte range in `header`.""" + a, b = moov_at + moov = header[a:b] + s0, s1 = _video_stbl(moov) + + def box(name): + return _find(moov, s0, s1, [name]) + + stsz = box(b"stsz") + fixed, n = _full(moov, stsz[0], ">II") + sizes = ( + np.full(n, fixed, np.int64) + if fixed + else np.frombuffer(moov, ">u4", n, stsz[0] + 12).astype(np.int64) + ) + if (co := box(b"stco")) is not None: + (k,) = _full(moov, co[0], ">I") + chunk_off = np.frombuffer(moov, ">u4", k, co[0] + 8).astype(np.int64) + else: + co = box(b"co64") + (k,) = _full(moov, co[0], ">I") + chunk_off = np.frombuffer(moov, ">u8", k, co[0] + 8).astype(np.int64) + stsc = box(b"stsc") + (m,) = _full(moov, stsc[0], ">I") + runs = np.frombuffer(moov, ">u4", 3 * m, stsc[0] + 8).reshape(m, 3).astype(np.int64) + per_chunk = np.empty(k, np.int64) + for j, (first, spc, _) in enumerate(runs): + last = runs[j + 1, 0] - 1 if j + 1 < m else k + per_chunk[first - 1 : last] = spc + chunk_of = np.repeat(np.arange(k), per_chunk)[:n] + first_in_chunk = np.r_[0, np.cumsum(per_chunk)[:-1]][chunk_of] + within = np.cumsum(sizes) - sizes + offsets = chunk_off[chunk_of] + within - within[first_in_chunk] + keys = np.arange(n) + if (ss := box(b"stss")) is not None: + (c,) = _full(moov, ss[0], ">I") + keys = np.frombuffer(moov, ">u4", c, ss[0] + 8).astype(np.int64) - 1 + keyframe_of = keys[np.searchsorted(keys, np.arange(n), side="right") - 1] + return Mp4Index(offsets, sizes, keyframe_of, header[:b]) + + +def read_index(read: Callable[[int, int], bytes], first: int = 1 << 16) -> Mp4Index: + """The index of an MP4 through `read(start, end)` (half-open byte range). A faststart + file costs one read (two if its moov exceeds `first` bytes); a moov after the media + costs two more, and beyond the first read the media is never fetched.""" + buf, base = read(0, first), 0 # buf holds file bytes [base, base + len(buf)) + head = b"" # boxes before the media, kept so the header parses as a file + pos = 0 + while True: + if pos + 16 > base + len(buf): + buf, base = read(pos, pos + max(16, first)), pos + size, typ = struct.unpack_from(">I4s", buf, pos - base) + hlen = 8 + if size == 1: + size, hlen = struct.unpack_from(">Q", buf, pos - base + 8)[0], 16 + if typ == b"moov": + if pos + size > base + len(buf): + buf += read(base + len(buf), pos + size) + moov = buf[pos - base : pos - base + size] + pre = head if base else buf[: pos - base] + return parse_index(pre + moov, (len(pre) + hlen, len(pre) + size)) + if typ == b"mdat" and not head: + head = buf[: pos - base] if base == 0 else b"" + pos += size + + +def ranges(index: Mp4Index, frames, gap: int = 0) -> list[tuple[int, int]]: + """Merged half-open byte ranges covering `frames` and the packets they decode from; + ranges closer than `gap` bytes are joined.""" + frames = np.unique(np.asarray(frames, np.int64)) + lo = index.offsets[index.keyframe_of[frames]] + hi = index.offsets[frames] + index.sizes[frames] + order = np.argsort(lo) + out: list[list[int]] = [] + for a, b in zip(lo[order].tolist(), hi[order].tolist()): + if out and a <= out[-1][1] + gap: + out[-1][1] = max(out[-1][1], b) + else: + out.append([a, b]) + return [(a, b) for a, b in out] + + +def _slice(fetched: dict[tuple[int, int], bytes], start: int, size: int) -> bytes: + for (a, b), data in fetched.items(): + if a <= start and start + size <= b: + return data[start - a : start - a + size] + raise KeyError(f"bytes {start}+{size} not fetched") + + +class Decoder: + """An AV1 decoder configured from an index's header; decodes frames from fetched bytes.""" + + def __init__(self, index: Mp4Index): + import av + import av.logging + + # FFmpeg logs "partial file" for a header without its media; expected here. + with av.logging.Capture(), av.open(io.BytesIO(index.header)) as c: + self._name = c.streams.video[0].codec_context.name + self._extradata = c.streams.video[0].codec_context.extradata + self.index = dataclasses.replace(index, header=b"") + + def _context(self): + import av + + ctx = av.CodecContext.create(self._name, "r") + ctx.extradata = self._extradata + # One context per frame: dav1d's thread pool costs more to start than it saves. + ctx.thread_count = 1 + return ctx + + def frames(self, frames, fetched: dict[tuple[int, int], bytes]) -> np.ndarray: + """(len(frames), H, W, 3) uint8 RGB, in the order asked.""" + import av + + idx = self.index + want = np.asarray(frames, np.int64) + got: dict[int, np.ndarray] = {} + for f in np.unique(want).tolist(): + if f in got: + continue + k = int(idx.keyframe_of[f]) + ctx = self._context() + shown = [] + for p in range(k, f + 1): + pkt = av.Packet(_slice(fetched, int(idx.offsets[p]), int(idx.sizes[p]))) + shown += ctx.decode(pkt) + shown += ctx.decode(None) + for j, fr in enumerate(shown): + got.setdefault(k + j, fr.to_ndarray(format="rgb24")) + return np.stack([got[f] for f in want.tolist()]) diff --git a/egomimic/rldb/lance/schema.py b/egomimic/rldb/lance/schema.py new file mode 100644 index 0000000000..ec4b3294d9 --- /dev/null +++ b/egomimic/rldb/lance/schema.py @@ -0,0 +1,179 @@ +"""Arrow schemas for the Lance episode tables (see docs/episode_representation_implementation.md).""" + +import uuid + +import pyarrow as pa + +RECORDING_NAMESPACE = uuid.UUID("6f1d4c3e-2b1a-5e7f-9c8d-0a1b2c3d4e5f") + + +def recording_id(lab: str, vendor_recording_id: str) -> str: + return str(uuid.uuid5(RECORDING_NAMESPACE, f"{lab}/{vendor_recording_id}")) + + +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) + + +HAND = pa.struct( + [ + ("wrist_pose", _vec(7)), + ("ee_pose", _vec(7)), + ("keypoints", _vec(63)), + # Pilot extras: raw Aria landmarks (pre-MANO) and MPS confidence (-1 = untracked). + ("aria_keypoints", _vec(63)), + ("confidence", pa.float32()), + ] +) + +HUMAN_STEPS = pa.schema( + [ + ("recording_id", pa.string()), + ("timestamp_ns", pa.int64()), + ("head_pose", _vec(7)), + ("left", HAND), + ("right", HAND), + ("left_valid", pa.bool_()), + ("right_valid", pa.bool_()), + ("head_pose_valid", pa.bool_()), + ("eye_gaze", _vec(3)), + ] +) + +# Robot arm, robot base frame: poses xyz (m) + quaternion wxyz, joints in rad, gripper in [0, 1]. +ARM = pa.struct( + [ + ("obs_ee_pose", _vec(7)), + ("obs_joints", _vec(6)), + ("obs_gripper", pa.float32()), + ("cmd_ee_pose", _vec(7)), + ("cmd_joints", _vec(6)), + ("cmd_gripper", pa.float32()), + ] +) + +# An arm the robot lacks is zeros with its flag False (Arrow's parquet reader rejects +# null fixed-size lists); an all-zero pose frame is forward-filled and flagged False. +EVA_STEPS = pa.schema( + [ + ("recording_id", pa.string()), + ("timestamp_ns", pa.int64()), + ("left", ARM), + ("right", ARM), + ("left_valid", pa.bool_()), + ("right_valid", pa.bool_()), + ] +) + +RECORDINGS = pa.schema( + [ + ("recording_id", pa.string()), + ("vendor_recording_id", pa.string()), + ("lab", pa.string()), + # Pilot extra: SQL `lab` is the collecting site for Aria, not the vendor. + ("site", pa.string()), + ("embodiment", pa.string()), + ("rig_name", pa.string()), + ("operator", pa.string()), + ("scene", pa.string()), + ("captured_at", pa.timestamp("ms", tz="UTC")), + ("fps", pa.float32()), + ("num_frames", pa.int64()), + ("video_path", pa.map_(pa.string(), pa.string())), + ("intrinsics", pa.map_(pa.string(), pa.list_(pa.float32()))), + ("extrinsics", pa.map_(pa.string(), pa.list_(pa.float32()))), + ("license", pa.string()), + # sha256 over the recording's steps rows and videos; cache keys use it instead of the table version. + ("content_digest", pa.string()), + ] +) + +TASKS = pa.schema( + [ + ("task_instance_id", pa.string()), + ("recording_id", pa.string()), + ("t_start_ns", pa.int64()), + ("t_end_ns", pa.int64()), + ("task_name", pa.string()), + ("task_description", pa.string()), + ("objects", pa.list_(pa.string())), + ] +) + +ALIASES = pa.schema( + [ + ("episode_hash", pa.string()), + ("recording_id", pa.string()), + ("t_start_ns", pa.int64()), + ("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()), + ] +) + +# One row per committer batch; `tag` names the version of every table the batch produced. +# `lab` is "annotations" for an annotation commit. `provenance` is JSON: {recording_id: +# what the converter recorded about it} for a delivery batch, the approver and source file +# for an annotation commit; null on rows written before the column existed. +CHANGES = pa.schema( + [ + ("tag", pa.string()), + ("lab", pa.string()), + ("committed_at", pa.timestamp("ms", tz="UTC")), + ("operator", pa.string()), + ("added", pa.list_(pa.string())), + ("replaced", pa.list_(pa.string())), + ("provenance", pa.string()), + ] +) diff --git a/egomimic/rldb/lance/timeline.py b/egomimic/rldb/lance/timeline.py new file mode 100644 index 0000000000..cbe3633dec --- /dev/null +++ b/egomimic/rldb/lance/timeline.py @@ -0,0 +1,62 @@ +"""Recording timeline and per-frame validity helpers shared by converters and the reader.""" + +import numpy as np + +# 30 cm in one frame is 9 m/s at 30 fps: a tracking jump, not motion (as scan_egomotion.py). +JUMP_M = 0.3 +FROZEN_RUN = 15 + + +def build_timeline(ts: np.ndarray, fps: float = 30.0): + """Constant-rate timeline over the RGB frames: dropped frames become filler rows. + + Returns (timestamps, source frame index per row, filler mask, first-row-after-gap mask). + """ + period = 1e9 / fps + dt = np.diff(ts) + missing = np.maximum(np.rint(dt / period).astype(np.int64) - 1, 0) + n = len(ts) + int(missing.sum()) + out_ts = np.empty(n, dtype=np.int64) + src = np.empty(n, dtype=np.int64) + filler = np.zeros(n, dtype=bool) + after_gap = np.zeros(n, dtype=bool) + r = 0 + for i in range(len(ts)): + out_ts[r], src[r] = ts[i], i + if i > 0 and missing[i - 1] > 0: + after_gap[r] = True + r += 1 + if i < len(missing) and missing[i] > 0: + k = missing[i] + out_ts[r : r + k] = ts[i] + np.rint(np.arange(1, k + 1) * period).astype( + np.int64 + ) + src[r : r + k] = i + filler[r : r + k] = True + r += k + return out_ts, src, filler, after_gap + + +def nearest_dt(query: np.ndarray, ref: np.ndarray) -> np.ndarray: + """|query - nearest(ref)| in ns; ref sorted.""" + j = np.clip(np.searchsorted(ref, query), 1, len(ref) - 1) + return np.minimum(np.abs(query - ref[j - 1]), np.abs(ref[j] - query)) + + +def pose_outliers(xyz: np.ndarray, ok: np.ndarray) -> np.ndarray: + """Rows that jump more than JUMP_M from the previous valid row, or sit in a frozen run.""" + bad = np.zeros(len(xyz), dtype=bool) + idx = np.flatnonzero(ok) + if len(idx) < 2: + return bad + step = np.linalg.norm(np.diff(xyz[idx], axis=0), axis=1) + bad[idx[1:]] = step > JUMP_M + # Hand results can repeat for a frame or two when MPS runs below the RGB rate, so + # only a long run of identical positions counts as frozen. + run_start = 0 + for j in range(1, len(idx) + 1): + if j == len(idx) or step[j - 1] != 0: + if j - run_start >= FROZEN_RUN: + bad[idx[run_start:j]] = True + run_start = j + return bad diff --git a/egomimic/rldb/leaf.py b/egomimic/rldb/leaf.py new file mode 100644 index 0000000000..e43b84aba9 --- /dev/null +++ b/egomimic/rldb/leaf.py @@ -0,0 +1,283 @@ +"""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 _hand_masks(self, idx: int) -> dict: + """Per-hand validity keys for sample ``idx`` (``action_hand_mask``, + ``proprio_hand_valid``); none when the store tracks no per-hand validity.""" + return {} + + def _proprio_history_spec(self) -> dict | None: + """The first non-camera keymap entry with ``history`` > 1, else None.""" + for spec in self.key_map.values(): + if spec.get("key_type") == "camera_keys": + continue + if int(spec.get("history") or 1) > 1: + return spec + return None + + def _proprio_history_mask(self, idx: int) -> np.ndarray | None: + """1.0 per real history frame, 0.0 per front-padded one, ``(K,)``; + ``None`` when the keymap has no proprio history key.""" + spec = self._proprio_history_spec() + if spec is None: + return None + return self._history_indices(idx, spec)[1] + + 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 _history_floor(self, idx: int) -> int: + """First frame history and lagged reads for ``idx`` may reach back to.""" + return 0 + + def _lagged_idx(self, idx: int, lag_s: float) -> int: + """Frame ``lag_s`` seconds before ``idx`` (at least one frame back), + clamped at the history floor -- so the first samples read the floor + frame, the same front padding a history window gets.""" + return max(self._history_floor(idx), 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 history floor (the episode start) are masked 0 + and repeat the oldest real step (the floor frame 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) + floor = self._history_floor(idx) + real = steps >= floor + fill = steps[real][0] if real.any() else floor + 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 _pad_sequences(self, data, horizon: int | None) -> dict: + """``_pad_to_horizon`` over a read result, in place. Keys are zarr keys.""" + if horizon is None: + return data + for k in data: + data[k] = self._pad_to_horizon(data[k], horizon) + return data + + 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] + + data.update(self._hand_masks(idx)) + + history_mask = self._proprio_history_mask(idx) + if history_mask is not None: + # Informational: the front padding already looks exactly like a + # history-dropout sample, so the stem needs nothing extra. + data["proprio_history_mask"] = history_mask + + 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/mcap/SPEC.md b/egomimic/rldb/mcap/SPEC.md new file mode 100644 index 0000000000..5f60e1cd3d --- /dev/null +++ b/egomimic/rldb/mcap/SPEC.md @@ -0,0 +1,295 @@ +# EgoVerse vendor MCAP spec 0.1 + +What a data vendor delivers to EgoVerse: one raw, self-contained MCAP recording per +episode, from which we build the Lance tables (`recordings`, `tasks`, `annotations`, +`steps`) and the training videos, or any other format later, without asking the vendor +anything. + +- **Raw, not processed.** Original sensor data at its own rates: no resampling, no + aligning streams to one clock rate, no undistorting or cropping video. We convert on + our side. +- **Self-contained.** Calibration, camera positions and episode info live inside the + file. No side JSONs, no constants in our converters. +- **Checked on arrival.** Every delivered episode is run through the validator + (`python -m egomimic.rldb.mcap.validate`), which lists what is missing or wrong. + +The machine-readable version of this document is [`spec.py`](spec.py); the message +definitions are [`proto/egoverse.proto`](proto/egoverse.proto) (generated from +[`schemas.py`](schemas.py)). The three change together. + +Spec 0.1 covers **human (egocentric) recordings** and **Eva robot recordings with +grippers**. Dexterous hands, tactile, humanoid / auxiliary joints (e.g. Sharpa) are not +in 0.1. + +## 1. What arrives per episode + +``` +episode_/ +├── episode.mcap # everything: video, sensors, poses, calibration, episode info +└── annotation.mcap # optional: /task spans and /subtask-annotation labels +``` + +- **An episode is one whole recording session, never a pre-cut clip.** It becomes one + row of `recordings`; the tasks performed in it become `tasks` rows. +- The label topics may instead be inline in `episode.mcap` (as ABC-130k does), but not in + both files. `annotation.mcap` holds nothing but label topics: no metadata records + (the episode info lives in `episode.mcap` only). +- MCAP with protobuf encoding; any chunk compression (zstd, lz4) is fine. + +## 2. Clock + +- **Every message's MCAP log time is its capture time in absolute Unix nanoseconds**, in + both files. Devices that stamp on their own clock (e.g. Aria device time) convert + with their device-to-UTC offset before writing. +- Within a topic, log times never decrease; each camera's strictly increase (every + frame has its own time). +- `start-time-unix` / `end-time-unix` in the episode info are **milliseconds** and must + match the recorded data within 1 s. +- **Steps clock.** The main camera (`/main-camera`, human) or top camera (`/top-camera`, + robot) defines the Lance `steps` rows: one row per decoded frame. Every other stream is + aligned to those frame times during conversion; wrist and SLAM cameras stay separate + videos and add no rows. + +## 3. Who, what, where: episode info + +One MCAP metadata record named **`session-metadata`**. All values are strings. + +| Key | Example | Meaning | +|---|---|---| +| `lab` | `aria` | Data source / capture family; the Lance partition (`steps/.lance`). Slug. | +| `site` | `rl2` | Organization that collected the recording. Slug. | +| `vendor` | `acme` | Organization or pipeline that delivered / processed it. Slug. | +| `vendor-recording-id` | `acme-rec-0001` | The vendor's own ID for this recording. **Stable across every re-export or correction** (new labels, fixed calibration, re-run poses). | +| `operator-id` | `9f2c…e41a` | Who performed it, as a SHA-256 hash (64 lowercase hex digits). Never a name or email. | +| `task` | `fold_clothes` | Task slug; reuse an existing one if it fits. | +| `instruction` | `fold the 2T shirt on the blue table` | Full task text. | +| `objects` | `baby_shirt_2t,basket` | Objects involved, comma-separated slugs. | +| `embodiment` | `human_bimanual` | One of `human_{bimanual,left_arm,right_arm}`, `eva_{bimanual,left_arm,right_arm}`. Must match the topics in the file. | +| `rig-name` | `aria_gen1` | Capture setup. | +| `device-serial` | `1WM103600M1292` | Physical unit; ties the episode to its calibration. | +| `scene` | `kitchen_a` | Where it was recorded. | +| `start-time-unix`, `end-time-unix` | `1736651249664` | Recording bounds, Unix **ms**. | +| `spec-version` | `0.1` | Version of this spec the file follows. | +| `license` | `CC-BY-4.0` | License the data is under. | +| `camera-timestamp-ref` | `mid` | *Human only.* Where in the exposure each frame is stamped: `start`, `mid` or `end`. | +| `shutter` | `rolling:5ms` | *Human only.* `global`, or `rolling:ms`. | +| `hand-joints` | see §6 | *Human only.* Names of the hand joints, in the order `HandKeypoints.positions` lists them. | + +Slugs are lowercase letters and digits, words joined by `_` or `-`. + +**IDs.** Our recording ID is +`egomimic.rldb.lance.schema.recording_id(lab, vendor-recording-id)` (UUIDv5 of +`/`); the validator prints it. Task and annotation IDs are +computed downstream from their content, so vendors send none. + +**Tasks.** Without `/task` topics, `task` / `instruction` / `objects` describe one task +spanning the whole recording: [first steps frame, last steps frame + one frame), as +every span is half-open. With `/task`, the spans are authoritative and these keys +only summarize the session (e.g. `task=laundry` over iron / fold / pack spans); no extra +task is made from them. + +## 4. Topics + +Message types are Foxglove's (`foxglove-schemas-protobuf`) where one exists, otherwise +`egoverse.*` from `proto/egoverse.proto`. Topic names, message types and required fields +are checked exactly. + +### Human / egocentric + +| Topic | Type | | Rate | Notes | +|---|---|---|---|---| +| `/main-camera` | `foxglove.CompressedVideo` | MUST | ≥ 30 fps | Head-mounted RGB; the steps clock. | +| `/main-camera-info` | `foxglove.CameraCalibration` | MUST | once | See §7. | +| `/slam-left-camera`, `/slam-right-camera` | `foxglove.CompressedVideo` | optional | ≥ 30 fps | Each with its `-info`. | +| `/left-wrist-camera`, `/right-wrist-camera` | `foxglove.CompressedVideo` | optional | ≥ 30 fps | Each with its `-info`. | +| `/imu` | `egoverse.Imu` | MUST | ≥ 200 Hz | Raw: gravity included, no fusion. | +| `/imu-info` | `egoverse.ImuInfo` | MUST | once | Noise model; datasheet values are fine. | +| `/head-pose` | `foxglove.PoseInFrame` | MUST | ≥ 30 Hz | `frame_id = "world"`, pose = T_world_head. | +| `/head-pose-quality` | `egoverse.TrackingQuality` | MUST | one per head pose | `OK`, `LOST`, `RELOCALIZED` (first pose after a SLAM reset) or `UNKNOWN`. | +| `/{left,right}-hand-keypoints` | `egoverse.HandKeypoints` | MUST per arm in use | tracker rate | See §6. | +| `/eye-gaze` | `egoverse.EyeGaze` | optional | tracker rate | Yaw / pitch / depth in a frame between the eyes (+x left, +y up, +z forward). | +| `/tf_static` | `foxglove.FrameTransforms` | MUST | once | See §5. | +| `/instruction` | `egoverse.Instruction` | MUST | once | Task text. | + +No arm, gripper or action topics: actions are derived from future hand and head poses. + +### Robot (Eva with grippers) + +`{side}` is `left` and/or `right`, as the embodiment uses. + +| Topic | Type | | Notes | +|---|---|---|---| +| `/top-camera` | `foxglove.CompressedVideo` | MUST | ≥ 30 fps; the steps clock. With `-info`. | +| `/{side}-wrist-camera` | `foxglove.CompressedVideo` | MUST | ≥ 30 fps. With `-info`. | +| `/{side}-arm-state` | `egoverse.ArmState` | MUST | Measured joints (names, position, velocity, torque) and `ee_pose`. | +| `/{side}-arm-action` | `egoverse.ArmAction` | MUST | The command sent at that time: target joints and `ee_pose`. | +| `/{side}-ee-state` | `egoverse.GripperState` | MUST | Measured opening, 0 = closed, 1 = fully open. | +| `/{side}-ee-action` | `egoverse.GripperState` | MUST | Commanded opening, same range. | +| `/tf_static` | `foxglove.FrameTransforms` | MUST | See §5. | +| `/instruction` | `egoverse.Instruction` | MUST | Task text. | + +`ArmState` / `ArmAction`: `frame_id` is the arm-base frame (`left_arm_base`, +`right_arm_base`); every joint array has one entry per `joint_names` entry, in that +order; joints in rad (m for prismatic), velocities rad/s, torques N·m; `ee_pose` is +T_base_ee as 16 numbers, a **row-major** 4×4. ABC-130k's own `RobotState` / +`GripperState` are converted to these types before validation. + +### Both profiles: tasks and labels (optional) + +| Topic | Type | Notes | +|---|---|---| +| `/task` | `egoverse.Task` | One message per task performed: `name` (slug), `description`, `objects`; starts at its log time, ends at `end_time_unix_ns` (after the start). | +| `/subtask-annotation` | `egoverse.Annotation` | A language label `data` over [log time, `end_time_unix_ns`); `end_time_unix_ns = 0` is a point event (one frame wide). `annotation_type` is `language` (the default when empty); `annotation_vendor` defaults to `vendor`. | + +Every label lies **inside one task**; without `/task`, inside the whole-recording task +made from the episode info. All spans lie inside the recorded data. Labels added or corrected after delivery go through +the Lance annotation review (`egomimic/rldb/lance/annotations.py`), not a re-export. + +## 5. Frames and conventions + +- Right-handed frames; metres, radians, seconds. +- Quaternions are Foxglove's named `x, y, z, w` fields and must be unit length. +- A transform or pose is **T_parent_child**: the child frame's pose in the parent + (Foxglove's `FrameTransform` / `PoseInFrame` convention). +- Every camera frame is a Foxglove optical frame: +x right, +y down, +z forward, in the + raw sensor orientation (see `upright_rotation_cw_deg`, §7). A camera's video, + `-info` and `-info-extra` all name that one frame. +- **Every frame any topic names must reach `world`** through `/tf_static` plus the + moving links the profile defines: + - **Human:** `world → head` moves via `/head-pose`. `/tf_static` links `head` to each + camera, the IMU, the gaze frame, and any frame hands are given in. The main camera + must be rigidly mounted on the head: a `/tf_static` path from `head` to its frame. + - **Robot:** `world → {side}_arm_base` and the top camera in `/tf_static`; + `{side}_arm_base → {side}_ee` moves via `/{side}-arm-state`; the wrist camera hangs + off `{side}_ee` in `/tf_static`. + +## 6. Hands + +`HandKeypoints`, one message per hand per tracker frame: + +| Field | | +|---|---| +| `frame_id` | Frame `positions` and `wrist_pose` are in (linked to `world`). | +| `positions` | `(x, y, z)` in m per `hand-joints` entry, flattened; or empty for devices without keypoints. | +| `confidence` | Whole-hand tracking confidence in [0, 1]. | +| `joint_confidence` | Optional: one value in [0, 1] per joint. | +| `wrist_pose` | Optional: T_frame_wrist, in the Aria MPS wrist convention below. Required when `positions` is empty. | +| `valid` | `false` when the hand is not tracked. **Never write zeros instead.** | + +**Wrist frame** (Aria MPS's, measured on real recordings): origin at the wrist; +y out of +the back of the hand (the palm faces −y); +z from the pinky side toward the index side; ++x = y × z, which points toward the fingers on the left hand and toward the forearm on +the right. + +**`hand-joints`** names the joints the device tracks, in its own order, from this +vocabulary: the 21 MANO joints `wrist`, `{thumb}_{cmc,mcp,ip,tip}`, +`{index,middle,ring,pinky}_{mcp,pip,dip,tip}`, plus `palm_center`. Each name at most +once; `wrist` is required. Joints with MANO names are what our MANO fit matches; extras +are kept but not fitted. A device that already outputs MANO sends the canonical order: + +``` +wrist,thumb_cmc,thumb_mcp,thumb_ip,thumb_tip,index_mcp,index_pip,index_dip,index_tip, +middle_mcp,middle_pip,middle_dip,middle_tip,ring_mcp,ring_pip,ring_dip,ring_tip, +pinky_mcp,pinky_pip,pinky_dip,pinky_tip +``` + +and a device with its own layout lists it as is, e.g. Aria MPS's 21 landmarks (no thumb +CMC, a palm center instead): + +``` +thumb_tip,index_tip,middle_tip,ring_tip,pinky_tip,wrist,thumb_mcp,thumb_ip, +index_mcp,index_pip,index_dip,middle_mcp,middle_pip,middle_dip,ring_mcp,ring_pip, +ring_dip,pinky_mcp,pinky_pip,pinky_dip,palm_center +``` + +## 7. Cameras + +- **Video.** H.264 or H.265, Annex B byte stream, `CompressedVideo.format` `h264` / + `h265`. **Each message is exactly one frame**, stamped with that frame's capture time; + the first message is a keyframe and carries SPS/PPS (repeat them at every keyframe). + Original resolution and orientation: no resizing, cropping, undistorting or rotating. +- **Calibration** `-info`: exactly one `foxglove.CameraCalibration` message, + finite values, principal point inside the image: + `width` / `height` equal the decoded frames; `K` has fx, fy > 0. Name the lens model in + `distortion_model` and give exactly its number of coefficients in `D` (an empty model + with no coefficients is a pinhole camera): + + | `distortion_model` | `K` | `D` | + |---|---|---| + | `plumb_bob` | fx, fy, cx, cy | k1, k2, p1, p2, k3 | + | `rational_polynomial` | fx, fy, cx, cy | k1, k2, p1, p2, k3, k4, k5, k6 | + | `kannala_brandt` | fx, fy, cx, cy | k1, k2, k3, k4 | + | `fisheye624` (Aria) | f, f, cx, cy | the 12 coefficients after f, cx, cy, in projectaria order (k0–k5, p0, p1, s0–s3) | + +- **Optional** `-info-extra`, `egoverse.CameraInfoExtra`, published once: + `valid_radius` (px from the principal point where the lens model holds; 0 = whole + image) and `upright_rotation_cw_deg` (0, 90, 180 or 270: the clockwise turn that makes + a frame upright, e.g. 90 for Aria's sideways RGB sensor). + +## 8. Please don't + +- Re-encode, resize, crop, rotate or undistort video. *(Devices that store frames as + images, e.g. Aria VRS, encode them once to H.264 / H.265 at full resolution.)* +- Pack several frames into one video message. +- Resample or align streams to one rate; we match by timestamp. +- Trim streams to the camera's span; deliver every sample. +- Write zeros when tracking is lost; set `valid = false`. +- Remove gravity from, filter or fuse the IMU data. +- Put real operator names or emails anywhere; hash them. +- Ship calibration as a separate file. + +## 9. How a delivery is checked + +`python -m egomimic.rldb.mcap.validate episode_/ [...]` prints `PASS` or every +problem per episode (with the episode's `recording_id`), and exits 1 if any fails. Every +problem fails the episode. + +| Group | Checks | +|---|---| +| everything there | Required topics present with the exact type and encoding; required fields non-empty on every message; every episode-info key present, slugs / integers / enums valid; `embodiment` matches the topics; `annotation.mcap` holds only label topics. | +| video usable | Codec h264 / h265; first message a keyframe; one frame per message; frames decode; size matches the calibration; ≥ 30 fps. | +| clocks | Unix ns (after 2015); never decreasing; no gap over 0.5 s; streams cover the recording to within 1 s; start / end times match the data to within 1 s; head pose and quality ≥ 30 Hz. | +| calibration | Every camera has a valid `-info` with a known lens model and its coefficient count; `-info-extra` values valid; a camera's video and calibration name the same frame; every named frame reaches `world`; a head-worn rig's main camera has a `/tf_static` path to `head`; `/tf_static` rotations unit and translations finite. | +| imu | ≥ 200 Hz; mean \|accel\| in the stillest 1 s within 0.5 m/s² of 9.81 (gravity not removed); noise values > 0. | +| poses | No NaN / Inf, unit rotations (head, wrist, `/tf_static`); hand positions match `hand-joints`; confidences in [0, 1]; no all-zero hands marked valid; gaze finite when valid; `ee_pose` a rigid 4×4; joint arrays match `joint_names`; grippers in [0, 1]. | +| tasks & labels | Spans end after they start (labels may be points); inside the recorded data; each label inside one task; `annotation_type` known. | + +Thresholds stated by the plan (30 fps, 200 Hz, 30 Hz, 9.81 m/s²) are fixed. The others +(5 % rate tolerance, 0.5 s gap, 1 s span tolerance, 0.5 m/s² gravity tolerance, 1e-3 +unit tolerance) are **provisional, to be tuned on real deliveries**; all live in +`spec.py`. Video is decoded for the first 60 messages of each camera; the converter then decodes +all of it and rejects an episode where any message is not exactly one frame. + +## 10. Where each Lance field comes from + +`python -m egomimic.scripts.mcap_process.mcap_to_lance --stage-to --lance episode_/...` +validates each episode, converts it and stages it for the committer +(`egomimic.scripts.lance_commit`, the only writer of the tables; `convert_and_stage` is the +same per episode, for a pipeline to call). Lance times stay absolute Unix ns, as delivered. + +| Lance | MCAP | +|---|---| +| `recordings.recording_id` | `recording_id(lab, vendor-recording-id)` | +| `recordings.vendor_recording_id`, `lab`, `site`, `embodiment`, `rig_name`, `scene`, `license` | episode info | +| `recordings.operator` | `operator-id` | +| `recordings.captured_at` | `start-time-unix` | +| `recordings.num_frames`, `fps` | Decoded steps-clock frames and their timestamps | +| `recordings.video_path` | One video per camera, one frame per steps row: the steps camera as `front_1_480` (480×640) and `front_1` (360×480) pinhole outputs, every other camera under its topic name (`/left-wrist-camera` → `left_wrist`) at its native size and focal | +| `recordings.intrinsics` / `extrinsics` | The pinhole K of each video (Aria's front outputs use `ARIA_INTRINSICS`; other sources keep their native field of view) / each camera's pose relative to `head_pose` (the front camera, upright), from `/tf_static` | +| `tasks.*` | `/task` messages, or the episode-info task keys over the whole recording | +| `annotations.*` | `/subtask-annotation` (type, text, vendor, span) | +| `steps.timestamp_ns` | Steps-clock frame times | +| `steps.head_pose`, `head_pose_valid` | `/head-pose` (× the camera's link), `/head-pose-quality`, timing gaps | +| `steps.{side}.keypoints`, `ee_pose` | `HandKeypoints.positions`, through the MANO fit when `hand-joints` is not canonical MANO | +| `steps.{side}.wrist_pose`, `confidence`, `{side}_valid` | `HandKeypoints.wrist_pose`, `confidence`, `valid` | +| `steps.eye_gaze` | `/eye-gaze` yaw, pitch, depth | +| Eva steps `obs_*` / `cmd_*` | `/{side}-arm-state` / `-arm-action` (`position`, `ee_pose`), `/{side}-ee-state` / `-ee-action` | + +## 11. Not in 0.1 + +Deliberately left out of this version, to revisit with real deliveries: + +- Graded SLAM quality (Aria MPS gives 0 / 0.5 / 1), beyond the `TrackingQuality` states. +- More than one IMU per device. +- Video from devices that store per-frame images: re-encode (as now) or accept them. +- Dexterous hands, tactile and auxiliary joints (Sharpa). diff --git a/egomimic/rldb/mcap/__init__.py b/egomimic/rldb/mcap/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/rldb/mcap/proto/egoverse.proto b/egomimic/rldb/mcap/proto/egoverse.proto new file mode 100644 index 0000000000..a6a4012fb2 --- /dev/null +++ b/egomimic/rldb/mcap/proto/egoverse.proto @@ -0,0 +1,137 @@ +// Generated by `python -m egomimic.rldb.mcap.schemas`; do not edit by hand. +// EgoVerse vendor spec 0.1: message types not covered by foxglove schemas. + +syntax = "proto3"; + +package egoverse; + +// A 3D vector. +message Vector3 { + double x = 1; + double y = 2; + double z = 3; +} + +// A rotation; must be unit length. +message Quaternion { + double x = 1; + double y = 2; + double z = 3; + double w = 4; +} + +// A child frame's pose in its parent (T_parent_child): m and a unit quaternion. +message Pose { + Vector3 position = 1; + Quaternion orientation = 2; +} + +// Optional per-camera facts foxglove.CameraCalibration has no field for, on the +// camera's -info-extra, published once next to -info. +message CameraInfoExtra { + string frame_id = 1; // the camera's frame, as in its -info + double valid_radius = 2; // px from the principal point where the lens model holds (fisheye); 0 = the whole image + int32 upright_rotation_cw_deg = 3; // clockwise turn (0, 90, 180, 270) that makes a frame upright, for sensors mounted sideways +} + +// One raw IMU sample, straight off the device: no gravity removal, no fusion. +message Imu { + string frame_id = 1; // IMU frame; must be linked in /tf_static + Vector3 linear_acceleration = 2; // m/s^2, gravity included + Vector3 angular_velocity = 3; // rad/s +} + +// IMU noise model, as SLAM/VIO tools expect it. Published once. Datasheet values +// are fine when the device does not calibrate them. +message ImuInfo { + string frame_id = 1; + double accel_noise_density = 2; // m/s^2/sqrt(Hz) + double accel_random_walk = 3; // m/s^3/sqrt(Hz) + double gyro_noise_density = 4; // rad/s/sqrt(Hz) + double gyro_random_walk = 5; // rad/s^2/sqrt(Hz) +} + +// State of the vendor's SLAM at this time, one per head pose. +message TrackingQuality { + enum State { + UNKNOWN = 0; + OK = 1; + LOST = 2; + RELOCALIZED = 3; + } + State state = 1; +} + +// One hand at one time: its joints and/or the wrist pose. Joints are in the order +// the episode's hand-joints key lists them (names from spec.HAND_JOINTS). Set +// valid = false when the hand is not tracked; never write zeros in its place. +// A device without keypoints leaves positions empty and sets wrist_pose. +message HandKeypoints { + string frame_id = 1; // frame positions and wrist_pose are in + repeated double positions = 2; // (x, y, z) in m per hand-joints entry, flattened; or empty + float confidence = 3; // whole-hand tracking confidence in [0, 1] + bool valid = 4; // hand tracked and its values usable + repeated float joint_confidence = 5; // optional: per-joint confidences in [0, 1], one per joint + Pose wrist_pose = 6; // optional: T_frame_wrist +} + +// Where the wearer looks, as in Aria MPS eye gaze. frame_id is centered between +// the eyes with +x to the wearer's left, +y up, +z forward; the gaze point is +// depth * normalize(tan(yaw), tan(pitch), 1). Set valid = false when gaze is not +// tracked; never write zeros in its place. +message EyeGaze { + string frame_id = 1; // gaze frame; linked in /tf_static + double yaw = 2; // rad, positive = looking left + double pitch = 3; // rad, positive = looking up + double depth = 4; // m, distance to the gaze point + bool valid = 5; // gaze tracked and values usable +} + +// Measured state of one arm. Joint arrays follow joint_names. ee_pose is +// T_base_ee: the end effector in the arm-base frame (frame_id), row-major 4x4. +message ArmState { + string frame_id = 1; // arm-base frame, e.g. "left_arm_base" + repeated string joint_names = 2; + repeated double position = 3; // rad (or m, prismatic) + repeated double velocity = 4; // rad/s (or m/s) + repeated double torque = 5; // N*m (or N) + repeated double ee_pose = 6; // 16 values, row-major +} + +// Command sent to one arm at this time (the target it was told to reach), +// same conventions as ArmState. +message ArmAction { + string frame_id = 1; // arm-base frame + repeated string joint_names = 2; + repeated double position = 3; // rad (or m, prismatic) + repeated double ee_pose = 4; // 16 values, row-major +} + +// Gripper opening, measured (ee-state) or commanded (ee-action). +message GripperState { + double position = 1; // 0 = closed, 1 = fully open +} + +// One task performed in the recording; a recording may hold several +// (e.g. iron, then fold, then pack). Starts at the message's log time. +message Task { + string name = 1; // task slug, e.g. "fold_clothes" + string description = 2; // free text for this instance + repeated string objects = 3; // object slugs involved + int64 end_time_unix_ns = 4; // task end, absolute Unix ns +} + +// A label over [start, end), start being the message's log time; it must lie +// inside one /task (or the recording, without /task). Field `data` matches +// ABC-130k's /subtask-annotation. +message Annotation { + string data = 1; // English text + int64 end_time_unix_ns = 2; // end, absolute Unix ns; 0 = point event (one frame wide) + string annotation_type = 3; // default "language" + string annotation_vendor = 4; // who wrote it; default the episode's vendor +} + +// Task text for the episode, e.g. "fold the shirt and put it in the basket". +message Instruction { + string data = 1; +} diff --git a/egomimic/rldb/mcap/schemas.py b/egomimic/rldb/mcap/schemas.py new file mode 100644 index 0000000000..dfbd92fd15 --- /dev/null +++ b/egomimic/rldb/mcap/schemas.py @@ -0,0 +1,320 @@ +"""EgoVerse's own MCAP message types (vendor spec 0.1). + +Only what Foxglove's schemas do not cover. Video, camera calibration, transforms and head +pose use `foxglove.CompressedVideo`, `foxglove.CameraCalibration`, `foxglove.FrameTransforms` +and `foxglove.PoseInFrame` (package `foxglove-schemas-protobuf`). Robot vendors use the +ArmState / ArmAction / GripperState below; ABC-130k's own robot types are normalized to +them before validation. + +The messages are declared once, here, and built into a descriptor pool at import time, so +no `protoc` step is needed. `render_proto()` prints the same definitions as the +`proto/egoverse.proto` file handed to vendors; a unit test keeps the two in sync. + +A message's time is its MCAP log time (absolute Unix ns), so none of them carry a timestamp; +spans (Task, Annotation) add only their end time. + + python -m egomimic.rldb.mcap.schemas > egomimic/rldb/mcap/proto/egoverse.proto +""" + +from dataclasses import dataclass, field + +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory + +PACKAGE = "egoverse" +PROTO_FILE = "egoverse.proto" + +_F = descriptor_pb2.FieldDescriptorProto +_SCALARS = { + "double": _F.TYPE_DOUBLE, + "float": _F.TYPE_FLOAT, + "int64": _F.TYPE_INT64, + "string": _F.TYPE_STRING, + "int32": _F.TYPE_INT32, + "bool": _F.TYPE_BOOL, +} + + +@dataclass(frozen=True) +class Field: + name: str + kind: str # a scalar above, an enum of the same message, or another message here + repeated: bool = False + doc: str = "" + + +@dataclass(frozen=True) +class Enum: + name: str + values: tuple[ + str, ... + ] # numbered 0, 1, ...; proto3 needs the 0 value to mean "unset" + + +@dataclass(frozen=True) +class Message: + name: str + doc: str + fields: tuple[Field, ...] + enums: tuple[Enum, ...] = field(default=()) + + +MESSAGES = ( + Message( + "Vector3", + "A 3D vector.", + (Field("x", "double"), Field("y", "double"), Field("z", "double")), + ), + Message( + "Quaternion", + "A rotation; must be unit length.", + ( + Field("x", "double"), + Field("y", "double"), + Field("z", "double"), + Field("w", "double"), + ), + ), + Message( + "Pose", + "A child frame's pose in its parent (T_parent_child): m and a unit quaternion.", + (Field("position", "Vector3"), Field("orientation", "Quaternion")), + ), + Message( + "CameraInfoExtra", + "Optional per-camera facts foxglove.CameraCalibration has no field for, on the\n" + "camera's -info-extra, published once next to -info.", + ( + Field("frame_id", "string", doc="the camera's frame, as in its -info"), + Field( + "valid_radius", + "double", + doc="px from the principal point where the lens model holds (fisheye); " + "0 = the whole image", + ), + Field( + "upright_rotation_cw_deg", + "int32", + doc="clockwise turn (0, 90, 180, 270) that makes a frame upright, for " + "sensors mounted sideways", + ), + ), + ), + Message( + "Imu", + "One raw IMU sample, straight off the device: no gravity removal, no fusion.", + ( + Field("frame_id", "string", doc="IMU frame; must be linked in /tf_static"), + Field("linear_acceleration", "Vector3", doc="m/s^2, gravity included"), + Field("angular_velocity", "Vector3", doc="rad/s"), + ), + ), + Message( + "ImuInfo", + "IMU noise model, as SLAM/VIO tools expect it. Published once. Datasheet values\n" + "are fine when the device does not calibrate them.", + ( + Field("frame_id", "string"), + Field("accel_noise_density", "double", doc="m/s^2/sqrt(Hz)"), + Field("accel_random_walk", "double", doc="m/s^3/sqrt(Hz)"), + Field("gyro_noise_density", "double", doc="rad/s/sqrt(Hz)"), + Field("gyro_random_walk", "double", doc="rad/s^2/sqrt(Hz)"), + ), + ), + Message( + "TrackingQuality", + "State of the vendor's SLAM at this time, one per head pose.", + (Field("state", "State"),), + enums=(Enum("State", ("UNKNOWN", "OK", "LOST", "RELOCALIZED")),), + ), + Message( + "HandKeypoints", + "One hand at one time: its joints and/or the wrist pose. Joints are in the order\n" + "the episode's hand-joints key lists them (names from spec.HAND_JOINTS). Set\n" + "valid = false when the hand is not tracked; never write zeros in its place.\n" + "A device without keypoints leaves positions empty and sets wrist_pose.", + ( + Field("frame_id", "string", doc="frame positions and wrist_pose are in"), + Field( + "positions", + "double", + repeated=True, + doc="(x, y, z) in m per hand-joints entry, flattened; or empty", + ), + Field( + "confidence", + "float", + doc="whole-hand tracking confidence in [0, 1]", + ), + Field("valid", "bool", doc="hand tracked and its values usable"), + Field( + "joint_confidence", + "float", + repeated=True, + doc="optional: per-joint confidences in [0, 1], one per joint", + ), + Field("wrist_pose", "Pose", doc="optional: T_frame_wrist"), + ), + ), + Message( + "EyeGaze", + "Where the wearer looks, as in Aria MPS eye gaze. frame_id is centered between\n" + "the eyes with +x to the wearer's left, +y up, +z forward; the gaze point is\n" + "depth * normalize(tan(yaw), tan(pitch), 1). Set valid = false when gaze is not\n" + "tracked; never write zeros in its place.", + ( + Field("frame_id", "string", doc="gaze frame; linked in /tf_static"), + Field("yaw", "double", doc="rad, positive = looking left"), + Field("pitch", "double", doc="rad, positive = looking up"), + Field("depth", "double", doc="m, distance to the gaze point"), + Field("valid", "bool", doc="gaze tracked and values usable"), + ), + ), + Message( + "ArmState", + "Measured state of one arm. Joint arrays follow joint_names. ee_pose is\n" + "T_base_ee: the end effector in the arm-base frame (frame_id), row-major 4x4.", + ( + Field("frame_id", "string", doc='arm-base frame, e.g. "left_arm_base"'), + Field("joint_names", "string", repeated=True), + Field("position", "double", repeated=True, doc="rad (or m, prismatic)"), + Field("velocity", "double", repeated=True, doc="rad/s (or m/s)"), + Field("torque", "double", repeated=True, doc="N*m (or N)"), + Field("ee_pose", "double", repeated=True, doc="16 values, row-major"), + ), + ), + Message( + "ArmAction", + "Command sent to one arm at this time (the target it was told to reach),\n" + "same conventions as ArmState.", + ( + Field("frame_id", "string", doc="arm-base frame"), + Field("joint_names", "string", repeated=True), + Field("position", "double", repeated=True, doc="rad (or m, prismatic)"), + Field("ee_pose", "double", repeated=True, doc="16 values, row-major"), + ), + ), + Message( + "GripperState", + "Gripper opening, measured (ee-state) or commanded (ee-action).", + (Field("position", "double", doc="0 = closed, 1 = fully open"),), + ), + Message( + "Task", + "One task performed in the recording; a recording may hold several\n" + "(e.g. iron, then fold, then pack). Starts at the message's log time.", + ( + Field("name", "string", doc='task slug, e.g. "fold_clothes"'), + Field("description", "string", doc="free text for this instance"), + Field("objects", "string", repeated=True, doc="object slugs involved"), + Field("end_time_unix_ns", "int64", doc="task end, absolute Unix ns"), + ), + ), + Message( + "Annotation", + "A label over [start, end), start being the message's log time; it must lie\n" + "inside one /task (or the recording, without /task). Field `data` matches\n" + "ABC-130k's /subtask-annotation.", + ( + Field("data", "string", doc="English text"), + Field( + "end_time_unix_ns", + "int64", + doc="end, absolute Unix ns; 0 = point event (one frame wide)", + ), + Field("annotation_type", "string", doc='default "language"'), + Field( + "annotation_vendor", + "string", + doc="who wrote it; default the episode's vendor", + ), + ), + ), + Message( + "Instruction", + 'Task text for the episode, e.g. "fold the shirt and put it in the basket".', + (Field("data", "string"),), + ), +) + + +def file_descriptor_proto() -> descriptor_pb2.FileDescriptorProto: + fp = descriptor_pb2.FileDescriptorProto( + name=PROTO_FILE, package=PACKAGE, syntax="proto3" + ) + for m in MESSAGES: + mp = fp.message_type.add(name=m.name) + for e in m.enums: + ep = mp.enum_type.add(name=e.name) + for number, value in enumerate(e.values): + ep.value.add(name=value, number=number) + enum_names = {e.name for e in m.enums} + for number, f in enumerate(m.fields, start=1): + label = _F.LABEL_REPEATED if f.repeated else _F.LABEL_OPTIONAL + fd = mp.field.add(name=f.name, number=number, label=label) + if f.kind in _SCALARS: + fd.type = _SCALARS[f.kind] + elif f.kind in enum_names: + fd.type = _F.TYPE_ENUM + fd.type_name = f".{PACKAGE}.{m.name}.{f.kind}" + else: + fd.type = _F.TYPE_MESSAGE + fd.type_name = f".{PACKAGE}.{f.kind}" + return fp + + +# Own pool, so importing this module never collides with another egoverse.proto. +_pool = descriptor_pool.DescriptorPool() +_pool.Add(file_descriptor_proto()) + + +def _message_class(name: str): + return message_factory.GetMessageClass( + _pool.FindMessageTypeByName(f"{PACKAGE}.{name}") + ) + + +Vector3 = _message_class("Vector3") +Quaternion = _message_class("Quaternion") +Pose = _message_class("Pose") +CameraInfoExtra = _message_class("CameraInfoExtra") +Imu = _message_class("Imu") +ImuInfo = _message_class("ImuInfo") +TrackingQuality = _message_class("TrackingQuality") +HandKeypoints = _message_class("HandKeypoints") +EyeGaze = _message_class("EyeGaze") +ArmState = _message_class("ArmState") +ArmAction = _message_class("ArmAction") +GripperState = _message_class("GripperState") +Task = _message_class("Task") +Annotation = _message_class("Annotation") +Instruction = _message_class("Instruction") + + +def render_proto() -> str: + """The `.proto` text for MESSAGES, as handed to vendors.""" + lines = [ + "// Generated by `python -m egomimic.rldb.mcap.schemas`; do not edit by hand.", + "// EgoVerse vendor spec 0.1: message types not covered by foxglove schemas.", + "", + 'syntax = "proto3";', + "", + f"package {PACKAGE};", + ] + for m in MESSAGES: + lines.append("") + lines += [f"// {line}" for line in m.doc.split("\n")] + lines.append(f"message {m.name} {{") + for e in m.enums: + lines.append(f" enum {e.name} {{") + lines += [f" {v} = {n};" for n, v in enumerate(e.values)] + lines.append(" }") + for number, f in enumerate(m.fields, start=1): + prefix = "repeated " if f.repeated else "" + comment = f" // {f.doc}" if f.doc else "" + lines.append(f" {prefix}{f.kind} {f.name} = {number};{comment}") + lines.append("}") + return "\n".join(lines) + "\n" + + +if __name__ == "__main__": + print(render_proto(), end="") diff --git a/egomimic/rldb/mcap/spec.py b/egomimic/rldb/mcap/spec.py new file mode 100644 index 0000000000..d28cf3a8c6 --- /dev/null +++ b/egomimic/rldb/mcap/spec.py @@ -0,0 +1,363 @@ +"""Vendor MCAP spec 0.1 in machine-readable form: what `validate.py` checks against. + +SPEC.md is the human-readable version of this file; change both together. Numbers the +plan states are used as-is; the ones it does not state are marked PROVISIONAL, to be tuned +on real deliveries. + +An episode is one whole recording session, never a pre-cut clip: it becomes one row of +the Lance `recordings` table, and the tasks inside it become `tasks` rows. + + episode_/ + ├── episode.mcap # everything: video, sensors, poses, calibration, episode info + └── annotation.mcap # optional: /task spans and /subtask-annotation labels + +Every message is stamped with absolute Unix time in nanoseconds (its MCAP log time), in +both files. Every camera is a topic ending in `-camera`; its calibration is on the sibling +`-info`. Each video message holds exactly one frame, stamped with its capture time. + +The profile's `step_clock` camera defines the Lance steps rows: one row per decoded frame +of it, every other stream aligned to those times during conversion. Other cameras stay +separate videos and add no rows. + +Who and where, as in the Lance pilot: `lab` is the data source / capture family and the +Lance partition (steps/.lance: aria, eva, mecka, scale), `site` the organization that +collected the recording (rl2, microagi), `vendor` who delivered or processed it (acme). +Hardware goes in rig-name / device-serial. All three are normalized slugs. + +IDs: Lance recording_id = recording_id(lab, vendor-recording-id) below, so +vendor-recording-id must stay the same across every re-export or correction of the +recording. Task and annotation IDs are computed downstream from their content. +""" + +import re +import uuid +from dataclasses import dataclass + +SPEC_VERSION = "0.1" +SUPPORTED_SPEC_VERSIONS = ("0.1",) + +EPISODE_FILE = "episode.mcap" +# Its topics may instead be inline in episode.mcap, as ABC-130k does. +ANNOTATION_FILE = "annotation.mcap" +# ABC's delivery format names its episode-info record this, with hyphenated keys. +METADATA_RECORD = "session-metadata" + +CAMERA_SUFFIX = "-camera" +CALIBRATION_SUFFIX = "-info" +# Optional egoverse.CameraInfoExtra (valid radius, upright rotation) next to the -info. +CAMERA_EXTRA_SUFFIX = "-info-extra" +UPRIGHT_ROTATIONS = (0, 90, 180, 270) +# CameraCalibration.distortion_model -> number of D coefficients. An empty model with no +# coefficients is a pinhole camera. +DISTORTION_MODELS = { + "plumb_bob": 5, + "rational_polynomial": 8, + "kannala_brandt": 4, + "fisheye624": 12, +} + +# Hand joint vocabulary. Vendors list the joints their device tracks, in the order they +# write them, in the "hand-joints" key. Names that are MANO joints are what the MANO fit +# (egomimic/scripts/aria_process/aria_utils.py) matches; extras are kept but not fitted. +# Canonical MANO order, as Human.FINGER_EDGES uses it: a vendor sending this exactly +# needs no fit. +MANO_JOINTS = ( + "wrist", + *( + f"{finger}_{joint}" + for finger, joints in ( + ("thumb", ("cmc", "mcp", "ip", "tip")), + ("index", ("mcp", "pip", "dip", "tip")), + ("middle", ("mcp", "pip", "dip", "tip")), + ("ring", ("mcp", "pip", "dip", "tip")), + ("pinky", ("mcp", "pip", "dip", "tip")), + ) + for joint in joints + ), +) +EXTRA_HAND_JOINTS = ("palm_center",) +HAND_JOINTS = MANO_JOINTS + EXTRA_HAND_JOINTS +# Shutter: global, or rolling with its readout time. +SHUTTER = re.compile(r"global|rolling:[0-9]+(\.[0-9]+)?ms") + +# Lowercase letters and digits, words joined by "_" or "-": safe in Lance paths and IDs. +SLUG = re.compile(r"[a-z0-9]+(?:[_-][a-z0-9]+)*") + +# Copied from egomimic/rldb/lance/schema.py (PR #716) so the validator needs no Lance +# install; a unit test checks the two agree once that module is on main. +RECORDING_NAMESPACE = uuid.UUID("6f1d4c3e-2b1a-5e7f-9c8d-0a1b2c3d4e5f") + + +def recording_id(lab: str, vendor_recording_id: str) -> str: + """The Lance recordings.recording_id this episode will get.""" + return str(uuid.uuid5(RECORDING_NAMESPACE, f"{lab}/{vendor_recording_id}")) + + +# Frames every vendor uses by these names in /tf_static and in poses. +WORLD_FRAME = "world" +HEAD_FRAME = "head" + +# ---- Thresholds stated in the plan ---- +MIN_VIDEO_FPS = 30.0 +MIN_IMU_HZ = 200.0 +MIN_HEAD_POSE_HZ = 30.0 +GRAVITY = 9.81 # m/s^2: mean |accel| of a raw IMU at rest +VIDEO_CODECS = ("h264", "h265") +GRIPPER_RANGE = (0.0, 1.0) # 0 = closed, 1 = open + +# ---- PROVISIONAL thresholds (not in the plan) ---- +RATE_TOL = ( + 0.05 # a measured rate may fall this fraction below its minimum (29.97 fps passes) +) +GRAVITY_TOL = 0.5 # m/s^2 allowed around GRAVITY in the stillest window +STILL_WINDOW_S = 1.0 # window used to find the IMU at rest +MAX_GAP_S = 0.5 # largest allowed gap between consecutive messages of one stream +SPAN_TOL_S = 1.0 # how far a stream may start late / end early vs. the episode +METADATA_TIME_TOL_S = 1.0 # start/end-time-unix vs. the data's actual span +UNIT_TOL = 1e-3 # |quaternion| - 1, and R^T R - I for rotation matrices +VIDEO_SAMPLE_FRAMES = 60 # frames decoded per camera to check the stream +MIN_UNIX_NS = 1_420_070_400 * 10**9 # 2015-01-01: anything earlier is not Unix time + +# Must match egomimic.rldb.embodiment.embodiment.EMBODIMENT (lowercased); a unit test +# checks. Listed here so the validator does not import torch. +EMBODIMENTS = ( + "human_right_arm", + "human_left_arm", + "human_bimanual", + "eva_right_arm", + "eva_left_arm", + "eva_bimanual", +) + + +@dataclass(frozen=True) +class Topic: + """One channel an episode may carry. + + `kind` picks the checks the validator runs on it. `schema` is the exact schema name + required, or None for any. `fields` must be non-empty on every message. `side` marks a + per-arm topic, dropped for single-arm embodiments. + """ + + name: str + kind: str + required: bool = True + schema: str | None = None + fields: tuple[str, ...] = () + min_hz: float | None = None + side: str | None = None + + +@dataclass(frozen=True) +class MovingLink: + """A parent -> child transform that changes over time, published on `topic` + instead of /tf_static (e.g. the head moving in the world).""" + + parent: str + child: str + topic: str + + +@dataclass(frozen=True) +class MetaKey: + """One episode-info key. MCAP metadata values are strings; `kind` is what the + string must parse as: "str", "slug" (SLUG), "int", "enum" (one of `values`), + "sha256" (64 lowercase hex digits), "shutter" (SHUTTER) or "hand_joints" + (comma-separated HAND_JOINTS names, each once, including wrist).""" + + name: str + kind: str = "str" + values: tuple[str, ...] = () + human_only: bool = False + + +@dataclass(frozen=True) +class Profile: + name: str + step_clock: str # camera whose frames define the Lance steps rows + topics: tuple[Topic, ...] + moving_links: tuple[MovingLink, ...] + + +_VIDEO = dict( + kind="video", + schema="foxglove.CompressedVideo", + min_hz=MIN_VIDEO_FPS, + fields=("data", "format", "frame_id"), +) +_TF = Topic( + "/tf_static", "tf_static", schema="foxglove.FrameTransforms", fields=("transforms",) +) +_INSTRUCTION = Topic( + "/instruction", "instruction", schema="egoverse.Instruction", fields=("data",) +) +# Optional in either profile. Without /task, the episode-info task / instruction / objects +# keys describe one task spanning the whole recording; with /task, the spans are +# authoritative and those keys only summarize the session (no task row is made from them). +_LABELS = ( + Topic("/task", "task", required=False, schema="egoverse.Task", fields=("name",)), + Topic( + "/subtask-annotation", + "annotation", + required=False, + schema="egoverse.Annotation", + fields=("data",), + ), +) +# Must match egomimic.rldb.lance.schema.ANNOTATION_TYPES (PR #717); empty means "language". +ANNOTATION_TYPES = ("language",) + + +def _per_side(make) -> tuple[Topic, ...]: + return tuple(make(side) for side in ("left", "right")) + + +HUMAN = Profile( + name="human", + step_clock="/main-camera", + topics=( + Topic("/main-camera", **_VIDEO), + Topic("/slam-left-camera", **_VIDEO, required=False), + Topic("/slam-right-camera", **_VIDEO, required=False), + *_per_side(lambda s: Topic(f"/{s}-wrist-camera", **_VIDEO, required=False)), + Topic( + "/imu", + "imu", + schema="egoverse.Imu", + min_hz=MIN_IMU_HZ, + fields=("frame_id", "linear_acceleration", "angular_velocity"), + ), + Topic("/imu-info", "imu_info", schema="egoverse.ImuInfo", fields=("frame_id",)), + Topic( + "/head-pose", + "pose", + schema="foxglove.PoseInFrame", + min_hz=MIN_HEAD_POSE_HZ, + fields=("frame_id", "pose"), + ), + Topic( + "/head-pose-quality", + "tracking_quality", + schema="egoverse.TrackingQuality", + min_hz=MIN_HEAD_POSE_HZ, + ), + *_per_side( + lambda s: Topic( + f"/{s}-hand-keypoints", + "hand_keypoints", + schema="egoverse.HandKeypoints", + fields=("frame_id",), + side=s, + ) + ), + Topic( + "/eye-gaze", + "eye_gaze", + required=False, + schema="egoverse.EyeGaze", + fields=("frame_id",), + ), + _TF, + _INSTRUCTION, + *_LABELS, + ), + moving_links=(MovingLink(WORLD_FRAME, HEAD_FRAME, "/head-pose"),), +) + +ROBOT = Profile( + name="robot", + step_clock="/top-camera", + topics=( + Topic("/top-camera", **_VIDEO), + *_per_side(lambda s: Topic(f"/{s}-wrist-camera", **_VIDEO, side=s)), + *_per_side( + lambda s: Topic( + f"/{s}-arm-state", + "arm", + schema="egoverse.ArmState", + side=s, + fields=( + "frame_id", + "joint_names", + "position", + "velocity", + "torque", + "ee_pose", + ), + ) + ), + *_per_side( + lambda s: Topic( + f"/{s}-arm-action", + "arm", + schema="egoverse.ArmAction", + side=s, + fields=("frame_id", "joint_names", "position", "ee_pose"), + ) + ), + *_per_side( + lambda s: Topic( + f"/{s}-ee-state", "gripper", schema="egoverse.GripperState", side=s + ) + ), + *_per_side( + lambda s: Topic( + f"/{s}-ee-action", "gripper", schema="egoverse.GripperState", side=s + ) + ), + _TF, + _INSTRUCTION, + *_LABELS, + ), + # The end effector moves relative to its arm base; the wrist camera hangs off it. + moving_links=tuple( + MovingLink(f"{s}_arm_base", f"{s}_ee", f"/{s}-arm-state") + for s in ("left", "right") + ), +) + +METADATA_KEYS = ( + MetaKey("lab", "slug"), # data source / Lance partition: aria, eva, mecka, scale + MetaKey("site", "slug"), # organization that collected it: rl2, microagi + MetaKey("operator-id", "sha256"), # SHA-256 of the operator, never a name + MetaKey("task"), + MetaKey("instruction"), + MetaKey("embodiment", "enum", EMBODIMENTS), + MetaKey("rig-name"), + MetaKey("device-serial"), + MetaKey("scene"), + MetaKey("objects"), # comma-separated + MetaKey("start-time-unix", "int"), # ms + MetaKey("end-time-unix", "int"), # ms + # The vendor's own ID for this recording; stable across re-exports and corrections. + MetaKey("vendor-recording-id"), + MetaKey("vendor", "slug"), # who delivered or processed it: acme + MetaKey("spec-version", "enum", SUPPORTED_SPEC_VERSIONS), + MetaKey("license"), # e.g. "CC-BY-4.0"; the Lance recordings.license column + # Where in the exposure frames are stamped. + MetaKey("camera-timestamp-ref", "enum", ("start", "mid", "end"), human_only=True), + MetaKey("shutter", "shutter", human_only=True), # "global" or "rolling:16ms" + # HandKeypoints.positions order, chosen by the vendor, e.g. ",".join(MANO_JOINTS) + MetaKey("hand-joints", "hand_joints", human_only=True), +) + + +def profile_for(embodiment: str) -> Profile: + return HUMAN if embodiment.startswith("human_") else ROBOT + + +def arm_sides(embodiment: str) -> tuple[str, ...]: + """Sides the embodiment uses: ("left", "right") for *_bimanual.""" + if embodiment.endswith("_left_arm"): + return ("left",) + if embodiment.endswith("_right_arm"): + return ("right",) + return ("left", "right") + + +def topics_for(embodiment: str) -> tuple[Topic, ...]: + """The profile's topics, minus per-side topics of an arm the embodiment lacks.""" + sides = arm_sides(embodiment) + return tuple( + t for t in profile_for(embodiment).topics if t.side is None or t.side in sides + ) diff --git a/egomimic/rldb/mcap/validate.py b/egomimic/rldb/mcap/validate.py new file mode 100644 index 0000000000..5256fbe6fb --- /dev/null +++ b/egomimic/rldb/mcap/validate.py @@ -0,0 +1,1214 @@ +"""Check a delivered episode against the vendor MCAP spec (spec.py, SPEC.md). + + python -m egomimic.rldb.mcap.validate episode_/ [episode_/ ...] + +Prints PASS, or every problem found, per episode; exits 1 if any episode fails. The +checks are the plan's six groups, plus the task / label spans the Lance tables need: + + 1. everything there required topics, types, non-empty fields, episode info + 2. video usable codec, starts on a keyframe, decodes, size matches calibration, fps + 3. clocks Unix ns, in order, no big gaps, streams cover the episode, rates + 4. calibration every camera has intrinsics and a path to the world frame + 5. imu raw (gravity still in it) and fast enough, noise specs present + 6. poses no NaNs, unit rotations, gripper in [0, 1], lost hands flagged not zeroed + 7. tasks & labels /task and /subtask-annotation spans are ordered, inside the data, + each label inside one task (the rules Lance's annotation review + applies, egomimic/rldb/lance/annotations.py in PR #717) + +The file is read once, in file order. Video is only decoded for its first +VIDEO_SAMPLE_FRAMES messages, each of which must decode to exactly one frame, so that +every frame has its own capture time (Lance needs one steps row per frame). annotation.mcap, when present, +is read too and may only hold the label topics. + +MCAP reading and packet-by-packet video decoding follow PR #606's +egomimic/scripts/abc_process/abc_to_zarr.py. +""" + +import argparse +import math +import re +import sys +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np + +from egomimic.rldb.mcap import spec + +GROUPS = { + "present": "everything there", + "video": "video usable", + "clocks": "clocks", + "calibration": "calibration", + "imu": "imu", + "poses": "poses", + "labels": "tasks & labels", +} +# Topics whose messages form a time series; the rest are published once or rarely. +STREAM_KINDS = { + "video", + "imu", + "pose", + "tracking_quality", + "hand_keypoints", + "arm", + "gripper", + "eye_gaze", +} + + +def _reach(edges: dict, start: str) -> set[str]: + reached, todo = {start}, [start] + while todo: + for nxt in edges[todo.pop()] - reached: + reached.add(nxt) + todo.append(nxt) + return reached + + +def frames_in_message(data: bytes, codec: str) -> int: + """Pictures (access units) in one video message; the spec wants exactly 1, so every + frame has its own capture time. Counted with the parser: the decoder silently drops + the extra frames of a packed message. Shared with the MCAP -> Lance converter.""" + import av + + parser = av.CodecContext.create("hevc" if codec == "h265" else "h264", "r") + with av.logging.Capture(): # lone P-frames log "non-existing PPS" + return len(parser.parse(data)) + len(parser.parse(None)) + + +# Topics allowed in annotation.mcap. +LABEL_TOPICS = {"/task", "/subtask-annotation"} +# Annex B NAL unit types that start a decodable picture (IDR / IRAP). +KEYFRAME_NALS = {"h264": {5}, "h265": set(range(16, 22))} + + +@dataclass(frozen=True) +class Problem: + group: str + topic: str | None + message: str + + def __str__(self) -> str: + where = f"{self.topic}: " if self.topic else "" + return f"[{GROUPS[self.group]}] {where}{self.message}" + + +@dataclass +class _Stream: + """Everything one pass keeps about one topic.""" + + schema: str + encoding: str + times: list[int] = field(default_factory=list) + msgs: list = field(default_factory=list) # decoded; video keeps only the sample + empty: Counter = field( + default_factory=Counter + ) # field -> messages where it is empty + missing_fields: set = field(default_factory=set) # fields the message type lacks + undecodable: bool = False + + @property + def n(self) -> int: + return len(self.times) + + def span_s(self) -> float: + return (max(self.times) - min(self.times)) / 1e9 if self.times else 0.0 + + def rate_hz(self) -> float: + span = self.span_s() + return (self.n - 1) / span if span > 0 else 0.0 + + +@dataclass +class _Episode: + path: Path + metadata: dict[str, dict[str, str]] # record name -> its key/values + streams: dict[str, _Stream] + + @property + def info(self) -> dict[str, str]: + return self.metadata.get(spec.METADATA_RECORD, {}) + + +# ---------------------------------------------------------------- reading + + +def _is_empty(msg, name: str) -> bool: + value = getattr(msg, name) + if isinstance(value, (str, bytes)) or hasattr(value, "__len__"): + return len(value) == 0 + try: + return not msg.HasField(name) # sub-message + except ValueError: + return False # proto3 scalar: zero is a legitimate value + + +def _read(path: Path) -> _Episode: + from mcap.reader import make_reader + from mcap_protobuf.decoder import DecoderFactory + + fields_of = {t.name: t.fields for p in (spec.HUMAN, spec.ROBOT) for t in p.topics} + video_topics = { + t.name for p in (spec.HUMAN, spec.ROBOT) for t in p.topics if t.kind == "video" + } + factory = DecoderFactory() + decoders: dict[int, object] = {} + streams: dict[str, _Stream] = {} + + with open(path, "rb") as f: + reader = make_reader(f) + metadata = {r.name: dict(r.metadata) for r in reader.iter_metadata()} + for schema, channel, message in reader.iter_messages(log_time_order=False): + topic = channel.topic + s = streams.get(topic) + if s is None: + s = streams[topic] = _Stream( + schema.name if schema else "", channel.message_encoding + ) + s.times.append(message.log_time) + if channel.id not in decoders: + decoders[channel.id] = factory.decoder_for( + channel.message_encoding, schema + ) + decode = decoders[channel.id] + if decode is None: + s.undecodable = True + continue + is_video = topic in video_topics or topic.endswith(spec.CAMERA_SUFFIX) + if is_video and len(s.msgs) >= spec.VIDEO_SAMPLE_FRAMES: + continue # video is checked on its sample; decoding all of it is costly + msg = decode(message.data) + for name in fields_of.get(topic, ()): + if not hasattr(msg, name): + s.missing_fields.add(name) + elif _is_empty(msg, name): + s.empty[name] += 1 + s.msgs.append(msg) + return _Episode(path, metadata, streams) + + +# ---------------------------------------------------------------- 1. everything there + + +def _embodiment(ep: _Episode) -> str: + """Declared embodiment, or a guess from the topics when it is missing/invalid.""" + declared = ep.info.get("embodiment", "") + if declared in spec.EMBODIMENTS: + return declared + return "human_bimanual" if "/main-camera" in ep.streams else "eva_bimanual" + + +def _int(value: str) -> int | None: + """The integer an episode-info value spells, or None (int() alone takes "²", " 5").""" + return int(value) if re.fullmatch(r"-?[0-9]+", value) else None + + +def _hand_joints_problem(value: str) -> str: + names = value.split(",") + unknown = [n for n in names if n not in spec.HAND_JOINTS] + if unknown: + return f"unknown joints {unknown}; use names from {list(spec.HAND_JOINTS)}" + if len(set(names)) != len(names): + return "lists a joint twice" + if "wrist" not in names: + return "must include wrist" + return "" + + +def _hand_joint_count(ep: _Episode) -> int: + """Joints per hand the episode declares; MANO's 21 when the key is unusable.""" + value = ep.info.get("hand-joints", "") + if value and not _hand_joints_problem(value): + return len(value.split(",")) + return len(spec.MANO_JOINTS) + + +def check_present(ep: _Episode, embodiment: str) -> list[Problem]: + out = [] + expected = spec.topics_for(embodiment) + for t in expected: + s = ep.streams.get(t.name) + if s is None: + if t.required: + out.append(Problem("present", t.name, "missing required topic")) + continue + if t.schema and s.schema != t.schema: + out.append( + Problem( + "present", t.name, f"type is {s.schema!r}, expected {t.schema!r}" + ) + ) + if s.encoding != "protobuf" or s.undecodable: + out.append( + Problem( + "present", + t.name, + f"cannot decode {s.encoding!r} messages; the spec uses protobuf", + ) + ) + for name in sorted(s.missing_fields): + out.append( + Problem("present", t.name, f"message type has no field {name!r}") + ) + for name, count in sorted(s.empty.items()): + out.append( + Problem("present", t.name, f"{name!r} empty on {count}/{s.n} messages") + ) + + # The embodiment must match the topics: no topics of the other profile or an unused arm. + allowed = {t.name for t in expected} + for p in (spec.HUMAN, spec.ROBOT): + for t in p.topics: + if t.name in ep.streams and t.name not in allowed: + out.append( + Problem( + "present", + t.name, + f"topic does not fit embodiment {embodiment!r}", + ) + ) + + human = spec.profile_for(embodiment) is spec.HUMAN + if spec.METADATA_RECORD not in ep.metadata: + found = ", ".join(sorted(ep.metadata)) or "none" + out.append( + Problem( + "present", + None, + f"no {spec.METADATA_RECORD!r} metadata record (found: {found})", + ) + ) + return out + for key in spec.METADATA_KEYS: + if key.human_only and not human: + continue + value = ep.info.get(key.name) + if value is None or value == "": + out.append(Problem("present", None, f"episode info {key.name!r} missing")) + elif key.kind == "sha256" and not re.fullmatch(r"[0-9a-f]{64}", value): + out.append( + Problem( + "present", + None, + f"episode info {key.name!r} is not a SHA-256 hex digest " + "(hash names and emails)", + ) + ) + elif key.kind == "shutter" and not spec.SHUTTER.fullmatch(value): + out.append( + Problem( + "present", + None, + f"episode info {key.name!r} = {value!r}, expected 'global' or " + "'rolling:ms'", + ) + ) + elif key.kind == "slug" and not spec.SLUG.fullmatch(value): + out.append( + Problem( + "present", + None, + f"episode info {key.name!r} = {value!r} is not a lowercase slug " + "(a-z, 0-9, words joined by _ or -)", + ) + ) + elif key.kind == "hand_joints" and (bad := _hand_joints_problem(value)): + out.append(Problem("present", None, f"episode info {key.name!r}: {bad}")) + elif key.kind == "int" and _int(value) is None: + out.append( + Problem( + "present", + None, + f"episode info {key.name!r} = {value!r} is not an integer", + ) + ) + elif key.kind == "enum" and value not in key.values: + out.append( + Problem( + "present", + None, + f"episode info {key.name!r} = {value!r}, expected one of {list(key.values)}", + ) + ) + return out + + +# ---------------------------------------------------------------- 2. video usable + + +def _codec(fmt: str) -> str | None: + fmt = fmt.lower() + return fmt if fmt in spec.VIDEO_CODECS else None + + +def _nal_types(data: bytes, codec: str) -> set[int]: + types, i = set(), 0 + while (j := data.find(b"\x00\x00\x01", i)) >= 0 and j + 3 < len(data): + b = data[j + 3] + types.add(b & 0x1F if codec == "h264" else (b >> 1) & 0x3F) + i = j + 3 + return types + + +def _calibration(ep: _Episode, camera: str): + s = ep.streams.get(camera + spec.CALIBRATION_SUFFIX) + return s.msgs[0] if s and s.msgs else None + + +def check_video(ep: _Episode) -> list[Problem]: + import av + + out = [] + for topic, s in sorted(ep.streams.items()): + if not topic.endswith(spec.CAMERA_SUFFIX) or not s.msgs: + continue + fmt = str(getattr(s.msgs[0], "format", "")) + codec = _codec(fmt) + if codec is None: + out.append( + Problem( + "video", + topic, + f"format {fmt!r}, expected one of {list(spec.VIDEO_CODECS)}", + ) + ) + continue + keyframes = [ + i + for i, msg in enumerate(s.msgs) + if _nal_types(bytes(msg.data), codec) & KEYFRAME_NALS[codec] + ] + if not keyframes: + out.append( + Problem( + "video", + topic, + f"no keyframe (Annex B IDR/IRAP) in the first {len(s.msgs)} messages", + ) + ) + continue + if keyframes[0] != 0: + out.append( + Problem( + "video", topic, "first message is not a keyframe (Annex B IDR/IRAP)" + ) + ) + + # Decode from the first keyframe, so a bad start is reported once, above. + sample = s.msgs[keyframes[0] :] + + packed = sum(frames_in_message(bytes(m.data), codec) != 1 for m in sample) + if packed: + out.append( + Problem( + "video", + topic, + f"{packed}/{len(sample)} sampled messages do not hold exactly one " + "frame; each message must be one frame at its capture time", + ) + ) + ctx = av.CodecContext.create("hevc" if codec == "h265" else "h264", "r") + frames, size = 0, None + try: + for msg in sample: + for frame in ctx.decode(av.Packet(bytes(msg.data))): + frames += 1 + size = size or (frame.width, frame.height) + for frame in ctx.decode(None): + frames += 1 + size = size or (frame.width, frame.height) + except av.error.FFmpegError as e: + out.append( + Problem( + "video", + topic, + f"decode error in the first {len(sample)} messages: {e}", + ) + ) + continue + if frames == 0: + out.append( + Problem( + "video", + topic, + f"no frames decoded from the first {len(sample)} messages", + ) + ) + continue + + calib = _calibration(ep, topic) + if calib is not None and size != (int(calib.width), int(calib.height)): + out.append( + Problem( + "video", + topic, + f"frames are {size[0]}x{size[1]}, calibration says {calib.width}x{calib.height}", + ) + ) + + fps = s.rate_hz() # meaningless for packed messages, reported above + if not packed and fps < spec.MIN_VIDEO_FPS * (1 - spec.RATE_TOL): + out.append( + Problem( + "video", topic, f"{fps:.1f} fps, need >= {spec.MIN_VIDEO_FPS:g}" + ) + ) + return out + + +# ---------------------------------------------------------------- 3. clocks + + +def _time_series(ep: _Episode, embodiment: str) -> dict[str, _Stream]: + """Cameras and the spec's time-series topics: what defines the recording's span.""" + expected = {t.name: t for t in spec.topics_for(embodiment)} + return { + topic: s + for topic, s in ep.streams.items() + if s.times + and ( + topic.endswith(spec.CAMERA_SUFFIX) + or (topic in expected and expected[topic].kind in STREAM_KINDS) + ) + } + + +def _span(streams: dict[str, _Stream]) -> tuple[int, int]: + return ( + min(min(s.times) for s in streams.values()), + max(max(s.times) for s in streams.values()), + ) + + +def check_clocks(ep: _Episode, embodiment: str) -> list[Problem]: + out = [] + expected = {t.name: t for t in spec.topics_for(embodiment)} + streams = _time_series(ep, embodiment) + for topic, s in sorted(ep.streams.items()): + if s.times and min(s.times) < spec.MIN_UNIX_NS: + out.append( + Problem( + "clocks", topic, "timestamps are before 2015: not absolute Unix ns" + ) + ) + for topic, s in sorted(streams.items()): + t = np.asarray(s.times, dtype=np.int64) + back = int((np.diff(t) < 0).sum()) + if back: + out.append( + Problem("clocks", topic, f"timestamps go backwards {back} times") + ) + same = int((np.diff(t) == 0).sum()) + if same and topic.endswith(spec.CAMERA_SUFFIX): + # Every frame needs its own time: Lance steps timestamps strictly increase. + out.append(Problem("clocks", topic, f"{same} frames repeat a timestamp")) + if len(t) > 1: + gap = float(np.diff(np.sort(t)).max()) / 1e9 + if gap > spec.MAX_GAP_S: + out.append( + Problem( + "clocks", topic, f"{gap:.2f} s gap, max {spec.MAX_GAP_S:g} s" + ) + ) + # Video rates are checked from decoded frames; IMU rate in its own group. + t_spec = expected.get(topic) + if t_spec and t_spec.min_hz and t_spec.kind not in ("video", "imu"): + if s.rate_hz() < t_spec.min_hz * (1 - spec.RATE_TOL): + out.append( + Problem( + "clocks", + topic, + f"{s.rate_hz():.1f} Hz, need >= {t_spec.min_hz:g}", + ) + ) + + if not streams: + return out + start, end = _span(streams) + for topic, s in sorted(streams.items()): + late, early = (min(s.times) - start) / 1e9, (end - max(s.times)) / 1e9 + if late > spec.SPAN_TOL_S or early > spec.SPAN_TOL_S: + out.append( + Problem( + "clocks", + topic, + f"covers only part of the episode (starts {late:.1f} s late, ends {early:.1f} s early)", + ) + ) + + for key, data_ns in (("start-time-unix", start), ("end-time-unix", end)): + value = ep.info.get(key, "") + if _int(value) is not None: + off = abs(_int(value) * 1e6 - data_ns) / 1e9 + if off > spec.METADATA_TIME_TOL_S: + out.append( + Problem( + "clocks", + None, + f"episode info {key!r} is {off:.1f} s from the data", + ) + ) + return out + + +# ---------------------------------------------------------------- 4. calibration + + +def _quat_ok(q) -> bool: + v = [q.x, q.y, q.z, q.w] + return ( + all(math.isfinite(c) for c in v) + and abs(math.sqrt(sum(c * c for c in v)) - 1) <= spec.UNIT_TOL + ) + + +def check_calibration(ep: _Episode, embodiment: str) -> list[Problem]: + out = [] + cameras = sorted(t for t in ep.streams if t.endswith(spec.CAMERA_SUFFIX)) + for cam in cameras: + calib = _calibration(ep, cam) + if calib is None: + out.append( + Problem( + "calibration", + cam, + f"no {cam + spec.CALIBRATION_SUFFIX} calibration", + ) + ) + continue + video = ep.streams.get(cam) + if video and video.msgs and hasattr(video.msgs[0], "frame_id"): + named = {str(m.frame_id) for m in video.msgs} + if named != {str(calib.frame_id)}: + out.append( + Problem( + "calibration", + cam, + f"video frames name {sorted(named)} but its calibration names " + f"{calib.frame_id!r}: they must be the camera's one frame", + ) + ) + info = ep.streams[cam + spec.CALIBRATION_SUFFIX] + if info.schema != "foxglove.CameraCalibration" or info.n != 1: + out.append( + Problem( + "calibration", + cam + spec.CALIBRATION_SUFFIX, + f"expected one foxglove.CameraCalibration, got {info.n} " + f"{info.schema or 'untyped'} message(s)", + ) + ) + K = list(getattr(calib, "K", [])) + D = list(getattr(calib, "D", [])) + if not all(math.isfinite(v) for v in K + D): + out.append(Problem("calibration", cam, "K or D has NaN / Inf")) + elif len(K) == 9 and not ( + 0 <= K[2] <= calib.width and 0 <= K[5] <= calib.height + ): + out.append( + Problem( + "calibration", + cam, + f"principal point ({K[2]:g}, {K[5]:g}) is outside the " + f"{calib.width}x{calib.height} image", + ) + ) + if len(K) != 9 or K[0] <= 0 or K[4] <= 0 or not calib.width or not calib.height: + out.append( + Problem( + "calibration", + cam, + "calibration has no valid K (fx, fy > 0) or size", + ) + ) + model, n_d = str(calib.distortion_model), len(calib.D) + if model == "" and n_d == 0: + pass # pinhole + elif model not in spec.DISTORTION_MODELS: + out.append( + Problem( + "calibration", + cam, + f"distortion_model {model!r}, expected one of " + f"{list(spec.DISTORTION_MODELS)}", + ) + ) + elif n_d != spec.DISTORTION_MODELS[model]: + out.append( + Problem( + "calibration", + cam, + f"{model} needs {spec.DISTORTION_MODELS[model]} D coefficients, " + f"got {n_d}", + ) + ) + + extra = ep.streams.get(cam + spec.CAMERA_EXTRA_SUFFIX) + if extra and extra.msgs: + m = extra.msgs[0] + if m.upright_rotation_cw_deg not in spec.UPRIGHT_ROTATIONS: + out.append( + Problem( + "calibration", + cam, + f"upright_rotation_cw_deg {m.upright_rotation_cw_deg}, expected " + f"one of {list(spec.UPRIGHT_ROTATIONS)}", + ) + ) + if not (math.isfinite(m.valid_radius) and m.valid_radius >= 0): + out.append( + Problem( + "calibration", cam, f"valid_radius {m.valid_radius} is not >= 0" + ) + ) + if m.frame_id != str(getattr(calib, "frame_id", m.frame_id)): + out.append( + Problem( + "calibration", + cam + spec.CAMERA_EXTRA_SUFFIX, + f"frame_id {m.frame_id!r} differs from the camera's -info", + ) + ) + + # Frame graph: /tf_static edges plus the moving links whose topic is present. + edges = defaultdict(set) + static = defaultdict(set) # /tf_static alone: what is rigidly mounted to what + tf = ep.streams.get("/tf_static") + bad_rot = bad_tr = 0 + for msg in tf.msgs if tf else (): + for x in msg.transforms: + for g in (edges, static): + g[x.parent_frame_id].add(x.child_frame_id) + g[x.child_frame_id].add(x.parent_frame_id) + bad_rot += not _quat_ok(x.rotation) + t = x.translation + bad_tr += not all(math.isfinite(v) for v in (t.x, t.y, t.z)) + if bad_rot: + out.append( + Problem( + "calibration", + "/tf_static", + f"{bad_rot} transforms have a non-unit rotation", + ) + ) + if bad_tr: + out.append( + Problem( + "calibration", + "/tf_static", + f"{bad_tr} transforms have a NaN / Inf translation", + ) + ) + sides = spec.arm_sides(embodiment) + for link in spec.profile_for(embodiment).moving_links: + if link.topic in ep.streams and ( + not link.parent.endswith("_arm_base") or link.parent.split("_")[0] in sides + ): + edges[link.parent].add(link.child) + edges[link.child].add(link.parent) + reached = _reach(edges, spec.WORLD_FRAME) + + # A head-worn rig's steps camera rides on the head: its frame must be rigidly + # mounted to the head (a /tf_static path), which is how the steps are placed. + profile = spec.profile_for(embodiment) + if any(link.child == spec.HEAD_FRAME for link in profile.moving_links): + clock = _calibration(ep, profile.step_clock) + fid = str(clock.frame_id) if clock is not None else "" + # (An unplaced frame is already reported below as having no path to world.) + if fid in reached and fid not in _reach(static, spec.HEAD_FRAME): + out.append( + Problem( + "calibration", + profile.step_clock, + f"frame {clock.frame_id!r} has no fixed path to " + f"{spec.HEAD_FRAME!r} via /tf_static: the steps camera must be " + "mounted on the head", + ) + ) + + # Cameras and the IMU must name their frame; every frame any topic names (hands, + # wrist poses, gaze, ...) must be placed in the world. + must_name = set(cameras) | {"/imu"} + for topic, s in sorted(ep.streams.items()): + skip = topic == "/tf_static" or topic.endswith( + (spec.CALIBRATION_SUFFIX, spec.CAMERA_EXTRA_SUFFIX) + ) + if skip or not s.msgs or not hasattr(s.msgs[0], "frame_id"): + continue # a camera's -info / -info-extra share its frame, checked above + named = {str(m.frame_id) for m in s.msgs} + if "" in named and topic in must_name: + out.append( + Problem("calibration", topic, "no frame_id, so it cannot be placed") + ) + for frame in sorted(named - reached - {""}): + out.append( + Problem( + "calibration", + topic, + f"frame {frame!r} has no path to {spec.WORLD_FRAME!r} via /tf_static", + ) + ) + head = ep.streams.get("/head-pose") + if ( + head + and head.msgs + and {str(m.frame_id) for m in head.msgs} != {spec.WORLD_FRAME} + ): + out.append( + Problem( + "calibration", + "/head-pose", + f"poses must be in frame {spec.WORLD_FRAME!r} (T_world_head)", + ) + ) + return out + + +# ---------------------------------------------------------------- 5. imu + + +def check_imu(ep: _Episode) -> list[Problem]: + out = [] + s = ep.streams.get("/imu") + if s and s.msgs: + if s.rate_hz() < spec.MIN_IMU_HZ * (1 - spec.RATE_TOL): + out.append( + Problem( + "imu", "/imu", f"{s.rate_hz():.0f} Hz, need >= {spec.MIN_IMU_HZ:g}" + ) + ) + acc = np.array( + [ + [ + m.linear_acceleration.x, + m.linear_acceleration.y, + m.linear_acceleration.z, + ] + for m in s.msgs + ] + ) + gyr = np.array( + [ + [m.angular_velocity.x, m.angular_velocity.y, m.angular_velocity.z] + for m in s.msgs + ] + ) + if not (np.isfinite(acc).all() and np.isfinite(gyr).all()): + out.append(Problem("imu", "/imu", "NaN/Inf samples")) + else: + # Stillest window = lowest mean rotation rate; a raw IMU there reads ~9.81 m/s^2. + w = max(1, min(len(acc), round(s.rate_hz() * spec.STILL_WINDOW_S))) + gyro_mag = np.convolve( + np.linalg.norm(gyr, axis=1), np.ones(w) / w, mode="valid" + ) + i = int(np.argmin(gyro_mag)) + still = float(np.linalg.norm(acc[i : i + w], axis=1).mean()) + if abs(still - spec.GRAVITY) > spec.GRAVITY_TOL: + out.append( + Problem( + "imu", + "/imu", + f"mean |accel| at rest = {still:.2f} m/s^2, expected ~{spec.GRAVITY} (gravity removed, or wrong units?)", + ) + ) + info = ep.streams.get("/imu-info") + if info and info.msgs: + names = ( + "accel_noise_density", + "accel_random_walk", + "gyro_noise_density", + "gyro_random_walk", + ) + bad = [n for n in names if not getattr(info.msgs[0], n, 0) > 0] + if bad: + out.append( + Problem( + "imu", "/imu-info", f"noise values must be > 0: {', '.join(bad)}" + ) + ) + return out + + +# ---------------------------------------------------------------- 6. poses + + +def _pose_ok(pose) -> bool: + p = pose.position + return _quat_ok(pose.orientation) and all(math.isfinite(v) for v in (p.x, p.y, p.z)) + + +def check_poses(ep: _Episode, embodiment: str) -> list[Problem]: + out = [] + for t in spec.topics_for(embodiment): + s = ep.streams.get(t.name) + if not s or not s.msgs: + continue + if t.kind == "pose": + bad = sum(not _pose_ok(m.pose) for m in s.msgs) + if bad: + out.append( + Problem( + "poses", + t.name, + f"{bad}/{s.n} poses have NaNs or a non-unit rotation", + ) + ) + elif t.kind == "hand_keypoints": + joints = _hand_joint_count(ep) + shape = sum(len(m.positions) not in (0, 3 * joints) for m in s.msgs) + if shape: + out.append( + Problem( + "poses", + t.name, + f"{shape}/{s.n} messages have neither {3 * joints} positions " + f"({joints} hand-joints x xyz) nor none", + ) + ) + empty = sum( + m.valid and not m.positions and not m.HasField("wrist_pose") + for m in s.msgs + ) + if empty: + out.append( + Problem( + "poses", + t.name, + f"{empty} valid messages have neither positions nor wrist_pose", + ) + ) + conf = sum( + not 0 <= m.confidence <= 1 + or len(m.joint_confidence) not in (0, joints) + or not all(0 <= c <= 1 for c in m.joint_confidence) + for m in s.msgs + ) + if conf: + out.append( + Problem( + "poses", + t.name, + f"{conf}/{s.n} messages have a confidence outside [0, 1] or " + f"joint_confidence not {joints} values", + ) + ) + wrist = sum( + m.valid and m.HasField("wrist_pose") and not _pose_ok(m.wrist_pose) + for m in s.msgs + ) + if wrist: + out.append( + Problem( + "poses", + t.name, + f"{wrist} valid wrist poses have NaNs or a non-unit rotation", + ) + ) + valid = [ + np.asarray(m.positions) + for m in s.msgs + if m.valid and len(m.positions) == 3 * joints + ] + nan = sum(not np.isfinite(p).all() for p in valid) + zero = sum(bool(np.isfinite(p).all() and not p.any()) for p in valid) + if nan: + out.append( + Problem( + "poses", t.name, f"{nan} valid messages have NaN/Inf positions" + ) + ) + if zero: + out.append( + Problem( + "poses", + t.name, + f"{zero} messages are all-zero with valid=true; lost tracking must be valid=false", + ) + ) + elif t.kind == "eye_gaze": + bad = sum( + m.valid and not all(math.isfinite(v) for v in (m.yaw, m.pitch, m.depth)) + for m in s.msgs + ) + if bad: + out.append( + Problem("poses", t.name, f"{bad} valid messages have NaN/Inf") + ) + elif t.kind == "arm": + names = sum( + len(m.joint_names) != len(m.position) + or (hasattr(m, "velocity") and len(m.velocity) != len(m.position)) + or (hasattr(m, "torque") and len(m.torque) != len(m.position)) + for m in s.msgs + ) + if names: + out.append( + Problem( + "poses", + t.name, + f"{names}/{s.n} messages' joint arrays do not match joint_names", + ) + ) + bad = 0 + for m in s.msgs: + pose = np.asarray(m.ee_pose, dtype=np.float64) + if ( + len(pose) != 16 + or not np.isfinite(pose).all() + or not np.isfinite(np.asarray(m.position)).all() + ): + bad += 1 + continue + T = pose.reshape(4, 4) + R = T[:3, :3] + if ( + np.abs(R.T @ R - np.eye(3)).max() > spec.UNIT_TOL + or abs(np.linalg.det(R) - 1) > spec.UNIT_TOL + or np.abs(T[3] - [0, 0, 0, 1]).max() > spec.UNIT_TOL + ): + bad += 1 + if bad: + out.append( + Problem( + "poses", + t.name, + f"{bad}/{s.n} poses are not a finite rigid 4x4 transform", + ) + ) + elif t.kind == "gripper": + lo, hi = spec.GRIPPER_RANGE + vals = np.array([m.position for m in s.msgs], dtype=np.float64) + bad = int((~np.isfinite(vals) | (vals < lo) | (vals > hi)).sum()) + if bad: + out.append( + Problem( + "poses", + t.name, + f"{bad}/{s.n} gripper values outside [{lo:g}, {hi:g}]", + ) + ) + return out + + +# ---------------------------------------------------------------- 7. tasks & labels + + +def _spans(s: _Stream | None) -> list[tuple[int, int]]: + """(start, end) per message; end 0 (unset) gives end == start: a point event for a + label (one frame wide in Lance), an error for a task.""" + if not s or not s.msgs: + return [] + return [ + (t, int(getattr(m, "end_time_unix_ns", 0)) or t) + for t, m in zip(s.times, s.msgs) + ] + + +def check_labels(ep: _Episode, embodiment: str) -> list[Problem]: + out = [] + streams = _time_series(ep, embodiment) + if not streams: + return out + start, end = _span(streams) + tol = spec.SPAN_TOL_S * 1e9 + tasks = _spans(ep.streams.get("/task")) + labels = _spans(ep.streams.get("/subtask-annotation")) + # A point label becomes one steps-camera frame wide in Lance (as the converter writes + # it), and Lance's annotation review wants the whole span inside a task. + cam = ep.streams.get(spec.profile_for(embodiment).step_clock) + cam_t = np.sort(np.asarray(cam.times if cam else [], dtype=np.int64)) + frame = float(np.median(np.diff(cam_t))) if len(cam_t) > 1 else 0.0 + wide = [(b, e if e > b else b + frame) for b, e in labels] + for topic, spans, point_ok in ( + ("/task", tasks, False), + ("/subtask-annotation", labels, True), + ): + backwards = sum(e < b or (e == b and not point_ok) for b, e in spans) + outside = sum(b < start - tol or e > end + tol for b, e in spans) + if backwards: + out.append( + Problem( + "labels", topic, f"{backwards} spans do not end after they start" + ) + ) + if outside: + out.append( + Problem( + "labels", + topic, + f"{outside}/{len(spans)} spans fall outside the recorded data", + ) + ) + if tasks: + orphans = sum(not any(tb <= b and e <= te for tb, te in tasks) for b, e in wide) + if orphans: + out.append( + Problem( + "labels", + "/subtask-annotation", + f"{orphans}/{len(labels)} labels do not lie inside one /task span", + ) + ) + else: + # The task made from the episode info: [first frame, last frame + one frame). + if len(cam_t) and labels: + first, end = int(cam_t[0]), int(cam_t[-1]) + frame + beyond = sum(b < first or e > end for b, e in wide) + if beyond: + out.append( + Problem( + "labels", + "/subtask-annotation", + f"{beyond}/{len(labels)} labels fall outside " + f"{spec.profile_for(embodiment).step_clock}'s frames (no /task)", + ) + ) + s = ep.streams.get("/subtask-annotation") + types = ( + {m.annotation_type for m in s.msgs if hasattr(m, "annotation_type")} + if s + else set() + ) + bad = sorted(types - {"", *spec.ANNOTATION_TYPES}) + if bad: + out.append( + Problem( + "labels", + "/subtask-annotation", + f"annotation_type {bad}, expected one of {list(spec.ANNOTATION_TYPES)}", + ) + ) + return out + + +# ---------------------------------------------------------------- entry points + + +@dataclass +class Result: + embodiment: str # declared, or guessed from the topics; "" if unreadable + recording_id: str # the Lance recording_id it will get; "" without lab / vendor id + problems: list[Problem] + + @property + def ok(self) -> bool: + return not self.problems + + +def validate_episode(episode: Path) -> Result: + """Validate one episode directory (or its episode.mcap).""" + episode = Path(episode) + path = episode / spec.EPISODE_FILE if episode.is_dir() else episode + if not path.exists(): + return Result( + "", "", [Problem("present", None, f"no {spec.EPISODE_FILE} in {episode}")] + ) + try: + ep = _read(path) + except Exception as e: # truncated or not an MCAP at all + return Result( + "", + "", + [ + Problem( + "present", None, f"cannot read {path.name}: {type(e).__name__}: {e}" + ) + ], + ) + problems = [] + annotation = path.with_name(spec.ANNOTATION_FILE) + if path.name == spec.EPISODE_FILE and annotation.exists(): + try: + extra = _read(annotation) + except Exception as e: + extra = None + problems.append( + Problem( + "present", + None, + f"cannot read {annotation.name}: {type(e).__name__}: {e}", + ) + ) + if extra and extra.metadata: + problems.append( + Problem( + "present", + None, + f"{annotation.name} carries metadata records {sorted(extra.metadata)}; " + f"episode info belongs in {path.name} only", + ) + ) + for topic, s in sorted(extra.streams.items()) if extra else (): + if topic not in LABEL_TOPICS: + problems.append( + Problem( + "present", + topic, + f"only {sorted(LABEL_TOPICS)} belong in {annotation.name}", + ) + ) + elif topic in ep.streams: + problems.append( + Problem( + "present", topic, f"in both {path.name} and {annotation.name}" + ) + ) + else: + ep.streams[topic] = s + embodiment = _embodiment(ep) + for group, check, args in ( + ("present", check_present, (ep, embodiment)), + ("video", check_video, (ep,)), + ("clocks", check_clocks, (ep, embodiment)), + ("calibration", check_calibration, (ep, embodiment)), + ("imu", check_imu, (ep,)), + ("poses", check_poses, (ep, embodiment)), + ("labels", check_labels, (ep, embodiment)), + ): + try: + problems += check(*args) + except Exception as e: # a topic of the wrong type the group's checks assume + problems.append( + Problem( + group, + None, + f"could not check: {type(e).__name__}: {e} (see 'type is' problems)", + ) + ) + lab, vid = ep.info.get("lab", ""), ep.info.get("vendor-recording-id", "") + rid = spec.recording_id(lab, vid) if lab and vid else "" + return Result(embodiment, rid, problems) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument( + "episodes", + nargs="+", + type=Path, + help="episode_/ dirs or episode.mcap files", + ) + args = p.parse_args(argv) + failed = 0 + for episode in args.episodes: + r = validate_episode(episode) + label = ( + f"{episode} (spec {spec.SPEC_VERSION}, embodiment {r.embodiment or '?'}, " + f"recording_id {r.recording_id or '?'})" + ) + if r.problems: + failed += 1 + print(f"FAIL {label}: {len(r.problems)} problem(s)") + for problem in r.problems: + print(f" {problem}") + else: + print(f"PASS {label}") + if len(args.episodes) > 1: + print(f"\n{len(args.episodes) - failed}/{len(args.episodes)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/egomimic/rldb/multi.py b/egomimic/rldb/multi.py new file mode 100644 index 0000000000..9409e45dbc --- /dev/null +++ b/egomimic/rldb/multi.py @@ -0,0 +1,1436 @@ +"""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 json +import logging +import math +import os +import random +import time +from pathlib import Path + +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, + _norm_center_scale, +) +from egomimic.utils.gpu_orphans import die_with_parent +from egomimic.utils.pose_utils import ( + bimanual_cartesian_layout, + bimanual_keypoint_layout, + hand_channel_mask, + 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 _bounds_check_channels(zarr_key: str, width: int) -> list[int] | None: + """Channels of a known bimanual layout that quantile bounds apply to + (translation, gripper, keypoints), or ``None`` to check the full vector. + Rotation channels (Euler or rot6d columns) are excluded.""" + if zarr_key in ("actions_cartesian", "observations.state.ee_pose"): + layout = bimanual_cartesian_layout(width) + if layout is not None: + return list(layout["xyz"]) + list(layout["grip"]) + elif zarr_key in ("actions_keypoints", "observations.state.keypoints"): + layout = bimanual_keypoint_layout(width) + if layout is not None: + return sorted(list(layout["wrist_xyz"]) + list(layout["keypoints"])) + return None + + +HAND_FLAGS = {"action_keys": "action_hand_mask", "proprio_keys": "proprio_hand_valid"} +# Norm-sample key holding the hand_keep mask of norm key k, as f"{k}{KEEP_SUFFIX}". +KEEP_SUFFIX = "#keep" + + +def hand_keep(data: dict, key_type: str | None, zarr_key: str, value): + """Bool mask, broadcastable to ``value``, of the channels of a tracked hand, from + the sample's (or batch's) ``action_hand_mask`` / ``proprio_hand_valid`` [left, + right] flags; None when the sample has no flags or ``value`` has no hand layout.""" + flags = data.get(HAND_FLAGS.get(key_type)) + if flags is None: + return None + flags = torch.as_tensor(flags) + width = value.shape[-1] + side = zarr_key.split(".", 1)[0] + if side in ("left", "right"): + channels = flags[..., [0 if side == "left" else 1] * width] + else: + channels = hand_channel_mask(flags, width) + if channels is None: + return None + # Flags are per sample (actions) or per history step (proprio); a key without + # that step axis takes the newest step. + while channels.ndim > value.ndim: + channels = channels[..., -1, :] + while channels.ndim < value.ndim: + channels = channels.unsqueeze(-2) + return channels > 0.5 + + +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.""" + for leaf in leaves: + for sub in ( + MultiDataset._iter_leaves(leaf) + if isinstance(leaf, MultiDataset) + else [leaf] + ): + if not getattr(sub, "fork_safe", True): + 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 + + 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.source_manifest: dict | None = None + 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) + + def episode(self, episode_hash: str): + """The leaf keyed `episode_hash`, or a Lance leaf narrowed to the aliased clip of + that hash; None when neither is here.""" + if episode_hash in self.datasets: + return self.datasets[episode_hash] + for leaf in self.datasets.values(): + if episode_hash in getattr(leaf, "clips", {}): + return leaf.clip(episode_hash) + return None + + @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 + # 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 + + # An untracked hand's channels hold fill values; the loss and the + # stats leave them out, so they cannot reject a sample either. + keep = None + if any(f in data for f in HAND_FLAGS.values()): + key_type = self.key_types.get(embodiment_id, {}).get(key_name) + keep = hand_keep(data, key_type, zarr_key, arr) + keep = ( + torch.ones_like(arr, dtype=torch.bool) + if keep is None + else torch.broadcast_to(keep.to(arr.device), arr.shape) + ) + if torch.any(~torch.isfinite(arr) & keep): + 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] + keep = keep[..., check_idx] + else: + arr_q = arr + + # Slack on the quantile bounds: a relative part (a fraction of + # each cell's own quantile range) plus a 1e-6 absolute floor. + # The bounds are per (timestep, channel) cell and the check is + # any-cell, so with 13 200 cells per keypoint chunk a strict + # compare against a subsampled stats pass rejects valid frames in + # bulk. Measured on the 2815-frame flagship episode through the + # real keypoints_wristframe_6d pipeline, norm_mode=zscore, stats + # from infer_norm_from_dataset(sample_frac=0.10) over its own + # frames, 403 probe frames, a frame counted as rejected if ANY of + # actions_keypoints / state.keypoints / state.ee_pose fails: + # slack 0 0.05 0.1 0.25 0.5 1.0 + # any 39.5% 17.9% 11.4% 5.2% 3.2% 1.0% + # The 132-channel action chunk drives it (39.0 % -> 3.2 %); the 8 + # checked ee_pose channels never exceed 2.7 %. Which frames the + # stats came from moves the absolute numbers, so re-measure the + # same way before comparing. Corrupt values + # (fill constants, wrong-frame data) sit orders of magnitude + # outside the range and are still caught. + tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low) + below = (arr_q < q_low - tol) & keep + above = (arr_q > q_high + tol) & keep + # 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. Which stat decides that is + # norm_mode's business -- std under zscore, the quantile range + # under quantile -- and every shipped config is zscore, so reading + # q99 - q1 here would exempt cells the normalizer still divides by. + 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: + _, scale = _norm_center_scale( + stats, self.norm_mode, torch.empty(0, dtype=torch.float32) + ) + except (KeyError, ValueError): + return None + return torch.broadcast_to(_degenerate(scale, self.norm_mode), shape) + + # Retries stay within the failing episode this many times, then widen to + # the full index space: a single wholly-bad episode must not exhaust the + # sampler and take the run down (NCCL-timeouts the other ranks). + GLOBAL_FALLBACK_ATTEMPTS = 25 + # Cap on consecutive bad samples, after which the worker raises. It is a + # systemic-failure detector, not a tolerance dial: `attempts` resets per + # __getitem__, so with p = the bad fraction of the whole index space the + # cap fires with probability p**(N - GLOBAL_FALLBACK_ATTEMPTS), a step + # function in p. At N = 100 the firing band sits at p >~ 0.8, so anything + # it can detect is systemic; at N = 1000 it needed ~99 % of the store to be + # bad, and got there via 975 sequential zarr/S3 reads per call per worker, + # which looks like a hang rather than a failure. + MAX_FALLBACK_ATTEMPTS = 100 + # Episodes tried per key_map/transform group before populate_from_datasets + # falls back to the raw key_map. + POPULATE_PROBES = 8 + + 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 | None, 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}" + ) + if attempts <= self.GLOBAL_FALLBACK_ATTEMPTS: + candidates = [ + c for c in self._global_indices_by_dataset[dataset_name] if c != idx + ] + else: + candidates = None + if candidates: + next_idx = random.choice(candidates) + else: + next_idx = random.randrange(len(self.index_map)) + next_dataset_name, next_local_idx = self.index_map[next_idx] + # Every retry used to log at WARNING: tens of thousands of lines per + # run. Keep the first few per process, then one in every 500. + self._retry_log_count = getattr(self, "_retry_log_count", 0) + 1 + if self._retry_log_count <= 20 or self._retry_log_count % 500 == 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 LocalEpisodeResolver can sync; other resolvers reject the kwarg. + extra = {"sync_from_s3": True} if sync_from_s3 else {} + resolved = resolver.resolve( + filters=filters, expected_embodiment=expected, **extra + ) + + pins = filters.episode_hashes if isinstance(filters, DatasetFilter) else () + # Lance keys leaves by recording_id; a pin may name the recording or an old clip. + keys = getattr(resolver, "dataset_keys", lambda ps: {p: p for p in ps})(pins) + dropped = sorted(p for p in pins if keys.get(p) not in 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)." + ) + + ds = cls(datasets=resolved, **kwargs) + manifest = getattr(resolver, "manifest", None) + # Which table versions the leaves came from; trainHydra writes it to the run dir. + ds.source_manifest = manifest() if manifest is not None else None + return ds + + # ===================================================================== + # 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, + ): + """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. + + ``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 not all( + not isinstance(ds, MultiDataset) + and getattr(ds, "episode_path", None) is not 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, + key_types={ + k: self.key_types.get(embodiment, {}).get(k) for k in zarr_keys + }, + seed=seed, + num_workers=num_workers, + ) + 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 + ) + keeps = {k: collected.pop(f"{k}{KEEP_SUFFIX}", None) for k in norm_keys} + 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, + self._keep_rows(collected[k], keeps.get(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) + flags = {f: batch[f][:take] for f in HAND_FLAGS.values() if f in batch} + 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] + keep = hand_keep( + flags, self.key_types.get(embodiment, {}).get(k), zarr_key, x + ) + if hasattr(x, "detach"): + x = x.detach().cpu().numpy() + collected[k].append(x) + collected.setdefault(f"{k}{KEEP_SUFFIX}", []).append( + None + if keep is None + else np.broadcast_to(keep.numpy(), x.shape).copy() + ) + cur += take + pbar.update(take) + return collected + + @staticmethod + def _keep_rows(chunks: list, keeps: list | None) -> np.ndarray | None: + """The hand_keep masks of the norm sample ``chunks``, concatenated (all True + for a chunk without one); None when no chunk has one.""" + if not keeps or all(m is None for m in keeps): + return None + return np.concatenate( + [ + np.ones(x.shape, bool) if m is None else m + for x, m in zip(chunks, keeps, strict=True) + ], + axis=0, + ) + + @staticmethod + def _drop_nonfinite_rows( + X: np.ndarray, key: str, keep: np.ndarray | None = None + ) -> np.ndarray: + """Drop norm samples (rows of X, shape (N, ...)) containing NaN/Inf. + + ``keep`` (X's shape) marks the cells of a tracked hand: the others are set + to NaN instead and left out of the stats per cell, so an untracked hand + neither drops the row nor shapes its channels' stats. + + 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. + """ + ok = np.isfinite(X) if keep is None else np.isfinite(X) | ~keep + finite = ok.reshape(X.shape[0], -1).all(axis=1) + n_bad = int((~finite).sum()) + if keep is not None: + X = np.where(keep, X, np.nan) + untracked = np.isnan(X[finite]).all(axis=0) + if finite.any() and untracked.any(): + raise ValueError( + f"[MultiDataset] key={key}: {int(untracked.sum())} cells have no " + "sample from a tracked hand" + ) + 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): + if np.isnan(X).any(): + # Untracked-hand cells (see _drop_nonfinite_rows); much slower than below. + return { + "mean": np.nanmean(X, axis=0), + "std": np.nanstd(X, axis=0), + "min": np.nanmin(X, axis=0), + "max": np.nanmax(X, axis=0), + "median": np.nanmedian(X, axis=0), + "quantile_1": np.nanpercentile(X, 1, axis=0), + "quantile_99": np.nanpercentile(X, 99, axis=0), + "quantile_0_01": np.nanpercentile(X, 0.01, axis=0, method="lower"), + "quantile_99_99": np.nanpercentile(X, 99.99, axis=0, method="higher"), + } + 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, so the extreme frames of the very episodes the stats were + # computed on were rejected at train time (3-16 % of a single + # episode in the overfit runs). + "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) + if key_type == "proprio_keys": + keep = hand_keep(data, key_type, zarr_key, out[zarr_key]) + if keep is not None: + # An untracked hand reads as the dataset mean pose instead of + # its fill values: 0 on normalized channels, the mean itself on + # pass-through ones (rot6d), where 0 is no rotation at all. + mean = torch.as_tensor(stats["mean"], dtype=out[zarr_key].dtype) + fill = self._apply_norm_one(mean, stats, zarr_key) + out[zarr_key] = torch.where(keep, out[zarr_key], fill) + 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 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 base.episode(h) is None] + 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.episode(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/norm/episode_norm_samples.py b/egomimic/rldb/norm/episode_norm_samples.py new file mode 100644 index 0000000000..dd473084fc --- /dev/null +++ b/egomimic/rldb/norm/episode_norm_samples.py @@ -0,0 +1,279 @@ +"""Per-episode norm-stat samples. + +Each episode's sampled proprio/action rows are cached on their own, keyed by +what a sample is (resolver keymap + transforms + reader code) but not by which +split the episode is in. A split's stats are computed from the union of its +episodes' rows, so a new split (20 % -> full, an operator hold-out, newly +collected data) only reads the episodes it has not seen before. + +Frames are taken on a power-of-two stride with a fixed per-episode phase, so +the rows at stride 2s are a subset of the rows at stride s: one cached file +serves every coarser stride, and a split only resamples an episode when it needs +it denser than it was ever cached. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import uuid +from collections.abc import Mapping +from pathlib import Path + +import numpy as np +import torch +from tqdm import tqdm + +from egomimic.rldb.multi import KEEP_SUFFIX, hand_keep, worker_context +from egomimic.rldb.norm import norm_cache +from egomimic.utils.gpu_orphans import die_with_parent + +logger = logging.getLogger(__name__) + +# Bump when the frame selection or the stored layout changes. +VERSION = 1 + + +def recipe_key(dataset_name: str, dataset_cfg) -> str: + """Identity of a sample: the resolver's keymap and transforms, the dataset + class it loads and the reader code. Filters, mode and split ratios only pick + episodes, so they are left out.""" + raw = norm_cache._to_dict(dataset_cfg) + resolver = {"resolver": raw.get("resolver") or {}} + inputs = { + "v": VERSION, + "dataset": dataset_name, + "resolver": norm_cache.recipe_inputs(resolver)["resolver"], + "code": norm_cache.code_hash(resolver), + } + blob = json.dumps(inputs, sort_keys=True, default=str) + return hashlib.sha256(blob.encode("utf-8")).hexdigest() + + +def cache_root(cache_dir, dataset_name: str, dataset_cfg) -> Path: + return ( + Path(cache_dir) + / "episodes" + / dataset_name + / recipe_key(dataset_name, dataset_cfg)[:16] + ) + + +def stride_for(total_frames: int, n_target: int) -> int: + """Largest power of two that still yields at least ``n_target`` frames.""" + s = 1 + while total_frames // (2 * s) >= max(n_target, 1): + s *= 2 + return s + + +def _phase(episode_hash: str) -> int: + return int(hashlib.sha256(episode_hash.encode("utf-8")).hexdigest()[:15], 16) + + +def frame_indices(episode_hash: str, n_frames: int, stride: int) -> np.ndarray: + return np.arange(_phase(episode_hash) % stride, n_frames, stride, dtype=np.int64) + + +def _file(root: Path, episode_hash: str, fingerprint: str, stride: int) -> Path: + return root / f"{episode_hash}.{fingerprint[:16]}.s{stride}.npz" + + +def find(root: Path, episode_hash: str, fingerprint: str, stride: int) -> Path | None: + """The coarsest cached file whose rows include every frame at ``stride``.""" + best = None + for p in root.glob(f"{episode_hash}.{fingerprint[:16]}.s*.npz"): + try: + s = int(p.name.rsplit(".s", 1)[1][: -len(".npz")]) + except ValueError: + continue + if stride % s == 0 and (best is None or s > best[0]): + best = (s, p) + return None if best is None else best[1] + + +def _write(path: Path, idx: list[int], rows: dict[str, list[np.ndarray]]) -> None: + keys = sorted(rows) + arrays = {"idx": np.asarray(idx, dtype=np.int64), "keys": np.asarray(keys)} + for i, k in enumerate(keys): + arrays[f"k{i}"] = np.stack(rows[k]) + # Several jobs (and every DDP rank) can share the cache, so each writer + # gets its own tmp file. + tmp = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") + try: + path.parent.mkdir(parents=True, exist_ok=True) + with open(tmp, "wb") as f: + np.savez(f, **arrays) + os.replace(tmp, path) + except OSError as e: + logger.warning("episode norm cache write failed (%s): %s", path, e) + tmp.unlink(missing_ok=True) + + +def _read(path: Path, rows: np.ndarray) -> dict[str, np.ndarray]: + with np.load(path) as z: + return {str(k): z[f"k{i}"][rows] for i, k in enumerate(z["keys"])} + + +def _to_numpy(x): + if hasattr(x, "detach"): + return x.detach().cpu().numpy() + return np.asarray(x) + + +class _Frames(torch.utils.data.Dataset): + def __init__(self, leaves, tasks, zarr_keys: Mapping[str, str], key_types=None): + self.leaves = leaves + self.tasks = tasks + self.zarr_keys = dict(zarr_keys) + self.key_types = dict(key_types or {}) + + def __len__(self): + return len(self.tasks) + + def __getitem__(self, i): + h, idx = self.tasks[i] + try: + data = self.leaves[h][idx] + except Exception as e: + # ZarrDataset already retried inside the episode; this one is bad. + logger.warning(f"norm sample failed ({type(e).__name__}: {e}) {h}[{idx}]") + return h, idx, None + row = {} + for k, z in self.zarr_keys.items(): + if z not in data: + continue + row[k] = _to_numpy(data[z]) + keep = hand_keep(data, self.key_types.get(k), z, data[z]) + if keep is not None: + row[f"{k}{KEEP_SUFFIX}"] = np.broadcast_to( + keep.numpy(), row[k].shape + ).copy() + return h, idx, row + + +def _as_list(batch): + return batch + + +def _sample_episodes( + root: Path, + leaves: Mapping, + fingerprints: Mapping[str, str], + zarr_keys: Mapping[str, str], + stride: int, + num_workers: int, + key_types: Mapping[str, str | None] | None = None, +) -> None: + tasks = [ + (h, int(i)) + for h in sorted(leaves) + for i in frame_indices(h, len(leaves[h]), stride) + ] + pending = {h: 0 for h in leaves} + for h, _ in tasks: + pending[h] += 1 + # An episode with no frames at this stride still gets an (empty) file. + for h in [h for h, n in pending.items() if n == 0]: + _write(_file(root, h, fingerprints[h], stride), [], {}) + del pending[h] + if not tasks: + return + + loader = torch.utils.data.DataLoader( + _Frames(leaves, tasks, zarr_keys, key_types), + batch_size=64, + num_workers=num_workers, + collate_fn=_as_list, + worker_init_fn=die_with_parent if num_workers > 0 else None, + multiprocessing_context=( + worker_context(leaves.values()) if num_workers > 0 else None + ), + ) + idx_buf: dict[str, list[int]] = {} + row_buf: dict[str, dict[str, list[np.ndarray]]] = {} + with tqdm(total=len(tasks), unit="sample", desc="norm samples") as pbar: + for batch in loader: + for h, idx, sample in batch: + if sample is not None: + idx_buf.setdefault(h, []).append(idx) + for k, v in sample.items(): + row_buf.setdefault(h, {}).setdefault(k, []).append(v) + pending[h] -= 1 + if pending[h] == 0: + _write( + _file(root, h, fingerprints[h], stride), + idx_buf.pop(h, []), + row_buf.pop(h, {}), + ) + pbar.update(len(batch)) + + +def collect( + root: Path, + leaves: Mapping, + zarr_keys: Mapping[str, str], + n_target: int, + *, + seed: int = 42, + num_workers: int = 4, + key_types: Mapping[str, str | None] | None = None, +) -> tuple[dict[str, list[np.ndarray]], dict]: + """``n_target`` rows per norm key (fewer if the episodes hold fewer frames), + drawn uniformly from ``leaves`` (episode hash -> ZarrDataset). Episodes + missing from ``root`` at the needed density are sampled and cached first.""" + root = Path(root) + fingerprints = {h: norm_cache.leaf_fingerprint(ds) for h, ds in leaves.items()} + lengths = {h: len(ds) for h, ds in leaves.items()} + stride = stride_for(sum(lengths.values()), n_target) + + files = {h: find(root, h, fingerprints[h], stride) for h in leaves} + missing = {h: leaves[h] for h, p in files.items() if p is None} + logger.info( + f"[norm samples] stride {stride}: {len(leaves) - len(missing)}/{len(leaves)} " + f"episodes cached under {root}, sampling {len(missing)}" + ) + if missing: + _sample_episodes( + root, missing, fingerprints, zarr_keys, stride, num_workers, key_types + ) + for h in missing: + files[h] = _file(root, h, fingerprints[h], stride) + + order = sorted(leaves) + at_stride = {} + for h in order: + phase = _phase(h) % stride + with np.load(files[h]) as z: + at_stride[h] = np.flatnonzero((z["idx"] - phase) % stride == 0) + total = sum(len(v) for v in at_stride.values()) + if total > n_target: + keep = np.zeros(total, bool) + keep[np.random.default_rng(seed).choice(total, n_target, replace=False)] = True + offset = 0 + for h in order: + n = len(at_stride[h]) + at_stride[h] = at_stride[h][keep[offset : offset + n]] + offset += n + + collected: dict[str, list] = {k: [] for k in zarr_keys} + for h in order: + if not len(at_stride[h]): + continue + rows = _read(files[h], at_stride[h]) + for k in zarr_keys: + if k in rows: + collected[k].append(rows[k]) + # Aligned with collected[k]: None for an episode without hand flags. + collected.setdefault(f"{k}{KEEP_SUFFIX}", []).append( + rows.get(f"{k}{KEEP_SUFFIX}") + ) + meta = { + "stride": stride, + "episodes": len(leaves), + "episodes_sampled": len(missing), + "frames": int(min(total, n_target)), + } + return collected, meta diff --git a/egomimic/rldb/zarr/norm_cache.py b/egomimic/rldb/norm/norm_cache.py similarity index 89% rename from egomimic/rldb/zarr/norm_cache.py rename to egomimic/rldb/norm/norm_cache.py index f706a62610..563e4bef77 100644 --- a/egomimic/rldb/zarr/norm_cache.py +++ b/egomimic/rldb/norm/norm_cache.py @@ -31,7 +31,8 @@ logger = logging.getLogger(__name__) -KEY_VERSION = 2 +# 3: rows come from the per-episode sample store (episode_norm_samples). +KEY_VERSION = 3 # Resolver fields that say where the data lives, not what it is. _LOCATION_FIELDS = ( "folder_path", @@ -43,8 +44,22 @@ # 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()`` (every ``EpisodeLeaf`` has one), else + ``episode_fingerprint`` of its zarr directory.""" + own = getattr(ds, "cache_fingerprint", None) + if own is not None: + return own() + return episode_fingerprint(getattr(ds, "episode_path", None)) def episode_fingerprint(episode_path) -> str | None: @@ -91,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() @@ -153,6 +168,8 @@ def cache_inputs( episodes: Mapping[str, str | None], dataset_cfg, sample_frac: float, + pool_horizon: bool = False, + max_samples: int | None = None, ) -> dict: """Everything the key hashes. ``episodes`` maps episode hash to its ``episode_fingerprint``. norm_mode is not an input: every stat is computed @@ -164,6 +181,8 @@ def cache_inputs( "episodes": sorted([str(h), fp] for h, fp in episodes.items()), "recipe": recipe_inputs(cfg), "sample_frac": float(sample_frac), + "pool_horizon": bool(pool_horizon), + "max_samples": None if max_samples is None else int(max_samples), "code": code_hash(cfg), } diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/transforms.py similarity index 74% rename from egomimic/rldb/zarr/action_chunk_transforms.py rename to egomimic/rldb/transforms.py index 62a0f91452..9150df551b 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/transforms.py @@ -33,6 +33,7 @@ _xyzwxyz_to_matrix, _xyzypr_to_matrix, _ypr_to_rot6d, + bimanual_keypoint_layout, wxyz_to_xyzw, xyzw_to_wxyz, ) @@ -139,6 +140,53 @@ def transform(self, batch: dict) -> dict: return batch +class InterpolatePadMask(Transform): + """Carry the dataset's action pad mask through the action interpolation. + + The interpolators resample a (T, D) chunk from ``np.linspace(0, 1, T)`` + onto ``np.linspace(0, 1, new_chunk_length)``; an output step is real only + if BOTH source frames it interpolates between are real (a step landing + exactly on a source frame inherits that frame). No-op when the mask key is + absent, and idempotent, so a mode that chains two builders + (``include_ee_pose``) still upsamples the mask once. + """ + + def __init__( + self, + new_chunk_length: int, + mask_key: str = "action_pad_mask", + stride: int = 1, + ): + if stride <= 0: + raise ValueError(f"stride must be positive, got {stride}") + self.new_chunk_length = new_chunk_length + self.mask_key = mask_key + self.stride = int(stride) + + def transform(self, batch: dict) -> dict: + mask = batch.get(self.mask_key) + if mask is None: + return batch + mask = np.asarray(mask, dtype=np.float32).reshape(-1) + if mask.shape[0] == self.new_chunk_length: + return batch # already at the model horizon + mask = mask[:: self.stride] + T = mask.shape[0] + if T < 2: + batch[self.mask_key] = np.full( + self.new_chunk_length, mask[0] if T else 0.0, dtype=np.float32 + ) + return batch + # Position of each output sample in source-frame index units. + pos = np.linspace(0, 1, self.new_chunk_length) * (T - 1) + nearest = np.rint(pos) + on_frame = np.abs(pos - nearest) < 1e-9 + lo = np.where(on_frame, nearest, np.floor(pos)).astype(int) + hi = np.where(on_frame, nearest, np.ceil(pos)).astype(int) + batch[self.mask_key] = ((mask[lo] > 0) & (mask[hi] > 0)).astype(np.float32) + return batch + + # --------------------------------------------------------------------------- # Coordinate Transforms # --------------------------------------------------------------------------- @@ -153,6 +201,7 @@ def __init__( extra_batch_key: dict = None, mode: Literal["xyz", "xyzwxyz", "xyzypr"] = "xyzwxyz", inverse: bool = True, + target_history: bool = False, ): """ args: @@ -160,6 +209,10 @@ def __init__( chunk_world: transformed_key_name: is_quat: if True, inputs are xyz + quat(wxyz); otherwise xyz + ypr. + target_history: the target frame is a PROPRIO key that may be read + with history, ``(K, D)`` with the current frame last. Take that + current frame; without this flag a 2-D target is an error (a + target frame is a single pose). """ self.target_world = target_world self.chunk_world = chunk_world @@ -167,6 +220,7 @@ def __init__( self.extra_batch_key = extra_batch_key self.mode = mode self.inverse = inverse + self.target_history = target_history def transform(self, batch): """ @@ -189,6 +243,11 @@ def transform(self, batch): chunk_world = np.asarray(batch[self.chunk_world]) chunk_world_shape = None + if self.target_history and target_world.ndim > 1: + # Proprio read with history: (K, D), current frame last. The chunk + # is always expressed in the CURRENT step's frame. + target_world = target_world[-1] + if chunk_world.ndim > 2: chunk_world_shape = chunk_world.shape chunk_world = chunk_world.reshape(-1, chunk_world_shape[-1]) @@ -240,7 +299,11 @@ def transform(self, batch): class QuaternionPoseToYPR(Transform): - """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr.""" + """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr. + + A leading time axis (``(K, 7)``, a proprio key read with history) is + treated as a batch axis and comes back as ``(K, 6)``. + """ def __init__(self, pose_key: str, output_key: str): self.pose_key = pose_key @@ -248,15 +311,17 @@ def __init__(self, pose_key: str, output_key: str): def transform(self, batch: dict) -> dict: pose = np.asarray(batch[self.pose_key]) - if pose.shape != (7,): + if pose.ndim > 2 or pose.shape[-1] != 7: raise ValueError( - f"QuaternionPoseToYPR expects shape (7,), got {pose.shape} for key " - f"'{self.pose_key}'" + f"QuaternionPoseToYPR expects shape (7,) or (K, 7), got " + f"{pose.shape} for key '{self.pose_key}'" ) - xyz = pose[:3] - xyzw = wxyz_to_xyzw(pose[3:7]) + flat = pose.reshape(-1, 7) + xyz = flat[:, :3] + xyzw = wxyz_to_xyzw(flat[:, 3:7]) ypr = R.from_quat(xyzw).as_euler("ZYX", degrees=False) - batch[self.output_key] = np.concatenate([xyz, ypr], axis=0) + out = np.concatenate([xyz, ypr], axis=-1) + batch[self.output_key] = out.reshape(*pose.shape[:-1], 6) return batch @@ -332,29 +397,71 @@ def __init__( pose_world: str, transformed_key_name: str, mode: Literal["xyzwxyz", "xyzypr", "xyz"] = "xyzwxyz", + target_history: bool = False, + per_step_target: bool = True, ): + """ + args: + per_step_target: put step k's pose in step k's OWN target frame + rather than every step in the current step's frame. On by + default, and inert unless the target itself was read with + history and the pose carries the matching axis -- which is the + only shape for which the other answer is wrong, not merely + different. + + For wrist-frame keypoints this is the difference between "where + the hand was, measured from today's wrist" and "what shape the + hand was in" -- only the latter is a state the norm stats + describe, and only it keeps keypoint 0 at the origin on every + step. Pass False to force the single-frame path; the ACTION + chunk does not come through here at all, since its leading axis + is the forward horizon and belongs in one current frame. + """ self.target_world = target_world self.pose_world = pose_world self.transformed_key_name = transformed_key_name self.mode = mode + self.target_history = target_history + self.per_step_target = per_step_target self._chunk_transform = ActionChunkCoordinateFrameTransform( target_world=target_world, chunk_world=pose_world, transformed_key_name=transformed_key_name, mode=mode, + target_history=target_history, ) - def transform(self, batch: dict) -> dict: - pose_world = np.asarray(batch[self.pose_world]) + def _one_step(self, target_pose, pose_world) -> np.ndarray: + """``pose_world`` in the single frame ``target_pose``.""" transformed = self._chunk_transform.transform( { - self.target_world: batch[self.target_world], + self.target_world: target_pose, self.pose_world: pose_world[None, :], } ) - batch[self.transformed_key_name] = np.asarray( - transformed[self.transformed_key_name] - )[0] + return np.asarray(transformed[self.transformed_key_name])[0] + + def transform(self, batch: dict) -> dict: + pose_world = np.asarray(batch[self.pose_world]) + target_world = np.asarray(batch[self.target_world]) + per_step = ( + self.per_step_target + and self.target_history + and target_world.ndim > 1 + and pose_world.ndim == target_world.ndim + 1 + and pose_world.shape[0] == target_world.shape[0] + ) + if per_step: + batch[self.transformed_key_name] = np.stack( + [ + self._one_step(target_world[k], pose_world[k]) + for k in range(target_world.shape[0]) + ] + ) + return batch + batch[self.transformed_key_name] = self._one_step( + batch[self.target_world], pose_world + ) return batch @@ -410,6 +517,10 @@ class CartesianYPRToRot6D(Transform): convention of the pi0.5 32D action blocks, so the resulting per-arm layout packs into them with no rotation math in the model. + Apply it to the proprio ee_pose as well as the actions: normalized YPR + saturates yaw/roll at +-pi, so per-dim normalization is only meaningful on + the continuous representation. + Input layouts (last dim), for action chunks ``(T, D)`` and single proprio poses ``(D,)`` alike: 12 -> [L xyz ypr, R xyz ypr] -> 18 [L xyz 6d, R xyz 6d] @@ -492,6 +603,74 @@ def transform(self, batch: dict) -> dict: return batch +class KeypointsYPRToRot6D(Transform): + """Convert a wrist-first bimanual keypoint vector from a ypr wrist rotation + to the continuous 6D one: per hand + ``[wrist xyz (3) | wrist ypr (3) | 21 keypoints (63)]`` (138 total) -> + ``[wrist xyz (3) | wrist rot6d (6) | 21 keypoints (63)]`` (144 total). + The keypoint blocks are copied through untouched. Works on action chunks + ``(T, 138)`` and single proprio vectors ``(138,)``. + """ + + def __init__( + self, action_key: str = "actions_keypoints", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or action_key + + def transform(self, batch: dict) -> dict: + arr, is_tensor = _as_array(batch[self.action_key]) + if arr.shape[-1] != 138: + raise ValueError( + f"KeypointsYPRToRot6D expects last-dim 138, got {arr.shape} for " + f"'{self.action_key}'" + ) + layout = bimanual_keypoint_layout(138) + per_hand, n_rot = layout["per_hand"], len(layout["rot"]) // 2 + blocks = [] + for hand in range(2): + o = hand * per_hand + xyz = len(layout["wrist_xyz"]) // 2 + blocks += [ + arr[..., o : o + xyz], + _ypr_to_rot6d(arr[..., o + xyz : o + xyz + n_rot]), + arr[..., o + xyz + n_rot : o + per_hand], + ] + batch[self.output_key] = _like_input(np.concatenate(blocks, axis=-1), is_tensor) + return batch + + +class KeypointsRot6DToYPR(Transform): + """Inverse of :class:`KeypointsYPRToRot6D`: 144 -> 138 (wrist rot6d -> + ypr via Gram-Schmidt, keypoints untouched).""" + + def __init__( + self, action_key: str = "actions_keypoints", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or action_key + + def transform(self, batch: dict) -> dict: + arr, is_tensor = _as_array(batch[self.action_key]) + layout = bimanual_keypoint_layout(arr.shape[-1]) + if layout is None or arr.shape[-1] != 144: + raise ValueError( + f"KeypointsRot6DToYPR expects last-dim 144, got {arr.shape} for " + f"'{self.action_key}'" + ) + per_hand = layout["per_hand"] + blocks = [] + for hand in range(2): + o = hand * per_hand + blocks += [ + arr[..., o : o + 3], + _rot6d_to_ypr(arr[..., o + 3 : o + 9]), + arr[..., o + 9 : o + per_hand], + ] + batch[self.output_key] = _like_input(np.concatenate(blocks, axis=-1), is_tensor) + return batch + + class RotateLocalFrame(Transform): """Right-multiply listed xyz+quat(wxyz) pose keys by a constant LOCAL rotation: ``R_new = R_old @ R_fix``. Relabels the pose's own axes without @@ -653,6 +832,32 @@ def transform(self, batch: dict) -> dict: return batch +class SelectCurrentStep(Transform): + """Drop a proprio key's observation-history axis, keeping the CURRENT step. + + A proprio key read with ``proprio_history`` K > 1 is ``(K, D)`` per sample + with the current frame last. The revert / viz pipelines want the single + current pose the forward pipeline used as the frame, so they start by + collapsing it. ``ndim`` is the flat per-sample rank of the key (1 for a + ``(D,)`` pose vector); a value already at that rank is left alone, so the + K = 1 path is bit-for-bit unchanged and the transform is idempotent. + """ + + def __init__(self, keys: list[str], ndim: int = 1): + self.keys = list(keys) + self.ndim = ndim + + def transform(self, batch: dict) -> dict: + for key in self.keys: + value = batch.get(key) + if value is None or not hasattr(value, "ndim"): + continue + while value.ndim > self.ndim: + value = value[..., -1, :] + batch[key] = value + return batch + + class ConcatKeys(Transform): def __init__(self, key_list, new_key_name, delete_old_keys=False): self.key_list = list(key_list) @@ -734,13 +939,22 @@ def transform(self, batch: dict) -> dict: class Reshape(Transform): + """Reshape a key to ``shape``. A value holding a whole multiple of + ``shape`` (a proprio key read with history, ``(K, ...)``) keeps its leading + time axis and is reshaped per step. + """ + def __init__(self, input_key: str, output_key: str, shape: tuple): self.input_key = input_key self.output_key = output_key self.shape = shape def transform(self, batch: dict) -> dict: - batch[self.output_key] = batch[self.input_key].reshape(*self.shape) + value = batch[self.input_key] + n = value.numel() if isinstance(value, torch.Tensor) else np.asarray(value).size + target = int(np.prod(self.shape)) + shape = self.shape if n == target else (-1, *self.shape) + batch[self.output_key] = value.reshape(*shape) return batch diff --git a/egomimic/rldb/zarr/zarr_dataset_action_expert.py b/egomimic/rldb/zarr/zarr_dataset_action_expert.py index d211ca9156..e507912b33 100644 --- a/egomimic/rldb/zarr/zarr_dataset_action_expert.py +++ b/egomimic/rldb/zarr/zarr_dataset_action_expert.py @@ -18,9 +18,9 @@ import torch from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.multi import MultiDataset 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 833f11741f..22a7a6377f 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -20,37 +20,39 @@ from __future__ import annotations -import copy 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 ( canonical_embodiment_name, - get_embodiment_id, ) - -# from action_chunk_transforms import Transform -from egomimic.rldb.filters import DatasetFilter -from egomimic.rldb.resolve_memo import memoized -from egomimic.utils.action_utils import ( - _apply_norm_one, - _apply_unnorm_one, +from egomimic.rldb.filters import DatasetFilter, PinError, _ensure_dataset_filter +from egomimic.rldb.leaf import EpisodeLeaf + +# Re-exported: Hydra `_target_`s (MultiDataset._from_resolver) and older imports +# name these here. +from egomimic.rldb.multi import ( # noqa: F401 + SEED, + EvenStrideDataset, + MultiDataset, + _IndexMap, + pinned_episode_subset, + split_dataset_names, + worker_context, ) +from egomimic.rldb.resolve_memo import memoized from egomimic.utils.env import load_env @@ -66,17 +68,8 @@ 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__) -SEED = 42 - def _episode_table() -> pd.DataFrame: return memoized( @@ -84,48 +77,6 @@ def _episode_table() -> pd.DataFrame: ) -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 @@ -156,38 +107,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], @@ -215,30 +134,27 @@ 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: 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 - ) + # First row per pinned hash, as dicts: a per-pin `.loc` took ~17 s for 41k + # pins. + by_hash = {} + if "episode_hash" in df.columns: + rows = df[df["episode_hash"].astype(str).isin(pins)] + by_hash = { + str(r["episode_hash"]): r + for r in rows.drop_duplicates("episode_hash").to_dict("records") + } problems: list[str] = [] for h in sorted(pins): - if by_hash is None or h not in by_hash.index: + if h not in by_hash: problems.append(f"{h}: not in app.episodes") continue - row = by_hash.loc[h] + row = by_hash[h] if not _is_missing_filter_value(row.get("is_deleted")) and bool( row["is_deleted"] ): @@ -482,7 +398,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( @@ -547,11 +463,16 @@ def _get_filtered_paths( logger.info("Episode table is empty.") return [] - mask = df.apply( - lambda row: filters.matches(_normalize_filter_row(row.to_dict())), - axis=1, - ) - matched = df.loc[mask] + if filters.episode_hashes: + # matches() rejects every other row anyway; the table has ~10x + # more rows than a full-lab pin list and 1e5x more than a val pin. + df = df[df["episode_hash"].astype(str).isin(filters.episode_hashes)] + # records, not a row-wise df.apply(row.to_dict()): ~20x faster on the + # full table, where the apply cost ~50 s per resolve. + mask = [ + filters.matches(_normalize_filter_row(row)) for row in df.to_dict("records") + ] + matched = df.loc[mask] if mask else df n_matched_sql = len(matched) matched = matched[ @@ -966,834 +887,7 @@ def resolve( return self.load(paths) -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") - - 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._global_indices_by_dataset = {n: [] for n in self.datasets} - for dataset_name, dataset in self.datasets.items(): - for local_idx in range(len(dataset)): - global_idx = len(self.index_map) - self.index_map.append((dataset_name, local_idx)) - self._global_indices_by_dataset[dataset_name].append(global_idx) - - 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 - # 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: - 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 - - below = arr < q_low - above = arr > q_high - 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.min().item():.4f}, {arr.max().item():.4f}]" - ) - return prefix - return None - - 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: - next_idx, attempts = self._next_after_failure( - idx, - dataset_name, - attempts, - reason=f"Sample failed ({type(e).__name__}: {e}) at " - f"{dataset_name}[{local_idx}]", - ) - idx = next_idx - 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: - next_idx, attempts = self._next_after_failure( - idx, - dataset_name, - attempts, - reason=violation, - ) - idx = next_idx - 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 _next_after_failure( - self, idx: int, dataset_name: str, attempts: int | None, *, reason: str - ) -> tuple[int, int]: - global_candidates = self._global_indices_by_dataset[dataset_name] - next_idx, attempts = get_fallback_idx( - idx=idx, - candidates=global_candidates, - _attempts=attempts, - max_attempts=len(global_candidates), - exhausted_error=( - f"Entire dataset bad (no valid indices): dataset={dataset_name}" - ), - ) - next_dataset_name, next_local_idx = self.index_map[next_idx] - logger.warning( - f"{reason} | attempt {attempts}, " - f"trying {next_dataset_name}[{next_local_idx}]" - ) - 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 leaf. ``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 - 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) - 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 - try: - sample = leaf[0] - if isinstance(sample, dict): - sample_keys = set(sample.keys()) - except Exception as e: - logger.warning( - f"[MultiDataset] Could not probe leaf for post-transform " - f"keys (emb={emb_id}): {e}. Falling back to raw key_map." - ) - - if sample_keys is None: - 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, - ): - 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 - - loader = torch.utils.data.DataLoader( - dataset, - batch_size=batch_size, - num_workers=num_workers, - shuffle=True, - generator=torch.Generator().manual_seed(seed), - ) - 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() - 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 - ) - stats_np = self._compute_stats_for_array(collected[k]) - 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, - } - 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), - "quantile_0_01": np.percentile(X, 0.01, axis=0), - "quantile_99_99": np.percentile(X, 99.99, axis=0), - } - - 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): - return _apply_norm_one(tensor, stats, self.norm_mode) - - def _apply_unnorm_one(self, tensor, stats): - return _apply_unnorm_one(tensor, stats, self.norm_mode) - - 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) - 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) - 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), - } - - @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", {})) - 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 set_data_schematic(self, data_schematic) -> None: - self.base.set_data_schematic(data_schematic) - self.data_schematic = data_schematic - - 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) - - -class ZarrDataset(torch.utils.data.Dataset): +class ZarrDataset(EpisodeLeaf): """ Base Zarr Dataset object, Just intializes as pass through to read from zarr episode """ @@ -1816,10 +910,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): """ @@ -1901,33 +992,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 _pad_sequences(self, data, horizon: int | None) -> dict: - if horizon is None: - return data - - # Note that k is zarr key - for k in data: - if isinstance(data[k], np.ndarray): - seq_len = data[k].shape[0] - if seq_len < horizon: - # Pad by repeating the last frame - pad_len = horizon - seq_len - last_frame = data[k][-1:] # Keep dims: (1, action_dim) - padding = np.repeat(last_frame, pad_len, axis=0) - data[k] = np.concatenate([data[k], padding], axis=0) - - return data - def __getitem__( self, idx: int, @@ -1961,6 +1025,36 @@ def _next(reason: str, key: str = "") -> int: return next_idx while True: + windows = self._read_windows(idx) + # One zarr call per key, not one per keymap entry: an action chunk + # and the proprio history of the same pose name the same key with + # different windows, and zarr's per-call metadata hop costs more + # than the handful of extra rows a merged read carries. Windows are + # merged only across rows in chunks the read fetches anyway (see + # ZarrEpisode.plan_reads), so two frames far apart on a JPEG key + # stay two rows. + intervals: dict[str, list[tuple[int, int]]] = {} + for zarr_key, window in windows.values(): + if isinstance(window, np.ndarray): + spans = [(int(i), int(i) + 1) for i in np.unique(window)] + else: + start, end = window + spans = [(start, start + 1 if end is None else end)] + intervals.setdefault(zarr_key, []).extend(spans) + groups = self.episode_reader.read_intervals(intervals) + + def _rows(zarr_key: str, start: int, stop: int): + for lo, rows in groups[zarr_key]: + if lo <= start and stop <= lo + len(rows): + return rows[start - lo : stop - lo] + # A StopIteration escaping __getitem__ reads as "epoch over" + # to the DataLoader; a short array must fail loudly. + raise ValueError( + f"ep={Path(self.episode_path).name} key={zarr_key}: " + f"window [{start}, {stop}) is past the rows read " + f"{[(lo, lo + len(rows)) for lo, rows in groups[zarr_key]]}" + ) + data = {} retry = False for k in self.key_map: @@ -1972,25 +1066,51 @@ def _next(reason: str, key: str = "") -> int: data[k] = self._annotation_text_for_frame(idx) continue - if horizon is not None: - end_idx = self._chunk_end_idx(idx, horizon, key_type) - read_interval = (idx, end_idx) + window = windows[k][1] + if isinstance(window, np.ndarray): + # (K,) rows; np.stack copies, so nothing aliases the read. + rows = [_rows(zarr_key, int(i), int(i) + 1)[0] for i in window] + chunk = rows if zarr_key in self._image_keys else np.stack(rows) else: - read_interval = (idx, None) - read_dict = {zarr_key: read_interval} - raw_data = self.episode_reader.read(read_dict) - self._pad_sequences(raw_data, horizon) # should be able to pad images - data[k] = raw_data[zarr_key] + start, end = window + chunk = _rows(zarr_key, start, start + 1 if end is None else end) + chunk = chunk[0] if end is None else chunk + if isinstance(chunk, np.ndarray): + # A slice of a shared read is a view; every keymap entry + # owned its own array before this coalescing and + # transforms are free to write into it. + chunk = chunk.copy() + data[k] = self._pad_to_horizon(chunk, horizon) if zarr_key in self._image_keys: - jpeg_bytes = data[k] try: - decoded = simplejpeg.decode_jpeg(jpeg_bytes, colorspace="RGB") + if isinstance(data[k], list): + # A history window, (K, H, W, 3): decode each + # distinct frame once (front padding repeats one). + window = windows[k][1] + unique = { + int(i): simplejpeg.decode_jpeg(b, colorspace="RGB") + for i, b in zip(window, data[k]) + } + decoded = np.stack([unique[int(i)] for i in window]) + else: + decoded = simplejpeg.decode_jpeg(data[k], colorspace="RGB") except Exception: idx = _next("JPEG decode failed", key=k) retry = True break - data[k] = np.transpose(decoded, (2, 0, 1)) / 255.0 + # float32, not numpy's float64 default for `uint8 / float`: + # process_batch_for_training casts to float32 on the GPU + # anyway, so the extra 4 bytes a pixel buy nothing and + # double the collate, the pin and the H2D copy. At 360x640 + # and batch 64 that is 338 MB a camera a batch instead of + # 169. The result is bit-identical: there are only 256 + # possible values and float32(u)/float32(255) equals + # float32(u / 255.0) for every one of them (see + # tests/unit/test_image_decode_dtype.py). + data[k] = np.moveaxis(decoded, -1, -3).astype( + np.float32 + ) / np.float32(255.0) elif zarr_key in self._json_keys: if isinstance(data[k], np.ndarray): data[k] = [self._decode_json_entry(v) for v in data[k]] @@ -1999,48 +1119,8 @@ def _next(reason: str, key: str = "") -> int: if retry: continue - 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) - # 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 - ) _ = origin # preserved for symmetry with prior API - return data + return self._finish(idx, data) class ZarrAnnotationCutoffDataset(ZarrDataset): @@ -2048,7 +1128,7 @@ class ZarrAnnotationCutoffDataset(ZarrDataset): Standard chunking from the start frame, but action reads stop at EOS+1 of the annotation span containing the start frame. The chunk is then padded out to - ``horizon`` via the base ``_pad_sequences`` (repeat-last), so frames beyond + ``horizon`` via the base ``_pad_to_horizon`` (repeat-last), so frames beyond EOS become the last action of the interval rather than crossing into the next annotation. @@ -2064,13 +1144,26 @@ def _build_frame_to_ann_end(self) -> dict[int, int]: annotation span. Annotations use half-open ``[start_idx, end_idx)``. """ mapping: dict[int, int] = {} + n_spans = 0 for ann in self._load_annotations(): start_idx = int(ann.get("start_idx", -1)) end_idx = int(ann.get("end_idx", -1)) if start_idx < 0 or end_idx <= start_idx: continue + n_spans += 1 for idx in range(start_idx, end_idx): mapping[idx] = end_idx + # Visibility into annotation-cutoff usage: spans/frames_covered of 0 + # means the cutoff is a no-op for this episode. One line per episode + # per rank is thousands of lines on a real split, so this is DEBUG and + # the resolver's kept N/M line carries the aggregate. + logger.debug( + "[AnnotationCutoff] ep=%s spans=%d frames_covered=%d/%d", + Path(self.episode_path).name, + n_spans, + len(mapping), + self.total_frames, + ) return mapping def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: @@ -2085,11 +1178,69 @@ def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> return min(end_idx, ann_end) +def _episode_has_annotation_spans(ds: "ZarrDataset") -> bool: + """True if the episode has at least one usable ``[start_idx, end_idx)`` span. + + Many Scale-"completed" episodes have an empty (or span-less) zarr + ``annotations`` array because the annotation-injection step lagged; the + AnnotationCutoff is a no-op for those, so they should be dropped when the + point of the run is to clamp chunks at annotation boundaries. + """ + try: + anns = ds._load_annotations() + except Exception: + return False + return any( + isinstance(a, dict) + and 0 <= int(a.get("start_idx", -1)) < int(a.get("end_idx", -1)) + for a in anns + ) + + class S3AnnotationCutoffEpisodeResolver(S3EpisodeResolver): - """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" + """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances. + + When ``require_annotations`` is set (default), episodes whose zarr + ``annotations`` array has no usable span are dropped; otherwise the + annotation cutoff would silently no-op on them. + """ _dataset_class = ZarrAnnotationCutoffDataset + def __init__(self, *args, require_annotations: bool = True, **kwargs): + super().__init__(*args, **kwargs) + self.require_annotations = require_annotations + + def resolve(self, filters=None, expected_embodiment=None): + datasets = super().resolve( + filters=filters, expected_embodiment=expected_embodiment + ) + if not self.require_annotations: + return datasets + kept = { + h: ds for h, ds in datasets.items() if _episode_has_annotation_spans(ds) + } + dropped = sorted(set(datasets) - set(kept)) + if dropped: + logger.warning( + "[AnnotationCutoff] dropped %d/%d episodes with no usable " + "annotation spans (e.g. %s)", + len(dropped), + len(datasets), + dropped[:5], + ) + logger.info( + "[AnnotationCutoff] kept %d/%d episodes with usable annotations", + len(kept), + len(datasets), + ) + if not kept: + raise ValueError( + "[AnnotationCutoff] no resolved episodes contain usable annotation " + "spans; check the filter / annotation injection for this dataset." + ) + return kept + class LocalAnnotationCutoffEpisodeResolver(LocalEpisodeResolver): """LocalEpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" @@ -2168,6 +1319,49 @@ def read( result[key] = data return result + @staticmethod + def plan_reads( + intervals: list[tuple[int, int]], chunk_len: int + ) -> list[tuple[int, int]]: + """Merge half-open ``[start, end)`` intervals on one key into read groups. + + Two merge only when every row between them lies in a chunk one of them + already touches, so the bridge costs no extra chunk fetch: the windows + on a numeric key (100-row chunks) collapse into one read, while two + frames far apart on a JPEG key (one frame per chunk) stay two rows + rather than every frame in between.""" + groups: list[list[int]] = [] + for start, end in sorted(intervals): + if groups and start // chunk_len <= (groups[-1][1] - 1) // chunk_len + 1: + groups[-1][1] = max(groups[-1][1], end) + else: + groups.append([start, end]) + return [(lo, hi) for lo, hi in groups] + + def read_intervals( + self, intervals: dict[str, list[tuple[int, int]]] + ) -> dict[str, list[tuple[int, np.ndarray]]]: + """One zarr call per key for all its ``[start, end)`` intervals. + + Returns ``{key: [(group_start, rows), ...]}``, one entry per group of + ``plan_reads``; several groups are fetched with a single orthogonal + index.""" + result = {} + for key, ivs in intervals.items(): + arr = self._store[key] + groups = self.plan_reads(ivs, arr.chunks[0]) + if len(groups) == 1: + lo, hi = groups[0] + result[key] = [(lo, arr[lo:hi])] + continue + rows = arr.oindex[np.concatenate([np.arange(lo, hi) for lo, hi in groups])] + out, offset = [], 0 + for lo, hi in groups: + out.append((lo, rows[offset : offset + hi - lo])) + offset += hi - lo + result[key] = out + return result + def _collect_keys(self) -> list[str]: """ Collect all array keys from the store. diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index fb252bb726..b6df442bd9 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -16,7 +16,11 @@ from egomimic.models.denoising_policy import DenoisingPolicy from egomimic.pl_utils.pl_model import ModelWrapper -from egomimic.rldb.embodiment.embodiment import get_embodiment +from egomimic.rldb.embodiment.embodiment import ( + EMBODIMENT, + Embodiment, + get_embodiment, +) from egomimic.rldb.embodiment.eva import Eva from egomimic.rldb.embodiment.human import Human from egomimic.robot.robot_utils import RateLoop @@ -154,9 +158,9 @@ def viz_rot_ee_pose(image, eepose, action_image_path, rot_image_path): LEFT_CAM_SERIAL = "" EMBODIMENT_MAP = { - "both": 8, - "left": 7, - "right": 6, + "both": EMBODIMENT.EVA_BIMANUAL.value, + "left": EMBODIMENT.EVA_LEFT_ARM.value, + "right": EMBODIMENT.EVA_RIGHT_ARM.value, } TEMP_DIR = "/home/robot/temp_dir" @@ -242,6 +246,7 @@ def __init__( resampled_action_len=None, debug=False, annotation_path=None, + transform_mode=None, ): super().__init__() self.arm = arm @@ -258,10 +263,25 @@ def __init__( self.debug_actions = None self.resampled_action_len = resampled_action_len self.debug = debug - # Inference on checkpoints trained before the 6D conversion. - self.transform_list = Eva.get_transform_list( - mode="cartesian_wristframe_ypr", allow_legacy_rotation=True - ) + self.transform_mode = self._checkpoint_transform_mode(transform_mode) + # Inference: a checkpoint recorded before the 6D conversion, or none at + # all, is exactly the case allow_legacy_rotation exists for. + if self.transform_mode is None: + # Legacy checkpoint: keep the historical input pipeline and use + # the predictions as-is. + self.transform_list = Eva.get_transform_list( + mode=self.LEGACY_TRANSFORM_MODE, allow_legacy_rotation=True + ) + self.revert_transform_list = None + else: + self.transform_list = Eva.get_transform_list( + mode=self.transform_mode, allow_legacy_rotation=True + ) + # Maps predictions (wrist frame and/or 6D rotation) back to the + # camera-frame xyz+ypr+gripper layout the command path expects. + self.revert_transform_list = Eva.get_revert_transform_list( + self.transform_mode + ) self.annotation = None self._tokenizer = None self.collate_fn = default_collate @@ -309,6 +329,63 @@ def _patch_checkpoint_paths(cls, ckpt_path): print(f"[rollout] Patched checkpoint saved to {patched_path}") return patched_path + # Input pipeline used for checkpoints whose config names no mode (the + # rollout hard-coded it before the mode was read from the checkpoint). + LEGACY_TRANSFORM_MODE = "cartesian_wristframe_ypr" + + def _checkpoint_transform_mode(self, override=None): + """Eva ``transform_list.mode`` the checkpoint was trained with, so the + rollout builds the same inputs and reverts the matching outputs; + ``None`` for a legacy 14-D ypr checkpoint. ``override`` wins.""" + from omegaconf import OmegaConf + + name = get_embodiment(self.embodiment_id).lower() + mode, source = override, "--transform-mode" + if mode is None: + cfg = self.policy._as_config( + getattr(self.policy.hparams, "config_tree", None) + ) + if cfg is not None: + mode = OmegaConf.select( + cfg, f"data_transform_modes.{name}", default=None + ) + source = "checkpoint" + if mode is None: + # Checkpoints from before trainHydra recorded the mode: only the + # legacy layout can be assumed; any other width is ambiguous. + width = self._checkpoint_action_width() + if width != 14: + raise ValueError( + f"[rollout] checkpoint records no {name} transform mode and its " + f"actions_cartesian width is {width}, not the legacy 14; pass " + "--transform-mode (e.g. cartesian_wristframe_6d)" + ) + mode, source = self.LEGACY_TRANSFORM_MODE, "14-D action width" + print(f"[rollout] Eva transform_list mode ({source}): {mode}") + return None if mode == self.LEGACY_TRANSFORM_MODE else mode + + def _checkpoint_action_width(self): + """Last dim of the checkpoint's Eva actions_cartesian, or None.""" + state = getattr(self.policy.hparams, "norm_stats_state", None) or {} + shapes = state.get("shapes", {}) + emb = shapes.get(self.embodiment_id, shapes.get(str(self.embodiment_id), {})) + shape = emb.get("actions_cartesian") + return None if shape is None else int(shape[-1]) + + def _revert_predictions(self, preds, obs_batch): + """(1, T, D) unnormalized predictions -> (1, T, 14) camera-frame + xyz+ypr+gripper per arm. ``obs_batch`` is the collated, not yet + normalized input batch (its proprio defines the wrist frame).""" + if self.revert_transform_list is None: + return preds + obs_key = "observations.state.ee_pose" + batch = { + "actions_cartesian": preds, + obs_key: torch.as_tensor(obs_batch[obs_key]).to(preds.device), + } + out = Embodiment.apply_transform(batch, self.revert_transform_list) + return torch.as_tensor(out["actions_cartesian"]) + def _load_policy(self): patched_path = self._patch_checkpoint_paths(self.policy_path) policy = ModelWrapper.load_from_checkpoint( @@ -372,6 +449,13 @@ def rollout_step(self, i, obs): for transform in self.transform_list: transform_list_batch = transform.transform(transform_list_batch) transform_list_batch = self.collate_fn([transform_list_batch]) + # process_batch_for_training normalizes in place; keep raw proprio + # for reverting the wrist-frame predictions. + raw_obs = { + "observations.state.ee_pose": transform_list_batch[ + "observations.state.ee_pose" + ].clone() + } if self.arm == "both": embodiment_name = "eva_bimanual" elif self.arm == "right": @@ -386,6 +470,7 @@ def rollout_step(self, i, obs): preds = self.policy.model.forward_eval(processed_batch)[ f"{embodiment_name}_actions_cartesian" ] + preds = self._revert_predictions(preds, raw_obs) self.actions = preds.detach().cpu().numpy().squeeze() self.debug_actions = self.actions.copy() if self.cartesian: @@ -589,6 +674,7 @@ def main( offline_debug=False, offline_episode_path=None, annotation_path=None, + transform_mode=None, ): if arms == "both": arms_list = ["right", "left"] @@ -620,6 +706,7 @@ def main( resampled_action_len=resampled_action_len, debug=debug, annotation_path=annotation_path, + transform_mode=transform_mode, ) elif dataset_path is not None: rollout_type = "replay" @@ -816,6 +903,12 @@ def build_arg_parser(description="Rollout robot model."): type=str, help="path to the annotation file", ) + parser.add_argument( + "--transform-mode", + type=str, + help="Eva transform_list mode the policy was trained with; only needed for " + "checkpoints that do not record it (default: read from the checkpoint)", + ) return parser @@ -833,6 +926,7 @@ def run_from_args(args): offline_debug=args.offline_debug, offline_episode_path=args.offline_episode_path, annotation_path=args.annotation_path, + transform_mode=args.transform_mode, ) diff --git a/egomimic/scripts/abc_process/__init__.py b/egomimic/scripts/abc_process/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/abc_process/viz_eva_episode.py b/egomimic/scripts/abc_process/viz_eva_episode.py new file mode 100644 index 0000000000..f60c1e1101 --- /dev/null +++ b/egomimic/scripts/abc_process/viz_eva_episode.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python +"""Visualize an EVA (ABC-130k) zarr episode -> mp4. + +Reads the camera frames straight out of the zarr store (no model / no GPU) and +writes a video. Optionally tiles multiple cameras side-by-side and burns in the +task description + frame index. + + ./emimic/bin/python egomimic/scripts/abc_process/viz_eva_episode.py \ + --episode 00022b62-0a07-4ac7-bb74-9b4da47f6c6e \ + --cams front_1 left_wrist right_wrist --out eva_ep.mp4 + +--episode accepts a bare hash (resolved under --folder, which defaults to the +dataset dir from hydra_configs/paths/default.yaml: $EGOVERSE_DATASET_DIR or +/coc/flash7/scratch/egoverseS3ZarrDataset) or a full .zarr path. + +--draw-axes overlays, on the front_1 tile, the per-arm EE pose triads AND the +world-frame origin triad (X=red, Y=green, Z=blue), projected with the SAME +math the training pipeline uses: world pose -> cam frame via +PoseCoordinateFrameTransform against EXTRINSICS[--extrinsics-key], then pixels +via INTRINSICS[--intrinsics-key]. The front_1 tile is stretched to 640x480 +first (the space that K and the abc_fold_viz solvePnP fit live in). Use +--still-frames to also dump overlaid PNGs of specific frames. +""" + +import argparse +import io +import os + +import numpy as np +import zarr + +from egomimic.utils.paths import default_dataset_dir + +DEFAULT_DATASET_DIR = default_dataset_dir() + + +def decode(elem): + """An image stored in the object array -> HxWx3 uint8. + + Frames are JPEG bytes wrapped in a 0-d object ndarray (sometimes nested), so + unwrap to the raw bytes and decode with PIL (imageio can't sniff a bare BytesIO). + """ + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +# Viz space the overlay lives in: front_1 stretched to 640x480, matching both +# the pipeline's _resize_image_keys and the clicker/solvePnP fit. +VIZ_W, VIZ_H = 640, 480 +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB + +# Viz-only calibrations, kept local to this script (training reads intrinsics +# per-episode from zarr.json and bakes Eva.EXTRINSICS; nothing else in the repo +# needs these). Keys that exist on the Eva class are added lazily from it in +# AxisOverlay (see _registries) so this module stays import-light. +# eva K: ABC-130k RealSense top camera (640x480 space) +# yam K: rectified stereo front cam, calibration-derived, 640x480 space +# abc_fold_viz: solvePnP world->cam fit over ABC fold-clothes clicks (one +# shared world frame for both arms); yam: per-arm base->cam hand-eye calib +# eva (extrinsics, from Eva.EXTRINSICS): the per-arm base->cam bake training +# uses; aria (intrinsics, from Eva.INTRINSICS): the class fallback K, +# 640x480 space +INTRINSICS = { + "eva": np.array( + [[436.26, 0.0, 310.1, 0.0], [0.0, 435.13, 241.93, 0.0], [0.0, 0.0, 1.0, 0.0]] + ), + "yam": np.array( + [ + [250.15807185802416, 0.0, 267.5409836065574, 0.0], + [0.0, 260.01789323969166, 152.72727272727272, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ), +} +_ABC_FOLD_VIZ = np.array( + [ + [-0.00540728648402, -0.91864395735192, 0.39504941573641, -0.07230219027034], + [-0.99924104589933, -0.01027605228132, -0.03757306135427, -0.23692835817344], + [0.03857581422213, -0.39495275966919, -0.91789118319482, 0.97783070025080], + [0.0, 0.0, 0.0, 1.0], + ] +) +EXTRINSICS = { + "abc_fold_viz": {"left": _ABC_FOLD_VIZ, "right": _ABC_FOLD_VIZ}, + "yam": { + "left": np.array( + [ + [0.0437291, -0.85821391, 0.5114261, 0.08690521], + [-0.99754573, -0.00948836, 0.06937217, -0.25327518], + [-0.05468356, -0.5132045, -0.85652253, 0.83357606], + [0.0, 0.0, 0.0, 1.0], + ] + ), + "right": np.array( + [ + [0.04251619, -0.84174054, 0.53820557, 0.04959916], + [-0.98836195, 0.04331659, 0.14582295, 0.27905066], + [-0.14605832, -0.53814174, -0.83010266, 0.87148985], + [0.0, 0.0, 0.0, 1.0], + ] + ), + }, +} +# Intrinsics keys whose K lives in the stretched 640x480 space (the pipeline's +# _resize_image_keys target); YAM's K is already in the stored image space. +_STRETCHED_INTRINSICS = {"eva", "aria"} + + +def _registries(): + """Local viz calibrations + the values that live on the Eva class.""" + from egomimic.rldb.embodiment.eva import Eva + + return ( + {**INTRINSICS, "aria": Eva.INTRINSICS}, + {**EXTRINSICS, "eva": Eva.EXTRINSICS}, + ) + + +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.transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _matrix_to_xyzwxyz, _xyzwxyz_to_matrix + + intrinsics, extrinsics = _registries() + self._to_mat = _xyzwxyz_to_matrix + self.K = intrinsics[intrinsics_key] + # EVA's K lives in the pipeline's stretched 640x480 space; YAM's K is + # already in the stored image space -> draw at native resolution there. + self.stretch = intrinsics_key in _STRETCHED_INTRINSICS + extr = extrinsics[extrinsics_key] + # cam<-world pose target per arm, exactly as the Eva transform list bakes it. + self.extr_pose = { + s: _matrix_to_xyzwxyz(np.asarray(extr[s], dtype=np.float64)[None])[0] + for s in ("left", "right") + } + # One shared world frame (e.g. abc_fold_viz) vs per-arm base frames (yam). + self.shared_base = np.allclose(extr["left"], extr["right"]) + self._t = PoseCoordinateFrameTransform( + target_world="extr", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self.key = extrinsics_key + + def cam_pose_matrix(self, world_pose_xyzwxyz, side: str) -> np.ndarray: + """World-frame pose (7,) -> 4x4 cam<-ee matrix (pipeline-identical).""" + out = self._t.transform( + { + "extr": self.extr_pose[side], + "pose": np.asarray(world_pose_xyzwxyz, dtype=np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, dtype=np.float64)[None])[0] + + def project(self, pts_cam: np.ndarray) -> np.ndarray: + """(N,3) cam-frame points -> (N,2) pixels; NaN where behind the camera.""" + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M_cam_pose, length: float, label: str, thickness=2): + """Draw an XYZ triad for a cam<-frame pose matrix onto img (RGB, 640x480).""" + import cv2 + + o = M_cam_pose[:3, 3] + pts = np.stack([o] + [o + M_cam_pose[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, front_native, lpose_w, rpose_w): + """front_1 native frame + world EE poses -> 640x480 RGB with triads. + + Returns (img, info) where info holds the projected EE origin pixels. + """ + import cv2 + + if self.stretch: + img = cv2.resize( + front_native, (VIZ_W, VIZ_H), interpolation=cv2.INTER_LINEAR + ) + else: + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = front_native.copy() + img = np.ascontiguousarray(img) + h = img.shape[0] + info = {} + # Base-frame triad(s): one shared world frame, or per-arm bases (yam). + ident = np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float64) + bases = ( + [("left", "world")] + if self.shared_base + else [("left", "L-base"), ("right", "R-base")] + ) + for side, lab in bases: + self.draw_triad( + img, self.cam_pose_matrix(ident, side), 0.15, lab, thickness=3 + ) + for side, pose_w, lab in ( + ("left", lpose_w, "L-ee"), + ("right", rpose_w, "R-ee"), + ): + # ABC episodes zero-pad the tail frames -> zero-norm quat; skip them. + if np.linalg.norm(np.asarray(pose_w, dtype=np.float64)[3:7]) < 1e-6: + info[side] = np.array([np.nan, np.nan]) + continue + M = self.cam_pose_matrix(pose_w, side) + self.draw_triad(img, M, 0.08, lab) + info[side] = self.project(M[:3, 3][None])[0] + cv2.putText( + img, + f"axes X=red Y=green Z=blue extr={self.key}", + (6, h - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img, info + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True, help="episode hash or full .zarr path") + ap.add_argument( + "--folder", + default=DEFAULT_DATASET_DIR, + help="zarr dataset dir a bare --episode hash is resolved under " + "(default: $EGOVERSE_DATASET_DIR or the lab dataset dir)", + ) + ap.add_argument( + "--cams", nargs="+", default=["front_1"], help="front_1 left_wrist right_wrist" + ) + ap.add_argument( + "--out", default=None, help="output mp4 (default: ./eva_.mp4)" + ) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + ap.add_argument( + "--draw-axes", + action="store_true", + help="overlay EE + world XYZ triads on front_1 (see module docstring)", + ) + ap.add_argument( + "--extrinsics-key", + default="abc_fold_viz", + help="EXTRINSICS key for the world->cam bake " + "(abc_fold_viz | yam | eva; default: abc_fold_viz)", + ) + ap.add_argument( + "--intrinsics-key", + default="eva", + help="INTRINSICS key for projection (eva | yam | aria)", + ) + ap.add_argument( + "--still-frames", + type=int, + nargs="*", + default=[], + help="also dump these frame indices as overlaid PNGs next to --out", + ) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + n = int(a.get("total_frames", 0)) + print( + f"[viz] {path}\n task={a.get('task_description') or a.get('task_name')!r} " + f"frames={n} fps={fps} cams={args.cams}" + ) + + overlay = None + lobs = robs = None + if args.draw_axes: + if "front_1" not in args.cams: + args.cams = ["front_1"] + args.cams + overlay = AxisOverlay(args.extrinsics_key, args.intrinsics_key) + lobs, robs = g["left.obs_ee_pose"], g["right.obs_ee_pose"] + print( + f"[viz] axis overlay on front_1: extrinsics_key={args.extrinsics_key}, " + f"K=INTRINSICS['{args.intrinsics_key}'], stretch={overlay.stretch}" + ) + + cam_arrs = {c: g[f"images.{c}"] for c in args.cams} + # Arrays are zero-padded past attrs total_frames -> cap at the real length. + nf = min( + next(iter(cam_arrs.values())).shape[0], n or 10**9, args.max_frames or 10**9 + ) + out = args.out or f"eva_{os.path.basename(path).replace('.zarr', '')}.mp4" + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + import imageio.v2 as imageio + + try: + from PIL import Image, ImageDraw + except Exception: + Image = None + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + tiles = [] + for c in args.cams: + f = decode(cam_arrs[c][i]) + if overlay is not None and c == "front_1": + f, info = overlay.overlay(f, np.asarray(lobs[i]), np.asarray(robs[i])) + if i in stills: + print( + f"[viz] frame {i}: EE pixels L={np.round(info['left'], 1)} " + f"R={np.round(info['right'], 1)}" + ) + tiles.append(f) + h = min(t.shape[0] for t in tiles) + tiles = ( + [ + t[:h] + if t.shape[0] == h + else np.asarray( + Image.fromarray(t).resize((int(t.shape[1] * h / t.shape[0]), h)) + ) + for t in tiles + ] + if Image + else tiles + ) + frame = np.concatenate(tiles, axis=1) + if Image: + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still_path = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still_path) + print(f"[viz] still -> {still_path}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/aria_process/aria_to_lance.py b/egomimic/scripts/aria_process/aria_to_lance.py new file mode 100644 index 0000000000..33a48bea3a --- /dev/null +++ b/egomimic/scripts/aria_process/aria_to_lance.py @@ -0,0 +1,288 @@ +"""Convert one raw Aria recording (VRS + MPS) into a delivery: per-frame parquet plus MP4s. + +Unlike aria_to_zarr, every RGB frame is kept; untracked hands, missing poses and timing +gaps become validity flags instead of deleted rows. With --stage-to the delivery is staged +for the committer (egomimic/scripts/lance_commit.py), which appends it to the Lance tables. +""" + +import argparse +import json +import logging +import time +from pathlib import Path + +import cv2 +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from projectaria_tools.core import calibration, data_provider, mps +from projectaria_tools.core.mps.utils import get_nearest_hand_tracking_result +from projectaria_tools.core.sensor_data import TimeDomain +from projectaria_tools.core.stream_id import StreamId + +from egomimic.rldb.embodiment.human import ARIA_INTRINSICS +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import HUMAN_STEPS, recording_id +from egomimic.rldb.lance.timeline import build_timeline, nearest_dt, pose_outliers +from egomimic.scripts.aria_process.aria_utils import ( + AriaVRSExtractor, + convert_aria_keypoints_to_mano, + mano_keypoints_to_cartesian, + slam_to_rgb, +) +from egomimic.utils.video_utils import open_av1_encoder + +log = logging.getLogger("aria_to_lance") + +FPS = 30.0 +RGB = StreamId("214-1") +SENTINEL = 1e8 +# A pose sample further than this from the frame time is treated as missing. +MATCH_TOL_NS = 50_000_000 +ARM_OF = { + "human_bimanual": "both", + "human_left_arm": "left", + "human_right_arm": "right", +} +VIDEOS = {"front_1": (360, 480), "front_1_480": (480, 640)} +IDENTITY_POSE = np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float32) + + +def make_undistort(provider, height=480, width=640): + """aria_utils.undistort_to_linear with the calibration built once.""" + label = provider.get_label_from_stream_id(RGB) + calib = provider.get_device_calibration().get_camera_calib(label) + focal = 133.25430222 * (height / 240) + linear = calibration.get_linear_camera_calibration( + height, width, focal, label, calib.get_transform_device_camera() + ) + + def f(raw): + return np.ascontiguousarray( + np.rot90(calibration.distort_by_calibration(raw, linear, calib), k=3) + ) + + return f + + +def hand_arrays(kp_rows: np.ndarray, side_offset: int): + """Split get_hand_keypoints rows (wrist pose 7 + 63 per hand) for one hand.""" + blk = kp_rows[:, side_offset : side_offset + 70] + return blk[:, :7], blk[:, 7:70] + + +def convert(args) -> dict: + t0 = time.time() + vrs = args.raw_dir / f"{args.vendor_recording_id}.vrs" + mps_dir = args.raw_dir / f"mps_{args.vendor_recording_id}_vrs" + out = args.out_dir / args.vendor_recording_id + out.mkdir(parents=True, exist_ok=True) + arm = ARM_OF[args.embodiment] + use = {"left": arm in ("left", "both"), "right": arm in ("right", "both")} + + provider = data_provider.create_vrs_data_provider(str(vrs)) + hands = mps.hand_tracking.read_hand_tracking_results( + str(mps_dir / "hand_tracking" / "hand_tracking_results.csv") + ) + traj = mps.read_closed_loop_trajectory( + str(mps_dir / "slam" / "closed_loop_trajectory.csv") + ) + gaze_csv = mps_dir / "eye_gaze" / "general_eye_gaze.csv" + gaze = mps.read_eyegaze(str(gaze_csv)) if gaze_csv.exists() else None + + rgb_ts = np.asarray( + provider.get_timestamps_ns(RGB, TimeDomain.DEVICE_TIME), dtype=np.int64 + ) + med_dt = float(np.median(np.diff(rgb_ts))) + nominal_fps = 1e9 / med_dt + if abs(nominal_fps - FPS) > 0.1 * FPS: + raise ValueError(f"RGB stream is {nominal_fps:.1f} fps, expected {FPS}") + stream_ts = {"rgb": list(rgb_ts)} + + # Same per-frame pose math as aria_to_zarr, on every RGB frame (no clean_data). + both = "both" + kp = AriaVRSExtractor.get_hand_keypoints(traj, stream_ts, hands, arm=both) + mps_ee = AriaVRSExtractor.get_ee_pose(traj, stream_ts, hands, arm=both) + head = AriaVRSExtractor.get_head_pose( + traj, slam_to_rgb(provider).inverse(), stream_ts + ) + eye = ( + AriaVRSExtractor.get_eye_gaze(gaze, stream_ts) + if gaze is not None + else np.full((len(rgb_ts), 3), SENTINEL * 10) + ) + + traj_ts = np.array( + [int(p.tracking_timestamp.total_seconds() * 1e9) for p in traj], dtype=np.int64 + ) + traj_graph = np.array([p.graph_uid for p in traj]) + hand_ts = np.array( + [int(h.tracking_timestamp.total_seconds() * 1e9) for h in hands], dtype=np.int64 + ) + pose_close = nearest_dt(rgb_ts, traj_ts) <= MATCH_TOL_NS + hand_close = nearest_dt(rgb_ts, hand_ts) <= MATCH_TOL_NS + graph = traj_graph[np.clip(np.searchsorted(traj_ts, rgb_ts), 0, len(traj_ts) - 1)] + graph_change = np.r_[False, graph[1:] != graph[:-1]] + + per_hand = {} + for side, off, ee_off, rhand in (("left", 0, 0, False), ("right", 70, 7, True)): + wrist, aria_kp = hand_arrays(kp, off) + tracked = use[side] & (wrist < SENTINEL).all(1) & hand_close & pose_close + conf = np.full(len(rgb_ts), -1.0, dtype=np.float32) + for i, t in enumerate(rgb_ts): + h = get_nearest_hand_tracking_result(hands, int(t)) + one = getattr(h, f"{side}_hand", None) if h is not None else None + if one is not None: + conf[i] = one.confidence + if args.mano and tracked.any(): + mano_kp = convert_aria_keypoints_to_mano( + np.where(tracked[:, None], aria_kp, 1e9), + is_rhand=rhand, + n_iters=args.mano_iters, + lr=0.05, + ) + ee = mano_keypoints_to_cartesian(mano_kp, is_rhand=rhand) + keypoints = mano_kp + else: + ee = mps_ee[:, ee_off : ee_off + 7] + keypoints = aria_kp + ok = tracked & np.isfinite(keypoints).all(1) & (np.abs(ee) < SENTINEL).all(1) + ok &= ~pose_outliers(wrist[:, :3], ok) + per_hand[side] = dict( + wrist=wrist, ee=ee, kp=keypoints, aria_kp=aria_kp, conf=conf, ok=ok + ) + + head_ok = (head < SENTINEL).all(1) & pose_close + head_ok &= ~pose_outliers(head[:, :3], head_ok) + + ts, src, filler, after_gap = build_timeline(rgb_ts) + brk = after_gap | (graph_change[src] & ~filler) + head_valid = head_ok[src] & ~filler & ~brk + n = len(ts) + + def col(x, valid, fill): + x = np.asarray(x, dtype=np.float32)[src] + x[~valid] = fill + return pa.FixedSizeListArray.from_arrays( + pa.array(x.reshape(-1)), x.shape[1] + ).cast(pa.list_(pa.float32(), x.shape[1])) + + cols = { + "recording_id": pa.array([args.recording_id] * n), + "timestamp_ns": pa.array(ts), + "head_pose": col(head, head_valid, IDENTITY_POSE), + } + valid = {} + for side, h in per_hand.items(): + v = h["ok"][src] & ~filler & ~brk + valid[side] = v + cols[side] = pa.StructArray.from_arrays( + [ + col(h["wrist"], v, IDENTITY_POSE), + col(h["ee"], v, IDENTITY_POSE), + col(h["kp"], v, 0.0), + col(h["aria_kp"], v, 0.0), + pa.array(np.where(filler, -1.0, h["conf"][src]).astype(np.float32)), + ], + fields=list(HUMAN_STEPS.field(side).type), + ) + cols["left_valid"] = pa.array(valid["left"]) + cols["right_valid"] = pa.array(valid["right"]) + cols["head_pose_valid"] = pa.array(head_valid) + eye_ok = (np.abs(eye) < SENTINEL).all(1)[src] & ~filler + cols["eye_gaze"] = col(eye, eye_ok, 0.0) + table = pa.table(cols, schema=HUMAN_STEPS) + pq.write_table(table, out / "steps.parquet", compression="zstd") + t_numeric = time.time() - t0 + + undistort = make_undistort(provider) + enc = { + k: open_av1_encoder(out / f"{k}.mp4", h, w, FPS, args.ffmpeg) + for k, (h, w) in VIDEOS.items() + } + last_src, frame = -1, None + for s in src: + if s != last_src: + frame = undistort( + provider.get_image_data_by_index(RGB, int(s))[0].to_numpy_array() + ) + assert frame.shape == (*VIDEOS["front_1_480"], 3), frame.shape + small = cv2.resize( + frame, + (VIDEOS["front_1"][1], VIDEOS["front_1"][0]), + interpolation=cv2.INTER_AREA, + ) + last_src = s + enc["front_1_480"].stdin.write(frame.tobytes()) + enc["front_1"].stdin.write(small.tobytes()) + for p in enc.values(): + p.stdin.close() + if p.wait() != 0: + raise RuntimeError("ffmpeg failed") + + k = np.asarray(ARIA_INTRINSICS, dtype=np.float64)[:3, :3] + scale = VIDEOS["front_1"][0] / 480 + k360 = k.copy() + k360[:2] *= scale + meta = { + "vendor_recording_id": args.vendor_recording_id, + "recording_id": args.recording_id, + "embodiment": args.embodiment, + "fps": FPS, + "rgb_fps_measured": nominal_fps, + "num_frames": n, + "rgb_frames": int(len(rgb_ts)), + "filler_frames": int(filler.sum()), + "gaps": int(after_gap.sum()), + "graph_changes": int(graph_change.sum()), + "t_start_ns": int(ts[0]), + # Half-open like every Lance interval: the last frame is inside its task and alias. + "t_end_ns": int(ts[-1]) + round(1e9 / FPS), + "mano": bool(args.mano), + "intrinsics": { + "front_1": k360.reshape(-1).tolist(), + "front_1_480": k.reshape(-1).tolist(), + }, + # head_pose is the front_1 camera pose, so the camera sits at the head frame origin. + "extrinsics": {"front_1": np.eye(4).reshape(-1).tolist()}, + "valid_frac": { + "left": float(valid["left"].mean()), + "right": float(valid["right"].mean()), + "head": float(head_valid.mean()), + }, + "pose_unmatched": int((~pose_close).sum()), + "hand_unmatched": int((~hand_close).sum()), + "seconds_numeric": round(t_numeric, 1), + "seconds_total": round(time.time() - t0, 1), + } + (out / "meta.json").write_text(json.dumps(meta, indent=1)) + return meta + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--raw-dir", type=Path, required=True) + ap.add_argument("--out-dir", type=Path, required=True) + ap.add_argument( + "--vendor-recording-id", required=True, help="epoch ms, the VRS file stem" + ) + ap.add_argument("--embodiment", required=True, choices=sorted(ARM_OF)) + ap.add_argument("--mano", action=argparse.BooleanOptionalAction, default=True) + ap.add_argument("--mano-iters", type=int, default=2400) + ap.add_argument("--ffmpeg", default="/coc/flash7/rco3/ffmpeg-7.1-shared/bin/ffmpeg") + ap.add_argument( + "--stage-to", help="staging root (s3:// or local) for the committer" + ) + args = ap.parse_args() + args.recording_id = recording_id("aria", args.vendor_recording_id) + logging.basicConfig(level=logging.INFO) + print(json.dumps(convert(args), indent=1)) + if args.stage_to: + commit.stage( + args.out_dir / args.vendor_recording_id, + commit.join(args.stage_to, "aria", args.vendor_recording_id), + ) + + +if __name__ == "__main__": + main() 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..6683bb5790 --- /dev/null +++ b/egomimic/scripts/aria_process/aria_to_lance_ray.py @@ -0,0 +1,149 @@ +"""Convert raw Aria recordings (VRS + MPS) into staged deliveries, one Ray task per recording. + +Workers run aria_to_lance.convert into `//`; the driver stages each +finished delivery under `/aria//` for the committer +(egomimic/scripts/lance_commit.py), which is the only Lance writer. 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 \\ + --stage-to s3://rldb/lance_staging --lance s3://rldb/lance +""" + +import argparse +import json +import logging +import time +import traceback +from pathlib import Path +from types import SimpleNamespace + +import ray + +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import recording_id +from egomimic.scripts.lance_commit import sql_rows +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]: + 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() + + +def stage_dir(args, d: Path, report: dict) -> None: + commit.stage(d, commit.join(args.stage_to, "aria", d.name)) + report["staged"].append(d.name) + + +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("--stage-to", required=True, help="staging root (s3:// or local)") + ap.add_argument( + "--lance", required=True, help="Lance root, read to skip committed recordings" + ) + 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( + "--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="/coc/flash7/rco3/ffmpeg-7.1-shared/bin/ffmpeg") + args = ap.parse_args() + logging.basicConfig(level=logging.INFO) + + have = commit.committed_ids(args.lance) + vids = [v for v in discover(args.raw_dir) if recording_id("aria", v) not in have] + delivered = {v for v in vids if (args.deliveries / v / "meta.json").exists()} - ( + set(vids) if args.reconvert else set() + ) + 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 = { + "staged": [], + "failed": {v: "no embodiment in SQL" for v in todo if v not in emb}, + } + 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()] + for v in sorted(delivered): + stage_dir(args, args.deliveries / v, report) + 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: + stage_dir(args, args.deliveries / vid, report) + log.info("%d/%d converted, %.0fs", done, len(emb), time.time() - t0) + 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"staged {len(report['staged'])}, failed {len(report['failed'])}; 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/data_visualization/inspector_lib/dataset_view.py b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py index 8fbe2a8b31..759ca7a757 100644 --- a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py +++ b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py @@ -429,7 +429,7 @@ def _world_keypoints_to_cam(seq, T_world_head): _2D_DOT_COLORS_BGR = { "left": (0, 120, 255), "right": (255, 80, 0), -} # mirrors _viz_keypoints +} # the inspector's own per-arm dot colours; the 2D overlay is the horizon trace _ARM_COLORS = { arm: f"rgb({bgr[2]},{bgr[1]},{bgr[0]})" # BGR -> RGB for arm, bgr in _2D_DOT_COLORS_BGR.items() @@ -491,9 +491,10 @@ def build_3d_figure(grp, frame: int, overlay: str, traj_window: int = 60): emb_cls, "FINGER_EDGE_RANGES", Human.FINGER_EDGE_RANGES if Human else [] ) # Map each edge index -> finger color (via FINGER_EDGE_RANGES), fall back gray. - # FINGER_COLORS are BGR tuples (they're handed to cv2 in the 2D skeleton draw - # in `_viz_keypoints`), so reverse BGR -> RGB to make the 3D edges render the - # SAME color cv2 paints in 2D. + # FINGER_COLORS are BGR tuples because cv2 consumes them in the 2D overlay, so + # reverse BGR -> RGB here. The 2D overlay is the horizon trace, which paints + # only the fingertips, so the colours agree at the tips and the 3D view is the + # only place a whole finger takes its finger's colour. edge_color_css = ["rgb(180,180,180)"] * len(edges) for name, lo, hi in edge_ranges: col = finger_colors.get(name) @@ -805,8 +806,9 @@ def _chunk(seq, arm): if all(np.allclose(p, 0) for p in parts): return _badge(img_rgb.copy(), "keypoint: no pts"), False, "off" # canonical keypoints (wrist_in_data=False) layout: [L 63, R 63]; - # .viz(mode="keypoints") supplies the MANO edges/colors/ - # edge_ranges. Only Human supports keypoints; if the resolved class + # .viz(mode="keypoints") draws the horizon trace: for this + # single frame that is the per-slot start markers, not a posed + # skeleton. Only Human supports keypoints; if the resolved class # doesn't, report via badge rather than reintroducing a low-level draw call. kp = np.concatenate(parts).reshape(1, -1) try: 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/eva_process/eva_to_lance.py b/egomimic/scripts/eva_process/eva_to_lance.py new file mode 100644 index 0000000000..6a483ec4b7 --- /dev/null +++ b/egomimic/scripts/eva_process/eva_to_lance.py @@ -0,0 +1,207 @@ +"""Convert one raw EVA demo (HDF5 + metadata JSON) into a delivery: per-frame parquet plus MP4s. + +Numerics go through the same EvaHD5Extractor and per-arm split as eva_to_zarr, so steps +columns equal the zarr keys. Frames where an arm's pose reads all zeros stay forward-filled +as in zarr, and are also flagged `_valid = False`. The HDF5 carries no per-frame +times, so frame i is at start + i / FPS. With --stage-to the delivery is staged for the +committer (egomimic/scripts/lance_commit.py), which appends it to the Lance tables. +""" + +import argparse +import json +import logging +import time +from pathlib import Path + +import cv2 +import h5py +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import ARM, EVA_STEPS, recording_id +from egomimic.scripts.eva_process.eva_to_zarr import ( + DATASET_KEY_MAPPINGS, + _split_per_arm, +) +from egomimic.scripts.eva_process.eva_utils import EvaHD5Extractor +from egomimic.utils.aws.aws_sql import timestamp_ms_to_episode_hash +from egomimic.utils.video_utils import open_av1_encoder + +log = logging.getLogger("eva_to_lance") + +FPS = 30.0 +ARM_OF = {"eva_bimanual": "both", "eva_left_arm": "left", "eva_right_arm": "right"} +# HDF5 camera -> delivery camera; 640x480 sources, encoded at short side 360. +CAMERAS = { + "front_img_1": "front_1", + "left_wrist_img": "left_wrist", + "right_wrist_img": "right_wrist", +} +SIZE = (360, 480) + + +def numeric_columns(raw: Path, arm: str) -> tuple[dict[str, np.ndarray], np.ndarray]: + """({zarr key: (T, d) array} of the robot's arms, (T, 2) [left, right] dropout), + with eva_to_zarr's per-arm math. + + Raw arrays are (T, 14) with the left arm in columns 0-6 and the right in 7-13 even + on a one-arm robot; eva_to_zarr reads a lone right arm from columns 0-6 (zeros), so + this splits both slots and keeps the robot's arms. + """ + feats = EvaHD5Extractor.process_episode(raw, arm=arm, with_images=False) + numeric = { + DATASET_KEY_MAPPINGS[k]: np.asarray(v) + for k, v in feats.items() + if k in DATASET_KEY_MAPPINGS and "images" not in k + } + widths = {k: v.shape[1] for k, v in numeric.items()} + if set(widths.values()) != {14}: + raise ValueError(f"expected (T, 14) arm arrays, got widths {widths}") + keep = ("left", "right") if arm == "both" else (arm,) + cols = { + k: v + for k, v in _split_per_arm(numeric, "both").items() + if k.split(".")[0] in keep + } + return cols, feats["metadata.pose_dropout"] + + +def steps_table(rid: str, ts: np.ndarray, cols: dict, valid: dict) -> pa.Table: + n = len(ts) + + def field(x, width): + x = np.asarray(x, np.float32).reshape(n, -1) + if width is None: + return pa.array(x[:, 0]) + return pa.FixedSizeListArray.from_arrays(pa.array(x.reshape(-1)), width) + + arms = {} + for side in ("left", "right"): + absent = f"{side}.obs_ee_pose" not in cols + arms[side] = pa.StructArray.from_arrays( + [ + field( + np.zeros((n, getattr(f.type, "list_size", 1))) + if absent + else cols[f"{side}.{f.name}"], + getattr(f.type, "list_size", None), + ) + for f in ARM + ], + fields=list(ARM), + ) + return pa.table( + { + "recording_id": pa.array([rid] * n), + "timestamp_ns": pa.array(ts), + **arms, + "left_valid": pa.array(valid["left"]), + "right_valid": pa.array(valid["right"]), + }, + schema=EVA_STEPS, + ) + + +def encode_videos(raw: Path, out: Path, n: int, ffmpeg: str) -> list[str]: + h, w = SIZE + done = [] + with h5py.File(raw, "r") as f: + imgs = f["observations/images"] + for src, cam in CAMERAS.items(): + if src not in imgs: + continue + if len(imgs[src]) != n: + raise ValueError(f"{src} has {len(imgs[src])} frames, numerics {n}") + enc = open_av1_encoder(out / f"{cam}.mp4", h, w, FPS, ffmpeg) + for a in range(0, n, 64): + for frame in imgs[src][a : a + 64]: + small = cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA) + enc.stdin.write(np.ascontiguousarray(small).tobytes()) + enc.stdin.close() + if enc.wait() != 0: + raise RuntimeError(f"ffmpeg failed on {cam}") + done.append(cam) + return done + + +def convert(args) -> dict: + t0 = time.time() + raw = args.raw_dir / f"{args.vendor_recording_id}.hdf5" + md = json.loads( + (args.raw_dir / f"{args.vendor_recording_id}_metadata.json").read_text() + ) + embodiment = md["robot_name"] + arm = ARM_OF[embodiment] + out = args.out_dir / args.vendor_recording_id + out.mkdir(parents=True, exist_ok=True) + rid = recording_id("eva", args.vendor_recording_id) + + cols, dropout = numeric_columns(raw, arm) + n = len(cols["left.obs_ee_pose" if arm != "right" else "right.obs_ee_pose"]) + valid = { + side: ~dropout[:, slot] & (arm in (side, "both")) + for slot, side in enumerate(("left", "right")) + } + start_ns = int(args.vendor_recording_id) * 1_000_000 + ts = start_ns + np.round(np.arange(n) * 1e9 / FPS).astype(np.int64) + pq.write_table( + steps_table(rid, ts, cols, valid), out / "steps.parquet", compression="zstd" + ) + t_numeric = time.time() - t0 + + cams = encode_videos(raw, out, n, args.ffmpeg) + k = np.asarray(Eva.INTRINSICS, np.float64)[:3, :3].copy() + k[:2] *= SIZE[0] / 480 + meta = { + "vendor_recording_id": args.vendor_recording_id, + "recording_id": rid, + "episode_hash": timestamp_ms_to_episode_hash(args.vendor_recording_id), + "embodiment": embodiment, + "site": md.get("lab"), + "operator": md.get("operator"), + "scene": md.get("scene"), + "task_name": md.get("task"), + "task_description": md.get("task_description"), + "objects": md.get("objects") or [], + "fps": FPS, + "num_frames": n, + "videos": cams, + "t_start_ns": int(ts[0]), + "t_end_ns": int(ts[-1]) + round(1e9 / FPS), + "intrinsics": {"front_1": k.reshape(-1).tolist()}, + # Same per-arm base-to-camera transforms eva_to_zarr stores. + "extrinsics": {a: m.reshape(-1).tolist() for a, m in Eva.EXTRINSICS.items()}, + "valid_frac": {s: float(v.mean()) for s, v in valid.items()}, + "seconds_numeric": round(t_numeric, 1), + "seconds_total": round(time.time() - t0, 1), + } + (out / "meta.json").write_text(json.dumps(meta, indent=1)) + return meta + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--raw-dir", type=Path, required=True) + ap.add_argument("--out-dir", type=Path, required=True) + ap.add_argument( + "--vendor-recording-id", required=True, help="epoch ms, the HDF5 file stem" + ) + ap.add_argument("--ffmpeg", default="/coc/flash7/rco3/ffmpeg-7.1-shared/bin/ffmpeg") + ap.add_argument( + "--stage-to", help="staging root (s3:// or local) for the committer" + ) + args = ap.parse_args() + logging.basicConfig(level=logging.INFO) + print(json.dumps(convert(args), indent=1)) + if args.stage_to: + commit.stage( + args.out_dir / args.vendor_recording_id, + commit.join(args.stage_to, "eva", args.vendor_recording_id), + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/eva_process/eva_utils.py b/egomimic/scripts/eva_process/eva_utils.py index 669025c7c1..04cc7b5d46 100644 --- a/egomimic/scripts/eva_process/eva_utils.py +++ b/egomimic/scripts/eva_process/eva_utils.py @@ -22,7 +22,7 @@ class EvaHD5Extractor: @staticmethod - def process_episode(episode_path, arm): + def process_episode(episode_path, arm, with_images=True): """ Extracts all feature keys from a given episode and returns as a dictionary Parameters @@ -49,7 +49,7 @@ def process_episode(episode_path, arm): with h5py.File(episode_path, "r") as episode: episode_name = Path(episode_path).stem - for camera in EvaHD5Extractor.get_cameras(episode): + for camera in EvaHD5Extractor.get_cameras(episode) if with_images else (): images = ( torch.from_numpy(episode["observations"]["images"][camera][:]) .permute(0, 3, 1, 2) @@ -73,8 +73,11 @@ def process_episode(episode_path, arm): episode_dt = datetime.datetime.fromtimestamp( timestamp_ms / 1000.0, tz=datetime.timezone.utc ) + # Per arm slot, frames clean_zero_data fills in any of the arrays it cleans. + dropout = np.zeros((len(episode_feats["obs_eepose"]), 2), bool) for key in ACTION_KEYS: if key in episode_feats: + dropout |= EvaHD5Extractor.zero_pose_frames(episode_feats[key]) episode_feats[key] = EvaHD5Extractor.clean_zero_data( episode_feats[key] ) @@ -91,6 +94,7 @@ def process_episode(episode_path, arm): else: value = EMBODIMENT.EVA_BIMANUAL.value + episode_feats["metadata.pose_dropout"] = dropout episode_feats["metadata.embodiment"] = np.full( (num_timesteps, 1), value, dtype=np.int32 ) @@ -148,6 +152,13 @@ def get_cmd_state(hdf5_data: h5py.File): states = [key for key in hdf5_data["/actions"]] return states + @staticmethod + def zero_pose_frames(data: np.ndarray) -> np.ndarray: + """(N, 2) per arm slot of a (N, 14) array: all six xyzypr values are zero.""" + return np.stack( + [(data[:, s] == 0).all(1) for s in (slice(0, 6), slice(7, 13))], axis=1 + ) + @staticmethod def clean_zero_data(data: np.ndarray) -> np.ndarray: """ @@ -164,11 +175,8 @@ def clean_zero_data(data: np.ndarray) -> np.ndarray: """ data = data.copy() - arm_pose_slices = [slice(0, 6), slice(7, 13)] # left, right - - for pose_slice in arm_pose_slices: - zero_mask = np.all(data[:, pose_slice] == 0, axis=1) # (N,) - + zero = EvaHD5Extractor.zero_pose_frames(data) + for pose_slice, zero_mask in zip((slice(0, 6), slice(7, 13)), zero.T): if not np.any(zero_mask): continue diff --git a/egomimic/scripts/format_bench/__init__.py b/egomimic/scripts/format_bench/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/format_bench/bench_ego.py b/egomimic/scripts/format_bench/bench_ego.py new file mode 100644 index 0000000000..826d58369d --- /dev/null +++ b/egomimic/scripts/format_bench/bench_ego.py @@ -0,0 +1,231 @@ +"""zarr / Lance loader bench: startup phases and steady-state throughput, one JSON line per run. + +Native loaders without transforms (LeRobot has no equivalent of our transforms, so all three +formats are measured on read + decode only): + zarr LocalEpisodeResolver over a directory of stores (local), or one ZarrDataset per store + under an s3:// prefix opened through obstore (s3patch) + zarr_fast the same stores through format_bench/zarr_fast.py: arrays opened once per worker, + shard indexes cached, a sample's (or a batch's) inner-chunk ranges read concurrently + lance LanceEpisodeResolver + MP4 decode (local), LanceStreamingResolver + MP4 ranges (s3://) +Patterns: + chunk1 Aria keypoints keymap: 30-step action chunks, proprio, 1 front frame + numeric same keymap without the camera + frames8 chunk1 with 8 consecutive front frames (camera history 8) + rdt RDT base's sample: chunk1 plus the front frame 0.1 s back (image_history_gap_s) +""" + +import argparse +import multiprocessing as mp +import os +import time +from pathlib import Path + +from egomimic.scripts.format_bench import ( # noqa: F401 (s3patch patches on import) + hooks, + s3patch, +) +from egomimic.scripts.format_bench.common import emit, evict, run_loader + +CONTEXT = {"zarr": "fork", "zarr_fast": "fork", "lance": "forkserver"} + + +def keymap(pattern: str) -> dict: + from egomimic.rldb.embodiment.human import Human + + km = Human.get_keymap( + keymap_mode="keypoints", + include_ee_pose=True, + image_history_gap_s=0.1 if pattern == "rdt" else None, + ) + if pattern == "numeric": + return {k: v for k, v in km.items() if v.get("key_type") != "camera_keys"} + if pattern == "frames8": + return { + k: {**v, "history": 8} if v.get("key_type") == "camera_keys" else v + for k, v in km.items() + } + return km + + +def zarr_s3_leaves(root: str, key_map: dict, embodiment: str | None) -> dict: + """One LIST for the store prefixes, then a ZarrDataset per store (the s3 analogue of the + local resolver's directory listing).""" + from egomimic.rldb.zarr.zarr_dataset_multi import ZarrDataset + + store = s3patch.obstore_s3(root) + names = sorted( + p.rstrip("/").split("/")[-1] + for p in store.list_with_delimiter()["common_prefixes"] + ) + leaves = {n: ZarrDataset(f"{root.rstrip('/')}/{n}", key_map, None) for n in names} + if embodiment: + leaves = {n: d for n, d in leaves.items() if d.embodiment.lower() == embodiment} + return leaves + + +def build(args, key_map): + from egomimic.rldb.filters import DatasetFilter + + # The training configs' filter; expected_embodiment alone only warns in LocalEpisodeResolver. + filters = ( + DatasetFilter( + filter_lambdas=[f"lambda row: row['embodiment'] == '{args.embodiment}'"] + ) + if args.embodiment + else None + ) + remote = args.root.startswith("s3://") + t0 = time.time() + if args.format in ("zarr", "zarr_fast"): + from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver + from egomimic.scripts.format_bench import zarr_fast + + fast = args.format == "zarr_fast" + if remote: + leaves = (zarr_fast.s3_leaves if fast else zarr_s3_leaves)( + args.root, key_map, args.embodiment + ) + else: + resolver = ( + zarr_fast.FastLocalEpisodeResolver if fast else LocalEpisodeResolver + ) + leaves = resolver( + Path(args.root), key_map=key_map, transform_list=None + ).resolve(filters=filters, expected_embodiment=args.embodiment) + else: + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.lance.dataset import ( + LanceEpisodeResolver, + LanceStreamingResolver, + ) + + has_cam = any(v.get("key_type") == "camera_keys" for v in key_map.values()) + transforms = ( + Human.get_transform_list(mode="keypoints_headframe_6d", stride=3) + if args.transforms + else None + ) + cams = {"images.front_1": args.camera} if has_cam else {} + if remote: + r = LanceStreamingResolver( + args.root, + key_map=key_map, + transform_list=transforms, + cameras=cams, + ) + else: + r = LanceEpisodeResolver( + args.root, key_map=key_map, transform_list=transforms, cameras=cams + ) + leaves = r.resolve(filters=filters, expected_embodiment=args.embodiment) + t1 = time.time() + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + ds = MultiDataset(datasets=leaves, mode="total") + t2 = time.time() + return ds, leaves, t1 - t0, t2 - t1 + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--format", choices=["zarr", "zarr_fast", "lance"], required=True) + ap.add_argument( + "--root", required=True, help="zarr store dir / Lance root, local or s3://" + ) + ap.add_argument( + "--pattern", choices=["chunk1", "numeric", "frames8", "rdt"], default="chunk1" + ) + ap.add_argument( + "--transforms", + action="store_true", + help="Lance only: apply the training transforms (keypoints_headframe_6d)", + ) + ap.add_argument("--camera", default="front_1_480") + ap.add_argument( + "--embodiment", default="human_bimanual", help="'' for every recording" + ) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--context", default=None) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=150) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument( + "--real", + type=Path, + help="first batch samples only these leaves (startup at scale)", + ) + ap.add_argument( + "--evict", nargs="*", default=[], help="paths to drop from the page cache first" + ) + ap.add_argument("--repeat", type=int, default=0) + ap.add_argument("--label", default="") + ap.add_argument("--out") + args = ap.parse_args() + args.embodiment = args.embodiment or None + + context = args.context or ( + "forkserver" if args.root.startswith("s3://") else CONTEXT[args.format] + ) + if context == "forkserver": + mp.set_forkserver_preload( + [ + "torch", + "egomimic.rldb.lance.dataset", + "egomimic.rldb.zarr.zarr_dataset_multi", + "egomimic.scripts.format_bench.s3patch", + "egomimic.scripts.format_bench.hooks", + "egomimic.scripts.format_bench.zarr_fast", + ] + ) + evicted = evict(*args.evict) if args.evict else 0 + km = keymap(args.pattern) + ds, leaves, resolve_s, construct_s = build(args, km) + indices = None + if args.real: + real = set(args.real.read_text().split()) + ranges = ds.index_map.ranges() + indices = [i for n in leaves if n in real for i in ranges[n]] + hook = None + if ( + args.format == "lance" + and not args.root.startswith("s3://") + and args.pattern != "numeric" + ): + hook = hooks.lance_decoders + res = run_loader( + ds, + workers=args.workers, + context=context, + batch_size=args.batch_size, + batches=args.batches, + warmup=args.warmup, + hook=hook, + indices=indices, + ) + emit( + dict( + label=args.label, + format=args.format, + location="s3" if args.root.startswith("s3://") else "local", + root=args.root, + pattern=args.pattern, + transforms=args.transforms, + workers=args.workers, + context=context, + batch_size=args.batch_size, + repeat=args.repeat, + episodes=len(leaves), + samples=len(ds), + evicted_gb=round(evicted / 1e9, 2), + resolve_s=round(resolve_s, 3), + construct_s=round(construct_s, 3), + **res, + startup_total_s=round(resolve_s + construct_s + res["first_batch_s"], 3), + ), + args.out, + ) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/bench_lerobot.py b/egomimic/scripts/format_bench/bench_lerobot.py new file mode 100644 index 0000000000..ebd102ba37 --- /dev/null +++ b/egomimic/scripts/format_bench/bench_lerobot.py @@ -0,0 +1,368 @@ +"""LeRobot v3.0 loader bench (LeRobot venv): same JSON fields as bench_ego.py. + +`LeRobotDataset` as shipped, filtered to the bimanual episodes with its own `episodes=` argument. +Phases: resolve = everything in __init__ except the HF dataset load (metadata, episode table, +cache checks); construct = `load_nested_dataset` (Arrow cache build when cold, memory-map when +warm). `--cold` empties HF_DATASETS_CACHE first. + +s3:// roots go through `S3LeRobotDataset`: LeRobot 0.6.1 only reads a local root or the HF Hub, +so the shim mirrors `meta/` locally (counted in resolve), loads data parquet from s3:// via fsspec, +and hands torchcodec s3:// video paths (LeRobot's decoder cache already opens through fsspec). +The endpoint comes from FSSPEC_S3_ENDPOINT_URL and AWS_* credentials. +""" + +import argparse +import os +import shutil +import time +from pathlib import Path + +from egomimic.scripts.format_bench import hooks +from egomimic.scripts.format_bench.common import emit, evict, run_loader + +HORIZON_KEYS = [ + f"{s}.{k}" + for s in ("left", "right") + for k in ("keypoints", "ee_pose", "wrist_pose") +] +CAM = "observation.images.front_1" +FPS = 30 + + +def delta_timestamps(pattern: str, packed: bool) -> dict: + keys = ["action"] if packed else HORIZON_KEYS + d = {k: [i / FPS for i in range(30)] for k in keys} + if pattern == "frames8": + d[CAM] = [i / FPS for i in range(-7, 1)] + if pattern == "rdt": + d[CAM] = [-3 / FPS, 0.0] + return d + + +def ds_features(root: Path) -> dict: + import json + + return json.loads((Path(root) / "meta/info.json").read_text())["features"] + + +def timed_hf_load(reader_cls, box: dict): + orig = reader_cls._load_hf_dataset + + def wrapped(self): + t = time.time() + out = orig(self) + box["construct_s"] = box.get("construct_s", 0.0) + time.time() - t + return out + + reader_cls._load_hf_dataset = wrapped + + +class S3: + root = None + local = None + + +def patch_s3(): + """Patches DatasetReader in place (LeRobotDataset imports the class directly), so forked + workers inherit it.""" + import datasets + import fsspec + import pyarrow.dataset as pa_ds + from lerobot.datasets import dataset_reader + from lerobot.datasets.feature_utils import get_hf_features_from_features + from lerobot.datasets.io_utils import hf_transform_to_torch + + # LeRobot's decoder cache opens videos with a bare fsspec.open, so s3fs defaults apply: a 50 MB + # readahead block per random seek, ~30 MB fetched for each ~70 KB frame. + fsspec.config.conf.setdefault("s3", {}).update( + default_cache_type="readahead", default_block_size=256 * 2**10 + ) + + def _load_hf_dataset(self): + fs = fsspec.filesystem("s3") + paths = sorted(fs.glob(f"{S3.root}/data/*/*.parquet")) + ds = datasets.Dataset.from_parquet( + [f"s3://{p}" for p in paths], + features=get_hf_features_from_features(self._meta.features), + # The parquet download keeps s3fs's default 50 MB block. + storage_options={**fs.storage_options, "default_block_size": 50 * 2**20}, + filters=None + if self.episodes is None + else pa_ds.field("episode_index").isin(self.episodes), + ) + ds.set_transform(hf_transform_to_torch) + return ds + + orig_decode = dataset_reader.decode_video_frames + + def decode_s3(video_path, *a, **kw): + rel = Path(video_path).relative_to(S3.local) + return orig_decode(f"{S3.root}/{rel}", *a, **kw) + + dataset_reader.DatasetReader._load_hf_dataset = _load_hf_dataset + dataset_reader.DatasetReader._check_cached_episodes_sufficient = ( + lambda self: self.hf_dataset is not None and len(self.hf_dataset) > 0 + ) + dataset_reader.decode_video_frames = decode_s3 + + +STREAM_COLUMNS = [ + "action", + "observation.state", + "timestamp", + "frame_index", + "episode_index", + "index", + "task_index", +] + + +class S3ParquetRows: + """Stand-in for `DatasetReader.hf_dataset` that reads rows from s3:// parquet per access, + so numerics stream like zarr and Lance instead of being downloaded at construct. + + The row -> (file, row group) map comes from meta/episodes, so the parent reads no parquet. + Each worker opens files lazily through pyarrow's S3 client (footer kept, column chunks + pre-buffered concurrently) and keeps the last FB_LR_RG_CACHE row groups (default 1, enough + for the item row, its action window and frame timestamps to share one fetch). Parquet's + smallest unit is a row group's column chunk: one episode, ~4 MB for `action`.""" + + def __init__(self, meta, episodes, s3_root: str, columns=STREAM_COLUMNS): + import numpy as np + + eps = meta.episodes + file_start = {} + for ep in eps: + f = (ep["data/chunk_index"], ep["data/file_index"]) + file_start[f] = min(file_start.get(f, 2**62), ep["dataset_from_index"]) + sel = sorted(range(len(eps)) if episodes is None else episodes) + abs_idx, row_of, self.file = [], [], [] + for e in sel: + ep = eps[e] + f = (ep["data/chunk_index"], ep["data/file_index"]) + r = np.arange(ep["dataset_from_index"], ep["dataset_to_index"]) + abs_idx.append(r) + row_of.append(r - file_start[f]) + self.file += [f] * len(r) + self.abs = np.concatenate(abs_idx) + self.row = np.concatenate(row_of) + self.root = s3_root + self.columns = columns + self.cache_size = int(os.environ.get("FB_LR_RG_CACHE", "1")) + self._fs = None + self._pf = {} + self._rg = {} + + def __len__(self): + return len(self.abs) + + @property + def data(self): + import pyarrow as pa + + return pa.table({"index": self.abs}) + + def __getstate__(self): + return { + k: v for k, v in self.__dict__.items() if k not in ("_fs", "_pf", "_rg") + } | { + "_fs": None, + "_pf": {}, + "_rg": {}, + } + + def _file(self, f): + import pyarrow.fs as pafs + import pyarrow.parquet as pq + + if self._fs is None: + ep = os.environ["FSSPEC_S3_ENDPOINT_URL"] + self._fs = pafs.S3FileSystem( + endpoint_override=ep.split("://")[1], + scheme=ep.split("://")[0], + access_key=os.environ["AWS_ACCESS_KEY_ID"], + secret_key=os.environ["AWS_SECRET_ACCESS_KEY"], + region=os.environ.get("AWS_REGION", "us-east-1"), + ) + if f not in self._pf: + path = f"{self.root[len('s3://'):]}/data/chunk-{f[0]:03d}/file-{f[1]:03d}.parquet" + pf = pq.ParquetFile(self._fs.open_input_file(path), pre_buffer=True) + md = pf.metadata + starts = [0] + for i in range(md.num_row_groups): + starts.append(starts[-1] + md.row_group(i).num_rows) + self._pf[f] = (pf, starts) + return self._pf[f] + + def _row_group(self, f, rg): + key = (f, rg) + if key in self._rg: + self._rg[key] = self._rg.pop(key) + return self._rg[key] + pf, _ = self._file(f) + t = pf.read_row_group(rg, columns=self.columns) + self._rg[key] = t + while len(self._rg) > self.cache_size: + self._rg.pop(next(iter(self._rg))) + return t + + def _rows(self, rel: list[int]): + import bisect + + import pyarrow as pa + + parts = [] + for i in rel: + f, r = self.file[i], int(self.row[i]) + _, starts = self._file(f) + rg = bisect.bisect_right(starts, r) - 1 + parts.append(self._row_group(f, rg).slice(r - starts[rg], 1)) + return pa.concat_tables(parts) + + def __getitem__(self, idx): + from lerobot.datasets.io_utils import hf_transform_to_torch + + if isinstance(idx, str): + raise TypeError("column access is not streamed; index rows first") + single = not isinstance(idx, list | tuple) + rel = [int(idx)] if single else [int(i) for i in idx] + batch = hf_transform_to_torch(self._rows(rel).to_pydict()) + return {k: v[0] for k, v in batch.items()} if single else batch + + +def patch_stream_parquet(): + """After patch_s3: numerics read from s3:// parquet per sample (S3ParquetRows).""" + from lerobot.datasets import dataset_reader + + def _load_hf_dataset(self): + return S3ParquetRows(self._meta, self.episodes, S3.root) + + dataset_reader.DatasetReader._load_hf_dataset = _load_hf_dataset + + +def mirror_meta(s3_root: str, local: Path): + import fsspec + + fs = fsspec.filesystem("s3") + if local.exists(): + shutil.rmtree(local) + fs.get(f"{s3_root}/meta", str(local / "meta"), recursive=True) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--root", required=True) + ap.add_argument( + "--pattern", choices=["chunk1", "numeric", "frames8", "rdt"], default="chunk1" + ) + ap.add_argument( + "--episodes", type=Path, help="episode indices to load (whitespace separated)" + ) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--context", default="fork") + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=150) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument( + "--real", type=Path, help="first batch samples only these episode indices" + ) + ap.add_argument("--cold", action="store_true", help="empty HF_DATASETS_CACHE first") + ap.add_argument( + "--stream-parquet", + action="store_true", + help="s3:// only: read parquet rows per sample instead of downloading them at construct", + ) + ap.add_argument("--evict", nargs="*", default=[]) + ap.add_argument("--repeat", type=int, default=0) + ap.add_argument("--label", default="") + ap.add_argument("--out") + args = ap.parse_args() + + from lerobot.datasets import dataset_reader + + cache = Path( + os.environ.get("HF_DATASETS_CACHE", Path(os.environ["HF_HOME"]) / "datasets") + ) + if args.cold and cache.exists(): + shutil.rmtree(cache) + evicted = evict(*args.evict) if args.evict else 0 + remote = args.root.startswith("s3://") + box = {} + episodes = ( + [int(x) for x in args.episodes.read_text().split()] if args.episodes else None + ) + + t0 = time.time() + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + if remote: + patch_s3() + if args.stream_parquet: + patch_stream_parquet() + S3.root, S3.local = args.root.rstrip("/"), Path(os.environ["FB_META_MIRROR"]) + mirror_meta(S3.root, S3.local) + root = S3.local + else: + root = Path(args.root) + timed_hf_load(dataset_reader.DatasetReader, box) + ds = LeRobotDataset( + "egoverse/aria_pilot", + root=root, + episodes=episodes, + delta_timestamps=delta_timestamps(args.pattern, "action" in ds_features(root)), + video_backend="torchcodec", + # float32 timestamps ~20 min into a 200 MB concatenated file are only good to ~1e-4 s, + # which trips the 1e-4 default. + tolerance_s=1e-3, + ) + if args.pattern == "numeric": + # LeRobot has no switch to skip video keys; hide them so get_item never decodes. + type(ds.meta).video_keys = property(lambda self: []) + init_s = time.time() - t0 + construct_s = box.get("construct_s", 0.0) + + indices = None + if args.real: + real = set(int(x) for x in args.real.read_text().split()) + ep_col = ds.hf_dataset.data.column("episode_index").to_numpy() + import numpy as np + + indices = np.flatnonzero(np.isin(ep_col, list(real))).tolist() + res = run_loader( + ds, + workers=args.workers, + context=args.context, + batch_size=args.batch_size, + batches=args.batches, + warmup=args.warmup, + hook=None if args.pattern == "numeric" or remote else hooks.lerobot_decoders, + indices=indices, + ) + emit( + dict( + label=args.label, + stream_parquet=args.stream_parquet, + format="lerobot" if "action" in ds_features(root) else "lerobot_unpacked", + location="s3" if remote else "local", + root=args.root, + pattern=args.pattern, + workers=args.workers, + context=args.context, + batch_size=args.batch_size, + repeat=args.repeat, + cold=args.cold, + episodes=ds.num_episodes, + samples=len(ds), + evicted_gb=round(evicted / 1e9, 2), + resolve_s=round(init_s - construct_s, 3), + construct_s=round(construct_s, 3), + **res, + startup_total_s=round(init_s + res["first_batch_s"], 3), + ), + args.out, + ) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/bench_lerobot_stream.py b/egomimic/scripts/format_bench/bench_lerobot_stream.py new file mode 100644 index 0000000000..fb05773e05 --- /dev/null +++ b/egomimic/scripts/format_bench/bench_lerobot_stream.py @@ -0,0 +1,193 @@ +"""StreamingLeRobotDataset bench (LeRobot venv): same JSON fields as bench_lerobot.py. + +LeRobot 0.6.1's streaming class as shipped, with three bench-only fixes it needs to read this +dataset at all: + data `load_dataset` reads `/data/*/*.parquet` (local or s3://) instead of a Hub repo, + filtered with parquet `filters` (its own `episodes=` is stored but never used) + timestamp `make_frame` seeks video to `index / fps`, the frame's position in the whole + dataset, but each concatenated mp4 restarts at 0; `index` is rewritten to the frame's + position in its own video file (from_timestamp * fps + frame_index) + workers every worker builds the same shard iterators, and inside a worker HF hands each + 1-shard iterator to worker 0 only (the rest yield nothing); each worker instead + streams its own episodes (episode % workers, a parquet filter that skips the other + row groups) through the stock iterator with HF's own worker split off +Shuffling is LeRobot's: per-worker shard sampling plus a `--buffer` sample shuffle buffer, filled +before the first sample comes out (counted in startup). +""" + +import argparse +import os +import time +from pathlib import Path + +import numpy as np + +from egomimic.scripts.format_bench import hooks +from egomimic.scripts.format_bench.bench_lerobot import ( + delta_timestamps, + ds_features, + mirror_meta, +) +from egomimic.scripts.format_bench.common import emit, evict, run_loader + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--root", required=True) + ap.add_argument( + "--pattern", choices=["chunk1", "numeric", "frames8"], default="chunk1" + ) + ap.add_argument( + "--episodes", type=Path, help="episode indices to stream (whitespace separated)" + ) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--context", default="fork") + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=150) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument( + "--buffer", + type=int, + default=None, + help="shuffle buffer (default 1000; 100 for frames8)", + ) + ap.add_argument("--evict", nargs="*", default=[]) + ap.add_argument("--repeat", type=int, default=0) + ap.add_argument("--label", default="") + ap.add_argument("--out") + args = ap.parse_args() + buffer = args.buffer or (100 if args.pattern == "frames8" else 1000) + + import fsspec + import torch + from lerobot.datasets import streaming_dataset as sd + + evicted = evict(*args.evict) if args.evict else 0 + remote = args.root.startswith("s3://") + root = args.root.rstrip("/") + episodes = ( + [int(x) for x in args.episodes.read_text().split()] if args.episodes else None + ) + + t0 = time.time() + if remote: + # Same s3fs settings as bench_lerobot.patch_s3: 256 KB readahead for video seeks, the + # default 50 MB block for parquet. + fsspec.config.conf.setdefault("s3", {}).update( + default_cache_type="readahead", default_block_size=256 * 2**10 + ) + fs = fsspec.filesystem("s3") + files = [f"s3://{p}" for p in sorted(fs.glob(f"{root}/data/*/*.parquet"))] + storage_options = { + "s3": {**fs.storage_options, "default_block_size": 50 * 2**20} + } + meta_root = Path(os.environ["FB_META_MIRROR"]) + mirror_meta(root, meta_root) + else: + files = sorted(str(p) for p in Path(root).glob("data/*/*.parquet")) + storage_options = {} + meta_root = Path(root) + pattern = args.pattern + orig_load = sd.load_dataset + + class Stream(sd.StreamingLeRobotDataset): + def stream(self, eps): + """The parquet rows of `eps`, `index` made relative to each row's video file.""" + hf = orig_load( + "parquet", + data_files=files, + split="train", + streaming=True, + filters=None if eps is None else [("episode_index", "in", list(eps))], + storage_options=storage_options, + ) + if pattern == "numeric": + return hf + (vk,) = self.meta.video_keys + e = self.meta.episodes + ep_idx = np.asarray(e["episode_index"]) + off = np.zeros(ep_idx.max() + 1, np.int64) + off[ep_idx] = np.round( + np.asarray(e[f"videos/{vk}/from_timestamp"]) * self.fps + ).astype(np.int64) + return hf.map( + lambda ep, fi: {"index": int(off[ep] + fi)}, + input_columns=["episode_index", "frame_index"], + ) + + def __iter__(self): + info = torch.utils.data.get_worker_info() + if info is not None: + eps = ( + episodes + if episodes is not None + else range(self.meta.total_episodes) + ) + self.hf_dataset = self.stream( + [e for e in eps if e % info.num_workers == info.id] + ) + self.num_shards = min(self.hf_dataset.num_shards, self._max_shards) + # HF splits shards across torch workers by itself; this worker's stream is + # already its own, so let HF iterate it whole. + torch.utils.data.get_worker_info = lambda: None + yield from super().__iter__() + + sd.load_dataset = lambda *a, **kw: orig_load( + "parquet", + data_files=files[:1], + split="train", + streaming=True, + storage_options=storage_options, + ) + ds = Stream( + "egoverse/aria_pilot", + root=meta_root, + delta_timestamps=delta_timestamps(pattern, "action" in ds_features(meta_root)), + tolerance_s=1e-3, + buffer_size=buffer, + ) + ds._max_shards = 16 + if remote: + ds.root = root # video paths are f"{root}/{relpath}", opened through fsspec + if pattern == "numeric": + type(ds.meta).video_keys = property(lambda self: []) + ds.hf_dataset = ds.stream(episodes) + init_s = time.time() - t0 + + res = run_loader( + ds, + workers=args.workers, + context=args.context, + batch_size=args.batch_size, + batches=args.batches, + warmup=args.warmup, + hook=None if args.pattern == "numeric" else hooks.lerobot_stream_decoders, + ) + n_eps = len(episodes) if episodes is not None else ds.meta.total_episodes + emit( + dict( + label=args.label, + format="lerobot_stream", + location="s3" if remote else "local", + root=args.root, + pattern=args.pattern, + workers=args.workers, + context=args.context, + batch_size=args.batch_size, + repeat=args.repeat, + episodes=n_eps, + buffer=buffer, + files=len(files), + evicted_gb=round(evicted / 1e9, 2), + resolve_s=round(init_s, 3), + construct_s=0.0, + **res, + startup_total_s=round(init_s + res["first_batch_s"], 3), + ), + args.out, + ) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/bench_update.py b/egomimic/scripts/format_bench/bench_update.py new file mode 100644 index 0000000000..b1feda674e --- /dev/null +++ b/egomimic/scripts/format_bench/bench_update.py @@ -0,0 +1,1016 @@ +"""Time whole-dataset edits in zarr vs LeRobot v3.1 vs Lance: replace head_pose, append a new +per-row column, rewrite every episode's annotations. + +`make` builds the same annotated rows in every format (no video, which none of these edits +touch); `update --op` applies one edit and times it end to end, then verifies it; `fit` projects +to a target row count. Roots are local paths or s3://bucket/prefix; s3 goes to FB_S3_ENDPOINT +(plain http, MinIO) with the R2_* credentials, as in s3patch.py. + +Rows are real: fake episodes are contiguous slices of the pilot recordings (Lance steps table), +with lengths cycling through the pilot zarr manifest's, so every format compresses like real +data and zarr has the real object count. The global row index g (`index` in LeRobot, +timestamp_ns // DT_NS in Lance, episode start + frame in zarr) maps back to the source row; edit +generation k writes source head_pose + k * 1e-3 and annotation text tagged with k. + +Annotations are one span per SPAN_FRAMES frames, stored as each format does: + 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 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 + append (its `features` attr lists every array): replace = 1 shard per episode, + append = array zarr.json + shard + group zarr.json, annotations = zarr.json + shard + lerobot every data parquet file in full (one row group per episode), plus meta/episodes + 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: overwrite annotations.lance +""" + +import argparse +import csv +import functools +import json +import os +import socket +import struct +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +DT_NS = 33_333_333 +FPS = 30 +CHUNK = 100 +SPAN_FRAMES = 300 +PARQUET_MB = 100 +CHUNKS_SIZE = 1000 +LANCE_ROWS_PER_FILE = 1024 * 1024 +STATS = ("min", "max", "mean", "std") +HANDS = ("wrist_pose", "ee_pose", "keypoints", "aria_keypoints") +# Flattened HUMAN_STEPS, as to_lerobot.py lays it out (copied: to_lerobot imports lerobot). +VECTORS = { + "head_pose": 7, + "eye_gaze": 3, + **{f"{s}.{f}": d for s in ("left", "right") for f, d in zip(HANDS, (7, 7, 63, 63))}, +} +SCALARS = ("left.confidence", "right.confidence") +BOOLS = ("left_valid", "right_valid", "head_pose_valid") +FORMATS = ("zarr", "lerobot", "lerobot_zstd", "lance") +OPS = ("replace", "append", "annotations") +COMPRESSION = {"lerobot": "snappy", "lerobot_zstd": "zstd"} +LANGUAGE_COLS = ["language_persistent", "language_events"] +# lerobot/datasets/language.py row types; tool_calls is its pa.string() fallback for pa.json_(). +LANG_PERSISTENT = pa.list_( + pa.struct( + [ + pa.field("role", pa.string(), nullable=False), + pa.field("content", pa.string()), + pa.field("style", pa.string()), + pa.field("timestamp", pa.float32(), nullable=False), + pa.field("camera", pa.string()), + pa.field("tool_calls", pa.list_(pa.string())), + ] + ) +) +LANG_EVENTS = pa.list_( + pa.struct( + [ + pa.field("role", pa.string(), nullable=False), + pa.field("content", pa.string()), + pa.field("style", pa.string()), + pa.field("camera", pa.string()), + pa.field("tool_calls", pa.list_(pa.string())), + ] + ) +) + + +# --------------------------------------------------------------------------- rows +def nested(t: pa.Table, dotted: str) -> pa.Array: + head, *rest = dotted.split(".") + col = t.column(head).combine_chunks() + for p in rest: + col = col.field(p) + return col + + +def vec(arr: pa.FixedSizeListArray, d: int) -> np.ndarray: + # .values keeps null-backed slots (flatten() drops them) but ignores the array offset. + v = np.asarray(arr.values) + return v[arr.offset * d : (arr.offset + len(arr)) * d].reshape(-1, d) + + +def fsl(a: np.ndarray) -> pa.Array: + return pa.FixedSizeListArray.from_arrays(pa.array(a.reshape(-1)), a.shape[1]) + + +def feat_name(k: int) -> str: + return f"feat_{k}" + + +def spans(ep: int, n: int, k: int) -> list[tuple[str, int, int]]: + """(text, start frame, end frame inclusive) per span, as ZarrWriter takes them.""" + return [ + ( + f"ep {ep} subtask {i} v{k}: pick up the folded shirt and place it on the stack", + i * SPAN_FRAMES, + min(n, (i + 1) * SPAN_FRAMES) - 1, + ) + for i in range(n // SPAN_FRAMES) + ] + + +class Layout: + """Fake episodes cut from source recordings: episode i has length ep_lens[i % len] (clipped + to its source recording, the last one to `rows`) and starts at a fixed offset inside + recording i % R.""" + + def __init__(self, steps: pa.Table, ep_lens: list[int], rows: int): + self.pool = {k: vec(nested(steps, k), d) for k, d in VECTORS.items()} + self.pool.update( + {k: pc.fill_null(nested(steps, k), 0.0).to_numpy() for k in SCALARS} + ) + self.pool.update( + { + k: pc.fill_null(nested(steps, k), False).to_numpy(zero_copy_only=False) + for k in BOOLS + } + ) + ids = steps.column("recording_id").to_numpy(zero_copy_only=False) + rec_starts = np.concatenate([[0], np.flatnonzero(ids[1:] != ids[:-1]) + 1]) + rec_lens = np.diff(np.append(rec_starts, len(ids))) + lens, src, total = [], [], 0 + while total < rows: + ep = len(lens) + r = ep % len(rec_lens) + n = int(min(ep_lens[ep % len(ep_lens)], rec_lens[r], rows - total)) + off = rec_starts[r] + (ep * 7919) % (rec_lens[r] - n + 1) + src.append(np.arange(off, off + n)) + lens.append(n) + total += n + self.lens = np.asarray(lens) + self.starts = np.concatenate([[0], np.cumsum(self.lens)[:-1]]) + self.src = np.concatenate(src) + + @property + def n_eps(self) -> int: + return len(self.lens) + + def g(self, ep: int) -> np.ndarray: + return np.arange(self.starts[ep], self.starts[ep] + self.lens[ep]) + + def episode(self, ep: int) -> dict[str, np.ndarray]: + g = self.g(ep) + cols = {k: v[self.src[g]] for k, v in self.pool.items()} + cols["g"] = g + return cols + + def new_head_pose(self, g: np.ndarray, k: int) -> np.ndarray: + return self.pool["head_pose"][self.src[g]] + np.float32(k * 1e-3) + + def spans(self, ep: int, k: int) -> list[tuple[str, int, int]]: + return spans(ep, int(self.lens[ep]), k) + + +def pilot_layout(pilot: Path, rows: int) -> Layout: + import lance + + steps = lance.dataset(str(pilot / "lance/steps/aria.lance")).to_table() + with open(pilot / "manifest.csv") as f: + ep_lens = [int(r["num_frames"]) for r in csv.DictReader(f)] + return Layout(steps, ep_lens, rows) + + +# --------------------------------------------------------------------------- storage +def is_s3(root: str) -> bool: + return root.startswith("s3://") + + +def creds() -> tuple[str, str, str]: + return ( + os.environ["FB_S3_ENDPOINT"], + os.environ["R2_ACCESS_KEY_ID"], + os.environ["R2_SECRET_ACCESS_KEY"], + ) + + +def arrow_fs(root: str) -> tuple[pafs.FileSystem, str]: + if not is_s3(root): + return pafs.LocalFileSystem(), str(Path(root).absolute()) + ep, key, secret = creds() + scheme, host = ep.split("://") + fs = pafs.S3FileSystem( + access_key=key, + secret_key=secret, + endpoint_override=host, + scheme=scheme, + region="us-east-1", + ) + return fs, root[len("s3://") :] + + +def lance_opts(root: str) -> dict | None: + if not is_s3(root): + return None + ep, key, secret = creds() + return { + "access_key_id": key, + "secret_access_key": secret, + "endpoint": ep, + "region": "us-east-1", + "allow_http": "true", + } + + +def obj_store(root: str): + from obstore.store import LocalStore, S3Store + + if not is_s3(root): + return LocalStore(prefix=Path(root).absolute(), mkdir=True) + ep, key, secret = creds() + bucket, prefix = root[len("s3://") :].split("/", 1) + return S3Store( + bucket, + prefix=prefix.rstrip("/"), + endpoint=ep, + access_key_id=key, + secret_access_key=secret, + region="us-east-1", + client_options={"allow_http": True}, + ) + + +def mkparent(fs, path: str) -> None: + if isinstance(fs, pafs.LocalFileSystem): + fs.create_dir(path.rsplit("/", 1)[0], recursive=True) + + +def write_json(fs, path: str, obj) -> None: + mkparent(fs, path) + with fs.open_output_stream(path) as f: + f.write(json.dumps(obj, indent=2).encode()) + + +def read_json(fs, path: str): + with fs.open_input_stream(path) as f: + return json.loads(f.read()) + + +def emit(rec: dict, out: str | None) -> None: + # common.emit imports torch, which takes minutes on a loaded node. + line = json.dumps( + {"node": socket.gethostname().split(".")[0], "t": time.time(), **rec} + ) + print(line, flush=True) + if out: + with open(out, "a") as f: + f.write(line + "\n") + + +def pmap(fn, items, workers: int) -> list: + with ThreadPoolExecutor(workers) as ex: + return list(ex.map(fn, items)) + + +# --------------------------------------------------------------------------- zarr +def ep_name(ep: int) -> str: + return f"ep-{ep:07d}.zarr" + + +def padded(hp: np.ndarray) -> np.ndarray: + """float64, zero-padded to a whole inner chunk, as ZarrWriter stores it.""" + out = np.zeros((-(-len(hp) // CHUNK) * CHUNK, hp.shape[1])) + out[: len(hp)] = hp + return out + + +def shard(chunks: list[bytes]) -> bytes: + """zarr 3 sharding_indexed with the default index codecs (bytes + crc32c, at the end): + zstd(0)-compressed inner chunks, then (offset, nbytes) uint64 pairs and their crc32c.""" + import google_crc32c + from numcodecs import Zstd + + z = Zstd(level=0) + enc = [z.encode(c) for c in chunks] + offsets = np.cumsum([0] + [len(e) for e in enc[:-1]]) + index = np.stack([offsets, [len(e) for e in enc]], 1).astype(" bytes: + a = np.ascontiguousarray(a, " list[bytes]: + return [ + json.dumps( + {"text": t, "start_idx": s, "end_idx": e}, + ensure_ascii=False, + separators=(",", ":"), + ).encode() + for t, s, e in sp + ] + + +def vlen_shard(items: list[bytes]) -> bytes: + """One inner chunk holding every item (ZarrWriter's annotations layout), vlen-bytes encoded.""" + body = struct.pack(" bytes: + """zarr.json of an array as ZarrWriter creates it, rendered by zarr once per shape.""" + import zarr + from zarr.core.dtype import VariableLengthBytes + + d = {} + edge = (max(1, shape[0]),) + zarr.create_array( + zarr.storage.MemoryStore(d), + shape=shape, + dtype=VariableLengthBytes() if vlen else "float64", + chunks=edge if vlen else (CHUNK, *shape[1:]), + shards=edge if vlen else shape, + ) + return d["zarr.json"].to_bytes() + + +def feature(arr: np.ndarray | None) -> dict: + if arr is None: + return { + "dtype": "json", + "shape": [0], + "names": ["json"], + "format": "annotation_v1", + } + return {"dtype": "float64", "shape": [arr.shape[1]], "names": ["dim_0"]} + + +def put_array(store, path: str, meta: bytes, data: bytes | None) -> None: + import obstore + + obstore.put(store, f"{path}/zarr.json", meta) + key = f"{path}/c/0" if path.endswith("annotations") else f"{path}/c/0/0" + if data is None: + # An empty array has no chunk object; S3 deletes are idempotent, local ones are not. + try: + obstore.delete(store, key) + except FileNotFoundError: + pass + else: + obstore.put(store, key, data) + + +def put_annotations(store, ep: int, sp: list[tuple[str, int, int]]) -> None: + items = annotation_items(sp) + put_array( + store, + f"{ep_name(ep)}/annotations", + array_meta((len(items),), True), + vlen_shard(items) if items else None, + ) + + +def make_zarr(root: str, lay: Layout, workers: int) -> None: + import zarr + + store = obj_store(root) + zstore = zarr.storage.ObjectStore(store) + + def one(ep): + hp = padded(lay.pool["head_pose"][lay.src[lay.g(ep)]]) + sp = lay.spans(ep, 0) + features = {"obs_head_pose": feature(hp), "annotations": feature(None)} + features["annotations"]["shape"] = [len(sp)] + zarr.open_group( + zstore, + path=ep_name(ep), + mode="w", + attributes={ + "embodiment": "human_bimanual", + "total_frames": int(lay.lens[ep]), + "fps": FPS, + "features": features, + }, + ) + put_array( + store, + f"{ep_name(ep)}/obs_head_pose", + array_meta(hp.shape, False), + numeric_shard(hp), + ) + put_annotations(store, ep, sp) + + pmap(one, range(lay.n_eps), workers) + + +def update_zarr(root: str, lay: Layout, op: str, k: int, workers: int) -> dict: + import obstore + + store = obj_store(root) + + def one(ep): + name = ep_name(ep) + if op == "annotations": + put_annotations(store, ep, lay.spans(ep, k)) + return 2 + hp = padded(lay.new_head_pose(lay.g(ep), k)) + if op == "replace": + obstore.put(store, f"{name}/obs_head_pose/c/0/0", numeric_shard(hp)) + return 1 + put_array( + store, + f"{name}/{feat_name(k)}", + array_meta(hp.shape, False), + numeric_shard(hp), + ) + group = json.loads(obstore.get(store, f"{name}/zarr.json").bytes().to_bytes()) + group["attributes"]["features"][feat_name(k)] = feature(hp) + obstore.put(store, f"{name}/zarr.json", json.dumps(group).encode()) + return 3 + + puts = pmap(one, range(lay.n_eps), workers) + return {"objects_written": sum(puts), "gets": lay.n_eps if op == "append" else 0} + + +def verify_zarr(root: str, lay: Layout, op: str, k: int) -> None: + import zarr + + store = zarr.storage.ObjectStore(obj_store(root)) + for ep in {0, lay.n_eps // 2, lay.n_eps - 1}: + g = zarr.open_group(store, path=ep_name(ep), mode="r") + if op == "annotations": + got = [json.loads(b) for b in g["annotations"][:]] + want = [json.loads(b) for b in annotation_items(lay.spans(ep, k))] + assert got == want, (ep, got[:1], want[:1]) + continue + key = "obs_head_pose" if op == "replace" else feat_name(k) + assert key in g.attrs["features"] + np.testing.assert_array_equal( + g[key][:], padded(lay.new_head_pose(lay.g(ep), k)) + ) + + +# --------------------------------------------------------------------------- lerobot +def language_columns( + sp: list[tuple[str, int, int]], n: int +) -> tuple[pa.Array, pa.Array]: + """The episode's spans broadcast to all n frames, and no events, as LeRobot's writer does.""" + rows = pa.array( + [ + { + "role": "assistant", + "content": t, + "style": "subtask", + "timestamp": s / FPS, + "camera": None, + "tool_calls": None, + } + for t, s, _ in sp + ], + LANG_PERSISTENT.value_type, + ) + m = len(rows) + persistent = pa.ListArray.from_arrays( + pa.array(np.arange(n + 1, dtype=np.int32) * m), + rows.take(pa.array(np.tile(np.arange(m), n))), + type=LANG_PERSISTENT, + ) + events = pa.ListArray.from_arrays( + pa.array(np.zeros(n + 1, np.int32)), + pa.array([], LANG_EVENTS.value_type), + type=LANG_EVENTS, + ) + return persistent, events + + +def lerobot_table(cols: dict, ep: int, sp: list) -> pa.Table: + n = len(cols["g"]) + t = {k: fsl(cols[k]) for k in VECTORS} + t.update({k: pa.array(cols[k]) for k in SCALARS + BOOLS}) + t["timestamp_ns"] = pa.array(cols["g"] * DT_NS) + t["timestamp"] = pa.array(np.arange(n, dtype=np.float32) / np.float32(FPS)) + t["frame_index"] = pa.array(np.arange(n, dtype=np.int64)) + t["episode_index"] = pa.array(np.full(n, ep, np.int64)) + t["index"] = pa.array(cols["g"]) + t["task_index"] = pa.array(np.zeros(n, np.int64)) + t["language_persistent"], t["language_events"] = language_columns(sp, n) + return pa.table(t) + + +def vec_stats(a: np.ndarray) -> dict[str, list[float]]: + a = a.astype(np.float64) + return { + "min": a.min(0).tolist(), + "max": a.max(0).tolist(), + "mean": a.mean(0).tolist(), + "std": a.std(0).tolist(), + } + + +def data_path(file_no: int) -> str: + c, f = divmod(file_no, CHUNKS_SIZE) + return f"data/chunk-{c:03d}/file-{f:03d}.parquet" + + +def episodes_meta(stats: list[dict], lens: np.ndarray, file_no: np.ndarray) -> pa.Table: + """meta/episodes rows as LeRobot v3 stores them: one stats// column per stat.""" + ends = np.cumsum(lens) + cols = { + "episode_index": np.arange(len(lens)), + "length": lens, + "dataset_from_index": ends - lens, + "dataset_to_index": ends, + "data/chunk_index": file_no // CHUNKS_SIZE, + "data/file_index": file_no % CHUNKS_SIZE, + } + t = {k: pa.array(np.asarray(v, np.int64)) for k, v in cols.items()} + for key in stats[0]: + for s in STATS: + t[f"stats/{key}/{s}"] = pa.array( + [st[key][s] for st in stats], pa.list_(pa.float64()) + ) + t[f"stats/{key}/count"] = pa.array( + [[int(n)] for n in lens], pa.list_(pa.int64()) + ) + return pa.table(t) + + +def aggregate(stats: list[dict], lens: np.ndarray, key: str) -> dict: + w = lens / lens.sum() + mean = w @ np.array([s[key]["mean"] for s in stats]) + m2 = w @ np.array( + [np.square(s[key]["std"]) + np.square(s[key]["mean"]) for s in stats] + ) + return { + "min": np.min([s[key]["min"] for s in stats], 0).tolist(), + "max": np.max([s[key]["max"] for s in stats], 0).tolist(), + "mean": mean.tolist(), + "std": np.sqrt(np.maximum(m2 - mean**2, 0)).tolist(), + } + + +def write_meta(fs, base: str, stats: list[dict], lens: np.ndarray, file_no: np.ndarray): + meta = episodes_meta(stats, lens, file_no) + rows_per_meta = 50_000 + written = 0 + for j in range(0, meta.num_rows, rows_per_meta): + p = f"{base}/meta/episodes/chunk-000/file-{j // rows_per_meta:03d}.parquet" + mkparent(fs, p) + pq.write_table( + meta.slice(j, rows_per_meta), p, filesystem=fs, compression="snappy" + ) + written += fs.get_file_info(p).size + write_json( + fs, f"{base}/meta/stats.json", {k: aggregate(stats, lens, k) for k in stats[0]} + ) + return written + + +def read_meta(fs, base: str) -> tuple[list[dict], np.ndarray, np.ndarray]: + meta = pq.read_table(f"{base}/meta/episodes/chunk-000", filesystem=fs) + keys = sorted( + {c.split("/")[1] for c in meta.column_names if c.startswith("stats/")} + ) + cols = { + (k, s): meta.column(f"stats/{k}/{s}").to_pylist() for k in keys for s in STATS + } + stats = [ + {k: {s: cols[k, s][i] for s in STATS} for k in keys} + for i in range(meta.num_rows) + ] + file_no = ( + meta.column("data/chunk_index").to_numpy() * CHUNKS_SIZE + + meta.column("data/file_index").to_numpy() + ) + return stats, meta.column("length").to_numpy(), file_no + + +def write_parquet(fs, path: str, tables, compression: str) -> None: + """One row group per episode, as LeRobot's writer leaves it.""" + mkparent(fs, path) + w = None + for t in tables: + if w is None: + w = pq.ParquetWriter( + path, + t.schema, + filesystem=fs, + compression=compression, + use_dictionary=True, + ) + w.write_table(t) + w.close() + + +def make_lerobot(root: str, lay: Layout, workers: int, compression: str) -> None: + fs, base = arrow_fs(root) + probe = pa.BufferOutputStream() + pq.write_table( + lerobot_table(lay.episode(0), 0, lay.spans(0, 0)), + probe, + compression=compression, + ) + rows_per_file = PARQUET_MB * 2**20 * lay.lens[0] / probe.getvalue().size + file_no, f, in_file = [], 0, 0 + for n in lay.lens: + if in_file and in_file + n > rows_per_file: + f, in_file = f + 1, 0 + file_no.append(f) + in_file += n + file_no = np.asarray(file_no) + + def one(f): + eps = np.flatnonzero(file_no == f) + cols = [lay.episode(ep) for ep in eps] + write_parquet( + fs, + f"{base}/{data_path(f)}", + (lerobot_table(c, ep, lay.spans(ep, 0)) for c, ep in zip(cols, eps)), + compression, + ) + return [{k: vec_stats(c[k]) for k in VECTORS} for c in cols] + + stats = [s for per in pmap(one, range(file_no[-1] + 1), workers) for s in per] + write_meta(fs, base, stats, lay.lens, file_no) + features = {k: {"dtype": "float32", "shape": [d]} for k, d in VECTORS.items()} + features.update({k: {"dtype": "language"} for k in LANGUAGE_COLS}) + write_json( + fs, + f"{base}/meta/info.json", + { + "codebase_version": "v3.1", + "total_episodes": lay.n_eps, + "total_frames": int(lay.lens.sum()), + "chunks_size": CHUNKS_SIZE, + "data_files_size_in_mb": PARQUET_MB, + "n_files": int(file_no[-1] + 1), + "compression": compression, + "features": features, + "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", + }, + ) + + +def update_lerobot(root: str, lay: Layout, op: str, k: int, workers: int) -> dict: + fs, base = arrow_fs(root) + info = read_json(fs, f"{base}/meta/info.json") + key = {"replace": "head_pose", "append": feat_name(k)}.get(op) + + def one(f): + p = f"{base}/{data_path(f)}" + size_in = fs.get_file_info(p).size + t = pq.read_table(p, filesystem=fs) + eps = t.column("episode_index").to_numpy() + cuts = np.flatnonzero(np.diff(eps)) + 1 + bounds = list(zip(np.r_[0, cuts], np.r_[cuts, len(eps)])) + per_ep = [] + if op == "annotations": + parts = [ + language_columns(lay.spans(int(eps[a]), k), b - a) for a, b in bounds + ] + for i, name in enumerate(LANGUAGE_COLS): + col = pa.chunked_array([part[i] for part in parts]) + t = t.set_column(t.schema.get_field_index(name), name, col) + else: + new = lay.new_head_pose(t.column("index").to_numpy(), k) + if op == "replace": + t = t.set_column(t.schema.get_field_index(key), key, fsl(new)) + else: + t = t.append_column(key, fsl(new)) + per_ep = [(int(eps[a]), vec_stats(new[a:b])) for a, b in bounds] + write_parquet( + fs, p, (t.slice(a, b - a) for a, b in bounds), info["compression"] + ) + return size_in, fs.get_file_info(p).size, per_ep + + res = pmap(one, range(info["n_files"]), workers) + rec = { + "objects_written": len(res), + "bytes_read": sum(r[0] for r in res), + "bytes_written": sum(r[1] for r in res), + } + if key is None: + return rec + stats, lens, file_no = read_meta(fs, base) + for _, _, per_ep in res: + for ep, st in per_ep: + stats[ep][key] = st + rec["bytes_written"] += write_meta(fs, base, stats, lens, file_no) + rec["objects_written"] += 2 + if op == "append": + info["features"][key] = {"dtype": "float32", "shape": [7]} + write_json(fs, f"{base}/meta/info.json", info) + rec["objects_written"] += 1 + return rec + + +def verify_lerobot(root: str, lay: Layout, op: str, k: int) -> None: + fs, base = arrow_fs(root) + info = read_json(fs, f"{base}/meta/info.json") + for f in {0, info["n_files"] - 1}: + t = pq.read_table(f"{base}/{data_path(f)}", filesystem=fs) + if op == "annotations": + eps = t.column("episode_index").to_numpy() + for ep in np.unique(eps): + first = int(np.flatnonzero(eps == ep)[0]) + rows = t.column("language_persistent")[first].as_py() + assert [r["content"] for r in rows] == [ + s[0] for s in lay.spans(int(ep), k) + ] + continue + key = "head_pose" if op == "replace" else feat_name(k) + np.testing.assert_array_equal( + vec(t.column(key).combine_chunks(), 7), + lay.new_head_pose(t.column("index").to_numpy(), k), + ) + if op != "annotations": + stats, _, _ = read_meta(fs, base) + want = vec_stats(lay.new_head_pose(lay.g(0), k)) + for s in STATS: + np.testing.assert_allclose(stats[0][key][s], want[s]) + + +# --------------------------------------------------------------------------- lance +def steps_uri(root: str) -> str: + return f"{root.rstrip('/')}/steps.lance" + + +def annotations_uri(root: str) -> str: + return f"{root.rstrip('/')}/annotations.lance" + + +def rec_id(ep: int) -> str: + return f"rec-{ep:07d}" + + +def steps_batch(cols: dict, recording_id: str) -> pa.RecordBatch: + """One HUMAN_STEPS batch from flattened columns (VECTORS/SCALARS/BOOLS keys plus g).""" + from egomimic.rldb.lance.schema import HUMAN_STEPS + + def hand(s): + arrs = [fsl(cols[f"{s}.{f}"]) for f in HANDS] + arrs.append(pa.array(cols[f"{s}.confidence"])) + return pa.StructArray.from_arrays(arrs, fields=list(HUMAN_STEPS.field(s).type)) + + n = len(cols["g"]) + return pa.RecordBatch.from_arrays( + [ + pa.array([recording_id] * n), + pa.array(cols["g"] * DT_NS), + fsl(cols["head_pose"]), + hand("left"), + hand("right"), + pa.array(cols["left_valid"]), + pa.array(cols["right_valid"]), + pa.array(cols["head_pose_valid"]), + fsl(cols["eye_gaze"]), + ], + schema=HUMAN_STEPS, + ) + + +def lance_annotations(lay: Layout, eps, k: int) -> pa.Table: + from egomimic.rldb.lance.schema import ANNOTATIONS + + 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 + + opts = lance_opts(root) + batches = (steps_batch(lay.episode(ep), rec_id(ep)) for ep in range(lay.n_eps)) + lance.write_dataset( + pa.RecordBatchReader.from_batches(HUMAN_STEPS, batches), + steps_uri(root), + schema=HUMAN_STEPS, + max_rows_per_file=rows_per_frag, + storage_options=opts, + ) + lance.write_dataset( + lance_annotations(lay, range(lay.n_eps), 0), + annotations_uri(root), + storage_options=opts, + ) + + +def lance_add(uri: str, opts, name: str, read: str, fn, workers: int): + """New column `name` = fn(batch of `read`) via one data file per fragment and one Merge.""" + import lance + + ds = lance.dataset(uri, storage_options=opts) + + def udf(batch: pa.RecordBatch) -> pa.RecordBatch: + return pa.RecordBatch.from_arrays([fn(batch.column(read))], names=[name]) + + res = pmap( + lambda frag: frag.merge_columns(udf, columns=[read]), + ds.get_fragments(), + workers, + ) + ds = lance.LanceDataset.commit( + uri, + lance.LanceOperation.Merge([f for f, _ in res], res[0][1]), + read_version=ds.version, + storage_options=opts, + ) + return ds, len(res) + + +def lance_replace(uri: str, opts, col: str, read: str, fn, workers: int) -> int: + ds, n = lance_add(uri, opts, f"{col}_new", read, fn, workers) + ds.drop_columns([col]) + ds.alter_columns({"path": f"{col}_new", "name": col}) + return n + + +def update_lance(root: str, lay: Layout, op: str, k: int, workers: int) -> dict: + opts = lance_opts(root) + if op == "annotations": + import lance + + ds = lance.write_dataset( + lance_annotations(lay, range(lay.n_eps), k), + annotations_uri(root), + mode="overwrite", + storage_options=opts, + ) + 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)) + + if op == "replace": + n = lance_replace( + steps_uri(root), opts, "head_pose", "timestamp_ns", hp, workers + ) + else: + _, n = lance_add( + steps_uri(root), opts, feat_name(k), "timestamp_ns", hp, workers + ) + return {"objects_written": n} + + +def verify_lance(root: str, lay: Layout, op: str, k: int) -> None: + import lance + + from egomimic.rldb.lance.schema import HUMAN_STEPS + + opts = lance_opts(root) + if op == "annotations": + eps = sorted({0, lay.n_eps // 2, lay.n_eps - 1}) + 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) + schema = {f.name: f.type for f in ds.schema} + assert all(schema[f.name] == f.type for f in HUMAN_STEPS) + key = "head_pose" if op == "replace" else feat_name(k) + n = ds.count_rows() + t = ds.take(sorted({0, n // 2, n - 1}), columns=["timestamp_ns", key]) + np.testing.assert_array_equal( + vec(t.column(key).combine_chunks(), 7), + lay.new_head_pose(t.column("timestamp_ns").to_numpy() // DT_NS, k), + ) + + +# --------------------------------------------------------------------------- cli +def fmt_root(root: str, fmt: str) -> str: + return f"{root.rstrip('/')}/{fmt}" + + +def make(fmt: str, root: str, lay: Layout, workers: int, rows_per_frag: int) -> None: + if fmt == "zarr": + make_zarr(root, lay, workers) + elif fmt == "lance": + make_lance(root, lay, rows_per_frag) + else: + make_lerobot(root, lay, workers, COMPRESSION[fmt]) + + +def update(fmt: str, root: str, lay: Layout, op: str, k: int, workers: int) -> dict: + fn = {"zarr": update_zarr, "lance": update_lance}.get(fmt, update_lerobot) + return fn(root, lay, op, k, workers) + + +def verify(fmt: str, root: str, lay: Layout, op: str, k: int) -> None: + fn = {"zarr": verify_zarr, "lance": verify_lance}.get(fmt, verify_lerobot) + fn(root, lay, op, k) + + +def cmd_make(a) -> None: + lay = pilot_layout(a.pilot, a.rows) + print(f"{lay.n_eps} episodes, median {int(np.median(lay.lens))} frames", flush=True) + for fmt in a.formats: + t = time.time() + make(fmt, fmt_root(a.root, fmt), lay, a.workers, a.rows_per_frag) + print(f"make {fmt}: {time.time() - t:.0f}s", flush=True) + + +def cmd_update(a) -> None: + lay = pilot_layout(a.pilot, a.rows) + r = fmt_root(a.root, a.fmt) + t = time.perf_counter() + rec = update(a.fmt, r, lay, a.op, a.k, a.workers) + sec = time.perf_counter() - t + verify(a.fmt, r, lay, a.op, a.k) + emit( + { + "bench": "update", + "op": a.op, + "format": a.fmt, + "rows": int(lay.lens.sum()), + "episodes": lay.n_eps, + "workers": a.workers, + "latency_ms": a.latency_ms, + "seconds": sec, + "rows_per_s": lay.lens.sum() / sec, + **rec, + }, + a.out, + ) + + +def cmd_fit(a) -> None: + """Rate-based projection from the largest run per (op, format, latency, workers).""" + runs = {} + for line in Path(a.results).read_text().splitlines(): + r = json.loads(line) + if r.get("bench") == "update": + key = (r.get("op", "replace"), r["format"], r["latency_ms"], r["workers"]) + runs.setdefault(key, []).append(r) + print( + f"| op | format | latency ms | workers | rows/s by size | hours at {a.target:.0e} rows |" + ) + print("|---|---|---|---|---|---|") + for key in sorted(runs): + rs = sorted(runs[key], key=lambda r: r["rows"]) + rates = ", ".join(f"{r['rows'] / 1e6:g}M: {r['rows_per_s']:,.0f}" for r in rs) + hours = a.target / rs[-1]["rows_per_s"] / 3600 + print(f"| {' | '.join(map(str, key))} | {rates} | {hours:,.2f} |") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + sub = ap.add_subparsers(dest="cmd", required=True) + for name in ("make", "update"): + p = sub.add_parser(name) + p.add_argument("--root", required=True) + p.add_argument("--rows", type=int, required=True) + p.add_argument( + "--pilot", type=Path, default=Path("/coc/flash7/scratch/lance_pilot") + ) + p.add_argument("--workers", type=int, default=16) + m = sub.choices["make"] + m.add_argument("--rows-per-frag", type=int, default=LANCE_ROWS_PER_FILE) + m.add_argument("--formats", nargs="+", default=list(FORMATS)) + u = sub.choices["update"] + u.add_argument("--fmt", required=True, choices=FORMATS) + u.add_argument("--op", required=True, choices=OPS) + u.add_argument( + "--k", type=int, required=True, help="edit generation; values encode it" + ) + u.add_argument("--latency-ms", type=float, default=0.0, help="label only") + u.add_argument("--out") + f = sub.add_parser("fit") + f.add_argument("results") + f.add_argument("--target", type=float, default=2e9) + a = ap.parse_args() + {"make": cmd_make, "update": cmd_update, "fit": cmd_fit}[a.cmd](a) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/common.py b/egomimic/scripts/format_bench/common.py new file mode 100644 index 0000000000..7b2c378f1e --- /dev/null +++ b/egomimic/scripts/format_bench/common.py @@ -0,0 +1,192 @@ +"""Loader timing shared by the zarr/Lance and LeRobot benches. Imports nothing from egomimic, +because the LeRobot bench runs in its own venv.""" + +import functools +import json +import multiprocessing as mp +import os +import socket +import subprocess +import threading +import time +from pathlib import Path + +import numpy as np +import psutil +import torch +from torch.utils.data import DataLoader, RandomSampler, SubsetRandomSampler + + +def git_sha() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + cwd=Path(__file__).parent, + text=True, + ).strip() + except Exception: + return "" + + +def emit(rec: dict, out: str | None): + rec = { + "node": socket.gethostname().split(".")[0], + "git": git_sha(), + "t": time.time(), + **rec, + } + line = json.dumps(rec) + print(line, flush=True) + if out: + with open(out, "a") as f: + f.write(line + "\n") + + +def evict(*roots) -> int: + """Drop our files from the page cache (no root needed). Returns bytes evicted.""" + total = 0 + for root in roots: + root = Path(root) + files = ( + [root] if root.is_file() else (p for p in root.rglob("*") if p.is_file()) + ) + for p in files: + try: + fd = os.open(p, os.O_RDONLY) + except OSError: + continue + try: + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + total += os.fstat(fd).st_size + finally: + os.close(fd) + return total + + +def tree_rss_mb() -> float: + proc = psutil.Process(os.getpid()) + total = proc.memory_info().rss + for c in proc.children(recursive=True): + try: + total += c.memory_info().rss + except psutil.NoSuchProcess: + pass + return total / 2**20 + + +class PeakRSS: + """Samples the process tree's RSS every `period` seconds on a thread.""" + + def __init__(self, period: float = 0.5): + self.peak, self.period, self._stop = 0.0, period, threading.Event() + self._t = threading.Thread(target=self._run, daemon=True) + + def _run(self): + while not self._stop.is_set(): + self.peak = max(self.peak, tree_rss_mb()) + self._stop.wait(self.period) + + def __enter__(self): + self._t.start() + return self + + def __exit__(self, *exc): + self._stop.set() + self._t.join() + self.peak = max(self.peak, tree_rss_mb()) + + +def collate(batch): + out = {} + for k in batch[0]: + vals = [b[k] for b in batch] + out[k] = torch.stack(vals) if isinstance(vals[0], torch.Tensor) else vals + return out + + +def _worker_init(worker_id, ready, counters, hook): + """Marks the worker ready (dataset unpickled) and installs the decoder-cache counter.""" + ready[worker_id] = time.time() + if hook: + hook(counters, worker_id) + + +def run_loader( + ds, + *, + workers: int, + context: str | None, + batch_size: int, + batches: int, + warmup: int, + hook=None, + seed: int = 0, + indices: list[int] | None = None, + collate_fn=collate, +) -> dict: + """Phase timings and steady-state throughput of a shuffled DataLoader over `ds`. `hook` is a + module-level function from `hooks.py`, called in each worker to count decoder-cache hits.""" + ctx = mp.get_context(context) if workers else mp.get_context() + ready = ctx.Array("d", max(workers, 1), lock=False) + # [hits, misses] per worker + counters = ctx.Array("q", 2 * max(workers, 1), lock=False) + g = torch.Generator().manual_seed(seed) + if isinstance(ds, torch.utils.data.IterableDataset): + # Streaming datasets shuffle themselves (shuffle buffer); DataLoader takes no sampler. + sampler = None + elif indices is None: + sampler = RandomSampler(ds, generator=g) + else: + sampler = SubsetRandomSampler(indices, generator=g) + loader = DataLoader( + ds, + batch_size=batch_size, + sampler=sampler, + num_workers=workers, + collate_fn=collate_fn, + prefetch_factor=2 if workers else None, + persistent_workers=workers > 0, + multiprocessing_context=ctx if workers else None, + worker_init_fn=functools.partial( + _worker_init, ready=ready, counters=counters, hook=hook + ) + if workers + else None, + ) + if not workers and hook: + hook(counters, 0) + with PeakRSS() as rss: + t0 = time.time() + it = iter(loader) + t_iter = time.time() + next(it) + t_first = time.time() + for _ in range(warmup): + next(it) + waits = [] + t1 = time.perf_counter() + for _ in range(batches): + tb = time.perf_counter() + next(it) + waits.append(time.perf_counter() - tb) + dt = max(time.perf_counter() - t1, 1e-9) + ready_t = [r for r in ready[:workers] if r > 0] + c = np.asarray(counters[:]).reshape(-1, 2) + hits, misses = int(c[:, 0].sum()), int(c[:, 1].sum()) + del it, loader + return dict( + iter_s=round(t_iter - t0, 3), + worker_spawn_s=round(max(ready_t) - t0, 3) if ready_t else 0.0, + first_batch_s=round(t_first - t0, 3), + samples_per_s=round(batches * batch_size / dt, 1) if batches else None, + batch_wait_p50_ms=round(float(np.percentile(waits, 50)) * 1e3, 1) + if waits + else None, + batch_wait_p95_ms=round(float(np.percentile(waits, 95)) * 1e3, 1) + if waits + else None, + peak_rss_mb=round(rss.peak), + decoder_hits=hits, + decoder_misses=misses, + decoder_hit_rate=round(hits / (hits + misses), 3) if hits + misses else None, + ) diff --git a/egomimic/scripts/format_bench/hooks.py b/egomimic/scripts/format_bench/hooks.py new file mode 100644 index 0000000000..0bda32fe9b --- /dev/null +++ b/egomimic/scripts/format_bench/hooks.py @@ -0,0 +1,46 @@ +"""Decoder-cache hit counters, installed inside each DataLoader worker. Module-level functions so +they pickle by reference under forkserver and spawn. FB_DECODER_CACHE=N shrinks Lance's per-worker +LRU (LeRobot reads LEROBOT_VIDEO_DECODER_CACHE_SIZE itself).""" + +import os + + +def lance_decoders(counters, wid): + from egomimic.rldb.lance import dataset as m + + if os.environ.get("FB_DECODER_CACHE"): + m._MAX_DECODERS = int(os.environ["FB_DECODER_CACHE"]) + + orig = m._decoder + + def counted(path): + counters[2 * wid + (0 if path in m._DECODERS else 1)] += 1 + return orig(path) + + m._decoder = counted + + +def lerobot_stream_decoders(counters, wid): + # StreamingLeRobotDataset builds its own VideoDecoderCache per iterator, so count at the class. + from lerobot.datasets.video_utils import VideoDecoderCache + + orig = VideoDecoderCache.get_decoder + + def counted(self, path): + counters[2 * wid + (0 if path in self else 1)] += 1 + return orig(self, path) + + VideoDecoderCache.get_decoder = counted + + +def lerobot_decoders(counters, wid): + from lerobot.datasets import video_utils as m + + cache = m._default_decoder_cache + orig = cache.get_decoder + + def counted(path): + counters[2 * wid + (0 if path in cache else 1)] += 1 + return orig(path) + + cache.get_decoder = counted diff --git a/egomimic/scripts/format_bench/make_fakes.py b/egomimic/scripts/format_bench/make_fakes.py new file mode 100644 index 0000000000..00a1c82dab --- /dev/null +++ b/egomimic/scripts/format_bench/make_fakes.py @@ -0,0 +1,307 @@ +"""Startup-scale datasets: the 48 real pilot episodes plus N fake bimanual episodes, per format. + +Each fake is `--frames` real rows sliced from a bimanual pilot episode and points at that +episode's real video, so every loader runs unmodified: + zarr a symlink to the source store (startup reads only attrs, so length does not matter) + lance recordings/aliases/tasks rows + contiguous steps rows; videos/ is a symlink + lerobot data parquet rows + episodes rows whose video from/to timestamps fall inside the + source episode's span of a real video file, rolled at LeRobot's file-size limits; + videos/ is a symlink +Also writes `real_.txt` (what the first batch may sample) and `lerobot_episodes.txt`. +""" + +import argparse +import json +import os +import shutil +import time +from pathlib import Path + +import lance +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +BIMANUAL = "human_bimanual" + + +def fake_name(i: int) -> str: + return f"fake-{i:07d}" + + +def sources( + n_fakes: int, src_lengths: list[int], frames: int, seed: int = 0 +) -> list[tuple[int, int]]: + """(source index, start frame) per fake, deterministic.""" + rng = np.random.default_rng(seed) + out = [] + for i in range(n_fakes): + s = i % len(src_lengths) + out.append((s, int(rng.integers(0, src_lengths[s] - frames)))) + return out + + +# --------------------------------------------------------------------------- zarr +def make_zarr(src: Path, out: Path, n: int, manifest: list[dict]) -> None: + out.mkdir(parents=True) + bi = [m["episode_hash"] for m in manifest if m["embodiment"] == BIMANUAL] + for m in manifest: + os.symlink(src / m["episode_hash"], out / m["episode_hash"]) + for i in range(n): + os.symlink(src / bi[i % len(bi)], out / fake_name(i)) + (out.parent / "real_zarr.txt").write_text(" ".join(bi) + "\n") + + +# --------------------------------------------------------------------------- lance +def make_lance(src: Path, out: Path, n: int, frames: int) -> None: + out.mkdir(parents=True) + os.symlink(src / "videos", out / "videos") + rec = lance.dataset(str(src / "recordings.lance")).to_table() + ali = lance.dataset(str(src / "aliases.lance")).to_table() + tasks = lance.dataset(str(src / "tasks.lance")).to_table() + steps = lance.dataset(str(src / "steps/aria.lance")).to_table() + + ids = steps.column("recording_id").to_numpy(zero_copy_only=False) + first = {} + for i, r in enumerate(ids): + first.setdefault(r, i) + recs = rec.to_pylist() + bi = sorted( + (r for r in recs if r["embodiment"] == BIMANUAL), + key=lambda r: first[r["recording_id"]], + ) + plan = sources(n, [r["num_frames"] for r in bi], frames) + fake_ids = [fake_name(i) for i in range(n)] + + def rows_like(tbl, key, src_ids, new_ids, overrides=None): + by_id = {r[key]: r for r in tbl.to_pylist()} + rows = [] + for s, f in zip(src_ids, new_ids): + r = dict(by_id[s]) + r[key] = f + r.update((overrides or {}).get(f, {})) + rows.append(r) + return pa.concat_tables([tbl, pa.Table.from_pylist(rows, schema=tbl.schema)]) + + src_of = [bi[s]["recording_id"] for s, _ in plan] + lance.write_dataset( + rows_like( + rec, + "recording_id", + src_of, + fake_ids, + {f: {"num_frames": frames} for f in fake_ids}, + ), + str(out / "recordings.lance"), + ) + a = rows_like(ali, "recording_id", src_of, fake_ids) + eh = a.column("episode_hash").to_pylist() + eh[len(ali) :] = fake_ids + a = a.set_column( + a.schema.get_field_index("episode_hash"), "episode_hash", pa.array(eh) + ) + lance.write_dataset(a, str(out / "aliases.lance")) + lance.write_dataset( + rows_like(tasks, "recording_id", src_of, fake_ids), str(out / "tasks.lance") + ) + + rid_col = steps.schema.get_field_index("recording_id") + + def step_batches(): + yield from steps.to_batches(max_chunksize=1 << 20) + per = max(1, (1 << 20) // frames) + for b0 in range(0, n, per): + idx = np.concatenate( + [ + np.arange( + first[bi[s]["recording_id"]] + f0, + first[bi[s]["recording_id"]] + f0 + frames, + ) + for s, f0 in plan[b0 : b0 + per] + ] + ) + t = steps.take(pa.array(idx)) + new = pa.array( + np.repeat(np.asarray(fake_ids[b0 : b0 + per], dtype=object), frames), + pa.string(), + ) + yield from t.set_column(rid_col, "recording_id", new).to_batches() + + lance.write_dataset( + pa.RecordBatchReader.from_batches(steps.schema, step_batches()), + str(out / "steps/aria.lance"), + max_rows_per_file=1 << 23, + ) + + a_all = ali.to_pylist() + real_bi = {r["recording_id"] for r in bi} + (out.parent / "real_lance.txt").write_text( + " ".join(r["episode_hash"] for r in a_all if r["recording_id"] in real_bi) + + "\n" + ) + + +# --------------------------------------------------------------------------- lerobot +def make_lerobot( + src: Path, out: Path, n: int, frames: int, bimanual_eps: list[int] +) -> None: + info = json.loads((src / "meta/info.json").read_text()) + fps = info["fps"] + (out / "meta").mkdir(parents=True) + os.symlink(src / "videos", out / "videos") + for f in ("tasks.parquet", "stats.json"): + shutil.copy(src / "meta" / f, out / "meta" / f) + data = pq.read_table(sorted((src / "data").glob("*/*.parquet"))) + eps = pq.read_table(sorted((src / "meta/episodes").glob("*/*.parquet"))) + n_real_eps, n_real_frames = eps.num_rows, data.num_rows + lengths = eps.column("length").to_pylist() + starts = eps.column("dataset_from_index").to_pylist() + plan = sources(n, [lengths[e] for e in bimanual_eps], frames) + data_mb, chunks_size = info["data_files_size_in_mb"], info["chunks_size"] + + # Data: real rows first, then fakes, rolled like DatasetWriter._save_episode_data. + chunk, file, writer, file_start, nxt = 0, 0, None, 0, 0 + loc = [] # (chunk, file) per episode, real then fake + + def path(c, f): + return out / info["data_path"].format(chunk_index=c, file_index=f) + + def write_episode(tbl): + nonlocal chunk, file, writer, file_start, nxt + if writer is not None: + size = os.path.getsize(path(chunk, file)) / 2**20 + in_file = nxt - file_start + if size + (size / in_file if in_file else 0) * tbl.num_rows >= data_mb: + writer.close() + writer = None + file += 1 + if file >= chunks_size: + chunk, file = chunk + 1, 0 + file_start = nxt + if writer is None: + path(chunk, file).parent.mkdir(parents=True, exist_ok=True) + writer = pq.ParquetWriter( + path(chunk, file), tbl.schema, compression="snappy", use_dictionary=True + ) + writer.write_table(tbl) + loc.append((chunk, file)) + nxt += tbl.num_rows + + for e in range(n_real_eps): + write_episode(data.slice(starts[e], lengths[e])) + ei, ii, fi, ti = ( + data.schema.get_field_index(c) + for c in ("episode_index", "index", "frame_index", "timestamp") + ) + for i, (s, f0) in enumerate(plan): + e = bimanual_eps[s] + t = data.slice(starts[e] + f0, frames) + t = t.set_column( + ei, "episode_index", pa.array(np.full(frames, n_real_eps + i, np.int64)) + ) + t = t.set_column( + ii, "index", pa.array(np.arange(nxt, nxt + frames, dtype=np.int64)) + ) + t = t.set_column(fi, "frame_index", pa.array(np.arange(frames, dtype=np.int64))) + t = t.set_column( + ti, + "timestamp", + pa.array(np.arange(frames, dtype=np.float32) / np.float32(fps)), + ) + write_episode(t) + writer.close() + + # Episodes metadata: copy each source row (stats included), rolled at the same limit. + vk = [k for k, v in info["features"].items() if v["dtype"] == "video"] + src_idx = np.asarray([bimanual_eps[s] for s, _ in plan], dtype=np.int64) + f0s = np.asarray([f0 for _, f0 in plan], dtype=np.float64) + fake = eps.take(pa.array(src_idx)) + + def put(tbl, name, values): + i = tbl.schema.get_field_index(name) + return tbl.set_column( + i, tbl.schema.field(i), pa.array(values, tbl.schema.field(i).type) + ) + + fake = put(fake, "episode_index", np.arange(n_real_eps, n_real_eps + n)) + fake = put(fake, "length", np.full(n, frames)) + fake = put(fake, "dataset_from_index", n_real_frames + np.arange(n) * frames) + fake = put(fake, "dataset_to_index", n_real_frames + (np.arange(n) + 1) * frames) + for key in vk: + start_ts = fake.column(f"videos/{key}/from_timestamp").to_numpy() + f0s / fps + fake = put(fake, f"videos/{key}/from_timestamp", start_ts) + fake = put(fake, f"videos/{key}/to_timestamp", start_ts + frames / fps) + allt = pa.concat_tables([eps, fake]) + allt = put(allt, "data/chunk_index", [c for c, _ in loc]) + allt = put(allt, "data/file_index", [f for _, f in loc]) + # Same rule as LeRobotDatasetMetadata._save_episode_metadata: roll at data_files_size_in_mb. + sink = pa.BufferOutputStream() + pq.write_table(allt.slice(0, 1000), sink, compression="snappy") + per_row_mb = sink.getvalue().size / 2**20 / min(1000, allt.num_rows) + rows_per_file = max(1, int(data_mb / per_row_mb)) + file_no = np.arange(allt.num_rows) // rows_per_file + allt = put(allt, "meta/episodes/chunk_index", file_no // chunks_size) + allt = put(allt, "meta/episodes/file_index", file_no % chunks_size) + for j in range(0, allt.num_rows, rows_per_file): + c, f = divmod(j // rows_per_file, chunks_size) + p = out / f"meta/episodes/chunk-{c:03d}/file-{f:03d}.parquet" + p.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(allt.slice(j, rows_per_file), p, compression="snappy") + + info["total_episodes"] = allt.num_rows + info["total_frames"] = n_real_frames + n * frames + info["splits"] = {"train": f"0:{allt.num_rows}"} + (out / "meta/info.json").write_text(json.dumps(info, indent=4)) + eps_list = bimanual_eps + list(range(n_real_eps, n_real_eps + n)) + (out.parent / "lerobot_episodes.txt").write_text( + " ".join(map(str, eps_list)) + "\n" + ) + (out.parent / "real_lerobot.txt").write_text( + " ".join(map(str, bimanual_eps)) + "\n" + ) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--pilot", type=Path, default=Path("/coc/flash7/scratch/lance_pilot") + ) + ap.add_argument("--lerobot", type=Path, required=True) + ap.add_argument("--lerobot-bimanual", type=Path, required=True) + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--n", type=int, required=True) + ap.add_argument("--frames", type=int, default=1000) + ap.add_argument("--formats", nargs="+", default=["zarr", "lance", "lerobot"]) + args = ap.parse_args() + + import csv + + manifest = list(csv.DictReader(open(args.pilot / "manifest.csv"))) + out = args.out + out.mkdir(parents=True, exist_ok=True) + for fmt in args.formats: + t = time.time() + if (out / fmt).exists(): + shutil.rmtree(out / fmt) + if fmt == "zarr": + make_zarr(args.pilot / "zarr", out / "zarr", args.n, manifest) + elif fmt == "lance": + make_lance(args.pilot / "lance", out / "lance", args.n, args.frames) + else: + bi = [int(x) for x in args.lerobot_bimanual.read_text().split()] + make_lerobot(args.lerobot, out / "lerobot", args.n, args.frames, bi) + print( + json.dumps( + dict( + format=fmt, + n=args.n, + frames=args.frames, + build_s=round(time.time() - t, 1), + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/s3patch.py b/egomimic/scripts/format_bench/s3patch.py new file mode 100644 index 0000000000..90e23c5a4a --- /dev/null +++ b/egomimic/scripts/format_bench/s3patch.py @@ -0,0 +1,83 @@ +"""Points the zarr and Lance loaders at a plain-http S3 endpoint (local MinIO), for the bench only. + +Active when FB_S3_ENDPOINT is set. Imported in the parent and preloaded in the forkserver, so +workers inherit the patched modules. + +- Lance: `storage_options` gains `allow_http` (plain-http endpoints), reading the same R2_* vars. +- zarr: `ZarrEpisode` opens s3:// paths through obstore, and reopens lazily in each worker rather + than pickling a live client (as logs/claude_scratch/zarr_r2_patch.py does for R2). +""" + +import os +from pathlib import Path + +ENDPOINT = os.environ.get("FB_S3_ENDPOINT") + + +def obstore_s3(uri: str): + from obstore.store import S3Store + + bucket, prefix = uri[len("s3://") :].split("/", 1) + return S3Store( + bucket, + prefix=prefix.rstrip("/"), + endpoint=ENDPOINT, + access_key_id=os.environ["R2_ACCESS_KEY_ID"], + secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"], + region="us-east-1", + client_options={"allow_http": True}, + ) + + +def _patch_lance(): + from egomimic.rldb.lance import dataset as m + + orig = m.storage_options + + def storage_options(uri): + opts = orig(uri) + if opts is not None: + opts = {**opts, "allow_http": "true", "region": "us-east-1"} + return opts + + m.storage_options = storage_options + + +def _patch_zarr(): + import zarr + + from egomimic.rldb.zarr import zarr_dataset_multi as zdm + + def _init(self, path): + self._path = path if str(path).startswith("s3://") else Path(path) + if str(path).startswith("s3://"): + store = zarr.storage.ObjectStore(obstore_s3(str(path)), read_only=True) + self._store = zarr.open_group(store=store, mode="r") + else: + self._store = zarr.open_group(str(path), mode="r") + self.metadata = dict(self._store.attrs) + self.keys = self.metadata["features"] + + def _lazy(self, name): + # Workers reopen a store on first use, not at unpickle: an eager reopen costs every worker + # one GET per episode (12 x 10k at startup) for stores the first batch never touches. + if name in ("_store", "metadata", "keys"): + _init(self, object.__getattribute__(self, "_path")) + return object.__getattribute__(self, name) + raise AttributeError(name) + + zdm.ZarrEpisode.__init__ = _init + zdm.ZarrEpisode.__getattr__ = _lazy + zdm.ZarrEpisode.__getstate__ = lambda self: {"_path": self._path} + zdm.ZarrEpisode.__setstate__ = lambda self, s: setattr(self, "_path", s["_path"]) + + +if ENDPOINT: + _patch_lance() + _patch_zarr() + +if os.environ.get("FB_NO_GETITEMS"): + # Ablation: per-sample reads for Lance too, to separate the batched fetch from the format. + from egomimic.rldb.zarr import zarr_dataset_multi as _zdm + + del _zdm.MultiDataset.__getitems__ diff --git a/egomimic/scripts/format_bench/storage.py b/egomimic/scripts/format_bench/storage.py new file mode 100644 index 0000000000..89767b3316 --- /dev/null +++ b/egomimic/scripts/format_bench/storage.py @@ -0,0 +1,135 @@ +"""Bytes on disk per format for the 48 pilot recordings, split into numeric / video-image / +metadata-index, total and per frame. One JSON line per format variant. + +zarr counts its trimmed episodes (manifest num_frames); Lance and LeRobot hold the full +recordings, so per-frame numbers are the comparable ones. Lance video is reported both as the +480p MP4s (local loader) and as the 480p 8-frame chunk table (object-storage loader). +""" + +import argparse +import csv +from pathlib import Path + +import lance + +from egomimic.scripts.format_bench.common import emit + + +def size(p: Path) -> int: + if p.is_file(): + return p.stat().st_size + return sum( + f.stat().st_size for f in p.rglob("*") if f.is_file() and not f.is_symlink() + ) + + +def lance_split(table: Path) -> tuple[int, int]: + """(data file bytes, manifest/version/transaction bytes) of one Lance table.""" + data = size(table / "data") if (table / "data").exists() else 0 + return data, size(table) - data + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--pilot", type=Path, default=Path("/coc/flash7/scratch/lance_pilot") + ) + ap.add_argument("--chunks", type=Path, required=True) + ap.add_argument( + "--lerobot", + type=Path, + nargs="+", + required=True, + help="one or more LeRobot roots", + ) + ap.add_argument("--camera", default="front_1_480") + ap.add_argument("--out") + args = ap.parse_args() + + manifest = list(csv.DictReader(open(args.pilot / "manifest.csv"))) + zarr_frames = sum(int(m["num_frames"]) for m in manifest) + num = img = meta = 0 + for m in manifest: + store = args.pilot / "zarr" / m["episode_hash"] + for child in store.iterdir(): + if child.name == "zarr.json": + meta += size(child) + elif child.name == "images.front_1": + img += size(child) + else: + num += size(child) + rows = [ + dict( + format="zarr", + variant="jpeg", + frames=zarr_frames, + numeric=num, + video=img, + metadata=meta, + ) + ] + + root = args.pilot / "lance" + rec = lance.dataset(str(root / "recordings.lance")).to_table().to_pylist() + lance_frames = sum(r["num_frames"] for r in rec) + s_data, s_meta = lance_split(root / "steps/aria.lance") + meta = s_meta + sum( + size(root / t) for t in ("recordings.lance", "aliases.lance", "tasks.lance") + ) + mp4 = sum(size(root / dict(r["video_path"])[args.camera]) for r in rec) + c_data, c_meta = lance_split(args.chunks) + rows.append( + dict( + format="lance", + variant="mp4", + frames=lance_frames, + numeric=s_data, + video=mp4, + metadata=meta, + ) + ) + rows.append( + dict( + format="lance", + variant="k8_chunks", + frames=lance_frames, + numeric=s_data, + video=c_data, + metadata=meta + c_meta, + ) + ) + + import json + + for lr in args.lerobot: + info = json.loads((lr / "meta/info.json").read_text()) + rows.append( + dict( + format="lerobot", + variant="packed" if "action" in info["features"] else "unpacked", + frames=info["total_frames"], + numeric=size(lr / "data"), + video=size(lr / "videos"), + metadata=size(lr / "meta"), + ) + ) + for r in rows: + r["total"] = r["numeric"] + r["video"] + r["metadata"] + for k in ("numeric", "video", "metadata", "total"): + r[f"{k}_mb"] = round(r[k] / 1e6, 1) + r[f"{k}_bytes_per_frame"] = round(r[k] / r["frames"], 1) + emit( + dict( + label="storage", + **{ + k: v + for k, v in r.items() + if k not in ("numeric", "video", "metadata", "total") + }, + ), + args.out, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/summarize.py b/egomimic/scripts/format_bench/summarize.py new file mode 100644 index 0000000000..d82a7290da --- /dev/null +++ b/egomimic/scripts/format_bench/summarize.py @@ -0,0 +1,354 @@ +"""results/*.jsonl -> markdown tables (medians over repeats).""" + +import argparse +import json +from collections import defaultdict +from pathlib import Path +from statistics import median + +FORMATS = [ + "zarr", + "zarr_fast", + "lance", + "lerobot", + "lerobot_stream", + "lerobot_unpacked", +] +PATTERNS = ["chunk1", "numeric", "frames8"] + + +def load(path: Path) -> list[dict]: + if not path.exists(): + return [] + rows = [] + for line in path.read_text().splitlines(): + try: + r = json.loads(line) + except json.JSONDecodeError: + continue + # Runs before the packed layout wrote format=lerobot for the one-column-per-key layout. + # The fake-length runs (rco3_fblen) are packed under a plain lerobot/ dir. + if ( + r.get("format") == "lerobot" + and "root" in r + and not any(s in r["root"] for s in ("lerobot_packed", "rco3_fblen")) + ): + r["format"] = "lerobot_unpacked" + rows.append(r) + return rows + + +def fmt(x, nd=1): + if x is None: + return "–" + if isinstance(x, float): + return f"{x:,.{nd}f}" + return f"{x:,}" + + +def table(header: list[str], rows: list[list]) -> str: + out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)] + out += ["| " + " | ".join(str(c) for c in r) + " |" for r in rows] + return "\n".join(out) + + +def throughput(rows: list[dict], title: str) -> str: + rows = [r for r in rows if r.get("label") == "tput"] + if not rows: + return f"## {title}\n\n(no results)\n" + by = defaultdict(list) + for r in rows: + by[(r["format"], r["pattern"], r["workers"])].append(r) + workers = sorted({r["workers"] for r in rows}) + best_rows, grid_rows = [], [] + for p in PATTERNS: + for f in FORMATS: + cells, best = [], None + for w in workers: + rs = by.get((f, p, w), []) + if not rs: + cells.append("–") + continue + sps = median(r["samples_per_s"] for r in rs) + cells.append(f"{sps:,.0f} (n={len(rs)})") + if best is None or sps > best[0]: + best = (sps, w, rs) + grid_rows.append([p, f, *cells]) + if best: + sps, w, rs = best + hit = [ + r["decoder_hit_rate"] + for r in rs + if r.get("decoder_hit_rate") is not None + ] + best_rows.append( + [ + p, + f, + f"{sps:,.0f}", + w, + fmt(median(r["batch_wait_p95_ms"] for r in rs), 0), + fmt(median(r["peak_rss_mb"] for r in rs) / 1024, 1), + fmt(median(hit), 2) if hit else "–", + ] + ) + return ( + f"## {title}\n\nBest worker count per format (median of repeats):\n\n" + + table( + [ + "pattern", + "format", + "samples/s", + "workers", + "p95 batch wait ms", + "peak RSS GB", + "decoder hit rate", + ], + best_rows, + ) + + "\n\nAll worker counts (median samples/s):\n\n" + + table(["pattern", "format", *[f"w={w}" for w in workers]], grid_rows) + + "\n" + ) + + +def nocache(rows: list[dict]) -> str: + rows = [r for r in rows if r.get("label") == "nocache"] + if not rows: + return "" + by = defaultdict(list) + for r in rows: + by[r["format"]].append(r) + out = [] + for f in FORMATS: + rs = by.get(f) + if rs: + out.append( + [ + f, + f"{median(r['samples_per_s'] for r in rs):,.0f}", + len(rs), + fmt(median(r["peak_rss_mb"] for r in rs) / 1024, 1), + ] + ) + return ( + "## Local, decoder cache size 1 (chunk1, 12 workers)\n\n" + "Approximates a 100k-episode run, where a worker almost never reuses an open video decoder.\n\n" + + table(["format", "samples/s", "n", "peak RSS GB"], out) + + "\n" + ) + + +def startup(rows: list[dict], title: str) -> str: + if not rows: + return f"## {title}\n\n(no results)\n" + by = defaultdict(list) + for r in rows: + variant = r["format"] + ( + "" + if not r["format"].startswith("lerobot") + else (" cold" if r.get("cold") else " warm") + ) + by[(r["label"], r["location"], variant, r["episodes"])].append(r) + out = [] + for (label, loc, variant, _), rs in sorted(by.items()): + m = { + k: median(r[k] for r in rs) + for k in ( + "resolve_s", + "construct_s", + "worker_spawn_s", + "startup_total_s", + "peak_rss_mb", + ) + } + m["first_only_s"] = median(r["first_batch_s"] - r["worker_spawn_s"] for r in rs) + out.append( + [ + label.replace("startup_", "").replace("_nfs", ""), + loc, + variant, + fmt(rs[0]["episodes"]), + fmt(m["resolve_s"], 2), + fmt(m["construct_s"], 2), + fmt(m["worker_spawn_s"], 2), + fmt(m["first_only_s"], 2), + fmt(m["startup_total_s"], 2), + fmt(m["peak_rss_mb"] / 1024, 1), + len(rs), + ] + ) + return ( + f"## {title}\n\n" + "resolve = finding episodes and building per-episode state; construct = MultiDataset / HF Arrow load; " + "spawn = DataLoader start until the last worker is ready; first batch = after workers are ready; " + "total = resolve + construct + time to first batch (each column is its own median). " + "Peak RSS sums RSS over the process tree, so pages shared between workers count once per worker.\n\n" + + table( + [ + "N fake", + "location", + "format", + "episodes", + "resolve s", + "construct s", + "spawn s", + "first batch s", + "total s", + "peak RSS GB", + "n", + ], + out, + ) + + "\n" + ) + + +def by_label(rows: list[dict], title: str, note: str = "") -> str: + """Median samples/s per (label, format, pattern, workers) for ablation-style runs.""" + if not rows: + return "" + by = defaultdict(list) + for r in rows: + by[(r["label"], r["format"], r["pattern"], r["workers"])].append(r) + out = [] + for (label, f, p, w), rs in sorted( + by.items(), key=lambda kv: (PATTERNS.index(kv[0][2]), kv[0]) + ): + hit = [ + r["decoder_hit_rate"] for r in rs if r.get("decoder_hit_rate") is not None + ] + out.append( + [ + p, + label, + f, + w, + f"{median(r['samples_per_s'] for r in rs):,.0f}", + fmt(median(hit), 2) if hit else "–", + len(rs), + ] + ) + return ( + f"## {title}\n\n{note}" + + table( + [ + "pattern", + "variant", + "format", + "workers", + "samples/s", + "decoder hit rate", + "n", + ], + out, + ) + + "\n" + ) + + +def storage(rows: list[dict]) -> str: + if not rows: + return "## Storage\n\n(no results)\n" + out = [ + [ + r["format"], + r["variant"], + fmt(r["frames"]), + fmt(r["numeric_mb"]), + fmt(r["video_mb"]), + fmt(r["metadata_mb"]), + fmt(r["total_mb"]), + fmt(r["numeric_bytes_per_frame"], 0), + fmt(r["video_bytes_per_frame"], 0), + fmt(r["total_bytes_per_frame"], 0), + ] + for r in rows + ] + return ( + "## Storage (48 pilot recordings)\n\n" + + table( + [ + "format", + "variant", + "frames", + "numeric MB", + "video/image MB", + "metadata MB", + "total MB", + "numeric B/frame", + "video B/frame", + "total B/frame", + ], + out, + ) + + "\n" + ) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--results", type=Path, required=True) + ap.add_argument("--out", type=Path, required=True) + args = ap.parse_args() + r = args.results + local = load(r / "local_throughput.jsonl") + s3 = load(r / "s3_throughput.jsonl") + parts = [ + "# LeRobot v3.0 vs zarr vs Lance\n", + storage(load(r / "storage.jsonl")), + throughput(local, "Throughput, local SATA SSD (walle)"), + nocache(local), + throughput( + s3, "Throughput, MinIO (S3) on clippy (SATA SSD), MinIO on 4 of the 16 CPUs" + ), + by_label( + [ + x + for x in s3 + if x["format"] in ("lance", "zarr_fast") + and x["workers"] == 8 + and x["label"] in ("tput", "per_sample") + ], + "Batched fetch vs per-sample reads over MinIO (8 workers)", + "per_sample removes MultiDataset.__getitems__, so the loader reads one sample at a time like zarr and LeRobot.\n\n", + ), + by_label( + [x for x in local if x["label"] == "per_sample"] + + [ + x + for x in local + if x["format"] == "zarr_fast" + and x["workers"] == 12 + and x["label"] == "tput" + and x["pattern"] != "frames8" + ], + "zarr_fast local: batched vs per-sample (12 workers)", + ), + by_label( + load(r / "cache_ablation.jsonl"), + "Lance decoder cache 16 vs 100 (local chunk1, 12 workers)", + ), + by_label( + load(r / "scale_throughput.jsonl"), + "Steady-state throughput over 10k episodes (cheetah local, 8 workers, samples drawn from all episodes)", + ), + startup( + load(r / "startup.jsonl"), + "Startup at scale (cheetah, SATA SSD; real + N fake episodes)", + ), + startup( + load(r / "startup_nfs.jsonl"), + "Startup at 100k from NFS (/coc/flash7/scratch, production check)", + ), + startup( + load(r / "startup_len.jsonl"), + "Startup at 10k vs fake episode length (kitt local; N fake = 10000_f)", + ), + ] + args.out.write_text("\n".join(parts)) + print(args.out.read_text()) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/to_lerobot.py b/egomimic/scripts/format_bench/to_lerobot.py new file mode 100644 index 0000000000..544f1ccbc2 --- /dev/null +++ b/egomimic/scripts/format_bench/to_lerobot.py @@ -0,0 +1,326 @@ +"""Lance pilot tables -> LeRobot v3.0 dataset, without re-encoding video. + +Runs in the LeRobot venv (no egomimic imports). Numerics come from the Lance steps table, so all +three formats hold the same values. The Lance MP4s already use LeRobot's default encode +(libsvtav1, g=2, crf=30, preset 12, yuv420p), so each episode's MP4 is appended to the current +video file with LeRobot's own `concatenate_video_files`, exactly as its writer does after encoding. +Data and video files roll over at LeRobot's size limits; episode metadata, tasks, info and stats +go through `LeRobotDatasetMetadata.save_episode`. +""" + +import argparse +import shutil +import time +from pathlib import Path + +import lance +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from lerobot.datasets.compute_stats import ( + auto_downsample_height_width, + compute_episode_stats, + get_feature_stats, + sample_indices, +) +from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata +from lerobot.datasets.io_utils import get_file_size_in_mb +from lerobot.datasets.utils import update_chunk_file_indices +from lerobot.datasets.video_utils import ( + concatenate_video_files, + get_video_duration_in_s, +) + +VIDEO_KEY = "observation.images.front_1" +VECTORS = { + "head_pose": 7, + "eye_gaze": 3, + **{ + f"{s}.{f}": d + for s in ("left", "right") + for f, d in ( + ("wrist_pose", 7), + ("ee_pose", 7), + ("keypoints", 63), + ("aria_keypoints", 63), + ) + }, +} +SCALARS = {"left.confidence": "float32", "right.confidence": "float32"} +BOOLS = ("left_valid", "right_valid", "head_pose_valid") +# Packed layout, as LeRobot datasets ship: one action and one state vector (the keymap's horizon +# keys); everything else stays its own observation.* column. +PACKED = [ + f"{s}.{f}" + for s in ("left", "right") + for f in ("keypoints", "ee_pose", "wrist_pose") +] +PACKED_DIM = sum(VECTORS[k] for k in PACKED) + + +def features(h: int, w: int, packed: bool) -> dict: + vec = {k: d for k, d in VECTORS.items() if not (packed and k in PACKED)} + f = {k: {"dtype": "float32", "shape": (d,), "names": None} for k, d in vec.items()} + if packed: + for k in ("action", "observation.state"): + f[k] = {"dtype": "float32", "shape": (PACKED_DIM,), "names": None} + f.update( + {k: {"dtype": t, "shape": (1,), "names": None} for k, t in SCALARS.items()} + ) + f.update({k: {"dtype": "bool", "shape": (1,), "names": None} for k in BOOLS}) + f["timestamp_ns"] = {"dtype": "int64", "shape": (1,), "names": None} + f[VIDEO_KEY] = { + "dtype": "video", + "shape": (h, w, 3), + "names": ["height", "width", "channels"], + } + return f + + +def column(tbl: pa.Table, dotted: str) -> np.ndarray: + path = dotted.split(".") + col = tbl.column(path[0]).combine_chunks() + for p in path[1:]: + col = col.field(p) + return col + + +def episode_table( + steps: pa.Table, ep: int, task_index: int, start: int, fps: float, packed: bool +) -> pa.Table: + n = steps.num_rows + cols = {} + if packed: + vec = np.concatenate( + [ + np.asarray(column(steps, k).flatten(), np.float32).reshape(n, -1) + for k in PACKED + ], + axis=1, + ) + arr = pa.FixedSizeListArray.from_arrays(pa.array(vec.reshape(-1)), PACKED_DIM) + cols["action"] = arr + cols["observation.state"] = arr + for k, d in VECTORS.items(): + if packed and k in PACKED: + continue + cols[k] = column(steps, k).cast(pa.list_(pa.float32(), d)) + for k in SCALARS: + cols[k] = column(steps, k).cast(pa.float32()) + for k in BOOLS: + cols[k] = column(steps, k) + cols["timestamp_ns"] = column(steps, "timestamp_ns") + cols["timestamp"] = pa.array(np.arange(n, dtype=np.float32) / np.float32(fps)) + cols["frame_index"] = pa.array(np.arange(n, dtype=np.int64)) + cols["episode_index"] = pa.array(np.full(n, ep, dtype=np.int64)) + cols["index"] = pa.array(np.arange(start, start + n, dtype=np.int64)) + cols["task_index"] = pa.array(np.full(n, task_index, dtype=np.int64)) + return pa.table(cols) + + +def numeric_stats(tbl: pa.Table, feats: dict) -> dict: + data = {} + for k in tbl.column_names: + if k not in feats or feats[k]["dtype"] in ("video", "string"): + continue + col = tbl.column(k).combine_chunks() + if pa.types.is_fixed_size_list(col.type) or pa.types.is_list(col.type): + arr = np.asarray(col.flatten()).reshape(len(tbl), -1) + else: + arr = np.asarray(col) + data[k] = arr.astype(np.float64) if arr.dtype == bool else arr + return compute_episode_stats(data, {k: feats[k] for k in data}) + + +def video_stats(mp4: Path, n: int) -> dict: + """LeRobot's image stats (sample_images + per-channel stats / 255) on frames decoded from the MP4.""" + from torchcodec.decoders import VideoDecoder + + dec = VideoDecoder(str(mp4), dimension_order="NCHW") + idx = sample_indices(min(n, dec.metadata.num_frames)) + imgs = [ + auto_downsample_height_width(f.numpy()) + for f in dec.get_frames_at(indices=idx).data + ] + st = get_feature_stats(np.stack(imgs), axis=(0, 2, 3), keepdims=True) + return { + k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in st.items() + } + + +class DataWriter: + """Mirrors `DatasetWriter._save_episode_data`: append episodes to one parquet file until the + next would push it past `data_files_size_in_mb`.""" + + def __init__(self, meta: LeRobotDatasetMetadata): + self.meta, self.chunk, self.file, self.writer = meta, 0, 0, None + self.file_start, self.next_index = 0, 0 + + def path(self) -> Path: + return self.meta.root / self.meta.data_path.format( + chunk_index=self.chunk, file_index=self.file + ) + + def write(self, tbl: pa.Table) -> dict: + if self.writer is not None: + size = get_file_size_in_mb(self.path()) + in_file = self.next_index - self.file_start + per_frame = size / in_file if in_file else 0 + if size + per_frame * tbl.num_rows >= self.meta.data_files_size_in_mb: + self.writer.close() + self.writer = None + self.chunk, self.file = update_chunk_file_indices( + self.chunk, self.file, self.meta.chunks_size + ) + self.file_start = self.next_index + if self.writer is None: + self.path().parent.mkdir(parents=True, exist_ok=True) + self.writer = pq.ParquetWriter( + self.path(), tbl.schema, compression="snappy", use_dictionary=True + ) + self.writer.write_table(tbl) + out = { + "data/chunk_index": self.chunk, + "data/file_index": self.file, + "dataset_from_index": self.next_index, + "dataset_to_index": self.next_index + tbl.num_rows, + } + self.next_index += tbl.num_rows + return out + + def close(self): + if self.writer is not None: + self.writer.close() + + +class VideoWriter: + """Mirrors `DatasetWriter._save_episode_video` for an already-encoded episode MP4.""" + + def __init__(self, meta: LeRobotDatasetMetadata, key: str, tmp: Path): + self.meta, self.key, self.tmp = meta, key, tmp + self.chunk, self.file, self.duration = 0, 0, None + + def path(self) -> Path: + return self.meta.root / self.meta.video_path.format( + video_key=self.key, chunk_index=self.chunk, file_index=self.file + ) + + def write(self, mp4: Path) -> dict: + ep_s = get_video_duration_in_s(mp4) + ep_mb = get_file_size_in_mb(mp4) + if ( + self.duration is not None + and get_file_size_in_mb(self.path()) + ep_mb + >= self.meta.video_files_size_in_mb + ): + self.chunk, self.file = update_chunk_file_indices( + self.chunk, self.file, self.meta.chunks_size + ) + self.duration = None + if self.duration is None: + self.path().parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(mp4, self.path()) + start = 0.0 + else: + start = self.duration + # concatenate_video_files reads its inputs before writing, but stage a copy anyway + # so the source table's MP4 is never an output. + staged = self.tmp / mp4.name + shutil.copyfile(mp4, staged) + concatenate_video_files([self.path(), staged], self.path()) + staged.unlink() + self.duration = start + ep_s + return { + f"videos/{self.key}/chunk_index": self.chunk, + f"videos/{self.key}/file_index": self.file, + f"videos/{self.key}/from_timestamp": start, + f"videos/{self.key}/to_timestamp": start + ep_s, + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--lance-root", type=Path, required=True) + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--camera", default="front_1_480") + ap.add_argument("--repo-id", default="egoverse/aria_pilot") + ap.add_argument( + "--limit", type=int, default=0, help="first N recordings only (smoke tests)" + ) + ap.add_argument( + "--packed", + action="store_true", + help="one action + one observation.state vector", + ) + args = ap.parse_args() + + root = args.lance_root + rec = lance.dataset(str(root / "recordings.lance")).to_table().to_pandas() + tasks = ( + lance.dataset(str(root / "tasks.lance")) + .to_table(columns=["recording_id", "task_name", "task_description"]) + .to_pandas() + ) + rec = rec.merge(tasks, on="recording_id") + steps_ds = lance.dataset(str(root / "steps/aria.lance")) + ids = ( + steps_ds.to_table(columns=["recording_id"]) + .column("recording_id") + .to_numpy(zero_copy_only=False) + ) + first_row = {} + for i, r in enumerate(ids): + first_row.setdefault(r, i) + rec["row_start"] = rec.recording_id.map(first_row) + rec = rec.sort_values("row_start").reset_index(drop=True) + if args.limit: + rec = rec.head(args.limit) + + from torchcodec.decoders import VideoDecoder + + first_mp4 = root / dict(rec.video_path[0])[args.camera] + m = VideoDecoder(str(first_mp4)).metadata + fps = int(round(float(rec.fps[0]))) + feats = features(m.height, m.width, args.packed) + + if args.out.exists(): + shutil.rmtree(args.out) + meta = LeRobotDatasetMetadata.create( + repo_id=args.repo_id, fps=fps, features=feats, root=args.out + ) + tmp = args.out.parent / f".{args.out.name}_tmp" + tmp.mkdir(parents=True, exist_ok=True) + data_w, video_w = DataWriter(meta), VideoWriter(meta, VIDEO_KEY, tmp) + t0 = time.perf_counter() + for ep, row in rec.iterrows(): + n = int(row.num_frames) + steps = steps_ds.take(list(range(int(row.row_start), int(row.row_start) + n))) + mp4 = root / dict(row.video_path)[args.camera] + nv = VideoDecoder(str(mp4)).metadata.num_frames + if nv != n: + raise ValueError(f"{row.recording_id}: {nv} video frames, {n} steps rows") + task = row.task_description or row.task_name + meta.save_episode_tasks([task]) + task_index = int(meta.get_task_index(task)) + tbl = episode_table(steps, ep, task_index, data_w.next_index, fps, args.packed) + md = data_w.write(tbl) + md.update(video_w.write(mp4)) + stats = numeric_stats(tbl, meta.features) + stats[VIDEO_KEY] = video_stats(mp4, n) + meta.save_episode(ep, n, [task], stats, md) + if ep == 0: + meta.update_video_info(VIDEO_KEY) + print( + f"ep {ep} {row.recording_id} {n} frames {time.perf_counter() - t0:.0f}s", + flush=True, + ) + data_w.close() + meta.finalize() + shutil.rmtree(tmp) + print( + f"done: {len(rec)} episodes, {data_w.next_index} frames, {time.perf_counter() - t0:.0f}s" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/format_bench/zarr_fast.py b/egomimic/scripts/format_bench/zarr_fast.py new file mode 100644 index 0000000000..979145d0aa --- /dev/null +++ b/egomimic/scripts/format_bench/zarr_fast.py @@ -0,0 +1,257 @@ +"""Bench-only zarr sampler over the stores as written (one shard per array): same samples as +ZarrDataset, fewer and concurrent reads. + +zarr-python 3 re-reads `/zarr.json` on every `group[key]` (no consolidated metadata) and the +whole shard index on every partial read (46 KB per JPEG frame), then fetches inner chunks one by +one. Here each worker opens every array once, keeps each shard index after its first read, and +issues every inner-chunk range of a sample, or of a whole DataLoader batch through +`MultiDataset.__getitems__`, in one asyncio.gather. Decoding stays byte-for-byte zarr's: the +inner codecs are checked against the two layouts ZarrWriter produces (bytes|vlen-bytes, then +optional zstd); anything else goes through stock zarr. +""" + +import asyncio + +import numcodecs +import numpy as np +from zarr.abc.store import RangeByteRequest, SuffixByteRequest +from zarr.codecs import BytesCodec, ShardingCodec, ZstdCodec +from zarr.codecs.vlen_utf8 import VLenBytesCodec +from zarr.core.buffer import default_buffer_prototype +from zarr.core.sync import sync + +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + ZarrDataset, + ZarrEpisode, +) + +_EMPTY = 2**64 - 1 +_VLEN = numcodecs.VLenBytes() + + +class _Layout: + """How to read one array's single shard without zarr's codec pipeline.""" + + __slots__ = ("path", "inner", "n_inner", "dtype", "vlen", "zstd", "index") + + def __init__(self, arr): + meta = arr.metadata + (shard,) = meta.codecs + if not isinstance(shard, ShardingCodec) or tuple( + meta.chunk_grid.chunk_shape + ) != tuple(arr.shape): + raise ValueError("not a single-shard array") + if shard.index_location.value != "end" or [ + type(c).__name__ for c in shard.index_codecs + ] != [ + "BytesCodec", + "Crc32cCodec", + ]: + raise ValueError("unexpected shard index codecs") + first, *rest = shard.codecs + if not isinstance(first, BytesCodec | VLenBytesCodec) or any( + not isinstance(c, ZstdCodec) for c in rest + ): + raise ValueError("unexpected inner codecs") + self.vlen = isinstance(first, VLenBytesCodec) + if ( + not self.vlen + and first.endian is not None + and first.endian.value != "little" + ): + raise ValueError("big-endian") + self.zstd = bool(rest) + self.inner = tuple(shard.chunk_shape) + self.n_inner = int(np.prod([-(-s // c) for s, c in zip(arr.shape, self.inner)])) + if len(self.inner) > 1 and self.inner[1:] != tuple(arr.shape[1:]): + raise ValueError("chunked along a non-leading axis") + self.dtype = None if self.vlen else arr.dtype + self.path = arr.store_path / meta.encode_chunk_key((0,) * len(arr.shape)) + self.index = None + + def decode(self, raw: bytes) -> np.ndarray: + if self.zstd: + raw = numcodecs.Zstd().decode(raw) + if self.vlen: + return _VLEN.decode(raw) + return np.frombuffer(raw, dtype=self.dtype).reshape(self.inner) + + +async def _get(path, byte_range): + buf = await path.get(prototype=default_buffer_prototype(), byte_range=byte_range) + return buf.to_bytes() + + +class FastZarrEpisode(ZarrEpisode): + """ZarrEpisode whose `read_intervals` reads only inner-chunk byte ranges, concurrently. + + Caches live in `__dict__` (ZarrEpisode uses __slots__), so the s3patch pickling, which keeps + only `_path`, gives each worker fresh ones.""" + + def _cache(self) -> dict: + c = self.__dict__.get("_fz") + if c is None: + c = self.__dict__["_fz"] = {"layouts": {}, "prefetched": None} + return c + + def layout(self, key: str): + layouts = self._cache()["layouts"] + if key not in layouts: + try: + layouts[key] = _Layout(self._store[key]) + except ValueError: + layouts[key] = None + return layouts[key] + + def requests(self, intervals: dict) -> tuple[list, list]: + """(indexes to load, (key, layout, inner chunk ids)) for `intervals`; keys without a + fast layout are left out and read by stock zarr.""" + need_index, reads = [], [] + for key, ivs in intervals.items(): + lay = self.layout(key) + if lay is None: + continue + ids = set() + for lo, hi in self.plan_reads(ivs, lay.inner[0]): + ids.update(range(lo // lay.inner[0], (hi - 1) // lay.inner[0] + 1)) + if lay.index is None: + need_index.append(lay) + reads.append((key, lay, sorted(ids))) + return need_index, reads + + def build(self, intervals: dict, reads: list, chunks: dict) -> dict: + """`read_intervals`' result from decoded inner chunks `{(id(layout), i): ndarray}`.""" + fast = {key for key, _, _ in reads} + result = {} + for key, lay, _ in reads: + c = lay.inner[0] + out = [] + for lo, hi in self.plan_reads(intervals[key], c): + first, last = lo // c, (hi - 1) // c + rows = np.concatenate( + [chunks[(id(lay), i)] for i in range(first, last + 1)] + ) + out.append((lo, rows[lo - first * c : hi - first * c])) + result[key] = out + rest = {k: v for k, v in intervals.items() if k not in fast} + if rest: + result.update(super().read_intervals(rest)) + return result + + def read_intervals(self, intervals: dict) -> dict: + cache = self._cache() + pre = cache["prefetched"] + if pre is not None and pre[0] == intervals: + cache["prefetched"] = None + return pre[1] + return fetch_many([(self, intervals)])[0] + + +def fetch_many(items: list) -> list: + """`read_intervals` for many (episode, intervals) at once: shard indexes first, then every + inner chunk, each phase one gather.""" + plans = [(ep, iv, *ep.requests(iv)) for ep, iv in items] + indexes = {id(lay): lay for _, _, need, _ in plans for lay in need}.values() + wanted = {} + for _, _, _, reads in plans: + for _, lay, ids in reads: + for i in ids: + wanted[(id(lay), i)] = (lay, i) + + async def run(): + if indexes: + raws = await asyncio.gather( + *( + _get(lay.path, SuffixByteRequest(lay.n_inner * 16 + 4)) + for lay in indexes + ) + ) + for lay, raw in zip(indexes, raws): + lay.index = np.frombuffer(raw[:-4], " dict: + """The per-key intervals `__getitem__` reads for frame `idx`.""" + out = {} + for zarr_key, window in self._read_windows(idx).values(): + if isinstance(window, np.ndarray): + spans = [(int(i), int(i) + 1) for i in np.unique(window)] + else: + start, end = window + spans = [(start, start + 1 if end is None else end)] + out.setdefault(zarr_key, []).extend(spans) + return out + + def plan_sample(self, i: int) -> dict: + return {"idx": i, "intervals": self.intervals(i)} + + @staticmethod + def fetch_planned(items: list) -> list: + return fetch_many( + [(leaf.episode_reader, plan["intervals"]) for leaf, plan in items] + ) + + def assemble_sample(self, plan: dict, fetched: dict): + self.episode_reader._cache()["prefetched"] = (plan["intervals"], fetched) + try: + return self[plan["idx"]] + finally: + self.episode_reader._cache()["prefetched"] = None + + +class FastLocalEpisodeResolver(LocalEpisodeResolver): + _dataset_class = FastZarrDataset + + +def s3_leaves(root: str, key_map: dict, embodiment: str | None) -> dict: + """`bench_ego.zarr_s3_leaves` with FastZarrDataset leaves.""" + from egomimic.scripts.format_bench import s3patch + + store = s3patch.obstore_s3(root) + names = sorted( + p.rstrip("/").split("/")[-1] + for p in store.list_with_delimiter()["common_prefixes"] + ) + leaves = { + n: FastZarrDataset(f"{root.rstrip('/')}/{n}", key_map, None) for n in names + } + if embodiment: + leaves = {n: d for n, d in leaves.items() if d.embodiment.lower() == embodiment} + return leaves diff --git a/egomimic/scripts/lance_annotations.py b/egomimic/scripts/lance_annotations.py new file mode 100644 index 0000000000..c59867814a --- /dev/null +++ b/egomimic/scripts/lance_annotations.py @@ -0,0 +1,257 @@ +"""Review and commit annotation change files, and build them from zarr episodes or the +Scale annotation bucket. + + review validate a change file against the tables and print the per-recording diff + stage the same, then put it under /annotations/ for the committer + (egomimic.scripts.lance_commit --lab annotations), which applies it as one + commit on annotations.lance, recorded in changes.lance and tagged + annotations- + from-zarr change file from zarr episodes' annotations, mapped onto the recording clock + through each episode's obs_rgb_timestamps_ns + from-bucket change file from _.json span lists, as + language_process/scale_to_bucket_annotation_parallel.py writes them + +from-bucket counts a span's frames from its alias's first recording frame. That is exact for +EVA, whose zarr episodes hold every recording frame, but not for Aria clips that dropped +frames: map those with from-zarr. + +`--lance` is a local directory or an s3:// prefix. See egomimic/rldb/lance/annotations.py for +the change file format. +""" + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +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 + +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 stage(file: Path, staging, approver: str) -> str: + """Upload a reviewed change file to `/annotations//`, meta.json last; `` + is the file's content hash, so staging the same file again changes nothing.""" + sid = hashlib.sha256(file.read_bytes()).hexdigest()[:16] + dst = commit.join(staging, "annotations", sid) + commit.put_file(file, commit.join(dst, "changes.parquet")) + meta = { + "approver": approver, + "source": file.name, + "rows": pq.read_metadata(file).num_rows, + "staged_at": time.time(), + } + commit.put_bytes(json.dumps(meta, indent=1).encode(), commit.join(dst, "meta.json")) + return dst + + +def cmd_stage(args): + rv = _review(args) + if rv.errors: + sys.exit("not staged: fix the errors above") + print(f"staged {args.file} as {stage(args.file, args.staging, args.approver)}") + + +def _aliases(root) -> dict[str, dict]: + rows = commit._open(commit.join(root, "aliases.lance")).to_table().to_pylist() + return {a["episode_hash"]: a for a in rows} + + +def _clocks(root, rids: set[str]) -> dict[str, tuple[np.ndarray, float]]: + """{recording_id: (sorted steps timestamp_ns, fps)}, one steps read per lab.""" + recs = commit._open(commit.join(root, "recordings.lance")).to_table( + columns=["recording_id", "lab", "fps"] + ) + out = {} + for lab in set(recs.column("lab").to_pylist()): + mine = [ + r for r in recs.to_pylist() if r["lab"] == lab and r["recording_id"] in rids + ] + if not mine: + continue + ids = ", ".join(commit._sql_str(r["recording_id"]) for r in mine) + t = commit._open(commit.steps_uri(root, lab)).to_table( + columns=["recording_id", "timestamp_ns"], filter=f"recording_id IN ({ids})" + ) + rid_col = np.asarray(t.column("recording_id").to_pylist()) + ts = t.column("timestamp_ns").to_numpy() + for r in mine: + out[r["recording_id"]] = ( + np.sort(ts[rid_col == r["recording_id"]]), + r["fps"], + ) + return out + + +def _add(rid, t0, t1, text, vendor) -> dict: + return dict( + op="add", + recording_id=rid, + t_start_ns=int(t0), + t_end_ns=int(t1), + annotation_type="language", + text=str(text), + annotation_vendor=vendor, + ) + + +def _write(rows, dropped, out): + pq.write_table(pa.Table.from_pylist(rows, schema=CHANGE_SCHEMA), out) + print( + f"wrote {len(rows)} annotations to {out}; dropped {sum(map(len, dropped.values()))}" + ) + for h, why in sorted(dropped.items()): + print(f" {h}: {len(why)} dropped, e.g. {why[0]}") + + +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_from_zarr(args): + found = {} + for h, a in _aliases(args.lance).items(): + path = next( + (p for p in (args.zarr / h, args.zarr / f"{h}.zarr") if p.exists()), None + ) + if path is not None: + found[h] = (a["recording_id"], path) + clocks = _clocks(args.lance, {rid for rid, _ in found.values()}) + rows, dropped = [], {} + for h, (rid, path) in found.items(): + anns, zts = _zarr_annotations(path) + if not anns: + continue + rts, fps = clocks[rid] + period = int(round(1e9 / fps)) + + 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 + if e < len(zts): + t1 = rts[last] + else: + t1 = rts[last + 1] if last + 1 < len(rts) else rts[-1] + period + rows.append(_add(rid, rts[i0], t1, d.get("text", ""), args.vendor)) + _write(rows, dropped, args.out) + + +def cmd_from_bucket(args): + suffix = f"_{args.key}.json" + aliases = _aliases(args.lance) + names = [n for n in commit.list_files(args.bucket) if n.endswith(suffix)] + hashes = {n[: -len(suffix)]: n for n in names} + dropped = {h: ["no alias"] for h in hashes if h not in aliases} + todo = {h: aliases[h] for h in hashes if h in aliases} + clocks = _clocks(args.lance, {a["recording_id"] for a in todo.values()}) + rows = [] + for h, a in sorted(todo.items()): + rts, _ = clocks[a["recording_id"]] + clip = rts[(rts >= a["t_start_ns"]) & (rts < a["t_end_ns"])] + spans = json.loads(commit.get_bytes(commit.join(args.bucket, hashes[h]))) + for d in spans: + s, e = int(d["start_idx"]), int(d["end_idx"]) + if not 0 <= s < e <= len(clip): + dropped.setdefault(h, []).append( + f"[{s}, {e}) is outside the alias's {len(clip)} frames" + ) + continue + t1 = clip[e] if e < len(clip) else a["t_end_ns"] + rows.append(_add(a["recording_id"], clip[s], t1, d["text"], args.vendor)) + _write(rows, dropped, args.out) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--lance", required=True, help="Lance root (dir or s3:// prefix)") + 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("stage", help="review, then stage a change file to commit") + p.add_argument("file", type=Path) + p.add_argument( + "--staging", required=True, help="staging root (dir or s3:// prefix)" + ) + p.add_argument("--approver", required=True) + p.set_defaults(fn=cmd_stage) + + p = sub.add_parser("from-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_from_zarr) + + p = sub.add_parser("from-bucket", help="change file from bucket span lists") + p.add_argument("--bucket", required=True, help="dir or s3:// prefix of the JSONs") + p.add_argument("--key", default="annotations", help="_.json") + 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_from_bucket) + + args = ap.parse_args() + args.fn(args) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_commit.py b/egomimic/scripts/lance_commit.py new file mode 100644 index 0000000000..3f4a267415 --- /dev/null +++ b/egomimic/scripts/lance_commit.py @@ -0,0 +1,463 @@ +"""The committer: the only writer of the Lance tables. + +Converters stage deliveries (steps.parquet, MP4s, meta.json last) under +`///` (commit.stage). Each run validates the complete ones, +commits them in batches (each batch tagged `-` on every table and recorded +in `changes`), and deletes their staged files. A delivery that fails validation gets +REJECTED.json (errors plus a digest of its files) and is skipped until its files change. +Aria recording and task metadata come from SQL; EVA's from the delivery's meta.json; a +delivery staged with commit.stage_delivery (meta.json "format": "delivery", e.g. vendor +MCAP) carries its own catalog rows, and its labels are synced right after its batch: the +label vendors' committed labels of that recording not among them are deleted. A delivery's +meta.json "provenance" is kept in its batch's `changes` row. + +`--lab annotations` commits annotation change files that `lance_annotations stage` put under +`/annotations//`, one commit and tag each, oldest first; each is reviewed again +against the tables as they are now and rejected like a bad delivery if it no longer applies. +See egomimic/rldb/lance/commit.py for the table operations. + + python -m egomimic.scripts.lance_commit --staging s3://rldb/lance_staging \\ + --lab eva --lance s3://rldb/lance --work-dir /coc/flash7/scratch/committer +""" + +import argparse +import hashlib +import json +import shutil +import tempfile +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from sqlalchemy import text + +from egomimic.rldb.lance import commit, mp4_ranges +from egomimic.rldb.lance.annotations import read_changes, review +from egomimic.rldb.lance.dataset import read_range +from egomimic.rldb.lance.schema import ANNOTATIONS as ANNOTATIONS_SCHEMA +from egomimic.rldb.lance.schema import task_instance_id +from egomimic.utils.aws.aws_sql import ( + create_default_engine, + episode_hash_to_timestamp_ms, + timestamp_ms_to_episode_hash, +) + +REJECTED = "REJECTED.json" +ANNOTATIONS = "annotations" + + +def sql_rows(hashes: list[str]) -> pd.DataFrame: + q = text( + """select episode_hash, operator, lab, embodiment, rig_name, scene, task, + task_description, objects, license + from app.episodes where episode_hash = any(:h)""" + ) + with create_default_engine().connect() as c: + return pd.read_sql(q, c, params={"h": hashes}).set_index("episode_hash") + + +def check_delivery(d: Path, meta: dict, steps: pa.Table) -> list[str]: + errs = [] + n = meta["num_frames"] + if steps.num_rows != n: + errs.append(f"steps rows {steps.num_rows} != num_frames {n}") + ts = steps.column("timestamp_ns").to_numpy() + if not (np.diff(ts) > 0).all(): + errs.append("timestamps not strictly increasing") + for cam in meta.get("videos", ("front_1", "front_1_480")): + frames = mp4_ranges.read_index( + lambda a, b, p=d / f"{cam}.mp4": read_range(str(p), a, b) + ).num_frames + if frames != n: + errs.append(f"{cam}.mp4 has {frames} frames, steps has {n}") + return errs + + +def text_or_none(v): + return None if v is None or (isinstance(v, float) and np.isnan(v)) else str(v) + + +def objects_list(v) -> list[str]: + if v is None: + return [] + items = list(v) if not isinstance(v, str) else v.strip("{}").split(",") + return [str(o).strip('"') for o in items if o and str(o) not in ("None", "NULL")] + + +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() + ] + sql = sql_rows(hashes) + + deliveries, rejected = [], {} + for d, m in sorted(metas.items(), key=lambda kv: kv[1]["recording_id"]): + steps = pq.read_table(d / "steps.parquet") + errs = check_delivery(d, m, steps) + h = timestamp_ms_to_episode_hash(m["vendor_recording_id"]) + if h not in sql.index: + errs.append("no SQL row") + if errs: + rejected[m["vendor_recording_id"]] = errs + continue + s = sql.loc[h] + rid = m["recording_id"] + recording = dict( + recording_id=rid, + vendor_recording_id=m["vendor_recording_id"], + lab="aria", + site=text_or_none(s.lab), + embodiment=text_or_none(s.embodiment), + rig_name=text_or_none(s.rig_name), + operator=text_or_none(s.operator), + scene=text_or_none(s.scene), + captured_at=pd.Timestamp( + episode_hash_to_timestamp_ms(h), unit="ms", tz="UTC" + ), + fps=m["fps"], + num_frames=m["num_frames"], + intrinsics=list(m["intrinsics"].items()), + extrinsics=list(m["extrinsics"].items()), + license=text_or_none(s.license), + ) + task = dict( + 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), + ) + alias = dict( + episode_hash=h, + recording_id=rid, + t_start_ns=m["t_start_ns"], + t_end_ns=m["t_end_ns"], + ) + videos = {cam: d / f"{cam}.mp4" for cam in ("front_1", "front_1_480")} + deliveries.append(commit.Delivery(recording, [task], alias, steps, videos)) + return deliveries, rejected + + +def load_eva_deliveries( + dirs: list[Path], +) -> tuple[list[commit.Delivery], dict[str, list[str]]]: + deliveries, rejected = [], {} + for d in sorted(dirs): + m = json.loads((d / "meta.json").read_text()) + steps = pq.read_table(d / "steps.parquet") + errs = check_delivery(d, m, steps) + if errs: + rejected[m["vendor_recording_id"]] = errs + continue + rid = m["recording_id"] + recording = dict( + recording_id=rid, + vendor_recording_id=m["vendor_recording_id"], + lab="eva", + site=m["site"], + embodiment=m["embodiment"], + rig_name="eva", + operator=m["operator"], + scene=m["scene"], + captured_at=pd.Timestamp( + int(m["vendor_recording_id"]), unit="ms", tz="UTC" + ), + fps=m["fps"], + num_frames=m["num_frames"], + intrinsics=list(m["intrinsics"].items()), + extrinsics=list(m["extrinsics"].items()), + license=None, + ) + task = dict( + task_instance_id=task_instance_id(rid, m["task_name"], m["t_start_ns"]), + recording_id=rid, + t_start_ns=m["t_start_ns"], + t_end_ns=m["t_end_ns"], + task_name=m["task_name"], + task_description=m["task_description"], + objects=m["objects"], + ) + alias = dict( + episode_hash=m["episode_hash"], + recording_id=rid, + t_start_ns=m["t_start_ns"], + t_end_ns=m["t_end_ns"], + ) + videos = {cam: d / f"{cam}.mp4" for cam in m["videos"]} + deliveries.append(commit.Delivery(recording, [task], alias, steps, videos)) + return deliveries, rejected + + +def load_described(dirs: list[Path]): + """Deliveries staged with their catalog rows; also {recording_id: (labels, label + vendors)} for the ones carrying labels.""" + deliveries, rejected, labels = [], {}, {} + for d in sorted(dirs): + m = json.loads((d / "meta.json").read_text()) + try: + delivery, rows = commit.load_staged_delivery(d) + errs = check_delivery(d, m, delivery.steps) + except Exception as e: + delivery, errs = None, [f"unreadable: {type(e).__name__}: {e}"] + if errs: + rejected[d.name] = errs + continue + deliveries.append(delivery) + if rows is not None: + labels[delivery.recording_id] = (rows, set(m.get("label_vendors", []))) + return deliveries, rejected, labels + + +def sync_labels(lance_root: str, labels: dict, operator, provenance) -> dict | None: + """Commit each recording's delivered labels in one annotation commit; the label + vendors' committed labels of that recording not among them are deleted.""" + uri = commit.annotations_uri(lance_root) + adds, stale = [], [] + for rid, (rows, vendors) in sorted(labels.items()): + new = set(rows.column("annotation_id").to_pylist()) + if commit._exists(uri): + old = commit._open(uri).to_table( + columns=["annotation_id", "annotation_vendor"], + filter=f"recording_id = {commit._sql_str(rid)}", + ) + stale += [ + r["annotation_id"] + for r in old.to_pylist() + if r["annotation_vendor"] in vendors and r["annotation_id"] not in new + ] + adds.append(rows) + table = pa.concat_tables(adds) if adds else ANNOTATIONS_SCHEMA.empty_table() + return commit.commit_annotations( + lance_root, table, stale, operator=operator, provenance=provenance + ) + + +LOADERS = {"aria": load_deliveries, "eva": load_eva_deliveries} + + +def _load(lab: str, dirs: list[Path]): + """LOADERS[lab] over `dirs`; a delivery the loader cannot even read is rejected alone.""" + try: + return LOADERS[lab](dirs) + except Exception: + deliveries, rejected = [], {} + for d in dirs: + try: + ok, bad = LOADERS[lab]([d]) + except Exception as e: + ok, bad = [], {d.name: [f"unreadable: {type(e).__name__}: {e}"]} + deliveries += ok + rejected.update(bad) + return deliveries, rejected + + +def files_digest(files: dict[str, tuple[int, str]]) -> str: + """Changes whenever a staged file is replaced (S3 ETag, or local size and mtime).""" + h = hashlib.sha256() + for name, (size, version) in sorted(files.items()): + if name != REJECTED: + h.update(f"{name}\0{size}\0{version}\0".encode()) + return h.hexdigest() + + +def pending(staging: str, lab: str, replace: bool = False) -> list[tuple[str, str]]: + """(vendor id, files digest) of complete staged deliveries not rejected as they are; + with `replace`, ones rejected only for being committed already come back.""" + out = [] + for vid in commit.list_dirs(commit.join(staging, lab)): + prefix = commit.join(staging, lab, vid) + files = commit.list_files(prefix) + if "meta.json" not in files: + continue + digest = files_digest(files) + if REJECTED in files: + with tempfile.TemporaryDirectory() as t: + commit.get_file(commit.join(prefix, REJECTED), Path(t) / REJECTED) + note = json.loads((Path(t) / REJECTED).read_text()) + if note["digest"] == digest and not (replace and note["committed"]): + continue + out.append((vid, digest)) + return out + + +def _reject(staging: str, lab: str, vid: str, errs, digest: str, committed=False): + note = { + "errors": errs, + "digest": digest, + "committed": committed, + "rejected_at": time.time(), + } + commit.put_bytes( + json.dumps(note, indent=1).encode(), commit.join(staging, lab, vid, REJECTED) + ) + + +def _unstage(staging: str, lab: str, vid: str) -> None: + prefix = commit.join(staging, lab, vid) + commit.remove([commit.join(prefix, name) for name in commit.list_files(prefix)]) + + +def commit_staged_annotations(staging: str, lance_root: str, work_dir: Path) -> dict: + """Commit every pending staged annotation change, oldest first, one commit each.""" + report = {"batches": [], "rejected": {}} + todo = [] + for sid, digest in pending(staging, ANNOTATIONS): + prefix = commit.join(staging, ANNOTATIONS, sid) + meta = json.loads(commit.get_bytes(commit.join(prefix, "meta.json"))) + todo.append((meta["staged_at"], sid, digest, meta)) + for _, sid, digest, meta in sorted(todo): + local = work_dir / ANNOTATIONS / sid + commit.get_file( + commit.join(staging, ANNOTATIONS, sid, "changes.parquet"), + local / "changes.parquet", + ) + rv = review(lance_root, read_changes(local / "changes.parquet")) + if rv.errors: + _reject(staging, ANNOTATIONS, sid, rv.errors, digest) + report["rejected"][sid] = rv.errors + else: + out = commit.commit_annotations( + lance_root, + rv.adds, + rv.deletes, + operator=meta["approver"], + provenance={k: meta[k] for k in ("approver", "source", "staged_at")}, + ) + if out: + report["batches"].append(out) + _unstage(staging, ANNOTATIONS, sid) + shutil.rmtree(local, ignore_errors=True) + return report + + +def commit_staged( + staging: str, + lab: str, + lance_root: str, + work_dir: Path, + batch_size: int = 64, + replace: bool = False, + operator: str | None = None, +) -> dict: + """Commit every pending staged delivery of `lab`, `batch_size` at a time; returns + {"batches": [commit_batch results], "rejected": {vendor id: errors}}.""" + if lab == ANNOTATIONS: + return commit_staged_annotations(staging, lance_root, work_dir) + report = {"batches": [], "rejected": {}} + todo = pending(staging, lab, replace) + have = commit.committed_ids(lance_root) + for b in range(0, len(todo), batch_size): + batch = dict(todo[b : b + batch_size]) + local = {} + for vid in batch: + src = commit.join(staging, lab, vid) + local[vid] = work_dir / lab / vid + for name in commit.list_files(src): + if name != REJECTED: + commit.get_file(commit.join(src, name), local[vid] / name) + described = { + vid: path + for vid, path in local.items() + if json.loads((path / "meta.json").read_text()).get("format") + == commit.DESCRIBED + } + rest = [path for vid, path in local.items() if vid not in described] + deliveries, rejected = _load(lab, rest) if rest else ([], {}) + more, bad, labels = load_described(list(described.values())) + deliveries += more + rejected.update(bad) + skipped = [] + if not replace: + skipped = [d for d in deliveries if d.recording_id in have] + for d in skipped: + deliveries.remove(d) + rejected[d.recording["vendor_recording_id"]] = [ + "already committed; rerun with --replace to replace it" + ] + provenance = {} + for d in deliveries: + meta = local[d.recording["vendor_recording_id"]] / "meta.json" + p = json.loads(meta.read_text()).get("provenance") + if p is not None: + provenance[d.recording_id] = p + result = commit.commit_batch( + lance_root, + deliveries, + lab, + replace_committed=replace, + operator=operator, + provenance=provenance or None, + ) + if result: + report["batches"].append(result) + have |= set(result["added"]) + done = set(result["added"]) | set(result["replaced"]) + mine = {rid: v for rid, v in labels.items() if rid in done} + if mine: + synced = sync_labels( + lance_root, mine, operator, {"labels_of": result["tag"]} + ) + if synced: + report["batches"].append(synced) + skipped_ids = {d.recording["vendor_recording_id"] for d in skipped} + for vid, errs in rejected.items(): + _reject(staging, lab, vid, errs, batch[vid], vid in skipped_ids) + report["rejected"].update(rejected) + for d in deliveries: + _unstage(staging, lab, d.recording["vendor_recording_id"]) + for path in local.values(): + shutil.rmtree(path, ignore_errors=True) + return report + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--staging", + required=True, + help="staging root (s3:// or local); deliveries are under ///", + ) + ap.add_argument("--lab", choices=sorted([*LOADERS, ANNOTATIONS]), required=True) + ap.add_argument("--lance", required=True, help="Lance root (s3:// or local)") + ap.add_argument( + "--work-dir", + type=Path, + required=True, + help="local scratch for downloaded deliveries", + ) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument( + "--replace", + action="store_true", + help="replace committed recordings with restaged deliveries", + ) + ap.add_argument("--operator", help="recorded in changes; defaults to $USER") + args = ap.parse_args() + report = commit_staged( + args.staging, + args.lab, + args.lance, + args.work_dir, + args.batch_size, + args.replace, + args.operator, + ) + for b in report["batches"]: + print(f"{b['tag']}: added {len(b['added'])}, replaced {len(b['replaced'])}") + for vid, errs in report["rejected"].items(): + print(f"rejected {vid}: {'; '.join(errs)}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/__init__.py b/egomimic/scripts/lance_pilot/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/lance_pilot/bench_decode.py b/egomimic/scripts/lance_pilot/bench_decode.py new file mode 100644 index 0000000000..a15cd6786a --- /dev/null +++ b/egomimic/scripts/lance_pilot/bench_decode.py @@ -0,0 +1,96 @@ +"""Per-frame decode cost, one thread, files preloaded into memory: zarr JPEG vs the Lance AV1 MP4s. + +jpeg simplejpeg on random zarr frames +av1_random get_frame_at on random indices (what the shuffle buffer does) +av1_sorted get_frames_at on the same indices, sorted +av1_sequential every frame of a contiguous range, in order +""" + +import argparse +import json +import random +import time +from pathlib import Path + +import numpy as np + +from egomimic.rldb.lance.dataset import open_table + + +def _timed(fn, n): + t0 = time.perf_counter() + fn() + return round(1000 * (time.perf_counter() - t0) / n, 3) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--hashes", nargs="+", required=True) + ap.add_argument( + "--zarr-root", type=Path, default=Path("/coc/flash7/scratch/lance_pilot/zarr") + ) + ap.add_argument( + "--lance-root", type=Path, default=Path("/coc/flash7/scratch/lance_pilot/lance") + ) + ap.add_argument("--n", type=int, default=500) + args = ap.parse_args() + + import simplejpeg + import zarr + from torchcodec.decoders import VideoDecoder + + rec = ( + open_table(str(args.lance_root / "recordings.lance")) + .to_table(columns=["recording_id", "video_path"]) + .to_pandas() + ) + ali = ( + open_table(str(args.lance_root / "aliases.lance")) + .to_table(columns=["recording_id", "episode_hash"]) + .to_pandas() + ) + videos = rec.merge(ali, on="recording_id").set_index("episode_hash").video_path + + for h in args.hashes: + rng = random.Random(0) + jpegs = zarr.open_group(str(args.zarr_root / h), mode="r")["images.front_1"][:] + idx = [rng.randrange(len(jpegs)) for _ in range(args.n)] + out = dict( + hash=h, + n=args.n, + jpeg_ms=_timed( + lambda: [ + simplejpeg.decode_jpeg(jpegs[i], colorspace="RGB") for i in idx + ], + args.n, + ), + ) + out["jpeg_kb"] = round(np.mean([len(jpegs[i]) for i in idx]) / 1024, 1) + for cam in ("front_1_480", "front_1"): + mp4 = (args.lance_root / dict(videos[h])[cam]).read_bytes() + dec = VideoDecoder( + mp4, + seek_mode="approximate", + dimension_order="NHWC", + num_ffmpeg_threads=1, + ) + nf = dec.metadata.num_frames + vidx = [rng.randrange(nf) for _ in range(args.n)] + res = f"{dec.metadata.height}p" + out[f"{res}_kb"] = round(len(mp4) / nf / 1024, 1) + out[f"{res}_random_ms"] = _timed( + lambda: [dec.get_frame_at(i) for i in vidx], args.n + ) + out[f"{res}_sorted_ms"] = _timed( + lambda: dec.get_frames_at(indices=sorted(vidx)), args.n + ) + s = rng.randrange(nf - args.n) + out[f"{res}_sequential_ms"] = _timed( + lambda: dec.get_frames_in_range(start=s, stop=s + args.n), args.n + ) + out["codec"] = dec.metadata.codec + print(json.dumps(out), flush=True) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/bench_loader.py b/egomimic/scripts/lance_pilot/bench_loader.py new file mode 100644 index 0000000000..0739cffeb9 --- /dev/null +++ b/egomimic/scripts/lance_pilot/bench_loader.py @@ -0,0 +1,209 @@ +"""Loader benchmark: zarr vs Lance on the same Aria recordings, through the real keymap and transforms. + +Prints one JSON line: resolve time, time to first batch, samples/s, worker startup, RSS. +""" + +import argparse +import json +import multiprocessing as mp +import os +import time +from pathlib import Path + +import psutil +import torch +from torch.utils.data import DataLoader, Dataset + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.utils.gpu_orphans import orphan_guarded + + +def collate(batch): + out = {} + for k in batch[0]: + vals = [b[k] for b in batch] + out[k] = torch.stack(vals) if isinstance(vals[0], torch.Tensor) else vals + return out + + +class PerSample(Dataset): + """Hides MultiDataset.__getitems__, so the DataLoader fetches one sample at a time.""" + + def __init__(self, ds): + self.ds = ds + + def __len__(self): + return len(self.ds) + + def __getitem__(self, i): + return self.ds[i] + + +def rss_mb(proc: psutil.Process) -> float: + total = proc.memory_info().rss + for c in proc.children(recursive=True): + try: + total += c.memory_info().rss + except psutil.NoSuchProcess: + pass + return total / 2**20 + + +def build(args): + key_map = Human.get_keymap(keymap_mode="keypoints", include_ee_pose=True) + if args.no_images: + key_map = { + k: v for k, v in key_map.items() if v.get("key_type") != "camera_keys" + } + transforms = Human.get_transform_list( + mode="keypoints_wristframe_6d", + include_ee_pose=True, + pad_proprio_gripper=True, + stride=3, + ) + hashes = ( + [h.strip() for h in Path(args.hashes).read_text().split() if h.strip()] + if args.hashes + else None + ) + filters = DatasetFilter(episode_hashes=hashes) if hashes else None + if args.format == "zarr": + resolver = LocalEpisodeResolver( + args.zarr_dir, key_map=key_map, transform_list=transforms + ) + leaves = resolver.resolve(filters=filters, expected_embodiment="human_bimanual") + else: + from egomimic.rldb.lance.dataset import ( + LanceEpisodeResolver, + LanceStreamingResolver, + ) + + kw = dict( + key_map=key_map, + transform_list=transforms, + cameras={} if args.no_images else {"images.front_1": args.camera}, + ) + resolver = ( + LanceStreamingResolver( + args.lance_root, + fetch_split=args.fetch_split or None, + overlap_reads=not args.no_overlap, + **kw, + ) + if args.stream + else LanceEpisodeResolver(args.lance_root, **kw) + ) + leaves = resolver.resolve(filters=filters, expected_embodiment="human_bimanual") + ds = MultiDataset(datasets=leaves, mode="total") + return (PerSample(ds) if args.per_sample else ds), len(leaves) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--format", choices=["zarr", "lance"], required=True) + ap.add_argument( + "--zarr-dir", type=Path, default=Path("/coc/flash7/scratch/lance_pilot/zarr") + ) + ap.add_argument( + "--lance-root", + default="/coc/flash7/scratch/lance_pilot/lance", + help="local path or s3:// URI", + ) + ap.add_argument("--camera", default="front_1", help="front_1 (360p) or front_1_480") + ap.add_argument( + "--no-images", action="store_true", help="drop camera keys: numerics only" + ) + ap.add_argument( + "--hashes", type=Path, help="whitespace-separated episode hashes to pin" + ) + ap.add_argument( + "--stream", + action="store_true", + help="read through LanceStreamingResolver: batched fetches, MP4 byte ranges", + ) + ap.add_argument( + "--per-sample", action="store_true", help="disable the batched __getitems__" + ) + ap.add_argument( + "--fetch-split", + type=int, + default=0, + help="samples per concurrent part; 0 = one", + ) + ap.add_argument( + "--no-overlap", action="store_true", help="steps take, then chunk reads" + ) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument( + "--context", + default=None, + help="fork | forkserver | spawn (default: fork for zarr, forkserver for lance)", + ) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=200) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--label", default="") + args = ap.parse_args() + context = args.context or ("fork" if args.format == "zarr" else "forkserver") + if context == "forkserver": + mp.set_forkserver_preload( + ["torch", "egomimic.rldb.lance.dataset", "egomimic.rldb.embodiment.human"] + ) + + t0 = time.perf_counter() + ds, n_leaves = build(args) + t_build = time.perf_counter() - t0 + + params = dict( + batch_size=args.batch_size, + shuffle=True, + num_workers=args.workers, + collate_fn=collate, + prefetch_factor=2 if args.workers else None, + persistent_workers=args.workers > 0, + multiprocessing_context=context if args.workers else None, + ) + loader = DataLoader(ds, **orphan_guarded(params)) + t1 = time.perf_counter() + it = iter(loader) + next(it) + t_first = time.perf_counter() - t1 + for _ in range(args.warmup): + next(it) + proc = psutil.Process(os.getpid()) + t2 = time.perf_counter() + for _ in range(args.batches): + next(it) + dt = time.perf_counter() - t2 + print( + json.dumps( + dict( + label=args.label, + format=args.format, + camera="none" + if args.no_images + else (args.camera if args.format == "lance" else "jpeg"), + context=context, + stream=args.stream, + fetch="per_sample" if args.per_sample else "batched", + fetch_split=args.fetch_split, + overlap=not args.no_overlap, + root="r2" if str(args.lance_root).startswith("s3://") else "local", + workers=args.workers, + batch_size=args.batch_size, + recordings=n_leaves, + samples=len(ds), + build_s=round(t_build, 2), + first_batch_s=round(t_first, 2), + samples_per_s=round(args.batches * args.batch_size / dt, 1), + rss_mb=round(rss_mb(proc)), + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/bench_ranges.py b/egomimic/scripts/lance_pilot/bench_ranges.py new file mode 100644 index 0000000000..bb74af9c35 --- /dev/null +++ b/egomimic/scripts/lance_pilot/bench_ranges.py @@ -0,0 +1,206 @@ +"""Random-sample read throughput with frames fetched from the MP4s by byte range. + +Each sample is one random 31-row steps window and the frame at its first row. Frames +come from the faststart MP4s through egomimic.rldb.lance.mp4_ranges: +each worker reads a recording's moov once (one range GET), then a batch issues one range +GET per frame (its keyframe through itself), concurrently, and decodes in memory. +Prints one JSON line per run. +""" + +import argparse +import concurrent.futures as cf +import json +import random +import time + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset + +from egomimic.rldb.lance import mp4_ranges +from egomimic.rldb.lance.dataset import _runs, _s3_client, _table, open_table + +H = 31 +COLS = [ + f"{s}.{k}" + for s in ("left", "right") + for k in ("ee_pose", "keypoints", "wrist_pose") +] + ["head_pose"] + + +def _flat(tbl, col): + path = col.split(".") + c = tbl.column(path[0]).combine_chunks() + for p in path[1:]: + c = c.field(p) + return np.asarray(c.flatten()).reshape(len(c), -1) + + +def _starts(uri): + ds = open_table(uri) + values, starts = _runs(ds.to_table(columns=["recording_id"]).column(0)) + return ds.version, dict(zip(values, starts.tolist())) + + +def run(args, ds): + """Times `ds` through a DataLoader; `args` carries the loader and report settings.""" + if args.workers == 0: + dl = DataLoader(ds, batch_size=args.batch_size, num_workers=0) + else: + dl = DataLoader( + ds, + batch_size=args.batch_size, + num_workers=args.workers, + multiprocessing_context="forkserver", + prefetch_factor=4, + persistent_workers=False, + ) + it = iter(dl) + t0 = time.perf_counter() + lat = [] + for _ in range(args.warmup): + next(it) + first = time.perf_counter() - t0 + t1 = time.perf_counter() + for _ in range(args.batches): + tb = time.perf_counter() + next(it) + lat.append(time.perf_counter() - tb) + secs = time.perf_counter() - t1 + out = dict( + label=args.label, + images=args.images, + workers=args.workers, + batch_size=args.batch_size, + root="r2" if args.lance_root.startswith("s3://") else "local", + warmup_s=round(first, 2), + samples_per_s=round(args.batches * args.batch_size / secs, 1), + batch_wait_p50_ms=round(1000 * float(np.percentile(lat, 50)), 1), + batch_wait_p95_ms=round(1000 * float(np.percentile(lat, 95)), 1), + ) + if args.workers == 0: + n = (args.warmup + args.batches) * args.batch_size + out.update( + {f"{k}_ms_per_sample": round(1000 * ds.timing[k] / n, 2) for k in ds.timing} + ) + return out + + +class Recs: + def __init__(self, root, hashes, camera): + rec = open_table(f"{root}/recordings.lance").to_table().to_pandas() + ali = open_table(f"{root}/aliases.lance").to_table().to_pandas() + rec = rec.merge(ali, on="recording_id").set_index("episode_hash").loc[hashes] + self.steps_uri = f"{root}/steps/aria.lance" + self.steps_version, steps = _starts(self.steps_uri) + self.items = [ + ( + steps[r.recording_id], + f"{root}/{dict(r.video_path)[camera]}", + int(r.num_frames), + ) + for r in rec.itertuples() + ] + + +def _read(url: str, a: int, b: int) -> bytes: + if not url.startswith("s3://"): + with open(url, "rb") as f: + f.seek(a) + return f.read(b - a) + bucket, key = url[len("s3://") :].split("/", 1) + return ( + _s3_client() + .get_object(Bucket=bucket, Key=key, Range=f"bytes={a}-{b - 1}")["Body"] + .read() + ) + + +class RangeDS(Dataset): + def __init__(self, recs, images, threads): + self.recs, self.images, self.threads = recs, images, threads + self.timing = {"take": 0.0, "index": 0.0, "get": 0.0, "decode": 0.0} + self._pool = None + self._dec: dict[str, mp4_ranges.Decoder] = {} + + def __len__(self): + return 10**9 + + def _draw(self): + s0, url, n = random.choice(self.recs.items) + f = random.randrange(n - H) + return s0 + f, url, f + + def _decoder(self, url): + if url not in self._dec: + idx = mp4_ranges.read_index(lambda a, b: _read(url, a, b)) + self._dec[url] = mp4_ranges.Decoder(idx) + return self._dec[url] + + def _fetch(self, draws): + if self._pool is None: + _s3_client() + self._pool = cf.ThreadPoolExecutor(self.threads) + r = self.recs + t = time.perf_counter() + rows = np.concatenate([np.arange(s, s + H) for s, _, _ in draws]) + uniq, inv = np.unique(rows, return_inverse=True) + tbl = _table(r.steps_uri, r.steps_version).take(uniq.tolist(), columns=COLS) + cols = {c: _flat(tbl, c)[inv].reshape(len(draws), H, -1) for c in COLS} + self.timing["take"] += time.perf_counter() - t + out = [ + {c: torch.from_numpy(cols[c][i]) for c in COLS} for i in range(len(draws)) + ] + if not self.images: + return out + t = time.perf_counter() + urls = sorted({u for _, u, _ in draws}) + decs = dict(zip(urls, self._pool.map(self._decoder, urls))) + self.timing["index"] += time.perf_counter() - t + t = time.perf_counter() + jobs = [ + (u, rg) + for u in urls + for rg in mp4_ranges.ranges( + decs[u].index, [f for _, uu, f in draws if uu == u] + ) + ] + got = self._pool.map(lambda j: _read(j[0], *j[1]), jobs) + fetched: dict[str, dict] = {u: {} for u in urls} + for (u, rg), data in zip(jobs, got): + fetched[u][rg] = data + self.timing["get"] += time.perf_counter() - t + t = time.perf_counter() + for o, (_, u, f) in zip(out, draws): + o["img"] = torch.from_numpy(decs[u].frames([f], fetched[u])[0]) + self.timing["decode"] += time.perf_counter() - t + return out + + def __getitem__(self, _): + return self._fetch([self._draw()])[0] + + def __getitems__(self, idxs): + return self._fetch([self._draw() for _ in idxs]) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--lance-root", required=True) + ap.add_argument("--hashes", required=True) + ap.add_argument("--camera", default="front_1") + ap.add_argument("--images", action="store_true") + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=60) + ap.add_argument("--warmup", type=int, default=5) + ap.add_argument("--label", default="") + args = ap.parse_args() + + hashes = open(args.hashes).read().split() + recs = Recs(args.lance_root.rstrip("/"), hashes, args.camera) + print(json.dumps(run(args, RangeDS(recs, args.images, args.threads))), flush=True) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/bench_stream.py b/egomimic/scripts/lance_pilot/bench_stream.py new file mode 100644 index 0000000000..cf0c92b3e2 --- /dev/null +++ b/egomimic/scripts/lance_pilot/bench_stream.py @@ -0,0 +1,167 @@ +"""Numeric-only read throughput: random-access `take` windows vs Lance's streaming loader. + + take one 31-row window per sample at a random eligible row (our reader's access pattern) + stream lance.torch.data.LanceDataset batches in table order (sequential scan) + window the stream cut into 31-row windows within each recording, through a shuffle + buffer: what a streaming training loader would hand the model + +Frames are not decoded here; the MP4 path is the same for all three. +""" + +import argparse +import json +import random +import time + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset, IterableDataset, get_worker_info + +from egomimic.rldb.lance.dataset import _table, open_table, storage_options + +COLS = [ + "left.ee_pose", + "left.keypoints", + "left.wrist_pose", + "right.ee_pose", + "right.keypoints", + "right.wrist_pose", + "head_pose", +] +H = 31 + + +def flat(tbl, col): + if col in tbl.schema.names: + c = tbl.column(col) + c = c.combine_chunks() if hasattr(c, "combine_chunks") else c + return np.asarray(c.flatten()).reshape(tbl.num_rows, -1) + path = col.split(".") + c = tbl.column(path[0]) + c = c.combine_chunks() if hasattr(c, "combine_chunks") else c + for p in path[1:]: + c = c.field(p) + return np.asarray(c.flatten()).reshape(len(c), -1) + + +class TakeWindows(Dataset): + def __init__(self, uri, version, n_rows, n): + self.uri, self.version, self.n_rows, self.n = uri, version, n_rows, n + + def __len__(self): + return self.n + + def __getitem__(self, i): + s = random.randrange(self.n_rows - H) + t = _table(self.uri, self.version).take(list(range(s, s + H)), columns=COLS) + return {c: torch.from_numpy(flat(t, c)) for c in COLS} + + +class StreamWindows(IterableDataset): + """Scans fragments (split across workers), cuts windows inside a recording, shuffles.""" + + def __init__(self, uri, version, buffer): + self.uri, self.version, self.buffer = uri, version, buffer + + def __iter__(self): + ds = _table(self.uri, self.version) + frags = ds.get_fragments() + w = get_worker_info() + if w is not None: + frags = frags[w.id :: w.num_workers] + rng = random.Random(0 if w is None else w.id) + buf = [] + for frag in frags: + for batch in ds.to_batches( + columns=["recording_id", *COLS], fragments=[frag], batch_size=4096 + ): + rid = batch.column(0).to_numpy(zero_copy_only=False) + arrs = {c: flat(batch, c) for c in COLS} + for s in range(0, len(rid) - H, 1): + if rid[s] != rid[s + H - 1]: + continue + buf.append({c: torch.from_numpy(arrs[c][s : s + H]) for c in COLS}) + if len(buf) >= self.buffer: + yield buf.pop(rng.randrange(len(buf))) + rng.shuffle(buf) + yield from buf + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--uri", required=True, help="steps table, local path or s3://") + ap.add_argument("--mode", choices=["take", "stream", "window"], required=True) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=300) + ap.add_argument("--buffer", type=int, default=10000) + ap.add_argument("--label", default="") + args = ap.parse_args() + ds = open_table(args.uri) + n_rows, version = ds.count_rows(), ds.version + ctx = "forkserver" if args.workers else None + + t0 = time.perf_counter() + if args.mode == "stream": + from lance.torch.data import LanceDataset + + it = iter( + LanceDataset( + args.uri, + batch_size=4096, + # Its projection only takes top-level fields, so whole hand structs are read. + columns=["left", "right", "head_pose"], + dataset_options={"storage_options": storage_options(args.uri)}, + batch_readahead=16, + to_tensor_fn=lambda batch, **_: batch, + ) + ) + next(it) + first = time.perf_counter() - t0 + rows, t1 = 0, time.perf_counter() + for batch in it: + rows += batch.num_rows + if rows >= args.batches * args.batch_size * H: + break + rate = rows / (time.perf_counter() - t1) + out = dict(rows_per_s=round(rate), samples_per_s_equiv=round(rate / H, 1)) + else: + data = ( + TakeWindows(args.uri, version, n_rows, 10**9) + if args.mode == "take" + else StreamWindows(args.uri, version, args.buffer) + ) + loader = DataLoader( + data, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.workers, + multiprocessing_context=ctx, + persistent_workers=args.workers > 0, + prefetch_factor=4 if args.workers else None, + ) + it = iter(loader) + next(it) + first = time.perf_counter() - t0 + t1 = time.perf_counter() + for _ in range(args.batches): + next(it) + rate = args.batches * args.batch_size / (time.perf_counter() - t1) + out = dict(samples_per_s=round(rate, 1), rows_per_s=round(rate * H)) + print( + json.dumps( + dict( + label=args.label, + mode=args.mode, + uri=args.uri.split("/")[0] or "local", + workers=args.workers, + first_s=round(first, 2), + **out, + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/bench_zarr_vs_lance_r2.py b/egomimic/scripts/lance_pilot/bench_zarr_vs_lance_r2.py new file mode 100644 index 0000000000..f5dc07d266 --- /dev/null +++ b/egomimic/scripts/lance_pilot/bench_zarr_vs_lance_r2.py @@ -0,0 +1,271 @@ +"""Zarr vs Lance read throughput from R2 or local disk, numerics with or without one 480p frame per sample. + + take one random 31-frame window per sample (today's random-access loader pattern) + window each worker reads whole recordings (prefetching the next), cuts every 31-frame window, + and serves them through a shuffle buffer; frames stay compressed in memory until served + +zarr reads `processed_v3/aria/.zarr` (sharded zarr v3, JPEG frames); Lance reads the pilot +tables and faststart MP4s under `--lance-root`. Window mode times one full pass over all recordings. +""" + +import argparse +import concurrent.futures as cf +import json +import os +import random +import time +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset, IterableDataset, get_worker_info + +from egomimic.rldb.lance.dataset import _decoder, _s3_client, _table, open_table + +H = 31 +ZARR_KEYS = [ + f"{s}.obs_{k}" + for s in ("left", "right") + for k in ("ee_pose", "keypoints", "wrist_pose") +] + ["obs_head_pose"] +LANCE_COLS = [ + f"{s}.{k}" + for s in ("left", "right") + for k in ("ee_pose", "keypoints", "wrist_pose") +] + ["head_pose"] +ZARR_IMG = "images.front_1" + + +def _zarr_group(h, root): + import zarr + + if not root.startswith("s3://"): + return zarr.open_group(f"{root}/{h}", mode="r") + from obstore.store import S3Store + + store = S3Store( + "rldb", + prefix=f"{root[len('s3://rldb/'):].rstrip('/')}/{h}.zarr", + endpoint=os.environ["R2_ENDPOINT_URL"], + access_key_id=os.environ["R2_ACCESS_KEY_ID"], + secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"], + region="auto", + ) + return zarr.open_group( + store=zarr.storage.ObjectStore(store, read_only=True), mode="r" + ) + + +def _jpeg(b): + import simplejpeg + + return torch.from_numpy(simplejpeg.decode_jpeg(b, colorspace="RGB")) + + +def _flat(tbl, col): + if col in tbl.schema.names: + return np.asarray(tbl.column(col).combine_chunks().flatten()).reshape( + tbl.num_rows, -1 + ) + path = col.split(".") + c = tbl.column(path[0]).combine_chunks() + for p in path[1:]: + c = c.field(p) + return np.asarray(c.flatten()).reshape(len(c), -1) + + +class Recs: + """Per-recording metadata, built in the parent and pickled to workers.""" + + def __init__(self, fmt, hashes, lance_root, zarr_root, cam): + self.fmt, self.lance_root, self.zarr_root = fmt, lance_root, zarr_root + rec = open_table(f"{lance_root}/recordings.lance").to_table().to_pandas() + ali = open_table(f"{lance_root}/aliases.lance").to_table().to_pandas() + rec = rec.merge(ali, on="recording_id").set_index("episode_hash").loc[hashes] + steps = open_table(f"{lance_root}/steps/aria.lance") + self.steps_uri, self.steps_version = ( + f"{lance_root}/steps/aria.lance", + steps.version, + ) + ids = ( + steps.to_table(columns=["recording_id"]) + .column(0) + .to_numpy(zero_copy_only=False) + ) + change = np.flatnonzero(np.r_[True, ids[1:] != ids[:-1]]) + starts = dict(zip(ids[change], change)) + self.items = [] + for h, row in rec.iterrows(): + self.items.append( + dict( + hash=h, + row_start=int(starts[row.recording_id]), + n=int(row.num_frames), + video=f"{lance_root}/{dict(row.video_path)[cam]}", + ) + ) + + +class TakeDS(Dataset): + def __init__(self, recs, images): + self.recs, self.images = recs, images + self._z = {} + + def __len__(self): + return 10**9 + + def _zarr(self, h): + if h not in self._z: + g = _zarr_group(h, self.recs.zarr_root) + self._z[h] = (g, g[ZARR_KEYS[0]].shape[0]) + return self._z[h] + + def __getitem__(self, _): + it = random.choice(self.recs.items) + if self.recs.fmt == "zarr": + g, n = self._zarr(it["hash"]) + s = random.randrange(n - H) + out = {k: torch.from_numpy(np.asarray(g[k][s : s + H])) for k in ZARR_KEYS} + if self.images: + out["img"] = _jpeg(g[ZARR_IMG][s : s + 1][0]) + return out + s = random.randrange(it["n"] - H) + t = _table(self.recs.steps_uri, self.recs.steps_version).take( + list(range(it["row_start"] + s, it["row_start"] + s + H)), + columns=LANCE_COLS, + ) + out = {c: torch.from_numpy(_flat(t, c)) for c in LANCE_COLS} + if self.images: + out["img"] = _decoder(it["video"]).get_frame_at(s).data + return out + + +class WindowDS(IterableDataset): + def __init__(self, recs, images, buffer): + self.recs, self.images, self.buffer = recs, images, buffer + + def _load(self, it): + if self.recs.fmt == "zarr": + g = _zarr_group(it["hash"], self.recs.zarr_root) + arrs = {k: np.asarray(g[k][:]) for k in ZARR_KEYS} + frames = g[ZARR_IMG][:] if self.images else None + return arrs, (lambda i: _jpeg(frames[i])) if self.images else None + from torchcodec.decoders import VideoDecoder + + n = it["n"] + t = _table(self.recs.steps_uri, self.recs.steps_version).to_table( + columns=LANCE_COLS, offset=it["row_start"], limit=n + ) + arrs = {c: _flat(t, c) for c in LANCE_COLS} + if not self.images: + return arrs, None + if it["video"].startswith("s3://"): + bucket, key = it["video"][len("s3://") :].split("/", 1) + mp4 = _s3_client().get_object(Bucket=bucket, Key=key)["Body"].read() + else: + mp4 = Path(it["video"]).read_bytes() + dec = VideoDecoder(mp4, seek_mode="approximate", dimension_order="NHWC") + # torchcodec does not hold a reference to the bytes; freeing them segfaults the next decode. + return arrs, lambda i, _keep=mp4: dec.get_frame_at(i).data + + def __iter__(self): + w = get_worker_info() + items = self.recs.items[w.id :: w.num_workers] if w else list(self.recs.items) + rng = random.Random(0 if w is None else w.id) + rng.shuffle(items) + buf = [] + with cf.ThreadPoolExecutor(1) as pool: + nxt = pool.submit(self._load, items[0]) if items else None + for i in range(len(items)): + arrs, frame = nxt.result() + nxt = ( + pool.submit(self._load, items[i + 1]) + if i + 1 < len(items) + else None + ) + n = len(next(iter(arrs.values()))) + for s in range(n - H): + buf.append((arrs, frame, s)) + if len(buf) >= self.buffer: + yield self._emit(buf.pop(rng.randrange(len(buf)))) + rng.shuffle(buf) + for e in buf: + yield self._emit(e) + + @staticmethod + def _emit(e): + arrs, frame, s = e + out = {k: torch.from_numpy(a[s : s + H]) for k, a in arrs.items()} + if frame is not None: + out["img"] = frame(s) + return out + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--format", choices=["zarr", "lance"], required=True) + ap.add_argument("--mode", choices=["take", "window"], required=True) + ap.add_argument("--images", action="store_true") + ap.add_argument("--hashes", type=Path, required=True) + ap.add_argument("--lance-root", required=True) + ap.add_argument( + "--zarr-root", + default="s3://rldb/processed_v3/aria", + help="s3://bucket/prefix of .zarr, or a local dir of /", + ) + ap.add_argument( + "--camera", + default="front_1_480", + help="Lance video: front_1_480 or front_1 (360p)", + ) + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--batches", type=int, default=60, help="take mode only") + ap.add_argument("--buffer", type=int, default=3000) + ap.add_argument("--label", default="") + args = ap.parse_args() + hashes = [h for h in args.hashes.read_text().split() if h] + recs = Recs(args.format, hashes, args.lance_root, args.zarr_root, args.camera) + ds = ( + TakeDS(recs, args.images) + if args.mode == "take" + else WindowDS(recs, args.images, args.buffer) + ) + loader = DataLoader( + ds, + batch_size=args.batch_size, + num_workers=args.workers, + multiprocessing_context="forkserver" if args.workers else None, + prefetch_factor=4 if args.workers else None, + ) + t0 = time.perf_counter() + it = iter(loader) + next(it) + first = time.perf_counter() - t0 + t1, n = time.perf_counter(), 0 + for i, b in enumerate(it): + n += len(b["obs_head_pose" if args.format == "zarr" else "head_pose"]) + if args.mode == "take" and i + 1 >= args.batches: + break + dt = time.perf_counter() - t1 + print( + json.dumps( + dict( + label=args.label, + format=args.format, + camera=args.camera if args.format == "lance" else "jpeg480", + mode=args.mode, + images=args.images, + workers=args.workers, + first_s=round(first, 1), + samples=n, + seconds=round(dt, 1), + samples_per_s=round(n / dt, 1), + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/check_eva_samples.py b/egomimic/scripts/lance_pilot/check_eva_samples.py new file mode 100644 index 0000000000..7bad404228 --- /dev/null +++ b/egomimic/scripts/lance_pilot/check_eva_samples.py @@ -0,0 +1,112 @@ +"""Check that the Lance reader serves the same EVA training samples as zarr. + +eva_to_zarr keeps every frame, so zarr frame i is Lance frame i. Numeric outputs after the +full transform list must match to float32 precision at anchors both sample; images are +compared against the zarr JPEGs downscaled to the 360p stream. +""" + +import argparse +from pathlib import Path + +import cv2 +import numpy as np +import torch + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.lance.dataset import LanceEpisodeResolver +from egomimic.rldb.zarr.zarr_dataset_multi import S3EpisodeResolver + +CAMERAS = { + "images.front_1": "front_1", + "images.left_wrist": "left_wrist", + "images.right_wrist": "right_wrist", +} + + +def _hwc(img: torch.Tensor) -> np.ndarray: + a = img.float().numpy() + return a if a.shape[-1] == 3 else a.transpose(1, 2, 0) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--zarr-dir", + type=Path, + default=Path("/coc/flash7/scratch/lance_pilot/eva_zarr"), + ) + ap.add_argument( + "--lance-root", + type=Path, + default=Path("/coc/flash7/scratch/lance_pilot/lance_eva"), + ) + ap.add_argument("--hashes", type=Path, required=True) + ap.add_argument("--embodiment", default="eva_bimanual") + ap.add_argument("--per-recording", type=int, default=20) + args = ap.parse_args() + + key_map = Eva.get_keymap(keymap_mode="cartesian") + tl = Eva.get_transform_list(mode="cartesian_wristframe_6d") + if args.embodiment != "eva_bimanual": + # The Eva transforms are bimanual only: compare a one-arm robot's raw arrays. + absent = "left." if args.embodiment == "eva_right_arm" else "right." + key_map = { + k: v + for k, v in key_map.items() + if not v["zarr_key"].startswith(absent) and v["key_type"] != "camera_keys" + } + tl = None + f = DatasetFilter(episode_hashes=args.hashes.read_text().split()) + zl = S3EpisodeResolver(args.zarr_dir, key_map=key_map, transform_list=tl).resolve( + filters=f, expected_embodiment=args.embodiment + ) + lance_res = LanceEpisodeResolver( + args.lance_root, + lab="eva", + key_map=key_map, + transform_list=tl, + cameras=CAMERAS, + require_hands_in_frame="none", + ) + keys = lance_res.dataset_keys(zl) + ll = lance_res.resolve(filters=f, expected_embodiment=args.embodiment) + + rng = np.random.default_rng(0) + worst: dict[str, float] = {} + img_err: dict[str, list] = {} + n_checked = 0 + for h, Z in zl.items(): + L = ll[keys[h]] + if L.total_frames != Z.total_frames: + print(f"{h}: zarr {Z.total_frames} frames, Lance {L.total_frames}") + continue + for i in rng.choice( + len(Z), size=min(args.per_recording, len(Z)), replace=False + ): + a, b = Z[int(i)], L[int(i)] + for k, v in a.items(): + if not isinstance(v, torch.Tensor) or k not in b: + continue + if k in key_map and key_map[k]["key_type"] == "camera_keys": + za, lb = _hwc(v), _hwc(b[k]) + small = cv2.resize( + za, lb.shape[1::-1], interpolation=cv2.INTER_AREA + ) + img_err.setdefault(k, []).append(float(np.abs(small - lb).mean())) + continue + if not v.is_floating_point(): + continue + d = float((v - b[k]).abs().nan_to_num(nan=float("inf")).max()) + worst[k] = max(worst.get(k, 0.0), d) + n_checked += 1 + + print(f"recordings {len(zl)}; samples checked {n_checked}") + for k, d in sorted(worst.items(), key=lambda kv: -kv[1]): + print(f" {k:40s} max |diff| {d:.2e}") + for k, e in img_err.items(): + print(f" {k:40s} image mean |diff| {np.mean(e):.4f} max {np.max(e):.4f}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/check_samples.py b/egomimic/scripts/lance_pilot/check_samples.py new file mode 100644 index 0000000000..0ca29323ca --- /dev/null +++ b/egomimic/scripts/lance_pilot/check_samples.py @@ -0,0 +1,120 @@ +"""Check that the Lance reader serves the same training samples as zarr where both can. + +Comparable anchors are zarr frames whose whole action chunk is contiguous in time (the +zarr converter deleted untracked frames, so elsewhere its "next 30 frames" skip ahead). +Numeric outputs after the full transform list must match to float32 precision; images +are compared on the 480p stream against the JPEGs. +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +import zarr + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.lance.dataset import LanceEpisodeResolver +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver + +HORIZON = 30 +PERIOD_NS = 1e9 / 30 + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--zarr-dir", type=Path, default=Path("/coc/flash7/scratch/lance_pilot/zarr") + ) + ap.add_argument( + "--lance-root", type=Path, default=Path("/coc/flash7/scratch/lance_pilot/lance") + ) + ap.add_argument("--hashes", type=Path, required=True) + ap.add_argument("--per-recording", type=int, default=20) + args = ap.parse_args() + + key_map = Human.get_keymap(keymap_mode="keypoints", include_ee_pose=True) + tl = Human.get_transform_list( + mode="keypoints_wristframe_6d", + include_ee_pose=True, + pad_proprio_gripper=True, + stride=3, + ) + hashes = args.hashes.read_text().split() + f = DatasetFilter(episode_hashes=hashes) + zl = LocalEpisodeResolver( + args.zarr_dir, key_map=key_map, transform_list=tl + ).resolve(filters=f) + ll = LanceEpisodeResolver( + args.lance_root, + key_map=key_map, + transform_list=tl, + cameras={"images.front_1": "front_1_480"}, + ).resolve(filters=f) + import lance + + steps = lance.dataset(str(args.lance_root / "steps" / "aria.lance")) + rng = np.random.default_rng(0) + worst: dict[str, float] = {} + img_err, n_checked, n_comparable, n_zarr = [], 0, 0, 0 + for h in hashes: + Z, L = zl[h], ll[h] + zts = np.asarray( + zarr.open(str(args.zarr_dir / h), mode="r")["obs_rgb_timestamps_ns"][ + : Z.total_frames + ] + ) + lts = ( + steps.take( + list(range(L.row_start, L.row_start + L.total_frames)), + columns=["timestamp_ns"], + ) + .column(0) + .to_numpy() + ) + n = len(zts) + n_zarr += n + contiguous = np.zeros(n, bool) + gap_ok = np.abs(np.diff(zts) - PERIOD_NS) < 0.5 * PERIOD_NS + c = np.r_[0, np.cumsum(gap_ok)] + for i in range(n): + end = min(i + HORIZON, n) - 1 + contiguous[i] = c[end] - c[i] == end - i + lance_frame = np.searchsorted(lts, zts) + local = {L.anchor(j): j for j in range(len(L))} + cands = [i for i in np.flatnonzero(contiguous) if int(lance_frame[i]) in local] + n_comparable += len(cands) + for i in ( + rng.choice(cands, size=min(args.per_recording, len(cands)), replace=False) + if cands + else [] + ): + a, b = Z[int(i)], L[local[int(lance_frame[i])]] + for k, v in a.items(): + if ( + not isinstance(v, torch.Tensor) + or k not in b + or v.dtype != torch.float32 + ): + continue + if k.startswith("observations.images"): + img_err.append(float((v - b[k]).abs().mean())) + continue + d = float((v - b[k]).abs().max()) + worst[k] = max(worst.get(k, 0.0), d) + n_checked += 1 + + print( + f"zarr anchors {n_zarr:,}; comparable (contiguous chunk, eligible in Lance) {n_comparable:,}; checked {n_checked}" + ) + for k, d in sorted(worst.items(), key=lambda kv: -kv[1]): + print(f" {k:40s} max |diff| {d:.2e}") + if img_err: + print( + f" image mean |diff| (AV1 480p vs JPEG): mean {np.mean(img_err):.4f} max {np.max(img_err):.4f}" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/duplicate.py b/egomimic/scripts/lance_pilot/duplicate.py new file mode 100644 index 0000000000..b656a7cebe --- /dev/null +++ b/egomimic/scripts/lance_pilot/duplicate.py @@ -0,0 +1,118 @@ +"""Duplicate the pilot recordings K times in both formats for loader scale tests. + +Each copy gets a new recording_id / episode hash, so neither format can share index +entries or decoders between copies. `--link` hardlinks files (index-scale test, near-zero +disk, page-cache-resident); without it, files are real copies (IO-scale test). +""" + +import argparse +import os +import shutil +from pathlib import Path + +import lance +import pandas as pd +import pyarrow as pa +import pyarrow.compute as pc + +from egomimic.rldb.lance.schema import recording_id + + +def place(src: Path, dst: Path, link: bool): + dst.parent.mkdir(parents=True, exist_ok=True) + if dst.exists(): + return + if link: + os.link(src, dst) + else: + shutil.copy2(src, dst) + + +def place_tree(src: Path, dst: Path, link: bool): + if dst.exists(): + return + shutil.copytree(src, dst, copy_function=os.link if link else shutil.copy2) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--src", type=Path, default=Path("/coc/flash7/scratch/lance_pilot")) + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--factor", type=int, required=True) + ap.add_argument("--hashes", type=Path, required=True) + ap.add_argument("--link", action="store_true") + ap.add_argument("--lance-only", action="store_true") + args = ap.parse_args() + + hashes = set(args.hashes.read_text().split()) + src_l = args.src / "lance" + rec = lance.dataset(str(src_l / "recordings.lance")).to_table().to_pandas() + tasks = lance.dataset(str(src_l / "tasks.lance")).to_table().to_pandas() + aliases = lance.dataset(str(src_l / "aliases.lance")).to_table().to_pandas() + aliases = aliases[aliases.episode_hash.isin(hashes)] + rec = rec[rec.recording_id.isin(aliases.recording_id)] + steps_ds = lance.dataset(str(src_l / "steps" / "aria.lance")) + steps = steps_ds.to_table( + filter=pc.field("recording_id").isin(list(rec.recording_id)) + ) + by_rid = aliases.set_index("recording_id").episode_hash + + out_l = args.out / "lance" + new_rec, new_tasks, new_alias = [], [], [] + steps_uri = str(out_l / "steps" / "aria.lance") + for k in range(args.factor): + rid_map = { + r: recording_id("aria", f"dup{k:04d}/{v}") + for r, v in zip(rec.recording_id, rec.vendor_recording_id) + } + for r in rec.itertuples(): + h = by_rid[r.recording_id] + nh = f"{h}-dup{k:04d}" + nr = rid_map[r.recording_id] + vids = [] + for cam, rel in r.video_path: + nrel = f"videos/{nr}/{cam}.mp4" + place(src_l / rel, out_l / nrel, args.link) + vids.append((cam, nrel)) + new_rec.append(r._asdict() | {"recording_id": nr, "video_path": vids}) + if not args.lance_only: + place_tree(args.src / "zarr" / h, args.out / "zarr" / nh, args.link) + t = tasks[tasks.recording_id == r.recording_id].iloc[0].to_dict() + new_tasks.append( + t + | { + "recording_id": nr, + "task_instance_id": f"{t['task_instance_id']}-{k}", + } + ) + a = aliases[aliases.recording_id == r.recording_id].iloc[0].to_dict() + new_alias.append(a | {"recording_id": nr, "episode_hash": nh}) + ids = pc.cast(steps.column("recording_id"), pa.string()).to_pylist() + copy = steps.set_column(0, "recording_id", pa.array([rid_map[i] for i in ids])) + lance.write_dataset( + copy, + steps_uri, + mode="append" if k else "overwrite", + max_rows_per_file=4 << 20, + ) + print(f"copy {k + 1}/{args.factor}", flush=True) + + for name, rows, src_name in ( + ("recordings", new_rec, "recordings"), + ("tasks", new_tasks, "tasks"), + ("aliases", new_alias, "aliases"), + ): + df = pd.DataFrame(rows).drop(columns=["Index"], errors="ignore") + schema = lance.dataset(str(src_l / f"{src_name}.lance")).schema + lance.write_dataset( + pa.Table.from_pandas(df[schema.names], schema=schema, preserve_index=False), + str(out_l / f"{name}.lance"), + mode="overwrite", + ) + print( + f"{len(new_rec)} recordings, {steps.num_rows * args.factor:,} frames -> {args.out}" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/report.py b/egomimic/scripts/lance_pilot/report.py new file mode 100644 index 0000000000..a945c87b2f --- /dev/null +++ b/egomimic/scripts/lance_pilot/report.py @@ -0,0 +1,175 @@ +"""Migration report: what each Aria delivery keeps, drops and changes relative to today's zarr.""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import zarr + + +def dir_bytes(p: Path) -> int: + return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) + + +def column(t, side, field=None): + c = t.column(side).combine_chunks() + if field is not None: + c = c.field(field) + return np.asarray(c.flatten()).reshape(t.num_rows, -1) + + +def quat_angle_deg(a, b): + dot = np.abs((a * b).sum(1)).clip(0, 1) + return np.degrees(2 * np.arccos(dot)) + + +def compare(d: Path, zdir: Path, emb: str) -> dict: + meta = json.loads((d / "meta.json").read_text()) + t = pq.read_table(d / "steps.parquet") + z = zarr.open(str(zdir), mode="r") + nz = int(z.attrs["total_frames"]) + ts = t.column("timestamp_ns").to_numpy() + zts = np.asarray(z["obs_rgb_timestamps_ns"][:nz]) + idx = np.searchsorted(ts, zts).clip(0, len(ts) - 1) + found = ts[idx] == zts + + lv, rv = t.column("left_valid").to_numpy(), t.column("right_valid").to_numpy() + hv = t.column("head_pose_valid").to_numpy() + sides = { + "human_bimanual": ("left", "right"), + "human_left_arm": ("left",), + "human_right_arm": ("right",), + }[emb] + any_hand = np.zeros(len(ts), bool) + for s in sides: + any_hand |= {"left": lv, "right": rv}[s] + in_zarr = np.zeros(len(ts), bool) + in_zarr[idx[found]] = True + + r = dict( + vendor_recording_id=meta["vendor_recording_id"], + embodiment=emb, + zarr_frames=nz, + zarr_frames_missing_in_lance=int((~found).sum()), + lance_frames=meta["num_frames"], + filler_frames=meta["filler_frames"], + gaps=meta["gaps"], + graph_changes=meta["graph_changes"], + left_valid=float(lv.mean()), + right_valid=float(rv.mean()), + head_valid=float(hv.mean()), + any_hand_valid=float(any_hand.mean()), + valid_not_in_zarr=int((any_hand & ~in_zarr).sum()), + zarr_kept_now_invalid=int((in_zarr & ~any_hand).sum()), + bytes_zarr=dir_bytes(zdir), + bytes_mp4_360=(d / "front_1.mp4").stat().st_size, + bytes_mp4_480=(d / "front_1_480.mp4").stat().st_size, + bytes_parquet=(d / "steps.parquet").stat().st_size, + seconds=meta["seconds_total"], + ) + zi = idx[found] + zrow = np.flatnonzero(found) + head = column(t, "head_pose")[zi] + ok = hv[zi] + r["head_pose_maxdiff"] = float( + np.abs(head[ok] - np.asarray(z["obs_head_pose"][:nz])[zrow][ok]).max(initial=0) + ) + for s in sides: + v = {"left": lv, "right": rv}[s][zi] + if not v.any(): + continue + zr = lambda k: np.asarray(z[f"{s}.{k}"][:nz])[zrow][v] # noqa: E731 + r[f"{s}_wrist_maxdiff"] = float( + np.abs(column(t, s, "wrist_pose")[zi][v] - zr("obs_wrist_pose")).max() + ) + r[f"{s}_aria_kp_maxdiff"] = float( + np.abs( + column(t, s, "aria_keypoints")[zi][v] - zr("obs_aria_keypoints") + ).max() + ) + kp_err = np.linalg.norm( + (column(t, s, "keypoints")[zi][v] - zr("obs_keypoints")).reshape(-1, 21, 3), + axis=-1, + ) + r[f"{s}_mano_kp_p99_mm"] = float(np.percentile(kp_err, 99) * 1e3) + ee, zee = column(t, s, "ee_pose")[zi][v], zr("obs_ee_pose") + r[f"{s}_ee_pos_p99_mm"] = float( + np.percentile(np.linalg.norm(ee[:, :3] - zee[:, :3], axis=1), 99) * 1e3 + ) + r[f"{s}_ee_rot_p99_deg"] = float( + np.percentile(quat_angle_deg(ee[:, 3:], zee[:, 3:]), 99) + ) + return r + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--root", type=Path, default=Path("/coc/flash7/scratch/lance_pilot") + ) + args = ap.parse_args() + m = pd.read_csv(args.root / "manifest.csv", dtype={"epoch_ms": str}) + rows, missing = [], [] + for rec in m.itertuples(): + d = args.root / "deliveries" / rec.epoch_ms + if not (d / "meta.json").exists(): + missing.append(rec.epoch_ms) + continue + rows.append( + compare(d, args.root / "zarr" / rec.episode_hash, rec.embodiment) + | {"site": rec.lab, "task": rec.task} + ) + df = pd.DataFrame(rows) + df.to_csv(args.root / "migration_report.csv", index=False) + + tot = df[ + [ + "zarr_frames", + "lance_frames", + "filler_frames", + "valid_not_in_zarr", + "zarr_kept_now_invalid", + ] + ].sum() + any_valid = (df.any_hand_valid * df.lance_frames).sum() + print(f"recordings {len(df)} (not converted: {missing or 'none'})") + print( + f"frames: zarr {tot.zarr_frames:,} lance {tot.lance_frames:,} ({tot.lance_frames / tot.zarr_frames:.2f}x) filler {tot.filler_frames:,}" + ) + print( + f"frames with a valid hand: {any_valid:,.0f} ({any_valid / tot.zarr_frames:.2f}x zarr)" + ) + print( + f"valid hand frames the zarr dropped: {tot.valid_not_in_zarr:,}; zarr frames now invalid: {tot.zarr_kept_now_invalid:,}" + ) + print( + f"recordings with gaps: {(df.gaps > 0).sum()}, with SLAM graph changes: {(df.graph_changes > 0).sum()}" + ) + print(f"zarr frames missing from lance: {df.zarr_frames_missing_in_lance.sum()}") + agree = [c for c in df.columns if c.endswith(("maxdiff", "_p99_mm", "_p99_deg"))] + print("agreement (max over recordings):") + print(df[agree].max().to_string()) + zb = df.bytes_zarr.sum() / tot.zarr_frames + lf = tot.lance_frames + print( + f"bytes/frame: zarr {zb / 1e3:.1f} KB | mp4 360 {df.bytes_mp4_360.sum() / lf / 1e3:.1f} KB" + f" | mp4 480 {df.bytes_mp4_480.sum() / lf / 1e3:.1f} KB | parquet {df.bytes_parquet.sum() / lf / 1e3:.2f} KB" + ) + print( + f"conversion: median {df.seconds.median():.0f} s/recording, total {df.seconds.sum() / 3600:.1f} GPU-h" + ) + print( + df.groupby("embodiment")[ + ["left_valid", "right_valid", "head_valid", "any_hand_valid"] + ] + .mean() + .round(3) + .to_string() + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/select_and_sync.py b/egomimic/scripts/lance_pilot/select_and_sync.py new file mode 100644 index 0000000000..a77b06b1e2 --- /dev/null +++ b/egomimic/scripts/lance_pilot/select_and_sync.py @@ -0,0 +1,145 @@ +"""Pick the Aria pilot recordings and sync their raw captures and current zarrs from R2.""" + +import argparse +import os +import subprocess +from pathlib import Path + +import pandas as pd +from sqlalchemy import text + +from egomimic.utils.aws.aws_sql import ( + create_default_engine, + episode_hash_to_timestamp_ms, +) + +RAW_PREFIX = "s3://rldb/raw_v2/aria" +# Semidense points and the open-loop trajectory are ~70% of MPS bytes and unused. +MPS_FILES = [ + "hand_tracking/*", + "eye_gaze/*", + "slam/closed_loop_trajectory.csv", + "slam/online_calibration.jsonl", + "slam/summary.json", + "vrs_health_check*.json", +] +ALWAYS = [ + "2025-09-20-17-47-54-000000", + "2025-10-29-22-06-00-411000", + "2025-11-24-20-24-32-080000", +] +SKIP_TASKS = {"debug", "interview"} + + +def load_rows() -> pd.DataFrame: + q = text( + """select episode_hash, operator, lab, embodiment, scene, task, num_frames, + zarr_processed_path + from app.episodes + where zarr_processed_path like :p and not is_deleted + order by episode_hash""" + ) + with create_default_engine().connect() as c: + df = pd.read_sql(q, c, params={"p": "%/aria/%"}) + df["epoch_ms"] = df.episode_hash.map(episode_hash_to_timestamp_ms) + return df + + +def select(df: pd.DataFrame, n: int, seed: int) -> pd.DataFrame: + """Round-robin over (site, task, embodiment) strata, preferring typical lengths.""" + lo, hi = df.num_frames.quantile([0.1, 0.9]) + pool = df[ + df.num_frames.between(lo, hi) + & ~df.episode_hash.isin(ALWAYS) + & ~df.task.isin(SKIP_TASKS) + ] + pool = pool.sample(frac=1.0, random_state=seed) + groups = [g for _, g in pool.groupby(["lab", "task", "embodiment"])] + groups.sort(key=len, reverse=True) + picked = list(df[df.episode_hash.isin(ALWAYS)].index) + i = 0 + while len(picked) < n and any(len(g) > i for g in groups): + for g in groups: + if len(picked) >= n: + break + if len(g) > i: + picked.append(g.index[i]) + i += 1 + return df.loc[picked].sort_values("episode_hash").reset_index(drop=True) + + +def s5cmd_lines(sel: pd.DataFrame, out: Path) -> list[str]: + lines = [] + for r in sel.itertuples(): + ms = r.epoch_ms + raw = out / "raw" + for name in (f"{ms}.vrs", f"{ms}.json", f"{ms}_metadata.json"): + lines.append(f'cp -n "{RAW_PREFIX}/{name}" "{raw}/{name}"') + for f in MPS_FILES: + sub = os.path.dirname(f) + dst = raw / f"mps_{ms}_vrs" / sub + lines.append(f'cp -n "{RAW_PREFIX}/mps_{ms}_vrs/{f}" "{dst}/"') + src = r.zarr_processed_path.rstrip("/") + if not src.startswith("s3://"): + src = f"s3://rldb/{src.lstrip('/')}" + lines.append(f'sync "{src}/*" "{out / "zarr" / r.episode_hash}/"') + return lines + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", type=Path, default=Path("/coc/flash7/scratch/lance_pilot")) + ap.add_argument("--n", type=int, default=50) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--numworkers", type=int, default=64) + ap.add_argument("--dry-run", action="store_true") + ap.add_argument( + "--shard", + type=int, + default=0, + help="download only recordings i with i %% nshards == shard", + ) + ap.add_argument("--nshards", type=int, default=1) + args = ap.parse_args() + + manifest = args.out / "manifest.csv" + # The manifest is the pilot set; select only once so reruns and shards agree. + if manifest.exists(): + sel = pd.read_csv(manifest) + else: + sel = select(load_rows(), args.n, args.seed) + args.out.mkdir(parents=True, exist_ok=True) + sel.to_csv(manifest, index=False) + print(sel.groupby(["lab", "embodiment"]).size().to_string()) + print( + f"{len(sel)} recordings, {sel.num_frames.sum():,} zarr frames, {sel.task.nunique()} tasks" + ) + + batch = args.out / f"s5cmd_batch_{args.shard}_of_{args.nshards}.txt" + shard = sel.iloc[args.shard :: args.nshards] + batch.write_text("\n".join(s5cmd_lines(shard, args.out)) + "\n") + if args.dry_run: + return + env = os.environ.copy() + env.update( + AWS_ACCESS_KEY_ID=env["R2_ACCESS_KEY_ID"], + AWS_SECRET_ACCESS_KEY=env["R2_SECRET_ACCESS_KEY"], + AWS_REGION="auto", + AWS_DEFAULT_REGION="auto", + ) + # R2 rejects the session token that ~/.egoverse_env also exports. + env.pop("AWS_SESSION_TOKEN", None) + cmd = [ + "s5cmd", + "--endpoint-url", + env["R2_ENDPOINT_URL"], + "--numworkers", + str(args.numworkers), + "run", + str(batch), + ] + subprocess.run(cmd, check=True, env=env) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/mcap_process/__init__.py b/egomimic/scripts/mcap_process/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/mcap_process/mcap_to_lance.py b/egomimic/scripts/mcap_process/mcap_to_lance.py new file mode 100644 index 0000000000..d8db8fc4bc --- /dev/null +++ b/egomimic/scripts/mcap_process/mcap_to_lance.py @@ -0,0 +1,973 @@ +"""Convert vendor MCAP episodes (spec 0.1, egomimic/rldb/mcap) into the Lance tables. + + python -m egomimic.scripts.mcap_process.mcap_to_lance --stage-to \ + --lance episode_/ ... + +Per episode: validate (any problem rejects it), check its labels against the tasks it +will have, build one steps row per frame of the steps-clock camera, write a video per +camera, and stage the recording, its tasks and its labels for the committer +(egomimic.scripts.lance_commit, the only writer of the tables) with +commit.stage_delivery. Everything comes from the MCAP itself: the episode info replaces +the SQL lookup of the Aria pilot. `--lance` is only read, to skip recordings already +committed; their labels are staged as an annotation change instead. + +The per-frame math follows aria_to_lance: nearest sample within MATCH_TOL_NS of each +frame (a true nearest search; projectaria's helper can land one sample off), filler rows +for dropped frames, head pose as the camera pose turned upright, hands in the world +frame, MANO keypoints from the Aria fit, validity from tracking, timing and pose jumps. +Timestamps stay absolute Unix ns, as delivered. Human embodiments only: Lance has no +robot steps schema yet. +""" + +from __future__ import annotations + +import argparse +import datetime +import json +import subprocess +import sys +import tempfile +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from scipy.spatial.transform import Rotation + +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.schema import ( + ANNOTATIONS, + HUMAN_STEPS, + annotation_id, + task_instance_id, +) +from egomimic.rldb.lance.timeline import build_timeline, pose_outliers +from egomimic.rldb.mcap import schemas, spec +from egomimic.rldb.mcap.validate import frames_in_message, validate_episode +from egomimic.utils.aws.aws_sql import timestamp_ms_to_episode_hash + +MATCH_TOL_NS = 50_000_000 # as aria_to_lance: further than this from a frame = missing +IDENTITY_POSE = np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float32) # xyz + wxyz +# The steps camera's pinhole outputs (h, w), as aria_to_lance. +FRONT = {"front_1": (360, 480), "front_1_480": (480, 640)} +# Aria's output focal at 480 rows (ARIA_INTRINSICS). Other sources keep their own field +# of view: their native focal, scaled to the output. +ARIA_FOCAL_480 = 133.25430222 * 2 +# Converter-side conventions shared with the Aria path (aria_utils). +T_ROT_CAM = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) +ENCODER_OPTS = { + "libsvtav1": ["-preset", "12", "-crf", "30"], # aria_to_lance's choice + "libaom-av1": ["-cpu-used", "8", "-crf", "30", "-b:v", "0"], + "libx264": ["-preset", "veryfast", "-crf", "23"], +} +# The only topics the conversion reads besides the cameras and their calibrations. +NEEDED = { + "/tf_static", + "/head-pose", + "/head-pose-quality", + "/left-hand-keypoints", + "/right-hand-keypoints", + "/eye-gaze", + "/task", + "/subtask-annotation", +} + + +class Rejected(Exception): + """A validated episode Lance would still refuse (e.g. labels outside its tasks).""" + + def __init__(self, problems: list[str]): + super().__init__("; ".join(problems)) + self.problems = problems + + +# ---------------------------------------------------------------- reading + + +@dataclass +class _Series: + t: np.ndarray # log times, Unix ns, sorted + msgs: list + + +@dataclass +class _Camera: + topic: str + t: np.ndarray # frame times, Unix ns, sorted + fmt: str # h264 / h265 + + +def _read(paths: list[Path]): + """The needed topics decoded, every camera's frame times and codec, and the episode + info. Cameras and info come from episode.mcap (paths[0]) only, as validated.""" + from mcap.reader import make_reader + from mcap_protobuf.decoder import DecoderFactory + + topics: dict[str, list] = defaultdict(list) + cameras: dict[str, _Camera] = {} + meta: dict[str, str] = {} + for k, path in enumerate(paths): + with open(path, "rb") as f: + r = make_reader(f, decoder_factories=[DecoderFactory()]) + summary = r.get_summary() + if summary is not None: + names = {c.topic for c in summary.channels.values()} + else: # unindexed file: one scan finds the topics + names = {ch.topic for _, ch, _ in r.iter_messages()} + wanted = sorted( + t + for t in names + if t in NEEDED + or t.endswith((spec.CALIBRATION_SUFFIX, spec.CAMERA_EXTRA_SUFFIX)) + ) + if wanted: + for _, ch, msg, dec in r.iter_decoded_messages( + topics=wanted, log_time_order=True + ): + topics[ch.topic].append((msg.log_time, dec)) + if k: + continue + meta = {m.name: dict(m.metadata) for m in r.iter_metadata()}.get( + spec.METADATA_RECORD, {} + ) + cams = sorted(t for t in names if t.endswith(spec.CAMERA_SUFFIX)) + times: dict[str, list[int]] = defaultdict(list) + for _, ch, msg in r.iter_messages(topics=cams) if cams else (): + times[ch.topic].append(msg.log_time) + for cam in cams: + _, _, _, first = next(r.iter_decoded_messages(topics=[cam])) + cameras[cam] = _Camera( + cam, np.sort(np.array(times[cam], dtype=np.int64)), first.format + ) + series = { + k: _Series(np.array([t for t, _ in v], dtype=np.int64), [m for _, m in v]) + for k, v in topics.items() + } + return series, cameras, meta + + +def _packets(path: Path, topic: str): + """One camera's frames as bytes, streamed in log-time order.""" + from mcap.reader import make_reader + from mcap_protobuf.decoder import DecoderFactory + + with open(path, "rb") as f: + r = make_reader(f, decoder_factories=[DecoderFactory()]) + for _, _, _, dec in r.iter_decoded_messages( + topics=[topic], log_time_order=True + ): + yield bytes(dec.data) + + +def _mat(position, orientation) -> np.ndarray: + T = np.eye(4) + o = orientation + T[:3, :3] = Rotation.from_quat([o.x, o.y, o.z, o.w]).as_matrix() + T[:3, 3] = (position.x, position.y, position.z) + return T + + +def _xyz_wxyz(T: np.ndarray) -> np.ndarray: + q = Rotation.from_matrix(T[:3, :3]).as_quat() # xyzw + return np.r_[T[:3, 3], q[3], q[:3]] + + +class _Frames: + """World poses of named frames: /tf_static plus the moving world -> head link.""" + + def __init__(self, series: dict[str, _Series]): + self.edges: dict[str, dict[str, np.ndarray]] = defaultdict(dict) + self._cache: dict[tuple[str, str], np.ndarray | None] = {} + for msg in series["/tf_static"].msgs: + for x in msg.transforms: + T = _mat(x.translation, x.rotation) + self.edges[x.parent_frame_id][x.child_frame_id] = T + self.edges[x.child_frame_id][x.parent_frame_id] = np.linalg.inv(T) + + def static(self, root: str, frame: str) -> np.ndarray | None: + """T_root_frame through /tf_static only, or None if unconnected (cached).""" + if (root, frame) not in self._cache: + self._cache[root, frame] = self._search(root, frame) + return self._cache[root, frame] + + def on_head(self, frame: str) -> bool: + """Whether placing `frame` in the world needs the head pose.""" + return ( + frame != spec.WORLD_FRAME + and self.static(spec.HEAD_FRAME, frame) is not None + ) + + def _search(self, root: str, frame: str) -> np.ndarray | None: + todo, seen = [(root, np.eye(4))], {root} + while todo: + name, T = todo.pop() + if name == frame: + return T + for nxt, step in self.edges[name].items(): + if nxt not in seen: + seen.add(nxt) + todo.append((nxt, T @ step)) + return None + + def world(self, frame: str, T_world_head: np.ndarray) -> np.ndarray: + """T_world_frame at a time whose head pose is T_world_head.""" + if frame == spec.WORLD_FRAME: + return np.eye(4) + T = self.static(spec.HEAD_FRAME, frame) + if T is not None: + return T_world_head @ T + T = self.static(spec.WORLD_FRAME, frame) + if T is None: + raise ValueError(f"frame {frame!r} is not linked to world or head") + return T + + +def _nearest(times: np.ndarray, t: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Index of the nearest sample to each time, and whether it is within MATCH_TOL_NS + (the index form of timeline.nearest_dt).""" + last = len(times) - 1 + j = np.clip(np.searchsorted(times, t), 0, last) + prev = np.clip(j - 1, 0, last) + pick = np.where(np.abs(times[prev] - t) <= np.abs(times[j] - t), prev, j) + return pick, np.abs(times[pick] - t) <= MATCH_TOL_NS + + +# ---------------------------------------------------------------- hands + + +def _aria_slots() -> list[str]: + """Joint name per slot of Aria's 21-landmark layout, which the MANO fit + (aria_utils.convert_aria_keypoints_to_mano) takes as input.""" + from egomimic.scripts.aria_process.aria_utils import ( + ARIA_TO_OTAHERI, + OTAHERI_TO_CANONICAL, + ) + + names = ["palm_center"] * 21 # the one slot the fit ignores + for aria, otaheri in ARIA_TO_OTAHERI.items(): + names[aria] = spec.MANO_JOINTS[OTAHERI_TO_CANONICAL.index(otaheri)] + return names + + +def _reorder(kp: np.ndarray, joints: list[str], slots: list[str]) -> np.ndarray: + """(T, J*3) in `joints` order -> (T, len(slots)*3) in `slots` order; absent joints NaN.""" + out = np.full((len(kp), len(slots), 3), np.nan) + src = kp.reshape(len(kp), len(joints), 3) + for i, name in enumerate(slots): + if name in joints: + out[:, i] = src[:, joints.index(name)] + return out.reshape(len(kp), -1) + + +def _hand( + series: _Series, + cam_t: np.ndarray, + T_world_head: list[np.ndarray], + frames: _Frames, + joints: list[str], + head_ok: np.ndarray, + *, + is_right: bool, + mano: bool, + mano_iters: int, +) -> dict[str, np.ndarray]: + """Per camera frame: wrist pose, ee pose, MANO keypoints, Aria keypoints, confidence, ok.""" + from egomimic.scripts.aria_process.aria_utils import ( + convert_aria_keypoints_to_mano, + mano_keypoints_to_cartesian, + ) + + n = len(cam_t) + pick, close = _nearest(series.t, cam_t) + raw = np.full((n, len(joints) * 3), np.nan) + wrist = np.full((n, 7), np.nan) + conf = np.full(n, -1.0, dtype=np.float32) + tracked = np.zeros(n, dtype=bool) + for i, (p, ok) in enumerate(zip(pick, close)): + m = series.msgs[p] + if not (ok and m.valid): + continue + if frames.on_head(m.frame_id) and not head_ok[i]: + continue # placed with a stale or lost head pose (aria_to_lance: pose_close) + T = frames.world(m.frame_id, T_world_head[i]) + conf[i] = m.confidence + tracked[i] = True + if len(m.positions): + pts = np.asarray(m.positions, dtype=np.float64).reshape(-1, 3) + raw[i] = (T[:3, :3] @ pts.T + T[:3, 3:4]).T.ravel() + if m.HasField("wrist_pose"): + Tw = T @ _mat(m.wrist_pose.position, m.wrist_pose.orientation) + Tw[:3, :3] = Tw[:3, :3] @ T_ROT_CAM + wrist[i] = _xyz_wxyz(Tw) + + has_kp = np.isfinite(raw).all(1) + aria_layout = joints == _aria_slots() + # MANO order by joint name; joints the device lacks stay NaN. + canonical = _reorder(raw, joints, list(spec.MANO_JOINTS)) + if tuple(joints) == spec.MANO_JOINTS: + keypoints = raw # already MANO: no fit + elif mano and has_kp.any(): + aria = _reorder(raw, joints, _aria_slots()) + aria = np.where(np.isfinite(aria) & (tracked & has_kp)[:, None], aria, 1e9) + keypoints = convert_aria_keypoints_to_mano( + aria, is_rhand=is_right, n_iters=mano_iters, lr=0.05 + ) + elif aria_layout: # --no-mano on Aria data, as aria_to_lance --no-mano + keypoints = raw + else: + keypoints = canonical + # The ee pose needs wrist + 4 MCPs: from the fit, or the vendor's joints by name. + ee_src = keypoints if (mano or tuple(joints) == spec.MANO_JOINTS) else canonical + ee = mano_keypoints_to_cartesian(np.nan_to_num(ee_src, nan=1e9), is_rhand=is_right) + ee_ok = (np.abs(ee) < 1e8).all(1) + + wrist_ok = np.isfinite(wrist).all(1) + if "wrist" in joints: # no wrist pose sent: the wrist joint, oriented as the ee + w = raw[:, 3 * joints.index("wrist") : 3 * joints.index("wrist") + 3] + fill = ~wrist_ok & ee_ok + wrist[fill] = np.c_[w[fill], ee[fill, 3:]] + no_kp = tracked & ~has_kp # wrist-only device: the wrist pose is the ee + ee[no_kp & wrist_ok] = wrist[no_kp & wrist_ok] + ee_ok |= no_kp & wrist_ok + + kp_ok = np.isfinite(keypoints).all(1) | no_kp + ok = tracked & np.isfinite(wrist).all(1) & kp_ok & ee_ok + ok &= ~pose_outliers(np.nan_to_num(wrist[:, :3]), ok) + return dict( + wrist=wrist, + ee=ee, + kp=keypoints, + # Lance's aria_keypoints holds Aria landmarks only, never another layout. + aria=raw if aria_layout else np.zeros((n, 63)), + conf=conf, + ok=ok, + ) + + +# ---------------------------------------------------------------- video + + +def _rot_z(deg_cw: int) -> np.ndarray: + """Camera-frame rotation matching a clockwise image turn (90 -> aria_utils.ROTATION_MATRIX).""" + return Rotation.from_euler("z", -deg_cw, degrees=True).as_matrix() + + +def _upright(extra) -> int: + return int(extra.upright_rotation_cw_deg) if extra else 0 + + +def _undistorter(calib, extra, size_hw: tuple[int, int], focal: float): + """raw frame -> upright pinhole frame of `size_hw` with `focal`, centred.""" + import cv2 + + h, w = size_hw + rot = _upright(extra) + pre_h, pre_w = (w, h) if rot % 180 else (h, w) # size before turning upright + model = (calib.distortion_model or "plumb_bob").lower() + k_rot = (-rot // 90) % 4 + + if model == "fisheye624": + import projectaria_tools.core.sophus as sp + from projectaria_tools.core import calibration + + params = np.array([calib.K[0], calib.K[2], calib.K[5], *calib.D]) + radius = float(extra.valid_radius) if extra and extra.valid_radius else None + src = calibration.CameraCalibration( + "camera", + calibration.CameraModelType.FISHEYE624, + params, + sp.SE3.from_matrix(np.eye(4)), + int(calib.width), + int(calib.height), + radius, + np.pi, + "", + ) + dst = calibration.get_linear_camera_calibration(pre_w, pre_h, focal, "camera") + + def remap(raw): + return calibration.distort_by_calibration(raw, dst, src) + + else: + K = np.asarray(calib.K, dtype=np.float64).reshape(3, 3) + D = np.asarray(calib.D, dtype=np.float64) + K_new = np.array( + [[focal, 0, (pre_w - 1) / 2], [0, focal, (pre_h - 1) / 2], [0, 0, 1]] + ) + if model == "kannala_brandt": + m1, m2 = cv2.fisheye.initUndistortRectifyMap( + K, D[:4], np.eye(3), K_new, (pre_w, pre_h), cv2.CV_32FC1 + ) + elif model in ("plumb_bob", "rational_polynomial"): + m1, m2 = cv2.initUndistortRectifyMap( + K, D, np.eye(3), K_new, (pre_w, pre_h), cv2.CV_32FC1 + ) + else: + raise ValueError(f"unsupported distortion_model {model!r}") + + def remap(raw): + return cv2.remap(raw, m1, m2, cv2.INTER_LINEAR) + + def f(raw): + return np.ascontiguousarray(np.rot90(remap(raw), k=k_rot)) + + return f + + +def _encoder(path: Path, h: int, w: int, ffmpeg: str, codec: str, fps: float): + cmd = [ + ffmpeg, "-hide_banner", "-loglevel", "error", "-y", + "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", f"{fps:g}", "-i", "-", + "-c:v", codec, "-g", "2", *ENCODER_OPTS.get(codec, []), "-pix_fmt", "yuv420p", + "-movflags", "+faststart", + str(path), + ] # fmt: skip + return subprocess.Popen(cmd, stdin=subprocess.PIPE) + + +def _write_camera(packets, cam: _Camera, rows, outputs, render, ffmpeg, codec, fps): + """Decode one camera once, streaming, and write one frame per steps row: rows[i] is + the camera frame shown on row i (non-decreasing; repeats for filler rows). + `render(raw)` gives one frame per (path, (h, w)) in `outputs`. Every message must be + exactly one frame: the validator checks a sample, this checks all of them.""" + import av + + ctx = av.CodecContext.create("hevc" if cam.fmt == "h265" else "h264", "r") + seen = 0 + + def frames(): + nonlocal seen + for data in packets: + seen += 1 + units = frames_in_message(data, cam.fmt) + if units != 1: + raise ValueError( + f"{cam.topic} message {seen} holds {units} frames, not 1" + ) + yield from ( + f.to_ndarray(format="rgb24") for f in ctx.decode(av.Packet(data)) + ) + yield from (f.to_ndarray(format="rgb24") for f in ctx.decode(None)) + + enc = [_encoder(p, h, w, ffmpeg, codec, fps) for p, (h, w) in outputs] + try: + decoded, have, shown = frames(), 0, None + for want in rows: + while have <= want: + raw = next(decoded, None) + if raw is None: + raise ValueError(f"{cam.topic} ends after {have} frames") + shown = render(raw) + have += 1 + for p, img in zip(enc, shown): + p.stdin.write(img.tobytes()) + total = have + sum(1 for _ in decoded) + if total != len(cam.t): + raise ValueError( + f"{cam.topic}: {len(cam.t)} messages decoded to {total} frames" + ) + for p in enc: + p.stdin.close() + if p.wait() != 0: + raise RuntimeError(f"ffmpeg ({codec}) failed") + finally: + for p in enc: + if p.poll() is None: + p.kill() + p.wait() + + +def _video_name(topic: str) -> str: + """/left-wrist-camera -> left_wrist (the zarr images.* names).""" + return topic.strip("/").removesuffix(spec.CAMERA_SUFFIX).replace("-", "_") + + +# ---------------------------------------------------------------- labels + + +def _label_rows(series, meta, rid: str, period: float) -> list[dict]: + """/subtask-annotation as Lance annotation rows; a point becomes one frame wide.""" + rows = [] + if "/subtask-annotation" not in series: + return rows + s = series["/subtask-annotation"] + for t0, m in zip(s.t, s.msgs): + row = dict( + op="add", + recording_id=rid, + t_start_ns=int(t0), + t_end_ns=int(m.end_time_unix_ns) or int(t0 + period), + annotation_type=m.annotation_type or "language", + text=m.data, + annotation_vendor=m.annotation_vendor or meta["vendor"], + quality_score=None, + graded_by=None, + rubric_version=None, + ) + row["annotation_id"] = annotation_id( + rid, row["annotation_vendor"], row["annotation_type"], + row["t_start_ns"], row["t_end_ns"], row["text"], + ) # fmt: skip + rows.append(row) + return rows + + +def _label_errors(rows: list[dict], spans: list[tuple[int, int]]) -> list[str]: + """Lance's own annotation review rule (annotations._add_errors) against these tasks.""" + from egomimic.rldb.lance.annotations import _add_errors + + if not rows: + return [] + tasks = {rows[0]["recording_id"]: spans} + return [ + f"label {r['text']!r} [{r['t_start_ns']}, {r['t_end_ns']}): {e}" + for r in rows + for e in _add_errors(r, tasks) + ] + + +def _label_changes(root, rid: str, rows: list[dict], vendors: set[str]) -> list[dict]: + """Change rows bringing a committed recording's labels to `rows`: adds, and deletes of + the vendors' committed labels not among them (a delivery's labels are authoritative + for its vendors).""" + uri = commit.annotations_uri(root) + new_ids = {r["annotation_id"] for r in rows} + stale = [] + if commit._exists(uri): + old = commit._open(uri).to_table( + columns=["annotation_id", "annotation_vendor"], + filter=f"recording_id = {commit._sql_str(rid)}", + ) + stale = [ + r["annotation_id"] + for r in old.to_pylist() + if r["annotation_vendor"] in vendors and r["annotation_id"] not in new_ids + ] + return [{**r, "op": "add"} for r in rows] + [ + {"op": "delete", "annotation_id": a} for a in stale + ] + + +# ---------------------------------------------------------------- convert + + +@dataclass +class Result: + """What one episode came to; what a watcher or dashboard records.""" + + episode: str + recording_id: str = "" + lab: str = "" + status: str = "" # rejected | failed | staged | skipped + problems: list[str] = field(default_factory=list) # why it was rejected + error: str = "" # why a validated episode could not be converted + num_frames: int = 0 + filler_frames: int = 0 + tasks: int = 0 + videos: list[str] = field(default_factory=list) + annotations_staged: int = 0 # label rows added or deleted, staged for the committer + valid_frac: dict[str, float] = field(default_factory=dict) + + +@dataclass +class Converted: + delivery: commit.Delivery + labels: list[dict] + vendor: str + stats: dict + + +def convert( + episode: Path, + out_dir: Path, + *, + mano: bool = True, + mano_iters: int = 2400, + ffmpeg: str = "ffmpeg", + codec: str = "libsvtav1", +) -> Converted: + """One validated episode -> its Lance rows and MP4s (under out_dir). Raises Rejected + when Lance would refuse its labels, before any video is written.""" + episode = Path(episode) + files = [episode / spec.EPISODE_FILE] + if (episode / spec.ANNOTATION_FILE).exists(): + files.append(episode / spec.ANNOTATION_FILE) + series, cameras, meta = _read(files) + embodiment = meta["embodiment"] + if spec.profile_for(embodiment) is not spec.HUMAN: + raise NotImplementedError( + "robot recordings: Lance has no robot steps schema yet" + ) + step_clock = spec.HUMAN.step_clock + cam_t = cameras[step_clock].t + period = float(np.median(np.diff(cam_t))) + fps = round(1e9 / period, 3) + ts, src, filler, after_gap = build_timeline(cam_t, fps=fps) + n = len(ts) + rid = spec.recording_id(meta["lab"], meta["vendor-recording-id"]) + + # Tasks and labels first: a label Lance would refuse rejects the episode up front. + if "/task" in series: + tasks = [ + (int(t0), int(m.end_time_unix_ns), m.name, m.description, list(m.objects)) + for t0, m in zip(series["/task"].t, series["/task"].msgs) + ] + else: # the episode info over the whole recording, half-open + objects = [o for o in meta.get("objects", "").split(",") if o] + tasks = [ + ( + int(ts[0]), + int(ts[-1] + period), + meta["task"], + meta["instruction"], + objects, + ) + ] + labels = _label_rows(series, meta, rid, period) + errors = _label_errors(labels, [(t0, t1) for t0, t1, *_ in tasks]) + if errors: + raise Rejected(errors) + + # Head: pose nearest each camera frame; SLAM resets break validity like gaps. + frames = _Frames(series) + hp = series["/head-pose"] + pick, pose_close = _nearest(hp.t, cam_t) + T_world_head = [_mat(hp.msgs[p].pose.position, hp.msgs[p].pose.orientation) for p in pick] # fmt: skip + calib = series[step_clock + spec.CALIBRATION_SUFFIX].msgs[0] + extra_s = series.get(step_clock + spec.CAMERA_EXTRA_SUFFIX) + extra = extra_s.msgs[0] if extra_s else None + R4 = np.eye(4) + R4[:3, :3] = _rot_z(_upright(extra)) + T_head_front = frames.static(spec.HEAD_FRAME, calib.frame_id) @ R4 + head = np.array([_xyz_wxyz(T @ T_head_front) for T in T_world_head]) + q = series["/head-pose-quality"] + qi, q_close = _nearest(q.t, cam_t) + state = np.array([q.msgs[i].state for i in qi]) + Q = schemas.TrackingQuality + head_ok = pose_close & q_close & (state != Q.LOST) + head_usable = head_ok.copy() # what hands placed via the head need + head_ok &= ~pose_outliers(head[:, :3], head_ok) + # A relocalization between a frame and the one before it breaks that frame. + reloc = q.t[[m.state == Q.RELOCALIZED for m in q.msgs]] + seen = np.searchsorted(reloc, cam_t, side="right") + reset = np.r_[False, seen[1:] > seen[:-1]] + brk = after_gap | (reset[src] & ~filler) + head_valid = head_ok[src] & ~filler & ~brk + + def col(x, valid, fill): + x = np.asarray(x, dtype=np.float32)[src] + x[~valid] = fill + x = np.nan_to_num(x) + return pa.FixedSizeListArray.from_arrays( + pa.array(x.reshape(-1)), x.shape[1] + ).cast(pa.list_(pa.float32(), x.shape[1])) + + cols = { + "recording_id": pa.array([rid] * n), + "timestamp_ns": pa.array(ts), + "head_pose": col(head, head_valid, IDENTITY_POSE), + } + joints = meta["hand-joints"].split(",") + valid = {} + for side in ("left", "right"): + topic = f"/{side}-hand-keypoints" + if topic in series: + h = _hand( + series[topic], cam_t, T_world_head, frames, joints, head_usable, + is_right=side == "right", mano=mano, mano_iters=mano_iters, + ) # fmt: skip + else: # single-arm embodiment + z7, z63 = np.zeros((len(cam_t), 7)), np.zeros((len(cam_t), 63)) + h = dict( + wrist=z7, ee=z7, kp=z63, aria=z63, + conf=np.full(len(cam_t), -1.0), ok=np.zeros(len(cam_t), dtype=bool), + ) # fmt: skip + v = h["ok"][src] & ~filler & ~brk + valid[side] = v + cols[side] = pa.StructArray.from_arrays( + [ + col(h["wrist"], v, IDENTITY_POSE), + col(h["ee"], v, IDENTITY_POSE), + col(h["kp"], v, 0.0), + col(h["aria"], v, 0.0), + pa.array(np.where(filler, -1.0, h["conf"][src]).astype(np.float32)), + ], + fields=list(HUMAN_STEPS.field(side).type), + ) + cols["left_valid"] = pa.array(valid["left"]) + cols["right_valid"] = pa.array(valid["right"]) + cols["head_pose_valid"] = pa.array(head_valid) + eye = np.zeros((len(cam_t), 3)) + eye_ok = np.zeros(len(cam_t), dtype=bool) + if "/eye-gaze" in series: + g = series["/eye-gaze"] + gi, g_close = _nearest(g.t, cam_t) + for i, (k, close) in enumerate(zip(gi, g_close)): + m = g.msgs[k] + eye[i] = (m.yaw, m.pitch, m.depth) + eye_ok[i] = close and m.valid and np.isfinite(eye[i]).all() + cols["eye_gaze"] = col(eye, eye_ok[src] & ~filler, 0.0) + steps = pa.table(cols, schema=HUMAN_STEPS) + + # Videos: the steps camera as the pinhole front outputs, every other camera as its + # own video aligned to the steps rows. Extrinsics are relative to head_pose (the + # front camera, upright). + import cv2 + + out = Path(out_dir) / rid + out.mkdir(parents=True, exist_ok=True) + intrinsics, extrinsics, videos = {}, {}, {} + big_h, big_w = FRONT["front_1_480"] + pre_w = big_h if _upright(extra) % 180 else big_w + aria = meta["lab"] == "aria" + focal = ARIA_FOCAL_480 if aria else float(calib.K[0]) * pre_w / float(calib.width) + undistort = _undistorter(calib, extra, (big_h, big_w), focal) + small_h, small_w = FRONT["front_1"] + + def render_front(raw): + big = undistort(raw) + return big, cv2.resize(big, (small_w, small_h), interpolation=cv2.INTER_AREA) + + outputs = [(out / f"{k}.mp4", FRONT[k]) for k in ("front_1_480", "front_1")] + _write_camera( + _packets(files[0], step_clock), cameras[step_clock], src, outputs, + render_front, ffmpeg, codec, fps, + ) # fmt: skip + for k, (vh, vw) in FRONT.items(): + f = focal * vh / big_h + # Aria keeps aria_to_lance's ARIA_INTRINSICS (centre at w/2, h/2). + cx, cy = (vw / 2, vh / 2) if aria else ((vw - 1) / 2, (vh - 1) / 2) + intrinsics[k] = [f, 0.0, cx, 0.0, f, cy, 0.0, 0.0, 1.0] + extrinsics[k] = np.eye(4).ravel().tolist() + videos[k] = out / f"{k}.mp4" + + for topic, cam in sorted(cameras.items()): + if topic == step_clock: + continue + name = _video_name(topic) + c = series[topic + spec.CALIBRATION_SUFFIX].msgs[0] + x_s = series.get(topic + spec.CAMERA_EXTRA_SUFFIX) + x = x_s.msgs[0] if x_s else None + h, w = ( + (int(c.width), int(c.height)) + if _upright(x) % 180 + else (int(c.height), int(c.width)) + ) + f = float(c.K[0]) # native size and focal: the camera keeps its field of view + rows, _ = _nearest(cam.t, ts) + render = _undistorter(c, x, (h, w), f) + _write_camera( + _packets(files[0], topic), cam, rows, [(out / f"{name}.mp4", (h, w))], + lambda raw, r=render: (r(raw),), ffmpeg, codec, fps, + ) # fmt: skip + Rc = np.eye(4) + Rc[:3, :3] = _rot_z(_upright(x)) + T_head_cam = frames.static(spec.HEAD_FRAME, c.frame_id) + intrinsics[name] = [f, 0.0, (w - 1) / 2, 0.0, f, (h - 1) / 2, 0.0, 0.0, 1.0] + if T_head_cam is not None: # a camera fixed in the world has no head extrinsic + extrinsics[name] = ( + (np.linalg.inv(T_head_front) @ T_head_cam @ Rc).ravel().tolist() + ) + videos[name] = out / f"{name}.mp4" + + # Catalog rows. + start_ms = int(meta["start-time-unix"]) + recording = dict( + recording_id=rid, + vendor_recording_id=meta["vendor-recording-id"], + lab=meta["lab"], + site=meta["site"], + embodiment=embodiment, + rig_name=meta["rig-name"], + operator=meta["operator-id"], + scene=meta["scene"], + captured_at=datetime.datetime.fromtimestamp( + start_ms / 1000, tz=datetime.timezone.utc + ), + fps=fps, + num_frames=n, + intrinsics=sorted(intrinsics.items()), + extrinsics=sorted(extrinsics.items()), + license=meta["license"], + ) + task_rows = [ + dict( + task_instance_id=task_instance_id(rid, name, t0), + recording_id=rid, + t_start_ns=t0, + t_end_ns=t1, + task_name=name, + task_description=desc, + objects=objs, + ) + for t0, t1, name, desc, objs in tasks + ] + alias = dict( + episode_hash=timestamp_ms_to_episode_hash(start_ms), + recording_id=rid, + t_start_ns=int(ts[0]), + t_end_ns=int(ts[-1] + period), # half-open: the last frame is inside + ) + stats = dict( + filler_frames=int(filler.sum()), + valid_frac={ + "head": float(head_valid.mean()), + "left": float(valid["left"].mean()), + "right": float(valid["right"].mean()), + }, + ) + delivery = commit.Delivery(recording, task_rows, alias, steps, videos) + return Converted(delivery, labels, meta["vendor"], stats) + + +def _labels_of_committed(episode: Path, root: Path, rid: str) -> tuple[list, str]: + """Label rows of an episode whose recording is already committed, with errors + against its committed tasks.""" + files = [episode / spec.EPISODE_FILE] + if (episode / spec.ANNOTATION_FILE).exists(): + files.append(episode / spec.ANNOTATION_FILE) + series, cameras, meta = _read(files) + period = float(np.median(np.diff(cameras[spec.HUMAN.step_clock].t))) + rows = _label_rows(series, meta, rid, period) + spans = [ + (t["t_start_ns"], t["t_end_ns"]) + for t in commit._open(commit.join(root, "tasks.lance")) + .to_table(filter=f"recording_id = {commit._sql_str(rid)}") + .to_pylist() + ] + errors = _label_errors(rows, spans) + if errors: + raise Rejected(errors) + return rows, meta["vendor"] + + +def convert_and_stage( + episode: Path, + staging, + lance_root, + *, + replace: bool = False, + work_dir: Path | None = None, + provenance: dict | None = None, + **convert_kw, +) -> Result: + """Validate, convert and stage one episode; safe to re-run. A recording already in + `lance_root` is skipped unless `replace` (the committer then needs --replace too), but + its labels are staged as an annotation change (labels delivered later).""" + from egomimic.scripts.lance_annotations import stage as stage_change + + res = Result(episode=str(episode)) + try: + v = validate_episode(Path(episode)) + except Exception as e: + res.status, res.error = "failed", f"{type(e).__name__}: {e}" + return res + res.recording_id = v.recording_id + if v.problems: + res.status = "rejected" + res.problems = [str(p) for p in v.problems] + return res + if v.recording_id in commit.committed_ids(lance_root) and not replace: + res.status = "skipped" + try: + rows, vendor = _labels_of_committed( + Path(episode), lance_root, v.recording_id + ) + except Rejected as e: + res.problems = e.problems + return res + vendors = {vendor, *(r["annotation_vendor"] for r in rows)} + changes = _label_changes(lance_root, v.recording_id, rows, vendors) + if changes: + with tempfile.TemporaryDirectory(dir=work_dir) as tmp: + f = Path(tmp) / f"{v.recording_id}-labels.parquet" + keys = list(dict.fromkeys(k for r in changes for k in r)) + pq.write_table( + pa.Table.from_pylist( + [{k: r.get(k) for k in keys} for r in changes] + ), + f, + ) + stage_change(f, staging, f"{vendor} delivery") + res.annotations_staged = len(changes) + return res + with tempfile.TemporaryDirectory(dir=work_dir) as tmp: + try: + c = convert(Path(episode), Path(tmp), **convert_kw) + except Rejected as e: # validated, but Lance would refuse it + res.status, res.problems = "rejected", e.problems + return res + except Exception as e: # unsupported embodiment / lens, broken video, ... + res.status, res.error = "failed", f"{type(e).__name__}: {e}" + return res + d = c.delivery + res.lab = d.recording["lab"] + labels = pa.Table.from_pylist( + [{k: r.get(k) for k in ANNOTATIONS.names} for r in c.labels], + schema=ANNOTATIONS, + ) + meta = { + "label_vendors": sorted( + {c.vendor, *(r["annotation_vendor"] for r in c.labels)} + ), + "provenance": provenance, + } + commit.stage_delivery( + d, + commit.join(staging, res.lab, d.recording["vendor_recording_id"]), + labels=labels, + meta=meta, + ) + res.status = "staged" + res.num_frames = d.recording["num_frames"] + res.tasks = len(d.tasks) + res.videos = sorted(d.videos) + res.filler_frames = c.stats["filler_frames"] + res.valid_frac = c.stats["valid_frac"] + res.annotations_staged = len(c.labels) + return res + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("episodes", nargs="+", type=Path, help="episode_/ dirs") + ap.add_argument( + "--stage-to", required=True, help="staging root (dir or s3:// prefix)" + ) + ap.add_argument( + "--lance", required=True, help="Lance root, read to skip committed recordings" + ) + ap.add_argument( + "--replace", action="store_true", help="restage committed 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("--codec", default="libsvtav1", choices=sorted(ENCODER_OPTS)) + ap.add_argument("--work-dir", type=Path, help="scratch for videos (default: tmp)") + args = ap.parse_args(argv) + + results = [] + for ep in args.episodes: + r = convert_and_stage( + ep, + args.stage_to, + args.lance, + replace=args.replace, + work_dir=args.work_dir, + mano=args.mano, + mano_iters=args.mano_iters, + ffmpeg=args.ffmpeg, + codec=args.codec, + ) + results.append(r) + print(json.dumps(r.__dict__)) + # Rejected / failed episodes and refused labels all need a look. + bad = [r for r in results if r.status in ("rejected", "failed") or r.problems] + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/egomimic/scripts/mecka_process/scan_egomotion.py b/egomimic/scripts/mecka_process/scan_egomotion.py new file mode 100644 index 0000000000..d5a0b9c2a3 --- /dev/null +++ b/egomimic/scripts/mecka_process/scan_egomotion.py @@ -0,0 +1,118 @@ +"""Scan local episodes' ``obs_head_pose`` (Mecka egomotion) for failed tracking. + +Writes a per-episode CSV and, with ``--blocklist``, the hashes whose status is +``constant`` (one pose for the whole episode: no tracking) or ``jump`` (a +single-frame head move over ``--jump-m``), in the format +``DatasetFilter(exclude_hashes=...)`` reads. + + python egomimic/scripts/mecka_process/scan_egomotion.py \\ + --dataset-dir $EGOVERSE_DATASET_DIR --csv scan.csv \\ + --blocklist egomimic/resources/blocklists/mecka_bad_egomotion.txt +""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +from collections import Counter +from datetime import date +from functools import partial +from multiprocessing import Pool +from pathlib import Path + +import numpy as np +import zarr + +BAD = ("constant", "jump") +FIELDS = [ + "episode", + "status", + "embodiment", + "frames", + "unique_quat", + "unique_xyz", + "xyz_std", + "max_step_m", + "nonfinite", +] + + +def scan_episode(root: Path, episode: str, jump_m: float) -> dict: + path = root / episode + row = {"episode": episode} + try: + attrs = json.loads((path / "zarr.json").read_text()).get("attributes", {}) + row["embodiment"] = attrs.get("embodiment") + if not (path / "obs_head_pose").is_dir(): + return {**row, "status": "no_head_pose"} + pose = np.asarray(zarr.open(str(path), mode="r")["obs_head_pose"][:], float) + except Exception as exc: + return {**row, "status": f"error:{type(exc).__name__}"} + xyz, quat = pose[:, :3], pose[:, 3:] + step = np.linalg.norm(np.diff(xyz, axis=0), axis=1) if len(pose) > 1 else [0.0] + row.update( + frames=len(pose), + unique_quat=len(np.unique(quat, axis=0)), + unique_xyz=len(np.unique(xyz, axis=0)), + xyz_std=float(xyz.std(0).max()), + max_step_m=float(np.max(step)), + nonfinite=int((~np.isfinite(pose)).any(1).sum()), + ) + if row["unique_quat"] == 1 and row["unique_xyz"] == 1: + status = "constant" + elif row["unique_quat"] == 1: + status = "constant_rot" + elif row["unique_xyz"] == 1: + status = "constant_xyz" + elif row["nonfinite"]: + status = "nonfinite" + elif row["max_step_m"] > jump_m: + status = "jump" + else: + status = "ok" + return {**row, "status": status} + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--dataset-dir", required=True, type=Path) + ap.add_argument("--csv", required=True, type=Path) + ap.add_argument("--blocklist", type=Path) + # 30 cm in one frame is 9 m/s at 30 fps: a tracking jump, not a head move. + ap.add_argument("--jump-m", type=float, default=0.3) + ap.add_argument("--workers", type=int, default=os.cpu_count()) + args = ap.parse_args() + + episodes = sorted( + d.name for d in args.dataset_dir.iterdir() if (d / "zarr.json").is_file() + ) + scan = partial(scan_episode, args.dataset_dir, jump_m=args.jump_m) + with Pool(args.workers) as pool: + rows = sorted( + pool.imap_unordered(scan, episodes, chunksize=32), + key=lambda r: r["episode"], + ) + + with args.csv.open("w", newline="") as f: + writer = csv.DictWriter(f, FIELDS) + writer.writeheader() + writer.writerows(rows) + print(Counter(r["status"] for r in rows)) + + if args.blocklist: + bad = [r for r in rows if r["status"] in BAD] + header = [ + "# Episodes with failed head tracking (obs_head_pose / vendor egomotion),", + "# from egomimic/scripts/mecka_process/scan_egomotion.py over the", + f"# {len(rows)} episodes in the local dataset dir on {date.today()}.", + "# constant: one pose for the whole episode; jump: a single-frame head", + f"# move over {args.jump_m} m. {len(bad)} episodes.", + ] + lines = [f"{r['episode']} # {r['status']}" for r in bad] + args.blocklist.write_text("\n".join(header + lines) + "\n") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/mecka_process/viz_mecka_episode.py b/egomimic/scripts/mecka_process/viz_mecka_episode.py new file mode 100644 index 0000000000..f117f7ce80 --- /dev/null +++ b/egomimic/scripts/mecka_process/viz_mecka_episode.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python +"""Visualize a Mecka (egocentric human) zarr episode -> mp4 with frame overlays. + +Egocentric analog of abc_process/viz_eva_episode.py --draw-axes: projects into +the moving head camera (per-frame ``obs_head_pose`` is the world->cam target, +exactly the headframe transform training uses) with MECKA_INTRINSICS. + +Overlays, per frame: + - L/R wrist triads (X=red Y=green Z=blue) from ``obs_ee_pose`` — shows the + hand-tracking wrist-frame convention (diagonal to hand geometry, mirrored + between hands, unlike EVA/YAM tool frames). + - hand keypoints (gray dots) + a WHITE line wrist-root -> middle-MCP = + "knuckle-forward", the physical reference to compare the triads against. + - the WORLD origin triad (thick, length 0.3m): origin sits on the floor under + the operator's head at episode start, Z up (gravity-aligned). + + ./emimic/bin/python egomimic/scripts/mecka_process/viz_mecka_episode.py \ + --episode 696d0f031244464a4aac1b8f --still-frames 150 450 750 + +--episode accepts a bare hash (resolved under --folder, which defaults to the +dataset dir from hydra_configs/paths/default.yaml: $EGOVERSE_DATASET_DIR or +/coc/flash7/scratch/egoverseS3ZarrDataset) or a full .zarr path. +""" + +import argparse +import io +import os + +import numpy as np +import zarr + +from egomimic.utils.paths import default_dataset_dir + +DEFAULT_DATASET_DIR = default_dataset_dir() + + +def decode(elem): + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB +# MediaPipe-style indices (verified on-data: kp0 roots all four finger chains, +# adjacent MCPs 5/9/13/17 are ~2.5cm apart). NOT Aria ordering. +WRIST_ROOT, MIDDLE_MCP = 0, 9 + + +class MeckaOverlay: + """World-frame geometry -> head-camera pixels via the training headframe math.""" + + 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.transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _xyzwxyz_to_matrix + + self.K = MECKA_INTRINSICS + self._to_mat = _xyzwxyz_to_matrix + self._t_pose = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self._t_xyz = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyz", + ) + + def cam_pose_matrix(self, head_pose, world_pose): + out = self._t_pose.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(world_pose, np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, np.float64)[None])[0] + + def cam_xyz(self, head_pose, world_xyz): + """(N,3) world points -> (N,3) cam-frame points.""" + return np.stack( + [ + self._t_xyz.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(p, np.float64), + } + )["out"] + for p in np.asarray(world_xyz, np.float64).reshape(-1, 3) + ] + ) + + def project(self, pts_cam): + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M, length, label, thickness=2): + import cv2 + + o = M[:3, 3] + pts = np.stack([o] + [o + M[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.4, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, img, head, lpose, rpose, lkp, rkp): + import cv2 + + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = np.ascontiguousarray(img.copy()) + # world origin triad (floor under head-at-start, Z up) + if np.linalg.norm(np.asarray(head)[3:7]) > 1e-6: + Mw = self.cam_pose_matrix(head, np.array([0, 0, 0, 1, 0, 0, 0.0])) + self.draw_triad(img, Mw, 0.3, "world", thickness=3) + for pose, kp, lab in ((lpose, lkp, "L-wrist"), (rpose, rkp, "R-wrist")): + if ( + np.linalg.norm(np.asarray(pose)[3:7]) < 1e-6 + or np.abs(kp).sum() < 1e-9 + ): + continue + # keypoints + knuckle-forward reference line + pc = self.cam_xyz(head, kp) + px = self.project(pc) + for p in px: + if not np.isnan(p).any(): + cv2.circle( + img, tuple(np.round(p).astype(int)), 2, (200, 200, 200), -1 + ) + a, b = px[WRIST_ROOT], px[MIDDLE_MCP] + if not (np.isnan(a).any() or np.isnan(b).any()): + cv2.line( + img, + tuple(np.round(a).astype(int)), + tuple(np.round(b).astype(int)), + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + # the wrist-frame triad the pipeline actually uses + M = self.cam_pose_matrix(head, pose) + self.draw_triad(img, M, 0.07, lab) + cv2.putText( + img, + "axes X=red Y=green Z=blue | white line = knuckle-forward", + (6, img.shape[0] - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True, help="episode hash or full .zarr path") + ap.add_argument( + "--folder", + default=DEFAULT_DATASET_DIR, + help="zarr dataset dir a bare --episode hash is resolved under " + "(default: $EGOVERSE_DATASET_DIR or the lab dataset dir)", + ) + ap.add_argument( + "--out", default=None, help="output mp4 (default: ./mecka__axes.mp4)" + ) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + ap.add_argument( + "--still-frames", + type=int, + nargs="*", + default=[], + help="also dump these frame indices as overlaid PNGs next to --out", + ) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + print(f"[viz] {path}\n task={(a.get('task_description') or '')[:80]!r} fps={fps}") + + imgs = g["images.front_1"] + head = np.asarray(g["obs_head_pose"]) + lp, rp = np.asarray(g["left.obs_ee_pose"]), np.asarray(g["right.obs_ee_pose"]) + lk = np.asarray(g["left.obs_keypoints"]).reshape(len(head), 21, 3) + rk = np.asarray(g["right.obs_keypoints"]).reshape(len(head), 21, 3) + nf = min(imgs.shape[0], len(head), args.max_frames or 10**9) + out = args.out or f"mecka_{os.path.basename(path).replace('.zarr', '')}_axes.mp4" + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + ov = MeckaOverlay() + import imageio.v2 as imageio + from PIL import Image, ImageDraw + + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + frame = ov.overlay(decode(imgs[i]), head[i], lp[i], rp[i], lk[i], rk[i]) + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still) + print(f"[viz] still -> {still}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py new file mode 100644 index 0000000000..d3edf831b8 --- /dev/null +++ b/egomimic/scripts/precompute_norm_stats.py @@ -0,0 +1,209 @@ +"""Precompute normalization stats on a CPU node, so GPU-node time isn't spent +computing them at training startup. + +Replicates trainHydra's norm loop EXACTLY (same hydra config + data config, +same keymap/transform, same sample_frac/seed) but builds NO model and NO +trainer: it only instantiates the train datasets (which s5cmd-syncs any +missing episodes from S3 as a side effect), infers shapes, computes norm +stats, and stores them twice: + + 1. an explicit file, ``/norm_stats/norm_stats.json``, for + ``norm_stats.precomputed_norm_path=/norm_stats``; + 2. the content-keyed cache under ``norm_stats.cache_dir`` (default + ``paths.cache_dir``), keyed exactly as trainHydra keys it, so a plain + training run on the same episodes + recipe hits the cache without any + override. + +The per-episode samples behind it also land in the cache (``episodes/``), so a +later split of the same recipe (a subset, a superset, an operator hold-out) +only samples the episodes it adds. + +The norm-mode keymap strips camera + annotation keys, so the stats pass reads +only the numeric proprio/action arrays: pure CPU work, no GPU, no JPEG decode +beyond one shape-inference sample. + +Usage (CPU node, repo root, emimic venv): + python egomimic/scripts/precompute_norm_stats.py \\ + --data mecka_all_6d --model pi0.5_bc_mecka_6d \\ + --sample-frac 0.1 --num-workers 30 --out /path/to/norm_stats/mecka_all_6d +""" + +from __future__ import annotations + +import argparse +import copy +import os +import time + +import hydra +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +import egomimic +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.norm import episode_norm_samples, norm_cache +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.utils.env import load_env + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--config-name", default="train_zarr_cartesian_pi") + ap.add_argument( + "--data", + required=True, + help="data config group; MUST match the training run's data config", + ) + ap.add_argument( + "--model", + required=True, + help="only needed so the config composes; the model is never built", + ) + ap.add_argument("--sample-frac", type=float, default=0.1) + ap.add_argument( + "--max-samples", + type=int, + default=None, + help="optional hard cap on collected samples: the (N, 100, D) float32 " + "action stack plus np.percentile's sort copy is ~2 x N x 100 x D x 4 " + "bytes of RAM, so an uncapped 0.1 frac of the full mecka set (~8.5M " + "frames) needs >100GB at D=18. Part of the cache key: a training run " + "hits this entry only with the same norm_stats.max_samples.", + ) + ap.add_argument("--num-workers", type=int, default=30) + ap.add_argument( + "--out", + required=True, + help="save_cache_dir; writes /norm_stats/norm_stats.json", + ) + ap.add_argument( + "--no-cache", + action="store_true", + help="only write the explicit file; skip the content-keyed cache entry", + ) + ap.add_argument( + "overrides", + nargs="*", + help="extra hydra overrides appended verbatim (e.g. " + "paths.dataset_dir=/path/to/zarr/mirror)", + ) + args = ap.parse_args() + + cfg_dir = os.path.join(os.path.dirname(egomimic.__file__), "hydra_configs") + GlobalHydra.instance().clear() + overrides = [ + f"data={args.data}", + f"model={args.model}", + f"norm_stats.sample_frac={args.sample_frac}", + f"norm_stats.max_samples={'null' if args.max_samples is None else args.max_samples}", + f"norm_stats.num_workers={args.num_workers}", + f"norm_stats.save_cache_dir={args.out}", + "norm_stats.precomputed_norm_path=null", + "seed=42", + *args.overrides, + ] + with initialize_config_dir(version_base=None, config_dir=cfg_dir): + cfg = compose(config_name=args.config_name, overrides=overrides) + + import lightning as L + + L.seed_everything(cfg.seed, workers=True) + load_env() + + norm_mode = OmegaConf.select(cfg, "norm_stats.norm_mode", default="quantile") + pool_horizon = bool(OmegaConf.select(cfg, "norm_stats.pool_horizon", default=False)) + cache_dir = None if args.no_cache else OmegaConf.select(cfg, "norm_stats.cache_dir") + + # Mirrors trainHydra: instantiate train datasets (resolver syncs from S3 + # as needed), then a stats-only MultiDataset computes the norm stats from + # a norm-mode (numerics-only) copy of each dataset. + train_datasets = {} + for dataset_name in cfg.data.train_datasets: + if cfg.data.train_datasets[dataset_name] is None: + continue + print(f"[precompute] dataset={dataset_name}: instantiating (syncs S3) ...") + train_datasets[dataset_name] = hydra.utils.instantiate( + cfg.data.train_datasets[dataset_name], dataset_name=dataset_name + ) + + norm_stats = MultiDataset(state={}, norm_mode=norm_mode) + norm_stats.populate_from_datasets(train_datasets) + + for dataset_name, dataset in train_datasets.items(): + print(f"[precompute] dataset={dataset_name}: inferring shapes ...") + norm_stats.infer_shapes_from_batch(dataset[0]) + + inst = copy.deepcopy(cfg.data.train_datasets[dataset_name]) + km = OmegaConf.to_container(inst.resolver.key_map, resolve=False) + km["norm_mode"] = True # strips image + annotation keys + if "proprio_history" in km: + # Current-step proprio stats even on a history (K > 1) config: every + # history step is normalized with the CURRENT-step stats, so a K > 1 + # run and the K = 1 baseline share one norm-stat file. Reading one + # frame is exactly `[..., -1, :]` of the window, minus the K-fold + # read. + km["proprio_history"] = 1 + inst.resolver.key_map = km + norm_dataset = hydra.utils.instantiate(inst, dataset_name=dataset_name) + + t0 = time.perf_counter() + norm_stats.infer_norm_from_dataset( + norm_dataset, + dataset_name, + sample_frac=args.sample_frac, + max_samples=args.max_samples, + num_workers=args.num_workers, + precomputed_norm_path=None, # force compute + pool_horizon=pool_horizon, + episode_cache=None + if not cache_dir + else episode_norm_samples.cache_root( + cache_dir, dataset_name, cfg.data.train_datasets[dataset_name] + ), + ) + print( + f"[precompute] {dataset_name}: norm computed in " + f"{time.perf_counter() - t0:.1f}s" + ) + + if cache_dir: + # Same key trainHydra derives, so its next run is a cache hit. + episodes = { + h: norm_cache.episode_fingerprint(getattr(ds, "episode_path", None)) + for h, ds in dataset.datasets.items() + } + inputs = norm_cache.cache_inputs( + dataset_name, + episodes, + cfg.data.train_datasets[dataset_name], + args.sample_frac, + pool_horizon, + args.max_samples, + ) + key = norm_cache.norm_cache_key(inputs) + emb = get_embodiment_id(dataset_name) + written = norm_cache.write_cached( + cache_dir, + dataset_name, + key, + inputs, + emb, + norm_stats.norm_stats[emb], + norm_stats._norm_run_metadata, + ) + print(f"[precompute] {dataset_name}: cache entry {written}") + + norm_stats.cache_stats(save_cache_dir=args.out) + out_dir = os.path.join(args.out, "norm_stats") + print("\nDONE. Use this in training:") + print(f" norm_stats.precomputed_norm_path={out_dir}") + if cache_dir: + print( + f"(or nothing: the content-keyed cache under {cache_dir} now holds " + "these stats for the same episodes + recipe)" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/tutorials/aria_to_mano_convert.py b/egomimic/scripts/tutorials/aria_to_mano_convert.py index 7047e06967..9d0aa527de 100644 --- a/egomimic/scripts/tutorials/aria_to_mano_convert.py +++ b/egomimic/scripts/tutorials/aria_to_mano_convert.py @@ -9,8 +9,8 @@ 4. Permute MANO's otaheri joint order (wrist, index*3, middle*3, pinky*3, ring*3, thumb*3, then 5 tips) to the canonical MANO ordering used by Human.FINGER_EDGES (wrist, thumb*4, index*4, middle*4, ring*4, pinky*4). -5. Render side-by-side: left half = Aria viz with Aria.FINGER_EDGES; right - half = fitted MANO viz with Human (=Mecka) FINGER_EDGES. +5. Render side-by-side: raw Aria keypoints on the left, the fitted MANO on + the right, both through Human.viz(mode="keypoints"). 6. Write mp4 to scratch/aria_to_mano.mp4. Prerequisites (not checked in): @@ -37,11 +37,7 @@ import numpy as np import torch -from egomimic.rldb.embodiment.human import ( - ARIA_FINGER_EDGE_RANGES, - ARIA_FINGER_EDGES, - Human, -) +from egomimic.rldb.embodiment.human import Human from egomimic.rldb.filters import DatasetFilter from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, S3EpisodeResolver from egomimic.scripts.aria_process.aria_utils import fit_mano_to_aria_batched @@ -78,16 +74,16 @@ def fetch_episode_loader(): print(f"Aria episode: {episode_hash} (of {len(aria_df)})") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=3) + transform_list = Human.get_transform_list( + mode="keypoints_headframe_ypr", stride=3, allow_legacy_rotation=True + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list ) 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) @@ -187,8 +183,6 @@ def render_side_by_side(rows, mano_left_canonical, mano_right_canonical): batch, mode="keypoints", viz_batch_key="actions_keypoints", - finger_edges=ARIA_FINGER_EDGES, - finger_edge_ranges=ARIA_FINGER_EDGE_RANGES, ) left_flat = mano_left_canonical[i].reshape(-1).numpy() right_flat = mano_right_canonical[i].reshape(-1).numpy() diff --git a/egomimic/scripts/tutorials/mecka_keypoints_viz.py b/egomimic/scripts/tutorials/mecka_keypoints_viz.py index 65786380ce..58cb6ab1da 100644 --- a/egomimic/scripts/tutorials/mecka_keypoints_viz.py +++ b/egomimic/scripts/tutorials/mecka_keypoints_viz.py @@ -34,7 +34,9 @@ def main() -> None: print(f"Using mecka episode: {episode_hash} (of {len(mecka_df)} candidates)") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=1) + transform_list = Human.get_transform_list( + mode="keypoints_headframe_ypr", stride=1, allow_legacy_rotation=True + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list @@ -42,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.ipynb b/egomimic/scripts/tutorials/zarr_data_viz.ipynb index 2f2e12d795..953fc1c802 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.ipynb +++ b/egomimic/scripts/tutorials/zarr_data_viz.ipynb @@ -291,7 +291,9 @@ ], "source": [ "key_map = Human.get_keymap(keymap_mode=\"keypoints\")\n", - "transform_list = Human.get_transform_list(mode=\"keypoints_headframe_ypr\", stride=3)\n", + "transform_list = Human.get_transform_list(\n", + " mode=\"keypoints_headframe_ypr\", stride=3, allow_legacy_rotation=True\n", + ")\n", "\n", "resolver = S3EpisodeResolver(\n", " TEMP_DIR,\n", 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 4b2b5a17b0..9bcba82a17 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -13,22 +13,25 @@ from lightning.fabric.utilities.cloud_io import get_filesystem from lightning.pytorch.loggers import Logger from lightning.pytorch.plugins.environments import SLURMEnvironment -from omegaconf import DictConfig, OmegaConf, open_dict +from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict from tabulate import tabulate import egomimic.utils.hydra_resolvers # noqa: F401 -- registers OmegaConf resolvers from egomimic.eval.eval import Eval 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 MultiDataset, PinError +from egomimic.utils import run_manifest 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 +from egomimic.utils.gpu_orphans import arm_rank_pdeathsig, reap_orphans 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) @@ -39,6 +42,21 @@ _PI_WEIGHT_KEY = "model.robomimic_model.config.pytorch_weight_path" +def _data_transform_modes(cfg: DictConfig) -> dict[str, str]: + """``{dataset name: transform_list.mode}`` for the train datasets that set one.""" + datasets = OmegaConf.select(cfg, "data.train_datasets", default=None) + if not isinstance(datasets, DictConfig): + return {} + modes = {} + for name in datasets: + mode = OmegaConf.select( + datasets, f"{name}.resolver.transform_list.mode", default=None + ) + if mode is not None: + modes[str(name)] = str(mode) + return modes + + def _build_model_config_tree(cfg: DictConfig) -> DictConfig: """The config tree ``ModelWrapper`` receives. Only this copy skips the PI base weights (see ``_weights_from_checkpoint``); ``cfg`` keeps the real path for @@ -51,6 +69,11 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: ): model_cfg.robomimic_model.norm_stats = None tree = OmegaConf.create({"model": model_cfg}) + # Saved with the checkpoint so the robot rollout rebuilds the same inputs + # and reverts the outputs; the data config itself is not saved. + modes = _data_transform_modes(cfg) + if modes: + tree.data_transform_modes = modes has_weights = OmegaConf.select(tree, _PI_WEIGHT_KEY, default=None) is not None if has_weights and _weights_from_checkpoint(cfg): log.info( @@ -61,6 +84,16 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: return tree +def _model_trainer_defaults(cfg: DictConfig) -> dict: + """Trainer kwargs a model config supplies: the gradient clip it was + published with (pi0.5: openpi's 1.0), unless ``trainer.gradient_clip_val`` + is set explicitly.""" + clip = cfg.model.get("gradient_clip_val") + if clip is None or cfg.trainer.get("gradient_clip_val") is not None: + return {} + return {"gradient_clip_val": clip} + + def _requeue_resume_path(cfg: DictConfig) -> Optional[str]: """``/last.ckpt`` when this process is a Slurm requeue (``SLURM_RESTART_COUNT`` > 0), else None. The dir is the ModelCheckpoint @@ -79,7 +112,12 @@ def _requeue_resume_path(cfg: DictConfig) -> Optional[str]: def _prepare_checkpoint_resume(cfg: DictConfig) -> None: """Settle ``cfg.ckpt_path`` before the model config tree is built: a requeued job resumes from ``last.ckpt`` if it exists; one preempted before its first - checkpoint keeps the launch-time ``ckpt_path`` (warning).""" + checkpoint keeps the launch-time ``ckpt_path`` (warning). + + ``last.ckpt`` is a symlink (``save_last: link``), so it can dangle (target + deleted) or be missing while checkpoints exist (``rsync`` without ``-l``, a + preemption between Lightning's unlink and relink). Both raise: falling back + would silently restart a run that has checkpoints.""" requeue = _requeue_resume_path(cfg) if requeue is None: return @@ -87,6 +125,26 @@ def _prepare_checkpoint_resume(cfg: DictConfig) -> None: log.info(f"Detected SLURM requeue — resuming from {requeue}") cfg.ckpt_path = requeue return + if os.path.islink(requeue): + raise FileNotFoundError( + f"SLURM requeue: {requeue} links to {os.readlink(requeue)!r}, which " + "does not exist. Refusing to restart a run that has checkpoints; " + "restore the target or relink last.ckpt to an existing checkpoint." + ) + ckpt_dir = os.path.dirname(requeue) + existing = [ + os.path.join(ckpt_dir, f) + for f in (os.listdir(ckpt_dir) if os.path.isdir(ckpt_dir) else []) + if f.endswith(".ckpt") + ] + if existing: + newest = max(existing, key=os.path.getmtime) + raise FileNotFoundError( + f"SLURM requeue: {requeue} is missing but {ckpt_dir} holds " + f"{len(existing)} checkpoint(s), newest {newest}. Refusing to restart " + f"from scratch; link it with `ln -s {os.path.basename(newest)} " + f"{requeue}` and requeue." + ) fallback = cfg.get("ckpt_path") log.warning( f"SLURM requeue detected but {requeue} does not exist; falling back to " @@ -108,7 +166,11 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") -def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> None: +def _log_dataset_frame_counts( + train_datasets: dict, + valid_datasets: dict, + unseen_op_valid_datasets: dict | None = None, +) -> None: rows = [] for name, ds in train_datasets.items(): rows.append(("train", name, len(ds))) @@ -122,6 +184,8 @@ def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> Non rows.append( ("TOTAL", "(valid)", sum(len(ds) for ds in valid_datasets.values())) ) + for name, ds in (unseen_op_valid_datasets or {}).items(): + rows.append(("unseen_op_valid", name, len(ds))) table = tabulate( rows, headers=["Split", "Dataset", "Frames"], @@ -184,6 +248,277 @@ def load_checkpoint( return pl_load(path, map_location=map_location, weights_only=weights_only) +def _resolve_mode(cfg: DictConfig) -> str: + """``mode`` key, else the legacy ``train`` / ``eval`` booleans.""" + if cfg.get("mode") is not None: + return cfg.mode + if cfg.get("train", False): + return "train" + if cfg.get("eval", False): + return "eval" + raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + + +def _train_viz_enabled(cfg: DictConfig) -> bool: + """The train_viz head (metrics + video on TRAIN data each validation) is on + by default for training runs with an evaluator; ``train_viz=false`` turns + it off, including the heads a data/top-level config declares itself.""" + return ( + bool(cfg.get("train_viz", True)) + and _resolve_mode(cfg) == "train" + and cfg.get("evaluator") is not None + ) + + +def _train_viz_datasets(cfg: DictConfig, train_datasets: dict, instantiate): + """Datasets for the train_viz loader and, when they are derived here, its + loader params (None = leave ``data.train_viz_dataloader_params`` to Hydra). + + A data config's own ``train_viz_datasets`` win. Otherwise the train split + itself is reused (same dataset objects, so no second resolve) through an + unshuffled loader with the valid loader's params.""" + if not _train_viz_enabled(cfg): + return {}, None + explicit = cfg.data.get("train_viz_datasets") + if explicit is not None: + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + }, None + valid_params = cfg.data.get("valid_dataloader_params") or {} + params = { + name: {**OmegaConf.to_container(valid_params[name]), "shuffle": False} + for name in train_datasets + if name in valid_params + } + return {name: train_datasets[name] for name in params}, params + + +def _build_train_viz_evaluator(cfg: DictConfig): + """``train_viz_evaluator`` if configured, else a TrainVizEvalVideo around a + second instance of the canonical evaluator (own frame buffers).""" + if not _train_viz_enabled(cfg): + return None + if cfg.get("train_viz_evaluator") is not None: + return hydra.utils.instantiate(cfg.train_viz_evaluator) + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + return TrainVizEvalVideo(hydra.utils.instantiate(cfg.evaluator)) + + +def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: + """Datasets for the third (unseen_op_valid) val loader: a data config's own + ``unseen_op_valid_datasets``, in training runs with an evaluator only (eval + mode validates ``valid_datasets`` alone).""" + explicit = cfg.data.get("unseen_op_valid_datasets") + if ( + explicit is None + or _resolve_mode(cfg) != "train" + or cfg.get("evaluator") is None + ): + return {} + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + } + + +def _trainer_world_size(cfg: DictConfig) -> int: + """``devices * num_nodes`` as the config declares it, 1 if it cannot say. + + DistributedSampler strides rather than chunks, so W ranks between them walk + the WHOLE split; the per-rank ``limit_val_batches`` window is therefore W + times wider than it looks from one rank. Eval mode runs on one device + (``train`` forces ``devices=1`` after the loaders are built), so it is 1. + """ + legacy = "train" if cfg.get("train") else "eval" if cfg.get("eval") else None + if (cfg.get("mode") or legacy) == "eval": # _resolve_mode, minus its raise + return 1 + devices = cfg.get("trainer", {}).get("devices", 1) + if isinstance(devices, (list, ListConfig)): + n = len(devices) + elif isinstance(devices, int) and devices > 0: + n = devices + else: # "auto", -1: resolved by lightning at runtime, unknown here + log.warning( + f"trainer.devices={devices!r} is not a count; treating the metric " + "loaders as single-rank, which under-scores the val split." + ) + n = 1 + nodes = cfg.get("trainer", {}).get("num_nodes", 1) + return n * (int(nodes) if isinstance(nodes, int) and nodes > 0 else 1) + + +def _metric_frames_per_episode( + cfg: DictConfig, + head: str, + n_episodes: int | None = None, + batch_size: int | None = None, +) -> int | None: + """Frames per episode to keep on this val head, or None for no subsampling. + + ``data.metric_frames_per_episode[head]`` is written for ONE rank -- it is + ``floor(limit_val_batches * batch_size / n_episodes)`` -- so it is scaled by + the world size here: the ranks stride through the subsampled set together + and each still reads at most ``limit_val_batches`` batches. Without this a + 4-GPU run scores a quarter of the frames the same config scores on 1 GPU, + which on the seen-val head is fewer than it scored before subsampling + existed at all. EvenStrideDataset keeps a whole episode when K exceeds its + length, so a split that fits entirely is not subsampled. + + ``auto`` computes that formula from the resolved split (``n_episodes``) and + the head's loader ``batch_size``, for splits defined by live SQL filters + whose episode count a hard-coded K would silently fall behind. With + ``limit_val_batches=0`` (validation off) it subsamples nothing. + """ + table = cfg.data.get("metric_frames_per_episode") + if table is None: + return None + k = table.get(head) + if k is None: + return None + if k == "auto": + limit = cfg.get("trainer", {}).get("limit_val_batches") + if limit == 0: # validation is off; nothing to budget + return None + if ( + isinstance(limit, bool) + or not isinstance(limit, int) + or limit <= 0 + or not batch_size + or not n_episodes + ): + raise ValueError( + f"data.metric_frames_per_episode.{head}=auto needs an int " + "trainer.limit_val_batches, the head's loader batch_size and a " + f"resolved split; got limit_val_batches={limit!r}, " + f"batch_size={batch_size!r}, episodes={n_episodes!r}" + ) + k = max(1, limit * int(batch_size) // n_episodes) + else: + k = int(k) + if k <= 0: + raise ValueError( + f"data.metric_frames_per_episode.{head} must be a positive int " + f"or 'auto', got {k}" + ) + return k * _trainer_world_size(cfg) + + +def _subsample_val_datasets( + cfg: DictConfig, head: str, datasets: dict, loader_params=None +) -> dict: + """Wrap each of this val head's datasets in ``EvenStrideDataset`` so the + ``limit_val_batches`` window spans EVERY episode of the split. + + The val loaders are unshuffled (so the overlay video is coherent), which + made the metric window the first ``limit_val_batches * batch_size * W`` + frames in episode-hash order: at W = 1 that is a third of the seen-val + split and effectively one operator of the unseen split (lane-0 coverage + report, 2026-09-15, measured single-rank). K evenly spaced frames per + episode gives every episode -- and so every operator -- equal weight at + the same batch count. + + ``loader_params`` is this head's ``{dataset_name: DataLoader kwargs}``, + read only for an ``auto`` K. No ``metric_frames_per_episode`` entry for this + head => datasets pass through untouched.""" + table = cfg.data.get("metric_frames_per_episode") + if table is None or table.get(head) is None or not datasets: + return datasets + wrapped = {} + for name, ds in datasets.items(): + if ds is None: + wrapped[name] = None + continue + params = (loader_params or {}).get(name) or {} + k = _metric_frames_per_episode( + cfg, head, n_episodes=len(ds.datasets), batch_size=params.get("batch_size") + ) + if k is None: + wrapped[name] = ds + continue + sub = EvenStrideDataset(ds, frames_per_episode=k) + log.info( + f"val head '{head}' dataset '{name}': EvenStrideDataset " + f"frames_per_episode={k} -> {len(sub)} / {len(ds)} frames over " + f"{len(ds.datasets)} episodes" + ) + wrapped[name] = sub + return wrapped + + +def _video_datasets(cfg: DictConfig, heads: dict) -> dict: + """``{head: {dataset_name: MultiDataset}}`` for the video-only val loaders. + + ``data.video_episodes[head]`` pins one episode hash per operator; each + pinned episode is taken out of that head's ALREADY resolved split (no second + SQL pull or path resolution -- see ``pinned_episode_subset``), contiguous + and unsubsampled so the overlay video stays watchable. + + A head with no pins gets no video loader, so a config without the key keeps + today's behaviour (metrics AND video on the one metric loader).""" + table = cfg.data.get("video_episodes") + if table is None: + return {} + out: dict = {} + for head, datasets in heads.items(): + pins = [str(h) for h in (table.get(head) or [])] + if not pins or not datasets: + continue + live = {name: ds for name, ds in datasets.items() if ds is not None} + unplaced = sorted( + h for h in pins if all(ds.episode(h) is None for ds in live.values()) + ) + if unplaced: + raise PinError( + f"{len(unplaced)} pinned video episode(s) for head '{head}' are in " + f"none of its datasets {sorted(live)}: {unplaced}. The pin must " + "name an episode of THIS head's split." + ) + built = {} + for name, ds in live.items(): + mine = [h for h in pins if ds.episode(h) is not None] + if not mine: + continue + built[name] = pinned_episode_subset(ds, mine, dataset_name=name) + log.info( + f"val head '{head}' video loader from dataset '{name}': " + f"{len(mine)} pinned episode(s), {len(built[name])} frames" + ) + out[head] = built + return out + + +def _require_capped_video_heads(model, datamodule) -> None: + """A pinned video loader renders on rank 0 alone while the other ranks wait + in the val reduction, so an uncapped render (``viz_max_batches=None``) is + bounded only by the 30-minute process-group timeout.""" + heads = model._val_heads() + for head, datasets in (getattr(datamodule, "video_datasets", None) or {}).items(): + evaluator = heads.get(head) + if ( + datasets + and getattr(evaluator, "viz_func", None) is not None + and evaluator.viz_max_batches is None + ): + raise ValueError( + f"val head '{head}' has pinned video episodes " + f"(data.video_episodes.{head}) but its evaluator sets no " + "viz_max_batches: rank 0 would render every video batch while " + "the other ranks wait on it. Set evaluator.viz_max_batches." + ) + + +def _build_unseen_op_valid_evaluator(cfg: DictConfig): + """The canonical evaluator (fresh instance, own frame buffers) wrapped to + log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + return TrainVizEvalVideo( + hydra.utils.instantiate(cfg.evaluator), prefix="unseen_op_valid" + ) + + @task_wrapper def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Trains the model. Can additionally evaluate on a testset, using best weights obtained during @@ -205,8 +540,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: load_env() - # One SQL pull / path resolution per dataset spec across train, valid and - # the norm-stat copies; dropped on exit so nothing outlives this run. + # One SQL pull / path resolution per dataset spec across train, valid, + # train_viz and the norm-stat copies; dropped on exit so nothing outlives + # this run. with resolve_once(): train_datasets = {} for dataset_name in cfg.data.train_datasets: @@ -220,12 +556,66 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: cfg.data.valid_datasets[dataset_name], dataset_name=dataset_name ) + train_viz_datasets, train_viz_params = _train_viz_datasets( + cfg, train_datasets, instantiate=_instantiate_dataset + ) + unseen_op_valid_datasets = _unseen_op_valid_datasets( + cfg, instantiate=_instantiate_dataset + ) + + # Split the val heads in two (lane A): a per-episode SUBSAMPLED metric + # loader that covers the whole split inside limit_val_batches, and a + # contiguous video-only loader over the pinned episodes. Build the video + # subsets from the full splits BEFORE wrapping, and let each head opt in + # independently (no keys in the data config => today's behaviour). + video_datasets = _video_datasets( + cfg, + { + "valid": valid_datasets, + "train_viz": train_viz_datasets, + "unseen_op_valid": unseen_op_valid_datasets, + }, + ) + valid_datasets = _subsample_val_datasets( + cfg, + "valid", + valid_datasets, + cfg.data.get("valid_dataloader_params"), + ) + train_viz_datasets = _subsample_val_datasets( + cfg, + "train_viz", + train_viz_datasets, + train_viz_params + if train_viz_params is not None + else cfg.data.get("train_viz_dataloader_params"), + ) + unseen_op_valid_datasets = _subsample_val_datasets( + cfg, + "unseen_op_valid", + unseen_op_valid_datasets, + cfg.data.get("unseen_op_valid_dataloader_params"), + ) + log.info(f"Instantiating datamodule <{cfg.data._target_}>") assert ( "MultiDataModuleWrapper" in cfg.data._target_ ), "cfg.data._target_ must be 'MultiDataModuleWrapper'" + datamodule_kwargs = dict( + train_datasets=train_datasets, valid_datasets=valid_datasets + ) + # Always passed, so train_viz=false also drops a data config's own + # train_viz_datasets. + datamodule_kwargs["train_viz_datasets"] = train_viz_datasets + if train_viz_params is not None: + datamodule_kwargs["train_viz_dataloader_params"] = train_viz_params + # Always passed too: hydra would otherwise instantiate the raw + # config nodes itself (eval mode included). + datamodule_kwargs["unseen_op_valid_datasets"] = unseen_op_valid_datasets + # Built above; hydra has nothing to instantiate for these. + datamodule_kwargs["video_datasets"] = video_datasets datamodule: LightningDataModule = hydra.utils.instantiate( - cfg.data, train_datasets=train_datasets, valid_datasets=valid_datasets + cfg.data, **datamodule_kwargs ) # Stats-only MultiDataset (no graph of its own; explicitly populated from @@ -236,9 +626,13 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ) norm_stats.populate_from_datasets(datamodule.train_datasets) - from egomimic.rldb.zarr import 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) + pool_horizon = bool( + OmegaConf.select(cfg, "norm_stats.pool_horizon", default=False) + ) explicit_path = OmegaConf.select( cfg, "norm_stats.precomputed_norm_path", default=None ) @@ -256,6 +650,13 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: # this remove annotation and image keys from the keymap km["norm_mode"] = True + if "proprio_history" in km: + # Current-step proprio stats even on a history (K > 1) config: + # every history step is normalized with the CURRENT-step, + # per-channel (D,) stats, so a K > 1 run and the K = 1 baseline + # share exactly one set of stats. Same rule as + # scripts/precompute_norm_stats.py. + km["proprio_history"] = 1 instantiate_copy.resolver.key_map = km norm_dataset = _instantiate_dataset( @@ -263,10 +664,13 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ) emb = get_embodiment_id(dataset_name) - key = inputs = cached = None + key = inputs = cached = episode_cache = None if explicit_path is None and use_cache and cache_dir: + episode_cache = episode_norm_samples.cache_root( + cache_dir, dataset_name, cfg.data.train_datasets[dataset_name] + ) episodes = { - h: norm_cache.episode_fingerprint(getattr(ds, "episode_path", None)) + h: norm_cache.leaf_fingerprint(ds) for h, ds in dataset.datasets.items() } inputs = norm_cache.cache_inputs( @@ -274,6 +678,8 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: episodes, cfg.data.train_datasets[dataset_name], sample_frac, + pool_horizon, + max_samples, ) key = norm_cache.norm_cache_key(inputs) cached = norm_cache.find_cached(cache_dir, dataset_name, key, emb) @@ -285,10 +691,13 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: norm_dataset, dataset_name, sample_frac=sample_frac, + max_samples=max_samples, num_workers=OmegaConf.select(cfg, "norm_stats.num_workers", default=4), precomputed_norm_path=explicit_path if explicit_path is not None else cached, + pool_horizon=pool_horizon, + episode_cache=episode_cache, ) if key is not None and cached is None: if norm_stats.norm_stats.get(emb): @@ -316,6 +725,16 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ds.set_norm_stats_from(norm_stats) for ds in datamodule.valid_datasets.values(): ds.set_norm_stats_from(norm_stats) + for ds in getattr(datamodule, "train_viz_datasets", {}).values(): + ds.set_norm_stats_from(norm_stats) + for ds in getattr(datamodule, "unseen_op_valid_datasets", {}).values(): + ds.set_norm_stats_from(norm_stats) + # The video subsets share leaves with their head's split but are separate + # MultiDatasets, so they need their own wiring (EvenStrideDataset forwards + # the call to the base it indexes into). + for head_datasets in getattr(datamodule, "video_datasets", {}).values(): + for ds in head_datasets.values(): + ds.set_norm_stats_from(norm_stats) _prepare_checkpoint_resume(cfg) @@ -324,22 +743,19 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: config_tree=_build_model_config_tree(cfg), norm_stats_state=norm_stats.to_state(), scheduler_interval=cfg.model.get("scheduler_interval", "step"), + enable_grad_norm=cfg.model.get("enable_grad_norm", True), ) - _log_dataset_frame_counts(datamodule.train_datasets, datamodule.valid_datasets) + _log_dataset_frame_counts( + datamodule.train_datasets, + datamodule.valid_datasets, + datamodule.unseen_op_valid_datasets, + ) log.info("Instantiating callbacks...") callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks")) - # Resolve mode: support both new `mode` key and legacy `train`/`eval` booleans - if cfg.get("mode") is not None: - mode = cfg.mode - elif cfg.get("train", False): - mode = "train" - elif cfg.get("eval", False): - mode = "eval" - else: - raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + mode = _resolve_mode(cfg) # In eval mode, apply trainer overrides from the eval object and disable logger if mode == "eval": @@ -368,7 +784,11 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: plugins.append(SLURMEnvironment(requeue_signal=signal.SIGUSR1)) print("SLURM REQUEUE ENABLED") trainer: Trainer = hydra.utils.instantiate( - cfg.trainer, callbacks=callbacks, logger=logger, plugins=plugins or None + cfg.trainer, + callbacks=callbacks, + logger=logger, + plugins=plugins or None, + **_model_trainer_defaults(cfg), ) object_dict = { @@ -385,13 +805,54 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: log_hyperparameters(object_dict) os.makedirs(os.path.join(trainer.default_root_dir, "videos"), exist_ok=True) + if trainer.is_global_zero: + manifest = run_manifest.run_manifest( + datamodule.train_datasets, datamodule.valid_datasets + ) + if mode == "train": + run_manifest.warn_on_split_drift( + cfg.get("ckpt_path"), manifest[run_manifest.SPLIT_FILE] + ) + run_manifest.write(trainer.default_root_dir, manifest) if mode == "train": if cfg.get("evaluator") is not None: eval_obj: Eval = hydra.utils.instantiate(cfg.evaluator) + # data.valid_prefix (e.g. seen_op_valid in the opsplit configs) + # renames the canonical head: `/Valid/...` metrics and + # `videos_/` instead of `Valid/...` and `videos/`. + valid_prefix = cfg.data.get("valid_prefix") + if valid_prefix: + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + eval_obj = TrainVizEvalVideo(eval_obj, prefix=valid_prefix) eval_obj.trainer = trainer eval_obj.model = model.model model.evaluator = eval_obj + train_viz_eval_obj = ( + _build_train_viz_evaluator(cfg) if datamodule.train_viz_datasets else None + ) + if train_viz_eval_obj is not None: + train_viz_eval_obj.trainer = trainer + train_viz_eval_obj.model = model.model + model.train_viz_evaluator = train_viz_eval_obj + if datamodule.unseen_op_valid_datasets: + unseen_eval_obj = _build_unseen_op_valid_evaluator(cfg) + unseen_eval_obj.trainer = trainer + unseen_eval_obj.model = model.model + model.unseen_op_valid_evaluator = unseen_eval_obj + _require_capped_video_heads(model, datamodule) + model.val_loader_names = datamodule.val_loader_names() + # Pre-fit baseline val. Skipped on requeues AND checkpoint resumes: + # trainer.validate here runs BEFORE fit restores ckpt_path weights, so + # on a resume it would score the un-resumed base model. + # (_prepare_checkpoint_resume already folded a requeue into ckpt_path.) + if cfg.get("val_at_start", False) and not cfg.get("ckpt_path"): + log.info( + "val_at_start: running validation at epoch 0 (pre-fit baseline; " + "the overlay video follows evaluator.viz_every_n_epochs)" + ) + trainer.validate(model=model, datamodule=datamodule) log.info("Starting training!") trainer.fit( model=model, @@ -403,6 +864,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: eval_obj.trainer = trainer eval_obj.model = model.model model.evaluator = eval_obj + # Without the names, loader indices fall back to [valid, train_viz] and + # the pinned video loader is misrouted. + model.val_loader_names = datamodule.val_loader_names() if hasattr(eval_obj, "run"): eval_obj.run(trainer, model, datamodule, cfg) @@ -451,6 +915,11 @@ def main(cfg: DictConfig) -> Optional[float]: # the job runs main(), so each job keys the cache on its own SLURM_JOB_ID. set_per_job_compile_cache_dir() + # Before CUDA init: a cancelled job can leave dataloader workers -- or whole + # ranks -- holding GPU contexts, and nothing else on this cluster sweeps them. + reap_orphans() + arm_rank_pdeathsig() + # apply extra utilities # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.) extras(cfg) diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 7a80fbd970..063410d292 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -1,16 +1,20 @@ """pi0.5 action I/O. Actions arrive from the data pipeline already in their model-native layout -and already normalized by the standard MultiDataset pipeline: xyz+rot6d -(+gripper) per arm, the ypr->6D conversion being the ``CartesianYPRToRot6D`` -data transform, with norm stats computed on 6D. +and already normalized by the standard MultiDataset pipeline: -The forward pass only *packs* the normalized vector into openpi's 32-slot -action vector (``to32_norm_6d``) and the eval path unpacks it +- cartesian: xyz+rot6d(+gripper) per arm (the ypr->6D conversion is the + ``CartesianYPRToRot6D`` data transform, norm stats are computed on 6D); +- hand keypoints: the 144-D wrist-first ``[wrist xyz | wrist rot6d | + 21 keypoints] x 2`` vector (``keypoints_*_6d`` modes). + +The forward pass only *packs* the normalized vector into openpi's action +vector (``to32_norm_6d`` builds the canonical 32-slot block layout, and the +PI algo pads that to ``model.action_dim``) and the eval path unpacks it (``from32_norm_6d``). No rotation math and no normalization happen here. """ -from typing import Any, Dict, Tuple +from typing import Any, Dict, Sequence, Tuple import torch @@ -44,15 +48,28 @@ def _ensure_bsd(x: torch.Tensor) -> torch.Tensor: return x -def _pad32(x: torch.Tensor) -> torch.Tensor: +def pad_to_width(x: torch.Tensor, width: int) -> torch.Tensor: + """Zero-pad the last dim of a (B,S,D) tensor up to ``width``. + + Raises when ``D > width``: a wider native action than the model's + ``action_dim`` would otherwise be truncated silently. + """ x = _ensure_bsd(x) B, S, D = x.shape - if D == 32: + if D == width: return x - if D < 32: - pad = torch.zeros(B, S, 32 - D, dtype=x.dtype, device=x.device) - return torch.cat([x, pad], dim=-1) - return x[..., :32] + if D > width: + raise ValueError( + f"Action has {D} dims but the model action_dim is {width}; set " + "model.robomimic_model.config.model.action_dim >= the widest " + "embodiment action" + ) + pad = torch.zeros(B, S, width - D, dtype=x.dtype, device=x.device) + return torch.cat([x, pad], dim=-1) + + +def _pad32(x: torch.Tensor) -> torch.Tensor: + return pad_to_width(x, 32) def _stat_tensor(stats: dict[str, Any], key: str, ref: torch.Tensor) -> torch.Tensor: @@ -60,6 +77,21 @@ def _stat_tensor(stats: dict[str, Any], key: str, ref: torch.Tensor) -> torch.Te return value.to(dtype=ref.dtype if ref.is_floating_point() else torch.float32) +# A channel whose stat range (std, max - min, or q99 - q1) is below this is +# treated as constant: it normalizes to 0 and unnormalizes to its centre. +# Structural zeros, not rare accidents: the t=0 cell of a wrist-frame action +# chunk is exactly the identity pose, and in a wrist-frame keypoint proprio the +# wrist keypoint sits at its own frame origin while kp9 defines the forward +# axis, so kp0 and kp9's off-axis components are zero at EVERY timestep (28 +# action cells and 10 proprio channels on a real 144-D episode). A bare +# ``+ 1e-6`` denominator scaled any off-convention value in them by 1e6. +# One absolute threshold across channels with different units is deliberate: +# 1e-4 is 0.1 mm of translation but 1/20000 of a rot6d channel's range. Same +# rule and the same constant as Diffusion Policy / robomimic; GR00T instead +# clamps the range, which keeps dividing. +NORM_MIN_RANGE = 1e-4 + + def _norm_center_scale( stats: dict[str, Any], norm_mode: str, ref: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -82,24 +114,48 @@ def _eps(norm_mode: str) -> float: return 1e-6 if norm_mode == "zscore" else 0.5e-6 +def _degenerate(scale: torch.Tensor, norm_mode: str) -> torch.Tensor: + full_range = scale if norm_mode == "zscore" else 2.0 * scale + return full_range < NORM_MIN_RANGE + + +def _passthrough( + out: torch.Tensor, tensor: torch.Tensor, channels: Sequence[int] | None +) -> torch.Tensor: + if not channels: + return out + out = out.clone() + out[..., list(channels)] = tensor[..., list(channels)] + return out + + def _apply_norm_one( tensor: torch.Tensor, stats: dict[str, Any], norm_mode: str, + identity_channels: Sequence[int] | None = None, ) -> torch.Tensor: - """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile).""" + """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile). + Channels with a range below ``NORM_MIN_RANGE`` map to 0; + ``identity_channels`` pass through untouched.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) - return (tensor - center) / (scale + _eps(norm_mode)) + out = (tensor - center) / (scale + _eps(norm_mode)) + out = torch.where(_degenerate(scale, norm_mode), torch.zeros_like(out), out) + return _passthrough(out, tensor, identity_channels) def _apply_unnorm_one( tensor: torch.Tensor, stats: dict[str, Any], norm_mode: str, + identity_channels: Sequence[int] | None = None, ) -> torch.Tensor: - """Inverse of :func:`_apply_norm_one`.""" + """Inverse of :func:`_apply_norm_one`; degenerate channels return their + centre (mean / midpoint) whatever the model predicted.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) - return tensor * (scale + _eps(norm_mode)) + center + out = tensor * (scale + _eps(norm_mode)) + center + out = torch.where(_degenerate(scale, norm_mode), center.expand_as(out), out) + return _passthrough(out, tensor, identity_channels) # ---------- rotation helpers (shared with the eval metrics) ---------- @@ -214,7 +270,7 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class RobotBimanualCartesianEuler(BaseActionConverter): +class RobotBimanualCartesian6D(BaseActionConverter): """ Native: (B,S,20) = [L xyz(3) 6d(6) g(1) | R xyz(3) 6d(6) g(1)] 32-pack: left block 0..9, right block 10..19, zero pad 20..31 @@ -245,7 +301,7 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class HumanBimanualCartesianEuler(BaseActionConverter): +class HumanBimanualCartesian6D(BaseActionConverter): """ Native: (B,S,18) = [L xyz(3) 6d(6) | R xyz(3) 6d(6)] (no gripper) 32-pack: left block 0..9 (g=0), right block 10..19 (g=0), zero pad 20..31 @@ -276,3 +332,33 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: L = actions32[..., 0:9] # drop left gripper slot at idx 9 R = actions32[..., 10:19] # drop right gripper slot at idx 19 return torch.cat([L, R], dim=-1) # (B,S,18) + + +class HumanBimanualKeypoints(BaseActionConverter): + """144-D wrist-first MANO keypoint action, per hand + ``[wrist xyz+rot6d in its own wrist frame (9) | 21 keypoints in that + wrist frame (63)]`` x {left, right} (``keypoints_wristframe_6d``), packed + identity-first into the model's action vector. The PI algo resizes + openpi's action projections to ``model.action_dim`` (>= 144) and pads the + vector up to it; ``from32_norm_6d`` slices the native width back out. + """ + + native_dim = 144 + + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + actions = _ensure_bsd(actions) + if actions.shape[-1] != self.native_dim: + raise ValueError( + f"HumanBimanualKeypoints: expected {self.native_dim}-dim, got " + f"{actions.shape[-1]}" + ) + return actions + + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + actions32 = _ensure_bsd(actions32) + if actions32.shape[-1] < self.native_dim: + raise ValueError( + f"HumanBimanualKeypoints: expected >={self.native_dim} dims, got " + f"{actions32.shape[-1]}" + ) + return actions32[..., : self.native_dim] diff --git a/egomimic/utils/compile_utils.py b/egomimic/utils/compile_utils.py new file mode 100644 index 0000000000..eb93e40e6f --- /dev/null +++ b/egomimic/utils/compile_utils.py @@ -0,0 +1,48 @@ +"""torch.compile wiring shared by the algos. + +``nn.Module.compile`` (not ``torch.compile(module)``) is what the algos call: +it swaps the module's call implementation IN PLACE, so the module keeps its +identity and -- the part that matters here -- its ``state_dict`` keys. Wrapping +a module in ``torch.compile`` instead returns an ``OptimizedModule`` whose +parameters are all prefixed ``_orig_mod.``, which silently breaks every +checkpoint written before the flag was turned on. + +The other half of the contract: ``Module.compile`` only routes ``__call__``. +A submodule the algo invokes through a named method (``stem.compute_latent``, +``head.compute_loss``) runs uncompiled no matter what, so only modules that are +actually *called* are worth handing over. +""" + +from __future__ import annotations + +import logging +from typing import Iterable, List, Optional + +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +def compile_modules( + modules: Iterable[tuple[str, Optional[nn.Module]]], + mode: Optional[str] = None, + dynamic: bool = False, +) -> List[str]: + """``module.compile(...)`` for each named module; returns the names done.""" + kwargs = {"dynamic": dynamic} + if mode: + kwargs["mode"] = mode + done = [] + for name, module in modules: + if not isinstance(module, nn.Module): + continue # a head with no sub-network, or a None slot + module.compile(**kwargs) + done.append(name) + if done: + logger.info( + "torch.compile (mode=%s, dynamic=%s): %s", + mode or "default", + dynamic, + ", ".join(done), + ) + return done diff --git a/egomimic/utils/ddp_hooks.py b/egomimic/utils/ddp_hooks.py new file mode 100644 index 0000000000..1d86790a77 --- /dev/null +++ b/egomimic/utils/ddp_hooks.py @@ -0,0 +1,12 @@ +"""DDP gradient-communication hooks, exposed as zero-argument factories so a +Hydra config can select one (``ddp_comm_hook: {_target_: ...}``). Lightning +logs the hook's ``__qualname__``, so the config must hand it the real function, +not a ``functools.partial`` (Hydra's ``_partial_``).""" + +from torch.distributed.algorithms.ddp_comm_hooks import default_hooks + + +def make_bf16_compress_hook(): + """All-reduce gradients in bf16 (half the bytes of fp32), decompress after. + Master weights and the optimizer stay fp32.""" + return default_hooks.bf16_compress_hook diff --git a/egomimic/utils/gpu_orphans.py b/egomimic/utils/gpu_orphans.py new file mode 100644 index 0000000000..92ec131ba3 --- /dev/null +++ b/egomimic/utils/gpu_orphans.py @@ -0,0 +1,175 @@ +"""Reclaim GPUs pinned by dataloader workers that outlived a cancelled job. + +A worker forked after CUDA init holds its rank's ``/dev/nvidia*`` fds, so the +rank's whole context stays resident after the rank is gone -- ``nvidia-smi`` then +reports the memory against a PID that is no longer in ``/proc``. The worker +cannot exit on its own (it blocks in ``pipe_write`` on a result pipe nobody +drains), and ``ProctrackType=proctrack/linuxproc`` walks the PPID tree, so +Slurm never sweeps it once it is reparented to init. + +``orphan_guarded`` / ``die_with_parent`` stop new leaks. This is the other half: a job +clears what an earlier job left on its node, before it touches CUDA itself. + +Nothing is imported from torch here -- this has to run before CUDA init. +""" + +import ctypes +import logging +import os +import signal + +from egomimic.utils.type_utils import str2bool + +logger = logging.getLogger(__name__) + +_DISABLE_ENV = "EGOMIMIC_REAP_GPU_ORPHANS" +_PR_SET_PDEATHSIG = 1 + + +def die_with_parent(_worker_id: int = 0) -> None: + """Ask the kernel to SIGKILL this dataloader worker when its rank dies. + + PyTorch's own parent check only runs between ``index_queue.get()`` timeouts, + which a worker blocked writing to a full result pipe never reaches. The + kernel does not care what the worker is blocked on. + """ + try: + ctypes.CDLL("libc.so.6", use_errno=True).prctl( + _PR_SET_PDEATHSIG, signal.SIGKILL + ) + except (OSError, AttributeError): # non-Linux, or no libc + logger.debug("PR_SET_PDEATHSIG unavailable; orphaned workers may leak GPUs") + return + # pdeathsig does not fire retroactively if the rank died during the fork. + if os.getppid() == 1: + # os._exit skips logging's flush; say why before a rank exits 1 silently. + os.write( + 2, + f"[gpu_orphans] pid {os.getpid()}: parent gone before " + "PR_SET_PDEATHSIG was armed; exiting\n".encode(), + ) + os._exit(1) + + +def orphan_guarded(params: dict) -> dict: + """``params`` with ``die_with_parent`` run before any caller ``worker_init_fn``.""" + params = dict(params) + if params.get("num_workers", 0) == 0: + return params + user_init = params.get("worker_init_fn") + if user_init is None: + params["worker_init_fn"] = die_with_parent + else: + + def _chained(worker_id: int, _user_init=user_init) -> None: + die_with_parent(worker_id) + _user_init(worker_id) + + params["worker_init_fn"] = _chained + return params + + +def _ppid(proc_root: str, pid: str) -> int | None: + try: + with open(f"{proc_root}/{pid}/status") as f: + for line in f: + if line.startswith("PPid:"): + return int(line.split()[1]) + except (OSError, ValueError): + pass + return None + + +def _holds_gpu(proc_root: str, pid: str) -> bool: + fd_dir = f"{proc_root}/{pid}/fd" + try: + fds = os.listdir(fd_dir) + except OSError: # gone, or not ours + return False + for fd in fds: + try: + if os.readlink(f"{fd_dir}/{fd}").startswith("/dev/nvidia"): + return True + except OSError: + continue + return False + + +def arm_rank_pdeathsig() -> bool: + """Have the kernel kill this process if its launcher dies first. + + Workers are not the only thing that leaks. A rank wedged in + ``torch.cuda.synchronize()`` inside inductor's cudagraph_trees does not + answer SIGTERM; Slurm hits ``UnkillableStepTimeout``, gives up, and the job + leaves the queue with its ranks still spinning on every GPU of the node. + + Only armed under Slurm, where the parent is slurmstepd or the launcher that + spawned this rank and never legitimately exits first. Off elsewhere, and off + for a ``nohup`` run (SIGHUP ignored): inside a ``salloc`` shell that is a run + meant to outlive the shell, which PR_SET_PDEATHSIG would kill. + """ + if "SLURM_JOB_ID" not in os.environ: + return False + if signal.getsignal(signal.SIGHUP) is signal.SIG_IGN: + logger.info("SIGHUP is ignored (nohup); not tying this rank to its parent") + return False + die_with_parent() + return True + + +def find_orphans(proc_root: str = "/proc", uid: int | None = None) -> list[int]: + """PIDs that are ours, reparented to init, and holding a GPU device open. + + A live job's workers always have a live parent, so the PPID-1 test alone + keeps this off anything still running -- ours or anyone else's. + """ + uid = os.getuid() if uid is None else uid + mine = {os.getpid(), os.getppid()} + orphans = [] + for pid in os.listdir(proc_root): + if not pid.isdigit() or int(pid) in mine: + continue + try: + if os.stat(f"{proc_root}/{pid}").st_uid != uid: + continue + except OSError: + continue + if _ppid(proc_root, pid) != 1: + continue + if _holds_gpu(proc_root, pid): + orphans.append(int(pid)) + return sorted(orphans) + + +def reap_orphans(dry_run: bool = False) -> list[int]: + """Kill this node's leaked GPU holders. Returns the PIDs signalled. + + Off outside Slurm: only there do we know every process of ours on the node + belongs to a job, so that a PPID-1 GPU holder can only be wreckage. + """ + if not str2bool(os.environ.get(_DISABLE_ENV, "1")): + return [] + if "SLURM_JOB_ID" not in os.environ: + return [] + if os.environ.get("SLURM_LOCALID", "0") != "0": # once per node + return [] + + orphans = find_orphans() + if not orphans: + return [] + logger.warning( + "Reaping %d orphaned GPU holder(s) left by an earlier job on %s: %s " + "(set %s=0 to disable)", + len(orphans), + os.uname().nodename, + orphans, + _DISABLE_ENV, + ) + if dry_run: + return orphans + for pid in orphans: + try: + os.kill(pid, signal.SIGKILL) + except OSError as e: + logger.warning("could not kill orphan %d: %s", pid, e) + return orphans diff --git a/egomimic/utils/metrics.py b/egomimic/utils/metrics.py index aa6c04756c..8b1ace86c0 100644 --- a/egomimic/utils/metrics.py +++ b/egomimic/utils/metrics.py @@ -1,33 +1,7 @@ -"""Distribution-distance metrics used by evaluation. - -Moved out of egomimicUtils; code unchanged. -""" - -import math +"""Distribution-distance metrics used by evaluation.""" import torch -# ---- moved from egomimicUtils.py (code unchanged) ---- - - -def reverse_kl_from_samples(pred_samples, targets): - M, B, T, D = pred_samples.shape - - TD = T * D - const = -0.5 * TD * math.log(2.0 * math.pi) - - A = pred_samples.permute(1, 0, 2, 3).reshape(B, M, TD) # (B,M,TD) - MU = targets.reshape(B, 1, TD) # (B,1,TD) - - d2 = torch.cdist(A, A).pow(2) # (B,M,M) - log_q_each = torch.logsumexp(const - 0.5 * d2, dim=-1) - math.log(M) # (B,M) - - d2p = ((A - MU) ** 2).sum(dim=-1) # (B,M) - log_p_each = const - 0.5 * d2p # (B,M) - - rkl_each = (log_q_each - log_p_each).mean(dim=-1) # (B,) - return rkl_each.mean() - def frechet_gaussian_over_time( pred: torch.Tensor, @@ -112,3 +86,39 @@ def frechet_gaussian_over_time( "max": dist.max().item(), } return dist + + +def dtw_distance(pred: torch.Tensor, tgt: torch.Tensor, normalize: bool = True): + """Batched dynamic-time-warping distance between (B, T1, D) and (B, T2, D) + trajectories. + + Local cost is the euclidean distance between D-dim frames; the classic + DTW recurrence runs as a wavefront over the T1+T2-1 anti-diagonals of the + cost grid so the whole batch advances in vectorized steps (no per-cell + python loop). Unlike paired MSE, DTW forgives temporal misalignment: a + correct trajectory executed slightly early/late scores near zero. + + Args: + pred/tgt: (B, T, D) action chunks (any device). + normalize: divide the accumulated path cost by (T1 + T2), giving a + per-step average frame distance comparable across chunk lengths. + Returns: + (B,) distances. + """ + B, T1, _ = pred.shape + _, T2, _ = tgt.shape + C = torch.cdist(pred.float(), tgt.float()) # (B, T1, T2) + acc = torch.full((B, T1 + 1, T2 + 1), float("inf"), device=C.device, dtype=C.dtype) + acc[:, 0, 0] = 0.0 + for k in range(2, T1 + T2 + 1): + i = torch.arange(max(1, k - T2), min(T1, k - 1) + 1, device=C.device) + j = k - i + prev = torch.minimum( + torch.minimum(acc[:, i - 1, j], acc[:, i, j - 1]), + acc[:, i - 1, j - 1], + ) + acc[:, i, j] = C[:, i - 1, j - 1] + prev + out = acc[:, T1, T2] + if normalize: + out = out / (T1 + T2) + return out diff --git a/egomimic/utils/paths.py b/egomimic/utils/paths.py new file mode 100644 index 0000000000..afed066cec --- /dev/null +++ b/egomimic/utils/paths.py @@ -0,0 +1,21 @@ +"""Path defaults for scripts that run outside hydra. + +A standalone script (a viz tool, a probe) still wants the cluster's dataset +mirror, and a second copy of the path silently drifts from the config the +training runs use. This reads the one hydra defines. +""" + +from __future__ import annotations + +from pathlib import Path + +from omegaconf import OmegaConf + +import egomimic.hydra_configs as _cfg_pkg + +_PATHS_YAML = Path(_cfg_pkg.__file__).parent / "paths" / "default.yaml" + + +def default_dataset_dir() -> str: + """``paths.dataset_dir`` from the hydra configs, env override included.""" + return str(OmegaConf.load(_PATHS_YAML).dataset_dir) diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index fa2095ae15..5cfc0163f4 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -227,6 +227,81 @@ def bimanual_cartesian_layout(width: int) -> dict | None: return BIMANUAL_CARTESIAN_LAYOUTS.get(int(width)) +# Wrist-first bimanual hand-keypoint layouts, per hand +# [wrist xyz (3) | wrist rot (3 ypr or 6 rot6d) | 21 MANO keypoints (63)] +# x {left, right}. ``wrist_xyz`` / ``rot`` index the wrist pose channels, +# ``keypoints`` the flattened (21, 3) keypoint block of each hand. +def _keypoint_layout(rot_width: int) -> dict: + per_hand = 3 + rot_width + 63 + wrist_xyz, rot, kp = [], [], [] + for hand in range(2): + o = hand * per_hand + wrist_xyz += list(range(o, o + 3)) + rot += list(range(o + 3, o + 3 + rot_width)) + kp += list(range(o + 3 + rot_width, o + per_hand)) + return { + "wrist_xyz": tuple(wrist_xyz), + "rot": tuple(rot), + "keypoints": tuple(kp), + "per_hand": per_hand, + } + + +BIMANUAL_KEYPOINT_LAYOUTS = { + 138: _keypoint_layout(3), # ypr wrist rotation + 144: _keypoint_layout(6), # rot6d wrist rotation (the default) +} + + +def bimanual_keypoint_layout(width: int) -> dict | None: + """Index layout for a wrist-first bimanual keypoint action/proprio vector. + + Returns a dict with ``wrist_xyz`` / ``rot`` / ``keypoints`` index tuples + and the per-hand block width, or ``None`` for an unrecognized width + (138 = ypr wrist, 144 = rot6d wrist). + + MANO keypoint 0 IS the wrist, so per hand the wrist position appears twice: + once as ``wrist_xyz`` and once as the first keypoint. Deliberate -- 21 + keypoints is the MANO interface every consumer reshapes to -- and it is a + duplicate, not a constant: the only place it degenerates is the wrist-frame + PROPRIO, where kp0 sits at its own frame origin and normalization treats it + as a constant channel (NORM_MIN_RANGE). + """ + return BIMANUAL_KEYPOINT_LAYOUTS.get(int(width)) + + +def hand_channel_mask(hands, width: int): + """Per-hand flags ``(..., 2)`` [left, right] spread over the channels of a bimanual + vector of ``width`` ([left | right] halves), ``(..., width)``; None when ``width`` + is not a known bimanual layout. Works on numpy arrays and tensors.""" + if ( + bimanual_cartesian_layout(width) is None + and bimanual_keypoint_layout(width) is None + ): + return None + half = int(width) // 2 + return hands[..., [0] * half + [1] * half] + + +# Keys whose last dim is one of the layouts above. +_CARTESIAN_KEYS = ("actions_cartesian", "observations.state.ee_pose") +_KEYPOINT_KEYS = ("actions_keypoints", "observations.state.keypoints") + + +def rot6d_channels(zarr_key: str, width: int) -> tuple[int, ...] | None: + """Indices of the rot6d channels of ``zarr_key`` at ``width``. + + ``None`` when the key or the width is not a known 6D layout, including the + ypr layouts: an Euler angle is not a rotation-matrix column and has none of + the properties that make these channels special. + """ + if zarr_key in _CARTESIAN_KEYS and int(width) in (18, 20): + return bimanual_cartesian_layout(width)["rot"] + if zarr_key in _KEYPOINT_KEYS and int(width) == 144: + return bimanual_keypoint_layout(width)["rot"] + return None + + def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray: """ args: @@ -327,10 +402,25 @@ def _split_action_pose(actions): return left_xyz, left_ypr, right_xyz, right_ypr -def _split_keypoints(keypoints, wrist_in_data: bool = False, is_quat: bool = True): +def _split_keypoints( + keypoints, + wrist_in_data: bool = False, + is_quat: bool = True, + is_rot6d: bool = False, +): + """Split a bimanual keypoint vector into its per-hand blocks. + + With ``wrist_in_data`` the layout is wrist-first per hand, + ``[wrist xyz | wrist rot | 21 keypoints]`` x 2, where the rotation block is + a quaternion (4, ``is_quat``), rot6d columns (6, ``is_rot6d``) or ypr (3); + returns the six blocks in order. Without it the vector is just + ``[left keypoints (63) | right keypoints (63)]`` and two blocks are returned. + """ if wrist_in_data: xyz_size = 3 - if is_quat: + if is_rot6d: + angle_size = 6 + elif is_quat: angle_size = 4 else: angle_size = 3 diff --git a/egomimic/utils/run_manifest.py b/egomimic/utils/run_manifest.py new file mode 100644 index 0000000000..e041241da1 --- /dev/null +++ b/egomimic/utils/run_manifest.py @@ -0,0 +1,79 @@ +"""Record which data a run trained on, next to its checkpoints: ``split.json`` +(train / valid dataset keys of each train dataset) and ``lance_versions.json`` +(the Lance table versions each dataset read).""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path + +log = logging.getLogger(__name__) + +SPLIT_FILE = "split.json" +LANCE_FILE = "lance_versions.json" + + +def run_manifest(train_datasets: dict, valid_datasets: dict) -> dict[str, dict]: + """{file name: contents} for ``write``; no Lance file without a Lance dataset.""" + split = { + name: { + "train": sorted(ds.train_collections), + "valid": sorted(ds.valid_collections), + } + for name, ds in train_datasets.items() + if hasattr(ds, "train_collections") + } + lance = { + f"{head}/{name}": ds.source_manifest + for head, datasets in (("train", train_datasets), ("valid", valid_datasets)) + for name, ds in datasets.items() + if getattr(ds, "source_manifest", None) + } + return {SPLIT_FILE: split, **({LANCE_FILE: lance} if lance else {})} + + +def write(root: str | Path, files: dict[str, dict]) -> None: + """Write each file into ``root``. A file already there with other contents (a + requeue or a resume into the same run dir) is kept, and the new one goes to + ``..json``, so the launch that trained the earlier checkpoints + keeps its record.""" + root = Path(root) + root.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + for name, contents in files.items(): + path = root / name + if path.exists(): + if json.loads(path.read_text(encoding="utf-8")) == contents: + continue + path = root / f"{path.stem}.{stamp}{path.suffix}" + path.write_text(json.dumps(contents, indent=1) + "\n", encoding="utf-8") + log.info("run manifest: wrote %s", path) + + +def warn_on_split_drift(ckpt_path: str | None, split: dict) -> list[str]: + """Warn about, and return, the keys the checkpoint's run validated on that this + run trains on: dataset keys added since then reshuffle ``split_dataset_names``. + Reads ``split.json`` of the run that wrote ``ckpt_path`` (``/checkpoints/``).""" + if not ckpt_path or not Path(ckpt_path).is_file(): + return [] + prev_file = Path(ckpt_path).resolve().parent.parent / SPLIT_FILE + if not prev_file.is_file(): + return [] + prev = json.loads(prev_file.read_text(encoding="utf-8")) + leaked = [] + for name, cur in split.items(): + moved = sorted(set(prev.get(name, {}).get("valid", [])) & set(cur["train"])) + if moved: + log.warning( + "dataset %r: %d episode(s) the checkpoint's run validated on are in " + "this run's train split (dataset keys changed since %s), e.g. %s; " + "its val metrics no longer measure held-out data", + name, + len(moved), + prev_file, + moved[:5], + ) + leaked += moved + return leaked diff --git a/egomimic/rldb/zarr/utils.py b/egomimic/utils/seed.py similarity index 100% rename from egomimic/rldb/zarr/utils.py rename to egomimic/utils/seed.py diff --git a/egomimic/utils/video_utils.py b/egomimic/utils/video_utils.py index 6042a6ae90..5cd1f5a1f3 100644 --- a/egomimic/utils/video_utils.py +++ b/egomimic/utils/video_utils.py @@ -114,3 +114,16 @@ def resize_video_thwc(video: np.ndarray) -> np.ndarray: for t in range(video.shape[0]): out[t] = cv2.resize(video[t], (640, 480), interpolation=cv2.INTER_AREA) return out + + +def open_av1_encoder(path: Path, h: int, w: int, fps: float, ffmpeg: str): + """ffmpeg reading rgb24 frames on stdin into the delivery-spec AV1 MP4.""" + cmd = [ + ffmpeg, "-hide_banner", "-loglevel", "error", "-y", + "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(int(fps)), "-i", "-", + "-c:v", "libsvtav1", "-g", "2", "-crf", "30", "-preset", "12", "-pix_fmt", "yuv420p", + # Index at the front, so a reader over HTTP range requests can seek without fetching the tail. + "-movflags", "+faststart", + str(path), + ] # fmt: skip + return subprocess.Popen(cmd, stdin=subprocess.PIPE) diff --git a/egomimic/utils/viz_utils.py b/egomimic/utils/viz_utils.py index 69b8d34fbf..0e3bf8c591 100644 --- a/egomimic/utils/viz_utils.py +++ b/egomimic/utils/viz_utils.py @@ -293,79 +293,128 @@ def _viz_gaze( ) -def _viz_keypoints( +# MANO slots per hand. The keypoint action layouts are +# ``[wrist xyz | wrist rot | 21 keypoints] x 2`` (rot6d 144-D, quat 140-D, +# ypr 138-D) or bare ``[21 keypoints] x 2`` (126-D). +_N_KP = 21 +_KEYPOINT_WIDTHS = {2 * (3 * _N_KP + pose) for pose in (0, 6, 7, 9)} +# Each finger runs base -> tip, so the tips are the last slot of each group of +# four (slot 0 is the wrist). +FINGERTIP_SLOTS = (4, 8, 12, 16, 20) + + +def _viz_keypoints_horizon_trace( image, actions, intrinsics, - edges, colors, - edge_ranges, - dot_color=None, + wrist_color=None, + fingertip_indices=FINGERTIP_SLOTS, + finger_names=("thumb", "index", "middle", "ring", "pinky"), **kwargs, ): - """Visualize all 21 MANO keypoints per hand, projected onto the image.""" + """Draw the horizon TRACE of the 5 fingertips + the wrist, per hand. + + One polyline per tracked slot across the whole chunk, with a filled circle + at the first drawable step and an open square at the last. The wrist comes + from the action's wrist xyz block when the layout has one (144 / 140 / 138) + and from MANO slot 0 otherwise (126-D). + + Args: + actions: ``(T, D)`` chunk, or ``(D,)`` for a single step. + colors: ``{finger_name: RGB}``, as ``Human.FINGER_COLORS``. + wrist_color: RGB for the wrist trace; ``None`` repeats Human.DOT_COLOR. + """ alpha = kwargs.get("alpha", 1.0) image = _prepare_viz_image(image) - base = image.copy() vis = base.copy() h, w = vis.shape[:2] - if actions.shape[-1] == 140: - _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( - actions, wrist_in_data=True - ) - elif actions.shape[-1] == 138: - _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( - actions, wrist_in_data=True, is_quat=False + actions = np.asarray(actions) + if actions.ndim == 1: + actions = actions[None] + if actions.ndim != 2 or actions.shape[-1] not in _KEYPOINT_WIDTHS: + raise ValueError( + f"keypoint trace expects (T, D) with D in {sorted(_KEYPOINT_WIDTHS)}, " + f"got {actions.shape}" ) + + if actions.shape[-1] == 2 * (3 * _N_KP + 9): + blocks = _split_keypoints(actions, wrist_in_data=True, is_rot6d=True) + elif actions.shape[-1] == 2 * (3 * _N_KP + 7): + blocks = _split_keypoints(actions, wrist_in_data=True) + elif actions.shape[-1] == 2 * (3 * _N_KP + 6): + blocks = _split_keypoints(actions, wrist_in_data=True, is_quat=False) else: - left_keypoints, right_keypoints = _split_keypoints(actions, wrist_in_data=False) - keypoints = {} - keypoints["left"] = left_keypoints.reshape(-1, 3) - keypoints["right"] = right_keypoints.reshape(-1, 3) - _default_dot_colors = {"left": (0, 120, 255), "right": (255, 80, 0)} - for hand in ("left", "right"): - hand_dot_color = ( - dot_color if dot_color is not None else _default_dot_colors[hand] + left_kps, right_kps = _split_keypoints(actions, wrist_in_data=False) + blocks = (None, None, left_kps, None, None, right_kps) + left_wrist, _, left_kps, right_wrist, _, right_kps = blocks + + left_kps = left_kps.reshape(-1, _N_KP, 3) + right_kps = right_kps.reshape(-1, _N_KP, 3) + if left_wrist is None: + left_wrist, right_wrist = left_kps[:, 0, :], right_kps[:, 0, :] + + if wrist_color is None: + # Human.DOT_COLOR's value, duplicated: human.py imports this module, + # so importing Human here would be a cycle. + wrist_color = (255, 165, 0) + + def _project_and_draw(points_cam, color, thickness=2, marker_size=4): + """``(T, 3)`` camera-frame points -> polyline + start / end markers.""" + points_cam = np.asarray(points_cam) + if points_cam.shape[0] == 0: + return + finite = np.isfinite(points_cam).all(axis=-1) & (np.abs(points_cam) < 1e8).all( + axis=-1 ) - kps_cam = keypoints[hand] - # Camera frame -> pixels - kps_px = cam_frame_to_cam_pixels(kps_cam, intrinsics) # (42, 3+) 21 per arm - - # Identify valid keypoints (z > 0 and in image bounds) - valid = kps_cam[:, 2] > 0.01 - valid &= (kps_px[:, 0] >= 0) & (kps_px[:, 0] < w) - valid &= (kps_px[:, 1] >= 0) & (kps_px[:, 1] < h) - - # Draw skeleton edges (colored by finger) - for finger, start, end in edge_ranges: - color = colors[finger] - for edge_idx in range(start, end): - i, j = edges[edge_idx] - if valid[i] and valid[j]: - p1 = (int(kps_px[i, 0]), int(kps_px[i, 1])) - p2 = (int(kps_px[j, 0]), int(kps_px[j, 1])) - cv2.line(vis, p1, p2, color, 2) - - # Draw keypoint dots on top - for k in range(21): - if valid[k]: - center = (int(kps_px[k, 0]), int(kps_px[k, 1])) - cv2.circle(vis, center, 4, hand_dot_color, -1) - cv2.circle(vis, center, 4, (255, 255, 255), 1) # white border - - # Label wrist - if valid[0]: - wrist_px = (int(kps_px[0, 0]) + 6, int(kps_px[0, 1]) - 6) - cv2.putText( + safe = np.where(finite[:, None], points_cam, [0.0, 0.0, -1.0]) + with np.errstate(divide="ignore", invalid="ignore"): + px = cam_frame_to_cam_pixels(safe, intrinsics) # (T, 3+) + valid = finite & (points_cam[:, 2] > 0.01) + valid &= (px[:, 0] >= 0) & (px[:, 0] < w) + valid &= (px[:, 1] >= 0) & (px[:, 1] < h) + pts = np.round(np.nan_to_num(px[:, :2])).astype(np.int32) + + # Connect consecutive drawable steps; a gap breaks the line rather + # than jumping across the frame. + prev = None + for t in range(pts.shape[0]): + if not valid[t]: + prev = None + continue + cur = (int(pts[t, 0]), int(pts[t, 1])) + if prev is not None: + cv2.line(vis, prev, cur, color, thickness, cv2.LINE_AA) + prev = cur + + first_t = next((t for t in range(pts.shape[0]) if valid[t]), None) + last_t = next((t for t in range(pts.shape[0] - 1, -1, -1) if valid[t]), None) + if first_t is not None: + cv2.circle( vis, - f"{hand[0].upper()}", - wrist_px, - cv2.FONT_HERSHEY_SIMPLEX, - 0.5, - hand_dot_color, - 2, + (int(pts[first_t, 0]), int(pts[first_t, 1])), + marker_size, + color, + -1, + cv2.LINE_AA, + ) + if last_t is not None and last_t != first_t: + ex, ey = int(pts[last_t, 0]), int(pts[last_t, 1]) + r = marker_size + 1 + cv2.rectangle( + vis, (ex - r, ey - r), (ex + r, ey + r), color, 2, cv2.LINE_AA + ) + + for wrist_xyz, kps in ((left_wrist, left_kps), (right_wrist, right_kps)): + _project_and_draw(wrist_xyz, wrist_color, thickness=2, marker_size=5) + for tip_idx, finger_name in zip(fingertip_indices, finger_names): + _project_and_draw( + kps[:, tip_idx, :], + colors.get(finger_name, (200, 200, 200)), + thickness=2, + marker_size=4, ) if alpha < 1.0: diff --git a/pi05.md b/pi05.md index ddcd12b643..e380422434 100644 --- a/pi05.md +++ b/pi05.md @@ -25,3 +25,120 @@ uv run --active examples/convert_jax_model_to_pytorch.py \ --checkpoint_dir ../../egomimic/algo/pi_checkpoints/pi05_base \ --output_path ../../egomimic/algo/pi_checkpoints/pi05_base_pytorch ``` + +## Starting from PaliGemma instead of the pi0.5 base + +`model/pi0.5_paligemma_init.yaml` is an overlay that swaps the pi0.5 base +checkpoint for the public PaliGemma release: the VLM prefix (SigLIP tower, +multimodal projector, Gemma-2B) is loaded from `paligemma-3b-pt-224`, and +everything pi0.5 adds on top — the 300M action expert, the action in/out +projections, the adaRMS time MLPs — trains from scratch. That is the ablation +that separates "what the VLM knows" from "what pi0.5's robot pretraining added". + +A ready-made arm ships for the mecka 6D configs: + +```bash +python egomimic/trainHydra.py --config-name=train_zarr_cartesian_pi \ + data=mecka_fold_freeform_opscale1_6d \ + model=pi0.5_bc_mecka_6d_paligemma_init \ + evaluator=eval_pi_wristframe_6d +``` + +Layer the overlay onto any other `pi0.5_*` config the same way (`defaults: +[pi0.5_bc_aria, pi0.5_paligemma_init, _self_]`), or set the two keys per run: + +```bash +model.robomimic_model.config.pytorch_weight_path=null \ +model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +``` + +Setting both paths raises: the pi0.5 base checkpoint already contains a trained +PaliGemma, so letting one silently win would make a run's provenance +unreadable from its config. + +### Getting the weights + +`google/paligemma-3b-pt-224` is license-gated. Either authenticate + +```bash +huggingface-cli login # with a token that has accepted the Gemma terms +``` + +or set `paligemma_allow_mirror: true` to let `egomimic/models/paligemma_init.py` +fall back to a third-party mirror on a 401/403 — it compares every shard's +sha256 against the canonical repo's file metadata (readable without a token) +before using it, so a mirror is transport only and cannot substitute different +weights. It is off by default so the weights of record come from the canonical +repo. `$EGOVERSE_PALIGEMMA_WEIGHTS` points the arm at a local snapshot +directory and skips the Hub entirely. + +Pre-fetch once before a multi-rank run, or every rank races on the same +download lock: + +```bash +python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +``` + +### What the loader checks + +The released checkpoint does not drop in unchanged, so the load is strict about +the two places it differs and fails loudly anywhere else: + +- the checkpoint's `vision_tower.vision_model.*` keeps its `vision_model.` hop + below transformers 5 and loses it from 5 on (`SiglipVisionModel` was + flattened); the loader maps onto whichever the model has. +- the release pads the token embedding from PaliGemma's 257152-token vocabulary + to 257216 rows (`` plus 63 unused). openpi builds the table at exactly + 257152, so `embed_tokens` / `lm_head` are truncated on their vocabulary axis; + the padding is at the end, so every real text token keeps its row. + +Any parameter the checkpoint does not cover, and any checkpoint tensor with no +home in the module tree, raises — a partial load that quietly leaves half the +VLM random is the failure mode this exists to prevent. + +### Hyperparameters + +The overlay changes the init source and nothing else: the optimizer stays +`pi0.5_base`'s fine-tuning recipe (constant 3e-5 after 2000 warmup steps). A +randomly initialized action expert may want a higher peak and a longer warmup — +that is a per-run decision, e.g. + +```bash +model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +``` + +## transformers 4.53 vs 5.x + +`external/openpi` pins `transformers==4.53.2` and its `transformers_replace` +patch is written against that version (`models/siglip/check.py` asserts it). +So does `pyproject.toml`, and the install steps copy that patch into the venv: +4.53.2 plus the patch is the supported setup, what CI installs and what the +shared `emimic` venv carries. Which transformers a run gets is whatever its venv +has, so check before comparing numbers across machines. + +A venv on transformers 5.x (the PaliGemma-init arm was first run on 5.17 with a +partial hand-port of the patch: the Gemma modules carry openpi's adaRMS changes, +but `modeling_paligemma.py` and `modeling_siglip.py` are effectively stock 5.x) +also works, but three of openpi's assumptions do not survive it: + +- `PaliGemmaModel.get_image_features` now returns a `BaseModelOutputWithPooling` + with the projected features on `pooler_output`; openpi's `embed_prefix` reads + `.shape` off the result. +- transformers 5 dropped the back-compat properties that forwarded + `PaliGemmaForConditionalGeneration.language_model` / `.vision_tower` / + `.multi_modal_projector` to the inner `PaliGemmaModel`; openpi reads all three + off the outer model. +- openpi's float32 keep-list names the SigLIP patch and position embeddings + `vision_tower.vision_model.embeddings.*`; with `SiglipVisionModel` flattened + nothing matches and they would train in bf16. + +`egomimic/models/openpi_compat.py` bridges all three, applied from +`egomimic.algo.pi`. Each shim is guarded on what it finds and stands down +entirely below transformers 5, so porting the fork's patch forward makes the +module a no-op — delete it then. + +On 5.x only the PaliGemma-init arm loads. The pi0.5 base checkpoint +(`pi05_base_pytorch/model.safetensors`) stores the nested +`vision_tower.vision_model.*` names and `safetensors.torch.load_model` is +strict, so the default init path stops at startup on missing SigLIP keys. +Train from the pi0.5 base on the pinned 4.53.2. diff --git a/pyproject.toml b/pyproject.toml index 9d1213ec3b..acc7877a05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ "tslearn", "typing_extensions", "pyarrow", + "pylance==13.0.0", "simplejpeg", "prettytable", "datasets==4.0.0", @@ -65,6 +66,10 @@ dependencies = [ "openai", "scikit-learn>=1.8.0", "scipy", + # Vendor MCAP deliveries (egomimic/rldb/mcap) + "mcap==1.5.0", + "mcap-protobuf-support==0.5.4", + "foxglove-schemas-protobuf==0.4.0", ] # Not declared: `openpi` (egomimic/algo/pi.py). The fork at external/openpi pins @@ -84,7 +89,12 @@ where = ["."] include = ["egomimic*"] [tool.setuptools.package-data] -egomimic = ["hydra_configs/**/*.yaml", "resources/*"] +egomimic = [ + "hydra_configs/**/*.yaml", + "resources/*", + "resources/blocklists/*.txt", + "rldb/mcap/proto/*.proto", +] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/reap_gpu_orphans.sh b/reap_gpu_orphans.sh new file mode 100755 index 0000000000..e277b6d006 --- /dev/null +++ b/reap_gpu_orphans.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# Find (and optionally kill) your own dataloader workers that outlived a cancelled +# job and are still pinning GPU memory. +# +# A worker forked after CUDA init holds its rank's /dev/nvidia* fds, so the rank's +# whole context stays resident even after the rank is gone -- nvidia-smi then +# reports memory against a PID that is no longer in /proc. ProctrackType on this +# cluster is proctrack/linuxproc, which tracks the PPID tree, so workers that got +# reparented to init are never swept. +# +# ./reap_gpu_orphans.sh gpu19 # report only +# ./reap_gpu_orphans.sh --kill gpu19 # reclaim +# ./reap_gpu_orphans.sh --kill gpu11 gpu19 +# +# Only processes that are yours, orphaned (PPID 1) and holding an nvidia fd are +# ever signalled, so this is safe to run against a node with live jobs on it. +set -uo pipefail + +KILL=0 +[[ ${1:-} == --kill ]] && { KILL=1; shift; } +if [[ $# -eq 0 ]]; then + echo "usage: $0 [--kill] [node...]" >&2 + exit 2 +fi + +PROBE=' +ME=$(id -u); victims=() +for d in /proc/[0-9]*; do + p=${d#/proc/} + [[ $(stat -c %u "$d" 2>/dev/null) == "$ME" ]] || continue + # PPid from status, as gpu_orphans._ppid reads it: in stat, field 2 is comm, + # which may contain spaces and shift every field after it. + [[ $(awk "/^PPid:/{print \$2}" "$d/status" 2>/dev/null) == 1 ]] || continue + ls -l "$d/fd" 2>/dev/null | grep -q /dev/nvidia || continue + victims+=("$p") +done +echo "node=$(hostname -s) orphans=${#victims[@]}" +[[ ${#victims[@]} -eq 0 ]] && exit 0 +for p in "${victims[@]}"; do + echo " $p $(cat /proc/$p/comm 2>/dev/null) rss=$(awk "/VmRSS/{print \$2}" /proc/$p/status 2>/dev/null)kB age=$(ps -o etime= -p $p 2>/dev/null | tr -d " ")" +done +echo "before: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" +if [[ $KILL == 1 ]]; then + kill -9 "${victims[@]}" 2>/dev/null + for i in 1 2 3 4 5 6; do + sleep 4 + n=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | wc -l) + [[ $n -eq 0 ]] && break + done + echo "after: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" + echo "compute apps left: $n" +fi +' + +probe() { # probe ; returns srun's status, not grep's + "$@" bash -c "KILL=$KILL; $PROBE" 2>&1 | grep -v '^srun: job' + return "${PIPESTATUS[0]}" +} + +rc=0 +for node in "$@"; do + echo "=== $node ===" + # --immediate everywhere: the partition is EXCLUSIVE, so a fresh allocation on a + # busy node queues forever, and a step inside a job whose training step already + # holds every CPU never starts. Better to say so than to hang. + job=$(squeue -h -u "$USER" -w "$node" -t R -o %i | head -1) + if [[ -n $job ]]; then + echo "(node held by your job $job; attaching to it)" + probe srun --overlap --immediate=20 --jobid="$job" -n1 -c 1 && continue + echo "(could not get a step in $job -- it is probably mid-training)" + fi + probe srun -p loaner -A loaner --nodelist="$node" -n1 -c 2 --mem=4G --overlap \ + --immediate=30 -t 5 && continue + echo "$node is busy; rerun when it is idle. A job starting there cleans it" >&2 + echo "automatically anyway (egomimic.utils.gpu_orphans.reap_orphans)." >&2 + rc=1 +done +exit "$rc" diff --git a/tests/fixtures/ddp_video_gate.py b/tests/fixtures/ddp_video_gate.py new file mode 100644 index 0000000000..75cd906285 --- /dev/null +++ b/tests/fixtures/ddp_video_gate.py @@ -0,0 +1,130 @@ +"""Run a 2-rank DDP fit where only rank 0 iterates the video loader. + +Launched as a subprocess by tests/unit/test_eval_loaders.py. Each rank writes +what it saw to ``/rank.json``; the test reads them back. Module scope +(not a closure) so the ddp launcher can re-import it. +""" + +from __future__ import annotations + +import json +import multiprocessing as mp +import sys +from pathlib import Path + +import numpy as np +import torch +from lightning import LightningModule, Trainer +from lightning.fabric.plugins.environments import LightningEnvironment + +from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + +HUMAN = "human_bimanual" + + +class _Episode: + """Counts reads in shared memory: the loader workers are other processes.""" + + def __init__(self, n: int): + self.n = n + self.reads = mp.Value("i", 0) + + def __len__(self): + return self.n + + def __getitem__(self, i): + with self.reads.get_lock(): + self.reads.value += 1 + return {"frame": np.float32(i)} + + +class _Module(LightningModule): + """Mimics ModelWrapper's collective pattern: sync_dist metrics off the + metric loader, nothing at all off the video loader.""" + + def __init__(self, pinned): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + self.pinned = pinned + self.video_steps: list[int] = [] + + def _val_heads(self): + return {"valid": _Viz(self)} + + def training_step(self, batch, batch_idx): + return self.layer(batch[HUMAN]["frame"].reshape(-1, 1)).sum() + + def validation_step(self, batch, batch_idx, dataloader_idx=0): + if dataloader_idx == 1: + self.video_steps.append(self.current_epoch) + return # the video loader logs nothing, as EvalVideo does + self.log("val/mse", torch.tensor(1.0), sync_dist=True, add_dataloader_idx=False) + + def on_validation_end(self): + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.barrier() + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + +class _Viz: + def __init__(self, module): + self.module = module + + def _should_viz(self) -> bool: + return self.module.current_epoch == 1 + + +def main(out_dir: str) -> None: + pinned, metric = _Episode(8), _Episode(8) + dm = MultiDataModuleWrapper( + train_datasets={ + HUMAN: MultiDataset(datasets={"ep0": _Episode(8)}, mode="total") + }, + valid_datasets={HUMAN: MultiDataset(datasets={"ep0": metric}, mode="total")}, + train_dataloader_params={HUMAN: {"batch_size": 2, "num_workers": 0}}, + # num_workers > 0 on purpose: with 0 a DataLoader prefetches nothing, so + # a gate that still builds an iterator would read no frame and look clean. + valid_dataloader_params={ + HUMAN: {"batch_size": 2, "num_workers": 2, "shuffle": False} + }, + video_datasets={ + "valid": {HUMAN: MultiDataset(datasets={"pin": pinned}, mode="total")} + }, + ) + model = _Module(pinned) + trainer = Trainer( + strategy="ddp", + # force the subprocess launcher: run under `srun pytest` lightning would + # otherwise detect SLURM and wait for ranks srun was never asked for + plugins=[LightningEnvironment()], + accelerator="cpu", + devices=2, + max_epochs=2, + limit_train_batches=2, + limit_val_batches=2, + check_val_every_n_epoch=1, + num_sanity_val_steps=2, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=out_dir, + ) + trainer.fit(model, datamodule=dm) + Path(out_dir, f"rank{trainer.global_rank}.json").write_text( + json.dumps( + { + "rank": trainer.global_rank, + "video_steps": model.video_steps, + "pinned_reads": pinned.reads.value, + "finished": True, + } + ) + ) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tests/fixtures/lance_tables.py b/tests/fixtures/lance_tables.py new file mode 100644 index 0000000000..534a30124f --- /dev/null +++ b/tests/fixtures/lance_tables.py @@ -0,0 +1,246 @@ +"""Synthetic Lance steps rows for unit tests.""" + +import json + +import lance +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from egomimic.rldb.lance.schema import ( + ALIASES, + ANNOTATIONS, + ARM, + EVA_STEPS, + HUMAN_STEPS, + RECORDINGS, + TASKS, + annotation_id, + recording_id, + task_instance_id, +) + +PERIOD_NS = 33_333_333 +EVA_DELIVERY_FRAMES = 6 + + +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): + return pa.FixedSizeListArray.from_arrays( + pa.array(np.asarray(x, np.float32).reshape(-1)), n + ) + + +def steps_table(rid, n, offset, valid=None): + t = np.arange(n, dtype=np.float32) + hand = lambda s: pa.StructArray.from_arrays( # noqa: E731 + [ + vec(np.c_[t + offset + s, np.zeros((n, 6))], 7), + vec(np.c_[t + offset + s, np.zeros((n, 6))], 7), + vec(np.zeros((n, 63)), 63), + vec(np.zeros((n, 63)), 63), + pa.array(np.ones(n, np.float32)), + ], + fields=list(HUMAN_STEPS.field("left").type), + ) + left_valid = np.ones(n, bool) + left_valid[4] = False + flags = { + "left_valid": left_valid, + "right_valid": np.ones(n, bool), + "head_pose_valid": np.ones(n, bool), + **(valid or {}), + } + return pa.table( + { + "recording_id": pa.array([rid] * n), + "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), + **{k: pa.array(v) for k, v in flags.items()}, + "eye_gaze": vec(np.zeros((n, 3)), 3), + }, + schema=HUMAN_STEPS, + ) + + +def eva_steps_table(rid, n, offset, arms=("left", "right"), valid=None): + """Robot rows: obs_ee_pose x = frame + offset (+0.5 left, +0.25 right), cmd x one + frame ahead; an arm not in `arms` is zeros and flagged False.""" + t = np.arange(n, dtype=np.float32) + + def arm(side, s): + on = float(side in arms) + pose = lambda x: vec( # noqa: E731 + on * np.c_[x, np.zeros((n, 3)), np.ones(n), np.zeros((n, 2))], 7 + ) + return pa.StructArray.from_arrays( + [ + pose(t + offset + s), + vec(np.zeros((n, 6)), 6), + pa.array(np.full(n, 0.5 * on, np.float32)), + pose(t + offset + s + 1), + vec(np.zeros((n, 6)), 6), + pa.array((on * t / n).astype(np.float32)), + ], + fields=list(ARM), + ) + + flags = {f"{a}_valid": np.full(n, a in arms) for a in ("left", "right")} + flags.update(valid or {}) + return pa.table( + { + "recording_id": pa.array([rid] * n), + "timestamp_ns": pa.array(np.arange(n, dtype=np.int64) * PERIOD_NS), + "left": arm("left", 0.5), + "right": arm("right", 0.25), + **{k: pa.array(v) for k, v in flags.items()}, + }, + schema=EVA_STEPS, + ) + + +def task(rid, f0, f1, name="fold_clothes", description="folding clothes"): + """A TASKS row over frames [f0, f1) of a `steps_table` recording.""" + t0 = f0 * PERIOD_NS + return dict( + task_instance_id=task_instance_id(rid, name, t0), + recording_id=rid, + t_start_ns=t0, + t_end_ns=f1 * PERIOD_NS, + task_name=name, + task_description=description, + objects=[], + ) + + +def alias(rid, f0, f1, episode_hash): + return dict( + episode_hash=episode_hash, + recording_id=rid, + t_start_ns=f0 * PERIOD_NS, + t_end_ns=f1 * PERIOD_NS, + ) + + +def write_catalog( + tmp_path, + annotations=lambda recs: [], + tasks=lambda recs: [task(rid, 0, n) for rid, _, n in recs], + aliases=lambda recs: [alias(rid, 0, n, h) for rid, h, n in recs], + valid=lambda rid, n: {}, + lab="aria", +): + """Steps, recordings, tasks, aliases and annotations tables for two recordings of + `lab` under ``tmp_path``; returns their sorted ``(recording_id, episode_hash, + num_frames)``. ``annotations(recs)``, ``tasks(recs)`` and ``aliases(recs)`` give + those rows (by default one task and one alias per whole recording); ``valid(rid, + n)`` overrides validity columns of a recording's steps. For "eva" the first + recording is bimanual and the second left-arm only; aria ones are bimanual.""" + recs = [ + (recording_id(lab, "1"), "2025-01-01-00-00-00-000000", 12), + (recording_id(lab, "2"), "2025-01-01-00-00-01-000000", 8), + ] + embodiment = {recs[0][0]: "human_bimanual", recs[1][0]: "human_bimanual"} + if lab == "eva": + embodiment = {recs[0][0]: "eva_bimanual", recs[1][0]: "eva_left_arm"} + recs.sort() + + def rows(i, rid, n): + if lab != "eva": + return steps_table(rid, n, 100 * i, valid(rid, n)) + arms = ("left", "right") if embodiment[rid] == "eva_bimanual" else ("left",) + return eva_steps_table(rid, n, 100 * i, arms, valid(rid, n)) + + steps = pa.concat_tables([rows(i, rid, n) for i, (rid, _, n) in enumerate(recs)]) + lance.write_dataset(steps, tmp_path / "steps" / f"{lab}.lance") + 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=lab, + site="rl2", + embodiment=embodiment[rid], + 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(tasks(recs), schema=TASKS), tmp_path / "tasks.lance" + ) + lance.write_dataset( + pa.Table.from_pylist(aliases(recs), schema=ALIASES), + tmp_path / "aliases.lance", + ) + lance.write_dataset( + pa.Table.from_pylist(annotations(recs), schema=ANNOTATIONS), + tmp_path / "annotations.lance", + ) + return recs + + +def eva_delivery(out, vid, frames_in_video=EVA_DELIVERY_FRAMES, offset=0): + """An eva_to_lance-shaped delivery directory for a bimanual demo of EVA_DELIVERY_FRAMES frames.""" + d = out / vid + d.mkdir(parents=True) + rid = recording_id("eva", vid) + pq.write_table( + eva_steps_table(rid, EVA_DELIVERY_FRAMES, offset), d / "steps.parquet" + ) + from fixtures.videos import write_av1 + + write_av1(d / "front_1.mp4", frames_in_video) + meta = { + "vendor_recording_id": vid, + "recording_id": rid, + "episode_hash": f"clip-{vid}", + "embodiment": "eva_bimanual", + "site": "rl2", + "operator": "op", + "scene": "1", + "task_name": "cup_on_saucer", + "task_description": "cup on saucer", + "objects": ["cup"], + "fps": 30.0, + "num_frames": EVA_DELIVERY_FRAMES, + "videos": ["front_1"], + "t_start_ns": 0, + "t_end_ns": EVA_DELIVERY_FRAMES * PERIOD_NS, + "intrinsics": {"front_1": [1.0] * 9}, + "extrinsics": {}, + } + (d / "meta.json").write_text(json.dumps(meta)) + return d diff --git a/tests/fixtures/mcap_episodes.py b/tests/fixtures/mcap_episodes.py new file mode 100644 index 0000000000..ada6bc673a --- /dev/null +++ b/tests/fixtures/mcap_episodes.py @@ -0,0 +1,603 @@ +"""Tiny synthetic vendor MCAP episodes (spec 0.1), valid unless a breakage is asked for. + +`write_mcap_episode` writes `episode_/episode.mcap` with every topic the spec +requires for the embodiment: ~2 s of 64x48 H.264 video, 200 Hz IMU, 30 Hz poses, etc. +It also writes the optional label topics (two /task spans with a /subtask-annotation +each, in annotation.mcap by default) and, for humans, /eye-gaze. +Each keyword-only switch breaks exactly one thing, so a test can assert the validator +reports that problem and nothing else. +""" + +from __future__ import annotations + +from collections.abc import Callable +from fractions import Fraction +from functools import lru_cache +from pathlib import Path + +import numpy as np +from foxglove_schemas_protobuf.CameraCalibration_pb2 import CameraCalibration +from foxglove_schemas_protobuf.CompressedVideo_pb2 import CompressedVideo +from foxglove_schemas_protobuf.FrameTransforms_pb2 import FrameTransforms +from foxglove_schemas_protobuf.PoseInFrame_pb2 import PoseInFrame +from mcap_protobuf.writer import Writer + +from egomimic.rldb.mcap import schemas, spec + +T0 = 1_760_000_000 * 10**9 # 2025-10-09, absolute Unix ns +DURATION_S = 2 +WIDTH, HEIGHT = 64, 48 +IMU_HZ = 200 +POSE_HZ = 30 +VIDEO_FPS = 30 + + +@lru_cache(maxsize=None) +def _h264_packets(n_frames: int) -> tuple[bytes, ...]: + """Annex B H.264, one frame per packet, keyframe every 2 frames, SPS/PPS repeated.""" + import av + + enc = av.CodecContext.create("libx264", "w") + enc.width, enc.height, enc.pix_fmt = WIDTH, HEIGHT, "yuv420p" + enc.time_base = Fraction(1, VIDEO_FPS) + enc.options = { + "bframes": "0", + "g": "2", + "repeat-headers": "1", + "tune": "zerolatency", + } + packets = [] + for i in range(n_frames): + img = np.full((HEIGHT, WIDTH, 3), i * 4 % 255, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(img, format="rgb24").reformat( + format="yuv420p" + ) + frame.pts = i + packets += [bytes(p) for p in enc.encode(frame)] + packets += [bytes(p) for p in enc.encode(None)] + return tuple(packets) + + +def _times(hz: float) -> list[int]: + return [T0 + round(i * 1e9 / hz) for i in range(int(DURATION_S * hz) + 1)] + + +def _sides(embodiment: str) -> tuple[str, ...]: + return spec.arm_sides(embodiment) + + +def _metadata(embodiment: str, declared: str, episode_id: str) -> dict[str, str]: + info = { + "lab": "aria", + "site": "rl2", + "operator-id": "ab" * 32, + "task": "fold_clothes", + "instruction": "fold the shirt", + "embodiment": declared, + "rig-name": "test_rig", + "device-serial": "SN0001", + "scene": "kitchen_A", + "objects": "shirt,basket", + "start-time-unix": str(T0 // 10**6), + "end-time-unix": str((T0 + DURATION_S * 10**9) // 10**6), + "vendor-recording-id": f"acme-rec-{episode_id}", + "vendor": "acme", + "spec-version": spec.SPEC_VERSION, + "license": "CC-BY-4.0", + } + if spec.profile_for(embodiment) is spec.HUMAN: + info["camera-timestamp-ref"] = "mid" + info["hand-joints"] = ",".join(spec.MANO_JOINTS) + info["shutter"] = "rolling:16ms" + return info + + +def _tf(links: list[tuple[str, str]], bad_rotation: bool) -> FrameTransforms: + tf = FrameTransforms() + for i, (parent, child) in enumerate(links): + x = tf.transforms.add(parent_frame_id=parent, child_frame_id=child) + x.translation.z = 0.1 + x.rotation.w = 2.0 if bad_rotation and i == 0 else 1.0 + return tf + + +def _hand_points(side: str, drift: float) -> list[float]: + """21 MANO joints (canonical order) of an open hand in front of the head: wrist, then + 4 joints per finger fanned out sideways, so wrist + knuckles span a plane.""" + y0 = 0.2 if side == "left" else -0.2 + pts = [(0.4 + drift, y0, 1.2)] + for f in range(5): # thumb .. pinky + for j in range(1, 5): + pts.append( + (0.4 + drift + 0.02 + 0.02 * j, y0 + 0.02 * (f - 2), 1.2 - 0.003 * j) + ) + return [c for p in pts for c in p] + + +def _rigid(z: float) -> list[float]: + T = np.eye(4) + T[:3, 3] = (0.3, 0.0, z) + return T.ravel().tolist() + + +def write_mcap_episode( + root: Path, + embodiment: str = "human_bimanual", + episode_id: str = "0", + *, + # 1. everything there + drop_topic: str | None = None, + drop_metadata_key: str | None = None, + declared_embodiment: str | None = None, + bad_site_slug: bool = False, + metadata_override: dict[str, str] | None = None, + unknown_hand_joint: bool = False, + empty_arm_pose: bool = False, + # 2. video + video_format: str = "h264", + skip_first_keyframe: bool = False, + calibration_width: int = WIDTH, + video_fps: int = VIDEO_FPS, + # 3. clocks + backwards_head_pose: bool = False, + hand_gap_s: float = 0.0, + device_clock_imu_info: bool = False, + # 4. calibration + unlink_main_camera: bool = False, + bad_tf_rotation: bool = False, + bad_upright_rotation: bool = False, + upright_rotation_cw_deg: int = 0, + drop_video_frame: int | None = None, + # 5. imu + imu_hz: int = IMU_HZ, + gravity_removed: bool = False, + # 6. poses + nan_head_pose: bool = False, + zeroed_left_hand: bool = False, + gripper_out_of_range: bool = False, + bad_arm_pose: bool = False, + bad_hand_confidence: bool = False, + hands_wrist_only: bool = False, + bad_wrist_pose: bool = False, + nan_eye_gaze: bool = False, + joint_names_mismatch: bool = False, + multi_frame_video: bool = False, + multi_frame_from: int = 0, + hand_frame: str = spec.WORLD_FRAME, + head_lost_from_s: float | None = None, + gaze_until_s: float | None = None, + # 7. tasks & labels + labels_in: str | None = "annotation", + task_past_end: bool = False, + label_outside_tasks: bool = False, + stray_annotation_topic: bool = False, + label_crosses_tasks: bool = False, + labels_without_tasks: bool = False, + label_past_video: bool = False, + label_at_last_frame: bool = False, + annotation_metadata: bool = False, + calibration_twice: bool = False, + principal_x: float = WIDTH / 2, + repeat_frame_time: bool = False, + tamper: Callable | None = None, + slam_camera: bool = False, + distortion_model: str = "", + distortion_coeffs: tuple[float, ...] = (), + bad_annotation_type: bool = False, +) -> Path: + """Write `root/episode_/episode.mcap` and return the episode dir. + + With no switches set, the episode passes the validator. Each switch breaks one + thing: + + drop_topic leave this topic out (e.g. "/imu-info", "/main-camera-info") + drop_metadata_key leave this episode-info key out + declared_embodiment write this as the "embodiment" key, while the topics + follow `embodiment` + bad_site_slug "site" is "RL2 Lab" (not a lowercase slug) + metadata_override episode-info values to set as given + unknown_hand_joint "hand-joints" names "thumb_base", not in the vocabulary + empty_arm_pose robot: one arm's state `pose` left empty (as in ABC's + release) + video_format CompressedVideo.format on every camera + skip_first_keyframe cameras start on a P-frame + calibration_width width in /main-camera-info (frames stay 64 wide) + video_fps frame rate of every camera + backwards_head_pose two /head-pose messages swap timestamps + hand_gap_s hole of this length in /left-hand-keypoints + device_clock_imu_info /imu-info stamped with device time (1 s), not Unix time + unlink_main_camera /tf_static has no transform for the main camera + bad_tf_rotation first /tf_static transform has a non-unit quaternion + bad_upright_rotation every camera's -info-extra says to turn it 45 degrees + upright_rotation_cw_deg valid: every camera's -info-extra turn (0/90/180/270) + drop_video_frame valid: this frame (odd index = a P-frame before a keyframe) is + missing from every camera, as a dropped frame would be + imu_hz /imu sample rate + gravity_removed /imu reads ~0 m/s^2 at rest + nan_head_pose one /head-pose position is NaN + zeroed_left_hand /left-hand-keypoints all zero with valid=true + gripper_out_of_range robot: one /-ee-state value is 1.5 + bad_arm_pose robot: one /-arm-state pose has a scaled rotation + bad_hand_confidence one /left-hand-keypoints confidence is 1.5 + hands_wrist_only valid: hands carry only wrist_pose, no keypoints + bad_wrist_pose one /left-hand-keypoints wrist_pose has a non-unit rotation + joint_names_mismatch robot: one arm's state lists 5 joint names for 6 joints + multi_frame_video /main-camera (or /top-camera) packs 2 frames per message + multi_frame_from ...only from this message on (past the validator's sample) + hand_frame valid: frame_id the hands are given in ("world" or "head") + head_lost_from_s valid: /head-pose-quality is LOST from this second on + gaze_until_s valid: /eye-gaze stops at this second + nan_eye_gaze one valid /eye-gaze message has a NaN yaw + labels_in where the label topics go: "annotation" (annotation.mcap), + "episode" (inline, as ABC does) or None (left out) + task_past_end the second /task ends 3 s after the recorded data + label_outside_tasks only the first /task is written, so the second label is + in no task + stray_annotation_topic annotation.mcap also carries /instruction + label_crosses_tasks the first label runs on into the second task + labels_without_tasks valid: labels but no /task (the episode info is the task) + label_past_video the point label sits after the last video frame + label_at_last_frame valid: the point label sits on the last video frame + annotation_metadata annotation.mcap carries its own session-metadata record + calibration_twice every camera's -info is published twice + principal_x cx in every camera's K (default: the image centre) + repeat_frame_time frame 21 of every camera repeats frame 20's timestamp + slam_camera valid: also a /slam-left-camera, linked to the head + tamper tamper(topic, msg, t) -> None | (msg, t), applied to every + message before it is written (both files): any breakage + distortion_model, distortion_coeffs every camera's -info lens model and D + bad_annotation_type the labels are typed "summary" (only "language" is known) + """ + human = spec.profile_for(embodiment) is spec.HUMAN + sides = _sides(embodiment) + out = Path(root) / f"episode_{episode_id}" + out.mkdir(parents=True, exist_ok=True) + + info = _metadata(embodiment, declared_embodiment or embodiment, episode_id) + if drop_metadata_key: + info.pop(drop_metadata_key) + if bad_site_slug: + info["site"] = "RL2 Lab" + info.update(metadata_override or {}) + if unknown_hand_joint: + info["hand-joints"] = info["hand-joints"].replace("thumb_cmc", "thumb_base") + + def keep(topic: str) -> bool: + return topic != drop_topic + + n_frames = DURATION_S * video_fps + packets = _h264_packets(n_frames) + if skip_first_keyframe: + packets = packets[1:] + video_times = [ + T0 + round((i + skip_first_keyframe) * 1e9 / video_fps) + for i in range(len(packets)) + ] + if repeat_frame_time: + video_times[21] = video_times[20] + if drop_video_frame is not None: + packets = packets[:drop_video_frame] + packets[drop_video_frame + 1 :] + del video_times[drop_video_frame] + + if human: + cameras = {"/main-camera": "main-camera"} + if slam_camera: + cameras["/slam-left-camera"] = "slam-left" + links = [("head", "main-camera"), ("head", "imu"), ("head", "et")] + if slam_camera: + links.append(("head", "slam-left")) + if unlink_main_camera: + links.remove(("head", "main-camera")) + else: + cameras = {"/top-camera": "top-camera"} + links = [("world", "top-camera")] + for s in sides: + cameras[f"/{s}-wrist-camera"] = f"{s}-wrist-camera" + links += [("world", f"{s}_arm_base"), (f"{s}_ee", f"{s}-wrist-camera")] + + with open(out / spec.EPISODE_FILE, "wb") as f, Writer(f) as w: + + def write(topic: str, msg, t: int) -> None: + if tamper and (changed := tamper(topic, msg, t)): + msg, t = changed + if keep(topic): + w.write_message(topic, msg, log_time=t, publish_time=t) + + # mcap_protobuf's Writer has no metadata call; its underlying mcap writer does. + w._writer.add_metadata(spec.METADATA_RECORD, info) + + for topic, frame_id in cameras.items(): + width = calibration_width if topic == "/main-camera" else WIDTH + write( + topic + spec.CALIBRATION_SUFFIX, + CameraCalibration( + frame_id=frame_id, + width=width, + height=HEIGHT, + K=[50, 0, principal_x, 0, 50, HEIGHT / 2, 0, 0, 1], + distortion_model=distortion_model, + D=list(distortion_coeffs), + ), + T0, + ) + if calibration_twice: + write( + topic + spec.CALIBRATION_SUFFIX, + CameraCalibration(frame_id=frame_id, width=width, height=HEIGHT), + T0 + 1, + ) + write( + topic + spec.CAMERA_EXTRA_SUFFIX, + schemas.CameraInfoExtra( + frame_id=frame_id, + valid_radius=0.0, + upright_rotation_cw_deg=45 + if bad_upright_rotation + else upright_rotation_cw_deg, + ), + T0, + ) + pairs = list(zip(packets, video_times)) + if multi_frame_video and topic in ("/main-camera", "/top-camera"): + k = multi_frame_from + pairs = pairs[:k] + [ + (b"".join(d for d, _ in pairs[i : i + 2]), pairs[i][1]) + for i in range(k, len(pairs), 2) + ] + for data, t in pairs: + write( + topic, + CompressedVideo(frame_id=frame_id, data=data, format=video_format), + t, + ) + + write("/tf_static", _tf(links, bad_tf_rotation), T0) + write("/instruction", schemas.Instruction(data=info["instruction"]), T0) + + if human: + _write_human( + write, + sides, + backwards_head_pose=backwards_head_pose, + hand_gap_s=hand_gap_s, + device_clock_imu_info=device_clock_imu_info, + imu_hz=imu_hz, + gravity_removed=gravity_removed, + nan_head_pose=nan_head_pose, + zeroed_left_hand=zeroed_left_hand, + bad_hand_confidence=bad_hand_confidence, + hands_wrist_only=hands_wrist_only, + bad_wrist_pose=bad_wrist_pose, + nan_eye_gaze=nan_eye_gaze, + hand_frame=hand_frame, + head_lost_from_s=head_lost_from_s, + gaze_until_s=gaze_until_s, + ) + else: + _write_robot( + write, + sides, + empty_arm_pose=empty_arm_pose, + gripper_out_of_range=gripper_out_of_range, + bad_arm_pose=bad_arm_pose, + joint_names_mismatch=joint_names_mismatch, + ) + if labels_in == "episode": + _write_labels( + write, + task_past_end, + label_outside_tasks, + label_crosses_tasks, + labels_without_tasks, + label_past_video, + label_at_last_frame, + bad_annotation_type, + ) + + if labels_in != "annotation": # rewriting an episode: no stale labels file + (out / spec.ANNOTATION_FILE).unlink(missing_ok=True) + if labels_in == "annotation": + with open(out / spec.ANNOTATION_FILE, "wb") as f, Writer(f) as w: + + def write_label(topic: str, msg, t: int) -> None: + if tamper and (changed := tamper(topic, msg, t)): + msg, t = changed + w.write_message(topic, msg, log_time=t, publish_time=t) + + _write_labels( + write_label, + task_past_end, + label_outside_tasks, + label_crosses_tasks, + labels_without_tasks, + label_past_video, + label_at_last_frame, + bad_annotation_type, + ) + if annotation_metadata: + w._writer.add_metadata(spec.METADATA_RECORD, {"lab": "other"}) + if stray_annotation_topic: + write_label("/instruction", schemas.Instruction(data="fold"), T0) + return out + + +def _write_labels( + write, + task_past_end: bool, + label_outside_tasks: bool, + label_crosses_tasks: bool, + labels_without_tasks: bool, + label_past_video: bool, + label_at_last_frame: bool, + bad_annotation_type: bool, +) -> None: + """Two back-to-back tasks, one ranged label inside the first and a point label + inside the second.""" + half = DURATION_S * 10**9 // 2 + end = T0 + 2 * half + (3 * 10**9 if task_past_end else 0) + tasks = [ + (T0, T0 + half, "fold_clothes", ["shirt"]), + (T0 + half, end, "pack_clothes", ["shirt", "basket"]), + ] + for start, stop, name, objects in tasks[ + : 0 if labels_without_tasks else 1 if label_outside_tasks else 2 + ]: + write( + "/task", + schemas.Task( + name=name, + description=name.replace("_", " "), + objects=objects, + end_time_unix_ns=stop, + ), + start, + ) + kind = "summary" if bad_annotation_type else "language" + first_end = T0 + half + half // 2 if label_crosses_tasks else T0 + half // 2 + write( + "/subtask-annotation", + schemas.Annotation( + data="pick up the shirt", + end_time_unix_ns=first_end, + annotation_type=kind, + annotation_vendor="acme", + ), + T0 + half // 5, + ) + write( + "/subtask-annotation", + schemas.Annotation(data="drop it in the basket", annotation_type=kind), + T0 + 2 * half - 10**6 + if label_past_video + else T0 + round((DURATION_S * VIDEO_FPS - 1) * 1e9 / VIDEO_FPS) + if label_at_last_frame + else T0 + half + half // 3, + ) + + +def _write_human( + write, + sides, + *, + backwards_head_pose, + hand_gap_s, + device_clock_imu_info, + imu_hz, + gravity_removed, + nan_head_pose, + zeroed_left_hand, + bad_hand_confidence, + hands_wrist_only, + bad_wrist_pose, + nan_eye_gaze, + hand_frame, + head_lost_from_s, + gaze_until_s, +) -> None: + write( + "/imu-info", + schemas.ImuInfo( + frame_id="imu", + accel_noise_density=1e-3, + accel_random_walk=1e-4, + gyro_noise_density=1e-4, + gyro_random_walk=1e-5, + ), + 10**9 if device_clock_imu_info else T0, + ) + for i, t in enumerate(_times(imu_hz)): + m = schemas.Imu(frame_id="imu") + m.linear_acceleration.z = 0.02 if gravity_removed else spec.GRAVITY + m.angular_velocity.x = 0.01 * np.sin(i / 20) + write("/imu", m, t) + + pose_times = _times(POSE_HZ) + if backwards_head_pose: + pose_times[10], pose_times[11] = pose_times[11], pose_times[10] + for i, t in enumerate(pose_times): + p = PoseInFrame(frame_id=spec.WORLD_FRAME) + p.pose.position.x = 0.001 * i # slow drift: real trackers never hold still + p.pose.position.z = 1.5 + p.pose.orientation.w = 1.0 + if nan_head_pose and i == 5: + p.pose.position.x = float("nan") + write("/head-pose", p, t) + + for i, t in enumerate(_times(POSE_HZ)): + if gaze_until_s is not None and t > T0 + gaze_until_s * 1e9: + break + g = schemas.EyeGaze(frame_id="et", yaw=0.1, pitch=-0.3, depth=0.8, valid=True) + if nan_eye_gaze and i == 5: + g.yaw = float("nan") + write("/eye-gaze", g, t) + + ok = schemas.TrackingQuality.OK + for i, t in enumerate(_times(POSE_HZ)): + lost = head_lost_from_s is not None and t >= T0 + head_lost_from_s * 1e9 + state = schemas.TrackingQuality.LOST if lost else ok + write("/head-pose-quality", schemas.TrackingQuality(state=state), t) + for s in sides: + in_gap = s == "left" and T0 + 0.5e9 < t < T0 + (0.5 + hand_gap_s) * 1e9 + if in_gap: + continue + zero = zeroed_left_hand and s == "left" + hand = _hand_points(s, drift=0.001 * i) + k = schemas.HandKeypoints( + frame_id=hand_frame, + positions=[] if hands_wrist_only else [0.0] * 63 if zero else hand, + confidence=1.5 + if bad_hand_confidence and s == "left" and t == T0 + else 0.9, + joint_confidence=[] if hands_wrist_only else [0.9] * 21, + valid=True, + ) + p = k.wrist_pose.position + p.x, p.y, p.z = hand[:3] # the wrist joint + k.wrist_pose.orientation.w = ( + 2.0 if bad_wrist_pose and s == "left" and t == T0 else 1.0 + ) + write(f"/{s}-hand-keypoints", k, t) + + +def _write_robot( + write, + sides, + *, + empty_arm_pose, + gripper_out_of_range, + bad_arm_pose, + joint_names_mismatch, +): + joints = [0.1, -0.2, 0.3, 0.0, 0.5, -0.1] + names = [f"joint{j}" for j in range(1, 7)] + for i, t in enumerate(_times(POSE_HZ)): + for s in sides: + pose = _rigid(0.2 + 0.001 * i) + if bad_arm_pose and s == sides[-1] and i == 5: + pose[0] = 2.0 # R[0, 0]: no longer a rotation + last = s == sides[-1] + write( + f"/{s}-arm-state", + schemas.ArmState( + frame_id=f"{s}_arm_base", + joint_names=names[:5] if joint_names_mismatch and last else names, + position=joints, + velocity=[0.0] * 6, + torque=[0.5] * 6, + ee_pose=[] if empty_arm_pose and last else pose, + ), + t, + ) + write( + f"/{s}-arm-action", + schemas.ArmAction( + frame_id=f"{s}_arm_base", + joint_names=names, + position=joints, + ee_pose=_rigid(0.2 + 0.001 * (i + 1)), + ), + t, + ) + g = 1.5 if gripper_out_of_range and last and i == 5 else 0.4 + write(f"/{s}-ee-state", schemas.GripperState(position=g), t) + write(f"/{s}-ee-action", schemas.GripperState(position=0.4), t) diff --git a/tests/fixtures/recipes.py b/tests/fixtures/recipes.py index bce90a939f..91de4221e0 100644 --- a/tests/fixtures/recipes.py +++ b/tests/fixtures/recipes.py @@ -22,13 +22,10 @@ class Recipe: def _pi(data: str, model: str, emb: str) -> Recipe: - return Recipe( - "train_zarr_cartesian_pi", - data, - model, - emb, - (f"+data.train_datasets.{emb}.resolver.key_map.annotation_key=annotations",), - ) + # data/eva.yaml already emits the continuous-6D (wrist-frame) layout Pi + # requires, the same mode cotrain_pi_base / eva_pi set. + extra = [f"+data.train_datasets.{emb}.resolver.key_map.annotation_key=annotations"] + return Recipe("train_zarr_cartesian_pi", data, model, emb, tuple(extra)) RECIPES: dict[tuple[str, str], Recipe] = { @@ -44,6 +41,11 @@ def _pi(data: str, model: str, emb: str) -> Recipe: ("scale", "hpt"): Recipe( "train_zarr_cartesian", "scale", "hpt_bc_flow_scale", "human_bimanual" ), + # RDT's shipped config is the 6D cartesian one; rdt_small_overrides points + # it at the keypoint action this data config emits. + ("mecka", "rdt"): Recipe( + "train_zarr_cartesian", "mecka", "rdt_bc_mecka_6d_dinov3_vitb", "human_bimanual" + ), # Pi trains on the same per-vendor data configs as HPT: datasets emit one # camera naming and the Pi wrapper renames onto openpi's slots. The plain # vendor configs load no language field (that is a data choice the *_lang @@ -124,6 +126,29 @@ def hpt_small_overrides(emb: str) -> list[str]: ] +def rdt_small_overrides(emb: str) -> list[str]: + """Tiny DiT, random-init tiny DINOv3 tower, no Qwen stem (network-free); + the 2-frame image history the shipped train config turns on.""" + rm = "model.robomimic_model" + enc = f"{rm}.encoder_specs.front_img_1" + return [ + f"{rm}.dims.{emb}.action=144", + f"{rm}.ac_keys.{emb}=actions_keypoints", + f"{rm}.width=32", + f"{rm}.trunk.depth=2", + f"{rm}.trunk.num_heads=4", + f"{rm}.head_specs.{emb}.num_inference_steps=2", + f"~{rm}.shared_stem_specs.annotation", + f"{rm}.shared_obs_keys=[front_img_1]", + f"{rm}.annotation_key=null", + f"{enc}.model_name=vit_small_patch16_dinov3", + f"{enc}.image_size=[32,48]", + f"+{enc}.pretrained=false", + f"+{enc}.tower_kwargs={{embed_dim: 32, depth: 1, num_heads: 2}}", + f"+data.train_datasets.{emb}.resolver.key_map.image_history_gap_s=0.1", + ] + + def pi_cpu_overrides() -> list[str]: return [ "model.robomimic_model.config.pytorch_weight_path=null", diff --git a/tests/fixtures/synthetic_episodes.py b/tests/fixtures/synthetic_episodes.py index b4317a057f..a63c8324ae 100644 --- a/tests/fixtures/synthetic_episodes.py +++ b/tests/fixtures/synthetic_episodes.py @@ -101,6 +101,16 @@ def write_episode( } # what the poses were composed through if v.has_head_pose: numeric["obs_head_pose"] = _pose(T, rng) + if v.embodiment == "human_bimanual": + # Hand keypoints (the default human action): a constant wrist pose per + # hand plus 21 MANO keypoints held at fixed offsets from it, so the + # wrist-frame keypoint action is constant in time like the poses. + for side in ("left", "right"): + wrist = _pose(T, rng) + offsets = rng.normal(0, 0.03, (1, 21, 3)) + kp = wrist[:, None, :3] + offsets # (T, 21, 3) world frame + numeric[f"{side}.obs_wrist_pose"] = wrist + numeric[f"{side}.obs_keypoints"] = kp.reshape(T, 63) return ZarrWriter.create_and_write( root / f"{vendor}_{seed:02d}.zarr", numeric_data=numeric, diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 1dc9db0db4..025692925f 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -14,9 +14,11 @@ import egomimic.trainHydra as train_hydra from fixtures.synthetic_episodes import write_episode -# Width of the concatenated proprio the Pi wrapper hands openpi: the -# 6D-encoded ee_pose, grip-padded to the robot 20-D layout on both. -STATE_DIM = {"eva_bimanual": 20, "human_bimanual": 20} +# Width of the concatenated proprio the Pi wrapper hands openpi: eva = the +# 6D-encoded 20-D ee_pose; human = the 144-D keypoint proprio + the 20-D +# grip-padded ee_pose the keypoint data configs also emit for the prompt +# (include_ee_pose + pad_proprio_gripper). +STATE_DIM = {"eva_bimanual": 20, "human_bimanual": 144 + 20} PI_TOKENIZER = ( "google/paligemma-3b-mix-224" # tokenizer_model_name in model/pi0.5_base.yaml ) @@ -71,7 +73,12 @@ def _hf_offline(monkeypatch) -> None: monkeypatch.setenv("HF_HUB_OFFLINE", "1") monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) - monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + # transformers <= 4 caches the flag in `_is_offline_mode`; 5.x reads it + # through `is_offline_mode()` instead. Whichever this version has. + if hasattr(transformers.utils.hub, "_is_offline_mode"): + monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + else: + monkeypatch.setattr(transformers.utils.hub, "is_offline_mode", lambda: True) def write_fixtures( @@ -131,6 +138,8 @@ class StubPI0(nn.Module): def __init__(self, config): super().__init__() self.config = config + # Like openpi: 32-wide regardless of config.action_dim. PI swaps them + # for action_dim-wide ones when action_dim != 32. self.action_in_proj = nn.Linear(32, 64) self.action_out_proj = nn.Linear(64, 32) @@ -153,11 +162,20 @@ def forward(self, observation, actions): k, mask, ) - assert actions.shape == (B, self.config.action_horizon, 32), actions.shape + action_dim = self.config.action_dim + assert actions.shape == ( + B, + self.config.action_horizon, + action_dim, + ), actions.shape + assert self.action_in_proj.in_features == action_dim + assert self.action_out_proj.out_features == action_dim pred = self.action_out_proj(torch.tanh(self.action_in_proj(actions))) return ((pred - actions) ** 2).mean(dim=(-1, -2)) @torch.no_grad() def sample_actions(self, device, observation, noise=None, num_steps=10): B = observation.state.shape[0] - return torch.zeros(B, self.config.action_horizon, 32, device=device) + return torch.zeros( + B, self.config.action_horizon, self.config.action_dim, device=device + ) diff --git a/tests/fixtures/videos.py b/tests/fixtures/videos.py new file mode 100644 index 0000000000..68ad29d1ce --- /dev/null +++ b/tests/fixtures/videos.py @@ -0,0 +1,24 @@ +"""Tiny AV1 MP4s for tests, encoded with PyAV (libaom).""" + +import numpy as np + + +def write_av1(path, n: int, faststart: bool = True, gop: int = 2) -> None: + """n distinct 64x48 frames (a white bar growing by 4 px per frame), keyframe every `gop`.""" + import av + + out = av.open( + str(path), "w", options={"movflags": "+faststart"} if faststart else {} + ) + s = out.add_stream("libaom-av1", rate=30) + s.width, s.height, s.pix_fmt = 64, 48, "yuv420p" + s.options = {"cpu-used": "8", "crf": "30"} + s.codec_context.gop_size = gop + for i in range(n): + img = np.zeros((48, 64, 3), np.uint8) + img[:, : min(64, 4 * (i + 1))] = 255 + for p in s.encode(av.VideoFrame.from_ndarray(img, format="rgb24")): + out.mux(p) + for p in s.encode(): + out.mux(p) + out.close() diff --git a/tests/integration/test_lance_r2.py b/tests/integration/test_lance_r2.py new file mode 100644 index 0000000000..70a8ac35a4 --- /dev/null +++ b/tests/integration/test_lance_r2.py @@ -0,0 +1,107 @@ +"""The committer on R2: concurrent Lance commits there are safe, and staged deliveries go +through commit_staged into an s3:// root that LanceStreamingResolver then reads. + +Needs EGOVERSE_TEST_R2_PREFIX (an s3:// scratch prefix; each run writes under a fresh +subprefix and deletes it) and R2_* credentials. CPU only. +""" + +import concurrent.futures as cf +import os +import uuid + +import pyarrow as pa +import pytest + +lance = pytest.importorskip("lance") +pytest.importorskip("av") + +from egomimic.rldb.lance import commit # noqa: E402 +from egomimic.rldb.lance.dataset import LanceStreamingResolver, open_table # noqa: E402 + + +@pytest.fixture +def prefix(): + base = os.environ.get("EGOVERSE_TEST_R2_PREFIX") + if not base: + pytest.skip("set EGOVERSE_TEST_R2_PREFIX") + p = commit.join(base, f"test-{uuid.uuid4().hex[:12]}") + yield p + commit.remove([commit.join(p, f) for f in commit.list_files(p)]) + + +def test_concurrent_appends_all_land(prefix): + uri = commit.join(prefix, "t.lance") + commit._append(uri, pa.table({"w": pa.array([], pa.int64())})) + writers = 6 + + def write(w): + commit._append(uri, pa.table({"w": [w] * 100})) + + with cf.ThreadPoolExecutor(writers) as ex: + list(ex.map(write, range(writers))) + ds = open_table(uri) + assert ds.count_rows() == writers * 100 + assert sorted(set(ds.to_table().column("w").to_pylist())) == list(range(writers)) + assert ds.version == writers + 1 + + +def test_staged_deliveries_commit_to_r2_and_stream(prefix, tmp_path): + from fixtures.lance_tables import eva_delivery + + from egomimic.scripts.lance_commit import commit_staged + + staging, root = commit.join(prefix, "staging"), commit.join(prefix, "lance") + for vid in ("100", "200"): + d = eva_delivery(tmp_path / vid, vid) + commit.stage(d, commit.join(staging, "eva", vid)) + report = commit_staged(staging, "eva", root, tmp_path / "work") + (batch,) = report["batches"] + assert len(batch["added"]) == 2 and report["rejected"] == {} + assert commit.list_files(commit.join(staging, "eva")) == {} + + from egomimic.rldb.lance.annotations import review + + rid = batch["added"][0] + t0 = next( + t + for t in open_table(commit.join(root, "tasks.lance")).to_table().to_pylist() + if t["recording_id"] == rid + ) + rv = review( + root, + pa.Table.from_pylist( + [ + dict( + op="add", + recording_id=rid, + t_start_ns=t0["t_start_ns"], + t_end_ns=t0["t_end_ns"], + annotation_type="language", + text="lift the cup", + annotation_vendor="scale", + ) + ] + ), + ) + assert not rv.errors + ann = commit.commit_annotations(root, rv.adds, rv.deletes) + assert ann["rows"] == {"added": 1, "deleted": 0} and ann["added"] == [rid] + + key_map = { + "img": {"key_type": "camera_keys", "zarr_key": "images.front_1"}, + "left.cmd_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.cmd_ee_pose", + "horizon": 2, + }, + } + leaves = LanceStreamingResolver( + root, + lab="eva", + key_map=key_map, + cameras={"images.front_1": "front_1"}, + tag=batch["tag"], + ).resolve() + leaf = next(iter(leaves.values())) + (sample,) = leaf.fetch_planned([(leaf, leaf.plan_sample(3))]) + assert sample[1]["img"].shape == (1, 48, 64, 3) diff --git a/tests/integration/test_lance_ranges.py b/tests/integration/test_lance_ranges.py new file mode 100644 index 0000000000..3884f2cc7c --- /dev/null +++ b/tests/integration/test_lance_ranges.py @@ -0,0 +1,58 @@ +"""Streaming leaves (batched MP4 byte-range reads) match the local per-sample MP4 path. + +Needs EGOVERSE_TEST_LANCE_ROOT (pilot tables). CPU only. +""" + +import os +import random +from pathlib import Path + +import pytest +import torch + +pytest.importorskip("lance") +pytest.importorskip("torchcodec") + +from egomimic.rldb.embodiment.human import Human # noqa: E402 +from egomimic.rldb.lance.dataset import ( # noqa: E402 + LanceEpisodeResolver, + LanceStreamingResolver, +) +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset # noqa: E402 + + +@pytest.fixture(scope="module") +def root(): + src = os.environ.get("EGOVERSE_TEST_LANCE_ROOT") + if not src: + pytest.skip("set EGOVERSE_TEST_LANCE_ROOT") + return Path(src) + + +def _dataset(root, resolver=LanceEpisodeResolver, **kw): + leaves = resolver( + root, + key_map=Human.get_keymap(keymap_mode="keypoints"), + transform_list=Human.get_transform_list( + mode="keypoints_headframe_6d", stride=3 + ), + cameras={"images.front_1": "front_1"}, + **kw, + ).resolve(expected_embodiment="human_bimanual") + return MultiDataset(datasets=leaves, mode="total") + + +def test_range_batch_matches_per_sample_mp4(root): + per_sample = _dataset(root) + batched = _dataset(root, LanceStreamingResolver) + assert len(batched) == len(per_sample) + rng = random.Random(0) + idxs = [rng.randrange(len(batched)) for _ in range(48)] + for i, got in zip(idxs, batched.__getitems__(idxs), strict=True): + want = per_sample[i] + assert got.keys() == want.keys() + for k, v in want.items(): + if isinstance(v, torch.Tensor): + assert torch.equal(got[k], v), (i, k) + else: + assert got[k] == v, (i, k) diff --git a/tests/integration/test_lance_versions.py b/tests/integration/test_lance_versions.py new file mode 100644 index 0000000000..30a2b02dfc --- /dev/null +++ b/tests/integration/test_lance_versions.py @@ -0,0 +1,268 @@ +"""Phase 7 of the Lance pilot: versioned commits on a copy of the pilot tables. + +v1 commit 24 recordings, tag t1; v2 append 10 one commit each; v3 replace one with +shifted poses; v4 edit a task's text; v5 add a column; v6 compact, tag t6, clean up. +Needs EGOVERSE_TEST_LANCE_ROOT (the pilot tables written by commit.py). CPU only. +""" + +import datetime +import json +import os +import random +import time +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pytest +import torch + +lance = pytest.importorskip("lance") + +from egomimic.rldb.lance import commit # noqa: E402 +from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 + +CAMERA = "front_1" +SAMPLES_PER_LEAF = 3 + + +@pytest.fixture(scope="module") +def pilot(): + root = os.environ.get("EGOVERSE_TEST_LANCE_ROOT") + if not root or not Path(root).exists(): + pytest.skip("set EGOVERSE_TEST_LANCE_ROOT to the pilot Lance tables") + return Path(root) + + +def _deliveries(src: Path) -> list[commit.Delivery]: + """Rebuild a delivery per bimanual pilot recording from the committed tables.""" + recs = lance.dataset(src / "recordings.lance").to_table().to_pylist() + tasks = { + r["recording_id"]: r + for r in lance.dataset(src / "tasks.lance").to_table().to_pylist() + } + aliases = { + r["recording_id"]: r + for r in lance.dataset(src / "aliases.lance").to_table().to_pylist() + } + steps = lance.dataset(src / "steps" / "aria.lance") + out = [] + for r in sorted(recs, key=lambda r: r["recording_id"]): + if r["embodiment"] != "human_bimanual": + continue + rid = r["recording_id"] + videos = {cam: src / p for cam, p in r["video_path"] if cam == CAMERA} + rec = {k: v for k, v in r.items() if k not in ("video_path", "content_digest")} + st = steps.to_table(filter=f"recording_id = '{rid}'") + out.append(commit.Delivery(rec, [tasks[rid]], aliases[rid], st, videos)) + return out + + +def _shift_head(d: commit.Delivery, dx: float) -> commit.Delivery: + head = d.steps.column("head_pose").combine_chunks() + vals = np.asarray(head.flatten()).reshape(-1, 7).copy() + vals[:, 0] += dx + new = pa.FixedSizeListArray.from_arrays(pa.array(vals.reshape(-1)), 7) + steps = d.steps.set_column( + d.steps.schema.get_field_index("head_pose"), "head_pose", new + ) + return commit.Delivery(d.recording, d.tasks, d.alias, steps, d.videos) + + +def _key_map(): + from egomimic.rldb.embodiment.human import Human + + return Human.get_keymap(keymap_mode="keypoints"), Human.get_transform_list( + mode="keypoints_headframe_6d", stride=3 + ) + + +def _resolver(root, **kw): + key_map, transforms = _key_map() + return LanceEpisodeResolver( + root, + key_map=key_map, + transform_list=transforms, + cameras={"images.front_1": CAMERA}, + **kw, + ) + + +def _samples(leaves, only=None): + out = {} + for h, leaf in leaves.items(): + if only is not None and h not in only: + continue + idx = np.linspace(0, len(leaf) - 1, SAMPLES_PER_LEAF).astype(int) + out[h] = [leaf[int(i)] for i in idx] + return out + + +def _same(a, b) -> bool: + if a.keys() != b.keys(): + return False + for k in a: + x, y = a[k], b[k] + if isinstance(x, torch.Tensor): + if x.dtype != y.dtype or x.shape != y.shape or not torch.equal(x, y): + return False + elif x != y: + return False + return True + + +def _all_same(s1, s2, hashes) -> list[str]: + return [ + h + for h in hashes + if not all(_same(a, b) for a, b in zip(s1[h], s2[h], strict=True)) + ] + + +def _row_count_ok(root, **kw) -> bool: + r = _resolver(root, **kw) + _, _, steps = r._open("steps/aria.lance") + _, _, recs = r._open("recordings.lance") + return ( + steps.count_rows() + == pc.sum(recs.to_table(columns=["num_frames"]).column(0)).as_py() + ) + + +def _take_rate(leaves, n=600, seed=0) -> float: + rng = random.Random(seed) + items = list(leaves.values()) + t0 = time.perf_counter() + for _ in range(n): + leaf = rng.choice(items) + s = rng.randrange(leaf.total_frames - 31) + leaf._read_numeric( + { + "obs_head_pose": np.arange(s, s + 31), + "left.obs_ee_pose": np.arange(s, s + 31), + } + ) + return n / (time.perf_counter() - t0) + + +def test_versioned_commits(pilot, tmp_path): + root = tmp_path / "lance" + ds = _deliveries(pilot) + assert len(ds) >= 25 + first, rest = ds[:24], ds[24:] + report = {} + + # v1 + commit.append(root, first) + commit.tag(root, "t1") + leaves1 = _resolver(root).resolve() + s1 = _samples(leaves1) + fp1 = {h: leaf.cache_fingerprint() for h, leaf in leaves1.items()} + assert _row_count_ok(root) + + # v2: one commit per recording, so many fragments + for d in rest: + commit.append(root, [d]) + steps_uri = commit.steps_uri(root, "aria") + report["v2_fragments"] = len(lance.dataset(steps_uri).get_fragments()) + assert report["v2_fragments"] == 1 + len(rest) + r2 = _resolver(root) + leaves2 = r2.resolve() + pinned2 = dict(r2.versions) + s2 = _samples(leaves2) + assert _all_same(s1, s2, s1) == [] + assert {h: leaves2[h].cache_fingerprint() for h in fp1} == fp1 + assert _row_count_ok(root) + + # v3: replace one recording with shifted head poses + victim = first[3] + vh = victim.alias["recording_id"] + commit.replace(root, _shift_head(victim, 0.05)) + leaves3 = _resolver(root).resolve() + s3 = _samples(leaves3) + untouched = [h for h in s1 if h != vh] + assert _all_same(s2, s3, untouched) == [] + assert _all_same(s2, s3, [vh]) == [vh] + assert leaves3[vh].row_start > leaves2[vh].row_start + assert {h: leaves3[h].cache_fingerprint() for h in untouched} == { + h: fp1[h] for h in untouched + } + assert leaves3[vh].cache_fingerprint() != fp1[vh] + assert _row_count_ok(root) + report["take_per_s_v2"] = _take_rate(leaves2) + report["take_per_s_v3"] = _take_rate(leaves3) + + # v4: edit one task's text + edited = first[5] + commit.edit( + root, + "tasks.lance", + pa.table( + { + "task_instance_id": [edited.tasks[0]["task_instance_id"]], + "task_description": ["edited text"], + } + ), + on="task_instance_id", + ) + leaves4 = _resolver(root).resolve() + eh = edited.alias["recording_id"] + assert leaves4[eh].task_text == "edited text" + assert leaves1[eh].task_text != "edited text" + assert _resolver(root, tag="t1").resolve()[eh].task_text == leaves1[eh].task_text + + # v5: add a column without rewriting existing data files + files_before = { + f.path + for frag in lance.dataset(steps_uri).get_fragments() + for f in frag.data_files() + } + commit.add_columns( + root, "steps/aria.lance", {"both_valid": "left_valid AND right_valid"} + ) + steps5 = lance.dataset(steps_uri) + files_after = {f.path for frag in steps5.get_fragments() for f in frag.data_files()} + assert files_before < files_after + assert len(files_after - files_before) == len(steps5.get_fragments()) + leaves5 = _resolver(root).resolve() + assert {h: leaves5[h].cache_fingerprint() for h in untouched} == { + h: fp1[h] for h in untouched + } + + # A reader pinned at v2 still serves v2 data after v3-v5. + assert _all_same(s2, _samples(leaves2), s2) == [] + assert ( + _all_same(s2, _samples(_resolver(root, versions=pinned2).resolve()), s2) == [] + ) + + # v6: compact, tag, clean up + commit.index(root) + commit.compact(root) + assert len(lance.dataset(steps_uri).get_fragments()) == 1 + commit.tag(root, "t6") + leaves6 = _resolver(root).resolve() + assert _all_same(s3, _samples(leaves6), s3) == [] + assert {h: leaves6[h].cache_fingerprint() for h in leaves6} == { + h: leaves3[h].cache_fingerprint() for h in leaves3 + } + assert ( + _all_same(s2, _samples(leaves2), s2) == [] + ) # compaction leaves old files in place + report["take_per_s_v6"] = _take_rate(leaves6) + + commit.cleanup(root, older_than=datetime.timedelta(0)) + report["versions_after_cleanup"] = { + str(Path(u).relative_to(root)): [ + v["version"] for v in lance.dataset(u).versions() + ] + for u in commit.table_uris(root) + } + t1 = _resolver(root, tag="t1").resolve() + assert _all_same(s1, _samples(t1), s1) == [] + assert {h: leaf.cache_fingerprint() for h, leaf in t1.items()} == fp1 + assert _row_count_ok(root, tag="t1") and _row_count_ok(root, tag="t6") + with pytest.raises(Exception): + _resolver(root, versions=pinned2).resolve() + + print("\nPHASE7 " + json.dumps(report)) diff --git a/tests/integration/test_mcap_abc.py b/tests/integration/test_mcap_abc.py new file mode 100644 index 0000000000..2fa1766e44 --- /dev/null +++ b/tests/integration/test_mcap_abc.py @@ -0,0 +1,46 @@ +"""Vendor MCAP validator on real episodes. + + EGOVERSE_TEST_ABC_DIR= pytest --integration tests/integration/test_mcap_abc.py + EGOVERSE_TEST_MCAP_DIR= pytest --integration tests/integration/test_mcap_abc.py + +ABC-130k predates the spec (own robot types, no /tf_static or full episode info), so its +episodes must be read end to end and reported, not pass. Episodes delivered to the spec +(e.g. Aria recordings converted to it) must pass. +""" + +import os +from pathlib import Path + +import pytest + +from egomimic.rldb.mcap import spec +from egomimic.rldb.mcap.validate import validate_episode + + +def _episodes(var: str) -> list[Path]: + root = os.environ.get(var) + if not root: + pytest.skip(f"set {var} to a directory of episode_*/ dirs") + eps = sorted(p for p in Path(root).glob("episode_*") if p.is_dir()) + if not eps: + pytest.skip(f"no episode_*/ dirs under {root}") + return eps + + +def test_abc_episodes_are_read_and_reported(): + for ep in _episodes("EGOVERSE_TEST_ABC_DIR"): + r = validate_episode(ep) + unreadable = [p for p in r.problems if p.message.startswith("cannot read")] + assert not unreadable, f"{ep}: {unreadable}" + assert r.problems, f"{ep}: ABC predates the spec, expected problems" + print(f"{ep.name}: {len(r.problems)} problems") + for p in r.problems: + print(f" {p}") + + +def test_spec_deliveries_pass(): + for ep in _episodes("EGOVERSE_TEST_MCAP_DIR"): + r = validate_episode(ep) + assert r.problems == [], f"{ep}:\n" + "\n".join(map(str, r.problems)) + assert r.recording_id, ep + assert (ep / spec.EPISODE_FILE).exists() diff --git a/tests/integration/test_mcap_to_lance_aria.py b/tests/integration/test_mcap_to_lance_aria.py new file mode 100644 index 0000000000..22dfb09d86 --- /dev/null +++ b/tests/integration/test_mcap_to_lance_aria.py @@ -0,0 +1,77 @@ +"""MCAP -> Lance on real spec deliveries, checked against Ryan's aria_to_lance. + + EGOVERSE_TEST_MCAP_DIR= + EGOVERSE_TEST_ARIA_LANCE_DIR= # optional + pytest --integration tests/integration/test_mcap_to_lance_aria.py + +Every episode converts, stages and commits. Where aria_to_lance output for the same recording is +given (run with --no-mano, as here), steps rows line up one to one on the same clock, and +the head pose matches wherever both picked the same pose sample: aria_to_lance uses +projectaria's search, which lands one sample off on about half the frames, while this +converter takes the true nearest sample. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import numpy as np +import pytest + +lance = pytest.importorskip("lance") +pq = pytest.importorskip("pyarrow.parquet") + +from egomimic.scripts.lance_commit import commit_staged # noqa: E402 +from egomimic.scripts.mcap_process.mcap_to_lance import convert_and_stage # noqa: E402 + + +def _codec() -> str: + out = subprocess.run( + [shutil.which("ffmpeg") or "ffmpeg", "-hide_banner", "-encoders"], + capture_output=True, + ).stdout.decode() + for c in ("libsvtav1", "libx264"): + if c in out: + return c + pytest.skip("ffmpeg has neither libsvtav1 nor libx264") + + +def test_real_episodes_convert_and_match_aria_to_lance(tmp_path): + root = os.environ.get("EGOVERSE_TEST_MCAP_DIR") + if not root: + pytest.skip("set EGOVERSE_TEST_MCAP_DIR to a directory of episode_*/ dirs") + eps = sorted(p for p in Path(root).glob("episode_*") if p.is_dir()) + ryan_root = os.environ.get("EGOVERSE_TEST_ARIA_LANCE_DIR") + codec = _codec() + for ep in eps: + r = convert_and_stage( + ep, tmp_path / "staging", tmp_path / "lance", codec=codec, mano=False + ) + assert r.status == "staged", r + commit_staged( + str(tmp_path / "staging"), r.lab, str(tmp_path / "lance"), tmp_path / "work" + ) + assert r.valid_frac["head"] > 0.9, r + if not ryan_root: + continue + ryan = Path(ryan_root) / ep.name.removeprefix("episode_") + if not (ryan / "steps.parquet").exists(): + continue + ours = ( + lance.dataset(tmp_path / "lance" / "steps" / f"{r.lab}.lance") + .to_table(filter=f"recording_id = '{r.recording_id}'") + .to_pydict() + ) + theirs = pq.read_table(ryan / "steps.parquet").to_pydict() + assert len(ours["timestamp_ns"]) == len(theirs["timestamp_ns"]) + offset = np.array(ours["timestamp_ns"]) - np.array(theirs["timestamp_ns"]) + assert np.ptp(offset) == 0 # same frames; ours Unix ns, theirs device ns + both = np.array(ours["head_pose_valid"]) & np.array(theirs["head_pose_valid"]) + a = np.array(ours["head_pose"])[both] + b = np.array(theirs["head_pose"])[both] + dp = np.linalg.norm(a[:, :3] - b[:, :3], axis=1) + # Same pose sample on about half the frames (exact), the neighbor on the rest + # (head moves < 1 mm between 1 kHz samples). + assert (dp < 1e-6).mean() > 0.4 + assert dp.max() < 2e-3 diff --git a/tests/integration/test_paligemma_init_real.py b/tests/integration/test_paligemma_init_real.py new file mode 100644 index 0000000000..e74135417f --- /dev/null +++ b/tests/integration/test_paligemma_init_real.py @@ -0,0 +1,253 @@ +"""The real PaliGemma release, loaded into the real openpi prefix. + +The unit tier proves the remapping and the coverage rules on a miniature module +tree; this is the part that can only be checked against the actual 3B +checkpoint -- that openpi's prefix and `paligemma-3b-pt-224` line up tensor for +tensor, that the values land unchanged, and that nothing outside the prefix +moves. It builds a 3.6B model and reads ~12 GB of weights, so it lives here. + + # the load itself: CPU, needs the checkpoint (or Hub access to fetch it) + python -m pytest --integration tests/integration/test_paligemma_init_real.py -q + # plus the flow-matching step, on a GPU node + python -m pytest --integration -m gpu tests/integration/test_paligemma_init_real.py -q +""" + +from __future__ import annotations + +import glob + +import pytest +import torch +from fixtures.train_harness import pi_unavailable +from safetensors import safe_open + +from egomimic.models.paligemma_init import ( + CANONICAL_REPO, + load_paligemma_weights, + resolve_paligemma_dir, +) + +pytestmark = pytest.mark.skipif( + pi_unavailable() is not None, reason=pi_unavailable() or "" +) + +PG = "paligemma_with_expert.paligemma." +EX = "paligemma_with_expert.gemma_expert." + +# These start random and must still be random afterwards: the whole point of +# this init is that pi0.5's own parameters are NOT pretrained. +FROM_SCRATCH = ( + EX + "model.layers.0.self_attn.q_proj.weight", + EX + "model.layers.17.mlp.up_proj.weight", + EX + "model.layers.0.input_layernorm.dense.weight", + EX + "model.norm.dense.weight", + "action_in_proj.weight", + "action_out_proj.weight", + "time_mlp_in.weight", + "time_mlp_out.weight", +) + +# module parameter -> the checkpoint key it must come from, one per remapping rule +PRETRAINED = { + PG + + "model.language_model.layers.0.self_attn.q_proj.weight": "language_model.model.layers.0.self_attn.q_proj.weight", + PG + + "model.language_model.layers.17.mlp.down_proj.weight": "language_model.model.layers.17.mlp.down_proj.weight", + PG + "model.language_model.norm.weight": "language_model.model.norm.weight", + PG + + "model.language_model.embed_tokens.weight": "language_model.model.embed_tokens.weight", + PG + + "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight": "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.weight", + PG + + "model.vision_tower.encoder.layers.26.mlp.fc2.bias": "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + PG + + "model.vision_tower.embeddings.patch_embedding.weight": "vision_tower.vision_model.embeddings.patch_embedding.weight", + PG + + "model.multi_modal_projector.linear.weight": "multi_modal_projector.linear.weight", +} + + +@pytest.fixture(scope="module") +def checkpoint_dir() -> str: + try: + return resolve_paligemma_dir(CANONICAL_REPO) + except Exception as exc: # gated and offline, or no verified mirror reachable + pytest.skip(f"cannot resolve {CANONICAL_REPO}: {type(exc).__name__}: {exc}") + + +def _flat(named) -> dict: + """Parameters under their transformers-5 names, which the tables above use: + below 5 the SigLIP tree keeps a ``vision_model.`` hop.""" + return { + k.replace("vision_tower.vision_model.", "vision_tower."): v for k, v in named + } + + +@pytest.fixture(scope="module") +def loaded(checkpoint_dir): + """The real pi0.5 chassis with only its VLM prefix initialized.""" + import openpi.models.pi0_config + import openpi.models_pytorch.pi0_pytorch as pi0t + + # Production reaches openpi only through `egomimic.algo.pi`, which applies + # the transformers-5 shims at import. This builds PI0Pytorch directly, so + # take the same entry point rather than an unpatched openpi. + import egomimic.algo.pi # noqa: F401 + + config = openpi.models.pi0_config.Pi0Config( + dtype="bfloat16", + action_dim=32, + action_horizon=100, + max_token_len=180, + paligemma_variant="gemma_2b", + action_expert_variant="gemma_300m", + pi05=True, + ) + torch.manual_seed(0) + model = pi0t.PI0Pytorch(config) + before = _flat((k, v.detach().clone()) for k, v in model.named_parameters()) + load_paligemma_weights(model, checkpoint_dir) + return model, before, _flat(model.named_parameters()) + + +def test_pi05s_own_parameters_are_left_at_their_init(loaded): + _, before, after = loaded + for name in FROM_SCRATCH: + assert torch.equal(after[name], before[name]), f"{name} was overwritten" + + +def test_the_prefix_moved_off_its_init(loaded): + _, before, after = loaded + for name in PRETRAINED: + assert not torch.equal(after[name], before[name]), f"{name} was not loaded" + + +def test_the_prefix_matches_the_checkpoint_bit_for_bit(loaded, checkpoint_dir): + _, _, after = loaded + wanted = set(PRETRAINED.values()) + found = {} + for shard in sorted(glob.glob(checkpoint_dir + "/*.safetensors")): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + if key in wanted: + found[key] = handle.get_tensor(key) + assert set(found) == wanted, f"checkpoint is missing {wanted - set(found)}" + + for name, key in PRETRAINED.items(): + param = after[name] + reference = found[key] + if reference.shape != param.shape: # padded vocabulary + reference = reference[: param.shape[0]] + assert torch.equal( + param, reference.to(param.dtype) + ), f"value mismatch at {name}" + + +def test_openpis_mixed_precision_policy_survives_the_load(loaded): + # The norms openpi keeps in fp32 must not be dragged to bf16 by the copy, + # and the matmul weights must not be widened. + _, _, after = loaded + assert ( + after[PG + "model.language_model.layers.0.input_layernorm.weight"].dtype + is torch.float32 + ) + assert after[PG + "model.language_model.norm.weight"].dtype is torch.float32 + assert ( + after[PG + "model.language_model.layers.0.self_attn.q_proj.weight"].dtype + is torch.bfloat16 + ) + assert ( + after[PG + "model.vision_tower.encoder.layers.0.mlp.fc1.weight"].dtype + is torch.bfloat16 + ) + # openpi's keep-list, on either transformers (openpi_compat on 5.x) + for name in ( + "patch_embedding.weight", + "patch_embedding.bias", + "position_embedding.weight", + ): + assert ( + after[PG + "model.vision_tower.embeddings." + name].dtype is torch.float32 + ) + + +# ------------------------------------------------------------------ a real step + + +def _observation(batch: int, device, max_token_len: int): + """A minimal openpi observation: three camera slots, the wrists masked out + the way human data leaves them, and a prompt of real token ids.""" + from egomimic.models.preprocess_pi_obs import PI_CAMERA_SLOTS, _SimpleObservation + + image = torch.rand(batch, 3, 224, 224, device=device) * 2 - 1 + present = { + "base_0_rgb": True, + "left_wrist_0_rgb": False, + "right_wrist_0_rgb": False, + } + mask = torch.ones(batch, max_token_len, dtype=torch.bool, device=device) + return _SimpleObservation( + images={slot: image.clone() for slot in PI_CAMERA_SLOTS}, + image_masks={ + slot: torch.full((batch,), present[slot], dtype=torch.bool, device=device) + for slot in PI_CAMERA_SLOTS + }, + state=torch.zeros(batch, 32, device=device), + tokenized_prompt=torch.randint( + 0, 1000, (batch, max_token_len), device=device, dtype=torch.long + ), + tokenized_prompt_mask=mask, + token_ar_mask=mask.clone(), + token_loss_mask=mask.clone(), + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_one_flow_matching_step_trains_both_halves(loaded): + """The point of the arm: a finite loss whose gradient reaches the frozen-in + PaliGemma prefix AND the action expert that started from noise. A load that + left half the model detached would still produce a loss here, so check the + gradients, not just the number.""" + model, _, _ = loaded + device = torch.device("cuda") + model = model.to(device) + model.train() + + observation = _observation(1, device, model.config.max_token_len) + actions = torch.randn( + 1, model.config.action_horizon, model.config.action_dim, device=device + ) + + torch.manual_seed(0) + losses = model.forward(observation, actions) + loss = ( + torch.stack(losses).mean() + if isinstance(losses, (list, tuple)) + else losses.mean() + ) + assert torch.isfinite(loss), f"non-finite flow-matching loss: {loss}" + + loss.backward() + probes = { + "prefix (pretrained)": model.paligemma_with_expert.paligemma.model.language_model.layers[ + 0 + ].self_attn.q_proj.weight, + "vision tower (pretrained)": model.paligemma_with_expert.paligemma.model.vision_tower.encoder.layers[ + 0 + ].self_attn.q_proj.weight, + "action expert (from scratch)": model.paligemma_with_expert.gemma_expert.model.layers[ + 0 + ].self_attn.q_proj.weight, + "action_out_proj (from scratch)": model.action_out_proj.weight, + } + for label, param in probes.items(): + assert param.grad is not None, f"no gradient reached {label}" + assert torch.isfinite(param.grad).all(), f"non-finite gradient at {label}" + assert param.grad.abs().sum() > 0, f"zero gradient at {label}" + + # and a step actually moves them + before = {k: v.detach().clone() for k, v in probes.items()} + torch.optim.SGD(model.parameters(), lr=1e-3).step() + for label, param in probes.items(): + assert not torch.equal(param, before[label]), f"{label} did not move" diff --git a/tests/integration/test_train_step_gpu.py b/tests/integration/test_train_step_gpu.py index 561eabd857..9b348a66a2 100644 --- a/tests/integration/test_train_step_gpu.py +++ b/tests/integration/test_train_step_gpu.py @@ -132,7 +132,12 @@ def test_hpt_real_recipe(tmp_path, monkeypatch, vendor): _vendor_filter(cfg, recipe.embodiment, vendor) spy = BatchKeySpy(monkeypatch, HPT) _run(cfg, out) - assert "actions_cartesian" in spy.keys[recipe.embodiment] + action_key = ( + "actions_cartesian" + if recipe.embodiment == "eva_bimanual" + else "actions_keypoints" # human data defaults to hand keypoints + ) + assert action_key in spy.keys[recipe.embodiment] @pytest.mark.parametrize("vendor", VENDOR_NAMES) diff --git a/tests/unit/test_action_chunk_transforms.py b/tests/unit/test_action_chunk_transforms.py index a9b9f9141c..260aee993c 100644 --- a/tests/unit/test_action_chunk_transforms.py +++ b/tests/unit/test_action_chunk_transforms.py @@ -8,10 +8,11 @@ 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, + PoseCoordinateFrameTransform, XYZWXYZ_to_XYZYPR, ) from egomimic.utils.pose_utils import _xyzwxyz_to_matrix @@ -344,6 +345,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: "InterpolatePose", "InterpolateLinear", "InterpolateLinear", + "InterpolatePadMask", "XYZWXYZ_to_XYZYPR", "ConcatKeys", "ConcatKeys", @@ -418,7 +420,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 8, + 9, "XYZWXYZ_to_XYZYPR", base_keys | {"left.cmd_ee_pose_camframe", "right.cmd_ee_pose_camframe"}, { @@ -430,7 +432,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 9, + 10, "ConcatKeys", { "actions_cartesian", @@ -451,7 +453,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 10, + 11, "ConcatKeys", { "actions_cartesian", @@ -468,7 +470,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 11, + 12, "DeleteKeys", {"actions_cartesian", "observations.state.ee_pose"}, { @@ -478,7 +480,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 12, + 13, "NumpyToTensor", {"actions_cartesian", "observations.state.ee_pose"}, { @@ -516,6 +518,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: "PoseCoordinateFrameTransform", "InterpolatePose", "InterpolatePose", + "InterpolatePadMask", "XYZWXYZ_to_XYZYPR", "ConcatKeys", "ConcatKeys", @@ -585,7 +588,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 6, + 7, "XYZWXYZ_to_XYZYPR", base_keys | { @@ -603,7 +606,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 7, + 8, "ConcatKeys", base_keys | { @@ -619,7 +622,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 8, + 9, "ConcatKeys", base_keys | {"actions_cartesian", "observations.state.ee_pose"}, { @@ -629,7 +632,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 9, + 10, "DeleteKeys", {"actions_cartesian", "observations.state.ee_pose"}, { @@ -637,3 +640,86 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: "observations.state.ee_pose": (12,), }, ) + + +# --------------------------------------------------------------------------- +# PoseCoordinateFrameTransform.per_step_target +# --------------------------------------------------------------------------- +def _wrist_history(steps: int) -> np.ndarray: + """A wrist moving +0.1 in x per step, oldest first, identity rotation.""" + poses = np.zeros((steps, 7), dtype=np.float64) + poses[:, 0] = 0.1 * np.arange(steps) + poses[:, 3] = 1.0 # w of wxyz + return poses + + +def _keypoints_on_the_wrist(wrist: np.ndarray) -> np.ndarray: + """(K, 2, 3): slot 0 sits exactly on that step's wrist, slot 1 0.4 above.""" + kp = np.zeros((wrist.shape[0], 2, 3), dtype=np.float64) + kp[:, 0, :] = wrist[:, :3] + kp[:, 1, :] = wrist[:, :3] + np.array([0.0, 0.4, 0.0]) + return kp + + +def _wristframe(wrist, keypoints, per_step_target): + tf = PoseCoordinateFrameTransform( + target_world="wrist", + pose_world="kp_head", + transformed_key_name="kp_wrist", + mode="xyz", + target_history=True, + per_step_target=per_step_target, + ) + return np.asarray(tf.transform({"wrist": wrist, "kp_head": keypoints})["kp_wrist"]) + + +def test_per_step_target_keeps_every_history_step_in_its_own_frame(): + """The whole point: keypoint 0 is the wrist, so it must be 0 at EVERY step. + + Measured from the current step's wrist instead, an older step carries the + wrist's own displacement -- which is what pushed palm-rigid MANO knuckles + outside their per-cell quantile bounds and starved the loader on retries. + """ + wrist = _wrist_history(4) + keypoints = _keypoints_on_the_wrist(wrist) + + current_frame = _wristframe(wrist, keypoints, per_step_target=False) + per_step = _wristframe(wrist, keypoints, per_step_target=True) + + # Slot 0 == that step's wrist, so per-step framing sends it to the origin. + np.testing.assert_allclose(per_step[:, 0, :], 0.0, atol=1e-12) + # Without it, the three older steps keep the wrist's -0.3/-0.2/-0.1 offset + # and only the current step lands at the origin. + np.testing.assert_allclose( + current_frame[:, 0, 0], [-0.3, -0.2, -0.1, 0.0], atol=1e-12 + ) + # The articulation channel is unchanged either way: 0.4 above the wrist. + for framed in (current_frame, per_step): + np.testing.assert_allclose(framed[:, 1, 1], 0.4, atol=1e-12) + + +def test_per_step_target_is_the_default(): + """Anything read with history through this transform wants per-step + framing; the flag is the escape hatch, not the opt-in.""" + import inspect + + from egomimic.rldb.transforms import ( + PoseCoordinateFrameTransform, + ) + + default = ( + inspect.signature(PoseCoordinateFrameTransform.__init__) + .parameters["per_step_target"] + .default + ) + assert default is True + + +def test_per_step_target_is_inert_without_a_history_axis(): + """Existing single-step call sites cannot change behaviour: the per-step + path needs a leading axis on the target, which a K=1 keymap never emits.""" + wrist = _wrist_history(1)[0] # (7,), the shape a no-history keymap gives + keypoints = _keypoints_on_the_wrist(_wrist_history(1))[0] # (2, 3) + off = _wristframe(wrist, keypoints, per_step_target=False) + on = _wristframe(wrist, keypoints, per_step_target=True) + np.testing.assert_array_equal(off, on) diff --git a/tests/unit/test_checkpoint_resume.py b/tests/unit/test_checkpoint_resume.py index 8c8e322e90..0e233e51ee 100644 --- a/tests/unit/test_checkpoint_resume.py +++ b/tests/unit/test_checkpoint_resume.py @@ -5,6 +5,7 @@ decision, and only when last.ckpt exists.""" import inspect +import os from pathlib import Path import pytest @@ -222,3 +223,120 @@ def test_load_checkpoint_weights_warns_on_unexpected_keys(tmp_path, caplog): load_checkpoint_weights(model, ckpt) assert "extra.buffer" in caplog.text assert torch.equal(model.weight, torch.ones(2, 2)) + + +# ------------------------------------------------ save_last writes a link, not 21 GB +def test_checkpoint_callback_links_last_instead_of_copying(compose_resolve, tmp_path): + """`save_last: true` wrote a second full copy of every checkpoint. "link" + symlinks it instead; save_top_k=-1 means the target is never collected.""" + import hydra + + cfg = _pi_cfg(compose_resolve, tmp_path) + assert cfg.callbacks.model_checkpoint.save_last == "link" + assert cfg.callbacks.model_checkpoint.save_top_k == -1, "the link's target" + + # lightning validates save_last at construction (it rejects "link" for a + # non-local dirpath), so instantiating is the real check + callback = hydra.utils.instantiate(cfg.callbacks.model_checkpoint) + assert callback.save_last == "link" + + +def test_requeue_resumes_through_a_symlinked_last_ckpt( + monkeypatch, compose_resolve, tmp_path +): + """What "link" leaves on disk: last.ckpt is a relative symlink to the real + epoch checkpoint, and the requeue path must still find it.""" + real = _touch(tmp_path / "checkpoints" / "epoch_epoch=49.ckpt") + link = tmp_path / "checkpoints" / "last.ckpt" + link.symlink_to(real.name) # relative, as ModelCheckpoint._link_checkpoint writes + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + train_hydra._prepare_checkpoint_resume(cfg) + assert Path(cfg.ckpt_path) == link + assert Path(cfg.ckpt_path).resolve() == real.resolve() + + +def test_requeue_with_a_dangling_last_ckpt_raises( + monkeypatch, compose_resolve, tmp_path +): + """os.path.isfile is False on a dangling link, which used to fall through to + "training starts over" on a run that has checkpoints.""" + link = tmp_path / "checkpoints" / "last.ckpt" + link.parent.mkdir(parents=True) + link.symlink_to("epoch_epoch=49.ckpt") + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + with pytest.raises(FileNotFoundError, match="does not exist"): + train_hydra._prepare_checkpoint_resume(cfg) + + +def test_requeue_with_checkpoints_but_no_last_ckpt_raises( + monkeypatch, compose_resolve, tmp_path +): + """rsync without -l, or a preemption between Lightning's unlink and relink, + leaves epoch checkpoints and no last.ckpt.""" + _touch(tmp_path / "checkpoints" / "epoch_epoch=49.ckpt") + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + with pytest.raises(FileNotFoundError, match="epoch_epoch=49.ckpt"): + train_hydra._prepare_checkpoint_resume(cfg) + + +@pytest.mark.parametrize("save_last", ["link", True]) +def test_save_last_link_writes_a_symlink_not_a_second_copy( + compose_resolve, tmp_path, save_last +): + """The point of the change, through a real fit: `true` writes last.ckpt as a + full second copy of the epoch checkpoint (~21 GB on Pi 0.5); "link" makes it + a symlink to the one already on disk.""" + import hydra + from lightning import LightningModule, Trainer + from torch.utils.data import DataLoader, TensorDataset + + class _M(LightningModule): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + + def training_step(self, batch, batch_idx): + return self.layer(batch[0]).sum() + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + cfg = _pi_cfg(compose_resolve, tmp_path) + ckpt_dir = tmp_path / "checkpoints" + callback = hydra.utils.instantiate( + cfg.callbacks.model_checkpoint, + dirpath=str(ckpt_dir), + every_n_epochs=1, + save_last=save_last, + ) + Trainer( + max_epochs=2, + accelerator="cpu", + devices=1, + limit_train_batches=1, + callbacks=[callback], + logger=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=str(tmp_path), + ).fit(_M(), DataLoader(TensorDataset(torch.zeros(2, 1)), batch_size=1)) + + last = ckpt_dir / "last.ckpt" + epochs = sorted(p.name for p in ckpt_dir.glob("epoch_*.ckpt")) + assert epochs == ["epoch_epoch=0.ckpt", "epoch_epoch=1.ckpt"], "save_top_k=-1" + assert last.is_file(), "either way the requeue path finds it" + + if save_last == "link": + assert last.is_symlink() + assert Path(os.readlink(last)) == Path("epoch_epoch=1.ckpt"), "relative" + assert last.resolve() == (ckpt_dir / "epoch_epoch=1.ckpt").resolve() + else: + # what we are getting rid of: a full second copy of the same checkpoint + assert not last.is_symlink() + assert last.stat().st_size == (ckpt_dir / "epoch_epoch=1.ckpt").stat().st_size diff --git a/tests/unit/test_chunk_renderer.py b/tests/unit/test_chunk_renderer.py new file mode 100644 index 0000000000..411cbb3965 --- /dev/null +++ b/tests/unit/test_chunk_renderer.py @@ -0,0 +1,114 @@ +"""The horizon trace is the only keypoint renderer, and it draws every frame. + +``Human.viz(mode="keypoints")`` has exactly one path: ``_viz_keypoints_horizon_trace``. +The static-skeleton renderer it used to fall back to is gone on purpose -- a +chunk that reached it rendered as a motionless overlay, which is what a +whole 50k-step campaign's validation videos turned out to be (2026-09-18). +""" + +import numpy as np +import pytest + +from egomimic.utils.viz_utils import _viz_keypoints_horizon_trace + +# px = 100 * x / z + 100, py = 100 * y / z + 100 on a 200x200 canvas. +K = np.array([[100, 0, 100, 0], [0, 100, 100, 0], [0, 0, 1, 0]]) + +BEHIND_CAMERA = [0.0, 0.0, -1.0] + +FINGERS = ("thumb", "index", "middle", "ring", "pinky") + + +def _trace(data, **kwargs): + return _viz_keypoints_horizon_trace( + np.zeros((200, 200, 3), np.uint8), + data, + K, + colors={name: (20, 255, 20) for name in FINGERS}, + **kwargs, + ) + + +def _sweep(length): + """A left hand whose slot 0 sweeps across the image, one x per frame. + + Slot 1 sits 0.4 above it; every other slot and the whole right hand stay + behind the camera. Returns a (T, 126) chunk. + """ + points = np.tile(np.array(BEHIND_CAMERA), (length, 2, 21, 1)) + xs = np.linspace(-0.7, 0.7, length) + points[:, 0, 0] = np.column_stack((xs, np.zeros(length), np.ones(length))) + points[:, 0, 1] = points[:, 0, 0] + [0.0, 0.4, 0.0] + return points.reshape(length, -1) + + +def test_horizon_trace_connects_every_step_of_the_chunk(): + """The whole sweep is a polyline, not just its endpoints.""" + image = _trace(_sweep(30)) + for x in np.linspace(-0.7, 0.7, 30): + assert image[100, round(100 + 100 * x)].any() + # the line fills the gaps between sampled steps too + assert image[100, 100].any() + + +def test_horizon_trace_accepts_the_126d_revert_output_and_144d_rot6d(): + """The evaluator hands it 126-D; raw wrist-frame actions are 144-D.""" + assert _trace(_sweep(8)).any() # 126-D, wrist falls back to MANO slot 0 + + hand = np.concatenate([np.zeros(9), np.tile(BEHIND_CAMERA, 21)]) + chunk = np.tile(np.concatenate([hand, hand]), (8, 1)) + chunk[:, 0:3] = np.column_stack( + (np.linspace(-0.5, 0.5, 8), np.zeros(8), np.ones(8)) + ) + assert chunk.shape[-1] == 144 + assert _trace(chunk).any() # wrist xyz block drives the trace + + +def test_a_lone_timestep_matches_its_own_flattened_form(): + """``(1, D)`` and ``(D,)`` are the same single-step trace, not two paths.""" + chunk = _sweep(1) + np.testing.assert_array_equal(_trace(chunk), _trace(chunk.reshape(-1))) + + +def test_non_finite_slots_are_skipped(): + """A NaN or absurd slot drops out of the trace; the finite ones still draw.""" + chunk = _sweep(3) + chunk[1, 0:3] = np.nan # frame 1's slot 0 + chunk[2, 0:3] = 1e12 + image = _trace(chunk) + # _sweep puts frame t's slot 0 at px = 100 * x + 100, i.e. 30 / 100 / 170. + assert image[100, 30].any() # frame 0 still drawn + assert not image[100, 100].any() # frame 1 was NaN + assert not image[100, 170].any() # frame 2 was 1e12 + + +def test_horizon_trace_rejects_a_bad_shape_and_an_unknown_width(): + with pytest.raises(ValueError, match="keypoint trace expects"): + _trace(np.zeros((2, 30, 126))) + with pytest.raises(ValueError, match="keypoint trace expects"): + _trace(np.zeros((4, 127))) + + +def test_human_viz_has_no_renderer_but_the_trace(monkeypatch): + """Every keypoint input shape routes to the trace -- there is no fallback.""" + import egomimic.rldb.embodiment.human as human_mod + from egomimic.rldb.embodiment.human import Human + + assert not hasattr(human_mod, "_viz_keypoints"), ( + "the static skeleton renderer is back; mode='keypoints' must have " + "exactly one path" + ) + + seen = [] + + def _record(**kw): + seen.append(kw["actions"].shape) + return np.zeros((2, 2, 3), np.uint8) + + monkeypatch.setattr(human_mod, "_viz_keypoints_horizon_trace", _record) + image = np.zeros((200, 200, 3), np.uint8) + Human.viz(image, _sweep(8), mode="keypoints", intrinsics=K) + Human.viz(image, _sweep(1), mode="keypoints", intrinsics=K) + Human.viz(image, _sweep(1).reshape(-1), mode="keypoints", intrinsics=K) + + assert seen == [(8, 126), (1, 126), (126,)] diff --git a/tests/unit/test_compile_cache.py b/tests/unit/test_compile_cache.py index b567423cd8..21b01b234c 100644 --- a/tests/unit/test_compile_cache.py +++ b/tests/unit/test_compile_cache.py @@ -32,12 +32,6 @@ def test_outside_slurm_keeps_torch_default(): assert env == {} -def test_explicit_value_is_respected(): - env = {KEY: "/scratch/ind", "SLURM_JOB_ID": "1"} - assert set_per_job_compile_cache_dir(env) == "/scratch/ind" - assert env[KEY] == "/scratch/ind" - - _LAUNCHER = """ import os, subprocess, sys from egomimic.utils.compile_cache import set_per_job_compile_cache_dir as s diff --git a/tests/unit/test_data_configs_compose.py b/tests/unit/test_data_configs_compose.py index 143bda48a8..88841b529c 100644 --- a/tests/unit/test_data_configs_compose.py +++ b/tests/unit/test_data_configs_compose.py @@ -9,6 +9,7 @@ import hydra import pytest +import yaml from hydra import compose, initialize_config_module from omegaconf import OmegaConf @@ -137,6 +138,35 @@ def test_no_data_config_uses_a_pi_keymap_mode(name): assert not str(mode).endswith("_pi"), (name, split, emb, mode) +# Aliases kept so configs saved by earlier runs still rebuild. +KNOWN_MODE_ONLY_CHILDREN = {"cotrain_pi_lang_wrist"} +MODE_KEYS = {"keymap_mode", "mode", "include_ee_pose"} + + +def _leaf_paths(node, prefix=()): + if isinstance(node, dict) and node: + for k, v in node.items(): + yield from _leaf_paths(v, (*prefix, k)) + else: + yield prefix + + +@pytest.mark.parametrize("name", DATA_CONFIGS) +def test_no_data_config_only_swaps_the_action_mode(name): + # A child that only changes keymap/transform mode (e.g. cam- vs wrist-frame) + # is a CLI override, not a data config; one data config per dataset slice. + raw = yaml.safe_load((DATA_DIR / f"{name}.yaml").read_text(encoding="utf-8")) + body = {k: v for k, v in (raw or {}).items() if k != "defaults"} + if not raw.get("defaults") or not body: + return + leaves = list(_leaf_paths(body)) + mode_only = all(path[-1] in MODE_KEYS for path in leaves) + if name in KNOWN_MODE_ONLY_CHILDREN: + assert mode_only, f"{name} is no longer mode-only; drop it from the allowlist" + else: + assert not mode_only, (name, leaves) + + def test_pi_train_config_defaults_to_all_eva_with_annotations(): # data=eva is a one-episode smoke config without annotations; Pi's default # must not silently train on that with an empty prompt. diff --git a/tests/unit/test_dataset_filter.py b/tests/unit/test_dataset_filter.py index fd0766ec45..6ab0d082c9 100644 --- a/tests/unit/test_dataset_filter.py +++ b/tests/unit/test_dataset_filter.py @@ -22,13 +22,6 @@ def test_dataset_filter_matches_rows_and_excludes_deleted_by_default() -> None: assert not filters.matches({"episode_hash": "episode-2"}) -def test_dataset_filter_empty_list_matches_all_non_deleted_rows() -> None: - filters = DatasetFilter() - - assert filters.matches({"episode_hash": "episode-1"}) - assert not filters.matches({"episode_hash": "episode-1", "is_deleted": True}) - - def test_dataset_filter_init_rejects_invalid_filter_and_prints_it(capsys) -> None: with pytest.raises(ValueError, match="Invalid filter"): DatasetFilter(filter_lambdas=["lambda row:"]) @@ -179,3 +172,36 @@ 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"} + + +def test_dataset_filter_exclude_hashes_never_match(tmp_path) -> None: + listed = DatasetFilter(exclude_hashes=["bad"]) + assert not listed.matches({"episode_hash": "bad"}) + assert listed.matches({"episode_hash": "good"}) + + path = tmp_path / "block.txt" + path.write_text("# header\nbad # constant\n\n") + from_file = DatasetFilter(exclude_hashes=str(path)) + assert from_file.exclude_hashes == frozenset({"bad"}) + + +def test_dataset_filter_exclude_hashes_by_blocklist_name() -> None: + filters = DatasetFilter(exclude_hashes="mecka_bad_egomotion") + assert len(filters.exclude_hashes) > 0 + with pytest.raises(FileNotFoundError, match="no_such_blocklist"): + DatasetFilter(exclude_hashes="no_such_blocklist") + + +def test_dataset_filter_drops_excluded_pins_but_not_all_of_them() -> None: + filters = DatasetFilter(episode_hashes=["a", "bad"], exclude_hashes=["bad"]) + assert filters.episode_hashes == frozenset({"a"}) + assert not filters.matches({"episode_hash": "bad"}) + # An empty pin set would mean "match everything". + with pytest.raises(ValueError, match="pinned episode"): + DatasetFilter(episode_hashes=["bad"], exclude_hashes=["bad"]) + + +def test_dataset_filter_cache_key_reflects_exclusions() -> None: + a = DatasetFilter(exclude_hashes=["x"]) + b = DatasetFilter(exclude_hashes=["y"]) + assert a.cache_key() != b.cache_key() diff --git a/tests/unit/test_device_propagation.py b/tests/unit/test_device_propagation.py new file mode 100644 index 0000000000..3103f70e76 --- /dev/null +++ b/tests/unit/test_device_propagation.py @@ -0,0 +1,50 @@ +"""The device the trainer actually runs on has to reach ``HPTModel``. + +``HPT`` builds its policy on CPU and leaves both devices unset; Lightning +learns the real one and ``ModelWrapper.on_fit_start`` / ``on_validation_start`` +write it to the algo. If that write stops at the algo, +``HPTModel.compute_loss`` builds its loss accumulators on the unset device +while ``HPT.compute_losses`` accumulates on the trainer's, and the run dies +with "Expected all tensors to be on the same device". +""" + +from __future__ import annotations + +import torch + +from egomimic.algo.hpt import HPT, HPTModel + + +def _algo_with_policy(construction_device: torch.device): + """An algo holding a policy, without running either ``__init__``. + + Building a real HPT means a full trunk, stems and heads; only the device + wiring is under test here, so both objects are created bare, exactly as + ``test_infra_port_wiring`` does for the stem registration. + """ + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict() + algo.device = construction_device + policy = HPTModel.__new__(HPTModel) + torch.nn.Module.__init__(policy) + policy.device = construction_device + algo.nets["policy"] = policy + return algo, policy + + +def test_assigning_the_algo_device_reaches_the_policy(): + algo, policy = _algo_with_policy(torch.device("cuda")) + + algo.device = torch.device("cpu") # what ModelWrapper.on_fit_start does + + assert policy.device == torch.device("cpu") + + +def test_algo_without_a_policy_still_takes_a_device(): + """The algo sets its device in ``__init__`` before the policy is built.""" + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict() + + algo.device = torch.device("cpu") + + assert algo.device == torch.device("cpu") diff --git a/tests/unit/test_dinov3_stem.py b/tests/unit/test_dinov3_stem.py new file mode 100644 index 0000000000..dfc1608478 --- /dev/null +++ b/tests/unit/test_dinov3_stem.py @@ -0,0 +1,101 @@ +"""``DINOv3Stem`` on a random-init tiny tower (no network, no weights).""" + +from __future__ import annotations + +import pytest +import torch + +from egomimic.models.hpt_nets import DINOv3Stem + +D, B = 32, 3 +TINY_TOWER = {"embed_dim": 32, "depth": 1, "num_heads": 2} + + +def _dino(**kwargs): + torch.manual_seed(0) + kwargs.setdefault("freeze_backbone", True) + return DINOv3Stem( + model_name="vit_small_patch16_dinov3", + output_dim=D, + image_size=[32, 48], + pretrained=False, + tower_kwargs=TINY_TOWER, + **kwargs, + ) + + +def test_dinov3_stem_returns_patch_tokens_only(): + stem = _dino() + assert stem.grid_size == (2, 3) + assert stem(torch.rand(B, 1, 1, 3, 90, 160)).shape == (B, 2 * 3, D) + # 2 frames on the time axis: past frame's patches first + assert stem(torch.rand(B, 2, 1, 3, 90, 160)).shape == (B, 2 * 2 * 3, D) + + +def test_dinov3_stem_frozen_tower(): + stem = _dino().train() + assert not stem.tower.training + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert all(p.grad is None for p in stem.tower.parameters()) + assert stem.proj.weight.grad is not None + assert stem.pretrained_submodules() == [stem.tower] + + +def test_dinov3_stem_finetuned_tower_gets_gradients(): + stem = _dino(freeze_backbone=False).train() + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert any(p.grad is not None for p in stem.tower.parameters()) + assert {id(p) for p in stem.backbone_parameters()} == { + id(p) for p in stem.tower.parameters() + } + + +def test_dinov3_stem_rejects_off_grid_size(): + with pytest.raises(ValueError, match="multiple"): + DINOv3Stem( + model_name="vit_small_patch16_dinov3", + image_size=250, + pretrained=False, + tower_kwargs=TINY_TOWER, + ) + + +def test_a_random_tower_is_not_verified(): + assert _dino().pretrained_reference_state_dict() is None + + +def test_the_verifier_checks_the_tower_against_its_cached_checkpoint( + tmp_path, monkeypatch +): + """The stem reads its timm checkpoint back from the HF cache, so #645's + hash check covers it like the Qwen stems.""" + import huggingface_hub + from safetensors.torch import save_file + + from egomimic.models.hpt_nets import verify_pretrained_weights + + stem = _dino() + path = tmp_path / "model.safetensors" + save_file({k: v.contiguous() for k, v in stem.tower.state_dict().items()}, path) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda *a, **k: str(path)) + stem.pretrained = True + + assert list(verify_pretrained_weights(stem, verbose=False)) == ["DINOv3Stem"] + with torch.no_grad(): + next(stem.tower.parameters()).add_(1.0) + with pytest.raises(RuntimeError, match="do not"): + verify_pretrained_weights(stem, verbose=False) + + +def test_an_uncached_checkpoint_is_skipped_not_fetched(monkeypatch): + import huggingface_hub + from huggingface_hub.errors import LocalEntryNotFoundError + + def offline(*a, **k): + assert k.get("local_files_only") is True + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", offline) + stem = _dino() + stem.pretrained = True + assert stem.pretrained_reference_state_dict() is None diff --git a/tests/unit/test_episode_norm_samples.py b/tests/unit/test_episode_norm_samples.py new file mode 100644 index 0000000000..da2c528c79 --- /dev/null +++ b/tests/unit/test_episode_norm_samples.py @@ -0,0 +1,148 @@ +"""Per-episode norm-sample cache: splits share episodes, strides nest.""" + +import numpy as np +import zarr +from fixtures.synthetic_episodes import write_episode +from omegaconf import OmegaConf + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.norm import episode_norm_samples as ens +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset + +T = 48 + + +def _episodes(tmp_path, n=3): + data = tmp_path / "data" + for i in range(n): + write_episode(data, "eva", T=T, seed=i) + resolver = LocalEpisodeResolver( + data, + key_map=Eva.get_keymap(keymap_mode="cartesian", norm_mode=True), + transform_list=Eva.get_transform_list(mode="cartesian_wristframe_6d"), + ) + return MultiDataset._from_resolver(resolver, mode="total").datasets + + +def _stats(leaves, cache=None, **kw): + ds = MultiDataset(datasets=dict(leaves), mode="total") + stats = MultiDataset(state={}, norm_mode="quantile") + stats.populate_from_datasets({"eva_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset( + ds, "eva_bimanual", num_workers=0, episode_cache=cache, **kw + ) + return next(iter(stats.norm_stats.values())), stats._norm_run_metadata + + +def _assert_same(a, b): + assert a.keys() == b.keys() + for key in a: + for name in a[key]: + np.testing.assert_allclose(a[key][name], b[key][name], rtol=1e-5, atol=1e-6) + + +def _spy(monkeypatch): + sampled = [] + orig = ens._sample_episodes + + def spy(root, leaves, *a, **k): + sampled.append(sorted(leaves)) + return orig(root, leaves, *a, **k) + + monkeypatch.setattr(ens, "_sample_episodes", spy) + return sampled + + +def test_every_frame_matches_the_whole_dataset_pass(tmp_path): + leaves = _episodes(tmp_path) + cached, meta = _stats(leaves, tmp_path / "cache", sample_frac=1.0) + legacy, _ = _stats(leaves, None, sample_frac=1.0) + _assert_same(cached, legacy) + assert meta["stride"] == 1 and meta["frames"] == 3 * T + + +def test_a_sub_split_reads_only_the_cache(tmp_path, monkeypatch): + leaves = _episodes(tmp_path) + sampled = _spy(monkeypatch) + _stats(leaves, tmp_path / "cache", sample_frac=1.0) + assert sampled == [sorted(leaves)] + + sub = {h: leaves[h] for h in sorted(leaves)[:2]} + from_cache, meta = _stats(sub, tmp_path / "cache", sample_frac=1.0) + assert len(sampled) == 1 and meta["episodes_sampled"] == 0 + fresh, _ = _stats(sub, tmp_path / "fresh", sample_frac=1.0) + _assert_same(from_cache, fresh) + + +def test_a_finer_file_serves_a_coarser_stride_but_not_the_reverse( + tmp_path, monkeypatch +): + leaves = _episodes(tmp_path) + sampled = _spy(monkeypatch) + _, meta = _stats(leaves, tmp_path / "cache", sample_frac=1.0) + _, meta = _stats(leaves, tmp_path / "cache", sample_frac=0.25) + assert meta["stride"] == 4 and len(sampled) == 1 + assert meta["frames"] == sum(len(ens.frame_indices(h, T, 4)) for h in leaves) + + _stats(leaves, tmp_path / "coarse", sample_frac=0.25) + _stats(leaves, tmp_path / "coarse", sample_frac=1.0) + assert len(sampled) == 3 + + +def test_max_samples_caps_rows(tmp_path): + leaves = _episodes(tmp_path) + _, meta = _stats(leaves, tmp_path / "cache", sample_frac=1.0, max_samples=50) + assert meta["frames"] == 50 and meta["stride"] == 2 + + +def test_rewritten_episode_is_resampled(tmp_path, monkeypatch): + leaves = _episodes(tmp_path) + sampled = _spy(monkeypatch) + _stats(leaves, tmp_path / "cache", sample_frac=1.0) + h = sorted(leaves)[0] + zarr.open_group(str(leaves[h].episode_path), mode="r+").attrs["note"] = "v2" + _stats(leaves, tmp_path / "cache", sample_frac=1.0) + assert sampled[1] == [h] + + +def test_frame_indices_nest_across_strides(): + for h in ("a", "696d9bf05c2b2b8722b2a08f"): + for s in (1, 2, 4, 8): + coarse = set(ens.frame_indices(h, 1000, 2 * s)) + assert coarse <= set(ens.frame_indices(h, 1000, s)) + + +def test_stride_is_the_largest_that_meets_the_target(): + assert ens.stride_for(11_520_569, 200_000) == 32 + assert ens.stride_for(100, 100) == 1 + assert ens.stride_for(100, 50) == 2 + assert ens.stride_for(100, 0) == 64 + + +def test_recipe_key_ignores_the_split_but_not_the_transforms(): + cfg = OmegaConf.create( + { + "_target_": "egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver", + "resolver": { + "_target_": "egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver", + "folder_path": "/data/a", + "transform_list": {"stride": 1}, + }, + "filters": {"episode_hashes": ["h1"]}, + "mode": "train", + } + ) + base = ens.recipe_key("human_bimanual", cfg) + split = OmegaConf.merge( + cfg, + { + "filters": {"episode_hashes": ["h1", "h2"]}, + "mode": "total", + "resolver": {"folder_path": "/elsewhere"}, + }, + ) + assert ens.recipe_key("human_bimanual", split) == base + strided = OmegaConf.merge(cfg, {"resolver": {"transform_list": {"stride": 3}}}) + assert ens.recipe_key("human_bimanual", strided) != base + assert ens.recipe_key("eva_bimanual", cfg) != base diff --git a/tests/unit/test_eva_to_lance.py b/tests/unit/test_eva_to_lance.py new file mode 100644 index 0000000000..ccd7cb0df9 --- /dev/null +++ b/tests/unit/test_eva_to_lance.py @@ -0,0 +1,82 @@ +import numpy as np +import pyarrow.parquet as pq +import pytest + +h5py = pytest.importorskip("h5py") + +from egomimic.rldb.lance.schema import EVA_STEPS # noqa: E402 +from egomimic.scripts.eva_process.eva_to_lance import ( # noqa: E402 + numeric_columns, + steps_table, +) +from egomimic.scripts.eva_process.eva_to_zarr import ( # noqa: E402 + _separate_numeric_and_image, + _split_per_arm, +) +from egomimic.scripts.eva_process.eva_utils import EvaHD5Extractor # noqa: E402 + +N = 10 + + +def _demo(tmp_path, name="1766709615098"): + """A bimanual demo whose right arm reads all zeros at frames 3-4 in every array.""" + rng = np.random.default_rng(0) + path = tmp_path / f"{name}.hdf5" + with h5py.File(path, "w") as f: + for group, key in ( + ("observations", "eepose"), + ("observations", "joints"), + ("actions", "eepose"), + ("actions", "joints"), + ): + a = rng.uniform(0.1, 0.5, (N, 14)).astype(np.float32) + a[:, 6] = a[:, 13] = np.linspace(0.002, 0.09, N) + a[3:5, 7:13] = 0 + f[f"{group}/{key}"] = a + f["observations/images/front_img_1"] = np.zeros((N, 4, 4, 3), np.uint8) + return path + + +def test_steps_match_the_zarr_arrays_and_flag_dropouts(tmp_path): + path = _demo(tmp_path) + cols, dropout = numeric_columns(path, "both") + want = _split_per_arm( + _separate_numeric_and_image(EvaHD5Extractor.process_episode(path, "both"))[0], + "both", + ) + assert set(want) <= set(cols) + for k in want: + np.testing.assert_array_equal(cols[k], want[k]) + assert ( + dropout[:, 0].sum() == 0 and dropout[3:5, 1].all() and dropout[:, 1].sum() == 2 + ) + # Gripper rescale at ingest: a pre-cutoff demo spans [0, 1]. + assert cols["left.obs_gripper"].min() == 0 and cols["left.obs_gripper"].max() == 1 + + valid = {"left": ~dropout[:, 0], "right": ~dropout[:, 1]} + tbl = steps_table("r", np.arange(N, dtype=np.int64), cols, valid) + assert tbl.schema == EVA_STEPS + assert tbl.column("right_valid").to_pylist() == [i not in (3, 4) for i in range(N)] + right = tbl.column("right").combine_chunks() + np.testing.assert_allclose( + np.stack(right.field("cmd_ee_pose").to_numpy(zero_copy_only=False)), + want["right.cmd_ee_pose"], + ) + np.testing.assert_allclose( + right.field("obs_gripper").to_numpy(), want["right.obs_gripper"][:, 0] + ) + + +def test_right_arm_robot_reads_the_right_slot_and_zeros_the_left(tmp_path): + path = _demo(tmp_path) + cols, dropout = numeric_columns(path, "right") + both, _ = numeric_columns(path, "both") + assert {k.split(".")[0] for k in cols} == {"right"} + np.testing.assert_array_equal(cols["right.obs_ee_pose"], both["right.obs_ee_pose"]) + valid = {"left": np.zeros(N, bool), "right": ~dropout[:, 1]} + tbl = steps_table("r", np.arange(N, dtype=np.int64), cols, valid) + pq.write_table(tbl, tmp_path / "steps.parquet") + tbl = pq.read_table(tmp_path / "steps.parquet") + left = tbl.column("left").combine_chunks() + assert not np.stack(left.field("obs_ee_pose").to_numpy(zero_copy_only=False)).any() + assert tbl.column("right_valid").to_pylist() == list(valid["right"]) diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py new file mode 100644 index 0000000000..25b742ffc7 --- /dev/null +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -0,0 +1,347 @@ +"""Eva wrist-frame defaults, horizon-pooled norm stats, per-segment L2 metrics.""" + +import numpy as np +import pytest +import torch +from fixtures.synthetic_episodes import write_episode +from hydra import compose, initialize_config_module + +from egomimic.eval.action_metrics import cartesian_metrics, keypoint_metrics +from egomimic.rldb.embodiment.eva import Eva +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" + + +def _compose(name, overrides=()): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + return compose(config_name=name, overrides=list(overrides)) + + +@pytest.mark.parametrize("top", ["train_zarr_cartesian", "train_zarr_cartesian_pi"]) +def test_eva_defaults_are_wrist_frame_6d(top): + cfg = _compose(top) + eva = cfg.data.train_datasets.eva_bimanual + assert eva.resolver.transform_list.mode == "cartesian_wristframe_6d" + assert cfg.evaluator.transform_lists.eva_bimanual._target_ == EVA_WRIST_REVERT + assert cfg.norm_stats.pool_horizon is False + + +def test_hpt_eva_dims_match_the_wrist_frame_6d_layout(): + cfg = _compose("train_zarr_cartesian") + dims = cfg.model.robomimic_model.dims.eva_bimanual + assert (dims.proprio, dims.action) == (20, 20) + + +def _eva_stats(tmp_path, pool_horizon): + for i in range(3): + write_episode(tmp_path, "eva", seed=i) + resolver = LocalEpisodeResolver( + tmp_path, + key_map=Eva.get_keymap(keymap_mode="cartesian", norm_mode=True), + transform_list=Eva.get_transform_list(mode="cartesian_wristframe_6d"), + ) + ds = MultiDataset._from_resolver(resolver, mode="total") + stats = MultiDataset(state={}, norm_mode="quantile") + stats.populate_from_datasets({"eva_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset( + ds, "eva_bimanual", sample_frac=1.0, num_workers=0, pool_horizon=pool_horizon + ) + return next(iter(stats.norm_stats.values())) + + +def test_pool_horizon_tiles_one_stat_per_channel(tmp_path): + per_t = _eva_stats(tmp_path / "a", pool_horizon=False) + pooled = _eva_stats(tmp_path / "b", pool_horizon=True) + act_t, act_p = per_t["actions_cartesian"], pooled["actions_cartesian"] + for name in act_t: + assert act_p[name].shape == act_t[name].shape + assert act_p[name].ndim == 2 + # every timestep carries the same per-channel value + np.testing.assert_array_equal( + act_p[name], act_p[name][:1].repeat(len(act_p[name]), 0) + ) + # every timestep has equal weight in the pool, so a pooled quantile lies + # between the smallest and largest per-timestep quantile of that channel + for name in ("quantile_1", "quantile_99", "median"): + lo = act_t[name].min(axis=0) - 1e-4 + hi = act_t[name].max(axis=0) + 1e-4 + assert np.all((act_p[name][0] >= lo) & (act_p[name][0] <= hi)), name + # proprio has no horizon axis: identical either way + for name, arr in per_t["observations.state.ee_pose"].items(): + np.testing.assert_allclose( + pooled["observations.state.ee_pose"][name], arr, rtol=1e-6 + ) + + +def test_pool_horizon_is_part_of_the_cache_key(): + args = ("eva_bimanual", {"h": "fp"}, {"resolver": {}}, 1.0) + assert norm_cache.norm_cache_key( + norm_cache.cache_inputs(*args, pool_horizon=True) + ) != norm_cache.norm_cache_key(norm_cache.cache_inputs(*args)) + + +def test_segment_l2_separates_early_from_late_error(): + B, T = 2, 100 + gt = torch.zeros(B, T, 20) + pred = gt.clone() + pred[:, :10, 0] = 0.03 # 3 cm on the left arm x, first 10 % only + m = cartesian_metrics(pred, gt, "p") + # mean over both arms: 3 cm on one of two + assert m["p_xyz_l2_early"].item() == pytest.approx(0.015) + assert m["p_xyz_l2_mid"].item() == pytest.approx(0.0) + assert m["p_xyz_l2_late"].item() == pytest.approx(0.0) + + +def test_segment_l2_on_keypoints(): + B, T = 1, 100 + gt = torch.zeros(B, T, 144) + pred = gt.clone() + pred[:, 50:, 9:72] = 0.02 # every left-hand keypoint off by (2,2,2) cm late + m = keypoint_metrics(pred, gt, "k") + expected = 0.5 * float(np.linalg.norm([0.02] * 3)) # one hand of two + assert m["k_kp_l2_late"].item() == pytest.approx(expected, rel=1e-5) + assert m["k_kp_l2_early"].item() == pytest.approx(0.0) + assert m["k_wrist_xyz_l2_late"].item() == pytest.approx(0.0) + + +# ------------------------------------------------- degenerate-range guard +@pytest.mark.parametrize("mode", ["zscore", "minmax", "quantile"]) +def test_norm_constant_channel_maps_to_zero_and_back_to_its_centre(mode): + from egomimic.utils.action_utils import ( + NORM_MIN_RANGE, + _apply_norm_one, + _apply_unnorm_one, + ) + + lo = np.array([0.0, -1.0, 2.0], dtype=np.float32) + hi = np.array([0.0, 1.0, 2.0 + NORM_MIN_RANGE / 4], dtype=np.float32) + stats = { + "mean": 0.5 * (lo + hi), + "std": hi - lo, + "min": lo, + "max": hi, + "quantile_1": lo, + "quantile_99": hi, + } + x = torch.tensor([[1e-3, 0.5, 2.0 + 1e-3]]) + n = _apply_norm_one(x, stats, mode) + # channels 0 and 2 are (near-)constant: 0 whatever the offset, not 1e3 + assert n[0, 0].item() == 0.0 and n[0, 2].item() == 0.0 + back = _apply_unnorm_one(torch.full_like(n, 0.7), stats, mode) + np.testing.assert_allclose(back[0, [0, 2]].numpy(), stats["mean"][[0, 2]]) + + # the regular channel keeps the historical formula exactly + if mode == "zscore": + ref = (0.5 - stats["mean"][1]) / (stats["std"][1] + 1e-6) + else: + ref = 2.0 * (0.5 - lo[1]) / (hi[1] - lo[1] + 1e-6) - 1.0 + assert n[0, 1].item() == pytest.approx(float(ref), rel=1e-6) + assert _apply_unnorm_one(n, stats, mode)[0, 1].item() == pytest.approx(0.5) + + +def test_multidataset_normalize_uses_the_guard(): + stats = {"quantile_1": np.zeros(2), "quantile_99": np.array([0.0, 1.0])} + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + n = md._apply_norm_one(torch.tensor([[5e-4, 0.5]]), stats) + assert n.tolist() == [[0.0, 0.0]] + + +# --- rot6d channels are never normalized ----------------------------------- + + +def _stats_for(width): + lo = -np.ones(width, dtype=np.float32) + hi = np.ones(width, dtype=np.float32) * 3.0 + return { + "mean": 0.5 * (lo + hi), + "std": hi - lo, + "min": lo, + "max": hi, + "quantile_1": lo, + "quantile_99": hi, + } + + +@pytest.mark.parametrize( + "zarr_key,width", + [ + ("actions_cartesian", 18), + ("actions_cartesian", 20), + ("observations.state.ee_pose", 20), + ("actions_keypoints", 144), + ("observations.state.keypoints", 144), + ], +) +def test_rot6d_channels_pass_through_normalization(zarr_key, width): + from egomimic.utils.action_utils import _apply_norm_one, _apply_unnorm_one + from egomimic.utils.pose_utils import rot6d_channels + + rot = list(rot6d_channels(zarr_key, width)) + other = [i for i in range(width) if i not in rot] + x = torch.arange(width, dtype=torch.float32).unsqueeze(0) / width + stats = _stats_for(width) + + n = _apply_norm_one(x, stats, "quantile", rot) + torch.testing.assert_close(n[:, rot], x[:, rot], atol=0, rtol=0) + assert not torch.allclose(n[:, other], x[:, other]) + back = _apply_unnorm_one(n, stats, "quantile", rot) + torch.testing.assert_close(back, x, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize( + "zarr_key,width", + [ + ("actions_cartesian", 12), # ypr: not a rot6d layout + ("actions_keypoints", 138), + ("observations.state.joint_positions", 20), # not a pose key + ("actions_cartesian", 7), # unknown width + ], +) +def test_no_rot6d_channels_outside_the_6d_layouts(zarr_key, width): + from egomimic.utils.pose_utils import rot6d_channels + + assert rot6d_channels(zarr_key, width) is None + + +def test_multidataset_leaves_the_rot6d_block_alone(): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + stats = _stats_for(18) + x = torch.arange(18, dtype=torch.float32).unsqueeze(0) / 18 + n = md._apply_norm_one(x, stats, "actions_cartesian") + rot = list(range(3, 9)) + list(range(12, 18)) + torch.testing.assert_close(n[:, rot], x[:, rot], atol=0, rtol=0) + # without the key there is no layout to consult, so nothing is skipped + assert not torch.allclose(md._apply_norm_one(x, stats)[:, rot], x[:, rot]) + + +def _state_with_stats(width, **extra): + return { + "norm_mode": "quantile", + "embodiments": [4], + "zarr_keys": {4: {"actions_cartesian": "actions_cartesian"}}, + "norm_stats": { + 4: {"actions_cartesian": {k: v for k, v in _stats_for(width).items()}} + }, + **extra, + } + + +def test_fresh_stats_pass_rot6d_through_and_save_the_flag(): + stats = MultiDataset(state={}, norm_mode="quantile") + assert stats.rot6d_passthrough + assert stats.to_state()["rot6d_passthrough"] is True + + +@pytest.mark.parametrize("saved,expect_identity", [(None, False), (True, True)]) +def test_state_round_trip_keeps_the_rot6d_convention(saved, expect_identity): + from egomimic.utils.pose_utils import rot6d_channels + + extra = {} if saved is None else {"rot6d_passthrough": saved} + restored = MultiDataset.from_state(_state_with_stats(20, **extra)) + x = torch.arange(20, dtype=torch.float32).unsqueeze(0) / 20 + rot = list(rot6d_channels("actions_cartesian", 20)) + n = restored._apply_norm_one(x, _stats_for(20), "actions_cartesian") + assert torch.equal(n[:, rot], x[:, rot]) == expect_identity + back = restored._apply_unnorm_one(n, _stats_for(20), "actions_cartesian") + torch.testing.assert_close(back, x, atol=1e-5, rtol=0) + + +def test_bounds_check_ignores_constant_cells(): + key = "actions_cartesian" + q1 = np.full((4, 14), -1.0, dtype=np.float32) + q99 = np.full((4, 14), 1.0, dtype=np.float32) + q1[0, :3] = q99[0, :3] = 0.0 # wrist-frame t=0: xyz exactly 0 + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = {0: {key: {"quantile_1": q1, "quantile_99": q99}}} + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + + arr = np.zeros((4, 14), dtype=np.float32) + arr[0, 0] = 1e-3 # off-convention offset at a constant cell: admitted + assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None + arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected + assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None + + +# ------------------------------------------------------- rollout reverts +def _rand_pose7(rng): + from scipy.spatial.transform import Rotation as R + + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() + return np.concatenate([rng.uniform(-0.5, 0.5, 3), q[[3, 0, 1, 2]]]) + + +def _rand_chunk7(rng, start, n=45): + from scipy.spatial.transform import Rotation as R + + out = np.zeros((n, 7)) + p, r = start[:3].copy(), R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + out[t] = np.concatenate([p, r.as_quat()[[3, 0, 1, 2]]]) + return out + + +# Compared against cartesian_6d, so that mode is the reference, not a case. +@pytest.mark.parametrize( + "mode", ["cartesian_wristframe_ypr", "cartesian_wristframe_6d"] +) +def test_eva_revert_for_rollout_recovers_camframe_ypr(mode): + """rollout.py feeds Eva.get_transform_list(mode) outputs to the model and + Eva.get_revert_transform_list(mode) must bring predictions back to the + same cam-frame xyz+ypr+gripper (14-D) whatever the training mode.""" + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.embodiment.embodiment import Embodiment + + rng = np.random.default_rng(3) + lobs, robs = _rand_pose7(rng), _rand_pose7(rng) + raw = { + "left.obs_ee_pose": lobs, + "right.obs_ee_pose": robs, + "left.cmd_ee_pose": _rand_chunk7(rng, lobs), + "right.cmd_ee_pose": _rand_chunk7(rng, robs), + "left.obs_gripper": np.array([0.2]), + "right.obs_gripper": np.array([0.8]), + "left.cmd_gripper": rng.uniform(0, 1, (45, 1)), + "right.cmd_gripper": rng.uniform(0, 1, (45, 1)), + } + + def run(m): + s = {k: v.copy() for k, v in raw.items()} + # the ypr half of this comparison is exactly what a legacy checkpoint + # feeds the rollout, which is the case allow_legacy_rotation covers + for t in Eva.get_transform_list(m, allow_legacy_rotation=True): + s = t.transform(s) + batch = { + "actions_cartesian": torch.as_tensor(s["actions_cartesian"])[None], + "observations.state.ee_pose": torch.as_tensor( + s["observations.state.ee_pose"] + )[None], + } + rev = Eva.get_revert_transform_list(m) + if rev is not None: + batch = Embodiment.apply_transform(batch, rev) + return np.asarray(batch["actions_cartesian"][0], dtype=np.float64) + + assert Eva.get_revert_transform_list("cartesian") is None + got, ref = run(mode), run("cartesian_6d") + assert got.shape == ref.shape == (100, 14) + for off in (0, 7): + np.testing.assert_allclose( + got[:, off : off + 3], ref[:, off : off + 3], atol=1e-5 + ) + np.testing.assert_allclose(got[:, off + 6], ref[:, off + 6], atol=1e-5) + Rg = R.from_euler("ZYX", got[:, off + 3 : off + 6]).as_matrix() + Rr = R.from_euler("ZYX", ref[:, off + 3 : off + 6]).as_matrix() + np.testing.assert_allclose(Rg, Rr, atol=1e-4) diff --git a/tests/unit/test_eval_determinism.py b/tests/unit/test_eval_determinism.py new file mode 100644 index 0000000000..d64a8a8fa8 --- /dev/null +++ b/tests/unit/test_eval_determinism.py @@ -0,0 +1,445 @@ +"""Deterministic evaluation (lane B). + +Validating the same checkpoint twice must give the same prompts, the same +sampled actions and the same metrics. Three pieces make that true: + +* ``HPT._build_prompts`` picks annotations with a per-pass seeded RNG whenever + the algo is in eval mode (``annotation_sampling_mode`` only steers training); +* ``HPT.forward_eval`` builds a ``torch.Generator`` seeded from + ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter`` and threads it + down to the denoising head's sampler, and forks/seeds the global RNG around + the BC val-loss call; ``ModelWrapper.on_validation_start`` rewinds the + counter, so pass N and pass N+1 replay the same draws. + +Plus the gradient-clipping change that rides along in this branch: the MAD +spike detector in ``ModelWrapper.on_after_backward`` is log-only (fixed +clipping is Lightning's ``trainer.gradient_clip_val``). + +Everything here is CPU-only and builds the smallest object that still runs the +real code path: the real ``FMPolicy`` Euler sampler, the real +``HPTModel.forward`` head dispatch and the real ``HPT.forward_eval``, with only +the stems/trunk (``forward_features``) and the norm stats stubbed out. +""" + +from __future__ import annotations + +import random +import types +from collections import deque +from pathlib import Path + +import hydra +import pytest +import torch +import torch.nn as nn + +import egomimic.hydra_configs +from egomimic.algo.hpt import EVAL_BASE_SEED, HPT, HPTModel +from egomimic.models.fm_policy import FMPolicy +from egomimic.pl_utils.pl_model import ModelWrapper + +DOMAIN = "eva_bimanual" +EMB_ID = 6 # EMBODIMENT.EVA_BIMANUAL; get_embodiment(6).lower() == DOMAIN +AC_KEY = "actions_cartesian" +ANNOTATION_KEY = "annotations" +BATCH = 4 +HORIZON = 3 +ACTION_DIM = 4 +FEAT_DIM = 5 + + +# --------------------------------------------------------------------------- +# Fixtures: a minimal but real HPT +# --------------------------------------------------------------------------- +class _TinyVelocity(nn.Module): + """Stands in for ConditionalUnet1D: v(x_t, t, cond), same call signature.""" + + def __init__(self, dim: int): + super().__init__() + self.lin = nn.Linear(dim, dim) + + def forward(self, x_t, t, global_cond): + return self.lin(x_t) * t.reshape(-1, 1, 1) + + +class _StubNormStats: + def unnormalize(self, predictions, embodiment_id): + return dict(predictions) + + +def _make_algo(annotation_key: str | None = ANNOTATION_KEY) -> HPT: + """A real HPT wired to a real FMPolicy head, without stems/trunk/dataset. + + ``HPT.__init__`` needs norm stats, configs and a dataset; every attribute + ``forward_eval`` / ``_build_prompts`` actually read is set here instead. + """ + torch.manual_seed(0) + head = FMPolicy( + model=_TinyVelocity(ACTION_DIM), + action_horizon=HORIZON, + infer_ac_dims={DOMAIN: ACTION_DIM}, + num_inference_steps=2, + ) + model = HPTModel.__new__(HPTModel) + nn.Module.__init__(model) + model.diffusion = True + model.shared_action = False + model.auxiliary_ac_keys = {} + model.heads = nn.ModuleDict({DOMAIN: head}) + model.device = torch.device("cpu") + features = ( + torch.arange(BATCH * FEAT_DIM, dtype=torch.float32).reshape(BATCH, FEAT_DIM) + / 10.0 + ) + model.forward_features = lambda domain, data: (features, None) + + algo = HPT.__new__(HPT) + algo.nets = nn.ModuleDict({"policy": model}) + algo.device = torch.device("cpu") + algo.norm_stats = _StubNormStats() + algo.camera_keys = {EMB_ID: []} + algo.proprio_keys = {EMB_ID: []} + algo.lang_keys = {EMB_ID: []} + algo.ac_keys = {EMB_ID: AC_KEY} + algo.auxiliary_ac_keys = {} + algo.shared_ac_key = None + algo.is_6dof = False + algo.freeze_repr = False + algo.freeze_depth = 8 + algo.eval_base_seed = EVAL_BASE_SEED + algo.annotation_key = annotation_key + algo.annotation_modality = "annotation" + algo.annotation_sampling_mode = "random" + algo.default_prompt = "do the task" + algo.train_image_augs = None + algo.eval_image_augs = None + algo.nets.eval() + return algo + + +def _make_batch() -> dict: + torch.manual_seed(1) + return { + EMB_ID: { + AC_KEY: torch.randn(BATCH, HORIZON, ACTION_DIM), + "pad_mask": torch.ones(BATCH, HORIZON, 1), + "embodiment": torch.tensor([EMB_ID]), + } + } + + +def _annotation_batch(n_options: int = 4) -> dict: + return { + ANNOTATION_KEY: [ + [f"sample{i}_option{j}" for j in range(n_options)] for i in range(BATCH) + ] + } + + +def _prompt_pass(algo: HPT, raw_batch: dict, n_batches: int = 3) -> list[list[str]]: + """One validation pass: N batches, counter advancing as forward_eval does.""" + algo.reset_eval_pass_counter() + out = [] + for _ in range(n_batches): + out.append(algo._build_prompts(raw_batch, BATCH)) + algo._eval_pass_counter += 1 # what forward_eval does after each batch + return out + + +# --------------------------------------------------------------------------- +# 1. Prompts +# --------------------------------------------------------------------------- +def test_eval_prompts_are_reproducible_across_passes(): + algo = _make_algo() + raw = _annotation_batch() + first = _prompt_pass(algo, raw) + second = _prompt_pass(algo, raw) + assert first == second + + +def test_eval_prompts_keep_variety_and_are_not_always_first(): + algo = _make_algo() + raw = _annotation_batch() + prompts = [p for batch in _prompt_pass(algo, raw, n_batches=5) for p in batch] + firsts = [options[0] for options in raw[ANNOTATION_KEY]] + assert any( + p not in firsts for p in prompts + ), f"eval prompts collapsed to sample[0]: {prompts}" + assert len(set(prompts)) > 1 + + +def test_eval_prompts_ignore_sampling_mode(): + """`annotation_sampling_mode` is an algo-level knob; eval overrides it.""" + algo = _make_algo() + raw = _annotation_batch() + algo.annotation_sampling_mode = "random" + as_random = _prompt_pass(algo, raw) + algo.annotation_sampling_mode = "first" + as_first = _prompt_pass(algo, raw) + assert as_random == as_first + assert as_first[0] != [options[0] for options in raw[ANNOTATION_KEY]] + + +def test_eval_prompts_fall_back_to_default_on_empty(): + algo = _make_algo() + raw = {ANNOTATION_KEY: [[], ["a", "b"], [], ["c"]]} + prompts = algo._build_prompts(raw, 4) + assert prompts[0] == algo.default_prompt + assert prompts[2] == algo.default_prompt + assert prompts[1] in {"a", "b"} and prompts[3] == "c" + + +def test_missing_annotation_key_uses_default_prompt(): + algo = _make_algo(annotation_key=None) + assert algo._build_prompts({}, 3) == [algo.default_prompt] * 3 + + +@pytest.mark.parametrize("mode", ["random", "first"]) +def test_train_prompt_behaviour_is_unchanged(mode): + """In train mode the parent's logic still applies verbatim.""" + algo = _make_algo() + algo.annotation_sampling_mode = mode + algo.nets.train() + raw = _annotation_batch() + + random.seed(12345) + got = algo._build_prompts(raw, BATCH) + + random.seed(12345) + expected = [] + for options in raw[ANNOTATION_KEY]: + if mode == "random": + expected.append(options[random.randint(0, len(options) - 1)]) + else: + expected.append(options[0]) + assert got == expected + + +# --------------------------------------------------------------------------- +# 2. Seeded flow sampling +# --------------------------------------------------------------------------- +def _eval_pass(algo: HPT, batch: dict, n_batches: int = 2) -> list[dict]: + algo.reset_eval_pass_counter() + return [algo.forward_eval(batch) for _ in range(n_batches)] + + +def _actions(preds: list[dict]) -> list[torch.Tensor]: + return [p[f"{DOMAIN}_{AC_KEY}"] for p in preds] + + +def test_two_validation_passes_give_bit_identical_actions(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(_actions(first), _actions(second)): + assert torch.equal(a, b) + + +def test_val_loss_is_identical_across_passes(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(first, second): + assert torch.equal(a[f"{DOMAIN}_loss"], b[f"{DOMAIN}_loss"]) + + +def test_batches_within_a_pass_use_different_noise(): + """Without the reset the counter has moved on, so the draws differ.""" + algo = _make_algo() + batch = _make_batch() + algo.reset_eval_pass_counter() + first = algo.forward_eval(batch) + second = algo.forward_eval(batch) # no reset: pass counter is now 1 + assert not torch.equal(first[f"{DOMAIN}_{AC_KEY}"], second[f"{DOMAIN}_{AC_KEY}"]) + + +def test_forward_eval_does_not_disturb_the_global_rng(): + """Eval must not shift the training RNG stream (fork_rng + generator).""" + algo = _make_algo() + batch = _make_batch() + torch.manual_seed(7) + before = torch.randn(3) + torch.manual_seed(7) + algo.reset_eval_pass_counter() + algo.forward_eval(batch) + after = torch.randn(3) + assert torch.equal(before, after) + + +def test_pass_seed_depends_on_rank_and_counter(monkeypatch): + algo = _make_algo() + algo.reset_eval_pass_counter() + monkeypatch.setenv("RANK", "0") + rank0 = algo._eval_pass_seed() + monkeypatch.setenv("RANK", "3") + rank3 = algo._eval_pass_seed() + assert rank0 != rank3 + algo._eval_pass_counter += 1 + assert algo._eval_pass_seed() == rank3 + 1 + + +def test_val_loss_noise_does_not_collide_with_the_next_batch(monkeypatch): + """A torch.Generator seeded with S draws the same stream as a global + manual_seed(S), so a +1 offset would give batch i's val loss the very noise + batch i+1 samples with.""" + from egomimic.algo.hpt import EVAL_LOSS_STRIDE + + assert EVAL_LOSS_STRIDE > 1000 + algo = _make_algo() + algo.reset_eval_pass_counter() + seeds = set() + for _ in range(4): + pass_seed = algo._eval_pass_seed() + seeds.add(pass_seed) + seeds.add(pass_seed + EVAL_LOSS_STRIDE) + algo._eval_pass_counter += 1 + assert len(seeds) == 8, "a val-loss seed repeats a sampling seed" + + +def test_eval_base_seed_is_a_real_knob(): + """A second eval pass can be made to draw DIFFERENT noise on purpose.""" + algo = _make_algo() + assert algo.eval_base_seed == EVAL_BASE_SEED + algo.reset_eval_pass_counter() + default = algo._eval_pass_seed() + algo.eval_base_seed = 12345 + assert algo._eval_pass_seed() == default + 12345 + + +def test_generator_is_threaded_to_the_head(): + """HPTModel.forward -> DenoisingPolicy.forward -> sample_action(generator).""" + algo = _make_algo() + model = algo.nets["policy"] + data = {"action": torch.zeros(BATCH, HORIZON, ACTION_DIM)} + gen = torch.Generator(device="cpu") + + gen.manual_seed(42) + a = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(42) + b = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(43) + c = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + assert torch.equal(a, b) + assert not torch.equal(a, c) + + +def test_on_validation_start_resets_the_pass_counter(): + algo = _make_algo() + algo.reset_eval_pass_counter() + algo._eval_pass_counter = 17 + wrapper = types.SimpleNamespace( + model=algo, + device=torch.device("cpu"), + _val_heads=lambda: {"valid": None}, + ) + ModelWrapper.on_validation_start(wrapper) + assert algo._eval_pass_counter == 0 + + +# --------------------------------------------------------------------------- +# 3. Log-only MAD gradient detector +# --------------------------------------------------------------------------- +class _GradHookStub: + """Just enough LightningModule for ModelWrapper.on_after_backward.""" + + def __init__(self, param: nn.Parameter, history): + self.enable_grad_norm = True + self.grad_norm_mad_scale = ModelWrapper.grad_norm_mad_scale + self.grad_norm_mad_min_count = ModelWrapper.grad_norm_mad_min_count + self.grad_norm_mad_window = ModelWrapper.grad_norm_mad_window + self.grad_norm_history = deque(history, maxlen=self.grad_norm_mad_window) + self._param = param + self.global_step = 4242 + self.trainer = types.SimpleNamespace(is_global_zero=True) + self.logged = {} + + def parameters(self): + return [self._param] + + def log(self, name, value, **kwargs): + self.logged[name] = value + + +def _spiking_stub(spike_norm: float = 50.0): + param = nn.Parameter(torch.zeros(4)) + grad = torch.full((4,), spike_norm / 2.0) # ||grad||_2 == spike_norm + param.grad = grad + history = [ + 0.9 if i % 2 else 1.1 for i in range(ModelWrapper.grad_norm_mad_min_count) + ] + return _GradHookStub(param, history), param + + +def test_on_after_backward_never_modifies_gradients(): + stub, param = _spiking_stub() + before = param.grad.clone() + ModelWrapper.on_after_backward(stub) + assert torch.equal(param.grad, before), "MAD detector rescaled the gradients" + + +def test_on_after_backward_still_flags_and_logs_the_spike(): + stub, _ = _spiking_stub() + ModelWrapper.on_after_backward(stub) + logged = stub.logged + assert logged["Train/policy_grad_norms_mad_flag"] == 1.0 + assert logged["Train/policy_grad_norms_raw"] == pytest.approx(50.0, rel=1e-5) + assert "Train/policy_grad_norms_mad_threshold" in logged + + +def test_flagged_steps_are_appended_to_the_history(): + """Log-only detector: every step goes into the rolling median/MAD window.""" + stub, _ = _spiking_stub() + n_before = len(stub.grad_norm_history) + ModelWrapper.on_after_backward(stub) + assert len(stub.grad_norm_history) == n_before + 1 + assert stub.grad_norm_history[-1] == pytest.approx(50.0, rel=1e-5) + + +# --------------------------------------------------------------------------- +# 4. Trainer config knob +# --------------------------------------------------------------------------- +def test_trainer_default_exposes_gradient_clip_val(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", []) + assert "gradient_clip_val" in cfg.trainer + assert cfg.trainer.gradient_clip_val is None # no clipping until chosen + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +# Derived from the config group, not listed: a new trainer config must not be +# able to join the repo without this knob being checked. +TRAINER_CONFIGS = sorted( + p.stem + for p in (Path(egomimic.hydra_configs.__file__).parent / "trainer").glob("*.yaml") +) + + +@pytest.mark.parametrize("trainer_cfg", TRAINER_CONFIGS) +def test_every_trainer_config_keeps_the_knob(trainer_cfg, compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", [f"trainer={trainer_cfg}"]) + assert cfg.trainer.gradient_clip_val is None + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +def test_gradient_clip_override_reaches_the_trainer(tmp_path, compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian", + [ + "trainer=default", + "trainer.gradient_clip_val=1.0", + "trainer.accelerator=cpu", + "trainer.precision=32", + "trainer.max_epochs=1", + "trainer.min_epochs=1", + f"paths.output_dir={tmp_path}", + ], + ) + trainer = hydra.utils.instantiate(cfg.trainer, logger=False, callbacks=None) + assert trainer.gradient_clip_val == 1.0 + assert trainer.gradient_clip_algorithm == "norm" + + +def test_lightning_module_does_not_override_gradient_clipping(): + """Nothing in ModelWrapper may intercept Lightning's clipping hook.""" + assert "configure_gradient_clipping" not in vars(ModelWrapper) diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py new file mode 100644 index 0000000000..1d8b92f22f --- /dev/null +++ b/tests/unit/test_eval_loaders.py @@ -0,0 +1,839 @@ +"""Lane A: split each val head into a subsampled METRIC loader and a +contiguous pinned-episode VIDEO loader. + +Before this change the unshuffled val loaders were truncated by +``limit_val_batches`` to their leading hash-sorted slice, so the metric saw a +third of the seen-op split and effectively one of the two held-out operators +(docs/hpt-experiments/supervisors/S1/reports/lane0-coverage.md). Now: + +* metric loader = ``EvenStrideDataset(split, frames_per_episode=K)``, K from + the data config, so every episode contributes and nothing renders; +* video loader = the head's split restricted to the pinned ``episode_hash``es, + contiguous, logging no metrics. + +A data config with neither key must behave exactly as before. +""" + +from __future__ import annotations + +import os +import re +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from lightning import LightningModule +from omegaconf import OmegaConf + +import egomimic.trainHydra as th +from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EvenStrideDataset, + MultiDataset, + PinError, + pinned_episode_subset, +) + +HUMAN = "human_bimanual" +FLAGSHIP = "train_zarr_mecka_flagship_6d_hpt" +LIMIT_VAL_BATCHES = 80 +BATCH = 64 + +# (data config, {head: K}, {head: [pinned hashes]}) as set on this branch. +CONFIG_EXPECTATIONS = { + "mecka_fold_flagship_opsplit_hpt_6d": ( + {"valid": 73, "train_viz": 3, "unseen_op_valid": 26}, + { + "valid": ["692e771c39719ab57395b408"], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), + "mecka_fold_flagship_topop_hpt_6d": ( + {"valid": 256, "train_viz": 13, "unseen_op_valid": 26}, + { + "valid": ["692e83ad567c97525f4011a5"], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), + "mecka_fold_flagship_top3_hpt_6d": ( + {"valid": 96, "train_viz": 4, "unseen_op_valid": 26}, + { + # one seen-val pin per train operator (683785ac / 690366b2 / + # 6903686e), 191 + 223 + 226 = 640f = the whole video budget + "valid": [ + "692ea44dc621d7f4aac3aae6", + "692eb08997ffd8ec90340290", + "692fe79f34a99e18e25289ad", + ], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), +} + +# Episode counts these K values were computed from (measured 2026-09-15). +SPLIT_EPISODES = { + "mecka_fold_flagship_opsplit_hpt_6d": { + "valid": 70, + "train_viz": 1346, + "unseen_op_valid": 191, + }, + "mecka_fold_flagship_topop_hpt_6d": { + "valid": 20, + "train_viz": 385, + "unseen_op_valid": 191, + }, + "mecka_fold_flagship_top3_hpt_6d": { + "valid": 53, + "train_viz": 1026, + "unseen_op_valid": 191, + }, +} + + +# --------------------------------------------------------------- fixtures +class _Episode: + """Minimal leaf dataset: n frames, each sample tagged with its frame idx.""" + + def __init__(self, n: int): + self.n = n + + def __len__(self): + return self.n + + def __getitem__(self, i): + return {"frame": np.float32(i)} + + +def _split(sizes: dict[str, int]) -> MultiDataset: + return MultiDataset( + datasets={h: _Episode(n) for h, n in sizes.items()}, mode="total" + ) + + +def _cfg(metric=None, video=None, **data): + node = dict(data) + if metric is not None: + node["metric_frames_per_episode"] = metric + if video is not None: + node["video_episodes"] = video + return OmegaConf.create({"data": node}) + + +# ------------------------------------------------------------ 1. config keys +@pytest.mark.parametrize("data_config", sorted(CONFIG_EXPECTATIONS)) +def test_flagship_configs_declare_k_and_pins(data_config, compose_resolve): + """Both flagship configs compose with the new keys, K per head matches + floor(5120 / episodes-in-split), and each head pins its video episodes.""" + expected_k, expected_pins = CONFIG_EXPECTATIONS[data_config] + cfg = compose_resolve(FLAGSHIP, [f"data={data_config}"]) + assert int(cfg.trainer.limit_val_batches) == LIMIT_VAL_BATCHES + assert int(cfg.data.valid_dataloader_params[HUMAN].batch_size) == BATCH + + k_table = OmegaConf.to_container(cfg.data.metric_frames_per_episode) + assert k_table == expected_k + pins = OmegaConf.to_container(cfg.data.video_episodes) + assert pins == expected_pins + + for head, k in k_table.items(): + n_eps = SPLIT_EPISODES[data_config][head] + assert k == LIMIT_VAL_BATCHES * BATCH // n_eps, (head, k, n_eps) + # the subsampled set still fits the limit_val_batches window + assert k * n_eps <= LIMIT_VAL_BATCHES * BATCH, (head, k, n_eps) + # the unseen head is the same split in every flagship config, so its pins + # and K are shared + assert k_table["unseen_op_valid"] == 26 + assert len(pins["unseen_op_valid"]) == 2, "one pin per held-out operator" + + # and trainHydra reads exactly these + for head, k in expected_k.items(): + assert th._metric_frames_per_episode(cfg, head) == k + + +# The 3 largest flagship operators (SQL app.episodes, 2026-09-15), by frames. +TOP3_OPERATORS = [ + "6903686e0e94ce070afd1f24", + "690366b20e94ce070afd1e8a", + "683785ac01ca734152093448", +] + + +def test_top3_config_trains_on_exactly_the_three_largest_operators(compose_resolve): + """mecka_fold_flagship_top3_hpt_6d composes under the flagship recipe, its + train filter names exactly the 3 top operators (none of them held out), the + seen-op val follows that filter, and both lane A tables cover all 3 heads.""" + cfg = compose_resolve(FLAGSHIP, ["data=mecka_fold_flagship_top3_hpt_6d"]) + parent = compose_resolve(FLAGSHIP, ["data=mecka_fold_flagship_opsplit_hpt_6d"]) + d, pd = cfg.data, parent.data + train = d.train_datasets[HUMAN] + + # the operator lambda names exactly the three ids, largest first + operator_lambda = train.filters.filter_lambdas[-1] + assert "row.get('operator'" in operator_lambda + assert re.findall(r"[0-9a-f]{24}", operator_lambda) == TOP3_OPERATORS + held_out = [str(o) for o in d.held_out_operators] + assert not set(TOP3_OPERATORS) & set(held_out), "a train operator is held out" + + # the first three (lab / task / flagship-path) lambdas are the parent's + assert list(train.filters.filter_lambdas[:3]) == list( + pd.train_datasets[HUMAN].filters.filter_lambdas[:3] + ) + + # seen-op val interpolates the train filters, so it follows this override + valid = d.valid_datasets[HUMAN] + assert list(valid.filters.filter_lambdas) == list(train.filters.filter_lambdas) + assert valid.valid_ratio == train.valid_ratio == 0.05 + assert (train.mode, valid.mode) == ("train", "valid") + assert d.valid_prefix == "seen_op_valid" + assert d.get("train_viz_datasets") is None + + # the unseen-operator head is untouched + unseen = d.unseen_op_valid_datasets[HUMAN] + assert list(unseen.filters.filter_lambdas) == list( + pd.unseen_op_valid_datasets[HUMAN].filters.filter_lambdas + ) + assert unseen.mode == "total" + + # lane A keys cover all three heads + heads = {"valid", "train_viz", "unseen_op_valid"} + k_table = OmegaConf.to_container(d.metric_frames_per_episode) + pins = OmegaConf.to_container(d.video_episodes) + assert set(k_table) == heads + assert set(pins) == heads + assert all(isinstance(v, int) and v > 0 for v in k_table.values()) + assert len(pins["valid"]) == 3, "one seen-val video pin per train operator" + assert len(set(pins["valid"])) == 3 + + +def test_config_without_the_keys_is_unchanged(compose_resolve): + """Backward compatibility: a data config that declares neither key gets the + old loader names and the unwrapped dataset objects.""" + cfg = compose_resolve("train_zarr_cartesian", ["data=mecka_all_6d"]) + assert cfg.data.get("metric_frames_per_episode") is None + assert cfg.data.get("video_episodes") is None + assert th._metric_frames_per_episode(cfg, "valid") is None + + split = _split({"ep0": 10, "ep1": 10}) + datasets = {HUMAN: split} + assert th._subsample_val_datasets(cfg, "valid", datasets) is datasets + assert th._video_datasets(cfg, {"valid": datasets}) == {} + + params = {HUMAN: {"batch_size": 2, "num_workers": 0}} + dm = MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets=datasets, + train_viz_dataloader_params=params, + unseen_op_valid_datasets=datasets, + unseen_op_valid_dataloader_params=params, + ) + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + assert dm.valid_datasets[HUMAN] is split + + +# ------------------------------------------------------------ 2. metric loaders +def test_metric_loader_is_even_stride_over_every_episode(): + sizes = {"ep0": 500, "ep1": 300, "ep2": 40} + split = _split(sizes) + cfg = _cfg(metric={"valid": 73}) + wrapped = th._subsample_val_datasets(cfg, "valid", {HUMAN: split})[HUMAN] + + assert isinstance(wrapped, EvenStrideDataset) + assert isinstance(wrapped, MultiDataset), "trainHydra's isinstance checks" + assert wrapped.frames_per_episode == 73 + assert wrapped.base is split + + per_episode = {name: 0 for name in sizes} + starts = {} + offset = 0 + for name, n in sizes.items(): + starts[name] = offset + offset += n + for global_idx in wrapped.indices: + name, _ = split.index_map[global_idx] + per_episode[name] += 1 + # K frames from each episode, or the whole episode when it is shorter + assert per_episode == {"ep0": 73, "ep1": 73, "ep2": 40} + assert len(wrapped) == sum(per_episode.values()) + assert len(wrapped) <= LIMIT_VAL_BATCHES * BATCH + # ... and the frames really are spread across each episode, not its head + ep0 = [ + split.index_map[g][1] for g in wrapped.indices if split.index_map[g][0] == "ep0" + ] + assert ep0[0] == 0 and ep0[-1] == 499 + + # samples still come through (the wrapper indexes into the base) + assert wrapped[0]["frame"] == split[wrapped.indices[0]]["frame"] + + +def test_metric_wrapper_forwards_norm_stats_to_the_base(): + """__getitem__ delegates to base, so the stats must land on the base or the + subsampled val loader would be unnormalized.""" + split = _split({"ep0": 10}) + wrapped = EvenStrideDataset(split, frames_per_episode=5) + source = MultiDataset(state={}) + source.norm_stats = {0: {"actions_cartesian": {"mean": np.zeros(3)}}} + wrapped.set_norm_stats_from(source) + assert split.norm_stats is source.norm_stats + + +def test_every_head_can_be_subsampled_independently(): + cfg = _cfg(metric={"valid": 5, "unseen_op_valid": 2}) + heads = {h: {HUMAN: _split({"ep0": 100})} for h in ("valid", "train_viz")} + heads["unseen_op_valid"] = {HUMAN: _split({"ep0": 100})} + assert len(th._subsample_val_datasets(cfg, "valid", heads["valid"])[HUMAN]) == 5 + assert ( + len( + th._subsample_val_datasets( + cfg, "unseen_op_valid", heads["unseen_op_valid"] + )[HUMAN] + ) + == 2 + ) + # train_viz has no entry -> untouched + assert ( + th._subsample_val_datasets(cfg, "train_viz", heads["train_viz"]) + is heads["train_viz"] + ) + with pytest.raises(ValueError, match="positive int"): + th._metric_frames_per_episode(_cfg(metric={"valid": 0}), "valid") + + +def test_k_scales_with_the_world_size(): + """K in the config is the single-rank value. DistributedSampler strides, + so W ranks between them read W times the frames in the same number of + batches each; without the scaling a 4-GPU run scores a quarter of what the + same config scores on one GPU.""" + cfg = _cfg(metric={"valid": 7}) + for devices, nodes, expected in ( + (1, 1, 7), + (4, 1, 28), + (8, 2, 112), + ([0, 1], 1, 14), + ("auto", 1, 7), # unknown at config time: stay single-rank + (-1, 1, 7), + ): + cfg.trainer = {"devices": devices, "num_nodes": nodes} + assert th._metric_frames_per_episode(cfg, "valid") == expected, devices + # a split shorter than K is kept whole, so scaling cannot truncate + cfg.trainer = {"devices": 4, "num_nodes": 1} + split = {HUMAN: _split({"ep0": 20})} + assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 + + +def test_eval_mode_is_single_rank(): + """Eval runs are forced onto one device after the loaders are built, so + their val split must be subsampled for one rank, not the training W.""" + cfg = _cfg(metric={"valid": 7}) + cfg.trainer = {"devices": 4, "num_nodes": 2} + assert th._metric_frames_per_episode(cfg, "valid") == 56 + cfg.mode = "eval" + assert th._metric_frames_per_episode(cfg, "valid") == 7 + del cfg["mode"] + cfg.eval = True + assert th._metric_frames_per_episode(cfg, "valid") == 7 + + +@pytest.mark.parametrize("trainer_cfg", ["ddp", "ddp_pi"]) +def test_ddp_devices_are_per_node(trainer_cfg, compose_resolve): + """Lightning reads devices per node; devices * num_nodes is the world.""" + cfg = compose_resolve( + "train_zarr_cartesian", + [ + f"trainer={trainer_cfg}", + "launch_params.gpus_per_node=8", + "launch_params.nodes=2", + ], + ) + assert cfg.trainer.devices == 8 and cfg.trainer.num_nodes == 2 + assert th._trainer_world_size(cfg) == 16 + + +def test_auto_k_is_derived_from_the_resolved_split(): + """`auto` = floor(limit_val_batches * batch_size / episodes), per dataset, + from the head's own loader batch size, then scaled by the world size.""" + cfg = _cfg(metric={"valid": "auto"}) + cfg.trainer = {"limit_val_batches": 4, "devices": 2} + split = {HUMAN: _split({"ep0": 100, "ep1": 100, "ep2": 100})} + params = {HUMAN: {"batch_size": 5}} + wrapped = th._subsample_val_datasets(cfg, "valid", split, params)[HUMAN] + assert wrapped.frames_per_episode == (4 * 5 // 3) * 2 + for limit in (1.0, None, True): + cfg.trainer.limit_val_batches = limit + with pytest.raises(ValueError, match="auto"): + th._subsample_val_datasets(cfg, "valid", split, params) + cfg.trainer.limit_val_batches = 0 # validation off, as the bench scripts run + assert th._subsample_val_datasets(cfg, "valid", split, params) == split + cfg.trainer.limit_val_batches = 4 + with pytest.raises(ValueError, match="batch_size"): + th._subsample_val_datasets(cfg, "valid", split, {}) + + +# ------------------------------------------------------------- 3. video loaders +def test_video_dataset_holds_only_the_pins_in_frame_order(): + split = _split({"ep0": 30, "ep1": 20, "ep2": 25}) + sub = pinned_episode_subset(split, ["ep2", "ep0"], dataset_name=HUMAN) + + assert list(sub.datasets) == ["ep2", "ep0"], "config order is playback order" + assert len(sub) == 55 + # contiguous: frame indices run 0..n-1 inside each episode, in order + by_episode: dict[str, list[int]] = {} + for name, local in sub.index_map: + by_episode.setdefault(name, []).append(local) + assert by_episode == {"ep2": list(range(25)), "ep0": list(range(30))} + + with pytest.raises(PinError, match="not in dataset"): + pinned_episode_subset(split, ["nope"], dataset_name=HUMAN) + + +def test_video_datasets_are_built_per_head_from_the_resolved_split(): + valid = _split({"a": 10, "b": 10}) + unseen = _split({"c": 10, "d": 10}) + cfg = _cfg(video={"valid": ["b"], "unseen_op_valid": ["c", "d"]}) + built = th._video_datasets( + cfg, + { + "valid": {HUMAN: valid}, + "train_viz": {HUMAN: _split({"t": 10})}, + "unseen_op_valid": {HUMAN: unseen}, + }, + ) + assert set(built) == {"valid", "unseen_op_valid"}, "no pins -> no video loader" + assert list(built["valid"][HUMAN].datasets) == ["b"] + assert list(built["unseen_op_valid"][HUMAN].datasets) == ["c", "d"] + # the video dataset reuses the SAME leaf objects: no second resolve + assert built["valid"][HUMAN].datasets["b"] is valid.datasets["b"] + + with pytest.raises(PinError, match="none of its datasets"): + th._video_datasets(_cfg(video={"valid": ["c"]}), {"valid": {HUMAN: valid}}) + + +def test_val_loader_names_append_the_video_loaders(): + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + datasets = {HUMAN: _split({"ep0": 4})} + dm = MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets=datasets, + train_viz_dataloader_params=params, + unseen_op_valid_datasets=datasets, + unseen_op_valid_dataloader_params=params, + video_datasets={ + "valid": {HUMAN: _split({"pin": 4})}, + "unseen_op_valid": {HUMAN: _split({"pin2": 4})}, + }, + ) + # metric loaders keep their historical indices; video loaders follow + assert dm.val_loader_names() == [ + "valid", + "train_viz", + "unseen_op_valid", + "valid_video", + "unseen_op_valid_video", + ] + loaders = dm.val_dataloader() + assert len(loaders) == 5 + # the video loader really iterates the pinned dataset, unshuffled + batch, _, _ = next(iter(loaders[3])) + assert [float(f) for f in batch[HUMAN]["frame"]] == [0.0, 1.0] + + # a head whose video datasets were dropped gets no name + dm.video_datasets = {} + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + + +# ----------------------------------------------------------------- 4. dispatch +class _StubEval: + """Records how validation_step called it, with EvalVideo's viz gating.""" + + def __init__(self, should_viz=True): + self.calls: list[tuple] = [] + self._viz = should_viz + + def _should_viz(self): + return self._viz + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + self.calls.append((batch_idx, dataloader_idx, mode)) + + +def _wrapper(names, heads): + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) + w.model = SimpleNamespace(process_batch_for_training=lambda b: b) + w.evaluator = heads.get("valid") + w.train_viz_evaluator = heads.get("train_viz") + w.unseen_op_valid_evaluator = heads.get("unseen_op_valid") + w.val_loader_names = names + return w + + +def test_validation_step_splits_metrics_from_video(): + valid, unseen = _StubEval(), _StubEval() + names = ["valid", "unseen_op_valid", "valid_video", "unseen_op_valid_video"] + w = _wrapper(names, {"valid": valid, "unseen_op_valid": unseen}) + + for idx in range(4): + w.validation_step({"k": idx}, 0, dataloader_idx=idx) + + assert valid.calls == [(0, 0, "metrics"), (0, 2, "video")] + assert unseen.calls == [(0, 1, "metrics"), (0, 3, "video")] + + +def test_non_video_epoch_skips_the_video_loader_before_the_forward(): + valid = _StubEval(should_viz=False) + processed = [] + w = _wrapper(["valid", "valid_video"], {"valid": valid}) + w.model = SimpleNamespace( + process_batch_for_training=lambda b: processed.append(b) or b + ) + w.validation_step({"k": 0}, 0, dataloader_idx=0) + w.validation_step({"k": 1}, 0, dataloader_idx=1) + assert valid.calls == [(0, 0, "metrics")], "video loader skipped" + assert processed == [{"k": 0}], "no forward pass for the skipped loader" + + +def test_head_without_a_video_loader_keeps_the_old_combined_call(): + """Evaluators that predate the split (and the stubs in older tests) are + called with no ``mode`` kwarg at all.""" + seen = [] + + class _Old: + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + seen.append((batch_idx, dataloader_idx)) + + w = _wrapper(["valid"], {"valid": _Old()}) + w.validation_step({"k": 0}, 7, dataloader_idx=0) + assert seen == [(7, 0)] + + +# ----------------------------------------------- EvalVideo's two modes +class _Trainer: + def __init__(self, epoch=0, max_epochs=1): + self.current_epoch = epoch + self.max_epochs = max_epochs + self.world_size = 1 + self.is_global_zero = True + self.loggers = [] + self.logged: list[dict] = [] + outer = self + + class _LM: + device = "cpu" + + def log_dict(self, metrics, **kw): + outer.logged.append(dict(metrics)) + + self.lightning_module = _LM() + + +def _eval_video(viz_max_batches=2): + from egomimic.eval.eval_video import EvalVideo + + class _Impl(EvalVideo): + def __init__(self, **kw): + super().__init__(**kw) + self.rendered: list[int] = [] + self.forwards = 0 + + def compute_metrics_and_viz(self, batch, do_viz=True): + self.forwards += 1 + if do_viz: + self.rendered.append(batch["i"]) + return {"Valid/mse": 1.0}, {HUMAN: np.zeros((1, 2, 2, 3), np.uint8)} + return {"Valid/mse": 1.0}, {} + + ev = _Impl(viz_every_n_epochs=1, viz_max_batches=viz_max_batches) + ev.trainer = _Trainer() + return ev + + +def test_metric_mode_never_renders_and_video_mode_never_logs(): + ev = _eval_video() + for i in range(3): + ev.on_validation_step({"i": i}, i, 0, mode="metrics") + assert ev.rendered == [], "no frame rendered on the metric loader" + assert ev.val_image_buffer == {}, "no frame buffered either" + assert len(ev.trainer.logged) == 3, "metrics logged for every metric batch" + + ev.trainer.logged.clear() + for i in range(4): + ev.on_validation_step({"i": i}, i, 1, mode="video") + assert ev.rendered == [0, 1], "rendered up to viz_max_batches" + assert ev.trainer.logged == [], "video loader logs no metric key" + # past viz_max_batches the video loader does not even run the forward + assert ev.forwards == 3 + 2 + + with pytest.raises(ValueError, match="unknown validation mode"): + ev.on_validation_step({"i": 0}, 0, 0, mode="bogus") + + +def test_default_mode_is_todays_behaviour(): + ev = _eval_video() + ev.on_validation_step({"i": 0}, 0, 0) + assert ev.rendered == [0] and len(ev.trainer.logged) == 1 + + +# --------------------------------------------- video writes happen on rank 0 only +def test_only_rank_zero_writes_the_video(monkeypatch, tmp_path): + """Every rank iterates every val batch and renders the same frames, so + without a guard N ranks race to write one mp4 path.""" + import egomimic.eval.eval_video as ev_mod + + written: list[str] = [] + monkeypatch.setattr( + ev_mod.tvio, "write_video", lambda path, *a, **kw: written.append(path) + ) + + buffers = {} + for rank_zero in (True, False): + ev = _eval_video(viz_max_batches=None) + ev.trainer.is_global_zero = rank_zero + monkeypatch.setattr(type(ev), "video_dir", lambda self: str(tmp_path)) + for i in range(3): + ev.on_validation_step({"i": i}, i, 0, mode="video") + ev.on_validation_end() + # rendering, buffering and counters are identical on every rank... + buffers[rank_zero] = (ev.rendered, dict(ev.val_counter), ev.val_image_buffer) + + assert buffers[True] == buffers[False], "only the write differs between ranks" + assert len(written) == 1, "one file, from rank 0" + + +# ------------------------------------------- video loaders are gated on non-viz +class _GateTrainer: + """Just enough trainer for MultiDataModuleWrapper._video_gate.""" + + def __init__(self, evaluators): + self.lightning_module = SimpleNamespace(_val_heads=lambda: evaluators) + self.is_global_zero = True + + +def _dm(): + """A datamodule with one metric head and its pinned video companion.""" + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + datasets = {HUMAN: _split({"ep0": 4})} + return MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + video_datasets={"valid": {HUMAN: _split({"pin": 4})}}, + ) + + +def test_video_loader_yields_nothing_when_the_gate_is_closed(): + from egomimic.pl_utils.pl_data_utils import _VizGatedLoader + + dm = _dm() + metric, video = dm.val_dataloader() + assert not isinstance(metric, _VizGatedLoader), "metric loaders run every pass" + assert isinstance(video, _VizGatedLoader) + + viz = _StubEval(should_viz=True) + dm.trainer = _GateTrainer({"valid": viz}) + assert len(list(iter(video))) == 2, "viz pass: the pinned episode is read" + + viz._viz = False + assert list(iter(video)) == [], "non-viz pass: not one batch is built" + # the length Lightning caches at setup does not move with the gate + assert len(video) == 2 + + +def test_gate_is_open_without_a_trainer_or_an_evaluator(): + dm = _dm() + _, video = dm.val_dataloader() + assert len(list(iter(video))) == 2, "no trainer attached -> gate open" + + dm.trainer = _GateTrainer({"valid": None}) + assert len(list(iter(video))) == 2, "head with no evaluator -> gate open" + + dm.trainer = _GateTrainer({"valid": object()}) + assert len(list(iter(video))) == 2, "evaluator without _should_viz -> gate open" + + +def test_gated_loader_length_survives_a_closed_gate_at_setup(): + """Lightning reads len() once, in setup_data, which can land on a non-viz + epoch: CombinedLoader.__len__ needs a live iterator the gate never made.""" + dm = _dm() + _, video = dm.val_dataloader() + dm.trainer = _GateTrainer({"valid": _StubEval(should_viz=False)}) + assert list(iter(video)) == [] + assert len(video) == 2 + + +def test_gated_loader_length_is_lightnings_own_and_opens_no_iterator(): + from lightning.pytorch.utilities.combined_loader import _MaxSizeCycle + + dm = _dm() + _, video = dm.val_dataloader() + video.limits = [1] + assert len(video) == len(_MaxSizeCycle(video.flattened, [1])) == 1 + assert video._iterator is None + + +def test_video_head_without_a_viz_cap_is_rejected(): + """Rank 0 renders a pinned video loader alone while the other ranks wait, so + an uncapped head is bounded only by the process-group timeout.""" + capped = SimpleNamespace(viz_func=print, viz_max_batches=10) + uncapped = SimpleNamespace(viz_func=print, viz_max_batches=None) + no_viz = SimpleNamespace(viz_func=None, viz_max_batches=None) + + def _run(heads, video): + model = SimpleNamespace(_val_heads=lambda: heads) + th._require_capped_video_heads(model, SimpleNamespace(video_datasets=video)) + + _run({"valid": capped}, {"valid": {HUMAN: object()}}) + _run({"valid": uncapped}, {"valid": {}}) + _run({"valid": uncapped}, {}) + _run({"valid": no_viz}, {"valid": {HUMAN: object()}}) + with pytest.raises(ValueError, match="viz_max_batches"): + _run({"valid": capped, "train_viz": uncapped}, {"train_viz": {HUMAN: object()}}) + + +# sanity steps on: the loaders (and their len()) are set up on a pass whose gate +# is already closed, which is what a real run does -- trainer/default.yaml sets no +# num_sanity_val_steps, so lightning's default of 2 applies. +@pytest.mark.parametrize("sanity_steps", [0, 2]) +def test_a_real_fit_builds_video_batches_only_on_viz_passes(tmp_path, sanity_steps): + """The gate against the real Lightning evaluation loop: a val dataloader + that yields nothing mid-sequence must not derail the pass, and the pinned + episode must not be touched at all on a non-viz epoch.""" + import torch + from lightning import Trainer + + class _Counting(_Episode): + def __init__(self, n): + super().__init__(n) + self.reads = 0 + + def __getitem__(self, i): + self.reads += 1 + return super().__getitem__(i) + + pinned, metric = _Counting(4), _Counting(4) + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + dm = MultiDataModuleWrapper( + train_datasets={HUMAN: _split({"ep0": 4})}, + valid_datasets={HUMAN: MultiDataset(datasets={"ep0": metric}, mode="total")}, + train_dataloader_params={HUMAN: {"batch_size": 2, "num_workers": 0}}, + valid_dataloader_params=params, + video_datasets={ + "valid": {HUMAN: MultiDataset(datasets={"pin": pinned}, mode="total")} + }, + ) + + class _Module(LightningModule): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + self.val_loaders_seen: list[tuple[int, int]] = [] + self.pinned_reads: dict[int, int] = {} + + # viz only on the second validation pass + def _val_heads(self): + return { + "valid": SimpleNamespace(_should_viz=lambda: self.current_epoch == 1) + } + + def training_step(self, batch, batch_idx): + return self.layer(batch[HUMAN]["frame"].reshape(-1, 1)).sum() + + def validation_step(self, batch, batch_idx, dataloader_idx=0): + if not self.trainer.sanity_checking: + self.val_loaders_seen.append((self.current_epoch, dataloader_idx)) + + def on_validation_epoch_end(self): + self.pinned_reads[self.current_epoch] = pinned.reads + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + model = _Module() + Trainer( + max_epochs=2, + accelerator="cpu", + devices=1, + limit_train_batches=2, + limit_val_batches=2, + check_val_every_n_epoch=1, + num_sanity_val_steps=sanity_steps, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=str(tmp_path), + ).fit(model, datamodule=dm) + + assert model.pinned_reads[0] == 0, "non-viz pass never touched the pinned episode" + assert model.pinned_reads[1] > 0, "viz pass still reads it" + # the metric loader is unaffected, and dataloader_idx stays stable + by_epoch = {e: [i for ep, i in model.val_loaders_seen if ep == e] for e in (0, 1)} + assert by_epoch[0] == [0, 0], "epoch 0: metric loader only" + assert by_epoch[1] == [0, 0, 1, 1], "epoch 1: metric loader then video loader" + + +def test_video_loader_is_closed_off_rank_zero(): + """Only rank 0 writes the overlay, and the val loaders carry no + DistributedSampler, so every other rank was decoding, forwarding and + rendering the same frames for nothing.""" + dm = _dm() + _, video = dm.val_dataloader() + viz = _StubEval(should_viz=True) + + dm.trainer = _GateTrainer({"valid": viz}) + dm.trainer.is_global_zero = True + assert len(list(iter(video))) == 2 + + dm.trainer.is_global_zero = False + assert list(iter(video)) == [], "rank 1 builds no video batch at all" + + +def test_two_rank_ddp_fit_survives_the_rank_divergence(tmp_path): + """The gate makes ranks iterate *different* numbers of val batches. Proven + against real DDP (gloo, 2 CPU processes) with ModelWrapper's collective + pattern -- sync_dist metrics off the metric loader, nothing off the video + one -- because a desync here would hang a campaign, not fail it.""" + import json + import subprocess + + script = Path(__file__).resolve().parents[1] / "fixtures" / "ddp_video_gate.py" + repo = Path(__file__).resolve().parents[2] + proc = subprocess.run( + [sys.executable, str(script), str(tmp_path)], + env={**os.environ, "PYTHONPATH": str(repo)}, + capture_output=True, + text=True, + timeout=600, + ) + assert proc.returncode == 0, proc.stderr[-3000:] + + ranks = { + int(p.stem.removeprefix("rank")): json.loads(p.read_text()) + for p in tmp_path.glob("rank*.json") + } + assert set(ranks) == {0, 1}, "both ranks finished; neither hung" + # epoch 0 renders nothing, epoch 1 does, and only on rank 0 + assert ranks[0]["video_steps"] == [1, 1] + assert ranks[0]["pinned_reads"] >= 4, "rank 0 read the frames it stepped" + # exactly 0, not "few": a gate that still builds an iterator makes the + # loader's workers prefetch, which is the cost this whole change is about + assert ranks[1]["video_steps"] == [] and ranks[1]["pinned_reads"] == 0 diff --git a/tests/unit/test_eval_video_upload.py b/tests/unit/test_eval_video_upload.py new file mode 100644 index 0000000000..145baabab6 --- /dev/null +++ b/tests/unit/test_eval_video_upload.py @@ -0,0 +1,61 @@ +"""Validation videos written by ``EvalVideo`` are also logged to wandb.""" + +from __future__ import annotations + +import types +from pathlib import Path + +import torch +from lightning.pytorch.loggers import WandbLogger + +from egomimic.eval.eval_video import EvalVideo + +HUMAN_BIMANUAL = 3 + + +class _Eval(EvalVideo): + def __init__(self, root): + super().__init__() + self._root = root + + def root_dir(self): + return self._root + + def compute_metrics_and_viz(self, batch, do_viz=True): + return {}, {} + + +def _trainer(loggers, epoch=99, step=1234): + return types.SimpleNamespace( + is_global_zero=True, current_epoch=epoch, global_step=step, loggers=loggers + ) + + +def _frames(): + return torch.randint(0, 255, (8, 32, 32, 3), dtype=torch.uint8) + + +def test_video_logged_to_wandb(tmp_path, monkeypatch): + monkeypatch.setenv("WANDB_MODE", "offline") + monkeypatch.setenv("WANDB_SILENT", "true") + logger = WandbLogger(save_dir=str(tmp_path), offline=True, project="t") + ev = _Eval(str(tmp_path / "run")) + ev.trainer = _trainer([logger]) + ev.val_counter[HUMAN_BIMANUAL] = 0 + + ev._write_video(HUMAN_BIMANUAL, _frames()) + run_dir = Path(logger.experiment.dir) + logger.experiment.finish() + + local = tmp_path / "run/videos/epoch_99/HUMAN_BIMANUAL/validation_video_0.mp4" + assert local.exists() + media = list((run_dir / "media/videos/videos").glob("HUMAN_BIMANUAL_*.mp4")) + assert len(media) == 1 + + +def test_no_wandb_logger_is_a_noop(tmp_path): + ev = _Eval(str(tmp_path)) + ev.trainer = _trainer([]) + ev.val_counter[HUMAN_BIMANUAL] = 0 + ev._write_video(HUMAN_BIMANUAL, _frames()) + assert (tmp_path / "videos/epoch_99/HUMAN_BIMANUAL/validation_video_0.mp4").exists() diff --git a/tests/unit/test_flagship_kp_configs.py b/tests/unit/test_flagship_kp_configs.py new file mode 100644 index 0000000000..05324f0cb8 --- /dev/null +++ b/tests/unit/test_flagship_kp_configs.py @@ -0,0 +1,79 @@ +"""The flagship keypoint HPT arm: recipe, data split and stem knobs compose.""" + +from __future__ import annotations + +import pytest + +RECIPE = "train_zarr_mecka_flagship_kp_hpt" + + +def test_kp_recipe_pairs_the_keypoint_data_evaluator_and_prompt(compose_resolve): + cfg = compose_resolve(RECIPE, []) + rm = cfg.model.robomimic_model + assert rm._target_ == "egomimic.algo.hpt.HPT" + assert rm.ac_keys.human_bimanual == "actions_keypoints" + assert rm.dims.human_bimanual.action == 144 + assert rm.dims.human_bimanual.proprio == 144 + # The flagship split carries no per-frame annotations, so every sample is + # meant to see one constant prompt; a comparison arm must set the same one. + assert rm.default_prompt == "fold clothes" + assert cfg.evaluator._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert ( + "keypoints_revert_6d_wristframe" + in cfg.evaluator.transform_lists.human_bimanual._target_ + ) + # The train-time overlay wraps its own evaluator under `base`, and renders + # the keypoint action through Human.viz(mode="keypoints"). + train_viz = cfg.train_viz_evaluator.base + assert train_viz._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert train_viz.viz_func.human_bimanual.mode == "keypoints" + assert train_viz.viz_func.human_bimanual.action_key == "actions_keypoints" + assert ( + train_viz.transform_lists.human_bimanual._target_ + == cfg.evaluator.transform_lists.human_bimanual._target_ + ) + resolver = cfg.data.train_datasets.human_bimanual.resolver + assert resolver.key_map.keymap_mode == "keypoints" + assert resolver.transform_list.mode == "keypoints_wristframe_6d" + assert resolver.transform_list.pad_proprio_gripper is False + + +def test_kp_top3_data_keeps_the_three_operators_and_the_keypoint_pipeline( + compose_resolve, +): + cfg = compose_resolve(RECIPE, ["data=mecka_fold_flagship_top3_hpt_keypoints"]) + lambdas = cfg.data.train_datasets.human_bimanual.filters.filter_lambdas + assert any("6903686e0e94ce070afd1f24" in fn for fn in lambdas) + assert ( + cfg.data.train_datasets.human_bimanual.resolver.transform_list.mode + == "keypoints_wristframe_6d" + ) + + +def test_hpt_kp_model_carries_the_history_and_backbone_lr_knobs(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", ["model=hpt_bc_keypoints_wrist_300M"]) + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert stem.history_len == 1 + assert stem.history_dropout == 0.2 + assert cfg.model.backbone_lr_scale == 1.0 + + +# The two knobs are one setting in two files: the keymap decides how many +# proprio steps a sample carries, the stem how many tokens consume them, and +# HPTModel.stem_process raises when they disagree. Composition alone cannot see +# it, so every keypoint data config this model is documented to pair with is +# checked here. +@pytest.mark.parametrize( + "data_cfg", + [ + "mecka_fold_flagship_opsplit_hpt_keypoints", + "mecka_fold_flagship_top3_hpt_keypoints", + ], +) +def test_keymap_proprio_history_matches_the_stem_that_consumes_it( + data_cfg, compose_resolve +): + cfg = compose_resolve(RECIPE, [f"data={data_cfg}"]) + key_map = cfg.data.train_datasets.human_bimanual.resolver.key_map + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert key_map.get("proprio_history", 1) == stem.history_len diff --git a/tests/unit/test_format_bench_update.py b/tests/unit/test_format_bench_update.py new file mode 100644 index 0000000000..b277d224a5 --- /dev/null +++ b/tests/unit/test_format_bench_update.py @@ -0,0 +1,156 @@ +"""Each edit of the update bench changes its target in every format and nothing else.""" + +from __future__ import annotations + +import csv +import hashlib +import json +from pathlib import Path + +import lance +import numpy as np +import pyarrow.parquet as pq +import pytest +import zarr + +from egomimic.rldb.lance.schema import HUMAN_STEPS +from egomimic.scripts.format_bench import bench_update as bu + +ROWS = 1300 + + +@pytest.fixture +def lay(tmp_path): + """Three source recordings; episode lengths that are not multiples of the zarr chunk, and + a 150-frame episode with no annotation span.""" + rng = np.random.default_rng(0) + batches = [] + for i, n in enumerate((300, 450, 700)): + cols = { + k: rng.standard_normal((n, d), dtype=np.float32) + for k, d in bu.VECTORS.items() + } + cols.update({k: rng.random(n, dtype=np.float32) for k in bu.SCALARS}) + cols.update({k: rng.random(n) > 0.1 for k in bu.BOOLS}) + cols["g"] = np.arange(n) + batches.append(bu.steps_batch(cols, f"src-{i}")) + pilot = tmp_path / "pilot" + lance.write_dataset( + batches, str(pilot / "lance/steps/aria.lance"), schema=HUMAN_STEPS + ) + with open(pilot / "manifest.csv", "w") as f: + w = csv.writer(f) + w.writerow(["episode_hash", "num_frames"]) + w.writerows([["a", 150], ["b", 330], ["c", 640]]) + lay = bu.pilot_layout(pilot, ROWS) + assert lay.lens.sum() == ROWS and lay.n_eps > 3 + assert min(len(lay.spans(ep, 0)) for ep in range(lay.n_eps)) == 0 + assert max(len(lay.spans(ep, 0)) for ep in range(lay.n_eps)) > 1 + return lay + + +def _digests(root: Path) -> dict[str, str]: + return { + str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(root.rglob("*")) + if p.is_file() + } + + +def test_zarr_shards_match_zarr(tmp_path, lay): + """Hand-built objects read back through zarr exactly as ZarrWriter would write them.""" + root = tmp_path / "zarr" + bu.make_zarr(str(root), lay, workers=4) + for ep in range(lay.n_eps): + g = zarr.open_group(str(root / bu.ep_name(ep)), mode="r") + hp = lay.pool["head_pose"][lay.src[lay.g(ep)]] + np.testing.assert_array_equal(g["obs_head_pose"][: len(hp)], hp) + assert not g["obs_head_pose"][len(hp) :].any() + want = [json.loads(b) for b in bu.annotation_items(lay.spans(ep, 0))] + assert [json.loads(b) for b in g["annotations"][:]] == want + assert sorted(g.keys()) == ["annotations", "obs_head_pose"] + + +@pytest.mark.parametrize("op", bu.OPS) +def test_zarr(tmp_path, lay, op): + root = tmp_path / "zarr" + bu.make_zarr(str(root), lay, workers=4) + target = {"replace": "obs_head_pose/c", "annotations": "annotations/"}.get(op) + before = _digests(root) + bu.update_zarr(str(root), lay, op, k=3, workers=4) + bu.verify_zarr(str(root), lay, op, k=3) + after = _digests(root) + for path, digest in before.items(): + if op == "append" and path.count("/") == 1 and path.endswith("zarr.json"): + continue + if target is None or target not in path: + assert after.get(path) == digest, path + if op == "append": + g = zarr.open_group(str(root / bu.ep_name(0)), mode="r") + assert set(g.attrs["features"]) == {"obs_head_pose", "annotations", "feat_3"} + + +@pytest.mark.parametrize("fmt", ["lerobot", "lerobot_zstd"]) +@pytest.mark.parametrize("op", bu.OPS) +def test_lerobot(tmp_path, lay, monkeypatch, op, fmt): + monkeypatch.setattr(bu, "PARQUET_MB", 0.2) + root = tmp_path / fmt + bu.make_lerobot(str(root), lay, workers=4, compression=bu.COMPRESSION[fmt]) + files = sorted((root / "data").rglob("*.parquet")) + assert len(files) > 1 + assert pq.ParquetFile(files[0]).metadata.row_group(0).column(0).compression == ( + bu.COMPRESSION[fmt].upper() + ) + target = {"replace": ["head_pose"], "annotations": bu.LANGUAGE_COLS}.get(op, []) + before = pq.read_table(files).drop_columns(target) + groups = [pq.ParquetFile(f).num_row_groups for f in files] + meta_before = pq.read_table(root / "meta/episodes") + bu.update_lerobot(str(root), lay, op, k=3, workers=4) + bu.verify_lerobot(str(root), lay, op, k=3) + after = pq.read_table(files) + if op == "append": + after = after.drop_columns(["feat_3"]) + assert after.drop_columns(target).equals(before) + assert [pq.ParquetFile(f).num_row_groups for f in files] == groups + meta = pq.read_table(root / "meta/episodes") + for c in meta_before.column_names: + changed = ( + op == "replace" and c.startswith("stats/head_pose/") and "count" not in c + ) + assert meta.column(c).equals(meta_before.column(c)) != changed, c + if op == "append": + assert "stats/feat_3/mean" in meta.column_names + info = json.loads((root / "meta/info.json").read_text()) + assert "feat_3" in info["features"] + + +@pytest.mark.parametrize("op", bu.OPS) +def test_lance(tmp_path, lay, op): + root = str(tmp_path / "lance") + bu.make_lance(root, lay, rows_per_frag=500) + steps = lance.dataset(bu.steps_uri(root)) + assert len(steps.get_fragments()) > 1 + steps_before = steps.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() + 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 anns_after.num_rows == anns_before.num_rows + return + 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) + if op == "replace": + assert steps_after.drop_columns(["head_pose"]).equals( + steps_before.drop_columns(["head_pose"]) + ) + else: + assert steps_after.drop_columns(["feat_3"]).equals(steps_before) diff --git a/tests/unit/test_format_bench_zarr_fast.py b/tests/unit/test_format_bench_zarr_fast.py new file mode 100644 index 0000000000..9d04c29cf8 --- /dev/null +++ b/tests/unit/test_format_bench_zarr_fast.py @@ -0,0 +1,123 @@ +"""The format bench's FastZarrDataset emits exactly ZarrDataset's samples, alone and batched. + +The episode is ramped (cell = frame * 1000 + column) and 160 frames long, so windows cross the +100-row inner-chunk boundary and a wrong offset changes numbers.""" + +from __future__ import annotations + +import numpy as np +import pytest +import simplejpeg +import torch + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, ZarrDataset +from egomimic.rldb.zarr.zarr_writer import ZarrWriter +from egomimic.scripts.format_bench import zarr_fast +from egomimic.scripts.format_bench.zarr_fast import FastZarrDataset + +T = 160 + + +def _ramp(width: int) -> np.ndarray: + return ( + np.arange(T, dtype=np.float64)[:, None] * 1000.0 + + np.arange(width, dtype=np.float64)[None, :] + ) + + +def _write(path, seed): + numeric = {"obs_head_pose": _ramp(7)} + for side, off in (("left", 0.0), ("right", 0.5)): + numeric[f"{side}.obs_wrist_pose"] = _ramp(7) + off + numeric[f"{side}.obs_keypoints"] = _ramp(63) + off + numeric[f"{side}.obs_ee_pose"] = _ramp(7) + off + rng = np.random.default_rng(seed) + return ZarrWriter.create_and_write( + path, + numeric_data=numeric, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + task_name="ramp", + intrinsics={"front_1": np.eye(3, 4)}, + ) + + +@pytest.fixture(scope="module") +def episodes(tmp_path_factory): + d = tmp_path_factory.mktemp("fast") + return [_write(d / f"ep{i}.zarr", i) for i in range(2)] + + +def _keymap(frames8: bool): + km = Human.get_keymap(keymap_mode="keypoints", include_ee_pose=True) + if frames8: + km = { + k: {**v, "history": 8} if v.get("key_type") == "camera_keys" else v + for k, v in km.items() + } + return km + + +def _same(a: dict, b: dict): + assert a.keys() == b.keys() + for k in a: + if isinstance(a[k], torch.Tensor): + assert a[k].dtype == b[k].dtype and torch.equal(a[k], b[k]), k + else: + assert a[k] == b[k], k + + +@pytest.mark.parametrize("frames8", [False, True]) +@pytest.mark.parametrize("idx", [0, 5, 70, 85, 99, 100, 130, T - 1]) +def test_getitem_matches_zarr_dataset(episodes, frames8, idx): + km = _keymap(frames8) + _same(FastZarrDataset(episodes[0], km)[idx], ZarrDataset(episodes[0], km)[idx]) + + +@pytest.mark.parametrize("frames8", [False, True]) +def test_getitems_batches_into_one_fetch_and_matches(episodes, frames8, monkeypatch): + km = _keymap(frames8) + fast = MultiDataset( + datasets={f"e{i}": FastZarrDataset(p, km) for i, p in enumerate(episodes)}, + mode="total", + ) + ref = MultiDataset( + datasets={f"e{i}": ZarrDataset(p, km) for i, p in enumerate(episodes)}, + mode="total", + ) + calls = [] + real = zarr_fast.fetch_many + monkeypatch.setattr( + zarr_fast, "fetch_many", lambda items: calls.append(len(items)) or real(items) + ) + idxs = [0, 3, 99, 150, T, T + 85, 2 * T - 1] + batch = fast.__getitems__(idxs) + assert calls == [len(idxs)], "one fetch for the whole batch" + for i, got in zip(idxs, batch): + _same(got, ref[i]) + + +def test_batched_jpeg_failure_falls_back_to_a_fresh_read(episodes, monkeypatch): + km = _keymap(False) + ds = FastZarrDataset(episodes[0], km) + multi = MultiDataset(datasets={"e0": ds}, mode="total") + real = simplejpeg.decode_jpeg + n = [] + + def fail_once(data, **kw): + n.append(1) + if len(n) == 1: + raise ValueError("corrupt jpeg") + return real(data, **kw) + + monkeypatch.setattr( + "egomimic.rldb.zarr.zarr_dataset_multi.simplejpeg.decode_jpeg", fail_once + ) + (sample,) = multi.__getitems__([60]) + frame = int(round(float(np.asarray(sample["obs_head_pose"])[0]) / 1000.0)) + assert frame != 60, "the retry moved to another frame and read it" + np.testing.assert_array_equal( + np.asarray(sample["obs_head_pose"]), _ramp(7)[frame].astype(np.float32) + ) + assert ds.episode_reader._cache()["prefetched"] is None diff --git a/tests/unit/test_hand_masks.py b/tests/unit/test_hand_masks.py new file mode 100644 index 0000000000..cfa69a748a --- /dev/null +++ b/tests/unit/test_hand_masks.py @@ -0,0 +1,226 @@ +"""An untracked hand (Lance `action_hand_mask` / `proprio_hand_valid`) is left out of +the loss, the norm stats and the bounds check, and reads as the mean pose in proprio.""" + +import numpy as np +import pytest +import torch +import torch.nn.functional as F + +from egomimic.algo.hpt import HPT +from egomimic.models.denoising_policy import DenoisingPolicy +from egomimic.rldb.embodiment.embodiment import EMBODIMENT, get_embodiment_id +from egomimic.rldb.multi import MultiDataset, hand_keep +from egomimic.utils.pose_utils import hand_channel_mask + +LEFT_ONLY = torch.tensor([[1.0, 0.0], [1.0, 0.0]]) + + +def test_hand_channel_mask_splits_bimanual_layouts_in_halves(): + for width in (14, 20, 144): + m = hand_channel_mask(np.array([1.0, 0.0]), width) + assert m[: width // 2].all() and not m[width // 2 :].any() + assert hand_channel_mask(np.array([1.0, 0.0]), 7) is None + + +class _FixedPred(DenoisingPolicy): + def __init__(self, pred): + torch.nn.Module.__init__(self) + self.action_horizon = pred.shape[1] + self.pooling = None + self.padding = None + self._pred = pred + + def predict(self, actions, global_cond): + return self._pred, actions + + +def test_denoising_loss_leaves_out_the_untracked_hand(): + g = torch.Generator().manual_seed(0) + target, pred = ( + torch.randn(2, 4, 144, generator=g), + torch.randn(2, 4, 144, generator=g), + ) + policy = _FixedPred(pred) + cond = torch.zeros(2, 3) + pad = torch.ones(2, 4, 1) + loss = policy.compute_loss( + cond, {"action": target, "pad_mask": pad, "action_hand_mask": LEFT_ONLY} + ) + torch.testing.assert_close(loss, F.mse_loss(pred[..., :72], target[..., :72])) + both = policy.compute_loss( + cond, {"action": target, "pad_mask": pad, "action_hand_mask": torch.ones(2, 2)} + ) + torch.testing.assert_close( + both, policy.compute_loss(cond, {"action": target, "pad_mask": pad}) + ) + + +def test_denoising_loss_refuses_a_hand_mask_it_cannot_place(): + policy = _FixedPred(torch.zeros(2, 4, 7)) + with pytest.raises(ValueError, match="bimanual layout"): + policy.compute_loss( + torch.zeros(2, 3), + {"action": torch.zeros(2, 4, 7), "action_hand_mask": LEFT_ONLY}, + ) + + +class _NoRename: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def test_hpt_process_batch_keeps_the_hand_mask(): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRename() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = False + emb = "human_bimanual" + algo.ac_keys = {get_embodiment_id(emb): "actions_keypoints"} + out = HPT.process_batch_for_training( + algo, + { + emb: { + "actions_keypoints": torch.zeros(2, 6, 144), + "action_hand_mask": LEFT_ONLY, + "proprio_hand_valid": LEFT_ONLY, + } + }, + )[get_embodiment_id(emb)] + torch.testing.assert_close(out["action_hand_mask"], LEFT_ONLY) + assert "proprio_hand_valid" not in out + + +def _pi(): + from egomimic.algo.pi import PI + from egomimic.utils.action_utils import ( + ConverterRegistry, + HumanBimanualCartesian6D, + HumanBimanualKeypoints, + ) + + pi = PI.__new__(PI) + pi.action_registry = ConverterRegistry() + pi._packed_widths = {} + human = EMBODIMENT.HUMAN_BIMANUAL.value + pi.action_registry.register(human, "actions_keypoints", HumanBimanualKeypoints()) + pi.action_registry.register(human, "actions_cartesian", HumanBimanualCartesian6D()) + return pi, human + + +def test_pi_loss_leaves_out_the_untracked_hand(): + pytest.importorskip("openpi") + pi, human = _pi() + losses = torch.zeros(2, 5, 160) + losses[..., 72:144] = 1.0 # error only on the right hand + kp = pi._reduce_loss( + losses, torch.zeros(2, 5, 144), human, "actions_keypoints", LEFT_ONLY + ) + assert kp.item() == 0.0 + plain = pi._reduce_loss(losses, torch.zeros(2, 5, 144), human, "actions_keypoints") + both = pi._reduce_loss( + losses, torch.zeros(2, 5, 144), human, "actions_keypoints", torch.ones(2, 2) + ) + torch.testing.assert_close(both, plain) + # Cartesian packs [L 9 | g | R 9 | g | pad 12]; the gripper slots stay in. + losses = torch.zeros(2, 5, 32) + losses[..., 0:10] = 1.0 + right_only = LEFT_ONLY.flip(-1) + cart = pi._reduce_loss( + losses, torch.zeros(2, 5, 18), human, "actions_cartesian", right_only + ) + assert abs(cart.item() - 1 / 23) < 1e-6 + + +def _stats_md(): + md = MultiDataset(state={}) + emb = get_embodiment_id("human_bimanual") + X = np.random.default_rng(0).normal(size=(64, 144)).astype(np.float32) + md.norm_mode = "zscore" + md.norm_stats = { + emb: { + "state": MultiDataset._compute_stats_for_array(X), + "act": MultiDataset._compute_stats_for_array(X[:, None, :]), + } + } + md.zarr_keys = { + emb: {"state": "observations.state.keypoints", "act": "actions_keypoints"} + } + md.key_types = {emb: {"state": "proprio_keys", "act": "action_keys"}} + return md, emb + + +def test_untracked_proprio_reads_as_the_mean_pose(): + md, emb = _stats_md() + data = { + "embodiment": emb, + "observations.state.keypoints": torch.full((144,), 5.0), + "proprio_hand_valid": torch.tensor([1.0, 0.0]), + } + out = md.normalize(data, emb)["observations.state.keypoints"] + assert (out[:72] != 0).all() + rot = list(range(75, 81)) + other = [c for c in range(72, 144) if c not in rot] + assert (out[other] == 0).all() + # rot6d passes through normalization: its neutral value is the mean, not 0. + mean = md.norm_stats[emb]["state"]["mean"] + np.testing.assert_allclose(out[rot].numpy(), mean[rot], rtol=1e-6) + + +def test_bounds_check_skips_the_untracked_hand(): + md, emb = _stats_md() + act = torch.zeros(1, 144) + act[0, 100] = 1e6 # an identity-pose fill value far outside the stats + data = {"embodiment": emb, "actions_keypoints": act} + assert md._check_bounds(dict(data), None, 0, "ep") is not None + data["action_hand_mask"] = torch.tensor([1.0, 0.0]) + assert md._check_bounds(data, None, 0, "ep") is None + + +def test_norm_stats_use_only_tracked_cells(): + rng = np.random.default_rng(1) + X = rng.normal(size=(40, 3, 144)).astype(np.float32) + flags = np.ones((40, 2), np.float32) + flags[:10, 1] = 0.0 + X[:10, :, 72:] = 1e6 # fill values of the untracked right hand + X[20, 0, 0] = np.nan # a genuinely bad sample is still dropped + keep = hand_keep( + {"action_hand_mask": torch.from_numpy(flags)}, + "action_keys", + "actions_keypoints", + torch.from_numpy(X), + ) + keep = np.broadcast_to(keep.numpy(), X.shape) + kept = MultiDataset._drop_nonfinite_rows(X, "actions_keypoints", keep) + assert kept.shape[0] == 39 + stats = MultiDataset._compute_stats_for_array(kept) + good = np.delete(X, 20, axis=0) + np.testing.assert_allclose( + stats["mean"][:, 72:], good[10:, :, 72:].mean(0), rtol=1e-5 + ) + np.testing.assert_allclose( + stats["mean"][:, :72], good[:, :, :72].mean(0), rtol=1e-5 + ) + assert all(np.isfinite(a).all() for a in stats.values()) + + +def test_norm_stats_refuse_a_cell_no_tracked_hand_reached(): + X = np.ones((4, 144), np.float32) + keep = np.ones_like(X, bool) + keep[:, 72:] = False + with pytest.raises(ValueError, match="tracked hand"): + MultiDataset._drop_nonfinite_rows(X, "k", keep) + + +def test_eval_takes_an_untracked_hands_gt_from_the_prediction(): + from egomimic.eval.action_metrics import fill_untracked_from_pred + + gt, pred = torch.zeros(2, 3, 144), torch.ones(2, 3, 144) + out = fill_untracked_from_pred( + {"actions_keypoints": gt, "action_hand_mask": LEFT_ONLY}, + "actions_keypoints", + pred, + )["actions_keypoints"] + assert (out[..., :72] == 0).all() and (out[..., 72:] == 1).all() + plain = {"actions_keypoints": gt} + assert fill_untracked_from_pred(plain, "actions_keypoints", pred) is plain diff --git a/tests/unit/test_hydra_configs.py b/tests/unit/test_hydra_configs.py index 0cdf0de78b..66e34e3856 100644 --- a/tests/unit/test_hydra_configs.py +++ b/tests/unit/test_hydra_configs.py @@ -31,13 +31,20 @@ def _options(group: str) -> list[str]: TRAIN_TOP_LEVEL = [c for c in TOP_LEVEL if c.startswith("train_")] -@pytest.mark.parametrize("config_name", TOP_LEVEL) +# train_* configs are composed (with more checks) by +# test_infra_port_wiring.py::test_human_action_width_matches_data_mode. +@pytest.mark.parametrize( + "config_name", [c for c in TOP_LEVEL if c not in TRAIN_TOP_LEVEL] +) def test_top_level_config_resolves(config_name, compose_resolve): compose_resolve(config_name, []) @pytest.mark.parametrize("data", _options("data")) -@pytest.mark.parametrize("top", TRAIN_TOP_LEVEL) +# train_zarr_cartesian x data is composed by test_data_filters_instantiate_to_dataset_filter. +@pytest.mark.parametrize( + "top", [t for t in TRAIN_TOP_LEVEL if t != "train_zarr_cartesian"] +) def test_data_option_resolves_under_every_train_config(top, data, compose_resolve): compose_resolve(top, [f"data={data}"]) diff --git a/tests/unit/test_image_augs.py b/tests/unit/test_image_augs.py new file mode 100644 index 0000000000..4854370d5f --- /dev/null +++ b/tests/unit/test_image_augs.py @@ -0,0 +1,109 @@ +"""BatchedColorJitter must match torchvision's per-image adjustments, keep +per-sample randomness, and hand anything it cannot batch back to the loop.""" + +import pytest +import torch +import torchvision.transforms.functional as F +from torchvision import transforms as T + +from egomimic.models.image_augs import BatchedColorJitter, PerSampleAugs + +IMAGENET = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + +@pytest.fixture +def images(): + torch.manual_seed(0) + return torch.rand(8, 3, 16, 16) + + +@pytest.mark.parametrize( + "kwargs, reference", + [ + ({"brightness": (0.7, 0.7)}, lambda x: F.adjust_brightness(x, 0.7)), + ({"contrast": (0.6, 0.6)}, lambda x: F.adjust_contrast(x, 0.6)), + ({"saturation": (1.3, 1.3)}, lambda x: F.adjust_saturation(x, 1.3)), + ({"hue": (0.04, 0.04)}, lambda x: F.adjust_hue(x, 0.04)), + ], +) +def test_matches_torchvision(images, kwargs, reference): + """A degenerate factor range pins every sample to one factor, so the batched + adjustment must equal torchvision's scalar one elementwise.""" + out = BatchedColorJitter(**kwargs)(images) + assert torch.allclose(out, reference(images), atol=1e-5) + + +def test_factors_are_per_sample(images): + jitter = BatchedColorJitter(brightness=(0.5, 1.5)) + ratios = (jitter(images) / images.clamp(min=1e-6)).flatten(1).median(dim=1).values + assert ratios.std() > 1e-3, "every sample got the same brightness factor" + + +def test_vectorizes_the_shipped_aug_list(): + augs = PerSampleAugs( + T.Compose( + [ + T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05), + IMAGENET, + ] + ) + ) + assert augs.vectorized is not None + assert isinstance(augs.vectorized[0], BatchedColorJitter) + + +def test_falls_back_for_unbatchable_transforms(images): + augs = PerSampleAugs(T.Compose([T.RandomHorizontalFlip(), IMAGENET])) + assert augs.vectorized is None + assert augs(images).shape == images.shape + + +def test_deterministic_aug_list_is_unchanged(images): + """Eval augs have no random member, so the batched path must be exact.""" + augs = PerSampleAugs(T.Compose([IMAGENET])) + assert augs.vectorized is not None + assert torch.allclose(augs(images), IMAGENET(images), atol=1e-6) + + +def test_output_stays_in_range(images): + augs = PerSampleAugs( + T.Compose( + [T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2)] + ) + ) + out = augs(images) + assert out.min() >= 0.0 and out.max() <= 1.0 + + +def test_factors_are_drawn_in_float32(): + """uniform_ into bfloat16 collapses 4096 draws onto a few dozen values, and + into uint8 it raises.""" + img = torch.zeros(4096, 3, 1, 1, dtype=torch.bfloat16) + f = BatchedColorJitter._factors((0.9, 1.1), 4096, img) + assert f.dtype == torch.bfloat16 + assert f.float().unique().numel() > 30 + BatchedColorJitter._factors( + (0.9, 1.1), 4, torch.zeros(4, 3, 1, 1, dtype=torch.uint8) + ) + + +def test_absent_camera_stays_zero_per_sample(): + """An all-zero sample is an absent camera: it must not be normalized to + -mean/std, even when the rest of its batch is present.""" + from types import SimpleNamespace + + from egomimic.algo.hpt import HPT + + stub = SimpleNamespace( + nets=SimpleNamespace(training=True), + train_image_augs=PerSampleAugs(T.Compose([T.ColorJitter(0.4), IMAGENET])), + eval_image_augs=None, + encoders={"front_img_1": object()}, + ) + torch.manual_seed(0) + images = torch.rand(3, 3, 8, 8) + images[1] = 0.0 + out = HPT._apply_image_augs(stub, images, "front_img_1") + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert not torch.equal(out[0], images[0]), "present samples are augmented" + assert torch.equal(HPT._apply_image_augs(stub, images, "not_an_encoder"), images) diff --git a/tests/unit/test_image_decode_dtype.py b/tests/unit/test_image_decode_dtype.py new file mode 100644 index 0000000000..31c405280a --- /dev/null +++ b/tests/unit/test_image_decode_dtype.py @@ -0,0 +1,28 @@ +"""Decoded frames leave the dataset as float32. + +`uint8 / 255.0` is float64 in numpy, which doubles what the collate, the +pinned-memory staging and the H2D copy move for every image in every batch -- +and the model casts to float32 on arrival regardless. +""" + +import numpy as np + + +def test_float32_division_is_bit_identical_to_float64(): + """The cheap dtype is also the same number: there are only 256 inputs, so + this is exhaustive, not a sample.""" + u = np.arange(256, dtype=np.uint8) + wide = (u / 255.0).astype(np.float32) + narrow = u.astype(np.float32) / np.float32(255.0) + assert np.array_equal(wide, narrow) + + +def test_decode_path_emits_float32(): + """Mirror of the expression in MultiDataset._read_keys' image branch.""" + decoded = np.random.randint(0, 256, size=(8, 8, 3), dtype=np.uint8) + out = np.transpose(decoded, (2, 0, 1)).astype(np.float32) / np.float32(255.0) + assert out.dtype == np.float32 + assert out.shape == (3, 8, 8) + assert np.array_equal( + out, (np.transpose(decoded, (2, 0, 1)) / 255.0).astype(np.float32) + ) diff --git a/tests/unit/test_image_history.py b/tests/unit/test_image_history.py new file mode 100644 index 0000000000..3bdaa6da67 --- /dev/null +++ b/tests/unit/test_image_history.py @@ -0,0 +1,102 @@ +"""Image history in the data path: ``Human.get_keymap(image_history_gap_s=...)`` +adds a ``_hist`` camera entry that ``ZarrDataset`` reads that many +seconds back (in the episode's own fps), and every sample carries ``fps``.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 90 +CAM = Human.VIZ_IMAGE_KEY +PAST = f"{CAM}_hist" + + +def _leaf(tmp_path, fps: int, **keymap_kwargs): + rng = np.random.default_rng(0) + t = np.arange(T, dtype=np.float64)[:, None] + pose = np.concatenate( + [t * 0.001 + np.zeros((T, 3)), np.tile([[1.0, 0.0, 0.0, 0.0]], (T, 1))], axis=1 + ) + ZarrWriter.create_and_write( + tmp_path / "ep.zarr", + numeric_data={ + "left.obs_ee_pose": pose, + "right.obs_ee_pose": pose, + "obs_head_pose": pose, + }, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + fps=fps, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + key_map = Human.get_keymap(keymap_mode="cartesian", **keymap_kwargs) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return next(iter(ds.datasets.values())) + + +@pytest.mark.parametrize("fps,gap_frames", [(30, 3), (10, 1), (60, 6)]) +def test_past_frame_is_gap_seconds_back_in_the_episodes_fps(tmp_path, fps, gap_frames): + leaf = _leaf(tmp_path, fps, image_history_gap_s=0.1) + idx = 8 + assert torch.equal(leaf[idx][PAST], leaf[idx - gap_frames][CAM]) + assert not torch.equal(leaf[idx][PAST], leaf[idx][CAM]) + assert leaf[idx]["fps"].item() == fps + + +def test_past_frame_clamps_to_the_episode_start(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.1) + for idx in (0, 1, 2): + assert torch.equal(leaf[idx][PAST], leaf[0][CAM]) + + +def test_a_gap_below_one_frame_still_reads_the_previous_frame(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.001) + assert torch.equal(leaf[5][PAST], leaf[4][CAM]) + + +def test_default_keymap_has_no_history_entry_but_samples_carry_fps(tmp_path): + leaf = _leaf(tmp_path, 30) + sample = leaf[3] + assert PAST not in sample and sample[CAM].shape == (3, 32, 32) + assert sample["fps"].dtype == torch.float32 + assert "proprio_history_mask" not in sample # a camera lag is not proprio history + + +def test_norm_mode_keymap_drops_the_history_camera_too(): + key_map = Human.get_keymap("cartesian", norm_mode=True, image_history_gap_s=0.1) + assert not [k for k in key_map if "img" in k] + + +@pytest.mark.parametrize("gap_s", [0.1, 0.5, 2.0]) +def test_the_past_frame_costs_one_jpeg_row_whatever_the_gap( + tmp_path, monkeypatch, gap_s +): + """The past frame shares the front camera's zarr key, so the read merges the + two windows. JPEGs are one frame per chunk: a plain union span would read + every frame between them (61 rows at 2 s and 30 fps) to use two.""" + leaf = _leaf(tmp_path, 30, image_history_gap_s=gap_s) + rows: list[dict] = [] + inner = ZarrEpisode.read_intervals + + def spy(self, intervals): + out = inner(self, intervals) + rows.append({k: sum(len(r) for _, r in g) for k, g in out.items()}) + return out + + monkeypatch.setattr(ZarrEpisode, "read_intervals", spy) + idx = 80 + sample = leaf[idx] + assert rows[-1]["images.front_1"] == 2 + assert torch.equal(sample[PAST], leaf[idx - round(gap_s * 30)][CAM]) diff --git a/tests/unit/test_image_memory.py b/tests/unit/test_image_memory.py new file mode 100644 index 0000000000..b899dda181 --- /dev/null +++ b/tests/unit/test_image_memory.py @@ -0,0 +1,126 @@ +"""Long-range image memory in the data path: ``Human.get_keymap(image_memory=N, +image_memory_stride_s=s)`` adds a ``_mem`` window of N frames s +seconds apart ending at the current frame, with a ``_mem_mask`` of the steps inside +the episode; ``history_stride_s`` spaces proprio history in seconds too.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 150 +CAM = Human.VIZ_IMAGE_KEY +MEM = f"{CAM}_mem" +MASK = f"{MEM}_mask" + + +def _leaf(tmp_path, fps: int = 30, **keymap_kwargs): + rng = np.random.default_rng(0) + t = np.arange(T, dtype=np.float64)[:, None] + pose = np.concatenate( + [t * 0.001 + np.zeros((T, 3)), np.tile([[1.0, 0.0, 0.0, 0.0]], (T, 1))], axis=1 + ) + ZarrWriter.create_and_write( + tmp_path / "ep.zarr", + numeric_data={ + "left.obs_ee_pose": pose, + "right.obs_ee_pose": pose, + "obs_head_pose": pose, + }, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + fps=fps, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + key_map = Human.get_keymap(keymap_mode="cartesian", **keymap_kwargs) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return next(iter(ds.datasets.values())) + + +@pytest.mark.parametrize("fps", [10, 30]) +def test_memory_frames_are_stride_seconds_apart_ending_at_the_current_frame( + tmp_path, fps +): + leaf = _leaf(tmp_path, fps, image_memory=4, image_memory_stride_s=0.5) + step = round(0.5 * fps) + idx = T - 1 + sample = leaf[idx] + assert sample[MEM].shape == (4, 3, 32, 32) + for j, frame in enumerate(sample[MEM]): + assert torch.equal(frame, leaf[idx - step * (3 - j)][CAM]) + np.testing.assert_array_equal(np.asarray(sample[MASK]), np.ones(4)) + + +def test_memory_front_pads_with_the_oldest_real_step_and_masks_the_rest(tmp_path): + leaf = _leaf(tmp_path, image_memory=4, image_memory_stride_s=1.0) + sample = leaf[35] # steps 5, 35 are real; -25, -55 are not + np.testing.assert_array_equal(np.asarray(sample[MASK]), [0, 0, 1, 1]) + for j, src in enumerate((5, 5, 5, 35)): + assert torch.equal(sample[MEM][j], leaf[src][CAM]) + + +def test_memory_at_the_first_frame_holds_only_the_current_step(tmp_path): + leaf = _leaf(tmp_path, image_memory=3, image_memory_stride_s=1.0) + sample = leaf[0] + np.testing.assert_array_equal(np.asarray(sample[MASK]), [0, 0, 1]) + assert sample[MEM].shape == (3, 3, 32, 32) + + +def test_a_one_frame_memory_keeps_its_window_axis(tmp_path): + sample = _leaf(tmp_path, image_memory=1)[40] + assert sample[MEM].shape == (1, 3, 32, 32) + assert sample[MASK].shape == (1,) + + +def test_memory_reads_and_decodes_each_distinct_frame_once(tmp_path, monkeypatch): + leaf = _leaf( + tmp_path, image_memory=10, image_memory_stride_s=1.0, image_history_gap_s=0.1 + ) + rows: list[dict] = [] + inner = ZarrEpisode.read_intervals + + def spy(self, intervals): + out = inner(self, intervals) + rows.append({k: sum(len(r) for _, r in g) for k, g in out.items()}) + return out + + monkeypatch.setattr(ZarrEpisode, "read_intervals", spy) + leaf[T - 1] + # current (also the newest memory step) + 0.1 s pair frame + 119 - 30k + assert rows[-1]["images.front_1"] == 1 + 1 + 4 + + +def test_default_keymap_has_no_memory(tmp_path): + sample = _leaf(tmp_path)[3] + assert MEM not in sample and MASK not in sample + + +def test_norm_mode_keymap_drops_the_memory_camera(): + key_map = Human.get_keymap("cartesian", norm_mode=True, image_memory=10) + assert not [k for k in key_map if "img" in k] + + +@pytest.mark.parametrize("fps", [10, 30]) +def test_proprio_history_stride_in_seconds_follows_the_episode_fps(tmp_path, fps): + leaf = _leaf(tmp_path, fps, proprio_history=3, history_stride_s=0.5) + step = round(0.5 * fps) + idx = T - 1 + got = np.asarray(leaf[idx]["right.obs_ee_pose"]) + want = np.stack( + [np.asarray(leaf[idx - step * j]["right.obs_ee_pose"])[-1] for j in (2, 1, 0)] + ) + np.testing.assert_array_equal(got, want) + np.testing.assert_array_equal( + np.asarray(leaf[step]["proprio_history_mask"]), [0, 1, 1] + ) diff --git a/tests/unit/test_index_map.py b/tests/unit/test_index_map.py new file mode 100644 index 0000000000..013ef5bb48 --- /dev/null +++ b/tests/unit/test_index_map.py @@ -0,0 +1,16 @@ +import pytest + +from egomimic.rldb.zarr.zarr_dataset_multi import _IndexMap + + +def test_index_map_matches_per_sample_list(): + datasets = {"a": [0] * 3, "empty": [], "b": [0] * 2} + im = _IndexMap(datasets) + expected = [("a", 0), ("a", 1), ("a", 2), ("b", 0), ("b", 1)] + assert len(im) == 5 + assert [im[i] for i in range(5)] == expected + assert list(im) == expected + assert im[-1] == ("b", 1) + assert im.ranges() == {"a": range(0, 3), "empty": range(3, 3), "b": range(3, 5)} + with pytest.raises(IndexError): + im[5] diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py new file mode 100644 index 0000000000..a150f562a7 --- /dev/null +++ b/tests/unit/test_infra_port_wiring.py @@ -0,0 +1,604 @@ +"""Wiring checks for the hand-keypoint default: recipes, stems, the mecka +left-wrist fix in keypoint modes, the annotation-cutoff span filter and the +train_viz second val loader.""" + +from __future__ import annotations + +import inspect +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from lightning import LightningModule +from scipy.spatial.transform import Rotation as R + +import egomimic +from egomimic.rldb.embodiment.human import Human + +CONFIG_DIR = Path(egomimic.__file__).parent / "hydra_configs" +TRAIN_TOP_LEVEL = sorted(p.stem for p in CONFIG_DIR.glob("train_*.yaml")) +HUMAN = "human_bimanual" + +# The all-cartesian cam-frame cotrain recipe is cotrain_pi_lang plus these +# overrides (model/pi0.5_cotrain_eva_aria_6d.yaml); there is no data config +# that differs from its parent only in the transform mode. +CAM_FRAME_COTRAIN = [ + "data.train_datasets.eva_bimanual.resolver.transform_list.mode=cartesian_6d", + f"data.train_datasets.{HUMAN}.resolver.key_map.keymap_mode=cartesian", + f"data.train_datasets.{HUMAN}.resolver.key_map.include_ee_pose=false", + f"data.train_datasets.{HUMAN}.resolver.transform_list.mode=cartesian_6d", + f"data.train_datasets.{HUMAN}.resolver.transform_list.include_ee_pose=false", +] + + +def _human_dataset_nodes(cfg): + for split in ( + "train_datasets", + "valid_datasets", + "train_viz_datasets", + "unseen_op_valid_datasets", + ): + node = (cfg.data.get(split) or {}).get(HUMAN) + if node is not None and node.get("resolver") is not None: + yield split, node.resolver + + +@pytest.mark.parametrize("top", TRAIN_TOP_LEVEL) +def test_human_action_width_matches_data_mode(top, compose_resolve): + """A model declaring the 144-D human action must be fed the + keypoints_*_6d data (and vice versa); a stale ypr/cartesian mode would + only fail at the first forward pass.""" + cfg = compose_resolve(top, []) + rm = cfg.model.robomimic_model + if HUMAN not in (rm.get("domains") or []): + pytest.skip(f"{top}: no human domain") + dims = (rm.get("dims") or {}).get(HUMAN) + ac_key = rm.ac_keys[HUMAN] + for split, resolver in _human_dataset_nodes(cfg): + mode = resolver.transform_list.mode + keymap_mode = resolver.key_map.keymap_mode + if ac_key == "actions_keypoints": + assert keymap_mode == "keypoints", (top, split, keymap_mode) + assert mode.startswith("keypoints_") and mode.endswith("_6d"), ( + top, + split, + mode, + ) + if dims is not None: + assert dims.action == 144, (top, split, dims) + else: + assert keymap_mode == "cartesian", (top, split, keymap_mode) + assert mode.startswith("cartesian"), (top, split, mode) + + +@pytest.mark.parametrize( + "model,data,extra", + [ + ("hpt_bc_flow_aria", "aria", []), + ("hpt_bc_flow_mecka", "mecka", []), + ("hpt_bc_flow_scale", "scale", []), + ("hpt_bc_flow_human", "human", []), + ("pi0.5_bc_aria", "aria", []), + ("pi0.5_bc_mecka", "mecka", []), + ("pi0.5_bc_scale", "scale", []), + ("pi0.5_cotrain_eva_aria", "cotrain_pi_base", []), + ("pi0.5_bc_mecka_6d", "mecka_all_6d", []), + ("pi0.5_cotrain_eva_aria_6d", "cotrain_pi_lang", CAM_FRAME_COTRAIN), + ("hpt_bc_mecka_6d_300M", "mecka_fold_flagship_opsplit_hpt_6d", []), + ( + "hpt_bc_keypoints_wrist_300M", + "mecka_fold_freeform_opsplit_hpt_keypoints", + [], + ), + ], +) +def test_vendor_pairings_agree_on_the_human_action(model, data, extra, compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian", [f"model={model}", f"data={data}", *extra] + ) + rm = cfg.model.robomimic_model + ac_key = rm.ac_keys[HUMAN] + resolver = cfg.data.train_datasets[HUMAN].resolver + mode = resolver.transform_list.mode + if ac_key == "actions_keypoints": + assert mode.endswith("_6d") and mode.startswith("keypoints_"), (model, mode) + if rm.get("dims"): + assert rm.dims[HUMAN].action == 144 + if "pi0.5" in model: + assert rm.config.model.action_dim == 144 + assert "Keypoints" in rm.action_converters.rules.HUMAN_BIMANUAL._target_ + # the pi0.5 prompt state is the cartesian ee_pose + assert resolver.key_map.get("include_ee_pose") is True + assert resolver.transform_list.get("include_ee_pose") is True + else: + assert mode.startswith("cartesian"), (model, mode) + if "pi0.5" in model: + assert mode.endswith("_6d"), (model, mode) + assert rm.config.model.action_dim == 32 + + +def test_hpt_stems_follow_the_human_action(compose_resolve): + """HPT human stems: keypoint models expose exactly ``state_keypoints`` + (a base's ``state_ee_pose`` is nulled out, not inherited).""" + from egomimic.algo.hpt import HPTModel + + for model in ("hpt_bc_flow_mecka", "hpt_cotrain_flow_seperate_head"): + cfg = compose_resolve("train_zarr_cartesian", [f"model={model}"]) + stems = cfg.model.robomimic_model.stem_specs[HUMAN] + live = {k for k, v in stems.items() if v is not None} + assert live == {"state_keypoints"}, (model, live) + cfg = compose_resolve( + "train_zarr_cartesian", ["model=hpt_cotrain_flow_shared_head"] + ) + stems = cfg.model.robomimic_model.stem_specs[HUMAN] + assert {k for k, v in stems.items() if v is not None} == {"state_ee_pose"} + + # the model side drops null stems instead of registering None + m = HPTModel.__new__(HPTModel) + torch.nn.Module.__init__(m) + m.stem_spec, m.modalities, m.stems = {}, {}, {} + m.init_domain_stem("human_bimanual", {"state_ee_pose": None, "x": object()}) + assert m.modalities["human_bimanual"] == ["x"] + assert set(m.stems) == {"human_bimanual_x"} + + +# ------------------------------------------------------ mecka left-wrist fix +def _pose(rng): + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() + return np.concatenate([rng.uniform(-1, 1, 3), q[[3, 0, 1, 2]]]) + + +def _chunk(rng, start, n): + out = np.zeros((n, 7)) + p, r = start[:3].copy(), R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + out[t] = np.concatenate([p, r.as_quat()[[3, 0, 1, 2]]]) + return out + + +def _rz180(pose7): + """What the fixed converter would have written: the same pose with its + local axes relabelled by Rz(180 deg).""" + out = np.array(pose7, dtype=np.float64, copy=True) + q = R.from_quat(out[..., [4, 5, 6, 3]]) * R.from_euler("z", np.pi) + out[..., 3:7] = q.as_quat()[..., [3, 0, 1, 2]] + return out + + +def _apply(tl, s): + s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in s.items()} + for t in tl: + s = t.transform(s) + return s + + +def test_fix_left_wrist_convention_equals_reconverting_in_keypoint_mode(): + rng = np.random.default_rng(31) + H = Human.ACTION_HORIZON + raw = {"obs_head_pose": _pose(rng)} + for side in ("left", "right"): + wrist = _pose(rng) + raw[f"{side}.obs_wrist_pose"] = wrist + raw[f"{side}.action_wrist_pose"] = _chunk(rng, wrist, H) + kp = raw[f"{side}.action_wrist_pose"][:, None, :3] + rng.uniform( + -0.1, 0.1, (H, 21, 3) + ) + raw[f"{side}.action_keypoints"] = kp.reshape(H, 63) + raw[f"{side}.obs_keypoints"] = kp[0].reshape(63) + ee = _pose(rng) + raw[f"{side}.obs_ee_pose"] = ee + raw[f"{side}.action_ee_pose"] = _chunk(rng, ee, H) + # "reconverted" twin: the fixed converter relabels the LEFT hand's axes on + # every pose it writes (wrist_pose and ee_pose share the rotation). + fixed = dict(raw) + for k in ( + "left.obs_wrist_pose", + "left.action_wrist_pose", + "left.obs_ee_pose", + "left.action_ee_pose", + ): + fixed[k] = _rz180(raw[k]) + + kw = dict(stride=1, include_ee_pose=True, pad_proprio_gripper=True) + with_flag = _apply( + Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True, **kw + ), + raw, + ) + reconverted = _apply( + Human.get_transform_list("keypoints_wristframe_6d", **kw), fixed + ) + assert set(with_flag) == set(reconverted) + for k in with_flag: + np.testing.assert_allclose(with_flag[k], reconverted[k], atol=1e-9, err_msg=k) + # and the flag really changes the left hand's frame + unfixed = _apply(Human.get_transform_list("keypoints_wristframe_6d", **kw), raw) + assert ( + np.abs(unfixed["actions_keypoints"] - with_flag["actions_keypoints"]).max() + > 1e-3 + ) + # the same equivalence without the ee_pose side (plain HPT keypoint data) + plain = {k: v for k, v in raw.items() if "ee_pose" not in k} + plain_fixed = {k: v for k, v in fixed.items() if "ee_pose" not in k} + a = _apply( + Human.get_transform_list( + "keypoints_wristframe_6d", stride=1, fix_left_wrist_convention=True + ), + plain, + ) + b = _apply( + Human.get_transform_list("keypoints_wristframe_6d", stride=1), plain_fixed + ) + for k in a: + np.testing.assert_allclose(a[k], b[k], atol=1e-9, err_msg=k) + + +# --------------------------------------------------- annotation-cutoff filter +def test_episode_has_annotation_spans(): + from egomimic.rldb.zarr.zarr_dataset_multi import _episode_has_annotation_spans + + def ds(anns): + return SimpleNamespace(_load_annotations=lambda: anns) + + assert _episode_has_annotation_spans( + ds([{"text": "a", "start_idx": 0, "end_idx": 5}]) + ) + assert not _episode_has_annotation_spans(ds([])) + assert not _episode_has_annotation_spans(ds([{"text": "a"}])) # span-less + assert not _episode_has_annotation_spans(ds([{"start_idx": 5, "end_idx": 5}])) + + def boom(): + raise OSError("corrupt") + + assert not _episode_has_annotation_spans(SimpleNamespace(_load_annotations=boom)) + + +def test_annotation_cutoff_resolver_defaults_to_requiring_spans(): + from egomimic.rldb.zarr.zarr_dataset_multi import S3AnnotationCutoffEpisodeResolver + + r = S3AnnotationCutoffEpisodeResolver.__new__(S3AnnotationCutoffEpisodeResolver) + sig = inspect.signature(S3AnnotationCutoffEpisodeResolver.__init__) + assert sig.parameters["require_annotations"].default is True + good = SimpleNamespace(_load_annotations=lambda: [{"start_idx": 0, "end_idx": 3}]) + bad = SimpleNamespace(_load_annotations=lambda: []) + r.require_annotations = True + # bypass the S3 base resolve: patch it on the instance's class chain + import egomimic.rldb.zarr.zarr_dataset_multi as m + + orig = m.S3EpisodeResolver.resolve + m.S3EpisodeResolver.resolve = lambda self, filters=None, expected_embodiment=None: { + "g": good, + "b": bad, + } + try: + kept = r.resolve(filters=None, expected_embodiment="human_bimanual") + assert set(kept) == {"g"} + r.require_annotations = False + assert set(r.resolve()) == {"g", "b"} + m.S3EpisodeResolver.resolve = ( + lambda self, filters=None, expected_embodiment=None: {"b": bad} + ) + r.require_annotations = True + with pytest.raises(ValueError, match="no resolved episodes"): + r.resolve() + finally: + m.S3EpisodeResolver.resolve = orig + + +# ------------------------------------------------------ train_viz second loader +class _Dicts: + def __init__(self, n, tag): + self.n, self.tag = n, tag + + def __len__(self): + return self.n + + def __getitem__(self, i): + return {"x": np.float32(i), "tag": [self.tag]} + + +def test_val_dataloader_returns_two_loaders_with_train_viz(): + from lightning.pytorch.utilities.combined_loader import CombinedLoader + + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + held_out_operators=["op1"], + ) + assert isinstance(dm.val_dataloader(), CombinedLoader) + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params={ + "human_bimanual": {**params["human_bimanual"], "shuffle": False} + }, + ) + loaders = dm.val_dataloader() + assert isinstance(loaders, list) and len(loaders) == 2 + batch, _, _ = next(iter(loaders[1])) + assert batch["human_bimanual"]["tag"] == [["viz"], ["viz"]] + + +def test_validation_step_routes_by_dataloader_idx(): + from egomimic.pl_utils.pl_model import ModelWrapper + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, batch, batch_idx, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) # global_rank etc. without a trainer + w.model = SimpleNamespace(process_batch_for_training=lambda b: {"processed": b}) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = _Ev("train_viz") + # Lightning hands the outer sequential loader's inner triple as `batch` + w.validation_step(({"k": 1}, 0, 0), 0, dataloader_idx=0) + w.validation_step(({"k": 2}, 0, 1), 3, dataloader_idx=1) + w.validation_step({"k": 3}, 1, dataloader_idx=0) + assert calls == [ + ("valid", {"processed": {"k": 1}}, 0, 0), + ("train_viz", {"processed": {"k": 2}}, 3, 1), + ("valid", {"processed": {"k": 3}}, 1, 0), + ] + w.train_viz_evaluator = None + w.validation_step({"k": 4}, 0, dataloader_idx=1) # no second head: ignored + assert len(calls) == 3 + + +def test_unseen_op_valid_third_loader_and_routing(): + """unseen_op_valid_datasets adds a loader after train_viz; without train_viz it + takes idx 1, and ModelWrapper routes by val_loader_names, not position.""" + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + from egomimic.pl_utils.pl_model import ModelWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + common = dict( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + unseen_op_valid_datasets={"human_bimanual": _Dicts(4, "unseen")}, + unseen_op_valid_dataloader_params=params, + ) + dm = MultiDataModuleWrapper( + **common, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params=params, + ) + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + loaders = dm.val_dataloader() + tags = [next(iter(ld))[0]["human_bimanual"]["tag"][0][0] for ld in loaders] + assert tags == ["valid", "viz", "unseen"] + + dm = MultiDataModuleWrapper(**common) + assert dm.val_loader_names() == ["valid", "unseen_op_valid"] + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) + w.model = SimpleNamespace(process_batch_for_training=lambda b: b) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = None + w.unseen_op_valid_evaluator = _Ev("unseen_op_valid") + w.val_loader_names = dm.val_loader_names() + w.validation_step({"k": 1}, 0, dataloader_idx=0) + w.validation_step({"k": 2}, 0, dataloader_idx=1) + assert calls == [("valid", 0), ("unseen_op_valid", 1)] + + +def test_flagship_opsplit_val_heads(compose_resolve): + """Flagship opsplit: valid (prefixed seen_op_valid) = seen operators' + held-out episodes (complement of train), train_viz = the train split (no + explicit datasets), unseen_op_valid = the held-out operators. The topop + twin inherits the same layout.""" + for data, train_op in ( + ("mecka_fold_flagship_opsplit_hpt_6d", "not in"), + ("mecka_fold_flagship_topop_hpt_6d", "== '6903686e0e94ce070afd1f24'"), + ): + cfg = compose_resolve("train_zarr_mecka_flagship_6d_hpt", [f"data={data}"]) + d = cfg.data + train, valid = d.train_datasets[HUMAN], d.valid_datasets[HUMAN] + unseen = d.unseen_op_valid_datasets[HUMAN] + assert d.get("train_viz_datasets") is None + assert d.valid_prefix == "seen_op_valid" + assert (train.mode, valid.mode, unseen.mode) == ("train", "valid", "total") + assert valid.valid_ratio == train.valid_ratio + assert list(valid.filters.filter_lambdas) == list(train.filters.filter_lambdas) + assert train_op in train.filters.filter_lambdas[-1] + ops = list(d.held_out_operators) + assert unseen.filters.filter_lambdas[-1].endswith(f" in {ops}") + + +def test_bounds_quantiles_include_the_observed_extremes(): + """With fewer than 10k samples a linear 0.01 / 99.99 percentile lands + strictly inside the sample range, so the stats' own extreme frames were + rejected at train time; the bounds quantiles must snap to samples.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + X = np.random.default_rng(0).normal(size=(500, 4)).astype(np.float32) + st = MultiDataset._compute_stats_for_array(X) + np.testing.assert_array_equal(st["quantile_0_01"], X.min(axis=0)) + np.testing.assert_array_equal(st["quantile_99_99"], X.max(axis=0)) + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = {0: {"actions_keypoints": st}} + md.zarr_keys = {0: {"actions_keypoints": "actions_keypoints"}} + md._warned_violations = set() + for row in (X.min(axis=0), X.max(axis=0)): + assert ( + md._check_bounds( + {"embodiment": 0, "actions_keypoints": row[None]}, None, 0, "ep" + ) + is None + ) + + +def test_pi_loss_is_reduced_over_the_packed_width(): + from egomimic.algo.pi import PI + from egomimic.rldb.embodiment.embodiment import EMBODIMENT + from egomimic.utils.action_utils import ( + ConverterRegistry, + HumanBimanualKeypoints, + RobotBimanualCartesian6D, + ) + + pi = PI.__new__(PI) + pi.action_registry = ConverterRegistry() + pi._packed_widths = {} + eva, human = EMBODIMENT.EVA_BIMANUAL.value, EMBODIMENT.HUMAN_BIMANUAL.value + pi.action_registry.register(eva, "actions_cartesian", RobotBimanualCartesian6D()) + pi.action_registry.register(human, "actions_keypoints", HumanBimanualKeypoints()) + losses = torch.zeros(2, 5, 144) + losses[..., 32:] = 1.0 # error only in the zero-padded slots + eva_loss = pi._reduce_loss(losses, torch.zeros(2, 5, 20), eva, "actions_cartesian") + assert eva_loss.item() == 0.0 + kp_loss = pi._reduce_loss( + losses, torch.zeros(2, 5, 144), human, "actions_keypoints" + ) + assert abs(kp_loss.item() - 112 / 144) < 1e-6 + # already-reduced losses pass through unchanged + assert ( + pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() + == 2.0 + ) + + +def test_pi_process_batch_keeps_the_per_sample_intrinsics(): + from egomimic.algo.pi import PI + from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + class _NormStats: + def zarr_key_to_keyname(self, key, embodiment_id): + return key if key == "actions_keypoints" else None + + pi = PI.__new__(PI) + pi.norm_stats = _NormStats() + pi.device = "cpu" + pi._build_prompts = lambda _batch, name, B: [""] * B + pi._tokenize_prompts = lambda prompts: {} + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + pi.ac_keys = {emb_id: "actions_keypoints"} + K = torch.eye(3).expand(2, 3, 3).clone() + out = PI.process_batch_for_training( + pi, + {emb: {"actions_keypoints": torch.zeros(2, 5, 144), "intrinsics": K}}, + )[emb_id] + torch.testing.assert_close(out["intrinsics"], K) + + +def test_bounds_check_has_relative_slack_but_catches_corrupt_values(): + """Per-cell bounds tolerate frames moderately beyond the stats sample's + range (valid extreme motion) and still reject values orders of magnitude + off (fill constants, wrong-frame data).""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + lo, hi = np.full(18, -1.0, np.float32), np.full(18, 1.0, np.float32) + md.norm_stats = {0: {"actions_cartesian": {"quantile_1": lo, "quantile_99": hi}}} + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + + def check(v): + arr = np.zeros((3, 18), np.float32) + arr[1, 0] = v # an xyz channel + return md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "e" + ) + + assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range + assert check(1.0 + 0.6 * 2.0) is not None + assert check(1e9) is not None + + +# ------------------------------------------------------ train_viz on by default +def _viz_cfg(**top): + from omegaconf import OmegaConf + + data = top.pop("data", {}) + cfg = { + "mode": "train", + "evaluator": { + "_target_": "egomimic.eval.eval_hpt.HPTEvalVideo", + "viz_every_n_epochs": 7, + }, + "data": { + "valid_dataloader_params": {HUMAN: {"batch_size": 4, "num_workers": 3}}, + **data, + }, + **top, + } + return OmegaConf.create(cfg) + + +def test_train_viz_defaults_to_the_train_datasets(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + viz, params = th._train_viz_datasets(_viz_cfg(), train, instantiate=None) + assert viz[HUMAN] is train[HUMAN], "reuses the train split, no second resolve" + # valid loader params, never shuffled (same leading slice every val) + assert params == {HUMAN: {"batch_size": 4, "num_workers": 3, "shuffle": False}} + + +def test_train_viz_explicit_datasets_and_opt_outs(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + explicit = _viz_cfg( + data={ + "train_viz_datasets": {HUMAN: {"tag": "explicit"}}, + "train_viz_dataloader_params": {HUMAN: {"batch_size": 32}}, + } + ) + viz, params = th._train_viz_datasets( + explicit, train, instantiate=lambda node, dataset_name: ("built", node.tag) + ) + assert viz == {HUMAN: ("built", "explicit")} + assert params is None, "the data config's own params are left to Hydra" + + for off in ( + _viz_cfg(train_viz=False), + _viz_cfg(mode="eval"), + _viz_cfg(evaluator=None), + ): + assert th._train_viz_datasets(off, train, instantiate=None) == ({}, None) + + +def test_train_viz_evaluator_wraps_the_canonical_evaluator(): + import egomimic.trainHydra as th + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + cfg = _viz_cfg() + ev = th._build_train_viz_evaluator(cfg) + assert isinstance(ev, TrainVizEvalVideo) + assert ev.viz_every_n_epochs == 7 + assert th._build_train_viz_evaluator(_viz_cfg(train_viz=False)) is None diff --git a/tests/unit/test_keypoints_6d.py b/tests/unit/test_keypoints_6d.py new file mode 100644 index 0000000000..103792683a --- /dev/null +++ b/tests/unit/test_keypoints_6d.py @@ -0,0 +1,276 @@ +"""Independent end-to-end checks of the 144-D wrist-frame hand-keypoint +pipeline (the default human action): + + raw world-frame keypoints + wrist / head poses + -> Human.get_transform_list("keypoints_wristframe_6d", include_ee_pose=True) + -> actions_keypoints (T, 144), observations.state.keypoints (144,), + observations.state.ee_pose (18 / 20,) + -> _build_human_keypoints_revert_6d_wristframe_transform_list (evaluator) + -> 126-D head-frame keypoints == independent inv(T_head) @ kp_world + +Ground truth is plain numpy/scipy SE(3) math sharing no code with the +pipeline. Interpolation endpoints are exact, so the chunk is compared at +t = 0 and t = -1 against raw frames 0 and horizon-1. +""" + +import numpy as np +import pytest +import torch +from scipy.spatial.transform import Rotation as R + +from egomimic.rldb.embodiment.embodiment import Embodiment +from egomimic.rldb.embodiment.human import ( + Human, + _build_human_keypoints_revert_6d_transform_list, + _build_human_keypoints_revert_6d_wristframe_transform_list, +) +from egomimic.rldb.transforms import ( + KeypointsRot6DToYPR, + KeypointsYPRToRot6D, +) +from egomimic.utils.pose_utils import bimanual_keypoint_layout + +HORIZON = Human.ACTION_HORIZON # raw frames per chunk (30) +CHUNK = 100 + + +def _rng(seed): + return np.random.default_rng(seed) + + +def _rand_pose(rng, scale=1.0): + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() # xyzw + return np.concatenate([rng.uniform(-scale, scale, 3), q[[3, 0, 1, 2]]]) + + +def _rand_chunk(rng, start, n): + out = np.zeros((n, 7)) + p = start[:3].copy() + r = R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t > 0: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + q = r.as_quat() + out[t] = np.concatenate([p, q[[3, 0, 1, 2]]]) + return out + + +def _T(p7): + M = np.eye(4) + M[:3, :3] = R.from_quat(p7[[4, 5, 6, 3]]).as_matrix() + M[:3, 3] = p7[:3] + return M + + +def _in_frame(T_frame, pts): + """Points (..., 3) world -> frame.""" + inv = np.linalg.inv(T_frame) + return pts @ inv[:3, :3].T + inv[:3, 3] + + +def _apply(transform_list, sample): + s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in sample.items()} + for t in transform_list: + s = t.transform(s) + return s + + +def _raw_sample(rng): + head = _rand_pose(rng) + raw = {"obs_head_pose": head} + world = {} + for side in ("left", "right"): + wrist_obs = _rand_pose(rng) + wrist_act = _rand_chunk(rng, wrist_obs, HORIZON) # (30, 7) world + offsets = rng.uniform(-0.1, 0.1, size=(HORIZON, 21, 3)) + kp_act = wrist_act[:, None, :3] + offsets # (30, 21, 3) world + kp_obs = kp_act[0] + ee_obs = _rand_pose(rng) + raw[f"{side}.obs_wrist_pose"] = wrist_obs + raw[f"{side}.action_wrist_pose"] = wrist_act + raw[f"{side}.obs_keypoints"] = kp_obs.reshape(63) + raw[f"{side}.action_keypoints"] = kp_act.reshape(HORIZON, 63) + raw[f"{side}.obs_ee_pose"] = ee_obs + raw[f"{side}.action_ee_pose"] = _rand_chunk(rng, ee_obs, HORIZON) + world[side] = (wrist_obs, wrist_act, kp_obs, kp_act) + return raw, head, world + + +def test_keypoint_transform_round_trips_138_144(): + rng = _rng(0) + x = np.concatenate( + [ + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 63)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 63)), + ], + axis=-1, + ) + assert x.shape == (5, 138) + y = KeypointsYPRToRot6D(action_key="k").transform({"k": x.copy()})["k"] + assert y.shape == (5, 144) + layout = bimanual_keypoint_layout(144) + # keypoint blocks pass through untouched + np.testing.assert_allclose(y[:, list(layout["keypoints"])][:, :63], x[:, 6:69]) + back = KeypointsRot6DToYPR(action_key="k").transform({"k": y})["k"] + np.testing.assert_allclose(back, x, atol=1e-9) + # a single proprio vector works the same way + single = KeypointsYPRToRot6D(action_key="k").transform({"k": x[0].copy()})["k"] + assert single.shape == (144,) + t = KeypointsYPRToRot6D(action_key="k").transform({"k": torch.from_numpy(x)})["k"] + assert isinstance(t, torch.Tensor) and t.shape == (5, 144) + + +@pytest.mark.parametrize("pad", [False, True]) +def test_keypoints_wristframe_6d_pipeline_matches_independent_math(pad): + rng = _rng(3 + int(pad)) + raw, head, world = _raw_sample(rng) + fwd = Human.get_transform_list( + "keypoints_wristframe_6d", + stride=1, + include_ee_pose=True, + pad_proprio_gripper=pad, + ) + out = _apply(fwd, raw) + assert set(out) == { + "actions_keypoints", + "observations.state.keypoints", + "observations.state.ee_pose", + }, sorted(out) + act = np.asarray(out["actions_keypoints"]) + obs = np.asarray(out["observations.state.keypoints"]) + ee = np.asarray(out["observations.state.ee_pose"]) + assert act.shape == (CHUNK, 144) and obs.shape == (144,) + assert ee.shape == ((20,) if pad else (18,)) + if pad: + assert ee[9] == 0.0 and ee[19] == 0.0 + + layout = bimanual_keypoint_layout(144) + per_hand = layout["per_hand"] + Th = _T(head) + for si, side in enumerate(("left", "right")): + wrist_obs, wrist_act, kp_obs, kp_act = world[side] + Tw = _T(wrist_obs) + o = si * per_hand + # proprio: wrist pose in HEAD frame (xyz + rot6d), keypoints in the + # wrist's own frame + Twh = np.linalg.inv(Th) @ Tw + np.testing.assert_allclose(obs[o : o + 3], Twh[:3, 3], atol=1e-9) + np.testing.assert_allclose(obs[o + 3 : o + 6], Twh[:3, 0], atol=1e-9) + np.testing.assert_allclose(obs[o + 6 : o + 9], Twh[:3, 1], atol=1e-9) + np.testing.assert_allclose( + obs[o + 9 : o + per_hand].reshape(21, 3), _in_frame(Tw, kp_obs), atol=1e-9 + ) + # action: wrist pose relative to the obs wrist (identity at t=0), + # keypoints in the obs-wrist frame; endpoints are exact under + # interpolation. + for t_chunk, t_raw in ((0, 0), (CHUNK - 1, HORIZON - 1)): + Ta = _T(wrist_act[t_raw]) + Trel = np.linalg.inv(Tw) @ Ta + row = act[t_chunk] + np.testing.assert_allclose(row[o : o + 3], Trel[:3, 3], atol=1e-6) + np.testing.assert_allclose(row[o + 3 : o + 6], Trel[:3, 0], atol=1e-6) + np.testing.assert_allclose(row[o + 6 : o + 9], Trel[:3, 1], atol=1e-6) + np.testing.assert_allclose( + row[o + 9 : o + per_hand].reshape(21, 3), + _in_frame(Tw, kp_act[t_raw]), + atol=1e-6, + ) + np.testing.assert_allclose(act[0, list(layout["wrist_xyz"])], 0.0, atol=1e-12) + + # evaluator revert: back to 126-D head-frame keypoints (batched) + rev = _build_human_keypoints_revert_6d_wristframe_transform_list() + reverted = Embodiment.apply_transform( + { + "actions_keypoints": torch.from_numpy(act[None]).float(), + "observations.state.keypoints": torch.from_numpy(obs[None]).float(), + }, + rev, + ) + kp_head = np.asarray(reverted["actions_keypoints"]) + assert kp_head.shape == (1, CHUNK, 126) + for si, side in enumerate(("left", "right")): + _, _, _, kp_act = world[side] + for t_chunk, t_raw in ((0, 0), (CHUNK - 1, HORIZON - 1)): + np.testing.assert_allclose( + kp_head[0, t_chunk, 63 * si : 63 * si + 63].reshape(21, 3), + _in_frame(Th, kp_act[t_raw]), + atol=1e-4, + ) + + +def test_keypoints_headframe_6d_pipeline_and_revert(): + rng = _rng(9) + raw, head, world = _raw_sample(rng) + for k in list(raw): + if "ee_pose" in k: + del raw[k] + fwd = Human.get_transform_list("keypoints_headframe_6d", stride=1) + out = _apply(fwd, raw) + assert set(out) == {"actions_keypoints", "observations.state.keypoints"} + act = np.asarray(out["actions_keypoints"]) + assert act.shape == (CHUNK, 144) + Th = _T(head) + layout = bimanual_keypoint_layout(144) + per_hand = layout["per_hand"] + for si, side in enumerate(("left", "right")): + _, wrist_act, _, kp_act = world[side] + o = si * per_hand + Tah = np.linalg.inv(Th) @ _T(wrist_act[0]) + np.testing.assert_allclose(act[0, o : o + 3], Tah[:3, 3], atol=1e-6) + np.testing.assert_allclose(act[0, o + 3 : o + 6], Tah[:3, 0], atol=1e-6) + np.testing.assert_allclose( + act[0, o + 9 : o + per_hand].reshape(21, 3), + _in_frame(Th, kp_act[0]), + atol=1e-6, + ) + rev = _build_human_keypoints_revert_6d_transform_list() + reverted = Embodiment.apply_transform( + { + "actions_keypoints": torch.from_numpy(act[None]).float(), + "observations.state.keypoints": torch.from_numpy( + np.asarray(out["observations.state.keypoints"])[None] + ).float(), + }, + rev, + ) + assert np.asarray(reverted["actions_keypoints"]).shape == (1, CHUNK, 138) + assert np.asarray(reverted["observations.state.keypoints"]).shape == (1, 138) + + +def test_keypoint_modes_without_ee_pose_emit_no_ee_pose(): + rng = _rng(12) + raw, _, _ = _raw_sample(rng) + for k in list(raw): + if "ee_pose" in k: + del raw[k] + for mode in ("keypoints_wristframe_6d", "keypoints_wristframe_ypr"): + out = _apply( + Human.get_transform_list(mode, stride=1, allow_legacy_rotation=True), raw + ) + assert set(out) == {"actions_keypoints", "observations.state.keypoints"} + ypr = _apply( + Human.get_transform_list( + "keypoints_wristframe_ypr", stride=1, allow_legacy_rotation=True + ), + raw, + ) + assert np.asarray(ypr["actions_keypoints"]).shape == (CHUNK, 138) + + +def test_keymap_include_ee_pose_adds_the_ee_pose_keys(): + km = Human.get_keymap("keypoints") + assert "left.obs_ee_pose" not in km + km = Human.get_keymap("keypoints", include_ee_pose=True) + for side in ("left", "right"): + assert km[f"{side}.obs_ee_pose"]["key_type"] == "proprio_keys" + assert km[f"{side}.action_ee_pose"]["key_type"] == "action_keys" + assert km[f"{side}.action_ee_pose"]["horizon"] == HORIZON + with pytest.raises(ValueError, match="include_ee_pose"): + Human.get_transform_list("cartesian_6d", include_ee_pose=True) + with pytest.raises(ValueError, match="pad_proprio_gripper"): + Human.get_transform_list("keypoints_wristframe_6d", pad_proprio_gripper=True) diff --git a/tests/unit/test_lance_annotations.py b/tests/unit/test_lance_annotations.py new file mode 100644 index 0000000000..8baa2ecab8 --- /dev/null +++ b/tests/unit/test_lance_annotations.py @@ -0,0 +1,427 @@ +import argparse +import json + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +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 +from egomimic.scripts.lance_annotations import cmd_stage, stage # noqa: E402 +from egomimic.scripts.lance_commit import REJECTED, commit_staged # 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): + return resolver( + root, key_map=KEY_MAP, cameras={}, annotations=ann or None + ).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][0]] + assert leaf[0]["annotations"] == ["folding clothes"] + + +def test_covering_annotations_in_start_order(annotated): + root, recs = annotated + leaf = _leaves(root, types=["language"])[recs[0][0]] + 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][0]] + assert other[0]["annotations"] == [] + + +def test_vendor_and_quality_filters(annotated): + root, recs = annotated + h = recs[0][0] + 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][0] + 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][0]] + # Eligible anchors are 0, 1 and 5.. (frame 4 is invalid); annotations cover [0, 9). + assert [leaves[recs[0][0]].anchor(i) for i in range(6)] == [0, 1, 5, 6, 7, 8] + assert len(leaves[recs[0][0]]) == 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) == { + "added": 1, + "deleted": 1, + } + ds = lance.dataset(uri) + assert ds.version == v + 1 + 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_from_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): + from egomimic.rldb.zarr.zarr_writer import ZarrWriter + from egomimic.scripts.lance_annotations import cmd_from_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)}, + ) + out = tmp_path / "changes.parquet" + cmd_from_zarr( + argparse.Namespace( + lance=str(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 + + +def test_from_bucket_counts_frames_from_the_alias_start(tmp_path): + from fixtures.lance_tables import alias + + from egomimic.scripts.lance_annotations import cmd_from_bucket + + root = tmp_path / "lance" + recs = write_catalog( + root, aliases=lambda recs: [alias(rid, 3, n, h) for rid, h, n in recs] + ) + (rid, h, n), (_, h2, _) = recs + bucket = tmp_path / "bucket" + bucket.mkdir() + spans = [ + {"text": "reach", "start_idx": 0, "end_idx": 2}, + {"text": "fold", "start_idx": 2, "end_idx": n - 3}, + {"text": "too long", "start_idx": 0, "end_idx": n}, + ] + (bucket / f"{h}_annotations.json").write_text(json.dumps(spans)) + (bucket / "unknown_annotations.json").write_text(json.dumps(spans)) + (bucket / f"{h2}_other.json").write_text(json.dumps(spans)) + out = tmp_path / "changes.parquet" + cmd_from_bucket( + argparse.Namespace( + lance=str(root), + bucket=str(bucket), + key="annotations", + vendor="scale", + out=out, + ) + ) + rows = pq.read_table(out).to_pylist() + assert {r["text"]: (r["t_start_ns"], r["t_end_ns"]) for r in rows} == { + "reach": (3 * PERIOD_NS, 5 * PERIOD_NS), + "fold": (5 * PERIOD_NS, n * PERIOD_NS), + } + assert {r["recording_id"] for r in rows} == {rid} + assert not review(root, pq.read_table(out)).errors + + +def test_commit_annotations_records_and_tags_the_change(annotated): + root, recs = annotated + rid = recs[0][0] + old = _rows(recs)[0] + rv = review( + root, + _changes( + [ + _change("delete", annotation_id=old["annotation_id"]), + _add(recs[1][0], 0, PERIOD_NS, text="new"), + ] + ), + ) + out = commit.commit_annotations(root, rv.adds, rv.deletes, operator="me") + assert out["tag"].startswith("annotations-") + assert out["rows"] == {"added": 1, "deleted": 1} + (row,) = lance.dataset(str(root / "changes.lance")).to_table().to_pylist() + assert (row["lab"], row["added"], row["replaced"], row["operator"]) == ( + "annotations", + [recs[1][0]], + [rid], + "me", + ) + ds = lance.dataset(str(root / "annotations.lance")) + assert ds.tags.list()[out["tag"]]["version"] == ds.version + assert commit.commit_annotations(root, rv.adds.slice(0, 0), []) is None + + +@pytest.fixture +def staged(annotated, tmp_path_factory): + root, recs = annotated + base = tmp_path_factory.mktemp("committer") + + def put(rows, name="change.parquet"): + f = base / name + pq.write_table(_changes(rows), f) + return stage(f, base / "staging", "me") + + def run(): + return commit_staged( + str(base / "staging"), "annotations", str(root), base / "work" + ) + + return root, recs, base / "staging", put, run + + +def test_staged_annotations_commit_through_the_committer(staged): + root, recs, staging, put, run = staged + rid = recs[1][0] + first = put([_add(rid, 0, PERIOD_NS, text="new")]) + assert put([_add(rid, 0, PERIOD_NS, text="new")]) == first + report = run() + (batch,) = report["batches"] + assert batch["tag"].startswith("annotations-") and batch["added"] == [rid] + (row,) = lance.dataset(str(root / "changes.lance")).to_table().to_pylist() + prov = json.loads(row["provenance"]) + assert (row["operator"], prov["approver"], prov["source"]) == ( + "me", + "me", + "change.parquet", + ) + assert commit.list_files(staging / "annotations") == {} + assert run() == {"batches": [], "rejected": {}} + + +def test_a_staged_change_that_no_longer_applies_is_rejected_until_restaged(staged): + root, recs, staging, put, run = staged + old = _rows(recs)[0] + put([_change("delete", annotation_id=old["annotation_id"])], "a.parquet") + second = put( + [ + _change("delete", annotation_id=old["annotation_id"]), + _add(recs[1][0], 0, PERIOD_NS, text="late"), + ], + "b.parquet", + ) + report = run() + assert [b["rows"] for b in report["batches"]] == [{"added": 0, "deleted": 1}] + (sid,) = report["rejected"] + assert second.endswith(sid) and "delete of" in report["rejected"][sid][0] + assert (staging / "annotations" / sid / REJECTED).exists() + assert run() == {"batches": [], "rejected": {}} + + +def test_stage_refuses_a_change_file_with_errors(annotated, tmp_path_factory): + root, recs = annotated + base = tmp_path_factory.mktemp("bad") + f = base / "bad.parquet" + pq.write_table(_changes([_add("no-such-recording", 0, PERIOD_NS)]), f) + args = argparse.Namespace( + lance=root, file=f, staging=base / "staging", approver="me" + ) + with pytest.raises(SystemExit, match="not staged"): + cmd_stage(args) + assert not (base / "staging").exists() + + +def test_a_changes_table_from_before_provenance_gains_the_column(annotated): + root, recs = annotated + old = commit.CHANGES.remove(commit.CHANGES.get_field_index("provenance")) + lance.write_dataset( + pa.Table.from_pylist( + [dict(tag="t0", lab="eva", operator="x", added=[], replaced=[])], + schema=old, + ), + str(root / "changes.lance"), + ) + rv = review(root, _changes([_add(recs[1][0], 0, PERIOD_NS, text="new")])) + commit.commit_annotations(root, rv.adds, rv.deletes, provenance={"approver": "me"}) + rows = lance.dataset(str(root / "changes.lance")).to_table().to_pylist() + assert [r["provenance"] for r in rows] == [None, '{"approver": "me"}'] diff --git a/tests/unit/test_lance_commit.py b/tests/unit/test_lance_commit.py new file mode 100644 index 0000000000..7e9aae655d --- /dev/null +++ b/tests/unit/test_lance_commit.py @@ -0,0 +1,216 @@ +import datetime + +import numpy as np +import pyarrow as pa +import pytest + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import eva_steps_table, 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 ( # noqa: E402 + TASKS, + recording_id, + task_instance_id, +) + +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 _delivery(tmp, i, n, offset=None, video=b"mp4"): + rid = recording_id("aria", str(i)) + h = f"2025-01-01-00-00-{i:02d}-000000" + src = tmp / "src" / f"{i}-{offset}" + src.mkdir(parents=True, exist_ok=True) + (src / "front_1.mp4").write_bytes(video) + return commit.Delivery( + recording=dict( + recording_id=rid, + vendor_recording_id=h, + lab="aria", + site="rl2", + embodiment="human_bimanual", + fps=30.0, + num_frames=n, + intrinsics=[], + extrinsics=[], + ), + tasks=[ + dict( + task_instance_id=f"t{i}", + recording_id=rid, + t_start_ns=0, + t_end_ns=n * 33_333_333, + task_name="fold", + task_description=f"task {i}", + objects=[], + ) + ], + alias=dict( + episode_hash=h, recording_id=rid, t_start_ns=0, t_end_ns=n * 33_333_333 + ), + steps=steps_table(rid, n, 100 * i if offset is None else offset), + videos={"front_1": src / "front_1.mp4"}, + ) + + +def _resolve(root, **kw): + return LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}, **kw).resolve() + + +def _head(leaf): + return float(leaf[0]["obs_head_pose"][0]) + + +def test_append_skips_committed(tmp_path): + root = tmp_path / "lance" + a, b = _delivery(tmp_path, 0, 10), _delivery(tmp_path, 1, 8) + assert commit.append(root, [a]) == [a.recording_id] + assert commit.append(root, [a, b]) == [b.recording_id] + assert commit.append(root, [a, b]) == [] + assert lance.dataset(root / "steps" / "aria.lance").count_rows() == 18 + assert len(lance.dataset(root / "steps" / "aria.lance").get_fragments()) == 2 + + +def test_replace_is_one_commit_and_keeps_old_versions_readable(tmp_path): + root = tmp_path / "lance" + ds = [_delivery(tmp_path, i, 10) for i in range(3)] + commit.append(root, ds) + commit.tag(root, "t1") + before = _resolve(root) + steps_v = lance.dataset(root / "steps" / "aria.lance").version + + new = _delivery(tmp_path, 0, 12, offset=1000, video=b"mp4-v2") + commit.replace(root, new) + assert lance.dataset(root / "steps" / "aria.lance").version == steps_v + 1 + + after = _resolve(root) + h0 = new.alias["recording_id"] + assert _head(after[h0]) == 1000 and len(after[h0]) == 12 - 3 + assert ( + after[h0].row_start == 20 + ) # moved to the end, contiguous and in timestamp order + for h in before: + if h != h0: + assert _head(after[h]) == _head(before[h]) + assert after[h].cache_fingerprint() == before[h].cache_fingerprint() + assert after[h0].cache_fingerprint() != before[h0].cache_fingerprint() + + old = _resolve(root, tag="t1") + assert _head(old[h0]) == 0 and old[h0].steps_version == steps_v + assert old[h0].content_digest == before[h0].content_digest + + +def test_edit_add_columns_tag_cleanup(tmp_path): + root = tmp_path / "lance" + ds = [_delivery(tmp_path, i, 10) for i in range(2)] + commit.append(root, ds) + commit.tag(root, "t1") + commit.replace(root, _delivery(tmp_path, 1, 10, offset=500, video=b"mp4-v2")) + + commit.edit( + root, + "tasks.lance", + pa.table({"task_instance_id": ["t0"], "task_description": ["fixed"]}), + on="task_instance_id", + ) + with pytest.raises(KeyError): + commit.edit( + root, + "tasks.lance", + pa.table({"task_instance_id": ["nope"], "task_description": ["x"]}), + on="task_instance_id", + ) + with pytest.raises(ValueError): + commit.edit( + root, + "steps/aria.lance", + pa.table({"recording_id": ["x"]}), + on="recording_id", + ) + leaves = _resolve(root) + h0, h1 = (d.alias["recording_id"] for d in ds) + assert leaves[h0].task_text == "fixed" and leaves[h1].task_text == "task 1" + + fp = {h: leaf.cache_fingerprint() for h, leaf in leaves.items()} + commit.add_columns( + root, "steps/aria.lance", {"both_valid": "left_valid AND right_valid"} + ) + commit.index(root) + commit.compact(root) + steps = lance.dataset(root / "steps" / "aria.lance") + assert len(steps.get_fragments()) == 1 + assert np.asarray(steps.to_table(columns=["both_valid"]).column(0)).sum() == 18 + assert {h: leaf.cache_fingerprint() for h, leaf in _resolve(root).items()} == fp + + videos = sorted(p.relative_to(root) for p in (root / "videos").rglob("*.mp4")) + assert len(videos) == 3 + commit.cleanup(root, older_than=datetime.timedelta(0)) + assert ( + sorted(p.relative_to(root) for p in (root / "videos").rglob("*.mp4")) == videos + ) # t1 still needs the old one + assert _head(_resolve(root, tag="t1")[h1]) == 100 + assert _head(_resolve(root)[h1]) == 500 + + lance.dataset(root / "recordings.lance").tags.delete("t1") + commit.cleanup(root, older_than=datetime.timedelta(0)) + assert len(list((root / "videos").rglob("*.mp4"))) == 2 + + +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) + + +def test_tag_covers_every_labs_steps_table(tmp_path): + root = tmp_path / "lance" + commit.append(root, [_delivery(tmp_path, 0, 10)]) + rid = recording_id("eva", "0") + lance.write_dataset(eva_steps_table(rid, 5, 0), root / "steps" / "eva.lance") + tagged = commit.tag(root, "t1") + assert {"steps/aria.lance", "steps/eva.lance"} <= set(tagged) + assert "t1" in lance.dataset(root / "steps" / "eva.lance").tags.list() diff --git a/tests/unit/test_lance_committer.py b/tests/unit/test_lance_committer.py new file mode 100644 index 0000000000..6355b4e43b --- /dev/null +++ b/tests/unit/test_lance_committer.py @@ -0,0 +1,112 @@ +import json + +import pytest + +lance = pytest.importorskip("lance") +pytest.importorskip("av") + +from fixtures.lance_tables import EVA_DELIVERY_FRAMES, eva_delivery # 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 # noqa: E402 +from egomimic.scripts.lance_commit import REJECTED, commit_staged # noqa: E402 + +N = EVA_DELIVERY_FRAMES + + +@pytest.fixture +def env(tmp_path): + staging, root = tmp_path / "staging", tmp_path / "lance" + built = iter(range(100)) + + def stage(vid, **kw): + src = eva_delivery(tmp_path / f"built{next(built)}", vid, **kw) + commit.stage(src, commit.join(staging, "eva", vid)) + + def run(**kw): + return commit_staged(str(staging), "eva", str(root), tmp_path / "work", **kw) + + return staging, root, stage, run + + +def _changes(root): + return lance.dataset(root / "changes.lance").to_table().to_pylist() + + +def test_staged_deliveries_commit_as_one_tagged_batch(env): + staging, root, stage, run = env + stage("100") + stage("200") + report = run() + (batch,) = report["batches"] + assert sorted(batch["added"]) == sorted( + recording_id("eva", v) for v in ("100", "200") + ) + assert batch["tag"].startswith("eva-") + assert {"steps/eva.lance", "recordings.lance", "changes.lance"} <= set( + batch["versions"] + ) + (row,) = _changes(root) + assert row["tag"] == batch["tag"] and sorted(row["added"]) == sorted(batch["added"]) + assert commit.list_files(staging / "eva") == {} + leaves = LanceEpisodeResolver( + root, lab="eva", key_map={}, cameras={}, tag=batch["tag"] + ).resolve() + assert len(leaves) == 2 + assert run() == {"batches": [], "rejected": {}} + + +def test_a_bad_delivery_is_skipped_until_its_files_change(env): + staging, root, stage, run = env + stage("100") + stage("200", frames_in_video=N - 2) + report = run() + assert [b["added"] for b in report["batches"]] == [[recording_id("eva", "100")]] + assert "front_1.mp4 has 4 frames" in report["rejected"]["200"][0] + note = json.loads((staging / "eva" / "200" / REJECTED).read_text()) + assert note["errors"] == report["rejected"]["200"] + + assert run() == {"batches": [], "rejected": {}} + assert len(_changes(root)) == 1 + + commit.remove([str(staging / "eva" / "200" / "front_1.mp4")]) + stage("200") + report = run() + assert [b["added"] for b in report["batches"]] == [[recording_id("eva", "200")]] + assert not (staging / "eva" / "200").exists() or not any( + (staging / "eva" / "200").iterdir() + ) + + +def test_restaging_a_committed_recording_needs_replace(env, tmp_path): + staging, root, stage, run = env + stage("100") + run() + commit.stage( + eva_delivery(tmp_path / "v2", "100", offset=50), + commit.join(staging, "eva", "100"), + ) + report = run() + assert ( + report["batches"] == [] and "already committed" in report["rejected"]["100"][0] + ) + report = run(replace=True) + assert [b["replaced"] for b in report["batches"]] == [[recording_id("eva", "100")]] + rows = lance.dataset(root / "steps" / "eva.lance").to_table().to_pylist() + assert len(rows) == N and rows[0]["left"]["obs_ee_pose"][0] == 50.5 + + +def test_delivery_provenance_is_kept_in_changes(env, tmp_path): + staging, root, stage, run = env + src = eva_delivery(tmp_path / "prov", "100") + meta = json.loads((src / "meta.json").read_text()) + meta["provenance"] = {"converter": "abc123"} + (src / "meta.json").write_text(json.dumps(meta)) + commit.stage(src, commit.join(staging, "eva", "100")) + stage("200") + run() + (row,) = _changes(root) + assert json.loads(row["provenance"]) == { + recording_id("eva", "100"): {"converter": "abc123"} + } diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py new file mode 100644 index 0000000000..27368f2f92 --- /dev/null +++ b/tests/unit/test_lance_dataset.py @@ -0,0 +1,492 @@ +import numpy as np +import pyarrow as pa +import pytest + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import alias, task, write_catalog # noqa: E402 + +from egomimic.rldb.filters import DatasetFilter, PinError # noqa: E402 +from egomimic.rldb.lance.dataset import ( # noqa: E402 + LanceEpisodeResolver, + LanceRecordingDataset, + LanceStreamingDataset, + LanceStreamingResolver, + _runs, + eligible_intervals, +) +from egomimic.rldb.lance.schema import task_instance_id # noqa: E402 + + +def test_eligible_intervals_requires_whole_chunk_valid(): + n = 10 + valid = { + "head": np.ones(n, bool), + "left": np.ones(n, bool), + "right": np.ones(n, bool), + } + valid["left"][5] = False + iv = eligible_intervals(valid, ("left", "right"), horizon=3) + # anchors 3, 4, 5 have frame 5 in [t, t+3); the tail chunk is clamped at n. + np.testing.assert_array_equal(iv, [[0, 3], [6, 10]]) + + +def test_eligible_intervals_ignores_unused_hand_and_needs_head(): + n = 6 + valid = { + "head": np.ones(n, bool), + "left": np.zeros(n, bool), + "right": np.ones(n, bool), + } + valid["head"][2] = False + iv = eligible_intervals(valid, ("right",), horizon=2) + np.testing.assert_array_equal(iv, [[0, 2], [3, 6]]) + + +def test_runs_merge_across_chunks_and_keep_repeats(): + ids = pa.chunked_array([["a", "a", "b"], ["b", "c"], [], ["c", "a"]]) + values, starts = _runs(ids) + assert values == ["a", "b", "c", "a"] + np.testing.assert_array_equal(starts, [0, 2, 4, 6]) + values, starts = _runs(pa.chunked_array([], pa.string())) + assert values == [] and len(starts) == 0 + + +@pytest.fixture +def catalog(tmp_path): + return tmp_path, write_catalog(tmp_path) + + +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"}, + "obs_head_pose": {"key_type": "proprio_keys", "zarr_key": "obs_head_pose"}, + "annotations": {"key_type": "annotation_keys", "zarr_key": "annotations"}, +} + + +def test_resolver_reads_windows_from_the_right_rows(catalog): + root, recs = catalog + leaves = LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}).resolve() + assert set(leaves) == {rid for rid, _, _ in recs} + for i, (rid, _, n) in enumerate(recs): + leaf = leaves[rid] + # left_valid is false at frame 4, so anchors 2..4 are dropped. + assert len(leaf) == n - 3 + assert [leaf.anchor(j) for j in range(3)] == [0, 1, 5] + s = leaf[2] + offset = 100 * i + np.testing.assert_allclose( + s["left.action_ee_pose"][:, 0].numpy(), offset + 0.5 + np.arange(5, 8) + ) + assert float(s["right.obs_ee_pose"][0]) == offset + 5.25 + assert float(s["obs_head_pose"][0]) == offset + 5 + assert s["annotations"] == ["folding clothes"] + assert s["action_pad_mask"].tolist() == [1, 1, 1] + tail = leaf[len(leaf) - 1] + assert tail["action_pad_mask"].tolist() == [1, 0, 0] + + +def test_resolver_applies_filters(catalog): + root, recs = catalog + pin = recs[0][1] + leaves = LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}).resolve( + filters=DatasetFilter( + filter_lambdas=["lambda row: row['lab'] == 'rl2'"], episode_hashes=[pin] + ) + ) + assert list(leaves) == [recs[0][0]] + + +def _same_sample(a, b): + assert a.keys() == b.keys() + for k in a: + if hasattr(a[k], "shape"): + assert a[k].shape == b[k].shape and bool((a[k] == b[k]).all()), k + else: + assert a[k] == b[k], k + + +def test_streaming_getitems_batches_reads_and_local_does_not(catalog, monkeypatch): + from egomimic.rldb.lance import dataset as lance_dataset + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + root, _ = catalog + local = MultiDataset( + datasets=LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}).resolve(), + mode="total", + ) + leaves = LanceStreamingResolver(root, key_map=KEY_MAP, cameras={}).resolve() + ds = MultiDataset(datasets=leaves, mode="total") + idxs = [0, len(ds) - 1, 3, 3, len(ds) // 2, 1] + want = [local[i] for i in idxs] + + takes = [] + real_fetch = lance_dataset._fetch_numeric + monkeypatch.setattr( + lance_dataset, + "_fetch_numeric", + lambda items: takes.append(len(items)) or real_fetch(items), + ) + for a, b in zip(ds.__getitems__(idxs), want, strict=True): + _same_sample(a, b) + assert takes == [len(idxs)] # one fetch for the whole batch, both recordings + + takes.clear() + for a, b in zip(local.__getitems__(idxs), want, strict=True): + _same_sample(a, b) + assert takes == [1] * len(idxs) + + def broken(items): + raise OSError("object store unavailable") + + monkeypatch.setattr(LanceStreamingDataset, "fetch_planned", staticmethod(broken)) + # Only the batched read fails; the per-sample fallback still reads the store. + monkeypatch.setattr( + LanceStreamingDataset, "__getitem__", LanceRecordingDataset.__getitem__ + ) + for a, b in zip(ds.__getitems__(idxs), want, strict=True): + _same_sample(a, b) + + +@pytest.mark.parametrize("overlap_reads", [True, False]) +def test_streaming_getitems_split_into_parts_matches( + catalog, monkeypatch, overlap_reads +): + from egomimic.rldb.lance import dataset as lance_dataset + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + root, _ = catalog + local = MultiDataset( + datasets=LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}).resolve(), + mode="total", + ) + ds = MultiDataset( + datasets=LanceStreamingResolver( + root, + fetch_split=2, + overlap_reads=overlap_reads, + key_map=KEY_MAP, + cameras={}, + ).resolve(), + mode="total", + ) + idxs = [0, len(ds) - 1, 3, 3, len(ds) // 2] + want = [local[i] for i in idxs] + takes = [] + real_fetch = lance_dataset._fetch_numeric + monkeypatch.setattr( + lance_dataset, + "_fetch_numeric", + lambda items: takes.append(len(items)) or real_fetch(items), + ) + for a, b in zip(ds.__getitems__(idxs), want, strict=True): + _same_sample(a, b) + assert sorted(takes) == [1, 2, 2] + + +def _all_valid(rid, n): + return {"left_valid": np.ones(n, bool)} + + +def _two_tasks(recs): + """The 12-frame recording: reach [0, 4), a gap, fold [6, 12); the 8-frame one: one + task, and two old clips aliased to its halves.""" + (a, _, _), (b, _, _) = _by_length(recs) + return [ + task(a, 0, 4, "reach", "reach"), + task(a, 6, 12, "fold", "fold"), + task(b, 0, 8), + ] + + +def _clips(recs): + (a, ha, _), (b, _, _) = _by_length(recs) + return [alias(a, 0, 12, ha), alias(b, 0, 4, "clip_b1"), alias(b, 4, 8, "clip_b2")] + + +def _by_length(recs): + return sorted(recs, key=lambda r: -r[2]) + + +@pytest.fixture +def segmented(tmp_path): + recs = write_catalog(tmp_path, tasks=_two_tasks, aliases=_clips, valid=_all_valid) + (a, _, _), (b, _, _) = _by_length(recs) + return tmp_path, a, b + + +def _frames(leaf): + return [leaf.anchor(i) for i in range(len(leaf))] + + +def _resolve(root, filters=None, **kw): + return LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}, **kw).resolve( + filters=filters + ) + + +def test_one_leaf_per_recording_sampling_only_its_tasks(segmented): + root, a, b = segmented + leaves = _resolve(root) + assert set(leaves) == {a, b} + assert _frames(leaves[a]) == [0, 1, 2, 3, *range(6, 12)] + assert _frames(leaves[b]) == list(range(8)) + assert leaves[a][1]["annotations"] == ["reach"] + assert leaves[a][4]["annotations"] == ["fold"] + + +def test_taskless_frames_only_when_asked(segmented): + root, a, _ = segmented + leaf = _resolve(root, include_taskless_frames=True)[a] + assert _frames(leaf) == list(range(12)) + assert leaf[4]["annotations"] == [] + gaps_only = _resolve( + root, + DatasetFilter(filter_lambdas=["lambda row: row['task'] is None"]), + include_taskless_frames=True, + ) + assert list(gaps_only) == [a] and _frames(gaps_only[a]) == [4, 5] + + +def test_filter_lambdas_run_per_task(segmented): + root, a, _ = segmented + leaves = _resolve( + root, DatasetFilter(filter_lambdas=["lambda row: row['task'] == 'fold'"]) + ) + assert list(leaves) == [a] and _frames(leaves[a]) == list(range(6, 12)) + tid = task_instance_id(a, "reach", 0) + leaves = _resolve( + root, + DatasetFilter( + filter_lambdas=[f"lambda row: row['task_instance_id'] != '{tid}'"] + ), + ) + assert _frames(leaves[a]) == list(range(6, 12)) + + +def test_clip_hash_pins_and_excludes_its_time_range(segmented): + root, a, b = segmented + pinned = _resolve(root, DatasetFilter(episode_hashes=["clip_b2"])) + assert list(pinned) == [b] and _frames(pinned[b]) == [4, 5, 6, 7] + both = _resolve(root, DatasetFilter(episode_hashes=["clip_b2", a])) + assert _frames(both[a]) == [0, 1, 2, 3, *range(6, 12)] + excluded = _resolve(root, DatasetFilter(exclude_hashes=["clip_b1", "zarr-only"])) + assert _frames(excluded[b]) == [4, 5, 6, 7] and a in excluded + with pytest.raises(PinError, match="nope"): + _resolve(root, DatasetFilter(episode_hashes=["nope"])) + + +def test_clip_pins_through_multidataset(segmented): + from egomimic.rldb.multi import MultiDataset, pinned_episode_subset + + root, _, b = segmented + resolver = LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}) + ds = MultiDataset._from_resolver( + resolver, filters=DatasetFilter(episode_hashes=["clip_b1"]), mode="total" + ) + assert list(ds.datasets) == [b] + whole = MultiDataset(datasets=resolver.resolve(), mode="total") + video = pinned_episode_subset(whole, ["clip_b2"]) + assert _frames(video.datasets["clip_b2"]) == [4, 5, 6, 7] + with pytest.raises(PinError): + pinned_episode_subset(whole, ["nope"]) + + +HISTORY_KEY_MAP = { + **KEY_MAP, + "obs_head_pose": { + "key_type": "proprio_keys", + "zarr_key": "obs_head_pose", + "history": 3, + }, +} + + +def _at(leaf, frame): + return leaf[_frames(leaf).index(frame)] + + +def test_action_chunks_stop_at_the_task_end_unless_asked(segmented): + root, a, _ = segmented + leaf = _resolve(root)[a] + assert _at(leaf, 2)["action_pad_mask"].tolist() == [1, 1, 0] + head = _at(leaf, 2)["left.action_ee_pose"][:, 0].tolist() + assert head == [2.5, 3.5, 3.5] + through = _resolve(root, action_clamp="recording")[a] + assert _at(through, 2)["action_pad_mask"].tolist() == [1, 1, 1] + gaps = _resolve(root, include_taskless_frames=True)[a] + assert _at(gaps, 4)["action_pad_mask"].tolist() == [1, 1, 0] + + +def test_history_floor_at_the_task_start_when_asked(segmented): + root, a, _ = segmented + reach_back = LanceEpisodeResolver( + root, key_map=HISTORY_KEY_MAP, cameras={} + ).resolve()[a] + s = _at(reach_back, 7) + assert s["obs_head_pose"][:, 0].tolist() == [5, 6, 7] + assert s["proprio_history_mask"].tolist() == [1, 1, 1] + floored = LanceEpisodeResolver( + root, key_map=HISTORY_KEY_MAP, cameras={}, history_floor="task" + ).resolve()[a] + s = _at(floored, 7) + assert s["obs_head_pose"][:, 0].tolist() == [6, 6, 7] + assert s["proprio_history_mask"].tolist() == [0, 1, 1] + with pytest.raises(ValueError, match="history_floor"): + LanceEpisodeResolver(root, history_floor="episode") + + +def _hands(n=8): + valid = {h: np.ones(n, bool) for h in ("head", "left", "right")} + valid["left"][2] = False + valid["right"][5] = False + return valid + + +@pytest.mark.parametrize( + "require, want", + [ + ("both", [[0, 1], [3, 4], [6, 8]]), + ("either", [[0, 8]]), + ("left", [[0, 1], [3, 8]]), + ("right", [[0, 4], [6, 8]]), + ("none", [[0, 8]]), + ], +) +def test_require_hands_in_frame_rules(require, want): + iv = eligible_intervals(_hands(), ("left", "right"), 2, require) + np.testing.assert_array_equal(iv, want) + + +def test_require_hands_needs_the_head_unless_none(): + valid = _hands() + valid["head"][7] = False + assert eligible_intervals(valid, ("left", "right"), 2, "either")[-1, 1] == 7 + assert eligible_intervals(valid, ("left", "right"), 2, "none")[-1, 1] == 8 + + +def test_require_a_hand_the_embodiment_lacks_is_an_error(): + with pytest.raises(ValueError, match="left"): + eligible_intervals(_hands(), ("right",), 2, "left") + + +def test_samples_carry_per_hand_masks(catalog): + root, recs = catalog + rid = recs[0][0] + both = _resolve(root)[rid] + s = both[0] + assert s["action_hand_mask"].tolist() == [1, 1] + assert s["proprio_hand_valid"].tolist() == [1, 1] + either = _resolve(root, require_hands_in_frame="either")[rid] + assert _frames(either) == list(range(recs[0][2])) + # left is untracked at frame 4: in the chunk of anchor 3, at the anchor 4. + assert _at(either, 3)["action_hand_mask"].tolist() == [0, 1] + assert _at(either, 3)["proprio_hand_valid"].tolist() == [1, 1] + assert _at(either, 4)["proprio_hand_valid"].tolist() == [0, 1] + assert _at(either, 5)["action_hand_mask"].tolist() == [1, 1] + with pytest.raises(ValueError, match="require_hands_in_frame"): + _resolve(root, require_hands_in_frame="at_least_one") + + +@pytest.mark.parametrize("episode_cache", [False, True]) +def test_norm_stats_leave_out_untracked_hand_chunks(catalog, tmp_path, episode_cache): + from egomimic.rldb.multi import MultiDataset + + root, recs = catalog + leaves = _resolve(root, require_hands_in_frame="either") + ds = MultiDataset(datasets=leaves, mode="total") + stats = MultiDataset(state={}, norm_mode="zscore") + stats.populate_from_datasets({"human_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset( + ds, + "human_bimanual", + sample_frac=1.0, + num_workers=0, + episode_cache=tmp_path / "norm" if episode_cache else None, + ) + emb = int(ds[0]["embodiment"]) + tracked = [ + s["left.action_ee_pose"].numpy() + for s in (ds[i] for i in range(len(ds))) + if s["action_hand_mask"][0] == 1 + ] + assert len(tracked) == len(ds) - 3 * len(recs) + np.testing.assert_allclose( + stats.norm_stats[emb]["left.action_ee_pose"]["mean"], + np.mean(tracked, axis=0), + rtol=1e-5, + ) + + +EVA_KEY_MAP = { + "left.cmd_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.cmd_ee_pose", + "horizon": 3, + }, + "left.cmd_gripper": { + "key_type": "action_keys", + "zarr_key": "left.cmd_gripper", + "horizon": 3, + }, + "right.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "right.obs_ee_pose"}, + "left.obs_gripper": {"key_type": "proprio_keys", "zarr_key": "left.obs_gripper"}, +} + + +def _eva(tmp_path, valid=lambda rid, n: {}): + recs = write_catalog(tmp_path, lab="eva", valid=valid) + (a, _, _), (b, _, _) = _by_length(recs) + return tmp_path, a, b + + +def _resolve_eva(root, key_map=EVA_KEY_MAP, **kw): + return LanceEpisodeResolver( + root, lab="eva", key_map=key_map, cameras={}, **kw + ).resolve() + + +def test_eva_reads_robot_columns_by_their_zarr_keys(tmp_path): + root, a, b = _eva(tmp_path) + leaf = _resolve_eva(root)[a] + assert _frames(leaf) == list(range(12)) + s = _at(leaf, 2) + offset = 100 * sorted([a, b]).index(a) + np.testing.assert_allclose( + s["left.cmd_ee_pose"][:, 0].numpy(), offset + 0.5 + np.arange(3, 6) + ) + assert s["left.cmd_gripper"].shape == (3, 1) + np.testing.assert_allclose( + s["left.cmd_gripper"][:, 0].numpy(), np.arange(2, 5) / 12 + ) + assert float(s["right.obs_ee_pose"][0]) == offset + 2.25 + assert float(s["left.obs_gripper"][0]) == 0.5 + assert s["action_hand_mask"].tolist() == [1, 1] + + +def test_eva_arm_dropouts_gate_anchors_and_mask_the_arm(tmp_path): + def dropout(rid, n): + v = np.ones(n, bool) + v[4] = False + return {"right_valid": v} if n == 12 else {} + + root, a, _ = _eva(tmp_path, valid=dropout) + assert _frames(_resolve_eva(root)[a]) == [0, 1, *range(5, 12)] + either = _resolve_eva(root, require_hands_in_frame="either")[a] + assert _at(either, 3)["action_hand_mask"].tolist() == [1, 0] + + +def test_eva_single_arm_robot_never_masks_the_missing_arm(tmp_path): + root, _, b = _eva(tmp_path) + key_map = {k: v for k, v in EVA_KEY_MAP.items() if not k.startswith("right.")} + leaf = _resolve_eva(root, key_map=key_map)[b] + assert _frames(leaf) == list(range(8)) + assert leaf[0]["action_hand_mask"].tolist() == [1, 1] + with pytest.raises(ValueError, match="require_hands_in_frame"): + _resolve_eva(root, key_map=key_map, require_hands_in_frame="right") diff --git a/tests/unit/test_lance_timeline.py b/tests/unit/test_lance_timeline.py new file mode 100644 index 0000000000..04567c1f91 --- /dev/null +++ b/tests/unit/test_lance_timeline.py @@ -0,0 +1,67 @@ +import numpy as np + +from egomimic.rldb.lance.timeline import ( + FROZEN_RUN, + build_timeline, + nearest_dt, + pose_outliers, +) + +P = 33_333_333 + + +def test_timeline_without_gaps_is_identity(): + ts = np.arange(10, dtype=np.int64) * P + out, src, filler, after_gap = build_timeline(ts) + np.testing.assert_array_equal(out, ts) + np.testing.assert_array_equal(src, np.arange(10)) + assert not filler.any() and not after_gap.any() + + +def test_timeline_fills_dropped_frames(): + ts = np.array([0, 1, 2, 5, 6], dtype=np.int64) * P + out, src, filler, after_gap = build_timeline(ts) + assert len(out) == 7 + np.testing.assert_array_equal(src, [0, 1, 2, 2, 2, 3, 4]) + np.testing.assert_array_equal(filler, [0, 0, 0, 1, 1, 0, 0]) + np.testing.assert_array_equal(after_gap, [0, 0, 0, 0, 0, 1, 0]) + assert (np.diff(out) > 0).all() + assert abs(out[3] - 3 * P) <= 1 + + +def test_timeline_tolerates_jitter(): + ts = np.array([0, P + 4_000_000, 2 * P - 3_000_000, 3 * P], dtype=np.int64) + _, _, filler, after_gap = build_timeline(ts) + assert not filler.any() and not after_gap.any() + + +def test_nearest_dt(): + ref = np.array([0, 100, 200], dtype=np.int64) + np.testing.assert_array_equal( + nearest_dt(np.array([0, 40, 160, 260]), ref), [0, 40, 40, 60] + ) + + +def test_pose_outliers_flags_jump_only_at_the_jump(): + xyz = np.zeros((6, 3)) + xyz[:, 0] = [0, 0.01, 0.02, 1.0, 1.01, 1.02] + bad = pose_outliers(xyz, np.ones(6, bool)) + np.testing.assert_array_equal(bad, [0, 0, 0, 1, 0, 0]) + + +def test_pose_outliers_skips_invalid_rows(): + xyz = np.zeros((4, 3)) + xyz[:, 0] = [0, 5.0, 0.01, 0.02] + ok = np.array([1, 0, 1, 1], bool) + assert not pose_outliers(xyz, ok).any() + + +def test_pose_outliers_frozen_run(): + n = FROZEN_RUN + 5 + xyz = np.cumsum(np.full((n, 3), 0.001), axis=0) + xyz[2 : 2 + FROZEN_RUN] = xyz[2] + bad = pose_outliers(xyz, np.ones(n, bool)) + assert bad[2 : 2 + FROZEN_RUN].all() and not bad[:2].any() + short = np.cumsum(np.full((n, 3), 0.001), axis=0) + short[3] = short[2] + assert not pose_outliers(short, np.ones(n, bool)).any() diff --git a/tests/unit/test_latent_dataset_kwargs.py b/tests/unit/test_latent_dataset_kwargs.py index 237f9715ad..f9f2526bc4 100644 --- a/tests/unit/test_latent_dataset_kwargs.py +++ b/tests/unit/test_latent_dataset_kwargs.py @@ -5,8 +5,6 @@ from __future__ import annotations -import inspect - import hydra import pytest from fixtures.synthetic_episodes import write_episode @@ -19,10 +17,6 @@ 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 = {} diff --git a/tests/unit/test_leaf_equivalence.py b/tests/unit/test_leaf_equivalence.py new file mode 100644 index 0000000000..845b539096 --- /dev/null +++ b/tests/unit/test_leaf_equivalence.py @@ -0,0 +1,89 @@ +"""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)] + }, + ) + 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)] + # Only Lance tracks per-hand validity; every action chunk sampled here + # has both hands (history can reach back to an untracked frame). + assert bool((a.pop("action_hand_mask") == 1).all()) + a.pop("proprio_hand_valid") + 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_loader_read_coalescing.py b/tests/unit/test_loader_read_coalescing.py new file mode 100644 index 0000000000..3c8929e044 --- /dev/null +++ b/tests/unit/test_loader_read_coalescing.py @@ -0,0 +1,342 @@ +"""``ZarrDataset.__getitem__`` reads each zarr key once, not once per keymap entry. + +The mecka keypoint keymap with ``proprio_history=4`` names ``left/right.obs_keypoints`` +and ``left/right.obs_wrist_pose`` twice each -- once as a forward action chunk +``(i, i+H)``, once as a backward proprio window ``(i-K+1, i+1)``. Those were ten +separate ``ZarrEpisode.read`` calls over six distinct keys; they are now one zarr call +per key (``ZarrEpisode.read_intervals``), sliced twice. Nothing the dataset emits may change. + +The episode here is RAMPED -- every cell is ``frame * 1000 + column`` -- so a slice +taken at the wrong offset into the coalesced span produces different numbers rather +than the same ones. ``tests/fixtures/synthetic_episodes.py`` deliberately holds each +numeric key constant in time (for the norm-stat quantile bounds), which would make an +off-by-one here invisible: the single most likely bug in this change. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import simplejpeg +import zarr + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + ZarrAnnotationCutoffDataset, + ZarrDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 160 +HORIZON = Human.ACTION_HORIZON +# half-open [start, end); ZarrAnnotationCutoffDataset clamps action chunks at the +# end of the span holding the start frame, so idx 60 and 130 both hit a cutoff +ANN_SPANS = ((0, 80), (80, T - 1)) + +# (label, keymap kwargs). Covers both action spaces, history on and off, and the +# ee_pose variant that gives four entries over two shared keys. +KEYMAPS = { + "keypoints_k4": {"keymap_mode": "keypoints", "proprio_history": 4}, + "keypoints_k1": {"keymap_mode": "keypoints", "proprio_history": 1}, + "keypoints_k4_ee": { + "keymap_mode": "keypoints", + "proprio_history": 4, + "include_ee_pose": True, + }, + "cartesian_k4": {"keymap_mode": "cartesian", "proprio_history": 4}, + "cartesian_k1": {"keymap_mode": "cartesian", "proprio_history": 1}, +} + + +def _keymap(label): + return Human.get_keymap( + has_head_pose=True, annotation_key="annotations", **KEYMAPS[label] + ) + + +def _ramp(width: int) -> np.ndarray: + """``[t, d] = t * 1000 + d`` -- every cell identifies its own frame.""" + return ( + np.arange(T, dtype=np.float64)[:, None] * 1000.0 + + np.arange(width, dtype=np.float64)[None, :] + ) + + +@pytest.fixture(scope="module") +def episode(tmp_path_factory): + numeric = {"obs_head_pose": _ramp(7)} + for side in ("left", "right"): + # distinct per side so a left/right mix-up is caught too + off = 0.0 if side == "left" else 0.5 + numeric[f"{side}.obs_wrist_pose"] = _ramp(7) + off + numeric[f"{side}.obs_keypoints"] = _ramp(63) + off + numeric[f"{side}.obs_ee_pose"] = _ramp(7) + off + rng = np.random.default_rng(0) + return ZarrWriter.create_and_write( + tmp_path_factory.mktemp("ramped") / "ramped.zarr", + numeric_data=numeric, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + task_name="ramp", + annotations=[(t, a, b) for t, (a, b) in zip(("fold", "put down"), ANN_SPANS)], + intrinsics={"front_1": np.eye(3, 4)}, + ) + + +def _ann_end(idx: int) -> int | None: + return next((e for s, e in ANN_SPANS if s <= idx < e), None) + + +def _expected_window(spec, idx, total, cutoff: bool): + """The read window for one keymap entry, written out independently of the + dataset: the spec this change must not move.""" + horizon, history = spec.get("horizon"), spec.get("history") + if horizon is not None: + end = min(idx + horizon, total) + if cutoff and spec.get("key_type") == "action_keys": + end = min(end, _ann_end(idx) or end) + return (idx, end) + if history is not None and int(history) > 1: + return (max(0, idx - int(history) + 1), idx + 1) + return (idx, None) + + +def _oracle(episode, key_map, idx, cutoff=False): + """One zarr slice per keymap entry -- the pre-coalescing read path.""" + store = zarr.open_group(str(episode), mode="r") + out = {} + for k, spec in key_map.items(): + if spec.get("key_type") == "annotation_keys": + continue + start, end = _expected_window(spec, idx, T, cutoff) + arr = store[spec["zarr_key"]] + chunk = arr[start:end] if end is not None else arr[start : start + 1][0] + horizon, history = spec.get("horizon"), spec.get("history") + if horizon is not None and chunk.shape[0] < horizon: + pad = np.repeat(chunk[-1:], horizon - chunk.shape[0], axis=0) + chunk = np.concatenate([chunk, pad], axis=0) + if horizon is None and history is not None and int(history) > 1: + if chunk.shape[0] < int(history): + pad = np.repeat(chunk[:1], int(history) - chunk.shape[0], axis=0) + chunk = np.concatenate([pad, chunk], axis=0) + out[k] = chunk + return out + + +# 0..3 clamp the backward window, 60 is interior, the rest truncate the forward chunk +@pytest.mark.parametrize("idx", [0, 1, 3, 60, 79, 130, T - HORIZON, T - 3, T - 1]) +@pytest.mark.parametrize("label", sorted(KEYMAPS)) +@pytest.mark.parametrize("cls", [ZarrDataset, ZarrAnnotationCutoffDataset]) +def test_getitem_matches_the_per_entry_read_path(episode, idx, label, cls): + key_map = _keymap(label) + sample = cls(episode, key_map=key_map, transform_list=None)[idx] + expected = _oracle(episode, key_map, idx, cutoff=cls is ZarrAnnotationCutoffDataset) + + # __getitem__ emits float32; compare bit-exactly in the dtype it emits + image_keys = {k for k, s in key_map.items() if s.get("key_type") == "camera_keys"} + for k, want in expected.items(): + if k in image_keys: + want = simplejpeg.decode_jpeg(want, colorspace="RGB") + want = np.transpose(want, (2, 0, 1)) / 255.0 + got = np.asarray(sample[k]) + assert got.shape == want.shape, k + np.testing.assert_array_equal(got, want.astype(np.float32), err_msg=k) + + +def test_shapes_are_the_keymap_contract(episode): + sample = ZarrDataset(episode, key_map=_keymap("keypoints_k4"))[60] + assert sample["left.action_keypoints"].shape == (HORIZON, 63) + assert sample["left.action_wrist_pose"].shape == (HORIZON, 7) + assert sample["left.obs_keypoints"].shape == (4, 63) + assert sample["left.obs_wrist_pose"].shape == (4, 7) + assert sample["obs_head_pose"].shape == (7,) + + +def _count_reads(monkeypatch) -> list[dict]: + """Per ``read_intervals`` call, ``{key: [(group_start, rows read), ...]}``.""" + calls: list[dict] = [] + inner = ZarrEpisode.read_intervals + + def spy(self, intervals): + out = inner(self, intervals) + calls.append({k: [(lo, len(rows)) for lo, rows in g] for k, g in out.items()}) + return out + + # ZarrEpisode has __slots__, so the counter goes on the class + monkeypatch.setattr(ZarrEpisode, "read_intervals", spy) + return calls + + +def test_one_read_per_distinct_zarr_key(episode, monkeypatch): + key_map = _keymap("keypoints_k4") + calls = _count_reads(monkeypatch) + + ZarrDataset(episode, key_map=key_map, transform_list=None)[60] + + assert len(calls) == 1, "one read call per __getitem__" + read_keys = set(calls[0]) + entries = [s for s in key_map.values() if s.get("key_type") != "annotation_keys"] + assert read_keys == {s["zarr_key"] for s in entries} + assert len(read_keys) == 6 < len(entries) == 10, "the four duplicates collapsed" + + +def test_rows_read_are_the_union_of_the_windows(episode, monkeypatch): + """Asserts rows READ, which the value oracle cannot see: a loader that read + the whole episode and sliced it would still match every value.""" + calls = _count_reads(monkeypatch) + ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=None)[60] + rows = calls[0] + assert rows["left.obs_keypoints"] == [ + (57, 3 + HORIZON) + ], "history + chunk, one group" + assert rows["obs_head_pose"] == [(60, 1)], "a key read once is not widened" + assert rows["images.front_1"] == [(60, 1)] + + +@pytest.mark.parametrize( + "intervals, chunk, groups", + [ + # numeric keys: 100-row chunks + ([(57, 61), (60, 160)], 100, [(57, 160)]), + ([(0, 1), (150, 151)], 100, [(0, 151)]), # adjacent chunks: no new fetch + ([(0, 1), (250, 251)], 100, [(0, 1), (250, 251)]), # chunk 1 untouched + ([(99, 100), (100, 101)], 100, [(99, 101)]), # straddles a boundary + # JPEG keys: one frame per chunk + ([(60, 61), (57, 58)], 1, [(57, 58), (60, 61)]), + ([(59, 60), (60, 61)], 1, [(59, 61)]), + ([(10, 20), (15, 25)], 1, [(10, 25)]), + ([(0, 1), (0, 1)], 1, [(0, 1)]), + ], +) +def test_plan_reads_bridges_only_chunks_already_fetched(intervals, chunk, groups): + assert ZarrEpisode.plan_reads(intervals, chunk) == groups + + +def test_far_apart_jpeg_frames_are_two_rows_not_the_gap(episode): + reader = ZarrEpisode(episode) + store = zarr.open_group(str(episode), mode="r") + got = reader.read_intervals({"images.front_1": [(100, 101), (10, 11)]}) + assert [(lo, len(rows)) for lo, rows in got["images.front_1"]] == [ + (10, 1), + (100, 1), + ] + for lo, rows in got["images.front_1"]: + assert rows[0] == store["images.front_1"][lo : lo + 1][0] + + +def test_union_does_not_depend_on_keymap_order(episode): + """The real keymaps happen to list the forward chunk before the backward + window, which would hide a union that just took the last entry's bounds.""" + pose = "left.obs_wrist_pose" + order = { + "proprio": {"key_type": "proprio_keys", "zarr_key": pose, "history": 4}, + "action": {"key_type": "action_keys", "zarr_key": pose, "horizon": 10}, + } + for key_map in (order, dict(reversed(order.items()))): + ds = ZarrDataset(episode, key_map=key_map, transform_list=None) + starts = [w[1][0] for w in ds._read_windows(60).values()] + assert min(starts) == 57 and max(starts) == 60 + sample = ds[60] + np.testing.assert_array_equal( + np.asarray(sample["proprio"]), _ramp(7)[57:61].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["action"]), _ramp(7)[60:70].astype(np.float32) + ) + + +def test_shared_span_slices_do_not_alias(episode): + """Two entries sliced out of one span must own their arrays: a transform is + handed the raw ndarrays and is free to write into them. Asserted at that + point, because __getitem__'s closing float64 -> float32 cast would sever the + aliasing before anything downstream could see it.""" + seen = {} + + class _Spy: + def transform(self, data): + seen["shares"] = np.shares_memory( + data["left.action_keypoints"], data["left.obs_keypoints"] + ) + data["left.action_keypoints"][0] += 1.0 # a transform writing in place + seen["proprio_last"] = data["left.obs_keypoints"][-1].copy() + return data + + ds = ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=[_Spy()]) + ds[60] + assert seen["shares"] is False, "the two windows must not share a buffer" + # frame 60 of the ramp, unperturbed by the write to the action chunk + np.testing.assert_array_equal(seen["proprio_last"], _ramp(63)[60]) + + +def test_single_frame_entry_is_sliced_at_the_span_offset(episode): + """A key read as a single frame AND as a window starting earlier: the + single-frame slice needs the span offset like any other. No shipped keymap + mixes the two on one key (the history keys are exactly the obs entries), so + nothing else would catch a `block[0]` here.""" + pose = "left.obs_wrist_pose" + key_map = { + "history": {"key_type": "proprio_keys", "zarr_key": pose, "history": 4}, + "current": {"key_type": "proprio_keys", "zarr_key": pose}, + } + sample = ZarrDataset(episode, key_map=key_map, transform_list=None)[60] + np.testing.assert_array_equal( + np.asarray(sample["current"]), _ramp(7)[60].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["history"]), _ramp(7)[57:61].astype(np.float32) + ) + + +def test_a_jpeg_failure_rereads_at_the_new_index(episode, monkeypatch): + """The decode retry picks a different frame and restarts the loop, so the + windows, spans and one read all have to be recomputed against the new idx -- + which the coalescing moved.""" + real = simplejpeg.decode_jpeg + calls: list[int] = [] + + def _fail_once(data, **kw): + calls.append(len(data)) + if len(calls) == 1: + raise ValueError("corrupt jpeg") + return real(data, **kw) + + monkeypatch.setattr( + "egomimic.rldb.zarr.zarr_dataset_multi.simplejpeg.decode_jpeg", _fail_once + ) + ds = ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=None) + sample = ds[60] + + assert len(calls) == 2, "one failure, then the retry's decode" + # every numeric key comes from the RETRY's frame, not frame 60 + frame = int(round(float(np.asarray(sample["obs_head_pose"])[0]) / 1000.0)) + assert frame != 60, "the retry moved to another frame" + np.testing.assert_array_equal( + np.asarray(sample["obs_head_pose"]), _ramp(7)[frame].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["left.obs_keypoints"]), + _ramp(63)[max(0, frame - 3) : frame + 1].astype(np.float32), + ) + + +def test_an_array_shorter_than_its_window_fails_loudly(tmp_path): + """A StopIteration out of __getitem__ would end the epoch silently.""" + numeric = {"obs_head_pose": _ramp(7)} + for side in ("left", "right"): + numeric[f"{side}.obs_wrist_pose"] = _ramp(7) + numeric[f"{side}.obs_keypoints"] = _ramp(63) + numeric[f"{side}.obs_ee_pose"] = _ramp(7) + path = ZarrWriter.create_and_write( + tmp_path / "short.zarr", + numeric_data=numeric, + image_data={"images.front_1": np.zeros((T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + task_name="ramp", + annotations=[("fold", 0, T - 1)], + intrinsics={"front_1": np.eye(3, 4)}, + ) + zarr.open_group(str(path), mode="r+")["left.obs_keypoints"].resize((130, 63)) + ds = ZarrDataset(path, key_map=_keymap("keypoints_k4"), transform_list=None) + with pytest.raises(ValueError, match="left.obs_keypoints"): + ds[115] diff --git a/tests/unit/test_loss_masking.py b/tests/unit/test_loss_masking.py new file mode 100644 index 0000000000..5197519d74 --- /dev/null +++ b/tests/unit/test_loss_masking.py @@ -0,0 +1,394 @@ +"""Episode-tail padding must not supervise "stop moving". + +The dataset emits a real ``action_pad_mask`` next to the action chunk, the +interpolators carry it to the model horizon, and ``DenoisingPolicy.loss_fn`` +drops the padded steps -- but only behind the ``use_pad_mask`` algo flag, so +the flag-off path is bit-for-bit today's behaviour. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from fixtures.synthetic_episodes import write_episode +from omegaconf import OmegaConf + +import egomimic +from egomimic.algo.hpt import HPT +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.zarr_dataset_multi import ( + LocalAnnotationCutoffEpisodeResolver, + LocalEpisodeResolver, + MultiDataset, +) + +T_FRAMES = 48 # write_episode default +HORIZON = Human.ACTION_HORIZON # 30 raw action frames +CHUNK = 100 # model horizon after interpolation + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _pad_mask_transform_cls(): + """Imported inside the tests so the module still collects (and every test + fails on its own assertion) before the transform exists.""" + from egomimic.rldb.transforms import InterpolatePadMask + + return InterpolatePadMask + + +def _leaf(ds: MultiDataset): + """The single per-episode ZarrDataset behind a one-episode MultiDataset.""" + return next(iter(ds.datasets.values())) + + +def _mecka_dataset(tmp_path, *, transforms, cutoff: bool = False, n: int = 1): + for i in range(n): + write_episode(tmp_path, "mecka", seed=i) + # No annotation_key: the cutoff reads the zarr's spans directly, and the + # default collate cannot stack the per-sample annotation string lists. + key_map = Human.get_keymap(keymap_mode="cartesian") + resolver_cls = ( + LocalAnnotationCutoffEpisodeResolver if cutoff else LocalEpisodeResolver + ) + resolver = resolver_cls(tmp_path, key_map=key_map, transform_list=transforms) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _expected_upsampled_mask(mask: np.ndarray, new_len: int) -> np.ndarray: + """Independent restatement of the rule: an output step is valid only if the + two source frames it interpolates between are both real; a step landing + exactly on a source frame inherits that frame.""" + T = mask.shape[0] + old = np.linspace(0, 1, T) + new = np.linspace(0, 1, new_len) + out = np.empty(new_len, dtype=np.float32) + for j, t in enumerate(new): + i = int(np.searchsorted(old, t)) + if i < T and np.isclose(old[i], t): + out[j] = mask[i] + else: + out[j] = mask[i - 1] * mask[i] + return out + + +# --------------------------------------------------------------------------- +# 1. dataset emits the real mask +# --------------------------------------------------------------------------- + + +def test_dataset_emits_action_pad_mask_for_a_tail_chunk(tmp_path): + ds = _mecka_dataset(tmp_path, transforms=[]) + leaf = _leaf(ds) + assert leaf.total_frames == T_FRAMES + + sample = leaf[T_FRAMES - 5] + mask = np.asarray(sample["action_pad_mask"]) + assert mask.shape == (HORIZON,) + expected = np.zeros(HORIZON, dtype=np.float32) + expected[:5] = 1.0 + np.testing.assert_array_equal(mask, expected) + + head = np.asarray(leaf[0]["action_pad_mask"]) + np.testing.assert_array_equal(head, np.ones(HORIZON, dtype=np.float32)) + + +def test_action_pad_mask_is_a_float32_tensor(tmp_path): + ds = _mecka_dataset(tmp_path, transforms=[]) + value = _leaf(ds)[0]["action_pad_mask"] + assert isinstance(value, torch.Tensor) + assert value.dtype == torch.float32 + + +def test_no_action_pad_mask_without_an_action_horizon(tmp_path): + """Keymaps with no horizoned ``action_keys`` entry must not grow a key.""" + write_episode(tmp_path, "mecka", seed=0) + key_map = Human.get_keymap(keymap_mode="cartesian") + for spec in key_map.values(): + if spec.get("key_type") == "action_keys": + spec["key_type"] = "proprio_keys" + spec.pop("horizon", None) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + assert "action_pad_mask" not in _leaf(ds)[0] + + +def test_mask_survives_bounds_check_and_normalization(tmp_path): + """``MultiDataset.__getitem__`` runs _check_bounds + normalize; the mask is + not in the norm stats and must come through untouched.""" + transforms = Human.get_transform_list(mode="cartesian_6d", stride=1) + ds = _mecka_dataset(tmp_path, transforms=transforms, n=3) + + stats = MultiDataset(state={}, norm_mode="quantile") + stats.populate_from_datasets({"human_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset(ds, "human_bimanual", sample_frac=1.0, num_workers=0) + emb = get_embodiment_id("human_bimanual") + assert "action_pad_mask" not in stats.norm_stats[emb] + + raw = _leaf(ds)[0]["action_pad_mask"].clone() + ds.set_norm_stats_from(stats) + normalized = ds[0] + assert "action_pad_mask" in normalized + assert normalized["action_pad_mask"].shape == (CHUNK,) + torch.testing.assert_close(normalized["action_pad_mask"], raw) + + +# --------------------------------------------------------------------------- +# 2. annotation cutoff +# --------------------------------------------------------------------------- + + +def test_annotation_cutoff_is_padding(tmp_path): + """Synthetic episodes carry spans [0, T//2) and [T//2, T-1); a chunk that + starts inside the first span is real only up to that span's end.""" + ds = _mecka_dataset(tmp_path, transforms=[], cutoff=True) + leaf = _leaf(ds) + ann_end = T_FRAMES // 2 + idx = ann_end - 4 + assert idx + HORIZON > ann_end # the cutoff, not the episode end, bites + + mask = np.asarray(leaf[idx]["action_pad_mask"]) + expected = np.zeros(HORIZON, dtype=np.float32) + expected[: ann_end - idx] = 1.0 + np.testing.assert_array_equal(mask, expected) + + +# --------------------------------------------------------------------------- +# 3. InterpolatePadMask +# --------------------------------------------------------------------------- + + +def test_interpolate_pad_mask_30_to_100_with_five_real_frames(): + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + batch = {"action_pad_mask": mask.copy()} + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch)[ + "action_pad_mask" + ] + + expected = _expected_upsampled_mask(mask, CHUNK) + assert out.shape == (CHUNK,) + assert out.dtype == np.float32 + np.testing.assert_array_equal(out, expected) + assert out.sum() == expected.sum() + # The tail is padding and the head is real: the mask is a prefix of ones. + assert out[0] == 1.0 and out[-1] == 0.0 + assert np.all(np.diff(out) <= 0) + + +def test_interpolate_pad_mask_all_real_stays_all_ones(): + batch = {"action_pad_mask": np.ones(HORIZON, dtype=np.float32)} + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch)[ + "action_pad_mask" + ] + np.testing.assert_array_equal(out, np.ones(CHUNK, dtype=np.float32)) + + +def test_interpolate_pad_mask_honours_stride(): + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:16] = 1.0 + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK, stride=3).transform( + {"action_pad_mask": mask.copy()} + )["action_pad_mask"] + np.testing.assert_array_equal(out, _expected_upsampled_mask(mask[::3], CHUNK)) + + +def test_interpolate_pad_mask_is_a_noop_without_the_key(): + batch = {"actions_cartesian": np.zeros((3, 4))} + assert _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch) == batch + + +def test_interpolate_pad_mask_is_idempotent(): + """``include_ee_pose`` chains two builders that both carry the mask.""" + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + t = _pad_mask_transform_cls()(new_chunk_length=CHUNK, stride=3) + once = t.transform({"action_pad_mask": mask.copy()}) + twice = t.transform(dict(once)) + np.testing.assert_array_equal(twice["action_pad_mask"], once["action_pad_mask"]) + + +def test_human_transform_lists_upsample_the_mask(): + modes = [ + "cartesian_6d", + "cartesian_wristframe_6d", + "keypoints_headframe_6d", + "keypoints_wristframe_6d", + ] + for mode in modes: + tl = Human.get_transform_list(mode=mode, stride=3) + assert any( + isinstance(t, _pad_mask_transform_cls()) for t in tl + ), f"{mode} does not carry the pad mask" + + +def test_keypoints_with_ee_pose_upsamples_the_mask_once(tmp_path): + """Two builders chain here; the mask must still come out (100,).""" + tl = Human.get_transform_list( + mode="keypoints_wristframe_6d", stride=1, include_ee_pose=True + ) + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + batch = {"action_pad_mask": mask.copy()} + for t in tl: + if isinstance(t, _pad_mask_transform_cls()): + batch = t.transform(batch) + assert batch["action_pad_mask"].shape == (CHUNK,) + np.testing.assert_array_equal( + batch["action_pad_mask"], _expected_upsampled_mask(mask, CHUNK) + ) + + +# --------------------------------------------------------------------------- +# 4. loss_fn +# --------------------------------------------------------------------------- + + +def _policy() -> DenoisingPolicy: + return DenoisingPolicy.__new__(DenoisingPolicy) + + +def test_loss_fn_without_mask_matches_mse_loss(): + g = torch.Generator().manual_seed(0) + pred = torch.randn(3, 7, 5, generator=g) + target = torch.randn(3, 7, 5, generator=g) + p = _policy() + torch.testing.assert_close(p.loss_fn(pred, target), F.mse_loss(pred, target)) + ones = torch.ones(3, 7, 1) + torch.testing.assert_close( + p.loss_fn(pred, target, pad_mask=ones), F.mse_loss(pred, target) + ) + + +def test_loss_fn_ignores_masked_steps(): + g = torch.Generator().manual_seed(1) + pred = torch.randn(2, 6, 4, generator=g) + target = torch.randn(2, 6, 4, generator=g) + mask = torch.ones(2, 6, 1) + mask[:, 4:] = 0.0 + p = _policy() + before = p.loss_fn(pred, target, pad_mask=mask) + + perturbed = pred.clone() + perturbed[:, 4:] += 100.0 + torch.testing.assert_close(p.loss_fn(perturbed, target, pad_mask=mask), before) + + # ... and it is exactly the mse over the kept steps. + torch.testing.assert_close(before, F.mse_loss(pred[:, :4], target[:, :4])) + + +def test_loss_fn_all_zero_mask_does_not_divide_by_zero(): + p = _policy() + loss = p.loss_fn( + torch.randn(2, 6, 4), torch.randn(2, 6, 4), pad_mask=torch.zeros(2, 6, 1) + ) + assert torch.isfinite(loss) and float(loss) == 0.0 + + +class _FixedPredPolicy(DenoisingPolicy): + """Minimal DenoisingPolicy: ``predict`` returns a canned prediction.""" + + def __init__(self, action_horizon: int, pred: torch.Tensor): + nn.Module.__init__(self) + self.action_horizon = action_horizon + self.pooling = None + self.padding = None + self._pred = pred + + def predict(self, actions, global_cond): + return self._pred, actions + + +def test_compute_loss_forwards_the_pad_mask(): + g = torch.Generator().manual_seed(2) + target = torch.randn(2, 6, 4, generator=g) + pred = torch.randn(2, 6, 4, generator=g) + policy = _FixedPredPolicy(6, pred) + mask = torch.ones(2, 6, 1) + mask[:, 4:] = 0.0 + global_cond = torch.zeros(2, 3) + loss = policy.compute_loss(global_cond, {"action": target, "pad_mask": mask}) + torch.testing.assert_close(loss, F.mse_loss(pred[:, :4], target[:, :4])) + + +# --------------------------------------------------------------------------- +# 5. process_batch_for_training + the config flag +# --------------------------------------------------------------------------- + + +class _NoRenameNormStats: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def _process(use_pad_mask: bool, action_pad_mask: torch.Tensor | None, S: int = 6): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRenameNormStats() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = use_pad_mask + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + algo.ac_keys = {emb_id: "actions_cartesian"} + inner = {"actions_cartesian": torch.zeros(2, S, 18)} + if action_pad_mask is not None: + inner["action_pad_mask"] = action_pad_mask + return HPT.process_batch_for_training(algo, {emb: inner})[emb_id] + + +def test_process_batch_keeps_ones_when_the_flag_is_off(): + mask = torch.ones(2, 6) + mask[:, 3:] = 0.0 + out = _process(False, mask) + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + assert "action_pad_mask" not in out + + +def test_process_batch_uses_the_dataset_mask_when_the_flag_is_on(): + mask = torch.ones(2, 6) + mask[:, 3:] = 0.0 + out = _process(True, mask) + assert out["pad_mask"].shape == (2, 6, 1) + torch.testing.assert_close(out["pad_mask"], mask[..., None]) + assert "action_pad_mask" not in out + + +def test_process_batch_falls_back_to_ones_without_the_dataset_key(): + out = _process(True, None) + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + + +def test_process_batch_raises_on_a_length_mismatch(): + with pytest.raises(ValueError, match="6"): + _process(True, torch.ones(2, 5)) + + +def test_use_pad_mask_defaults_to_false(): + import inspect + + assert inspect.signature(HPT.__init__).parameters["use_pad_mask"].default is False + + +# The two flagship recipes: the ones a run would flip the knob on. Every other +# HPT config takes the HPT.__init__ default, which the test above pins. +@pytest.mark.parametrize( + "yaml_name", + ["hpt_bc_mecka_6d_300M", "hpt_bc_keypoints_wrist_300M"], +) +def test_algo_yaml_declares_use_pad_mask_off(yaml_name): + from pathlib import Path + + path = ( + Path(egomimic.__file__).parent / "hydra_configs" / "model" / f"{yaml_name}.yaml" + ) + cfg = OmegaConf.load(path) + assert cfg.robomimic_model.use_pad_mask is False diff --git a/tests/unit/test_mcap_to_lance.py b/tests/unit/test_mcap_to_lance.py new file mode 100644 index 0000000000..25d6aeb902 --- /dev/null +++ b/tests/unit/test_mcap_to_lance.py @@ -0,0 +1,293 @@ +"""MCAP -> Lance converter (egomimic/scripts/mcap_process) on tiny synthetic episodes.""" + +import shutil +import subprocess + +import numpy as np +import pytest +from fixtures.mcap_episodes import T0, write_mcap_episode + +from egomimic.rldb.lance import commit +from egomimic.rldb.mcap import spec + +lance = pytest.importorskip("lance") + + +def _ffmpeg_has(codec: str) -> bool: + ffmpeg = shutil.which("ffmpeg") + if not ffmpeg: + return False + out = subprocess.run([ffmpeg, "-hide_banner", "-encoders"], capture_output=True) + return codec in out.stdout.decode() + + +if not _ffmpeg_has("libx264"): + pytest.skip("needs ffmpeg with libx264", allow_module_level=True) + +from egomimic.scripts.lance_commit import commit_staged # noqa: E402 +from egomimic.scripts.mcap_process.mcap_to_lance import ( # noqa: E402 + convert_and_stage, +) + + +def commit_all(tmp_path, replace=False): + """Run the committer over every lab with staged work, annotations last.""" + staging = tmp_path / "staging" + labs = sorted(p.name for p in staging.iterdir()) if staging.exists() else [] + for lab in sorted(labs, key=lambda x: x == "annotations"): + commit_staged( + str(staging), + lab, + str(tmp_path / "lance"), + tmp_path / "work", + replace=replace, + ) + + +def land(tmp_path, ep, replace=False, convert=None): + r = (convert or convert_and_stage)( + ep, + tmp_path / "staging", + tmp_path / "lance", + codec="libx264", + mano=False, + replace=replace, + ) + commit_all(tmp_path, replace) + return r + + +def run(tmp_path, embodiment="human_bimanual", episode_id="0", **kw): + ep = write_mcap_episode(tmp_path / "eps", embodiment, episode_id, **kw) + return land(tmp_path, ep) + + +def table(tmp_path, name): + return lance.dataset(tmp_path / "lance" / name).to_table().to_pylist() + + +def test_episode_lands_in_every_table(tmp_path): + r = run(tmp_path) + assert r.status == "staged" and r.problems == [] + assert r.recording_id == spec.recording_id("aria", "acme-rec-0") + assert (r.num_frames, r.tasks, r.annotations_staged) == (60, 2, 2) + assert r.valid_frac == {"head": 1.0, "left": 1.0, "right": 1.0} + + (rec,) = table(tmp_path, "recordings.lance") + assert (rec["lab"], rec["site"], rec["license"]) == ("aria", "rl2", "CC-BY-4.0") + assert rec["operator"] == "ab" * 32 and rec["num_frames"] == 60 + assert int(rec["captured_at"].timestamp() * 1000) == T0 // 10**6 + assert {k for k, _ in rec["video_path"]} == {"front_1", "front_1_480"} + + tasks = sorted(table(tmp_path, "tasks.lance"), key=lambda t: t["t_start_ns"]) + assert [t["task_name"] for t in tasks] == ["fold_clothes", "pack_clothes"] + assert tasks[0]["t_start_ns"] == T0 and tasks[1]["objects"] == ["shirt", "basket"] + labels = table(tmp_path, "annotations.lance") + assert {a["text"] for a in labels} == {"pick up the shirt", "drop it in the basket"} + assert {a["annotation_vendor"] for a in labels} == {"acme"} + point = next(a for a in labels if a["text"] == "drop it in the basket") + assert 0 < point["t_end_ns"] - point["t_start_ns"] <= 34_000_000 # one frame wide + (alias,) = table(tmp_path, "aliases.lance") + assert alias["recording_id"] == r.recording_id + # Half-open: the alias covers the last frame. + assert alias["t_end_ns"] > max( + s["timestamp_ns"] for s in table(tmp_path, "steps/aria.lance") + ) + assert commit.list_files(tmp_path / "staging" / "aria") == {} + + steps = table(tmp_path, "steps/aria.lance") + ts = np.array([s["timestamp_ns"] for s in steps]) + assert ts[0] == T0 and (np.diff(ts) > 0).all() # absolute Unix ns + # Head pose = world->head (x drifts 1 mm/frame, z 1.5) @ head->camera (z + 0.1). + np.testing.assert_allclose( + steps[5]["head_pose"], [0.005, 0, 1.6, 1, 0, 0, 0], atol=1e-5 + ) + # Canonical MANO joints pass through, in the world frame. + assert steps[5]["left"]["keypoints"][:3] == pytest.approx( + [0.405, 0.2, 1.2], abs=1e-5 + ) + + from torchcodec.decoders import VideoDecoder + + for cam, rel in rec["video_path"]: + meta = VideoDecoder(str(tmp_path / "lance" / rel), seek_mode="exact").metadata + assert meta.num_frames == 60, cam + + +def test_rerun_skips_and_replace_keeps_one_copy(tmp_path): + assert run(tmp_path).status == "staged" + assert run(tmp_path).status == "skipped" + r = land(tmp_path, tmp_path / "eps" / "episode_0", replace=True) + assert r.status == "staged" + assert len(table(tmp_path, "recordings.lance")) == 1 + assert len(table(tmp_path, "tasks.lance")) == 2 + assert len(table(tmp_path, "annotations.lance")) == 2 + assert len(table(tmp_path, "steps/aria.lance")) == 60 + + +def test_replace_drops_the_vendors_removed_labels(tmp_path): + run(tmp_path) + write_mcap_episode(tmp_path / "eps", labels_in=None) # re-export without labels + r = land(tmp_path, tmp_path / "eps" / "episode_0", replace=True) + assert r.status == "staged" + assert table(tmp_path, "annotations.lance") == [] + # Without /task the episode info is one task over the whole recording. + (task,) = table(tmp_path, "tasks.lance") + assert task["task_name"] == "fold_clothes" and task["t_start_ns"] == T0 + + +def test_dropped_frame_becomes_an_invalid_filler_row(tmp_path): + r = run(tmp_path, drop_video_frame=11) + assert (r.num_frames, r.filler_frames) == (60, 1) + steps = table(tmp_path, "steps/aria.lance") + filler = steps[11] + assert not ( + filler["head_pose_valid"] or filler["left_valid"] or filler["right_valid"] + ) + assert filler["left"]["confidence"] == -1.0 + assert steps[12]["head_pose_valid"] is False # first frame after the gap + assert steps[13]["head_pose_valid"] is True + + +def test_single_arm_and_wrist_only_hands(tmp_path): + r = run(tmp_path, "human_right_arm") + assert r.valid_frac == {"head": 1.0, "left": 0.0, "right": 1.0} + r = run(tmp_path, episode_id="w", hands_wrist_only=True) + assert r.valid_frac["left"] == 1.0 + steps = [ + s + for s in table(tmp_path, "steps/aria.lance") + if s["recording_id"] == r.recording_id + ] + assert steps[0]["left"]["ee_pose"] == steps[0]["left"]["wrist_pose"] + + +def test_upright_rotation_turns_the_head_pose(tmp_path): + run(tmp_path, upright_rotation_cw_deg=90) + (s,) = [x for x in table(tmp_path, "steps/aria.lance")][5:6] + # Rz(-90 deg) as wxyz, as aria_utils.ROTATION_MATRIX for Aria's sideways sensor. + np.testing.assert_allclose( + np.abs(s["head_pose"][3:]), [np.sqrt(0.5), 0, 0, np.sqrt(0.5)], atol=1e-5 + ) + + +def test_failing_episode_is_rejected_and_nothing_committed(tmp_path): + r = run(tmp_path, imu_hz=100) + assert r.status == "rejected" + assert r.problems == ["[imu] /imu: 100 Hz, need >= 200"] + assert not (tmp_path / "lance" / "recordings.lance").exists() + + +def test_robot_episodes_wait_for_a_steps_schema(tmp_path): + r = run(tmp_path, "eva_bimanual") + assert r.status == "failed" and "no robot steps schema" in r.error + + +def test_hands_need_a_usable_head_pose(tmp_path): + # Hands given in the head frame are placed with the head pose: lost head -> invalid. + r = run(tmp_path, hand_frame="head", head_lost_from_s=1.0) + steps = table(tmp_path, "steps/aria.lance") + late = [s for s in steps if s["timestamp_ns"] >= T0 + 1_100_000_000] + assert late and not any(s["head_pose_valid"] or s["left_valid"] for s in late) + assert r.valid_frac["left"] == pytest.approx(r.valid_frac["head"], abs=0.02) + # Hands given in the world frame do not depend on it. + r = run(tmp_path, episode_id="w", head_lost_from_s=1.0) + assert r.valid_frac["left"] == 1.0 and r.valid_frac["head"] < 0.6 + + +def test_gaze_beyond_the_match_window_is_invalid(tmp_path): + run(tmp_path, gaze_until_s=1.5) + steps = table(tmp_path, "steps/aria.lance") + late = [s for s in steps if s["timestamp_ns"] > T0 + 1_600_000_000] + early = [s for s in steps if s["timestamp_ns"] < T0 + 1_400_000_000] + assert late and all(s["eye_gaze"] == [0, 0, 0] for s in late) + assert all(s["eye_gaze"] != [0, 0, 0] for s in early) + + +def test_packed_frame_past_the_validator_sample_fails(tmp_path): + # 120 messages at 60 fps; the validator samples the first 60, frames pack from 100. + r = run(tmp_path, video_fps=60, multi_frame_video=True, multi_frame_from=100) + assert r.problems == [] + assert r.status == "failed" and "holds 2 frames, not 1" in r.error + assert not (tmp_path / "lance" / "recordings.lance").exists() + + +def test_cli_keeps_going_past_a_failed_episode(tmp_path, capsys): + from egomimic.scripts.mcap_process.mcap_to_lance import main + + robot = write_mcap_episode(tmp_path / "eps", "eva_bimanual", "robot") + human = write_mcap_episode(tmp_path / "eps", "human_bimanual", "human") + args = [ + "--stage-to", + str(tmp_path / "staging"), + "--lance", + str(tmp_path / "lance"), + "--codec", + "libx264", + "--no-mano", + ] + assert main([str(robot), str(human), *args]) == 1 + out = capsys.readouterr().out + assert '"status": "failed"' in out and '"status": "staged"' in out + assert not (tmp_path / "lance" / "recordings.lance").exists() + commit_all(tmp_path) + assert len(table(tmp_path, "recordings.lance")) == 1 + + +def test_every_camera_becomes_a_video_with_its_calibration(tmp_path): + r = run(tmp_path, slam_camera=True) + assert r.videos == ["front_1", "front_1_480", "slam_left"] + (rec,) = table(tmp_path, "recordings.lance") + intr, extr = dict(rec["intrinsics"]), dict(rec["extrinsics"]) + # Native size and focal (64x48, f = 50), centred. + assert intr["slam_left"] == pytest.approx([50, 0, 31.5, 0, 50, 23.5, 0, 0, 1]) + # Both cameras hang off the head at the same offset: identity relative to front_1. + np.testing.assert_allclose( + np.reshape(extr["slam_left"], (4, 4)), np.eye(4), atol=1e-9 + ) + from torchcodec.decoders import VideoDecoder + + rel = dict(rec["video_path"])["slam_left"] + meta = VideoDecoder(str(tmp_path / "lance" / rel), seek_mode="exact").metadata + assert (meta.num_frames, meta.height, meta.width) == (60, 48, 64) + + +def test_non_aria_sources_keep_their_field_of_view(tmp_path): + run(tmp_path, metadata_override={"lab": "mecka"}) + (rec,) = table(tmp_path, "recordings.lance") + # f = 50 px at 64 px wide -> 500 px at 640 wide; Aria would use 266.5. + assert dict(rec["intrinsics"])["front_1_480"][0] == pytest.approx(500) + assert len(table(tmp_path, "steps/mecka.lance")) == 60 + + +def test_labels_lance_would_refuse_reject_before_any_commit(tmp_path, monkeypatch): + import egomimic.scripts.mcap_process.mcap_to_lance as m + from egomimic.rldb.mcap.validate import Result + + def late_point(topic, msg, t): # 10 ms before task 1 ends: Lance makes it 1 frame + if topic == "/subtask-annotation" and not msg.end_time_unix_ns: + return msg, T0 + 10**9 - 10_000_000 + + ep = write_mcap_episode(tmp_path / "eps", tamper=late_point) + # Even if a validator let it through, the converter checks before committing. + monkeypatch.setattr( + m, "validate_episode", lambda e: Result("human_bimanual", "x", []) + ) + r = land(tmp_path, ep, convert=m.convert_and_stage) + assert r.status == "rejected" and "not inside any task" in r.problems[0] + assert not (tmp_path / "staging").exists() + + +def test_labels_delivered_after_the_commit_are_picked_up(tmp_path): + r = run(tmp_path, labels_in=None) + assert (r.status, r.annotations_staged) == ("staged", 0) + write_mcap_episode(tmp_path / "eps", labels_without_tasks=True) # labels arrive + r = land(tmp_path, tmp_path / "eps" / "episode_0") + assert (r.status, r.annotations_staged) == ("skipped", 2) + assert len(table(tmp_path, "annotations.lance")) == 2 + (row,) = [ + c + for c in table(tmp_path, "changes.lance") + if c["lab"] == "annotations" and c["added"] + ] + assert row["operator"] == "acme delivery" diff --git a/tests/unit/test_mcap_validate.py b/tests/unit/test_mcap_validate.py new file mode 100644 index 0000000000..b4cca2cf93 --- /dev/null +++ b/tests/unit/test_mcap_validate.py @@ -0,0 +1,284 @@ +"""Vendor MCAP validator (egomimic/rldb/mcap) on tiny synthetic episodes. + +Valid episodes of every profile pass; each breakage in fixtures.mcap_episodes produces +exactly its expected problems and nothing else. +""" + +from pathlib import Path + +import pytest +from fixtures.mcap_episodes import T0, write_mcap_episode + +from egomimic.rldb.mcap import schemas, spec +from egomimic.rldb.mcap.validate import main, validate_episode + +PROTO = Path(schemas.__file__).parent / "proto" / "egoverse.proto" + + +def problems(tmp_path, embodiment="human_bimanual", **kw) -> list[tuple]: + r = validate_episode(write_mcap_episode(tmp_path, embodiment, **kw)) + return [(p.group, p.topic, p.message) for p in r.problems] + + +@pytest.mark.parametrize( + "embodiment, kw", + [ + ("human_bimanual", {}), + ("human_right_arm", {}), + ("human_left_arm", {}), + ("eva_bimanual", {}), + ("eva_left_arm", {}), + ("human_bimanual", {"labels_in": "episode"}), + ("eva_bimanual", {"labels_in": None}), + ("human_bimanual", {"hands_wrist_only": True}), + ("human_bimanual", {"drop_topic": "/main-camera-info-extra"}), + ("human_bimanual", {"labels_without_tasks": True}), + ("human_bimanual", {"labels_without_tasks": True, "label_at_last_frame": True}), + ( + "human_bimanual", + {"distortion_model": "fisheye624", "distortion_coeffs": (0.1,) * 12}, + ), + ], +) +def test_valid_episode_passes(tmp_path, embodiment, kw): + r = validate_episode(write_mcap_episode(tmp_path, embodiment, **kw)) + assert r.problems == [] + assert r.embodiment == embodiment + assert r.recording_id == spec.recording_id("aria", "acme-rec-0") + + +# (embodiment, breakage, [(group, topic, substring of the message), ...]) +BREAKAGES = [ + # 1. everything there + ("human_bimanual", {"drop_topic": "/imu-info"}, [("present", "/imu-info", "missing required topic")]), + ("human_bimanual", {"drop_metadata_key": "vendor"}, [("present", None, "'vendor' missing")]), + ("human_bimanual", {"drop_metadata_key": "site"}, [("present", None, "'site' missing")]), + ("human_bimanual", {"drop_metadata_key": "license"}, [("present", None, "'license' missing")]), + ("human_bimanual", {"drop_metadata_key": "vendor-recording-id"}, [("present", None, "'vendor-recording-id' missing")]), + ("human_bimanual", {"drop_metadata_key": "hand-joints"}, [("present", None, "'hand-joints' missing")]), + ("human_bimanual", {"bad_site_slug": True}, [("present", None, "not a lowercase slug")]), + ("human_bimanual", {"metadata_override": {"operator-id": "jane.doe@gatech.edu"}}, [("present", None, "not a SHA-256 hex digest")]), + ("human_bimanual", {"metadata_override": {"shutter": "rolling"}}, [("present", None, "expected 'global' or 'rolling:ms'")]), + ("human_bimanual", {"metadata_override": {"camera-timestamp-ref": "middle"}}, [("present", None, "'camera-timestamp-ref' = 'middle'")]), + ("human_bimanual", {"annotation_metadata": True}, [("present", None, "annotation.mcap carries metadata records")]), + ("human_bimanual", {"metadata_override": {"start-time-unix": "--1760000000000"}}, [("present", None, "is not an integer")]), + ("human_bimanual", {"metadata_override": {"end-time-unix": "1760000002000²"}}, [("present", None, "is not an integer")]), + ("human_bimanual", {"unknown_hand_joint": True}, [("present", None, "unknown joints ['thumb_base']")]), + ("human_bimanual", {"declared_embodiment": "robot_dog"}, [("present", None, "'embodiment' = 'robot_dog'")]), + ("human_bimanual", {"declared_embodiment": "human_right_arm"}, [("present", "/left-hand-keypoints", "does not fit embodiment")]), + ("human_bimanual", {"stray_annotation_topic": True}, [("present", "/instruction", "belong in annotation.mcap")]), + ("eva_bimanual", {"empty_arm_pose": True}, [ + ("present", "/right-arm-state", "'ee_pose' empty on 61/61"), + ("poses", "/right-arm-state", "not a finite rigid 4x4"), + ]), + # 2. video + ("human_bimanual", {"video_format": "vp9"}, [("video", "/main-camera", "format 'vp9'")]), + ("human_bimanual", {"skip_first_keyframe": True}, [("video", "/main-camera", "not a keyframe")]), + ("human_bimanual", {"calibration_width": 128}, [("video", "/main-camera", "calibration says 128x48")]), + ("human_bimanual", {"video_fps": 15}, [("video", "/main-camera", "15.0 fps")]), + ("human_bimanual", {"multi_frame_video": True}, [("video", "/main-camera", "exactly one frame")]), + ("eva_bimanual", {"multi_frame_video": True}, [("video", "/top-camera", "exactly one frame")]), + # 3. clocks + ("human_bimanual", {"backwards_head_pose": True}, [("clocks", "/head-pose", "backwards 1 times")]), + ("human_bimanual", {"hand_gap_s": 0.8}, [("clocks", "/left-hand-keypoints", "0.80 s gap")]), + ("human_bimanual", {"device_clock_imu_info": True}, [("clocks", "/imu-info", "before 2015")]), + ("human_bimanual", {"repeat_frame_time": True}, [("clocks", "/main-camera", "1 frames repeat a timestamp")]), + # 4. calibration + ("human_bimanual", {"drop_topic": "/main-camera-info"}, [("calibration", "/main-camera", "no /main-camera-info")]), + ("human_bimanual", {"unlink_main_camera": True}, [("calibration", "/main-camera", "no path to 'world'")]), + ("human_bimanual", {"bad_tf_rotation": True}, [("calibration", "/tf_static", "non-unit rotation")]), + ("human_bimanual", {"bad_upright_rotation": True}, [("calibration", "/main-camera", "upright_rotation_cw_deg 45")]), + ("human_bimanual", {"calibration_twice": True}, [("calibration", "/main-camera-info", "got 2 foxglove.CameraCalibration")]), + ("human_bimanual", {"principal_x": 500.0}, [("calibration", "/main-camera", "principal point (500, 24) is outside")]), + ("human_bimanual", {"distortion_model": "equidistant", "distortion_coeffs": (0.1,) * 4}, [("calibration", "/main-camera", "distortion_model 'equidistant'")]), + ("human_bimanual", {"distortion_model": "fisheye624", "distortion_coeffs": (0.1,) * 5}, [("calibration", "/main-camera", "needs 12 D coefficients, got 5")]), + # 5. imu + ("human_bimanual", {"imu_hz": 100}, [("imu", "/imu", "100 Hz")]), + ("human_bimanual", {"gravity_removed": True}, [("imu", "/imu", "gravity removed")]), + # 6. poses + ("human_bimanual", {"nan_head_pose": True}, [("poses", "/head-pose", "1/61 poses")]), + ("human_bimanual", {"zeroed_left_hand": True}, [("poses", "/left-hand-keypoints", "all-zero with valid=true")]), + ("human_bimanual", {"bad_hand_confidence": True}, [("poses", "/left-hand-keypoints", "confidence outside [0, 1]")]), + ("human_bimanual", {"bad_wrist_pose": True}, [("poses", "/left-hand-keypoints", "wrist poses")]), + ("human_bimanual", {"nan_eye_gaze": True}, [("poses", "/eye-gaze", "NaN/Inf")]), + ("eva_bimanual", {"gripper_out_of_range": True}, [("poses", "/right-ee-state", "outside [0, 1]")]), + ("eva_bimanual", {"bad_arm_pose": True}, [("poses", "/right-arm-state", "1/61 poses")]), + ("eva_bimanual", {"joint_names_mismatch": True}, [("poses", "/right-arm-state", "do not match joint_names")]), + # 7. tasks & labels + ("human_bimanual", {"task_past_end": True}, [("labels", "/task", "outside the recorded data")]), + ("eva_bimanual", {"label_outside_tasks": True}, [("labels", "/subtask-annotation", "inside one /task")]), + ("human_bimanual", {"label_crosses_tasks": True}, [("labels", "/subtask-annotation", "inside one /task")]), + ("human_bimanual", {"bad_annotation_type": True}, [("labels", "/subtask-annotation", "['summary']")]), + ("human_bimanual", {"labels_without_tasks": True, "label_past_video": True}, [("labels", "/subtask-annotation", "outside /main-camera's frames")]), +] # fmt: skip + + +@pytest.mark.parametrize( + "embodiment, kw, expected", + BREAKAGES, + ids=[f"{e}-{next(iter(kw))}={next(iter(kw.values()))}" for e, kw, _ in BREAKAGES], +) +def test_breakage_reports_exactly_its_problem(tmp_path, embodiment, kw, expected): + found = problems(tmp_path, embodiment, **kw) + assert len(found) == len(expected), found + for (group, topic, message), (g, t, sub) in zip(sorted(found), sorted(expected)): + assert (group, topic) == (g, t), found + assert sub in message, found + + +def test_cli_pass_fail_and_summary(tmp_path, capsys): + good = write_mcap_episode(tmp_path, episode_id="good") + bad = write_mcap_episode(tmp_path, episode_id="bad", imu_hz=100) + assert main([str(good)]) == 0 + assert main([str(good), str(bad)]) == 1 + out = capsys.readouterr().out + assert f"PASS {good}" in out + assert f"FAIL {bad}" in out and "[imu] /imu: 100 Hz" in out + assert "1/2 passed" in out + assert spec.recording_id("aria", "acme-rec-good") in out + + +def test_missing_and_unreadable_episodes(tmp_path): + r = validate_episode(tmp_path / "episode_nope") + assert [p.message for p in r.problems] == [ + f"no episode.mcap in {tmp_path / 'episode_nope'}" + ] + junk = tmp_path / "episode_junk" + junk.mkdir() + (junk / "episode.mcap").write_bytes(b"not an mcap") + (p,) = validate_episode(junk).problems + assert p.message.startswith("cannot read episode.mcap") + + +def test_unreadable_annotation_file(tmp_path): + ep = write_mcap_episode(tmp_path, labels_in=None) + (ep / spec.ANNOTATION_FILE).write_bytes(b"garbage") + (p,) = validate_episode(ep).problems + assert p.message.startswith("cannot read annotation.mcap") + + +def test_proto_file_matches_schemas(): + assert PROTO.read_text() == schemas.render_proto(), ( + "regenerate: python -m egomimic.rldb.mcap.schemas > " + str(PROTO) + ) + + +def test_embodiments_match_enum(): + from egomimic.rldb.embodiment.embodiment import EMBODIMENT + + assert set(spec.EMBODIMENTS) == {e.name.lower() for e in EMBODIMENT} + + +def test_mano_joints_match_human_finger_edges(): + # Canonical MANO: wrist, then 4 joints per finger, thumb to pinky. + assert len(spec.MANO_JOINTS) == 21 and spec.MANO_JOINTS[0] == "wrist" + assert spec.MANO_JOINTS[1:5] == ("thumb_cmc", "thumb_mcp", "thumb_ip", "thumb_tip") + assert spec.MANO_JOINTS[17:] == ("pinky_mcp", "pinky_pip", "pinky_dip", "pinky_tip") + + +def test_lance_ids_and_annotation_types_match_lance_schema(): + lance_schema = pytest.importorskip("egomimic.rldb.lance.schema") + assert spec.recording_id("aria", "x") == lance_schema.recording_id("aria", "x") + if hasattr(lance_schema, "ANNOTATION_TYPES"): + assert spec.ANNOTATION_TYPES == lance_schema.ANNOTATION_TYPES + + +def _late_hand_frame(topic, msg, t): + if topic == "/left-hand-keypoints" and t > T0 + 10**9: + msg.frame_id = "tracker" # message 0 is fine; later ones are not linked + return msg, t + + +def _head_in_odom(topic, msg, t): + if topic == "/head-pose": + msg.frame_id = "odom" + return msg, t + if ( + topic == "/tf_static" + ): # odom does reach world, but /head-pose must be T_world_head + x = msg.transforms.add(parent_frame_id="world", child_frame_id="odom") + x.rotation.w = 1.0 + return msg, t + + +def _nan_static_translation(topic, msg, t): + if topic == "/tf_static": + msg.transforms[0].translation.x = float("nan") + return msg, t + + +def _main_camera_fixed_in_world(topic, msg, t): + """A consistent camera frame, rigidly mounted to the world instead of the head.""" + if topic.startswith("/main-camera"): + msg.frame_id = "room_cam" + return msg, t + if topic == "/tf_static": + x = msg.transforms.add(parent_frame_id="world", child_frame_id="room_cam") + x.rotation.w = 1.0 + return msg, t + + +def _point_label_at_task_end(topic, msg, t): + if topic == "/subtask-annotation" and not msg.end_time_unix_ns: + return msg, T0 + 10**9 - 10_000_000 # 10 ms before task 1 ends: not a frame + + +@pytest.mark.parametrize( + "tamper, expected", + [ + ( + _late_hand_frame, + ("calibration", "/left-hand-keypoints", "'tracker' has no path"), + ), + (_head_in_odom, ("calibration", "/head-pose", "must be in frame 'world'")), + ( + _nan_static_translation, + ("calibration", "/tf_static", "1 transforms have a NaN / Inf translation"), + ), + ( + _main_camera_fixed_in_world, + ("calibration", "/main-camera", "no fixed path to 'head'"), + ), + ( + _point_label_at_task_end, + ("labels", "/subtask-annotation", "inside one /task"), + ), + ], +) +def test_tampered_episode_reports_exactly_its_problem(tmp_path, tamper, expected): + found = problems(tmp_path, tamper=tamper) + assert len(found) == 1, found + group, topic, message = found[0] + assert (group, topic) == expected[:2] and expected[2] in message, found + + +def test_calibration_naming_another_frame_than_the_video_is_caught(tmp_path): + """The calibration (and -info-extra) moved to `world`, the video frames not: the + camera has two frames, and the steps camera is no longer on the head.""" + + def calib_in_world(topic, msg, t): + if topic in ("/main-camera-info", "/main-camera-info-extra"): + msg.frame_id = "world" + return msg, t + + found = problems(tmp_path, tamper=calib_in_world) + assert ("calibration", "/main-camera") == found[0][:2] + assert any("video frames name ['main_cam']" in m or "video frames name" in m + for *_, m in found), found # fmt: skip + assert any("no fixed path to 'head'" in m for *_, m in found), found + + +def test_wrong_message_type_is_reported_not_raised(tmp_path): + from egomimic.rldb.mcap import schemas + + def imu_as_instruction(topic, msg, t): + if topic == "/imu": + return schemas.Instruction(data="not an imu"), t + + found = problems(tmp_path, tamper=imu_as_instruction) + assert ( + "present", + "/imu", + "type is 'egoverse.Instruction', expected 'egoverse.Imu'", + ) in found + assert any(g == "imu" and "could not check" in m for g, _, m in found), found diff --git a/tests/unit/test_model_dims.py b/tests/unit/test_model_dims.py index c68f993d2b..259ed791be 100644 --- a/tests/unit/test_model_dims.py +++ b/tests/unit/test_model_dims.py @@ -37,21 +37,29 @@ def _target(name): if _target(p.stem) == "egomimic.algo.hpt.HPT" and p.stem != "egobridge" ) -COTRAIN = {"eva_bimanual": (14, 14), "human_bimanual": (12, 12)} +COTRAIN = {"eva_bimanual": (14, 14), "human_bimanual": (20, 18)} # One shared head: human actions are zero-padded to its width. -SHARED = (14, {"eva_bimanual": (14, 14), "human_bimanual": (12, 14)}) +SHARED = (20, {"eva_bimanual": (14, 20), "human_bimanual": (20, 20)}) # {config: (action_width or None, {emb: (proprio, action)})}; None = no shared # head, so no action_width key. EXPECTED = { "hpt_bc_flow_eva": (None, {"eva_bimanual": (20, 20)}), - "hpt_bc_flow_aria": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_human": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_mecka": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_scale": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_keypoints_base": (None, {"human_bimanual": (138, 138)}), + # Human data defaults to the 144-D wrist-frame hand-keypoint action. + "hpt_bc_flow_aria": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_human": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_mecka": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_scale": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_human_cartesian": (None, {"human_bimanual": (20, 18)}), + "hpt_bc_keypoints_base": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_keypoints_wrist_300M": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_mecka_6d_300M": (None, {"human_bimanual": (20, 18)}), "hpt_cotrain_enc_dec_base": (None, COTRAIN), - "hpt_cotrain_flow_seperate_head": (None, COTRAIN), + # Separate heads: the human head follows the keypoint default. + "hpt_cotrain_flow_seperate_head": ( + None, + {"eva_bimanual": (14, 14), "human_bimanual": (144, 144)}, + ), "hpt_cotrain_flow_shared_head": SHARED, "hpt_cotrain_mecka_flow_shared_head": SHARED, "hpt_cotrain_scale_flow_shared_head": SHARED, @@ -71,6 +79,8 @@ def _width_leaves(rm): should read: ``.proprio``, ``.action`` or ``action_width``.""" for emb, stems in (rm.get("stem_specs") or {}).items(): for key, stem in (stems or {}).items(): + if stem is None: + continue # a child config dropping a base stem if key.startswith("state_"): path = f"stem_specs.{emb}.{key}.input_dim" yield path, stem["input_dim"], f"{emb}.proprio" diff --git a/tests/unit/test_mp4_ranges.py b/tests/unit/test_mp4_ranges.py new file mode 100644 index 0000000000..2c3d7bc227 --- /dev/null +++ b/tests/unit/test_mp4_ranges.py @@ -0,0 +1,60 @@ +import numpy as np +import pytest + +av = pytest.importorskip("av") + +from fixtures.videos import write_av1 # noqa: E402 + +from egomimic.rldb.lance import mp4_ranges # noqa: E402 + +N = 12 + + +def _reference(path): + with av.open(str(path)) as c: + packets = [(p.pos, p.size, p.is_keyframe) for p in c.demux(video=0) if p.size] + with av.open(str(path)) as c: + frames = [f.to_ndarray(format="rgb24") for f in c.decode(video=0)] + return packets, np.stack(frames) + + +@pytest.fixture(params=[True, False], ids=["faststart", "moov_at_end"]) +def mp4(tmp_path, request): + try: + write_av1(tmp_path / "v.mp4", N, faststart=request.param) + except (av.FFmpegError, ValueError) as e: + pytest.skip(f"no AV1 encoder: {e}") + return tmp_path / "v.mp4" + + +def test_index_matches_the_demuxer_and_reads_no_media(mp4): + data = mp4.read_bytes() + reads = [] + index = mp4_ranges.read_index( + lambda a, b: reads.append((a, b)) or data[a:b], first=64 + ) + packets, _ = _reference(mp4) + assert index.num_frames == N + assert index.offsets.tolist() == [p[0] for p in packets] + assert index.sizes.tolist() == [p[1] for p in packets] + keys = [i for i, p in enumerate(packets) if p[2]] + assert index.keyframe_of.tolist() == [ + max(k for k in keys if k <= i) for i in range(N) + ] + media = (min(index.offsets), max(index.offsets + index.sizes)) + # Past the first read, only box headers and the moov are fetched. + assert not any(a < media[1] and b > media[0] + 16 for a, b in reads[1:]) + + +def test_frames_decode_from_fetched_ranges_only(mp4): + data = mp4.read_bytes() + index = mp4_ranges.read_index(lambda a, b: data[a:b]) + _, want = _reference(mp4) + frames = [7, 0, 3, 3, N - 1] + rg = mp4_ranges.ranges(index, frames) + assert rg == sorted(rg) and all(b <= a2 for (_, b), (a2, _) in zip(rg, rg[1:])) + fetched = {r: data[r[0] : r[1]] for r in rg} + got = mp4_ranges.Decoder(index).frames(frames, fetched) + np.testing.assert_array_equal(got, want[frames]) + with pytest.raises(KeyError): + mp4_ranges.Decoder(index).frames([5], fetched) diff --git a/tests/unit/test_norm_cache.py b/tests/unit/test_norm_cache.py index 8af26e89b3..0ee527ff5a 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 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 ZarrWriter @@ -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) @@ -290,6 +291,14 @@ def test_train_writes_then_reuses_norm_cache(tmp_path, monkeypatch) -> None: lambda self, *a, **k: seen.append(k.get("precomputed_norm_path")) or orig(self, *a, **k), ) + sampled = [] + orig_sample = episode_norm_samples._sample_episodes + monkeypatch.setattr( + episode_norm_samples, + "_sample_episodes", + lambda root, leaves, *a, **k: sampled.append(sorted(leaves)) + or orig_sample(root, leaves, *a, **k), + ) train_hydra.train(_cfg(tmp_path, cache_dir, data, out, hashes)) files = list((cache_dir / "norm_stats" / "human_bimanual").glob("*.json")) @@ -305,11 +314,13 @@ def test_train_writes_then_reuses_norm_cache(tmp_path, monkeypatch) -> None: train_hydra.train(_cfg(tmp_path, cache_dir, data, out, hashes)) assert seen[2] is None assert len(list((cache_dir / "norm_stats" / "human_bimanual").glob("*.json"))) == 2 + assert sampled == [sorted(hashes), [hashes[0]]] # only the re-export resampled train_hydra.train(_cfg(tmp_path, cache_dir, data, out, hashes[:2])) assert ( len(list((cache_dir / "norm_stats" / "human_bimanual").glob("*.json"))) == 3 ) # new episode set, new key + assert len(sampled) == 2 # ...served from the per-episode samples def test_common_overrides_keep_the_cache_out_of_the_checkout(tmp_path) -> None: diff --git a/tests/unit/test_openpi_compat.py b/tests/unit/test_openpi_compat.py new file mode 100644 index 0000000000..951a90fa6a --- /dev/null +++ b/tests/unit/test_openpi_compat.py @@ -0,0 +1,115 @@ +"""openpi pins transformers 4.53 and ships a patch written against it; a venv +on transformers 5 needs these shims. The shims in egomimic/models/openpi_compat.py bridge the +gap and must disappear on their own once the fork's patch is ported forward.""" + +import pytest +import torch + +pytest.importorskip("openpi.models_pytorch.gemma_pytorch") + +from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel # noqa: E402 + +from egomimic.models import openpi_compat # noqa: E402 + + +@pytest.fixture(autouse=True) +def _restore_embed_image(): + original = PaliGemmaWithExpertModel.embed_image + cast = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + applied = openpi_compat._APPLIED + yield + PaliGemmaWithExpertModel.embed_image = original + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = cast + openpi_compat._APPLIED = applied + + +class _Outputs: + """What transformers 5's ``get_image_features`` hands back: the vision + output object with the projected features on ``pooler_output``.""" + + def __init__(self, features): + self.last_hidden_state = torch.zeros_like(features) + self.pooler_output = features + + +class _Prefix: + def __init__(self, result): + self._result = result + + class _Inner: + def __init__(self, result): + self._result = result + + def get_image_features(self, image): + return self._result + + @property + def model(self): + return self._Inner(self._result) + + +def _embed(result): + fake = object.__new__(PaliGemmaWithExpertModel) + fake.paligemma = _Prefix(result) + return PaliGemmaWithExpertModel.embed_image(fake, torch.zeros(1, 3, 224, 224)) + + +def test_a_wrapped_output_is_unwrapped_to_its_features(): + # Without this, openpi's embed_prefix reads `.shape` off the output object. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(_Outputs(features)), features) + + +def test_a_plain_tensor_passes_straight_through(): + # openpi's own patched transformers already returns the tensor; the shim + # must not change that behaviour. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(features), features) + + +def test_apply_is_a_no_op_the_second_time(caplog): + openpi_compat._APPLIED = False + openpi_compat.apply() + first = PaliGemmaWithExpertModel.embed_image + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is first + + +def test_the_shim_stands_down_on_the_transformers_openpi_pins(monkeypatch): + import transformers + + original = PaliGemmaWithExpertModel.embed_image + monkeypatch.setattr(transformers, "__version__", "4.53.2") + openpi_compat._APPLIED = False + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is original + + +def _siglip_embeddings(nested: bool) -> torch.nn.Module: + root = torch.nn.Module() + tower = torch.nn.Module() + parent = tower + if nested: + tower.vision_model = torch.nn.Module() + parent = tower.vision_model + parent.embeddings = torch.nn.Module() + parent.embeddings.patch_embedding = torch.nn.Conv2d(3, 4, 2) + parent.embeddings.position_embedding = torch.nn.Embedding(4, 4) + parent.encoder = torch.nn.Linear(4, 4) + root.vision_tower = tower + return root + + +@pytest.mark.parametrize("nested", [True, False]) +def test_the_siglip_embeddings_stay_float32(nested): + """openpi's keep-list matches the nested 4.x names only; with the shim the + flattened 5.x names are kept too, and everything else still goes bf16.""" + openpi_compat._patch_fp32_keep_list() + module = _siglip_embeddings(nested) + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params(module, "bfloat16") + dtypes = {n: p.dtype for n, p in module.named_parameters()} + for name, dtype in dtypes.items(): + expected = torch.float32 if "embeddings" in name else torch.bfloat16 + assert dtype == expected, name diff --git a/tests/unit/test_opscale_ladder.py b/tests/unit/test_opscale_ladder.py new file mode 100644 index 0000000000..f058274a06 --- /dev/null +++ b/tests/unit/test_opscale_ladder.py @@ -0,0 +1,299 @@ +"""The freeform fold-clothes OPERATOR-SCALING ladder holds its shape. + +mecka_fold_freeform_opscale{1,2,4} vary ONE thing: how many operators the +training data comes from. Each arm trains on ALL the annotated data its +operators have, minus the shared held-out val episodes. Everything that would +confound the comparison -- the seen-op val, the unseen-op val, the pipeline -- +must be identical across the three rungs, and the operator sets must be nested. + +The train split is an OPERATOR FILTER, not an episode pin, so these tests drive +the real DatasetFilter with synthetic rows rather than counting hashes. +""" + +from pathlib import Path + +import hydra +import pytest +import yaml +from hydra import compose, initialize_config_module +from omegaconf import OmegaConf + +import egomimic.hydra_configs as _cfg_pkg +from egomimic.rldb.filters import DatasetFilter + +DATA_DIR = Path(_cfg_pkg.__file__).parent / "data" +ARMS = (1, 2, 4) +VARIANTS = ("6d", "hpt_keypoints") + +# The 4 train operators, ranked by ANNOTATED hours (SQL app.episodes, 2026-09-17). +# Arm k trains on the first k of these. NOT the same order as raw hours. +OPERATORS = [ + "6964d38183a9fdf2d8663792", + "68e0b875c33a1abcb8fc55b9", + "695d09fc83a9fdf2d84d9c11", + "69439ad5a2a8b5aee76602f5", +] +# The only episodes withheld from training: 3 episodes of operator #1. +SEEN_VAL = [ + "696bcbe0e13d37dd4293f58b", + "696bcbe86377d0871b6a9253", + "696bcc71ed8ed158122ec794", +] +# Episodes per arm as measured 2026-09-17. The filter is LIVE, so these are +# documentation, not a guarantee; the metric budget is derived from the resolved +# split at runtime (`auto`), so nothing has to be kept in step with them. +TRAIN_EPISODES = {1: 9, 2: 24, 4: 48} +SEEN_VAL_EPISODES = 3 +UNSEEN_OPERATORS = 31 +SPLITS = ("train_datasets", "valid_datasets", "unseen_op_valid_datasets") + +# A row that clears every non-operator lambda: mecka, freeform, fold-clothes, +# annotated, and not one of the withheld episodes. +BASE_ROW = { + "lab": "mecka", + "task": "folding_clothes", + "zarr_processed_path": "s3://rldb/processed_v3/mecka/freeform/abc.zarr", + "episode_hash": "an_episode_not_in_the_val_set", + "segments": [{"label": "fold shirt", "start_seconds": 0, "end_seconds": 9}], +} + + +def _data(name: str): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose(config_name="train_zarr_cartesian", overrides=[f"data={name}"]) + cfg._set_flag("allow_objects", True) + OmegaConf.resolve(cfg.data) + return cfg.data + + +def _name(k: int, variant: str) -> str: + return f"mecka_fold_freeform_opscale{k}_{variant}" + + +def _filter(data, split: str) -> DatasetFilter: + node = data[split].human_bimanual + return DatasetFilter( + filter_lambdas=list(node.filters.filter_lambdas), + episode_hashes=list(node.filters.get("episode_hashes") or []), + ) + + +@pytest.fixture(scope="module") +def ladder(): + return {(k, v): _data(_name(k, v)) for k in ARMS for v in VARIANTS} + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_accepts_exactly_the_first_k_operators(ladder, k, variant): + """Nesting, stated as behaviour: arm k takes operators 1..k and no others, + so each rung only ADDS an operator to the one below it.""" + f = _filter(ladder[(k, variant)], "train_datasets") + for i, op in enumerate(OPERATORS): + row = {**BASE_ROW, "operator": op} + assert f.matches(row) is (i < k), f"operator #{i + 1} in the k={k} arm" + assert not f.matches({**BASE_ROW, "operator": "some_other_operator"}) + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_withholds_the_seen_val_episodes(ladder, k, variant): + """The seen-op val is the ONLY thing kept out of training, so the exclusion + must be present in every arm -- including k=1, whose single operator owns + all three of those episodes.""" + f = _filter(ladder[(k, variant)], "train_datasets") + for h in SEEN_VAL: + row = {**BASE_ROW, "operator": OPERATORS[0], "episode_hash": h} + assert not f.matches(row), f"k={k}: train accepted seen-val episode {h}" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_seen_val_pin_matches_what_train_excludes(ladder, k, variant): + """The val pin and the train exclusion are written separately; if they drift, + an episode is either trained on AND validated, or dropped entirely.""" + data = ladder[(k, variant)] + pin = list(data.valid_datasets.human_bimanual.filters.episode_hashes) + assert sorted(pin) == sorted(SEEN_VAL) + assert len(pin) == SEEN_VAL_EPISODES + train = _filter(data, "train_datasets") + # every pinned episode is refused by train, and a non-pinned one is not + assert all( + not train.matches({**BASE_ROW, "operator": OPERATORS[0], "episode_hash": h}) + for h in pin + ) + assert train.matches({**BASE_ROW, "operator": OPERATORS[0]}) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_val_sets_are_identical_across_the_ladder(ladder, variant): + seen = { + k: list( + ladder[(k, variant)].valid_datasets.human_bimanual.filters.episode_hashes + ) + for k in ARMS + } + unseen = {k: list(ladder[(k, variant)].held_out_operators) for k in ARMS} + assert len({tuple(v) for v in seen.values()}) == 1, "seen-op val differs by rung" + assert len({tuple(v) for v in unseen.values()}) == 1, "unseen-op list differs" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_operators_are_never_in_the_unseen_split(ladder, k, variant): + data = ladder[(k, variant)] + assert len(data.held_out_operators) == UNSEEN_OPERATORS + assert not (set(OPERATORS) & set(data.held_out_operators)) + # and the unseen filter refuses a train operator outright + unseen = _filter(data, "unseen_op_valid_datasets") + assert not unseen.matches({**BASE_ROW, "operator": OPERATORS[0]}) + assert unseen.matches({**BASE_ROW, "operator": data.held_out_operators[0]}) + + +def test_unseen_operators_match_the_opsplit_config(): + """Reusing mecka_fold_freeform_opsplit_6d's held-out list verbatim is what + keeps this ladder's unseen_op_valid/ numbers comparable with the existing + freeform opsplit checkpoints. Do not re-cut it.""" + raw = yaml.safe_load( + (DATA_DIR / "mecka_fold_freeform_opsplit_6d.yaml").read_text(encoding="utf-8") + ) + assert sorted(_data(_name(4, "6d")).held_out_operators) == sorted( + raw["held_out_operators"] + ) + + +# --- language annotations ------------------------------------------------- +# 112 of the 462 freeform episodes have segments = NULL -> an empty `annotations` +# track -> _annotation_text_for_frame returns [] for every frame -> _build_prompts +# substitutes `default_prompt` ("" unless overridden). Every split of every rung +# must filter those out, or a third of training runs on an empty instruction. + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize("split", SPLITS) +def test_every_split_requires_language_annotations(ladder, k, variant, split): + data = ladder[(k, variant)] + f = _filter(data, split) + operator = ( + data.held_out_operators[0] + if split == "unseen_op_valid_datasets" + else OPERATORS[0] + ) + good = {**BASE_ROW, "operator": operator} + if split == "valid_datasets": + good["episode_hash"] = SEEN_VAL[0] # that split is pinned by hash + for bad in (None, []): + assert not f.matches( + {**good, "segments": bad} + ), f"{split}: accepted an episode with segments={bad!r}" + assert f.matches(good), f"{split}: guard rejects an annotated row" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_annotation_key_is_wired(ladder, k, variant): + """Without annotation_key the keymap emits no annotation entry at all, and + the prompt is `default_prompt` everywhere regardless of the guard.""" + km = ladder[(k, variant)].train_datasets.human_bimanual.resolver.key_map + assert km.annotation_key == "annotations" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_no_rung_asks_for_proprio_history_alone(ladder, k, variant): + """proprio_history is half a setting: the keymap decides how many steps a + sample carries and the model's stem how many tokens consume them, and + HPTModel.stem_process raises when they disagree. Every model in the repo is + history_len 1, so a rung shipping K > 1 is a rung nothing can train -- which + is what happened when the ladder outlived the one model that wanted K=4.""" + km = ladder[(k, variant)].train_datasets.human_bimanual.resolver.key_map + assert km.get("proprio_history", 1) == 1 + + +# --- pipeline and loaders ------------------------------------------------- + + +@pytest.mark.parametrize("k", ARMS) +def test_keypoint_twin_shares_its_parents_split(ladder, k): + """The HPT twin may change only the keymap/transform; if its filters drift + from the 6d parent the two model families stop being comparable.""" + for split in SPLITS: + assert list( + ladder[(k, "hpt_keypoints")][split].human_bimanual.filters.filter_lambdas + ) == list(ladder[(k, "6d")][split].human_bimanual.filters.filter_lambdas), split + + +@pytest.mark.parametrize("k", ARMS) +def test_pipeline_per_variant(ladder, k): + for variant, keymap, mode, pad in ( + ("6d", "cartesian", "cartesian_wristframe_6d", True), + ("hpt_keypoints", "keypoints", "keypoints_wristframe_6d", False), + ): + res = ladder[(k, variant)].train_datasets.human_bimanual.resolver + assert res.key_map.keymap_mode == keymap + assert res.transform_list.mode == mode + assert res.transform_list.pad_proprio_gripper is pad + # processed_v3 mecka zarrs have the LEFT wrist frame double-mirrored. + # Instantiated, not just read off the node: a renamed kwarg would keep + # the assertion green while the config could no longer be built. + assert res.transform_list.fix_left_wrist_convention is True + assert hydra.utils.instantiate(res.transform_list) + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_filters_are_the_whole_split(ladder, k, variant): + """`mode: total` everywhere: the filters ARE the split, so no seeded episode + split may run on top of them and quietly drop a fraction of the data.""" + data = ladder[(k, variant)] + for split in SPLITS: + assert data[split].human_bimanual.mode == "total", split + assert data.valid_prefix == "seen_op_valid" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_metric_budget_is_derived_from_the_resolved_split(ladder, k, variant): + """Every head is `auto`, so K follows the live filters; a hard-coded K would + push K x episodes past the limit_val_batches window as SQL grows. At the + documented counts `auto` reproduces floor(5120 / episodes).""" + from egomimic import trainHydra as th + + data = ladder[(k, variant)] + assert dict(data.metric_frames_per_episode) == { + "valid": "auto", + "unseen_op_valid": "auto", + "train_viz": "auto", + } + cfg = OmegaConf.create( + { + "data": {"metric_frames_per_episode": dict(data.metric_frames_per_episode)}, + "trainer": {"limit_val_batches": 80, "devices": 1}, + } + ) + bs = data.valid_dataloader_params.human_bimanual.batch_size + for head, n in (("train_viz", TRAIN_EPISODES[k]), ("valid", SEEN_VAL_EPISODES)): + kk = th._metric_frames_per_episode(cfg, head, n_episodes=n, batch_size=bs) + assert kk == 80 * bs // n + assert kk * n <= 80 * bs + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_video_pins_belong_to_their_own_split(ladder, k, variant): + """A video pin is taken out of that head's already-resolved split; a pin from + elsewhere would make the loader re-resolve or render nothing.""" + data = ladder[(k, variant)] + vids = data.video_episodes + assert set(vids) == {"valid", "unseen_op_valid", "train_viz"} + assert set(vids.valid) <= set(SEEN_VAL) + assert len(vids.unseen_op_valid) == 1 and len(vids.train_viz) == 1 + # the train_viz pin is an operator #1 episode, so every arm -- including + # k=1 -- must actually accept it, which is why the children can inherit it + train = _filter(data, "train_datasets") + assert train.matches( + {**BASE_ROW, "operator": OPERATORS[0], "episode_hash": vids.train_viz[0]} + ), f"k={k}: the inherited train_viz video pin is not in this arm's split" diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py index dd035fc212..bcca6ae48a 100644 --- a/tests/unit/test_packaging.py +++ b/tests/unit/test_packaging.py @@ -77,3 +77,8 @@ def test_wheel_ships_robot_resources(wheel_files): # XML relative to egomimic.__file__. assert "egomimic/resources/model_x5.xml" in wheel_files assert "egomimic/resources/model_arx.urdf" in wheel_files + + +def test_wheel_ships_mcap_proto(wheel_files): + # The .proto handed to vendors; egomimic.rldb.mcap.schemas renders the same text. + assert "egomimic/rldb/mcap/proto/egoverse.proto" in wheel_files diff --git a/tests/unit/test_paligemma_init.py b/tests/unit/test_paligemma_init.py new file mode 100644 index 0000000000..61e9f0d3ed --- /dev/null +++ b/tests/unit/test_paligemma_init.py @@ -0,0 +1,457 @@ +"""PaliGemma-only init: the VLM prefix comes from the public PaliGemma release +and everything pi0.5 adds on top starts from its PyTorch init. + +No network and no 3B model: the loader is exercised against a miniature module +tree, and against a tiny real ``PaliGemmaForConditionalGeneration`` built by +whichever transformers is installed (the vision tree differs across 4.x and +5.x), with safetensors written on the fly; the download resolution is driven +through stubs. The one thing that needs the real thing -- +that the released checkpoint covers openpi's prefix exactly -- is asserted by +``load_paligemma_weights`` itself (it raises on any missing or unexpected +tensor), so a rename upstream cannot pass silently. +""" + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +from egomimic.models import paligemma_init as pgi + +CANONICAL = pgi.CANONICAL_REPO + + +# ---------------------------------------------------------------- key remapping + + +# Spelled exactly as google/paligemma-3b-pt-224 ships them. +@pytest.mark.parametrize( + ("checkpoint_key", "module_name"), + [ + ( + "language_model.model.layers.0.self_attn.q_proj.weight", + "model.language_model.layers.0.self_attn.q_proj.weight", + ), + ( + "language_model.model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ), + ("language_model.model.norm.weight", "model.language_model.norm.weight"), + ("language_model.lm_head.weight", "lm_head.weight"), + # transformers >= 5 flattened SiglipVisionModel: the `.vision_model.` hop + # in the checkpoint has no counterpart in the module tree. + ( + "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + "model.vision_tower.encoder.layers.26.mlp.fc2.bias", + ), + ( + "vision_tower.vision_model.embeddings.patch_embedding.weight", + "model.vision_tower.embeddings.patch_embedding.weight", + ), + ( + "multi_modal_projector.linear.bias", + "model.multi_modal_projector.linear.bias", + ), + ], +) +def test_checkpoint_keys_map_onto_the_module_tree(checkpoint_key, module_name): + assert pgi._remap_key(checkpoint_key) == module_name + + +def test_unknown_prefix_is_left_alone(): + assert pgi._remap_key("action_in_proj.weight") == "action_in_proj.weight" + + +# ------------------------------------------------------------- vocab truncation + + +def test_padded_vocabulary_is_truncated_to_the_model_width(): + # The release pads 257152 -> 257216; the extra rows are + padding. + tensor, param = torch.arange(12.0).reshape(6, 2), torch.zeros(4, 2) + fitted = pgi._fit("model.language_model.embed_tokens.weight", tensor, param) + assert fitted.shape == param.shape + assert torch.equal(fitted, tensor[:4]) # padding is at the END + + +def test_truncation_is_not_applied_to_other_tensors(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit( + "model.vision_tower.encoder.layers.0.mlp.fc1.weight", + torch.zeros(6, 2), + torch.zeros(4, 2), + ) + + +def test_a_narrower_checkpoint_is_never_padded_out(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit("lm_head.weight", torch.zeros(3, 2), torch.zeros(4, 2)) + + +# ------------------------------------------------------------------ init source + + +def test_pi05_base_wins_when_only_it_is_set(): + assert pgi.select_init_source("/ckpt/pi05_base_pytorch", None) == "pi05_base" + + +def test_paligemma_when_the_base_checkpoint_is_off(): + assert pgi.select_init_source(None, CANONICAL) == "paligemma" + + +def test_neither_is_reported_as_no_weights(): + assert pgi.select_init_source(None, None) == "none" + + +def test_both_at_once_is_refused(): + # The pi0.5 base already contains a trained PaliGemma; silently letting one + # win would make the run's provenance unreadable from its config. + with pytest.raises(ValueError, match="both set"): + pgi.select_init_source("/ckpt/pi05_base_pytorch", CANONICAL) + + +# ------------------------------------------------------------ the loader itself + + +class _Prefix(nn.Module): + """Shaped like the part of ``PaliGemmaForConditionalGeneration`` the loader + walks: a ``model.`` subtree, a tied ``lm_head`` and a derived buffer.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, vocab=4, width=3): + super().__init__() + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.embed_tokens = nn.Embedding(vocab, width) + self.model.vision_tower = nn.Module() + self.model.vision_tower.post_layernorm = nn.Linear(width, width) + self.lm_head = nn.Linear(width, vocab, bias=False) + self.lm_head.weight = self.model.language_model.embed_tokens.weight # tied + # config-derived, absent from every released checkpoint + self.model.language_model.register_buffer("inv_freq", torch.ones(width)) + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.paligemma_with_expert = nn.Module() + self.paligemma_with_expert.paligemma = _Prefix() + self.action_in_proj = nn.Linear(3, 3) + + +def _write_checkpoint(tmp_path, tensors, shards=2): + items = list(tensors.items()) + size = -(-len(items) // shards) + for i in range(shards): + chunk = dict(items[i * size : (i + 1) * size]) + if chunk: + save_file(chunk, str(tmp_path / f"model-{i}.safetensors")) + return str(tmp_path) + + +def _full_checkpoint(vocab=4, width=3): + torch.manual_seed(1) + return { + "language_model.model.embed_tokens.weight": torch.randn(vocab, width), + "vision_tower.vision_model.post_layernorm.weight": torch.randn(width, width), + "vision_tower.vision_model.post_layernorm.bias": torch.randn(width), + } + + +def test_the_prefix_loads_and_nothing_else_moves(tmp_path): + model = _Model() + before = model.action_in_proj.weight.detach().clone() + tensors = _full_checkpoint() + directory = _write_checkpoint(tmp_path, tensors) + + assert pgi.load_paligemma_weights(model, directory) == directory + + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.model.language_model.embed_tokens.weight, + tensors["language_model.model.embed_tokens.weight"], + ) + assert torch.equal( + prefix.model.vision_tower.post_layernorm.bias, + tensors["vision_tower.vision_model.post_layernorm.bias"], + ) + # everything outside the prefix is untouched + assert torch.equal(model.action_in_proj.weight, before) + + +def test_the_tied_head_follows_the_embedding(tmp_path): + # The checkpoint carries no lm_head; tying must still leave it consistent. + model = _Model() + tensors = _full_checkpoint() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, tensors)) + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.lm_head.weight, prefix.model.language_model.embed_tokens.weight + ) + + +def test_a_derived_buffer_is_not_required(tmp_path): + model = _Model() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, _full_checkpoint())) + assert torch.equal( + model.paligemma_with_expert.paligemma.model.language_model.inv_freq, + torch.ones(3), + ) + + +def test_a_partial_checkpoint_is_refused(tmp_path): + # Silently leaving part of the VLM at its random init is the failure this + # whole module exists to avoid. + tensors = _full_checkpoint() + del tensors["vision_tower.vision_model.post_layernorm.weight"] + with pytest.raises(RuntimeError, match="missing"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_unknown_tensor_is_refused(tmp_path): + tensors = _full_checkpoint() + tensors["vision_tower.vision_model.renamed_upstream.weight"] = torch.zeros(3) + with pytest.raises(RuntimeError, match="unexpected"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_empty_directory_is_refused(tmp_path): + with pytest.raises(FileNotFoundError, match="safetensors"): + pgi.load_paligemma_weights(_Model(), str(tmp_path)) + + +def test_a_padded_vocabulary_checkpoint_loads(tmp_path): + # The real release is 64 rows wider than openpi's table. + tensors = _full_checkpoint() + wide = torch.randn(6, 3) + tensors["language_model.model.embed_tokens.weight"] = wide + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +# ------------------------------------------------------- download / mirror path + + +def _tiny_paligemma(): + from transformers import PaliGemmaConfig, PaliGemmaForConditionalGeneration + + config = PaliGemmaConfig( + text_config={ + "model_type": "gemma", + "num_hidden_layers": 1, + "hidden_size": 32, + "intermediate_size": 32, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 32, + "vocab_size": 300, + }, + vision_config={ + "num_hidden_layers": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_attention_heads": 1, + "image_size": 28, + "patch_size": 14, + "projection_dim": 32, + }, + image_token_index=299, + projection_dim=32, + ) + return PaliGemmaForConditionalGeneration(config) + + +def _release_key(name: str) -> str: + """The module-tree name as google/paligemma-3b-pt-224 spells it.""" + for old, new in ( + ("model.language_model.", "language_model.model."), + ("lm_head.", "language_model.lm_head."), + ("model.vision_tower.vision_model.", "vision_tower.vision_model."), + ("model.vision_tower.", "vision_tower.vision_model."), + ("model.multi_modal_projector.", "multi_modal_projector."), + ): + if name.startswith(old): + return new + name[len(old) :] + return name + + +def test_a_real_paligemma_loads_on_this_transformers(tmp_path): + """The miniature tree above is spelled the 5.x way; a real model is spelled + whichever way the installed transformers spells it, which is how a 4.53 + venv came to fail 21 SigLIP tensors while the key tests all passed.""" + source = _tiny_paligemma() + released = { + _release_key(n): p.detach().clone().contiguous() + for n, p in source.named_parameters() + } + target = _tiny_paligemma() + model = nn.Module() + model.paligemma_with_expert = nn.Module() + model.paligemma_with_expert.paligemma = target + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, released)) + for name, param in target.named_parameters(): + assert torch.equal(param, dict(source.named_parameters())[name]), name + + +def test_a_nested_vision_tree_is_found(): + own = {"model.vision_tower.vision_model.post_layernorm.bias": None} + key = "vision_tower.vision_model.post_layernorm.bias" + assert ( + pgi._module_name(key, own) + == "model.vision_tower.vision_model.post_layernorm.bias" + ) + assert pgi._module_name(key, {}) == "model.vision_tower.post_layernorm.bias" + + +def test_only_the_shards_are_downloaded(): + assert pgi._ALLOW_PATTERNS == ["*.safetensors"] + + +def test_a_local_directory_is_used_as_is(tmp_path): + assert pgi.resolve_paligemma_dir(str(tmp_path)) == str(tmp_path) + + +def _gated_error(repo_id): + """The 401 huggingface_hub raises for a license-gated repo.""" + import requests + from huggingface_hub.errors import GatedRepoError + + response = requests.Response() + response.status_code = 401 + return GatedRepoError(f"401 for {repo_id}", response=response) + + +def _stub_hub(monkeypatch, *, gated, digests, downloads): + import huggingface_hub + + def snapshot_download(repo_id, **_): + if repo_id in gated: + raise _gated_error(repo_id) + return downloads[repo_id] + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + monkeypatch.setattr(pgi, "_shard_digests", lambda repo: digests[repo]) + monkeypatch.setattr(pgi, "_verify_shards", lambda directory, expected: None) + + +def test_the_canonical_repo_is_preferred(monkeypatch, tmp_path): + _stub_hub( + monkeypatch, gated=set(), digests={}, downloads={CANONICAL: str(tmp_path)} + ) + assert pgi.resolve_paligemma_dir(CANONICAL) == str(tmp_path) + + +def test_a_gated_canonical_repo_falls_back_to_a_matching_mirror(monkeypatch, tmp_path): + good = pgi.MIRROR_REPOS[0] + digests = { + CANONICAL: {"model.safetensors": "abc"}, + good: {"model.safetensors": "abc"}, + } + digests.update({m: {"model.safetensors": "abc"} for m in pgi.MIRROR_REPOS}) + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + assert pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) == str(tmp_path) + + +def test_the_mirror_is_opt_in(monkeypatch, tmp_path): + from huggingface_hub.errors import GatedRepoError + + digests = {r: {"model.safetensors": "abc"} for r in (CANONICAL, *pgi.MIRROR_REPOS)} + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir(CANONICAL) + + +def test_a_rate_limit_is_not_a_reason_to_use_a_mirror(monkeypatch): + import huggingface_hub + import requests + from huggingface_hub.errors import HfHubHTTPError + + response = requests.Response() + response.status_code = 429 + + def snapshot_download(repo_id, **_): + raise HfHubHTTPError("429", response=response) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + with pytest.raises(HfHubHTTPError): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_a_mirror_whose_shards_differ_is_rejected(monkeypatch, tmp_path): + # A mirror is transport only: different bytes means it is not PaliGemma. + digests = {CANONICAL: {"model.safetensors": "abc"}} + digests.update({m: {"model.safetensors": "TAMPERED"} for m in pgi.MIRROR_REPOS}) + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(RuntimeError, match="no verified mirror"): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_the_mirror_fallback_can_be_switched_off(monkeypatch, tmp_path): + from huggingface_hub.errors import GatedRepoError + + _stub_hub(monkeypatch, gated={CANONICAL}, digests={}, downloads={}) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=False) + + +def test_a_non_canonical_repo_never_falls_back(monkeypatch, tmp_path): + # Only the gated canonical repo has verified stand-ins; anything else the + # user names must fail loudly rather than silently load other weights. + from huggingface_hub.errors import GatedRepoError + + _stub_hub( + monkeypatch, gated={"someone/private-paligemma"}, digests={}, downloads={} + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir("someone/private-paligemma", allow_mirror=True) + + +# -------------------------------------------------------------------- the config + + +def test_the_overlay_flips_the_init_source_and_keeps_its_parent(compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d_paligemma_init"] + ) + model = cfg.model.robomimic_model + assert model.config.pytorch_weight_path is None + assert model.config.paligemma_weight_path == CANONICAL + assert ( + pgi.select_init_source( + model.config.pytorch_weight_path, model.config.paligemma_weight_path + ) + == "paligemma" + ) + # the two keys are all the overlay touches: pi0.5_bc_mecka_6d's own choices + # (cartesian 6D action, no Embodiment prompt block) must survive it + assert model.ac_keys.human_bimanual == "actions_cartesian" + assert model.embodiment_label is False + assert model.config.model.action_dim == 32 + assert ( + model.action_converters.rules.HUMAN_BIMANUAL._target_ + == "egomimic.utils.action_utils.HumanBimanualCartesian6D" + ) + + +def test_the_base_pi_config_still_starts_from_pi05_base(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d"]) + config = cfg.model.robomimic_model.config + assert config.paligemma_weight_path is None + assert ( + pgi.select_init_source(config.pytorch_weight_path, config.paligemma_weight_path) + == "pi05_base" + ) diff --git a/tests/unit/test_per_sample_augs.py b/tests/unit/test_per_sample_augs.py new file mode 100644 index 0000000000..9aa678de15 --- /dev/null +++ b/tests/unit/test_per_sample_augs.py @@ -0,0 +1,133 @@ +"""Per-sample image augmentation (ideas 0.10). + +torchvision's random transforms draw one parameter set per call, so the old +inline ``self.train_image_augs(_data)`` on the stacked (B, 3, H, W) batch gave +every image in the batch the same jitter. These tests pin the per-sample +behaviour of ``PerSampleAugs`` and of ``HPT._apply_image_augs``, and pin that +the deterministic eval path is unchanged. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torchvision.transforms import ColorJitter, Compose, Normalize + +from egomimic.models.image_augs import PerSampleAugs + +MEAN = [0.485, 0.456, 0.406] +STD = [0.229, 0.224, 0.225] + + +def _identical_batch(n=64, mean=0.5): + """``n`` copies of one image whose per-image mean is exactly ``mean``.""" + return torch.full((n, 3, 16, 16), float(mean)) + + +def _per_image_means(x): + return x.flatten(1).mean(dim=1) + + +def test_per_sample_augs_jitters_each_image_independently(): + """A batch of identical images must come out with per-image brightness + spread; the same ColorJitter called on the whole batch must not.""" + torch.manual_seed(0) + jitter = ColorJitter(brightness=0.5) + images = _identical_batch() + + per_sample_std = _per_image_means(PerSampleAugs(jitter)(images)).std().item() + batched_std = _per_image_means(jitter(images)).std().item() + print( + f"per-image brightness std: per-sample={per_sample_std:.6f} " + f"batched={batched_std:.6f}" + ) + + # brightness=0.5 scales each image by U(0.5, 1.5), so the per-image means + # of a mean-0.5 batch are U(0.25, 0.75): std = 0.5 / sqrt(12) ~ 0.144. + # 0.01 is ~14x below that and ~0 for the batched call. + assert per_sample_std > 0.01 + assert batched_std < 1e-6 + + +def test_per_sample_augs_matches_batched_for_deterministic_transform(): + """A deterministic Compose through PerSampleAugs is exactly the batched + result, so wrapping a Normalize-only aug list is a no-op.""" + augs = Compose([Normalize(mean=MEAN, std=STD)]) + images = torch.rand(8, 3, 16, 16) + torch.testing.assert_close( + PerSampleAugs(augs)(images), augs(images), rtol=0, atol=0 + ) + + +def test_per_sample_augs_preserves_shape_dtype_device_and_grad(): + images = torch.rand(4, 3, 16, 16, dtype=torch.float32, requires_grad=True) + out = PerSampleAugs(Compose([Normalize(mean=MEAN, std=STD)]))(images) + + assert out.shape == images.shape + assert out.dtype == images.dtype + assert out.device == images.device + out.sum().backward() + assert images.grad is not None and torch.isfinite(images.grad).all() + + +def _hpt_stub(train_image_augs=None, eval_image_augs=None, training=True): + """A bare HPT carrying only what ``_apply_image_augs`` reads; constructing + a real HPT would build the trunk and stems.""" + from egomimic.algo.hpt import HPT + + algo = HPT.__new__(HPT) + algo.nets = nn.ModuleDict() + algo.nets.train(training) + algo.encoders = {"front_img_1": {}} + algo.train_image_augs = ( + PerSampleAugs(train_image_augs) if train_image_augs is not None else None + ) + algo.eval_image_augs = eval_image_augs + return algo + + +def test_apply_image_augs_is_per_sample_in_train_mode(): + torch.manual_seed(0) + algo = _hpt_stub(train_image_augs=ColorJitter(brightness=0.5), training=True) + out = algo._apply_image_augs(_identical_batch(), "front_img_1") + assert _per_image_means(out).std().item() > 0.01 + + +def test_apply_image_augs_eval_mode_equals_batched_normalize(): + eval_augs = Compose([Normalize(mean=MEAN, std=STD)]) + algo = _hpt_stub( + train_image_augs=ColorJitter(brightness=0.5), + eval_image_augs=eval_augs, + training=False, + ) + images = torch.rand(8, 3, 16, 16) + torch.testing.assert_close( + algo._apply_image_augs(images, "front_img_1"), eval_augs(images), rtol=0, atol=0 + ) + + +def test_apply_image_augs_passes_through_uncoded_camera(): + """A camera with no encoder is left alone in either mode.""" + images = torch.rand(4, 3, 16, 16) + for training in (True, False): + algo = _hpt_stub( + train_image_augs=ColorJitter(brightness=0.5), + eval_image_augs=Compose([Normalize(mean=MEAN, std=STD)]), + training=training, + ) + out = algo._apply_image_augs(images, "wrist_img_1") + torch.testing.assert_close(out, images, rtol=0, atol=0) + + +def test_missing_train_augs_pass_through(): + """__init__ allows train_image_augs=None; the helper must not call it.""" + from egomimic.algo.hpt import HPT + + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict({"policy": torch.nn.Identity()}) + algo.nets.train() + algo.train_image_augs = None + algo.eval_image_augs = None + algo.encoders = {"front_img_1": object()} + images = torch.rand(2, 3, 8, 8) + assert algo._apply_image_augs(images, "front_img_1") is images diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 89c40d9279..3e03f36acc 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -5,23 +5,30 @@ 20D 6D, with gripper) and human bimanual (12D ypr / 18D 6D, no gripper) layouts, plus the proprio ee_pose: the 6D transform modes convert the proprio too (a single pose vector, same per-arm layout as one action row), and the 6D revert -lists convert it back before the eef-frame revert reads it. +lists convert it back before the eef-frame revert reads it. Also the pieces +that ride along with the 6D pipeline: the rotation-free bounds check, the +mecka left-wrist fix, embodiment aliasing, the sampler fallback and the eval +helpers. """ import math +from types import SimpleNamespace import numpy as np import pytest import torch -from egomimic.rldb.zarr.action_chunk_transforms import ( +from egomimic.rldb.transforms import ( CartesianRot6DToYPR, CartesianYPRToRot6D, ) from egomimic.utils.action_utils import ( BaseActionConverter, - RobotBimanualCartesianEuler, + HumanBimanualKeypoints, + RobotBimanualCartesian6D, + pad_to_width, ) +from egomimic.utils.pose_utils import _ypr_to_rot6d def _eva_ypr_chunk(T: int = 5) -> np.ndarray: @@ -88,7 +95,7 @@ def test_transform_preserves_tensor_type(): def test_robot_bimanual_norm_6d_pack_round_trips(): - converter = RobotBimanualCartesianEuler() + converter = RobotBimanualCartesian6D() six6d = torch.from_numpy( CartesianYPRToRot6D().transform({"actions_cartesian": _eva_ypr_chunk()})[ "actions_cartesian" @@ -101,6 +108,29 @@ def test_robot_bimanual_norm_6d_pack_round_trips(): torch.testing.assert_close(decoded, six6d, atol=1e-6, rtol=1e-6) +def test_keypoint_converter_is_identity_and_pads_to_model_width(): + # The 144-D keypoint action is packed identity-first; PI pads it (and the + # cartesian 32-slot layout) up to model.action_dim and the decode slices + # the native width back out of the wider vector. + conv = HumanBimanualKeypoints() + kp = torch.randn(2, 4, 144) + packed = conv.to32_norm_6d(kp) + torch.testing.assert_close(packed, kp) + torch.testing.assert_close(conv.from32_norm_6d(pad_to_width(packed, 160)), kp) + with pytest.raises(ValueError, match="expected 144-dim"): + conv.to32_norm_6d(torch.zeros(1, 1, 138)) + + eva32 = RobotBimanualCartesian6D().to32_norm_6d(torch.ones(1, 3, 20)) + wide = pad_to_width(eva32, 144) + assert wide.shape[-1] == 144 and torch.all(wide[..., 32:] == 0) + torch.testing.assert_close( + RobotBimanualCartesian6D().from32_norm_6d(wide)[..., :20], + torch.ones(1, 3, 20), + ) + with pytest.raises(ValueError, match="action_dim is 32"): + pad_to_width(kp, 32) + + @pytest.mark.parametrize( "chunk_fn,ypr_dim,six_dim", [(_eva_ypr_chunk, 14, 20), (_aria_ypr_chunk, 12, 18)], @@ -184,6 +214,222 @@ def test_6d_revert_lists_revert_proprio(): }, f"{build.__name__} must revert both action and proprio to ypr" +def _bounds_check_dataset(key: str, width: int): + """Minimal MultiDataset shell exposing _check_bounds with +-1 quantile + bounds on ``key`` for embodiment 0.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, -1.0, dtype=np.float32), + "quantile_99": np.full(width, 1.0, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + return md + + +@pytest.mark.parametrize("key", ["actions_cartesian", "observations.state.ee_pose"]) +@pytest.mark.parametrize("width,rot_idx,xyz_idx", [(14, 3, 0), (20, 4, 0), (18, 5, 9)]) +def test_bounds_check_ignores_rotation_channels(key, width, rot_idx, xyz_idx): + # Rotation channels (Euler wraps at +-pi; 6D columns are ~[-1, 1]) must be + # excluded from quantile bounds checking, while translation/gripper + # channels are still checked and NaN/Inf still rejects the full vector. + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + + arr[2, rot_idx] = 50.0 # far outside +-1, but a rotation channel + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + + bad = arr.copy() + bad[2, xyz_idx] = 50.0 # translation channel out of bounds -> violation + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + nan = arr.copy() + nan[2, rot_idx] = np.nan # NaN anywhere (even rotation) -> violation + assert md._check_bounds({"embodiment": 0, key: nan}, None, 0, "ep") is not None + + +@pytest.mark.parametrize("key", ["actions_keypoints", "observations.state.keypoints"]) +@pytest.mark.parametrize("width,rot_idx,kp_idx", [(138, 4, 10), (144, 78, 90)]) +def test_bounds_check_ignores_wrist_rotation_in_keypoint_layout( + key, width, rot_idx, kp_idx +): + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + arr[1, rot_idx] = 50.0 # wrist rotation channel: not bounds-checked + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[1, kp_idx] = 50.0 # a keypoint coordinate: checked + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +def test_bounds_check_full_vector_for_other_keys(): + # Keys without a known layout (or unrecognized widths) keep the + # full-vector check. + md = _bounds_check_dataset("some_other_key", 20) + arr = np.zeros((5, 20), dtype=np.float32) + arr[2, 4] = 50.0 + assert ( + md._check_bounds({"embodiment": 0, "some_other_key": arr}, None, 0, "ep") + is not None + ) + + md16 = _bounds_check_dataset("actions_cartesian", 16) + arr16 = np.zeros((5, 16), dtype=np.float32) + arr16[2, 4] = 50.0 + assert ( + md16._check_bounds({"embodiment": 0, "actions_cartesian": arr16}, None, 0, "ep") + is not None + ) + + +def test_constant_cell_exemption_follows_norm_mode(): + """A channel the zscore normalizer still divides by must still be + bounds-checked, even though its quantile range is zero: 0 in 99.6 % of + samples and 1.0 in the rest gives q99 - q1 = 0 but std = 0.063.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + stats = { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": np.zeros(18, dtype=np.float32), + "mean": np.full(18, 0.004, dtype=np.float32), + "std": np.full(18, 0.063, dtype=np.float32), + } + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 50.0 # corrupt value in a quantile-collapsed channel + + for mode, rejected in (("quantile", False), ("zscore", True)): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = mode + md.norm_stats = {0: {"actions_cartesian": stats}} + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + got = md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep" + ) + assert (got is not None) is rejected, mode + + +def test_bounds_check_tolerates_roundoff_on_collapsed_bounds(): + """Wrist-frame t=0 cells have bounds [0, 0]; values there must not reject + (normalize() maps such constant cells to 0), while other cells are still + checked.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + q99 = np.ones(18, dtype=np.float32) + q99[0] = 0.0 # channel 0 collapsed to [0, 0] + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + "actions_cartesian": { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": q99, + } + } + } + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 1e-9 # a roundoff-scale xyz value at a [0, 0] bound + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 0] = 1e-3 # off-convention offset at the constant cell: admitted + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 1] = 50.0 # a real violation on a regular cell is still caught + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is not None + ) + + +def test_bounds_check_warns_once_on_stat_shape_mismatch(caplog): + md = _bounds_check_dataset("actions_cartesian", 18) + arr = np.zeros((5, 20), dtype=np.float32) # stats are 18-wide + with caplog.at_level("WARNING"): + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 1, "ep") + is None + ) + msgs = [r.message for r in caplog.records if "bounds check skipped" in r.message] + assert len(msgs) == 1, msgs + + +def test_rotate_local_frame_flips_left_wrist_convention(): + # Right-multiplying by Rz(180 deg) must flip the pose's own x/y axes, keep + # z (knuckle-forward) and the position, skip zero-quat padding rows, and + # handle both (7,) poses and (T, 7) chunks. + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.transforms import RotateLocalFrame + + rng = np.random.default_rng(4) + q = R.random(3, random_state=5) + chunk = np.zeros((4, 7)) + chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) + chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded + + t = RotateLocalFrame(keys=["k"]) + out = t.transform({"k": chunk.copy()})["k"] + + np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged + np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched + R_old = q.as_matrix() + R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() + np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip + np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip + np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept + + single = t.transform({"k": chunk[0].copy()})["k"] + np.testing.assert_allclose(single, out[0], atol=1e-12) + + +def test_fix_left_wrist_convention_flag_prepends_correction(): + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.transforms import RotateLocalFrame + + tl = Human.get_transform_list( + "cartesian_wristframe_6d", stride=1, fix_left_wrist_convention=True + ) + assert isinstance(tl[0], RotateLocalFrame) + assert set(tl[0].keys) == {"left.action_ee_pose", "left.obs_ee_pose"} + # default off: other vendors' data must be untouched + tl_off = Human.get_transform_list("cartesian_wristframe_6d", stride=1) + assert not isinstance(tl_off[0], RotateLocalFrame) + # keypoints modes correct the wrist_pose keys their frames are built on + # (the converter wrote the same double-mirrored rotation into both), plus + # the ee_pose keys when those are built too. + tl_kp = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True + ) + assert isinstance(tl_kp[0], RotateLocalFrame) + assert set(tl_kp[0].keys) == {"left.action_wrist_pose", "left.obs_wrist_pose"} + tl_kp_ee = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True, include_ee_pose=True + ) + assert set(tl_kp_ee[0].keys) == { + "left.action_wrist_pose", + "left.obs_wrist_pose", + "left.action_ee_pose", + "left.obs_ee_pose", + } + + def test_vendor_embodiment_names_collapse_to_human(): # Mirror episodes written by the vendor-split registry carry names like # MECKA_BIMANUAL in their zarr metadata; locally all human demo data is @@ -221,7 +467,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, ) @@ -243,7 +489,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, @@ -252,30 +498,432 @@ def test_human_6d_reverts_unpad_proprio(): assert any(isinstance(t, UnpadGripperZeros) for t in build()), build.__name__ -def test_rotate_local_frame_flips_left_wrist_convention(): - # Right-multiplying by Rz(180 deg) must flip the pose's own x/y axes, keep - # z (knuckle-forward) and the position, skip zero-quat padding rows, and - # handle both (7,) poses and (T, 7) chunks. +def _fallback_dataset(): + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.index_map = [("bad", i) for i in range(10)] + [("good", i) for i in range(1000)] + md._global_indices_by_dataset = { + "bad": list(range(10)), + "good": list(range(10, 1010)), + } + return md + + +def test_fallback_widens_to_global_after_local_attempts(): + # A wholly-bad episode must not exhaust the sampler: retries stay inside + # the failing episode for GLOBAL_FALLBACK_ATTEMPTS, then widen to the full + # index space. + md = _fallback_dataset() + attempts = None + seen_local, seen_global = set(), set() + idx = 0 + for _ in range(md.GLOBAL_FALLBACK_ATTEMPTS): + idx, attempts = md._next_after_failure(0, "bad", attempts, reason="r") + seen_local.add(md.index_map[idx][0]) + assert seen_local == {"bad"}, "early retries must stay within the episode" + + for _ in range(50): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="r") + seen_global.add(md.index_map[idx][0]) + assert "good" in seen_global, "post-threshold retries must sample globally" + + +def test_fallback_cap_raises(): + # The cap is a systemic-failure detector: everything it can reach has + # already been drawn uniformly from the whole index space and failed. + md = _fallback_dataset() + attempts, idx, n = None, 0, 0 + with pytest.raises(RuntimeError, match="consecutive bad samples"): + for n in range(1, md.MAX_FALLBACK_ATTEMPTS + 2): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="boom") + assert n == md.MAX_FALLBACK_ATTEMPTS + + +def test_video_fps_is_the_source_rate_regardless_of_world_size(): + """Val loaders are CombinedLoaders that Lightning never shards, so rank 0 + renders every frame: playback must stay at the source rate (an earlier + world_size division made 8-GPU videos 8x slow motion).""" + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub.__new__(_Stub) + for world in (1, 2, 8): + ev.trainer = SimpleNamespace(world_size=world) + assert ev._video_fps() == 30, (world, ev._video_fps()) + ev.trainer = SimpleNamespace(world_size=None) + assert ev._video_fps() == 30 + assert ev._video_fps(source_fps=10) == 10 + + +def test_viz_gate_follows_lightning_epoch_convention(): + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub(viz_every_n_epochs=20) + # Lightning validates at current_epoch 19, 39, ... for check_val_every_n_epoch=20 + for epoch, expected in [(0, False), (19, True), (20, False), (39, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch) + assert ev._should_viz() is expected, epoch + ev.viz_every_n_epochs = 0 + assert ev._should_viz() is False + + +def test_viz_gate_always_renders_the_final_epoch(): + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub(viz_every_n_epochs=100) + # trainer=debug: 4 epochs, val at current_epoch 1 and 3 -> video only at 3 + for epoch, expected in [(1, False), (3, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch, max_epochs=4) + assert ev._should_viz() is expected, epoch + # eval mode validates once at epoch 0 with max_epochs=1 + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=1) + assert ev._should_viz() is True + # unbounded runs (max_epochs None / -1) only follow the interval + for max_epochs in (None, -1): + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=max_epochs) + assert ev._should_viz() is False + # viz_every_n_epochs=0 still turns video off entirely + ev.viz_every_n_epochs = 0 + ev.trainer = SimpleNamespace(current_epoch=3, max_epochs=4) + assert ev._should_viz() is False + + +def test_frechet_helper(): + from egomimic.utils.metrics import frechet_gaussian_over_time + + B, T, D = 4, 10, 6 + torch.manual_seed(0) + pred = torch.randn(B, T, D) + fd = frechet_gaussian_over_time(pred, pred.clone()) + assert fd.shape == (B,) + assert fd.max() < 1e-2, "identical distributions must score ~0" + fd_shift = frechet_gaussian_over_time(pred + 5.0, pred) + assert (fd_shift > fd).all(), "mean shift must increase the distance" + + +def test_reverse_kl_is_gone(): + import egomimic.utils.metrics as metrics + + assert not hasattr(metrics, "reverse_kl_from_samples") + + +def test_train_viz_wrapper_prefixes_and_forwards(): + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + from egomimic.eval.eval_video import EvalVideo + + class _Base(EvalVideo): + def __init__(self): + super().__init__(viz_func={}, transform_lists={}, viz_every_n_epochs=7) + self.seen = None + + def compute_metrics_and_viz(self, batch, do_viz=True): + self.seen = (self.model, do_viz) + return {"Valid/x": 1.0}, {} + + base = _Base() + tv = TrainVizEvalVideo(base) + algo = object() + tv.model = algo # property setter forwards to base too + metrics, _ = tv.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"train_viz/Valid/x"}, metrics + assert base.seen == (algo, False) + assert tv.viz_every_n_epochs == 7, "wrapper inherits the base viz gate" + tv.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert base.trainer is tv.trainer + assert tv.video_dir().endswith("videos_train_viz") + + op = TrainVizEvalVideo(base, prefix="unseen_op_valid") + metrics, _ = op.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"unseen_op_valid/Valid/x"}, metrics + op.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert op.video_dir().endswith("videos_unseen_op_valid") + + +def test_dtw_distance_matches_bruteforce_and_tolerates_shift(): + from egomimic.utils.metrics import dtw_distance + + def _dtw_ref(x, y): + t1, t2 = len(x), len(y) + cost = np.linalg.norm(x[:, None, :] - y[None, :, :], axis=-1) + acc = np.full((t1 + 1, t2 + 1), np.inf) + acc[0, 0] = 0.0 + for i in range(1, t1 + 1): + for j in range(1, t2 + 1): + acc[i, j] = cost[i - 1, j - 1] + min( + acc[i - 1, j], acc[i, j - 1], acc[i - 1, j - 1] + ) + return acc[t1, t2] + + torch.manual_seed(3) + pred = torch.randn(3, 9, 4) + tgt = torch.randn(3, 9, 4) + got = dtw_distance(pred, tgt, normalize=False) + for b in range(3): + ref = _dtw_ref(pred[b].numpy(), tgt[b].numpy()) + assert abs(got[b].item() - ref) < 1e-4, (b, got[b].item(), ref) + + # identical trajectories -> 0 + assert dtw_distance(pred, pred.clone()).max() < 1e-6 + + # a time-shifted copy of the same smooth trajectory: DTW must forgive the + # shift (score far below the paired per-step distance of the shifted pair) + t = torch.linspace(0, 6.28, 40) + traj = torch.stack([torch.sin(t), torch.cos(t)], dim=-1)[None] # (1, 40, 2) + shifted = torch.roll(traj, shifts=3, dims=1) + dtw_shift = dtw_distance(traj, shifted, normalize=False).item() + paired = (traj - shifted).norm(dim=-1).sum().item() + assert dtw_shift < 0.5 * paired, (dtw_shift, paired) + + +def test_split_mse_is_stateless_and_matches_manual(): + from egomimic.eval.action_metrics import _paired_mse, _split_mse + + rng = np.random.default_rng(21) + pred = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + gt = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + xyz_idx = [0, 1, 2, 9, 10, 11] + rot_idx = [i for i in range(18) if i not in xyz_idx] + xyz, rot = _split_mse(pred, gt) + torch.testing.assert_close( + xyz, (pred[..., xyz_idx] - gt[..., xyz_idx]).pow(2).mean() + ) + torch.testing.assert_close( + rot, (pred[..., rot_idx] - gt[..., rot_idx]).pow(2).mean() + ) + # stateless: a second, unrelated call is unaffected by the first + a, b = torch.zeros(2, 3, 18), torch.ones(2, 3, 18) + torch.testing.assert_close(_split_mse(a, b)[0], torch.tensor(1.0)) + torch.testing.assert_close(_paired_mse(a, b), torch.tensor(1.0)) + assert _split_mse(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) == (None, None) + + +def test_rot_geodesic_error_matches_angle_and_survives_gimbal_lock(): from scipy.spatial.transform import Rotation as R - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + from egomimic.eval.action_metrics import _rot_geodesic_error + + rng = np.random.default_rng(22) + ypr = rng.uniform(-1.0, 1.0, size=(6, 5, 3)) + theta = 0.3 + # rotate every pose by theta about a random axis -> geodesic error == theta + axes = rng.normal(size=(6, 5, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + Rg = R.from_euler("ZYX", ypr.reshape(-1, 3)) + Rp = R.from_rotvec(theta * axes.reshape(-1, 3)) * Rg + ypr_p = Rp.as_euler("ZYX").reshape(6, 5, 3) + + # 12-dim ypr layout (both arms the same pose) + gt12 = torch.from_numpy(np.concatenate([ypr, ypr], -1)).float() + gt12 = torch.cat( + [torch.zeros(6, 5, 3), gt12[..., :3], torch.zeros(6, 5, 3), gt12[..., 3:]], -1 + ) + pr12 = torch.from_numpy(np.concatenate([ypr_p, ypr_p], -1)).float() + pr12 = torch.cat( + [torch.zeros(6, 5, 3), pr12[..., :3], torch.zeros(6, 5, 3), pr12[..., 3:]], -1 + ) + assert abs(_rot_geodesic_error(pr12, gt12).item() - theta) < 1e-4 + assert _rot_geodesic_error(gt12, gt12).item() < 1e-5 + + # 18-dim 6D layout, same rotations -> same answer + def to18(y): + six = _ypr_to_rot6d(y) + arm = np.concatenate([np.zeros(y.shape[:-1] + (3,)), six], -1) + return torch.from_numpy(np.concatenate([arm, arm], -1)).float() + + assert abs(_rot_geodesic_error(to18(ypr_p), to18(ypr)).item() - theta) < 1e-4 + assert _rot_geodesic_error(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) is None + + # Gimbal lock: pitch ~ pi/2, pred = gt rotated 0.004 rad about local y. + # A ypr MSE (even wrap-aware) explodes because yaw and roll trade off; + # the geodesic error reports the true 0.004. + gt = np.array([0.3, np.pi / 2 - 0.002, 0.2]) + Rgt = R.from_euler("ZYX", gt) + Rpr = Rgt * R.from_rotvec([0.0, 0.004, 0.0]) + pr = Rpr.as_euler("ZYX") + v_gt = torch.tensor([[0, 0, 0, *gt, 0, 0, 0, *gt]], dtype=torch.float32) + v_pr = torch.tensor([[0, 0, 0, *pr, 0, 0, 0, *pr]], dtype=torch.float32) + geo = _rot_geodesic_error(v_pr, v_gt).item() + assert abs(geo - 0.004) < 1e-3, geo + euler_mse = (v_pr - v_gt).pow(2).mean().item() + assert euler_mse > 0.1, euler_mse # the Euler metric is fooled here + + +# ------------------------------------------------ PIEvalVideo metric contract +def _rand_6d_chunks(seed, B=3, T=6): + """(pred, gt) native 18D human chunks with a small, known perturbation, + plus the per-(batch, arm) rotation matrices so a test can re-express both + in another frame. 6D columns are built directly from rotation matrices + (col 0 | col 1), matching ``_reconstruct_R_from_cols``.""" + from scipy.spatial.transform import Rotation as R - rng = np.random.default_rng(4) - q = R.random(3, random_state=5) - chunk = np.zeros((4, 7)) - chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) - chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded + rng = np.random.default_rng(seed) + xyz_g = rng.uniform(-1.0, 1.0, size=(B, T, 2, 3)) + R_g = R.random(B * T * 2, random_state=seed).as_matrix().reshape(B, T, 2, 3, 3) + xyz_p = xyz_g + 0.05 * rng.normal(size=xyz_g.shape) + R_p = ( + ( + R.from_rotvec(0.1 * rng.normal(size=(B * T * 2, 3))) + * R.from_matrix(R_g.reshape(-1, 3, 3)) + ) + .as_matrix() + .reshape(B, T, 2, 3, 3) + ) + return (xyz_p, R_p), (xyz_g, R_g) + + +def _chunk18(xyz, Rm): + six = np.concatenate([Rm[..., :, 0], Rm[..., :, 1]], -1) # (B,T,2,6) + arm = np.concatenate([xyz, six], -1) # (B,T,2,9) + return torch.from_numpy(arm.reshape(arm.shape[0], arm.shape[1], 18)).float() + + +def _run_pi_eval(pred, gt, ac_key="actions_cartesian"): + from egomimic.eval.eval_pi import PIEvalVideo + from egomimic.rldb.embodiment.embodiment import EMBODIMENT, get_embodiment + + eid = EMBODIMENT.HUMAN_BIMANUAL.value + name = get_embodiment(eid).lower() + + class _NormStats: + def unnormalize(self, batch, embodiment_id): + return batch + + class _Algo: + ac_keys = {eid: ac_key} + norm_stats = _NormStats() + + def forward_eval(self, batch): + return {f"{name}_loss": torch.tensor(0.5), f"{name}_{ac_key}": pred} + + ev = PIEvalVideo(viz_func={}, transform_lists={}, viz_every_n_epochs=1) + ev.model = _Algo() + batch = {eid: {ac_key: gt, "embodiment": torch.tensor([eid])}} + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=False) + return f"Valid/{name}_{ac_key}_", f"Valid/{name}_loss", metrics, images + + +def test_pi_eval_metric_set_and_identity(): + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(40) + gt = _chunk18(xyz_g, R_g) + prefix, loss_key, metrics, images = _run_pi_eval(gt.clone(), gt) + expected = { + prefix + k + for k in ( + "xyz_paired_mse_avg", + "xyz_final_mse_avg", + "rot6d_paired_mse_avg", + "rot_err_deg_avg", + "rot_err_deg_final", + "xyz_dtw_avg", + "xyz_frechet_gauss_avg", + "xyz_l2_early", + "xyz_l2_mid", + "xyz_l2_late", + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + assert images == {}, "do_viz=False must not render" + assert float(metrics[prefix + "xyz_paired_mse_avg"]) < 1e-10 + assert float(metrics[prefix + "rot_err_deg_avg"]) < 1e-3 + assert float(metrics[prefix + "xyz_dtw_avg"]) < 1e-6 + assert float(metrics[prefix + "xyz_l2_late"]) < 1e-5 + assert float(metrics[prefix + "xyz_frechet_gauss_avg"]) < 1e-2 + + +def test_pi_eval_metrics_are_invariant_to_a_shared_rigid_transform(): + # The evaluator logs no cam-frame twins of the native metrics: the revert + # applies the SAME rigid transform (the observed wrist pose) to prediction + # and gt, and every kept metric is invariant to that. Re-express pred and + # gt in a random frame per (batch, arm) and check. + from scipy.spatial.transform import Rotation as R - t = RotateLocalFrame(keys=["k"]) - out = t.transform({"k": chunk.copy()})["k"] + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(41) + B = xyz_g.shape[0] + R_w = R.random(B * 2, random_state=7).as_matrix().reshape(B, 1, 2, 3, 3) + t_w = np.random.default_rng(8).uniform(-2, 2, size=(B, 1, 2, 3)) - np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged - np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched - R_old = q.as_matrix() - R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() - np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip - np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip - np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept + def move(xyz, Rm): + return np.einsum("btaij,btaj->btai", R_w, xyz) + t_w, R_w @ Rm - single = t.transform({"k": chunk[0].copy()})["k"] - np.testing.assert_allclose(single, out[0], atol=1e-12) + prefix, _, native, _ = _run_pi_eval(_chunk18(xyz_p, R_p), _chunk18(xyz_g, R_g)) + prefix, _, moved, _ = _run_pi_eval( + _chunk18(*move(xyz_p, R_p)), _chunk18(*move(xyz_g, R_g)) + ) + for key, atol in ( + ("xyz_paired_mse_avg", 1e-5), + ("xyz_final_mse_avg", 1e-5), + ("rot_err_deg_avg", 1e-3), + ("rot_err_deg_final", 1e-3), + ("xyz_dtw_avg", 1e-4), + ("xyz_frechet_gauss_avg", 1e-3), + ("xyz_l2_early", 1e-5), + ("xyz_l2_late", 1e-5), + ): + a, b = float(native[prefix + key]), float(moved[prefix + key]) + assert abs(a - b) < atol, (key, a, b) + # sanity: the perturbation is visible in the headline metrics + assert float(native[prefix + "rot_err_deg_avg"]) > 1.0 + assert float(native[prefix + "xyz_paired_mse_avg"]) > 1e-4 + + +def test_pi_eval_keypoint_metrics(): + from scipy.spatial.transform import Rotation as R + + rng = np.random.default_rng(50) + B, T = 2, 5 + wrist_xyz = rng.uniform(-1, 1, size=(B, T, 2, 3)) + ypr = rng.uniform(-1, 1, size=(B, T, 2, 3)) + kp = rng.uniform(-0.1, 0.1, size=(B, T, 2, 63)) + + def chunk(wx, y, k): + hand = np.concatenate([wx, _ypr_to_rot6d(y), k], -1) # (B,T,2,72) + return torch.from_numpy(hand.reshape(B, T, 144)).float() + + gt = chunk(wrist_xyz, ypr, kp) + theta = 0.2 + axes = rng.normal(size=(B * T * 2, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + ypr_p = ( + (R.from_rotvec(theta * axes) * R.from_euler("ZYX", ypr.reshape(-1, 3))) + .as_euler("ZYX") + .reshape(B, T, 2, 3) + ) + pred = chunk(wrist_xyz + 0.1, ypr_p, kp + 0.02) + prefix, loss_key, metrics, _ = _run_pi_eval(pred, gt, ac_key="actions_keypoints") + expected = { + prefix + k + for k in ( + "kp_l2_avg", + "kp_l2_final", + "wrist_xyz_paired_mse_avg", + "wrist_xyz_final_mse_avg", + "wrist_rot_err_deg_avg", + "wrist_rot_err_deg_final", + "kp_dtw_avg", + *( + f"{m}_l2_{seg}" + for m in ("kp", "wrist_xyz") + for seg in ("early", "mid", "late") + ), + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + # every keypoint is off by 0.02 in each axis -> L2 = 0.02 * sqrt(3) + assert abs(float(metrics[prefix + "kp_l2_avg"]) - 0.02 * math.sqrt(3)) < 1e-5 + assert abs(float(metrics[prefix + "wrist_xyz_paired_mse_avg"]) - 0.01) < 1e-6 + assert ( + abs(float(metrics[prefix + "wrist_rot_err_deg_avg"]) - math.degrees(theta)) + < 1e-2 + ) diff --git a/tests/unit/test_populate_probe.py b/tests/unit/test_populate_probe.py new file mode 100644 index 0000000000..491dd2470f --- /dev/null +++ b/tests/unit/test_populate_probe.py @@ -0,0 +1,60 @@ +"""populate_from_datasets probes one episode per key_map/transform group.""" + +from fixtures.synthetic_episodes import write_episode + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, + ZarrDataset, +) + + +def _dataset(tmp_path, n=4): + for i in range(n): + write_episode(tmp_path, "eva", seed=i) + resolver = LocalEpisodeResolver( + tmp_path, + key_map=Eva.get_keymap(keymap_mode="cartesian"), + transform_list=Eva.get_transform_list(mode="cartesian_wristframe_6d"), + ) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _count_probes(monkeypatch, fail_first=0): + calls = [] + orig = ZarrDataset.__getitem__ + + def spy(self, idx, *a, **k): + calls.append(self.episode_path) + if len(calls) <= fail_first: + raise RuntimeError("bad episode") + return orig(self, idx, *a, **k) + + monkeypatch.setattr(ZarrDataset, "__getitem__", spy) + return calls + + +def test_one_probe_per_group_gives_the_per_leaf_inventory(tmp_path, monkeypatch): + ds = _dataset(tmp_path) + expected = { + k for k in ds[0] if k in ds.datasets[next(iter(ds.datasets))].key_map + } | {"actions_cartesian", "observations.state.ee_pose"} + calls = _count_probes(monkeypatch) + stats = MultiDataset(state={}) + stats.populate_from_datasets({"eva_bimanual": ds}) + assert len(calls) == 1 + (emb,) = stats.key_types + assert expected <= set(stats.key_types[emb]) + assert stats.key_types[emb]["actions_cartesian"] == "action_keys" + assert stats.zarr_keys[emb]["actions_cartesian"] == "actions_cartesian" + + +def test_a_failed_probe_tries_the_next_episode(tmp_path, monkeypatch): + ds = _dataset(tmp_path) + calls = _count_probes(monkeypatch, fail_first=2) + stats = MultiDataset(state={}) + stats.populate_from_datasets({"eva_bimanual": ds}) + assert len(calls) == 3 and len(set(calls)) == 3 + (emb,) = stats.key_types + assert stats.zarr_keys[emb]["actions_cartesian"] == "actions_cartesian" diff --git a/tests/unit/test_proprio_history.py b/tests/unit/test_proprio_history.py new file mode 100644 index 0000000000..e12c6f1332 --- /dev/null +++ b/tests/unit/test_proprio_history.py @@ -0,0 +1,898 @@ +"""Proprio observation history: the last K proprio vectors, one token per step. + +K = 1 must stay bit-for-bit today's pipeline -- same sample keys, same shapes, +same stem tokens, same sinusoidal PE -- so every check below has a K = 1 twin. +For K > 1 the dataset reads a backward window (front-padded with the first real +frame so the current frame is always ``[-1]``), the cartesian transforms treat +the leading axis as batch, the baseline's per-channel proprio stats broadcast +over it, and the stem replaces the input-side sinusoid with a learned +time-step embedding added after the projection. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +import egomimic +from egomimic.algo.hpt import HPT, HPTModel +from egomimic.models.hpt_nets import MLPPolicyStem +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T_FRAMES = 24 +HORIZON = Human.ACTION_HORIZON +CHUNK = 100 +PROPRIO_KEYS = ("left.obs_ee_pose", "right.obs_ee_pose") + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _pose_series(T: int, offset: float) -> np.ndarray: + """(T, 7) xyz+quat(wxyz) whose translation is a distinct ramp per frame. + + The shared synthetic fixture holds every pose CONSTANT over time (so the + per-cell quantile bounds stay usable); history needs frames that differ, + so this writes its own episode with an identity rotation and a ramp. + """ + t = np.arange(T, dtype=np.float64)[:, None] + xyz = offset + t * np.array([[0.001, 0.002, 0.003]]) + quat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]]), (T, 1)) + return np.concatenate([xyz, quat], axis=1) + + +def _write_varying_episode(root, *, T: int = T_FRAMES, seed: int = 0): + rng = np.random.default_rng(seed) + H = W = 64 + K = np.array( + [[100.0, 0.0, W / 2, 0.0], [0.0, 100.0, H / 2, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + numeric = { + "left.obs_ee_pose": _pose_series(T, 0.10 + seed), + "right.obs_ee_pose": _pose_series(T, 0.20 + seed), + "obs_head_pose": _pose_series(T, 0.30 + seed), + } + images = {"images.front_1": rng.integers(0, 255, (T, H, W, 3), dtype=np.uint8)} + return ZarrWriter.create_and_write( + root / f"varying_{seed:02d}.zarr", + numeric_data=numeric, + image_data=images, + embodiment="human_bimanual", + fps=30, + task_name="synthetic", + intrinsics={"front_1": K}, + ) + + +def _leaf(ds: MultiDataset): + return next(iter(ds.datasets.values())) + + +def _dataset( + tmp_path, + *, + proprio_history: int = 1, + history_stride: int = 1, + transforms=None, + n: int = 1, +): + for i in range(n): + _write_varying_episode(tmp_path, seed=i) + key_map = Human.get_keymap( + keymap_mode="cartesian", + proprio_history=proprio_history, + history_stride=history_stride, + ) + resolver = LocalEpisodeResolver( + tmp_path, key_map=key_map, transform_list=transforms or [] + ) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _raw_poses(tmp_path, key: str) -> np.ndarray: + """The written (T, 7) array for ``key``, straight off disk.""" + import zarr + + store = zarr.open_group(str(next(tmp_path.glob("*.zarr"))), mode="r") + return np.asarray(store[key][:]) + + +# --------------------------------------------------------------------------- +# 1. dataset: backward window, front padding, mask +# --------------------------------------------------------------------------- + + +def test_history_window_is_the_last_k_frames(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[10][key]) + assert got.shape == (4, 7) + np.testing.assert_allclose(got, raw[7:11], rtol=0, atol=1e-6) + # The current frame is always last. + np.testing.assert_allclose(got[-1], raw[10], rtol=0, atol=1e-6) + + +def test_history_stride_spaces_the_window_out(tmp_path): + """K = 3, s = 5 reads frames idx-10, idx-5, idx -- 0.33 s of past at 30 fps + instead of 0.07 s, with the current frame still last.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[12][key]) + assert got.shape == (3, 7) + np.testing.assert_allclose(got, raw[[2, 7, 12]], rtol=0, atol=1e-6) + np.testing.assert_array_equal( + np.asarray(leaf[12]["proprio_history_mask"]), np.ones(3, dtype=np.float32) + ) + + +def test_strided_history_front_pads_at_the_episode_start(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + raw = _raw_poses(tmp_path, PROPRIO_KEYS[0]) + # idx 7: frames 7 and 2 are real, idx-10 would be negative + got = np.asarray(leaf[7][PROPRIO_KEYS[0]]) + np.testing.assert_allclose(got, raw[[2, 2, 7]], rtol=0, atol=1e-6) + np.testing.assert_array_equal( + np.asarray(leaf[7]["proprio_history_mask"]), + np.array([0, 1, 1], dtype=np.float32), + ) + # idx 3: only the current frame is real + np.testing.assert_array_equal( + np.asarray(leaf[3]["proprio_history_mask"]), + np.array([0, 0, 1], dtype=np.float32), + ) + + +def test_stride_one_is_the_contiguous_window(tmp_path): + a = _leaf(_dataset(tmp_path, proprio_history=4)) + b = _leaf(_dataset(tmp_path, proprio_history=4, history_stride=1)) + for key in PROPRIO_KEYS: + np.testing.assert_array_equal(np.asarray(a[10][key]), np.asarray(b[10][key])) + + +def test_history_mask_is_all_ones_away_from_the_episode_start(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + mask = np.asarray(leaf[10]["proprio_history_mask"]) + assert mask.shape == (4,) + np.testing.assert_array_equal(mask, np.ones(4, dtype=np.float32)) + + +def test_history_front_pads_by_repeating_the_first_real_frame(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[1][key]) + assert got.shape == (4, 7) + expected = np.stack([raw[0], raw[0], raw[0], raw[1]]) + np.testing.assert_allclose(got, expected, rtol=0, atol=1e-6) + + # NOTE: the brief's example mask for idx=1 reads [0, 0, 0, 1]; the rule it + # states ("1 = real, 0 = front padding") gives [0, 0, 1, 1] -- frames 0 and + # 1 are both real reads, only the two leading copies of frame 0 are + # padding. The rule wins; the mask is informational only (it is dropped in + # process_batch_for_training and the padded steps already look exactly like + # a history-dropout sample). + mask = np.asarray(leaf[1]["proprio_history_mask"]) + np.testing.assert_array_equal(mask, np.array([0, 0, 1, 1], dtype=np.float32)) + + np.testing.assert_array_equal( + np.asarray(leaf[0]["proprio_history_mask"]), + np.array([0, 0, 0, 1], dtype=np.float32), + ) + + +def test_history_mask_is_float32(tmp_path): + value = _leaf(_dataset(tmp_path, proprio_history=4))[3]["proprio_history_mask"] + assert isinstance(value, torch.Tensor) + assert value.dtype == torch.float32 + + +def test_k1_sample_is_bit_for_bit_todays_sample(tmp_path): + """The K = 1 keymap must produce exactly what a keymap with no + ``proprio_history`` argument at all produces.""" + _write_varying_episode(tmp_path, seed=0) + baseline = Human.get_keymap(keymap_mode="cartesian") + with_arg = Human.get_keymap(keymap_mode="cartesian", proprio_history=1) + assert with_arg == baseline + assert not any("history" in spec for spec in with_arg.values()) + + def _sample(key_map): + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return _leaf(ds)[5] + + a, b = _sample(baseline), _sample(with_arg) + assert set(a) == set(b) + assert "proprio_history_mask" not in a + for key in PROPRIO_KEYS: + assert tuple(a[key].shape) == (7,) + for key, value in a.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(value, b[key]) + + +def test_keys_without_history_keep_the_single_frame_read(tmp_path): + """Only the proprio keys the model consumes get a time axis: the head pose + (the frame everything is expressed in) and the action chunks do not.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + sample = leaf[10] + assert tuple(sample["obs_head_pose"].shape) == (7,) + assert tuple(sample["left.action_ee_pose"].shape) == (HORIZON, 7) + + +# --------------------------------------------------------------------------- +# 2. transforms: a leading time axis is a batch axis +# --------------------------------------------------------------------------- + + +def _wristframe_batch(poses: dict[str, np.ndarray]) -> dict: + rng = np.random.default_rng(0) + chunk = _pose_series(HORIZON, 0.4) + batch = { + "obs_head_pose": _pose_series(1, 0.3)[0], + "left.action_ee_pose": chunk.copy(), + "right.action_ee_pose": chunk.copy() + 0.01, + "action_pad_mask": np.ones(HORIZON, dtype=np.float32), + } + del rng + batch.update({k: v.copy() for k, v in poses.items()}) + return batch + + +def _run_wristframe(batch: dict) -> dict: + for t in Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ): + batch = t.transform(batch) + return batch + + +def test_wristframe_chain_maps_k_by_7_proprio_to_k_by_20(): + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + out = _run_wristframe(_wristframe_batch(hist)) + state = np.asarray(out["observations.state.ee_pose"]) + assert state.shape == (K, 20) + + current = {key: value[-1] for key, value in hist.items()} + out1 = _run_wristframe(_wristframe_batch(current)) + state1 = np.asarray(out1["observations.state.ee_pose"]) + assert state1.shape == (20,) + np.testing.assert_allclose(state[-1], state1, rtol=1e-6, atol=1e-8) + + +def test_wristframe_actions_are_unchanged_by_the_proprio_time_axis(): + """The action chunk is expressed in the CURRENT step's wrist frame, so it + must not move when history is added next to it.""" + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + current = {key: value[-1] for key, value in hist.items()} + a = _run_wristframe(_wristframe_batch(hist))["actions_cartesian"] + b = _run_wristframe(_wristframe_batch(current))["actions_cartesian"] + np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-6, atol=1e-8) + + +def test_dataset_and_transforms_compose_to_a_k_by_20_state(tmp_path): + transforms = Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ) + hist = _leaf(_dataset(tmp_path, proprio_history=4, transforms=transforms)) + base = _leaf(_dataset(tmp_path, proprio_history=1, transforms=transforms)) + for idx in (0, 3, 11): + h = hist[idx]["observations.state.ee_pose"] + b = base[idx]["observations.state.ee_pose"] + assert tuple(h.shape) == (4, 20) + assert tuple(b.shape) == (20,) + torch.testing.assert_close(h[-1], b, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# 3. normalization / bounds over the time axis +# --------------------------------------------------------------------------- + + +def _norm_shell(key: str, width: int, *, lo=-1.0, hi=1.0): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, lo, dtype=np.float32), + "quantile_99": np.full(width, hi, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md.key_types = {0: {key: "proprio_keys"}} + md._warned_violations = set() + return md + + +def test_per_channel_stats_broadcast_over_the_time_axis(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20, lo=-2.0, hi=4.0) + stats = md.norm_stats[0][key] + g = torch.Generator().manual_seed(0) + arr = torch.randn(4, 20, generator=g) + out = md._apply_norm_one(arr, stats) + assert out.shape == (4, 20) + for k in range(4): + torch.testing.assert_close(out[k], md._apply_norm_one(arr[k], stats)) + + +def test_bounds_check_catches_a_violation_in_a_past_step(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20) + arr = np.zeros((4, 20), dtype=np.float32) + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[0, 0] = 50.0 # oldest history step, translation channel + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +# --------------------------------------------------------------------------- +# 4. the stem +# --------------------------------------------------------------------------- + + +def _stems(D: int, out: int, **kwargs): + torch.manual_seed(0) + a = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], **kwargs) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16]) + return a, b + + +def test_history_stem_shape_and_zero_init_matches_the_baseline_stem(): + D, out, B, K = 5, 7, 3, 3 + hist, base = _stems(D, out, history_len=K) + x = torch.randn(B, K, 1, D) + y = hist(x) + assert y.shape == (B, K, 1, out) + # time_embed is zero at init, so every step is today's projection. + torch.testing.assert_close(y, base(x)) + for k in range(K): + torch.testing.assert_close(y[:, k], base(x[:, k])) + + +def test_time_embed_is_learned_zero_init_and_per_step(): + D, out, K = 5, 7, 3 + hist, _ = _stems(D, out, history_len=K) + assert hist.time_embed.shape == (K, out) + assert hist.time_embed.requires_grad + torch.testing.assert_close(hist.time_embed, torch.zeros(K, out)) + + with torch.no_grad(): + hist.time_embed.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + x = torch.randn(2, K, 1, D) + y = hist(x) + assert not torch.allclose(y[:, 0], y[:, 1]) + # ... and the shift is exactly the per-step embedding. + hist.time_embed.data.zero_() + y0 = hist(x) + hist.time_embed.data.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + torch.testing.assert_close( + hist(x) - y0, + torch.arange(K, dtype=torch.float32)[None, :, None, None].expand(2, K, 1, out), + ) + + +def test_history_len_one_stem_has_no_time_embed_and_is_todays_stem(): + D, out = 5, 7 + one, base = _stems(D, out, history_len=1) + assert not hasattr(one, "time_embed") + assert dict(one.named_parameters()).keys() == dict(base.named_parameters()).keys() + x = torch.randn(4, 1, 1, D) + torch.testing.assert_close(one(x), base(x)) + + +def test_history_dropout_replaces_the_past_with_the_current_step_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + hist, _ = _stems(D, out, history_len=K, history_dropout=1.0) + x = torch.randn(B, K, 1, D) + hist.train() + y = hist(x) + # Every step equals the current one: exactly an episode-start sample after + # front padding. + for k in range(K): + torch.testing.assert_close( + y[:, k], y[:, -1] - hist.time_embed[-1] + hist.time_embed[k] + ) + hist.eval() + torch.testing.assert_close(y, hist(_current_only(x))) + + +def _current_only(x: torch.Tensor) -> torch.Tensor: + out = x.clone() + out[:, :-1] = out[:, -1:] + return out + + +def test_history_dropout_is_a_noop_in_eval_mode(): + D, out, B, K = 5, 7, 6, 4 + drop, _ = _stems(D, out, history_len=K, history_dropout=1.0) + torch.manual_seed(0) + plain = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + drop.eval() + plain.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(drop(x), plain(x)) + assert not torch.allclose(drop(x), drop(_current_only(x))) + + +def test_history_dropout_zero_is_a_noop_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + a, _ = _stems(D, out, history_len=K, history_dropout=0.0) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + a.train() + b.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(a(x), b(x)) + + +def test_history_stem_rejects_a_wrong_length_history(): + hist, _ = _stems(5, 7, history_len=3) + with pytest.raises(ValueError, match="3"): + hist(torch.randn(2, 4, 1, 5)) + + +# --------------------------------------------------------------------------- +# 5. stem_process: the input-side sinusoid goes once history_len > 1 +# --------------------------------------------------------------------------- + + +class _EchoStem: + """Stem stub whose ``compute_latent`` hands back what stem_process fed it.""" + + def __init__(self, history_len: int = 1): + self.history_len = history_len + self.seen = None + + def compute_latent(self, x): + self.seen = x.clone() + return x + + +def _stem_process(history_len: int, data_shape): + model = HPTModel.__new__(HPTModel) + model.modalities = {"human_bimanual": ["state_ee_pose"]} + model.shared_keys = [] + model.encoders = {} + stem = _EchoStem(history_len) + model.stems = {"human_bimanual_state_ee_pose": stem} + model.stem_spec = { + "human_bimanual": { + "state_ee_pose": OmegaConf.create( + {"specs": {"random_horizon_masking": False}} + ) + } + } + model.train_mode = False + torch.manual_seed(0) + x = torch.randn(*data_shape) + model.stem_process("human_bimanual", {"state_ee_pose": x.clone()}) + return x, stem.seen + + +def test_stem_process_skips_the_sinusoid_for_a_history_stem(): + x, seen = _stem_process(4, (2, 4, 20)) + torch.testing.assert_close(seen, x) + + +def test_stem_process_keeps_the_sinusoid_for_a_history_len_one_stem(): + x, seen = _stem_process(1, (2, 1, 20)) + from egomimic.utils.tensor_utils import get_sinusoid_encoding_table + + pe = get_sinusoid_encoding_table(0, 1, 20).to(x) + torch.testing.assert_close(seen, x + pe.view(1, 1, 20)) + assert not torch.allclose(seen, x) + + +# --------------------------------------------------------------------------- +# 6. HPT batch plumbing +# --------------------------------------------------------------------------- + + +class _NoRenameNormStats: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def _hpt_with_stem(history_len: int): + algo = HPT.__new__(HPT) + policy = HPTModel.__new__(HPTModel) + policy.stems = {"human_bimanual_state_ee_pose": _EchoStem(history_len)} + algo.nets = {"policy": policy} + algo.is_6dof = True + algo.shared_ac_key = None + algo.annotation_modality = "annotation" + algo.eval_image_augs = None + return algo + + +def _to_hpt(algo, proprio: torch.Tensor): + batch = { + "observations.state.ee_pose": proprio, + "pad_mask": torch.ones(proprio.shape[0], 5, 1), + "embodiment": torch.tensor([0]), + "actions_cartesian": torch.zeros(proprio.shape[0], 5, 18), + } + return HPT._robomimic_to_hpt_data( + algo, + batch, + cam_keys=[], + proprio_keys=["observations.state.ee_pose"], + lang_keys=[], + ac_key="actions_cartesian", + domain="human_bimanual", + ) + + +def test_two_dim_proprio_is_still_unsqueezed(): + algo = _hpt_with_stem(1) + data = _to_hpt(algo, torch.randn(3, 20)) + assert data["state_ee_pose"].shape == (3, 1, 20) + + +def test_three_dim_proprio_passes_through_without_unsqueeze(): + algo = _hpt_with_stem(4) + x = torch.randn(3, 4, 20) + data = _to_hpt(algo, x) + assert data["state_ee_pose"].shape == (3, 4, 20) + torch.testing.assert_close(data["state_ee_pose"], x) + + +def test_history_len_is_read_from_the_domain_not_by_suffix(): + """A cotrain run has one state_ee_pose stem per domain; a suffix match + returns whichever comes first in the dict, so the human side would be + checked against eva's history_len.""" + algo = _hpt_with_stem(3) + algo.nets["policy"].stems = { + "eva_bimanual_state_ee_pose": _EchoStem(1), + "human_bimanual_state_ee_pose": _EchoStem(3), + } + data = _to_hpt(algo, torch.randn(3, 3, 20)) + assert data["state_ee_pose"].shape == (3, 3, 20) + + +def test_history_len_is_read_from_a_built_policys_module_dict(): + """A built policy holds its stems in an nn.ModuleDict, which has no .get.""" + import types + + algo = _hpt_with_stem(3) + algo.nets["policy"] = types.SimpleNamespace( + stems=torch.nn.ModuleDict( + {"human_bimanual_state_ee_pose": MLPPolicyStem(20, 8, history_len=3)} + ) + ) + data = _to_hpt(algo, torch.randn(3, 3, 20)) + assert data["state_ee_pose"].shape == (3, 3, 20) + with pytest.raises(ValueError, match="state_ee_pose"): + _to_hpt(algo, torch.randn(3, 4, 20)) + + +def test_history_length_mismatch_raises_naming_both(): + algo = _hpt_with_stem(2) + with pytest.raises(ValueError, match="state_ee_pose"): + _to_hpt(algo, torch.randn(3, 4, 20)) + + +def _process(action_pad_mask=None, history_mask=None, S: int = 6): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRenameNormStats() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = False + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + algo.ac_keys = {emb_id: "actions_cartesian"} + inner = { + "actions_cartesian": torch.zeros(2, S, 18), + "observations.state.ee_pose": torch.zeros(2, 4, 20), + } + if action_pad_mask is not None: + inner["action_pad_mask"] = action_pad_mask + if history_mask is not None: + inner["proprio_history_mask"] = history_mask + return HPT.process_batch_for_training(algo, {emb: inner})[emb_id] + + +def test_process_batch_drops_the_history_mask(): + out = _process(history_mask=torch.ones(2, 4)) + assert "proprio_history_mask" not in out + assert out["observations.state.ee_pose"].shape == (2, 4, 20) + + +def test_process_batch_is_unchanged_without_the_history_mask(): + out = _process() + assert "proprio_history_mask" not in out + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + + +# --------------------------------------------------------------------------- +# 7. configs +# --------------------------------------------------------------------------- + +RECIPE = "train_zarr_mecka_flagship_6d_hpt" +DATA = "data=mecka_fold_flagship_opsplit_hpt_6d" +KEYMAP = "data.train_datasets.human_bimanual.resolver.key_map" +STEM = "model.robomimic_model.stem_specs.human_bimanual.state_ee_pose" + + +def test_flagship_ships_k_three_stride_five(compose_resolve): + cfg = compose_resolve(RECIPE, [DATA]) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 3 + assert km.history_stride == 5 + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose + assert stem.history_len == km.proprio_history, "data K and stem K must agree" + assert stem.history_dropout == 0.2 + + +def test_k_one_baseline_is_still_reachable(compose_resolve): + cfg = compose_resolve( + RECIPE, + [ + DATA, + f"{KEYMAP}.proprio_history=1", + f"{KEYMAP}.history_stride=1", + f"{STEM}.history_len=1", + ], + ) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 1 and km.history_stride == 1 + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 1 + ) + + +def test_k_two_override_reaches_every_loader(compose_resolve): + cfg = compose_resolve( + RECIPE, [DATA, f"{KEYMAP}.proprio_history=2", f"{STEM}.history_len=2"] + ) + data = cfg.data + for group in ("train_datasets", "valid_datasets", "unseen_op_valid_datasets"): + km = data[group]["human_bimanual"]["resolver"]["key_map"] + assert km["proprio_history"] == 2, group + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 2 + ) + + +def test_k_two_keymap_marks_only_the_model_proprio_keys(): + km = Human.get_keymap(keymap_mode="cartesian", proprio_history=2) + assert {k for k, v in km.items() if "history" in v} == set(PROPRIO_KEYS) + assert all(km[k]["history"] == 2 for k in PROPRIO_KEYS) + + +# --------------------------------------------------------------------------- +# 8. the eval path: the revert transforms and the evaluator take the current +# step, so val metrics and the overlay video work at K > 1 +# --------------------------------------------------------------------------- + +REVERT_BUILDERS = [ + "_build_human_cartesian_revert_6d_wristframe_transform_list", + "_build_human_cartesian_revert_6d_transform_list", +] + + +def _revert_inputs(B: int = 2, K: int = 3): + """``(actions, current_proprio, history_proprio)`` whose last history step + is exactly the 2-D proprio.""" + g = torch.Generator().manual_seed(7) + actions = torch.randn(B, 100, 18, generator=g) + current = torch.randn(B, 20, generator=g) + history = torch.randn(B, K, 20, generator=g) + history[:, -1] = current + return actions, current, history + + +@pytest.mark.parametrize("builder_name", REVERT_BUILDERS) +def test_revert_transform_list_is_blind_to_proprio_history(builder_name): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.embodiment.embodiment import Embodiment + + transform_list = getattr(human_mod, builder_name)() + actions, current, history = _revert_inputs() + + def _run(proprio): + return Embodiment.apply_transform( + { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + }, + transform_list, + ) + + flat, hist = _run(current), _run(history) + assert flat["observations.state.ee_pose"].shape == (2, 12) + assert hist["observations.state.ee_pose"].shape == (2, 12) + for key in ("actions_cartesian", "observations.state.ee_pose"): + torch.testing.assert_close(hist[key], flat[key], rtol=1e-5, atol=1e-6) + + +def test_select_current_step_leaves_a_flat_value_alone_and_is_idempotent(): + from egomimic.rldb.transforms import SelectCurrentStep + + t = SelectCurrentStep(keys=["observations.state.ee_pose", "missing"]) + flat = np.arange(20, dtype=np.float32) + out = t.transform({"observations.state.ee_pose": flat}) + assert out["observations.state.ee_pose"] is flat # untouched, not a copy + + hist = np.stack([flat - 1, flat, flat + 1]) + once = t.transform({"observations.state.ee_pose": hist.copy()}) + np.testing.assert_array_equal(once["observations.state.ee_pose"], flat + 1) + twice = t.transform(dict(once)) + np.testing.assert_array_equal(twice["observations.state.ee_pose"], flat + 1) + + +def test_every_human_revert_list_selects_the_current_proprio_step(): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.transforms import SelectCurrentStep + + builders = { + "_build_human_cartesian_revert_6d_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_6d_wristframe_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_eef_frame_transform_list": "observations.state.ee_pose", + "_build_human_keypoints_revert_6d_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_6d_wristframe_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_eef_frame_transform_list": "observations.state.keypoints", + } + for name, obs_key in builders.items(): + tl = getattr(human_mod, name)() + assert isinstance(tl[0], SelectCurrentStep), name + assert tl[0].keys == [obs_key], name + + +# --- the evaluator itself ----------------------------------------------- + + +class _IdentityNormStats: + def unnormalize(self, batch, embodiment_id): + return dict(batch) + + +class _StubAlgo: + """Just enough of HPT for HPTEvalVideo.compute_metrics_and_viz.""" + + def __init__(self, emb_id, preds): + self.norm_stats = _IdentityNormStats() + self.ac_keys = {emb_id: "actions_cartesian"} + self.shared_ac_key = None + self.auxiliary_ac_keys = {} + self._preds = preds + + def forward_eval(self, batch): + return self._preds + + +def _eval_once(proprio): + """Run the real HPTEvalVideo over one batch, returning its metrics and the + batch the viz function was handed.""" + from egomimic.eval.eval_hpt import HPTEvalVideo + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + + emb_id = get_embodiment_id("human_bimanual") + emb_name = "human_bimanual" + # Action horizon 1: the evaluator's paired/final MSE does + # ``pred[:, -1].cpu()``, which is a real (contiguous) copy for the CUDA + # tensors a run produces but a non-contiguous view for the CPU tensors a + # test hands it, and torchmetrics' ``view(-1)`` rejects that. S = 1 keeps + # the slice contiguous; the proprio path under test is unaffected. + g = torch.Generator().manual_seed(11) + actions = torch.randn(2, 1, 18, generator=g) + preds = { + f"{emb_name}_actions_cartesian": torch.randn(2, 1, 18, generator=g), + f"{emb_name}_loss": torch.tensor(0.5), + } + batch = { + emb_id: { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + "embodiment": torch.tensor([emb_id, emb_id]), + } + } + + seen = {} + + def _viz(predictions, viz_batch): + seen["batch"] = viz_batch + seen["preds"] = predictions + return np.zeros((2, 4, 4, 3), dtype=np.uint8) + + ev = HPTEvalVideo.__new__(HPTEvalVideo) + ev._replay_now = False + ev.model = _StubAlgo(emb_id, preds) + ev.viz_func = {emb_name: _viz} + ev.transform_lists = { + emb_name: _build_human_cartesian_revert_6d_wristframe_transform_list() + } + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=True) + return metrics, images, seen + + +def test_evaluator_metrics_and_viz_are_blind_to_proprio_history(): + _, current, history = _revert_inputs() + m_flat, im_flat, seen_flat = _eval_once(current) + m_hist, im_hist, seen_hist = _eval_once(history) + + assert set(m_flat) == set(m_hist) + cam_keys = [k for k in m_flat if "_cam_" in k] + assert cam_keys, "the cam-frame metrics must be produced" + for key in m_flat: + torch.testing.assert_close( + torch.as_tensor(m_flat[key]).float(), + torch.as_tensor(m_hist[key]).float(), + rtol=1e-5, + atol=1e-6, + ) + # The viz function sees the reverted, current-step proprio either way. + assert seen_flat["batch"]["observations.state.ee_pose"].shape == (2, 12) + torch.testing.assert_close( + seen_hist["batch"]["observations.state.ee_pose"], + seen_flat["batch"]["observations.state.ee_pose"], + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + seen_hist["preds"]["human_bimanual_actions_cartesian"], + seen_flat["preds"]["human_bimanual_actions_cartesian"], + rtol=1e-5, + atol=1e-6, + ) + assert set(im_flat) == set(im_hist) + + +TRAIN_CONFIGS = sorted( + p.stem + for p in (Path(egomimic.__file__).parent / "hydra_configs").glob("train_*.yaml") +) + + +@pytest.mark.parametrize("top", TRAIN_CONFIGS) +def test_every_recipe_agrees_with_its_stem_on_K(top, compose_resolve): + """The data's `proprio_history` and the stem's `history_len` must match, or + process_batch_for_training raises at the first training step. A recipe that + inherits a data config from one model and overrides the model is exactly + how they drift apart, and nothing else composes both halves.""" + cfg = compose_resolve(top, []) + stem_specs = (cfg.model.robomimic_model.get("stem_specs") or {}).get( + "human_bimanual" + ) or {} + for split in ("train_datasets", "valid_datasets"): + node = (cfg.data.get(split) or {}).get("human_bimanual") + if node is None or node.get("resolver") is None: + continue + k = int(node.resolver.key_map.get("proprio_history", 1) or 1) + for name, stem in stem_specs.items(): + if stem is None or "history_len" not in stem: + continue + assert int(stem.history_len) == k, (top, split, name, k, stem.history_len) diff --git a/tests/unit/test_rdt.py b/tests/unit/test_rdt.py new file mode 100644 index 0000000000..06bb103ce5 --- /dev/null +++ b/tests/unit/test_rdt.py @@ -0,0 +1,669 @@ +"""RDT on the HPT chassis: the shared DiT backbone and per-embodiment denoiser, +``RDTModel``'s condition streams (random-init tiny DINOv3 tower, no network) +and the algo's image-history batch layout.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from torchvision import transforms as T + +from egomimic.algo.rdt import RDT, RDTModel +from egomimic.models.fm_policy import FMPolicy +from egomimic.models.hpt_nets import DINOv3Stem, MLPPolicyStem, PolicyStem +from egomimic.models.rdt_nets import ( + RDTBackbone, + RDTConditions, + RDTDenoiser, + sincos_2d, +) + +D, ACT, SEQ, B = 32, 6, 5, 3 +DOMAIN = "human_bimanual" +TINY_TOWER = {"embed_dim": 32, "depth": 1, "num_heads": 2} + + +def _backbone(depth=4): + torch.manual_seed(0) + return RDTBackbone(hidden_dim=D, depth=depth, n_heads=4).eval() + + +def _denoiser(**kwargs): + torch.manual_seed(0) + net = RDTDenoiser(act_dim=ACT, act_seq=SEQ, hidden_dim=D, **kwargs) + # the zero-init output layer would hide every upstream effect + nn.init.normal_(net.ffn_final.fc2.weight, std=0.1) + return net.eval() + + +def _cond(backbone=None, lang_len=4, n_state=1): + torch.manual_seed(1) + return RDTConditions( + img=torch.randn(B, 7, D), + freq=torch.full((B,), 30.0), + backbone=backbone or _backbone(), + state=torch.randn(B, n_state, D) if n_state else None, + lang=torch.randn(B, lang_len, D), + lang_mask=torch.ones(B, lang_len, dtype=torch.bool), + ) + + +def test_denoiser_shape_and_zero_init(): + net = RDTDenoiser(act_dim=ACT, act_seq=SEQ, hidden_dim=D) + out = net(torch.randn(B, SEQ, ACT), torch.rand(B), _cond()) + assert out.shape == (B, SEQ, ACT) + assert torch.count_nonzero(out) == 0 + + +def test_denoiser_ignores_masked_language_tokens(): + net, x, t = _denoiser(), torch.randn(B, SEQ, ACT), torch.rand(B) + cond = _cond() + cond.lang_mask[:, -2:] = False + ref = net(x, t, cond) + cond.lang[:, -2:] += 10.0 + assert torch.allclose(net(x, t, cond), ref, atol=1e-6) + cond.lang[:, 0] += 10.0 + assert not torch.allclose(net(x, t, cond), ref, atol=1e-4) + + +def test_denoiser_uses_every_condition_time_and_frequency(): + net, x, t = _denoiser(), torch.randn(B, SEQ, ACT), torch.rand(B) + backbone = _backbone() + ref = net(x, t, _cond(backbone)) + for field in ("img", "state", "lang", "freq"): + cond = _cond(backbone) + setattr(cond, field, getattr(cond, field) + 1.0) + assert not torch.allclose(net(x, t, cond), ref, atol=1e-4), field + assert not torch.allclose(net(x, t * 0.5, _cond(backbone)), ref, atol=1e-4) + + +def test_denoisers_of_different_widths_share_one_backbone(): + """Co-training: per-embodiment adaptors, one set of DiT blocks.""" + backbone = _backbone() + wide = RDTDenoiser(act_dim=ACT + 3, act_seq=SEQ, hidden_dim=D) + cond = _cond(backbone) + assert _denoiser()(torch.randn(B, SEQ, ACT), torch.rand(B), cond).shape[-1] == ACT + assert wide(torch.randn(B, SEQ, ACT + 3), torch.rand(B), cond).shape[-1] == ACT + 3 + assert not any("blocks" in n for n, _ in wide.named_parameters()) + + +def test_denoiser_without_language_or_state(): + net = _denoiser(n_state_tokens=0) + cond = RDTConditions( + img=torch.randn(B, 7, D), freq=torch.full((B,), 30.0), backbone=_backbone() + ) + assert net(torch.randn(B, SEQ, ACT), torch.rand(B), cond).shape == (B, SEQ, ACT) + + +def test_denoiser_rejects_wrong_state_token_count(): + with pytest.raises(ValueError, match="n_state_tokens"): + _denoiser()(torch.randn(B, SEQ, ACT), torch.rand(B), _cond(n_state=2)) + + +def test_fm_head_samples_through_the_conditions_container(): + head = FMPolicy( + model=_denoiser(), + action_horizon=SEQ, + infer_ac_dims={DOMAIN: ACT}, + num_inference_steps=2, + ) + assert head((_cond(), DOMAIN)).shape == (B, SEQ, ACT) + + +def test_sincos_2d_is_resolution_agnostic_at_the_corners(): + small, large = sincos_2d(D, (2, 3)), sincos_2d(D, (14, 25)) + assert small.shape == (6, D) and large.shape == (350, D) + for a, b in ((0, 0), (2, 24), (3, 325), (5, 349)): # the four corners + assert torch.allclose(small[a], large[b], atol=1e-6) + # vertical neighbours share their column half + assert torch.equal(large[0, D // 2 :], large[25, D // 2 :]) + + +def _dino(**kwargs): + torch.manual_seed(0) + kwargs.setdefault("freeze_backbone", True) + return DINOv3Stem( + model_name="vit_small_patch16_dinov3", + output_dim=D, + image_size=[32, 48], + pretrained=False, + tower_kwargs=TINY_TOWER, + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# RDTModel +# -------------------------------------------------------------------------- + + +class _StubText(PolicyStem): + """A per-token text stem: prompt i has i + 1 real tokens, left-padded.""" + + def __init__(self): + super().__init__() + self.embed = nn.Embedding(8, D) + + def forward_with_mask(self, prompts): + n = len(prompts) + ids = torch.arange(n)[None, :].expand(n, -1) + mask = ids >= (n - 1 - torch.arange(n))[:, None] + return self.embed(ids), mask + + +def _model( + cond_drop=None, + text=True, + history_len=1, + frames=2, + frame_dropout=0.0, + memory=None, + short_dropout=0.0, +): + torch.manual_seed(0) + model = RDTModel( + embed_dim=D, + depth=2, + num_heads=4, + cond_drop=cond_drop, + image_history=frames, + image_history_dropout=frame_dropout, + memory=memory, + short_dropout=short_dropout, + ) + shared = {"front_img_1": MLPPolicyStem(input_dim=D, output_dim=D, widths=[D])} + if text: + shared["annotation"] = _StubText() + model.init_domain_stem("shared", shared) + model.shared_keys = list(shared) + model.init_domain_stem( + DOMAIN, + { + "state_ee_pose": MLPPolicyStem( + input_dim=4, output_dim=D, widths=[D], history_len=history_len + ) + }, + ) + model.init_domain_head( + DOMAIN, + FMPolicy( + model=RDTDenoiser( + act_dim=ACT, + act_seq=SEQ, + hidden_dim=D, + n_state_tokens=history_len, + n_memory_tokens=memory["frames"] if memory else 0, + ), + action_horizon=SEQ, + infer_ac_dims={DOMAIN: ACT}, + num_inference_steps=2, + ), + ) + model.init_encoder("front_img_1", _dino()) + model.auxiliary_ac_keys = {} + model.diffusion = True + model.device = torch.device("cpu") + model.finalize_modules() + return model + + +def _data(history_len=1, text=True, frames=2, memory_frames=0): + torch.manual_seed(2) + data = { + "front_img_1": torch.rand(B, frames, 1, 3, 32, 48), + "state_ee_pose": torch.randn(B, history_len, 4), + "action": torch.randn(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "fps": torch.full((B,), 30.0), + } + if text: + data["annotation"] = ["a", "b", "c"] + if memory_frames: + data["memory"] = torch.rand(B, memory_frames, 3, 32, 48) + data["memory_mask"] = torch.ones(B, memory_frames, dtype=torch.bool) + return data + + +ALL_DROPS = {"state": 1.0, "lang": 1.0, "img": 1.0} + + +def test_model_streams_and_shapes(): + model = _model(history_len=2).eval() + cond, _ = model.forward_features(DOMAIN, _data(history_len=2)) + assert cond.img.shape == (B, 2 * 6, D) # 2 frames x 6 patches + assert cond.state.shape == (B, 2, D) + assert cond.lang.shape == (B, B, D) + assert cond.lang_mask.sum(1).tolist() == [1, 2, 3] + assert cond.freq.tolist() == [30.0] * B + assert cond.backbone is model.trunk["trunk"] + + +def test_model_shares_the_dit_in_the_trunk_not_the_head(): + names = [n for n, _ in _model().named_parameters()] + assert any(n.startswith("trunk.trunk.blocks.") for n in names) + assert not any(n.startswith("heads.") and ".blocks." in n for n in names) + assert not any(".tokens" in n or "cross_attention" in n for n in names) + + +def test_model_keeps_rdt_init_after_hpt_xavier_pass(): + net = _model().heads[DOMAIN].model + assert torch.count_nonzero(net.ffn_final.fc2.weight) == 0 + assert net.x_pos_embed.abs().sum() > 0 + + +def test_model_loss_reaches_every_trainable_parameter(): + """Plain DDP crashes on a parameter that gets no gradient.""" + model = _model(cond_drop={"state": 0.5, "lang": 0.5}, frame_dropout=0.5).train() + nn.init.normal_(model.heads[DOMAIN].model.ffn_final.fc2.weight, std=0.1) + model.compute_loss({"domain": DOMAIN, "data": _data()}).backward() + missing = [ + n for n, p in model.named_parameters() if p.requires_grad and p.grad is None + ] + assert not missing, missing + + +def test_model_without_text_stem_and_single_frame(): + model = _model(cond_drop=ALL_DROPS, text=False, frames=1).train() + data = _data(text=False, frames=1) + cond, _ = model.forward_features(DOMAIN, data) + assert cond.lang is None and cond.img.shape == (B, 6, D) + model.compute_loss({"domain": DOMAIN, "data": data}).backward() + out = model.eval().forward(DOMAIN, _data(text=False, frames=1)) + assert out[DOMAIN].shape == (B, SEQ, ACT) + + +def test_model_rejects_a_frame_count_the_config_did_not_declare(): + with pytest.raises(ValueError, match="image_history"): + _model(frames=2).forward_features(DOMAIN, _data(frames=1)) + + +def test_condition_dropout_is_train_only_and_never_blinds_one_camera(): + model = _model(cond_drop=ALL_DROPS) + assert not hasattr(model, "null_img") # one camera: img dropout is inert + seen = [] + stem = model.stems["shared_annotation"] + real = stem.forward_with_mask + stem.forward_with_mask = lambda prompts: ( + seen.append(list(prompts)), + real(prompts), + )[1] + + ref, _ = model.eval().forward_features(DOMAIN, _data()) + assert seen[-1] == ["a", "b", "c"] + cond, _ = model.train().forward_features(DOMAIN, _data()) + assert seen[-1] == ["", "", ""] # a dropped prompt is the empty prompt + assert torch.allclose(cond.state, model.null_state.expand_as(cond.state)) + assert torch.allclose(cond.img, ref.img) # tower frozen + in eval: same tokens + + +def test_image_history_dropout_replaces_the_past_frame_by_the_current(): + model = _model(frame_dropout=1.0).train() + with torch.no_grad(): + model.frame_embed.zero_() + cond, _ = model.forward_features(DOMAIN, _data()) + past, current = cond.img[:, :6], cond.img[:, 6:] + assert torch.allclose(past, current, atol=1e-6) + cond, _ = model.eval().forward_features(DOMAIN, _data()) + assert not torch.allclose(cond.img[:, :6], cond.img[:, 6:], atol=1e-4) + + +def test_frame_and_camera_tags_start_at_zero(): + model = _model() + assert torch.count_nonzero(model.frame_embed) == 0 + assert all(torch.count_nonzero(p) == 0 for p in model.camera_embed.values()) + + +def test_model_needs_fps(): + data = _data() + del data["fps"] + with pytest.raises(ValueError, match="fps"): + _model().forward_features(DOMAIN, data) + + +# -------------------------------------------------------------------------- +# RDT algo +# -------------------------------------------------------------------------- + + +def test_rdt_rejects_hpt_trunk_features(): + with pytest.raises(ValueError, match="HPT-only"): + RDT(None, {}, None, None, ot=True) + + +def _pairing_algo(augs): + from egomimic.models.image_augs import PerSampleAugs + + algo = RDT.__new__(RDT) + algo.nets = nn.ModuleDict().train() + algo.encoders = {"front_img_1": None} + algo.train_image_augs = PerSampleAugs(augs) + algo.eval_image_augs = None + algo.annotation_modality, algo.is_6dof, algo.shared_ac_key = ( + "annotation", + True, + None, + ) + return algo + + +class _Shift(nn.Module): # a "random" aug with no batched form: one draw per call + def forward(self, image): + return image + torch.rand(()) + + +@pytest.mark.parametrize("vectorized", [True, False]) +def test_algo_pairs_history_frames_under_one_augmentation(vectorized): + """(past, current) reach the model as one (B, 2, 1, 3, H, W) input, jittered + with the same per-sample draw, on the batched and the per-image aug path.""" + from torchvision.transforms import ColorJitter + + algo = _pairing_algo(ColorJitter(0.4, 0.4, 0.4, 0.1) if vectorized else _Shift()) + assert (algo.train_image_augs.vectorized is not None) == vectorized + cam = "observations.images.front_img_1" + frame = torch.rand(1, 3, 8, 8).expand(B, -1, -1, -1).contiguous() + batch = { + cam: frame, + f"{cam}_hist": frame.clone(), + "actions": torch.zeros(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((B,), 30.0), + } + data = algo._robomimic_to_hpt_data(batch, [cam, f"{cam}_hist"], [], [], "actions") + frames = data["front_img_1"] + assert frames.shape == (B, 2, 1, 3, 8, 8) + assert "front_img_1_hist" not in data + # identical inputs: the pair stays identical, the samples do not + assert torch.allclose(frames[:, 0], frames[:, 1], atol=1e-6) + assert not torch.allclose(frames[0], frames[1], atol=1e-4) + assert data["fps"].tolist() == [30.0] * B + + +def test_algo_keeps_past_and_current_in_order(): + algo = _pairing_algo(nn.Identity()) + cam = "observations.images.front_img_1" + batch = { + cam: torch.ones(B, 3, 8, 8), + f"{cam}_hist": torch.zeros(B, 3, 8, 8), + "actions": torch.zeros(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((B,), 30.0), + } + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert frames[:, 0].sum() == 0 and frames[:, 1].min() == 1 # current frame last + + +def _pair_batch(current, past): + cam = "observations.images.front_img_1" + n = current.shape[0] + return cam, { + cam: current, + f"{cam}_hist": past, + "actions": torch.zeros(n, SEQ, ACT), + "pad_mask": torch.ones(n, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((n,), 30.0), + } + + +@pytest.mark.parametrize("vectorized", [True, False]) +def test_a_frames_contrast_mean_is_its_own(vectorized): + """Side by side, contrast took one mean over both frames, so the past + frame's pixels depended on the current frame's.""" + import torchvision.transforms.functional as F + from torchvision.transforms import ColorJitter + + jitter = ColorJitter(contrast=(0.5, 0.5)) + algo = _pairing_algo(jitter if vectorized else T.Compose([jitter, _Shift()])) + past = torch.rand(B, 3, 8, 8) + outs = [] + for current in (torch.zeros(B, 3, 8, 8) + 0.9, torch.rand(B, 3, 8, 8)): + cam, batch = _pair_batch(current, past.clone()) + outs.append( + algo._robomimic_to_hpt_data(batch, [cam, f"{cam}_hist"], [], [], "actions")[ + "front_img_1" + ][:, 0, 0] + ) + reference = F.adjust_contrast(past, 0.5) + for out in outs: + # _Shift adds one random offset per sample on the per-image path + offset = (out - reference).flatten(1) + assert torch.allclose(offset, offset[:, :1].expand_as(offset), atol=1e-5) + if vectorized: + assert torch.allclose(out, reference, atol=1e-5) + + +def test_a_resize_runs_per_frame(): + algo = _pairing_algo(T.Compose([T.ColorJitter(0.4), T.Resize(4)])) + frame = torch.rand(1, 3, 8, 8).expand(B, -1, -1, -1).contiguous() + cam, batch = _pair_batch(frame, frame.clone()) + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert frames.shape == (B, 2, 1, 3, 4, 4) + assert torch.allclose(frames[:, 0], frames[:, 1], atol=1e-6) + + +def test_an_absent_past_frame_is_not_normalized(): + algo = _pairing_algo( + T.Compose([T.ColorJitter(0.4), T.Normalize([0.5] * 3, [0.25] * 3)]) + ) + cam, batch = _pair_batch(torch.rand(B, 3, 8, 8), torch.zeros(B, 3, 8, 8)) + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert torch.equal(frames[:, 0], torch.zeros_like(frames[:, 0])) + assert frames[:, 1].min() < 0, "the present frame is normalized" + + +def test_compile_traces_the_dit_through_the_heads(): + """RDT's trunk is not a compile target of its own (no_trunk); each head's + denoiser calls it, so its blocks land in the head's graph.""" + import types + + graphs = [] + + def backend(gm, _inputs): + graphs.append(gm) + return gm.forward + + model = _model().train() + original = nn.Module.compile + + def compile_with(self, **kwargs): + kwargs.pop("mode", None) + return original(self, backend=backend, **kwargs) + + torch._dynamo.reset() + try: + nn.Module.compile = compile_with + RDT.compile_for_training(types.SimpleNamespace(nets={"policy": model})) + torch.manual_seed(0) + model.compute_loss({"domain": DOMAIN, "data": _data()}).backward() + finally: + nn.Module.compile = original + torch._dynamo.reset() + assert any("blocks" in g.code and "backbone" in g.code for g in graphs) + assert all( + p.grad is not None + for n, p in model.named_parameters() + if "trunk.trunk.blocks" in n + ) + + +# -------------------------------------------------------------------------- +# Long-range memory +# -------------------------------------------------------------------------- + +MEM_FRAMES = 3 + + +def _memory(**kwargs): + return { + "encoder": _dino(pool="cls", freeze_backbone=False), + "frames": MEM_FRAMES, + "stride_s": 1.0, + **kwargs, + } + + +def _mem_cond(backbone=None): + cond = _cond(backbone) + torch.manual_seed(3) + cond.memory = torch.randn(B, 4, D) + cond.memory_mask = torch.ones(B, 4, dtype=torch.bool) + return cond + + +def test_denoiser_ignores_masked_memory_tokens_and_uses_the_rest(): + net, x, t = _denoiser(n_memory_tokens=4), torch.randn(B, SEQ, ACT), torch.rand(B) + cond = _mem_cond() + cond.memory_mask[:, :2] = False + ref = net(x, t, cond) + cond.memory[:, :2] += 10.0 + assert torch.allclose(net(x, t, cond), ref, atol=1e-5) + cond.memory[:, 3] += 1.0 + assert not torch.allclose(net(x, t, cond), ref, atol=1e-4) + + +def test_a_denoiser_without_memory_is_unchanged_by_the_new_arguments(): + old = _denoiser() + new = _denoiser(n_memory_tokens=0) + x, t = torch.randn(B, SEQ, ACT), torch.rand(B) + assert torch.equal(old(x, t, _cond()), new(x, t, _cond())) + + +def test_denoiser_rejects_wrong_memory_token_count(): + with pytest.raises(ValueError, match="n_memory_tokens"): + _denoiser(n_memory_tokens=3)( + torch.randn(B, SEQ, ACT), torch.rand(B), _mem_cond() + ) + + +def test_model_memory_tokens_and_mask(): + model = _model(memory=_memory()).eval() + data = _data(memory_frames=MEM_FRAMES) + data["memory_mask"][0, 0] = False + cond, _ = model.forward_features(DOMAIN, data) + assert cond.memory.shape == (B, MEM_FRAMES, D) + assert cond.memory_mask[0].tolist() == [False] + [True] * (MEM_FRAMES - 1) + assert cond.memory_mask[1:].all() + out = model.forward(DOMAIN, data) + assert out[DOMAIN].shape == (B, SEQ, ACT) + + +def test_memory_frames_are_tagged_by_age(): + model = _model(memory=_memory()).eval() + data = _data(memory_frames=MEM_FRAMES) + data["memory"][:] = data["memory"][:, :1].clone() # three identical frames + cond, _ = model.forward_features(DOMAIN, data) + assert not torch.allclose(cond.memory[:, 0], cond.memory[:, 1], atol=1e-4) + + +def test_short_dropout_is_train_only_and_nulls_the_whole_short_stream(): + model = _model(memory=_memory(), short_dropout=1.0) + data = _data(memory_frames=MEM_FRAMES) + cond, _ = model.train().forward_features(DOMAIN, data) + assert torch.equal(cond.img, model.null_short.expand_as(cond.img)) + assert cond.memory_mask.all() + cond, _ = model.eval().forward_features(DOMAIN, data) + assert not torch.equal(cond.img, model.null_short.expand_as(cond.img)) + + +def test_no_short_dropout_adds_no_parameter(): + assert not hasattr(_model(memory=_memory()), "null_short") + + +@pytest.mark.parametrize("fuse", [False, True]) +def test_memory_loss_reaches_every_trainable_parameter(fuse): + history = MEM_FRAMES + 1 if fuse else 1 + model = _model( + history_len=history, memory=_memory(fuse_proprio=fuse), short_dropout=0.5 + ).train() + nn.init.normal_(model.heads[DOMAIN].model.ffn_final.fc2.weight, std=0.1) + data = _data(history_len=history, memory_frames=MEM_FRAMES) + model.compute_loss({"domain": DOMAIN, "data": data}).backward() + missing = [ + n for n, p in model.named_parameters() if p.requires_grad and p.grad is None + ] + assert not missing, missing + assert all(p.requires_grad for p in model.long_memory.encoder.tower.parameters()) + + +def test_fused_memory_reads_the_proprio_of_its_own_step(): + history = MEM_FRAMES + 1 + model = _model(history_len=history, memory=_memory(fuse_proprio=True)).eval() + data = _data(history_len=history, memory_frames=MEM_FRAMES) + ref, _ = model.forward_features(DOMAIN, data) + data["state_ee_pose"][:, 1] += 1.0 # the current step is last, so 1 <-> 0 + cond, _ = model.forward_features(DOMAIN, data) + diff = (cond.memory - ref.memory).abs().reshape(B, MEM_FRAMES, -1).amax(-1) + assert (diff[:, 0] > 1e-4).all() and (diff[:, 1:] < 1e-6).all() + + +def test_the_memory_tower_trains_in_the_backbone_lr_group(): + from types import SimpleNamespace + + from egomimic.pl_utils.pl_model import ModelWrapper + + model = _model(memory=_memory()) + root = nn.Module() + root.nets = nn.ModuleDict({"policy": model}) + stub = SimpleNamespace(model=root, trainer=SimpleNamespace(model=root)) + backbone, rest = ModelWrapper._backbone_param_groups(stub, 1.0, 0.1) + assert backbone["lr"] == pytest.approx(0.1) + tower = {id(p) for p in model.long_memory.encoder.tower.parameters()} + assert {id(p) for p in backbone["params"]} == tower + assert not tower & {id(p) for p in rest["params"]} + + +def test_fused_memory_ignores_proprio_steps_older_than_its_window(): + history = MEM_FRAMES + 1 + model = _model(history_len=history, memory=_memory(fuse_proprio=True)).eval() + data = _data(history_len=history, memory_frames=MEM_FRAMES) + ref, _ = model.forward_features(DOMAIN, data) + data["state_ee_pose"][:, 0] += 1.0 + cond, _ = model.forward_features(DOMAIN, data) + assert torch.equal(cond.memory, ref.memory) + + +def test_fused_memory_needs_the_proprio_on_its_grid(): + model = _model(history_len=1, memory=_memory(fuse_proprio=True)).eval() + with pytest.raises(ValueError, match="fuse_proprio"): + model.forward_features(DOMAIN, _data(memory_frames=MEM_FRAMES)) + + +def test_memory_config_and_batch_must_agree(): + with pytest.raises(ValueError, match="image_memory"): + _model(memory=_memory()).forward_features(DOMAIN, _data()) + with pytest.raises(ValueError, match="image_memory"): + _model().forward_features(DOMAIN, _data(memory_frames=MEM_FRAMES)) + with pytest.raises(ValueError, match="frames"): + _model(memory=_memory()).forward_features(DOMAIN, _data(memory_frames=2)) + + +def test_memory_encoder_must_pool_each_frame_to_one_token(): + model = _model(memory=_memory(encoder=_dino(freeze_backbone=False))) + with pytest.raises(ValueError, match="pool: cls"): + model.forward_features(DOMAIN, _data(memory_frames=MEM_FRAMES)) + + +def test_algo_hands_the_memory_window_over_unaugmented(): + algo = _pairing_algo(_Shift()) + window = torch.rand(B, 4, 3, 8, 8) + mask = torch.tensor([[0.0, 1.0, 1.0, 1.0]] * B) + cam, batch = _pair_batch(torch.rand(B, 3, 8, 8), torch.rand(B, 3, 8, 8)) + batch[f"{cam}_mem"] = window + batch[f"{cam}_mem_mask"] = mask + data = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist", f"{cam}_mem"], [], [], "actions" + ) + assert torch.equal(data["memory"], window) + assert data["memory_mask"].dtype == torch.bool + assert data["memory_mask"].tolist() == [[False, True, True, True]] * B + assert data["front_img_1"].shape == (B, 2, 1, 3, 8, 8) + assert "front_img_1_mem" not in data diff --git a/tests/unit/test_rdt_mem_configs.py b/tests/unit/test_rdt_mem_configs.py new file mode 100644 index 0000000000..cfabc58c12 --- /dev/null +++ b/tests/unit/test_rdt_mem_configs.py @@ -0,0 +1,49 @@ +"""The RDT memory-ablation arms: each model tree resolves on its own (``ModelWrapper`` +builds the model from the model config alone, so ``${data.*}`` there fails only +at train time), and the model's history / memory sizes match the keymap's.""" + +from __future__ import annotations + +import pytest +from omegaconf import OmegaConf + +from egomimic.trainHydra import _build_model_config_tree + +ARMS = ("base", "proprio", "vision", "both") + + +def _compose(arm, compose_resolve): + cfg = compose_resolve( + f"train_zarr_mecka_20pct_6d_rdt_mem_{arm}", [], keep_hydra=True + ) + return OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]) + + +@pytest.mark.parametrize("arm", ARMS) +def test_model_tree_resolves_without_the_rest_of_the_config(arm, compose_resolve): + tree = _build_model_config_tree(_compose(arm, compose_resolve)) + tree._set_flag("allow_objects", True) + OmegaConf.resolve(tree) + + +@pytest.mark.parametrize("arm", ARMS) +def test_model_sizes_match_the_keymap(arm, compose_resolve): + cfg = _compose(arm, compose_resolve) + OmegaConf.resolve(cfg) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + model = cfg.model.robomimic_model + denoiser = model.head_specs.human_bimanual.model + history = km.get("proprio_history", 1) + assert model.stem_specs.human_bimanual.state_ee_pose.history_len == history + assert denoiser.n_state_tokens == history + memory = model.trunk.get("memory") + if memory is None: + assert not km.get("image_memory") and not denoiser.get("n_memory_tokens") + return + assert memory.frames == km.image_memory + assert memory.stride_s == km.image_memory_stride_s + assert denoiser.n_memory_tokens == memory.frames + assert memory.encoder.pool == "cls" + if memory.fuse_proprio: + assert history >= memory.frames + assert km.history_stride_s == memory.stride_s diff --git a/tests/unit/test_replay_viz.py b/tests/unit/test_replay_viz.py new file mode 100644 index 0000000000..6f770c0677 --- /dev/null +++ b/tests/unit/test_replay_viz.py @@ -0,0 +1,188 @@ +"""GT / prediction replay video (``egomimic/eval/replay_viz.py``).""" + +from __future__ import annotations + +import types + +import numpy as np +from scipy.spatial.transform import Rotation + +from egomimic.eval.eval_video import EvalVideo +from egomimic.eval.replay_viz import ReplayClip, infer_stride, render_replay, rigid_fit +from egomimic.rldb.embodiment.human import Human + + +def _clip(n, horizon, seed=0): + rng = np.random.default_rng(seed) + gt = rng.normal(size=(n, horizon, 126)) * 0.05 + np.tile([0.0, 0.0, 0.5], 42) + gt = gt.astype(np.float32) + return ReplayClip( + images=np.zeros((n, 90, 160, 3), np.uint8), + gt=gt, + pred=gt + 0.01, + intrinsics=[None] * n, + ) + + +def test_rigid_fit_recovers_transform(): + rng = np.random.default_rng(0) + src = rng.normal(size=(42, 3)) + R0 = Rotation.random(random_state=1).as_matrix() + t0 = np.array([0.1, -0.2, 0.3]) + R, t = rigid_fit(src, src @ R0.T + t0) + np.testing.assert_allclose(R, R0, atol=1e-9) + np.testing.assert_allclose(t, t0, atol=1e-9) + + +def test_rigid_fit_ignores_invalid_points_and_degenerates_to_identity(): + src = np.full((42, 3), np.nan) + R, t = rigid_fit(src, src) + np.testing.assert_array_equal(R, np.eye(3)) + np.testing.assert_array_equal(t, np.zeros(3)) + + +def test_render_replays_each_complete_chunk_twice(): + horizon = 10 + frames, used = render_replay(_clip(25, horizon), Human, mode="keypoints", stride=1) + # Two complete chunks (anchors 0 and 10); frames 20-24 wait for the next clip. + assert used == 2 * horizon + assert frames.shape == (2 * 2 * horizon, 90, 160, 3) + assert frames.dtype == np.uint8 + + +def test_render_with_stride_spans_stride_frames_per_step(): + horizon, stride = 5, 3 + frames, used = render_replay( + _clip(16, horizon), Human, mode="keypoints", stride=stride + ) + assert used == horizon * stride + assert len(frames) == 2 * horizon * stride + + +def test_short_clip_renders_nothing_and_consumes_nothing(): + frames, used = render_replay(_clip(9, 10), Human, mode="keypoints", stride=1) + assert used == 0 + assert len(frames) == 0 + + +def test_tail_and_concat_round_trip(): + clip = _clip(12, 4) + joined = ReplayClip.concat([clip.tail(8), clip.tail(10)]) + assert len(joined) == 6 + np.testing.assert_array_equal( + joined.gt, np.concatenate([clip.gt[8:], clip.gt[10:]]) + ) + assert joined.intrinsics == clip.intrinsics[8:] + clip.intrinsics[10:] + + +def _moving_camera_clip(stride, n=60, horizon=12, seed=0): + """GT chunks of a hand moving in the world, seen by a rotating camera.""" + rng = np.random.default_rng(seed) + steps = n + horizon * stride + world = np.cumsum(rng.normal(scale=0.01, size=(steps, 42, 3)), axis=0) + world += rng.normal(scale=0.05, size=(42, 3)) + [0.0, 0.0, 0.5] + cams = Rotation.from_rotvec(np.cumsum(rng.normal(scale=0.02, size=(steps, 3)), 0)) + gt = np.stack( + [ + np.stack( + [ + cams[t].inv().apply(world[t + k * stride]).reshape(-1) + for k in range(horizon) + ] + ) + for t in range(n) + ] + ).astype(np.float32) + return ReplayClip(np.zeros((n, 90, 160, 3), np.uint8), gt, gt, [None] * n) + + +def test_infer_stride_recovers_the_frames_per_step(): + for stride in (1, 2, 3): + assert infer_stride(_moving_camera_clip(stride)) == stride + + +def _resampled_clip(stride, n=60, horizon=100, seed=0): + """As _moving_camera_clip, but with a hand and camera that move + continuously, so a chunk step can fall between frames.""" + rng = np.random.default_rng(seed) + base = rng.normal(scale=0.05, size=(42, 3)) + [0.0, 0.0, 0.5] + amp, freq, phase = rng.normal(scale=0.05, size=(3, 42, 3)) + rot_amp = rng.normal(scale=0.2, size=3) + + def world(t): + return base + amp * np.sin(0.2 * t * (1 + freq) + phase) + + def cam(t): + return Rotation.from_rotvec(rot_amp * np.sin(0.05 * t)) + + gt = np.stack( + [ + np.stack( + [ + cam(t).inv().apply(world(t + k * stride)).reshape(-1) + for k in range(horizon) + ] + ) + for t in range(n) + ] + ).astype(np.float32) + return ReplayClip(np.zeros((n, 90, 160, 3), np.uint8), gt, gt, [None] * n) + + +def test_infer_stride_recovers_a_resampled_chunk(): + # 30 raw frames resampled to 100 steps (mecka, action stride 1). + assert infer_stride(_resampled_clip(29 / 99)) == 29 / 99 + + +def test_render_with_fractional_stride_plays_in_real_time(): + stride, horizon = 29 / 99, 100 + frames, used = render_replay( + _resampled_clip(stride, n=70, horizon=horizon), + Human, + mode="keypoints", + stride=stride, + ) + # Anchors 0 and 30: each chunk covers frames a..a+29, once as GT, once as pred. + assert used == 60 + assert len(frames) == 2 * 2 * 30 + + +def test_infer_stride_needs_keypoints(): + clip = _clip(20, 10) + clip.gt = clip.gt[..., :12] + assert infer_stride(clip) is None + + +class _Buffering(EvalVideo): + def __init__(self, images_fn): + super().__init__(viz_max_batches=100) + self._images_fn = images_fn + + def compute_metrics_and_viz(self, batch, do_viz=True): + return {}, ({3: self._images_fn()} if do_viz else {}) + + +def test_overlay_buffer_counts_frames_not_pixel_rows(monkeypatch): + written = [] + monkeypatch.setattr( + "egomimic.eval.eval_video.tvio.write_video", + lambda path, frames, **kw: written.append(len(frames)), + ) + ev = _Buffering(lambda: np.zeros((8, 360, 640, 3), np.uint8)) + ev.trainer = types.SimpleNamespace( + is_global_zero=True, + current_epoch=0, + max_epochs=1, + global_step=0, + loggers=[], + default_root_dir="/tmp/unused", + lightning_module=types.SimpleNamespace( + device="cpu", log_dict=lambda *a, **k: None + ), + ) + monkeypatch.setattr(ev, "_upload_video", lambda key, path: None) + monkeypatch.setattr("os.makedirs", lambda *a, **k: None) + for i in range(10): + ev.on_validation_step({}, i, mode="both") + ev.on_validation_end() + assert written == [80] diff --git a/tests/unit/test_rollout_transform_mode.py b/tests/unit/test_rollout_transform_mode.py new file mode 100644 index 0000000000..793612eed7 --- /dev/null +++ b/tests/unit/test_rollout_transform_mode.py @@ -0,0 +1,71 @@ +"""The robot rollout rebuilds a checkpoint's Eva input pipeline from the +transform mode trainHydra records in the model config tree, and refuses to +guess for an unrecorded non-legacy action width.""" + +from types import SimpleNamespace + +import pytest + +import egomimic.trainHydra as train_hydra +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.robot.rollout import EMBODIMENT_MAP, PolicyRollout + + +@pytest.fixture(autouse=True) +def no_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + monkeypatch.delenv("SLURM_RESTART_COUNT", raising=False) + + +def test_config_tree_records_the_eva_transform_mode(compose_resolve, tmp_path): + cfg = compose_resolve( + "train_zarr_cartesian", ["data=eva", f"paths.output_dir={tmp_path}"] + ) + tree = train_hydra._build_model_config_tree(cfg) + assert tree.data_transform_modes.eva_bimanual == ( + cfg.data.train_datasets.eva_bimanual.resolver.transform_list.mode + ) + + +def _rollout(config_tree=None, action_width=None): + emb = EMBODIMENT_MAP["both"] + shapes = ( + {} + if action_width is None + else {emb: {"actions_cartesian": (100, action_width)}} + ) + policy = SimpleNamespace( + hparams=SimpleNamespace( + config_tree=config_tree, norm_stats_state={"shapes": shapes} + ), + _as_config=ModelWrapper._as_config, + ) + r = PolicyRollout.__new__(PolicyRollout) + r.arm, r.embodiment_id, r.policy = "both", emb, policy + return r + + +def test_recorded_mode_is_used(): + tree = { + "model": {}, + "data_transform_modes": {"eva_bimanual": "cartesian_wristframe_6d"}, + } + assert _rollout(tree, 20)._checkpoint_transform_mode() == "cartesian_wristframe_6d" + + +def test_override_wins_over_the_checkpoint(): + tree = {"model": {}, "data_transform_modes": {"eva_bimanual": "cartesian_6d"}} + r = _rollout(tree, 20) + assert r._checkpoint_transform_mode("cartesian_wristframe_6d") == ( + "cartesian_wristframe_6d" + ) + + +def test_unrecorded_legacy_width_uses_the_legacy_pipeline(): + assert _rollout({"model": {}}, 14)._checkpoint_transform_mode() is None + + +@pytest.mark.parametrize("width", [20, None]) +def test_unrecorded_other_width_refuses_to_guess(width): + with pytest.raises(ValueError, match="--transform-mode"): + _rollout({"model": {}}, width)._checkpoint_transform_mode() diff --git a/tests/unit/test_run_manifest.py b/tests/unit/test_run_manifest.py new file mode 100644 index 0000000000..a4647b79f0 --- /dev/null +++ b/tests/unit/test_run_manifest.py @@ -0,0 +1,70 @@ +import json + +import pytest + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import write_catalog # noqa: E402 + +from egomimic.rldb.lance import commit # noqa: E402 +from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 +from egomimic.rldb.multi import MultiDataset # noqa: E402 +from egomimic.utils import run_manifest # noqa: E402 + +KEY_MAP = { + "left.action_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.obs_ee_pose", + "horizon": 3, + }, +} + + +def _dataset(root, **kw): + resolver = LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}, **kw) + return MultiDataset._from_resolver(resolver, mode="total", valid_ratio=0.5) + + +def test_manifest_names_the_split_and_the_versions_read(tmp_path, caplog): + root = tmp_path / "lance" + recs = write_catalog(root) + with caplog.at_level("WARNING"): + ds = _dataset(root) + assert "untagged" in caplog.text + files = run_manifest.run_manifest({"human_bimanual": ds}, {"human_bimanual": ds}) + split = files[run_manifest.SPLIT_FILE]["human_bimanual"] + assert sorted(split["train"] + split["valid"]) == sorted(r for r, _, _ in recs) + read = files[run_manifest.LANCE_FILE]["train/human_bimanual"] + assert read["versions"]["steps/aria.lance"] == 1 + assert read["tags"]["recordings.lance"] == [] + + commit.tag(root, "t1") + caplog.clear() + with caplog.at_level("WARNING"): + tagged = _dataset(root, tag="t1") + assert "untagged" not in caplog.text + assert tagged.source_manifest["tags"]["steps/aria.lance"] == ["t1"] + + +def test_write_keeps_the_first_launchs_record(tmp_path): + run_manifest.write(tmp_path, {"split.json": {"a": 1}}) + run_manifest.write(tmp_path, {"split.json": {"a": 1}}) + assert [p.name for p in tmp_path.iterdir()] == ["split.json"] + run_manifest.write(tmp_path, {"split.json": {"a": 2}}) + assert json.loads((tmp_path / "split.json").read_text()) == {"a": 1} + assert len(list(tmp_path.glob("split.*.json"))) == 1 + + +def test_resume_warns_when_old_val_episodes_now_train(tmp_path, caplog): + run = tmp_path / "run" + (run / "checkpoints").mkdir(parents=True) + ckpt = run / "checkpoints" / "last.ckpt" + ckpt.write_bytes(b"") + run_manifest.write(run, {"split.json": {"h": {"train": ["a"], "valid": ["b"]}}}) + same = {"h": {"train": ["a", "c"], "valid": ["b"]}} + assert run_manifest.warn_on_split_drift(str(ckpt), same) == [] + moved = {"h": {"train": ["a", "b"], "valid": ["c"]}} + with caplog.at_level("WARNING"): + assert run_manifest.warn_on_split_drift(str(ckpt), moved) == ["b"] + assert "validated on" in caplog.text + assert run_manifest.warn_on_split_drift(None, moved) == [] diff --git a/tests/unit/test_text_encoder_fixes.py b/tests/unit/test_text_encoder_fixes.py new file mode 100644 index 0000000000..fda2736059 --- /dev/null +++ b/tests/unit/test_text_encoder_fixes.py @@ -0,0 +1,388 @@ +"""Lane C: the HPT text-encoder fixes. + +Three defects, one test file: + +1. ``HPTModel.finalize_modules`` used ``self.apply(self._init_weights)``, which + xavier-randomised the *pretrained* HF encoders inside the text stems (and + would do the same to any future pretrained stem). It now skips anything a + module declares through ``PretrainedWeights``. +2. ``HPT.__init__`` called ``finalize_modules`` twice. +3. ``QwenPerTokenEncoder.compute_latent`` never passed the tokenizer attention + mask to its cross-attention, so padded positions - which are the learned + constant ``proj.bias``, not zero, because the zeroing happens before the + biased ``proj`` - got real attention weight. + +Plus the new guard rail: the pretrained weights are hashed against their +on-disk snapshot once the model is built, and a mismatch fails the run. +""" + +from __future__ import annotations + +import glob +import inspect +import os + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.models.hpt_nets import ( + PolicyStem, + PretrainedWeights, + QwenPerTokenEncoder, + apply_skipping_pretrained, + hash_state_dict, + verify_pretrained_weights, +) + +EMBED_DIM = 32 + + +def _cross_attn_specs(latent: int = 4) -> OmegaConf: + return OmegaConf.create( + { + "random_horizon_masking": False, + "cross_attn": { + "crossattn_latent": latent, + "crossattn_heads": 2, + "crossattn_dim_head": 8, + "crossattn_modality_dropout": 0.0, + "modality_embed_dim": EMBED_DIM, + }, + } + ) + + +# -------------------------------------------------------------------------- +# a tiny stand-in for a pretrained stem +# -------------------------------------------------------------------------- + + +class FakePretrainedStem(PretrainedWeights, nn.Module): + """Carries the marker, an "encoder" that must never be touched, and a + ``proj`` that HPT owns and therefore must still be initialised.""" + + _hpt_pretrained_attrs = ("encoder",) + + def __init__(self) -> None: + super().__init__() + self.encoder = nn.Sequential(nn.Linear(8, 8), nn.LayerNorm(8), nn.Linear(8, 8)) + self.proj = nn.Linear(8, EMBED_DIM) + with torch.no_grad(): # distinctive "checkpoint" values, no default zeros + for param in self.encoder.parameters(): + param.normal_(mean=1.0, std=1.0) + # pretend these came off disk + self._reference = { + k: v.detach().clone() for k, v in self.pretrained_state_dict().items() + } + + def pretrained_reference_state_dict(self): + return {k: v.clone() for k, v in self._reference.items()} + + +def _small_hpt_model(stems: dict): + from egomimic.algo.hpt import HPTModel + + model = HPTModel( + embed_dim=EMBED_DIM, + num_blocks=1, + num_heads=2, + observation_horizon=1, + action_horizon=2, + token_postprocessing="action_token", + ) + model.init_domain_stem("d", stems) + return model + + +# -------------------------------------------------------------------------- +# 1 / 2. the init pass skips pretrained weights and nothing else +# -------------------------------------------------------------------------- + + +def test_apply_skipping_pretrained_matches_apply_when_nothing_is_pretrained(): + """Modules HPT builds itself must be initialised exactly as before: with no + pretrained markers in the tree the new traversal is bit-identical to + ``nn.Module.apply`` (same post-order, same RNG draws).""" + from egomimic.algo.hpt import HPTModel + + def build(): + torch.manual_seed(0) + return nn.Sequential( + nn.Linear(8, 8), + nn.LayerNorm(8), + nn.Sequential(nn.Linear(8, 4), nn.Linear(4, 4)), + ) + + def init(module): # HPTModel._init_weights only ever looks at its argument + HPTModel._init_weights(None, module) + + reference, candidate = build(), build() + torch.manual_seed(1234) + reference.apply(init) + torch.manual_seed(1234) + apply_skipping_pretrained(candidate, init) + + ref_state, cand_state = reference.state_dict(), candidate.state_dict() + assert set(ref_state) == set(cand_state) + for key in ref_state: + assert torch.equal(ref_state[key], cand_state[key]), key + + +def test_finalize_modules_leaves_pretrained_weights_untouched(): + torch.manual_seed(0) + stem = FakePretrainedStem() + before_pretrained = { + k: v.detach().clone() for k, v in stem.pretrained_state_dict().items() + } + before_proj = stem.proj.weight.detach().clone() + + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + + after = stem.pretrained_state_dict() + for key, value in before_pretrained.items(): + assert torch.equal(value, after[key]), f"pretrained tensor {key} was modified" + + # ... while the module HPT owns *was* re-initialised (xavier + zero bias). + assert not torch.equal(before_proj, stem.proj.weight) + assert torch.equal(stem.proj.bias, torch.zeros_like(stem.proj.bias)) + + +def test_finalize_modules_initialises_everything_hpt_builds(): + """Every Linear/LayerNorm outside a pretrained subtree gets the xavier + treatment (bias 0 / LayerNorm 1,0), trunk and head included.""" + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + + skipped = {id(m) for m in stem.encoder.modules()} + seen = 0 + for module in model.modules(): + if id(module) in skipped: + continue + if isinstance(module, nn.Linear) and module.bias is not None: + assert torch.equal(module.bias, torch.zeros_like(module.bias)) + seen += 1 + elif isinstance(module, nn.LayerNorm): + assert torch.equal(module.weight, torch.ones_like(module.weight)) + assert torch.equal(module.bias, torch.zeros_like(module.bias)) + seen += 1 + assert seen > 5, "expected the trunk/stem modules to be visited" + + # the pretrained LayerNorm keeps its (non-default) loaded values + ln = stem.encoder[1] + assert not torch.equal(ln.bias, torch.zeros_like(ln.bias)) + + +# -------------------------------------------------------------------------- +# 3. finalize_modules is called once per HPT.__init__ +# -------------------------------------------------------------------------- + + +def test_finalize_modules_called_once_in_hpt_init(): + """Counted by walking HPT.__init__'s AST, not its source text: a comment + mentioning the call would satisfy a substring count.""" + import ast + import textwrap + + from egomimic.algo.hpt import HPT + + tree = ast.parse(textwrap.dedent(inspect.getsource(HPT.__init__))) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "finalize_modules" + ] + assert len(calls) == 1, ( + "HPT.__init__ must call finalize_modules exactly once; a second call " + "re-runs the init pass and re-draws the shared action tokens" + ) + + +# -------------------------------------------------------------------------- +# 4. the weight hash check +# -------------------------------------------------------------------------- + + +def test_verify_pretrained_weights_passes_on_a_clean_model(): + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + checked = verify_pretrained_weights(model, verbose=False) + assert checked, "the marked stem should have been hashed" + ((name, digest),) = checked.items() + assert "annotation" in name + assert digest == hash_state_dict(stem.pretrained_reference_state_dict()) + + +def test_verify_pretrained_weights_raises_when_a_weight_is_perturbed(): + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + with torch.no_grad(): + stem.encoder[0].weight[0, 0] += 1e-3 + with pytest.raises(RuntimeError, match="do not match their checkpoint"): + verify_pretrained_weights(model, verbose=False) + + +# -------------------------------------------------------------------------- +# 5. the per-token text stem passes the attention mask +# -------------------------------------------------------------------------- + + +class StubPerTokenEncoder(QwenPerTokenEncoder): + """``QwenPerTokenEncoder`` with the HF model swapped for fixed features, so + the mask plumbing is tested without RoPE/position effects.""" + + def __init__(self, hidden: torch.Tensor, mask: torch.Tensor, specs) -> None: + PolicyStem.__init__(self, specs=specs) + self.hidden_size = hidden.shape[-1] + self.output_dim = EMBED_DIM + self.freeze_encoder = False # no HF encoder to keep in eval mode + self._snapshot_dir = "" + self._load_dtype = torch.float32 + self.proj = nn.Linear(self.hidden_size, EMBED_DIM) + self.stub = (hidden, mask) + + def _encode(self, prompts): + return self.stub + + +def _padded_stub(seq_short: int = 3, seq_long: int = 7, hidden_dim: int = 16): + torch.manual_seed(7) + hidden = torch.randn(2, seq_long, hidden_dim) + mask = torch.ones(2, seq_long, dtype=torch.long) + mask[0, : seq_long - seq_short] = 0 # left padding, as the tokenizer pads + return hidden, mask + + +def test_per_token_stem_masks_padding(): + hidden, mask = _padded_stub() + specs = _cross_attn_specs() + stem = StubPerTokenEncoder(hidden, mask, specs) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + + with torch.no_grad(): + batched = stem.compute_latent(None) + + # the same (short) prompt encoded on its own: no padding at all + real = mask[0].bool() + stem.stub = (hidden[:1, real], mask[:1, real]) + alone = stem.compute_latent(None) + + # what the parent did: cross-attention over every position, mask ignored + stem.stub = (hidden, mask) + feat = stem(None) + tokens = stem.tokens.repeat(feat.shape[0], 1, 1) + unmasked = stem.cross_attention(tokens, feat) + + masked_diff = (batched[0] - alone[0]).abs().max().item() + unmasked_diff = (unmasked[0] - alone[0]).abs().max().item() + print( + f"\n[mask] max|short-in-padded-batch - short-alone|: " + f"with mask {masked_diff:.3e} / without mask (parent) {unmasked_diff:.3e}" + ) + assert masked_diff < 1e-5, masked_diff + assert ( + unmasked_diff > 1e-3 + ), "the stub is degenerate: padding made no difference even unmasked" + + +# -------------------------------------------------------------------------- +# the same, with the real Qwen encoder (skipped when the snapshot is absent) +# -------------------------------------------------------------------------- + + +def _qwen_snapshot() -> str | None: + cache = os.environ.get("HF_HUB_CACHE") or os.path.join( + os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")), "hub" + ) + hits = sorted( + glob.glob( + os.path.join(cache, "models--Qwen--Qwen3-Embedding-0.6B", "snapshots", "*") + ) + ) + return hits[-1] if hits else None + + +QWEN_SNAPSHOT = _qwen_snapshot() +requires_qwen = pytest.mark.skipif( + QWEN_SNAPSHOT is None, reason="Qwen3-Embedding-0.6B snapshot not in the HF cache" +) + + +@pytest.fixture(scope="module") +def qwen_stem(): + if QWEN_SNAPSHOT is None: + pytest.skip("Qwen3-Embedding-0.6B snapshot not in the HF cache") + os.environ.setdefault("HF_HUB_OFFLINE", "1") + specs = _cross_attn_specs() + stem = QwenPerTokenEncoder( + model_name=QWEN_SNAPSHOT, max_length=32, output_dim=EMBED_DIM, specs=specs + ) + stem.init_cross_attn(specs.cross_attn) + return stem + + +@requires_qwen +def test_finalize_modules_keeps_the_real_qwen_encoder_bit_identical(qwen_stem): + before = hash_state_dict(qwen_stem.pretrained_state_dict()) + model = _small_hpt_model({"annotation": qwen_stem}) + model.finalize_modules() + after = hash_state_dict(qwen_stem.pretrained_state_dict()) + reference = hash_state_dict(qwen_stem.pretrained_reference_state_dict()) + print( + f"\n[qwen] sha256 after finalize_modules: {after}\n[qwen] snapshot: {reference}" + ) + assert after == before + assert after == reference + # and the check that runs at train start agrees, including after the upcast + model.float() + verify_pretrained_weights(model, verbose=False) + + +@requires_qwen +def test_verify_pretrained_weights_catches_a_perturbed_qwen_weight(qwen_stem): + param = next(iter(qwen_stem.encoder.parameters())) + original = param.detach().clone() + try: + with torch.no_grad(): + param[0, 0] += 1.0 + with pytest.raises(RuntimeError, match="do not match their checkpoint"): + verify_pretrained_weights(qwen_stem, verbose=False) + finally: + with torch.no_grad(): + param.copy_(original) + + +@requires_qwen +def test_real_qwen_per_token_stem_masks_padding(qwen_stem): + stem = qwen_stem + stem.eval() + prompts = ["pick up the cup", "pick up the red cup and place it on the blue plate"] + with torch.no_grad(): + feat, mask = stem.forward_with_mask(prompts) + latent = stem.compute_latent(prompts) + one_token = stem.tokens.repeat(1, 1, 1) + # reference: the short row's real positions only, same features + alone = stem.cross_attention(one_token, feat[0][mask[0]].unsqueeze(0)) + unmasked = stem.cross_attention(one_token, feat[:1]) + + masked_diff = (latent[0] - alone[0]).abs().max().item() + unmasked_diff = (unmasked[0] - alone[0]).abs().max().item() + n_pad = int((~mask[0]).sum()) + print( + f"\n[qwen-mask] {n_pad} padded positions; max|padded-row - real-tokens-only|: " + f"with mask {masked_diff:.3e} / without mask (parent) {unmasked_diff:.3e}" + ) + assert n_pad > 0, "prompts must differ in length for this test to mean anything" + assert masked_diff < 1e-5, masked_diff + assert unmasked_diff > 1e-3, unmasked_diff diff --git a/tests/unit/test_text_feature_cache.py b/tests/unit/test_text_feature_cache.py new file mode 100644 index 0000000000..db2678bf10 --- /dev/null +++ b/tests/unit/test_text_feature_cache.py @@ -0,0 +1,421 @@ +"""Lane C step 2: the frozen-text-feature cache on the Qwen stems. + +The annotation vocabulary is tiny compared with the number of training steps, +so a frozen text encoder recomputes the same ~600M-parameter forward over and +over. ``_Qwen3BaseEncoder`` now keeps a lazy in-memory cache of the per-token +hidden states, keyed by the exact prompt string, and re-assembles the batch +(left-padded, exactly like the tokenizer) from the cached rows. + +The cache is only legal when the encoder is frozen; a trainable encoder must +stay online (its outputs change every step and need a graph). +""" + +from __future__ import annotations + +import glob +import os +import time + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.models.hpt_nets import PolicyStem, QwenPerTokenEncoder + +EMBED_DIM = 32 +HIDDEN_DIM = 16 + + +def _cross_attn_specs(latent: int = 4) -> OmegaConf: + return OmegaConf.create( + { + "random_horizon_masking": False, + "cross_attn": { + "crossattn_latent": latent, + "crossattn_heads": 2, + "crossattn_dim_head": 8, + "crossattn_modality_dropout": 0.0, + "modality_embed_dim": EMBED_DIM, + }, + } + ) + + +# -------------------------------------------------------------------------- +# a stub encoder: deterministic per-prompt features, counts its forwards +# -------------------------------------------------------------------------- + +PROMPT_LENGTHS = {"fold the towel": 2, "pick up the cup": 3, "wipe the table": 5} + + +class CountingStubEncoder(QwenPerTokenEncoder): + """``QwenPerTokenEncoder`` with the tokenizer + HF model replaced by a fixed + prompt -> per-token-features table, so the cache can be tested without + loading Qwen. ``_encode_online`` is what the cache is supposed to avoid; + every call is counted, along with how many prompts it saw.""" + + def __init__( + self, + specs, + freeze: bool = True, + cache_text_features: bool = True, + text_cache_max_entries: int = 8192, + lengths: dict | None = None, + ) -> None: + PolicyStem.__init__(self, specs=specs) + self.hidden_size = HIDDEN_DIM + self.output_dim = EMBED_DIM + self.freeze_encoder = freeze + self._snapshot_dir = "" + self._load_dtype = torch.float32 + self.proj = nn.Linear(HIDDEN_DIM, EMBED_DIM) + self.encoder = nn.Identity() # never called; `train()` keeps it in eval + self._init_text_feature_cache(cache_text_features, text_cache_max_entries) + + generator = torch.Generator().manual_seed(0) + self._table = { + prompt: torch.randn(length, HIDDEN_DIM, generator=generator) + for prompt, length in (lengths or PROMPT_LENGTHS).items() + } + self.calls = [] # one entry per _encode_online call: the prompt list + + def _encode_online(self, prompts): + self.calls.append(list(prompts)) + rows = [self._table[p] for p in prompts] + longest = max(row.shape[0] for row in rows) + hidden = torch.zeros(len(rows), longest, HIDDEN_DIM) + mask = torch.zeros(len(rows), longest, dtype=torch.long) + for i, row in enumerate(rows): # left padding, as padding_side="left" + hidden[i, longest - row.shape[0] :] = row + mask[i, longest - row.shape[0] :] = 1 + return hidden, mask + + +def _stub(**kwargs) -> CountingStubEncoder: + specs = _cross_attn_specs() + stem = CountingStubEncoder(specs, **kwargs) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + return stem + + +A, B, C = list(PROMPT_LENGTHS) + + +# -------------------------------------------------------------------------- +# 1. misses are encoded once, in one sub-batch of unique strings +# -------------------------------------------------------------------------- + + +def test_cache_encodes_each_unique_prompt_exactly_once(): + stem = _stub() + with torch.no_grad(): + stem._encode([A, B]) + stem._encode([B, C, A]) + stem._encode([C]) + stem._encode([A, A, B, C]) + + # one call per batch that had misses, each carrying only the unique misses + assert stem.calls == [[A, B], [C]], stem.calls + stats = stem.text_feature_cache_stats() + assert stats["entries"] == 3 + assert stats["misses"] == 3 + assert stats["hits"] == 2 + 3 + 1 + 4 - 3 # every prompt seen, minus misses + + +def test_cached_output_equals_the_uncached_output(): + """Bit-identical, since the stub features do not depend on the batch.""" + batches = [[A, B], [B, C, A], [C], [A, A, C]] + cached, online = _stub(), _stub(cache_text_features=False) + online.load_state_dict(cached.state_dict()) # same proj / cross-attn weights + with torch.no_grad(): + for prompts in batches: + hid_c, mask_c = cached._encode(prompts) + hid_o, mask_o = online._encode(prompts) + assert hid_c.shape == hid_o.shape + assert torch.equal(mask_c.bool(), mask_o.bool()) + assert torch.equal(hid_c, hid_o), prompts + # and through the whole stem, cross-attention included + assert torch.equal( + cached.compute_latent(prompts), online.compute_latent(prompts) + ) + # the uncached stem re-encodes every batch (twice here: _encode + + # compute_latent), the cached one only ever saw the misses + assert online.calls == [b for b in batches for _ in range(2)] + assert cached.calls == [[A, B], [C]] + + +# -------------------------------------------------------------------------- +# 2. the re-assembled batch keeps the tokenizer's left padding +# -------------------------------------------------------------------------- + + +def test_reassembled_batch_is_left_padded(): + stem = _stub() + prompts = [C, A, B] # lengths 5, 2, 3 + with torch.no_grad(): + stem._encode([A]) # warm one entry so the batch mixes hits and misses + hidden, mask = stem._encode(prompts) + + lengths = [PROMPT_LENGTHS[p] for p in prompts] + longest = max(lengths) + assert hidden.shape == (len(prompts), longest, HIDDEN_DIM) + assert mask.shape == (len(prompts), longest) + for i, length in enumerate(lengths): + pad = longest - length + assert mask[i, :pad].sum() == 0, "padding must be on the LEFT" + assert mask[i, pad:].min() == 1 + assert torch.equal(hidden[i, pad:], stem._table[prompts[i]]) + assert torch.equal(hidden[i, :pad], torch.zeros(pad, HIDDEN_DIM)) + + +def test_cached_tensors_are_detached(): + stem = _stub(freeze=True) + with torch.no_grad(): + hidden, _ = stem._encode([A, B]) + assert not hidden.requires_grad + for entry in stem._text_feature_cache.values(): + assert not entry.requires_grad + assert entry.grad_fn is None + + +# -------------------------------------------------------------------------- +# 3. bypasses +# -------------------------------------------------------------------------- + + +def test_trainable_encoder_bypasses_the_cache(): + stem = _stub(freeze=False) + stem._encode([A, B]) + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats()["entries"] == 0 + + +def test_cache_text_features_false_bypasses_the_cache(): + stem = _stub(cache_text_features=False) + with torch.no_grad(): + stem._encode([A, B]) + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats()["entries"] == 0 + + +def test_clear_text_feature_cache(): + stem = _stub() + with torch.no_grad(): + stem._encode([A, B]) + stem.clear_text_feature_cache() + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats() == {"entries": 2, "hits": 0, "misses": 4} + + +# -------------------------------------------------------------------------- +# 4. LRU eviction +# -------------------------------------------------------------------------- + + +def test_lru_eviction_at_max_entries(): + stem = _stub(text_cache_max_entries=2) + with torch.no_grad(): + stem._encode([A, B]) # cache: A, B + stem._encode([A]) # hit; A becomes most-recent, B is the LRU victim + stem._encode([C]) # evicts B + assert set(stem._text_feature_cache) == {A, C} + stem._encode([A, C]) # both still cached: no new forward + assert stem.calls == [[A, B], [C]] + stem._encode([B]) # B was evicted -> re-encoded + assert stem.calls == [[A, B], [C], [B]] + assert stem.text_feature_cache_stats()["entries"] == 2 + + +def test_cache_never_exceeds_max_entries(): + lengths = {f"prompt {i}": 1 + (i % 4) for i in range(20)} + stem = _stub(text_cache_max_entries=5, lengths=lengths) + with torch.no_grad(): + for prompt in lengths: + stem._encode([prompt]) + assert len(stem._text_feature_cache) <= 5 + + +# -------------------------------------------------------------------------- +# 5. the real Qwen encoder (skipped when the snapshot is absent) +# -------------------------------------------------------------------------- + + +def _qwen_snapshot() -> str | None: + cache = os.environ.get("HF_HUB_CACHE") or os.path.join( + os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")), "hub" + ) + hits = sorted( + glob.glob( + os.path.join(cache, "models--Qwen--Qwen3-Embedding-0.6B", "snapshots", "*") + ) + ) + return hits[-1] if hits else None + + +QWEN_SNAPSHOT = _qwen_snapshot() +requires_qwen = pytest.mark.skipif( + QWEN_SNAPSHOT is None, reason="Qwen3-Embedding-0.6B snapshot not in the HF cache" +) + +QWEN_PROMPTS = [ + "fold the towel", + "pick up the red cup and place it on the blue plate", + "wipe", + "open the drawer and take out the blue sponge, then close the drawer", +] + + +def _build_qwen_stems(dtype: str = "float16"): + """Two stems off the same snapshot: one cached, one always online.""" + if QWEN_SNAPSHOT is None: + pytest.skip("Qwen3-Embedding-0.6B snapshot not in the HF cache") + os.environ.setdefault("HF_HUB_OFFLINE", "1") + specs = _cross_attn_specs() + built = [] + for cache_text_features in (True, False): + stem = QwenPerTokenEncoder( + model_name=QWEN_SNAPSHOT, + max_length=32, + output_dim=EMBED_DIM, + dtype=dtype, + cache_text_features=cache_text_features, + specs=specs, + ) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + built.append(stem) + # identical cross-attention / projection weights on both stems + built[1].load_state_dict(built[0].state_dict()) + return tuple(built) + + +@pytest.fixture(scope="module") +def qwen_stems(): + return _build_qwen_stems("float16") + + +@pytest.fixture(scope="module") +def qwen_stems_fp32(): + return _build_qwen_stems("float32") + + +@requires_qwen +def test_real_qwen_cached_matches_online(qwen_stems): + cached, online = qwen_stems + forwards = [] + original = type(cached)._encode_online + + def counting(self, prompts): + forwards.append(list(prompts)) + return original(self, prompts) + + with torch.no_grad(): + hid_on, mask_on = online._encode(QWEN_PROMPTS) + lat_on = online.compute_latent(QWEN_PROMPTS) + + # Warm one prompt on its own first, so the batch below is a mix of a + # hit (encoded with NO padding) and misses (encoded padded to each + # other) - the case where cached and online padding really differ. + cached.clear_text_feature_cache() + cached._encode([QWEN_PROMPTS[0]]) + cached._encode_online = counting.__get__(cached, type(cached)) + hid_c1, mask_c1 = cached._encode(QWEN_PROMPTS) + lat_c1 = cached.compute_latent(QWEN_PROMPTS) + first_calls = list(forwards) + n_first = len(forwards) + forwards.clear() + lat_c2 = cached.compute_latent(QWEN_PROMPTS) # fully warm + + assert n_first == 1, "the misses must go through the encoder in ONE sub-batch" + assert first_calls[0] == QWEN_PROMPTS[1:], "only the misses, unique, in order" + assert forwards == [], "a warm cache must perform zero encoder forwards" + assert torch.equal(mask_c1.bool(), mask_on.bool()) + assert hid_c1.shape == hid_on.shape + assert torch.equal(lat_c1, lat_c2), "the cache must be deterministic once warm" + + # Only the REAL token positions are comparable: the cache pads with zeros + # while the online path leaves the encoder's output for the pad tokens + # there. Both are masked out downstream (forward_with_mask), so this is the + # tensor the rest of the stem actually reads. + real = mask_on.bool() + hidden_diff = (hid_c1[real].float() - hid_on[real].float()).abs().max().item() + latent_diff = (lat_c1 - lat_on).abs().max().item() + print( + f"\n[qwen-cache] max|cached - online| hidden {hidden_diff:.3e} " + f"latent {latent_diff:.3e} (entries {cached.text_feature_cache_stats()})" + ) + # The cache encodes a miss with different left padding than the online + # batch would, and Qwen's RoPE positions are absolute, so fp16 rounding + # differs. The right yardstick is how much the ONLINE path itself moves + # under the same change of padding: encode one prompt alone and compare it + # with its row in the padded batch. + with torch.no_grad(): + solo = online._encode([QWEN_PROMPTS[1]])[0][0].float() + control = (solo - hid_on[1][mask_on[1].bool()].float()).abs().max().item() + print(f"[qwen-cache] online-vs-online padding noise (fp16): {control:.3e}") + assert hidden_diff <= max(2 * control, 1e-6), (hidden_diff, control) + # absolute ceilings: hidden states run to |x| ~ 27 here, so fp16's ~1e-3 + # relative precision is worth a few 1e-2 in absolute terms. + assert hidden_diff < 1e-1, hidden_diff + assert latent_diff < 5e-3, latent_diff + + +@requires_qwen +def test_real_qwen_cached_matches_online_in_fp32(qwen_stems_fp32): + """The same comparison without fp16 rounding: the cache is then exact to + within ordinary float noise (the brief's 1e-2 / 1e-3 bounds).""" + cached, online = qwen_stems_fp32 + with torch.no_grad(): + hid_on, mask_on = online._encode(QWEN_PROMPTS) + lat_on = online.compute_latent(QWEN_PROMPTS) + cached.clear_text_feature_cache() + cached._encode([QWEN_PROMPTS[0]]) # mixed hit / miss padding again + hid_c, _ = cached._encode(QWEN_PROMPTS) + lat_c = cached.compute_latent(QWEN_PROMPTS) + real = mask_on.bool() + hidden_diff = (hid_c[real] - hid_on[real]).abs().max().item() + latent_diff = (lat_c - lat_on).abs().max().item() + print( + f"\n[qwen-cache-fp32] max|cached - online| hidden {hidden_diff:.3e} " + f"latent {latent_diff:.3e}" + ) + assert hidden_diff < 1e-2, hidden_diff + assert latent_diff < 1e-3, latent_diff + + +@requires_qwen +def test_real_qwen_cache_is_faster_when_warm(qwen_stems): + cached, online = qwen_stems + cached.clear_text_feature_cache() + with torch.no_grad(): + cached.compute_latent(QWEN_PROMPTS) # warm + start = time.perf_counter() + cached.compute_latent(QWEN_PROMPTS) + warm = time.perf_counter() - start + start = time.perf_counter() + online.compute_latent(QWEN_PROMPTS) + live = time.perf_counter() - start + print(f"\n[qwen-cache] warm {warm * 1e3:.1f} ms vs online {live * 1e3:.1f} ms") + assert warm < live + + +def test_a_batch_wider_than_the_cache_still_serves_every_prompt(): + """The LRU trim runs after the miss pass, so a batch with more unique + prompts than the capacity evicts its own earliest entries; reading them + back must not KeyError.""" + stem = _stub(text_cache_max_entries=2) + with torch.no_grad(): + hidden, mask = stem._encode([A, B, C]) + online_hidden, online_mask = stem._encode_online([A, B, C]) + assert hidden.shape[0] == 3 and mask.shape[0] == 3 + assert len(stem._text_feature_cache) <= 2 + for i in range(3): + real, real_on = mask[i].bool(), online_mask[i].bool() + torch.testing.assert_close( + hidden[i][real], online_hidden[i][real_on], atol=0, rtol=0 + ) diff --git a/tests/unit/test_train_perf_wiring.py b/tests/unit/test_train_perf_wiring.py new file mode 100644 index 0000000000..b55bcd4e7e --- /dev/null +++ b/tests/unit/test_train_perf_wiring.py @@ -0,0 +1,87 @@ +"""The torch.compile knob: off unless asked, in-place so checkpoints survive.""" + +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.utils.compile_utils import compile_modules + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(4, 4) + + def forward(self, x): + return self.lin(x) + + +def test_compile_keeps_state_dict_keys(): + """nn.Module.compile is in-place; torch.compile(module) would prefix every + key with `_orig_mod.` and break checkpoints written before the flag.""" + module = _Tiny() + before = list(module.state_dict()) + compile_modules([("tiny", module)]) + assert list(module.state_dict()) == before + + +def test_compile_modules_skips_non_modules_and_reports_names(): + """A head without a sub-network hands back None (or something that is not a + Module); compiling it must be a skip, not a crash.""" + assert compile_modules([("a", _Tiny()), ("b", None), ("c", object())]) == ["a"] + + +class _SpyAlgo: + def __init__(self): + self.calls = [] + self.nets = nn.ModuleDict({"policy": _Tiny()}) + + def compile_for_training(self, mode=None, dynamic=False): + self.calls.append((mode, dynamic)) + return ["policy"] + + +def _wrapper_with(compile_cfg, monkeypatch): + algo = _SpyAlgo() + monkeypatch.setattr(ModelWrapper, "_instantiate_model", lambda *a, **k: algo) + model = {"robomimic_model": {}} + if compile_cfg is not None: + model["compile"] = compile_cfg + ModelWrapper(config_tree=OmegaConf.create({"model": model}), norm_stats_state=None) + return algo + + +def test_compile_is_off_unless_asked(monkeypatch): + assert _wrapper_with(None, monkeypatch).calls == [] + assert _wrapper_with({"enabled": False, "mode": None}, monkeypatch).calls == [] + + +def test_compile_forwards_mode_and_dynamic(monkeypatch): + cfg = {"enabled": True, "mode": "reduce-overhead", "dynamic": True} + assert _wrapper_with(cfg, monkeypatch).calls == [("reduce-overhead", True)] + + +def test_enable_grad_norm_reaches_the_wrapper(): + """It used to be dropped between the config and ModelWrapper, so a model's + `enable_grad_norm: false` did nothing.""" + import inspect + + from egomimic import trainHydra + + src = inspect.getsource(trainHydra) + assert 'enable_grad_norm=cfg.model.get("enable_grad_norm", True)' in src + + +def test_pi_clips_like_openpi_unless_the_trainer_says_otherwise(compose_resolve): + from egomimic.trainHydra import _model_trainer_defaults + + cfg = compose_resolve("train_zarr_cartesian", ["model=pi0.5_base"]) + assert cfg.model.gradient_clip_val == 1.0 + assert cfg.model.get("enable_grad_norm", True) is True + assert _model_trainer_defaults(cfg) == {"gradient_clip_val": 1.0} + + cfg.trainer.gradient_clip_val = 0.5 + assert _model_trainer_defaults(cfg) == {} + + hpt = compose_resolve("train_zarr_cartesian", []) + assert _model_trainer_defaults(hpt) == {} diff --git a/tests/unit/test_train_step.py b/tests/unit/test_train_step.py index e5fbe91141..75a08e4ebe 100644 --- a/tests/unit/test_train_step.py +++ b/tests/unit/test_train_step.py @@ -9,12 +9,14 @@ import os import pytest +import torch from fixtures.recipes import ( RECIPES, common_overrides, cpu_trainer_overrides, hpt_small_overrides, pi_cpu_overrides, + rdt_small_overrides, ) from fixtures.train_harness import ( STATE_DIM, @@ -43,6 +45,26 @@ def _run(cfg) -> float: return loss +def _expected_keys(embodiment: str) -> set[str]: + """Batch keys every recipe must hand the algo: eva trains on the cartesian + action, human on the wrist-frame hand-keypoint action (plus the cartesian + ee_pose proprio the keypoint data configs emit for Pi's prompt).""" + if embodiment == "eva_bimanual": + return { + "actions_cartesian", + "observations.state.ee_pose", + Eva.VIZ_IMAGE_KEY, + "observations.images.right_wrist_img", + "observations.images.left_wrist_img", + } + return { + "actions_keypoints", + "observations.state.keypoints", + "observations.state.ee_pose", + Human.VIZ_IMAGE_KEY, + } + + @pytest.mark.parametrize("vendor", VENDOR_NAMES) def test_hpt_train_step(tmp_path, monkeypatch, vendor): hermetic_env(monkeypatch) @@ -65,19 +87,86 @@ def test_hpt_train_step(tmp_path, monkeypatch, vendor): spy = BatchKeySpy(monkeypatch, HPT) _run(cfg) seen = spy.keys[recipe.embodiment] - is_eva = recipe.embodiment == "eva_bimanual" - front = Eva.VIZ_IMAGE_KEY if is_eva else Human.VIZ_IMAGE_KEY - expected = {"actions_cartesian", "observations.state.ee_pose", front} - if is_eva: - expected |= { - "observations.images.right_wrist_img", - "observations.images.left_wrist_img", - } + expected = _expected_keys(recipe.embodiment) assert ( expected <= seen ), f"{vendor}/hpt missing {expected - seen}; saw {sorted(seen)}" +def test_hpt_builds_the_policy_on_cpu_while_gpus_are_visible(tmp_path, monkeypatch): + """Under submitit every rank sees all 8 GPUs, so an algo that picks its own + construction device puts all 8 policies on cuda:0 and OOMs at 3B. Lightning + owns the move, so nothing may leave CPU in ``__init__``.""" + hermetic_env(monkeypatch) + recipe = RECIPES[("aria", "hpt")] + data, out, hashes = write_fixtures(tmp_path, "aria") + cfg = compose_recipe( + recipe, + 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, + ) + + built = {} + init = HPT.__init__ + + def spy(self, *args, **kwargs): + # Only __init__ sees the GPUs: faking them for the whole run would send + # the rest of the stack down CUDA paths this CPU tier cannot take. + with monkeypatch.context() as gpus_visible: + gpus_visible.setattr(torch.cuda, "is_available", lambda: True) + gpus_visible.setattr(torch.cuda, "device_count", lambda: 8) + init(self, *args, **kwargs) + built["param_devices"] = {p.device.type for p in self.nets.parameters()} + built["algo_device"] = self.device + built["policy_device"] = self.nets["policy"].device + + monkeypatch.setattr(HPT, "__init__", spy) + _run(cfg) + + assert built["param_devices"] == {"cpu"} + assert built["algo_device"] is None + assert built["policy_device"] is None + + +def test_rdt_train_step(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + recipe = RECIPES[("mecka", "rdt")] + data, out, hashes = write_fixtures(tmp_path, "mecka") + cfg = compose_recipe( + recipe, + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + + cpu_trainer_overrides(STEPS) + + rdt_small_overrides(recipe.embodiment), + out, + ) + spy = BatchKeySpy(monkeypatch, HPT) + _run(cfg) + expected = { + "actions_keypoints", + "observations.state.ee_pose", + Human.VIZ_IMAGE_KEY, + f"{Human.VIZ_IMAGE_KEY}_hist", + "fps", + } + assert expected <= spy.keys[recipe.embodiment] + + @pytest.mark.parametrize("vendor", VENDOR_NAMES) def test_pi_train_step(tmp_path, monkeypatch, vendor): reason = pi_unavailable() @@ -118,13 +207,7 @@ def test_pi_train_step(tmp_path, monkeypatch, vendor): # The wrapper sees dataset camera names; StubPI0 asserts it was handed # openpi's slot names, which is the remap working end to end. seen = spy.keys[recipe.embodiment] - front = Eva.VIZ_IMAGE_KEY if is_eva else Human.VIZ_IMAGE_KEY - expected = {"actions_cartesian", "observations.state.ee_pose", front, "annotations"} - if is_eva: - expected |= { - "observations.images.left_wrist_img", - "observations.images.right_wrist_img", - } + expected = _expected_keys(recipe.embodiment) | {"annotations"} assert ( expected <= seen ), f"{vendor}/pi missing {expected - seen}; saw {sorted(seen)}" diff --git a/tests/unit/test_vision_backbone.py b/tests/unit/test_vision_backbone.py new file mode 100644 index 0000000000..e8d643d888 --- /dev/null +++ b/tests/unit/test_vision_backbone.py @@ -0,0 +1,199 @@ +"""The frozen-backbone eval-mode fix and the separate backbone LR group. + +Covers ``ResNet`` (unchanged forward, BatchNorm in eval when frozen, +``backbone_parameters``), ``ModelWrapper.configure_optimizers`` param groups, +and the flagship model config. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torchvision +from omegaconf import OmegaConf +from torch import nn + +from egomimic.models.hpt_nets import ResNet +from egomimic.pl_utils.pl_model import ModelWrapper + +# -------------------------------------------------------------------------- +# ResNet +# -------------------------------------------------------------------------- + + +def _resnet_stem(**kwargs): + torch.manual_seed(0) + return ResNet(output_dim=16, **kwargs) + + +def test_resnet_forward_matches_reference_implementation(): + """The baseline image path is bit-identical: the stem's output equals the + torchvision backbone followed by the stem's projection.""" + stem = _resnet_stem().eval() + x = torch.rand(2, 1, 1, 3, 96, 128) + with torch.no_grad(): + out = stem(x) + + backbone = nn.Sequential( + *list(torchvision.models.resnet18(weights="DEFAULT").children())[:-2] + ).eval() + with torch.no_grad(): + feat = backbone(x.view(-1, 3, 96, 128)) + expected = stem.proj(feat.view(2, feat.shape[1], -1).transpose(1, 2)) + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +def test_resnet_frozen_backbone_batchnorm_stays_in_eval(): + """0.15: a frozen ResNet must not update / use batch statistics.""" + stem = _resnet_stem(freeze_backbone=True) + stem.train() + assert stem.training + bns = [ + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ] + assert bns + assert all(not m.training for m in bns) + assert all(not p.requires_grad for p in stem.net.parameters()) + # an unfrozen stem keeps the default behaviour + trainable = _resnet_stem() + trainable.train() + assert all( + m.training + for m in trainable.net.modules() + if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + + +def test_resnet_frozen_backbone_running_stats_do_not_drift(): + stem = _resnet_stem(freeze_backbone=True) + stem.train() + bn = next( + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + before = bn.running_mean.clone() + with torch.no_grad(): + stem(torch.rand(2, 1, 1, 3, 96, 128) * 5.0) + assert torch.equal(bn.running_mean, before) + + +def test_resnet_backbone_parameters_are_the_net(): + stem = _resnet_stem() + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + assert all(id(p) not in backbone for p in stem.proj.parameters()) + + +def test_resnet_backbone_parameters_with_multiple_copies(): + stem = ResNet(output_dim=16, num_of_copy=2) + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + + +# -------------------------------------------------------------------------- +# configure_optimizers +# -------------------------------------------------------------------------- + + +class _StubEncoder(nn.Module): + """Encoder with a pretrained-style backbone plus a trainable projection.""" + + def __init__(self, frozen: bool = False): + super().__init__() + self.net = nn.Linear(4, 4) + self.proj = nn.Linear(4, 4) + if frozen: + for p in self.net.parameters(): + p.requires_grad = False + + def backbone_parameters(self): + return list(self.net.parameters()) + + +def _stub_wrapper(backbone_lr_scale=None, frozen=False): + policy = nn.Module() + policy.encoders = nn.ModuleDict( + { + "front_img_1": _StubEncoder(frozen=frozen), + "other": nn.Linear(4, 4), # no backbone_parameters -> main group + } + ) + policy.trunk = nn.Linear(4, 4) + root = nn.Module() + root.nets = nn.ModuleDict({"policy": policy}) + + model_cfg = { + "optimizer": { + "_target_": "torch.optim.AdamW", + "_partial_": True, + "lr": 5e-5, + "weight_decay": 0.0001, + }, + "scheduler": None, + } + if backbone_lr_scale is not None: + model_cfg["backbone_lr_scale"] = backbone_lr_scale + config_tree = OmegaConf.create({"model": model_cfg}) + # A LightningModule cannot be built without a trainer / datamodule here, so + # call the real methods against the attributes they actually read. + stub = SimpleNamespace( + hparams=SimpleNamespace(config_tree=config_tree), + _as_config=ModelWrapper._as_config, + model=root, + trainer=SimpleNamespace(model=root), + ) + stub._backbone_param_groups = ModelWrapper._backbone_param_groups.__get__(stub) + return stub, root + + +@pytest.mark.parametrize("scale", [None, 1.0]) +def test_configure_optimizers_single_group_by_default(scale): + stub, _ = _stub_wrapper(backbone_lr_scale=scale) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 1 + assert groups[0]["lr"] == pytest.approx(5e-5) + + +def test_configure_optimizers_backbone_group(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 2 + by_lr = {round(g["lr"], 12): g for g in groups} + assert set(by_lr) == {round(5e-5, 12), round(5e-6, 12)} + + encoder = root.nets["policy"].encoders["front_img_1"] + backbone_ids = {id(p) for p in encoder.backbone_parameters()} + got_backbone = {id(p) for p in by_lr[round(5e-6, 12)]["params"]} + main = {id(p) for p in by_lr[round(5e-5, 12)]["params"]} + assert got_backbone == backbone_ids + assert not (got_backbone & main) + trainable = {id(p) for p in root.parameters() if p.requires_grad} + assert got_backbone | main == trainable + + +def test_configure_optimizers_skips_frozen_backbone_params(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1, frozen=True) + out = ModelWrapper.configure_optimizers(stub) + every = [p for g in out["optimizer"].param_groups for p in g["params"]] + assert all(p.requires_grad for p in every) + assert {id(p) for p in every} == { + id(p) for p in root.parameters() if p.requires_grad + } + + +# -------------------------------------------------------------------------- +# Configs +# -------------------------------------------------------------------------- + +FLAGSHIP = "train_zarr_mecka_flagship_6d_hpt" + + +def test_base_model_config_declares_unit_backbone_lr_scale(compose_resolve): + cfg = compose_resolve(FLAGSHIP, []) + assert cfg.model.backbone_lr_scale == 1.0 + encoder = cfg.model.robomimic_model.encoder_specs.front_img_1 + assert encoder._target_ == "egomimic.models.hpt_nets.ResNet" diff --git a/tests/unit/test_worker_orphan_guard.py b/tests/unit/test_worker_orphan_guard.py new file mode 100644 index 0000000000..27fce77db4 --- /dev/null +++ b/tests/unit/test_worker_orphan_guard.py @@ -0,0 +1,217 @@ +"""Dataloader workers must die with their rank. + +Workers are forked after CUDA init, so they inherit the rank's ``/dev/nvidia*`` +fds. One survivor keeps every GPU context on the node alive, and +``proctrack/linuxproc`` will not sweep it because it has been reparented to init. +""" + +import os +import subprocess +import sys +import textwrap +import time + +import pytest + +from egomimic.utils import gpu_orphans +from egomimic.utils.gpu_orphans import ( + arm_rank_pdeathsig, + die_with_parent, + find_orphans, + orphan_guarded, + reap_orphans, +) + +# A worker that inherits both ends of its result pipe never gets an EPIPE when +# the rank dies; it blocks in pipe_write, so PyTorch's own ppid check -- which +# only runs between index_queue timeouts -- is never reached. +RANK = textwrap.dedent( + """ + import os, sys, time + sys.path.insert(0, {root!r}) + from egomimic.utils.gpu_orphans import die_with_parent + guard = sys.argv[1] == "guard" + r, w = os.pipe() + if os.fork() == 0: + if guard: + die_with_parent(0) + with open(sys.argv[2], "w") as f: + f.write(str(os.getpid())) + while True: + os.write(w, b"x" * 65536) + time.sleep(600) + """ +) + + +def _alive(pid: int) -> bool: + return os.path.isdir(f"/proc/{pid}") + + +def _pid_written(pidfile) -> bool: + return pidfile.exists() and pidfile.read_text().strip().isdigit() + + +# The wait channel of a writer blocked on a full pipe: split out of pipe_wait in +# Linux 5.5, so a 5.4 kernel (sky2) still reports the old name. +_PIPE_BLOCKED = {"pipe_write", "pipe_wait"} + + +def _run_rank(tmp_path, mode: str) -> int: + """Start a rank, SIGKILL it once its worker is wedged, return the worker pid.""" + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + pidfile = tmp_path / f"pid_{mode}" + rank = subprocess.Popen( + [sys.executable, "-c", RANK.format(root=root), mode, str(pidfile)] + ) + try: + deadline = time.time() + 180 # the rank imports torch first + while time.time() < deadline and not _pid_written(pidfile): + time.sleep(0.2) + assert _pid_written(pidfile), "worker never started" + worker = int(pidfile.read_text()) + deadline = time.time() + 30 + wchan = "" + while time.time() < deadline: + wchan = open(f"/proc/{worker}/wchan").read() + if wchan in _PIPE_BLOCKED: + break + time.sleep(0.2) + else: + os.kill(worker, 9) # an unguarded worker would outlive the test + pytest.skip(f"this kernel reports the blocked worker as {wchan!r}") + finally: + rank.kill() + rank.wait() + return worker + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_unguarded_worker_outlives_its_rank(tmp_path): + worker = _run_rank(tmp_path, "noguard") + time.sleep(3) + try: + assert _alive(worker), "baseline no longer reproduces the leak" + finally: + os.kill(worker, 9) + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_guarded_worker_dies_with_its_rank(tmp_path): + worker = _run_rank(tmp_path, "guard") + deadline = time.time() + 15 + while time.time() < deadline and _alive(worker): + time.sleep(0.2) + if _alive(worker): + os.kill(worker, 9) + pytest.fail("worker survived its rank despite PR_SET_PDEATHSIG") + + +def test_orphan_guard_runs_before_a_caller_init(): + calls = [] + params = orphan_guarded( + {"num_workers": 2, "worker_init_fn": lambda i: calls.append(("user", i))} + ) + params["worker_init_fn"](3) + assert calls == [("user", 3)] + + +def test_orphan_guard_is_a_noop_without_workers(): + assert orphan_guarded({"num_workers": 0}) == {"num_workers": 0} + + +def test_orphan_guard_installs_the_default(): + assert orphan_guarded({"num_workers": 4})["worker_init_fn"] is die_with_parent + + +def _fake_proc(root, pid: int, *, ppid: int, fds: list[str]) -> None: + d = root / str(pid) + (d / "fd").mkdir(parents=True) + (d / "status").write_text(f"Name:\tpt_data_worker\nPPid:\t{ppid}\n") + for i, target in enumerate(fds): + os.symlink(target, d / "fd" / str(i)) + + +def test_find_orphans_picks_only_ours_orphaned_and_on_a_gpu(tmp_path): + uid = os.getuid() + gpu = ["/dev/nvidia0", "/dev/nvidiactl"] + _fake_proc(tmp_path, 100, ppid=1, fds=gpu) # the leak + _fake_proc(tmp_path, 101, ppid=42, fds=gpu) # live job's worker + _fake_proc(tmp_path, 102, ppid=1, fds=["/dev/null"]) # no GPU + assert find_orphans(str(tmp_path), uid=uid) == [100] + + +def test_find_orphans_skips_other_users(tmp_path): + uid = os.getuid() + _fake_proc(tmp_path, 200, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid + 1) == [] + + +def test_find_orphans_tolerates_a_process_exiting_mid_scan(tmp_path): + uid = os.getuid() + (tmp_path / "300").mkdir() # no status/fd: raced away + _fake_proc(tmp_path, 301, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid) == [301] + + +def test_reap_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert reap_orphans() == [] + + +def test_reap_runs_once_per_node(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "3") + assert reap_orphans() == [] + + +@pytest.mark.parametrize( + "value, on", + [("0", False), ("false", False), ("1", True), ("true", True), ("yes", True)], +) +def test_reap_honours_the_kill_switch(monkeypatch, value, on): + """Any truthy spelling keeps it on. Never the real /proc: a regression here + would SIGKILL the tester's own orphans.""" + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "0") + monkeypatch.setenv("EGOMIMIC_REAP_GPU_ORPHANS", value) + monkeypatch.setattr(gpu_orphans, "find_orphans", lambda: [4242]) + assert reap_orphans(dry_run=True) == ([4242] if on else []) + + +def test_reap_rejects_a_misspelled_switch(monkeypatch): + import argparse + + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("EGOMIMIC_REAP_GPU_ORPHANS", "of") + monkeypatch.setattr(gpu_orphans, "find_orphans", lambda: [4242]) + with pytest.raises(argparse.ArgumentTypeError): + reap_orphans(dry_run=True) + + +def test_rank_pdeathsig_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert arm_rank_pdeathsig() is False + + +def test_rank_pdeathsig_arms_under_slurm(monkeypatch): + import ctypes + + monkeypatch.setenv("SLURM_JOB_ID", "1") + try: + assert arm_rank_pdeathsig() is True + finally: + # do not leave pytest itself armed + ctypes.CDLL("libc.so.6").prctl(1, 0) + + +def test_rank_pdeathsig_leaves_a_nohup_run_alone(monkeypatch): + """nohup inside salloc means "outlive this shell"; arming would kill it.""" + import signal + + monkeypatch.setenv("SLURM_JOB_ID", "1") + previous = signal.signal(signal.SIGHUP, signal.SIG_IGN) + try: + assert arm_rank_pdeathsig() is False + finally: + signal.signal(signal.SIGHUP, previous) diff --git a/tests/unit/test_wrist6d_roundtrip.py b/tests/unit/test_wrist6d_roundtrip.py index ff98b4857b..749977caac 100644 --- a/tests/unit/test_wrist6d_roundtrip.py +++ b/tests/unit/test_wrist6d_roundtrip.py @@ -22,20 +22,24 @@ 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, ) from egomimic.utils.action_utils import ( - HumanBimanualCartesianEuler, - RobotBimanualCartesianEuler, + HumanBimanualCartesian6D, + RobotBimanualCartesian6D, _apply_norm_one, _apply_unnorm_one, _matrix_to_ypr, _reconstruct_R_from_cols, _ypr_to_matrix, ) -from egomimic.utils.pose_utils import _rot6d_to_ypr, _ypr_to_rot6d +from egomimic.utils.pose_utils import ( + _rot6d_to_ypr, + _ypr_to_rot6d, + rot6d_channels, +) T = 100 @@ -245,15 +249,17 @@ def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): st_act = _quantile_stats(act6, 18) st_obs = _quantile_stats(obs6, 20) a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() - a_n = _apply_norm_one(a_t, st_act, "quantile") - o_n = _apply_norm_one(o_t, st_obs, "quantile") - conv = HumanBimanualCartesianEuler() + act_rot = rot6d_channels("actions_cartesian", a_t.shape[-1]) + obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) + a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) + o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) + conv = HumanBimanualCartesian6D() a32 = conv.to32_norm_6d(a_n) assert a32.shape == (B, T, 32) assert torch.all(a32[..., [9, 19]] == 0) and torch.all(a32[..., 20:] == 0) torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) - a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") - o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile", act_rot) + o_un = _apply_unnorm_one(o_n, st_obs, "quantile", obs_rot) torch.testing.assert_close(a_un, a_t, atol=1e-5, rtol=0) torch.testing.assert_close(o_un, o_t, atol=1e-5, rtol=0) @@ -268,7 +274,9 @@ def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): _assert_pose12_close(out["observations.state.ee_pose"], gt_obs, atol=1e-4) # (5) a noisy (non-orthonormal) prediction still reverts to finite poses - a_noisy = _apply_unnorm_one(a_n + 0.05 * torch.randn_like(a_n), st_act, "quantile") + a_noisy = _apply_unnorm_one( + a_n + 0.05 * torch.randn_like(a_n), st_act, "quantile", act_rot + ) out_n = Embodiment.apply_transform( {"actions_cartesian": a_noisy, "observations.state.ee_pose": o_un}, rev ) @@ -336,13 +344,15 @@ def test_eva_wristframe_6d_pipeline_round_trips_to_camframe(): st_act, st_obs = _quantile_stats(act6, 20), _quantile_stats(obs6, 20) a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() - a_n = _apply_norm_one(a_t, st_act, "quantile") - o_n = _apply_norm_one(o_t, st_obs, "quantile") - conv = RobotBimanualCartesianEuler() + act_rot = rot6d_channels("actions_cartesian", a_t.shape[-1]) + obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) + a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) + o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) + conv = RobotBimanualCartesian6D() a32 = conv.to32_norm_6d(a_n) torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) - a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") - o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile", act_rot) + o_un = _apply_unnorm_one(o_n, st_obs, "quantile", obs_rot) rev = _build_eva_cartesian_revert_6d_wristframe_transform_list() out = Embodiment.apply_transform( diff --git a/uv.lock b/uv.lock index 13130da7b9..4adf54eb7f 100644 --- a/uv.lock +++ b/uv.lock @@ -6,10 +6,10 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'emscripten'", "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version < '3.12' and sys_platform == 'win32'", "python_full_version < '3.12' and sys_platform == 'emscripten'", @@ -1030,6 +1030,7 @@ dependencies = [ { name = "datasets" }, { name = "dm-control" }, { name = "einops" }, + { name = "foxglove-schemas-protobuf" }, { name = "geomloss" }, { name = "gpustat" }, { name = "h5py" }, @@ -1038,6 +1039,8 @@ dependencies = [ { name = "ipython" }, { name = "lightning" }, { name = "matplotlib" }, + { name = "mcap" }, + { name = "mcap-protobuf-support" }, { name = "mediapy" }, { name = "mink" }, { name = "mujoco" }, @@ -1053,6 +1056,7 @@ dependencies = [ { name = "projectaria-tools", extra = ["all"] }, { name = "psycopg", extra = ["binary"] }, { name = "pyarrow" }, + { name = "pylance" }, { name = "pynvml" }, { name = "pyquaternion" }, { name = "pytest" }, @@ -1092,6 +1096,7 @@ requires-dist = [ { name = "datasets", specifier = "==4.0.0" }, { name = "dm-control", specifier = "==1.0.8" }, { name = "einops" }, + { name = "foxglove-schemas-protobuf", specifier = "==0.4.0" }, { name = "geomloss" }, { name = "gpustat" }, { name = "h5py" }, @@ -1100,6 +1105,8 @@ requires-dist = [ { name = "ipython" }, { name = "lightning" }, { name = "matplotlib" }, + { name = "mcap", specifier = "==1.5.0" }, + { name = "mcap-protobuf-support", specifier = "==0.5.4" }, { name = "mediapy" }, { name = "mink", specifier = "==0.0.13" }, { name = "mujoco", specifier = "==3.4.0" }, @@ -1114,6 +1121,7 @@ requires-dist = [ { name = "projectaria-tools", extras = ["all"], specifier = "==2.0.0" }, { name = "psycopg", extras = ["binary"] }, { name = "pyarrow" }, + { name = "pylance", specifier = "==13.0.0" }, { name = "pynvml" }, { name = "pyquaternion" }, { name = "pytest" }, @@ -1243,6 +1251,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/ba/56147c165442cc5ba7e82ecf301c9a68353cede498185869e6e02b4c264f/fonttools-4.62.1-py3-none-any.whl", hash = "sha256:7487782e2113861f4ddcc07c3436450659e3caa5e470b27dc2177cade2d8e7fd", size = 1152647, upload-time = "2026-03-13T13:54:22.735Z" }, ] +[[package]] +name = "foxglove-schemas-protobuf" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/d8/7ca265b4721e2ae912cd0fa5cc96d90357b5ae1ae15c5dff653bf9f22ec5/foxglove_schemas_protobuf-0.4.0.tar.gz", hash = "sha256:3fc9610d7ec6cc4b4029f746ab9699a4ad396bdd7551b99d2a37509d13ef3cc2", size = 31026, upload-time = "2026-07-29T20:46:47.162Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/ff/67f49b46e40dca6fb5e5077c3f665c07155f5a878cf349185957d075e999/foxglove_schemas_protobuf-0.4.0-py3-none-any.whl", hash = "sha256:e6a35df574f4c10ea455e8660abd09af028ea00d90e01c78f8696972e2decc95", size = 98809, upload-time = "2026-07-29T20:46:45.863Z" }, +] + [[package]] name = "fqdn" version = "1.5.1" @@ -1864,7 +1884,7 @@ name = "jinxed" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ansicon", marker = "sys_platform == 'win32'" }, + { name = "ansicon" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cb/eb/8821ce6e7386e96355f2c6be83944925b4a0870572896fd33a4e61b8aa5a/jinxed-2.0.0.tar.gz", hash = "sha256:64b960b8f8d9966e1b2e7cb57ea2b1aa0a8d7e68045b96dc7ef58201e43a1209", size = 127118, upload-time = "2026-05-08T21:25:25.917Z" } wheels = [ @@ -2377,6 +2397,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/0f/13f0d54305e66c14c90512f3682f713273ec9aa94d107be7947157b37a74/labmaze-1.0.6-cp312-cp312-win_amd64.whl", hash = "sha256:5af997598cc46b1929d1c5a1febc32fd56c75874fe481a2a5982c65cee8450c9", size = 4811813, upload-time = "2023-10-04T17:20:30.837Z" }, ] +[[package]] +name = "lance-namespace" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lance-namespace-urllib3-client" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bf/93/da5f7fcac690db9b282a3439ed9e34960c147619a0d6e1f4eb8cd240e7a5/lance_namespace-0.11.1.tar.gz", hash = "sha256:f67cfbbe0647b7cb42f23b673e7edf8a75b7d8a047265a916492f8d247ee1bc2", size = 11631, upload-time = "2026-08-18T17:40:06.294Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/bc/601f2b3cc4cfa0070d858a33223bc823fffdd7981a25c45984a5216ca952/lance_namespace-0.11.1-py3-none-any.whl", hash = "sha256:07643fce9a42ad4d58cc8bf91e3f592bc7f4cbd8d0ad5233223506debf67551c", size = 13507, upload-time = "2026-08-18T17:40:03.561Z" }, +] + +[[package]] +name = "lance-namespace-urllib3-client" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dateutil" }, + { name = "typing-extensions" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/c5/2bdd0ff98b469894c8a73be809d26ffdad5402517b0e5f9e758026cba29e/lance_namespace_urllib3_client-0.11.1.tar.gz", hash = "sha256:145a9e9424d7597487249b5b95ee274423bf2910e1a9160b6a07b676b61ea46a", size = 237345, upload-time = "2026-08-18T17:40:07.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/eaaefd55d1190291207049fedc6b3eb22b506e57d6de91bae46bbaaa9c60/lance_namespace_urllib3_client-0.11.1-py3-none-any.whl", hash = "sha256:36537f529294da6d884ba0fe783704483f0a75463497c7705fd083a4d0257990", size = 406311, upload-time = "2026-08-18T17:40:04.842Z" }, +] + [[package]] name = "lark" version = "1.3.1" @@ -2549,6 +2596,54 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/44/3ee09a5b60cb44c4f2fbc1c9015cfd6ff5afc08f991cab295d3024dcbf2d/lxml-6.1.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:7da13bb6fbadfafb474e0226a30570a3445cfd47c86296f2446dafbd77079ace", size = 3508860, upload-time = "2026-04-18T04:32:48.619Z" }, ] +[[package]] +name = "lz4" +version = "4.4.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/51/f1b86d93029f418033dddf9b9f79c8d2641e7454080478ee2aab5123173e/lz4-4.4.5.tar.gz", hash = "sha256:5f0b9e53c1e82e88c10d7c180069363980136b9d7a8306c4dca4f760d60c39f0", size = 172886, upload-time = "2025-11-03T13:02:36.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/5b/6edcd23319d9e28b1bedf32768c3d1fd56eed8223960a2c47dacd2cec2af/lz4-4.4.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d6da84a26b3aa5da13a62e4b89ab36a396e9327de8cd48b436a3467077f8ccd4", size = 207391, upload-time = "2025-11-03T13:01:36.644Z" }, + { url = "https://files.pythonhosted.org/packages/34/36/5f9b772e85b3d5769367a79973b8030afad0d6b724444083bad09becd66f/lz4-4.4.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:61d0ee03e6c616f4a8b69987d03d514e8896c8b1b7cc7598ad029e5c6aedfd43", size = 207146, upload-time = "2025-11-03T13:01:37.928Z" }, + { url = "https://files.pythonhosted.org/packages/04/f4/f66da5647c0d72592081a37c8775feacc3d14d2625bbdaabd6307c274565/lz4-4.4.5-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:33dd86cea8375d8e5dd001e41f321d0a4b1eb7985f39be1b6a4f466cd480b8a7", size = 1292623, upload-time = "2025-11-03T13:01:39.341Z" }, + { url = "https://files.pythonhosted.org/packages/85/fc/5df0f17467cdda0cad464a9197a447027879197761b55faad7ca29c29a04/lz4-4.4.5-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:609a69c68e7cfcfa9d894dc06be13f2e00761485b62df4e2472f1b66f7b405fb", size = 1279982, upload-time = "2025-11-03T13:01:40.816Z" }, + { url = "https://files.pythonhosted.org/packages/25/3b/b55cb577aa148ed4e383e9700c36f70b651cd434e1c07568f0a86c9d5fbb/lz4-4.4.5-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:75419bb1a559af00250b8f1360d508444e80ed4b26d9d40ec5b09fe7875cb989", size = 1368674, upload-time = "2025-11-03T13:01:42.118Z" }, + { url = "https://files.pythonhosted.org/packages/fb/31/e97e8c74c59ea479598e5c55cbe0b1334f03ee74ca97726e872944ed42df/lz4-4.4.5-cp311-cp311-win32.whl", hash = "sha256:12233624f1bc2cebc414f9efb3113a03e89acce3ab6f72035577bc61b270d24d", size = 88168, upload-time = "2025-11-03T13:01:43.282Z" }, + { url = "https://files.pythonhosted.org/packages/18/47/715865a6c7071f417bef9b57c8644f29cb7a55b77742bd5d93a609274e7e/lz4-4.4.5-cp311-cp311-win_amd64.whl", hash = "sha256:8a842ead8ca7c0ee2f396ca5d878c4c40439a527ebad2b996b0444f0074ed004", size = 99491, upload-time = "2025-11-03T13:01:44.167Z" }, + { url = "https://files.pythonhosted.org/packages/14/e7/ac120c2ca8caec5c945e6356ada2aa5cfabd83a01e3170f264a5c42c8231/lz4-4.4.5-cp311-cp311-win_arm64.whl", hash = "sha256:83bc23ef65b6ae44f3287c38cbf82c269e2e96a26e560aa551735883388dcc4b", size = 91271, upload-time = "2025-11-03T13:01:45.016Z" }, + { url = "https://files.pythonhosted.org/packages/1b/ac/016e4f6de37d806f7cc8f13add0a46c9a7cfc41a5ddc2bc831d7954cf1ce/lz4-4.4.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:df5aa4cead2044bab83e0ebae56e0944cc7fcc1505c7787e9e1057d6d549897e", size = 207163, upload-time = "2025-11-03T13:01:45.895Z" }, + { url = "https://files.pythonhosted.org/packages/8d/df/0fadac6e5bd31b6f34a1a8dbd4db6a7606e70715387c27368586455b7fc9/lz4-4.4.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6d0bf51e7745484d2092b3a51ae6eb58c3bd3ce0300cf2b2c14f76c536d5697a", size = 207150, upload-time = "2025-11-03T13:01:47.205Z" }, + { url = "https://files.pythonhosted.org/packages/b7/17/34e36cc49bb16ca73fb57fbd4c5eaa61760c6b64bce91fcb4e0f4a97f852/lz4-4.4.5-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7b62f94b523c251cf32aa4ab555f14d39bd1a9df385b72443fd76d7c7fb051f5", size = 1292045, upload-time = "2025-11-03T13:01:48.667Z" }, + { url = "https://files.pythonhosted.org/packages/90/1c/b1d8e3741e9fc89ed3b5f7ef5f22586c07ed6bb04e8343c2e98f0fa7ff04/lz4-4.4.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c3ea562c3af274264444819ae9b14dbbf1ab070aff214a05e97db6896c7597e", size = 1279546, upload-time = "2025-11-03T13:01:50.159Z" }, + { url = "https://files.pythonhosted.org/packages/55/d9/e3867222474f6c1b76e89f3bd914595af69f55bf2c1866e984c548afdc15/lz4-4.4.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24092635f47538b392c4eaeff14c7270d2c8e806bf4be2a6446a378591c5e69e", size = 1368249, upload-time = "2025-11-03T13:01:51.273Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e7/d667d337367686311c38b580d1ca3d5a23a6617e129f26becd4f5dc458df/lz4-4.4.5-cp312-cp312-win32.whl", hash = "sha256:214e37cfe270948ea7eb777229e211c601a3e0875541c1035ab408fbceaddf50", size = 88189, upload-time = "2025-11-03T13:01:52.605Z" }, + { url = "https://files.pythonhosted.org/packages/a5/0b/a54cd7406995ab097fceb907c7eb13a6ddd49e0b231e448f1a81a50af65c/lz4-4.4.5-cp312-cp312-win_amd64.whl", hash = "sha256:713a777de88a73425cf08eb11f742cd2c98628e79a8673d6a52e3c5f0c116f33", size = 99497, upload-time = "2025-11-03T13:01:53.477Z" }, + { url = "https://files.pythonhosted.org/packages/6a/7e/dc28a952e4bfa32ca16fa2eb026e7a6ce5d1411fcd5986cd08c74ec187b9/lz4-4.4.5-cp312-cp312-win_arm64.whl", hash = "sha256:a88cbb729cc333334ccfb52f070463c21560fca63afcf636a9f160a55fac3301", size = 91279, upload-time = "2025-11-03T13:01:54.419Z" }, + { url = "https://files.pythonhosted.org/packages/2f/46/08fd8ef19b782f301d56a9ccfd7dafec5fd4fc1a9f017cf22a1accb585d7/lz4-4.4.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:6bb05416444fafea170b07181bc70640975ecc2a8c92b3b658c554119519716c", size = 207171, upload-time = "2025-11-03T13:01:56.595Z" }, + { url = "https://files.pythonhosted.org/packages/8f/3f/ea3334e59de30871d773963997ecdba96c4584c5f8007fd83cfc8f1ee935/lz4-4.4.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:b424df1076e40d4e884cfcc4c77d815368b7fb9ebcd7e634f937725cd9a8a72a", size = 207163, upload-time = "2025-11-03T13:01:57.721Z" }, + { url = "https://files.pythonhosted.org/packages/41/7b/7b3a2a0feb998969f4793c650bb16eff5b06e80d1f7bff867feb332f2af2/lz4-4.4.5-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:216ca0c6c90719731c64f41cfbd6f27a736d7e50a10b70fad2a9c9b262ec923d", size = 1292136, upload-time = "2025-11-03T13:02:00.375Z" }, + { url = "https://files.pythonhosted.org/packages/89/d1/f1d259352227bb1c185288dd694121ea303e43404aa77560b879c90e7073/lz4-4.4.5-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:533298d208b58b651662dd972f52d807d48915176e5b032fb4f8c3b6f5fe535c", size = 1279639, upload-time = "2025-11-03T13:02:01.649Z" }, + { url = "https://files.pythonhosted.org/packages/d2/fb/ba9256c48266a09012ed1d9b0253b9aa4fe9cdff094f8febf5b26a4aa2a2/lz4-4.4.5-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:451039b609b9a88a934800b5fc6ee401c89ad9c175abf2f4d9f8b2e4ef1afc64", size = 1368257, upload-time = "2025-11-03T13:02:03.35Z" }, + { url = "https://files.pythonhosted.org/packages/a5/6d/dee32a9430c8b0e01bbb4537573cabd00555827f1a0a42d4e24ca803935c/lz4-4.4.5-cp313-cp313-win32.whl", hash = "sha256:a5f197ffa6fc0e93207b0af71b302e0a2f6f29982e5de0fbda61606dd3a55832", size = 88191, upload-time = "2025-11-03T13:02:04.406Z" }, + { url = "https://files.pythonhosted.org/packages/18/e0/f06028aea741bbecb2a7e9648f4643235279a770c7ffaf70bd4860c73661/lz4-4.4.5-cp313-cp313-win_amd64.whl", hash = "sha256:da68497f78953017deb20edff0dba95641cc86e7423dfadf7c0264e1ac60dc22", size = 99502, upload-time = "2025-11-03T13:02:05.886Z" }, + { url = "https://files.pythonhosted.org/packages/61/72/5bef44afb303e56078676b9f2486f13173a3c1e7f17eaac1793538174817/lz4-4.4.5-cp313-cp313-win_arm64.whl", hash = "sha256:c1cfa663468a189dab510ab231aad030970593f997746d7a324d40104db0d0a9", size = 91285, upload-time = "2025-11-03T13:02:06.77Z" }, + { url = "https://files.pythonhosted.org/packages/49/55/6a5c2952971af73f15ed4ebfdd69774b454bd0dc905b289082ca8664fba1/lz4-4.4.5-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:67531da3b62f49c939e09d56492baf397175ff39926d0bd5bd2d191ac2bff95f", size = 207348, upload-time = "2025-11-03T13:02:08.117Z" }, + { url = "https://files.pythonhosted.org/packages/4e/d7/fd62cbdbdccc35341e83aabdb3f6d5c19be2687d0a4eaf6457ddf53bba64/lz4-4.4.5-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a1acbbba9edbcbb982bc2cac5e7108f0f553aebac1040fbec67a011a45afa1ba", size = 207340, upload-time = "2025-11-03T13:02:09.152Z" }, + { url = "https://files.pythonhosted.org/packages/77/69/225ffadaacb4b0e0eb5fd263541edd938f16cd21fe1eae3cd6d5b6a259dc/lz4-4.4.5-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:a482eecc0b7829c89b498fda883dbd50e98153a116de612ee7c111c8bcf82d1d", size = 1293398, upload-time = "2025-11-03T13:02:10.272Z" }, + { url = "https://files.pythonhosted.org/packages/c6/9e/2ce59ba4a21ea5dc43460cba6f34584e187328019abc0e66698f2b66c881/lz4-4.4.5-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e099ddfaa88f59dd8d36c8a3c66bd982b4984edf127eb18e30bb49bdba68ce67", size = 1281209, upload-time = "2025-11-03T13:02:12.091Z" }, + { url = "https://files.pythonhosted.org/packages/80/4f/4d946bd1624ec229b386a3bc8e7a85fa9a963d67d0a62043f0af0978d3da/lz4-4.4.5-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2af2897333b421360fdcce895c6f6281dc3fab018d19d341cf64d043fc8d90d", size = 1369406, upload-time = "2025-11-03T13:02:13.683Z" }, + { url = "https://files.pythonhosted.org/packages/02/a2/d429ba4720a9064722698b4b754fb93e42e625f1318b8fe834086c7c783b/lz4-4.4.5-cp313-cp313t-win32.whl", hash = "sha256:66c5de72bf4988e1b284ebdd6524c4bead2c507a2d7f172201572bac6f593901", size = 88325, upload-time = "2025-11-03T13:02:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/4b/85/7ba10c9b97c06af6c8f7032ec942ff127558863df52d866019ce9d2425cf/lz4-4.4.5-cp313-cp313t-win_amd64.whl", hash = "sha256:cdd4bdcbaf35056086d910d219106f6a04e1ab0daa40ec0eeef1626c27d0fddb", size = 99643, upload-time = "2025-11-03T13:02:15.978Z" }, + { url = "https://files.pythonhosted.org/packages/77/4d/a175459fb29f909e13e57c8f475181ad8085d8d7869bd8ad99033e3ee5fa/lz4-4.4.5-cp313-cp313t-win_arm64.whl", hash = "sha256:28ccaeb7c5222454cd5f60fcd152564205bcb801bd80e125949d2dfbadc76bbd", size = 91504, upload-time = "2025-11-03T13:02:17.313Z" }, + { url = "https://files.pythonhosted.org/packages/63/9c/70bdbdb9f54053a308b200b4678afd13efd0eafb6ddcbb7f00077213c2e5/lz4-4.4.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:c216b6d5275fc060c6280936bb3bb0e0be6126afb08abccde27eed23dead135f", size = 207586, upload-time = "2025-11-03T13:02:18.263Z" }, + { url = "https://files.pythonhosted.org/packages/b6/cb/bfead8f437741ce51e14b3c7d404e3a1f6b409c440bad9b8f3945d4c40a7/lz4-4.4.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c8e71b14938082ebaf78144f3b3917ac715f72d14c076f384a4c062df96f9df6", size = 207161, upload-time = "2025-11-03T13:02:19.286Z" }, + { url = "https://files.pythonhosted.org/packages/e7/18/b192b2ce465dfbeabc4fc957ece7a1d34aded0d95a588862f1c8a86ac448/lz4-4.4.5-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:9b5e6abca8df9f9bdc5c3085f33ff32cdc86ed04c65e0355506d46a5ac19b6e9", size = 1292415, upload-time = "2025-11-03T13:02:20.829Z" }, + { url = "https://files.pythonhosted.org/packages/67/79/a4e91872ab60f5e89bfad3e996ea7dc74a30f27253faf95865771225ccba/lz4-4.4.5-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b84a42da86e8ad8537aabef062e7f661f4a877d1c74d65606c49d835d36d668", size = 1279920, upload-time = "2025-11-03T13:02:22.013Z" }, + { url = "https://files.pythonhosted.org/packages/f1/01/d52c7b11eaa286d49dae619c0eec4aabc0bf3cda7a7467eb77c62c4471f3/lz4-4.4.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bba042ec5a61fa77c7e380351a61cb768277801240249841defd2ff0a10742f", size = 1368661, upload-time = "2025-11-03T13:02:23.208Z" }, + { url = "https://files.pythonhosted.org/packages/f7/da/137ddeea14c2cb86864838277b2607d09f8253f152156a07f84e11768a28/lz4-4.4.5-cp314-cp314-win32.whl", hash = "sha256:bd85d118316b53ed73956435bee1997bd06cc66dd2fa74073e3b1322bd520a67", size = 90139, upload-time = "2025-11-03T13:02:24.301Z" }, + { url = "https://files.pythonhosted.org/packages/18/2c/8332080fd293f8337779a440b3a143f85e374311705d243439a3349b81ad/lz4-4.4.5-cp314-cp314-win_amd64.whl", hash = "sha256:92159782a4502858a21e0079d77cdcaade23e8a5d252ddf46b0652604300d7be", size = 101497, upload-time = "2025-11-03T13:02:25.187Z" }, + { url = "https://files.pythonhosted.org/packages/ca/28/2635a8141c9a4f4bc23f5135a92bbcf48d928d8ca094088c962df1879d64/lz4-4.4.5-cp314-cp314-win_arm64.whl", hash = "sha256:d994b87abaa7a88ceb7a37c90f547b8284ff9da694e6afcfaa8568d739faf3f7", size = 93812, upload-time = "2025-11-03T13:02:26.133Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -2711,6 +2806,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" }, ] +[[package]] +name = "mcap" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lz4" }, + { name = "zstandard" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/26/16/2ac1cbeaa524453d616050db9fda14010eda507fd8130e0676879bb8faf9/mcap-1.5.0.tar.gz", hash = "sha256:9c385cc5e5a6bccff4aa0c6814b305faa2c3239ece46778661eaa6d809b1768a", size = 23110, upload-time = "2026-09-24T21:45:44.714Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/59/caa4d60a58f129659ea5403e6fb52808f87005dfcbaa03e84f0e07f0033e/mcap-1.5.0-py3-none-any.whl", hash = "sha256:44ba129d381abdca474fbf5bcee036db87d7075d9059b139bf2fee2975c8bd52", size = 21021, upload-time = "2026-09-24T21:45:43.51Z" }, +] + +[[package]] +name = "mcap-protobuf-support" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mcap" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cc/8e/98af824adce25082f1c7fb8e8aab48b1678a370f97633a55bfb06c40e3fd/mcap_protobuf_support-0.5.4.tar.gz", hash = "sha256:3af2de2c2dbda9d1dbb1a526d2ad0cf3604ab44947dd1c99a52a2f990b1eb308", size = 7059, upload-time = "2025-12-24T21:22:57.716Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/74/ee/cb4beace2bbaf3877848dbc555366c5a72ec1dabba5d620bf746143d91c4/mcap_protobuf_support-0.5.4-py3-none-any.whl", hash = "sha256:072f83b7b147b5faa25df21c175cbf6e53bf959f8a6aea534167055bc4526649", size = 7286, upload-time = "2025-12-24T21:22:56.393Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -3309,7 +3430,7 @@ name = "nvidia-cudnn-cu12" version = "9.5.1.17" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/2a/78/4535c9c7f859a64781e43c969a3a7e84c54634e319a996d43ef32ce46f83/nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2", size = 570988386, upload-time = "2024-10-25T19:54:26.39Z" }, @@ -3320,7 +3441,7 @@ name = "nvidia-cufft-cu12" version = "11.3.0.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8f/16/73727675941ab8e6ffd86ca3a4b7b47065edcca7a997920b831f8147c99d/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5", size = 200221632, upload-time = "2024-11-20T17:41:32.357Z" }, @@ -3349,9 +3470,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f0/6e/c2cf12c9ff8b872e92b4a5740701e51ff17689c4d726fca91875b07f655d/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c", size = 158229790, upload-time = "2024-11-20T17:43:43.211Z" }, @@ -3363,7 +3484,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/06/1e/b8b7c2f4099a37b96af5c9bb158632ea9e5d9d27d7391d7eb8fc45236674/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73", size = 216561367, upload-time = "2024-11-20T17:44:54.824Z" }, @@ -3572,7 +3693,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "ptyprocess" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -3707,7 +3828,7 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12'" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/ae/677391c5e75858b5496dcdf4146952646c13a8b31c04090fa98c37a04198/positional_encodings-6.0.3.tar.gz", hash = "sha256:c144ad52e78f09d16effcbfabdf3e84050b56a37e342248476e18bbffd747b24", size = 6626, upload-time = "2024-01-04T18:17:44.133Z" } wheels = [ @@ -3716,7 +3837,7 @@ wheels = [ [package.optional-dependencies] pytorch = [ - { name = "torch", marker = "python_full_version < '3.12'" }, + { name = "torch" }, ] [[package]] @@ -3728,14 +3849,14 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'emscripten'", "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12'" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/04/93a78a33088fe9115602883a39f8395cdb5561388a24e65f57c11920c638/positional_encodings-6.0.4.tar.gz", hash = "sha256:435db9596a73759e7caa8f677ebf6d2ee4d39187112244e4d8b6d7b62bc119ad", size = 6814, upload-time = "2024-10-23T07:33:48.273Z" } wheels = [ @@ -3744,7 +3865,7 @@ wheels = [ [package.optional-dependencies] pytorch = [ - { name = "torch", marker = "python_full_version >= '3.12'" }, + { name = "torch" }, ] [[package]] @@ -4294,6 +4415,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, ] +[[package]] +name = "pylance" +version = "13.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lance-namespace" }, + { name = "numpy" }, + { name = "pyarrow" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/12/fa8b39d84bfac672fd44d1369b31069021829e585ca565d49d33cedc90a1/pylance-13.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:38cbe8d204785e697909e8faca91adb3b00a2f6a865c425987cb232d0865c010", size = 74507904, upload-time = "2026-10-07T07:00:00.369Z" }, + { url = "https://files.pythonhosted.org/packages/27/e1/0399a1dc66664ed6d66fc4cd7fdaeb457c4dd5b2ea536241643388a888fb/pylance-13.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:22a37a0af446964e79cfd6046f9ea2735fcb42f8020eab63947b2bdca04989c2", size = 78945779, upload-time = "2026-10-07T07:03:56.389Z" }, + { url = "https://files.pythonhosted.org/packages/7d/72/7ba2a773a9fc3be815f39fae39f1f26b87e1a82aa7adf7cc748cb294ea36/pylance-13.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2c94649a35161c6100ed822ac022756bfefb840b7e01cf491505c5245186f7b4", size = 83169863, upload-time = "2026-10-07T07:19:35.436Z" }, + { url = "https://files.pythonhosted.org/packages/03/48/81ffda7fb308a87e81416f5f5ca67507abd8f0470c7820f68e9b32881260/pylance-13.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:22e47adcc2c7300ff876fb398ee9c625973932163ec055adae14871f8a09d7b0", size = 78969946, upload-time = "2026-10-07T07:05:09.987Z" }, + { url = "https://files.pythonhosted.org/packages/77/4c/8734e6c12500521cc92e3594715cb5d58cd770938127458e4b7252a17a92/pylance-13.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a54496c4e6c01a8c3d49479fce98474c95a265515091d173aa9e16df271d837d", size = 83164662, upload-time = "2026-10-07T07:17:16.085Z" }, + { url = "https://files.pythonhosted.org/packages/ee/39/7ff19ec586f460f7851f96e07ee20e67815806c19eccc08155aa110a7d0d/pylance-13.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:8a340dcf750171dd6386db0b6ed303bb22594e1c1b266b7fe149ea1684bf4a9a", size = 89711200, upload-time = "2026-10-07T07:07:25.895Z" }, +] + [[package]] name = "pynvml" version = "13.0.1" @@ -5641,7 +5780,7 @@ name = "triton" version = "3.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "setuptools", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "setuptools" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937, upload-time = "2025-05-29T23:39:44.182Z" }, @@ -6163,3 +6302,77 @@ sdist = { url = "https://files.pythonhosted.org/packages/30/21/093488dfc7cc8964d wheels = [ { url = "https://files.pythonhosted.org/packages/08/8a/0861bec20485572fbddf3dfba2910e38fe249796cb73ecdeb74e07eeb8d3/zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc", size = 10378, upload-time = "2026-04-13T23:21:45.386Z" }, ] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/83/c3ca27c363d104980f1c9cee1101cc8ba724ac8c28a033ede6aab89585b1/zstandard-0.25.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:933b65d7680ea337180733cf9e87293cc5500cc0eb3fc8769f4d3c88d724ec5c", size = 795254, upload-time = "2025-09-14T22:16:26.137Z" }, + { url = "https://files.pythonhosted.org/packages/ac/4d/e66465c5411a7cf4866aeadc7d108081d8ceba9bc7abe6b14aa21c671ec3/zstandard-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a3f79487c687b1fc69f19e487cd949bf3aae653d181dfb5fde3bf6d18894706f", size = 640559, upload-time = "2025-09-14T22:16:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/12/56/354fe655905f290d3b147b33fe946b0f27e791e4b50a5f004c802cb3eb7b/zstandard-0.25.0-cp311-cp311-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:0bbc9a0c65ce0eea3c34a691e3c4b6889f5f3909ba4822ab385fab9057099431", size = 5348020, upload-time = "2025-09-14T22:16:29.523Z" }, + { url = "https://files.pythonhosted.org/packages/3b/13/2b7ed68bd85e69a2069bcc72141d378f22cae5a0f3b353a2c8f50ef30c1b/zstandard-0.25.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:01582723b3ccd6939ab7b3a78622c573799d5d8737b534b86d0e06ac18dbde4a", size = 5058126, upload-time = "2025-09-14T22:16:31.811Z" }, + { url = "https://files.pythonhosted.org/packages/c9/dd/fdaf0674f4b10d92cb120ccff58bbb6626bf8368f00ebfd2a41ba4a0dc99/zstandard-0.25.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5f1ad7bf88535edcf30038f6919abe087f606f62c00a87d7e33e7fc57cb69fcc", size = 5405390, upload-time = "2025-09-14T22:16:33.486Z" }, + { url = "https://files.pythonhosted.org/packages/0f/67/354d1555575bc2490435f90d67ca4dd65238ff2f119f30f72d5cde09c2ad/zstandard-0.25.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:06acb75eebeedb77b69048031282737717a63e71e4ae3f77cc0c3b9508320df6", size = 5452914, upload-time = "2025-09-14T22:16:35.277Z" }, + { url = "https://files.pythonhosted.org/packages/bb/1f/e9cfd801a3f9190bf3e759c422bbfd2247db9d7f3d54a56ecde70137791a/zstandard-0.25.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9300d02ea7c6506f00e627e287e0492a5eb0371ec1670ae852fefffa6164b072", size = 5559635, upload-time = "2025-09-14T22:16:37.141Z" }, + { url = "https://files.pythonhosted.org/packages/21/88/5ba550f797ca953a52d708c8e4f380959e7e3280af029e38fbf47b55916e/zstandard-0.25.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:bfd06b1c5584b657a2892a6014c2f4c20e0db0208c159148fa78c65f7e0b0277", size = 5048277, upload-time = "2025-09-14T22:16:38.807Z" }, + { url = "https://files.pythonhosted.org/packages/46/c0/ca3e533b4fa03112facbe7fbe7779cb1ebec215688e5df576fe5429172e0/zstandard-0.25.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:f373da2c1757bb7f1acaf09369cdc1d51d84131e50d5fa9863982fd626466313", size = 5574377, upload-time = "2025-09-14T22:16:40.523Z" }, + { url = "https://files.pythonhosted.org/packages/12/9b/3fb626390113f272abd0799fd677ea33d5fc3ec185e62e6be534493c4b60/zstandard-0.25.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6c0e5a65158a7946e7a7affa6418878ef97ab66636f13353b8502d7ea03c8097", size = 4961493, upload-time = "2025-09-14T22:16:43.3Z" }, + { url = "https://files.pythonhosted.org/packages/cb/d3/23094a6b6a4b1343b27ae68249daa17ae0651fcfec9ed4de09d14b940285/zstandard-0.25.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:c8e167d5adf59476fa3e37bee730890e389410c354771a62e3c076c86f9f7778", size = 5269018, upload-time = "2025-09-14T22:16:45.292Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a7/bb5a0c1c0f3f4b5e9d5b55198e39de91e04ba7c205cc46fcb0f95f0383c1/zstandard-0.25.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:98750a309eb2f020da61e727de7d7ba3c57c97cf6213f6f6277bb7fb42a8e065", size = 5443672, upload-time = "2025-09-14T22:16:47.076Z" }, + { url = "https://files.pythonhosted.org/packages/27/22/503347aa08d073993f25109c36c8d9f029c7d5949198050962cb568dfa5e/zstandard-0.25.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:22a086cff1b6ceca18a8dd6096ec631e430e93a8e70a9ca5efa7561a00f826fa", size = 5822753, upload-time = "2025-09-14T22:16:49.316Z" }, + { url = "https://files.pythonhosted.org/packages/e2/be/94267dc6ee64f0f8ba2b2ae7c7a2df934a816baaa7291db9e1aa77394c3c/zstandard-0.25.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:72d35d7aa0bba323965da807a462b0966c91608ef3a48ba761678cb20ce5d8b7", size = 5366047, upload-time = "2025-09-14T22:16:51.328Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a3/732893eab0a3a7aecff8b99052fecf9f605cf0fb5fb6d0290e36beee47a4/zstandard-0.25.0-cp311-cp311-win32.whl", hash = "sha256:f5aeea11ded7320a84dcdd62a3d95b5186834224a9e55b92ccae35d21a8b63d4", size = 436484, upload-time = "2025-09-14T22:16:55.005Z" }, + { url = "https://files.pythonhosted.org/packages/43/a3/c6155f5c1cce691cb80dfd38627046e50af3ee9ddc5d0b45b9b063bfb8c9/zstandard-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:daab68faadb847063d0c56f361a289c4f268706b598afbf9ad113cbe5c38b6b2", size = 506183, upload-time = "2025-09-14T22:16:52.753Z" }, + { url = "https://files.pythonhosted.org/packages/8c/3e/8945ab86a0820cc0e0cdbf38086a92868a9172020fdab8a03ac19662b0e5/zstandard-0.25.0-cp311-cp311-win_arm64.whl", hash = "sha256:22a06c5df3751bb7dc67406f5374734ccee8ed37fc5981bf1ad7041831fa1137", size = 462533, upload-time = "2025-09-14T22:16:53.878Z" }, + { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, + { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, + { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, + { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, + { url = "https://files.pythonhosted.org/packages/35/0b/8df9c4ad06af91d39e94fa96cc010a24ac4ef1378d3efab9223cc8593d40/zstandard-0.25.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ec996f12524f88e151c339688c3897194821d7f03081ab35d31d1e12ec975e94", size = 795735, upload-time = "2025-09-14T22:17:26.042Z" }, + { url = "https://files.pythonhosted.org/packages/3f/06/9ae96a3e5dcfd119377ba33d4c42a7d89da1efabd5cb3e366b156c45ff4d/zstandard-0.25.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a1a4ae2dec3993a32247995bdfe367fc3266da832d82f8438c8570f989753de1", size = 640440, upload-time = "2025-09-14T22:17:27.366Z" }, + { url = "https://files.pythonhosted.org/packages/d9/14/933d27204c2bd404229c69f445862454dcc101cd69ef8c6068f15aaec12c/zstandard-0.25.0-cp313-cp313-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:e96594a5537722fdfb79951672a2a63aec5ebfb823e7560586f7484819f2a08f", size = 5343070, upload-time = "2025-09-14T22:17:28.896Z" }, + { url = "https://files.pythonhosted.org/packages/6d/db/ddb11011826ed7db9d0e485d13df79b58586bfdec56e5c84a928a9a78c1c/zstandard-0.25.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bfc4e20784722098822e3eee42b8e576b379ed72cca4a7cb856ae733e62192ea", size = 5063001, upload-time = "2025-09-14T22:17:31.044Z" }, + { url = "https://files.pythonhosted.org/packages/db/00/87466ea3f99599d02a5238498b87bf84a6348290c19571051839ca943777/zstandard-0.25.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:457ed498fc58cdc12fc48f7950e02740d4f7ae9493dd4ab2168a47c93c31298e", size = 5394120, upload-time = "2025-09-14T22:17:32.711Z" }, + { url = "https://files.pythonhosted.org/packages/2b/95/fc5531d9c618a679a20ff6c29e2b3ef1d1f4ad66c5e161ae6ff847d102a9/zstandard-0.25.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:fd7a5004eb1980d3cefe26b2685bcb0b17989901a70a1040d1ac86f1d898c551", size = 5451230, upload-time = "2025-09-14T22:17:34.41Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/e3678b4e776db00f9f7b2fe58e547e8928ef32727d7a1ff01dea010f3f13/zstandard-0.25.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8e735494da3db08694d26480f1493ad2cf86e99bdd53e8e9771b2752a5c0246a", size = 5547173, upload-time = "2025-09-14T22:17:36.084Z" }, + { url = "https://files.pythonhosted.org/packages/4e/d5/ba05ed95c6b8ec30bd468dfeab20589f2cf709b5c940483e31d991f2ca58/zstandard-0.25.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3a39c94ad7866160a4a46d772e43311a743c316942037671beb264e395bdd611", size = 5046736, upload-time = "2025-09-14T22:17:37.891Z" }, + { url = "https://files.pythonhosted.org/packages/50/d5/870aa06b3a76c73eced65c044b92286a3c4e00554005ff51962deef28e28/zstandard-0.25.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:172de1f06947577d3a3005416977cce6168f2261284c02080e7ad0185faeced3", size = 5576368, upload-time = "2025-09-14T22:17:40.206Z" }, + { url = "https://files.pythonhosted.org/packages/5d/35/398dc2ffc89d304d59bc12f0fdd931b4ce455bddf7038a0a67733a25f550/zstandard-0.25.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3c83b0188c852a47cd13ef3bf9209fb0a77fa5374958b8c53aaa699398c6bd7b", size = 4954022, upload-time = "2025-09-14T22:17:41.879Z" }, + { url = "https://files.pythonhosted.org/packages/9a/5c/36ba1e5507d56d2213202ec2b05e8541734af5f2ce378c5d1ceaf4d88dc4/zstandard-0.25.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:1673b7199bbe763365b81a4f3252b8e80f44c9e323fc42940dc8843bfeaf9851", size = 5267889, upload-time = "2025-09-14T22:17:43.577Z" }, + { url = "https://files.pythonhosted.org/packages/70/e8/2ec6b6fb7358b2ec0113ae202647ca7c0e9d15b61c005ae5225ad0995df5/zstandard-0.25.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:0be7622c37c183406f3dbf0cba104118eb16a4ea7359eeb5752f0794882fc250", size = 5433952, upload-time = "2025-09-14T22:17:45.271Z" }, + { url = "https://files.pythonhosted.org/packages/7b/01/b5f4d4dbc59ef193e870495c6f1275f5b2928e01ff5a81fecb22a06e22fb/zstandard-0.25.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:5f5e4c2a23ca271c218ac025bd7d635597048b366d6f31f420aaeb715239fc98", size = 5814054, upload-time = "2025-09-14T22:17:47.08Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e5/fbd822d5c6f427cf158316d012c5a12f233473c2f9c5fe5ab1ae5d21f3d8/zstandard-0.25.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4f187a0bb61b35119d1926aee039524d1f93aaf38a9916b8c4b78ac8514a0aaf", size = 5360113, upload-time = "2025-09-14T22:17:48.893Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e0/69a553d2047f9a2c7347caa225bb3a63b6d7704ad74610cb7823baa08ed7/zstandard-0.25.0-cp313-cp313-win32.whl", hash = "sha256:7030defa83eef3e51ff26f0b7bfb229f0204b66fe18e04359ce3474ac33cbc09", size = 436936, upload-time = "2025-09-14T22:17:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/d9/82/b9c06c870f3bd8767c201f1edbdf9e8dc34be5b0fbc5682c4f80fe948475/zstandard-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:1f830a0dac88719af0ae43b8b2d6aef487d437036468ef3c2ea59c51f9d55fd5", size = 506232, upload-time = "2025-09-14T22:17:50.402Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/60c3c01243bb81d381c9916e2a6d9e149ab8627c0c7d7abb2d73384b3c0c/zstandard-0.25.0-cp313-cp313-win_arm64.whl", hash = "sha256:85304a43f4d513f5464ceb938aa02c1e78c2943b29f44a750b48b25ac999a049", size = 462671, upload-time = "2025-09-14T22:17:51.533Z" }, + { url = "https://files.pythonhosted.org/packages/3d/5c/f8923b595b55fe49e30612987ad8bf053aef555c14f05bb659dd5dbe3e8a/zstandard-0.25.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e29f0cf06974c899b2c188ef7f783607dbef36da4c242eb6c82dcd8b512855e3", size = 795887, upload-time = "2025-09-14T22:17:54.198Z" }, + { url = "https://files.pythonhosted.org/packages/8d/09/d0a2a14fc3439c5f874042dca72a79c70a532090b7ba0003be73fee37ae2/zstandard-0.25.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:05df5136bc5a011f33cd25bc9f506e7426c0c9b3f9954f056831ce68f3b6689f", size = 640658, upload-time = "2025-09-14T22:17:55.423Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/8b6b71b1ddd517f68ffb55e10834388d4f793c49c6b83effaaa05785b0b4/zstandard-0.25.0-cp314-cp314-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:f604efd28f239cc21b3adb53eb061e2a205dc164be408e553b41ba2ffe0ca15c", size = 5379849, upload-time = "2025-09-14T22:17:57.372Z" }, + { url = "https://files.pythonhosted.org/packages/a4/86/a48e56320d0a17189ab7a42645387334fba2200e904ee47fc5a26c1fd8ca/zstandard-0.25.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223415140608d0f0da010499eaa8ccdb9af210a543fac54bce15babbcfc78439", size = 5058095, upload-time = "2025-09-14T22:17:59.498Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ad/eb659984ee2c0a779f9d06dbfe45e2dc39d99ff40a319895df2d3d9a48e5/zstandard-0.25.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2e54296a283f3ab5a26fc9b8b5d4978ea0532f37b231644f367aa588930aa043", size = 5551751, upload-time = "2025-09-14T22:18:01.618Z" }, + { url = "https://files.pythonhosted.org/packages/61/b3/b637faea43677eb7bd42ab204dfb7053bd5c4582bfe6b1baefa80ac0c47b/zstandard-0.25.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ca54090275939dc8ec5dea2d2afb400e0f83444b2fc24e07df7fdef677110859", size = 6364818, upload-time = "2025-09-14T22:18:03.769Z" }, + { url = "https://files.pythonhosted.org/packages/31/dc/cc50210e11e465c975462439a492516a73300ab8caa8f5e0902544fd748b/zstandard-0.25.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e09bb6252b6476d8d56100e8147b803befa9a12cea144bbe629dd508800d1ad0", size = 5560402, upload-time = "2025-09-14T22:18:05.954Z" }, + { url = "https://files.pythonhosted.org/packages/c9/ae/56523ae9c142f0c08efd5e868a6da613ae76614eca1305259c3bf6a0ed43/zstandard-0.25.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a9ec8c642d1ec73287ae3e726792dd86c96f5681eb8df274a757bf62b750eae7", size = 4955108, upload-time = "2025-09-14T22:18:07.68Z" }, + { url = "https://files.pythonhosted.org/packages/98/cf/c899f2d6df0840d5e384cf4c4121458c72802e8bda19691f3b16619f51e9/zstandard-0.25.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a4089a10e598eae6393756b036e0f419e8c1d60f44a831520f9af41c14216cf2", size = 5269248, upload-time = "2025-09-14T22:18:09.753Z" }, + { url = "https://files.pythonhosted.org/packages/1b/c0/59e912a531d91e1c192d3085fc0f6fb2852753c301a812d856d857ea03c6/zstandard-0.25.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:f67e8f1a324a900e75b5e28ffb152bcac9fbed1cc7b43f99cd90f395c4375344", size = 5430330, upload-time = "2025-09-14T22:18:11.966Z" }, + { url = "https://files.pythonhosted.org/packages/a0/1d/7e31db1240de2df22a58e2ea9a93fc6e38cc29353e660c0272b6735d6669/zstandard-0.25.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:9654dbc012d8b06fc3d19cc825af3f7bf8ae242226df5f83936cb39f5fdc846c", size = 5811123, upload-time = "2025-09-14T22:18:13.907Z" }, + { url = "https://files.pythonhosted.org/packages/f6/49/fac46df5ad353d50535e118d6983069df68ca5908d4d65b8c466150a4ff1/zstandard-0.25.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4203ce3b31aec23012d3a4cf4a2ed64d12fea5269c49aed5e4c3611b938e4088", size = 5359591, upload-time = "2025-09-14T22:18:16.465Z" }, + { url = "https://files.pythonhosted.org/packages/c2/38/f249a2050ad1eea0bb364046153942e34abba95dd5520af199aed86fbb49/zstandard-0.25.0-cp314-cp314-win32.whl", hash = "sha256:da469dc041701583e34de852d8634703550348d5822e66a0c827d39b05365b12", size = 444513, upload-time = "2025-09-14T22:18:20.61Z" }, + { url = "https://files.pythonhosted.org/packages/3a/43/241f9615bcf8ba8903b3f0432da069e857fc4fd1783bd26183db53c4804b/zstandard-0.25.0-cp314-cp314-win_amd64.whl", hash = "sha256:c19bcdd826e95671065f8692b5a4aa95c52dc7a02a4c5a0cac46deb879a017a2", size = 516118, upload-time = "2025-09-14T22:18:17.849Z" }, + { url = "https://files.pythonhosted.org/packages/f0/ef/da163ce2450ed4febf6467d77ccb4cd52c4c30ab45624bad26ca0a27260c/zstandard-0.25.0-cp314-cp314-win_arm64.whl", hash = "sha256:d7541afd73985c630bafcd6338d2518ae96060075f9463d7dc14cfb33514383d", size = 476940, upload-time = "2025-09-14T22:18:19.088Z" }, +]