From bbe528bbd42baeb6bf0589c1302f71c3a45ddcdc Mon Sep 17 00:00:00 2001 From: ElmoPA Date: Sun, 18 Oct 2026 14:04:23 -0400 Subject: [PATCH 1/4] Add modular pixel video training and resilient log mirroring Implement image-only discrete and continuous LIBERO video recipes with separate corruption, denoising and loss stages, axial RoPE, AdaLN-Zero, packed attention and distributed validation/resume support. Keep the frozen-image state/control readout configurable and deferred. Preserve upstream review ancestry and the legacy checkpoint fixture; record frozen running-source identities. Retain transient mirror failures for retry without terminating training, while requiring a final flush. Validation: 540 changed-training tests and 2 legacy parity checks passed; 5 GPU-gated tests skipped locally. Ruff, shell syntax and diff checks pass. --- docs/PIXEL_VIDEO_PILOTS.md | 97 ++ docs/WORLD_MODEL.md | 124 +++ egomimic/eval/dino_proprio_eval.py | 13 + egomimic/eval/image_proprio_eval.py | 246 +++++ egomimic/eval/pixel_diagnostics.py | 435 ++++++++ .../data/dino_proprio/prepared.yaml | 18 + .../data/pixel_proprio/prepared.yaml | 22 + .../dino_proprio/direct_action.yaml | 39 + .../experiment/dino_proprio/visual_state.yaml | 47 + .../pixel_proprio/readout_continuous.yaml | 31 + .../pixel_proprio/readout_discrete.yaml | 18 + .../pixel_proprio/video_continuous.yaml | 10 + .../pixel_proprio/video_discrete.yaml | 16 + .../experiment/pixel_proprio/video_only.yaml | 61 ++ .../model/dino_proprio/base.yaml | 39 + .../model/pixel_proprio/base.yaml | 49 + .../model/pixel_proprio/readout.yaml | 106 ++ .../model/pixel_proprio/video.yaml | 67 ++ .../hydra_configs/pixel_execution/ddp4.yaml | 18 + .../pixel_execution/ddp4_prefetch.yaml | 13 + .../flash3_bf16_mixed.yaml | 7 + .../flash4_bf16_mixed.yaml | 10 + .../world_model_attention/native.yaml | 1 + .../world_model_execution/batched_eager.yaml | 7 + .../packed_regional_compile.yaml | 8 + .../regional_compile.yaml | 10 + egomimic/models/attention_execution.py | 268 +++++ egomimic/models/compilation.py | 48 + egomimic/models/dino_proprio.py | 28 + egomimic/models/image_proprio.py | 952 +++++++++++++++++ egomimic/models/pixel_diffusion.py | 198 ++++ egomimic/models/world_model.py | 234 +++-- egomimic/pipeline/stages_dino_proprio.py | 13 + egomimic/pipeline/stages_image_proprio.py | 86 ++ egomimic/pipeline/stages_pixel_readout.py | 350 +++++++ egomimic/pipeline/stages_pixel_video.py | 137 +++ egomimic/pl_utils/world_model_behavior.py | 104 +- egomimic/rldb/dino_proprio.py | 611 +++++++++++ egomimic/rldb/dino_proprio_datamodule.py | 212 ++++ egomimic/rldb/pixel_proprio.py | 890 ++++++++++++++++ egomimic/rldb/pixel_proprio_datamodule.py | 363 +++++++ egomimic/trainHydra.py | 45 +- scripts/train/log_mirror.sh | 78 ++ .../world_model/benchmark_pixel_execution.py | 655 ++++++++++++ .../export_libero_action_sidecars.py | 406 +++++++ .../world_model/initialize_pixel_readout.py | 496 +++++++++ scripts/world_model/prepare_pixel_proprio.py | 715 +++++++++++++ tests/test_attention_execution.py | 384 +++++++ tests/test_dino_proprio.py | 361 +++++++ tests/test_dino_proprio_data.py | 452 ++++++++ tests/test_dino_proprio_integration.py | 139 +++ tests/test_image_proprio.py | 319 ++++++ tests/test_image_proprio_adaln.py | 183 ++++ tests/test_image_proprio_attention.py | 199 ++++ tests/test_image_proprio_distributed.py | 130 +++ tests/test_image_proprio_image_only.py | 135 +++ ...image_proprio_image_only_initialization.py | 150 +++ tests/test_image_proprio_integration.py | 187 ++++ tests/test_image_proprio_mixed_precision.py | 106 ++ tests/test_image_proprio_no_phase.py | 93 ++ tests/test_image_proprio_positions.py | 197 ++++ ...test_image_proprio_validation_lifecycle.py | 130 +++ tests/test_libero_action_sidecar_export.py | 291 ++++++ tests/test_log_mirror.py | 199 ++++ tests/test_pixel_attention_profiles.py | 186 ++++ tests/test_pixel_compile_ddp.py | 96 ++ tests/test_pixel_diffusion_components.py | 273 +++++ tests/test_pixel_execution_benchmark.py | 304 ++++++ tests/test_pixel_execution_model.py | 170 +++ tests/test_pixel_loader_execution.py | 364 +++++++ .../test_pixel_noise_conditioning_profiles.py | 269 +++++ tests/test_pixel_packed_execution.py | 464 ++++++++ tests/test_pixel_prediction_diagnostics.py | 228 ++++ tests/test_pixel_proprio_data.py | 989 ++++++++++++++++++ tests/test_pixel_proprio_preparation.py | 498 +++++++++ tests/test_pixel_readout_graph.py | 459 ++++++++ tests/test_pixel_readout_initialization.py | 366 +++++++ tests/test_pixel_readout_modularity.py | 218 ++++ tests/test_pixel_rgb_frame_reads.py | 144 +++ tests/test_pixel_token_execution_profiles.py | 104 ++ tests/test_pixel_video_graph.py | 246 +++++ tests/test_train_hydra_torch_elastic.py | 97 ++ tests/test_world_model_distributed_rng.py | 260 +++++ 83 files changed, 17718 insertions(+), 73 deletions(-) create mode 100644 docs/PIXEL_VIDEO_PILOTS.md create mode 100644 egomimic/eval/dino_proprio_eval.py create mode 100644 egomimic/eval/image_proprio_eval.py create mode 100644 egomimic/eval/pixel_diagnostics.py create mode 100644 egomimic/hydra_configs/data/dino_proprio/prepared.yaml create mode 100644 egomimic/hydra_configs/data/pixel_proprio/prepared.yaml create mode 100644 egomimic/hydra_configs/experiment/dino_proprio/direct_action.yaml create mode 100644 egomimic/hydra_configs/experiment/dino_proprio/visual_state.yaml create mode 100644 egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml create mode 100644 egomimic/hydra_configs/experiment/pixel_proprio/readout_discrete.yaml create mode 100644 egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml create mode 100644 egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml create mode 100644 egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml create mode 100644 egomimic/hydra_configs/model/dino_proprio/base.yaml create mode 100644 egomimic/hydra_configs/model/pixel_proprio/base.yaml create mode 100644 egomimic/hydra_configs/model/pixel_proprio/readout.yaml create mode 100644 egomimic/hydra_configs/model/pixel_proprio/video.yaml create mode 100644 egomimic/hydra_configs/pixel_execution/ddp4.yaml create mode 100644 egomimic/hydra_configs/pixel_execution/ddp4_prefetch.yaml create mode 100644 egomimic/hydra_configs/world_model_attention/flash3_bf16_mixed.yaml create mode 100644 egomimic/hydra_configs/world_model_attention/flash4_bf16_mixed.yaml create mode 100644 egomimic/hydra_configs/world_model_execution/batched_eager.yaml create mode 100644 egomimic/hydra_configs/world_model_execution/packed_regional_compile.yaml create mode 100644 egomimic/hydra_configs/world_model_execution/regional_compile.yaml create mode 100644 egomimic/models/attention_execution.py create mode 100644 egomimic/models/compilation.py create mode 100644 egomimic/models/dino_proprio.py create mode 100644 egomimic/models/image_proprio.py create mode 100644 egomimic/models/pixel_diffusion.py create mode 100644 egomimic/pipeline/stages_dino_proprio.py create mode 100644 egomimic/pipeline/stages_image_proprio.py create mode 100644 egomimic/pipeline/stages_pixel_readout.py create mode 100644 egomimic/pipeline/stages_pixel_video.py create mode 100644 egomimic/rldb/dino_proprio.py create mode 100644 egomimic/rldb/dino_proprio_datamodule.py create mode 100644 egomimic/rldb/pixel_proprio.py create mode 100644 egomimic/rldb/pixel_proprio_datamodule.py create mode 100644 scripts/train/log_mirror.sh create mode 100644 scripts/world_model/benchmark_pixel_execution.py create mode 100644 scripts/world_model/export_libero_action_sidecars.py create mode 100644 scripts/world_model/initialize_pixel_readout.py create mode 100644 scripts/world_model/prepare_pixel_proprio.py create mode 100644 tests/test_attention_execution.py create mode 100644 tests/test_dino_proprio.py create mode 100644 tests/test_dino_proprio_data.py create mode 100644 tests/test_dino_proprio_integration.py create mode 100644 tests/test_image_proprio.py create mode 100644 tests/test_image_proprio_adaln.py create mode 100644 tests/test_image_proprio_attention.py create mode 100644 tests/test_image_proprio_distributed.py create mode 100644 tests/test_image_proprio_image_only.py create mode 100644 tests/test_image_proprio_image_only_initialization.py create mode 100644 tests/test_image_proprio_integration.py create mode 100644 tests/test_image_proprio_mixed_precision.py create mode 100644 tests/test_image_proprio_no_phase.py create mode 100644 tests/test_image_proprio_positions.py create mode 100644 tests/test_image_proprio_validation_lifecycle.py create mode 100644 tests/test_libero_action_sidecar_export.py create mode 100644 tests/test_log_mirror.py create mode 100644 tests/test_pixel_attention_profiles.py create mode 100644 tests/test_pixel_compile_ddp.py create mode 100644 tests/test_pixel_diffusion_components.py create mode 100644 tests/test_pixel_execution_benchmark.py create mode 100644 tests/test_pixel_execution_model.py create mode 100644 tests/test_pixel_loader_execution.py create mode 100644 tests/test_pixel_noise_conditioning_profiles.py create mode 100644 tests/test_pixel_packed_execution.py create mode 100644 tests/test_pixel_prediction_diagnostics.py create mode 100644 tests/test_pixel_proprio_data.py create mode 100644 tests/test_pixel_proprio_preparation.py create mode 100644 tests/test_pixel_readout_graph.py create mode 100644 tests/test_pixel_readout_initialization.py create mode 100644 tests/test_pixel_readout_modularity.py create mode 100644 tests/test_pixel_rgb_frame_reads.py create mode 100644 tests/test_pixel_token_execution_profiles.py create mode 100644 tests/test_pixel_video_graph.py create mode 100644 tests/test_train_hydra_torch_elastic.py create mode 100644 tests/test_world_model_distributed_rng.py diff --git a/docs/PIXEL_VIDEO_PILOTS.md b/docs/PIXEL_VIDEO_PILOTS.md new file mode 100644 index 000000000..7526828a0 --- /dev/null +++ b/docs/PIXEL_VIDEO_PILOTS.md @@ -0,0 +1,97 @@ +# Images-only pixel video pilots + +The initial experiment uses LIBERO gripper RGB only: two cameras at 160×160, +with 16×16 RGB patches (100 tokens per camera). There is no pretrained image +encoder, point track input, proprioception conditioning, state head, or action +head in stage1. Camera IDs, spatial positions, and the selected time encoding +remain. Each stage1 model has one image expert, width 768, depth 8, six attention +heads, and 88,178,688 parameters. The current model uses per-frame AdaLN-Zero +noise conditioning and spatial/temporal RoPE (64 time, 32 x, 32 y channels per +head); it has no additive XY or learned phase embedding. + +The two recipes keep the same frozen 450-training/50-validation episode split, +with 45/5 episodes per task. Discrete training selects a complete 64-frame +window and selects 16 timestamps at stride 4, producing 3,200 tokens. Continuous +training samples 2–min(64, episode length) timestamps across the whole episode, +producing at most 12,800 tokens. This caps the number of selected timestamps, +not their temporal span. The unchanged model capacity ceiling remains 25,600 +tokens; deferred stage2 retains its separately configured 128-timestamp cap. +Both cameras are retained at each selected +timestamp. Training anchors are uniform over eligible recorded frames. Equal +optimizer-step budgets therefore do not imply equal token or attention work. + +The discrete and continuous full-training recipes use AdamW at LR 5e-5, +weight decay 0.01, and 300,000 optimizer steps each. The first major quality +review is at 100,000 steps; it does not automatically stop either run. +Only these two recipe defaults change the total budget. The shared pilot base, +explicit smoke step overrides, and deferred stage2 budgets remain unchanged. +The allocation profile `+pixel_execution=ddp4` runs each pilot on four GPUs on +one node. Its reusable default is batch 1 per GPU; the selected full runs use +the explicit override `data.batch_size=8`, retaining global batch 32 with no +gradient accumulation or learning-rate scaling. Each arm uses its own four-GPU +process group; its scheduler allocation and node are recorded independently. +DDP uses one image expert with all trainable parameters used; automatic +sampler replacement is disabled because the pixel adapter owns sharding. +The explicit TorchElastic environment reads ranks from each run's `torchrun` +process group, while Slurm IDs remain available for artifact provenance. +Validation covers two fixed anchors from each of the 50 held-out +episodes every 250 steps, with identical 16-timestamp evaluation grids in both arms. +The 100 global validation batches are divided into four contiguous shards of +25 without padding or duplicated examples. Sampling noise uses global batch +IDs so rank assignment does not change it. Metrics sum squared errors and +valid-target counts across ranks before dividing; RGB diagnostics are written +only by rank zero, whose shard contains both forecast and completion examples. +All scheduled checkpoints and `last` are retained: 1,200 scheduled checkpoints +per run at steps 250 through 300,000, including the 100,000-step review point. +FP32 weights and two AdamW moment tensors alone use approximately 1.06 GB per +checkpoint for this model, or 1.27 TB per run and 2.54 TB for both runs at this +retention. These are decimal raw-tensor estimates; additional checkpoint +metadata, the separate `last` files, and diagnostic artifacts require more +space. Validation reports held-out +image MSE and RGB diagnostics; it is not a policy-success evaluation. + +The default execution profile is `world_model_attention=flash4_bf16_mixed`, +with Lightning `trainer.precision=bf16-mixed`. Standard autocast uses BF16 for +eligible operations, including the linear layers and FA4 attention. Stored +parameters and AdamW state remain FP32, image MSE is evaluated in FP32, and +time coordinates retain FP32 precision. Some residuals and other activations +remain FP32; this is not an all-BF16 activation graph. TF32 is disabled, +queries are unchunked, and whole-block activation checkpointing is off. +The earlier `flash4_bf16` profile retains attention-only BF16 with +`trainer.precision=32-true`; `native` and `unchunked_fp32` also retain `32-true`. +FA4 requires the separately verified runtime documented in `WORLD_MODEL.md`. +CPU integration tests select `native` and do not certify CUDA capacity or +performance. + +`data.image_only=true` emits only the image stream and does not load measured +state/action arrays or pose-normalization statistics. The model uses +`phase_a_condition_on_proprio=false`, `pose_dims=null`, and `action_dims=null`. +The reusable prepared manifest may still retain state/action provenance for +later work; those fields are not stage1 inputs. + +Stage2 is deferred until review of the generated videos. Its separate model +may introduce proprioception plus measured-state and control-command MSE +heads, initialized from the audited stage1 image weights. Selecting these +stage1 recipes does not activate or schedule stage2 training. + +## Source and publication provenance + +The reviewed training tree is derived from local snapshot +`252275a903b9f444a297dda303fbb6a1ca9ddf93`. The publication branch is a genuine +Git child of `cb95574e7f3dcdec8c0b0374aecbb1fefaf5f616` +([PR #261](https://github.com/GaTech-RL2/EgoVerse-graph/pull/261)), retaining its +parent relationship to +[PR #260](https://github.com/GaTech-RL2/EgoVerse-graph/pull/260). Consolidation +preserves the upstream `tests/fixtures/legacy_hpt_trunk/trunk.pth` fixture, +which was absent from the isolated local snapshot. The original local history +and the frozen experiment source remain available under their recorded IDs. + +The running experiments and their checkpoints retain exact source +`3785f94392b9d900e21d9948fbc5bf94a6d975b8`, with explicit run overrides for +300,000 steps, batch 8 per rank on four ranks, and the continuous arm's +64-timestamp sampling cap. Publishing this branch does not replace that source +or relabel existing checkpoint provenance. As recorded on 2026-10-18, the +continuous arm uses allocation `411469` on `gpu19`, while the discrete resume +uses allocation `411503` on `gpu21`; they are separate four-GPU allocations. +Subsequent resumes must retain their own exact source, checkpoint, and +allocation receipts rather than infer identity from this publication branch. diff --git a/docs/WORLD_MODEL.md b/docs/WORLD_MODEL.md index 41781012c..c1ef57942 100644 --- a/docs/WORLD_MODEL.md +++ b/docs/WORLD_MODEL.md @@ -479,3 +479,127 @@ do not prove unchanged convergence or final task quality. A requested future full run must use its own resolved-source/config/runtime optimizer, validation and checkpoint compatibility preflight. No existing canceled run is resumed by selecting or adding these profiles. + +## Pixel-only distributed loader execution + +The optional `+pixel_execution=ddp4_prefetch` profile applies to the new +`PixelProprioDataModule` and pixel video recipes. It inherits four-rank DDP, +enables two loader workers per rank, pinned host memory, persistent workers, +and two prefetched batches per worker. It explicitly fixes +`data.validation_batch_size=1`. Training batch size remains one until an +allocated GPU benchmark selects a larger `data.batch_size`; the global training +batch is that value times four. Attention backend and compilation choices remain +separate execution settings. This profile is not a launch or performance claim. + +For the frozen 450/50 gripper split, validation still covers two deterministic +anchors in each of 50 episodes: 100 global batches and 25 contiguous batches per +rank, with no dropped or duplicated examples. Global validation IDs preserve +the same frame choices, observation masks, and sampling-noise seeds as the +unsharded evaluation. Training batch changes do not alter this validation grid. +The loader rejects a validation batch count that cannot divide evenly across +the configured ranks. Unattached inspection and strict reload see all 100 +validation batches. + +Checkpointed training draws use an immutable `(anchor, consumed cursor, global +draw slot)` key. Workers derive frame selection and conditioning masks from +that key, so prefetched but unconsumed batches are discarded safely on restart. +Only acknowledged batches advance the saved cursor. Worker count, pinning, +persistence, and prefetch depth are operational settings recorded in the +resolved config; they do not change source draw identity. Explicit train/val +batch sizes and the optional work-assignment strategy are checkpoint-bound. +Changing those identities is not an exact resume. Historical configs omit the +new optional identity fields and retain their previous data-context bindings. + +`data.balance_training_work=true` enables deterministic work redistribution +inside each global training batch. It previews the exact sampled frame count +from the same RNG helper used by the loader, without reading image data, and +estimates attention work as squared image-token count. Longest-first assignment +keeps exactly the configured number of samples on each rank and is accepted +only when it lowers the largest rank's estimated total. It never changes the +global batch's selected clips, frame indices, conditioning masks, or draw keys, +and never moves a clip across optimizer updates. Batch-one and equal-length +discrete batches retain their original assignment. The estimate is not measured +GPU time: FFNs, padding, transfers, and synchronization can affect throughput. +Reassigned clips receive different rank-seeded diffusion-noise realizations; +their noise distribution is unchanged, but per-clip noise equality with an +unbalanced run is not claimed. + +CPU tests exercise two workers with training batch two, discard prefetched +queues, and compare resumed versus uninterrupted frames, values, timestamps, +and masks exactly, with both balancing settings and changed worker counts. +They also verify full 100-anchor coverage across four rank shards and reject +incompatible batching/work-assignment checkpoint identities. Lightning's default +batch transfer already requests nonblocking CUDA copies for these nested tensor +containers; pinned allocation and useful transfer overlap require CUDA proof. +GPU throughput, memory, the selected larger batch, and an actual optimized +train/validation/checkpoint smoke remain pending measurements. Existing +zero-worker and four-rank baseline receipts do not certify this new profile. + +## Modular pixel state/command readout + +The deferred `pixel_proprio/readout_continuous` and `readout_discrete` recipes +keep one registered model at pipeline stage index 1. `PixelReadoutModelStage` +owns its parameters, applies the explicit `frozen_image` or `joint` policy, and +publishes the transient `pixel_video/denoiser` reference. The stable +`pixel_denoiser` and `image_proprio` aliases both select this owner; strict +initialization and tensor checkpoints retain `stages.1.model.*` names. The +legacy `PixelReadoutDenoiserStage` still implements its previous combined +ownership/prediction behavior for older serialized configurations. + +The high-noise image objective is an independent branch: `PixelVideoNoisingStage` +followed by `PixelReadoutVideoDenoiserStage` and `PixelReadoutImageLossStage`. +Removing that branch does not remove the shared model or readout gradients. +When editing the stage list, preserve the owner's index 1 if reusing its exact +checkpoint namespace; for example, the retained low-noise noising node can +occupy index 0. The default recipe still executes the high-noise branch, keeping +the existing losses and random-draw order. Under `frozen_image` its image MSE +contributes to reporting but has no trainable dependency. + +Both training and inference now supply the tensor port +`pixel_readout/image_prediction` to `PixelStateActionReadoutStage(image_key=...)`. +The default training producer remains `PixelReadoutImageDenoiserStage`, fed by +the separate corruption node with maximum strength 0.2. The inference producer +is `PixelReadoutSampledImageStage`, using the configured 20-step sampler and +`modes: [inference]`. The readout itself does not choose how images are generated. + +For a future experiment, a sampled training producer can replace the low-noise +noising/prediction branch using the same port: + +```yaml +_target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutSampledImageStage +modes: [train] +prediction_key: pixel_readout/image_prediction +gradient_policy: stop +sampler: + _target_: egomimic.models.pixel_diffusion.LinearDDIMImageSampler + steps: 20 +``` + +This is an available extension point, not a selected training algorithm. The +built-in sampler and this provider stop gradients through generated images. +Frozen-image readout training is supported. A future joint-training recipe +would receive image gradients only through separate image objectives or Phase B, +not through this sampled producer. The low-noise producer retains its original +differentiable behavior. Image, measured-state and controller-command losses +remain separate configured nodes. + +The current image architecture uses AdaLN-Zero and spatial-temporal RoPE with +time/x/y dimensions `[64,32,32]` per 128-channel head and a `[10,10]` patch grid. +There is no learned XY projection or phase embedding; camera embeddings remain. +Discrete time uses sampled-position indices and continuous time uses recorded +seconds, an intentional comparison difference. The image expert has 88,178,688 +parameters. Stage-two expansion requires identical noise/position settings and +geometry, and freezes all copied image parameters, including AdaLN modules. + +Frozen image parameters preserve image-only Phase A and its sampler. Joint +attention during Phase B can still change image activations through state/query +tokens; the exposed image prediction is the Phase-A result. The default readout +training images come from a one-pass estimate of low-noise ground-truth future +images, whereas inference uses full sampled images. That distribution gap and +the eventual stage-two training strategy remain unresolved. No stage-two GPU +launch or certification is implied by these extension points. + +CPU tests verify default output, gradient and RNG parity; strict legacy/new +tensor checkpoint compatibility; removal of the high-noise branch; sampled +provider substitution with unchanged model weights; explicit stopped gradients; +and frozen readout initialization under both clock and position variants. diff --git a/egomimic/eval/dino_proprio_eval.py b/egomimic/eval/dino_proprio_eval.py new file mode 100644 index 000000000..53b5f31df --- /dev/null +++ b/egomimic/eval/dino_proprio_eval.py @@ -0,0 +1,13 @@ +"""Compatibility evaluator for original DINO recipe prediction keys.""" + +from egomimic.eval.eval import EvaluationDataRequirements +from egomimic.eval.image_proprio_eval import ImageProprioEval + + +class DinoProprioEval(ImageProprioEval): + prediction_key = "dino_proprio_prediction" + + def data_requirements(self): + # Preserve the archived DINO datamodule contract. Pixel diagnostics use + # the generic evaluator and its richer metadata declaration. + return EvaluationDataRequirements(required_keys=("world", "embodiment")) diff --git a/egomimic/eval/image_proprio_eval.py b/egomimic/eval/image_proprio_eval.py new file mode 100644 index 000000000..a14904704 --- /dev/null +++ b/egomimic/eval/image_proprio_eval.py @@ -0,0 +1,246 @@ +"""Held-out prediction metrics, kept separate for states and action commands.""" + +import torch + +from egomimic.eval.eval import Eval, EvaluationDataRequirements + + +class ImageProprioEval(Eval): + prediction_key = "image_proprio_prediction" + + def __init__( + self, objective="visual_proprio", sample_seed=42, pixel_diagnostics=None + ): + if objective not in ( + "visual_only", + "visual_proprio", + "visual_state_action", + "direct_action", + ): + raise ValueError("Unknown image/proprio evaluation objective") + self.objective, self.sample_seed = objective, int(sample_seed) + self.totals = {} + self._completed_validation_step = None + self._validation_step = None + self._validation_batches = 0 + self.global_validation_batches = 0 + self._global_totals = {} + self._metrics_logged = False + self._validation_failed = False + self._validation_sanity = False + self.pixel_diagnostics = None + if pixel_diagnostics is not None: + from egomimic.eval.pixel_diagnostics import PixelPredictionDiagnostics + + self.pixel_diagnostics = PixelPredictionDiagnostics( + **dict(pixel_diagnostics) + ) + + def bind_data_context(self, *, normalizer): + if normalizer is not None: + raise ValueError("Image/proprio evaluator expects data-owned normalization") + + def data_requirements(self): + keys = ("world", "embodiment", "task_kind") + if self.pixel_diagnostics is not None: + keys += ("pixel_metadata",) + return EvaluationDataRequirements(required_keys=keys) + + def trainer_overrides(self): + return {} + + def on_validation_start(self): + self._completed_validation_step = None + self._validation_step = int(self.model.global_step) + self._validation_batches = 0 + self.global_validation_batches = 0 + self._global_totals = {} + self._metrics_logged = False + self._validation_failed = False + self._validation_sanity = bool( + getattr(getattr(self.model, "trainer", None), "sanity_checking", False) + ) + self.totals = {} + if self.pixel_diagnostics is not None: + self.pixel_diagnostics.start() + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + self._completed_validation_step = None + self._metrics_logged = False + try: + output = self._evaluate_batch(batch, batch_idx, dataloader_idx) + except Exception: + self._validation_failed = True + raise + self._validation_batches += 1 + return output + + @staticmethod + def _sampling_index(batch, batch_idx): + """Use the adapter's global batch identity across validation shards.""" + indexed = [item.get("validation_batch_index") for item in batch.values()] + if all(value is None for value in indexed): + return int(batch_idx) + if any(value is None for value in indexed): + raise ValueError("Global validation batch identities must cover every source") + indices = [] + for item, values in zip(batch.values(), indexed): + if not isinstance(values, (tuple, list)) or len(values) != len(item["task_kind"]): + raise ValueError("One global validation batch identity is required per sample") + if any(type(value) is not int or value < 0 for value in values): + raise ValueError("Global validation batch identities must be nonnegative integers") + indices.extend(values) + if len(set(indices)) != 1: + raise ValueError("A validation batch must share one global batch identity") + return indices[0] + + def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0): + devices = ( + [self.model.device.index or 0] if self.model.device.type == "cuda" else [] + ) + with torch.random.fork_rng(devices=devices): + torch.manual_seed( + self.sample_seed + self._sampling_index(batch, batch_idx) + + 100000 * dataloader_idx + ) + output = self.model.forward_eval(batch) + names = { + "visual_only": ("image",), + "visual_proprio": ("image", "proprio"), + "visual_state_action": ("image", "proprio", "action"), + "direct_action": ("action",), + }[self.objective] + for source, item in batch.items(): + kinds = item["task_kind"] + for name in names: + stream = item["world"][name] + mask = ( + stream["valid"] + & stream.get("target_valid", stream["valid"]) + & ~stream["observed"] + ) + target = torch.where(mask[..., None], stream["value"], 0) + prediction = output[source][self.prediction_key][name] + if prediction.shape != target.shape: + raise ValueError(f"Prediction/target shape mismatch for {name}") + error = (prediction.float() - target.float()).square().mean(-1) + if not torch.isfinite(error[mask]).all(): + raise ValueError(f"Non-finite held-out {name} prediction/target") + metric = ( + "action_command_mse" + if name == "action" + else f"{name}_normalized_mse" + ) + if len(kinds) != error.shape[0]: + raise ValueError("One task kind is required per validation sample") + for kind in sorted(set(kinds)): + if kind not in ("forecast", "completion", "direct_action"): + raise ValueError(f"Unknown validation task kind: {kind}") + rows = torch.tensor( + [value == kind for value in kinds], device=mask.device + ) + selected = mask & rows[:, None] + numerator = error.masked_fill(~selected, 0).sum().double() + denominator = selected.sum().double() + key = f"val/{source}/{kind}/{metric}" + old = self.totals.get( + key, (numerator.new_zeros(()), denominator.new_zeros(())) + ) + self.totals[key] = old[0] + numerator, old[1] + denominator + if self.pixel_diagnostics is not None: + self.pixel_diagnostics.record( + model=self.model, + source=source, + item=item, + prediction=output[source][self.prediction_key], + objective=self.objective, + batch_idx=batch_idx, + dataloader_idx=dataloader_idx, + ) + return output + + def _reduce_validation_totals(self): + """Sum target-weighted errors, including keys absent on some ranks.""" + self.global_validation_batches = self._validation_batches + if not (torch.distributed.is_available() and torch.distributed.is_initialized()): + return self.totals + keys_by_rank = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(keys_by_rank, sorted(self.totals)) + keys = sorted({key for rank_keys in keys_by_rank for key in rank_keys}) + totals = torch.zeros(len(keys), 2, dtype=torch.float64, device=self.model.device) + for index, key in enumerate(keys): + if key in self.totals: + numerator, denominator = self.totals[key] + totals[index, 0] = numerator + totals[index, 1] = denominator + valid = (torch.isfinite(totals).all() & (totals >= 0).all()).to(torch.int64) + torch.distributed.all_reduce(valid, op=torch.distributed.ReduceOp.MIN) + if not valid.item(): + raise ValueError("Invalid accumulated validation metric on a distributed rank") + if keys: + torch.distributed.all_reduce(totals, op=torch.distributed.ReduceOp.SUM) + batches = torch.tensor(self._validation_batches, dtype=torch.int64, device=self.model.device) + torch.distributed.all_reduce(batches, op=torch.distributed.ReduceOp.SUM) + self.global_validation_batches = int(batches.item()) + return {key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys)} + + def log_epoch_metrics(self): + self._metrics_logged = False + try: + self._global_totals = self._reduce_validation_totals() + for key, (numerator, denominator) in self._global_totals.items(): + if ( + not torch.isfinite(numerator) + or not torch.isfinite(denominator) + or numerator < 0 + or denominator < 0 + ): + raise ValueError(f"Invalid accumulated validation metric: {key}") + if denominator > 0: + metric = (numerator / denominator).float() + if not torch.isfinite(metric): + raise ValueError(f"Non-finite validation metric: {key}") + self.model.log(key, metric, sync_dist=False) + except Exception: + self._validation_failed = True + raise + self._metrics_logged = True + + def has_completed_validation(self, global_step): + """Prove a successful validation/publication at this exact optimizer step.""" + return bool( + self._completed_validation_step is not None + and self._completed_validation_step == int(global_step) + ) + + def on_validation_end(self): + self._completed_validation_step = None + publication_error = None + if self.pixel_diagnostics is not None: + try: + self.pixel_diagnostics.log_to_wandb(self.model) + except Exception as error: + self._validation_failed = True + publication_error = error + # Publication must finish before advertising completion. A zero-target + # or sanity pass cannot suppress the maintained runner's final check. + complete = ( + self.global_validation_batches > 0 + and self._metrics_logged + and not self._validation_failed + and not self._validation_sanity + and not getattr( + getattr(self.model, "trainer", None), "sanity_checking", False + ) + and self._validation_step == int(self.model.global_step) + and any(denominator > 0 for _, denominator in self._global_totals.values()) + ) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + success = torch.tensor(int(complete), dtype=torch.int64, device=self.model.device) + torch.distributed.all_reduce(success, op=torch.distributed.ReduceOp.MIN) + complete = bool(success.item()) + if publication_error is not None: + raise publication_error + if complete: + self._completed_validation_step = self._validation_step + return None diff --git a/egomimic/eval/pixel_diagnostics.py b/egomimic/eval/pixel_diagnostics.py new file mode 100644 index 000000000..ef1a8f329 --- /dev/null +++ b/egomimic/eval/pixel_diagnostics.py @@ -0,0 +1,435 @@ +"""Bounded, timestamp-faithful RGB prediction diagnostics for patch models.""" + +from __future__ import annotations + +import json +import math +import re +from fractions import Fraction +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw + +from egomimic.eval.artifact_paths import ( + artifact_destination, + artifact_execution_identity, +) + + +def _component(value): + value = str(value) + if not re.fullmatch(r"[A-Za-z0-9_.-]+", value) or value in (".", ".."): + raise ValueError("Artifact path component must be a simple identifier") + return value + + +def unpatchify_rgb(patches, image_size, patch_size): + """Invert row-major/HWC patchification without clipping predictions.""" + height, width = image_size + gy, gx = height // patch_size, width // patch_size + if patches.shape != (gy * gx, 3 * patch_size**2): + raise ValueError("Pixel patch shape differs from declared image geometry") + return ( + patches.reshape(gy, gx, patch_size, patch_size, 3) + .transpose(0, 2, 1, 3, 4) + .reshape(height, width, 3) + ) + + +def pixel_frame_records(stream, prediction, metadata, row): + """Validate exact physical metadata and restore each camera's RGB frames.""" + size, patch = metadata["image_size"], metadata["patch_size"] + if ( + len(size) != 2 + or any(type(x) is not int or x < 1 for x in size) + or type(patch) is not int + or patch < 1 + or any(x % patch for x in size) + or metadata.get("normalization") != "uint8_to_minus_one_one" + or metadata.get("patch_order") != "row_major_hwc" + ): + raise ValueError( + "Pixel diagnostics require explicit RGB patch geometry and normalization" + ) + times = np.asarray(metadata["timestamps_seconds"], np.float64) + frames = np.asarray(metadata["frame_indices"]) + cameras = metadata["camera_names"] + origin = float(metadata["time_origin_seconds"]) + if ( + times.ndim != 1 + or len(times) == 0 + or not np.isfinite(times).all() + or not (np.diff(times) > 0).all() + or not math.isfinite(origin) + or frames.shape != times.shape + or frames.dtype.kind not in "iu" + or (frames < 0).any() + or not (np.diff(frames) > 0).all() + or not cameras + or len(set(cameras)) != len(cameras) + or any(not isinstance(x, str) or not x for x in cameras) + ): + raise ValueError( + "Pixel metadata requires ordered physical times, frame indices and camera names" + ) + arrays = {key: value[row].detach().cpu().numpy() for key, value in stream.items()} + predicted = ( + None if prediction is None else prediction[row].detach().float().cpu().numpy() + ) + if predicted is not None and predicted.shape != arrays["value"].shape: + raise ValueError("Pixel prediction and target shapes differ") + valid = arrays["valid"] + index, camera = arrays["index"], arrays["camera"] + if ( + (index[valid] < 0).any() + or (index[valid] >= len(times)).any() + or not np.equal(index[valid], np.round(index[valid])).all() + or (camera[valid] < 0).any() + or (camera[valid] >= len(cameras)).any() + ): + raise ValueError("Pixel tokens do not match metadata frames/cameras") + result = {} + gy, gx = size[0] // patch, size[1] // patch + for camera_id, camera_name in enumerate(cameras): + sequence = [] + for ordinal, (timestamp, frame_index) in enumerate(zip(times, frames)): + selected = np.flatnonzero( + valid & (index == ordinal) & (camera == camera_id) + ) + if len(selected) != gy * gx: + raise ValueError( + "Pixel diagnostic requires one complete grid per frame/camera" + ) + xy = arrays["xy"][selected] + grid = np.rint((xy + 1) * np.array([gx, gy]) / 2 - 0.5).astype(int) + expected_xy = (grid + 0.5) / np.array([gx, gy]) * 2 - 1 + order = grid[:, 1] * gx + grid[:, 0] + if ( + not np.isfinite(xy).all() + or not np.allclose(xy, expected_xy, atol=1e-5) + or (grid < 0).any() + or (grid >= np.array([gx, gy])).any() + or not np.array_equal(np.sort(order), np.arange(gy * gx)) + ): + raise ValueError( + "Pixel spatial coordinates do not form the declared grid" + ) + selected = selected[np.argsort(order)] + if not np.allclose( + arrays["time"][selected], timestamp - origin, atol=1e-4, rtol=1e-5 + ): + raise ValueError( + "Pixel token time does not match recorded physical timestamp" + ) + if not arrays.get("target_valid", valid)[selected].all(): + raise ValueError("RGB diagnostics require recorded target pixels") + target = arrays["value"][selected] + estimate = None if predicted is None else predicted[selected] + if not np.isfinite(target).all() or ( + estimate is not None and not np.isfinite(estimate).all() + ): + raise ValueError("Non-finite RGB diagnostic target/prediction") + if (np.abs(target) > 1.00001).any(): + raise ValueError("Target pixels are not normalized to [-1,1]") + observed = arrays["observed"][selected] + status = ( + "observed" + if observed.all() + else "query" + if not observed.any() + else "mixed" + ) + details = { + "frame_index": int(frame_index), + "timestamp_seconds": float(timestamp), + "relative_time_seconds": float(timestamp - origin), + "observation": status, + "observed_patches": observed.tolist(), + "query_pixel_mse_rgb01": None, + "prediction_outside_rgb_range_fraction": None, + } + if estimate is not None: + if (~observed).any(): + details["query_pixel_mse_rgb01"] = float( + np.square((estimate[~observed] - target[~observed]) / 2).mean() + ) + details["prediction_outside_rgb_range_fraction"] = float( + (np.abs(estimate) > 1).mean() + ) + sequence.append( + { + "metadata": details, + "target": (unpatchify_rgb(target, size, patch) + 1) / 2, + "prediction": None + if estimate is None + else (unpatchify_rgb(estimate, size, patch) + 1) / 2, + } + ) + result[camera_id] = {"name": camera_name, "frames": sequence} + return result + + +def _panel(frame, camera_name, *, direct=False): + target, prediction, metadata = ( + frame["target"], + frame["prediction"], + frame["metadata"], + ) + height, width = target.shape[:2] + scale = 2 + columns = 1 if direct else 3 + panel = Image.new( + "RGB", (max(480, columns * width * scale), height * scale + 58), (18, 21, 28) + ) + draw = ImageDraw.Draw(panel) + error = metadata["query_pixel_mse_rgb01"] + label = ( + "context only; no RGB prediction" + if direct + else f"{metadata['observation']} | query MSE={error:.6f}" + if error is not None + else "observed; excluded from query metrics" + ) + draw.text( + (5, 3), + f"{camera_name} | frame {metadata['frame_index']} | t={metadata['timestamp_seconds']:.6f}s", + fill="white", + ) + draw.text( + (5, 19), + label, + fill=(120, 220, 170) + if metadata["observation"] == "observed" + else (255, 205, 110), + ) + arrays, labels = [target], ["Recorded RGB"] + if not direct: + difference = np.abs(prediction - target).mean(-1) + heat = np.stack( + (difference, np.zeros_like(difference), np.zeros_like(difference)), -1 + ) + arrays += [prediction, heat] + labels += ["Predicted RGB (display clipped)", "Absolute error (RGB [0,1])"] + for i, (array, title) in enumerate(zip(arrays, labels)): + image = Image.fromarray(np.rint(np.clip(array, 0, 1) * 255).astype(np.uint8)) + image = image.resize((width * scale, height * scale), Image.Resampling.NEAREST) + panel.paste(image, (i * width * scale, 58)) + draw.text((i * width * scale + 5, 39), title, fill="white") + return panel + + +def write_timestamp_video(path, panels, timestamps, nominal_fps=10): + """Store every sampled frame with physical-time PTS, without uniform retiming.""" + import av + + clock = np.asarray(timestamps, np.float64) + ticks = np.rint((clock - clock[0]) * 1_000_000).astype(np.int64) + if len(ticks) > 1 and not (np.diff(ticks) > 0).all(): + raise ValueError("Video timestamps must be distinct at microsecond precision") + time_base = Fraction(1, 1_000_000) + with av.open(str(path), mode="w") as container: + stream = container.add_stream( + "libx264", rate=Fraction(nominal_fps).limit_denominator() + ) + stream.pix_fmt = "yuv420p" + stream.time_base = stream.codec_context.time_base = time_base + stream.options = {"crf": "18", "preset": "veryfast", "bf": "0"} + for panel, tick in zip(panels, ticks): + array = np.asarray(panel, np.uint8) + stream.height, stream.width = array.shape[:2] + frame = av.VideoFrame.from_ndarray(array, format="rgb24") + frame.pts, frame.time_base = int(tick), time_base + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + +class PixelPredictionDiagnostics: + def __init__( + self, + artifact_dir, + max_samples_per_source_kind=1, + playback_fps=10, + write_video=True, + log_wandb=True, + contact_sheet_frames=8, + ): + if not artifact_dir: + raise ValueError("Pixel diagnostics require an explicit artifact_dir") + if ( + type(max_samples_per_source_kind) is not int + or max_samples_per_source_kind < 1 + or type(contact_sheet_frames) is not int + or contact_sheet_frames < 1 + ): + raise ValueError("Pixel diagnostic limits must be positive integers") + if not math.isfinite(playback_fps) or playback_fps <= 0: + raise ValueError("playback_fps must be positive") + if type(write_video) is not bool or type(log_wandb) is not bool: + raise ValueError("Pixel output switches must be boolean") + self.root = Path(artifact_dir) + self.limit, self.contact_sheet_frames = ( + max_samples_per_source_kind, + contact_sheet_frames, + ) + self.playback_fps, self.write_video, self.log_wandb = ( + playback_fps, + write_video, + log_wandb, + ) + self.start() + + def start(self): + self.counts, self.written = {}, [] + + def record( + self, *, model, source, item, prediction, objective, batch_idx, dataloader_idx + ): + if model.global_rank != 0: + return + if len(item["pixel_metadata"]) != len(item["task_kind"]): + raise ValueError("Every pixel sample requires its own metadata") + for row, (metadata, kind) in enumerate( + zip(item["pixel_metadata"], item["task_kind"]) + ): + key = (source, kind, dataloader_idx) + if self.counts.get(key, 0) >= self.limit: + continue + direct = objective == "direct_action" + frames = pixel_frame_records( + item["world"]["image"], + None if direct else prediction["image"], + metadata, + row, + ) + root = ( + self.root + / f"loader-{dataloader_idx}" + / _component(source) + / _component(kind) + ) + execution = artifact_execution_identity() + path = ( + artifact_destination( + root, + execution, + epoch=model.current_epoch, + global_step=model.global_step, + rank=model.global_rank, + batch_idx=batch_idx, + ).with_suffix("") + / f"sample-{row}" + ) + path.mkdir(parents=True, exist_ok=False) + report = { + "schema": "pixel-prediction-diagnostics-v1", + "objective": objective, + "source": source, + "task_kind": kind, + "pixel_metadata": metadata, + "execution": execution, + "global_step": int(model.global_step), + "video_timing": "variable PTS from recorded seconds; no uniform-frame retiming", + "last_frame_hold_seconds": 1 / self.playback_fps, + "pixel_error_space": "RGB [0,1], unclipped predictions; observed patches excluded", + "display": "2x nearest-neighbor; predictions clipped only for display; error is absolute RGB error, red channel", + "cameras": {}, + } + for name in ("proprio", "action"): + if name not in prediction: + continue + stream = item["world"][name] + valid = ( + stream["valid"][row] + & stream.get("target_valid", stream["valid"])[row] + & ~stream["observed"][row] + ) + report[name] = { + "prediction": prediction[name][row][valid].detach().cpu().tolist(), + "target": stream["value"][row][valid].detach().cpu().tolist(), + "relative_time_seconds": stream["time"][row][valid] + .detach() + .cpu() + .tolist(), + "squared_error": ( + prediction[name][row][valid] - stream["value"][row][valid] + ) + .square() + .detach() + .cpu() + .tolist(), + "value_space": "normalized_measured_state" + if name == "proprio" + else "native_action_commands", + "note": "RGB is observed context only; this model does not predict images" + if direct + else "Readout after generated RGB; observed state entries excluded", + } + for camera_id, camera in frames.items(): + sequence = camera["frames"] + selected = np.unique( + np.linspace( + 0, + len(sequence) - 1, + min(self.contact_sheet_frames, len(sequence)), + dtype=int, + ) + ) + thumbnails = [ + _panel(sequence[i], camera["name"], direct=direct) for i in selected + ] + sheet = Image.new( + "RGB", (thumbnails[0].width, sum(im.height for im in thumbnails)) + ) + for offset, thumbnail in enumerate(thumbnails): + sheet.paste(thumbnail, (0, offset * thumbnail.height)) + sheet_path = path / f"camera-{camera_id}.png" + sheet.save(sheet_path) + report["cameras"][str(camera_id)] = { + "name": camera["name"], + "frames": [f["metadata"] for f in sequence], + "contact_sheet_frame_indices": [ + int(metadata["frame_indices"][i]) for i in selected + ], + "contact_sheet": sheet_path.name, + "video": None, + } + self.written.append((source, kind, "image", sheet_path)) + if self.write_video and not direct: + video_path = path / f"camera-{camera_id}.mp4" + write_timestamp_video( + video_path, + (_panel(f, camera["name"]) for f in sequence), + metadata["timestamps_seconds"], + self.playback_fps, + ) + report["cameras"][str(camera_id)]["video"] = video_path.name + self.written.append((source, kind, "video", video_path)) + (path / "diagnostic.json").write_text( + json.dumps(report, indent=2, allow_nan=False) + "\n" + ) + self.counts[key] = self.counts.get(key, 0) + 1 + + def log_to_wandb(self, model): + if not self.log_wandb or model.global_rank != 0 or not self.written: + return + import wandb + from lightning.pytorch.loggers import WandbLogger + + for logger in model.trainer.loggers: + if not isinstance(logger, WandbLogger): + continue + payload = {} + for source, kind, media, path in self.written: + key = f"val_pixel/{source}/{kind}/{media}" + value = ( + wandb.Image(str(path)) + if media == "image" + else wandb.Video(str(path), format="mp4") + ) + payload.setdefault(key, []).append(value) + # Lightning records trainer/global_step as the metric axis. W&B's + # internal history counter also includes validation/media rows. + logger.log_metrics(payload, step=int(model.global_step)) diff --git a/egomimic/hydra_configs/data/dino_proprio/prepared.yaml b/egomimic/hydra_configs/data/dino_proprio/prepared.yaml new file mode 100644 index 000000000..5c9a41ca5 --- /dev/null +++ b/egomimic/hydra_configs/data/dino_proprio/prepared.yaml @@ -0,0 +1,18 @@ +_target_: egomimic.rldb.dino_proprio_datamodule.DinoProprioDataModule +manifest_path: ??? +task: visual_diffusion +action_manifest_path: null +batch_size: 1 +num_workers: 0 +mode: continuous +seed: 42 +min_frames: 4 +max_frames: 4 +discrete_frames: 4 +stride: 1 +forecasting_probability: 0.5 +max_past_seconds: null +max_future_seconds: null +checkpointable_sampling: true +validation_batches_per_embodiment: null +validation_anchors_per_episode: 2 diff --git a/egomimic/hydra_configs/data/pixel_proprio/prepared.yaml b/egomimic/hydra_configs/data/pixel_proprio/prepared.yaml new file mode 100644 index 000000000..87214178b --- /dev/null +++ b/egomimic/hydra_configs/data/pixel_proprio/prepared.yaml @@ -0,0 +1,22 @@ +_target_: egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule +manifest_path: ??? +task: visual_diffusion +batch_size: 1 +num_workers: 0 +mode: continuous +seed: 42 +min_frames: 2 +max_frames: 128 +sampling_strategy: full_episode_uniform +discrete_frames: 4 +stride: 1 +forecasting_probability: 0.5 +max_past_seconds: null +max_future_seconds: null +checkpointable_sampling: true +validation_batches_per_embodiment: null +validation_anchors_per_episode: 2 +# Compare the same held-out RGB targets/masks; training sampling remains distinct. +validation_sampling_strategy: discrete_window +validation_discrete_frames: 16 +validation_stride: 4 diff --git a/egomimic/hydra_configs/experiment/dino_proprio/direct_action.yaml b/egomimic/hydra_configs/experiment/dino_proprio/direct_action.yaml new file mode 100644 index 000000000..eec22a598 --- /dev/null +++ b/egomimic/hydra_configs/experiment/dino_proprio/direct_action.yaml @@ -0,0 +1,39 @@ +# @package _global_ +# Requires verified action commands aligned to the current rendered state. +defaults: + - /experiment/dino_proprio/visual_state + - _self_ + +name: dino_proprio_direct_action_pilot +data: + task: direct_action + action_manifest_path: null # Required by the direct-action data constructor. +model: + pipeline: + stages: + - _target_: egomimic.pipeline.stages_dino_proprio.DinoProprioStage + warmup_steps: 0 + pose_ramp_steps: 0 + image_weight: 0.0 + pose_weight: 1.0 + sampler_steps: 20 + model: + _target_: egomimic.models.dino_proprio.DinoProprioModel + objective: direct_action + pose_dims: + libero_gripper: 9 + # Expected LIBERO controller width; loader must verify actual labels. + action_dims: + libero_gripper: 7 + image_dim: 768 + width: 768 + depth: 8 + heads: 6 + mlp_ratio: 4 + max_cameras: 2 + max_tokens: 3202 + time_mode: ${data.mode} + attention_query_chunk_size: 512 + activation_checkpointing: true +evaluator: + objective: direct_action diff --git a/egomimic/hydra_configs/experiment/dino_proprio/visual_state.yaml b/egomimic/hydra_configs/experiment/dino_proprio/visual_state.yaml new file mode 100644 index 000000000..7e5e43a8d --- /dev/null +++ b/egomimic/hydra_configs/experiment/dino_proprio/visual_state.yaml @@ -0,0 +1,47 @@ +# @package _global_ +# Small gripper-only pilot. This predicts measured states, not control commands. +defaults: + - override /model: dino_proprio/base + - override /data: dino_proprio/prepared + - override /trainer: default + +name: dino_proprio_visual_state_pilot +description: dino_proprio_only_s42 +seed: 42 +norm_stats: null +normalizer: null +inference_config: + enabled: false +evaluator: + _target_: egomimic.eval.dino_proprio_eval.DinoProprioEval + objective: visual_proprio + sample_seed: 42 +trainer: + accelerator: gpu + precision: 32-true + devices: 1 + min_epochs: 0 + max_epochs: -1 + limit_train_batches: 1.0 + limit_val_batches: 1.0 + val_check_interval: 250 + check_val_every_n_epoch: null + max_steps: 1000 + num_sanity_val_steps: 0 + log_every_n_steps: 1 +callbacks: + model_checkpoint: + filename: "epoch_{epoch:04d}-step_{step:09d}" + every_n_epochs: 0 + every_n_train_steps: 250 + save_on_train_epoch_end: false + save_top_k: -1 + save_last: true +runtime: + strict_fp32: true +logger: + wandb: + project: mot-world-model + entity: rl2-group + group: dino-proprio-minimal-bc + tags: [dino-proprio-only, no-point-tracks, pilot] diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml new file mode 100644 index 000000000..22fe4780e --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml @@ -0,0 +1,31 @@ +# @package _global_ +# Deferred stage 2. Supply the strictly expanded tensor artifact and its SHA256. +defaults: + - /experiment/pixel_proprio/video_continuous + - override /model: pixel_proprio/readout + - _self_ + +name: pixel_proprio_readout_continuous +description: gripper_pixel160_continuous_frozen_image_measured9_control7_s42 +pixel_readout: + initial_weights: ??? + initial_weights_sha256: ??? + training_mode: frozen_image # Freeze all image weights, including noise conditioning. + pose_dims: {libero_gripper: 9} + action_dims: {libero_gripper: 7} + max_tokens: 25856 # 128 * (200 image + 1 measured-state + 1 command) tokens. + image_loss_weight: 1.0 + measured_state_loss_weight: 1.0 + control_command_loss_weight: 1.0 +data: + image_only: false + task: visual_state_action + max_frames: 128 # Preserve deferred stage2 sampling independently of stage1. +trainer: + max_steps: 1000 # Deferred readout keeps its original pilot budget. +evaluator: + objective: visual_state_action +logger: + wandb: + group: pixel-proprio-readout + tags: [pixels-from-scratch, frozen-image-expert, state-and-command-readout, no-point-tracks, gripper-only] diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/readout_discrete.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/readout_discrete.yaml new file mode 100644 index 000000000..62e8e8ee9 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/readout_discrete.yaml @@ -0,0 +1,18 @@ +# @package _global_ +defaults: + - /experiment/pixel_proprio/readout_continuous + - _self_ + +name: pixel_proprio_readout_discrete +description: gripper_pixel160_discrete16_stride4_frozen_image_measured36_control28_s42 +data: + mode: discrete + sampling_strategy: discrete_window + discrete_frames: 16 + stride: 4 +pixel_readout: + pose_dims: {libero_gripper: 36} # Four chronological normalized 9D measured states. + action_dims: {libero_gripper: 28} # Four chronological raw 7D command vectors. + max_tokens: 3232 # 16 * (200 image + 1 measured-state + 1 command) tokens. +pixel_video: + max_tokens: 3200 diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml new file mode 100644 index 000000000..c18a53122 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_continuous.yaml @@ -0,0 +1,10 @@ +# @package _global_ +defaults: + - /experiment/pixel_proprio/video_only + - _self_ +name: pixel_proprio_video_continuous_pilot +description: gripper_pixel160_patch16_continuous_full_episode_s42 +trainer: + max_steps: 300000 +data: + max_frames: 64 # Uniform timestamp selection still spans the full episode. diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml new file mode 100644 index 000000000..3e38073e5 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_discrete.yaml @@ -0,0 +1,16 @@ +# @package _global_ +defaults: + - /experiment/pixel_proprio/video_only + - _self_ + +name: pixel_proprio_video_discrete_pilot +description: gripper_pixel160_patch16_discrete16_stride4_s42 +trainer: + max_steps: 300000 +data: + mode: discrete + sampling_strategy: discrete_window + discrete_frames: 16 + stride: 4 +pixel_video: + max_tokens: 3200 # 16 * 2 * 100 image patches. diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml new file mode 100644 index 000000000..63141a7c9 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml @@ -0,0 +1,61 @@ +# @package _global_ +# First diagnostic pilot: RGB prediction conditioned only on observed images. +# No action head/loss or second-stage state readout is instantiated or trained. +defaults: + - /world_model_attention: flash4_bf16_mixed + - override /model: pixel_proprio/video + - override /data: pixel_proprio/prepared + - override /trainer: default + +name: pixel_proprio_video_only_pilot +description: pixel160_patch16_full_episode_s42 +seed: 42 +pixel_video: + max_tokens: 25600 # 128 * 2 * 100 image patches. +data: + image_only: true +norm_stats: null +normalizer: null +val_at_end: true +inference_config: + enabled: false +evaluator: + _target_: egomimic.eval.image_proprio_eval.ImageProprioEval + objective: visual_only + sample_seed: 42 + pixel_diagnostics: + artifact_dir: ${paths.output_dir}/pixel_diagnostics + max_samples_per_source_kind: 1 + playback_fps: 20 + write_video: true + log_wandb: true +trainer: + accelerator: gpu + precision: ${oc.select:world_model_attention.trainer_precision,32-true} + devices: 1 + min_epochs: 0 + max_epochs: -1 + limit_train_batches: 1.0 + limit_val_batches: 1.0 + val_check_interval: 250 + check_val_every_n_epoch: null + max_steps: 1000 + num_sanity_val_steps: 0 + log_every_n_steps: 1 +callbacks: + model_checkpoint: + filename: "epoch_{epoch:04d}-step_{step:09d}" + every_n_epochs: 0 + every_n_train_steps: 250 + save_on_train_epoch_end: false + save_top_k: -1 + save_last: true +runtime: + strict_fp32: ${world_model_attention.strict_fp32} + disable_tf32: ${world_model_attention.disable_tf32} +logger: + wandb: + project: mot-world-model + entity: rl2-group + group: pixel-proprio-minimal-bc + tags: [pixels-from-scratch, video-only, no-point-tracks, gripper-only, pilot] diff --git a/egomimic/hydra_configs/model/dino_proprio/base.yaml b/egomimic/hydra_configs/model/dino_proprio/base.yaml new file mode 100644 index 000000000..19caa62fa --- /dev/null +++ b/egomimic/hydra_configs/model/dino_proprio/base.yaml @@ -0,0 +1,39 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 0.0001 + weight_decay: 0.01 +scheduler: null +training_behavior: + _target_: egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior +train_log_on_step: true +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stage_ids: + dino_proprio: 0 + stages: + - _target_: egomimic.pipeline.stages_dino_proprio.DinoProprioStage + warmup_steps: 100 + pose_ramp_steps: 100 + image_weight: 1.0 + pose_weight: 1.0 + sampler_steps: 20 + model: + _target_: egomimic.models.dino_proprio.DinoProprioModel + objective: visual_proprio + pose_dims: + libero_gripper: 9 + action_dims: null + image_dim: 768 + width: 768 + depth: 8 + heads: 6 + mlp_ratio: 4 + max_cameras: 2 + max_tokens: 12804 + time_mode: ${data.mode} + attention_query_chunk_size: 512 + activation_checkpointing: true + pose_conditioning: low_noise + pose_noise_max: 0.2 diff --git a/egomimic/hydra_configs/model/pixel_proprio/base.yaml b/egomimic/hydra_configs/model/pixel_proprio/base.yaml new file mode 100644 index 000000000..5d752aa62 --- /dev/null +++ b/egomimic/hydra_configs/model/pixel_proprio/base.yaml @@ -0,0 +1,49 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +enable_grad_norm: false # Disable inherited median/MAD adaptive gradient clipping. +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 0.00005 + weight_decay: 0.01 +scheduler: null +training_behavior: + _target_: egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior +train_log_on_step: true +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stage_ids: + image_proprio: 0 + stages: + - _target_: egomimic.pipeline.stages_image_proprio.ImageProprioStage + warmup_steps: 0 + pose_ramp_steps: 0 + image_weight: 1.0 + pose_weight: 0.0 + sampler_steps: 20 + model: + _target_: egomimic.models.image_proprio.ImageProprioModel + objective: visual_only + phase_a_condition_on_proprio: false + pose_dims: null + action_dims: null + image_dim: 768 # 16 * 16 * 3 raw RGB values; no image encoder. + width: 768 + depth: 8 + heads: 6 + mlp_ratio: 4 + max_cameras: 2 + use_camera_embedding: true + use_modality_embedding: false + use_embodiment_embedding: false + max_tokens: 25600 # 128 * 2 * 100 image patches; no proprio stream. + time_mode: ${data.mode} + attention_backend: ${oc.select:world_model_attention.backend,sdpa} + attention_dtype: ${oc.select:world_model_attention.dtype,native} + strict_fp32: ${runtime.strict_fp32} + attention_query_chunk_size: ${oc.select:world_model_attention.query_chunk_size,0} + activation_checkpointing: ${oc.select:world_model_attention.activation_checkpointing,false} + allow_patch_compression: false + loss_reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + compile_mode: ${oc.select:world_model_execution.compile_mode,none} + compile_dynamic: ${oc.select:world_model_execution.compile_dynamic,false} + compile_backend: ${oc.select:world_model_execution.compile_backend,inductor} diff --git a/egomimic/hydra_configs/model/pixel_proprio/readout.yaml b/egomimic/hydra_configs/model/pixel_proprio/readout.yaml new file mode 100644 index 000000000..65f098273 --- /dev/null +++ b/egomimic/hydra_configs/model/pixel_proprio/readout.yaml @@ -0,0 +1,106 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +enable_grad_norm: false +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 0.00005 + weight_decay: 0.01 +scheduler: null +training_behavior: + _target_: egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior +train_log_on_step: true +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + initialization: + - stage_id: image_proprio + module_path: model + source: ${pixel_readout.initial_weights} + sha256: ${pixel_readout.initial_weights_sha256} + stage_ids: + pixel_noising: 0 + pixel_denoiser: 1 + pixel_prediction: 2 + pixel_loss: 3 + readout_noising: 4 + readout_denoiser: 5 + pixel_sampler: 6 + state_action_readout: 7 + measured_state_loss: 8 + control_command_loss: 9 + image_proprio: 1 + stages: + - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage + noising: + _target_: egomimic.models.pixel_diffusion.LinearGaussianImageNoising + maximum: 1.0 + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutModelStage + training_mode: ${pixel_readout.training_mode} + model: + _target_: egomimic.models.image_proprio.ImageProprioModel + objective: visual_state_action + phase_a_condition_on_proprio: false + pose_dims: ${pixel_readout.pose_dims} + action_dims: ${pixel_readout.action_dims} + image_dim: 768 + width: 768 + depth: 8 + heads: 6 + mlp_ratio: 4 + max_cameras: 2 + use_camera_embedding: true + use_modality_embedding: false + use_embodiment_embedding: false + use_phase_embedding: false + noise_conditioning: adaln_zero + position_encoding: spatiotemporal_rope + rope_axis_dims: [64, 32, 32] + spatial_grid_size: [10, 10] + max_tokens: ${pixel_readout.max_tokens} + time_mode: ${data.mode} + attention_backend: ${oc.select:world_model_attention.backend,sdpa} + attention_dtype: ${oc.select:world_model_attention.dtype,native} + strict_fp32: ${runtime.strict_fp32} + attention_query_chunk_size: ${oc.select:world_model_attention.query_chunk_size,0} + activation_checkpointing: ${oc.select:world_model_attention.activation_checkpointing,false} + allow_patch_compression: false + loss_reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + token_execution: ${oc.select:world_model_execution.token_execution,padded} + compile_mode: ${oc.select:world_model_execution.compile_mode,none} + compile_dynamic: ${oc.select:world_model_execution.compile_dynamic,false} + compile_backend: ${oc.select:world_model_execution.compile_backend,inductor} + # Optional high-noise prediction/loss branch; it does not own model weights. + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutVideoDenoiserStage + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutImageLossStage + weight: ${pixel_readout.image_loss_weight} + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutNoisingStage + noising: + _target_: egomimic.models.pixel_diffusion.LinearGaussianImageNoising + maximum: 0.2 + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutImageDenoiserStage + prediction_key: pixel_readout/image_prediction + # Training keeps the existing low-noise provider above. A future sampled + # training provider can use the same tensor port with explicit modes: [train]. + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutSampledImageStage + modes: [inference] + prediction_key: pixel_readout/image_prediction + gradient_policy: stop + sampler: + _target_: egomimic.models.pixel_diffusion.LinearDDIMImageSampler + steps: 20 + - _target_: egomimic.pipeline.stages_pixel_readout.PixelStateActionReadoutStage + image_key: pixel_readout/image_prediction + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutMSEStage + stream: proprio + weight: ${pixel_readout.measured_state_loss_weight} + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + - _target_: egomimic.pipeline.stages_pixel_readout.PixelReadoutMSEStage + stream: action + weight: ${pixel_readout.control_command_loss_weight} + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} diff --git a/egomimic/hydra_configs/model/pixel_proprio/video.yaml b/egomimic/hydra_configs/model/pixel_proprio/video.yaml new file mode 100644 index 000000000..fbc5d875b --- /dev/null +++ b/egomimic/hydra_configs/model/pixel_proprio/video.yaml @@ -0,0 +1,67 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +enable_grad_norm: false +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 0.00005 + weight_decay: 0.01 +scheduler: null +training_behavior: + _target_: egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior +train_log_on_step: true +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stage_ids: + pixel_noising: 0 + pixel_denoiser: 1 + pixel_loss: 2 + pixel_sampler: 3 + image_proprio: 1 # Existing stable initialization name, same registered model. + stages: + - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoNoisingStage + noising: + _target_: egomimic.models.pixel_diffusion.LinearGaussianImageNoising + maximum: 1.0 + - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoDenoiserStage + model: + _target_: egomimic.models.image_proprio.ImageProprioModel + objective: visual_only + phase_a_condition_on_proprio: false + pose_dims: null + action_dims: null + image_dim: 768 # 16 * 16 * 3 raw RGB values; no image encoder. + width: 768 + depth: 8 + heads: 6 + mlp_ratio: 4 + max_cameras: 2 + use_camera_embedding: true + use_modality_embedding: false + use_embodiment_embedding: false + use_phase_embedding: false + noise_conditioning: adaln_zero + position_encoding: spatiotemporal_rope + rope_axis_dims: [64, 32, 32] + spatial_grid_size: [10, 10] + max_tokens: ${pixel_video.max_tokens} + time_mode: ${data.mode} + attention_backend: ${oc.select:world_model_attention.backend,sdpa} + attention_dtype: ${oc.select:world_model_attention.dtype,native} + strict_fp32: ${runtime.strict_fp32} + attention_query_chunk_size: ${oc.select:world_model_attention.query_chunk_size,0} + activation_checkpointing: ${oc.select:world_model_attention.activation_checkpointing,false} + allow_patch_compression: false + loss_reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + token_execution: ${oc.select:world_model_execution.token_execution,padded} + compile_mode: ${oc.select:world_model_execution.compile_mode,none} + compile_dynamic: ${oc.select:world_model_execution.compile_dynamic,false} + compile_backend: ${oc.select:world_model_execution.compile_backend,inductor} + - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoLossStage + weight: 1.0 + loss: + _target_: egomimic.models.pixel_diffusion.PlainPixelMSE + reduction: ${oc.select:world_model_execution.loss_reduction,token_mean} + - _target_: egomimic.pipeline.stages_pixel_video.PixelVideoSamplerStage + sampler: + _target_: egomimic.models.pixel_diffusion.LinearDDIMImageSampler + steps: 20 diff --git a/egomimic/hydra_configs/pixel_execution/ddp4.yaml b/egomimic/hydra_configs/pixel_execution/ddp4.yaml new file mode 100644 index 000000000..4c34190a2 --- /dev/null +++ b/egomimic/hydra_configs/pixel_execution/ddp4.yaml @@ -0,0 +1,18 @@ +# @package _global_ +# Two independent runs can each use four GPUs on one allocated eight-GPU node. +# Rank-aware pixel samplers own training draws and the nonduplicated val shards. +data: + batch_size: 1 + distributed_world_size: 4 +model: + training_behavior: + distributed_noise_seed: ${seed} +trainer: + strategy: ddp + devices: 4 + num_nodes: 1 + accumulate_grad_batches: 1 + use_distributed_sampler: false +runtime: + cluster_environment: torch_elastic + slurm_requeue_owner: none diff --git a/egomimic/hydra_configs/pixel_execution/ddp4_prefetch.yaml b/egomimic/hydra_configs/pixel_execution/ddp4_prefetch.yaml new file mode 100644 index 000000000..a79b5d686 --- /dev/null +++ b/egomimic/hydra_configs/pixel_execution/ddp4_prefetch.yaml @@ -0,0 +1,13 @@ +# @package _global_ +# Keep held-out anchor/noise identity fixed while benchmarking train batch size. +defaults: + - ddp4 + - _self_ + +data: + validation_batch_size: 1 + num_workers: 2 + pin_memory: true + persistent_workers: true + prefetch_factor: 2 + balance_training_work: true diff --git a/egomimic/hydra_configs/world_model_attention/flash3_bf16_mixed.yaml b/egomimic/hydra_configs/world_model_attention/flash3_bf16_mixed.yaml new file mode 100644 index 000000000..007d48322 --- /dev/null +++ b/egomimic/hydra_configs/world_model_attention/flash3_bf16_mixed.yaml @@ -0,0 +1,7 @@ +# @package _global_ +defaults: + - flash4_bf16_mixed + - _self_ +# Optional separately pinned Hopper backend; missing dependency fails closed. +world_model_attention: + backend: flash3 diff --git a/egomimic/hydra_configs/world_model_attention/flash4_bf16_mixed.yaml b/egomimic/hydra_configs/world_model_attention/flash4_bf16_mixed.yaml new file mode 100644 index 000000000..ebb677992 --- /dev/null +++ b/egomimic/hydra_configs/world_model_attention/flash4_bf16_mixed.yaml @@ -0,0 +1,10 @@ +# @package _global_ +defaults: + - flash4_bf16 + - _self_ +# Standard autocast: FP32 stored parameters/AdamW state, BF16 eligible operations. +# FA4 Q/K/V remain BF16; FP32 MSE and time coordinates are model-owned contracts. +world_model_attention: + trainer_precision: bf16-mixed +trainer: + precision: ${world_model_attention.trainer_precision} diff --git a/egomimic/hydra_configs/world_model_attention/native.yaml b/egomimic/hydra_configs/world_model_attention/native.yaml index 631deb1c2..712638c7a 100644 --- a/egomimic/hydra_configs/world_model_attention/native.yaml +++ b/egomimic/hydra_configs/world_model_attention/native.yaml @@ -7,3 +7,4 @@ world_model_attention: activation_checkpointing: false strict_fp32: false disable_tf32: false + trainer_precision: 32-true diff --git a/egomimic/hydra_configs/world_model_execution/batched_eager.yaml b/egomimic/hydra_configs/world_model_execution/batched_eager.yaml new file mode 100644 index 000000000..6b918edcc --- /dev/null +++ b/egomimic/hydra_configs/world_model_execution/batched_eager.yaml @@ -0,0 +1,7 @@ +# @package _global_ +# Preserve the previous B1/DDP objective when combining variable-length clips. +world_model_execution: + loss_reduction: mean_per_example + compile_mode: none + compile_dynamic: false + compile_backend: inductor diff --git a/egomimic/hydra_configs/world_model_execution/packed_regional_compile.yaml b/egomimic/hydra_configs/world_model_execution/packed_regional_compile.yaml new file mode 100644 index 000000000..ad28df2b2 --- /dev/null +++ b/egomimic/hydra_configs/world_model_execution/packed_regional_compile.yaml @@ -0,0 +1,8 @@ +# @package _global_ +defaults: + - regional_compile + - _self_ +# Opt in to valid-token execution through image-only Phase A. Ragged inputs +# require the selected FlashAttention3/4 backend; dense batches keep their path. +world_model_execution: + token_execution: packed_valid diff --git a/egomimic/hydra_configs/world_model_execution/regional_compile.yaml b/egomimic/hydra_configs/world_model_execution/regional_compile.yaml new file mode 100644 index 000000000..f1c2b22e6 --- /dev/null +++ b/egomimic/hydra_configs/world_model_execution/regional_compile.yaml @@ -0,0 +1,10 @@ +# @package _global_ +defaults: + - batched_eager + - _self_ +# Compile tensor regions with fullgraph=True; optional attention kernels remain +# explicit eager calls. Dynamic shapes suit variable-length continuous clips. +# Fixed-window discrete runs may override compile_dynamic=false. +world_model_execution: + compile_mode: default + compile_dynamic: true diff --git a/egomimic/models/attention_execution.py b/egomimic/models/attention_execution.py new file mode 100644 index 000000000..eb056c6a2 --- /dev/null +++ b/egomimic/models/attention_execution.py @@ -0,0 +1,268 @@ +"""Reusable sequence layout and explicit optional fused-attention kernels. + +Layout construction is an eager boundary: inspect masks once, then reuse the +same gather indices and sequence boundaries across layers and sampling steps. +No backend or layout changes model parameters or permits cross-example keys. +""" + +import importlib +import math +from collections.abc import Sequence +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable + +import torch + + +@dataclass(frozen=True) +class AttentionLayout: + kind: str + batch_size: int + query_tokens: int + key_tokens: int + key_mask: torch.Tensor + query_indices: torch.Tensor | None = None + key_indices: torch.Tensor | None = None + cu_seqlens_q: torch.Tensor | None = None + cu_seqlens_k: torch.Tensor | None = None + max_seqlen_q: int = 0 + max_seqlen_k: int = 0 + + +def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_valid=None): + """Prepare dense or packed attention without assuming prefix-only padding. + + ``query_valid=None`` retains every query, including for query chunks or + cross attention. Supply the full self-attention validity mask to pack Q as + well as K/V. ``all_valid=True`` is a trusted, already-validated caller + contract and avoids reading CUDA mask values entirely. + """ + if valid.ndim != 2 or valid.dtype != torch.bool: + raise ValueError("Attention validity must be boolean [B,N]") + batch, key_tokens = valid.shape + if batch < 1 or key_tokens < 1: + raise ValueError("Attention requires nonempty batch and key dimensions") + if query_valid is not None: + if (query_valid.ndim != 2 or query_valid.dtype != torch.bool + or query_valid.shape[0] != batch or query_valid.device != valid.device): + raise ValueError("Query validity must be boolean [B,Q] on the key device") + if query_tokens is not None and query_tokens != query_valid.shape[1]: + raise ValueError("Query token count differs from query validity") + query_tokens = query_valid.shape[1] + query_tokens = key_tokens if query_tokens is None else query_tokens + if type(query_tokens) is not int or query_tokens < 1: + raise ValueError("Query token count must be positive") + if all_valid is not None and type(all_valid) is not bool: + raise ValueError("all_valid must be a boolean or None") + if all_valid is None: + all_valid = bool(valid.all()) and ( + query_valid is None or bool(query_valid.all()) + ) + # Clone masks so later batch mutations cannot alter this prepared layout. + key_mask = valid[:, None, None, :].clone() + if all_valid: + return AttentionLayout("dense", batch, query_tokens, key_tokens, key_mask) + key_counts = valid.sum(-1, dtype=torch.int32) + query_counts = ( + query_valid.sum(-1, dtype=torch.int32) if query_valid is not None + else torch.full((batch,), query_tokens, device=valid.device, dtype=torch.int32) + ) + # This one eager transfer also verifies the optional kernel preconditions. + counts = torch.stack((query_counts, key_counts)).cpu().tolist() + if min(counts[1]) < 1: + raise ValueError("Fused attention requires at least one valid key per example") + if min(counts[0]) < 1: + raise ValueError("Packed attention requires at least one valid query per example") + zero = key_counts.new_zeros(1) + return AttentionLayout( + "varlen", batch, query_tokens, key_tokens, key_mask, + query_indices=(query_valid.reshape(-1).nonzero(as_tuple=True)[0] + if query_valid is not None else None), + key_indices=valid.reshape(-1).nonzero(as_tuple=True)[0], + cu_seqlens_q=torch.cat((zero, query_counts.cumsum(0, dtype=torch.int32))), + cu_seqlens_k=torch.cat((zero, key_counts.cumsum(0, dtype=torch.int32))), + max_seqlen_q=max(counts[0]), max_seqlen_k=max(counts[1]), + ) + + +def prepare_temporal_rope(times, head_dim): + """Return reusable FP32 (cos, sin) for times[B,N] and even head width.""" + if times.ndim != 2 or type(head_dim) is not int or head_dim < 2 or head_dim % 2: + raise ValueError("RoPE requires [B,N] times and a positive even head width") + half = head_dim // 2 + freq = torch.exp(-math.log(10000) * torch.arange( + half, device=times.device, dtype=torch.float32 + ) / half) + phase = times.float()[:, None, :, None] * freq + return phase.cos(), phase.sin() + + +def resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size): + """Validate time/x/y channel allocation and explicit image patch grid H/W.""" + if type(head_dim) is not int or head_dim < 6 or head_dim % 2: + raise ValueError("Spatial-temporal RoPE requires an even head width of at least 6") + if axis_dims is None: + if head_dim % 8: + raise ValueError("Automatic RoPE axis dimensions require head width divisible by 8") + axis_dims = (head_dim // 2, head_dim // 4, head_dim // 4) + if (not isinstance(axis_dims, Sequence) or isinstance(axis_dims, (str, bytes)) or len(axis_dims) != 3 + or any(type(value) is not int or value < 2 or value % 2 for value in axis_dims) + or sum(axis_dims) != head_dim): + raise ValueError("rope_axis_dims must be positive-even time/x/y dimensions summing to head width") + if (not isinstance(spatial_grid_size, Sequence) or isinstance(spatial_grid_size, (str, bytes)) or len(spatial_grid_size) != 2 + or any(type(value) is not int or value < 1 for value in spatial_grid_size)): + raise ValueError("spatial_grid_size must explicitly name positive integer patch-grid height/width") + return tuple(axis_dims), tuple(spatial_grid_size) + + +def prepare_spatiotemporal_rope(times, xy, head_dim, *, axis_dims=None, spatial_grid_size=None): + """FP32 axial rotations for time plus normalized image-patch center XY. + + Normalized centers scale by(grid_width/2,grid_height/2), producing centered + patch units: a10-wide grid has x=-4.5,...,4.5 with adjacent spacing1. + ``xy=None`` gives nonimage tokens zero spatial phases while preserving the + same temporal frequencies as image tokens in joint attention. + + The existing rotation pairs the first and second halves of each head. + Therefore each half orders its coordinates as[time,x,y], and this helper + concatenates axis frequency vectors in that order. A64/32/32 allocation + means32 time,16 x and16 y *pairs*, not a contiguous-slice rotation. + """ + axes, grid = resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size) + if times.ndim != 2 or times.dtype == torch.bool or times.is_complex(): + raise ValueError("Spatial-temporal RoPE requires real-valued times[B,N]") + if xy is not None and (xy.shape != (*times.shape, 2) or xy.device != times.device + or not xy.is_floating_point()): + raise ValueError("Spatial RoPE coordinates must be floating-point[B,N,2] on the time device") + positions = [times.float()] + positions.extend((torch.zeros_like(times, dtype=torch.float32),) * 2 if xy is None else + (xy[..., 0].float() * (grid[1] / 2), xy[..., 1].float() * (grid[0] / 2))) + phases = [] + for coordinate, dimension in zip(positions, axes): + pairs = dimension // 2 + frequency = torch.exp(-math.log(10000) * torch.arange( + pairs, device=times.device, dtype=torch.float32 + ) / pairs) + phases.append(coordinate[:, None, :, None] * frequency) + phase = torch.cat(phases, dim=-1) + return phase.cos(), phase.sin() + + +def apply_temporal_rope(x, rope): + cos, sin = rope + left, right = x.float().chunk(2, dim=-1) + return torch.cat((left * cos - right * sin, left * sin + right * cos), -1).to(x.dtype) + + +@dataclass(frozen=True) +class FlashAttentionKernels: + backend: str + dense: Callable + varlen: Callable + + +@lru_cache(maxsize=16) +def load_flash_attention_kernels(backend, device): + """Resolve a declared CUDA backend once; missing kernels never fall back.""" + device = torch.device(device) + if backend not in ("flash3", "flash4"): + raise ValueError("Expected flash3 or flash4 backend") + if device.type != "cuda": + raise RuntimeError(f"{backend} attention requires a CUDA Hopper/Blackwell GPU") + if torch.cuda.get_device_capability(device)[0] < 9: + raise RuntimeError(f"{backend} requires Hopper or newer CUDA hardware") + if backend == "flash4": + from packaging.version import Version + if Version(torch.__version__.split("+")[0]) < Version("2.9"): + raise RuntimeError("flash4 requires the separately verified Torch>=2.9 runtime") + module_names = (("flash_attn.cute",) if backend == "flash4" + else ("flash_attn_3.flash_attn_interface", "flash_attn_interface")) + try: + for index, module_name in enumerate(module_names): + try: + module = importlib.import_module(module_name) + break + except ModuleNotFoundError as error: + # Older official FA3 builds expose a top-level alias. Only + # missing module/package names permit that compatibility path; + # an installed backend with a missing dependency must fail. + if (index == len(module_names) - 1 + or error.name not in {module_name, module_name.split(".")[0]}): + raise + dense, varlen = module.flash_attn_func, module.flash_attn_varlen_func + except (ImportError, AttributeError) as error: + raise RuntimeError( + f"Install the verified optional {backend} runtime ({module_name}); no backend fallback" + ) from error + if not callable(dense) or not callable(varlen): + raise RuntimeError(f"Unsupported {backend} dense/varlen API") + return FlashAttentionKernels(backend, dense, varlen) + + +def _output_tensor(result, expected): + # FA4 b35 returns (out, optional_lse); FA3 returns out when LSE is disabled. + if isinstance(result, tuple): + if len(result) != 2: + raise RuntimeError("Unsupported fused-attention result tuple") + result = result[0] + if (not torch.is_tensor(result) or result.shape != expected.shape + or result.dtype != expected.dtype or result.device != expected.device): + raise RuntimeError("Unsupported fused-attention output shape/dtype/device") + return result + + +def execute_flash_attention(query, key, value, layout, kernels): + """Execute prepared full attention; arguments use [B,H,N,D] layout.""" + if any(x.dtype != torch.bfloat16 for x in (query, key, value)): + raise ValueError("Fused attention requires bfloat16 Q/K/V") + b, h, nq, d = query.shape + if (b != layout.batch_size or nq != layout.query_tokens + or key.shape != (b, h, layout.key_tokens, d) or value.shape != key.shape + or key.device != query.device or value.device != query.device + or layout.key_mask.device != query.device): + raise ValueError("Prepared attention layout differs from Q/K/V shape/device") + q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value)) + if layout.kind == "dense": + out = _output_tensor(kernels.dense(q, k, v, causal=False), q) + return out.transpose(1, 2) + if layout.kind != "varlen": + raise ValueError("Unknown attention layout kind") + q = q.reshape(-1, h, d) + if layout.query_indices is not None: + q = q.index_select(0, layout.query_indices) + k = k.reshape(-1, h, d).index_select(0, layout.key_indices) + v = v.reshape(-1, h, d).index_select(0, layout.key_indices) + out = _output_tensor(kernels.varlen( + q, k, v, cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), q) + if layout.query_indices is not None: + out = out.new_zeros(b * nq, h, d).index_copy(0, layout.query_indices, out) + return out.reshape(b, nq, h, d).transpose(1, 2) + + +def execute_packed_flash_attention(query, key, value, layout, kernels): + """Attend already-packed [total_tokens,H,D] tensors without repadding. + + The layout retains the original examples' boundaries. It must never be + rebuilt as one sequence merely because tensor regions use a leading 1. + """ + if layout.kind != "varlen": + raise ValueError("Direct packed attention requires a varlen layout") + if any(x.ndim != 3 or x.dtype != torch.bfloat16 for x in (query, key, value)): + raise ValueError("Packed fused attention requires bfloat16 [tokens,H,D] Q/K/V") + query_count = (layout.query_indices.numel() if layout.query_indices is not None + else layout.batch_size * layout.query_tokens) + if (query.shape[0] != query_count + or key.shape != (layout.key_indices.numel(), *query.shape[1:]) + or value.shape != key.shape or key.device != query.device + or value.device != query.device or layout.key_mask.device != query.device): + raise ValueError("Prepared packed attention layout differs from Q/K/V shape/device") + return _output_tensor(kernels.varlen( + query.contiguous(), key.contiguous(), value.contiguous(), + cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), query) diff --git a/egomimic/models/compilation.py b/egomimic/models/compilation.py new file mode 100644 index 000000000..a7ff7d458 --- /dev/null +++ b/egomimic/models/compilation.py @@ -0,0 +1,48 @@ +"""Explicit compilation of tensor-only regions, without eager fallback.""" + +from functools import wraps + +import torch +from torch._functorch import config as _aot_config + + +def compile_region(function, *, mode="none", dynamic=False, backend="inductor"): + if mode not in ("none", "default", "reduce-overhead", "max-autotune"): + raise ValueError("Unknown compile mode") + if type(dynamic) is not bool: + raise ValueError("compile dynamic must be a boolean") + if not isinstance(backend, str) or not backend: + raise ValueError("compile backend must be a nonempty string") + if mode == "none": + return function + if torch._dynamo.config.suppress_errors: + raise RuntimeError("Compilation requires suppress_errors=False; no eager fallback") + compiled = torch.compile( + function, backend=backend, mode=mode, dynamic=dynamic, fullgraph=True + ) + + @wraps(function) + def execute(*args, **kwargs): + # Compilation is lazy, so the policy must enclose invocation, including + # retracing, rather than only the torch.compile call above. Lightning's + # standard AMP backward runs outside the forward autocast context. + # These are already bounded tensor regions inside the ordinary DDP + # wrapper. DDPOptimizer must not repartition each region against the + # whole model's gradient buckets (Torch 2.9 can emit invalid SymInt + # subgraph outputs). This affects compilation only; DDP's reducer and + # gradient synchronization remain active. Restore the caller's policy. + with torch._dynamo.config.patch(optimize_ddp=False): + if hasattr(_aot_config, "backward_pass_autocast"): + with _aot_config.patch(backward_pass_autocast="off"): + return compiled(*args, **kwargs) + # Older runtimes can still compile ordinary FP32 regions, but cannot + # express the required AOT backward policy for our CPU/CUDA mixed mode. + if torch.is_autocast_enabled("cpu") or torch.is_autocast_enabled("cuda"): + raise RuntimeError( + "Compiled autocast requires a PyTorch runtime with " + "torch._functorch.config.backward_pass_autocast; " + "no unsafe mixed-precision fallback" + ) + return compiled(*args, **kwargs) + + return execute diff --git a/egomimic/models/dino_proprio.py b/egomimic/models/dino_proprio.py new file mode 100644 index 000000000..444f10ea6 --- /dev/null +++ b/egomimic/models/dino_proprio.py @@ -0,0 +1,28 @@ +"""Compatibility class for the original unlaunched DINO/proprio recipes. + +The feature-agnostic implementation lives in :mod:`image_proprio`. Its default +width guard and identity embeddings retain the original model's defaults. +""" + +from egomimic.models.image_proprio import MODALITIES, ImageProprioModel + + +class DinoProprioModel(ImageProprioModel): + def __init__( + self, + *args, + use_camera_embedding=True, + use_modality_embedding=True, + use_embodiment_embedding=True, + **kwargs, + ): + super().__init__( + *args, + use_camera_embedding=use_camera_embedding, + use_modality_embedding=use_modality_embedding, + use_embodiment_embedding=use_embodiment_embedding, + **kwargs, + ) + + +__all__ = ["MODALITIES", "DinoProprioModel"] diff --git a/egomimic/models/image_proprio.py b/egomimic/models/image_proprio.py new file mode 100644 index 000000000..909aef38d --- /dev/null +++ b/egomimic/models/image_proprio.py @@ -0,0 +1,952 @@ +"""Minimal two-stream image-patch/proprio models; no point-track dependency. + +``visual_proprio`` predicts clean image-patch values under linear Gaussian +corruption, then reads out withheld proprioceptive states in a separate phase. +``direct_action`` regresses recorded action commands from clean observed image +patches and observed proprioception. Action labels are never model inputs. +``visual_only`` predicts image patches, without state/action heads, withheld +proprio queries, or a readout phase. With phase_a_condition_on_proprio=False it +constructs only an image expert and takes only images. The same flag keeps +phase A image-only when later adding a proprio expert for phase B readout. +``visual_state_action`` adds typed state and command queries in phase B, sharing +the existing proprio branch. Both readouts use MSE; action labels stay loss-only. + +Prepared streams contain value[B,N,D], time[B,N], valid[B,N], observed[B,N]. +Image tokens additionally carry xy[B,N,2] and camera[B,N]. Physical time and +noise level are separate conditioning variables. No image observed/query flag +is embedded. Target availability is used only by the loss. Image values may be +prepared features or flattened RGB patches; the model adds no image encoder. +Camera identity conditioning is enabled by default to distinguish multiple views. +Modality and embodiment identity embeddings are disabled by default; private +modality weights and the dimension-specific proprio adapters remain in use. +Channel compression requires an explicit constructor opt-in. That opt-in does +not add a wide denoising head or remove the resulting low-rank output constraint. +Attention execution is independently selectable: native SDPA remains the default, +while explicit BF16 attention changes only Q/K/V and attention output arithmetic. +The caller owns trainer precision and the TF32 policy for the remaining model. +Learned phase offsets are optional; A/B labels always select control flow. +The optional spatial-temporal RoPE mode replaces additive XY with axial query/key +rotations in centered patch units. Camera conditioning remains separate. State +and command tokens share image temporal frequencies and have zero spatial phases. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass + +import torch +from torch import nn +from torch.utils.checkpoint import checkpoint + +from egomimic.models.attention_execution import ( + AttentionLayout, + load_flash_attention_kernels, + prepare_attention_layout, + prepare_spatiotemporal_rope, + prepare_temporal_rope, + resolve_spatiotemporal_rope_geometry, +) +from egomimic.models.compilation import compile_region +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, + PlainPixelMSE, +) +from egomimic.models.world_model import MoTBlock + +MODALITIES = ("image", "proprio") + + +@dataclass(frozen=True) +class _PackedImageExecution: + layout: AttentionLayout + image: Mapping[str, torch.Tensor] + valid: torch.Tensor + rope: tuple[torch.Tensor, torch.Tensor] + + +@dataclass(frozen=True) +class _NoiseFrameMap: + representatives: torch.Tensor + token_groups: torch.Tensor + + +@dataclass(frozen=True) +class _ConditionedImageExecution: + attention: object + noise_frames: _NoiseFrameMap + + +class ImageProprioModel(nn.Module): + def __init__( + self, + pose_dims: Mapping[str, int] | None = None, + action_dims: Mapping[str, int] | None = None, + objective="visual_proprio", + image_dim=768, + width=768, + depth=8, + heads=6, + mlp_ratio=4, + max_cameras=3, + max_tokens=76816, + time_mode="continuous", + attention_query_chunk_size=0, + activation_checkpointing=False, + pose_conditioning="low_noise", + pose_noise_max=0.2, + allow_patch_compression=False, + use_camera_embedding=True, + use_modality_embedding=False, + use_embodiment_embedding=False, + attention_backend="sdpa", + attention_dtype="native", + strict_fp32=False, + phase_a_condition_on_proprio=True, + loss_reduction="token_mean", + compile_mode="none", + compile_dynamic=False, + compile_backend="inductor", + token_execution="padded", + use_phase_embedding=True, + noise_conditioning="additive", + position_encoding="linear_xy_temporal_rope", + rope_axis_dims=None, + spatial_grid_size=None, + ): + super().__init__() + if type(strict_fp32) is not bool: + raise ValueError("strict_fp32 must be a boolean") + if strict_fp32 and attention_dtype != "native": + raise ValueError( + "Explicit BF16 attention conflicts with strict_fp32; " + "select the mixed-attention execution profile" + ) + if objective not in ( + "visual_only", + "visual_proprio", + "visual_state_action", + "direct_action", + ): + raise ValueError( + "objective must be visual_only, visual_proprio, visual_state_action or direct_action" + ) + if type(phase_a_condition_on_proprio) is not bool: + raise ValueError("phase_a_condition_on_proprio must be bool") + if objective == "direct_action" and not phase_a_condition_on_proprio: + raise ValueError("Direct action BC requires its observed proprio context") + self.phase_a_condition_on_proprio = phase_a_condition_on_proprio + self.has_proprio_branch = ( + objective != "visual_only" or phase_a_condition_on_proprio + ) + self.modalities = MODALITIES if self.has_proprio_branch else ("image",) + if not self.has_proprio_branch: + if pose_dims not in (None, {}): + raise ValueError("Image-only visual_only must omit pose_dims") + if use_embodiment_embedding: + raise ValueError("Image-only visual_only has no embodiment embedding") + if type(image_dim) is not int or image_dim < 1: + raise ValueError("image_dim must be a positive integer") + if type(allow_patch_compression) is not bool: + raise ValueError("allow_patch_compression must be bool") + for name, value in ( + ("use_camera_embedding", use_camera_embedding), + ("use_modality_embedding", use_modality_embedding), + ("use_embodiment_embedding", use_embodiment_embedding), + ("use_phase_embedding", use_phase_embedding), + ): + if type(value) is not bool: + raise ValueError(f"{name} must be bool") + if type(width) is not int or width < 1: + raise ValueError("width must be a positive integer") + if width < image_dim and not allow_patch_compression: + raise ValueError( + "Hidden width must be at least image_dim unless " + "allow_patch_compression=True is explicitly selected" + ) + if type(depth) is not int or depth < 1: + raise ValueError("depth must be a positive integer") + if type(heads) is not int or heads < 1: + raise ValueError("heads must be a positive integer") + if type(mlp_ratio) is not int or mlp_ratio < 1: + raise ValueError("mlp_ratio must be a positive integer") + if type(max_tokens) is not int or max_tokens < 1: + raise ValueError("max_tokens must be a positive integer") + if type(activation_checkpointing) is not bool: + raise ValueError("activation_checkpointing must be bool") + if time_mode not in ("continuous", "discrete"): + raise ValueError("time_mode must be continuous or discrete") + if pose_conditioning not in ("clean", "low_noise"): + raise ValueError("pose_conditioning must be clean or low_noise") + if not math.isfinite(pose_noise_max) or not 0 < pose_noise_max <= 0.2: + raise ValueError("pose_noise_max must be in (0, 0.2]") + self.pose_dims = ( + self._dimensions(pose_dims, "pose_dims") if self.has_proprio_branch else {} + ) + self.action_dims = ( + self._dimensions(action_dims, "action_dims") + if objective in ("direct_action", "visual_state_action") + else {} + ) + if ( + objective in ("direct_action", "visual_state_action") + and self.action_dims.keys() != self.pose_dims.keys() + ): + raise ValueError("action_dims and pose_dims must name the same embodiments") + if ( + objective not in ("direct_action", "visual_state_action") + and action_dims is not None + ): + raise ValueError("This objective has no action head; omit action_dims") + self.objective = objective + if noise_conditioning not in ("additive", "adaln_zero"): + raise ValueError("noise_conditioning must be additive or adaln_zero") + if noise_conditioning == "adaln_zero" and objective == "direct_action": + raise ValueError("Direct action BC has no diffusion noise conditioning") + self.noise_conditioning = noise_conditioning + if position_encoding not in ("linear_xy_temporal_rope", "spatiotemporal_rope"): + raise ValueError("position_encoding must be linear_xy_temporal_rope or spatiotemporal_rope") + self.position_encoding = position_encoding + if position_encoding == "spatiotemporal_rope": + if width % heads: + raise ValueError("Model width must be divisible by attention heads") + self.rope_axis_dims, self.spatial_grid_size = resolve_spatiotemporal_rope_geometry( + width // heads, rope_axis_dims, spatial_grid_size + ) + else: + if rope_axis_dims is not None or spatial_grid_size is not None: + raise ValueError("Spatial RoPE geometry requires position_encoding='spatiotemporal_rope'") + self.rope_axis_dims = self.spatial_grid_size = None + self.image_dim, self.width = image_dim, width + self.allow_patch_compression = allow_patch_compression + self.use_camera_embedding = use_camera_embedding + self.use_modality_embedding = use_modality_embedding + self.use_embodiment_embedding = use_embodiment_embedding + self.use_phase_embedding = use_phase_embedding + self.max_tokens, self.time_mode = max_tokens, time_mode + self.activation_checkpointing = activation_checkpointing + self.attention_backend = attention_backend + self.attention_dtype = attention_dtype + self.strict_fp32 = strict_fp32 + if loss_reduction not in ("token_mean", "mean_per_example"): + raise ValueError("loss_reduction must be token_mean or mean_per_example") + self.loss_reduction = loss_reduction + self.compile_mode, self.compile_dynamic = compile_mode, compile_dynamic + self.compile_backend = compile_backend + if token_execution not in ("padded", "packed_valid"): + raise ValueError("token_execution must be padded or packed_valid") + self.token_execution = token_execution + self.pose_conditioning, self.pose_noise_max = pose_conditioning, pose_noise_max + self.image_in = nn.Linear(image_dim, width) + if self.has_proprio_branch: + self.pose_in = nn.ModuleDict( + {k: nn.Linear(d, width) for k, d in self.pose_dims.items()} + ) + if position_encoding == "linear_xy_temporal_rope": + self.xy = nn.Linear(2, width) + if use_camera_embedding: + self.camera = nn.Embedding(max_cameras, width) + if use_modality_embedding: + self.modality = nn.Parameter( + torch.randn(len(self.modalities), width) * 0.02 + ) + if use_embodiment_embedding: + self.embodiment = nn.Embedding(len(self.pose_dims), width) + self.embodiment_ids = {k: i for i, k in enumerate(self.pose_dims)} + roles = ( + 1 + if objective == "visual_only" + else 3 + if objective == "visual_state_action" + else 2 + ) + if self.has_proprio_branch: + self.proprio_role = nn.Embedding(roles, width) + if objective != "visual_only": + self.query = nn.Parameter(torch.randn(width) * 0.02) + if objective == "visual_state_action": + self.action_query = nn.Parameter(torch.randn(width) * 0.02) + if objective != "direct_action": + self.image_out = nn.Linear(width, image_dim) + if objective in ("visual_proprio", "visual_state_action"): + self.pose_out = nn.ModuleDict( + {k: nn.Linear(width, d) for k, d in self.pose_dims.items()} + ) + self.noise_level = nn.Sequential( + nn.Linear(256 if noise_conditioning == "adaln_zero" else 1, width), + nn.SiLU(), nn.Linear(width, width) + ) + if use_phase_embedding: + self.phase = nn.Embedding(1 if objective == "visual_only" else 2, width) + if objective in ("direct_action", "visual_state_action"): + self.action_out = nn.ModuleDict( + {k: nn.Linear(width, d) for k, d in self.action_dims.items()} + ) + self.blocks = nn.ModuleList( + MoTBlock( + width, + heads, + mlp_ratio, + attention_query_chunk_size, + modalities=self.modalities, + attention_backend=attention_backend, + attention_dtype=attention_dtype, + image_adaln=noise_conditioning == "adaln_zero", + ) + for _ in range(depth) + ) + self.norm = nn.ModuleList([ + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if noise_conditioning == "adaln_zero" else nn.LayerNorm(width) + ]) + if self.has_proprio_branch: + self.norm.append( + nn.Identity() if objective == "visual_only" else nn.LayerNorm(width) + ) + if noise_conditioning == "adaln_zero": + self.final_adaln = nn.Sequential(nn.SiLU(), nn.Linear(width, 2 * width)) + nn.init.zeros_(self.final_adaln[-1].weight) + nn.init.zeros_(self.final_adaln[-1].bias) + nn.init.zeros_(self.image_out.weight) + nn.init.zeros_(self.image_out.bias) + for layer in (self.noise_level[0], self.noise_level[2]): + nn.init.normal_(layer.weight, std=0.02) + nn.init.zeros_(layer.bias) + # Compile pure tensor regions, keeping validation, packing, and optional + # attention-library dispatch outside graph capture. Never register an + # OptimizedModule: historical checkpoint parameter names stay intact. + for block in self.blocks: + block.configure_compilation( + mode=compile_mode, dynamic=compile_dynamic, backend=compile_backend + ) + compiled_names = ("_embed_image", "_decode_image") + if noise_conditioning == "adaln_zero": + compiled_names += ("_embed_noise_frames",) + for name in compiled_names: + object.__setattr__( + self, + name + "_execute", + compile_region( + getattr(self, name), + mode=compile_mode, + dynamic=compile_dynamic, + backend=compile_backend, + ), + ) + + @staticmethod + def _dimensions(dims, name): + if ( + not isinstance(dims, Mapping) + or not dims + or any( + not isinstance(k, str) or not k or type(d) is not int or d < 1 + for k, d in dims.items() + ) + ): + raise ValueError( + f"{name} requires positive integer dimensions by embodiment" + ) + return dict(dims) + + def validate(self, batch, embodiment): + if self.has_proprio_branch and embodiment not in self.pose_dims: + raise ValueError(f"Unknown embodiment {embodiment!r}") + if not self.has_proprio_branch and set(batch) != {"image"}: + raise ValueError("Image-only visual_only accepts exactly the image stream") + if "track" in batch: + raise ValueError("Image-patch/proprio models do not accept a track stream") + dimensions = {"image": self.image_dim} + if self.has_proprio_branch: + dimensions["proprio"] = self.pose_dims[embodiment] + if self.objective in ("direct_action", "visual_state_action"): + dimensions["action"] = self.action_dims[embodiment] + total, batch_size = 0, None + checks = [] + for name, dim in dimensions.items(): + stream = batch[name] + value = stream["value"] + if ( + value.ndim != 3 + or value.shape[-1] != dim + or not value.is_floating_point() + ): + raise ValueError(f"{name}: invalid value shape or dtype") + if batch_size is not None and batch_size != value.shape[0]: + raise ValueError("Mixed batch sizes") + batch_size = value.shape[0] + total += value.shape[1] + for key in ("time", "valid", "observed"): + if stream[key].shape != value.shape[:2]: + raise ValueError(f"{name}: {key} shape mismatch") + valid, observed = stream["valid"], stream["observed"] + if valid.dtype != torch.bool or observed.dtype != torch.bool: + raise ValueError("Masks must be boolean") + checks.append(((observed & ~valid).any(), "Padding cannot be observed")) + if "target_valid" in stream and ( + stream["target_valid"].shape != valid.shape + or stream["target_valid"].dtype != torch.bool + ): + raise ValueError("Target validity must be a boolean token mask") + checks.append( + ( + ~torch.isfinite(stream["time"].masked_fill(~valid, 0)).all(), + "Non-finite timestamp", + ) + ) + checks.append( + ( + ~torch.isfinite(torch.where(observed[..., None], value, 0)).all(), + "Non-finite observed value", + ) + ) + if name == "action": + checks.append( + (observed.any(), "Action labels must never be marked observed") + ) + if self.time_mode == "discrete": + if "index" not in stream or stream["index"].shape != valid.shape: + raise ValueError("Discrete mode requires explicit grid indices") + indices = stream["index"].masked_fill(~valid, 0) + checks.append( + ( + (~torch.isfinite(indices) | (indices != indices.round())).any(), + "Discrete indices must be finite integers", + ) + ) + if name == "image": + if stream["xy"].shape != (*valid.shape, 2): + raise ValueError("Reference image coordinates must match tokens") + checks.append( + ( + ~torch.isfinite( + torch.where(valid[..., None], stream["xy"], 0) + ).all(), + "Non-finite reference image coordinates", + ) + ) + if self.use_camera_embedding: + ids = stream["camera"] + if ids.shape != valid.shape or ids.dtype != torch.long: + raise ValueError("Camera IDs must be integer token IDs") + checks.append( + ( + ( + ((ids < 0) | (ids >= self.camera.num_embeddings)) + & valid + ).any(), + "Camera ID out of range", + ) + ) + if total > self.max_tokens: + raise ValueError(f"Token budget exceeded: {total} > {self.max_tokens}") + # A query without any clean observation is not a conditional BC example. + conditioning = MODALITIES if self.phase_a_condition_on_proprio else ("image",) + context = torch.cat([batch[n]["observed"] for n in conditioning], 1) + checks.append( + ( + ~context.any(1).all(), + "Every sample requires an observed image or proprio token", + ) + ) + # One device-to-host synchronization for all value checks. Boolean + # indexing/nonzero used here previously introduced additional syncs. + for failed, (_, message) in zip( + torch.stack([item[0] for item in checks]).tolist(), checks + ): + if failed: + raise ValueError(message) + + def _clock(self, stream, valid): + key = "time" if self.time_mode == "continuous" else "index" + return stream[key].masked_fill(~valid, 0) + + def forward( + self, + batch, + embodiment, + image_values=None, + levels=None, + phase="A", + *, + _validated=False, + _execution=None, + ): + """Forward one phase. Withheld labels never enter an embedding. + + Phase A excludes withheld proprio queries from joint attention entirely. + Phase B adds those queries, whose values are learned constants. In the + direct-action variant only observed context and action queries are used. + """ + if not _validated: + self.validate(batch, embodiment) + if self.objective != "direct_action": + image = batch["image"] + if image_values is None or levels is None: + raise ValueError( + "Visual phases require image_values and noise levels" + ) + if ( + image_values.shape != image["value"].shape + or levels.shape != image["valid"].shape + ): + raise ValueError("Visual values/levels must match image tokens") + bad_levels = ( + ~torch.isfinite(levels) | (levels < 0) | (levels > 1) + ) & image["valid"] + if bad_levels.any(): + raise ValueError("Noise levels must be finite in [0,1]") + return self._forward_validated( + batch, embodiment, image_values, levels, phase, execution=_execution + ) + + def _prepare_rope(self, times, xy=None): + if self.position_encoding == "spatiotemporal_rope": + return prepare_spatiotemporal_rope( + times, xy, self.blocks[0].head_dim, + axis_dims=self.rope_axis_dims, spatial_grid_size=self.spatial_grid_size, + ) + return prepare_temporal_rope(times, self.blocks[0].head_dim) + + def _prepare_execution(self, times, valid, image_xy=None): + joint_valid = torch.cat(valid, 1) + layout = None + fused_backends = {block.attention_backend for block in self.blocks} & { + "flash3", + "flash4", + } + if fused_backends: + # Respect explicit block execution overrides used by the offline + # weights-preservation proof; model config identity stays recorded. + for backend in fused_backends: + load_flash_attention_kernels(backend, joint_valid.device) + layout = prepare_attention_layout(joint_valid, query_valid=joint_valid) + if self.position_encoding == "spatiotemporal_rope" and image_xy is None: + raise ValueError("Image spatial-temporal attention requires patch coordinates") + safe_xy = None if image_xy is None else torch.where(valid[0][..., None], image_xy, 0) + rope = tuple(self._prepare_rope(t, safe_xy if index == 0 else None) + for index, t in enumerate(times)) + return layout, rope + + @staticmethod + def _prepare_noise_frame_map(image): + """Map same-example timestamp/observed-role tokens to one noise condition. + + Camera and patch identities are deliberately excluded. Grouping is eager, + reused by all blocks and cached with attention metadata during sampling. + Noise values are checked on each call and never cached across sampler steps. + """ + valid = image["valid"] + groups = torch.zeros_like(valid, dtype=torch.long) + representatives, offset = [], 0 + for batch_index in range(valid.shape[0]): + positions = valid[batch_index].nonzero(as_tuple=False).flatten() + if positions.numel() == 0: + continue + keys = torch.stack((image["time"][batch_index, positions].double(), + image["observed"][batch_index, positions].double()), dim=1) + unique, inverse = torch.unique(keys, dim=0, return_inverse=True) + first = torch.full((unique.shape[0],), positions.numel(), device=positions.device, dtype=torch.long) + first.scatter_reduce_(0, inverse, torch.arange(positions.numel(), device=positions.device), reduce="amin") + representatives.append(batch_index * valid.shape[1] + positions[first]) + groups[batch_index, positions] = inverse + offset + offset += unique.shape[0] + indices = (torch.cat(representatives) if representatives else + torch.zeros(1, device=valid.device, dtype=torch.long)) + return _NoiseFrameMap(indices, groups) + + def _embed_noise_frames(self, levels): + # Embedding convention only: tau remains the original linear-corruption + # strength in [0,1]. Multiplying by1000 does not discretize/change it. + frequency = torch.exp(-math.log(10000) * torch.arange( + 128, device=levels.device, dtype=torch.float32 + ) / 128) + angles = (1000 * levels.float())[:, None] * frequency[None] + return self.noise_level(torch.cat((angles.cos(), angles.sin()), dim=-1)) + + def _image_noise_condition(self, image, levels, frame_map): + effective = levels.masked_fill(image["observed"] | ~image["valid"], 0) + grouped = effective.reshape(-1).index_select(0, frame_map.representatives) + reconstructed = grouped.index_select(0, frame_map.token_groups.reshape(-1)).reshape_as(effective) + if ((effective != reconstructed) & image["valid"]).any(): + raise ValueError("Image noise levels must agree within each timestamp and observed/query group") + return self._embed_noise_frames_execute(grouped), frame_map.token_groups + + def _prepare_image_execution(self, image, valid): + execution = self._prepare_image_attention_execution(image, valid) + if self.noise_conditioning == "adaln_zero": + return _ConditionedImageExecution(execution, self._prepare_noise_frame_map(image)) + return execution + + def _prepare_image_attention_execution(self, image, valid): + """Prepare reusable image-only metadata before any learned projection.""" + times = self._clock(image, valid) + if self.token_execution == "padded": + return self._prepare_execution([times], [valid], image["xy"]) + layout = prepare_attention_layout(valid, query_valid=valid) + backends = {block.attention_backend for block in self.blocks} + for backend in backends & {"flash3", "flash4"}: + load_flash_attention_kernels(backend, valid.device) + if layout.kind == "dense": + # Equal, fully valid batches already have no padding overhead and + # retain the existing dense kernel and tensor-region shapes. + return ( + layout if backends & {"flash3", "flash4"} else None, + (self._prepare_rope(times, image["xy"]),), + ) + if not backends <= {"flash3", "flash4"}: + raise ValueError("Ragged packed_valid execution requires flash3 or flash4") + if any(block.attention_query_chunk_size for block in self.blocks): + raise ValueError("Packed token execution requires unchunked attention") + indices = layout.key_indices + metadata_keys = ("xy", "observed", "camera") if self.use_camera_embedding else ("xy", "observed") + packed_image = { + key: image[key].flatten(0, 1).index_select(0, indices).unsqueeze(0) + for key in metadata_keys + } + packed_times = times.reshape(-1).index_select(0, indices).unsqueeze(0) + packed_valid = valid.new_ones((1, indices.numel())) + return _PackedImageExecution( + layout, packed_image, packed_valid, + self._prepare_rope(packed_times, packed_image["xy"]), + ) + + def _forward_packed_image(self, values, levels, embodiment, execution, image_condition=None): + indices = execution.layout.key_indices + packed_values = values.flatten(0, 1).index_select(0, indices).unsqueeze(0) + packed_levels = levels.reshape(-1).index_select(0, indices).unsqueeze(0) + if image_condition is not None: + condition, frame_groups = image_condition + image_condition = (condition, frame_groups.reshape(-1).index_select(0, indices).unsqueeze(0)) + stream = self._embed_image_execute( + execution.image, packed_values, packed_levels, execution.valid, "A", False + ) + if self.use_modality_embedding: + stream = stream + self.modality[0] + if self.use_embodiment_embedding: + stream = stream + self.embodiment.weight[self.embodiment_ids[embodiment]] + for block in self.blocks: + if self.activation_checkpointing and torch.is_grad_enabled(): + stream = checkpoint( + block.forward_packed, stream, execution.valid, + attention_layout=execution.layout, rope=execution.rope, + image_condition=image_condition, + use_reentrant=False, + ) + else: + stream = block.forward_packed( + stream, execution.valid, + attention_layout=execution.layout, rope=execution.rope, + image_condition=image_condition, + ) + predicted = self._decode_image_execute(stream, execution.valid, image_condition).squeeze(0) + batch, tokens = execution.layout.batch_size, execution.layout.key_tokens + return predicted.new_zeros(batch * tokens, self.image_dim).index_copy( + 0, indices, predicted + ).reshape(batch, tokens, self.image_dim) + + def _embed_image(self, image, values, levels, image_valid, phase, direct): + image_x = self.image_in(torch.where(image_valid[..., None], values, 0)) + if self.position_encoding == "linear_xy_temporal_rope": + image_x = image_x + self.xy(torch.where(image_valid[..., None], image["xy"], 0)) + if self.use_camera_embedding: + image_x = image_x + self.camera( + image["camera"].masked_fill(~image_valid, 0) + ) + if not direct and self.noise_conditioning == "additive": + effective_levels = levels.masked_fill(image["observed"] | ~image_valid, 0) + image_x = image_x + self.noise_level( + effective_levels[..., None].to(image_x.dtype) + ) + if not direct and self.use_phase_embedding: + image_x = image_x + self.phase.weight[int(phase == "B")] + return image_x + + def _decode_image(self, stream, valid, image_condition=None): + normalized = self.norm[0](stream) + if self.noise_conditioning == "adaln_zero": + condition, frame_groups = image_condition + shift, scale = self.final_adaln(condition).chunk(2, dim=-1) + normalized = normalized * (1 + MoTBlock._frame_values(scale, frame_groups)) + MoTBlock._frame_values(shift, frame_groups) + return self.image_out(normalized).masked_fill(~valid[..., None], 0) + + def _forward_validated( + self, + batch, + embodiment, + image_values=None, + levels=None, + phase="A", + execution=None, + ): + if phase not in ("A", "B"): + raise ValueError("phase must be A or B") + if self.objective == "visual_only" and phase != "A": + raise ValueError("visual_only has no proprio readout phase") + direct = self.objective == "direct_action" + combined = self.objective == "visual_state_action" + include_pose = self.has_proprio_branch and ( + direct or phase == "B" or self.phase_a_condition_on_proprio + ) + image = batch["image"] + image_valid = image["valid"] & (image["observed"] if direct else True) + if direct: + if image_values is not None or levels is not None: + raise ValueError("Direct action BC takes only clean observed features") + values = image["value"] + else: + values = torch.where( + image["observed"][..., None], image["value"], image_values + ) + frame_map = None + if self.noise_conditioning == "adaln_zero" and phase == "A" and not include_pose and execution is None: + execution = self._prepare_image_execution(image, image_valid) + if isinstance(execution, _ConditionedImageExecution): + frame_map = execution.noise_frames + execution = execution.attention + image_condition = None + if self.noise_conditioning == "adaln_zero": + frame_map = frame_map or self._prepare_noise_frame_map(image) + image_condition = self._image_noise_condition(image, levels, frame_map) + # Phase B retains the existing multi-expert route. Packing here changes + # execution only; no token, camera, or per-example boundary is removed. + if self.token_execution == "packed_valid" and phase == "A" and not include_pose: + execution = (self._prepare_image_execution(image, image_valid) + if execution is None else execution) + if isinstance(execution, _PackedImageExecution): + return {"image": self._forward_packed_image( + values, levels, embodiment, execution, image_condition + )} + image_x = self._embed_image_execute( + image, values, levels, image_valid, phase, direct + ) + if include_pose: + pose = batch["proprio"] + pose_valid = pose["valid"] & ( + pose["observed"] if direct or phase == "A" else True + ) + observed_pose = pose["observed"] & pose_valid + pose_value = torch.where(observed_pose[..., None], pose["value"], 0) + pose_x = self.pose_in[embodiment](pose_value) + if self.objective == "visual_only": + pose_x = pose_x + self.proprio_role.weight[0] + else: + pose_x = torch.where(observed_pose[..., None], pose_x, self.query) + pose_x = pose_x + self.proprio_role((~observed_pose).long()) + pose_time = self._clock(pose, pose_valid) + action_start = pose_x.shape[1] + if direct or (combined and phase == "B"): + action = batch["action"] + # The shape and query metadata are needed, but action values are not. + query = (self.action_query if combined else self.query).expand( + *action["valid"].shape, -1 + ) + query = query + self.proprio_role.weight[2 if combined else 1] + pose_x = torch.cat((pose_x, query), 1) + pose_valid = torch.cat((pose_valid, action["valid"]), 1) + pose_time = torch.cat((pose_time, self._clock(action, action["valid"])), 1) + if not direct and self.use_phase_embedding: + phase_embedding = self.phase.weight[int(phase == "B")] + if include_pose: + pose_x = pose_x + phase_embedding + valid = [image_valid] + times = [self._clock(image, image_valid)] + streams = [image_x] + if include_pose: + valid.append(pose_valid) + times.append(pose_time) + streams.append(pose_x) + if self.use_modality_embedding: + streams = [x + self.modality[i] for i, x in enumerate(streams)] + if self.use_embodiment_embedding: + embodiment_x = self.embodiment.weight[self.embodiment_ids[embodiment]] + streams = [x + embodiment_x for x in streams] + streams = [ + x.masked_fill(~valid[i][..., None], 0) for i, x in enumerate(streams) + ] + attention_layout, rope = ( + execution + if execution is not None + else self._prepare_execution(times, valid, image["xy"]) + ) + for block in self.blocks: + if self.activation_checkpointing and torch.is_grad_enabled(): + streams = checkpoint( + block, + streams, + times, + valid, + attention_layout=attention_layout, + rope=rope, + image_condition=image_condition, + use_reentrant=False, + ) + else: + streams = block( + streams, times, valid, attention_layout=attention_layout, rope=rope, + image_condition=image_condition, + ) + image_result = ( + None if direct else self._decode_image_execute(streams[0], image_valid, image_condition) + ) + if include_pose: + streams[1] = self.norm[1](streams[1]) + if direct: + action = self.action_out[embodiment](streams[1][:, action_start:]) + return { + "action": action.masked_fill(~batch["action"]["valid"][..., None], 0) + } + result = {"image": image_result} + if self.objective == "visual_only" or not include_pose: + return result + result.update( + { + "proprio": self.pose_out[embodiment]( + streams[1][:, :action_start] + ).masked_fill(~pose_valid[:, :action_start, None], 0), + } + ) + if combined and phase == "B": + result["action"] = self.action_out[embodiment]( + streams[1][:, action_start:] + ).masked_fill(~batch["action"]["valid"][..., None], 0) + return result + + @staticmethod + def frame_noise_levels(stream, maximum=1.0): + """One scalar per physical timestamp; observed tokens stay at zero.""" + return LinearGaussianImageNoising.frame_noise_levels(stream, maximum) + + def _noisy_images(self, image, clean, maximum=1.0): + corruption = LinearGaussianImageNoising(maximum)(image, clean=clean) + return corruption["image_values"], corruption["noise_levels"] + + def _loss(self, prediction, target, mask, *, check_target=True): + return PlainPixelMSE(self.loss_reduction)( + prediction, target, mask, check_target=check_target + ) + + def training_losses( + self, batch, embodiment, pose_weight=1.0, image_weight=1.0, action_weight=1.0 + ): + self.validate(batch, embodiment) + for name, weight in ( + ("pose_weight", pose_weight), + ("image_weight", image_weight), + ("action_weight", action_weight), + ): + if not math.isfinite(weight) or weight < 0: + raise ValueError(f"{name} must be finite and nonnegative") + if self.objective == "direct_action": + prediction = self(batch, embodiment, _validated=True)["action"] + stream = batch["action"] + mask = stream["valid"] & stream.get("target_valid", stream["valid"]) + loss = self._loss(prediction, stream["value"], mask) + return {"action": loss, "total": loss} + image = batch["image"] + clean, image_mask = LinearGaussianImageNoising.prepare_target(image) + noisy, levels = self._noisy_images(image, clean) + prediction = self(batch, embodiment, noisy, levels, phase="A", _validated=True) + losses = { + "image": self._loss( + prediction["image"], + clean, + image_mask, + check_target=False, + ) + } + if self.objective == "visual_only": + losses["total"] = image_weight * losses["image"] + return losses + pose = batch["proprio"] + pose_mask = ( + pose["valid"] & ~pose["observed"] & pose.get("target_valid", pose["valid"]) + ) + combined = self.objective == "visual_state_action" + action_mask = None + if combined: + action = batch["action"] + action_mask = action["valid"] & action.get("target_valid", action["valid"]) + if (pose_weight > 0 and pose_mask.any()) or ( + combined and action_weight > 0 and action_mask.any() + ): + if self.pose_conditioning == "low_noise": + # A separate low-noise prediction aligns the readout with this + # trajectory. Do not train pose readout on arbitrary high-noise + # x0 predictions. No detach: state loss reaches the image head. + low_noisy, low_levels = self._noisy_images( + image, clean, self.pose_noise_max + ) + estimated = self( + batch, embodiment, low_noisy, low_levels, phase="A", _validated=True + )["image"] + pose_images = torch.where( + image["observed"][..., None], clean, estimated + ) + else: + pose_images = clean + readout = self( + batch, + embodiment, + pose_images, + torch.zeros_like(levels), + phase="B", + _validated=True, + ) + losses["pose"] = ( + self._loss(readout["proprio"], pose["value"], pose_mask) + if pose_weight > 0 + else readout["proprio"].sum() * 0 + ) + if combined: + losses["action"] = ( + self._loss(readout["action"], action["value"], action_mask) + if action_weight > 0 + else readout["action"].sum() * 0 + ) + else: + losses["pose"] = prediction.get("proprio", prediction["image"]).sum() * 0 + if combined: + losses["action"] = losses["image"] * 0 + losses["total"] = image_weight * losses["image"] + pose_weight * losses["pose"] + if combined: + losses["total"] = losses["total"] + action_weight * losses["action"] + return losses + + @torch.no_grad() + def sample(self, batch, embodiment, steps=20, generator=None): + """Sample images from noise, then read states; or one direct BC pass. + + For linear corruption, epsilon=(z_s-(1-s)*x0)/s. Implement the DDIM + update using the bounded next_s/s ratio, never division at s=0. All + recorded observations are clamped exactly throughout sampling. + """ + self.validate(batch, embodiment) + if self.objective == "direct_action": + return self(batch, embodiment, _validated=True) + current = LinearDDIMImageSampler(steps)( + self, batch, embodiment, generator=generator, validated=True + ) + if self.objective == "visual_only": + return {"image": current} + readout = self( + batch, + embodiment, + current, + torch.zeros_like(batch["image"]["time"], dtype=torch.float32), + phase="B", + _validated=True, + ) + pose = readout["proprio"] + pose = torch.where( + batch["proprio"]["observed"][..., None], batch["proprio"]["value"], pose + ) + result = { + "image": current, + "proprio": pose.masked_fill(~batch["proprio"]["valid"][..., None], 0), + } + if self.objective == "visual_state_action": + result["action"] = readout["action"] + return result diff --git a/egomimic/models/pixel_diffusion.py b/egomimic/models/pixel_diffusion.py new file mode 100644 index 000000000..b8f162d22 --- /dev/null +++ b/egomimic/models/pixel_diffusion.py @@ -0,0 +1,198 @@ +"""Replaceable pixel noising, objective, and image sampling components. + +These objects own no trainable parameters. The graph registers the denoiser +once and passes it to the sampler; archived model convenience methods delegate +to the same implementations. Diffusion strength and physical time remain +separate inputs, including at the loss boundary for future objective choices. +""" + +import math + +import torch + +from egomimic.models.world_model import masked_mse + + +class LinearGaussianImageNoising: + """Linear corruption with one uniform strength per queried timestamp. + + All cameras and patches at a physical timestamp share a strength. Gaussian + noise remains independent per channel. Draw order is timestamp strengths for + each example, then one full-shaped Gaussian draw, including observed/padded + positions whose noise is subsequently discarded. + """ + + def __init__(self, maximum=1.0): + self._check_maximum(maximum) + self.maximum = maximum + + @staticmethod + def _check_maximum(maximum): + if not math.isfinite(maximum) or not 0 <= maximum <= 1: + raise ValueError("maximum noise strength must be finite in [0,1]") + + @staticmethod + def prepare_target(image): + """Return recorded clean targets and the queried supervision mask.""" + known = image["valid"] & image.get("target_valid", image["valid"]) + if (image["valid"] & ~known).any(): + raise ValueError( + "Visual training requires recorded features for every valid image token" + ) + clean = torch.where(known[..., None], image["value"], 0) + if not torch.isfinite(clean).all(): + raise ValueError("Non-finite image target") + return clean, known & ~image["observed"] + + @staticmethod + def frame_noise_levels(stream, maximum=1.0): + LinearGaussianImageNoising._check_maximum(maximum) + levels = torch.zeros_like(stream["time"], dtype=torch.float32) + for b in range(levels.shape[0]): + selected = stream["valid"][b] & ~stream["observed"][b] + if selected.any(): + times, inverse = torch.unique( + stream["time"][b][selected], return_inverse=True + ) + draws = torch.rand(times.numel(), device=levels.device) * maximum + levels[b, selected] = draws[inverse] + return levels + + def __call__(self, image, *, clean=None, maximum=None): + maximum = self.maximum if maximum is None else maximum + self._check_maximum(maximum) + if clean is None: + clean, mask = self.prepare_target(image) + else: + # Internal legacy low-noise passes provide an already checked target. + if clean.shape != image["value"].shape: + raise ValueError("Clean image target must match image token shape") + mask = ( + image["valid"] + & image.get("target_valid", image["valid"]) + & ~image["observed"] + ) + levels = self.frame_noise_levels(image, maximum) + strength = levels[..., None].to(clean.dtype) + noisy = (1 - strength) * clean + strength * torch.randn_like(clean) + noisy = noisy.masked_fill(~image["valid"][..., None], 0) + return { + "image_values": noisy, + "noise_levels": levels, + "target": clean, + "mask": mask, + "metadata": { + "timestamps": image["time"], + "valid": image["valid"], + "observed": image["observed"], + }, + } + + +class PlainPixelMSE: + """FP32 plain x0 MSE; noise/time inputs are explicit and unweighted. + + The configurable loss boundary deliberately accepts noise levels and metadata + without using them. A later objective can consume these inputs without changing + the noising stage or denoiser. No SNR weighting is applied here. + """ + + def __init__(self, reduction="mean_per_example"): + if reduction not in ("token_mean", "mean_per_example"): + raise ValueError("Unknown masked MSE reduction") + self.reduction = reduction + + def __call__( + self, + prediction, + target, + mask, + noise_levels=None, + metadata=None, + *, + check_target=True, + ): + return masked_mse( + prediction, + target, + mask.float(), + reduction=self.reduction, + check_target=check_target, + ) + + +class LinearDDIMImageSampler: + """Integrate x0 predictions, then return images without state readout. + + The denoiser is the graph's existing registered model, not a copied submodule. + Its image-only attention layout and FP32 RoPE tables are reused across steps. + With standard BF16 mixed precision, recorded patches/state are FP32 and every + prediction is promoted to that state dtype before integration arithmetic. + """ + + def __init__(self, steps=20): + self._check_steps(steps) + self.steps = steps + + @staticmethod + def _check_steps(steps): + if type(steps) is not int or steps < 1: + raise ValueError("steps must be a positive integer") + + @torch.no_grad() + def __call__( + self, + denoiser, + batch, + embodiment, + *, + steps=None, + generator=None, + validated=False, + ): + steps = self.steps if steps is None else steps + if not validated: + denoiser.validate(batch, embodiment) + self._check_steps(steps) + image = batch["image"] + valid, observed = image["valid"], image["observed"] + anchor = torch.where(observed[..., None], image["value"], 0) + noise = torch.randn( + image["value"].shape, + device=anchor.device, + dtype=anchor.dtype, + generator=generator, + ) + current = torch.where(observed[..., None], anchor, noise).masked_fill( + ~valid[..., None], 0 + ) + execution = None + if not denoiser.has_proprio_branch or ( + ( + denoiser.token_execution == "packed_valid" + or denoiser.noise_conditioning == "adaln_zero" + ) + and not denoiser.phase_a_condition_on_proprio + ): + execution = denoiser._prepare_image_execution(image, valid) + for j in range(steps, 0, -1): + level, next_level = j / steps, (j - 1) / steps + levels = torch.full_like( + image["time"], level, dtype=torch.float32 + ).masked_fill(observed | ~valid, 0) + x0 = denoiser( + batch, + embodiment, + current, + levels, + phase="A", + _validated=True, + _execution=execution, + )["image"].to(current.dtype) + candidate = (1 - next_level) * x0 + (next_level / level) * ( + current - (1 - level) * x0 + ) + current = torch.where(observed[..., None], anchor, candidate).masked_fill( + ~valid[..., None], 0 + ) + return current diff --git a/egomimic/models/world_model.py b/egomimic/models/world_model.py index b9def84db..58679be61 100644 --- a/egomimic/models/world_model.py +++ b/egomimic/models/world_model.py @@ -16,15 +16,32 @@ from torch.nn import functional as F from torch.utils.checkpoint import checkpoint +from egomimic.models.attention_execution import ( + apply_temporal_rope, + execute_flash_attention, + execute_packed_flash_attention, + load_flash_attention_kernels, + prepare_attention_layout, +) +from egomimic.models.compilation import compile_region + MODALITIES = ("image", "track", "proprio") -def masked_mse(pred, target, weight): +def masked_mse(pred, target, weight, *, reduction="token_mean", check_target=True): + if reduction not in ("token_mean", "mean_per_example"): + raise ValueError("Unknown masked MSE reduction") active = weight > 0 - if not torch.isfinite(target[active]).all(): + if check_target and not torch.isfinite(torch.where(active[..., None], target, 0)).all(): raise ValueError("Non-finite supervised target") safe_target = torch.where(active[..., None], target, torch.zeros_like(target)) error = (pred.float() - safe_target.float()).square().mean(-1) + if reduction == "mean_per_example": + # Match the mean of independent batch-one losses, including empty + # examples' differentiable zero. Padding and clip length add no weight. + numerator = (error * weight.float()).flatten(1).sum(1) + denominator = weight.float().flatten(1).sum(1) + return (numerator / torch.where(denominator > 0, denominator, 1.0)).mean() total_weight = weight.float().sum() # Confidence weights can sum to less than one. Preserve a weighted mean; # use one only for an empty objective so its value and gradients stay zero. @@ -55,51 +72,113 @@ def temporal_rope(x, times): class MoTBlock(nn.Module): """Private modality projections/FFNs, one joint attention operation.""" - def __init__(self, width, heads, mlp_ratio=4, attention_query_chunk_size=0, - attention_backend="sdpa", attention_dtype="native"): + def __init__( + self, + width, + heads, + mlp_ratio=4, + attention_query_chunk_size=0, + modalities=MODALITIES, + attention_backend="sdpa", + attention_dtype="native", + image_adaln=False, + ): super().__init__() if width % heads or (width // heads) % 2: raise ValueError("RoPE requires an even attention head width") self.heads = heads + self.head_dim = width // heads if type(attention_query_chunk_size) is not int or attention_query_chunk_size < 0: raise ValueError("attention_query_chunk_size must be a nonnegative integer") self.attention_query_chunk_size = attention_query_chunk_size - if attention_backend not in ("sdpa", "efficient", "flash4"): - raise ValueError("attention_backend must be sdpa, efficient or flash4") + if attention_backend not in ("sdpa", "efficient", "flash3", "flash4"): + raise ValueError("attention_backend must be sdpa, efficient, flash3 or flash4") if attention_dtype not in ("native", "bfloat16"): raise ValueError("attention_dtype must be native or bfloat16") - if attention_backend == "flash4" and attention_dtype != "bfloat16": - raise ValueError("flash4 requires explicit bfloat16 attention") + if attention_backend in ("flash3", "flash4") and attention_dtype != "bfloat16": + raise ValueError(f"{attention_backend} requires explicit bfloat16 attention") self.attention_backend = attention_backend self.attention_dtype = attention_dtype - self.norm1 = nn.ModuleList(nn.LayerNorm(width) for _ in MODALITIES) - self.qkv = nn.ModuleList(nn.Linear(width, 3 * width) for _ in MODALITIES) - self.out = nn.ModuleList(nn.Linear(width, width) for _ in MODALITIES) - self.norm2 = nn.ModuleList(nn.LayerNorm(width) for _ in MODALITIES) + if type(image_adaln) is not bool: + raise ValueError("image_adaln must be bool") + self.image_adaln = image_adaln + self.norm1 = nn.ModuleList( + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if image_adaln and i == 0 else nn.LayerNorm(width) + for i, _ in enumerate(modalities) + ) + self.qkv = nn.ModuleList(nn.Linear(width, 3 * width) for _ in modalities) + self.out = nn.ModuleList(nn.Linear(width, width) for _ in modalities) + self.norm2 = nn.ModuleList( + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if image_adaln and i == 0 else nn.LayerNorm(width) + for i, _ in enumerate(modalities) + ) self.ffn = nn.ModuleList( nn.Sequential( nn.Linear(width, width * mlp_ratio), nn.GELU(), nn.Linear(width * mlp_ratio, width), ) - for _ in MODALITIES + for _ in modalities ) + if image_adaln: + self.adaln_image = nn.Sequential(nn.SiLU(), nn.Linear(width, 6 * width)) + nn.init.zeros_(self.adaln_image[-1].weight) + nn.init.zeros_(self.adaln_image[-1].bias) + self.configure_compilation() + + def configure_compilation(self, mode="none", dynamic=False, backend="inductor"): + names = ("_project_inputs", "_finish_outputs") + if self.image_adaln: + names += ("_image_modulation",) + for name in names: + object.__setattr__(self, name + "_executor", compile_region( + getattr(self, name), mode=mode, dynamic=dynamic, backend=backend + )) + + def _image_modulation(self, condition): + # Project once per physical-frame/role group, not once per patch. + return self.adaln_image(condition).chunk(6, dim=-1) - def forward(self, streams, times, valid): - sizes = [s.shape[1] for s in streams] + @staticmethod + def _frame_values(values, frame_indices): + return values.index_select(0, frame_indices.reshape(-1)).reshape( + *frame_indices.shape, values.shape[-1] + ) + + def _project_inputs(self, streams, times, rope, image_modulation=None, frame_indices=None): qs, ks, vs = [], [], [] for i, x in enumerate(streams): b, n, d = x.shape + normalized = self.norm1[i](x) + if i == 0 and self.image_adaln: + shift, scale = (self._frame_values(value, frame_indices) for value in image_modulation[:2]) + normalized = normalized * (1 + scale) + shift q, k, v = ( - self.qkv[i](self.norm1[i](x)) + self.qkv[i](normalized) .reshape(b, n, 3, self.heads, d // self.heads) .permute(2, 0, 3, 1, 4) ) - qs.append(temporal_rope(q, times[i])) - ks.append(temporal_rope(k, times[i])) + qs.append(temporal_rope(q, times[i]) if rope is None else apply_temporal_rope(q, rope[i])) + ks.append(temporal_rope(k, times[i]) if rope is None else apply_temporal_rope(k, rope[i])) vs.append(v) - mask = torch.cat(valid, 1)[:, None, None, :] - query, key, value = torch.cat(qs, 2), torch.cat(ks, 2), torch.cat(vs, 2) + return torch.cat(qs, 2), torch.cat(ks, 2), torch.cat(vs, 2) + + def forward(self, streams, times, valid, *, attention_layout=None, rope=None, image_condition=None): + sizes = [s.shape[1] for s in streams] + modulation = None + frame_indices = None + if self.image_adaln: + if image_condition is None: + raise ValueError("Image AdaLN requires frame conditioning") + condition, frame_indices = image_condition + modulation = self._image_modulation_executor(condition) + query, key, value = self._project_inputs_executor(streams, times, rope, modulation, frame_indices) + else: + query, key, value = self._project_inputs_executor(streams, times, rope) + mask = (torch.cat(valid, 1)[:, None, None, :] if attention_layout is None + else attention_layout.key_mask) chunk = self.attention_query_chunk_size if chunk and query.shape[2] > chunk: # Split only queries. Every row retains all keys/values and the same @@ -115,16 +194,75 @@ def forward(self, streams, times, valid): parts.append(part) result = torch.cat(parts, dim=2) else: - result = self._apply_attention(query, key, value, mask) + # Keep the historical four-argument call for native SDPA overrides. + result = (self._apply_attention(query, key, value, mask) + if attention_layout is None or self.attention_backend not in ("flash3", "flash4") + else self._apply_attention( + query, key, value, mask, attention_layout=attention_layout)) + if self.image_adaln: + return self._finish_outputs_executor(streams, result, valid, sizes, modulation, frame_indices) + return self._finish_outputs_executor(streams, result, valid, sizes) + + def _finish_outputs(self, streams, result, valid, sizes, image_modulation=None, frame_indices=None): result = result.transpose(1, 2).flatten(2).split(sizes, 1) outputs = [] for i, (x, y) in enumerate(zip(streams, result)): - x = x + self.out[i](y) - x = x + self.ffn[i](self.norm2[i](x)) + if i == 0 and self.image_adaln: + gate_attention, shift, scale, gate_mlp = ( + self._frame_values(value, frame_indices) for value in image_modulation[2:] + ) + x = x + gate_attention * self.out[i](y) + x = x + gate_mlp * self.ffn[i](self.norm2[i](x) * (1 + scale) + shift) + else: + x = x + self.out[i](y) + x = x + self.ffn[i](self.norm2[i](x)) outputs.append(x.masked_fill(~valid[i][..., None], 0)) return outputs - def _apply_attention(self, query, key, value, mask): + def forward_packed(self, stream, valid, *, attention_layout, rope, image_condition=None): + """Run expert zero on packed image tokens, preserving example boundaries. + + Pure tensor regions reuse the same modules and compilation policy as + padded execution. Only optional attention dispatch remains outside them. + """ + if self.attention_backend not in ("flash3", "flash4"): + raise ValueError("Packed token execution requires flash3 or flash4") + if self.attention_dtype != "bfloat16": + raise ValueError("Packed token execution requires explicit bfloat16 attention") + if self.attention_query_chunk_size: + raise ValueError("Packed token execution requires unchunked attention") + if stream.ndim != 3 or stream.shape[0] != 1 or valid.shape != stream.shape[:2]: + raise ValueError("Packed expert input must be [1,total_tokens,width]") + modulation = None + frame_indices = None + if self.image_adaln: + if image_condition is None: + raise ValueError("Image AdaLN requires frame conditioning") + condition, frame_indices = image_condition + modulation = self._image_modulation_executor(condition) + query, key, value = self._project_inputs_executor([stream], None, (rope,), modulation, frame_indices) + else: + query, key, value = self._project_inputs_executor([stream], None, (rope,)) + original_dtype = query.dtype + kernels = load_flash_attention_kernels(self.attention_backend, query.device) + with torch.autocast(device_type=query.device.type, enabled=False): + query, key, value = ( + x.squeeze(0).transpose(0, 1).to(torch.bfloat16) + for x in (query, key, value) + ) + result = execute_packed_flash_attention( + query, key, value, attention_layout, kernels + ).to(original_dtype) + result = result.transpose(0, 1).unsqueeze(0) + if self.image_adaln: + return self._finish_outputs_executor( + [stream], result, [valid], [stream.shape[1]], modulation, frame_indices + )[0] + return self._finish_outputs_executor( + [stream], result, [valid], [stream.shape[1]] + )[0] + + def _apply_attention(self, query, key, value, mask, *, attention_layout=None): # Preserve historical default arithmetic and autocast behavior exactly. if self.attention_backend == "sdpa" and self.attention_dtype == "native": return self._attend(query, key, value, mask) @@ -132,8 +270,12 @@ def _apply_attention(self, query, key, value, mask): with torch.autocast(device_type=query.device.type, enabled=False): if self.attention_dtype == "bfloat16": query, key, value = (x.to(torch.bfloat16) for x in (query, key, value)) - if self.attention_backend == "flash4": - result = self._flash4_attend(query, key, value, mask) + if self.attention_backend in ("flash3", "flash4"): + kernels = load_flash_attention_kernels(self.attention_backend, query.device) + layout = attention_layout or prepare_attention_layout( + mask[:, 0, 0], query_tokens=query.shape[2] + ) + result = execute_flash_attention(query, key, value, layout, kernels) elif self.attention_backend == "efficient": if query.device.type != "cuda": raise RuntimeError("efficient attention requires CUDA; no math fallback") @@ -146,40 +288,9 @@ def _apply_attention(self, query, key, value, mask): @staticmethod def _flash4_attend(query, key, value, mask): - # Batch-one support is explicit until packed variable-length batches are - # validated. Removing padded keys preserves full joint attention, - # including noncontiguous padding; no queries or valid keys are dropped. - if query.device.type != "cuda": - raise RuntimeError("flash4 attention requires a CUDA Hopper/Blackwell GPU") - if query.shape[0] != 1 or mask.dtype != torch.bool or mask.shape != (1, 1, 1, key.shape[2]): - raise ValueError("flash4 currently requires batch1 and a boolean key-padding mask") - if any(x.dtype != torch.bfloat16 for x in (query, key, value)): - raise ValueError("flash4 requires bfloat16 Q/K/V") - from packaging.version import Version - if Version(torch.__version__.split("+")[0]) < Version("2.9"): - raise RuntimeError("flash4 profile requires the separately verified Torch>=2.9 runtime; see docs/WORLD_MODEL.md") - if torch.cuda.get_device_capability(query.device)[0] < 9: - raise RuntimeError("flash4 profile requires Hopper or newer CUDA hardware") - try: - from flash_attn.cute import flash_attn_func - except (ImportError, AttributeError) as error: - raise RuntimeError("Install the verified optional FlashAttention4 runtime; no silent backend fallback") from error - valid_keys = mask[0, 0, 0].nonzero(as_tuple=True)[0] - if valid_keys.numel() == 0: - raise ValueError("flash4 attention requires at least one valid key") - key = key.index_select(2, valid_keys) - value = value.index_select(2, valid_keys) - q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value)) - result = flash_attn_func(q, k, v, causal=False, return_lse=False) - # The tested 4.0.0b35 API returns (output, optional_lse). Reject unknown - # containers/shapes instead of assuming the README illustration's API. - if isinstance(result, tuple): - if len(result) != 2: - raise RuntimeError("Unsupported FlashAttention4 result tuple") - result = result[0] - if not torch.is_tensor(result) or result.shape != q.shape or result.dtype != q.dtype or result.device != q.device: - raise RuntimeError("Unsupported FlashAttention4 output shape/dtype/device") - return result.transpose(1, 2) + kernels = load_flash_attention_kernels("flash4", query.device) + layout = prepare_attention_layout(mask[:, 0, 0], query_tokens=query.shape[2]) + return execute_flash_attention(query, key, value, layout, kernels) @staticmethod def _attend(query, key, value, mask): @@ -277,7 +388,8 @@ def __init__( ) self.blocks = nn.ModuleList( MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size, - attention_backend, attention_dtype) + attention_backend=attention_backend, + attention_dtype=attention_dtype) for _ in range(depth) ) # One private image adapter per embodiment per block. Zero disables the diff --git a/egomimic/pipeline/stages_dino_proprio.py b/egomimic/pipeline/stages_dino_proprio.py new file mode 100644 index 000000000..4464f824d --- /dev/null +++ b/egomimic/pipeline/stages_dino_proprio.py @@ -0,0 +1,13 @@ +"""Preserve the original DINO recipe's graph keys with the generic stage.""" + +from egomimic.pipeline.stages_image_proprio import ImageProprioStage + + +class DinoProprioStage(ImageProprioStage): + loss_key = "loss/dino_proprio" + prediction_key = "dino_proprio_prediction" + log_prefix = "dino_proprio" + writes_by_mode = { + "train": ("loss/dino_proprio",), + "inference": ("dino_proprio_prediction",), + } diff --git a/egomimic/pipeline/stages_image_proprio.py b/egomimic/pipeline/stages_image_proprio.py new file mode 100644 index 000000000..3d9fa1e5a --- /dev/null +++ b/egomimic/pipeline/stages_image_proprio.py @@ -0,0 +1,86 @@ +"""Two-modality visual-state and direct-action experiments.""" + +from egomimic.pipeline.core import Stage, resolve_homogeneous_scalar + + +class ImageProprioStage(Stage): + loss_key = "loss/image_proprio" + prediction_key = "image_proprio_prediction" + log_prefix = "image_proprio" + reads_by_mode = { + "train": ("world", "embodiment", "optimizer_step"), + "inference": ("world", "embodiment"), + } + writes_by_mode = { + "train": ("loss/image_proprio",), + "inference": ("image_proprio_prediction",), + } + + def __init__( + self, + model, + warmup_steps=100, + pose_ramp_steps=100, + image_weight=1.0, + pose_weight=1.0, + sampler_steps=20, + action_weight=1.0, + ): + super().__init__() + if ( + min(warmup_steps, pose_ramp_steps, image_weight, pose_weight, action_weight) + < 0 + ): + raise ValueError("Schedules and loss weights must be nonnegative") + if type(sampler_steps) is not int or sampler_steps < 1: + raise ValueError("sampler_steps must be a positive integer") + self.model = model + self.warmup_steps, self.pose_ramp_steps = warmup_steps, pose_ramp_steps + self.image_weight, self.pose_weight = image_weight, pose_weight + self.action_weight = action_weight + self.sampler_steps = sampler_steps + + def pose_scale(self, step): + return self.pose_weight * self.readout_ramp(step) + + def readout_ramp(self, step): + if step < self.warmup_steps: + return 0.0 + if self.pose_ramp_steps == 0: + return 1.0 + return min(1.0, (step - self.warmup_steps) / self.pose_ramp_steps) + + def execute(self, batch, *, mode): + embodiment = str( + resolve_homogeneous_scalar(batch["embodiment"], label="embodiment") + ) + if mode == "train": + scale = self.pose_scale(int(batch["optimizer_step"])) + action_scale = self.action_weight * self.readout_ramp( + int(batch["optimizer_step"]) + ) + losses = self.model.training_losses( + batch["world"], + embodiment, + pose_weight=scale, + image_weight=self.image_weight, + action_weight=action_scale, + ) + batch[self.loss_key] = losses["total"] + batch.update( + {f"log/{self.log_prefix}_{k}": v.detach() for k, v in losses.items()} + ) + image = batch["world"]["image"] + batch[f"log/{self.log_prefix}_image_tokens"] = image["valid"].sum().float() + batch[f"log/{self.log_prefix}_image_query_tokens"] = ( + image["valid"] & ~image["observed"] + ).sum().float() + if self.model.objective in ("visual_proprio", "visual_state_action"): + batch["log/pose_weight"] = losses["total"].new_tensor(scale) + if self.model.objective == "visual_state_action": + batch["log/action_weight"] = losses["total"].new_tensor(action_scale) + else: + batch[self.prediction_key] = self.model.sample( + batch["world"], embodiment, steps=self.sampler_steps + ) + return batch diff --git a/egomimic/pipeline/stages_pixel_readout.py b/egomimic/pipeline/stages_pixel_readout.py new file mode 100644 index 000000000..b945e97a9 --- /dev/null +++ b/egomimic/pipeline/stages_pixel_readout.py @@ -0,0 +1,350 @@ +"""Explicit state/command readout graph sharing one pixel diffusion model. + +The initial readout phase freezes every copied image parameter. Legacy models +with learned phase offsets freeze that entire embedding too; new configurations +omit it. State/control adapters, queries, expert and heads remain trainable. +""" + +import math + +import torch + +from egomimic.pipeline.core import Stage +from egomimic.pipeline.stages_pixel_video import ( + PixelVideoDenoiserStage, + PixelVideoLossStage, + PixelVideoNoisingStage, + _embodiment, +) + + +def configure_readout_trainability(model, training_mode): + """Partition every parameter explicitly; never freeze only an embedding row.""" + if training_mode not in ("frozen_image", "joint"): + raise ValueError("Readout training_mode must be frozen_image or joint") + if model.objective != "visual_state_action" or model.phase_a_condition_on_proprio: + raise ValueError("Readout requires visual_state_action with image-only Phase A") + image_modules = [ + model.image_in, + model.image_out, + model.noise_level, + model.norm[0], + ] + if hasattr(model, "xy"): + image_modules.append(model.xy) + if model.noise_conditioning == "adaln_zero": + image_modules.append(model.final_adaln) + if model.use_phase_embedding: + image_modules.append(model.phase) + if model.use_camera_embedding: + image_modules.append(model.camera) + if model.use_embodiment_embedding: + image_modules.append(model.embodiment) + readout_modules = [ + model.pose_in, + model.pose_out, + model.action_out, + model.proprio_role, + model.norm[1], + ] + for block in model.blocks: + if model.noise_conditioning == "adaln_zero": + image_modules.append(block.adaln_image) + for name in ("norm1", "qkv", "out", "norm2", "ffn"): + branches = getattr(block, name) + if len(branches) != 2: + raise ValueError( + "Readout requires exactly image and shared proprio experts" + ) + image_modules.append(branches[0]) + readout_modules.append(branches[1]) + image_ids = {id(p) for module in image_modules for p in module.parameters()} + if model.use_modality_embedding: + image_ids.add(id(model.modality)) + readout_ids = {id(p) for module in readout_modules for p in module.parameters()} + readout_ids.update((id(model.query), id(model.action_query))) + parameters = dict(model.named_parameters()) + if image_ids & readout_ids or image_ids | readout_ids != { + id(p) for p in parameters.values() + }: + raise ValueError( + "Readout trainability partition has unknown or shared parameters" + ) + for parameter in parameters.values(): + parameter.requires_grad_( + training_mode == "joint" or id(parameter) in readout_ids + ) + return { + "training_mode": training_mode, + "frozen_parameters": tuple( + name for name, p in parameters.items() if not p.requires_grad + ), + "trainable_parameters": tuple( + name for name, p in parameters.items() if p.requires_grad + ), + "phase_embedding_present": model.use_phase_embedding, + "phase_embedding_frozen": model.use_phase_embedding + and training_mode == "frozen_image", + "noise_conditioning": model.noise_conditioning, + } + + +class PixelReadoutDenoiserStage(PixelVideoDenoiserStage): + """Legacy combined model owner and high-noise training prediction stage.""" + + def __init__(self, model, training_mode="frozen_image"): + super().__init__(model) + self.trainability = configure_readout_trainability(model, training_mode) + + +class PixelReadoutModelStage(Stage): + """Register/freeze the sole model without choosing any forward branch.""" + + writes = ("pixel_video/denoiser",) + + def __init__(self, model, training_mode="frozen_image"): + super().__init__() + self.model = model + self.trainability = configure_readout_trainability(model, training_mode) + + def execute(self, batch, *, mode): + batch[self.writes[0]] = self.model + return batch + + +class PixelReadoutVideoDenoiserStage(Stage): + """Optional high-noise image prediction, independent of model ownership.""" + + train_only = True + reads = ( + "world", + "embodiment", + "pixel_video/denoiser", + "pixel_video/noisy_image", + "pixel_video/noise_levels", + ) + writes = ("pixel_video/prediction",) + + def execute(self, batch, *, mode): + batch[self.writes[0]] = batch["pixel_video/denoiser"]( + batch["world"], + _embodiment(batch), + batch["pixel_video/noisy_image"], + batch["pixel_video/noise_levels"], + phase="A", + )["image"] + return batch + + +class PixelReadoutImageLossStage(PixelVideoLossStage): + """Image objective in a multi-loss graph; generic Train/loss is the total.""" + + def execute(self, batch, *, mode): + batch = super().execute(batch, mode=mode) + # The base image-only node labels its sole loss as total. Here the two + # readout objectives are separate nodes, so that partial total is invalid. + batch.pop("log/image_proprio_total", None) + return batch + + +class PixelReadoutNoisingStage(PixelVideoNoisingStage): + """Independent low-noise corruption for paired-trajectory training.""" + + writes = ( + "pixel_readout/noisy_image", + "pixel_readout/noise_levels", + "pixel_readout/target", + "pixel_readout/loss_mask", + "pixel_readout/noise_metadata", + ) + + +class PixelReadoutImageDenoiserStage(Stage): + """Estimate aligned images using the same phase-A weights, without detaching.""" + + train_only = True + reads = ( + "world", + "embodiment", + "pixel_video/denoiser", + "pixel_readout/noisy_image", + "pixel_readout/noise_levels", + "pixel_readout/target", + ) + writes = ("pixel_readout/image_prediction",) + + def __init__(self, prediction_key="pixel_readout/image_prediction"): + super().__init__() + if not isinstance(prediction_key, str) or not prediction_key: + raise ValueError("Image prediction key must be a nonempty string") + self.writes = (prediction_key,) + + def execute(self, batch, *, mode): + image = batch["world"]["image"] + prediction = batch["pixel_video/denoiser"]( + batch["world"], + _embodiment(batch), + batch["pixel_readout/noisy_image"], + batch["pixel_readout/noise_levels"], + phase="A", + )["image"] + batch[self.writes[0]] = torch.where( + image["observed"][..., None], batch["pixel_readout/target"], prediction + ).masked_fill(~image["valid"][..., None], 0) + return batch + + +class PixelReadoutSampledImageStage(Stage): + """Exchangeable sampled-image provider with explicit stopped gradients. + + The built-in sampler does not backpropagate through generation. Frozen-image + training can use this provider; joint training would need separate image + objectives or the Phase-B path for gradients to the image parameters. + """ + + reads = ("world", "embodiment", "pixel_video/denoiser") + + def __init__( + self, + sampler, + modes=("inference",), + prediction_key="pixel_readout/image_prediction", + gradient_policy="stop", + ): + super().__init__() + if isinstance(modes, str): + raise ValueError("Provider modes must be a sequence of train/inference") + modes = tuple(modes) + if ( + not modes + or len(set(modes)) != len(modes) + or not set(modes) + <= { + "train", + "inference", + } + ): + raise ValueError("Provider modes must select train and/or inference once") + if not isinstance(prediction_key, str) or not prediction_key: + raise ValueError("Image prediction key must be a nonempty string") + if gradient_policy != "stop": + raise ValueError( + "Sampled image provider only supports gradient_policy=stop" + ) + self.sampler = sampler + self.modes, self.gradient_policy = modes, gradient_policy + self.train_only = "inference" not in modes + self.inference_only = "train" not in modes + self.writes = (prediction_key,) + + def execute(self, batch, *, mode): + with torch.no_grad(): + image = self.sampler( + batch["pixel_video/denoiser"], batch["world"], _embodiment(batch) + ) + if not isinstance(image, torch.Tensor): + raise TypeError("Image providers must produce a tensor") + batch[self.writes[0]] = image.detach() + return batch + + +class PixelStateActionReadoutStage(Stage): + """Phase B uses observed states and typed state/command queries, never labels.""" + + reads_by_mode = { + "train": ( + "world", + "embodiment", + "pixel_video/denoiser", + "pixel_readout/image_prediction", + ), + "inference": ( + "world", + "embodiment", + "pixel_video/denoiser", + "pixel_video/sampled_prediction", + ), + } + writes_by_mode = { + "train": ("pixel_readout/prediction",), + "inference": ("image_proprio_prediction",), + } + + def __init__(self, image_key=None): + super().__init__() + if image_key is not None: + if not isinstance(image_key, str) or not image_key: + raise ValueError("Readout image key must be a nonempty string") + self.reads_by_mode = { + mode: ("world", "embodiment", "pixel_video/denoiser", image_key) + for mode in ("train", "inference") + } + self.image_key = image_key + + def execute(self, batch, *, mode): + world = batch["world"] + image = ( + batch[self.image_key] + if self.image_key is not None + else batch["pixel_readout/image_prediction"] + if mode == "train" + else batch["pixel_video/sampled_prediction"]["image"] + ) + prediction = batch["pixel_video/denoiser"]( + world, + _embodiment(batch), + image, + torch.zeros_like(world["image"]["time"], dtype=torch.float32), + phase="B", + ) + if mode == "train": + batch["pixel_readout/prediction"] = prediction + else: + pose = world["proprio"] + observed = pose["valid"] & pose["observed"] + state = torch.where( + observed[..., None], pose["value"], prediction["proprio"] + ) + batch["image_proprio_prediction"] = { + "image": image, + "proprio": state.masked_fill(~pose["valid"][..., None], 0), + "action": prediction["action"], + } + return batch + + +class PixelReadoutMSEStage(Stage): + """One explicit supervised objective for measured state or actual commands.""" + + train_only = True + reads = ("world", "pixel_readout/prediction") + + def __init__(self, loss, stream, weight=1.0): + super().__init__() + if stream not in ("proprio", "action"): + raise ValueError("Readout MSE stream must be proprio or action") + if not math.isfinite(weight) or weight < 0: + raise ValueError("Readout loss weight must be finite and nonnegative") + self.loss, self.stream, self.weight = loss, stream, weight + name = "measured_state" if stream == "proprio" else "control_command" + self.writes = (f"loss/{name}",) + self.log_key = f"log/{name}_mse" + + def execute(self, batch, *, mode): + stream = batch["world"][self.stream] + mask = stream["valid"] & stream.get("target_valid", stream["valid"]) + if self.stream == "proprio": + mask = mask & ~stream["observed"] + loss = self.loss( + batch["pixel_readout/prediction"][self.stream], + stream["value"], + mask, + metadata={ + "timestamps": stream["time"], + "valid": stream["valid"], + "observed": stream["observed"], + }, + ) + batch[self.writes[0]] = self.weight * loss + batch[self.log_key] = loss.detach() + return batch diff --git a/egomimic/pipeline/stages_pixel_video.py b/egomimic/pipeline/stages_pixel_video.py new file mode 100644 index 000000000..d631ae7c8 --- /dev/null +++ b/egomimic/pipeline/stages_pixel_video.py @@ -0,0 +1,137 @@ +"""Replaceable image-only diffusion stages with one registered denoiser.""" + +import math + +from egomimic.pipeline.core import Stage, resolve_homogeneous_scalar + + +def _embodiment(batch): + return str(resolve_homogeneous_scalar(batch["embodiment"], label="embodiment")) + + +class PixelVideoNoisingStage(Stage): + """Draw corruption independently of the denoiser and training objective.""" + + train_only = True + reads = ("world",) + writes = ( + "pixel_video/noisy_image", + "pixel_video/noise_levels", + "pixel_video/target", + "pixel_video/loss_mask", + "pixel_video/noise_metadata", + ) + + def __init__(self, noising): + super().__init__() + self.noising = noising + + def execute(self, batch, *, mode): + result = self.noising(batch["world"]["image"]) + for destination, source in zip( + self.writes, ("image_values", "noise_levels", "target", "mask", "metadata") + ): + batch[destination] = result[source] + return batch + + +class PixelVideoDenoiserStage(Stage): + """Own the only denoiser parameters and expose them to the sampler.""" + + reads_by_mode = { + "train": ( + "world", + "embodiment", + "pixel_video/noisy_image", + "pixel_video/noise_levels", + ), + "inference": (), + } + writes_by_mode = { + "train": ("pixel_video/prediction", "pixel_video/denoiser"), + "inference": ("pixel_video/denoiser",), + } + + def __init__(self, model): + super().__init__() + if ( + model.objective not in ("visual_only", "visual_state_action") + or model.phase_a_condition_on_proprio + ): + raise ValueError("Pixel video denoising requires an image-only Phase A") + self.model = model + + def execute(self, batch, *, mode): + # A transient reference shares the sole registered model with any + # configured sampler/readout branch, without a second parameter copy. + batch["pixel_video/denoiser"] = self.model + if mode == "train": + batch["pixel_video/prediction"] = self.model( + batch["world"], + _embodiment(batch), + batch["pixel_video/noisy_image"], + batch["pixel_video/noise_levels"], + phase="A", + )["image"] + return batch + + +class PixelVideoLossStage(Stage): + """Configure the objective separately, retaining temporal/noise metadata.""" + + train_only = True + reads = ( + "pixel_video/prediction", + "pixel_video/target", + "pixel_video/loss_mask", + "pixel_video/noise_levels", + "pixel_video/noise_metadata", + ) + writes = ("loss/pixel_video",) + + def __init__(self, loss, weight=1.0): + super().__init__() + if not math.isfinite(weight) or weight < 0: + raise ValueError("Pixel loss weight must be finite and nonnegative") + self.loss, self.weight = loss, weight + + def execute(self, batch, *, mode): + loss = self.loss( + batch["pixel_video/prediction"], + batch["pixel_video/target"], + batch["pixel_video/loss_mask"], + noise_levels=batch["pixel_video/noise_levels"], + metadata=batch["pixel_video/noise_metadata"], + check_target=False, + ) + batch["loss/pixel_video"] = self.weight * loss + batch["log/image_proprio_image"] = loss.detach() + batch["log/image_proprio_total"] = batch["loss/pixel_video"].detach() + metadata = batch["pixel_video/noise_metadata"] + batch["log/image_proprio_image_tokens"] = metadata["valid"].sum().float() + batch["log/image_proprio_image_query_tokens"] = ( + (metadata["valid"] & ~metadata["observed"]).sum().float() + ) + return batch + + +class PixelVideoSamplerStage(Stage): + """Run a replaceable inference integrator against the shared denoiser.""" + + inference_only = True + reads = ("world", "embodiment", "pixel_video/denoiser") + + def __init__(self, sampler, prediction_key="image_proprio_prediction"): + super().__init__() + if not isinstance(prediction_key, str) or not prediction_key: + raise ValueError("Sampler prediction key must be a nonempty string") + self.sampler = sampler + self.prediction_key = prediction_key + self.writes = (prediction_key,) + + def execute(self, batch, *, mode): + image = self.sampler( + batch["pixel_video/denoiser"], batch["world"], _embodiment(batch) + ) + batch[self.prediction_key] = {"image": image} + return batch diff --git a/egomimic/pl_utils/world_model_behavior.py b/egomimic/pl_utils/world_model_behavior.py index d1c36f783..74a79090f 100644 --- a/egomimic/pl_utils/world_model_behavior.py +++ b/egomimic/pl_utils/world_model_behavior.py @@ -7,8 +7,75 @@ from egomimic.pl_utils.training_behavior import TrainingBehavior +_DISTRIBUTED_RNG_SCHEMA = "mot_rng_state.distributed.v1" + + +def _distributed_world(): + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_world_size(), torch.distributed.get_rank() + return 1, 0 + + +def _capture_rng_state(*, current_cuda_only=False): + cuda = [] + if torch.cuda.is_available(): + cuda = ([torch.cuda.get_rng_state()] if current_cuda_only + else torch.cuda.get_rng_state_all()) + state = { + "torch_cpu": torch.get_rng_state(), "torch_cuda": cuda, + "numpy": np.random.get_state(), "python": random.getstate(), + } + if current_cuda_only: + state["cuda_scope"] = "current_device" + return state + + +def _rank_rng_state(saved): + """Select at the first train step, after distributed initialization.""" + world_size, rank = _distributed_world() + if "schema" not in saved: + if world_size != 1: + raise ValueError("Exact MoT DDP resume requires per-rank RNG state") + return saved + if saved.get("schema") != _DISTRIBUTED_RNG_SCHEMA: + raise ValueError("Unknown MoT RNG checkpoint schema") + if saved.get("world_size") != world_size: + raise ValueError("Exact MoT resume requires the saved RNG world size") + ranks = saved.get("ranks") + if not isinstance(ranks, list) or len(ranks) != world_size: + raise ValueError("MoT RNG checkpoint must contain one state per rank") + state = ranks[rank] + if not isinstance(state, dict) or state.get("cuda_scope") != "current_device": + raise ValueError("Distributed MoT RNG state must bind the current CUDA device") + if len(state.get("torch_cuda", [])) > 1: + raise ValueError("Distributed MoT RNG state contains multiple CUDA devices") + return state + + +def _restore_rng_state(state): + torch.set_rng_state(state["torch_cpu"].cpu()) + if state["torch_cuda"]: + if not torch.cuda.is_available(): + raise ValueError("Cannot restore CUDA RNG on a CPU training runtime") + if state.get("cuda_scope") == "current_device": + torch.cuda.set_rng_state(state["torch_cuda"][0].cpu()) + else: + torch.cuda.set_rng_state_all([value.cpu() for value in state["torch_cuda"]]) + np.random.set_state(state["numpy"]) + random.setstate(state["python"]) + class WorldModelTrainingBehavior(TrainingBehavior): + def __init__(self, distributed_noise_seed=None): + super().__init__() + if distributed_noise_seed is not None and ( + type(distributed_noise_seed) is not int + or not 0 <= distributed_noise_seed < 2**32 + ): + raise ValueError("distributed_noise_seed must be a uint32 integer or None") + self.distributed_noise_seed = distributed_noise_seed + self._training_rng_initialized = False + def on_validation_epoch_end(self): if self.context.evaluator is not None: self.context.evaluator.log_epoch_metrics() @@ -16,14 +83,21 @@ def on_validation_epoch_end(self): def training_step(self, batch, batch_idx): pending = getattr(self, "_pending_rng", None) if pending is not None: - torch.set_rng_state(pending["torch_cpu"].cpu()) - if pending["torch_cuda"]: - if not torch.cuda.is_available(): - raise ValueError("Cannot restore CUDA RNG on a CPU training runtime") - torch.cuda.set_rng_state_all([state.cpu() for state in pending["torch_cuda"]]) - np.random.set_state(pending["numpy"]) - random.setstate(pending["python"]) + _restore_rng_state(_rank_rng_state(pending)) self._pending_rng = None + self._training_rng_initialized = True + elif not self._training_rng_initialized: + world_size, rank = _distributed_world() + if self.distributed_noise_seed is not None and world_size > 1: + # Model construction remains identically seeded on every rank. + # Fresh training then gets distinct streams; resume above wins. + seed = (self.distributed_noise_seed + rank) % (2**32) + torch.random.default_generator.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) # Current rank's device only. + np.random.seed(seed) + random.seed(seed) + self._training_rng_initialized = True data = self.context.trainer.datamodule if getattr(data, "checkpointable_sampling", False): cursors = {values["sampling_cursor"] for values in batch.values()} @@ -37,10 +111,20 @@ def training_step(self, batch, batch_idx): return self.context._default_training_step(batch, batch_idx) def on_save_checkpoint(self, checkpoint): + world_size, _ = _distributed_world() + local = _capture_rng_state(current_cuda_only=world_size > 1) + if world_size == 1: + checkpoint["mot_rng_state"] = local + return + # Lightning calls model checkpoint hooks on every rank before its + # strategy limits the file write to rank zero. This collective must + # remain outside a rank-zero guard. + ranks = [None] * world_size + torch.distributed.all_gather_object(ranks, local) checkpoint["mot_rng_state"] = { - "torch_cpu": torch.get_rng_state(), - "torch_cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [], - "numpy": np.random.get_state(), "python": random.getstate(), + "schema": _DISTRIBUTED_RNG_SCHEMA, + "world_size": world_size, + "ranks": ranks, } def on_load_checkpoint(self, checkpoint): diff --git a/egomimic/rldb/dino_proprio.py b/egomimic/rldb/dino_proprio.py new file mode 100644 index 000000000..09fe4c304 --- /dev/null +++ b/egomimic/rldb/dino_proprio.py @@ -0,0 +1,611 @@ +"""DINO and measured-proprioception caches, with optional verified action labels. + +Point tracks are neither required nor loaded. Action commands are never inferred +from measured poses: direct BC requires a separate, SHA-bound alignment receipt. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from egomimic.rldb.world_model import ( + _bounded_time_sampling, + _physical_times, + open_prepared_episode, + sample_times, +) +from egomimic.rldb.world_model_preparation import ( + ArtifactVerificationMemo, + training_identities, +) + + +def _sha(value): + return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None + + +def _coordinate_schema(schema, label): + dimension = schema.get("dimension") + fields = schema.get("fields", []) + if ( + type(dimension) is not int + or dimension < 1 + or len(fields) != dimension + or any(not isinstance(field, str) or not field for field in fields) + or len(set(fields)) != dimension + ): + raise ValueError( + f"{label} requires explicit unique fields for every coordinate" + ) + + +def load_dino_proprio_normalization(manifest, root, verifier): + """Use only image/pose statistics, including from historical three-stream files.""" + contract = manifest["normalization"] + if ( + contract.get("scope") != "train_only" + or not _sha(contract.get("sha256")) + or type(contract.get("arrays_normalized")) is not bool + ): + raise ValueError("Require train-only normalization provenance") + if contract["arrays_normalized"] and not contract.get("path"): + return None + if not contract.get("path"): + raise ValueError( + "Raw caches require bound image/proprio normalization statistics" + ) + path = root / contract["path"] + verifier.verify(path, contract["sha256"]) + saved = json.loads(path.read_text()) + if ( + saved.get("schema") + not in ("mot-world-normalization-v1", "dino-proprio-normalization-v1") + or saved.get("scope") != "train_only" + or saved.get("train_episodes") != training_identities(manifest["episodes"]) + ): + raise ValueError( + "Normalization statistics do not match training episode identities" + ) + arrays = {} + required = {"image"} | {f"pose/{name}" for name in manifest["pose_schemas"]} + for key in required: + item = saved.get("statistics", {}).get(key) + if item is None: + raise ValueError(f"Missing normalization statistics: {key}") + mean, std = ( + np.asarray(item["mean"], np.float32), + np.asarray(item["std"], np.float32), + ) + if ( + mean.ndim != 1 + or not len(mean) + or mean.shape != std.shape + or not np.isfinite(mean).all() + or not np.isfinite(std).all() + or np.any(std <= 0) + or item.get("count", 0) < 1 + ): + raise ValueError(f"Invalid normalization statistics: {key}") + if ( + key.startswith("pose/") + and len(mean) != manifest["pose_schemas"][key[5:]]["dimension"] + ): + raise ValueError(f"Normalization schema width mismatch: {key}") + arrays[key] = mean, std + return arrays if not contract["arrays_normalized"] else None + + +def collate_dino_proprio(samples): + """Pad only the declared streams; action values remain label-only tensors.""" + if not samples or len({s["embodiment"] for s in samples}) != 1: + raise ValueError("Batch must be nonempty and homogeneous in embodiment") + names = set(samples[0]["world"]) + if names not in ({"image", "proprio"}, {"image", "proprio", "action"}): + raise ValueError("Expected DINO/proprio with optional action labels") + if any(set(s["world"]) != names for s in samples): + raise ValueError("Cannot mix training objectives in a batch") + result = {} + for name in sorted(names): + size = max(len(s["world"][name]["value"]) for s in samples) + fields = set(samples[0]["world"][name]) + if any(set(s["world"][name]) != fields for s in samples): + raise ValueError("Inconsistent stream fields") + result[name] = {} + for key in fields: + values = [] + for sample in samples: + value = sample["world"][name][key] + padded = torch.zeros((size, *value.shape[1:]), dtype=value.dtype) + padded[: len(value)] = value + values.append(padded) + result[name][key] = torch.stack(values) + embodiment = samples[0]["embodiment"] + kinds = [sample.get("task_kind") for sample in samples] + allowed = {"direct_action"} if "action" in names else {"forecast", "completion"} + if any(kind not in allowed for kind in kinds): + raise ValueError( + "Every sample requires its explicit observation-mask task kind" + ) + item = {"embodiment": embodiment, "world": result, "task_kind": kinds} + if any("sampling_cursor" in s for s in samples): + cursors = {s.get("sampling_cursor") for s in samples} + if None in cursors or len(cursors) != 1: + raise ValueError("A resumed batch must have one source-draw cursor") + item["sampling_cursor"] = cursors.pop() + return {embodiment: item} + + +class EvenlySpacedEpisodeBatchSampler: + """Fixed held-out anchors spanning every episode, with homogeneous batches.""" + + def __init__(self, dataset, anchors_per_episode, batch_size=1): + if getattr(dataset, "split", None) != "val": + raise ValueError("Episode-spanning anchor subsets are for validation only") + if type(anchors_per_episode) is not int or anchors_per_episode < 1: + raise ValueError( + "validation_anchors_per_episode must be a positive integer" + ) + if type(batch_size) is not int or batch_size < 1: + raise ValueError("Positive batch size required") + self.batch_size = batch_size + self.groups = {} + offset = 0 + for episode, count in zip(dataset.episodes, dataset.counts): + selected = np.linspace( + 0, count - 1, min(anchors_per_episode, count), dtype=np.int64 + ) + self.groups.setdefault(episode["embodiment"], []).extend( + (selected + offset).tolist() + ) + offset += count + + def __len__(self): + return sum( + (len(indices) + self.batch_size - 1) // self.batch_size + for indices in self.groups.values() + ) + + def __iter__(self): + for indices in self.groups.values(): + for start in range(0, len(indices), self.batch_size): + yield indices[start : start + self.batch_size] + + +class PreparedDinoProprioDataset(Dataset): + """DINO + proprio diffusion/readout, or same-frame causal action BC. + + Direct BC reads exactly one observed frame and its aligned action command. + The visual objective retains the configured recorded-time sampler and + observation masks. A global anchor index makes source draws frame-weighted. + """ + + def __init__( + self, + manifest_path, + split="train", + task="visual_diffusion", + action_manifest_path=None, + mode="continuous", + seed=42, + min_frames=2, + max_frames=4, + discrete_frames=4, + stride=1, + forecasting_probability=0.5, + max_past_seconds=None, + max_future_seconds=None, + verification_memo=None, + ): + if task not in ("visual_diffusion", "direct_action"): + raise ValueError("Unknown DINO/proprio objective") + if task == "direct_action" and not action_manifest_path: + raise ValueError( + "Direct action BC requires a verified action alignment manifest" + ) + if split not in ("train", "val"): + raise ValueError("Explicit train/val split required") + self.path, self.task = Path(manifest_path), task + self.manifest = json.loads(self.path.read_text()) + manifest = self.manifest + if manifest.get("schema") not in ( + "mot-world-prepared-v1", + "dino-proprio-prepared-v1", + ): + raise ValueError("Unsupported DINO/proprio prepared schema") + for key in ("dino_identity", "normalization", "pose_schemas", "episodes"): + if not manifest.get(key): + raise ValueError(f"Missing {key}") + identity = manifest["dino_identity"] + for key in ("model", "revision", "image_size", "patch_grid", "feature_layer"): + if not identity.get(key): + raise ValueError(f"Missing DINO identity field {key}") + for key in ("image_size", "patch_grid"): + value = identity[key] + if ( + not isinstance(value, list) + or len(value) != 2 + or any(type(x) is not int or x < 1 for x in value) + ): + raise ValueError(f"DINO {key} requires two positive integers") + for schema in manifest["pose_schemas"].values(): + _coordinate_schema(schema, "Pose schema") + self.verifier = verification_memo or ArtifactVerificationMemo() + self.normalization = load_dino_proprio_normalization( + manifest, self.path.parent, self.verifier + ) + ids, paths = set(), set() + for ep in manifest["episodes"]: + key = (ep["embodiment"], ep["episode_id"]) + physical = (self.path.parent / ep["path"]).resolve() + if key in ids or physical in paths: + raise ValueError( + "Duplicate episode identity or physical path across splits" + ) + ids.add(key) + paths.add(physical) + if ( + ep.get("frame") != "fixed_reference" + or ep.get("pose_kind") != "measured_pose" + or ep["embodiment"] not in manifest["pose_schemas"] + ): + raise ValueError( + "Require measured proprioception in a declared fixed reference frame" + ) + if ( + ep.get("split") not in ("train", "val") + or type(ep.get("frames")) is not int + or ep["frames"] < 1 + ): + raise ValueError("Explicit split and positive frame count required") + if not _sha(ep.get("sha256")): + raise ValueError("Require feature-cache SHA256") + self.verifier.verify( + physical, + ep["sha256"], + trusted_fingerprint=ep.get("verified_fingerprint"), + ) + self.action_path = Path(action_manifest_path) if action_manifest_path else None + self.action_manifest, self.actions, self.action_alignment = None, {}, {} + if task == "direct_action": + self._load_actions() + self.episodes = [ep for ep in manifest["episodes"] if ep["split"] == split] + if not self.episodes: + raise ValueError("No episodes in selected split") + self.mode, self.seed, self.split = mode, seed, split + self.sampling = dict( + mode=mode, + min_frames=min_frames, + max_frames=max_frames, + discrete_frames=discrete_frames, + stride=stride, + ) + bounded = _bounded_time_sampling(max_past_seconds, max_future_seconds) + if ( + mode not in ("continuous", "discrete") + or min_frames < 2 + or max_frames < min_frames + or stride < 1 + ): + raise ValueError("Invalid time sampler") + if bounded and mode != "continuous": + raise ValueError("Physical-time bounds apply only to continuous sampling") + if bounded: + self.sampling.update( + max_past_seconds=max_past_seconds, max_future_seconds=max_future_seconds + ) + if not 0 <= forecasting_probability <= 1: + raise ValueError("Invalid forecasting probability") + self.forecasting_probability = forecasting_probability + self.anchor_indices = [] + for ep in self.episodes: + if task == "direct_action": + anchors = np.arange(ep["frames"]) + elif bounded: + with open_prepared_episode(self.path.parent / ep["path"]) as saved: + times = _physical_times(saved["time"][:], ep["frames"]) + lo = np.searchsorted(times, times - max_past_seconds, side="left") + hi = np.searchsorted(times, times + max_future_seconds, side="right") + anchors = np.flatnonzero(hi - lo >= min_frames) + else: + count = ep["frames"] - ( + (discrete_frames - 1) * stride if mode == "discrete" else 0 + ) + anchors = ( + np.arange(max(0, count)) + if ep["frames"] >= min_frames + else np.array([], np.int64) + ) + self.anchor_indices.append(anchors) + self.counts = [len(a) for a in self.anchor_indices] + if any(c < 1 for c in self.counts): + raise ValueError("Episode too short for temporal recipe") + self.ends = np.cumsum(self.counts) + self.rng = self.worker_seed = None + + def _load_actions(self): + self.action_manifest = json.loads(self.action_path.read_text()) + action = self.action_manifest + if action.get("schema") != "dino-proprio-actions-v1": + raise ValueError("Unsupported action alignment schema") + schemas = action.get("action_schemas", {}) + for schema in schemas.values(): + _coordinate_schema(schema, "Action schema") + if ( + schema.get("semantics") != "action_command" + or schema.get("value_space") != "native_commands" + ): + raise ValueError("Action schema must declare native action commands") + for item in action.get("episodes", []): + key = (item["embodiment"], item["episode_id"]) + if key in self.actions: + raise ValueError("Duplicate action episode identity") + self.actions[key] = item + expected = { + (ep["embodiment"], ep["episode_id"]) for ep in self.manifest["episodes"] + } + if set(self.actions) != expected: + raise ValueError( + "Action manifest must match every prepared episode exactly" + ) + for ep in self.manifest["episodes"]: + item = self.actions[(ep["embodiment"], ep["episode_id"])] + source = ep.get("native_source", {}) + proof = item.get("alignment", {}) + if ( + ep["embodiment"] not in schemas + or item.get("feature_sha256") != ep["sha256"] + or not _sha(source.get("source_demo_sha256")) + or item.get("source_demo_sha256") != source["source_demo_sha256"] + or not source.get("episode") + or item.get("source_episode") != source["episode"] + or proof.get("relation") != "action_t_applied_to_state_t" + or proof.get("cache_frame_mapping") != "explicit_source_indices" + or not _sha(proof.get("producer_sha256")) + or not _sha(proof.get("source_alignment_sha256")) + or not proof.get("source_alignment_path") + ): + raise ValueError( + "Action alignment proof must bind source commands to exact feature frames" + ) + # A digest string is not evidence of source alignment. Require the + # actual bound receipt and inspect its source identities and mapping. + # This authenticates the exporter-recorded mapping, not the truth of + # an unknown controller convention; that remains an upstream audit. + receipt_path = self.action_path.parent / proof["source_alignment_path"] + self.verifier.verify(receipt_path, proof["source_alignment_sha256"]) + receipt = json.loads(receipt_path.read_text()) + basis = receipt.get("basis", {}) + if ( + receipt.get("schema") != "dino-proprio-source-alignment-v1" + or any( + receipt.get(k) != item[k] + for k in ("feature_sha256", "source_demo_sha256", "source_episode") + ) + or receipt.get("relation") != proof["relation"] + or receipt.get("producer_sha256") != proof["producer_sha256"] + or any( + not isinstance(basis.get(k), str) or not basis[k].strip() + for k in ( + "cache_index_source", + "action_index_source", + "state_action_relation_source", + ) + ) + ): + raise ValueError( + "Action source alignment receipt has incomplete or conflicting provenance" + ) + self.action_alignment[(ep["embodiment"], ep["episode_id"])] = receipt + self.verifier.verify( + self.action_path.parent / item["path"], + item["sha256"], + trusted_fingerprint=item.get("verified_fingerprint"), + ) + # Fail during dataset setup, before optimization, for any malformed + # alignment or invalid command. These arrays are tiny relative to DINO. + with open_prepared_episode(self.path.parent / ep["path"]) as cache: + time = cache["time"][:] + self._read_actions(ep, np.arange(ep["frames"]), time) + + def _read_actions(self, ep, indices, cache_time): + item = self.actions[(ep["embodiment"], ep["episode_id"])] + receipt = self.action_alignment[(ep["embodiment"], ep["episode_id"])] + schema = self.action_manifest["action_schemas"][ep["embodiment"]] + with open_prepared_episode(self.action_path.parent / item["path"]) as data: + names = ( + "action", + "time", + "valid", + "cache_frame_index", + "source_state_index", + "source_action_index", + ) + if any(name not in data for name in names): + raise ValueError("Action sidecar missing explicit alignment arrays") + n = ep["frames"] + if data["action"].shape != (n, schema["dimension"]) or any( + data[k].shape != (n,) for k in names[1:] + ): + raise ValueError("Action alignment array shape mismatch") + for name in names[3:]: + recorded = np.asarray(receipt.get(name)) + if ( + np.asarray(data[name][:]).dtype.kind not in "iu" + or recorded.dtype.kind not in "iu" + or recorded.shape != (n,) + ): + raise ValueError("Source alignment indices must be integers") + if not np.array_equal(data[name][:], recorded): + raise ValueError( + "Action alignment arrays differ from the bound source receipt" + ) + state, command = ( + data["source_state_index"][:], + data["source_action_index"][:], + ) + if ( + not np.array_equal(data["cache_frame_index"][:], np.arange(n)) + or not np.array_equal(state, command) + or np.any(state < 0) + or not (np.diff(state) > 0).all() + or not np.array_equal(data["time"][:], cache_time) + ): + raise ValueError("Action/state/cache-frame alignment mismatch") + value = np.asarray(data["action"][indices], np.float32) + valid = np.asarray(data["valid"][indices]) + if ( + valid.dtype != np.bool_ + or not valid.all() + or not np.isfinite(value).all() + ): + raise ValueError( + "Direct action BC requires finite valid commands at every cache frame" + ) + return value + + def __len__(self): + return int(self.ends[-1]) + + def __getitem__(self, index): + draw = None + if isinstance(index, tuple): + index, cursor, slot = index + draw = int(cursor), int(slot) + if not 0 <= index < len(self): + raise IndexError(index) + worker_seed = torch.initial_seed() + if self.rng is None or self.worker_seed != worker_seed: + self.rng = np.random.default_rng((worker_seed + self.seed) % 2**63) + self.worker_seed = worker_seed + rng = ( + self.rng + if self.split == "train" + else np.random.default_rng(self.seed + index) + ) + if draw is not None: + if self.split != "train": + raise ValueError("Checkpointed draw keys apply only to training") + rng = np.random.default_rng(np.random.SeedSequence([self.seed, *draw])) + ei = int(np.searchsorted(self.ends, index, side="right")) + ep = self.episodes[ei] + offset = 0 if ei == 0 else int(self.ends[ei - 1]) + anchor = int(self.anchor_indices[ei][index - offset]) + with open_prepared_episode(self.path.parent / ep["path"]) as data: + if any( + k not in data + for k in ("time", "image", "image_xy", "pose", "pose_valid") + ): + raise ValueError("Prepared cache missing required DINO/proprio arrays") + full_time = data["time"][:] + _physical_times(full_time, ep["frames"]) + indices = ( + np.array([anchor]) + if self.task == "direct_action" + else sample_times( + ep["frames"], + rng, + anchor=anchor, + timestamps=full_time, + **self.sampling, + ) + ) + time = np.asarray(full_time[indices], np.float64) + origin = 0 if self.task == "direct_action" else int(rng.integers(len(time))) + time = (time - time[origin]).astype(np.float32) + image, pose = ( + np.asarray(data[k][indices], np.float32) for k in ("image", "pose") + ) + if ( + image.ndim != 4 + or pose.ndim != 2 + or data["image"].shape[0] != ep["frames"] + or data["pose"].shape[0] != ep["frames"] + or pose.shape[-1] + != self.manifest["pose_schemas"][ep["embodiment"]]["dimension"] + ): + raise ValueError("Expected image[T,V,P,D] and measured pose[T,Dm]") + f, cameras, patches, _ = image.shape + if patches != np.prod(self.manifest["dino_identity"]["patch_grid"]): + raise ValueError("DINO patch grid differs from cached image tokens") + if data["image_xy"].shape != (patches, 2) or data["pose_valid"].shape != ( + ep["frames"], + ): + raise ValueError( + "Prepared geometry/mask shapes differ from token arrays" + ) + pose_valid = np.asarray(data["pose_valid"][indices], bool) + if not np.isfinite(image).all() or not np.isfinite(pose[pose_valid]).all(): + raise ValueError("Non-finite valid DINO/proprio values") + if self.normalization is not None: + im, scale = self.normalization["image"] + pm, pscale = self.normalization[f"pose/{ep['embodiment']}"] + if image.shape[-1] != len(im): + raise ValueError("Normalization schema width mismatch: image") + image, pose = (image - im) / scale, (pose - pm) / pscale + # Absent measurements cannot contaminate token projections/losses. + pose[~pose_valid] = 0 + if self.task == "direct_action": + task_kind = "direct_action" + observed = {name: np.ones(f, bool) for name in ("image", "proprio")} + elif rng.random() < self.forecasting_probability: + task_kind = "forecast" + known = np.arange(f) < int(rng.integers(1, f)) + observed = { + "image": known.copy(), + "proprio": known & (rng.random(f) < 0.8), + } + else: + task_kind = "completion" + observed = {name: rng.random(f) < 0.5 for name in ("image", "proprio")} + if not observed["image"].any(): + observed["image"][rng.integers(f)] = True + if observed["image"].all(): + observed["image"][rng.integers(f)] = False + + def base(value, repeats, observed_frames): + n = f * repeats + return dict( + value=value.reshape(n, -1), + time=np.repeat(time, repeats), + index=np.repeat(np.arange(f), repeats).astype(np.float32), + valid=np.ones(n, bool), + target_valid=np.ones(n, bool), + observed=np.repeat(observed_frames, repeats), + ) + + streams = { + "image": base(image, cameras * patches, observed["image"]), + "proprio": base(pose, 1, observed["proprio"]), + } + xy = np.asarray(data["image_xy"][:], np.float32) + if not np.isfinite(xy).all(): + raise ValueError("Non-finite image patch coordinates") + streams["image"].update( + xy=np.tile(xy, (f * cameras, 1)), + camera=np.tile(np.repeat(np.arange(cameras), patches), f), + ) + streams["proprio"]["target_valid"] = pose_valid + streams["proprio"]["observed"] &= pose_valid + if self.task == "direct_action": + action = self._read_actions(ep, indices, full_time) + streams["action"] = base(action, 1, np.zeros(f, bool)) + sample = { + "embodiment": ep["embodiment"], + "task_kind": task_kind, + "world": { + name: { + key: torch.from_numpy(np.array(value, copy=True)) + for key, value in stream.items() + } + for name, stream in streams.items() + }, + } + if draw is not None: + sample["sampling_cursor"] = draw[0] + 1 + return sample diff --git a/egomimic/rldb/dino_proprio_datamodule.py b/egomimic/rldb/dino_proprio_datamodule.py new file mode 100644 index 000000000..9abef00a3 --- /dev/null +++ b/egomimic/rldb/dino_proprio_datamodule.py @@ -0,0 +1,212 @@ +"""DINO/proprio objectives through the shared prepared-data training boundary.""" + +import hashlib + +import torch +from torch.utils.data import DataLoader + +from egomimic.pl_utils.data_context import DataContext +from egomimic.rldb.dino_proprio import ( + EvenlySpacedEpisodeBatchSampler, + PreparedDinoProprioDataset, + collate_dino_proprio, +) +from egomimic.rldb.world_model import HomogeneousWorldBatchSampler +from egomimic.rldb.world_model_datamodule import WorldModelDataModule +from egomimic.rldb.world_model_preparation import ArtifactVerificationMemo + + +class DinoProprioDataModule(WorldModelDataModule): + """Reuse the checkpointable source sampler, never the three-stream dataset.""" + + def __init__( + self, + manifest_path, + task="visual_diffusion", + action_manifest_path=None, + batch_size=1, + num_workers=0, + mode="continuous", + seed=42, + min_frames=2, + max_frames=4, + discrete_frames=4, + stride=1, + forecasting_probability=0.5, + max_past_seconds=None, + max_future_seconds=None, + checkpointable_sampling=False, + validation_batches_per_embodiment=None, + validation_anchors_per_episode=None, + ): + if task not in ("visual_diffusion", "direct_action"): + raise ValueError("Unknown DINO/proprio objective") + if task == "direct_action" and not action_manifest_path: + raise ValueError( + "Direct action BC requires a verified action alignment manifest" + ) + if task == "visual_diffusion" and action_manifest_path is not None: + raise ValueError("Visual/proprio prediction does not consume action labels") + if validation_anchors_per_episode is not None: + if ( + type(validation_anchors_per_episode) is not int + or validation_anchors_per_episode < 1 + ): + raise ValueError( + "validation_anchors_per_episode must be a positive integer or None" + ) + if validation_batches_per_embodiment is not None: + raise ValueError( + "Per-episode validation anchors and prefix batch caps are mutually exclusive" + ) + super().__init__( + manifest_path, + batch_size, + num_workers, + mode, + seed, + min_frames, + max_frames, + discrete_frames, + stride, + forecasting_probability, + max_past_seconds, + max_future_seconds, + checkpointable_sampling, + validation_batches_per_embodiment, + ) + self.validation_anchors_per_episode = validation_anchors_per_episode + self.settings.update(task=task) + if action_manifest_path is not None: + self.settings["action_manifest_path"] = str(action_manifest_path) + + def setup(self, stage=None): + if not self.datasets: + verifier = ArtifactVerificationMemo() + self.datasets = { + split: PreparedDinoProprioDataset( + self.manifest_path, + split=split, + verification_memo=verifier, + **self.settings, + ) + for split in ("train", "val") + } + + def prepare_context(self, *, mode, normalization, normalizer, restored_state=None): + del mode + if normalizer is not None or normalization not in (None, "prepared"): + raise ValueError("DINO/proprio caches carry bound train-only normalization") + self.setup() + dataset = self.datasets["train"] + state = { + "kind": "dino-proprio-prepared-v1", + "manifest": dataset.manifest, + "manifest_sha256": hashlib.sha256(dataset.path.read_bytes()).hexdigest(), + "sampling": self.settings, + } + if dataset.action_manifest is not None: + state["action_manifest"] = dataset.action_manifest + state["action_manifest_sha256"] = hashlib.sha256( + dataset.action_path.read_bytes() + ).hexdigest() + if self.checkpointable_sampling: + state["sampling_resume"] = { + "schema": "mot-sampling-state-v1", + "batch_size": self.batch_size, + } + if self.validation_batches_per_embodiment is not None: + state["validation_sampling"] = { + "batches_per_embodiment": self.validation_batches_per_embodiment + } + if self.validation_anchors_per_episode is not None: + state["validation_sampling"] = { + "anchors_per_episode": self.validation_anchors_per_episode, + "selection": "deterministic_evenly_spaced_eligible_anchors", + } + if restored_state is not None and restored_state != state: + raise ValueError("DINO/proprio checkpoint/data context mismatch") + return DataContext( + normalizer=None, + sample_schema={ + "world": ( + "image_proprio_action_labels" + if dataset.task == "direct_action" + else "image_proprio" + ) + }, + validation_groups=tuple( + sorted({ep["embodiment"] for ep in dataset.episodes}) + ), + state=state, + ) + + def _loader(self, split): + self.setup() + if self.trainer is not None and self.trainer.world_size != 1: + raise ValueError( + "Distributed source sharding is not implemented; use one device" + ) + if split == "val" and self.validation_anchors_per_episode is not None: + sampler = EvenlySpacedEpisodeBatchSampler( + self.datasets[split], + self.validation_anchors_per_episode, + self.batch_size, + ) + else: + sampler = HomogeneousWorldBatchSampler( + self.datasets[split], + self.batch_size, + self.settings["seed"], + shuffle=split == "train", + checkpointable=self.checkpointable_sampling and split == "train", + validation_batches_per_embodiment=( + self.validation_batches_per_embodiment if split == "val" else None + ), + ) + if split == "train": + self.train_sampler = sampler + sampler.consumed_batches = self.consumed_batches + generator = ( + torch.Generator().manual_seed(self.settings["seed"]) + if self.checkpointable_sampling + else None + ) + return DataLoader( + self.datasets[split], + batch_sampler=sampler, + collate_fn=collate_dino_proprio, + num_workers=self.num_workers, + generator=generator, + ) + + def state_dict(self): + state = super().state_dict() + if state: + state["validation_anchors_per_episode"] = ( + self.validation_anchors_per_episode + ) + if state and "action_manifest_path" in self.settings: + from pathlib import Path + + state["action_manifest_sha256"] = hashlib.sha256( + Path(self.settings["action_manifest_path"]).read_bytes() + ).hexdigest() + return state + + def load_state_dict(self, state_dict): + if ( + self.checkpointable_sampling + and state_dict.get("validation_anchors_per_episode") + != self.validation_anchors_per_episode + ): + raise ValueError( + "DINO/proprio sampling restore mismatch: validation_anchors_per_episode" + ) + if state_dict.get("action_manifest_sha256") != self.state_dict().get( + "action_manifest_sha256" + ): + raise ValueError( + "DINO/proprio sampling restore mismatch: action_manifest_sha256" + ) + super().load_state_dict(state_dict) diff --git a/egomimic/rldb/pixel_proprio.py b/egomimic/rldb/pixel_proprio.py new file mode 100644 index 000000000..21f386ef3 --- /dev/null +++ b/egomimic/rldb/pixel_proprio.py @@ -0,0 +1,890 @@ +"""Native RGB patches with optional proprioception and verified command labels. + +RGB remains uint8 in its prepared immutable cache. Selected frames are converted +to [-1, 1] and rearranged into full raw-pixel patches at load time. This adapter +does not resize images, average patches, or consume encoder features or tracks. +""" + +from __future__ import annotations + +import json +import re +from contextlib import ExitStack +from dataclasses import dataclass +from itertools import islice +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from egomimic.rldb.world_model import ( + HomogeneousWorldBatchSampler, + _physical_times, + open_prepared_episode, +) +from egomimic.rldb.world_model_preparation import ( + ArtifactVerificationMemo, + training_identities, +) + + +def collate_pixel_proprio(samples): + """Pad declared streams independently, preserving per-sample task labels.""" + if not samples or len({s["embodiment"] for s in samples}) != 1: + raise ValueError("Batch must be nonempty and homogeneous in embodiment") + names = set(samples[0]["world"]) + if names not in ({"image"}, {"image", "proprio"}, {"image", "proprio", "action"}): + raise ValueError("Expected image stream with optional proprio/action labels") + if any(set(s["world"]) != names for s in samples): + raise ValueError("Cannot mix objectives in a batch") + result = {} + for name in sorted(names): + size = max(len(s["world"][name]["value"]) for s in samples) + fields = set(samples[0]["world"][name]) + if any(set(s["world"][name]) != fields for s in samples): + raise ValueError("Inconsistent stream fields") + result[name] = {} + for key in fields: + values = [] + for sample in samples: + value = sample["world"][name][key] + target = torch.zeros((size, *value.shape[1:]), dtype=value.dtype) + target[: len(value)] = value + values.append(target) + result[name][key] = torch.stack(values) + kinds = [s.get("task_kind") for s in samples] + allowed = {"forecast", "completion", "direct_action"} + if any(kind not in allowed for kind in kinds): + raise ValueError( + "Every sample requires its explicit observation-mask task kind" + ) + embodiment = samples[0]["embodiment"] + item = dict(embodiment=embodiment, world=result, task_kind=kinds) + if any("pixel_metadata" in s for s in samples): + if not all("pixel_metadata" in s for s in samples): + raise ValueError( + "Pixel provenance metadata must be present for every sample" + ) + item["pixel_metadata"] = [s["pixel_metadata"] for s in samples] + if any("sampling_cursor" in s for s in samples): + cursors = {s.get("sampling_cursor") for s in samples} + if None in cursors or len(cursors) != 1: + raise ValueError("A resumed batch must have one source-draw cursor") + item["sampling_cursor"] = cursors.pop() + if any("validation_batch_index" in s for s in samples): + indices = [s.get("validation_batch_index") for s in samples] + if ( + any(type(index) is not int or index < 0 for index in indices) + or len(set(indices)) != 1 + ): + raise ValueError("A validation batch requires one global batch index") + item["validation_batch_index"] = indices + return {embodiment: item} + + +@dataclass(frozen=True) +class PixelValidationIndex: + """Keep the original validation draw identity after rank sharding.""" + + index: int + batch_index: int + + +class DistributedPixelTrainingBatchSampler(HomogeneousWorldBatchSampler): + """Split each deterministic global draw into rank-local homogeneous batches. + + The inherited source distribution remains sampling with replacement. Ranks + consume distinct draw slots, not copies of rank zero's draw. A slot may + independently draw the same anchor as another slot, as in the original + sampler, but its frame/mask RNG remains independently keyed by global slot. + """ + + def __init__( + self, + dataset, + batch_size, + seed, + *, + world_size, + rank, + balance_training_work=False, + ): + _validate_rank(world_size, rank) + self.world_size, self.rank = world_size, rank + self.per_rank_batch_size = batch_size + self.balance_training_work = balance_training_work + super().__init__( + dataset, + batch_size * world_size, + seed, + shuffle=True, + checkpointable=True, + ) + + def __iter__(self): + start = self.rank * self.per_rank_batch_size + for batch in super().__iter__(): + if self.balance_training_work: + yield balance_global_pixel_draws( + self.dataset, batch, self.per_rank_batch_size, self.world_size + )[self.rank] + else: + yield batch[start : start + self.per_rank_batch_size] + + +def balance_global_pixel_draws(dataset, draws, batch_size, world_size): + """Assign the same global draw keys to fixed-size ranks by estimated work. + + Longest-first assignment is used only when it lowers the largest rank's + summed squared token count. In particular B=1 and equal-length discrete + batches retain their original assignment. No sample is added, removed, or + moved to another optimizer update. Model noise remains rank-seeded, so a + reassigned clip has the same source frames/masks but different model noise. + """ + if len(draws) != batch_size * world_size: + raise ValueError("Work balancing requires one complete global batch") + original = [draws[i * batch_size : (i + 1) * batch_size] for i in range(world_size)] + if batch_size == 1 or world_size == 1: + return original + work = {draw: dataset.predicted_token_work(draw[0], draw[1:]) for draw in draws} + assigned = [[] for _ in range(world_size)] + totals = [0] * world_size + for draw in sorted(draws, key=lambda draw: (-work[draw], draw[2])): + rank = min( + (rank for rank in range(world_size) if len(assigned[rank]) < batch_size), + key=lambda rank: (totals[rank], rank), + ) + assigned[rank].append(draw) + totals[rank] += work[draw] + if max(totals) >= max(sum(work[draw] for draw in batch) for batch in original): + return original + return [sorted(batch, key=lambda draw: draw[2]) for batch in assigned] + + +def _validate_rank(world_size, rank): + if ( + type(world_size) is not int + or world_size < 1 + or type(rank) is not int + or not 0 <= rank < world_size + ): + raise ValueError("Invalid distributed sampling topology") + + +class DistributedPixelValidationBatchSampler: + """Contiguous exact shards, without dropped or padded validation examples. + + Equal local lengths prevent DDP forward collectives from hanging at epoch + boundaries. An incompatible total fails explicitly instead of duplicating + validation samples. The unsharded loader is still used for offline audits. + """ + + def __init__(self, sampler, *, world_size, rank): + _validate_rank(world_size, rank) + self.sampler, self.world_size, self.rank = sampler, world_size, rank + self.global_batches = len(sampler) + if not self.global_batches or self.global_batches % world_size: + raise ValueError( + "Validation batch count must be positive and divisible by world_size " + "for exact distributed coverage without padding or dropping" + ) + + def __len__(self): + return self.global_batches // self.world_size + + def __iter__(self): + start = self.rank * len(self) + for batch_index, batch in enumerate( + islice(iter(self.sampler), start, start + len(self)), start=start + ): + yield [PixelValidationIndex(int(index), batch_index) for index in batch] + + +class EvenlySpacedEpisodeBatchSampler: + """Deterministic held-out anchor coverage for every episode and embodiment.""" + + def __init__(self, dataset, anchors_per_episode, batch_size=1): + if dataset.split != "val": + raise ValueError("Episode-spanning subsets are for validation only") + if ( + type(anchors_per_episode) is not int + or anchors_per_episode < 1 + or type(batch_size) is not int + or batch_size < 1 + ): + raise ValueError("Positive integer anchor and batch counts required") + self.groups, self.batch_size = {}, batch_size + offset = 0 + for ep, count in zip(dataset.episodes, dataset.counts): + selected = np.linspace( + 0, count - 1, min(anchors_per_episode, count), dtype=np.int64 + ) + self.groups.setdefault(ep["embodiment"], []).extend( + (selected + offset).tolist() + ) + offset += count + + def __len__(self): + return sum( + (len(v) + self.batch_size - 1) // self.batch_size + for v in self.groups.values() + ) + + def __iter__(self): + for indices in self.groups.values(): + for start in range(0, len(indices), self.batch_size): + yield indices[start : start + self.batch_size] + + +def sample_full_episode_times(length, rng, min_frames=2, max_frames=128, anchor=None): + """Uniform frame-count draw and uniform other frames over the whole episode. + + The globally sampled anchor is always retained. No fixed time window or + nearby-frame preference is imposed. Recorded frames are not interpolated. + """ + if ( + any(type(v) is not int for v in (length, min_frames, max_frames)) + or min_frames < 2 + or max_frames < min_frames + or length < min_frames + ): + raise ValueError("Invalid full-episode frame count") + if anchor is None: + anchor = int(rng.integers(length)) + if not 0 <= anchor < length: + raise ValueError("Invalid full-episode anchor") + count = int(rng.integers(min_frames, min(max_frames, length) + 1)) + candidates = np.arange(length) + other = rng.choice(candidates[candidates != anchor], count - 1, replace=False) + return np.sort(np.r_[anchor, other]).astype(np.int64) + + +def sample_discrete_window(length, rng, frames, stride, anchor): + """A complete raw-frame window containing anchor, grouped without overlap. + + Images use each group's first frame. Each proprio/action token contains all + ``stride`` consecutive recorded states/commands in that group. + """ + if ( + any(type(v) is not int for v in (length, frames, stride, anchor)) + or frames < 2 + or stride < 1 + or length < frames * stride + or not 0 <= anchor < length + ): + raise ValueError("Invalid complete discrete window") + span = frames * stride + lower, upper = max(0, anchor - span + 1), min(anchor, length - span) + start = int(rng.integers(lower, upper + 1)) + return np.arange(start, start + span, dtype=np.int64).reshape(frames, stride) + + +def _sha(value): + return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None + + +def _coordinate_schema(schema, label): + dimension, fields = schema.get("dimension"), schema.get("fields", []) + if ( + type(dimension) is not int + or dimension < 1 + or len(fields) != dimension + or any(not isinstance(v, str) or not v for v in fields) + or len(set(fields)) != dimension + ): + raise ValueError( + f"{label} requires an explicit unique field for every coordinate" + ) + + +def read_selected_rgb_frames(rgb, indices): + """Read complete selected frames without HDF5's costly point selection. + + Prepared HDF5 RGB is chunked by frame and camera. Basic integer reads keep + those chunks intact; stacking preserves the exact requested frame order and + dtype. NPZ members are NumPy arrays and retain their existing indexing path. + """ + if isinstance(rgb, np.ndarray): + return rgb[indices] + import h5py + + if not isinstance(rgb, h5py.Dataset): + raise TypeError("Expected an HDF5 RGB dataset or a NumPy RGB array") + indices = np.asarray(indices) + if indices.ndim != 1 or (indices.size and indices.dtype.kind not in "iu"): + raise TypeError("Selected RGB frames must be a one-dimensional integer array") + if not indices.size: + return np.empty((0, *rgb.shape[1:]), dtype=rgb.dtype) + return np.stack([rgb[int(index)] for index in indices], axis=0) + + +def patchify_rgb(rgb, patch_size): + """Return [T,V,patches,3*P*P] with row-major patches and HWC within each patch.""" + if rgb.dtype != np.uint8 or rgb.ndim != 5 or rgb.shape[-1] != 3: + raise ValueError("Pixel input must be uint8 RGB[T,V,H,W,3]") + if type(patch_size) is not int or patch_size < 1: + raise ValueError("patch_size must be a positive integer") + frames, cameras, height, width, channels = rgb.shape + if height % patch_size or width % patch_size: + raise ValueError("Image dimensions must be divisible by patch_size") + patches = rgb.reshape( + frames, + cameras, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + channels, + ) + patches = patches.transpose(0, 1, 2, 4, 3, 5, 6).reshape( + frames, cameras, height // patch_size * (width // patch_size), 3 * patch_size**2 + ) + return patches.astype(np.float32) / np.float32(127.5) - np.float32(1) + + +def load_pose_normalization(manifest, root, verifier): + contract = manifest.get("normalization", {}) + if ( + contract.get("scope") != "train_only" + or contract.get("arrays_normalized") is not False + or not contract.get("path") + or not _sha(contract.get("sha256")) + ): + raise ValueError( + "Pixel caches require bound train-only normalization of raw poses" + ) + path = root / contract["path"] + verifier.verify(path, contract["sha256"]) + saved = json.loads(path.read_text()) + required = {f"pose/{emb}" for emb in manifest["pose_schemas"]} + if ( + saved.get("schema") != "pixel-proprio-normalization-v1" + or saved.get("scope") != "train_only" + or saved.get("train_episodes") != training_identities(manifest["episodes"]) + or set(saved.get("statistics", {})) != required + ): + raise ValueError( + "Pixel normalization must contain only poses fitted to exact training episodes" + ) + result = {} + for key, item in saved["statistics"].items(): + mean, std = ( + np.asarray(item["mean"], np.float32), + np.asarray(item["std"], np.float32), + ) + dim = manifest["pose_schemas"][key[5:]]["dimension"] + if ( + mean.shape != (dim,) + or std.shape != (dim,) + or item.get("count", 0) < 1 + or not np.isfinite(mean).all() + or not np.isfinite(std).all() + or np.any(std <= 0) + ): + raise ValueError(f"Invalid pose normalization: {key}") + result[key] = mean, std + return result + + +class PreparedPixelProprioDataset(Dataset): + """Image-only/conditioned visual prediction or current-frame causal action BC. + + Image-only batches retain the same frozen manifest and RGB verification but + never open the optional pose/command caches or pose normalization artifact. + """ + + def __init__( + self, + manifest_path, + split="train", + task="visual_diffusion", + mode="continuous", + seed=42, + min_frames=2, + max_frames=128, + discrete_frames=4, + stride=1, + forecasting_probability=0.5, + max_past_seconds=None, + max_future_seconds=None, + verification_memo=None, + sampling_strategy="full_episode_uniform", + image_only=False, + ): + if task not in ( + "visual_diffusion", + "visual_state_action", + "direct_action", + ) or split not in ( + "train", + "val", + ): + raise ValueError( + "Require a declared pixel objective and explicit train/val split" + ) + if type(image_only) is not bool or (image_only and task != "visual_diffusion"): + raise ValueError("image_only requires visual_diffusion and a boolean flag") + self.image_only = image_only + self.path, self.task, self.split, self.seed, self.mode = ( + Path(manifest_path), + task, + split, + seed, + mode, + ) + self.manifest = json.loads(self.path.read_text()) + manifest = self.manifest + if manifest.get("schema") != "pixel-proprio-prepared-v1": + raise ValueError("Unsupported pixel/proprio prepared schema") + identity = manifest.get("pixel_identity", {}) + size, patch = identity.get("image_size"), identity.get("patch_size") + if ( + not isinstance(size, list) + or len(size) != 2 + or any(type(v) is not int or v < 1 for v in size) + or type(patch) is not int + or patch < 1 + or any(v % patch for v in size) + or identity.get("channels") != 3 + or identity.get("patch_order") != "row_major_hwc" + or identity.get("normalization") != "uint8_to_minus_one_one" + ): + raise ValueError( + "Require explicit divisible RGB dimensions, patch size/order, and pixel normalization" + ) + self.image_size, self.patch_size = tuple(size), patch + self.grid = tuple(v // patch for v in size) + self.image_dim = 3 * patch**2 + yy, xx = np.meshgrid( + (np.arange(self.grid[0]) + 0.5) / self.grid[0] * 2 - 1, + (np.arange(self.grid[1]) + 0.5) / self.grid[1] * 2 - 1, + indexing="ij", + ) + self.image_xy = np.stack((xx, yy), -1).reshape(-1, 2).astype(np.float32) + if not manifest.get("episodes"): + raise ValueError("Require explicit episode identities") + if not image_only: + if not manifest.get("pose_schemas"): + raise ValueError("Require measured pose schemas") + for schema in manifest["pose_schemas"].values(): + _coordinate_schema(schema, "Pose schema") + self.verifier = verification_memo or ArtifactVerificationMemo() + ids, paths, rgbs = set(), set(), set() + self.action_alignment = {} + for ep in manifest["episodes"]: + key = (ep["embodiment"], ep["episode_id"]) + rgb = (self.path.parent / ep["rgb_path"]).resolve() + if key in ids or rgb in rgbs: + raise ValueError( + "Duplicate pixel episode identity or physical source across splits" + ) + ids.add(key) + rgbs.add(rgb) + if ( + ep.get("split") not in ("train", "val") + or type(ep.get("frames")) is not int + or ep["frames"] < 1 + or not ep.get("cameras") + or len(set(ep["cameras"])) != len(ep["cameras"]) + or not _sha(ep.get("rgb_sha256")) + ): + raise ValueError( + "Pixel episode requires exact RGB identities, cameras and split" + ) + if not image_only: + path = (self.path.parent / ep["path"]).resolve() + if path in paths: + raise ValueError("Duplicate measured-pose source across splits") + paths.add(path) + if ( + ep["embodiment"] not in manifest["pose_schemas"] + or ep.get("frame") != "fixed_reference" + or ep.get("pose_kind") != "measured_pose" + or not _sha(ep.get("sha256")) + ): + raise ValueError("Require exact measured-pose identity and frame") + self.verifier.verify( + path, + ep["sha256"], + trusted_fingerprint=ep.get("verified_fingerprint"), + ) + self.verifier.verify( + rgb, + ep["rgb_sha256"], + trusted_fingerprint=ep.get("rgb_verified_fingerprint"), + ) + if task in ("direct_action", "visual_state_action"): + self._validate_actions(ep) + self.normalization = ( + {} + if image_only + else load_pose_normalization(manifest, self.path.parent, self.verifier) + ) + self.episodes = [ep for ep in manifest["episodes"] if ep["split"] == split] + if not self.episodes: + raise ValueError("No episodes in selected split") + if ( + (sampling_strategy, mode) + not in ( + ("full_episode_uniform", "continuous"), + ("discrete_window", "continuous"), + ("discrete_window", "discrete"), + ) + or max_past_seconds is not None + or max_future_seconds is not None + ): + raise ValueError( + "Pixel sampling requires continuous full_episode_uniform or discrete_window and no time bounds" + ) + if ( + type(min_frames) is not int + or type(max_frames) is not int + or min_frames < 2 + or max_frames < min_frames + ): + raise ValueError("Invalid full-episode frame count") + if not 0 <= forecasting_probability <= 1: + raise ValueError("Invalid forecasting probability") + if ( + type(discrete_frames) is not int + or discrete_frames < 2 + or type(stride) is not int + or stride < 1 + ): + raise ValueError("Invalid discrete frame count or stride") + if task == "direct_action" and sampling_strategy != "full_episode_uniform": + raise ValueError( + "Direct action baseline requires single-frame continuous observations" + ) + self.forecasting_probability = forecasting_probability + self.sampling = dict(min_frames=min_frames, max_frames=max_frames) + self.sampling_strategy = sampling_strategy + self.discrete_frames, self.stride = discrete_frames, stride + self.chunk_size = stride if mode == "discrete" else 1 + self.anchor_indices = [] + for ep in self.episodes: + if task == "direct_action": + anchors = np.arange(ep["frames"]) + else: + anchors = ( + np.arange(ep["frames"]) + if ep["frames"] + >= ( + discrete_frames * stride + if sampling_strategy == "discrete_window" + else min_frames + ) + else np.array([], np.int64) + ) + self.anchor_indices.append(anchors) + self.counts = [len(indices) for indices in self.anchor_indices] + if any(c < 1 for c in self.counts): + raise ValueError("Episode too short for temporal recipe") + self.ends = np.cumsum(self.counts) + self.rng = self.worker_seed = None + + def _validate_actions(self, ep): + schema = self.manifest.get("action_schemas", {}).get(ep["embodiment"]) + provenance, source = ( + ep.get("action_provenance", {}), + ep.get("native_source", {}), + ) + if ( + schema is None + or not provenance.get("path") + or not _sha(provenance.get("sha256")) + ): + raise ValueError( + "Direct action BC requires verified real command labels for every embodiment/episode" + ) + _coordinate_schema(schema, "Action schema") + if ( + schema.get("semantics") != "action_command" + or schema.get("value_space") != "native_commands" + ): + raise ValueError("Action schemas must declare native action commands") + path = self.path.parent / provenance["path"] + self.verifier.verify(path, provenance["sha256"]) + proof = json.loads(path.read_text()) + if ( + proof.get("schema") != "pixel-proprio-action-alignment-v1" + or proof.get("relation") != "action_t_applied_to_state_t" + or not _sha(source.get("source_demo_sha256")) + or proof.get("source_demo_sha256") != source["source_demo_sha256"] + or not source.get("episode") + or proof.get("source_episode") != source["episode"] + or any( + not _sha(proof.get(k)) + for k in ( + "producer_sha256", + "raw_export_metadata_sha256", + "raw_payload_sha256", + ) + ) + or any( + not isinstance(proof.get("basis", {}).get(k), str) + or not proof["basis"][k].strip() + for k in ( + "cache_index_source", + "action_index_source", + "state_action_relation_source", + ) + ) + ): + raise ValueError( + "Incomplete or conflicting pixel action alignment provenance" + ) + n = ep["frames"] + indices = { + k: np.asarray(proof.get(k)) + for k in ("cache_frame_index", "source_state_index", "source_action_index") + } + if ( + any(v.shape != (n,) or v.dtype.kind not in "iu" for v in indices.values()) + or not np.array_equal(indices["cache_frame_index"], np.arange(n)) + or not np.array_equal( + indices["source_state_index"], indices["source_action_index"] + ) + or np.any(indices["source_state_index"] < 0) + or not (np.diff(indices["source_state_index"]) > 0).all() + ): + raise ValueError("Pixel action/state/cache index alignment mismatch") + with open_prepared_episode(self.path.parent / ep["path"]) as data: + if ( + "action" not in data + or "action_valid" not in data + or data["action"].shape != (n, schema["dimension"]) + or data["action_valid"].shape != (n,) + ): + raise ValueError("Pixel cache lacks declared real action arrays") + values, valid = data["action"][:], data["action_valid"][:] + if ( + valid.dtype != np.bool_ + or not valid.all() + or not np.isfinite(values).all() + ): + raise ValueError("Direct action labels must all be finite and valid") + self.action_alignment[(ep["embodiment"], ep["episode_id"])] = proof + + def __len__(self): + return int(self.ends[-1]) + + def _episode_anchor(self, index): + if not 0 <= index < len(self): + raise IndexError(index) + ei = int(np.searchsorted(self.ends, index, side="right")) + offset = 0 if ei == 0 else int(self.ends[ei - 1]) + return self.episodes[ei], int(self.anchor_indices[ei][index - offset]) + + def _sample_frame_chunks(self, length, anchor, rng): + """One RNG contract shared by actual data reads and metadata-only work preview.""" + if self.task == "direct_action": + return np.array([[anchor]]) + if self.sampling_strategy == "discrete_window": + chunks = sample_discrete_window( + length, rng, self.discrete_frames, self.stride, anchor + ) + return chunks[:, :1] if self.mode == "continuous" else chunks + return sample_full_episode_times(length, rng, anchor=anchor, **self.sampling)[ + :, None + ] + + def predicted_token_work(self, index, draw): + """Pure preview of image-attention work for the exact checkpointed draw. + + The preview reads manifest metadata only and never advances a worker's + RNG. Squared image-token count estimates attention cost; it is not a + measured GPU time or a reason to change the selected frames. + """ + if ( + self.split != "train" + or len(draw) != 2 + or any(type(value) is not int or value < 0 for value in draw) + ): + raise ValueError("Work preview requires a training cursor/global-slot draw") + episode, anchor = self._episode_anchor(index) + rng = np.random.default_rng(np.random.SeedSequence([self.seed, *draw])) + frames = len(self._sample_frame_chunks(episode["frames"], anchor, rng)) + return (frames * len(episode["cameras"]) * self.grid[0] * self.grid[1]) ** 2 + + def __getitem__(self, index): + draw = None + validation_batch_index = None + if isinstance(index, PixelValidationIndex): + if self.split != "val" or index.batch_index < 0: + raise ValueError("Global validation indices apply only to validation") + validation_batch_index, index = index.batch_index, index.index + if isinstance(index, tuple): + index, cursor, slot = index + draw = int(cursor), int(slot) + ep, anchor = self._episode_anchor(index) + worker_seed = torch.initial_seed() + if self.rng is None or self.worker_seed != worker_seed: + self.rng = np.random.default_rng((worker_seed + self.seed) % 2**63) + self.worker_seed = worker_seed + rng = ( + self.rng + if self.split == "train" + else np.random.default_rng(self.seed + index) + ) + if draw is not None: + if self.split != "train": + raise ValueError("Checkpointed draw keys apply only to training") + rng = np.random.default_rng(np.random.SeedSequence([self.seed, *draw])) + with ExitStack() as stack: + rgb_data = stack.enter_context( + open_prepared_episode(self.path.parent / ep["rgb_path"]) + ) + if any(k not in rgb_data for k in ("rgb", "time")): + raise ValueError("Missing required raw RGB arrays") + full_time = rgb_data["time"][:] + _physical_times(full_time, ep["frames"]) + if not self.image_only: + data = stack.enter_context( + open_prepared_episode(self.path.parent / ep["path"]) + ) + if any(k not in data for k in ("time", "pose", "pose_valid")): + raise ValueError("Missing required measured proprioception arrays") + if not np.array_equal(full_time, data["time"][:]): + raise ValueError( + "RGB and measured-proprioception timestamps disagree" + ) + if rgb_data["rgb"].dtype != np.uint8 or rgb_data["rgb"].shape != ( + ep["frames"], + len(ep["cameras"]), + *self.image_size, + 3, + ): + raise ValueError( + "Native RGB shape/dtype differs from declared pixel geometry/cameras" + ) + chunk_indices = self._sample_frame_chunks(ep["frames"], anchor, rng) + indices = chunk_indices[:, 0] + raw_indices = chunk_indices.reshape(-1) + time = np.asarray(full_time[indices], np.float64) + origin = 0 if self.task == "direct_action" else int(rng.integers(len(time))) + physical_times = time.tolist() + physical_origin = float(time[origin]) + time = (time - time[origin]).astype(np.float32) + image = patchify_rgb( + read_selected_rgb_frames(rgb_data["rgb"], indices), self.patch_size + ) + f, cameras, patches, _ = image.shape + if not self.image_only: + dim = self.manifest["pose_schemas"][ep["embodiment"]]["dimension"] + if data["pose"].shape != (ep["frames"], dim) or data[ + "pose_valid" + ].shape != (ep["frames"],): + raise ValueError( + "Measured proprioception shape differs from schema" + ) + pose, pose_valid = ( + np.asarray(data["pose"][raw_indices], np.float32), + data["pose_valid"][raw_indices], + ) + if ( + pose_valid.dtype != np.bool_ + or not np.isfinite(pose[pose_valid]).all() + ): + raise ValueError( + "Measured proprioception validity or values are invalid" + ) + mean, std = self.normalization[f"pose/{ep['embodiment']}"] + pose = (pose - mean) / std + pose[~pose_valid] = 0 + pose = pose.reshape(f, self.chunk_size * dim) + pose_valid = pose_valid.reshape(f, self.chunk_size).all(axis=1) + # Preserve the established RNG consumption even when proprio is + # omitted: completion's image-mask repairs and later draws must + # select identical RGB frames/masks across conditioning variants. + if self.task == "direct_action": + task_kind = "direct_action" + observed = {name: np.ones(f, bool) for name in ("image", "proprio")} + elif rng.random() < self.forecasting_probability: + task_kind = "forecast" + known = np.arange(f) < int(rng.integers(1, f)) + observed = { + "image": known.copy(), + "proprio": known & (rng.random(f) < 0.8), + } + else: + task_kind = "completion" + observed = {name: rng.random(f) < 0.5 for name in ("image", "proprio")} + if not observed["image"].any(): + observed["image"][rng.integers(f)] = True + if observed["image"].all(): + observed["image"][rng.integers(f)] = False + + def base(value, repeats, known): + n = f * repeats + return dict( + value=value.reshape(n, -1), + time=np.repeat(time, repeats), + index=np.repeat(np.arange(f), repeats).astype(np.float32), + valid=np.ones(n, bool), + target_valid=np.ones(n, bool), + observed=np.repeat(known, repeats), + ) + + streams = {"image": base(image, cameras * patches, observed["image"])} + streams["image"].update( + xy=np.tile(self.image_xy, (f * cameras, 1)), + camera=np.tile(np.repeat(np.arange(cameras), patches), f), + ) + if not self.image_only: + streams["proprio"] = base(pose, 1, observed["proprio"]) + streams["proprio"]["target_valid"] = pose_valid + streams["proprio"]["observed"] &= pose_valid + if self.task in ("direct_action", "visual_state_action"): + action = np.asarray(data["action"][raw_indices], np.float32).reshape( + f, -1 + ) + streams["action"] = base(action, 1, np.zeros(f, bool)) + if task_kind == "forecast": + # Past command queries can see later observed context. They + # remain real valid labels, but are not causal BC targets. + streams["action"]["target_valid"] &= ~observed["image"] + sample = dict( + embodiment=ep["embodiment"], + task_kind=task_kind, + pixel_metadata=dict( + episode_id=ep["episode_id"], + camera_names=list(ep["cameras"]), + frame_indices=indices.tolist(), + timestamps_seconds=physical_times, + time_origin_seconds=physical_origin, + chunk_size=self.chunk_size, + sampling_strategy=self.sampling_strategy, + image_size=list(self.image_size), + patch_size=self.patch_size, + normalization="uint8_to_minus_one_one", + patch_order="row_major_hwc", + ), + world={ + name: { + key: torch.from_numpy(np.array(value, copy=True)) + for key, value in stream.items() + } + for name, stream in streams.items() + }, + ) + if not self.image_only: + sample["pixel_metadata"].update( + pose_base_dimension=dim, + proprio_frame_indices=chunk_indices.tolist(), + proprio_timestamps_seconds=full_time[chunk_indices].tolist(), + ) + if draw is not None: + sample["sampling_cursor"] = draw[0] + 1 + if validation_batch_index is not None: + sample["validation_batch_index"] = validation_batch_index + return sample + + +__all__ = [ + "PreparedPixelProprioDataset", + "collate_pixel_proprio", + "patchify_rgb", + "EvenlySpacedEpisodeBatchSampler", +] diff --git a/egomimic/rldb/pixel_proprio_datamodule.py b/egomimic/rldb/pixel_proprio_datamodule.py new file mode 100644 index 000000000..838287deb --- /dev/null +++ b/egomimic/rldb/pixel_proprio_datamodule.py @@ -0,0 +1,363 @@ +"""Configurable native-pixel/proprio data through the shared training boundary.""" + +import hashlib + +import torch +from torch.utils.data import DataLoader + +from egomimic.pl_utils.data_context import DataContext +from egomimic.rldb.pixel_proprio import ( + DistributedPixelTrainingBatchSampler, + DistributedPixelValidationBatchSampler, + EvenlySpacedEpisodeBatchSampler, + PreparedPixelProprioDataset, + collate_pixel_proprio, +) +from egomimic.rldb.world_model import HomogeneousWorldBatchSampler +from egomimic.rldb.world_model_datamodule import WorldModelDataModule +from egomimic.rldb.world_model_preparation import ArtifactVerificationMemo + + +class PixelProprioDataModule(WorldModelDataModule): + def __init__( + self, + manifest_path, + task="visual_diffusion", + batch_size=1, + num_workers=0, + mode="continuous", + seed=42, + min_frames=2, + max_frames=128, + discrete_frames=4, + stride=1, + forecasting_probability=0.5, + max_past_seconds=None, + max_future_seconds=None, + checkpointable_sampling=False, + validation_batches_per_embodiment=None, + validation_anchors_per_episode=None, + sampling_strategy="full_episode_uniform", + validation_sampling_strategy=None, + validation_discrete_frames=None, + validation_stride=None, + image_only=False, + distributed_world_size=1, + validation_batch_size=None, + pin_memory=False, + persistent_workers=False, + prefetch_factor=None, + balance_training_work=False, + ): + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be a positive integer") + if type(num_workers) is not int or num_workers < 0: + raise ValueError("num_workers must be a nonnegative integer") + if validation_batch_size is not None and ( + type(validation_batch_size) is not int or validation_batch_size < 1 + ): + raise ValueError("validation_batch_size must be a positive integer or None") + if type(pin_memory) is not bool or type(persistent_workers) is not bool: + raise ValueError("pin_memory and persistent_workers must be booleans") + if persistent_workers and num_workers == 0: + raise ValueError("persistent_workers requires num_workers > 0") + if prefetch_factor is not None and ( + type(prefetch_factor) is not int or prefetch_factor < 1 or num_workers == 0 + ): + raise ValueError( + "prefetch_factor requires a positive integer and num_workers > 0" + ) + if type(distributed_world_size) is not int or distributed_world_size < 1: + raise ValueError("distributed_world_size must be a positive integer") + if distributed_world_size > 1 and not checkpointable_sampling: + raise ValueError( + "Distributed pixel loading requires checkpointable_sampling" + ) + if type(balance_training_work) is not bool or ( + balance_training_work and not checkpointable_sampling + ): + raise ValueError( + "balance_training_work requires a boolean and checkpointable_sampling" + ) + self.balance_training_work = balance_training_work + self.distributed_world_size = distributed_world_size + if task not in ("visual_diffusion", "visual_state_action", "direct_action"): + raise ValueError("Unknown pixel/proprio objective") + if type(image_only) is not bool or (image_only and task != "visual_diffusion"): + raise ValueError("image_only requires visual_diffusion and a boolean flag") + if validation_anchors_per_episode is not None: + if ( + type(validation_anchors_per_episode) is not int + or validation_anchors_per_episode < 1 + ): + raise ValueError( + "validation_anchors_per_episode must be a positive integer or None" + ) + if validation_batches_per_embodiment is not None: + raise ValueError( + "Per-episode validation anchors and prefix batch caps are mutually exclusive" + ) + super().__init__( + manifest_path, + batch_size, + # The legacy parent forbids workers for exact resume. Pixel draws + # instead key every frame/mask RNG by (cursor, global draw slot), so + # worker scheduling and discarded prefetched batches cannot change + # the first unconsumed draw. Leave the old dataset guard untouched. + 0, + mode, + seed, + min_frames, + max_frames, + discrete_frames, + stride, + forecasting_probability, + max_past_seconds, + max_future_seconds, + checkpointable_sampling, + validation_batches_per_embodiment, + ) + self.num_workers = num_workers + self.pin_memory = pin_memory + self.persistent_workers = persistent_workers + self.prefetch_factor = prefetch_factor + self._explicit_validation_batch_size = validation_batch_size is not None + self.validation_batch_size = ( + batch_size if validation_batch_size is None else validation_batch_size + ) + self.settings.update(task=task, sampling_strategy=sampling_strategy) + # Preserve historical conditioned checkpoint contexts when disabled. + if image_only: + self.settings["image_only"] = True + self.validation_anchors_per_episode = validation_anchors_per_episode + self.validation_settings = { + key: value + for key, value in dict( + sampling_strategy=validation_sampling_strategy, + discrete_frames=validation_discrete_frames, + stride=validation_stride, + ).items() + if value is not None + } + + def setup(self, stage=None): + if not self.datasets: + verifier = ArtifactVerificationMemo() + self.datasets = { + split: PreparedPixelProprioDataset( + self.manifest_path, + split=split, + verification_memo=verifier, + **( + self.settings + | (self.validation_settings if split == "val" else {}) + ), + ) + for split in ("train", "val") + } + + def configure_evaluation(self, requirements): + if requirements.complete_episodes or set(requirements.required_keys) - { + "world", + "embodiment", + "task_kind", + "pixel_metadata", + }: + raise ValueError("Unsupported pixel/proprio evaluation requirements") + + def prepare_context(self, *, mode, normalization, normalizer, restored_state=None): + del mode + if normalizer is not None or normalization not in (None, "prepared"): + raise ValueError( + "Pixel adapter owns train-only pose normalization and fixed RGB scaling" + ) + self.setup() + dataset = self.datasets["train"] + state = dict( + kind="pixel-image-prepared-v1" + if dataset.image_only + else "pixel-proprio-prepared-v1", + manifest=dataset.manifest, + manifest_sha256=hashlib.sha256(dataset.path.read_bytes()).hexdigest(), + sampling=self.settings, + validation_overrides=self.validation_settings, + ) + if self.checkpointable_sampling: + state["sampling_resume"] = { + "schema": "mot-sampling-state-v1", + "batch_size": self.batch_size, + } + if self.distributed_world_size > 1: + state["distributed_sampling"] = self._distributed_sampling_identity() + if self._explicit_validation_batch_size: + state["batching"] = self._batching_identity() + if self.balance_training_work: + state["training_work_balance"] = self._work_balance_identity() + if self.validation_batches_per_embodiment is not None: + state["validation_sampling"] = { + "batches_per_embodiment": self.validation_batches_per_embodiment + } + if self.validation_anchors_per_episode is not None: + state["validation_sampling"] = { + "anchors_per_episode": self.validation_anchors_per_episode, + "selection": "deterministic_evenly_spaced_eligible_anchors", + } + if restored_state is not None and restored_state != state: + raise ValueError("Pixel/proprio checkpoint/data context mismatch") + return DataContext( + normalizer=None, + sample_schema={ + "world": "image" + if dataset.image_only + else "image_proprio_action_labels" + if dataset.task in ("direct_action", "visual_state_action") + else "image_proprio" + }, + validation_groups=tuple( + sorted({ep["embodiment"] for ep in dataset.episodes}) + ), + state=state, + ) + + def _loader(self, split): + self.setup() + rank = 0 + if self.trainer is not None: + if self.trainer.world_size != self.distributed_world_size: + raise ValueError( + "Trainer world_size differs from bound distributed_world_size" + ) + rank = self.trainer.global_rank + batch_size = self.batch_size if split == "train" else self.validation_batch_size + if split == "val" and self.validation_anchors_per_episode is not None: + sampler = EvenlySpacedEpisodeBatchSampler( + self.datasets[split], + self.validation_anchors_per_episode, + batch_size, + ) + elif split == "train" and self.distributed_world_size > 1: + sampler = DistributedPixelTrainingBatchSampler( + self.datasets[split], + batch_size, + self.settings["seed"], + world_size=self.distributed_world_size, + rank=rank, + balance_training_work=self.balance_training_work, + ) + else: + sampler = HomogeneousWorldBatchSampler( + self.datasets[split], + batch_size, + self.settings["seed"], + shuffle=split == "train", + checkpointable=self.checkpointable_sampling and split == "train", + validation_batches_per_embodiment=self.validation_batches_per_embodiment + if split == "val" + else None, + ) + # Before Trainer attachment, inspection/reload must see all held-out + # anchors. Runtime sharding preserves their original global IDs. + if ( + split == "val" + and self.trainer is not None + and self.distributed_world_size > 1 + ): + sampler = DistributedPixelValidationBatchSampler( + sampler, world_size=self.distributed_world_size, rank=rank + ) + if split == "train": + self.train_sampler = sampler + sampler.consumed_batches = self.consumed_batches + generator = ( + torch.Generator().manual_seed(self.settings["seed"]) + if self.checkpointable_sampling + else None + ) + return DataLoader( + self.datasets[split], + batch_sampler=sampler, + collate_fn=collate_pixel_proprio, + num_workers=self.num_workers, + generator=generator, + pin_memory=self.pin_memory, + persistent_workers=self.persistent_workers, + prefetch_factor=self.prefetch_factor, + ) + + def state_dict(self): + state = super().state_dict() + if state: + state["validation_anchors_per_episode"] = ( + self.validation_anchors_per_episode + ) + state["validation_overrides"] = self.validation_settings + if self.distributed_world_size > 1: + state["distributed_sampling"] = self._distributed_sampling_identity() + if self._explicit_validation_batch_size: + state["batching"] = self._batching_identity() + if self.balance_training_work: + state["training_work_balance"] = self._work_balance_identity() + return state + + def _batching_identity(self): + # Only explicit split-specific batching adds a new identity field. + # Historical configs keep their exact bound data/checkpoint contexts. + # Worker/pinning/prefetch options are recorded in the resolved config, + # but do not change a stateless source draw or its data identity. + return { + "schema": "pixel-split-batching-v1", + "train_batch_size": self.batch_size, + "validation_batch_size": self.validation_batch_size, + } + + def _distributed_sampling_identity(self): + return { + "schema": "pixel-distributed-sampling-v1", + "world_size": self.distributed_world_size, + "per_rank_batch_size": self.batch_size, + "global_batch_size": self.batch_size * self.distributed_world_size, + "training_sharding": "global_draw_rank_slots", + "validation_sharding": "contiguous_global_batches_no_padding", + } + + def _work_balance_identity(self): + return { + "schema": "pixel-global-batch-work-v1", + "strategy": "longest_first_fixed_rank_capacity_if_lower_maximum", + "estimate": "sampled_image_tokens_squared", + "preserved": "global_draw_keys_frames_and_conditioning_masks", + "model_noise": "rank_seeded_realization_changes_for_reassigned_clips", + } + + def load_state_dict(self, state_dict): + expected_balance = ( + self._work_balance_identity() if self.balance_training_work else None + ) + if state_dict.get("training_work_balance") != expected_balance: + raise ValueError("Pixel sampling restore mismatch: training work balance") + expected_batching = ( + self._batching_identity() if self._explicit_validation_batch_size else None + ) + if ( + self.checkpointable_sampling + and state_dict.get("batching") != expected_batching + ): + raise ValueError( + "Pixel sampling restore mismatch: train/validation batching" + ) + expected_topology = ( + self._distributed_sampling_identity() + if self.distributed_world_size > 1 + else None + ) + if state_dict.get("distributed_sampling") != expected_topology: + raise ValueError("Pixel sampling restore mismatch: distributed topology") + if self.checkpointable_sampling and ( + state_dict.get("validation_anchors_per_episode") + != self.validation_anchors_per_episode + or state_dict.get("validation_overrides", {}) != self.validation_settings + ): + raise ValueError( + "Pixel sampling restore mismatch: validation_anchors_per_episode" + ) + super().load_state_dict(state_dict) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index b1e605c32..648fe201c 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -14,7 +14,10 @@ from lightning.fabric.utilities.cloud_io import _load as pl_load from lightning.fabric.utilities.cloud_io import get_filesystem from lightning.pytorch.loggers import Logger -from lightning.pytorch.plugins.environments import SLURMEnvironment +from lightning.pytorch.plugins.environments import ( + SLURMEnvironment, + TorchElasticEnvironment, +) from omegaconf import DictConfig, OmegaConf, open_dict from tabulate import tabulate @@ -160,7 +163,13 @@ def _instantiate_trainer_plugins(cfg: DictConfig) -> List[Any]: plugins: List[Any] = [] if cfg.get("mmap_checkpoint", True): plugins.append(MmapCheckpointIO()) - if os.environ.get("SLURM_JOB_ID"): + environment = OmegaConf.select(cfg, "runtime.cluster_environment", default="auto") + if environment not in {"auto", "torch_elastic"}: + raise ValueError("runtime.cluster_environment must be auto or torch_elastic") + if environment == "torch_elastic": + plugins.append(_torch_elastic_environment(cfg)) + print("TORCH ELASTIC ENVIRONMENT ENABLED") + elif os.environ.get("SLURM_JOB_ID"): slurm_environment = _slurm_environment(cfg) plugins.append(slurm_environment) print( @@ -170,6 +179,38 @@ def _instantiate_trainer_plugins(cfg: DictConfig) -> List[Any]: return plugins +def _torch_elastic_environment(cfg: DictConfig) -> TorchElasticEnvironment: + """Use torchrun ranks even when a Slurm step owns the parent process.""" + owner = OmegaConf.select(cfg, "runtime.slurm_requeue_owner", default="lightning") + if owner != "none": + raise ValueError("torch_elastic requires runtime.slurm_requeue_owner=none") + _slurm_auto_requeue(cfg) # Reject inherited, conflicting runner ownership. + if not TorchElasticEnvironment.detect(): + raise RuntimeError("torch_elastic requires a process launched by torchrun") + integers = {} + for name in ("RANK", "LOCAL_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "GROUP_RANK"): + value = os.environ.get(name, "") + try: + integers[name] = int(value) + except ValueError as exc: + raise ValueError(f"torch_elastic requires integer {name}") from exc + world, local_world = integers["WORLD_SIZE"], integers["LOCAL_WORLD_SIZE"] + if (world < 1 or local_world < 1 or world % local_world + or not 0 <= integers["RANK"] < world + or not 0 <= integers["LOCAL_RANK"] < local_world + or not 0 <= integers["GROUP_RANK"] < world // local_world): + raise ValueError("torch_elastic has inconsistent world size or rank variables") + if not os.environ.get("MASTER_ADDR"): + raise ValueError("torch_elastic requires MASTER_ADDR") + try: + port = int(os.environ.get("MASTER_PORT", "")) + except ValueError as exc: + raise ValueError("torch_elastic requires integer MASTER_PORT") from exc + if not 1 <= port <= 65535: + raise ValueError("torch_elastic MASTER_PORT must be between 1 and 65535") + return TorchElasticEnvironment() + + def _build_model_config_tree(cfg: DictConfig) -> DictConfig: model_cfg = OmegaConf.create(OmegaConf.to_container(cfg.model, resolve=True)) config_tree = {"model": model_cfg} diff --git a/scripts/train/log_mirror.sh b/scripts/train/log_mirror.sh new file mode 100644 index 000000000..f461b6621 --- /dev/null +++ b/scripts/train/log_mirror.sh @@ -0,0 +1,78 @@ +#!/bin/bash +# Source this reusable helper from a training launcher on Linux. +# Requires bash, flock, GNU timeout, rsync, and standard POSIX file utilities. +# +# mirror_once LOCAL_LOG_DIR DESTINATION [ATTEMPTS=3] +# Final flush: bounded retries, nonzero on failure. The caller owns its exit policy. +# watch_log_mirror CHILD_PID LOCAL_LOG_DIR DESTINATION ADMISSION_FILE [INTERVAL=60] +# Periodic best-effort copies after admission. Never signals the training process. +# Run in the background; reap it, then call mirror_once after training exits. +# +# Local files and existing destination history are retained. W&B spool is copied; +# only the disposable root wandb-cache directory and internal lock/temp files are +# excluded. Health is updated atomically locally after each transfer, so a copied +# health receipt describes the preceding attempt, not proof of this transfer. + +_mirror_health() { + local source=${1:?} code=${2:?} previous="$1/.mirror-health.json" + local last_success=0 failures=0 now temporary + if [[ -f "$previous" ]]; then + last_success=$(sed -n 's/.*"last_success_unix":\([0-9][0-9]*\).*/\1/p' "$previous") + failures=$(sed -n 's/.*"consecutive_failures":\([0-9][0-9]*\).*/\1/p' "$previous") + fi + [[ "$last_success" =~ ^[0-9]+$ ]] || last_success=0 + [[ "$failures" =~ ^[0-9]+$ ]] || failures=0 + now=$(date +%s) + if [[ "$code" == 0 ]]; then last_success=$now; failures=0; else failures=$((failures + 1)); fi + temporary=$(mktemp "$source/.mirror-health.tmp.XXXXXX") || return + if ! printf '{"schema":"training-log-mirror-health-v1","checked_at_unix":%s,"last_success_unix":%s,"consecutive_failures":%s,"last_exit_code":%s}\n' \ + "$now" "$last_success" "$failures" "$code" > "$temporary"; then + rm -f "$temporary" + return 1 + fi + mv -f "$temporary" "$previous" +} + +# Default is the final flush: at most three bounded attempts, requiring exit 0. +# The periodic watcher passes attempts=1 and retries on its next interval. +mirror_once() ( + local source=${1:?} destination=${2:?} attempts=${3:-3} + local attempt code=1 + [[ -d "$source" ]] || { printf 'TRAINING LOG MIRROR ERROR: local log directory missing\n' >&2; return 2; } + [[ "$attempts" =~ ^[123]$ ]] || { printf 'TRAINING LOG MIRROR ERROR: attempts must be 1..3\n' >&2; return 2; } + # The lock belongs to this subshell for the entire copy/retry/health update. + # It is outside the transfer and cannot be copied into its own destination. + exec 9> "$source/.mirror.lock" || return + flock --wait 330 9 || { printf 'TRAINING LOG MIRROR WARNING: lock wait expired\n' >&2; return 75; } + for ((attempt=1; attempt<=attempts; attempt++)); do + code=0 + # Bound mkdir as well as rsync: either can stall on the shared filesystem. + timeout --kill-after=15s 300s bash -c ' + mkdir -p "$2" && exec rsync -a --no-owner --no-group --delay-updates \ + --exclude=/wandb-cache/ --exclude=/.mirror.lock --exclude=/.mirror-health.tmp.\* \ + -- "$1/" "$2/" + ' training-log-mirror "$source" "$destination" || code=$? + if ! _mirror_health "$source" "$code"; then + printf 'TRAINING LOG MIRROR ERROR: could not write local health receipt\n' >&2 + return 74 + fi + [[ "$code" == 0 ]] && return 0 + printf 'TRAINING LOG MIRROR WARNING: attempt %s/%s exited %s; local files retained\n' "$attempt" "$attempts" "$code" >&2 + [[ "$attempt" == "$attempts" ]] || sleep 2 + done + return "$code" +) + +watch_log_mirror() { + local child=${1:?} source=${2:?} destination=${3:?} admission=${4:?} + # Transient periodic errors do not signal the trainer or poison future attempts. + local interval=${5:-60} + while kill -0 "$child" 2>/dev/null; do + sleep "$interval" + kill -0 "$child" 2>/dev/null || break + [[ -f "$admission" ]] || continue + if ! mirror_once "$source" "$destination" 1; then + printf 'TRAINING LOG MIRROR WARNING: periodic copy incomplete; retrying next interval, training continues\n' >&2 + fi + done +} diff --git a/scripts/world_model/benchmark_pixel_execution.py b/scripts/world_model/benchmark_pixel_execution.py new file mode 100644 index 000000000..d809df63a --- /dev/null +++ b/scripts/world_model/benchmark_pixel_execution.py @@ -0,0 +1,655 @@ +"""Bounded single-GPU benchmark of exact image-only pixel execution candidates. + +Inputs are synthetic and device-resident. This measures neither a data loader nor +NCCL, and does not replace real-data train/validation/checkpoint smoke tests. +Every candidate runs in a fresh process; missing kernels, OOMs, correctness +failures and compilation failures are reported, never replaced by a fallback. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import importlib.metadata +import json +import math +import os +import signal +import statistics +import subprocess +import sys +import time +import traceback +from contextlib import contextmanager +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +PARAMETERS = 88_178_688 +BACKENDS = ("flash4", "flash3", "sdpa") +SCENARIOS = ("fixed16", "mixed", "worst128") +COMPILE_MODES = ("none", "default", "reduce-overhead", "max-autotune") +TOKEN_EXECUTIONS = ("padded", "packed_valid") +MIXED_FRAMES = (2, 16, 64, 128) + + +class CorrectnessError(RuntimeError): + pass + + +class UnsupportedBackend(RuntimeError): + pass + + +def require(condition, message): + if not condition: + raise CorrectnessError(message) + + +def write_once(path, value): + with Path(path).open("x") as stream: + stream.write(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") + + +def digest(path): + with Path(path).open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def frame_schedule(scenario, batch_size): + """A complete mixed cycle gives every batch size the same frame histogram.""" + if scenario == "fixed16": + return [(16,) * batch_size] + if scenario == "worst128": + return [(128,) * batch_size] + if scenario != "mixed": + raise ValueError(f"Unknown scenario: {scenario}") + cycle = len(MIXED_FRAMES) // math.gcd(batch_size, len(MIXED_FRAMES)) + return [tuple(MIXED_FRAMES[(step * batch_size + row) % len(MIXED_FRAMES)] + for row in range(batch_size)) for step in range(cycle)] + + +def complete_cycles(steps, schedule): + return math.ceil(steps / len(schedule)) * len(schedule) + + +def model_arguments(backend, scenario, *, compile_mode="none", compile_dynamic=False, + token_execution="padded"): + if backend not in BACKENDS or scenario not in SCENARIOS or token_execution not in TOKEN_EXECUTIONS: + raise ValueError("Unknown execution candidate") + return dict( + objective="visual_only", phase_a_condition_on_proprio=False, + pose_dims=None, action_dims=None, image_dim=768, width=768, + depth=8, heads=6, mlp_ratio=4, max_cameras=2, + max_tokens=3200 if scenario == "fixed16" else 25600, + time_mode="discrete" if scenario == "fixed16" else "continuous", + use_camera_embedding=True, use_modality_embedding=False, + use_embodiment_embedding=False, allow_patch_compression=False, + use_phase_embedding=False, noise_conditioning="adaln_zero", + position_encoding="spatiotemporal_rope", rope_axis_dims=(64, 32, 32), + spatial_grid_size=(10, 10), + attention_backend=backend, attention_dtype="bfloat16", strict_fp32=False, + attention_query_chunk_size=0, activation_checkpointing=False, + loss_reduction="mean_per_example", compile_mode=compile_mode, + compile_dynamic=compile_dynamic, compile_backend="inductor", token_execution=token_execution, + ) + + +def make_batch(frames, *, seed, device, image_size=160, patch_size=16, cameras=2, + frame_seconds=0.05): + """True RGB patch geometry; only batch padding is invalid.""" + if not frames or any(type(length) is not int or length < 2 for length in frames): + raise ValueError("Every synthetic episode needs at least two frames") + if image_size % patch_size or image_size < patch_size: + raise ValueError("Image size must be divisible by patch size") + side = image_size // patch_size + patches, dimension = side * side, patch_size * patch_size * 3 + per_frame, tokens = cameras * patches, max(frames) * cameras * patches + generator = torch.Generator().manual_seed(seed) + value = torch.zeros(len(frames), tokens, dimension, dtype=torch.float32) + valid = torch.zeros(len(frames), tokens, dtype=torch.bool) + observed = torch.zeros_like(valid) + clock = torch.zeros(len(frames), tokens, dtype=torch.float32) + index = torch.zeros_like(clock) + camera = torch.arange(cameras).repeat_interleave(patches).repeat(max(frames)) + axis = (torch.arange(side, dtype=torch.float32) + 0.5) * 2 / side - 1 + y, x = torch.meshgrid(axis, axis, indexing="ij") + xy = torch.stack((x.flatten(), y.flatten()), -1).repeat(max(frames) * cameras, 1) + for row, length in enumerate(frames): + size = length * per_frame + value[row, :size] = torch.rand(size, dimension, generator=generator) * 2 - 1 + valid[row, :size] = True + observed[row, :max(1, length // 2) * per_frame] = True + indices = torch.arange(length, dtype=torch.float32).repeat_interleave(per_frame) + index[row, :size] = indices + clock[row, :size] = indices * frame_seconds + return {"image": {key: tensor.to(device) for key, tensor in dict( + value=value, valid=valid, observed=observed, time=clock, index=index, + camera=camera[None].expand(len(frames), -1).clone(), + xy=xy[None].expand(len(frames), -1, -1).clone(), + ).items()}} + + +def single_example(batch, row): + indices = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] + return {"image": {key: value[row:row + 1].index_select(1, indices) + for key, value in batch["image"].items()}} + + +def fixed_corruption(batch, seed): + image = batch["image"] + generator = torch.Generator(device=image["value"].device).manual_seed(seed) + noise = torch.randn(image["value"].shape, dtype=torch.float32, + device=image["value"].device, generator=generator) + levels = (0.15 + 0.7 * torch.remainder(image["index"], 7) / 6).masked_fill( + image["observed"] | ~image["valid"], 0 + ) + noisy = ((1 - levels[..., None]) * image["value"] + levels[..., None] * noise) + return noisy.masked_fill(~image["valid"][..., None], 0), levels + + +def example_mean_loss(prediction, batch): + image = batch["image"] + mask = image["valid"] & ~image["observed"] + counts = mask.sum(1) + require(bool((counts > 0).all()), "Every example needs a query target") + error = (prediction.float() - image["value"].float()).square().mean(-1) + return (error.masked_fill(~mask, 0).sum(1) / counts).mean() + + +@contextmanager +def execution_context(backend, device): + # The SDPA candidate explicitly enables only its fused Flash kernel. + # A backend that cannot handle the mask/shape raises instead of using math. + with torch.autocast(device.type, dtype=torch.bfloat16): + if backend == "sdpa": + from torch.nn.attention import SDPBackend, sdpa_kernel + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + yield + else: + yield + + +def relative_error(actual, expected): + error = (actual.float() - expected.float()).double() + return dict(squared_error=float(error.square().sum()), + squared_reference=float(expected.double().square().sum()), + max_abs=float(error.abs().max())) + + +@contextmanager +def nontrivial_conditioning_probe(model, seed): + """Exercise attention and noise conditioning beyond DiT's zero initialization. + + Only correctness probes use these deterministic nonzero weights. Restore + every changed parameter and RNG before timing actual zero-initialized + training; otherwise zero gates/head would make parity checks vacuous. + """ + if getattr(model, "noise_conditioning", "additive") != "adaln_zero": + yield {"nonzero_adaln_probe": False} + return + selected = {name: parameter for name, parameter in model.named_parameters() + if name.startswith(("image_out.", "final_adaln.")) + or ".adaln_image." in name} + require(selected and any("adaln_image" in name for name in selected), + "AdaLN correctness probe could not find modulation parameters") + original = {name: parameter.detach().clone() for name, parameter in selected.items()} + device = next(model.parameters()).device + devices = [device.index] if device.type == "cuda" else [] + try: + with torch.random.fork_rng(devices=devices), torch.no_grad(): + # Seed only the probed device; do not disturb other CUDA generators. + torch.random.default_generator.manual_seed(seed) + if devices: + torch.cuda.default_generators[device.index].manual_seed(seed) + for parameter in selected.values(): + parameter.normal_(mean=0.0, std=0.02) + yield {"nonzero_adaln_probe": True, "seed": seed, + "restores_training_initialization": True} + finally: + with torch.no_grad(): + for name, parameter in selected.items(): + parameter.copy_(original[name]) + model.zero_grad(set_to_none=True) + + +def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): + """Compare a batch to independent, unpadded B1 calls at fixed noise levels.""" + device = batch["image"]["value"].device + noisy, levels = fixed_corruption(batch, seed) + model.zero_grad(set_to_none=True) + with execution_context(backend, device): + prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] + loss = example_mean_loss(prediction, batch) + require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), "Nonfinite/non-FP32 probe loss") + actual = prediction.detach() + batched_gradients = None + if gradients: + loss.backward() + require(all(p.grad is not None for p in model.parameters()), "Unused trainable parameter") + batched_gradients = {name: p.grad.detach().cpu().clone() for name, p in model.named_parameters()} + del prediction, loss + model.zero_grad(set_to_none=True) + errors = [] + batch_size = actual.shape[0] + for row in range(batch_size): + selected = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] + independent = single_example(batch, row) + with execution_context(backend, device): + predicted = model(independent, "libero_gripper", + noisy[row:row + 1].index_select(1, selected), + levels[row:row + 1].index_select(1, selected), phase="A")["image"] + individual_loss = example_mean_loss(predicted, independent) / batch_size + errors.append(relative_error(actual[row:row + 1].index_select(1, selected), predicted.detach())) + if gradients: + individual_loss.backward() + del predicted, individual_loss + output_relative = math.sqrt(sum(v["squared_error"] for v in errors) / + max(sum(v["squared_reference"] for v in errors), 1e-30)) + require(math.isfinite(output_relative) and output_relative <= tolerance, + f"Batched output differs from independent B1: relative L2={output_relative}") + report = {"examples": batch_size, "output_relative_l2": output_relative, + "output_max_abs": max(v["max_abs"] for v in errors)} + if gradients: + error_sum, reference_sum = 0.0, 0.0 + for name, parameter in model.named_parameters(): + require(parameter.grad is not None and parameter.grad.dtype == torch.float32, + "Missing/non-FP32 reference gradient") + measured = relative_error(batched_gradients[name], parameter.grad.detach().cpu()) + error_sum += measured["squared_error"] + reference_sum += measured["squared_reference"] + gradient_relative = math.sqrt(error_sum / max(reference_sum, 1e-30)) + require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Batched gradient differs from mean independent B1: relative L2={gradient_relative}") + report["gradient_relative_l2"] = gradient_relative + model.zero_grad(set_to_none=True) + # Perturb every other example, preserving row zero and its noise exactly. + if batch_size > 1: + changed = {"image": {key: value.clone() for key, value in batch["image"].items()}} + changed["image"]["value"][1:] *= -0.7 + perturbed = noisy.clone() + perturbed[1:] *= -0.7 + with torch.no_grad(), execution_context(backend, device): + # Inductor may select different kernels for grad-enabled and + # no-grad execution. Compare intervention/control under the same + # mode and graph; separately bound the mode's BF16 rounding drift. + control = model(batch, "libero_gripper", noisy, levels, phase="A")["image"][0].clone() + independent = model(changed, "libero_gripper", perturbed, levels, phase="A")["image"][0] + mode_error = relative_error(control, actual[0]) + mode_relative = math.sqrt(mode_error["squared_error"] / max(mode_error["squared_reference"], 1e-30)) + require(math.isfinite(mode_relative) and mode_relative <= tolerance, + f"Grad/no-grad output differs beyond BF16 tolerance: relative L2={mode_relative}") + report["grad_vs_no_grad_relative_l2"] = mode_relative + error = relative_error(independent, control) + leak = math.sqrt(error["squared_error"] / max(error["squared_reference"], 1e-30)) + require(math.isfinite(leak) and leak <= 1e-6, + f"Cross-example attention leakage detected: same-mode relative L2={leak}") + report["cross_example_relative_l2"] = leak + else: + report["cross_example_check"] = "not_applicable_batch_one" + return report + + +def check_eager_compiled(eager, compiled, batch, *, backend, seed, tolerance): + """Require state-identical eager/compiled AMP forward and FP32 backward parity.""" + require(eager.token_execution == compiled.token_execution, + "Eager/compiled token execution differs") + reference_state, candidate_state = eager.state_dict(), compiled.state_dict() + require(reference_state.keys() == candidate_state.keys(), "Eager/compiled state keys differ") + require(all(torch.equal(value, candidate_state[name]) for name, value in reference_state.items()), + "Eager/compiled initial state differs") + require(all(p.dtype == torch.float32 for model in (eager, compiled) for p in model.parameters()), + "Eager/compiled stored parameters must be FP32") + noisy, levels = fixed_corruption(batch, seed) + device = batch["image"]["value"].device + + def forward_backward(model): + model.zero_grad(set_to_none=True) + with execution_context(backend, device): + prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] + loss = example_mean_loss(prediction, batch) + require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), + "Eager/compiled loss must be finite FP32") + # Match standard mixed precision: backward is outside forward autocast. + loss.backward() + require(all(p.grad is not None and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() for p in model.parameters()), + "Eager/compiled gradient must be present and finite FP32") + return prediction.detach().cpu(), float(loss.detach()) + + try: + reference_output, reference_loss = forward_backward(eager) + reference_gradients = {name: p.grad.detach().cpu().clone() for name, p in eager.named_parameters()} + eager.zero_grad(set_to_none=True) + output, loss = forward_backward(compiled) + measured = relative_error(output, reference_output) + output_relative = math.sqrt(measured["squared_error"] / max(measured["squared_reference"], 1e-30)) + loss_relative = abs(loss - reference_loss) / max(abs(reference_loss), 1e-30) + require(math.isfinite(output_relative) and output_relative <= tolerance, + f"Compiled output differs from eager: relative L2={output_relative}") + require(math.isfinite(loss_relative) and loss_relative <= tolerance, + f"Compiled loss differs from eager: relative error={loss_relative}") + errors = [relative_error(p.grad.detach().cpu(), reference_gradients[name]) + for name, p in compiled.named_parameters()] + gradient_relative = math.sqrt(sum(value["squared_error"] for value in errors) / + max(sum(value["squared_reference"] for value in errors), 1e-30)) + require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Compiled gradient differs from eager: relative L2={gradient_relative}") + return {"output_relative_l2": output_relative, "output_max_abs": measured["max_abs"], + "loss_relative_error": loss_relative, "eager_loss": reference_loss, + "compiled_loss": loss, "gradient_relative_l2": gradient_relative, + "parameters_and_gradients_dtype": "float32", "backward_autocast": "off", + "same_state_verified": True, "backend": backend, + "token_execution": compiled.token_execution} + finally: + eager.zero_grad(set_to_none=True) + compiled.zero_grad(set_to_none=True) + + +def compiler_counters(): + from torch._dynamo.utils import counters + return {group: {str(k): int(v) for k, v in values.items()} for group, values in counters.items()} + + +def counter_total(counters, group, key=None): + values = counters.get(group, {}) + return sum(values.values()) if key is None else values.get(key, 0) + + +def backend_inventory(backend): + module_names = {"flash4": ("flash_attn.cute.interface",), + "flash3": ("flash_attn_3.flash_attn_interface", "flash_attn_interface"), + "sdpa": ("torch",)}[backend] + try: + for index, module_name in enumerate(module_names): + try: + module = importlib.import_module(module_name) + break + except ModuleNotFoundError as error: + # Match the model's official FA3 package/legacy alias support; + # an installed backend with a missing dependency must fail. + if (index == len(module_names) - 1 + or error.name not in {module_name, module_name.split(".")[0]}): + raise + except ImportError as error: + raise UnsupportedBackend(f"Required backend import failed: {module_name}: {error}") from error + module_path = Path(module.__file__) + versions = {} + for name in ("torch", "flash-attn-4", "flash-attn-3", "flash-attn", "nvidia-cutlass-dsl"): + try: + versions[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + versions[name] = None + return {"module": module_name, "path": str(module_path), "sha256": digest(module_path), + "versions": versions, "sdpa_kernel": "FLASH_ATTENTION_ONLY" if backend == "sdpa" else None} + + +def timed_step(model, batch, optimizer, backend, device): + optimizer.zero_grad(set_to_none=True) + with execution_context(backend, device): + losses = model.training_losses(batch, "libero_gripper", pose_weight=0.0) + require(losses["total"].dtype == losses["image"].dtype == torch.float32, "MSE must remain FP32") + losses["total"].backward() + optimizer.step() + return losses["total"].detach() + + +def run_candidate(request, destination): + from egomimic.models.image_proprio import ImageProprioModel + + if not torch.cuda.is_available(): + raise UnsupportedBackend("A scheduled H100 CUDA device is required; no CPU benchmark fallback") + device = torch.device("cuda", request["device"]) + torch.cuda.set_device(device) + properties = torch.cuda.get_device_properties(device) + if "H100" not in properties.name: + raise UnsupportedBackend(f"This benchmark binds H100 hardware, got {properties.name}") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + torch.manual_seed(request["seed"]) + inventory = backend_inventory(request["backend"]) + arguments = model_arguments(request["backend"], request["scenario"], + compile_mode=request["compile_mode"], compile_dynamic=request["compile_dynamic"], + token_execution=request.get("token_execution", "padded")) + model = ImageProprioModel(**arguments).to(device).train() + require(sum(p.numel() for p in model.parameters()) == PARAMETERS, "Exact architecture parameter count changed") + require(all(p.dtype == torch.float32 for p in model.parameters()), "Stored parameters must be FP32") + before_counters = compiler_counters() + schedule = frame_schedule(request["scenario"], request["batch_size"]) + batches = [make_batch(frames, seed=request["seed"] + index, device=device, + frame_seconds=0.2 if request["scenario"] == "fixed16" else 0.05) + for index, frames in enumerate(schedule)] + with nontrivial_conditioning_probe(model, request["seed"] + 1000) as probe_record: + probe = make_batch((2, 3) if request["batch_size"] > 1 else (2,), seed=request["seed"], device=device) + correctness = {"bounded_gradient_probe": check_outputs( + model, probe, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"], gradients=True)} + if request["compile_mode"] != "none": + # This is a separate eager model, not another invocation of the same + # compiled callable. Keep synthetic noise RNG unchanged by construction. + with torch.random.fork_rng(devices=[device.index]): + eager = ImageProprioModel(**{**arguments, "compile_mode": "none"}).to(device).train() + eager.load_state_dict(model.state_dict(), strict=True) + correctness["compiled_vs_eager"] = check_eager_compiled( + eager, model, probe, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"]) + del eager + del probe + correctness["conditioning_probe"] = probe_record + correctness["full_shape_outputs"] = [] + with torch.no_grad(): + for batch in batches: + correctness["full_shape_outputs"].append(check_outputs( + model, batch, backend=request["backend"], seed=request["seed"], + tolerance=request["tolerance"], gradients=False)) + torch.cuda.empty_cache() # Release unused correctness-probe allocator blocks. + optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01) + warmup = complete_cycles(request["warmup"], schedule) + steps = complete_cycles(request["steps"], schedule) + warm_start = time.perf_counter() + for step in range(warmup): + timed_step(model, batches[step % len(batches)], optimizer, request["backend"], device) + torch.cuda.synchronize(device) + warmup_seconds = time.perf_counter() - warm_start + warm_counters = compiler_counters() + if request["compile_mode"] != "none": + require(counter_total(warm_counters, "stats", "unique_graphs") > counter_total(before_counters, "stats", "unique_graphs"), + "Requested compilation captured no tensor regions") + require(counter_total(warm_counters, "graph_break") == counter_total(before_counters, "graph_break"), + "A compiled tensor region broke its fullgraph contract") + torch.cuda.reset_peak_memory_stats(device) + walls, cuda_times, loss_values = [], [], [] + valid_tokens = 0 + for step in range(steps): + batch = batches[step % len(batches)] + torch.cuda.synchronize(device) + start_event, end_event = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start = time.perf_counter() + start_event.record() + loss = timed_step(model, batch, optimizer, request["backend"], device) + end_event.record() + torch.cuda.synchronize(device) + walls.append(time.perf_counter() - start) + cuda_times.append(start_event.elapsed_time(end_event) / 1000) + loss_values.append(float(loss)) + valid_tokens += sum(schedule[step % len(schedule)]) * 200 + after_counters = compiler_counters() + require(all(math.isfinite(value) for value in loss_values), "Nonfinite measured training loss") + require(counter_total(after_counters, "stats", "unique_graphs") == counter_total(warm_counters, "stats", "unique_graphs"), + "New compiled graphs appeared during steady-state timing") + require(all(p.grad is not None and p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() + for p in model.parameters()), "Missing, nonfinite or non-FP32 measured gradient") + require(all(state[key].dtype == torch.float32 and torch.isfinite(state[key]).all() + for state in optimizer.state.values() for key in ("exp_avg", "exp_avg_sq")), + "Nonfinite or non-FP32 AdamW moments") + require(all(p.dtype == torch.float32 and torch.isfinite(p).all() for p in model.parameters()), + "Nonfinite or non-FP32 optimized parameters") + result = { + "status": "PASS", "candidate": request, "parameters_total_trainable": PARAMETERS, + "model_arguments": arguments, "backend_inventory": inventory, + "gpu": {"name": properties.name, "uuid": str(properties.uuid), "total_memory_bytes": properties.total_memory}, + "correctness": correctness, "frame_schedule": schedule, + "warmup_steps": warmup, "timed_steps": steps, "warmup_seconds": warmup_seconds, + "wall_seconds": walls, "cuda_event_seconds": cuda_times, + "median_wall_seconds": statistics.median(walls), "median_cuda_seconds": statistics.median(cuda_times), + "examples_per_second": steps * request["batch_size"] / sum(walls), + "valid_tokens_per_second": valid_tokens / sum(walls), "valid_tokens_measured": valid_tokens, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), + "resident_input_bytes": sum(value.numel() * value.element_size() + for batch in batches for value in batch["image"].values()), + "losses": loss_values, "compiler_counters_before": before_counters, + "compiler_counters_after_warmup": warm_counters, "compiler_counters_after_timing": after_counters, + "compile_scope": "declared_fullgraph_tensor_regions; attention_adapter_and_validation_eager", + "compiled_backward_autocast": "off" if request["compile_mode"] != "none" else None, + "measurement_scope": "single_gpu_device_resident_synthetic_model; no_loader_H2D_NCCL_or_checkpoint_IO", + } + if request["profile"]: + trace = destination / "trace.json" + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, profile_memory=True) as profiler: + with torch.profiler.record_function("synthetic_device_resident_optimizer_step"): + timed_step(model, batches[0], optimizer, request["backend"], device) + torch.cuda.synchronize(device) + profiler.export_chrome_trace(str(trace)) + result["profile"] = {"path": str(trace), "sha256": digest(trace)} + return result + + +def worker(request_path): + request_path = Path(request_path) + request = json.loads(request_path.read_text()) + try: + result = run_candidate(request, request_path.parent) + except torch.cuda.OutOfMemoryError as error: + result = {"status": "OOM", "error": str(error), "candidate": request} + except UnsupportedBackend as error: + result = {"status": "UNSUPPORTED", "error": str(error), "candidate": request} + except CorrectnessError as error: + result = {"status": "CORRECTNESS_FAILED", "error": str(error), "candidate": request} + except Exception as error: + # Backend/API/compiler failures are never retried with another kernel. + unsupported = (isinstance(error, RuntimeError) and any( + phrase in str(error) for phrase in ("No available kernel", "No viable backend") + )) or (isinstance(error, ValueError) and + str(error) == "Ragged packed_valid execution requires flash3 or flash4") + result = {"status": "UNSUPPORTED" if unsupported else "ERROR", "error": f"{type(error).__name__}: {error}", + "traceback": traceback.format_exc(), "candidate": request} + write_once(request_path.parent / "result.json", result) + + +def csv_choices(value, choices): + values = value.split(",") + if not values or len(values) != len(set(values)) or set(values) - set(choices): + raise argparse.ArgumentTypeError(f"Choose distinct values from {choices}") + return values + + +def run_worker_process(command, folder, timeout): + """Bound compiler subprocesses as well as the GPU-owning worker on timeout.""" + with (folder / "stdout.log").open("x") as out, (folder / "stderr.log").open("x") as err: + process = subprocess.Popen(command, stdout=out, stderr=err, start_new_session=True) + try: + return process.wait(timeout=timeout), False + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass # The bounded child exited between timeout and kill. + process.wait() + return process.returncode, True + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path) + parser.add_argument("--backends", type=lambda s: csv_choices(s, BACKENDS), default=list(BACKENDS)) + parser.add_argument("--scenarios", type=lambda s: csv_choices(s, SCENARIOS), default=list(SCENARIOS)) + parser.add_argument("--batch-sizes", type=lambda s: [int(v) for v in csv_choices(s, ("1", "2", "4", "8"))], default=[1, 2, 4, 8]) + parser.add_argument("--compile-mode", choices=COMPILE_MODES, default="none") + parser.add_argument("--compile-dynamic", action="store_true") + parser.add_argument("--token-execution", choices=TOKEN_EXECUTIONS, default="padded", + help="Keep padded tensors, or pack valid image tokens across the image-only expert") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--steps", type=int, default=10) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--tolerance", type=float, default=0.03) + parser.add_argument("--worker-timeout", type=int, default=1800) + parser.add_argument("--python-bootstrap", type=Path, + help="Run every child through this reviewed runtime bootstrap (.pth hooks included)") + parser.add_argument("--profile-candidate", help="Exact candidate ID such as flash4-mixed-b4-none") + parser.add_argument("--dry-run", action="store_true", help="Print the immutable candidate plan without CUDA work") + parser.add_argument("--worker-request", type=Path, help=argparse.SUPPRESS) + args = parser.parse_args() + if args.worker_request: + worker(args.worker_request) + return + if min(args.warmup, args.steps, args.worker_timeout) < 1 or args.device < 0 or not 0 < args.tolerance < 1: + parser.error("Positive bounded warmup/steps/timeout/tolerance and nonnegative device required") + candidates = [] + for backend in args.backends: + for scenario in args.scenarios: + for batch_size in args.batch_sizes: + identity = f"{backend}-{scenario}-b{batch_size}-{args.compile_mode}" + if args.token_execution != "padded": + identity += f"-{args.token_execution}" + candidates.append(dict(id=identity, backend=backend, scenario=scenario, + batch_size=batch_size, compile_mode=args.compile_mode, + compile_dynamic=args.compile_dynamic, token_execution=args.token_execution, + seed=args.seed, + device=args.device, warmup=args.warmup, steps=args.steps, + tolerance=args.tolerance, profile=args.profile_candidate == identity)) + if args.profile_candidate and not any(c["profile"] for c in candidates): + parser.error("--profile-candidate must identify an enabled candidate") + if args.dry_run: + print(json.dumps({"parameters_total_trainable": PARAMETERS, "candidates": candidates}, indent=2)) + return + if args.output is None: + parser.error("--output is required") + bootstrap = None + if args.python_bootstrap: + bootstrap_path = args.python_bootstrap.resolve(strict=True) + bootstrap = {"path": str(bootstrap_path), "sha256": digest(bootstrap_path)} + args.output.mkdir(parents=True, exist_ok=False) + source_files = [Path(__file__), ROOT / "egomimic/models/image_proprio.py", + ROOT / "egomimic/models/pixel_diffusion.py", + ROOT / "egomimic/models/world_model.py", ROOT / "egomimic/models/attention_execution.py", + ROOT / "egomimic/models/compilation.py"] + source = {str(path.relative_to(ROOT)): digest(path) for path in source_files if path.exists()} + write_once(args.output / "plan.json", {"schema": "pixel-execution-benchmark-v1", "candidates": candidates, + "source_sha256": source, "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "python_bootstrap": bootstrap, "python_executable": sys.executable, + "scope": "synthetic_single_gpu_execution_only; not_training_readiness_or_policy_quality"}) + results = [] + for candidate in candidates: + folder = args.output / candidate["id"] + folder.mkdir() + request_path = folder / "request.json" + write_once(request_path, candidate) + command = [sys.executable] + if args.python_bootstrap: + command.append(str(args.python_bootstrap.resolve())) + command += [str(Path(__file__).resolve()), "--worker-request", str(request_path.resolve())] + returncode, timed_out = run_worker_process(command, folder, args.worker_timeout) + if timed_out: + result = {"status": "TIMEOUT", "seconds": args.worker_timeout, "candidate": candidate} + elif returncode != 0 or not (folder / "result.json").exists(): + result = {"status": "PROCESS_ERROR", "returncode": returncode, "candidate": candidate} + else: + result = json.loads((folder / "result.json").read_text()) + results.append(result) + print(json.dumps({"candidate": candidate["id"], "status": result["status"]}), flush=True) + write_once(args.output / "report.json", {"schema": "pixel-execution-benchmark-v1", "source_sha256": source, + "status": "COMPLETED", "results": results, + "scope": "Only PASS candidates have measured and checked synthetic throughput; real-data DDP smoke still required"}) + + +if __name__ == "__main__": + main() diff --git a/scripts/world_model/export_libero_action_sidecars.py b/scripts/world_model/export_libero_action_sidecars.py new file mode 100644 index 000000000..28d8f4ac6 --- /dev/null +++ b/scripts/world_model/export_libero_action_sidecars.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +"""Publish immutable LIBERO action sidecars after exact source-row/time checks. + +CPU-only. Run where the prepared caches and their native alignment artifacts +exist. The raw-export bundle records the separate source/controller audit; this +tool verifies that evidence and binds its commands to actual cached frames. +""" + +# ruff: noqa: E402 +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import shutil +import sys +from pathlib import Path + +import h5py +import numpy as np + +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from egomimic.rldb.world_model_preparation import ArtifactVerificationMemo + + +def sha256(path): + with Path(path).open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def _json(path): + return json.loads(Path(path).read_text()) + + +def _write_json(path, value): + with Path(path).open("x") as stream: + json.dump(value, stream, indent=2, sort_keys=True, allow_nan=False) + stream.write("\n") + + +def _inside(root, name): + candidate = (root / name).resolve() + if not candidate.is_relative_to(root.resolve()): + raise ValueError("Raw-export artifact paths must stay inside their bundle") + return candidate + + +def _source_hashes(path): + result = {} + for line in path.read_text().splitlines(): + digest, name = line.split(maxsplit=1) + name = name.lstrip(" *") + if not re.fullmatch(r"[0-9a-f]{64}", digest) or name in result: + raise ValueError("Invalid or duplicate raw source hash evidence") + result[name] = digest + return result + + +def _verified_bundle(raw_dir, manifest_path, verifier): + metadata_path = raw_dir / "metadata.json" + metadata = _json(metadata_path) + if ( + metadata.get("schema") != "mot-libero-raw-command-export-v1" + or metadata.get("status") != "PASS" + ): + raise ValueError("Require completed raw command export evidence") + if metadata.get("source_manifest_sha256") != sha256(manifest_path): + raise ValueError("Raw export is bound to a different prepared manifest") + evidence = {"metadata.json": metadata_path} + for key in ("raw_source_hashes", "headers", "controller_metadata", "validation"): + artifact = metadata[key] + path = _inside(raw_dir, artifact["path"]) + verifier.verify(path, artifact["sha256"]) + evidence[artifact["path"]] = path + script = raw_dir / "export-commands.sh" + verifier.verify(script, metadata["export_script_sha256"]) + evidence[script.name] = script + producer = metadata["producer_convention"] + expected_url = ( + f"https://github.com/{producer.get('repository')}/blob/" + f"{producer.get('commit')}/{producer.get('path')}" + ) + if ( + producer.get("repository") != "Lifelong-Robot-Learning/LIBERO" + or not re.fullmatch(r"[0-9a-f]{40}", producer.get("commit", "")) + or producer.get("path") != "scripts/create_dataset.py" + or producer.get("relation") != "action_t_applied_to_state_t" + or producer.get("url") != expected_url + ): + raise ValueError( + "Missing pinned official LIBERO state/action convention evidence" + ) + source = _inside(raw_dir, producer["local_path"]) + verifier.verify(source, producer["sha256"]) + evidence[producer["local_path"]] = source + lines = producer.get("evidence_lines", []) + if not lines or any( + type(i) is not int or not 1 <= i <= len(source.read_text().splitlines()) + for i in lines + ): + raise ValueError("Producer convention evidence lines are invalid") + schema = metadata["action_schema"] + if ( + schema.get("dimension") != 7 + or schema.get("fields") + != [ + "delta_x", + "delta_y", + "delta_z", + "delta_rx", + "delta_ry", + "delta_rz", + "gripper_command", + ] + or schema.get("semantics") != "action_command" + or schema.get("value_space") != "native_commands" + or schema.get("controller_type") != "OSC_POSE" + or schema.get("control_delta") is not True + ): + raise ValueError( + "Require explicit seven-coordinate native delta OSC_POSE command schema" + ) + hashes = _source_hashes(evidence[metadata["raw_source_hashes"]["path"]]) + controllers = _json(evidence[metadata["controller_metadata"]["path"]]) + by_source = {item["source_demo"]: item for item in controllers} + if len(by_source) != len(controllers) or set(by_source) != set(hashes): + raise ValueError( + "Controller evidence must cover each hashed raw demonstration once" + ) + for item in controllers: + config = item["controller_configs"] + if ( + config.get("type") != schema["controller_type"] + or config.get("control_delta") is not True + or item.get("control_freq") != schema["control_frequency_hz"] + or not np.array_equal( + np.broadcast_to(config["input_min"], (7,)), schema["input_min"] + ) + or not np.array_equal( + np.broadcast_to(config["input_max"], (7,)), schema["input_max"] + ) + or config["output_min"] != schema["pose_output_min"] + or config["output_max"] != schema["pose_output_max"] + ): + raise ValueError("Raw controller evidence disagrees with command schema") + payload = _inside(raw_dir, metadata["payload_path"]) + verifier.verify(payload, metadata["payload_sha256"]) + if ( + payload.stat().st_size != metadata["payload_nbytes"] + or metadata["payload_nbytes"] != metadata["expected_payload_nbytes"] + ): + raise ValueError("Raw binary payload size mismatch") + evidence[metadata["payload_path"]] = payload + return metadata, evidence, hashes, payload + + +def _frame_binding(episode, root, verifier): + feature = (root / episode["path"]).resolve() + verifier.verify( + feature, + episode["sha256"], + trusted_fingerprint=episode.get("verified_fingerprint"), + ) + receipt_path = feature.with_suffix(feature.suffix + ".json") + verifier.verify(receipt_path, episode["receipt_sha256"]) + receipt = _json(receipt_path) + if ( + receipt.get("schema") != "mot-world-feature-cache-v1" + or receipt.get("status") != "PASS" + or receipt.get("output_sha256") != episode["sha256"] + or receipt.get("native_source") != episode["native_source"] + or receipt.get("episode_id") != episode["episode_id"] + or receipt.get("embodiment") != episode["embodiment"] + ): + raise ValueError("Feature receipt disagrees with prepared episode") + inputs = receipt["request"]["inputs"] + native = Path(receipt["native_cache"]) + verifier.verify(native / "manifest.json", inputs["native_manifest_sha256"]) + native_manifest = _json(native / "manifest.json") + if ( + native_manifest.get("schema") != "libero-same-recorded-state-v1" + or any(native_manifest.get(k) != v for k, v in episode["native_source"].items()) + or native_manifest.get("artifacts", {}).get("aligned.npz") + != inputs["native_aligned_sha256"] + ): + raise ValueError("Native alignment manifest identity mismatch") + aligned = native / "aligned.npz" + verifier.verify(aligned, inputs["native_aligned_sha256"]) + with np.load(aligned, allow_pickle=False) as data: + rows, native_time = data["source_row"], data["time"] + with h5py.File(feature, "r") as data: + feature_time = data["time"][:] + n = episode["frames"] + if ( + rows.shape != (n,) + or rows.dtype.kind not in "iu" + or not np.array_equal(rows, np.arange(n)) + or native_time.shape != (n,) + or feature_time.shape != (n,) + or not np.array_equal(native_time, feature_time) + or not np.isfinite(feature_time).all() + or not (np.diff(feature_time) > 0).all() + ): + raise ValueError( + "Native source_row or feature/native timestamp alignment mismatch" + ) + return ( + rows.astype(np.int64), + feature_time.astype(np.float64), + { + "feature_receipt_sha256": episode["receipt_sha256"], + "native_manifest_sha256": inputs["native_manifest_sha256"], + "native_aligned_sha256": inputs["native_aligned_sha256"], + "native_cache": str(native), + }, + ) + + +def export_sidecars(manifest_path, raw_export_dir, output_dir): + manifest_path, raw_dir, output = map( + Path, (manifest_path, raw_export_dir, output_dir) + ) + if output.exists(): + raise FileExistsError(f"Refusing to replace immutable sidecar output: {output}") + verifier = ArtifactVerificationMemo() + manifest = _json(manifest_path) + metadata, evidence, raw_hashes, payload = _verified_bundle( + raw_dir, manifest_path, verifier + ) + raw_episodes = {(e["embodiment"], e["episode_id"]): e for e in metadata["episodes"]} + expected = {(e["embodiment"], e["episode_id"]) for e in manifest["episodes"]} + if len(raw_episodes) != len(metadata["episodes"]) or set(raw_episodes) != expected: + raise ValueError("Raw export must match every prepared episode exactly once") + # Validate the complete small binary and all cache bindings before publishing. + values, ranges = [], [] + schema = metadata["action_schema"] + for episode in manifest["episodes"]: + raw = raw_episodes[(episode["embodiment"], episode["episode_id"])] + n = episode["frames"] + if ( + episode["embodiment"] != "libero_gripper" + or raw["frames"] != n + or raw["split"] != episode["split"] + or raw["feature_sha256"] != episode["sha256"] + or raw["native_source"] != episode["native_source"] + or raw_hashes.get(raw["native_source"]["source_demo"]) + != raw["native_source"]["source_demo_sha256"] + ): + raise ValueError( + "Raw export episode or original demonstration SHA mismatch" + ) + arrays = [] + for name, shape in (("action", [n, 7]), ("state_time", [n])): + offset, size = raw[f"{name}_offset"], raw[f"{name}_nbytes"] + if ( + raw[f"{name}_dtype"] != " metadata["payload_nbytes"] + ): + raise ValueError("Invalid raw command binary dtype/shape/range") + ranges.append((offset, offset + size)) + arrays.append( + np.fromfile( + payload, dtype=" np.asarray(schema["input_max"])).any() + ): + raise ValueError("Non-finite or out-of-controller-bound commands") + rows, cached_time, binding = _frame_binding( + episode, manifest_path.parent, verifier + ) + if not np.array_equal(times, cached_time): + raise ValueError( + "Raw states[:,0] timestamps differ from exact cached feature frames" + ) + values.append((episode, actions.astype(np.float32), cached_time, rows, binding)) + ordered = sorted(ranges) + if ( + not ordered + or ordered[0][0] != 0 + or ordered[-1][1] != metadata["payload_nbytes"] + or any(left[1] != right[0] for left, right in zip(ordered, ordered[1:])) + ): + raise ValueError( + "Raw binary ranges must cover the payload without gaps or overlaps" + ) + output.mkdir(parents=True, exist_ok=False) + provenance = output / "provenance" + provenance.mkdir() + for name, source in evidence.items(): + target = provenance / name + target.parent.mkdir(parents=True, exist_ok=True) + with source.open("rb") as src, target.open("xb") as dst: + shutil.copyfileobj(src, dst) + producer_sha = sha256(__file__) + metadata_sha = sha256(raw_dir / "metadata.json") + output_episodes = [] + for number, (episode, actions, times, rows, binding) in enumerate(values): + name = f"episode-{number:05d}" + sidecar = output / f"{name}.h5" + with h5py.File(sidecar, "x") as data: + for key, value in dict( + action=actions, + time=times, + valid=np.ones(len(rows), bool), + cache_frame_index=np.arange(len(rows)), + source_state_index=rows, + source_action_index=rows, + ).items(): + data.create_dataset(key, data=value) + native = episode["native_source"] + identity = dict( + feature_sha256=episode["sha256"], + source_demo_sha256=native["source_demo_sha256"], + source_episode=native["episode"], + ) + receipt_path = output / f"{name}.alignment.json" + receipt = dict( + schema="dino-proprio-source-alignment-v1", + **identity, + relation="action_t_applied_to_state_t", + producer_sha256=producer_sha, + basis={ + "cache_index_source": f"SHA-verified native aligned.npz source_row; {binding}", + "action_index_source": f"Raw {native['source_demo']} data/{native['episode']}/actions, SHA-verified export metadata {metadata_sha}", + "state_action_relation_source": json.dumps( + metadata["producer_convention"], sort_keys=True + ), + }, + cache_frame_index=list(range(len(rows))), + source_state_index=rows.tolist(), + source_action_index=rows.tolist(), + raw_export_metadata_sha256=metadata_sha, + raw_payload_sha256=metadata["payload_sha256"], + feature_source_binding=binding, + producer_convention=metadata["producer_convention"], + ) + _write_json(receipt_path, receipt) + output_episodes.append( + dict( + embodiment=episode["embodiment"], + episode_id=episode["episode_id"], + path=sidecar.name, + sha256=sha256(sidecar), + **identity, + alignment={ + "relation": "action_t_applied_to_state_t", + "cache_frame_mapping": "explicit_source_indices", + "producer_sha256": producer_sha, + "source_alignment_path": receipt_path.name, + "source_alignment_sha256": sha256(receipt_path), + }, + ) + ) + result = dict( + schema="dino-proprio-actions-v1", + action_schemas={"libero_gripper": schema}, + episodes=output_episodes, + source_manifest_sha256=sha256(manifest_path), + raw_export_metadata_sha256=metadata_sha, + raw_payload_sha256=metadata["payload_sha256"], + exporter_sha256=producer_sha, + producer_convention=metadata["producer_convention"], + normalization="none; native controller commands preserved as float32", + ) + _write_json(output / "manifest.json", result) + return { + "status": "PASS", + "manifest": str(output / "manifest.json"), + "manifest_sha256": sha256(output / "manifest.json"), + "episodes": len(values), + "frames": sum(len(v[2]) for v in values), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--raw-export-dir", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + print( + json.dumps( + export_sidecars(args.manifest, args.raw_export_dir, args.output_dir) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/world_model/initialize_pixel_readout.py b/scripts/world_model/initialize_pixel_readout.py new file mode 100644 index 000000000..d5166034b --- /dev/null +++ b/scripts/world_model/initialize_pixel_readout.py @@ -0,0 +1,496 @@ +#!/usr/bin/env python3 +"""Explicit, offline visual-only -> state/action readout weights migration. + +This writes a plain tensor dictionary for PipelineAlgo.initialization. It never +resumes an optimizer or edits a training configuration. Model arguments must be +resolved JSON mappings, not Hydra experiment names. For a Lightning checkpoint, +declare its exact state dictionary/config paths and namespace; nothing is guessed. +Loading a pickle-based Lightning checkpoint additionally requires the explicit +trusted flag, after its SHA256 has been checked. Tensor-only sources use the safe +weights-only loader, with their architecture explicitly supplied by the caller. +""" + +# ruff: noqa: E402 +from __future__ import annotations + +import argparse +import copy +import hashlib +import inspect +import json +import sys +from collections.abc import Mapping +from pathlib import Path + +import torch + +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from egomimic.models.image_proprio import ImageProprioModel + +MODEL_TARGET = "egomimic.models.image_proprio.ImageProprioModel" +EXPANDED_ROWS = {"proprio_role.weight": (1, 3), "phase.weight": (1, 2)} + + +def sha256_file(path): + with Path(path).open("rb") as handle: + return hashlib.file_digest(handle, "sha256").hexdigest() + + +def resolve_model_config(config): + """Bind the exact class signature; reject unknown Hydra/model arguments.""" + if not isinstance(config, Mapping): + raise TypeError("Model configuration must be an argument mapping") + config = copy.deepcopy(dict(config)) + if config.pop("_target_", MODEL_TARGET) != MODEL_TARGET: + raise ValueError("Only ImageProprioModel is supported") + signature = inspect.signature(ImageProprioModel) + bound = signature.bind(**config) + bound.apply_defaults() + # Reject unresolved interpolation and non-JSON objects before model creation. + resolved = json.loads(json.dumps(dict(bound.arguments), allow_nan=False)) + if "${" in json.dumps(resolved): + raise ValueError("Model configuration must have resolved interpolations") + return resolved + + +def validate_transition(source_config, target_config): + source = resolve_model_config(source_config) + target = resolve_model_config(target_config) + if source["objective"] != "visual_only": + raise ValueError("Source objective must be visual_only") + if target["objective"] != "visual_state_action": + raise ValueError("Target objective must be visual_state_action") + if source["action_dims"] is not None: + raise ValueError("Visual-only source must not declare action dimensions") + if not isinstance(target["action_dims"], Mapping): + raise ValueError("Target requires explicit action dimensions") + image_only = not source["phase_a_condition_on_proprio"] + target_poses = ImageProprioModel._dimensions(target["pose_dims"], "pose_dims") + if image_only: + if source["pose_dims"] not in (None, {}): + raise ValueError("Image-only source must omit pose_dims") + elif list(source["pose_dims"].items()) != list(target_poses.items()): + # Embodiment order affects embedding IDs and deterministic initialization. + raise ValueError("Source and target pose_dims and order must match") + if list(target["action_dims"]) != list(target_poses): + raise ValueError("Action dimensions must use the same embodiment order") + allowed = {"objective", "action_dims", "max_tokens"} + if image_only: + allowed.add("pose_dims") + changed = [ + key for key in source if key not in allowed and source[key] != target[key] + ] + if changed: + raise ValueError(f"Source/target model configuration mismatch: {changed}") + if target["max_tokens"] < source["max_tokens"]: + raise ValueError("Target token capacity must not decrease") + return source, target + + +def _select(payload, path): + for key in path: + # JSON path components can be object keys or integer sequence indices. + payload = payload[key] + return payload + + +def _plain_config(value): + from omegaconf import OmegaConf + + return ( + OmegaConf.to_container(value, resolve=True) + if OmegaConf.is_config(value) + else value + ) + + +def _validate_state(state, expected): + if set(state) != set(expected): + raise ValueError( + "Source namespace mismatch: " + f"missing={sorted(set(expected) - set(state))}, " + f"unexpected={sorted(set(state) - set(expected))}" + ) + for name, value in state.items(): + if ( + not isinstance(value, torch.Tensor) + or value.shape != expected[name].shape + or value.dtype != expected[name].dtype + ): + raise ValueError(f"Source tensor shape/dtype mismatch: {name}") + if not torch.isfinite(value).all(): + raise ValueError(f"Source tensor is nonfinite: {name}") + + +def transfer_shared_weights(source, target): + """Copy every source tensor; initialize only this exact readout expansion.""" + if source.use_phase_embedding != target.use_phase_embedding: + raise ValueError("Source/target use_phase_embedding must match") + if source.noise_conditioning != target.noise_conditioning: + raise ValueError("Source/target noise_conditioning must match") + for field in ("position_encoding", "rope_axis_dims", "spatial_grid_size"): + if getattr(source, field) != getattr(target, field): + raise ValueError(f"Source/target {field} must match") + source_state, target_state = source.state_dict(), target.state_dict() + added = {"query", "action_query", "norm.1.weight", "norm.1.bias"} + for head in ("pose_out", "action_out"): + for embodiment in target.pose_dims: + added.update(f"{head}.{embodiment}.{field}" for field in ("weight", "bias")) + expanded = { + name: shape for name, shape in EXPANDED_ROWS.items() if name in source_state + } + if not source.has_proprio_branch: + added.add("proprio_role.weight") + for embodiment in target.pose_dims: + added.update( + f"pose_in.{embodiment}.{field}" for field in ("weight", "bias") + ) + for index, block in enumerate(target.blocks): + for branch in ("norm1", "qkv", "out", "norm2", "ffn"): + added.update( + f"blocks.{index}.{branch}.1.{name}" + for name, _ in getattr(block, branch)[1].named_parameters() + ) + if source.use_modality_embedding: + expanded["modality"] = (1, 2) + if ( + set(source_state) - set(target_state) + or set(target_state) - set(source_state) != added + ): + raise ValueError("Unexpected parameter namespace in readout transition") + for name, value in source_state.items(): + expected = target_state[name] + if name in expanded: + old_rows, new_rows = expanded[name] + if value.shape != (old_rows, source.width) or expected.shape != ( + new_rows, + source.width, + ): + raise ValueError(f"Unexpected expanded embedding shape: {name}") + if value.dtype != expected.dtype: + raise ValueError(f"Shared tensor dtype mismatch: {name}") + expected[:old_rows].copy_(value) + else: + if value.shape != expected.shape or value.dtype != expected.dtype: + raise ValueError(f"Shared tensor shape/dtype mismatch: {name}") + expected.copy_(value) + target.load_state_dict(target_state, strict=True) + return { + "copied_tensors": sorted(set(source_state) - set(expanded)), + "expanded_embeddings": { + key: {"copied_rows": [0], "new_rows": list(range(1, rows[1]))} + for key, rows in expanded.items() + }, + "new_tensors": sorted(added), + "new_parameter_initialization": "ImageProprioModel constructor at declared seed", + } + + +def _probe_batch(source, target, embodiment, completion): + """Small deterministic input with observed/query images and withheld labels.""" + image_observed = torch.tensor([[True, False, False, True if completion else False]]) + pose_observed = ( + torch.tensor([[False, True]]) if completion else torch.tensor([[True, False]]) + ) + + def stream(tokens, dimensions, observed): + return { + "value": torch.linspace(-0.8, 0.9, tokens * dimensions).reshape( + 1, tokens, dimensions + ), + "time": torch.linspace(-0.6, 0.7, tokens).reshape(1, tokens), + "index": torch.arange(tokens).reshape(1, tokens), + "valid": torch.ones(1, tokens, dtype=torch.bool), + "observed": observed, + } + + batch = { + "image": stream(4, source.image_dim, image_observed), + "proprio": stream(2, target.pose_dims[embodiment], pose_observed), + "action": stream( + 2, target.action_dims[embodiment], torch.zeros(1, 2, dtype=torch.bool) + ), + } + batch["image"]["xy"] = torch.tensor( + [[[-0.5, -0.5], [0.5, -0.5], [-0.5, 0.5], [0.5, 0.5]]] + ) + cameras = source.camera.num_embeddings if source.use_camera_embedding else 1 + batch["image"]["camera"] = torch.arange(4).remainder(cameras).reshape(1, 4) + # Neither missing state labels nor command labels may influence phase A. + batch["proprio"]["value"][~pose_observed] = float("nan") + batch["action"]["value"].fill_(float("nan")) + values = torch.linspace(0.7, -0.8, 4 * source.image_dim).reshape( + 1, 4, source.image_dim + ) + levels = torch.tensor([[0.0, 0.25, 0.75, 0.5]]) + return batch, values, levels + + +@torch.no_grad() +def verify_phase_a(source, target, cpu_attention_override=None): + """Require bit-exact CPU output; an optional native override is disclosed.""" + if cpu_attention_override not in (None, "native"): + raise ValueError("cpu_attention_override must be None or native") + blocks = [block for model in (source, target) for block in model.blocks] + execution = [(block.attention_backend, block.attention_dtype) for block in blocks] + if cpu_attention_override is None and any( + backend != "sdpa" for backend, _ in execution + ): + raise ValueError( + "GPU attention needs an explicit cpu_attention_override='native' " + "for this CPU weights-preservation proof" + ) + source.eval() + target.eval() + checks = [] + try: + if cpu_attention_override == "native": + for block in blocks: + block.attention_backend, block.attention_dtype = "sdpa", "native" + for embodiment in target.pose_dims: + for completion in (False, True): + batch, values, levels = _probe_batch( + source, target, embodiment, completion + ) + source_batch = ( + batch if source.has_proprio_branch else {"image": batch["image"]} + ) + original = source(source_batch, embodiment, values, levels, phase="A")[ + "image" + ] + migrated = target(batch, embodiment, values, levels, phase="A")["image"] + if not torch.isfinite(original).all() or not torch.equal( + original, migrated + ): + difference = (original - migrated).abs().max().item() + raise ValueError( + f"Phase A preservation failed: {embodiment}, max_abs={difference}" + ) + checks.append( + { + "embodiment": embodiment, + "mask": "completion" if completion else "forecast", + "max_abs_error": 0.0, + } + ) + finally: + for block, (backend, dtype) in zip(blocks, execution): + block.attention_backend, block.attention_dtype = backend, dtype + return { + "device": "cpu", + "dtype": "float32", + "bit_exact": True, + "checks": checks, + "cpu_attention_override": cpu_attention_override, + "declared_attention": sorted( + {f"{backend}/{dtype}" for backend, dtype in execution} + ), + "gpu_kernel_verified": False, + } + + +def initialize_readout( + *, + source_checkpoint, + source_sha256, + source_config, + target_config, + output_path, + receipt_path, + seed, + source_prefix="", + state_dict_path=(), + source_config_path=(), + trusted_lightning_checkpoint=False, + stage_id="image_proprio", + module_path="model", + cpu_attention_override=None, +): + """Create a fresh stage-two tensor artifact, never an optimizer resume. + + Full checkpoints must bind their model configuration using an explicit path. + A bare tensor dictionary has no architecture metadata; in that case the + receipt explicitly records that the caller supplied the model declaration. + """ + if type(seed) is not int or not 0 <= seed < 2**63: + raise ValueError("Seed must be an integer in [0, 2**63)") + source_config, target_config = validate_transition(source_config, target_config) + source_path = Path(source_checkpoint).resolve() + output_path, receipt_path = ( + Path(output_path).resolve(), + Path(receipt_path).resolve(), + ) + if output_path == receipt_path or source_path in (output_path, receipt_path): + raise ValueError("Input, output and receipt paths must be distinct") + if output_path.exists() or receipt_path.exists(): + raise FileExistsError( + "Initialization artifacts are immutable; choose new paths" + ) + if trusted_lightning_checkpoint and not source_config_path: + raise ValueError( + "Trusted Lightning input requires a source_config_path binding" + ) + # Hash and load the same open file. A replacement at this path cannot swap + # the verified source underneath torch.load; also reject in-place mutation. + with source_path.open("rb") as handle: + digest = hashlib.file_digest(handle, "sha256").hexdigest() + if digest != source_sha256: + raise ValueError("Source checkpoint SHA256 mismatch") + handle.seek(0) + # No automatic unsafe-pickle fallback. + payload = torch.load( + handle, map_location="cpu", weights_only=not trusted_lightning_checkpoint + ) + handle.seek(0) + if hashlib.file_digest(handle, "sha256").hexdigest() != digest: + raise ValueError("Source checkpoint changed while loading") + if source_config_path: + recorded = resolve_model_config( + _plain_config(_select(payload, source_config_path)) + ) + if recorded != source_config or list( + (recorded["pose_dims"] or {}).items() + ) != list((source_config["pose_dims"] or {}).items()): + raise ValueError( + "Checkpoint model configuration does not match declared source" + ) + selected = _select(payload, state_dict_path) + if not isinstance(selected, Mapping) or any( + not isinstance(key, str) for key in selected + ): + raise TypeError("State dictionary path must select a string-keyed mapping") + if not source_config_path and (state_dict_path or source_prefix): + raise ValueError( + "Nested/namespaced checkpoint input requires source_config_path" + ) + state = { + key[len(source_prefix) :]: value + for key, value in selected.items() + if key.startswith(source_prefix) + } + # fork_rng prevents a utility call from changing the caller's initialization RNG. + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(seed) + source = ImageProprioModel(**source_config).cpu().float() + _validate_state(state, source.state_dict()) + source.load_state_dict(state, strict=True) + torch.random.default_generator.manual_seed(seed) + target = ImageProprioModel(**target_config).cpu().float() + transfer = transfer_shared_weights(source, target) + preservation = verify_phase_a(source, target, cpu_attention_override) + weights = { + key: value.detach().cpu().clone() for key, value in target.state_dict().items() + } + for path in (output_path, receipt_path): + path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("xb") as handle: + torch.save(weights, handle) + output_sha = sha256_file(output_path) + receipt = { + "schema": "pixel-readout-initialization-v1", + "status": "PASS", + "source": { + "path": str(source_path), + "sha256": digest, + "state_dict_path": list(state_dict_path), + "source_prefix": source_prefix, + "model_config_path": list(source_config_path), + "config_binding": "checkpoint" + if source_config_path + else "caller_declared_tensor_state", + "trusted_pickle": bool(trusted_lightning_checkpoint), + }, + "source_model": source_config, + "target_model": target_config, + "seed": seed, + "transfer": transfer, + "phase_a_preservation": preservation, + "output": { + "path": str(output_path), + "sha256": output_sha, + "tensor_count": len(weights), + "format": "plain_tensor_state_dict", + }, + "implementation": { + "utility_sha256": sha256_file(__file__), + "model_sha256": sha256_file(inspect.getfile(ImageProprioModel)), + "torch_version": str(torch.__version__), + }, + "optimizer_state_transferred": False, + "training_step_transferred": False, + "initialization": [ + { + "stage_id": stage_id, + "module_path": module_path, + "source": str(output_path), + "sha256": output_sha, + } + ], + } + with receipt_path.open("x") as handle: + json.dump(receipt, handle, indent=2, allow_nan=False) + handle.write("\n") + return receipt + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-checkpoint", required=True) + parser.add_argument("--source-sha256", required=True) + parser.add_argument( + "--source-config", required=True, help="Resolved model arguments JSON" + ) + parser.add_argument( + "--target-config", required=True, help="Resolved model arguments JSON" + ) + parser.add_argument("--output", required=True) + parser.add_argument("--receipt", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--source-prefix", default="") + parser.add_argument( + "--state-dict-path", + type=json.loads, + default=[], + help='JSON path, e.g. ["state_dict"]', + ) + parser.add_argument( + "--source-config-path", + type=json.loads, + default=[], + help="JSON path into checkpoint for exact model arguments", + ) + parser.add_argument("--trusted-lightning-checkpoint", action="store_true") + parser.add_argument("--stage-id", default="image_proprio") + parser.add_argument("--module-path", default="model") + parser.add_argument( + "--cpu-attention-override", + choices=["native"], + help="Explicit native SDPA CPU weight-preservation proof; does not verify GPU kernels", + ) + args = parser.parse_args() + receipt = initialize_readout( + source_checkpoint=args.source_checkpoint, + source_sha256=args.source_sha256, + source_config=json.loads(Path(args.source_config).read_text()), + target_config=json.loads(Path(args.target_config).read_text()), + output_path=args.output, + receipt_path=args.receipt, + seed=args.seed, + source_prefix=args.source_prefix, + state_dict_path=args.state_dict_path, + source_config_path=args.source_config_path, + trusted_lightning_checkpoint=args.trusted_lightning_checkpoint, + stage_id=args.stage_id, + module_path=args.module_path, + cpu_attention_override=args.cpu_attention_override, + ) + print( + json.dumps({"status": receipt["status"], "output": receipt["output"]}, indent=2) + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/world_model/prepare_pixel_proprio.py b/scripts/world_model/prepare_pixel_proprio.py new file mode 100644 index 000000000..cf91bdb95 --- /dev/null +++ b/scripts/world_model/prepare_pixel_proprio.py @@ -0,0 +1,715 @@ +#!/usr/bin/env python3 +"""Prepare immutable RGB pixels and same-state measured proprio, without an encoder. + +CPU preparation preserves every recorded timestamp, camera and source identity. +Splits are preserved unless an explicit balanced derivation is requested. +Native RGB is reduced by exact integer area averaging with ties-to-even +uint8 rounding; no crop, frame resampling, encoder or track arrays are used. +Only training episodes fit pose statistics. Action commands are optional and +require a separately audited export. Run on an allocated cluster CPU node. +""" + +# ruff: noqa: E402 +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import shutil +import sys +from pathlib import Path + +import h5py +import numpy as np + +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from egomimic.rldb.world_model_preparation import ( + ArtifactVerificationMemo, + StreamingMoments, + training_identities, +) +from scripts.world_model.export_libero_action_sidecars import ( + _json, + _verified_bundle, + _write_json, + sha256, +) + + +def _key(episode): + return episode["embodiment"], episode["episode_id"] + + +def _index(episodes): + result = {_key(episode): episode for episode in episodes} + if len(result) != len(episodes): + raise ValueError("Duplicate embodiment/episode identity") + return result + + +def _json_sha256(value): + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def derive_split(episodes, validation_episodes_per_task=None, split_seed=42): + """Derive balanced gripper splits without mutating source assignments. + + A task is the verified native demonstration filename without ``_demo``; + its episode_id must independently equal that task plus native demo index. + Existing held-out episodes always remain held out. Remaining choices use + SHA256 ranks, independent of source ordering and NumPy/RNG versions. + """ + if type(split_seed) is not int or split_seed < 0: + raise ValueError("split_seed must be a nonnegative integer") + original = _index(episodes) + if not original or any(ep["split"] not in ("train", "val") for ep in episodes): + raise ValueError("Require original explicit source splits") + result = {key: dict(ep, original_split=ep["split"]) for key, ep in original.items()} + policy = {"kind": "preserved_source_episode_assignments"} + if validation_episodes_per_task is not None: + if ( + type(validation_episodes_per_task) is not int + or validation_episodes_per_task < 1 + ): + raise ValueError( + "validation_episodes_per_task must be a positive integer or None" + ) + groups, task_sources = {}, {} + for key, ep in original.items(): + native = ep["native_source"] + source = Path(native["source_demo"]) + task = source.stem.removesuffix("_demo") + source_episode = native["episode"] + source_sha = native["source_demo_sha256"] + if ( + key[0] != "libero_gripper" + or source.suffix not in (".hdf5", ".h5") + or not source.stem.endswith("_demo") + or not task + or re.fullmatch(r"demo_\d+", source_episode) is None + or ep["episode_id"] != f"{task}-{source_episode}" + or re.fullmatch(r"[0-9a-f]{64}", source_sha) is None + ): + raise ValueError( + "Balanced split requires matching native LIBERO task filename and episode identity" + ) + if task in task_sources and task_sources[task] != source_sha: + raise ValueError( + "One task filename refers to conflicting native source hashes" + ) + task_sources[task] = source_sha + groups.setdefault(task, []).append(key) + result[key]["task_id"] = task + task_counts = {} + for task, keys in sorted(groups.items()): + retained = {key for key in keys if original[key]["split"] == "val"} + if len( + retained + ) > validation_episodes_per_task or validation_episodes_per_task >= len( + keys + ): + raise ValueError( + "Requested balance must retain existing validation and nonempty training per task" + ) + candidates = sorted( + (key for key in keys if key not in retained), + key=lambda key: (_json_sha256([split_seed, task, key[0], key[1]]), key), + ) + validation = retained | set( + candidates[: validation_episodes_per_task - len(retained)] + ) + for key in keys: + result[key]["split"] = "val" if key in validation else "train" + task_counts[task] = dict( + train=len(keys) - len(validation), + val=len(validation), + retained_original_val=len(retained), + source_demo_sha256=task_sources[task], + ) + policy = dict( + kind="balanced_native_libero_tasks_v1", + seed=split_seed, + validation_episodes_per_task=validation_episodes_per_task, + task_identity="native source filename without _demo, cross-checked against episode_id and native demo index", + selection="retain_original_validation_then_lowest_sha256_of_compact_json_[seed,task,embodiment,episode_id]", + task_counts=task_counts, + ) + + def identities(values): + return sorted( + [ + ep["embodiment"], + ep["episode_id"], + ep["native_source"]["source_demo_sha256"], + ep["native_source"]["episode"], + ep["frames"], + ep["split"], + ] + for ep in values + ) + + original_ids, derived_ids = ( + identities(original.values()), + identities(result.values()), + ) + policy.update( + original_identities=original_ids, + derived_identities=derived_ids, + original_identities_sha256=_json_sha256(original_ids), + derived_identities_sha256=_json_sha256(derived_ids), + counts={ + emb: { + split: sum( + ep["embodiment"] == emb and ep["split"] == split + for ep in result.values() + ) + for split in ("train", "val") + } + for emb in sorted({ep["embodiment"] for ep in episodes}) + }, + ) + return result, policy + + +def _raw_actions(raw_dir, source_manifest, verifier): + """Validate the entire small command export, independent of feature caches.""" + metadata, evidence, hashes, payload = _verified_bundle( + raw_dir, source_manifest, verifier + ) + reference = _index(_json(source_manifest)["episodes"]) + raw = _index(metadata["episodes"]) + if set(raw) != set(reference): + raise ValueError("Raw command export must exactly cover its source manifest") + ranges, values = [], {} + schema = metadata["action_schema"] + for key, item in raw.items(): + episode, n = reference[key], item["frames"] + if ( + key[0] != "libero_gripper" + or type(n) is not int + or n < 1 + or n != episode["frames"] + or item["split"] != episode["split"] + or item["native_source"] != episode["native_source"] + or hashes.get(item["native_source"]["source_demo"]) + != item["native_source"]["source_demo_sha256"] + ): + raise ValueError("Raw source identity does not match source split manifest") + arrays = [] + for name, shape in (("action", [n, 7]), ("state_time", [n])): + offset, size = item[f"{name}_offset"], item[f"{name}_nbytes"] + if ( + item[f"{name}_dtype"] != " metadata["payload_nbytes"] + ): + raise ValueError("Invalid raw command binary dtype, shape or range") + ranges.append((offset, offset + size)) + arrays.append( + np.fromfile(payload, " 0).all() + or (actions < np.asarray(schema["input_min"])).any() + or (actions > np.asarray(schema["input_max"])).any() + ): + raise ValueError("Invalid raw commands or state times") + values[key] = (actions.astype(np.float32), times, item) + ranges.sort() + if ( + not ranges + or ranges[0][0] != 0 + or ranges[-1][1] != metadata["payload_nbytes"] + or any(a[1] != b[0] for a, b in zip(ranges, ranges[1:])) + ): + raise ValueError("Raw payload ranges must be contiguous without overlap") + return metadata, evidence, values + + +def resize_rgb_area(rgb, image_size): + """Exact integer area mean in RGB, rounded to nearest uint8 with ties to even.""" + if rgb.ndim < 3 or rgb.shape[-1] != 3 or rgb.dtype != np.uint8: + raise ValueError("Area resize requires uint8 RGB[...,H,W,3]") + height, width = image_size + source_h, source_w = rgb.shape[-3:-1] + if ( + height < 1 + or width < 1 + or source_h % height + or source_w % width + or height > source_h + or width > source_w + ): + raise ValueError( + "Area resize requires integer downsampling factors with no upscaling" + ) + fy, fx = source_h // height, source_w // width + if fy == fx == 1: + return rgb.copy() + cells = rgb.reshape(*rgb.shape[:-3], height, fy, width, fx, 3) + summed = cells.sum(axis=(-4, -2), dtype=np.uint64) + return np.rint(summed / (fy * fx)).astype(np.uint8) + + +def _native_binding(episode, locations, pose_schema, image_size, verifier): + native = Path(locations["native_cache"]).resolve(strict=True) + rgb_path = Path(locations["rgb_cache"]).resolve(strict=True) + receipt_path = Path(locations["rgb_receipt"]).resolve(strict=True) + receipt_sha = sha256(receipt_path) + receipt = _json(receipt_path) + height, width = receipt.get("height"), receipt.get("width") + if ( + receipt.get("schema") != "libero-native-rgb-v1" + or receipt.get("status") != "PASS" + or receipt.get("native_render") is not True + or receipt.get("native_geometry_replay_exact") is not True + or type(height) is not int + or type(width) is not int + or height < image_size[0] + or width < image_size[1] + or height % image_size[0] + or width % image_size[1] + or receipt.get("frames") != episode["frames"] + ): + raise ValueError( + "Require successful native RGB receipt compatible with integer area downsampling" + ) + native_manifest_path = native / "manifest.json" + verifier.verify(native_manifest_path, receipt["source_manifest_sha256"]) + native_manifest = _json(native_manifest_path) + source = episode["native_source"] + cameras = receipt["cameras"] + if ( + native_manifest.get("schema") != "libero-same-recorded-state-v1" + or any(native_manifest.get(k) != v for k, v in source.items()) + or native_manifest.get("pose_fields") != pose_schema["fields"] + or native_manifest.get("cameras") != cameras + or not isinstance(cameras, list) + or not cameras + or len(set(cameras)) != len(cameras) + or native_manifest.get("artifacts", {}).get("aligned.npz") + != receipt["source_aligned_sha256"] + ): + raise ValueError("Native measured-pose/source/camera schema mismatch") + aligned_path = native / "aligned.npz" + aligned_fingerprint = verifier.verify( + aligned_path, + receipt["source_aligned_sha256"], + trusted_fingerprint=locations.get("aligned_verified_fingerprint"), + ) + # Load no tracks, original low-resolution RGB, or feature arrays. + with np.load(aligned_path, allow_pickle=False) as aligned: + pose = aligned["pose"].astype(np.float32) + time = aligned["time"] + rows = aligned["source_row"] + state_time = aligned["state"][:, 0] + n, d = episode["frames"], pose_schema["dimension"] + if ( + pose.shape != (n, d) + or not np.isfinite(pose).all() + or time.shape != (n,) + or not np.isfinite(time).all() + or not (np.diff(time) > 0).all() + or not np.array_equal(time, state_time) + or rows.shape != (n,) + or rows.dtype.kind not in "iu" + or not np.array_equal(rows, np.arange(n)) + ): + raise ValueError("Native pose/state-time/source-row alignment mismatch") + rgb_fingerprint = verifier.verify( + rgb_path, + receipt["rgb_cache_sha256"], + trusted_fingerprint=locations.get("rgb_verified_fingerprint"), + ) + with h5py.File(rgb_path, "r") as rgb: + if ( + rgb.attrs.get("schema") != "libero-native-rgb-v1" + or not bool(rgb.attrs.get("native_render", False)) + or rgb["rgb"].dtype != np.dtype("uint8") + or rgb["rgb"].shape != (n, len(cameras), height, width, 3) + or not np.array_equal(rgb["time"][:], time) + ): + raise ValueError("Native RGB shape/type/time differs from measured proprio") + # Receipts are small; also rule out mutation during the long data hash. + verifier.verify(receipt_path, receipt_sha) + binding = dict( + source_rgb_path=str(rgb_path), + source_rgb_sha256=receipt["rgb_cache_sha256"], + source_image_size=[height, width], + rgb_receipt_path=str(receipt_path), + rgb_receipt_sha256=receipt_sha, + native_cache=str(native), + native_manifest_sha256=receipt["source_manifest_sha256"], + native_aligned_sha256=receipt["source_aligned_sha256"], + ) + if rgb_fingerprint: + binding["source_rgb_verified_fingerprint"] = rgb_fingerprint + if aligned_fingerprint: + binding["native_aligned_verified_fingerprint"] = aligned_fingerprint + return pose, time.astype(np.float64), rows.astype(np.int64), cameras, binding + + +def _write_resized_rgb(path, time, cameras, binding, image_size, batch_size, verifier): + source_path = Path(binding["source_rgb_path"]) + verifier.verify( + source_path, + binding["source_rgb_sha256"], + trusted_fingerprint=binding.get("source_rgb_verified_fingerprint"), + ) + before = verifier.identity(source_path) + with h5py.File(source_path, "r") as source, h5py.File(path, "x") as target: + target.attrs["schema"] = "pixel-proprio-rgb-v1" + target.attrs["resize"] = "integer_area_mean_round_ties_to_even" + target.attrs["source_sha256"] = binding["source_rgb_sha256"] + target.create_dataset("time", data=time) + target.create_dataset( + "cameras", data=np.asarray(cameras, dtype=h5py.string_dtype("utf-8")) + ) + output = target.create_dataset( + "rgb", + shape=(len(time), len(cameras), *image_size, 3), + dtype=np.uint8, + chunks=(1, 1, *image_size, 3), + compression="lzf", + ) + for start in range(0, len(time), batch_size): + stop = min(start + batch_size, len(time)) + output[start:stop] = resize_rgb_area(source["rgb"][start:stop], image_size) + if verifier.identity(source_path) != before: + raise ValueError("Native RGB source changed during pixel preparation") + return sha256(path) + + +def prepare_pixels( + source_manifest, + native_spec, + output_dir, + *, + image_size, + patch_size, + raw_export_dir=None, + raw_source_manifest=None, + std_floor=1e-5, + resize_batch_size=4, + embodiments=None, + validation_episodes_per_task=None, + split_seed=42, +): + source_manifest, native_spec, output = map( + lambda p: Path(p).resolve(), (source_manifest, native_spec, output_dir) + ) + if output.exists(): + raise FileExistsError(f"Preserve immutable pixel preparation output: {output}") + if ( + len(image_size) != 2 + or any(type(v) is not int or v <= 0 for v in image_size) + or type(patch_size) is not int + or patch_size <= 0 + or any(v % patch_size for v in image_size) + or not np.isfinite(std_floor) + or std_floor <= 0 + or type(resize_batch_size) is not int + or resize_batch_size < 1 + ): + raise ValueError( + "Require patch-divisible image dimensions, positive batch size and std floor" + ) + source_sha, locations_sha = sha256(source_manifest), sha256(native_spec) + source, spec = _json(source_manifest), _json(native_spec) + all_episodes, locations = _index(source["episodes"]), _index(spec["episodes"]) + if not all_episodes: + raise ValueError("Require source episodes") + available_embodiments = {key[0] for key in all_episodes} + schemas = { + emb: source["pose_schemas"][emb] for emb in sorted(available_embodiments) + } + source_keys = set() + for ep in all_episodes.values(): + if ( + ep["split"] not in {"train", "val"} + or type(ep["frames"]) is not int + or ep["frames"] < 1 + ): + raise ValueError("Invalid source episode frames/split") + source_key = ( + ep["embodiment"], + ep["native_source"]["source_demo_sha256"], + ep["native_source"]["episode"], + ) + if source_key in source_keys: + raise ValueError("Duplicate native demonstration episode in source split") + source_keys.add(source_key) + for emb in available_embodiments: + schema = schemas[emb] + if ( + type(schema["dimension"]) is not int + or schema["dimension"] < 1 + or len(schema["fields"]) != schema["dimension"] + or len(set(schema["fields"])) != schema["dimension"] + or {ep["split"] for ep in all_episodes.values() if ep["embodiment"] == emb} + != {"train", "val"} + ): + raise ValueError( + "Require complete pose schema and nonempty train/val split per embodiment" + ) + selected = set(embodiments) if embodiments is not None else available_embodiments + if ( + not selected + or not selected.issubset(available_embodiments) + or (embodiments is not None and len(selected) != len(embodiments)) + ): + raise ValueError( + "Select nonempty unique embodiments present in the frozen source manifest" + ) + embodiments = selected + episodes = {key: ep for key, ep in all_episodes.items() if key[0] in selected} + schemas = {emb: schemas[emb] for emb in sorted(selected)} + if not set(episodes).issubset(locations): + raise ValueError("Native locations must cover each selected source episode") + verifier = ArtifactVerificationMemo() + raw_metadata, raw_evidence, raw_values = None, {}, {} + raw_reference = ( + Path(raw_source_manifest).resolve() if raw_source_manifest else source_manifest + ) + if raw_export_dir: + raw_metadata, raw_evidence, raw_values = _raw_actions( + Path(raw_export_dir).resolve(), raw_reference, verifier + ) + if not set(raw_values).issubset(episodes): + raise ValueError( + "Action export contains episodes absent from selected pixel corpus" + ) + elif raw_source_manifest: + raise ValueError("raw_source_manifest requires raw_export_dir") + # Validate policy metadata before long file reads; apply only after every + # native/raw alignment below has been checked against its original split. + derived_episodes, split_provenance = derive_split( + list(episodes.values()), validation_episodes_per_task, split_seed + ) + values, moments = [], {emb: StreamingMoments() for emb in embodiments} + for key, ep in episodes.items(): + pose, time, rows, cameras, binding = _native_binding( + ep, locations[key], schemas[key[0]], image_size, verifier + ) + actions = None + if key in raw_values: + actions, raw_times, raw = raw_values[key] + if ( + raw["native_source"] != ep["native_source"] + or raw["split"] != ep["split"] + or not np.array_equal(raw_times, time) + ): + raise ValueError( + "Raw action source/split/state time differs from native RGB/proprio" + ) + values.append((ep, pose, time, rows, cameras, binding, actions)) + verifier.verify(source_manifest, source_sha) + verifier.verify(native_spec, locations_sha) + values = [(derived_episodes[_key(ep)], *rest) for ep, *rest in values] + for ep, pose, *_ in values: + if ep["split"] == "train": + moments[ep["embodiment"]].update(pose) + # Publish only after every input has passed the CPU validation above. + output.mkdir(parents=True, exist_ok=False) + provenance = output / "provenance" + provenance.mkdir() + split_path = provenance / "split.json" + _write_json(split_path, split_provenance) + for name, path in ( + ("source-manifest.json", source_manifest), + ("native-locations.json", native_spec), + ): + with path.open("rb") as src, (provenance / name).open("xb") as dst: + shutil.copyfileobj(src, dst) + for name, path in raw_evidence.items(): + target = provenance / "raw-commands" / name + target.parent.mkdir(parents=True, exist_ok=True) + with path.open("rb") as src, target.open("xb") as dst: + shutil.copyfileobj(src, dst) + producer_sha = sha256(__file__) + raw_metadata_sha = ( + sha256(Path(raw_export_dir) / "metadata.json") if raw_export_dir else None + ) + result_episodes = [] + for index, (ep, pose, time, rows, cameras, binding, actions) in enumerate(values): + name = f"episode-{index:05d}" + path = output / f"{name}.h5" + rgb_path = output / f"{name}.rgb.h5" + rgb_sha = _write_resized_rgb( + rgb_path, time, cameras, binding, image_size, resize_batch_size, verifier + ) + with h5py.File(path, "x") as data: + data.attrs["schema"] = "pixel-proprio-metadata-v1" + arrays = dict(pose=pose, pose_valid=np.ones(len(time), bool), time=time) + if actions is not None: + arrays.update(action=actions, action_valid=np.ones(len(time), bool)) + for field, value in arrays.items(): + data.create_dataset(field, data=value) + item = dict( + embodiment=ep["embodiment"], + episode_id=ep["episode_id"], + split=ep["split"], + original_split=ep["original_split"], + frames=len(time), + path=path.name, + sha256=sha256(path), + frame="fixed_reference", + pose_kind="measured_pose", + native_source=ep["native_source"], + cameras=cameras, + rgb_path=rgb_path.name, + rgb_sha256=rgb_sha, + **binding, + ) + if "task_id" in ep: + item["task_id"] = ep["task_id"] + if actions is not None: + native = ep["native_source"] + receipt = dict( + schema="pixel-proprio-action-alignment-v1", + source_demo_sha256=native["source_demo_sha256"], + source_episode=native["episode"], + relation="action_t_applied_to_state_t", + producer_sha256=producer_sha, + basis={ + "cache_index_source": "SHA-verified native source_row equals frame index; RGB time, pose time and state[:,0] are exactly equal", + "action_index_source": f"Raw {native['source_demo']} data/{native['episode']}/actions; verified export {raw_metadata_sha}", + "state_action_relation_source": json.dumps( + raw_metadata["producer_convention"], sort_keys=True + ), + }, + cache_frame_index=list(range(len(time))), + source_state_index=rows.tolist(), + source_action_index=rows.tolist(), + raw_export_metadata_sha256=raw_metadata_sha, + raw_payload_sha256=raw_metadata["payload_sha256"], + producer_convention=raw_metadata["producer_convention"], + pixel_source_binding=binding, + ) + receipt_path = output / f"{name}.alignment.json" + _write_json(receipt_path, receipt) + item["action_provenance"] = dict( + path=receipt_path.name, sha256=sha256(receipt_path) + ) + result_episodes.append(item) + # Cache exact content proofs for the newly finalized files when their ctime + # is old enough. The final manifest binds these same-device fast paths. + for item in result_episodes: + for field, digest_field, fingerprint_field in ( + ("path", "sha256", "verified_fingerprint"), + ("rgb_path", "rgb_sha256", "rgb_verified_fingerprint"), + ): + proof = verifier.verify(output / item[field], item[digest_field]) + if proof is not None: + item[fingerprint_field] = proof + stats = dict( + schema="pixel-proprio-normalization-v1", + scope="train_only", + train_episodes=training_identities(result_episodes), + statistics={ + f"pose/{emb}": moments[emb].result(std_floor) for emb in sorted(embodiments) + }, + std_floor=std_floor, + ) + stats_path = output / "normalization.json" + _write_json(stats_path, stats) + manifest = dict( + schema="pixel-proprio-prepared-v1", + episodes=result_episodes, + pose_schemas=schemas, + pixel_identity=dict( + image_size=list(image_size), + patch_size=patch_size, + channels=3, + patch_order="row_major_hwc", + normalization="uint8_to_minus_one_one", + resize="integer_area_mean_round_ties_to_even", + resize_numpy_version=np.__version__, + crop="none", + temporal_resampling="none", + ), + normalization=dict( + scope="train_only", + arrays_normalized=False, + path=stats_path.name, + sha256=sha256(stats_path), + ), + source_manifest_sha256=source_sha, + native_locations_sha256=locations_sha, + selected_embodiments=sorted(embodiments), + split_provenance={ + **split_provenance, + "path": str(split_path.relative_to(output)), + "sha256": sha256(split_path), + }, + preparer_sha256=producer_sha, + ) + if raw_metadata: + manifest.update( + action_schemas={"libero_gripper": raw_metadata["action_schema"]}, + raw_export_metadata_sha256=raw_metadata_sha, + raw_payload_sha256=raw_metadata["payload_sha256"], + ) + verifier.verify(source_manifest, source_sha) + verifier.verify(native_spec, locations_sha) + _write_json(output / "manifest.json", manifest) + return dict( + status="PASS", + manifest=str(output / "manifest.json"), + manifest_sha256=sha256(output / "manifest.json"), + episodes=len(values), + frames=sum(len(item[2]) for item in values), + image_size=list(image_size), + patch_size=patch_size, + rgb_copied=True, + resize="integer_area_mean_round_ties_to_even", + action_episodes=sum(item[-1] is not None for item in values), + split_counts=split_provenance["counts"], + derived_split_identities_sha256=split_provenance["derived_identities_sha256"], + split_provenance_sha256=sha256(split_path), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-manifest", required=True, type=Path) + parser.add_argument("--native-spec", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument( + "--image-size", required=True, type=int, nargs=2, metavar=("HEIGHT", "WIDTH") + ) + parser.add_argument("--patch-size", required=True, type=int) + parser.add_argument( + "--embodiments", + required=True, + nargs="+", + choices=("libero_gripper", "libero_sharpa"), + ) + parser.add_argument("--raw-export-dir", type=Path) + parser.add_argument("--raw-source-manifest", type=Path) + parser.add_argument("--std-floor", type=float, default=1e-5) + parser.add_argument("--resize-batch-size", type=int, default=4) + parser.add_argument("--validation-episodes-per-task", type=int) + parser.add_argument("--split-seed", type=int, default=42) + args = parser.parse_args() + print(json.dumps(prepare_pixels(**vars(args))), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/test_attention_execution.py b/tests/test_attention_execution.py new file mode 100644 index 000000000..f9a58931a --- /dev/null +++ b/tests/test_attention_execution.py @@ -0,0 +1,384 @@ +"""Packed attention isolation, regional compilation, and optional CUDA parity.""" + +import copy +import os + +import pytest +import torch +from torch.nn import functional as F + +from egomimic.models.attention_execution import ( + FlashAttentionKernels, + apply_temporal_rope, + execute_flash_attention, + load_flash_attention_kernels, + prepare_attention_layout, + prepare_temporal_rope, +) +from egomimic.models.compilation import compile_region +from egomimic.models.world_model import MoTBlock, temporal_rope + + +def _dense(q, k, v, *, causal): + assert causal is False + return F.scaled_dot_product_attention( + q.float().transpose(1, 2), k.float().transpose(1, 2), v.float().transpose(1, 2) + ).transpose(1, 2).to(q.dtype) + + +def _varlen(q, k, v, *, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal): + assert cu_seqlens_q.dtype == cu_seqlens_k.dtype == torch.int32 + qs, ks = cu_seqlens_q.tolist(), cu_seqlens_k.tolist() + assert max_seqlen_q == max(b - a for a, b in zip(qs, qs[1:])) + assert max_seqlen_k == max(b - a for a, b in zip(ks, ks[1:])) + return torch.cat([ + _dense(q[qa:qb][None], k[ka:kb][None], v[ka:kb][None], causal=causal)[0] + for qa, qb, ka, kb in zip(qs, qs[1:], ks, ks[1:]) + ]) + + +def _kernels(backend="flash4"): + def dense(*args, **kwargs): + out = _dense(*args, **kwargs) + return (out, None) if backend == "flash4" else out + + def varlen(*args, **kwargs): + out = _varlen(*args, **kwargs) + return (out, None) if backend == "flash4" else out + + return FlashAttentionKernels(backend, dense, varlen) + + +@pytest.mark.parametrize("backend", ["flash3", "flash4"]) +@pytest.mark.parametrize("packing", ["dense", "keys", "queries_and_keys"]) +def test_batched_outputs_gradients_and_example_isolation(backend, packing): + torch.manual_seed(43) + b, h, nq, nk, d = 3, 2, 5, 7, 8 + q, k, v = [x.requires_grad_() for x in ( + torch.randn(b, h, nq, d).bfloat16(), + torch.randn(b, h, nk, d).bfloat16(), + torch.randn(b, h, nk, d).bfloat16(), + )] + valid = torch.ones(b, nk, dtype=torch.bool) + qvalid = torch.ones(b, nq, dtype=torch.bool) + if packing != "dense": + valid[0, [1, 4, 6]] = False + valid[1, [0, 3]] = False + if packing == "queries_and_keys": + qvalid[0, [1, 3]] = False + qvalid[2, 4] = False + layout = prepare_attention_layout(valid, query_valid=qvalid, query_tokens=nq) + out = execute_flash_attention(q, k, v, layout, _kernels(backend)) + expected = F.scaled_dot_product_attention( + q.float(), k.float(), v.float(), attn_mask=valid[:, None, None] + ).to(q.dtype).masked_fill(~qvalid[:, None, :, None], 0) + torch.testing.assert_close(out, expected, atol=.008, rtol=.02) + actual_grads = torch.autograd.grad(out.float().square().sum(), (q, k, v), retain_graph=True) + reference_grads = torch.autograd.grad(expected.float().square().sum(), (q, k, v)) + for actual, reference in zip(actual_grads, reference_grads): + torch.testing.assert_close(actual, reference, atol=.015, rtol=.03) + assert torch.isfinite(actual).all() + assert actual_grads[1].transpose(1, 2)[~valid].count_nonzero() == 0 + assert actual_grads[0].transpose(1, 2)[~qvalid].count_nonzero() == 0 + changed = [x.detach().clone() for x in (q, k, v)] + for x in changed: + x[1:] *= 100 + isolated = execute_flash_attention(*changed, layout, _kernels(backend)) + assert torch.equal(out[0], isolated[0]) + + +def test_key_only_layout_retains_all_queries_with_different_lengths(): + q = torch.randn(2, 2, 3, 8).bfloat16() + k, v = (torch.randn(2, 2, 7, 8).bfloat16() for _ in range(2)) + valid = torch.tensor([[True, False, True, False, True, False, True], [True] * 7]) + layout = prepare_attention_layout(valid, query_tokens=3) + assert layout.query_indices is None + out = execute_flash_attention(q, k, v, layout, _kernels()) + expected = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), + attn_mask=valid[:, None, None]).to(q.dtype) + torch.testing.assert_close(out, expected, atol=.008, rtol=.02) + + +def test_verified_dense_layout_does_not_read_mask_values(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("Fully valid path must not inspect CUDA mask values") + valid = torch.ones(4, 11, dtype=torch.bool) + monkeypatch.setattr(torch.Tensor, "all", forbidden) + monkeypatch.setattr(torch.Tensor, "nonzero", forbidden) + monkeypatch.setattr(torch.Tensor, "item", forbidden) + layout = prepare_attention_layout(valid, query_valid=valid, all_valid=True) + assert layout.kind == "dense" and layout.key_indices is None + + +def test_layout_owns_mask_and_handles_interior_holes(): + valid = torch.tensor([[True, False, True], [False, True, True]]) + layout = prepare_attention_layout(valid, query_valid=valid) + valid[:] = False + assert layout.key_mask.sum() == 4 + assert layout.key_indices.tolist() == [0, 2, 4, 5] + assert layout.cu_seqlens_k.tolist() == [0, 2, 4] + + +@pytest.mark.parametrize("valid,qvalid,match", [ + (torch.zeros(2, 3, dtype=torch.bool), None, "valid key"), + (torch.ones(2, 3, dtype=torch.bool), torch.zeros(2, 3, dtype=torch.bool), "valid query"), + (torch.ones(2, 3), None, "boolean"), +]) +def test_bad_layout_fails_closed(valid, qvalid, match): + with pytest.raises(ValueError, match=match): + prepare_attention_layout(valid, query_valid=qvalid) + + +def test_precomputed_rope_matches_historical_values_and_gradients(): + x = torch.randn(2, 3, 7, 8, requires_grad=True) + times = torch.randn(2, 7) + rope = prepare_temporal_rope(times, 8) + assert all(value.dtype == torch.float32 for value in rope) + actual, expected = apply_temporal_rope(x, rope), temporal_rope(x, times) + assert torch.equal(actual, expected) + assert torch.equal(torch.autograd.grad(actual.square().sum(), x)[0], + torch.autograd.grad(expected.square().sum(), x)[0]) + + +def test_mot_reuses_layout_and_rope_without_mask_discovery(monkeypatch): + block = MoTBlock(16, 2, modalities=("image", "proprio"), + attention_backend="flash4", attention_dtype="bfloat16") + streams = [torch.randn(2, 5, 16, requires_grad=True), torch.randn(2, 3, 16, requires_grad=True)] + times = [torch.randn(2, 5), torch.randn(2, 3)] + valid = [torch.tensor([[1, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.bool), + torch.tensor([[0, 1, 0], [1, 0, 1]], dtype=torch.bool)] + joined = torch.cat(valid, 1) + layout = prepare_attention_layout(joined, query_valid=joined) + ropes = tuple(prepare_temporal_rope(t, 8) for t in times) + reference = copy.deepcopy(block) + reference.attention_backend = "sdpa" + monkeypatch.setattr("egomimic.models.world_model.load_flash_attention_kernels", lambda *args: _kernels()) + monkeypatch.setattr(torch.Tensor, "nonzero", lambda *args, **kwargs: pytest.fail("repeated packing")) + actual = block(streams, times, valid, attention_layout=layout, rope=ropes) + expected = reference(streams, times, valid) + for a, e, mask in zip(actual, expected, valid): + torch.testing.assert_close(a, e, atol=.01, rtol=.02) + assert a[~mask].count_nonzero() == 0 + sum(x.square().sum() for x in actual).backward() + for x, mask in zip(streams, valid): + assert torch.isfinite(x.grad).all() and x.grad[~mask].count_nonzero() == 0 + + +def test_query_chunks_preserve_generic_batched_nonprefix_mask_semantics(monkeypatch): + torch.manual_seed(53) + full = MoTBlock(16, 2, modalities=("image",), + attention_backend="flash4", attention_dtype="bfloat16") + chunked = copy.deepcopy(full) + chunked.attention_query_chunk_size = 3 + monkeypatch.setattr("egomimic.models.world_model.load_flash_attention_kernels", lambda *args: _kernels()) + values = torch.randn(2, 7, 16, requires_grad=True) + times = torch.randn(2, 7) + valid = torch.tensor([[1, 0, 1, 1, 0, 1, 1], [1, 1, 1, 0, 1, 1, 0]], dtype=torch.bool) + layout = prepare_attention_layout(valid, query_valid=valid) + a, = full([values], [times], [valid], attention_layout=layout) + b, = chunked([values], [times], [valid], attention_layout=layout) + torch.testing.assert_close(a, b, rtol=.02, atol=.008) + ga, = torch.autograd.grad(a.square().sum(), values) + gb, = torch.autograd.grad(b.square().sum(), values) + torch.testing.assert_close(ga, gb, rtol=.03, atol=.015) + + +@pytest.mark.parametrize("backend", ["flash3", "flash4"]) +def test_missing_cuda_never_falls_back(backend): + with pytest.raises(RuntimeError, match="CUDA"): + load_flash_attention_kernels(backend, torch.device("cpu")) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_flash3_import_uses_official_namespace_or_legacy_alias(monkeypatch, legacy): + from types import SimpleNamespace + + calls = [] + def load(name): + calls.append(name) + if legacy and name == "flash_attn_3.flash_attn_interface": + raise ModuleNotFoundError("legacy install", name="flash_attn_3") + return SimpleNamespace(flash_attn_func=_dense, flash_attn_varlen_func=_varlen) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0)) + monkeypatch.setattr("egomimic.models.attention_execution.importlib.import_module", load) + load_flash_attention_kernels.cache_clear() + try: + result = load_flash_attention_kernels("flash3", torch.device("cuda:0")) + assert result.backend == "flash3" + expected = ["flash_attn_3.flash_attn_interface"] + assert calls == expected + (["flash_attn_interface"] if legacy else []) + finally: + load_flash_attention_kernels.cache_clear() + + +def test_installed_backend_missing_dependency_does_not_try_alias(monkeypatch): + calls = [] + def load(name): + calls.append(name) + raise ModuleNotFoundError("broken installed extension", name="flash_attn_3_cuda") + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0)) + monkeypatch.setattr("egomimic.models.attention_execution.importlib.import_module", load) + load_flash_attention_kernels.cache_clear() + try: + with pytest.raises(RuntimeError, match="no backend fallback"): + load_flash_attention_kernels("flash3", torch.device("cuda:0")) + assert calls == ["flash_attn_3.flash_attn_interface"] + finally: + load_flash_attention_kernels.cache_clear() + + +def test_regional_compile_preserves_checkpoint_keys_and_backward(): + torch.manual_seed(19) + eager = MoTBlock(16, 2, modalities=("image",)) + compiled = copy.deepcopy(eager) + compiled.configure_compilation(mode="default", dynamic=True, backend="eager") + assert list(eager.state_dict()) == list(compiled.state_dict()) + for n in (5, 8): + x = torch.randn(2, n, 16, requires_grad=True) + times = torch.randn(2, n) + valid = torch.ones(2, n, dtype=torch.bool) + rope = (prepare_temporal_rope(times, 8),) + a, = eager([x], [times], [valid], rope=rope) + b, = compiled([x], [times], [valid], rope=rope) + assert torch.equal(a, b) + assert torch.equal(torch.autograd.grad(a.square().sum(), x)[0], + torch.autograd.grad(b.square().sum(), x)[0]) + compiled.load_state_dict(eager.state_dict(), strict=True) + + +def test_compile_fails_instead_of_falling_back(monkeypatch): + import torch._dynamo + with torch._dynamo.config.patch(suppress_errors=True): + with pytest.raises(RuntimeError, match="no eager fallback"): + compile_region(lambda x: x + 1, mode="default") + with pytest.raises(ValueError, match="compile mode"): + compile_region(lambda x: x, mode="invalid") + + +def test_compiled_backward_autocast_policy_covers_lazy_invocation_and_restores(monkeypatch): + from contextlib import contextmanager + from types import SimpleNamespace + + import egomimic.models.compilation as compilation + + config = SimpleNamespace(backward_pass_autocast="same_as_forward") + calls = [] + + @contextmanager + def patch(**settings): + previous = config.backward_pass_autocast + config.backward_pass_autocast = settings["backward_pass_autocast"] + try: + yield + finally: + config.backward_pass_autocast = previous + + config.patch = patch + monkeypatch.setattr(compilation, "_aot_config", config) + + def lazy_compile(function, **options): + assert config.backward_pass_autocast == "same_as_forward" + assert options["fullgraph"] + + def lazy_call(value): + calls.append(config.backward_pass_autocast) + assert config.backward_pass_autocast == "off" + if value == "fail": + raise RuntimeError("tracing failed") + return function(value) + return lazy_call + + monkeypatch.setattr(torch, "compile", lazy_compile) + execute = compile_region(lambda value: value + 1, mode="default") + assert execute(2) == 3 + assert config.backward_pass_autocast == "same_as_forward" + with pytest.raises(RuntimeError, match="tracing failed"): + execute("fail") + assert config.backward_pass_autocast == "same_as_forward" + assert calls == ["off", "off"] + + +def test_missing_aot_autocast_policy_fails_closed_only_for_compiled_mixed_mode(monkeypatch): + from types import SimpleNamespace + + import egomimic.models.compilation as compilation + + monkeypatch.setattr(compilation, "_aot_config", SimpleNamespace()) + monkeypatch.setattr(torch, "compile", lambda function, **_: function) + def function(value): + return value + 1 + execute = compile_region(function, mode="default") + assert execute(torch.ones(1)).item() == 2 + with torch.autocast("cpu", dtype=torch.bfloat16): + with pytest.raises(RuntimeError, match="no unsafe mixed-precision fallback"): + execute(torch.ones(1)) + assert compile_region(function, mode="none") is function + assert function(torch.ones(1)).item() == 2 + + +def test_aot_eager_bf16_forward_backward_matches_standard_amp(): + from torch._functorch import config + if not hasattr(config, "backward_pass_autocast"): + pytest.skip("This PyTorch lacks the required compiled AMP backward policy") + torch.manual_seed(19) + eager = MoTBlock(32, 4, modalities=("image",), attention_dtype="bfloat16") + compiled = copy.deepcopy(eager) + compiled.configure_compilation(mode="default", backend="aot_eager") + times = torch.randn(2, 7) + valid = torch.ones(2, 7, dtype=torch.bool) + rope = (prepare_temporal_rope(times, 8),) + x = torch.randn(2, 7, 32, requires_grad=True) + copied = x.detach().clone().requires_grad_() + before = config.backward_pass_autocast + with torch.autocast("cpu", dtype=torch.bfloat16): + expected, = eager([x], [times], [valid], rope=rope) + actual, = compiled([copied], [times], [valid], rope=rope) + expected_loss = expected.float().square().mean() + actual_loss = actual.float().square().mean() + assert config.backward_pass_autocast == before + expected_loss.backward() + actual_loss.backward() + torch.testing.assert_close(actual, expected, rtol=.02, atol=.01) + for a, b in [(copied.grad, x.grad)] + [ + (a.grad, b.grad) for a, b in zip(compiled.parameters(), eager.parameters()) + ]: + assert a.dtype == b.dtype == torch.float32 + assert torch.isfinite(a).all() + torch.testing.assert_close(a, b, rtol=.02, atol=.002) + assert (a - b).norm() / b.norm().clamp_min(1e-8) < .02 + + +@pytest.mark.skipif(os.environ.get("EGOMIMIC_RUN_FUSED_ATTENTION_TESTS") != "1", + reason="Explicit opt-in real CUDA attention proof") +@pytest.mark.parametrize("backend", ["flash3", "flash4"]) +@pytest.mark.parametrize("dense", [True, False]) +def test_real_cuda_batched_attention_matches_independent_gradients(backend, dense): + if not torch.cuda.is_available(): + pytest.fail("GPU proof requested without CUDA") + device = torch.device("cuda", torch.cuda.current_device()) + kernels = load_flash_attention_kernels(backend, device) + torch.manual_seed(71) + q, k, v = [torch.randn(3, 6, 97, 128, device=device, dtype=torch.bfloat16, + requires_grad=True) for _ in range(3)] + valid = torch.ones(3, 97, device=device, dtype=torch.bool) + if not dense: + valid[0, 1::3] = False + valid[1, 1::2] = False + layout = prepare_attention_layout(valid, query_valid=valid) + actual = execute_flash_attention(q, k, v, layout, kernels) + reference = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), + attn_mask=valid[:, None, None]).masked_fill(~valid[:, None, :, None], 0) + torch.testing.assert_close(actual.float(), reference, rtol=.03, atol=.02) + ga = torch.autograd.grad(actual.float().square().sum(), (q, k, v), retain_graph=True) + gr = torch.autograd.grad(reference.square().sum(), (q, k, v)) + for a, r in zip(ga, gr): + assert torch.isfinite(a).all() + assert (a.float() - r.float()).norm() / r.float().norm().clamp_min(1e-8) < .04 + for i in range(3): + one = prepare_attention_layout(valid[i:i+1], query_valid=valid[i:i+1]) + independent = execute_flash_attention(q[i:i+1], k[i:i+1], v[i:i+1], one, kernels) + torch.testing.assert_close(actual[i:i+1], independent, rtol=.03, atol=.02) + independent_grads = torch.autograd.grad(independent.float().square().sum(), (q, k, v)) + for full, single in zip(ga, independent_grads): + relative = (full[i].float() - single[i].float()).norm() + assert relative / full[i].float().norm().clamp_min(1e-8) < .04 diff --git a/tests/test_dino_proprio.py b/tests/test_dino_proprio.py new file mode 100644 index 000000000..5ae18a909 --- /dev/null +++ b/tests/test_dino_proprio.py @@ -0,0 +1,361 @@ +"""CPU contracts for two-stream models; these are not model-quality tests.""" + +import copy + +import pytest +import torch + +from egomimic.models.dino_proprio import MODALITIES, DinoProprioModel +from egomimic.models.world_model import MoTBlock + + +def make_batch(direct=False): + torch.manual_seed(13) + image = { + "value": torch.randn(2, 6, 8), + "time": torch.tensor([[0.0, 0.0, 1.0, 1.0, 2.0, 2.0]]).expand(2, -1).clone(), + "index": torch.tensor([[0.0, 0.0, 1.0, 1.0, 2.0, 2.0]]).expand(2, -1).clone(), + "valid": torch.ones(2, 6, dtype=torch.bool), + "observed": torch.tensor([[True, True, False, False, False, False]]) + .expand(2, -1) + .clone(), + "xy": torch.randn(2, 6, 2), + "camera": torch.tensor([[0, 1, 0, 1, 0, 1]]).expand(2, -1).clone(), + "target_valid": torch.ones(2, 6, dtype=torch.bool), + } + pose = { + "value": torch.randn(2, 3, 3), + "time": torch.tensor([[0.0, 1.0, 2.0]]).expand(2, -1).clone(), + "index": torch.tensor([[0.0, 1.0, 2.0]]).expand(2, -1).clone(), + "valid": torch.ones(2, 3, dtype=torch.bool), + "observed": torch.tensor([[True, False, False]]).expand(2, -1).clone(), + "target_valid": torch.ones(2, 3, dtype=torch.bool), + } + batch = {"image": image, "proprio": pose} + if direct: + batch["action"] = { + "value": torch.randn(2, 2, 4), + "time": torch.tensor([[0.0, 0.2]]).expand(2, -1).clone(), + "index": torch.tensor([[0.0, 1.0]]).expand(2, -1).clone(), + "valid": torch.ones(2, 2, dtype=torch.bool), + "observed": torch.zeros(2, 2, dtype=torch.bool), + "target_valid": torch.ones(2, 2, dtype=torch.bool), + } + return batch + + +def make_model(direct=False, **kwargs): + settings = dict( + image_dim=8, width=16, depth=2, heads=2, max_cameras=2, max_tokens=32 + ) + settings.update(kwargs) + return DinoProprioModel( + pose_dims={"gripper": 3}, + action_dims={"gripper": 4} if direct else None, + objective="direct_action" if direct else "visual_proprio", + **settings, + ) + + +def test_two_branches_and_width_guard_without_track_parameters(): + assert MODALITIES == ("image", "proprio") + for direct in (False, True): + model = make_model(direct) + assert all(len(block.qkv) == 2 for block in model.blocks) + assert all( + "track" not in name and "point" not in name + for name, _ in model.named_parameters() + ) + assert model.image_in.in_features == 8 + assert model.image_in.out_features >= 8 + if direct: + assert not hasattr(model, "image_out") + assert not hasattr(model, "pose_out") + assert not hasattr(model, "noise_level") + else: + assert model.noise_level[0].in_features == 1 + assert model.image_out.out_features == 8 + with pytest.raises(ValueError, match="at least image_dim"): + make_model(width=4) + # Existing checkpoints still construct three private branches by default. + assert len(MoTBlock(16, 2).qkv) == 3 + + +def test_world_phase_a_ignores_withheld_pose_values_times_and_presence(): + model, batch = make_model().eval(), make_batch() + levels = torch.full_like(batch["image"]["time"], 0.5) + images = torch.randn_like(batch["image"]["value"]) + before = model(batch, "gripper", images, levels, phase="A") + changed = copy.deepcopy(batch) + withheld = ~changed["proprio"]["observed"] + changed["proprio"]["value"][withheld] = float("nan") + changed["proprio"]["time"][withheld] = 999 + changed["proprio"]["target_valid"][:] = False + after = model(changed, "gripper", images, levels, phase="A") + torch.testing.assert_close(before["image"], after["image"], rtol=0, atol=0) + assert torch.count_nonzero(after["proprio"][withheld]) == 0 + for key, value in changed["proprio"].items(): + changed["proprio"][key] = value[:, :1] + without_queries = model(changed, "gripper", images, levels, phase="A") + torch.testing.assert_close( + before["image"], without_queries["image"], rtol=1e-6, atol=1e-6 + ) + + +def test_world_phase_b_pose_targets_are_loss_only(): + model, batch = make_model().eval(), make_batch() + images = torch.randn_like(batch["image"]["value"]) + levels = torch.zeros_like(batch["image"]["time"]) + before = model(batch, "gripper", images, levels, phase="B") + batch["proprio"]["value"][~batch["proprio"]["observed"]] = float("nan") + after = model(batch, "gripper", images, levels, phase="B") + for name in MODALITIES: + torch.testing.assert_close(before[name], after[name], rtol=0, atol=0) + + +def test_visual_observation_status_is_not_a_learned_condition(): + model, batch = make_model().eval(), make_batch() + images = batch["image"]["value"].clone() + levels = torch.zeros_like(batch["image"]["time"]) + before = model(batch, "gripper", images, levels, phase="A")["image"] + batch["image"]["observed"][:] = True + after = model(batch, "gripper", images, levels, phase="A")["image"] + torch.testing.assert_close(before, after, rtol=0, atol=0) + + +def test_joint_attention_couples_image_and_proprio_gradients(): + model, batch = make_model(), make_batch() + images = batch["image"]["value"].clone().requires_grad_() + batch["proprio"]["value"].requires_grad_() + result = model( + batch, "gripper", images, torch.ones_like(batch["image"]["time"]), phase="A" + ) + result["image"].square().sum().backward() + pose_grad = batch["proprio"]["value"].grad + assert pose_grad[batch["proprio"]["observed"]].abs().sum() > 0 + assert pose_grad[~batch["proprio"]["observed"]].abs().sum() == 0 + assert all(block.qkv[1].weight.grad.abs().sum() > 0 for block in model.blocks) + model.zero_grad(set_to_none=True) + images = batch["image"]["value"].clone().requires_grad_() + result = model( + batch, "gripper", images, torch.zeros_like(batch["image"]["time"]), phase="B" + ) + result["proprio"][:, 1:].square().sum().backward() + assert images.grad[~batch["image"]["observed"]].abs().sum() > 0 + assert all(block.qkv[0].weight.grad.abs().sum() > 0 for block in model.blocks) + + +def test_image_loss_is_clean_feature_mse_and_pose_warmup_skips_readout(monkeypatch): + model, batch = make_model(), make_batch() + forwards = [] + original = model.forward + + def capture(*args, **kwargs): + result = original(*args, **kwargs) + forwards.append((kwargs["phase"], result)) + return result + + monkeypatch.setattr(model, "forward", capture) + losses = model.training_losses(batch, "gripper", pose_weight=0) + assert len(forwards) == 1 and forwards[0][0] == "A" + mask = batch["image"]["valid"] & ~batch["image"]["observed"] + expected = ( + (forwards[0][1]["image"][mask] - batch["image"]["value"][mask]).square().mean() + ) + torch.testing.assert_close(losses["image"], expected) + assert losses["pose"].item() == 0 + torch.testing.assert_close(losses["total"], losses["image"]) + + +def test_low_noise_pose_pass_and_gradient_through_image_head(monkeypatch): + model, batch = make_model(pose_noise_max=0.15), make_batch() + forwards = [] + original = model.forward + + def capture(*args, **kwargs): + forwards.append((kwargs["phase"], args[3].detach().clone())) + return original(*args, **kwargs) + + monkeypatch.setattr(model, "forward", capture) + losses = model.training_losses(batch, "gripper", image_weight=0) + assert [phase for phase, _ in forwards] == ["A", "A", "B"] + assert forwards[1][1].max() <= 0.15 + assert torch.count_nonzero(forwards[2][1]) == 0 + losses["total"].backward() + assert model.image_out.weight.grad.abs().sum() > 0 + assert model.pose_out["gripper"].weight.grad.abs().sum() > 0 + + +def test_clean_pose_conditioning_is_explicit_and_skips_auxiliary_denoising(monkeypatch): + model, batch = make_model(pose_conditioning="clean"), make_batch() + phases = [] + original = model.forward + + def capture(*args, **kwargs): + phases.append(kwargs["phase"]) + if kwargs["phase"] == "B": + torch.testing.assert_close(args[2], batch["image"]["value"]) + return original(*args, **kwargs) + + monkeypatch.setattr(model, "forward", capture) + model.training_losses(batch, "gripper") + assert phases == ["A", "B"] + + +@pytest.mark.parametrize("steps", [1, 2, 20]) +def test_sampling_finite_preserves_observations_and_ignores_targets(steps): + model, batch = make_model().eval(), make_batch() + before = model.sample( + batch, "gripper", steps=steps, generator=torch.Generator().manual_seed(7) + ) + for name in MODALITIES: + assert torch.isfinite(before[name]).all() + observed = batch[name]["observed"] + torch.testing.assert_close( + before[name][observed], batch[name]["value"][observed], rtol=0, atol=0 + ) + batch[name]["value"][~observed] = float("nan") + batch[name]["target_valid"][:] = False + after = model.sample( + batch, "gripper", steps=steps, generator=torch.Generator().manual_seed(7) + ) + for name in MODALITIES: + torch.testing.assert_close(before[name], after[name], rtol=0, atol=0) + + +def test_all_queries_masked_or_absent_produces_finite_zero_losses(): + world, batch = make_model(), make_batch() + for name in MODALITIES: + batch[name]["observed"][:] = True + losses = world.training_losses(batch, "gripper") + assert all(torch.isfinite(loss) and loss.item() == 0 for loss in losses.values()) + losses["total"].backward() + for empty in (False, True): + direct, batch = make_model(True), make_batch(True) + if empty: + for key, value in batch["action"].items(): + batch["action"][key] = value[:, :0] + else: + batch["action"]["target_valid"][:] = False + batch["action"]["value"][:] = float("nan") + losses = direct.training_losses(batch, "gripper") + assert all( + torch.isfinite(loss) and loss.item() == 0 for loss in losses.values() + ) + losses["total"].backward() + + +def test_direct_action_targets_and_unobserved_context_never_enter_model(): + model, batch = make_model(True).eval(), make_batch(True) + before = model.sample(batch, "gripper")["action"] + for name in MODALITIES: + withheld = ~batch[name]["observed"] + batch[name]["value"][withheld] = float("nan") + batch[name]["time"][withheld] = 12345 + batch["action"]["value"][:] = float("nan") + batch["action"]["target_valid"][:] = False + after = model.sample(batch, "gripper")["action"] + torch.testing.assert_close(before, after, rtol=0, atol=0) + + +def test_direct_action_uses_both_observed_modalities_and_action_query_time(): + model, batch = make_model(True), make_batch(True) + for name in MODALITIES: + batch[name]["value"].requires_grad_() + batch["action"]["value"].requires_grad_() + prediction = model(batch, "gripper")["action"] + prediction.square().sum().backward() + assert batch["action"]["value"].grad is None + for name in MODALITIES: + grad = batch[name]["value"].grad + assert grad[batch[name]["observed"]].abs().sum() > 0 + assert grad[~batch[name]["observed"]].abs().sum() == 0 + batch["action"]["time"] += 1 + assert not torch.allclose(prediction, model(batch, "gripper")["action"]) + losses = model.training_losses(batch, "gripper") + assert set(losses) == {"action", "total"} + expected = ( + (model(batch, "gripper")["action"] - batch["action"]["value"]).square().mean() + ) + torch.testing.assert_close(losses["total"], expected) + + +def test_frame_noise_and_padding_are_consistent_and_target_masks_loss_only(): + batch = make_batch() + levels = DinoProprioModel.frame_noise_levels(batch["image"]) + assert torch.equal(levels[:, 0::2], levels[:, 1::2]) + assert torch.count_nonzero(levels[batch["image"]["observed"]]) == 0 + for name in MODALITIES: + batch[name]["valid"][:, -1] = False + batch[name]["observed"][:, -1] = False + batch[name]["value"][:, -1] = float("nan") + batch[name]["time"][:, -1] = float("nan") + batch["image"]["xy"][:, -1] = float("nan") + batch["image"]["camera"][:, -1] = -999 + model = make_model() + losses = model.training_losses(batch, "gripper") + assert all(torch.isfinite(v) for v in losses.values()) + sampled = model.sample(batch, "gripper", steps=1) + assert all(torch.count_nonzero(v[:, -1]) == 0 for v in sampled.values()) + + +@pytest.mark.parametrize("direct", [False, True]) +def test_checkpointed_chunked_attention_matches_full_attention(direct): + model = make_model(direct) + chunked = make_model( + direct, attention_query_chunk_size=3, activation_checkpointing=True + ) + chunked.load_state_dict(model.state_dict()) + batch = make_batch(direct) + args = ( + (batch, "gripper") + if direct + else ( + batch, + "gripper", + batch["image"]["value"], + torch.zeros_like(batch["image"]["time"]), + ) + ) + first, second = model(*args), chunked(*args) + for name in first: + torch.testing.assert_close(first[name], second[name], rtol=1e-5, atol=1e-6) + sum(v.square().sum() for v in first.values()).backward() + sum(v.square().sum() for v in second.values()).backward() + for first_param, second_param in zip(model.parameters(), chunked.parameters()): + if first_param.grad is not None: + torch.testing.assert_close( + first_param.grad, second_param.grad, rtol=1e-4, atol=1e-5 + ) + + +def test_rejects_tracks_observed_action_labels_missing_targets_and_invalid_controls(): + model, batch = make_model(), make_batch() + batch["track"] = {} + with pytest.raises(ValueError, match="track stream"): + model.training_losses(batch, "gripper") + del batch["track"] + batch["image"]["target_valid"][:, -1] = False + with pytest.raises(ValueError, match="recorded features"): + model.training_losses(batch, "gripper") + with pytest.raises(ValueError, match="positive integer"): + model.sample(make_batch(), "gripper", steps=0) + with pytest.raises(ValueError, match="0.2"): + make_model(pose_noise_max=0.3) + direct, batch = make_model(True), make_batch(True) + batch["action"]["observed"][:] = True + with pytest.raises(ValueError, match="never be marked observed"): + direct.sample(batch, "gripper") + + +def test_discrete_requires_indices_and_changes_clock_not_noise(): + model, batch = make_model(time_mode="discrete").eval(), make_batch() + images, levels = batch["image"]["value"], torch.zeros_like(batch["image"]["time"]) + before = model(batch, "gripper", images, levels, phase="B") + for name in MODALITIES: + batch[name]["time"] *= 100 + after = model(batch, "gripper", images, levels, phase="B") + torch.testing.assert_close(before["proprio"], after["proprio"], rtol=0, atol=0) + del batch["image"]["index"] + with pytest.raises(ValueError, match="explicit grid indices"): + model(batch, "gripper", images, levels) diff --git a/tests/test_dino_proprio_data.py b/tests/test_dino_proprio_data.py new file mode 100644 index 000000000..42ed85108 --- /dev/null +++ b/tests/test_dino_proprio_data.py @@ -0,0 +1,452 @@ +"""CPU contracts for two-stream data and fail-closed, aligned command labels.""" + +import copy +import hashlib +import json + +import h5py +import numpy as np +import pytest +import torch + +from egomimic.rldb.dino_proprio import ( + EvenlySpacedEpisodeBatchSampler, + PreparedDinoProprioDataset, + collate_dino_proprio, +) +from egomimic.rldb.dino_proprio_datamodule import DinoProprioDataModule + + +def digest(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def save_json(path, data): + path.write_text(json.dumps(data)) + return path + + +def prepared(tmp_path, *, raw=False): + episodes = [] + for split in ("train", "val"): + path = tmp_path / f"{split}.h5" + with h5py.File(path, "w") as data: + data["time"] = 1e9 + np.arange(6) * 0.125 + data["image"] = np.arange(6 * 2 * 3 * 8, dtype=np.float32).reshape( + 6, 2, 3, 8 + ) + data["image_xy"] = np.zeros((3, 2), np.float32) + data["pose"] = np.arange(6 * 4, dtype=np.float32).reshape(6, 4) + data["pose_valid"] = np.ones(6, bool) + episodes.append( + dict( + embodiment="gripper", + episode_id=split, + split=split, + path=path.name, + frames=6, + sha256=digest(path), + frame="fixed_reference", + pose_kind="measured_pose", + native_source={ + "episode": f"demo_{split}", + "source_demo_sha256": "a" * 64, + }, + ) + ) + manifest = dict( + schema="dino-proprio-prepared-v1", + episodes=episodes, + dino_identity={ + "model": "fixture", + "revision": "fixture", + "image_size": [16, 48], + "patch_grid": [1, 3], + "feature_layer": "last", + }, + pose_schemas={ + "gripper": {"dimension": 4, "fields": [f"p{i}" for i in range(4)]} + }, + normalization={ + "scope": "train_only", + "sha256": "0" * 64, + "arrays_normalized": not raw, + }, + ) + if raw: + norm = dict( + schema="mot-world-normalization-v1", + scope="train_only", + train_episodes=[["gripper", "train", episodes[0]["sha256"]]], + statistics={ + "image": {"mean": [1.0] * 8, "std": [2.0] * 8, "count": 36}, + "pose/gripper": {"mean": [3.0] * 4, "std": [4.0] * 4, "count": 6}, + # Deliberately unusable: unused legacy stats must not be required/parsed. + "track": {"not": "a normalizer"}, + }, + ) + normpath = save_json(tmp_path / "normalization.json", norm) + manifest["normalization"].update(path=normpath.name, sha256=digest(normpath)) + return save_json(tmp_path / "manifest.json", manifest) + + +def actions(path): + manifest = json.loads(path.read_text()) + episodes = [] + for ep in manifest["episodes"]: + target = path.parent / f"actions-{ep['split']}.h5" + with ( + h5py.File(path.parent / ep["path"], "r") as cache, + h5py.File(target, "w") as data, + ): + data["time"] = cache["time"][:] + data["action"] = np.arange(18, dtype=np.float32).reshape(6, 3) + 100 + data["valid"] = np.ones(6, bool) + data["cache_frame_index"] = np.arange(6) + data["source_state_index"] = np.arange(6) * 2 + data["source_action_index"] = np.arange(6) * 2 + receipt = save_json( + path.parent / f"alignment-{ep['split']}.json", + dict( + schema="dino-proprio-source-alignment-v1", + feature_sha256=ep["sha256"], + source_demo_sha256=ep["native_source"]["source_demo_sha256"], + source_episode=ep["native_source"]["episode"], + relation="action_t_applied_to_state_t", + producer_sha256="b" * 64, + basis={ + "cache_index_source": "synthetic fixture frame map", + "action_index_source": "synthetic fixture command map", + "state_action_relation_source": "synthetic fixture definition", + }, + cache_frame_index=list(range(6)), + source_state_index=list(range(0, 12, 2)), + source_action_index=list(range(0, 12, 2)), + ), + ) + episodes.append( + dict( + embodiment=ep["embodiment"], + episode_id=ep["episode_id"], + path=target.name, + sha256=digest(target), + feature_sha256=ep["sha256"], + source_demo_sha256=ep["native_source"]["source_demo_sha256"], + source_episode=ep["native_source"]["episode"], + alignment={ + "relation": "action_t_applied_to_state_t", + "cache_frame_mapping": "explicit_source_indices", + "producer_sha256": "b" * 64, + "source_alignment_path": receipt.name, + "source_alignment_sha256": digest(receipt), + }, + ) + ) + return save_json( + path.parent / "actions.json", + dict( + schema="dino-proprio-actions-v1", + episodes=episodes, + action_schemas={ + "gripper": { + "dimension": 3, + "fields": ["dx", "dy", "grip"], + "semantics": "action_command", + "value_space": "native_commands", + } + }, + ), + ) + + +def test_two_stream_cache_and_legacy_normalizer_never_require_track_arrays( + tmp_path, monkeypatch +): + path = prepared(tmp_path, raw=True) + old = h5py.Group.__getitem__ + + def guarded(group, name): + assert not str(name).startswith("track") + return old(group, name) + + monkeypatch.setattr(h5py.Group, "__getitem__", guarded) + dataset = PreparedDinoProprioDataset(path, split="val", min_frames=4, max_frames=4) + sample = dataset[0] + assert set(sample["world"]) == {"image", "proprio"} + assert set(dataset.normalization) == {"image", "pose/gripper"} + assert sample["world"]["image"]["value"].shape == (24, 8) + assert sample["world"]["proprio"]["value"].shape == (4, 4) + torch.testing.assert_close( + sample["world"]["image"]["value"][0], (torch.arange(8) - 1.0) / 2 + ) + assert sample["world"]["image"]["observed"].any() + assert (~sample["world"]["image"]["observed"]).any() + + +def test_forecast_mask_never_observes_future_and_padding_is_not_supervision(tmp_path): + path = prepared(tmp_path) + dataset = PreparedDinoProprioDataset(path, split="val", forecasting_probability=1.0) + sample = dataset[0] + image, pose = sample["world"]["image"], sample["world"]["proprio"] + last_image = image["time"][image["observed"]].max() + assert (image["time"][~image["observed"]] > last_image).all() + assert (pose["time"][pose["observed"]] <= last_image).all() + shorter = copy.deepcopy(sample) + for stream in shorter["world"].values(): + for key in stream: + stream[key] = stream[key][:-1] + batch = collate_dino_proprio([sample, shorter])["gripper"]["world"] + for stream in batch.values(): + assert not stream["valid"][1, -1] + assert not stream["target_valid"][1, -1] + assert not stream["observed"][1, -1] + + +def test_direct_action_requires_real_labels_and_preserves_native_command_values( + tmp_path, +): + path = prepared(tmp_path, raw=True) + with pytest.raises(ValueError, match="verified action alignment"): + PreparedDinoProprioDataset(path, task="direct_action") + apath = actions(path) + sample = PreparedDinoProprioDataset( + path, split="val", task="direct_action", action_manifest_path=apath + )[2] + assert set(sample["world"]) == {"image", "proprio", "action"} + image, pose, command = ( + sample["world"][key] for key in ("image", "proprio", "action") + ) + assert image["value"].shape == (6, 8) + assert pose["value"].shape == (1, 4) + assert image["observed"].all() and pose["observed"].all() + assert not command["observed"].any() + assert sample["task_kind"] == "direct_action" + assert collate_dino_proprio([sample])["gripper"]["task_kind"] == ["direct_action"] + torch.testing.assert_close(command["value"], torch.tensor([[106.0, 107.0, 108.0]])) + assert all((s["time"] == 0).all() for s in (image, pose, command)) + assert not torch.equal(command["value"], pose["value"][:, :3]) + + +def test_task_kind_survives_mixed_observation_regimes_in_batch_order(tmp_path): + path = prepared(tmp_path) + forecast = PreparedDinoProprioDataset( + path, split="val", forecasting_probability=1.0 + )[0] + completion = PreparedDinoProprioDataset( + path, split="val", forecasting_probability=0.0 + )[0] + assert forecast["task_kind"] == "forecast" + assert completion["task_kind"] == "completion" + assert collate_dino_proprio([completion, forecast])["gripper"]["task_kind"] == [ + "completion", + "forecast", + ] + + +def test_validation_subset_spans_every_episode_and_keeps_batches_homogeneous(tmp_path): + from egomimic.rldb.world_model import HomogeneousWorldBatchSampler + + path = prepared(tmp_path) + manifest = json.loads(path.read_text()) + original = manifest["episodes"][1] + for label, embodiment in (("second", "gripper"), ("third", "other")): + cache = path.parent / f"{label}.h5" + cache.write_bytes((path.parent / original["path"]).read_bytes()) + episode = copy.deepcopy(original) + episode.update(episode_id=label, embodiment=embodiment, path=cache.name) + manifest["episodes"].append(episode) + manifest["pose_schemas"]["other"] = copy.deepcopy( + manifest["pose_schemas"]["gripper"] + ) + save_json(path, manifest) + data = DinoProprioDataModule( + path, + batch_size=2, + validation_anchors_per_episode=2, + checkpointable_sampling=True, + ) + context = data.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + sampler = data.val_dataloader().batch_sampler + assert isinstance(sampler, EvenlySpacedEpisodeBatchSampler) + assert list(sampler) == [[0, 5], [6, 11], [12, 17]] + assert len(sampler) == 3 + assert list(sampler) == list(data.val_dataloader().batch_sampler) + for indices in sampler: + assert len({data.datasets["val"][i]["embodiment"] for i in indices}) == 1 + assert isinstance( + data.train_dataloader().batch_sampler, HomogeneousWorldBatchSampler + ) + assert context.state["validation_sampling"] == { + "anchors_per_episode": 2, + "selection": "deterministic_evenly_spaced_eligible_anchors", + } + assert data.state_dict()["validation_anchors_per_episode"] == 2 + changed = DinoProprioDataModule( + path, + batch_size=2, + validation_anchors_per_episode=3, + checkpointable_sampling=True, + ) + with pytest.raises(ValueError, match="checkpoint/data context mismatch"): + changed.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + with pytest.raises(ValueError, match="validation_anchors_per_episode"): + changed.load_state_dict(data.state_dict()) + with pytest.raises(ValueError, match="mutually exclusive"): + DinoProprioDataModule( + path, validation_anchors_per_episode=2, validation_batches_per_embodiment=2 + ) + + +def test_validation_subset_uses_all_anchors_when_episode_is_short(tmp_path): + path = prepared(tmp_path) + dataset = PreparedDinoProprioDataset(path, split="val") + sampler = EvenlySpacedEpisodeBatchSampler( + dataset, anchors_per_episode=10, batch_size=4 + ) + assert list(sampler) == [[0, 1, 2, 3], [4, 5]] + with pytest.raises(ValueError, match="validation only"): + EvenlySpacedEpisodeBatchSampler( + PreparedDinoProprioDataset(path), anchors_per_episode=2 + ) + + +@pytest.mark.parametrize( + "damage", + [ + "command_index", + "time", + "feature_sha", + "semantics", + "source_sha", + "missing_episode", + "receipt_provenance", + ], +) +def test_action_alignment_fails_before_training(tmp_path, damage): + path = prepared(tmp_path) + apath = actions(path) + action = json.loads(apath.read_text()) + ep = action["episodes"][0] + if damage in ("command_index", "time"): + cache = apath.parent / ep["path"] + with h5py.File(cache, "a") as data: + name = "source_action_index" if damage == "command_index" else "time" + data[name][0] += 1 + ep["sha256"] = digest(cache) + elif damage == "feature_sha": + ep["feature_sha256"] = "d" * 64 + elif damage == "semantics": + action["action_schemas"]["gripper"]["semantics"] = "measured_pose" + elif damage == "source_sha": + ep["source_demo_sha256"] = "d" * 64 + elif damage == "receipt_provenance": + target = apath.parent / ep["alignment"]["source_alignment_path"] + receipt = json.loads(target.read_text()) + receipt["basis"].pop("state_action_relation_source") + save_json(target, receipt) + ep["alignment"]["source_alignment_sha256"] = digest(target) + else: + action["episodes"].pop() + save_json(apath, action) + with pytest.raises(ValueError, match="alignment|native action|every prepared"): + PreparedDinoProprioDataset( + path, task="direct_action", action_manifest_path=apath + ) + + +def test_feature_and_normalization_checksum_are_enforced(tmp_path): + path = prepared(tmp_path, raw=True) + manifest = json.loads(path.read_text()) + target = path.parent / manifest["episodes"][0]["path"] + with h5py.File(target, "a") as data: + data["pose"][0, 0] += 1 + with pytest.raises(ValueError, match="Artifact hash mismatch"): + PreparedDinoProprioDataset(path) + + +def test_checkpointable_sampling_context_and_action_binding(tmp_path): + path = prepared(tmp_path) + apath = actions(path) + kwargs = dict( + manifest_path=path, + task="direct_action", + action_manifest_path=apath, + checkpointable_sampling=True, + batch_size=1, + ) + data = DinoProprioDataModule(**kwargs) + context = data.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert context.state["kind"] == "dino-proprio-prepared-v1" + assert context.sample_schema == {"world": "image_proprio_action_labels"} + first = next(iter(data.train_dataloader()))["gripper"] + data.acknowledge_batch(first["sampling_cursor"]) + expected = next(iter(data.train_dataloader()))["gripper"] + state = data.state_dict() + resumed = DinoProprioDataModule(**kwargs) + resumed.load_state_dict(state) + actual = next(iter(resumed.train_dataloader()))["gripper"] + assert actual["sampling_cursor"] == expected["sampling_cursor"] == 2 + for name in actual["world"]: + for key in actual["world"][name]: + torch.testing.assert_close( + actual["world"][name][key], expected["world"][name][key], rtol=0, atol=0 + ) + changed = json.loads(apath.read_text()) + changed["note"] = "different immutable manifest" + save_json(apath, changed) + with pytest.raises(ValueError, match="action_manifest_sha256"): + DinoProprioDataModule(**kwargs).load_state_dict(state) + + +@pytest.mark.parametrize( + "task,objective", + [("visual_diffusion", "visual_proprio"), ("direct_action", "direct_action")], +) +def test_dataset_batch_passes_stage_training_and_inference(tmp_path, task, objective): + from egomimic.models.dino_proprio import DinoProprioModel + from egomimic.pipeline.stages_dino_proprio import DinoProprioStage + + path = prepared(tmp_path, raw=True) + apath = actions(path) if task == "direct_action" else None + dataset = PreparedDinoProprioDataset( + path, + split="val", + task=task, + action_manifest_path=apath, + forecasting_probability=1.0, + ) + item = collate_dino_proprio([dataset[0]])["gripper"] + item["optimizer_step"] = 1 + model = DinoProprioModel( + objective=objective, + pose_dims={"gripper": 4}, + action_dims={"gripper": 3} if task == "direct_action" else None, + image_dim=8, + width=24, + depth=1, + heads=3, + max_cameras=2, + max_tokens=100, + ) + stage = DinoProprioStage(model, warmup_steps=0, pose_ramp_steps=0, sampler_steps=2) + trained = stage.execute(copy.deepcopy(item), mode="train") + loss = trained["loss/dino_proprio"] + assert torch.isfinite(loss) and loss > 0 + loss.backward() + assert model.image_in.weight.grad is not None + inferred = stage.execute(copy.deepcopy(item), mode="inference")[ + "dino_proprio_prediction" + ] + expected = ("action",) if task == "direct_action" else ("image", "proprio") + for name in expected: + assert inferred[name].shape == item["world"][name]["value"].shape + assert torch.isfinite(inferred[name]).all() diff --git a/tests/test_dino_proprio_integration.py b/tests/test_dino_proprio_integration.py new file mode 100644 index 000000000..97a79846c --- /dev/null +++ b/tests/test_dino_proprio_integration.py @@ -0,0 +1,139 @@ +"""Exercise composed configs through optimization and scheduled validation. + +Tiny synthetic fixtures test framework wiring, not LIBERO experiment quality. +""" + +from pathlib import Path + +import pytest +import torch +from hydra import compose, initialize_config_dir +from hydra.utils import instantiate +from lightning.pytorch import Trainer +from lightning.pytorch.callbacks import ModelCheckpoint +from omegaconf import open_dict + +from egomimic.pl_utils.pl_model import ModelWrapper +from tests.test_dino_proprio_data import actions, prepared + + +def configuration(arm): + directory = Path(__file__).parents[1] / "egomimic/hydra_configs" + with initialize_config_dir(config_dir=str(directory), version_base="1.2"): + return compose( + config_name="train_zarr_cartesian", + overrides=[f"+experiment=dino_proprio/{arm}"], + ) + + +@pytest.mark.parametrize( + "arm,parameters", + [ + ("visual_state", 115206921), + ("direct_action", 114021127), + ], +) +def test_exact_pilot_configs_dimensions_and_optimizer(arm, parameters): + cfg = configuration(arm) + stage = cfg.model.pipeline.stages[0] + with torch.device("meta"): + model = instantiate(stage.model) + assert sum(p.numel() for p in model.parameters()) == parameters + assert sum(p.numel() for p in model.parameters() if p.requires_grad) == parameters + assert model.image_dim == model.width == 768 + assert len(model.blocks) == 8 + assert len(model.blocks[0].qkv) == 2 + assert not any("track" in key for key in model.state_dict()) + assert cfg.model.optimizer.lr == 1e-4 + assert cfg.model.optimizer.weight_decay == 0.01 + assert cfg.model.scheduler is None + assert cfg.trainer.max_steps == 1000 + assert cfg.callbacks.model_checkpoint.every_n_train_steps == 250 + assert cfg.callbacks.model_checkpoint.save_top_k == -1 + assert cfg.data.min_frames == cfg.data.max_frames == 4 + assert cfg.data.max_past_seconds is cfg.data.max_future_seconds is None + if arm == "direct_action": + assert cfg.data.action_manifest_path is None + assert not hasattr(model, "image_out") + else: + assert model.pose_conditioning == "low_noise" + + +@pytest.mark.parametrize("arm", ["visual_state", "direct_action"]) +def test_optimizer_scheduled_validation_and_checkpoint(tmp_path, arm): + cfg = configuration(arm) + manifest = prepared(tmp_path) + with open_dict(cfg): + cfg.data.manifest_path = str(manifest) + if arm == "direct_action": + cfg.data.action_manifest_path = str(actions(manifest)) + stage = cfg.model.pipeline.stages[0] + stage.warmup_steps = 0 + stage.pose_ramp_steps = 0 + stage.sampler_steps = 2 + stage.model.pose_dims = {"gripper": 4} + stage.model.action_dims = {"gripper": 3} if arm == "direct_action" else None + stage.model.image_dim = 8 + stage.model.width = 12 + stage.model.depth = 1 + stage.model.heads = 3 + stage.model.attention_query_chunk_size = 4 + stage.model.max_tokens = 100 + data = instantiate(cfg.data) + context = data.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + evaluator = instantiate(cfg.evaluator) + model = ModelWrapper(config_tree=cfg, evaluator=evaluator, enable_grad_norm=False) + model.data_context = context + context.bind(model.model, evaluator) + data.configure_evaluation(evaluator.data_requirements()) + before = next(model.parameters()).detach().clone() + callback = ModelCheckpoint( + dirpath=tmp_path / "checkpoints", + every_n_train_steps=2, + save_top_k=-1, + save_last=True, + save_on_train_epoch_end=False, + ) + trainer = Trainer( + accelerator="cpu", + devices=1, + precision="32-true", + logger=False, + max_steps=3, + max_epochs=-1, + min_epochs=0, + limit_val_batches=1, + val_check_interval=2, + check_val_every_n_epoch=None, + num_sanity_val_steps=0, + enable_progress_bar=False, + enable_model_summary=False, + callbacks=[callback], + default_root_dir=tmp_path, + ) + trainer.fit(model, datamodule=data) + assert trainer.global_step == 3 + assert data.consumed_batches == 3 + assert not torch.equal(before, next(model.parameters()).detach()) + expected = ( + ["image_normalized_mse", "proprio_normalized_mse"] + if arm == "visual_state" + else ["action_command_mse"] + ) + for metric in expected: + found = [ + v + for k, v in trainer.callback_metrics.items() + if k.startswith("val/gripper/") and k.endswith("/" + metric) + ] + assert found and all(torch.isfinite(value) for value in found) + assert callback.last_model_path + checkpoint = torch.load(callback.last_model_path, weights_only=False) + assert checkpoint["global_step"] == 2 + assert checkpoint["data_context"] == context.snapshot() + assert "mot_rng_state" in checkpoint + optimizer = trainer.optimizers[0] + assert optimizer.param_groups[0]["lr"] == 1e-4 + assert optimizer.param_groups[0]["betas"] == (0.9, 0.999) diff --git a/tests/test_image_proprio.py b/tests/test_image_proprio.py new file mode 100644 index 000000000..b596f4018 --- /dev/null +++ b/tests/test_image_proprio.py @@ -0,0 +1,319 @@ +"""Image-patch core compatibility and explicit compression contracts.""" + +import copy + +import pytest +import torch + +from egomimic.models.dino_proprio import DinoProprioModel +from egomimic.models.image_proprio import ImageProprioModel + + +def raw_patch_batch(direct): + # Two 2x2 RGB patches per timestamp, scaled to [-1,1]. These are raw + # channel values, without a pretrained feature extractor or learned encoder. + values = torch.linspace(-1, 1, 4 * 12).reshape(1, 4, 12) + batch = { + "image": { + "value": values, + "time": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), + "valid": torch.ones(1, 4, dtype=torch.bool), + "observed": torch.tensor([[True, True, False, False]]), + "xy": torch.tensor([[[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [1.0, 0.0]]]), + "camera": torch.zeros(1, 4, dtype=torch.long), + }, + "proprio": { + "value": torch.tensor([[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]]), + "time": torch.tensor([[0.0, 1.0]]), + "valid": torch.ones(1, 2, dtype=torch.bool), + "observed": torch.tensor([[True, False]]), + }, + } + if direct: + batch["action"] = { + "value": torch.tensor([[[0.2, -0.3]]]), + "time": torch.zeros(1, 1), + "valid": torch.ones(1, 1, dtype=torch.bool), + "observed": torch.zeros(1, 1, dtype=torch.bool), + } + return batch + + +def model_arguments(direct): + return dict( + pose_dims={"robot": 3}, + action_dims={"robot": 2} if direct else None, + objective="direct_action" if direct else "visual_proprio", + image_dim=12, + depth=1, + heads=2, + max_cameras=1, + max_tokens=8, + ) + + +def test_legacy_class_retains_identity_embeddings_and_parameter_compatibility(): + assert issubclass(DinoProprioModel, ImageProprioModel) + previous_name = DinoProprioModel(width=16, **model_arguments(False)) + current_name = ImageProprioModel( + width=16, + use_modality_embedding=True, + use_embodiment_embedding=True, + **model_arguments(False), + ) + assert previous_name.use_camera_embedding + assert previous_name.use_modality_embedding + assert previous_name.use_embodiment_embedding + current_name.load_state_dict(previous_name.state_dict(), strict=True) + assert not current_name.allow_patch_compression + + +def test_pixel_default_keeps_camera_but_removes_modality_and_embodiment_ids(): + arguments = model_arguments(False) + arguments["max_cameras"] = 2 + model = ImageProprioModel(width=16, **arguments) + assert model.use_camera_embedding + assert not model.use_modality_embedding and not model.use_embodiment_embedding + assert not any( + hasattr(model, key) for key in ("modality", "embodiment", "embodiment_ids") + ) + assert not any( + key.startswith(("modality", "embodiment.")) for key in model.state_dict() + ) + assert all(len(block.qkv) == 2 for block in model.blocks) + batch = raw_patch_batch(False) + images = torch.randn_like(batch["image"]["value"]) + levels = torch.zeros_like(batch["image"]["time"]) + expected = model(batch, "robot", images, levels) + batch["image"]["camera"].fill_(1) + changed = model(batch, "robot", images, levels) + assert not torch.equal(expected["image"], changed["image"]) + model.training_losses(batch, "robot")["total"].backward() + assert model.camera.weight.grad.abs().sum() > 0 + for flag in ( + "use_camera_embedding", + "use_modality_embedding", + "use_embodiment_embedding", + ): + for value in (0, 1, None, "false"): + with pytest.raises(ValueError, match=f"{flag} must be bool"): + ImageProprioModel(width=16, **{flag: value}, **arguments) + + +def test_camera_conditioning_can_be_disabled_without_hidden_slot_encoding(): + model = ImageProprioModel( + width=16, use_camera_embedding=False, **model_arguments(False) + ) + assert not hasattr(model, "camera") + batch = raw_patch_batch(False) + images = torch.randn_like(batch["image"]["value"]) + levels = torch.zeros_like(batch["image"]["time"]) + expected = model(batch, "robot", images, levels) + batch["image"]["camera"] = torch.tensor([[100, -20, 9, 7]]) + changed = model(batch, "robot", images, levels) + del batch["image"]["camera"] + absent = model(batch, "robot", images, levels) + for name in expected: + torch.testing.assert_close(expected[name], changed[name], rtol=0, atol=0) + torch.testing.assert_close(expected[name], absent[name], rtol=0, atol=0) + + +def test_pixel_has_no_embodiment_id_offset_when_adapter_weights_match(): + arguments = model_arguments(False) + arguments["pose_dims"] = {"robot": 3, "another": 3} + model = ImageProprioModel(width=16, **arguments) + model.pose_in["another"].load_state_dict(model.pose_in["robot"].state_dict()) + model.pose_out["another"].load_state_dict(model.pose_out["robot"].state_dict()) + batch = raw_patch_batch(False) + images = torch.randn_like(batch["image"]["value"]) + levels = torch.zeros_like(batch["image"]["time"]) + expected = model(batch, "robot", images, levels, phase="B") + changed = model(batch, "another", images, levels, phase="B") + for name in expected: + torch.testing.assert_close(expected[name], changed[name], rtol=0, atol=0) + + +@pytest.mark.parametrize("direct", [False, True]) +def test_patch_compression_requires_explicit_boolean_opt_in(direct): + for kwargs in ({}, {"allow_patch_compression": False}): + with pytest.raises(ValueError, match="allow_patch_compression=True"): + ImageProprioModel(width=8, **model_arguments(direct), **kwargs) + for bad_value in (1, "true", None): + with pytest.raises(ValueError, match="must be bool"): + ImageProprioModel( + width=8, + allow_patch_compression=bad_value, + **model_arguments(direct), + ) + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.parametrize("compress", [False, True]) +def test_raw_rgb_patch_objectives_backward_and_sampling(direct, compress): + model = ImageProprioModel( + width=8 if compress else 16, + allow_patch_compression=compress, + **model_arguments(direct), + ) + assert model.image_in.in_features == 12 + assert model.image_in.out_features == (8 if compress else 16) + assert not any(isinstance(module, torch.nn.Conv2d) for module in model.modules()) + batch = raw_patch_batch(direct) + losses = model.training_losses(batch, "robot") + assert all(torch.isfinite(value) for value in losses.values()) + losses["total"].backward() + assert model.image_in.weight.grad.abs().sum() > 0 + sampled = model.sample(batch, "robot", steps=2) + assert all(torch.isfinite(value).all() for value in sampled.values()) + if direct: + assert sampled["action"].shape == (1, 1, 2) + assert not hasattr(model, "image_out") + else: + assert model.image_out.out_features == 12 + assert model.image_out.in_features == (8 if compress else 16) + assert model.image_out.weight.grad.abs().sum() > 0 + assert sampled["image"].shape == (1, 4, 12) + for name in ("image", "proprio"): + observed = batch[name]["observed"] + torch.testing.assert_close( + sampled[name][observed], batch[name]["value"][observed], rtol=0, atol=0 + ) + + +def test_visual_only_has_no_state_action_heads_queries_or_second_phase(monkeypatch): + arguments = model_arguments(False) + arguments["objective"] = "visual_only" + model = ImageProprioModel(width=16, **arguments) + assert not any(hasattr(model, name) for name in ("pose_out", "action_out", "query")) + batch = raw_patch_batch(False) + batch["proprio"]["value"][~batch["proprio"]["observed"]] = float("nan") + phases = [] + original = model.forward + + def capture(*args, **kwargs): + phases.append(kwargs.get("phase", "A")) + return original(*args, **kwargs) + + monkeypatch.setattr(model, "forward", capture) + losses = model.training_losses(batch, "robot", pose_weight=1) + assert set(losses) == {"image", "total"} + assert phases == ["A"] + losses["total"].backward() + assert model.pose_in["robot"].weight.grad.abs().sum() > 0 + phases.clear() + sampled = model.sample(batch, "robot", steps=2) + assert phases == ["A", "A"] + assert set(sampled) == {"image"} + assert torch.isfinite(sampled["image"]).all() + with pytest.raises(ValueError, match="no proprio readout phase"): + model( + batch, "robot", batch["image"]["value"], batch["image"]["time"], phase="B" + ) + + +def combined_model_and_batch(**kwargs): + arguments = model_arguments(True) + arguments["objective"] = "visual_state_action" + model = ImageProprioModel(width=16, **arguments, **kwargs) + batch = raw_patch_batch(True) + batch["action"] = { + "value": torch.tensor([[[0.2, -0.3], [-0.4, 0.5]]]), + "time": torch.tensor([[0.0, 1.0]]), + "valid": torch.ones(1, 2, dtype=torch.bool), + "observed": torch.zeros(1, 2, dtype=torch.bool), + } + return model, batch + + +def test_joint_state_action_readout_is_two_branches_and_never_embeds_labels(): + model, batch = combined_model_and_batch() + assert all(len(block.qkv) == 2 for block in model.blocks) + assert model.proprio_role.num_embeddings == 3 + assert model.action_query is not model.query + images = torch.randn_like(batch["image"]["value"]) + levels = torch.zeros_like(batch["image"]["time"]) + changed = copy.deepcopy(batch) + changed["action"]["value"].fill_(float("nan")) + changed["proprio"]["value"][~changed["proprio"]["observed"]] = float("nan") + expected = model(batch, "robot", images, levels, phase="B") + actual = model(changed, "robot", images, levels, phase="B") + assert set(actual) == {"image", "proprio", "action"} + for name in expected: + torch.testing.assert_close(expected[name], actual[name], rtol=0, atol=0) + # Action query metadata is absent in phase A as well as action labels. + before = model(batch, "robot", images, levels, phase="A") + changed["action"]["time"] += 100 + changed["action"]["valid"].fill_(False) + after = model(changed, "robot", images, levels, phase="A") + assert "action" not in before + for name in before: + torch.testing.assert_close(before[name], after[name], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "conditioning,phases", [("low_noise", ["A", "A", "B"]), ("clean", ["A", "B"])] +) +def test_combined_readout_shares_one_aligned_image_estimate( + monkeypatch, conditioning, phases +): + model, batch = combined_model_and_batch(pose_conditioning=conditioning) + calls = [] + original = model.forward + + def capture(*args, **kwargs): + calls.append(kwargs.get("phase", "A")) + return original(*args, **kwargs) + + monkeypatch.setattr(model, "forward", capture) + losses = model.training_losses( + batch, "robot", image_weight=0, pose_weight=2, action_weight=3 + ) + assert calls == phases + assert set(losses) == {"image", "pose", "action", "total"} + torch.testing.assert_close( + losses["total"], 2 * losses["pose"] + 3 * losses["action"] + ) + losses["total"].backward() + for head in (model.pose_out["robot"], model.action_out["robot"]): + assert head.weight.grad.abs().sum() > 0 + assert model.action_query.grad.abs().sum() > 0 + assert model.query.grad.abs().sum() > 0 + if conditioning == "low_noise": + assert model.image_out.weight.grad.abs().sum() > 0 + calls.clear() + sample = model.sample(batch, "robot", steps=2) + assert calls == ["A", "A", "B"] + assert all(torch.isfinite(value).all() for value in sample.values()) + for name in ("image", "proprio"): + observed = batch[name]["observed"] + torch.testing.assert_close( + sample[name][observed], batch[name]["value"][observed], rtol=0, atol=0 + ) + + +def test_combined_readout_zero_weights_and_missing_targets(monkeypatch): + model, batch = combined_model_and_batch() + calls = [] + original = model.forward + + def capture(*args, **kwargs): + calls.append(kwargs.get("phase", "A")) + return original(*args, **kwargs) + + monkeypatch.setattr(model, "forward", capture) + losses = model.training_losses(batch, "robot", pose_weight=0, action_weight=0) + assert calls == ["A"] + assert losses["pose"] == losses["action"] == 0 + # Action supervision can train by itself when withheld state labels are absent. + calls.clear() + batch["proprio"]["target_valid"] = batch["proprio"]["observed"].clone() + batch["proprio"]["value"][~batch["proprio"]["observed"]] = float("nan") + losses = model.training_losses(batch, "robot") + assert calls == ["A", "A", "B"] + assert losses["pose"] == 0 and torch.isfinite(losses["total"]) + calls.clear() + batch["action"]["target_valid"] = torch.zeros_like(batch["action"]["valid"]) + batch["action"]["value"].fill_(float("nan")) + losses = model.training_losses(batch, "robot") + assert calls == ["A"] + assert torch.isfinite(losses["total"]) diff --git a/tests/test_image_proprio_adaln.py b/tests/test_image_proprio_adaln.py new file mode 100644 index 000000000..cbc47d282 --- /dev/null +++ b/tests/test_image_proprio_adaln.py @@ -0,0 +1,183 @@ +"""DiT-style image noise conditioning with per-frame shared modulation.""" + +import copy +from contextlib import nullcontext + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from scripts.world_model.benchmark_pixel_execution import fixed_corruption +from tests.test_pixel_packed_execution import ( + assert_gradient_parity, + batch, + reference_kernels, # noqa: F401 -- explicitly shared CPU kernel fixture +) +from tests.test_pixel_packed_execution import ( + model as packed_model, +) + + +@pytest.fixture(autouse=True) +def one_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def model(**kwargs): + settings = dict(noise_conditioning="adaln_zero", use_phase_embedding=False) + settings.update(kwargs) + return packed_model(backend="sdpa", attention_dtype="native", **settings) + + +def activate_conditioning(network): + with torch.no_grad(): + network.image_out.weight.normal_(std=0.08) + network.image_out.bias.normal_(std=0.02) + network.final_adaln[-1].weight.normal_(std=0.1) + for block in network.blocks: + block.adaln_image[-1].weight.normal_(std=0.1) + block.adaln_image[-1].bias.normal_(std=0.05) + + +def test_additive_default_preserves_seeded_state_rng_and_outputs(): + torch.manual_seed(12) + original = model(noise_conditioning="additive") + rng = torch.get_rng_state().clone() + torch.manual_seed(12) + explicit = packed_model(backend="sdpa", attention_dtype="native", use_phase_embedding=False) + assert torch.equal(rng, torch.get_rng_state()) + assert all(torch.equal(value, explicit.state_dict()[name]) for name, value in original.state_dict().items()) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 8) + assert torch.equal(original(world, "libero_gripper", noisy, levels)["image"], + explicit(world, "libero_gripper", noisy, levels)["image"]) + + +def test_standard_zero_initialization_and_exact_parameter_count(): + with torch.device("meta"): + full = ImageProprioModel(objective="visual_only", phase_a_condition_on_proprio=False, + max_cameras=2, use_phase_embedding=False, noise_conditioning="adaln_zero") + assert sum(parameter.numel() for parameter in full.parameters()) == 88_180_992 + network = model() + assert not network.norm[0].elementwise_affine + assert all(not block.norm1[0].elementwise_affine and not block.norm2[0].elementwise_affine for block in network.blocks) + world = batch() + noisy, levels = fixed_corruption(world, 7) + prediction = network(world, "libero_gripper", noisy, levels)["image"] + assert prediction.count_nonzero() == 0 + network.training_losses(world, "libero_gripper")["total"].backward() + assert all(parameter.grad is not None and torch.isfinite(parameter.grad).all() for parameter in network.parameters()) + assert network.image_out.weight.grad.abs().sum() > 0 + assert network.blocks[0].qkv[0].weight.grad.count_nonzero() == 0 + + +def test_frame_modulation_is_shared_but_examples_and_levels_remain_distinct(): + network = model() + activate_conditioning(network) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 11) + frames = network._prepare_noise_frame_map(world["image"]) + assert frames.representatives.numel() == 5 #2 and3frames, each8camera/patchtokens. + shapes = [] + handles = [block.adaln_image[-1].register_forward_pre_hook(lambda _, args: shapes.append(args[0].shape)) + for block in network.blocks] + first = network(world, "libero_gripper", noisy, levels)["image"] + for handle in handles: + handle.remove() + assert shapes == [torch.Size([5, 16])] * 2 + changed = levels.clone() + changed[1] = changed[1] * 0.5 + second = network(world, "libero_gripper", noisy, changed)["image"] + assert torch.equal(first[0], second[0]) + assert not torch.allclose(first[1], second[1], atol=1e-7, rtol=1e-7) + # Input embeddings no longer receive additive noise vectors. + image = world["image"] + assert torch.equal(network._embed_image(image, noisy, levels, image["valid"], "A", False), + network._embed_image(image, noisy, changed, image["valid"], "A", False)) + query = image["valid"] & ~image["observed"] + row, token = query.nonzero()[0].tolist() + invalid = levels.clone() + invalid[row, token] += 0.01 + with pytest.raises(ValueError, match="agree within each timestamp"): + network(world, "libero_gripper", noisy, invalid) + + +def test_sampler_caches_only_frame_mapping_and_recomputes_noise_embedding(monkeypatch): + network = model() + activate_conditioning(network) + world = batch() + calls, levels_seen = [], [] + original = network._prepare_noise_frame_map + + def record_map(image): + calls.append(True) + return original(image) + + monkeypatch.setattr(network, "_prepare_noise_frame_map", record_map) + original_embed = network._embed_noise_frames_execute + + def record_embed(levels): + levels_seen.append(levels.clone()) + return original_embed(levels) + + object.__setattr__(network, "_embed_noise_frames_execute", record_embed) + output = network.sample(world, "libero_gripper", steps=4, generator=torch.Generator().manual_seed(4))["image"] + assert len(calls) == 1 and len(levels_seen) == 4 + assert [float(level.max()) for level in levels_seen] == [1.0, 0.75, 0.5, 0.25] + assert all((level == 0).sum() == 2 for level in levels_seen) + assert torch.isfinite(output).all() + observed = world["image"]["observed"] + assert torch.equal(output[observed], world["image"]["value"][observed]) + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("time_mode", ["continuous", "discrete"]) +@pytest.mark.usefixtures("reference_kernels") +def test_nonzero_packed_conditioning_matches_padded_outputs_and_gradients(mixed, time_mode): + arguments = dict(noise_conditioning="adaln_zero", use_phase_embedding=False, time_mode=time_mode) + padded = packed_model(**arguments) + activate_conditioning(padded) + packed = packed_model(token_execution="packed_valid", **arguments) + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 19) + outputs = [] + for network in (padded, packed): + with torch.autocast("cpu", dtype=torch.bfloat16) if mixed else nullcontext(): + prediction = network(world, "libero_gripper", noisy, levels)["image"] + loss = prediction.float().square().sum() + loss.backward() + outputs.append(prediction.detach()) + assert outputs[0].abs().max() > 0.01 + torch.testing.assert_close(outputs[1], outputs[0], rtol=0.03, atol=0.003) + assert_gradient_parity(packed, padded) + assert packed.noise_level[0].weight.grad.abs().sum() > 0 + changed = copy.deepcopy(world) + changed["image"]["value"][1] *= -3 + altered = noisy.clone() + altered[1] *= 2 + with torch.no_grad(): + expected = packed(world, "libero_gripper", noisy, levels)["image"][0] + isolated = packed(changed, "libero_gripper", altered, levels)["image"][0] + assert torch.equal(expected, isolated) + + +def test_fullgraph_tensor_regions_match_eager_with_nonzero_conditioning(): + eager = model() + activate_conditioning(eager) + compiled = model(compile_mode="default", compile_backend="aot_eager", compile_dynamic=True) + compiled.load_state_dict(eager.state_dict(), strict=True) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 23) + outputs = [] + for network in (eager, compiled): + prediction = network(world, "libero_gripper", noisy, levels)["image"] + prediction.square().mean().backward() + outputs.append(prediction.detach()) + torch.testing.assert_close(outputs[0], outputs[1], rtol=1e-5, atol=1e-6) + for name, parameter in eager.named_parameters(): + torch.testing.assert_close(parameter.grad, dict(compiled.named_parameters())[name].grad, + rtol=2e-5, atol=2e-6, msg=name) diff --git a/tests/test_image_proprio_attention.py b/tests/test_image_proprio_attention.py new file mode 100644 index 000000000..46ea095f8 --- /dev/null +++ b/tests/test_image_proprio_attention.py @@ -0,0 +1,199 @@ +"""CPU execution/masking contracts; real Flash4 capacity needs a GPU smoke.""" + +import copy + +import pytest +import torch + +from egomimic.models.dino_proprio import DinoProprioModel +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.world_model import MoTBlock + + +def model(**kwargs): + config = dict( + pose_dims={"gripper": 3}, + objective="visual_only", + image_dim=8, + width=24, + depth=2, + heads=3, + max_cameras=2, + max_tokens=32, + ) + config.update(kwargs) + return ImageProprioModel(**config) + + +def example(): + torch.manual_seed(31) + image = { + "value": torch.randn(1, 6, 8), + "time": torch.tensor([[-1., -1., 0., 0., 2., 2.]]), + "valid": torch.tensor([[True, True, False, True, True, True]]), + "observed": torch.tensor([[True, True, False, False, False, False]]), + "xy": torch.randn(1, 6, 2), + "camera": torch.tensor([[0, 1, 0, 1, 0, 1]]), + } + proprio = { + "value": torch.randn(1, 4, 3), + "time": torch.tensor([[-1., -.5, 0., 2.]]), + "valid": torch.ones(1, 4, dtype=torch.bool), + "observed": torch.tensor([[True, False, True, False]]), + } + return {"image": image, "proprio": proprio} + + +def forward(network, batch): + return network( + batch, "gripper", batch["image"]["value"], + torch.zeros_like(batch["image"]["time"]), phase="A", + )["image"] + + +def test_native_default_matches_historical_attention_and_strict_state_loading(): + torch.manual_seed(8) + default = model() + explicit = model(attention_backend="sdpa", attention_dtype="native", strict_fp32=True) + historical_attention = model() + explicit.load_state_dict(default.state_dict(), strict=True) + historical_attention.load_state_dict(default.state_dict(), strict=True) + # The old block called _attend directly. Bypassing only the new dispatcher + # checks that native execution keeps precisely that arithmetic. + for block in historical_attention.blocks: + block._apply_attention = block._attend + assert list(default.state_dict()) == list(explicit.state_dict()) + assert all(len(block.qkv) == 2 for block in explicit.blocks) + assert not any("attention_backend" in key for key in explicit.state_dict()) + batch = example() + expected = forward(historical_attention, batch) + assert torch.equal(forward(default, batch), expected) + assert torch.equal(forward(explicit, batch), expected) + # Execution controls do not add or reshape checkpoint tensors. + for settings in ( + {"attention_dtype": "bfloat16"}, + {"attention_backend": "flash4", "attention_dtype": "bfloat16"}, + ): + alternate = model(**settings) + alternate.load_state_dict(default.state_dict(), strict=True) + assert list(alternate.state_dict()) == list(default.state_dict()) + + +@pytest.mark.parametrize( + "settings,match", + [ + ({"strict_fp32": "true"}, "boolean"), + ({"strict_fp32": 1}, "boolean"), + ({"strict_fp32": True, "attention_dtype": "bfloat16"}, "conflicts"), + ({"attention_backend": "flash4"}, "explicit bfloat16"), + ({"attention_backend": "unknown"}, "attention_backend"), + ({"attention_dtype": "float16"}, "attention_dtype"), + ], +) +def test_explicit_precision_guards(settings, match): + with pytest.raises(ValueError, match=match): + model(**settings) + + +@pytest.mark.parametrize("dtype", ["native", "bfloat16"]) +def test_two_stream_holes_have_no_label_leakage_or_invalid_value_gradients(dtype): + network = model(attention_dtype=dtype) + batch = example() + batch["image"]["value"].requires_grad_() + batch["proprio"]["value"].requires_grad_() + result = forward(network, batch) + changed = copy.deepcopy(batch) + hidden = ~changed["proprio"]["observed"] + changed["proprio"]["value"] = changed["proprio"]["value"].detach() + changed["proprio"]["value"][hidden] = float("nan") + changed["proprio"]["time"][hidden] = 999 + assert torch.equal(forward(network, changed), result) + result.square().sum().backward() + pose_grad = batch["proprio"]["value"].grad + assert pose_grad[hidden].count_nonzero() == 0 + assert pose_grad[~hidden].abs().sum() > 0 + assert batch["image"]["value"].grad[~batch["image"]["valid"]].count_nonzero() == 0 + assert all(torch.isfinite(p.grad).all() for p in network.parameters() if p.grad is not None) + assert all(p.dtype == torch.float32 for p in network.parameters()) + + # Removing invalid keys (including interior holes) preserves valid outputs. + compact = copy.deepcopy(changed) + for name, keep in ( + ("image", batch["image"]["valid"][0]), + ("proprio", batch["proprio"]["observed"][0]), + ): + compact[name] = {key: value[:, keep] for key, value in compact[name].items()} + torch.testing.assert_close( + forward(network, compact), result[:, batch["image"]["valid"][0]], + rtol=1e-5, atol=1e-5, + ) + + +def test_bf16_is_confined_to_attention_with_fp32_outputs_and_gradients(monkeypatch): + network = model(attention_dtype="bfloat16") + calls = [] + for block in network.blocks: + original = block._attend + + def record(q, k, v, mask, original=original): + calls.append((q.dtype, k.dtype, v.dtype, mask.detach().clone())) + return original(q, k, v, mask) + + monkeypatch.setattr(block, "_attend", record) + batch = example() + result = forward(network, batch) + assert result.dtype == torch.float32 + result.square().mean().backward() + expected_keys = torch.cat((batch["image"]["valid"], batch["proprio"]["observed"]), 1) + assert len(calls) == len(network.blocks) + for q_dtype, k_dtype, v_dtype, mask in calls: + assert q_dtype == k_dtype == v_dtype == torch.bfloat16 + assert torch.equal(mask[:, 0, 0], expected_keys) + assert all( + p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() + for p in network.parameters() if p.grad is not None + ) + + +def test_unchunked_noncheckpointed_model_does_not_recompute(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("Checkpointing must be disabled by both execution controls") + + monkeypatch.setattr("egomimic.models.image_proprio.checkpoint", forbidden) + monkeypatch.setattr("egomimic.models.world_model.checkpoint", forbidden) + network = model( + attention_query_chunk_size=0, activation_checkpointing=False, + attention_dtype="bfloat16", + ) + forward(network, example()).square().mean().backward() + + +def test_flash4_model_does_not_silently_fallback_on_cpu(): + network = model(attention_backend="flash4", attention_dtype="bfloat16") + with pytest.raises(RuntimeError, match="CUDA"): + forward(network, example()) + + +def test_legacy_dino_wrapper_keeps_native_execution_defaults(): + network = DinoProprioModel( + pose_dims={"gripper": 3}, image_dim=8, width=16, depth=1, heads=2, + ) + assert network.attention_backend == "sdpa" + assert network.attention_dtype == "native" + assert not network.strict_fp32 + assert network.use_camera_embedding and network.use_modality_embedding + assert network.use_embodiment_embedding + + +def test_generic_single_modality_block_has_valid_backward(): + block = MoTBlock(24, 3, modalities=("image",), attention_dtype="bfloat16") + values = torch.randn(1, 5, 24, requires_grad=True) + times = torch.tensor([[-1., 0., .5, 1., 2.]]) + valid = torch.tensor([[True, False, True, True, True]]) + output, = block([values], [times], [valid]) + assert len(block.qkv) == 1 + assert output.shape == values.shape and output.dtype == torch.float32 + assert output[~valid].count_nonzero() == 0 + output.square().mean().backward() + assert torch.isfinite(values.grad).all() and values.grad[valid].abs().sum() > 0 + assert values.grad[~valid].count_nonzero() == 0 diff --git a/tests/test_image_proprio_distributed.py b/tests/test_image_proprio_distributed.py new file mode 100644 index 000000000..67b76a0aa --- /dev/null +++ b/tests/test_image_proprio_distributed.py @@ -0,0 +1,130 @@ +"""CPU Gloo proves distributed metric weighting and validation completion.""" + +import json +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + +from egomimic.eval.image_proprio_eval import ImageProprioEval +from egomimic.eval.pixel_diagnostics import PixelPredictionDiagnostics +from tests.test_pixel_prediction_diagnostics import batch_fixture, fake_model + + +def _distributed_metrics_worker(rank, rendezvous, output): + torch.set_num_threads(1) + torch.distributed.init_process_group( + "gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2, + timeout=timedelta(seconds=30), + ) + try: + logged = {} + model = SimpleNamespace( + device=torch.device("cpu"), global_step=4, global_rank=rank, + trainer=SimpleNamespace(sanity_checking=False), + log=lambda key, value, **kwargs: logged.update({key: float(value)}), + ) + evaluator = ImageProprioEval(objective="visual_only") + evaluator.model = model + + def prepare(): + evaluator.pixel_diagnostics = None + evaluator.on_validation_start() + evaluator._validation_batches = (60, 40)[rank] + # An unweighted average of rank means is wrong: (2 + 6)/2 != 5. + evaluator.totals = { + "forecast": (torch.tensor((2.0, 18.0)[rank]), torch.tensor((1.0, 3.0)[rank])) + } + if rank == 1: + evaluator.totals["completion"] = (torch.tensor(10.0), torch.tensor(2.0)) + + prepare() + evaluator.log_epoch_metrics() + evaluator.on_validation_end() + assert logged == {"forecast": 5.0, "completion": 5.0} + assert evaluator.global_validation_batches == 100 + assert evaluator._validation_batches == (60, 40)[rank] + assert evaluator.has_completed_validation(4) + # Repeated metric publication must not multiply already-reduced totals. + evaluator.log_epoch_metrics() + assert evaluator.global_validation_batches == 100 + assert logged == {"forecast": 5.0, "completion": 5.0} + + prepare() + evaluator.log_epoch_metrics() + + def publish(_): + if rank == 0: + raise RuntimeError("rank-zero artifact publication failed") + + evaluator.pixel_diagnostics = SimpleNamespace(log_to_wandb=publish) + try: + evaluator.on_validation_end() + except RuntimeError: + assert rank == 0 + assert not evaluator.has_completed_validation(4) + + prepare() + if rank == 0: + # A valid global sum must not hide an invalid local contribution. + evaluator.totals["forecast"] = (torch.tensor(-2.0), torch.tensor(1.0)) + with pytest.raises(ValueError, match="distributed rank"): + evaluator.log_epoch_metrics() + evaluator.on_validation_end() + assert not evaluator.has_completed_validation(4) + Path(output, f"rank-{rank}.json").write_text(json.dumps({"status": "PASS"})) + finally: + torch.distributed.destroy_process_group() + + +@pytest.mark.skipif(not torch.distributed.is_gloo_available(), reason="Gloo unavailable") +def test_distributed_metrics_reduce_sums_and_global_completion(tmp_path): + mp.spawn(_distributed_metrics_worker, args=(str(tmp_path / "rendezvous"), str(tmp_path)), + nprocs=2, join=True) + for rank in range(2): + assert json.loads((tmp_path / f"rank-{rank}.json").read_text()) == {"status": "PASS"} + + +def test_nonzero_rank_does_not_publish_rgb_artifacts(tmp_path): + batch, _ = batch_fixture() + model, prediction, _ = fake_model(batch) + model.global_rank = 1 + diagnostics = PixelPredictionDiagnostics(artifact_dir=tmp_path, log_wandb=False) + diagnostics.record(model=model, source="robot", item=batch["robot"], + prediction=prediction, objective="visual_only", batch_idx=0, dataloader_idx=0) + assert diagnostics.written == [] and not list(tmp_path.iterdir()) + + +def test_global_validation_identity_preserves_sampling_noise_across_shards(): + batch, _ = batch_fixture() + model, prediction, _ = fake_model(batch) + draws = [] + + def forward_eval(_): + draws.append(torch.rand(4)) + return {"robot": {"image_proprio_prediction": prediction}} + + model.forward_eval = forward_eval + evaluator = ImageProprioEval(objective="visual_only") + evaluator.model = model + evaluator.on_validation_start() + before = torch.get_rng_state().clone() + evaluator.on_validation_step(batch, 29) + batch["robot"]["validation_batch_index"] = [29] + evaluator.on_validation_step(batch, 4) + assert torch.equal(draws[0], draws[1]) + assert torch.equal(torch.get_rng_state(), before) + batch["robot"]["validation_batch_index"] = [30] + evaluator.on_validation_step(batch, 4) + assert not torch.equal(draws[0], draws[2]) + + +@pytest.mark.parametrize("invalid", [[], [-1], [True], [1.0], [1, 2]]) +def test_invalid_global_validation_identity_is_rejected(invalid): + batch, _ = batch_fixture() + batch["robot"]["validation_batch_index"] = invalid + with pytest.raises(ValueError, match="validation batch identit"): + ImageProprioEval._sampling_index(batch, 0) diff --git a/tests/test_image_proprio_image_only.py b/tests/test_image_proprio_image_only.py new file mode 100644 index 000000000..bc96d6d64 --- /dev/null +++ b/tests/test_image_proprio_image_only.py @@ -0,0 +1,135 @@ +"""Image-only stage one and image-only phase A of the deferred joint readout.""" + +import copy + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from tests.test_image_proprio_attention import example + + +def image_model(**kwargs): + settings = dict( + objective="visual_only", phase_a_condition_on_proprio=False, + image_dim=8, width=24, heads=3, depth=2, max_cameras=2, max_tokens=32, + ) + settings.update(kwargs) + return ImageProprioModel(**settings) + + +def joint_model(**kwargs): + return image_model( + objective="visual_state_action", pose_dims={"gripper": 3}, + action_dims={"gripper": 2}, **kwargs, + ) + + +def joint_batch(): + batch = example() + batch["action"] = { + "value": torch.randn(1, 4, 2), + "time": batch["proprio"]["time"].clone(), + "valid": torch.ones(1, 4, dtype=torch.bool), + "observed": torch.zeros(1, 4, dtype=torch.bool), + } + return batch + + +def predict(model, batch, phase="A"): + return model( + batch, "gripper", batch["image"]["value"], + torch.zeros_like(batch["image"]["time"]), phase=phase, + ) + + +def test_image_only_has_exact_parameters_and_no_proprio_modules(): + with torch.device("meta"): + full = image_model(image_dim=768, width=768, heads=6, depth=8) + assert sum(parameter.numel() for parameter in full.parameters()) == 58_482_432 + assert full.pose_dims == full.action_dims == {} + assert full.modalities == ("image",) + assert not full.has_proprio_branch + assert len(full.norm) == 1 + assert all(len(block.qkv) == 1 for block in full.blocks) + for name in ("pose_in", "proprio_role", "pose_out", "action_out", "query", "action_query"): + assert not hasattr(full, name) + assert not any("pose" in key or "proprio" in key for key in full.state_dict()) + + +def test_image_only_trains_and_samples_without_state_data(): + model = image_model() + batch = {"image": example()["image"]} + losses = model.training_losses(batch, "gripper") + assert set(losses) == {"image", "total"} + losses["total"].backward() + assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) + query = ~batch["image"]["observed"] + batch["image"]["value"][query] = float("nan") + prediction = model.sample(batch, "gripper", steps=20)["image"] + assert torch.isfinite(prediction).all() + observed = batch["image"]["observed"] + assert torch.equal(prediction[observed], batch["image"]["value"][observed]) + assert prediction[~batch["image"]["valid"]].count_nonzero() == 0 + + +def test_image_only_rejects_extra_streams_or_unusable_conditioning(): + model = image_model() + with pytest.raises(ValueError, match="exactly the image stream"): + predict(model, example()) + batch = {"image": example()["image"]} + batch["image"]["observed"].zero_() + with pytest.raises(ValueError, match="observed"): + predict(model, batch) + with pytest.raises(ValueError, match="omit pose_dims"): + image_model(pose_dims={"gripper": 3}) + with pytest.raises(ValueError, match="no embodiment embedding"): + image_model(use_embodiment_embedding=True) + with pytest.raises(ValueError, match="must be bool"): + image_model(phase_a_condition_on_proprio=0) + + +def test_joint_phase_a_never_uses_even_observed_proprio_or_action_queries(): + model, batch = joint_model(), joint_batch() + expected = predict(model, batch) + assert set(expected) == {"image"} + changed = copy.deepcopy(batch) + changed["proprio"]["value"] *= 1000 + changed["proprio"]["time"] += 123 + changed["action"]["value"].fill_(float("nan")) + assert torch.equal(expected["image"], predict(model, changed)["image"]) + expected["image"].square().mean().backward() + assert all(p.grad is None for p in model.pose_in.parameters()) + assert model.proprio_role.weight.grad is None + for block in model.blocks: + assert block.qkv[0].weight.grad.abs().sum() > 0 + assert block.qkv[1].weight.grad is None + + +def test_joint_phase_b_uses_new_branch_but_withheld_labels_are_loss_only(): + model, batch = joint_model(), joint_batch() + before = predict(model, batch, phase="B") + assert set(before) == {"image", "proprio", "action"} + changed = copy.deepcopy(batch) + changed["proprio"]["value"][~changed["proprio"]["observed"]] = float("nan") + changed["action"]["value"].fill_(float("nan")) + after = predict(model, changed, phase="B") + assert all(torch.equal(before[key], after[key]) for key in before) + (before["proprio"].square().mean() + before["action"].square().mean()).backward() + assert model.pose_in["gripper"].weight.grad.abs().sum() > 0 + assert model.pose_out["gripper"].weight.grad.abs().sum() > 0 + assert model.action_out["gripper"].weight.grad.abs().sum() > 0 + assert all(block.qkv[1].weight.grad.abs().sum() > 0 for block in model.blocks) + + +@pytest.mark.parametrize("readout_weight", [0., 1.]) +def test_joint_training_supports_zero_or_active_readout_after_image_only_phase_a(readout_weight): + model = joint_model() + losses = model.training_losses( + joint_batch(), "gripper", pose_weight=readout_weight, action_weight=readout_weight, + ) + assert set(losses) == {"image", "pose", "action", "total"} + assert all(torch.isfinite(value) for value in losses.values()) + losses["total"].backward() + assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None) + assert (model.action_out["gripper"].weight.grad is None) == (readout_weight == 0) diff --git a/tests/test_image_proprio_image_only_initialization.py b/tests/test_image_proprio_image_only_initialization.py new file mode 100644 index 000000000..127ed243f --- /dev/null +++ b/tests/test_image_proprio_image_only_initialization.py @@ -0,0 +1,150 @@ +"""Image-only -> two-expert readout migration with explicit CPU execution proof.""" + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from scripts.world_model.initialize_pixel_readout import ( + initialize_readout, + sha256_file, + transfer_shared_weights, + validate_transition, + verify_phase_a, +) +from tests.test_image_proprio_image_only import joint_batch + + +def configs(mode="continuous", modality=False, backend="sdpa"): + source = dict( + objective="visual_only", phase_a_condition_on_proprio=False, + image_dim=12, width=16, heads=2, depth=2, max_cameras=2, max_tokens=12, + time_mode=mode, use_modality_embedding=modality, + attention_backend=backend, + attention_dtype="bfloat16" if backend == "flash4" else "native", + ) + target = dict( + source, objective="visual_state_action", max_tokens=16, + pose_dims={"gripper": 9 if mode == "continuous" else 36}, + action_dims={"gripper": 7 if mode == "continuous" else 28}, + ) + return source, target + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("modality", [False, True]) +def test_image_weights_copy_exactly_and_added_branch_keeps_seeded_initialization(mode, modality): + source_config, target_config = configs(mode, modality) + validate_transition(source_config, target_config) + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + initialized = {key: value.clone() for key, value in target.state_dict().items()} + transfer = transfer_shared_weights(source, target) + assert "proprio_role.weight" in transfer["new_tensors"] + assert "pose_in.gripper.weight" in transfer["new_tensors"] + assert "blocks.0.qkv.1.weight" in transfer["new_tensors"] + for key in transfer["new_tensors"]: + assert torch.equal(target.state_dict()[key], initialized[key]) + for key, value in source.state_dict().items(): + actual = target.state_dict()[key] + if key in transfer["expanded_embeddings"]: + actual = actual[:value.shape[0]] + assert torch.equal(value, actual) + proof = verify_phase_a(source, target) + assert proof["bit_exact"] and len(proof["checks"]) == 2 + assert not proof["gpu_kernel_verified"] + + +def test_transition_cannot_enable_phase_a_proprio_or_relabel_source_architecture(): + source, target = configs() + with pytest.raises(ValueError, match="configuration mismatch"): + validate_transition(source, {**target, "phase_a_condition_on_proprio": True}) + with pytest.raises(ValueError, match="omit pose_dims"): + validate_transition({**source, "pose_dims": {"gripper": 9}}, target) + with pytest.raises(ValueError, match="configuration mismatch"): + validate_transition(source, {**target, "image_dim": 8}) + + +def test_full_twenty_step_generation_and_image_features_are_preserved(): + source_config, target_config = configs() + source_config["image_dim"] = target_config["image_dim"] = 8 + target_config["pose_dims"] = {"gripper": 3} + target_config["action_dims"] = {"gripper": 2} + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + transfer_shared_weights(source, target) + batch = joint_batch() + image_features, target_features = [], [] + handles = [ + source.blocks[-1].register_forward_hook( + lambda module, args, output: image_features.append(output[0].clone()) + ), + target.blocks[-1].register_forward_hook( + lambda module, args, output: target_features.append(output[0].clone()) + ), + ] + try: + original = source.sample( + {"image": batch["image"]}, "gripper", steps=20, + generator=torch.Generator().manual_seed(17), + ) + migrated = target.sample( + batch, "gripper", steps=20, generator=torch.Generator().manual_seed(17), + ) + finally: + for handle in handles: + handle.remove() + assert torch.equal(original["image"], migrated["image"]) + assert len(image_features) == 20 and len(target_features) == 21 + assert all(torch.equal(a, b) for a, b in zip(image_features, target_features[:20])) + + +def test_fa4_cpu_probe_requires_declared_override_and_restores_model_execution(): + source_config, target_config = configs(backend="flash4") + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + transfer_shared_weights(source, target) + with pytest.raises(ValueError, match="explicit cpu_attention_override"): + verify_phase_a(source, target) + proof = verify_phase_a(source, target, cpu_attention_override="native") + assert proof["bit_exact"] and proof["cpu_attention_override"] == "native" + assert proof["declared_attention"] == ["flash4/bfloat16"] + assert not proof["gpu_kernel_verified"] + assert all(block.attention_backend == "flash4" for model in (source, target) for block in model.blocks) + with torch.no_grad(): + target.image_out.bias.add_(1) + with pytest.raises(ValueError, match="preservation failed"): + verify_phase_a(source, target, cpu_attention_override="native") + assert all(block.attention_backend == "flash4" for model in (source, target) for block in model.blocks) + + +def test_offline_fa4_artifact_binds_original_configs_and_records_only_cpu_proof(tmp_path): + source_config, target_config = configs(backend="flash4") + source = ImageProprioModel(**source_config) + path = tmp_path / "source.pt" + torch.save(source.state_dict(), path) + kwargs = dict( + source_checkpoint=path, source_sha256=sha256_file(path), + source_config=source_config, target_config=target_config, + output_path=tmp_path / "target.pt", receipt_path=tmp_path / "receipt.json", + seed=42, + ) + with pytest.raises(ValueError, match="explicit cpu_attention_override"): + initialize_readout(**kwargs) + assert not kwargs["output_path"].exists() + rng_before = torch.random.get_rng_state().clone() + receipt = initialize_readout(**kwargs, cpu_attention_override="native") + assert torch.equal(rng_before, torch.random.get_rng_state()) + assert receipt["source_model"]["attention_backend"] == "flash4" + assert receipt["target_model"]["attention_backend"] == "flash4" + assert receipt["phase_a_preservation"]["cpu_attention_override"] == "native" + assert not receipt["phase_a_preservation"]["gpu_kernel_verified"] + target = ImageProprioModel(**target_config) + target.load_state_dict(torch.load(kwargs["output_path"], weights_only=True), strict=True) + assert verify_phase_a(source, target, "native")["bit_exact"] + assert not receipt["optimizer_state_transferred"] + assert not receipt["training_step_transferred"] + + +def test_unknown_added_parameter_cannot_be_silently_initialized(): + source_config, target_config = configs() + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + target.extra_parameter = torch.nn.Parameter(torch.ones(1)) + with pytest.raises(ValueError, match="Unexpected parameter namespace"): + transfer_shared_weights(source, target) diff --git a/tests/test_image_proprio_integration.py b/tests/test_image_proprio_integration.py new file mode 100644 index 000000000..fc149a5ca --- /dev/null +++ b/tests/test_image_proprio_integration.py @@ -0,0 +1,187 @@ +"""Generic pixel pipeline reaches optimization, RGB validation and checkpoints.""" + +import json + +import pytest +import torch +from hydra.utils import instantiate +from lightning.pytorch import Trainer +from lightning.pytorch.callbacks import ModelCheckpoint +from omegaconf import open_dict + +from egomimic.pl_utils.pl_model import ModelWrapper +from tests.test_dino_proprio_integration import configuration +from tests.test_image_proprio_validation_lifecycle import needs_final_validation +from tests.test_pixel_attention_profiles import pixel_stage, recipe +from tests.test_pixel_proprio_data import pixel_manifest + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize( + "objective", ["visual_only", "visual_state_action", "images_only"] +) +def test_pixel_pipeline_optimization_validation_and_checkpoint( + tmp_path, mode, objective +): + images_only = objective == "images_only" + cfg = recipe(mode, "native") if images_only else configuration("visual_state") + objective = "visual_only" if images_only else objective + manifest = pixel_manifest( + tmp_path, include_sharpa=False, frames=12, actions=not images_only + ) + if images_only: + # The actual image-only loader must not depend on measured states or + # their normalizer, even when the reusable source manifest records them. + metadata = json.loads(manifest.read_text()) + for episode in metadata["episodes"]: + (manifest.parent / episode["path"]).unlink() + (manifest.parent / metadata["normalization"]["path"]).unlink() + chunk = 2 if mode == "discrete" else 1 + with open_dict(cfg): + cfg.data = { + "_target_": "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule", + "manifest_path": str(manifest), + "task": "visual_state_action" + if objective == "visual_state_action" + else "visual_diffusion", + "mode": mode, + "sampling_strategy": "discrete_window" + if mode == "discrete" + else "full_episode_uniform", + "min_frames": 4, + "max_frames": 4, + "discrete_frames": 4, + "stride": chunk, + "batch_size": 1, + "num_workers": 0, + "checkpointable_sampling": True, + "validation_anchors_per_episode": 2, + "image_only": images_only, + } + if images_only: + stage = pixel_stage(cfg) + pixel_stage(cfg, "pixel_sampler").sampler.steps = 2 + else: + stage = cfg.model.pipeline.stages[0] + cfg.model.pipeline.stage_ids = {"image_proprio": 0} + stage._target_ = "egomimic.pipeline.stages_image_proprio.ImageProprioStage" + stage.warmup_steps = stage.pose_ramp_steps = 0 + stage.sampler_steps = 2 + stage.model._target_ = "egomimic.models.image_proprio.ImageProprioModel" + stage.model.objective = objective + stage.model.pose_dims = None if images_only else {"libero_gripper": 9 * chunk} + stage.model.action_dims = ( + {"libero_gripper": 7 * chunk} + if objective == "visual_state_action" + else None + ) + stage.model.image_dim = 12 + stage.model.width = 16 + stage.model.depth = 1 + stage.model.heads = 2 + if images_only: + # Keep the production three-axis position encoding, with dimensions + # matching this tiny model and the fixture's 4px / 2px patch grid. + stage.model.rope_axis_dims = [4, 2, 2] + stage.model.spatial_grid_size = [2, 2] + stage.model.max_tokens = 100 + stage.model.attention_query_chunk_size = 0 if images_only else 16 + cfg.evaluator = { + "_target_": "egomimic.eval.image_proprio_eval.ImageProprioEval", + "objective": objective, + "pixel_diagnostics": { + "artifact_dir": str(tmp_path / "pixel_diagnostics"), + "write_video": True, + "log_wandb": False, + }, + } + data = instantiate(cfg.data) + context = data.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + evaluator = instantiate(cfg.evaluator) + model = ModelWrapper(config_tree=cfg, evaluator=evaluator, enable_grad_norm=False) + model.data_context = context + context.bind(model.model, evaluator) + data.configure_evaluation(evaluator.data_requirements()) + if images_only: + for loader in (data.train_dataloader(), data.val_dataloader()): + batch = next(iter(loader)) + assert batch + assert all(set(item["world"]) == {"image"} for item in batch.values()) + network = model.model.pipeline.stage_by_id("pixel_denoiser").model + assert network.modalities == ("image",) + assert network.pose_dims == {} and network.action_dims == {} + assert all(len(block.qkv) == 1 for block in network.blocks) + before = next(model.parameters()).detach().clone() + callback = ModelCheckpoint( + dirpath=tmp_path / "checkpoints", + every_n_train_steps=2, + save_top_k=-1, + save_last=True, + save_on_train_epoch_end=False, + ) + trainer = Trainer( + accelerator="cpu", + devices=1, + precision="32-true", + logger=False, + max_steps=2, + max_epochs=-1, + min_epochs=0, + limit_val_batches=2, + val_check_interval=2, + check_val_every_n_epoch=None, + num_sanity_val_steps=0, + enable_progress_bar=False, + enable_model_summary=False, + callbacks=[callback], + default_root_dir=tmp_path, + ) + trainer.fit(model, datamodule=data) + assert trainer.global_step == data.consumed_batches == 2 + # The scheduled terminal pass already published immutable RGB artifacts. + # This is the same guard used by trainHydra's val_at_end path. + assert evaluator.has_completed_validation(trainer.global_step) is True + assert needs_final_validation(evaluator, trainer.global_step) is False + if needs_final_validation(evaluator, trainer.global_step): + trainer.validate(model, datamodule=data) + assert needs_final_validation(evaluator, trainer.global_step + 1) is True + assert not torch.equal(before, next(model.parameters()).detach()) + metrics = { + k: v for k, v in trainer.callback_metrics.items() if k.startswith("val/") + } + assert metrics and all(torch.isfinite(value) for value in metrics.values()) + assert any(k.endswith("image_normalized_mse") for k in metrics) + if objective == "visual_only": + assert all(k.endswith("image_normalized_mse") for k in metrics) + else: + assert any(k.endswith("proprio_normalized_mse") for k in metrics) + assert any(k.endswith("action_command_mse") for k in metrics) + checkpoint = torch.load(callback.last_model_path, weights_only=False) + assert checkpoint["global_step"] == 2 + assert checkpoint["data_context"] == context.snapshot() + assert "mot_rng_state" in checkpoint + if images_only: + assert trainer.optimizers[0].param_groups[0]["lr"] == 5e-5 + restored = ModelWrapper( + config_tree=cfg, evaluator=instantiate(cfg.evaluator), enable_grad_norm=False + ) + restored.load_state_dict(checkpoint["state_dict"], strict=True) + for key, value in restored.state_dict().items(): + assert torch.equal(value, checkpoint["state_dict"][key]) + reports = list((tmp_path / "pixel_diagnostics").rglob("diagnostic.json")) + assert reports + for path in reports: + report = json.loads(path.read_text()) + assert report["objective"] == objective + assert report["pixel_metadata"]["chunk_size"] == chunk + assert set(report["cameras"]) == {"0", "1"} + for camera in report["cameras"].values(): + assert (path.parent / camera["video"]).is_file() + assert (path.parent / camera["contact_sheet"]).is_file() + if objective == "visual_only": + assert "proprio" not in report and "action" not in report + else: + assert report["action"]["value_space"] == "native_action_commands" + assert report["proprio"]["value_space"] == "normalized_measured_state" diff --git a/tests/test_image_proprio_mixed_precision.py b/tests/test_image_proprio_mixed_precision.py new file mode 100644 index 000000000..5a8d5b899 --- /dev/null +++ b/tests/test_image_proprio_mixed_precision.py @@ -0,0 +1,106 @@ +"""CPU autocast analogues; CUDA/Flash4 execution is a separate GPU smoke.""" + +import pytest +import torch + +from tests.test_image_proprio_attention import example +from tests.test_image_proprio_image_only import image_model + + +def test_mixed_training_keeps_loss_timestamps_parameters_and_adamw_state_fp32(): + model = image_model(attention_dtype="bfloat16") + batch = {"image": example()["image"]} + # Nearby physical timestamps must remain distinct under mixed precision. + batch["image"]["time"] = torch.tensor([[30.010, 30.010, 30.020, 30.020, 30.030, 30.030]]) + times = batch["image"]["time"].clone() + optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) + observed_dtypes = {} + + def capture(name): + def record(module, inputs, output): + observed_dtypes[name] = output.dtype + return record + + handles = [ + model.blocks[0].qkv[0].register_forward_hook(capture("qkv")), + model.blocks[0].ffn[0][0].register_forward_hook(capture("ffn")), + model.image_out.register_forward_hook(capture("prediction")), + ] + try: + with torch.autocast("cpu", dtype=torch.bfloat16): + losses = model.training_losses(batch, "gripper") + assert set(observed_dtypes.values()) == {torch.bfloat16} + assert all(loss.dtype == torch.float32 and torch.isfinite(loss) for loss in losses.values()) + losses["total"].backward() + assert all( + p.dtype == torch.float32 and p.grad is not None + and p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() + for p in model.parameters() + ) + before = model.image_out.weight.detach().clone() + optimizer.step() + assert not torch.equal(before, model.image_out.weight) + for state in optimizer.state.values(): + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + with torch.autocast("cpu", dtype=torch.bfloat16): + predicted = model.sample(batch, "gripper", steps=20)["image"] + finally: + for handle in handles: + handle.remove() + assert torch.equal(times, batch["image"]["time"]) + assert batch["image"]["time"].dtype == torch.float32 + assert predicted.dtype == torch.float32 and torch.isfinite(predicted).all() + observed = batch["image"]["observed"] + assert torch.equal(predicted[observed], batch["image"]["value"][observed]) + + +@pytest.mark.parametrize("steps", [3, 20]) +def test_every_mixed_sampler_update_uses_fp32_arithmetic_before_clamping(monkeypatch, steps): + model = image_model(attention_dtype="bfloat16") + batch = {"image": example()["image"]} + image = batch["image"] + valid, observed = image["valid"], image["observed"] + fixed_prediction = torch.linspace(-.73, .91, image["value"].numel()).reshape_as( + image["value"] + ).bfloat16() + actual_inputs, actual_levels = [], [] + + def predict(batch, embodiment, image_values, levels, phase, **execution): + assert phase == "A" + actual_inputs.append(image_values.clone()) + actual_levels.append(levels.clone()) + return {"image": fixed_prediction} + + monkeypatch.setattr(model, "forward", predict) + with torch.autocast("cpu", dtype=torch.bfloat16): + actual = model.sample( + batch, "gripper", steps=steps, generator=torch.Generator().manual_seed(101) + )["image"] + anchor = torch.where(observed[..., None], image["value"], 0) + current = torch.where( + observed[..., None], anchor, + torch.randn(image["value"].shape, generator=torch.Generator().manual_seed(101)), + ).masked_fill(~valid[..., None], 0) + prediction_fp32 = fixed_prediction.float() + rounded_path_differs = False + for index, j in enumerate(range(steps, 0, -1)): + assert actual_inputs[index].dtype == torch.float32 + assert actual_levels[index].dtype == torch.float32 + torch.testing.assert_close(actual_inputs[index], current, rtol=0, atol=0) + level, next_level = j / steps, (j - 1) / steps + expected = (1 - next_level) * prediction_fp32 + (next_level / level) * ( + current - (1 - level) * prediction_fp32 + ) + # This reference distinguishes FP32 integration from the old BF16 + # scalar products; checking only the final x0 would miss the rounding. + rounded = (1 - next_level) * fixed_prediction + (next_level / level) * ( + current - (1 - level) * fixed_prediction + ) + rounded_path_differs |= not torch.equal(expected, rounded) + current = torch.where(observed[..., None], anchor, expected).masked_fill( + ~valid[..., None], 0 + ) + assert rounded_path_differs + torch.testing.assert_close(actual, current, rtol=0, atol=0) + assert torch.equal(actual[observed], image["value"][observed]) diff --git a/tests/test_image_proprio_no_phase.py b/tests/test_image_proprio_no_phase.py new file mode 100644 index 000000000..8fd0bc552 --- /dev/null +++ b/tests/test_image_proprio_no_phase.py @@ -0,0 +1,93 @@ +"""A/B routing without learned phase offsets, including strict readout expansion.""" + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.pipeline.stages_pixel_readout import configure_readout_trainability +from scripts.world_model.initialize_pixel_readout import ( + transfer_shared_weights, + validate_transition, + verify_phase_a, +) +from tests.test_image_proprio_image_only import image_model, joint_batch, predict +from tests.test_image_proprio_image_only_initialization import configs + + +@pytest.fixture(autouse=True) +def one_thread(): + before = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(before) + + +def test_legacy_default_has_identical_seeded_state_and_strict_reload(): + torch.manual_seed(37) + original = image_model() + torch.manual_seed(37) + explicit = image_model(use_phase_embedding=True) + assert original.use_phase_embedding and "phase.weight" in original.state_dict() + assert all(torch.equal(value, explicit.state_dict()[key]) for key, value in original.state_dict().items()) + explicit.load_state_dict(original.state_dict(), strict=True) + with pytest.raises(ValueError, match="use_phase_embedding must be bool"): + image_model(use_phase_embedding=0) + + +def test_no_phase_has_exact_parameter_delta_and_all_stage_one_gradients(): + with torch.device("meta"): + model = image_model(image_dim=768, width=768, heads=6, depth=8, use_phase_embedding=False) + assert sum(p.numel() for p in model.parameters()) == 58_481_664 + assert not hasattr(model, "phase") + assert not any("phase" in key for key in model.state_dict()) + model = image_model(use_phase_embedding=False) + batch = {"image": joint_batch()["image"]} + model.training_losses(batch, "gripper")["total"].backward() + assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) + with pytest.raises(ValueError, match="no proprio readout phase"): + predict(model, batch, phase="B") + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_no_phase_readout_warm_start_preserves_images_and_routes_state_commands(mode): + source_config, target_config = configs(mode) + for config in (source_config, target_config): + config.update(use_phase_embedding=False, image_dim=8) + target_config.update(pose_dims={"gripper": 3}, action_dims={"gripper": 2}) + validate_transition(source_config, target_config) + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + transfer = transfer_shared_weights(source, target) + assert "phase.weight" not in transfer["expanded_embeddings"] + assert not hasattr(source, "phase") and not hasattr(target, "phase") + assert verify_phase_a(source, target)["bit_exact"] + clone = ImageProprioModel(**target_config) + clone.load_state_dict(target.state_dict(), strict=True) + policy = configure_readout_trainability(target, "frozen_image") + assert not policy["phase_embedding_present"] and not policy["phase_embedding_frozen"] + assert sum(p.numel() for p in target.parameters() if not p.requires_grad) == sum(p.numel() for p in source.parameters()) + batch = joint_batch() + if mode == "discrete": + for stream in batch.values(): + stream["index"] = torch.arange(stream["valid"].shape[1]).expand_as(stream["valid"]) + original = source.sample({"image": batch["image"]}, "gripper", steps=20, + generator=torch.Generator().manual_seed(19)) + actual = target.sample(batch, "gripper", steps=20, generator=torch.Generator().manual_seed(19)) + assert torch.equal(original["image"], actual["image"]) + assert set(actual) == {"image", "proprio", "action"} + assert all(torch.isfinite(value).all() for value in actual.values()) + image = batch["image"] + levels = torch.zeros_like(image["time"]) + # Control-flow labels do not introduce a hidden phase offset in either image route. + a = target._embed_image(image, image["value"], levels, image["valid"], "A", False) + b = target._embed_image(image, image["value"], levels, image["valid"], "B", False) + assert torch.equal(a, b) + + +def test_initializer_refuses_to_drop_or_add_phase_parameters_implicitly(): + source_config, target_config = configs() + target_config["use_phase_embedding"] = False + with pytest.raises(ValueError, match="configuration mismatch.*use_phase_embedding"): + validate_transition(source_config, target_config) + source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + with pytest.raises(ValueError, match="use_phase_embedding must match"): + transfer_shared_weights(source, target) diff --git a/tests/test_image_proprio_positions.py b/tests/test_image_proprio_positions.py new file mode 100644 index 000000000..3e966f0fb --- /dev/null +++ b/tests/test_image_proprio_positions.py @@ -0,0 +1,197 @@ +"""Axial time/XY RoPE geometry, legacy identity and packed execution parity.""" + +from contextlib import nullcontext + +import pytest +import torch +from omegaconf import OmegaConf + +from egomimic.models.attention_execution import ( + apply_temporal_rope, + prepare_spatiotemporal_rope, + resolve_spatiotemporal_rope_geometry, +) +from egomimic.models.image_proprio import ImageProprioModel +from scripts.world_model.benchmark_pixel_execution import fixed_corruption +from tests.test_image_proprio_adaln import activate_conditioning +from tests.test_pixel_packed_execution import ( + assert_gradient_parity, + batch, + model, + reference_kernels, # noqa: F401 -- shared explicit CPU attention fixture +) + + +@pytest.fixture(autouse=True) +def one_thread(): + old = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(old) + + +def spatial_model(**kwargs): + settings = dict(noise_conditioning="adaln_zero", use_phase_embedding=False, + position_encoding="spatiotemporal_rope", rope_axis_dims=[4, 2, 2], + spatial_grid_size=[2, 2]) + settings.update(kwargs) + return model(**settings) + + +def test_axes_use_centered_patch_units_and_correct_half_split_pairs(): + times = torch.tensor([[0.0, 1.0]]) + xy = torch.tensor([[[-0.5, -0.5], [0.5, 0.5]]]) + rope = prepare_spatiotemporal_rope(times, xy, 8, axis_dims=[4, 2, 2], spatial_grid_size=[2, 2]) + expected_phase = torch.tensor([[[[0.0, 0.0, -0.5, -0.5], [1.0, 0.01, 0.5, 0.5]]]]) + torch.testing.assert_close(rope[0], expected_phase.cos()) + torch.testing.assert_close(rope[1], expected_phase.sin()) + basis = torch.zeros(1, 1, 2, 8) + basis[..., 2] = 1 # x cosine coordinate pairs with index6, not adjacent3. + rotated = apply_temporal_rope(basis, rope) + torch.testing.assert_close(rotated[..., 2], expected_phase[..., 2].cos()) + torch.testing.assert_close(rotated[..., 6], expected_phase[..., 2].sin()) + assert rotated[..., [0, 1, 3, 4, 5, 7]].count_nonzero() == 0 + + +def test_joint_coordinate_translation_preserves_attention_dot_products(): + torch.manual_seed(6) + q, k = torch.randn(2, 3, 5, 16), torch.randn(2, 3, 5, 16) + times, xy = torch.randn(2, 5), torch.randn(2, 5, 2) + arguments = dict(head_dim=16, axis_dims=[8, 4, 4], spatial_grid_size=[10, 10]) + original = prepare_spatiotemporal_rope(times, xy, **arguments) + shifted = prepare_spatiotemporal_rope(times + 13, xy + torch.tensor([0.2, -0.3]), **arguments) + def logits(table): + return apply_temporal_rope(q, table) @ apply_temporal_rope(k, table).transpose(-1, -2) + + torch.testing.assert_close(logits(original), logits(shifted), rtol=1e-5, atol=1e-5) + with torch.autocast("cpu", dtype=torch.bfloat16): + mixed = prepare_spatiotemporal_rope(times, xy, **arguments) + assert all(value.dtype == torch.float32 for value in mixed) + assert all(torch.equal(a, b) for a, b in zip(original, mixed)) + + +def test_nonimage_tokens_share_temporal_frequencies_with_zero_spatial_rotation(): + times = torch.tensor([[0.2, 2.1]]) + kwargs = dict(head_dim=128, axis_dims=[64, 32, 32], spatial_grid_size=[10, 10]) + image = prepare_spatiotemporal_rope(times, torch.ones(1, 2, 2), **kwargs) + pose = prepare_spatiotemporal_rope(times, None, **kwargs) + assert all(torch.equal(a[..., :32], b[..., :32]) for a, b in zip(image, pose)) + assert torch.equal(pose[0][..., 32:], torch.ones_like(pose[0][..., 32:])) + assert pose[1][..., 32:].count_nonzero() == 0 + + +def test_hydra_geometry_and_discrete_integer_clock_match_float_tables(): + axes, grid = OmegaConf.create([4, 2, 2]), OmegaConf.create([2, 2]) + assert resolve_spatiotemporal_rope_geometry(8, axes, grid) == ((4, 2, 2), (2, 2)) + times = torch.tensor([[0, 1, 3]], dtype=torch.int64) + xy = torch.zeros(1, 3, 2) + expected = prepare_spatiotemporal_rope(times.float(), xy, 8, axis_dims=axes, spatial_grid_size=grid) + actual = prepare_spatiotemporal_rope(times, xy, 8, axis_dims=axes, spatial_grid_size=grid) + assert all(torch.equal(a, b) for a, b in zip(expected, actual)) + network = spatial_model(backend="sdpa", attention_dtype="native", time_mode="discrete", + rope_axis_dims=axes, spatial_grid_size=grid) + world = batch() + world["image"]["index"] = world["image"]["index"].long() + noisy, levels = fixed_corruption(world, 8) + assert torch.isfinite(network(world, "libero_gripper", noisy, levels)["image"]).all() + + +@pytest.mark.parametrize("axes,grid", [([4, 3, 1], [2, 2]), ([4, 2, 4], [2, 2]), ([4, 2, 2], None), ([4, 2, 2], [0, 2])]) +def test_explicit_spatial_geometry_is_validated(axes, grid): + with pytest.raises(ValueError): + resolve_spatiotemporal_rope_geometry(8, axes, grid) + + +def test_new_parameter_count_and_legacy_rng_state_compatibility(): + with torch.device("meta"): + full = ImageProprioModel(objective="visual_only", phase_a_condition_on_proprio=False, + noise_conditioning="adaln_zero", use_phase_embedding=False, max_cameras=2, + position_encoding="spatiotemporal_rope", rope_axis_dims=[64, 32, 32], spatial_grid_size=[10, 10]) + assert sum(parameter.numel() for parameter in full.parameters()) == 88_178_688 + assert not hasattr(full, "xy") and hasattr(full, "camera") + assert not any(name.startswith("xy.") for name in full.state_dict()) + torch.manual_seed(29) + default = model() + rng = torch.get_rng_state().clone() + torch.manual_seed(29) + explicit = model(position_encoding="linear_xy_temporal_rope", rope_axis_dims=None, spatial_grid_size=None) + assert torch.equal(rng, torch.get_rng_state()) + assert all(torch.equal(value, explicit.state_dict()[name]) for name, value in default.state_dict().items()) + assert resolve_spatiotemporal_rope_geometry(8, None, [2, 2])[0] == (4, 2, 2) + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("time_mode", ["continuous", "discrete"]) +@pytest.mark.usefixtures("reference_kernels") +def test_nonzero_spatial_packed_outputs_and_all_gradients_match_padded(mixed, time_mode): + padded = spatial_model(time_mode=time_mode) + activate_conditioning(padded) + packed = spatial_model(time_mode=time_mode, token_execution="packed_valid") + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 20) + outputs = [] + for network in (padded, packed): + with torch.autocast("cpu", dtype=torch.bfloat16) if mixed else nullcontext(): + output = network(world, "libero_gripper", noisy, levels)["image"] + loss = output.float().square().sum() + loss.backward() + outputs.append(output.detach()) + assert outputs[0].abs().max() > 0.01 + torch.testing.assert_close(outputs[0], outputs[1], rtol=0.03, atol=0.003) + assert_gradient_parity(packed, padded) + + +def test_sampler_reuses_tables_and_phase_b_nonimage_axes_are_aligned(monkeypatch): + network = spatial_model(backend="sdpa", attention_dtype="native") + activate_conditioning(network) + world = batch() + calls = [] + original = network._prepare_rope + + def prepare(times, xy=None): + result = original(times, xy) + calls.append(result) + return result + + monkeypatch.setattr(network, "_prepare_rope", prepare) + output = network.sample(world, "libero_gripper", steps=4)["image"] + assert len(calls) == 1 and torch.isfinite(output).all() + times = torch.tensor([[1.0, 2.0]]) + _, rope = network._prepare_execution([times, times], [torch.ones(1, 2, dtype=torch.bool)] * 2, + torch.ones(1, 2, 2)) + assert all(torch.equal(a[..., :2], b[..., :2]) for a, b in zip(*rope)) + assert rope[1][1][..., 2:].count_nonzero() == 0 + + +def test_spatial_fullgraph_tensor_regions_match_eager(): + eager = spatial_model(backend="sdpa", attention_dtype="native") + activate_conditioning(eager) + compiled = spatial_model(backend="sdpa", attention_dtype="native", compile_mode="default", + compile_backend="aot_eager", compile_dynamic=True) + compiled.load_state_dict(eager.state_dict(), strict=True) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 9) + results = [] + for network in (eager, compiled): + result = network(world, "libero_gripper", noisy, levels)["image"] + result.square().mean().backward() + results.append(result.detach()) + torch.testing.assert_close(*results, rtol=1e-5, atol=1e-6) + for name, parameter in eager.named_parameters(): + torch.testing.assert_close(parameter.grad, dict(compiled.named_parameters())[name].grad, + rtol=2e-5, atol=2e-6, msg=name) + + +def test_full_image_model_has_no_additive_xy_and_keeps_relative_time_origin(): + network = spatial_model(backend="sdpa", attention_dtype="native") + activate_conditioning(network) + world = batch(nonprefix=True) + noisy, levels = fixed_corruption(world, 31) + expected = network(world, "libero_gripper", noisy, levels)["image"] + shifted = {"image": {key: value.clone() for key, value in world["image"].items()}} + shifted["image"]["time"] += 10 + shifted["image"]["xy"] += torch.tensor([0.2, -0.4]) + actual = network(shifted, "libero_gripper", noisy, levels)["image"] + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + assert network.camera.weight.numel() > 0 and not hasattr(network, "xy") diff --git a/tests/test_image_proprio_validation_lifecycle.py b/tests/test_image_proprio_validation_lifecycle.py new file mode 100644 index 000000000..91d8b8c40 --- /dev/null +++ b/tests/test_image_proprio_validation_lifecycle.py @@ -0,0 +1,130 @@ +"""Only successful final-step validation can suppress duplicate publication.""" + +import ast +from pathlib import Path + +import pytest +import torch + +from egomimic.eval.image_proprio_eval import ImageProprioEval +from tests.test_pixel_prediction_diagnostics import batch_fixture, fake_model + + +# Execute the maintained runner guard directly without unrelated optional +# runtime imports (the local test environment does not include tabulate). +def _runner_guard(): + path = Path(__file__).parents[1] / "egomimic/trainHydra.py" + node = next( + item + for item in ast.parse(path.read_text()).body + if isinstance(item, ast.FunctionDef) and item.name == "_needs_final_validation" + ) + namespace = {} + exec( + compile(ast.Module(body=[node], type_ignores=[]), str(path), "exec"), namespace + ) + return namespace["_needs_final_validation"] + + +needs_final_validation = _runner_guard() + + +def evaluator_fixture(tmp_path=None): + batch, _ = batch_fixture() + model, prediction, _ = fake_model(batch) + diagnostics = ( + None + if tmp_path is None + else { + "artifact_dir": str(tmp_path), + "write_video": False, + "log_wandb": False, + } + ) + evaluator = ImageProprioEval(objective="visual_only", pixel_diagnostics=diagnostics) + evaluator.model = model + evaluator.on_validation_start() + return evaluator, batch, model, prediction + + +def finish(evaluator, batch): + evaluator.on_validation_step(batch, 0) + evaluator.log_epoch_metrics() + evaluator.on_validation_end() + + +def test_completed_pass_suppresses_only_same_step_and_start_resets_it(tmp_path): + evaluator, batch, model, _ = evaluator_fixture(tmp_path) + assert evaluator.has_completed_validation(model.global_step) is False + assert needs_final_validation(evaluator, model.global_step) is True + finish(evaluator, batch) + assert evaluator.has_completed_validation(model.global_step) is True + assert needs_final_validation(evaluator, model.global_step) is False + assert needs_final_validation(evaluator, model.global_step + 1) is True + evaluator.on_validation_start() + assert needs_final_validation(evaluator, model.global_step) is True + with pytest.raises(FileExistsError): + evaluator.on_validation_step(batch, 0) + evaluator.log_epoch_metrics() + evaluator.on_validation_end() + assert needs_final_validation(evaluator, model.global_step) is True + + +@pytest.mark.parametrize( + "case", ["empty", "zero_targets", "sanity", "unlogged", "changed_step"] +) +def test_incomplete_pass_keeps_final_validation(case): + evaluator, batch, model, _ = evaluator_fixture() + if case == "sanity": + model.trainer.sanity_checking = True + evaluator.on_validation_start() + if case == "zero_targets": + batch["robot"]["world"]["image"]["observed"].fill_(True) + if case != "empty": + evaluator.on_validation_step(batch, 0) + if case != "unlogged": + evaluator.log_epoch_metrics() + if case == "changed_step": + model.global_step += 1 + evaluator.on_validation_end() + assert evaluator.has_completed_validation(model.global_step) is False + assert needs_final_validation(evaluator, model.global_step) is True + + +def test_failed_artifact_publication_cannot_mark_complete(tmp_path, monkeypatch): + evaluator, batch, model, _ = evaluator_fixture(tmp_path) + evaluator.on_validation_step(batch, 0) + evaluator.log_epoch_metrics() + + def fail(_): + assert evaluator.has_completed_validation(model.global_step) is False + raise RuntimeError("publication failed") + + monkeypatch.setattr(evaluator.pixel_diagnostics, "log_to_wandb", fail) + with pytest.raises(RuntimeError, match="publication failed"): + evaluator.on_validation_end() + assert needs_final_validation(evaluator, model.global_step) is True + + +@pytest.mark.parametrize("case", ["prediction", "accumulator", "metric_publication"]) +def test_nonfinite_or_failed_metric_pass_cannot_mark_complete(case): + evaluator, batch, model, prediction = evaluator_fixture() + if case == "prediction": + prediction["image"].fill_(float("nan")) + with pytest.raises(ValueError, match="Non-finite"): + evaluator.on_validation_step(batch, 0) + else: + evaluator.on_validation_step(batch, 0) + if case == "accumulator": + key = next(iter(evaluator.totals)) + evaluator.totals[key] = (torch.tensor(float("inf")), torch.tensor(1.0)) + else: + + def fail(*args, **kwargs): + raise RuntimeError("metric publication failed") + + model.log = fail + with pytest.raises((ValueError, RuntimeError)): + evaluator.log_epoch_metrics() + evaluator.on_validation_end() + assert needs_final_validation(evaluator, model.global_step) is True diff --git a/tests/test_libero_action_sidecar_export.py b/tests/test_libero_action_sidecar_export.py new file mode 100644 index 000000000..b1b8d6a26 --- /dev/null +++ b/tests/test_libero_action_sidecar_export.py @@ -0,0 +1,291 @@ +"""Exporter integrity, source-row alignment, and loader interoperability.""" + +import hashlib +import json +from pathlib import Path + +import h5py +import numpy as np +import pytest + +from egomimic.rldb.dino_proprio import PreparedDinoProprioDataset +from scripts.world_model.export_libero_action_sidecars import export_sidecars + + +def digest(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def write_json(path, value): + path.write_text(json.dumps(value)) + return path + + +def bundle(tmp_path, *, bad_row=False, time_offset=0): + raw = tmp_path / "raw" + raw.mkdir() + source_name, source_sha = "/fixture/original.hdf5", "a" * 64 + episodes, exports, payload = [], [], bytearray() + for i, split in enumerate(("train", "val")): + times = np.arange(4, dtype=np.float64) * 0.05 + native = tmp_path / f"native-{i}" + native.mkdir() + rows = np.arange(4) + if bad_row: + rows += 1 + np.savez(native / "aligned.npz", source_row=rows, time=times) + native_source = dict( + source_demo=source_name, source_demo_sha256=source_sha, episode=f"demo_{i}" + ) + native_manifest = write_json( + native / "manifest.json", + dict( + schema="libero-same-recorded-state-v1", + **native_source, + artifacts={"aligned.npz": digest(native / "aligned.npz")}, + ), + ) + feature = tmp_path / f"feature-{i}.h5" + with h5py.File(feature, "w") as data: + data["time"] = times + data["image"] = np.zeros((4, 2, 1, 8), np.float32) + data["image_xy"] = np.zeros((1, 2), np.float32) + data["pose"] = np.zeros((4, 9), np.float32) + data["pose_valid"] = np.ones(4, bool) + receipt = write_json( + feature.with_suffix(".h5.json"), + dict( + schema="mot-world-feature-cache-v1", + status="PASS", + output_sha256=digest(feature), + native_source=native_source, + episode_id=f"episode_{i}", + embodiment="libero_gripper", + native_cache=str(native), + request={ + "inputs": { + "native_manifest_sha256": digest(native_manifest), + "native_aligned_sha256": digest(native / "aligned.npz"), + } + }, + ), + ) + ep = dict( + embodiment="libero_gripper", + episode_id=f"episode_{i}", + split=split, + frames=4, + path=feature.name, + sha256=digest(feature), + receipt_sha256=digest(receipt), + native_source=native_source, + frame="fixed_reference", + pose_kind="measured_pose", + ) + episodes.append(ep) + actions = np.arange(28, dtype=" 0 + assert health["consecutive_failures"] == 0 + assert (tmp_path / "attempts").read_text() == "2" + assert not list(source.glob(".mirror-health.tmp.*")) + + +def test_final_failure_is_nonzero_and_health_records_all_failures(tmp_path): + source, destination, env = fixture(tmp_path, [20, 24, 124]) + result = run_mirror(source, destination, env) + assert result.returncode == 124 + health = json.loads((source / ".mirror-health.json").read_text()) + assert health["last_exit_code"] == 124 + assert health["last_success_unix"] == 0 + assert health["consecutive_failures"] == 3 + assert (tmp_path / "attempts").read_text() == "3" + + +def test_periodic_failures_and_recovery_never_kill_healthy_trainer(tmp_path): + source, destination, env = fixture(tmp_path, [20, 24, 124, 0]) + admission = tmp_path / "admitted" + admission.touch() + failure = tmp_path / "MIRROR_FAILED" + child = subprocess.Popen(["/bin/sleep", "30"]) + args = [str(child.pid), str(source), str(destination), str(admission), ".01"] + watcher = subprocess.Popen( + command("watch_log_mirror " + " ".join(map(shlex.quote, args))), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + receipt = source / ".mirror-health.json" + if receipt.exists(): + health = json.loads(receipt.read_text()) + if health["last_success_unix"] > 0: + break + assert child.poll() is None + time.sleep(0.01) + else: + pytest.fail("Periodic watcher never recovered") + assert child.poll() is None + assert not failure.exists() + assert health["consecutive_failures"] == 0 + child.terminate() + child.wait(timeout=3) + assert watcher.wait(timeout=3) == 0 + _, errors = watcher.communicate() + assert "training continues" in errors + finally: + for process in (watcher, child): + if process.poll() is None: + process.terminate() + process.wait(timeout=3) + + +def test_transfer_excludes_only_disposable_cache_lock_and_temp_receipts(tmp_path): + actual = shutil.which("rsync") + if actual is None: + pytest.skip("Local rsync unavailable") + source, destination, env = fixture(tmp_path, [0]) + env["MIRROR_TEST_REAL_RSYNC"] = actual + (source / "wandb-cache").mkdir() + (source / "wandb-cache" / "disposable").write_text("cache") + (source / "wandb").mkdir() + (source / "wandb" / "run.wandb").write_text("important spool") + (source / "metrics.csv").write_text("step,loss\n12001,0.2\n") + (source / ".mirror-health.tmp.orphan").write_text("partial") + destination.mkdir() + (destination / "historical.txt").write_text("retained") + result = run_mirror(source, destination, env) + assert result.returncode == 0, result.stderr + assert (destination / "metrics.csv").read_bytes() == ( + source / "metrics.csv" + ).read_bytes() + assert (destination / "wandb" / "run.wandb").read_text() == "important spool" + assert (destination / "historical.txt").read_text() == "retained" + assert not (destination / "wandb-cache").exists() + assert not (destination / ".mirror.lock").exists() + assert not (destination / ".mirror-health.tmp.orphan").exists() + # A second final flush updates the same CSV rather than ignoring it. + (source / "metrics.csv").write_text("step,loss\n12001,0.2\n12002,0.19\n") + assert run_mirror(source, destination, env).returncode == 0 + assert (destination / "metrics.csv").read_bytes() == ( + source / "metrics.csv" + ).read_bytes() + + +def test_concurrent_calls_are_serialized(tmp_path): + source, destination, env = fixture(tmp_path, [0]) + env["MIRROR_TEST_COPY_SECONDS"] = ".12" + args = command( + f"mirror_once {shlex.quote(str(source))} {shlex.quote(str(destination))} 1" + ) + first = subprocess.Popen(args, env=env) + second = subprocess.Popen(args, env=env) + assert first.wait(timeout=5) == 0 + assert second.wait(timeout=5) == 0 + assert (tmp_path / "events").read_text().splitlines() == [ + "start", + "end", + "start", + "end", + ] + assert ( + json.loads((source / ".mirror-health.json").read_text())["consecutive_failures"] + == 0 + ) + + +def test_shell_syntax(): + subprocess.run(["/bin/bash", "-n", str(COMMON)], check=True) diff --git a/tests/test_pixel_attention_profiles.py b/tests/test_pixel_attention_profiles.py new file mode 100644 index 000000000..f8131bef4 --- /dev/null +++ b/tests/test_pixel_attention_profiles.py @@ -0,0 +1,186 @@ +"""Compose actual pixel recipes without loading data or the optional CUDA runtime.""" + +from pathlib import Path + +import pytest +import torch +from hydra import compose, initialize_config_dir +from hydra.utils import instantiate +from omegaconf import OmegaConf, open_dict + +CONFIG_ROOT = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" +EXECUTION_KEYS = { + "attention_backend", + "attention_dtype", + "strict_fp32", + "attention_query_chunk_size", + "activation_checkpointing", +} + + +def recipe(arm, profile=None, execution=None): + overrides = [ + f"+experiment=pixel_proprio/video_{arm}", + "hydra/launcher=basic", + "data.manifest_path=/fixture/prepared.json", + "paths.output_dir=/fixture/output", + ] + if profile is not None: + overrides.append(f"world_model_attention={profile}") + if execution is not None: + overrides.append(f"+pixel_execution={execution}") + with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): + return compose(config_name="train_zarr_cartesian", overrides=overrides) + + +def scientific_stage(cfg): + stage = OmegaConf.to_container(cfg.model.pipeline, resolve=True) + model = stage["stages"][stage["stage_ids"]["pixel_denoiser"]]["model"] + for key in EXECUTION_KEYS: + model.pop(key) + return stage + + +def pixel_stage(cfg, name="pixel_denoiser"): + return cfg.model.pipeline.stages[cfg.model.pipeline.stage_ids[name]] + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +def test_pixel_default_is_explicit_flash4_profile(arm): + default, explicit = recipe(arm), recipe(arm, "flash4_bf16_mixed") + assert OmegaConf.to_container(default, resolve=False) == OmegaConf.to_container( + explicit, resolve=False + ) + model = pixel_stage(default).model + assert model.attention_backend == "flash4" + assert model.attention_dtype == "bfloat16" + assert not model.strict_fp32 and not default.runtime.strict_fp32 + assert default.runtime.disable_tf32 + assert default.trainer.precision == "bf16-mixed" + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +@pytest.mark.parametrize( + "profile,backend,dtype,strict,disable_tf32,trainer_precision", + [ + ("native", "sdpa", "native", False, False, "32-true"), + ("unchunked_fp32", "efficient", "native", True, True, "32-true"), + ("flash4_bf16", "flash4", "bfloat16", False, True, "32-true"), + ("flash4_bf16_mixed", "flash4", "bfloat16", False, True, "bf16-mixed"), + ], +) +def test_pixel_profile_override_preserves_science( + arm, profile, backend, dtype, strict, disable_tf32, trainer_precision +): + cfg, baseline = recipe(arm, profile), recipe(arm, "native") + model = pixel_stage(cfg).model + assert model.attention_backend == backend + assert model.attention_dtype == dtype + assert model.strict_fp32 is strict and cfg.runtime.strict_fp32 is strict + assert cfg.runtime.disable_tf32 is disable_tf32 + assert model.attention_query_chunk_size == 0 + assert not model.activation_checkpointing + assert scientific_stage(cfg) == scientific_stage(baseline) + for section in ("data", "callbacks"): + assert OmegaConf.to_container(cfg[section], resolve=True) == OmegaConf.to_container( + baseline[section], resolve=True + ) + trainer = OmegaConf.to_container(cfg.trainer, resolve=True) + baseline_trainer = OmegaConf.to_container(baseline.trainer, resolve=True) + assert trainer.pop("precision") == trainer_precision + baseline_trainer.pop("precision") + assert trainer == baseline_trainer + assert cfg.trainer.max_steps == 300000 and cfg.trainer.val_check_interval == 250 + assert cfg.model.optimizer == baseline.model.optimizer + assert cfg.model.optimizer.lr == 5e-5 and cfg.model.optimizer.weight_decay == 0.01 + assert cfg.model.scheduler is None and not cfg.model.enable_grad_norm + assert model.objective == "visual_only" and model.action_dims is None + assert cfg.data.image_only and not model.phase_a_condition_on_proprio + assert model.pose_dims is None + assert model.image_dim == model.width == 768 + assert model.depth == 8 and model.heads == 6 and model.mlp_ratio == 4 + assert model.max_cameras == 2 and model.use_camera_embedding + assert not model.use_modality_embedding and not model.use_embodiment_embedding + assert not model.use_phase_embedding + assert model.noise_conditioning == "adaln_zero" + assert not model.allow_patch_compression + assert pixel_stage(cfg, "pixel_loss").weight == 1 + assert pixel_stage(cfg, "pixel_sampler").sampler.steps == 20 + assert all("warmup_steps" not in stage and "pose_ramp_steps" not in stage + for stage in cfg.model.pipeline.stages) + assert cfg.data.mode == arm and model.time_mode == arm + if arm == "discrete": + assert cfg.data.sampling_strategy == "discrete_window" + assert cfg.data.discrete_frames == 16 and cfg.data.stride == 4 + assert model.max_tokens == 3200 + else: + assert cfg.data.sampling_strategy == "full_episode_uniform" + assert cfg.data.min_frames == 2 and cfg.data.max_frames == 64 + # Retain the existing model ceiling; stage1 sampling supplies <=12,800. + assert model.max_tokens == 25600 + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +def test_composed_image_only_model_has_one_expert_and_no_readouts(arm): + cfg = recipe(arm) + with torch.device("meta"): + model = instantiate(pixel_stage(cfg).model) + assert model.modalities == ("image",) + assert model.pose_dims == {} and model.action_dims == {} + assert sum(p.numel() for p in model.parameters()) == 88_178_688 + assert not hasattr(model, "phase") + assert all(p.dtype == torch.float32 for p in model.parameters()) + assert all(len(block.qkv) == 1 for block in model.blocks) + assert not any( + key.startswith(("pose_in.", "pose_out.", "action_out.", "proprio_role.")) + for key in model.state_dict() + ) + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +def test_pixel_native_precision_fallback_remains_cpu_compatible(arm): + cfg = recipe(arm, "native") + with open_dict(cfg): + del cfg.world_model_attention.trainer_precision + assert cfg.trainer.precision == "32-true" + + +@pytest.mark.parametrize("name", ["video_only", "video_discrete", "readout_continuous", "readout_discrete"]) +def test_continuous_stage_one_cap_does_not_change_other_recipe_defaults(name): + with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): + cfg = compose(config_name="train_zarr_cartesian", overrides=[ + f"+experiment=pixel_proprio/{name}", "hydra/launcher=basic", + ]) + assert cfg.data.min_frames == 2 and cfg.data.max_frames == 128 + assert cfg.data.validation_sampling_strategy == "discrete_window" + assert cfg.data.validation_discrete_frames == 16 + assert cfg.data.validation_stride == 4 + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +def test_four_gpu_profile_changes_only_execution_and_global_batch(arm): + baseline, cfg = recipe(arm), recipe(arm, execution="ddp4") + assert scientific_stage(cfg) == scientific_stage(baseline) + assert cfg.model.optimizer == baseline.model.optimizer + assert cfg.trainer.precision == "bf16-mixed" + assert cfg.trainer.strategy == "ddp" + assert cfg.trainer.devices == cfg.data.distributed_world_size == 4 + assert cfg.trainer.num_nodes == cfg.data.batch_size == 1 + assert cfg.trainer.accumulate_grad_batches == 1 + assert cfg.trainer.use_distributed_sampler is False + assert cfg.data.checkpointable_sampling is True + assert cfg.runtime.cluster_environment == "torch_elastic" + assert cfg.runtime.slurm_requeue_owner == "none" + actual_data = OmegaConf.to_container(cfg.data, resolve=True) + original_data = OmegaConf.to_container(baseline.data, resolve=True) + actual_data.pop("distributed_world_size") + original_data.pop("distributed_world_size", None) + assert actual_data == original_data + actual_trainer = OmegaConf.to_container(cfg.trainer, resolve=True) + original_trainer = OmegaConf.to_container(baseline.trainer, resolve=True) + for key in ("strategy", "devices", "num_nodes", "accumulate_grad_batches", "use_distributed_sampler"): + actual_trainer.pop(key, None) + original_trainer.pop(key, None) + assert actual_trainer == original_trainer + assert cfg.callbacks == baseline.callbacks + assert recipe(arm, "native", "ddp4").trainer.precision == "32-true" diff --git a/tests/test_pixel_compile_ddp.py b/tests/test_pixel_compile_ddp.py new file mode 100644 index 000000000..fdf605381 --- /dev/null +++ b/tests/test_pixel_compile_ddp.py @@ -0,0 +1,96 @@ +"""Regional compilation must preserve ordinary distributed gradient reduction.""" + +import copy +import json +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + +from egomimic.models.compilation import compile_region + + +class RegionModel(torch.nn.Module): + def __init__(self, *, compiled=False): + super().__init__() + self.norm = torch.nn.LayerNorm(4) + self.linear = torch.nn.Linear(4, 3) + self.execute = compile_region( + self.region, mode="default" if compiled else "none", dynamic=True, + backend="pixel_region_ddp_test", + ) + + def region(self, value): + return self.linear(self.norm(value)).square() + + def forward(self, value): + return self.execute(value) + + +def _input(rank): + return torch.linspace(-1, 2, (rank + 2) * 3 * 4).reshape(rank + 2, 3, 4) + + +def _worker(rank, directory): + torch.set_num_threads(1) + from torch._dynamo.backends.registry import lookup_backend, register_backend + + policies = [] + backend = lookup_backend("aot_eager") + + def checked_backend(graph, inputs): + policies.append(torch._dynamo.config.optimize_ddp) + return backend(graph, inputs) + + register_backend(checked_backend, name="pixel_region_ddp_test") + dist.init_process_group( + "gloo", init_method=Path(directory, "rendezvous").as_uri(), rank=rank, world_size=2, + ) + try: + torch.manual_seed(42) + network = RegionModel(compiled=True) + reference = RegionModel() + reference.load_state_dict(copy.deepcopy(network.state_dict()), strict=True) + expected = sum(reference(_input(r)).mean() for r in range(2)) / 2 + expected.backward() + wrapped = torch.nn.parallel.DistributedDataParallel(network) + original_policy = torch._dynamo.config.optimize_ddp + loss = wrapped(_input(rank)).mean() + loss.backward() + assert policies and all(value is False for value in policies) + assert torch._dynamo.config.optimize_ddp == original_policy + for actual, wanted in zip(network.parameters(), reference.parameters()): + torch.testing.assert_close(actual.grad, wanted.grad, rtol=1e-5, atol=1e-6) + torch.optim.SGD(network.parameters(), lr=0.01).step() + state = torch.cat([parameter.detach().reshape(-1) for parameter in network.parameters()]) + ranks = [torch.empty_like(state) for _ in range(2)] + dist.all_gather(ranks, state) + assert torch.equal(ranks[0], ranks[1]) + Path(directory, f"rank{rank}.json").write_text(json.dumps({ + "rank": rank, "gradient_average_verified": True, + "compiled_regions": len(policies), "identical_updated_weights": True, + })) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="Gloo is unavailable") +def test_regional_compile_preserves_two_rank_ddp_gradient_average(tmp_path): + torch.multiprocessing.spawn(_worker, args=(str(tmp_path),), nprocs=2, join=True) + assert all(json.loads((tmp_path / f"rank{rank}.json").read_text())[ + "gradient_average_verified" + ] for rank in range(2)) + + +def test_region_restores_ddp_policy_after_backend_failure(monkeypatch): + def fail(*args, **kwargs): + assert torch._dynamo.config.optimize_ddp is False + raise RuntimeError("test compiler failure") + + monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: fail) + original = torch._dynamo.config.optimize_ddp + region = compile_region(lambda x: x + 1, mode="default") + with pytest.raises(RuntimeError, match="test compiler failure"): + region(torch.ones(1)) + assert torch._dynamo.config.optimize_ddp == original diff --git a/tests/test_pixel_diffusion_components.py b/tests/test_pixel_diffusion_components.py new file mode 100644 index 000000000..0857d2bf0 --- /dev/null +++ b/tests/test_pixel_diffusion_components.py @@ -0,0 +1,273 @@ +"""Independent pre-extraction arithmetic/RNG references for graph components.""" + +import copy + +import pytest +import torch + +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, + PlainPixelMSE, +) +from tests.test_image_proprio_image_only import joint_batch, joint_model +from tests.test_pixel_execution_model import batch, network + + +def previous_noising(image, clean, maximum=1.0): + # Frozen arithmetic from the pre-extraction model; deliberately independent + # of all new components so draw-order changes cannot validate themselves. + levels = torch.zeros_like(image["time"], dtype=torch.float32) + for row in range(levels.shape[0]): + selected = image["valid"][row] & ~image["observed"][row] + if selected.any(): + times, inverse = torch.unique( + image["time"][row][selected], return_inverse=True + ) + draws = torch.rand(times.numel(), device=levels.device) * maximum + levels[row, selected] = draws[inverse] + strength = levels[..., None].to(clean.dtype) + noisy = (1 - strength) * clean + strength * torch.randn_like(clean) + return noisy.masked_fill(~image["valid"][..., None], 0), levels + + +def previous_loss(prediction, target, mask, reduction): + weight = mask.float() + target = torch.where((weight > 0)[..., None], target, torch.zeros_like(target)) + error = (prediction.float() - target.float()).square().mean(-1) + if reduction == "mean_per_example": + numerator = (error * weight).flatten(1).sum(1) + denominator = weight.flatten(1).sum(1) + return (numerator / torch.where(denominator > 0, denominator, 1.0)).mean() + denominator = weight.sum() + return (error * weight).sum() / torch.where(denominator > 0, denominator, 1.0) + + +@torch.no_grad() +def previous_sample(model, data, steps, generator): + model.validate(data, "gripper") + image = data["image"] + valid, observed = image["valid"], image["observed"] + anchor = torch.where(observed[..., None], image["value"], 0) + noise = torch.randn( + image["value"].shape, + device=anchor.device, + dtype=anchor.dtype, + generator=generator, + ) + current = torch.where(observed[..., None], anchor, noise).masked_fill( + ~valid[..., None], 0 + ) + execution = None + if not model.has_proprio_branch: + execution = model._prepare_execution([model._clock(image, valid)], [valid]) + for j in range(steps, 0, -1): + level, next_level = j / steps, (j - 1) / steps + levels = torch.full_like(image["time"], level, dtype=torch.float32).masked_fill( + observed | ~valid, 0 + ) + x0 = model( + data, + "gripper", + current, + levels, + phase="A", + _validated=True, + _execution=execution, + )["image"].to(current.dtype) + candidate = (1 - next_level) * x0 + (next_level / level) * ( + current - (1 - level) * x0 + ) + current = torch.where(observed[..., None], anchor, candidate).masked_fill( + ~valid[..., None], 0 + ) + return current + + +@pytest.mark.parametrize("maximum", [0.0, 0.2, 1.0]) +@pytest.mark.parametrize("all_observed", [False, True]) +def test_noising_exact_legacy_draw_order_observed_padding_and_metadata( + maximum, all_observed +): + image = batch()["image"] + image["value"][~image["valid"]] = float("nan") + if all_observed: + image["observed"] = image["valid"].clone() + clean = torch.where(image["valid"][..., None], image["value"], 0) + torch.manual_seed(91) + expected, expected_levels = previous_noising(image, clean, maximum) + expected_rng = torch.get_rng_state().clone() + torch.manual_seed(91) + actual = LinearGaussianImageNoising(maximum)(image) + assert torch.equal(actual["image_values"], expected) + assert torch.equal(actual["noise_levels"], expected_levels) + assert torch.equal(torch.get_rng_state(), expected_rng) + assert torch.equal(actual["target"], clean) + assert torch.equal( + actual["image_values"][image["observed"]], image["value"][image["observed"]] + ) + assert not actual["image_values"][~image["valid"]].any() + assert actual["metadata"]["timestamps"] is image["time"] + assert torch.equal(actual["mask"], image["valid"] & ~image["observed"]) + assert actual["noise_levels"].dtype == torch.float32 + + +def test_noising_rejects_missing_or_nonfinite_recorded_query_targets(): + image = batch()["image"] + image["target_valid"] = image["valid"].clone() + image["target_valid"][0, 2] = False + with pytest.raises(ValueError, match="recorded features"): + LinearGaussianImageNoising()(image) + image["target_valid"][0, 2] = True + image["value"][0, 2] = float("nan") + with pytest.raises(ValueError, match="Non-finite image target"): + LinearGaussianImageNoising()(image) + + +@pytest.mark.parametrize("reduction", ["token_mean", "mean_per_example"]) +def test_loss_preserves_fp32_arithmetic_gradients_and_empty_clips(reduction): + target = torch.randn(2, 4, 3) + mask = torch.tensor([[True, False, True, True], [False] * 4]) + target[~mask] = float("nan") + predicted = torch.randn_like(target).bfloat16().requires_grad_() + expected = previous_loss(predicted, target, mask, reduction) + # Noise/time metadata must not accidentally turn plain MSE into a weighting. + actual = PlainPixelMSE(reduction)( + predicted, + target, + mask, + noise_levels=torch.full((2, 4), 0.9), + metadata={"timestamps": torch.arange(4)}, + ) + assert actual.dtype == torch.float32 + assert torch.equal(actual, expected) + expected_grad = torch.autograd.grad(expected, predicted, retain_graph=True)[0] + actual_grad = torch.autograd.grad(actual, predicted)[0] + assert torch.equal(actual_grad, expected_grad) + assert not actual_grad[1].any() + + +@pytest.mark.parametrize("reduction", ["token_mean", "mean_per_example"]) +def test_separate_components_equal_legacy_training_loss_gradients_and_rng(reduction): + torch.manual_seed(7) + legacy = network() + legacy.loss_reduction = reduction + separate = copy.deepcopy(legacy) + data = batch() + image = data["image"] + torch.manual_seed(111) + clean = torch.where(image["valid"][..., None], image["value"], 0) + noisy, levels = previous_noising(image, clean) + expected_prediction = legacy(data, "gripper", noisy, levels)["image"] + expected = previous_loss( + expected_prediction, clean, image["valid"] & ~image["observed"], reduction + ) + expected_rng = torch.get_rng_state().clone() + expected.backward() + torch.manual_seed(111) + corruption = LinearGaussianImageNoising()(image) + actual_prediction = separate( + data, "gripper", corruption["image_values"], corruption["noise_levels"] + )["image"] + actual = PlainPixelMSE(reduction)( + actual_prediction, + corruption["target"], + corruption["mask"], + corruption["noise_levels"], + corruption["metadata"], + ) + assert torch.equal(actual, expected) + assert torch.equal(torch.get_rng_state(), expected_rng) + actual.backward() + for a, b in zip(legacy.parameters(), separate.parameters()): + assert torch.equal(a.grad, b.grad) + torch.manual_seed(111) + wrapped = separate.training_losses(data, "gripper")["total"] + assert torch.equal(wrapped, expected) + assert torch.equal(torch.get_rng_state(), expected_rng) + assert legacy.state_dict().keys() == separate.state_dict().keys() + + +@pytest.mark.parametrize("steps", [1, 3, 20]) +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("explicit_generator", [False, True]) +def test_sampler_component_is_bit_exact_to_previous_integration( + steps, mixed, explicit_generator +): + model = network(attention_dtype="bfloat16" if mixed else "native") + data = batch() + image = data["image"] + image["value"][~image["observed"]] = float("nan") + + def generator(): + return torch.Generator().manual_seed(101) if explicit_generator else None + + torch.manual_seed(83) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=mixed): + expected = previous_sample(model, data, steps, generator()) + expected_rng = torch.get_rng_state().clone() + torch.manual_seed(83) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=mixed): + actual = LinearDDIMImageSampler(steps)( + model, data, "gripper", generator=generator() + ) + assert actual.dtype == torch.float32 + assert torch.equal(actual, expected) + assert torch.equal(torch.get_rng_state(), expected_rng) + torch.manual_seed(83) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=mixed): + wrapped = model.sample(data, "gripper", steps=steps, generator=generator())[ + "image" + ] + assert torch.equal(wrapped, expected) + assert torch.equal(torch.get_rng_state(), expected_rng) + assert torch.isfinite(actual).all() + assert torch.equal(actual[image["observed"]], image["value"][image["observed"]]) + + +def test_stage_two_wrapper_keeps_phase_b_readout_and_observed_states(): + model, data = joint_model(), joint_batch() + current = previous_sample(model, data, 3, torch.Generator().manual_seed(23)) + with torch.no_grad(): + expected = model( + data, "gripper", current, torch.zeros_like(data["image"]["time"]), phase="B" + ) + actual = model.sample( + data, "gripper", steps=3, generator=torch.Generator().manual_seed(23) + ) + assert torch.equal(actual["image"], current) + assert torch.equal(actual["action"], expected["action"]) + pose = data["proprio"] + expected_pose = torch.where( + pose["observed"][..., None], pose["value"], expected["proprio"] + ) + assert torch.equal( + actual["proprio"], expected_pose.masked_fill(~pose["valid"][..., None], 0) + ) + + +def test_sampler_prepares_once_and_components_register_no_parameters(monkeypatch): + model, data = network(), batch() + before = tuple(model.state_dict()) + calls = {"validate": 0, "prepare": 0} + validate, prepare = model.validate, model._prepare_execution + + def checked(*args, **kwargs): + calls["validate"] += 1 + return validate(*args, **kwargs) + + def prepared(*args, **kwargs): + calls["prepare"] += 1 + return prepare(*args, **kwargs) + + monkeypatch.setattr(model, "validate", checked) + monkeypatch.setattr(model, "_prepare_execution", prepared) + LinearDDIMImageSampler(20)(model, data, "gripper") + assert calls == {"validate": 1, "prepare": 1} + assert tuple(model.state_dict()) == before + for component in ( + LinearGaussianImageNoising(), + PlainPixelMSE(), + LinearDDIMImageSampler(), + ): + assert not isinstance(component, torch.nn.Module) diff --git a/tests/test_pixel_execution_benchmark.py b/tests/test_pixel_execution_benchmark.py new file mode 100644 index 000000000..987877a0f --- /dev/null +++ b/tests/test_pixel_execution_benchmark.py @@ -0,0 +1,304 @@ +"""Benchmark protocol checks; these CPU tests make no H100 speed claim.""" + +import json +import sys +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from scripts.world_model import benchmark_pixel_execution as benchmark + + +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8]) +def test_mixed_schedule_matches_frame_histogram_across_batch_sizes(batch_size): + schedule = benchmark.frame_schedule("mixed", batch_size) + histogram = Counter(frame for batch in schedule for frame in batch) + assert set(histogram) == {2, 16, 64, 128} + assert len(set(histogram.values())) == 1 + assert all(len(batch) == batch_size for batch in schedule) + assert benchmark.complete_cycles(10, schedule) % len(schedule) == 0 + assert benchmark.complete_cycles(10, schedule) >= 10 + + +@pytest.mark.parametrize("scenario,capacity", [("fixed16", 3200), ("mixed", 25600), ("worst128", 25600)]) +def test_benchmark_binds_exact_architecture_and_explicit_loss(scenario, capacity): + arguments = benchmark.model_arguments("flash4", scenario) + with torch.device("meta"): + model = ImageProprioModel(**arguments) + assert sum(p.numel() for p in model.parameters()) == 88_178_688 + assert not model.use_phase_embedding + assert model.noise_conditioning == "adaln_zero" + assert model.position_encoding == "spatiotemporal_rope" + assert not hasattr(model, "xy") + assert not any(name.startswith("phase.") for name in model.state_dict()) + assert model.max_tokens == capacity + assert model.loss_reduction == "mean_per_example" + assert model.pose_dims == model.action_dims == {} + assert model.modalities == ("image",) + assert all(p.dtype == torch.float32 for p in model.parameters()) + + +def small_batch(): + return benchmark.make_batch((2, 3), seed=17, device=torch.device("cpu"), image_size=4, patch_size=2) + + +def small_model(): + return ImageProprioModel( + image_dim=12, width=16, depth=1, heads=2, max_cameras=2, max_tokens=32, + objective="visual_only", phase_a_condition_on_proprio=False, + attention_backend="sdpa", attention_dtype="bfloat16", + loss_reduction="mean_per_example", + ) + + +def test_geometry_padding_and_unpadded_independent_examples(): + batch = small_batch() + image = batch["image"] + assert image["value"].shape == (2, 24, 12) + assert image["valid"].sum(1).tolist() == [16, 24] + assert image["observed"].sum(1).tolist() == [8, 8] + assert image["time"].dtype == torch.float32 + assert image["value"].dtype == torch.float32 + assert image["camera"].unique().tolist() == [0, 1] + single = benchmark.single_example(batch, 0)["image"] + assert single["value"].shape == (1, 16, 12) + assert single["valid"].all() + torch.testing.assert_close(single["value"], image["value"][0:1, :16]) + + +def test_mean_per_example_is_not_padding_or_token_weighted(): + batch = small_batch() + batch["image"]["value"].zero_() + batch["image"]["value"][1] = 1 + loss = benchmark.example_mean_loss(torch.zeros_like(batch["image"]["value"]), batch) + assert loss == 0.5 # Unequal query-token counts must not weight example 1 twice. + + +def test_batched_probe_checks_outputs_gradients_and_cross_example_isolation(): + # Flash4's context means standard BF16 autocast. The tiny CPU model uses + # native SDPA explicitly; hardware/kernel certification happens on CUDA. + result = benchmark.check_outputs(small_model(), small_batch(), backend="flash4", + seed=42, tolerance=0.03, gradients=True) + assert result["output_relative_l2"] <= 0.03 + assert result["gradient_relative_l2"] <= 0.03 + assert result["cross_example_relative_l2"] <= 1e-6 + + +def test_measured_training_loss_matches_explicit_mean_independent_example_contract(monkeypatch): + model, batch = small_model(), small_batch() + noisy, levels = benchmark.fixed_corruption(batch, 42) + monkeypatch.setattr(model, "_noisy_images", lambda *args, **kwargs: (noisy, levels)) + with benchmark.execution_context("flash4", torch.device("cpu")): + prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] + explicit = benchmark.example_mean_loss(prediction, batch) + losses = model.training_losses(batch, "libero_gripper", pose_weight=0.0) + assert losses["total"].dtype == torch.float32 + torch.testing.assert_close(losses["total"], explicit, rtol=0, atol=0) + + +def test_eager_compiled_gate_accepts_equal_execution_and_rejects_bad_backward(): + # Exercise the comparison gate without claiming CPU execution certifies + # CUDA Inductor. The scheduled harness supplies an actual compiled model. + eager, candidate = small_model(), small_model() + candidate.load_state_dict(eager.state_dict(), strict=True) + result = benchmark.check_eager_compiled(eager, candidate, small_batch(), + backend="flash4", seed=42, tolerance=0.03) + assert result["same_state_verified"] + assert result["output_relative_l2"] == result["gradient_relative_l2"] == 0 + assert result["loss_relative_error"] == 0 + hooks = [p.register_hook(lambda gradient: gradient * 2) for p in candidate.parameters()] + with pytest.raises(benchmark.CorrectnessError, match="Compiled gradient differs from eager"): + benchmark.check_eager_compiled(eager, candidate, small_batch(), + backend="flash4", seed=42, tolerance=0.03) + assert all(p.grad is None for model in (eager, candidate) for p in model.parameters()) + for hook in hooks: + hook.remove() + + +def test_eager_compiled_gate_rejects_unmatched_weights(): + with pytest.raises(benchmark.CorrectnessError, match="initial state differs"): + benchmark.check_eager_compiled(small_model(), small_model(), small_batch(), + backend="flash4", seed=42, tolerance=0.03) + + +def test_probe_rejects_cross_example_interference(): + model = small_model() + forward = model.forward + + def leaking_forward(batch, *args, **kwargs): + result = forward(batch, *args, **kwargs) + result["image"] = result["image"] + batch["image"]["value"].mean() * 10 + return result + + model.forward = leaking_forward + with pytest.raises(benchmark.CorrectnessError, match="independent B1|leakage"): + benchmark.check_outputs(model, small_batch(), backend="flash4", seed=42, + tolerance=0.03, gradients=True) + + +def test_leakage_gate_does_not_confuse_grad_mode_rounding_with_interference(): + model = small_model() + forward = model.forward + + def mode_dependent_forward(*args, **kwargs): + result = forward(*args, **kwargs) + if not torch.is_grad_enabled(): + result["image"] = result["image"] + 0.001 + return result + + model.forward = mode_dependent_forward + result = benchmark.check_outputs(model, small_batch(), backend="flash4", seed=42, + tolerance=0.03, gradients=True) + assert 1e-6 < result["grad_vs_no_grad_relative_l2"] < 0.03 + assert result["cross_example_relative_l2"] == 0 + + +def test_cpu_worker_is_explicitly_unsupported_without_fallback(tmp_path, monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + request = tmp_path / "request.json" + benchmark.write_once(request, {"backend": "flash4"}) + benchmark.worker(request) + result = json.loads((tmp_path / "result.json").read_text()) + assert result["status"] == "UNSUPPORTED" + assert "no CPU benchmark fallback" in result["error"] + with pytest.raises(FileExistsError): + benchmark.worker(request) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_fa3_inventory_matches_canonical_or_absent_package_alias(tmp_path, monkeypatch, legacy): + module_file = tmp_path / "interface.py" + module_file.write_text("# exact fake backend source\n") + calls = [] + + def import_module(name): + calls.append(name) + if legacy and name == "flash_attn_3.flash_attn_interface": + raise ModuleNotFoundError("No canonical package", name="flash_attn_3") + return SimpleNamespace(__file__=str(module_file)) + + monkeypatch.setattr(benchmark.importlib, "import_module", import_module) + inventory = benchmark.backend_inventory("flash3") + expected = ["flash_attn_3.flash_attn_interface"] + (["flash_attn_interface"] if legacy else []) + assert calls == expected + assert inventory["module"] == expected[-1] + assert inventory["sha256"] == benchmark.digest(module_file) + + +def test_fa3_installed_but_broken_dependency_never_uses_alias(monkeypatch): + calls = [] + + def import_module(name): + calls.append(name) + raise ModuleNotFoundError("Missing installed backend dependency", name="flash_attn_3_cuda") + + monkeypatch.setattr(benchmark.importlib, "import_module", import_module) + with pytest.raises(benchmark.UnsupportedBackend, match="Required backend import failed"): + benchmark.backend_inventory("flash3") + assert calls == ["flash_attn_3.flash_attn_interface"] + + +def test_worker_subprocess_and_timeout_are_bounded(tmp_path): + successful = tmp_path / "success" + successful.mkdir() + code, timed_out = benchmark.run_worker_process([sys.executable, "-c", "print('done')"], successful, 10) + assert code == 0 and not timed_out + assert (successful / "stdout.log").read_text().strip() == "done" + timeout = tmp_path / "timeout" + timeout.mkdir() + code, timed_out = benchmark.run_worker_process([sys.executable, "-c", "import time; time.sleep(30)"], timeout, 0.1) + assert timed_out and code != 0 + + +@pytest.mark.parametrize("token_execution,suffix", [("padded", ""), ("packed_valid", "-packed_valid")]) +def test_token_execution_cli_is_explicit_and_keeps_legacy_padded_ids(monkeypatch, capsys, token_execution, suffix): + monkeypatch.setattr(sys, "argv", [ + "benchmark_pixel_execution.py", "--dry-run", "--backends", "flash3", + "--scenarios", "mixed", "--batch-sizes", "2", "--compile-mode", "default", + "--compile-dynamic", "--token-execution", token_execution, + ]) + benchmark.main() + plan = json.loads(capsys.readouterr().out) + candidate, = plan["candidates"] + assert candidate["id"] == "flash3-mixed-b2-default" + suffix + assert candidate["token_execution"] == token_execution + assert candidate["compile_mode"] == "default" and candidate["compile_dynamic"] + arguments = benchmark.model_arguments( + candidate["backend"], candidate["scenario"], compile_mode=candidate["compile_mode"], + compile_dynamic=candidate["compile_dynamic"], token_execution=candidate["token_execution"], + ) + assert arguments["token_execution"] == token_execution + # The compiled/eager gate changes compilation only, not token layout. + assert {**arguments, "compile_mode": "none"}["token_execution"] == token_execution + + +def test_token_execution_cli_defaults_to_archived_padded_id(monkeypatch, capsys): + monkeypatch.setattr(sys, "argv", [ + "benchmark_pixel_execution.py", "--dry-run", "--backends", "flash4", + "--scenarios", "fixed16", "--batch-sizes", "1", + ]) + benchmark.main() + candidate, = json.loads(capsys.readouterr().out)["candidates"] + assert candidate["id"] == "flash4-fixed16-b1-none" + assert candidate["token_execution"] == "padded" + assert benchmark.model_arguments("flash4", "fixed16")["token_execution"] == "padded" + + +def test_token_execution_invalid_choice_fails_closed(): + with pytest.raises(ValueError, match="Unknown execution candidate"): + benchmark.model_arguments("flash4", "mixed", token_execution="automatic") + + +def test_compiled_reference_cannot_change_token_execution(): + eager, compiled = small_model(), small_model() + compiled.load_state_dict(eager.state_dict(), strict=True) + compiled.token_execution = "packed_valid" + with pytest.raises(benchmark.CorrectnessError, match="token execution differs"): + benchmark.check_eager_compiled(eager, compiled, small_batch(), backend="flash4", seed=42, tolerance=.03) + + +def test_worker_reports_unsupported_ragged_backend_without_fallback(tmp_path, monkeypatch): + def unsupported(*args): + raise ValueError("Ragged packed_valid execution requires flash3 or flash4") + monkeypatch.setattr(benchmark, "run_candidate", unsupported) + request = tmp_path / "request.json" + benchmark.write_once(request, {"backend": "sdpa", "token_execution": "packed_valid"}) + benchmark.worker(request) + result = json.loads((tmp_path / "result.json").read_text()) + assert result["status"] == "UNSUPPORTED" + assert result["candidate"]["token_execution"] == "packed_valid" + + +def test_adaln_correctness_probe_is_nontrivial_and_restores_state_and_rng(): + model = ImageProprioModel( + image_dim=12, width=16, depth=1, heads=2, max_cameras=2, max_tokens=32, + objective="visual_only", phase_a_condition_on_proprio=False, + use_phase_embedding=False, noise_conditioning="adaln_zero", + attention_backend="sdpa", attention_dtype="bfloat16", + loss_reduction="mean_per_example", + ) + batch = small_batch() + noisy, levels = benchmark.fixed_corruption(batch, 42) + state = {name: value.clone() for name, value in model.state_dict().items()} + rng = torch.get_rng_state().clone() + assert not model(batch, "libero_gripper", noisy, levels)["image"].any() + with benchmark.nontrivial_conditioning_probe(model, 19) as receipt: + assert receipt["nonzero_adaln_probe"] + first = model(batch, "libero_gripper", noisy, levels)["image"] + changed_levels = (1 - levels).masked_fill(batch["image"]["observed"], 0) + second = model(batch, "libero_gripper", noisy, changed_levels)["image"] + assert first.abs().sum() > 0 + assert not torch.allclose(first, second) + result = benchmark.check_outputs(model, batch, backend="flash4", seed=42, + tolerance=0.03, gradients=True) + assert result["gradient_relative_l2"] < 0.03 + assert torch.equal(rng, torch.get_rng_state()) + assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) + assert all(parameter.grad is None for parameter in model.parameters()) + with pytest.raises(RuntimeError, match="intentional probe failure"): + with benchmark.nontrivial_conditioning_probe(model, 20): + raise RuntimeError("intentional probe failure") + assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) diff --git a/tests/test_pixel_execution_model.py b/tests/test_pixel_execution_model.py new file mode 100644 index 000000000..a5aeb7364 --- /dev/null +++ b/tests/test_pixel_execution_model.py @@ -0,0 +1,170 @@ +"""Batch-independent objective and regional compilation contracts.""" + +import copy + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.world_model import masked_mse + + +def network(**kwargs): + return ImageProprioModel( + objective="visual_only", + phase_a_condition_on_proprio=False, + image_dim=12, + width=24, + heads=3, + depth=2, + max_cameras=2, + loss_reduction="mean_per_example", + **kwargs, + ) + + +def batch(): + torch.manual_seed(19) + valid = torch.tensor([[True] * 6, [True, True, False, True, False, False]]) + return { + "image": { + "value": torch.randn(2, 6, 12), + "time": torch.tensor([[-1.0, -1.0, 0.0, 0.0, 1.0, 1.0]] * 2), + "valid": valid, + "observed": torch.tensor( + [ + [True, True, False, False, False, False], + [True, False, False, False, False, False], + ] + ), + "xy": torch.randn(2, 6, 2), + "camera": torch.tensor([[0, 1, 0, 1, 0, 1]] * 2), + } + } + + +def test_per_example_loss_matches_independent_clips_and_gradients(): + pred = torch.tensor( + [[[1.0], [3.0], [5.0]], [[7.0], [0.0], [0.0]]], requires_grad=True + ) + target = torch.zeros_like(pred) + mask = torch.tensor([[1.0, 1.0, 1.0], [1.0, 0.0, 0.0]]) + actual = masked_mse(pred, target, mask, reduction="mean_per_example") + expected = torch.stack( + [masked_mse(p[None], t[None], m[None]) for p, t, m in zip(pred, target, mask)] + ).mean() + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + torch.autograd.grad(actual, pred, retain_graph=True)[0], + torch.autograd.grad(expected, pred)[0], + ) + assert actual != masked_mse(pred, target, mask) + + +def test_per_example_empty_clip_has_differentiable_zero_weight(): + pred = torch.ones(2, 3, 2, requires_grad=True) + target = torch.zeros_like(pred) + target[1] = float("nan") + mask = torch.tensor([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) + loss = masked_mse(pred, target, mask, reduction="mean_per_example") + assert loss == 0.5 + loss.backward() + assert torch.isfinite(pred.grad).all() and not pred.grad[1].any() + + +def test_batched_image_model_matches_independent_outputs_and_mean_gradients(): + torch.manual_seed(3) + net = network() + independent = copy.deepcopy(net) + data = batch() + values = torch.randn_like(data["image"]["value"]) + levels = torch.full_like(data["image"]["time"], 0.4) + prediction = net(data, "gripper", values, levels)["image"] + expected = [] + for b in range(2): + single = {"image": {k: v[b : b + 1] for k, v in data["image"].items()}} + expected.append( + independent(single, "gripper", values[b : b + 1], levels[b : b + 1])[ + "image" + ] + ) + torch.testing.assert_close(prediction, torch.cat(expected), rtol=1e-5, atol=1e-6) + mask = data["image"]["valid"] & ~data["image"]["observed"] + net._loss(prediction, data["image"]["value"], mask).backward() + torch.stack( + [ + independent._loss( + expected[b], data["image"]["value"][b : b + 1], mask[b : b + 1] + ) + for b in range(2) + ] + ).mean().backward() + for a, b in zip(net.parameters(), independent.parameters()): + if a.grad is not None: + torch.testing.assert_close(a.grad, b.grad, rtol=2e-5, atol=2e-6) + + +def test_regional_compile_matches_eager_and_preserves_checkpoint_names(): + torch.manual_seed(3) + eager = network() + compiled = network( + compile_mode="default", compile_backend="eager", compile_dynamic=True + ) + compiled.load_state_dict(eager.state_dict(), strict=True) + assert eager.state_dict().keys() == compiled.state_dict().keys() + data = batch() + values = torch.randn_like(data["image"]["value"]) + levels = torch.full_like(data["image"]["time"], 0.4) + expected = eager(data, "gripper", values, levels)["image"] + actual = compiled(data, "gripper", values, levels)["image"] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual.square().mean().backward() + expected.square().mean().backward() + for a, b in zip(eager.parameters(), compiled.parameters()): + if a.grad is not None: + torch.testing.assert_close(a.grad, b.grad, rtol=0, atol=0) + assert not any("_orig_mod" in key for key in compiled.state_dict()) + + +def test_sampler_validates_and_prepares_once_without_changing_observations(monkeypatch): + net = network() + data = batch() + calls = {"validate": 0, "prepare": 0} + validate, prepare = net.validate, net._prepare_execution + + def checked(*args, **kwargs): + calls["validate"] += 1 + return validate(*args, **kwargs) + + def prepared(*args, **kwargs): + calls["prepare"] += 1 + return prepare(*args, **kwargs) + + monkeypatch.setattr(net, "validate", checked) + monkeypatch.setattr(net, "_prepare_execution", prepared) + out = net.sample(data, "gripper", steps=4)["image"] + assert calls == {"validate": 1, "prepare": 1} + known = data["image"]["observed"] + assert torch.equal(out[known], data["image"]["value"][known]) + + +def test_training_validates_once(monkeypatch): + net = network() + original = net.validate + calls = [] + + def checked(*args, **kwargs): + calls.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(net, "validate", checked) + assert torch.isfinite(net.training_losses(batch(), "gripper")["total"]) + assert len(calls) == 1 + + +@pytest.mark.parametrize("field", ["time", "xy"]) +def test_value_validation_still_rejects_nonfinite_valid_metadata(field): + data = batch() + data["image"][field][0, 0] = float("nan") + with pytest.raises(ValueError, match="Non-finite"): + network().training_losses(data, "gripper") diff --git a/tests/test_pixel_loader_execution.py b/tests/test_pixel_loader_execution.py new file mode 100644 index 000000000..1b6d84a59 --- /dev/null +++ b/tests/test_pixel_loader_execution.py @@ -0,0 +1,364 @@ +"""Prefetch-safe exact draws, independent validation batching and rank work balance.""" + +import copy +import json +import os +import sys +from itertools import islice +from types import SimpleNamespace + +import pytest +import torch + +from egomimic.rldb.pixel_proprio import ( + DistributedPixelTrainingBatchSampler, + PreparedPixelProprioDataset, + balance_global_pixel_draws, +) +from egomimic.rldb.pixel_proprio_datamodule import PixelProprioDataModule +from egomimic.rldb.world_model import HomogeneousWorldBatchSampler +from tests.test_pixel_attention_profiles import recipe +from tests.test_pixel_proprio_data import ( + _assert_identical_pixel_item, + _fifty_validation_episode_manifest, + pixel_manifest, +) + + +def close_iterator(iterator): + # Tests intentionally abandon prefetched samples, as a stopped process does. + if hasattr(iterator, "_shutdown_workers"): + iterator._shutdown_workers() + + +def _protect_spawn_tracker_fd(worker_id): + """Keep macOS Torch's exec'd shared-memory manager from retaining this pipe. + + Python spawn passes its tracker FD as inheritable. Torch's file_system + transport can inherit it again into a manager that outlives the worker, + creating a tracker/manager wait cycle during interpreter shutdown. This + test-only process hygiene does not change dataset sampling or IPC tensors. + Linux's default fork loader does not need this workaround. + """ + del worker_id + if sys.platform == "darwin": + from multiprocessing import resource_tracker + + os.set_inheritable(resource_tracker.getfd(), False) + + +@pytest.mark.parametrize("balance", [False, True]) +def test_b2_two_worker_prefetch_resume_replays_only_unconsumed_draws(tmp_path, balance): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=80) + manifest = json.loads(path.read_text()) + for episode in manifest["episodes"]: + (path.parent / episode["path"]).unlink() + (path.parent / manifest["normalization"]["path"]).unlink() + kwargs = dict( + manifest_path=path, + image_only=True, + checkpointable_sampling=True, + batch_size=2, + validation_batch_size=1, + distributed_world_size=4, + validation_anchors_per_episode=2, + num_workers=2, + prefetch_factor=3, + persistent_workers=True, + balance_training_work=balance, + ) + module = PixelProprioDataModule(**kwargs) + module.trainer = SimpleNamespace(world_size=4, global_rank=1) + loader = module.train_dataloader() + loader.multiprocessing_context = "spawn" + loader.worker_init_fn = _protect_spawn_tracker_fd + rng = torch.get_rng_state().clone() + iterator = iter(loader) + try: + for cursor in (1, 2): + item = next(iterator)["libero_gripper"] + assert item["sampling_cursor"] == cursor + module.acknowledge_batch(cursor) + assert iterator._send_idx > 2 # Worker queue has advanced past consumed data. + saved = copy.deepcopy(module.state_dict()) + assert saved["consumed_batches"] == 2 + expected = [next(iterator)["libero_gripper"] for _ in range(3)] + # Iterator construction/prefetch must not consume model Torch RNG. + assert torch.equal(torch.get_rng_state(), rng) + finally: + close_iterator(iterator) + for workers in (2, 0): + restored = PixelProprioDataModule( + **{ + **kwargs, + "num_workers": workers, + "prefetch_factor": 2 if workers else None, + "persistent_workers": bool(workers), + } + ) + restored.load_state_dict(saved) + restored.trainer = SimpleNamespace(world_size=4, global_rank=1) + loader = restored.train_dataloader() + if workers: + loader.multiprocessing_context = "spawn" + loader.worker_init_fn = _protect_spawn_tracker_fd + iterator = iter(loader) + try: + for cursor, want in enumerate(expected, start=3): + actual = next(iterator)["libero_gripper"] + assert actual["sampling_cursor"] == cursor + _assert_identical_pixel_item(want, actual) + restored.acknowledge_batch(cursor) + finally: + close_iterator(iterator) + assert restored.consumed_batches == 5 + + +def test_train_b2_prefetch_keeps_all_100_validation_anchors_once_across_four_ranks( + tmp_path, +): + path = _fifty_validation_episode_manifest(tmp_path) + baseline = PixelProprioDataModule( + path, + image_only=True, + checkpointable_sampling=True, + validation_anchors_per_episode=2, + ) + expected = [batch["libero_gripper"] for batch in baseline.val_dataloader()] + visited, episodes = [], [] + for rank in range(4): + module = PixelProprioDataModule( + path, + image_only=True, + checkpointable_sampling=True, + batch_size=2, + validation_batch_size=1, + distributed_world_size=4, + validation_anchors_per_episode=2, + num_workers=2, + prefetch_factor=2, + persistent_workers=True, + balance_training_work=True, + ) + assert len(module.val_dataloader()) == 100 # Unattached inspection is global. + module.trainer = SimpleNamespace(world_size=4, global_rank=rank) + loader = module.val_dataloader() + loader.worker_init_fn = _protect_spawn_tracker_fd + assert len(loader) == 25 + iterator = iter(loader) + try: + for batch in iterator: + item = batch["libero_gripper"] + (index,) = item["validation_batch_index"] + visited.append(index) + episodes.append(item["pixel_metadata"][0]["episode_id"]) + _assert_identical_pixel_item(expected[index], item) + finally: + close_iterator(iterator) + assert visited == list(range(100)) + assert len(set(episodes)) == 50 + assert all(episodes.count(episode) == 2 for episode in set(episodes)) + + +def test_loader_execution_options_preserve_legacy_context_but_explicit_batches_bind( + tmp_path, +): + path = pixel_manifest(tmp_path, actions=False) + kwargs = dict(manifest_path=path, image_only=True, checkpointable_sampling=True) + legacy = PixelProprioDataModule(**kwargs) + context = legacy.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + state = legacy.state_dict() + assert ( + "batching" not in context.state and "training_work_balance" not in context.state + ) + changed_runtime = PixelProprioDataModule( + **kwargs, + num_workers=2, + pin_memory=True, + persistent_workers=True, + prefetch_factor=3, + ) + changed_runtime.load_state_dict(state) + assert changed_runtime.state_dict() == state + assert ( + changed_runtime.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ).state + == context.state + ) + loader = changed_runtime.train_dataloader() + assert loader.num_workers == 2 and loader.prefetch_factor == 3 + assert loader.pin_memory and loader.persistent_workers + explicit = PixelProprioDataModule(**kwargs, batch_size=2, validation_batch_size=1) + explicit_context = explicit.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert explicit_context.state["batching"] == { + "schema": "pixel-split-batching-v1", + "train_batch_size": 2, + "validation_batch_size": 1, + } + for invalid in ( + PixelProprioDataModule(**kwargs, batch_size=2), + PixelProprioDataModule(**kwargs, batch_size=2, validation_batch_size=2), + ): + with pytest.raises(ValueError, match="batching"): + invalid.load_state_dict(explicit.state_dict()) + with pytest.raises(ValueError, match="context mismatch"): + invalid.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=explicit_context.state, + ) + + +@pytest.mark.parametrize( + "options", + [ + {"batch_size": False}, + {"batch_size": 0}, + {"validation_batch_size": 0}, + {"validation_batch_size": True}, + {"num_workers": -1}, + {"num_workers": True}, + {"persistent_workers": True}, + {"prefetch_factor": 2}, + {"num_workers": 2, "prefetch_factor": 0}, + {"num_workers": 2, "prefetch_factor": True}, + {"pin_memory": "yes"}, + {"persistent_workers": 1}, + {"balance_training_work": 1}, + {"balance_training_work": True}, + ], +) +def test_loader_options_fail_closed(options): + with pytest.raises(ValueError): + PixelProprioDataModule("unused.json", **options) + + +def test_longest_first_fixed_capacity_preserves_draws_and_reduces_worst_work(): + costs = [100, 81, 64, 49, 16, 9, 4, 1] + source = SimpleNamespace(predicted_token_work=lambda index, draw: costs[index]) + draws = [(i, 7, i) for i in range(8)] + balanced = balance_global_pixel_draws(source, draws, 2, 4) + assert all(len(rank) == 2 for rank in balanced) + assert sorted(draw for rank in balanced for draw in rank) == draws + assert max(sum(costs[draw[0]] for draw in rank) for rank in balanced) == 101 + assert ( + max(sum(costs[draw[0]] for draw in draws[i : i + 2]) for i in range(0, 8, 2)) + == 181 + ) + assert balance_global_pixel_draws(source, draws, 2, 4) == balanced + assert balance_global_pixel_draws(source, draws[:4], 1, 4) == [ + [draw] for draw in draws[:4] + ] + equal = SimpleNamespace(predicted_token_work=lambda index, draw: 256) + assert balance_global_pixel_draws(equal, draws, 2, 4) == [ + draws[i : i + 2] for i in range(0, 8, 2) + ] + with pytest.raises(ValueError, match="complete global batch"): + balance_global_pixel_draws(source, draws[:-1], 2, 4) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_work_preview_shares_exact_frame_rng_and_balancing_preserves_global_examples( + tmp_path, mode +): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=80) + data = PreparedPixelProprioDataset( + path, + image_only=True, + mode=mode, + discrete_frames=16, + stride=4, + sampling_strategy="full_episode_uniform" + if mode == "continuous" + else "discrete_window", + ) + global_draws = list( + islice(HomogeneousWorldBatchSampler(data, 8, 42, checkpointable=True), 8) + ) + assert data.rng is None + work = { + draw: data.predicted_token_work(draw[0], draw[1:]) + for batch in global_draws + for draw in batch + } + assert data.rng is None # Preview does not read pixels or initialize worker RNG. + ranked = [] + for rank in range(4): + sampler = DistributedPixelTrainingBatchSampler( + data, 2, 42, world_size=4, rank=rank, balance_training_work=True + ) + ranked.append(list(islice(sampler, 8))) + sampler.consumed_batches = 3 + assert list(islice(sampler, 5)) == ranked[-1][3:] + reassigned = False + for step, original in enumerate(global_draws): + combined = [draw for rank in range(4) for draw in ranked[rank][step]] + assert sorted(combined) == sorted(original) + reassigned |= combined != original + max_after = max( + sum(work[draw] for draw in ranked[rank][step]) for rank in range(4) + ) + max_before = max( + sum(work[draw] for draw in original[i : i + 2]) for i in range(0, 8, 2) + ) + assert max_after <= max_before + for draw in original: + expected, actual = ( + data[draw], + data[next(key for key in combined if key == draw)], + ) + assert work[draw] == len(actual["world"]["image"]["value"]) ** 2 + assert actual["pixel_metadata"] == expected["pixel_metadata"] + for field in actual["world"]["image"]: + assert torch.equal( + actual["world"]["image"][field], expected["world"]["image"][field] + ) + assert reassigned is (mode == "continuous") + + +def test_work_assignment_choice_is_context_and_resume_bound(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False) + kwargs = dict( + manifest_path=path, + image_only=True, + checkpointable_sampling=True, + batch_size=2, + validation_batch_size=1, + distributed_world_size=4, + ) + balanced = PixelProprioDataModule(**kwargs, balance_training_work=True) + state = balanced.state_dict() + context = balanced.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert "rank_seeded" in context.state["training_work_balance"]["model_noise"] + unbalanced = PixelProprioDataModule(**kwargs) + with pytest.raises(ValueError, match="work balance"): + unbalanced.load_state_dict(state) + with pytest.raises(ValueError, match="context mismatch"): + unbalanced.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + + +@pytest.mark.parametrize("arm", ["discrete", "continuous"]) +def test_prefetch_execution_profile_keeps_validation_batch_fixed(arm): + cfg = recipe(arm, execution="ddp4_prefetch") + assert cfg.data.batch_size == cfg.data.validation_batch_size == 1 + assert cfg.data.distributed_world_size == 4 + assert cfg.data.num_workers == cfg.data.prefetch_factor == 2 + assert cfg.data.pin_memory and cfg.data.persistent_workers + assert cfg.data.balance_training_work and cfg.data.checkpointable_sampling + assert not cfg.trainer.use_distributed_sampler diff --git a/tests/test_pixel_noise_conditioning_profiles.py b/tests/test_pixel_noise_conditioning_profiles.py new file mode 100644 index 000000000..28e82f178 --- /dev/null +++ b/tests/test_pixel_noise_conditioning_profiles.py @@ -0,0 +1,269 @@ +"""Pilot AdaLN selection and strict preservation when adding frozen readouts.""" + +import copy + +import pytest +import torch + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.pixel_diffusion import LinearDDIMImageSampler +from egomimic.pipeline.core import Pipeline, sum_losses +from egomimic.pipeline.stages_pixel_readout import configure_readout_trainability +from scripts.world_model.initialize_pixel_readout import ( + initialize_readout, + resolve_model_config, + sha256_file, + transfer_shared_weights, + validate_transition, + verify_phase_a, +) +from tests.test_pixel_readout_graph import IDS, fixture_batch, stages +from tests.test_pixel_token_execution_profiles import configuration, denoiser + + +@pytest.fixture(autouse=True) +def single_cpu_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +@pytest.mark.parametrize("stage", ["video", "readout"]) +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +@pytest.mark.parametrize("profile", [None, "packed_regional_compile"]) +def test_current_pilots_select_adaln_without_other_identity_conditioning( + stage, arm, profile +): + model = denoiser(configuration(stage, arm, profile)) + assert model["noise_conditioning"] == "adaln_zero" + assert model["use_phase_embedding"] is False + assert model["use_modality_embedding"] is False + assert model["use_embodiment_embedding"] is False + assert model["use_camera_embedding"] is True + assert model["time_mode"] == arm + assert model["position_encoding"] == "spatiotemporal_rope" + assert model["rope_axis_dims"] == [64, 32, 32] + assert model["spatial_grid_size"] == [10, 10] + + +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_readout_requires_exact_noise_conditioning_and_legacy_default(arm): + source = denoiser(configuration("video", arm, "packed_regional_compile")) + target = denoiser(configuration("readout", arm, "packed_regional_compile")) + left, right = validate_transition(source, target) + assert left["noise_conditioning"] == right["noise_conditioning"] == "adaln_zero" + del source["noise_conditioning"] + assert resolve_model_config(source)["noise_conditioning"] == "additive" + with pytest.raises(ValueError, match="configuration mismatch.*noise_conditioning"): + validate_transition(source, target) + del target["noise_conditioning"] + left, right = validate_transition(source, target) + assert left["noise_conditioning"] == right["noise_conditioning"] == "additive" + + +def model_arguments(mode="continuous"): + return dict( + objective="visual_only", + phase_a_condition_on_proprio=False, + pose_dims=None, + action_dims=None, + image_dim=12, + width=16, + depth=2, + heads=2, + mlp_ratio=2, + max_cameras=2, + max_tokens=16, + time_mode=mode, + use_phase_embedding=False, + noise_conditioning="adaln_zero", + loss_reduction="mean_per_example", + ) + + +def readout_arguments(source): + return { + **source, + "objective": "visual_state_action", + "pose_dims": {"robot": 3}, + "action_dims": {"robot": 2}, + } + + +def _trained_image_source(arguments): + """Nonzero gates/modulations/head prevent a vacuous zero-output proof.""" + with torch.random.fork_rng(devices=[]): + torch.manual_seed(349) + model = ImageProprioModel(**arguments) + with torch.no_grad(): + for name, parameter in model.named_parameters(): + if any(part in name for part in ("noise_level", "adaln", "image_out")): + parameter.normal_(mean=0, std=0.05) + return model + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("spatial_rope", [False, True]) +def test_adaln_checkpoint_expansion_freezes_all_image_parameters( + tmp_path, mode, spatial_rope +): + source_arguments = model_arguments(mode) + if spatial_rope: + source_arguments.update( + position_encoding="spatiotemporal_rope", + rope_axis_dims=[4, 2, 2], + spatial_grid_size=[2, 2], + ) + target_arguments = readout_arguments(source_arguments) + source = _trained_image_source(source_arguments) + source_path = tmp_path / "source.pt" + target_path = tmp_path / "readout.pt" + torch.save(source.state_dict(), source_path) + receipt = initialize_readout( + source_checkpoint=source_path, + source_sha256=sha256_file(source_path), + source_config=source_arguments, + target_config=target_arguments, + output_path=target_path, + receipt_path=tmp_path / "receipt.json", + seed=42, + ) + assert receipt["phase_a_preservation"]["bit_exact"] + assert not receipt["optimizer_state_transferred"] + assert not receipt["training_step_transferred"] + assert receipt["output"]["sha256"] == sha256_file(target_path) + assert receipt["source_model"]["noise_conditioning"] == "adaln_zero" + assert receipt["target_model"]["noise_conditioning"] == "adaln_zero" + conditioning = { + name for name in source.state_dict() if "adaln" in name or "noise_level" in name + } + assert any(name.startswith("final_adaln.") for name in conditioning) + assert any(".adaln_image." in name for name in conditioning) + assert conditioning <= set(receipt["transfer"]["copied_tensors"]) + target = ImageProprioModel(**target_arguments) + target.load_state_dict(torch.load(target_path, weights_only=True), strict=True) + graph = Pipeline(stages(target), stage_ids=IDS) + policy = graph.stage_by_id("pixel_denoiser").trainability + assert policy["noise_conditioning"] == "adaln_zero" + assert not policy["phase_embedding_present"] + assert hasattr(target, "xy") is not spatial_rope + assert conditioning <= set(policy["frozen_parameters"]) + assert set(source.state_dict()) <= set(policy["frozen_parameters"]) + before = { + name: parameter.detach().clone() + for name, parameter in target.named_parameters() + } + batch = fixture_batch() + sampler = LinearDDIMImageSampler(3) + torch.manual_seed(78) + expected_images = sampler(source, {"image": batch["world"]["image"]}, "robot") + # Check unknown pixels too: observed clean values would be nonzero even if + # the image head were still at its constructor's zero initialization. + unknown = batch["world"]["image"]["valid"] & ~batch["world"]["image"]["observed"] + assert expected_images[unknown].abs().sum() > 0 + optimizer = torch.optim.AdamW(graph.parameters(), lr=5e-5, weight_decay=0.01) + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + output = graph.execute(batch, mode="train") + assert not output["loss/pixel_video"].requires_grad + loss = sum_losses(output) + assert loss.requires_grad and torch.isfinite(loss) + loss.backward() + for name, parameter in target.named_parameters(): + if parameter.requires_grad: + assert ( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + ), name + else: + assert parameter.grad is None, name + optimizer.step() + for name in policy["frozen_parameters"]: + assert torch.equal(target.state_dict()[name], before[name]), name + for name in ("pose_out.robot.weight", "action_out.robot.weight"): + assert not torch.equal(target.state_dict()[name], before[name]), name + assert verify_phase_a(source, target)["bit_exact"] + torch.manual_seed(78) + actual_images = sampler(target, batch["world"], "robot") + torch.testing.assert_close(actual_images, expected_images, rtol=0, atol=0) + restored = Pipeline(stages(ImageProprioModel(**target_arguments)), stage_ids=IDS) + restored.load_state_dict(graph.state_dict(), strict=True) + assert restored.stage_by_id("pixel_denoiser").trainability == policy + + +def test_direct_transfer_rejects_noise_mode_drift_and_joint_mode_unfreezes_adaln(): + arguments = model_arguments() + source = ImageProprioModel(**arguments) + target_arguments = readout_arguments(arguments) + additive_target = ImageProprioModel( + **{**target_arguments, "noise_conditioning": "additive"} + ) + with pytest.raises(ValueError, match="noise_conditioning must match"): + transfer_shared_weights(source, additive_target) + target = ImageProprioModel(**target_arguments) + transfer_shared_weights(source, target) + policy = configure_readout_trainability(target, "frozen_image") + assert policy["frozen_parameters"] + policy = configure_readout_trainability(target, "joint") + assert not policy["frozen_parameters"] + assert all(parameter.requires_grad for parameter in target.parameters()) + # A fresh strict load retains its explicitly selected noise architecture. + restored = ImageProprioModel(**copy.deepcopy(target_arguments)) + restored.load_state_dict(target.state_dict(), strict=True) + + +def test_padded_stage_two_sampler_caches_frame_groups_once(monkeypatch): + model = _trained_image_source(readout_arguments(model_arguments())) + original = model._prepare_image_execution + calls = [] + + def prepare(image, valid): + calls.append(valid.shape) + return original(image, valid) + + monkeypatch.setattr(model, "_prepare_image_execution", prepare) + image = LinearDDIMImageSampler(3)(model, fixture_batch()["world"], "robot") + assert torch.isfinite(image).all() + assert calls == [torch.Size([2, 4])] + + +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_readout_initialization_binds_position_encoding_and_geometry(arm): + source = denoiser(configuration("video", arm, "packed_regional_compile")) + target = denoiser(configuration("readout", arm, "packed_regional_compile")) + validate_transition(source, target) + for field, value in ( + ("position_encoding", "linear_xy_temporal_rope"), + ("rope_axis_dims", [32, 48, 48]), + ("spatial_grid_size", [20, 20]), + ): + changed = {**target, field: value} + with pytest.raises(ValueError, match=f"configuration mismatch.*{field}"): + validate_transition(source, changed) + for cfg in (source, target): + for field in ("position_encoding", "rope_axis_dims", "spatial_grid_size"): + cfg.pop(field) + left, right = validate_transition(source, target) + assert ( + left["position_encoding"] + == right["position_encoding"] + == "linear_xy_temporal_rope" + ) + assert left["rope_axis_dims"] is right["rope_axis_dims"] is None + assert left["spatial_grid_size"] is right["spatial_grid_size"] is None + + +@pytest.mark.parametrize( + "field,value", [("rope_axis_dims", [2, 2, 4]), ("spatial_grid_size", [4, 4])] +) +def test_direct_tensor_transfer_rejects_same_shape_positional_drift(field, value): + source_args = { + **model_arguments(), + "position_encoding": "spatiotemporal_rope", + "rope_axis_dims": [4, 2, 2], + "spatial_grid_size": [2, 2], + } + source = ImageProprioModel(**source_args) + target = ImageProprioModel(**{**readout_arguments(source_args), field: value}) + with pytest.raises(ValueError, match=f"{field} must match"): + transfer_shared_weights(source, target) diff --git a/tests/test_pixel_packed_execution.py b/tests/test_pixel_packed_execution.py new file mode 100644 index 000000000..90d68acef --- /dev/null +++ b/tests/test_pixel_packed_execution.py @@ -0,0 +1,464 @@ +"""Packed image-expert parity using explicit CPU reference attention kernels. + +These tests exercise the actual packing/layout/transformer/scatter path, while +replacing optional CUDA kernel loading with isolated SDPA math references. They +make no claim about actual FA3/FA4 kernels or CUDA Inductor performance. +""" + +from contextlib import nullcontext + +import pytest +import torch +from torch.nn import functional as F + +from egomimic.models.attention_execution import FlashAttentionKernels +from egomimic.models.image_proprio import ImageProprioModel +from scripts.world_model import benchmark_pixel_execution as benchmark + + +@pytest.fixture(autouse=True) +def one_cpu_thread(): + original = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(original) + + +@pytest.fixture +def reference_kernels(monkeypatch): + calls = [] + + def dense(q, k, v, *, causal): + assert causal is False + return ( + F.scaled_dot_product_attention( + q.float().transpose(1, 2), + k.float().transpose(1, 2), + v.float().transpose(1, 2), + dropout_p=0.0, + ) + .transpose(1, 2) + .to(q.dtype) + ) + + def load(backend, device): + assert backend in ("flash3", "flash4") + assert torch.device(device).type == "cpu" + + def dense_reference(q, k, v, **kwargs): + calls.append(("dense", tuple(q.shape))) + output = dense(q, k, v, **kwargs) + return (output, None) if backend == "flash4" else output + + def packed_reference( + q, k, v, *, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal + ): + assert cu_seqlens_q.dtype == cu_seqlens_k.dtype == torch.int32 + qs, ks = cu_seqlens_q.tolist(), cu_seqlens_k.tolist() + assert max_seqlen_q == max(b - a for a, b in zip(qs, qs[1:])) + assert max_seqlen_k == max(b - a for a, b in zip(ks, ks[1:])) + assert qs[-1] == len(q) and ks[-1] == len(k) + calls.append(("varlen", tuple(q.shape), tuple(qs), tuple(ks))) + output = torch.cat( + [ + dense( + q[qa:qb][None], k[ka:kb][None], v[ka:kb][None], causal=causal + )[0] + for qa, qb, ka, kb in zip(qs, qs[1:], ks, ks[1:]) + ] + ) + return (output, None) if backend == "flash4" else output + + return FlashAttentionKernels(backend, dense_reference, packed_reference) + + monkeypatch.setattr( + "egomimic.models.image_proprio.load_flash_attention_kernels", load + ) + monkeypatch.setattr( + "egomimic.models.world_model.load_flash_attention_kernels", load + ) + return calls + + +def model(token_execution="padded", backend="flash3", **overrides): + arguments = dict( + objective="visual_only", + phase_a_condition_on_proprio=False, + image_dim=12, + width=16, + depth=2, + heads=2, + max_tokens=32, + mlp_ratio=2, + attention_backend=backend, + attention_dtype="bfloat16", + loss_reduction="mean_per_example", + token_execution=token_execution, + ) + arguments.update(overrides) + return ImageProprioModel(**arguments) + + +def batch(nonprefix=False, all_valid=False): + result = benchmark.make_batch( + (3, 3) if all_valid else (2, 3), + seed=57, + device=torch.device("cpu"), + image_size=4, + patch_size=2, + ) + image = result["image"] + if nonprefix: + permutation = torch.tensor( + [ + 0, + 16, + 1, + 17, + 2, + 18, + 3, + 19, + 4, + 20, + 5, + 21, + 6, + 22, + 7, + 23, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + ] + ) + for key, value in image.items(): + image[key] = value.index_select(1, permutation) + # Invalid padding must be eliminated before it can contaminate features. + image["value"][~image["valid"]] = float("nan") + image["xy"][~image["valid"]] = float("nan") + image["time"][~image["valid"]] = float("nan") + image["camera"][~image["valid"]] = -999 + return result + + +def relative_l2(actual, reference): + actual, reference = actual.float(), reference.float() + return float((actual - reference).norm() / reference.norm().clamp_min(1e-8)) + + +def assert_gradient_parity(actual, reference): + assert ( + dict(actual.named_parameters()).keys() + == dict(reference.named_parameters()).keys() + ) + errors, scales = 0.0, 0.0 + for name, parameter in actual.named_parameters(): + expected = dict(reference.named_parameters())[name] + assert parameter.grad is not None and expected.grad is not None, name + assert parameter.grad.dtype == expected.grad.dtype == torch.float32 + assert ( + torch.isfinite(parameter.grad).all() and torch.isfinite(expected.grad).all() + ), name + torch.testing.assert_close( + parameter.grad, expected.grad, rtol=0.04, atol=0.002, msg=name + ) + errors += float((parameter.grad - expected.grad).double().square().sum()) + scales += float(expected.grad.double().square().sum()) + assert (errors / max(scales, 1e-30)) ** 0.5 < 0.03 + + +@pytest.mark.parametrize("backend", ["flash3", "flash4"]) +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("time_mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("nonprefix", [False, True]) +def test_packed_outputs_all_parameter_gradients_and_projection_work_match( + reference_kernels, backend, mixed, nonprefix, time_mode +): + torch.manual_seed(71) + padded, packed = ( + model(backend=backend, time_mode=time_mode), + model("packed_valid", backend, time_mode=time_mode), + ) + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix) + noisy, levels = benchmark.fixed_corruption(world, 19) + shapes = [] + hooks = [ + module.register_forward_pre_hook( + lambda _, args: shapes.append(tuple(args[0].shape)) + ) + for module in ( + packed.image_in, + packed.image_out, + *(block.qkv[0] for block in packed.blocks), + ) + ] + outputs = [] + for network in (padded, packed): + context = ( + torch.autocast("cpu", dtype=torch.bfloat16) if mixed else nullcontext() + ) + with context: + predicted = network(world, "libero_gripper", noisy, levels)["image"] + loss = benchmark.example_mean_loss(predicted, world) + assert loss.dtype == torch.float32 and torch.isfinite(loss) + loss.backward() + outputs.append(predicted.detach()) + for hook in hooks: + hook.remove() + assert len(shapes) == 4 + assert all(shape[:2] == (1, int(world["image"]["valid"].sum())) for shape in shapes) + valid = world["image"]["valid"] + assert relative_l2(outputs[1][valid], outputs[0][valid]) < 0.03 + assert not outputs[1][~valid].any() + assert_gradient_parity(packed, padded) + assert all(call[0] == "varlen" for call in reference_kernels) + assert all(call[2] == call[3] == (0, 16, 40) for call in reference_kernels) + + +@pytest.mark.parametrize("backend", ["flash3", "flash4"]) +def test_full_valid_batch_retains_existing_dense_path( + reference_kernels, monkeypatch, backend +): + padded, packed = model(backend=backend), model("packed_valid", backend) + packed.load_state_dict(padded.state_dict(), strict=True) + for block in packed.blocks: + monkeypatch.setattr( + block, + "forward_packed", + lambda *a, **kw: pytest.fail("unneeded packed path"), + ) + world = batch(all_valid=True) + noisy, levels = benchmark.fixed_corruption(world, 29) + actual = packed(world, "libero_gripper", noisy, levels)["image"] + expected = padded(world, "libero_gripper", noisy, levels)["image"] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert len(reference_kernels) == 4 and all( + call[0] == "dense" for call in reference_kernels + ) + + +def test_packing_does_not_change_full_shaped_noising_or_rng_order( + reference_kernels, monkeypatch +): + padded, packed = model(), model("packed_valid") + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix=True) + captures = [] + for network in (padded, packed): + original = network._noisy_images + + def record(*args, _original=original, **kwargs): + noisy, levels = _original(*args, **kwargs) + captures.append((noisy.detach().clone(), levels.detach().clone())) + return noisy, levels + + monkeypatch.setattr(network, "_noisy_images", record) + results, states = [], [] + for network in (padded, packed): + torch.manual_seed(91) + results.append(network.training_losses(world, "libero_gripper")) + states.append(torch.get_rng_state().clone()) + assert torch.equal(states[0], states[1]) + assert captures[0][0].shape == world["image"]["value"].shape + assert all(torch.equal(left, right) for left, right in zip(*captures)) + torch.testing.assert_close( + results[0]["total"], results[1]["total"], rtol=0.03, atol=1e-6 + ) + + +def test_packed_benchmark_independent_b1_and_same_mode_isolation(reference_kernels): + result = benchmark.check_outputs( + model("packed_valid"), + batch(nonprefix=True), + backend="flash3", + seed=27, + tolerance=0.03, + gradients=True, + ) + assert result["output_relative_l2"] <= 0.03 + assert result["gradient_relative_l2"] <= 0.03 + assert result["cross_example_relative_l2"] == 0 + + +def test_phase_a_sampler_matches_and_reuses_one_packed_layout( + reference_kernels, monkeypatch +): + padded, packed = model(), model("packed_valid") + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix=True) + calls = [] + original = packed._prepare_image_execution + + def prepare(*args, **kwargs): + result = original(*args, **kwargs) + calls.append(result) + return result + + monkeypatch.setattr(packed, "_prepare_image_execution", prepare) + torch.manual_seed(102) + expected = padded.sample(world, "libero_gripper", steps=3)["image"] + reference_rng = torch.get_rng_state().clone() + torch.manual_seed(102) + actual = packed.sample(world, "libero_gripper", steps=3)["image"] + assert len(calls) == 1 + assert torch.equal(torch.get_rng_state(), reference_rng) + assert relative_l2(actual, expected) < 0.03 + observed = world["image"]["observed"] + assert torch.equal(actual[observed], world["image"]["value"][observed]) + assert not actual[~world["image"]["valid"]].any() + assert actual.dtype == torch.float32 + + +def stage_two_batch(): + world = batch(nonprefix=True) + world["proprio"] = { + "value": torch.tensor( + [[[0.1, 0.2, -0.3], [0.4, -0.5, 0.6]], [[-0.2, 0.3, 0.1], [0.5, 0.1, -0.4]]] + ), + "time": torch.tensor([[0.0, 0.1], [0.0, 0.1]]), + "index": torch.tensor([[0, 2], [0, 2]]), + "valid": torch.ones(2, 2, dtype=torch.bool), + "observed": torch.tensor([[True, False], [True, False]]), + } + world["action"] = { + **{ + key: value.clone() + for key, value in world["proprio"].items() + if key != "value" + }, + "value": torch.zeros(2, 2, 2), + "observed": torch.zeros(2, 2, dtype=torch.bool), + } + return world + + +def test_stage_two_keeps_phase_b_shared_expert_path(reference_kernels, monkeypatch): + arguments = dict( + objective="visual_state_action", + pose_dims={"libero_gripper": 3}, + action_dims={"libero_gripper": 2}, + ) + padded, packed = model(**arguments), model("packed_valid", **arguments) + packed.load_state_dict(padded.state_dict(), strict=True) + world = stage_two_batch() + noisy, levels = benchmark.fixed_corruption(world, 11) + for block in packed.blocks: + monkeypatch.setattr( + block, + "forward_packed", + lambda *a, **kw: pytest.fail("Phase B must retain joint attention"), + ) + before = tuple(packed.state_dict()) + actual = packed(world, "libero_gripper", noisy, levels, phase="B") + expected = padded(world, "libero_gripper", noisy, levels, phase="B") + assert set(actual) == {"image", "proprio", "action"} + for key in actual: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + assert tuple(packed.state_dict()) == before == tuple(padded.state_dict()) + assert all(len(block.qkv) == 2 for block in packed.blocks) + + +def test_stage_two_observed_proprio_phase_a_keeps_joint_path( + reference_kernels, monkeypatch +): + arguments = dict( + objective="visual_state_action", + phase_a_condition_on_proprio=True, + pose_dims={"libero_gripper": 3}, + action_dims={"libero_gripper": 2}, + ) + padded, packed = model(**arguments), model("packed_valid", **arguments) + packed.load_state_dict(padded.state_dict(), strict=True) + for block in packed.blocks: + monkeypatch.setattr( + block, + "forward_packed", + lambda *a, **kw: pytest.fail("Observed-proprio Phase A is not image-only"), + ) + world = stage_two_batch() + noisy, levels = benchmark.fixed_corruption(world, 31) + actual = packed(world, "libero_gripper", noisy, levels, phase="A") + expected = padded(world, "libero_gripper", noisy, levels, phase="A") + for key in actual: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + + +def test_state_keys_parameter_count_and_default_execution_are_unchanged(): + arguments = benchmark.model_arguments("flash3", "mixed") + assert arguments["token_execution"] == "padded" + with torch.device("meta"): + padded = ImageProprioModel(**arguments) + packed = ImageProprioModel(**{**arguments, "token_execution": "packed_valid"}) + assert ( + padded.token_execution == "padded" and packed.token_execution == "packed_valid" + ) + assert tuple(packed.state_dict()) == tuple(padded.state_dict()) + assert ( + sum(p.numel() for p in padded.parameters()) + == sum(p.numel() for p in packed.parameters()) + == 88_178_688 + ) + assert all(parameter.dtype == torch.float32 for parameter in packed.parameters()) + + +def test_unsupported_ragged_backend_or_chunks_fail_closed(reference_kernels): + world = batch(nonprefix=True) + noisy, levels = benchmark.fixed_corruption(world, 37) + with pytest.raises(ValueError, match="flash3 or flash4"): + model("packed_valid", backend="sdpa")(world, "libero_gripper", noisy, levels) + with pytest.raises(ValueError, match="unchunked"): + model("packed_valid", attention_query_chunk_size=2)( + world, "libero_gripper", noisy, levels + ) + # Fully valid native SDPA is still the declared dense route, never a fallback. + dense = batch(all_valid=True) + noisy, levels = benchmark.fixed_corruption(dense, 37) + result = model("packed_valid", backend="sdpa")( + dense, "libero_gripper", noisy, levels + ) + assert torch.isfinite(result["image"]).all() + + +def test_packed_regional_compile_matches_eager_with_cpu_graph_backend( + reference_kernels, +): + torch.manual_seed(117) + eager = model("packed_valid") + compiled = model( + "packed_valid", + compile_mode="default", + compile_dynamic=True, + compile_backend="eager", + ) + compiled.load_state_dict(eager.state_dict(), strict=True) + world = batch(nonprefix=True) + noisy, levels = benchmark.fixed_corruption(world, 71) + actual = compiled(world, "libero_gripper", noisy, levels)["image"] + expected = eager(world, "libero_gripper", noisy, levels)["image"] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + benchmark.example_mean_loss(actual, world).backward() + benchmark.example_mean_loss(expected, world).backward() + assert_gradient_parity(compiled, eager) + + +def test_packed_optional_camera_embedding_does_not_require_camera_metadata( + reference_kernels, +): + padded, packed = ( + model(use_camera_embedding=False), + model("packed_valid", use_camera_embedding=False), + ) + packed.load_state_dict(padded.state_dict(), strict=True) + world = batch(nonprefix=True) + del world["image"]["camera"] + noisy, levels = benchmark.fixed_corruption(world, 78) + actual = packed(world, "libero_gripper", noisy, levels)["image"] + expected = padded(world, "libero_gripper", noisy, levels)["image"] + assert relative_l2(actual, expected) < 0.03 diff --git a/tests/test_pixel_prediction_diagnostics.py b/tests/test_pixel_prediction_diagnostics.py new file mode 100644 index 000000000..44dc45f5c --- /dev/null +++ b/tests/test_pixel_prediction_diagnostics.py @@ -0,0 +1,228 @@ +"""RGB reconstruction, mask semantics, metadata binding and real-time videos.""" + +import copy +import json +from types import SimpleNamespace + +import av +import numpy as np +import pytest +import torch + +from egomimic.eval.image_proprio_eval import ImageProprioEval +from egomimic.eval.pixel_diagnostics import ( + PixelPredictionDiagnostics, + pixel_frame_records, + unpatchify_rgb, + write_timestamp_video, +) +from egomimic.rldb.pixel_proprio import patchify_rgb + + +def test_wandb_media_uses_optimizer_axis_not_history_counter(tmp_path, monkeypatch): + import wandb + from lightning.pytorch.loggers import WandbLogger + + calls = [] + experiment = SimpleNamespace(log=lambda *args, **kwargs: calls.append((args, kwargs))) + logger = WandbLogger(experiment=experiment, offline=True) + monkeypatch.setattr(wandb, "Image", lambda path: ("image", path)) + diagnostic = PixelPredictionDiagnostics(tmp_path, log_wandb=True) + diagnostic.written = [("robot", "forecast", "image", tmp_path / "example.png")] + model = SimpleNamespace(global_rank=0, global_step=250, + trainer=SimpleNamespace(loggers=[logger])) + diagnostic.log_to_wandb(model) + assert len(calls) == 1 + args, kwargs = calls[0] + assert "step" not in kwargs + assert args[0]["trainer/global_step"] == 250 + assert "val_pixel/robot/forecast/image" in args[0] + + +def batch_fixture(direct=False): + frames = 1 if direct else 3 + rgb = np.arange(frames * 2 * 4 * 4 * 3, dtype=np.uint8).reshape(frames, 2, 4, 4, 3) + patches = patchify_rgb(rgb, 2) + clock = np.array([10.1, 10.3, 11.05])[:frames] + origin = clock[0] if direct else clock[1] + xy = np.array([[-0.5, -0.5], [0.5, -0.5], [-0.5, 0.5], [0.5, 0.5]], np.float32) + known = np.arange(frames) == 0 + image = { + "value": torch.from_numpy(patches.reshape(1, frames * 8, 12)), + "time": torch.from_numpy(np.repeat(clock - origin, 8).astype(np.float32)[None]), + "index": torch.from_numpy( + np.repeat(np.arange(frames), 8).astype(np.float32)[None] + ), + "camera": torch.from_numpy(np.tile(np.repeat(np.arange(2), 4), frames)[None]), + "xy": torch.from_numpy(np.tile(xy, (frames * 2, 1))[None]), + "valid": torch.ones(1, frames * 8, dtype=torch.bool), + "observed": torch.from_numpy(np.repeat(known, 8)[None]), + } + proprio = { + "value": torch.ones(1, frames, 3), + "time": torch.from_numpy((clock - origin).astype(np.float32)[None]), + "valid": torch.ones(1, frames, dtype=torch.bool), + "observed": torch.from_numpy(known[None]), + } + metadata = { + "episode_id": "example-episode", + "camera_names": ["agentview", "wrist"], + "frame_indices": [2, 9, 18][:frames], + "timestamps_seconds": clock.tolist(), + "time_origin_seconds": float(origin), + "image_size": [4, 4], + "patch_size": 2, + "normalization": "uint8_to_minus_one_one", + "patch_order": "row_major_hwc", + } + item = { + "world": {"image": image, "proprio": proprio}, + "embodiment": "robot", + "task_kind": ["direct_action" if direct else "forecast"], + "pixel_metadata": [metadata], + } + if direct: + item["world"]["action"] = { + "value": torch.tensor([[[0.1, 0.2]]]), + "time": torch.zeros(1, 1), + "valid": torch.ones(1, 1, dtype=torch.bool), + "observed": torch.zeros(1, 1, dtype=torch.bool), + } + return {"robot": item}, rgb + + +def fake_model(batch, direct=False): + world = batch["robot"]["world"] + prediction = ( + {"action": world["action"]["value"] + 0.5} + if direct + else { + "image": world["image"]["value"] + + (~world["image"]["observed"])[..., None] * 0.2 + } + ) + + def forward_eval(_): + torch.rand(3) # Sampling must not advance the training RNG. + return {"robot": {"image_proprio_prediction": prediction}} + + logged = {} + model = SimpleNamespace( + device=torch.device("cpu"), + current_epoch=1, + global_step=7, + global_rank=0, + trainer=SimpleNamespace(loggers=[]), + forward_eval=forward_eval, + log=lambda k, v, **_: logged.update({k: float(v)}), + ) + return model, prediction, logged + + +def test_pixel_grid_roundtrip_and_reordered_tokens_use_true_coordinates(): + batch, rgb = batch_fixture() + item = batch["robot"] + image = item["world"]["image"] + for frame in range(3): + for camera in range(2): + patches = patchify_rgb(rgb, 2)[frame, camera] + np.testing.assert_allclose( + (unpatchify_rgb(patches, [4, 4], 2) + 1) / 2, + rgb[frame, camera] / 255, + atol=1e-7, + ) + permutation = torch.randperm(image["value"].shape[1]) + shuffled = {k: v[:, permutation] for k, v in image.items()} + result = pixel_frame_records( + shuffled, shuffled["value"], item["pixel_metadata"][0], 0 + ) + assert list(result) == [0, 1] + for camera, value in result.items(): + for i, frame in enumerate(value["frames"]): + np.testing.assert_allclose(frame["target"], rgb[i, camera] / 255, atol=1e-7) + assert frame["metadata"]["frame_index"] == [2, 9, 18][i] + + +def test_metadata_time_and_spatial_grid_mismatch_fail_closed(): + batch, _ = batch_fixture() + item = batch["robot"] + stream = item["world"]["image"] + metadata = copy.deepcopy(item["pixel_metadata"][0]) + metadata["timestamps_seconds"][1] += 0.01 + with pytest.raises(ValueError, match="recorded physical timestamp"): + pixel_frame_records(stream, stream["value"], metadata, 0) + stream["xy"][:, 1] = stream["xy"][:, 0] + with pytest.raises(ValueError, match="declared grid"): + pixel_frame_records(stream, stream["value"], item["pixel_metadata"][0], 0) + + +def test_video_preserves_uneven_physical_timestamp_intervals(tmp_path): + from PIL import Image + + path = tmp_path / "uneven.mp4" + panels = [Image.new("RGB", (64, 64), (i * 40, 0, 0)) for i in range(3)] + write_timestamp_video(path, panels, [100.0, 100.2, 100.95], nominal_fps=10) + with av.open(path) as video: + times = [float(frame.pts * frame.time_base) for frame in video.decode(video=0)] + np.testing.assert_allclose(times, [0.0, 0.2, 0.95], atol=1e-5) + + +def test_visual_only_metrics_contact_sheets_errors_and_immutable_files( + tmp_path, monkeypatch +): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + batch, _ = batch_fixture() + model, _, metrics = fake_model(batch) + evaluator = ImageProprioEval( + objective="visual_only", + pixel_diagnostics={ + "artifact_dir": str(tmp_path), + "write_video": True, + "log_wandb": False, + }, + ) + evaluator.model = model + evaluator.on_validation_start() + rng = torch.get_rng_state().clone() + evaluator.on_validation_step(batch, 0) + assert torch.equal(torch.get_rng_state(), rng) + evaluator.log_epoch_metrics() + assert list(metrics) == ["val/robot/forecast/image_normalized_mse"] + assert metrics["val/robot/forecast/image_normalized_mse"] == pytest.approx(0.04) + reports = list(tmp_path.rglob("diagnostic.json")) + assert len(reports) == 1 + report = json.loads(reports[0].read_text()) + assert report["objective"] == "visual_only" + assert "action" not in report + for camera in report["cameras"].values(): + assert camera["frames"][0]["observation"] == "observed" + assert camera["frames"][0]["query_pixel_mse_rgb01"] is None + assert camera["frames"][1]["query_pixel_mse_rgb01"] == pytest.approx(0.01) + assert (reports[0].parent / camera["contact_sheet"]).is_file() + assert (reports[0].parent / camera["video"]).is_file() + evaluator.on_validation_step(batch, 1) # Bounded diagnostics, all metrics continue. + assert len(list(tmp_path.rglob("diagnostic.json"))) == 1 + evaluator.on_validation_start() + with pytest.raises(FileExistsError): + evaluator.on_validation_step(batch, 0) + + +def test_direct_action_context_never_masquerades_as_pixel_prediction(tmp_path): + batch, _ = batch_fixture(direct=True) + model, _, metrics = fake_model(batch, direct=True) + evaluator = ImageProprioEval( + objective="direct_action", + pixel_diagnostics={ + "artifact_dir": str(tmp_path), + "log_wandb": False, + }, + ) + evaluator.model = model + evaluator.on_validation_start() + evaluator.on_validation_step(batch, 0) + evaluator.log_epoch_metrics() + assert metrics == {"val/robot/direct_action/action_command_mse": 0.25} + report = json.loads(next(tmp_path.rglob("diagnostic.json")).read_text()) + assert report["action"]["squared_error"][0] == pytest.approx([0.25, 0.25]) + assert all(value["video"] is None for value in report["cameras"].values()) + assert not list(tmp_path.rglob("*.mp4")) diff --git a/tests/test_pixel_proprio_data.py b/tests/test_pixel_proprio_data.py new file mode 100644 index 000000000..b973693fb --- /dev/null +++ b/tests/test_pixel_proprio_data.py @@ -0,0 +1,989 @@ +"""Raw-pixel ordering, full-episode sampling, provenance and two-stream behavior.""" + +import copy +import hashlib +import json +import shutil +from types import SimpleNamespace + +import h5py +import numpy as np +import pytest +import torch + +from egomimic.rldb.pixel_proprio import ( + DistributedPixelTrainingBatchSampler, + DistributedPixelValidationBatchSampler, + EvenlySpacedEpisodeBatchSampler, + PreparedPixelProprioDataset, + collate_pixel_proprio, + patchify_rgb, + sample_discrete_window, + sample_full_episode_times, +) +from egomimic.rldb.pixel_proprio_datamodule import PixelProprioDataModule +from egomimic.rldb.world_model import HomogeneousWorldBatchSampler + + +def digest(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def save_json(path, value): + path.write_text(json.dumps(value)) + return path + + +def pixel_manifest(tmp_path, *, include_sharpa=True, actions=True, frames=6): + schemas = { + "libero_gripper": {"dimension": 9, "fields": [f"gripper_{i}" for i in range(9)]} + } + if include_sharpa: + schemas["libero_sharpa"] = { + "dimension": 29, + "fields": [f"hand_{i}" for i in range(29)], + } + episodes = [] + for embodiment, schema in schemas.items(): + for split in ("train", "val"): + name = f"{embodiment}-{split}" + rgb_path, path = tmp_path / f"{name}-rgb.h5", tmp_path / f"{name}.h5" + time = 1e9 + np.arange(frames) * 2.0 + rgb = ( + (np.arange(frames * 2 * 4 * 4 * 3) % 256) + .astype(np.uint8) + .reshape(frames, 2, 4, 4, 3) + ) + with h5py.File(rgb_path, "w") as data: + data["rgb"] = rgb + data["time"] = time + with h5py.File(path, "w") as data: + data["time"] = time + data["pose"] = np.arange( + frames * schema["dimension"], dtype=np.float32 + ).reshape(frames, -1) + data["pose_valid"] = np.ones(frames, bool) + if actions and embodiment == "libero_gripper": + data["action"] = ( + np.arange(frames * 7, dtype=np.float32).reshape(frames, 7) / 100 + ) + data["action_valid"] = np.ones(frames, bool) + ep = dict( + embodiment=embodiment, + episode_id=name, + split=split, + frames=frames, + path=path.name, + sha256=digest(path), + rgb_path=rgb_path.name, + rgb_sha256=digest(rgb_path), + cameras=["agentview", "wrist"], + frame="fixed_reference", + pose_kind="measured_pose", + native_source={"source_demo_sha256": "a" * 64, "episode": name}, + ) + if actions and embodiment == "libero_gripper": + proof = save_json( + tmp_path / f"{name}-alignment.json", + dict( + schema="pixel-proprio-action-alignment-v1", + relation="action_t_applied_to_state_t", + source_demo_sha256="a" * 64, + source_episode=name, + producer_sha256="b" * 64, + raw_export_metadata_sha256="c" * 64, + raw_payload_sha256="d" * 64, + cache_frame_index=list(range(frames)), + source_state_index=list(range(frames)), + source_action_index=list(range(frames)), + basis={ + "cache_index_source": "synthetic fixture mapping", + "action_index_source": "synthetic fixture commands", + "state_action_relation_source": "synthetic fixture relation", + }, + ), + ) + ep["action_provenance"] = {"path": proof.name, "sha256": digest(proof)} + episodes.append(ep) + norm = save_json( + tmp_path / "normalization.json", + dict( + schema="pixel-proprio-normalization-v1", + scope="train_only", + train_episodes=sorted( + [ + [ep["embodiment"], ep["episode_id"], ep["sha256"]] + for ep in episodes + if ep["split"] == "train" + ] + ), + statistics={ + f"pose/{emb}": { + "mean": [1.0] * schema["dimension"], + "std": [2.0] * schema["dimension"], + "count": frames, + } + for emb, schema in schemas.items() + }, + ), + ) + manifest = dict( + schema="pixel-proprio-prepared-v1", + pose_schemas=schemas, + episodes=episodes, + pixel_identity={ + "image_size": [4, 4], + "patch_size": 2, + "channels": 3, + "patch_order": "row_major_hwc", + "normalization": "uint8_to_minus_one_one", + }, + normalization={ + "scope": "train_only", + "arrays_normalized": False, + "path": norm.name, + "sha256": digest(norm), + }, + ) + if actions: + manifest["action_schemas"] = { + "libero_gripper": { + "dimension": 7, + "fields": [f"command_{i}" for i in range(7)], + "semantics": "action_command", + "value_space": "native_commands", + } + } + return save_json(tmp_path / "manifest.json", manifest) + + +def test_patch_order_retains_every_raw_hwc_value_and_each_camera(): + rgb = np.arange(1 * 2 * 4 * 4 * 3, dtype=np.uint8).reshape(1, 2, 4, 4, 3) + patches = patchify_rgb(rgb, 2) + assert patches.shape == (1, 2, 4, 12) and patches.dtype == np.float32 + for camera in range(2): + for row in range(2): + for col in range(2): + expected = rgb[ + 0, camera, row * 2 : (row + 1) * 2, col * 2 : (col + 1) * 2 + ].reshape(-1) + np.testing.assert_allclose( + patches[0, camera, row * 2 + col], expected / 127.5 - 1, atol=1e-7 + ) + rgb = np.zeros((1, 2, 160, 160, 3), np.uint8) + rgb[0, 1] = 255 + patches = patchify_rgb(rgb, 16) + assert patches.shape == (1, 2, 100, 768) + assert np.all(patches[0, 0] == -1) and np.all(patches[0, 1] == 1) + + +def test_full_episode_uniform_sampling_clips_maximum_without_time_window(): + all_frames = sample_full_episode_times( + 128, np.random.default_rng(1), 128, 128, anchor=37 + ) + np.testing.assert_array_equal(all_frames, np.arange(128)) + clipped = sample_full_episode_times(5, np.random.default_rng(2), 5, 128, anchor=2) + np.testing.assert_array_equal(clipped, np.arange(5)) + seen = set() + rng = np.random.default_rng(21) + for _ in range(200): + selected = sample_full_episode_times(20, rng, 2, 2, anchor=0) + assert len(selected) == len(set(selected)) == 2 and selected[0] == 0 + seen.add(selected[1]) + assert seen == set(range(1, 20)) + + +def test_pixels_and_both_measured_pose_schemas_need_no_encoder_tracks_or_commands( + tmp_path, monkeypatch +): + path = pixel_manifest(tmp_path, actions=False) + original = h5py.Group.__getitem__ + + def guard(group, key): + assert key not in ("image", "dino", "track", "action", "action_valid") + return original(group, key) + + monkeypatch.setattr(h5py.Group, "__getitem__", guard) + dataset = PreparedPixelProprioDataset( + path, split="val", min_frames=6, max_frames=128 + ) + assert set(dataset.normalization) == {"pose/libero_gripper", "pose/libero_sharpa"} + first, second = dataset[0], dataset[6] + assert first["world"]["image"]["value"].shape == (48, 12) + assert first["world"]["proprio"]["value"].shape == (6, 9) + assert second["world"]["proprio"]["value"].shape == (6, 29) + torch.testing.assert_close( + first["world"]["proprio"]["value"][0], (torch.arange(9) - 1) / 2 + ) + assert torch.unique(first["world"]["image"]["time"]).numel() == 6 + assert ( + first["world"]["image"]["time"].max() - first["world"]["image"]["time"].min() + == 10 + ) + + +@pytest.mark.parametrize("image_only", [False, True]) +def test_forecast_and_completion_metadata_padding_and_no_future_observations( + tmp_path, image_only +): + path = pixel_manifest(tmp_path, actions=False) + forecast = PreparedPixelProprioDataset( + path, + split="val", + forecasting_probability=1.0, + min_frames=6, + max_frames=6, + image_only=image_only, + )[0] + completion = PreparedPixelProprioDataset( + path, + split="val", + forecasting_probability=0.0, + min_frames=2, + max_frames=2, + image_only=image_only, + )[0] + image = forecast["world"]["image"] + cutoff = image["time"][image["observed"]].max() + assert (image["time"][~image["observed"]] > cutoff).all() + if image_only: + assert set(forecast["world"]) == set(completion["world"]) == {"image"} + else: + pose = forecast["world"]["proprio"] + assert (pose["time"][pose["observed"]] <= cutoff).all() + batch = collate_pixel_proprio([forecast, completion])["libero_gripper"] + assert batch["task_kind"] == ["forecast", "completion"] + assert batch["pixel_metadata"] == [ + forecast["pixel_metadata"], + completion["pixel_metadata"], + ] + metadata = forecast["pixel_metadata"] + assert metadata["timestamps_seconds"] == (1e9 + np.arange(6) * 2).tolist() + assert metadata["frame_indices"] == list(range(6)) + assert metadata["camera_names"] == ["agentview", "wrist"] + for stream in batch["world"].values(): + assert ( + not stream["valid"][1, -1] + and not stream["observed"][1, -1] + and not stream["target_valid"][1, -1] + ) + + +def test_direct_bc_is_one_observed_anchor_and_real_unscaled_command(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False) + sample = PreparedPixelProprioDataset(path, task="direct_action", split="val")[3] + assert sample["task_kind"] == "direct_action" + image, pose, action = ( + sample["world"][name] for name in ("image", "proprio", "action") + ) + assert image["value"].shape == (8, 12) and pose["value"].shape == (1, 9) + assert ( + image["observed"].all() + and pose["observed"].all() + and not action["observed"].any() + ) + torch.testing.assert_close( + action["value"], torch.arange(21, 28)[None].float() / 100 + ) + assert all( + torch.equal(s["time"], torch.zeros_like(s["time"])) + for s in (image, pose, action) + ) + + +def test_direct_mixed_corpus_fails_closed_when_sharpa_commands_are_unverified(tmp_path): + path = pixel_manifest(tmp_path) + with pytest.raises(ValueError, match="verified real command labels"): + PreparedPixelProprioDataset(path, task="direct_action") + + +@pytest.mark.parametrize( + "damage", + ["rgb_checksum", "rgb_time", "action_proof", "image_statistics", "time_window"], +) +def test_pixel_provenance_normalization_and_sampling_fail_closed(tmp_path, damage): + path = pixel_manifest(tmp_path, include_sharpa=False) + manifest = json.loads(path.read_text()) + episode = manifest["episodes"][1] + kwargs = {} + if damage in ("rgb_checksum", "rgb_time"): + rgb_path = path.parent / episode["rgb_path"] + with h5py.File(rgb_path, "a") as data: + if damage == "rgb_checksum": + data["rgb"][0, 0, 0, 0, 0] += 1 + else: + data["time"][0] += 0.25 + if damage == "rgb_time": + episode["rgb_sha256"] = digest(rgb_path) + elif damage == "action_proof": + proof_path = path.parent / episode["action_provenance"]["path"] + proof = json.loads(proof_path.read_text()) + proof["source_action_index"][0] += 1 + save_json(proof_path, proof) + episode["action_provenance"]["sha256"] = digest(proof_path) + kwargs["task"] = "direct_action" + elif damage == "image_statistics": + norm_path = path.parent / manifest["normalization"]["path"] + norm = json.loads(norm_path.read_text()) + norm["statistics"]["image"] = {"mean": [0], "std": [1], "count": 1} + save_json(norm_path, norm) + manifest["normalization"]["sha256"] = digest(norm_path) + else: + kwargs.update(max_past_seconds=3, max_future_seconds=3) + save_json(path, manifest) + with pytest.raises(ValueError): + PreparedPixelProprioDataset(path, split="val", **kwargs)[0] + + +@pytest.mark.parametrize("image_only", [False, True]) +def test_full_episode_sampler_and_heldout_coverage_survive_exact_resume( + tmp_path, image_only +): + path = pixel_manifest(tmp_path, actions=False) + kwargs = dict( + manifest_path=path, + batch_size=1, + checkpointable_sampling=True, + validation_anchors_per_episode=2, + min_frames=2, + max_frames=128, + image_only=image_only, + ) + data = PixelProprioDataModule(**kwargs) + context = data.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert context.state["sampling"]["sampling_strategy"] == "full_episode_uniform" + assert context.state["kind"] == ( + "pixel-image-prepared-v1" if image_only else "pixel-proprio-prepared-v1" + ) + assert list(data.val_dataloader().batch_sampler) == [[0], [5], [6], [11]] + first = next(iter(data.train_dataloader())) + data.acknowledge_batch(next(iter(first.values()))["sampling_cursor"]) + state = data.state_dict() + expected = next(iter(data.train_dataloader())) + restored = PixelProprioDataModule(**kwargs) + restored.load_state_dict(state) + actual = next(iter(restored.train_dataloader())) + assert set(actual) == set(expected) + key = next(iter(actual)) + assert actual[key]["task_kind"] == expected[key]["task_kind"] + for name in actual[key]["world"]: + for field in actual[key]["world"][name]: + torch.testing.assert_close( + actual[key]["world"][name][field], + expected[key]["world"][name][field], + rtol=0, + atol=0, + ) + changed = PixelProprioDataModule(**{**kwargs, "max_frames": 127}) + with pytest.raises(ValueError, match="settings"): + changed.load_state_dict(state) + changed = PixelProprioDataModule(**{**kwargs, "image_only": not image_only}) + with pytest.raises(ValueError, match="settings"): + changed.load_state_dict(state) + with pytest.raises(ValueError, match="context mismatch"): + changed.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + + +@pytest.mark.parametrize("direct", [False, True]) +def test_pixel_batches_train_and_inference_ignore_withheld_values(tmp_path, direct): + from egomimic.models.image_proprio import ImageProprioModel + + path = pixel_manifest(tmp_path, include_sharpa=False) + task = "direct_action" if direct else "visual_diffusion" + dataset = PreparedPixelProprioDataset( + path, + task=task, + split="val", + min_frames=2, + max_frames=2, + forecasting_probability=1.0, + ) + batch = collate_pixel_proprio([dataset[0]])["libero_gripper"]["world"] + model = ImageProprioModel( + pose_dims={"libero_gripper": 9}, + action_dims={"libero_gripper": 7} if direct else None, + objective="direct_action" if direct else "visual_proprio", + image_dim=12, + width=24, + depth=1, + heads=3, + max_cameras=2, + max_tokens=100, + ) + losses = model.training_losses(batch, "libero_gripper") + assert torch.isfinite(losses["total"]) and losses["total"] > 0 + losses["total"].backward() + assert ( + model.image_in.weight.grad is not None and model.image_in.weight.grad.norm() > 0 + ) + changed = copy.deepcopy(batch) + for stream in changed.values(): + stream["value"][~stream["observed"]] = float("nan") + torch.manual_seed(7) + expected = model.sample(batch, "libero_gripper", steps=2) + torch.manual_seed(7) + actual = model.sample(changed, "libero_gripper", steps=2) + for name in expected: + torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0) + + +def test_discrete_window_contains_anchor_and_all_four_state_command_rows(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False, frames=75) + rng = np.random.default_rng(3) + for anchor in (0, 30, 74): + selected = sample_discrete_window(75, rng, 16, 4, anchor) + assert selected.shape == (16, 4) and anchor in selected + np.testing.assert_array_equal(np.diff(selected.reshape(-1)), np.ones(63)) + dataset = PreparedPixelProprioDataset( + path, + split="val", + task="visual_state_action", + mode="discrete", + sampling_strategy="discrete_window", + discrete_frames=16, + stride=4, + forecasting_probability=1, + ) + sample = dataset[74] + meta, world = sample["pixel_metadata"], sample["world"] + rows = np.asarray(meta["proprio_frame_indices"]) + assert rows.shape == (16, 4) and rows[-1, -1] == 74 + np.testing.assert_array_equal(meta["frame_indices"], rows[:, 0]) + assert world["image"]["value"].shape == (16 * 8, 12) + assert world["proprio"]["value"].shape == (16, 36) + assert world["action"]["value"].shape == (16, 28) + assert not world["action"]["observed"].any() + expected_pose = (np.arange(75 * 9).reshape(75, 9)[rows].reshape(16, 36) - 1) / 2 + np.testing.assert_allclose(world["proprio"]["value"].numpy(), expected_pose) + expected_action = ( + np.arange(75 * 7).reshape(75, 7).astype(np.float32)[rows].reshape(16, 28) / 100 + ) + np.testing.assert_allclose(world["action"]["value"].numpy(), expected_action) + known_images = world["image"]["observed"].reshape(16, -1).all(1).numpy() + known_pose = world["proprio"]["observed"].numpy() + cutoff = rows[known_images][-1, -1] + assert (rows[known_pose] <= cutoff).all() + assert (rows[~known_images] > cutoff).all() + assert not known_pose[~known_images].any() + first_unknown_image = np.asarray(meta["frame_indices"])[~known_images][0] + assert rows[known_pose].max() < first_unknown_image + assert world["action"]["valid"].all() + np.testing.assert_array_equal( + world["action"]["target_valid"].numpy(), ~known_images + ) + assert sample["task_kind"] == "forecast" + assert collate_pixel_proprio([sample])["libero_gripper"]["task_kind"] == [ + "forecast" + ] + + +@pytest.mark.parametrize("image_only", [False, True]) +def test_shared_validation_grid_masks_and_requirements_across_discrete_and_continuous( + tmp_path, + image_only, +): + from egomimic.pl_utils.data_context import EvaluationDataRequirements + + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=75) + shared = dict( + manifest_path=path, + validation_anchors_per_episode=2, + validation_sampling_strategy="discrete_window", + validation_discrete_frames=16, + validation_stride=4, + forecasting_probability=1, + checkpointable_sampling=True, + image_only=image_only, + ) + discrete = PixelProprioDataModule( + **shared, + mode="discrete", + sampling_strategy="discrete_window", + discrete_frames=16, + stride=4, + ) + continuous = PixelProprioDataModule(**shared) + for data in (discrete, continuous): + data.configure_evaluation( + EvaluationDataRequirements( + required_keys=("world", "embodiment", "task_kind", "pixel_metadata") + ) + ) + with pytest.raises(ValueError, match="Unsupported"): + data.configure_evaluation( + EvaluationDataRequirements(complete_episodes=True) + ) + with pytest.raises(ValueError, match="Unsupported"): + data.configure_evaluation( + EvaluationDataRequirements(required_keys=("missing",)) + ) + for left, right in zip(discrete.val_dataloader(), continuous.val_dataloader()): + left, right = left["libero_gripper"], right["libero_gripper"] + assert ( + left["pixel_metadata"][0]["frame_indices"] + == right["pixel_metadata"][0]["frame_indices"] + ) + assert ( + left["pixel_metadata"][0]["timestamps_seconds"] + == right["pixel_metadata"][0]["timestamps_seconds"] + ) + for key in left["world"]["image"]: + torch.testing.assert_close( + left["world"]["image"][key], + right["world"]["image"][key], + rtol=0, + atol=0, + ) + if image_only: + assert set(left["world"]) == set(right["world"]) == {"image"} + else: + torch.testing.assert_close( + left["world"]["proprio"]["observed"], + right["world"]["proprio"]["observed"], + ) + assert left["world"]["proprio"]["value"].shape[-1] == 36 + assert right["world"]["proprio"]["value"].shape[-1] == 9 + context = continuous.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert context.state["validation_overrides"]["stride"] == 4 + changed = PixelProprioDataModule(**{**shared, "validation_stride": 3}) + with pytest.raises(ValueError, match="context mismatch"): + changed.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + with pytest.raises(ValueError, match="sampling restore"): + changed.load_state_dict(continuous.state_dict()) + + +@pytest.mark.parametrize("discrete", [False, True]) +def test_forecast_prefix_command_labels_do_not_contribute_to_loss_or_metrics( + tmp_path, discrete +): + from types import SimpleNamespace + + from egomimic.eval.image_proprio_eval import ImageProprioEval + from egomimic.models.image_proprio import ImageProprioModel + + path = pixel_manifest(tmp_path, include_sharpa=False, frames=75) + settings = dict( + task="visual_state_action", + split="val", + forecasting_probability=1, + min_frames=4, + max_frames=4, + ) + if discrete: + settings.update( + mode="discrete", + sampling_strategy="discrete_window", + discrete_frames=16, + stride=4, + ) + sample = PreparedPixelProprioDataset(path, **settings)[0] + batch = collate_pixel_proprio([sample]) + world = batch["libero_gripper"]["world"] + targets = world["action"]["target_valid"] + assert targets.any() and (~targets).any() + assert world["action"]["valid"].all() and not world["action"]["observed"].any() + changed = copy.deepcopy(batch) + changed["libero_gripper"]["world"]["action"]["value"][~targets] = 1e6 + chunk = 4 if discrete else 1 + model = ImageProprioModel( + pose_dims={"libero_gripper": 9 * chunk}, + action_dims={"libero_gripper": 7 * chunk}, + objective="visual_state_action", + image_dim=12, + width=24, + depth=1, + heads=3, + max_cameras=2, + max_tokens=400, + ) + torch.manual_seed(19) + expected = model.training_losses(world, "libero_gripper") + torch.manual_seed(19) + actual = model.training_losses(changed["libero_gripper"]["world"], "libero_gripper") + torch.testing.assert_close(actual["action"], expected["action"], rtol=0, atol=0) + torch.testing.assert_close(actual["total"], expected["total"], rtol=0, atol=0) + + def evaluate(value): + prediction = { + name: torch.zeros_like(stream["value"]) for name, stream in world.items() + } + evaluator = ImageProprioEval(objective="visual_state_action") + evaluator.model = SimpleNamespace( + device=torch.device("cpu"), + global_step=0, + forward_eval=lambda _: { + "libero_gripper": {"image_proprio_prediction": prediction} + }, + ) + evaluator.on_validation_start() + evaluator.on_validation_step(value, 0) + return evaluator.totals["val/libero_gripper/forecast/action_command_mse"] + + numerator, denominator = evaluate(batch) + changed_numerator, changed_denominator = evaluate(changed) + assert denominator == targets.sum() == changed_denominator + torch.testing.assert_close(numerator, changed_numerator, rtol=0, atol=0) + completion = PreparedPixelProprioDataset( + path, **{**settings, "forecasting_probability": 0} + )[0] + assert completion["task_kind"] == "completion" + assert completion["world"]["action"]["target_valid"].all() + + +def test_visual_only_discrete_observed_chunks_stop_before_first_unknown_image(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=75) + sample = PreparedPixelProprioDataset( + path, + split="val", + mode="discrete", + sampling_strategy="discrete_window", + discrete_frames=16, + stride=4, + forecasting_probability=1, + )[0] + world, meta = sample["world"], sample["pixel_metadata"] + assert set(world) == {"image", "proprio"} + known_image = world["image"]["observed"].reshape(16, -1).all(1).numpy() + known_pose = world["proprio"]["observed"].numpy() + first_unknown_image = np.asarray(meta["frame_indices"])[~known_image][0] + rows = np.asarray(meta["proprio_frame_indices"]) + assert (rows[known_pose] < first_unknown_image).all() + assert not known_pose[~known_image].any() + + +def test_image_only_needs_no_pose_command_or_normalization_artifacts( + tmp_path, monkeypatch +): + path = pixel_manifest(tmp_path) + manifest = json.loads(path.read_text()) + manifest_digest = digest(path) + for ep in manifest["episodes"]: + (path.parent / ep["path"]).unlink() + if "action_provenance" in ep: + (path.parent / ep["action_provenance"]["path"]).unlink() + (path.parent / manifest["normalization"]["path"]).unlink() + original = h5py.Group.__getitem__ + + def guard(group, key): + assert key in ("rgb", "time"), f"Image-only path opened {key}" + return original(group, key) + + monkeypatch.setattr(h5py.Group, "__getitem__", guard) + module = PixelProprioDataModule( + path, + image_only=True, + validation_anchors_per_episode=2, + checkpointable_sampling=True, + ) + context = module.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert context.sample_schema == {"world": "image"} + assert context.state["manifest"] == manifest + assert context.state["manifest_sha256"] == manifest_digest + assert context.state["sampling"]["image_only"] is True + assert module.datasets["train"].normalization == {} + assert list(module.val_dataloader().batch_sampler) == [[0], [5], [6], [11]] + for batch in [next(iter(module.train_dataloader())), *module.val_dataloader()]: + item = next(iter(batch.values())) + assert set(item["world"]) == {"image"} + assert torch.isfinite(item["world"]["image"]["value"]).all() + assert not any( + "proprio" in key or "pose" in key for key in item["pixel_metadata"][0] + ) + module.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + + +@pytest.mark.parametrize("discrete", [False, True]) +@pytest.mark.parametrize("forecasting_probability", [0.0, 1.0]) +def test_image_only_preserves_rgb_sampling_masks_and_relative_time( + tmp_path, + discrete, + forecasting_probability, +): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=75) + settings = dict(forecasting_probability=forecasting_probability) + if discrete: + settings.update( + mode="discrete", + sampling_strategy="discrete_window", + discrete_frames=16, + stride=4, + ) + else: + # Two-frame completion exercises all-known/all-unknown image-mask repairs. + settings.update(min_frames=2, max_frames=2) + for split in ("train", "val"): + legacy = PreparedPixelProprioDataset(path, split=split, **settings) + image_only = PreparedPixelProprioDataset( + path, split=split, image_only=True, **settings + ) + for index in range(20): + expected, actual = legacy[index], image_only[index] + assert set(actual["world"]) == {"image"} + assert actual["task_kind"] == expected["task_kind"] + for field in expected["world"]["image"]: + torch.testing.assert_close( + actual["world"]["image"][field], + expected["world"]["image"][field], + rtol=0, + atol=0, + ) + for field in ( + "frame_indices", + "timestamps_seconds", + "time_origin_seconds", + "camera_names", + ): + assert ( + actual["pixel_metadata"][field] == expected["pixel_metadata"][field] + ) + + +@pytest.mark.parametrize("damage", ["checksum", "timestamps", "duplicate"]) +def test_image_only_still_validates_rgb_identity_and_time(tmp_path, damage): + path = pixel_manifest(tmp_path, include_sharpa=False) + manifest = json.loads(path.read_text()) + ep = manifest["episodes"][1] + if damage == "duplicate": + ep["rgb_path"] = manifest["episodes"][0]["rgb_path"] + else: + rgb_path = path.parent / ep["rgb_path"] + with h5py.File(rgb_path, "a") as data: + if damage == "checksum": + data["rgb"][0, 0, 0, 0, 0] += 1 + else: + data["time"][1] = data["time"][0] + if damage == "timestamps": + ep["rgb_sha256"] = digest(rgb_path) + save_json(path, manifest) + with pytest.raises(ValueError): + PreparedPixelProprioDataset(path, split="val", image_only=True)[0] + + +@pytest.mark.parametrize("task", ["direct_action", "visual_state_action"]) +def test_image_only_rejects_label_requiring_tasks(tmp_path, task): + path = pixel_manifest(tmp_path, include_sharpa=False) + with pytest.raises(ValueError, match="image_only"): + PreparedPixelProprioDataset(path, task=task, image_only=True) + with pytest.raises(ValueError, match="image_only"): + PixelProprioDataModule(path, task=task, image_only=True) + + +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_four_rank_training_is_partition_of_one_global_draw_and_resumes( + tmp_path, batch_size +): + path = pixel_manifest(tmp_path, actions=False, frames=20) + dataset = PreparedPixelProprioDataset(path, image_only=True) + global_sampler = HomogeneousWorldBatchSampler( + dataset, batch_size=4 * batch_size, seed=42, checkpointable=True + ) + global_draws = list(global_sampler) + ranked = [] + for rank in range(4): + sampler = DistributedPixelTrainingBatchSampler( + dataset, batch_size, 42, world_size=4, rank=rank + ) + ranked.append(list(sampler)) + assert len(sampler) == len(global_sampler) + sampler.consumed_batches = 1 + assert list(sampler) == ranked[-1][1:] + for step, expected in enumerate(global_draws): + actual = [draw for rank in range(4) for draw in ranked[rank][step]] + assert actual == expected + # Distinct deterministic RNG slots; replacement may draw the same anchor. + assert [draw[2] for draw in actual] == list(range(4 * batch_size)) + assert {draw[1] for draw in actual} == {step} + assert ( + len( + { + dataset.episodes[ + int(np.searchsorted(dataset.ends, draw[0], side="right")) + ]["embodiment"] + for draw in actual + } + ) + == 1 + ) + + +def _fifty_validation_episode_manifest(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False) + manifest = json.loads(path.read_text()) + template = next(ep for ep in manifest["episodes"] if ep["split"] == "val") + original_rgb = tmp_path / template["rgb_path"] + for index in range(1, 50): + episode = copy.deepcopy(template) + episode["episode_id"] = f"validation-episode-{index}" + episode["rgb_path"] = f"validation-episode-{index}-rgb.h5" + shutil.copyfile(original_rgb, tmp_path / episode["rgb_path"]) + manifest["episodes"].append(episode) + return save_json(path, manifest) + + +def _assert_identical_pixel_item(expected, actual): + assert actual["task_kind"] == expected["task_kind"] + assert actual["pixel_metadata"] == expected["pixel_metadata"] + assert set(actual["world"]) == set(expected["world"]) == {"image"} + for field in actual["world"]["image"]: + torch.testing.assert_close( + actual["world"]["image"][field], + expected["world"]["image"][field], + rtol=0, + atol=0, + ) + + +def test_four_rank_validation_exactly_covers_100_global_draws_and_preserves_masks( + tmp_path, +): + path = _fifty_validation_episode_manifest(tmp_path) + kwargs = dict( + manifest_path=path, + image_only=True, + checkpointable_sampling=True, + validation_anchors_per_episode=2, + distributed_world_size=4, + ) + audit = PixelProprioDataModule(**kwargs) + context = audit.prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + assert context.state["distributed_sampling"]["global_batch_size"] == 4 + assert len(audit.val_dataloader()) == 100 # Unattached audit/reload is global. + expected = [batch["libero_gripper"] for batch in audit.val_dataloader()] + visited, episodes = [], [] + for rank in range(4): + module = PixelProprioDataModule(**kwargs) + module.trainer = SimpleNamespace(world_size=4, global_rank=rank) + loader = module.val_dataloader() + assert len(loader) == 25 + assert ( + module.prepare_context( + mode="train", normalization="prepared", normalizer=None + ).state + == context.state + ) + kinds = set() + for batch in loader: + item = batch["libero_gripper"] + (global_index,) = item["validation_batch_index"] + visited.append(global_index) + episodes.append(item["pixel_metadata"][0]["episode_id"]) + kinds.update(item["task_kind"]) + _assert_identical_pixel_item(expected[global_index], item) + if rank == 0: + assert kinds == {"forecast", "completion"} + assert visited == list(range(100)) # No drops, duplicates, or reordered seeds. + assert len(set(episodes)) == 50 + assert all(episodes.count(episode) == 2 for episode in set(episodes)) + + +def test_four_rank_checkpoint_cursor_restores_each_local_draw_and_binds_topology( + tmp_path, +): + path = pixel_manifest(tmp_path, actions=False, frames=20) + kwargs = dict( + manifest_path=path, + image_only=True, + checkpointable_sampling=True, + validation_anchors_per_episode=2, + distributed_world_size=4, + ) + states = [] + for rank in range(4): + data = PixelProprioDataModule(**kwargs) + data.trainer = SimpleNamespace(world_size=4, global_rank=rank) + consumed = next(iter(data.train_dataloader())) + item = next(iter(consumed.values())) + assert item["sampling_cursor"] == 1 + data.acknowledge_batch(item["sampling_cursor"]) + state = data.state_dict() + states.append(state) + expected = next(iter(data.train_dataloader())) + restored = PixelProprioDataModule(**kwargs) + restored.load_state_dict(state) + restored.trainer = SimpleNamespace(world_size=4, global_rank=rank) + actual = next(iter(restored.train_dataloader())) + assert actual.keys() == expected.keys() + key = next(iter(actual)) + _assert_identical_pixel_item(expected[key], actual[key]) + assert actual[key]["sampling_cursor"] == 2 + with pytest.raises(ValueError, match="Nonconsecutive"): + restored.acknowledge_batch(3) + assert all(state == states[0] for state in states) # One global cursor. + single = PixelProprioDataModule(**{**kwargs, "distributed_world_size": 1}) + with pytest.raises(ValueError, match="distributed topology"): + single.load_state_dict(states[0]) + with pytest.raises(ValueError, match="distributed topology"): + PixelProprioDataModule(**kwargs).load_state_dict(single.state_dict()) + context = PixelProprioDataModule(**kwargs).prepare_context( + mode="train", normalization="prepared", normalizer=None + ) + with pytest.raises(ValueError, match="context mismatch"): + single.prepare_context( + mode="train", + normalization="prepared", + normalizer=None, + restored_state=context.state, + ) + + +def test_distributed_pixel_loading_rejects_ambiguous_or_lossy_topologies(tmp_path): + path = pixel_manifest(tmp_path, actions=False) + with pytest.raises(ValueError, match="checkpointable_sampling"): + PixelProprioDataModule(path, distributed_world_size=4) + with pytest.raises(ValueError, match="positive integer"): + PixelProprioDataModule(path, distributed_world_size=True) + data = PixelProprioDataModule( + path, image_only=True, checkpointable_sampling=True, distributed_world_size=4 + ) + data.trainer = SimpleNamespace(world_size=1, global_rank=0) + with pytest.raises(ValueError, match="world_size"): + data.train_dataloader() + dataset = PreparedPixelProprioDataset(path, split="val", image_only=True) + base = EvenlySpacedEpisodeBatchSampler(dataset, 2) + with pytest.raises(ValueError, match="divisible"): + DistributedPixelValidationBatchSampler(base, world_size=3, rank=0) + with pytest.raises(ValueError, match="topology"): + DistributedPixelTrainingBatchSampler(dataset, 1, 42, world_size=4, rank=4) + + +def test_validation_global_batch_index_is_shared_and_preserved_when_collating(tmp_path): + path = pixel_manifest(tmp_path, include_sharpa=False, actions=False, frames=12) + data = PreparedPixelProprioDataset(path, split="val", image_only=True) + base = EvenlySpacedEpisodeBatchSampler(data, 8, batch_size=2) + shards = [ + DistributedPixelValidationBatchSampler(base, world_size=2, rank=rank) + for rank in range(2) + ] + for rank, shard in enumerate(shards): + for offset, indices in enumerate(shard): + item = collate_pixel_proprio([data[index] for index in indices])[ + "libero_gripper" + ] + assert item["validation_batch_index"] == [rank * 2 + offset] * 2 + sample = data[0] + sample["validation_batch_index"] = 0 + with pytest.raises(ValueError, match="global batch index"): + collate_pixel_proprio([sample, data[1]]) diff --git a/tests/test_pixel_proprio_preparation.py b/tests/test_pixel_proprio_preparation.py new file mode 100644 index 000000000..d718248bf --- /dev/null +++ b/tests/test_pixel_proprio_preparation.py @@ -0,0 +1,498 @@ +"""End-to-end pixel preparation without feature or track dependencies.""" + +import copy +import json +from pathlib import Path + +import h5py +import numpy as np +import pytest + +from egomimic.rldb.pixel_proprio import PreparedPixelProprioDataset +from scripts.world_model.export_libero_action_sidecars import sha256 +from scripts.world_model.prepare_pixel_proprio import ( + derive_split, + prepare_pixels, + resize_rgb_area, +) + + +def save_json(path, value): + path.write_text(json.dumps(value)) + return path + + +def native_episode(root, episode, fields, *, pose_offset=0, bad_state_time=False): + name = episode["episode_id"] + native = root / name + native.mkdir() + n = episode["frames"] + time = 0.25 + np.arange(n) * 0.05 + pose = pose_offset + np.arange(n)[:, None] + np.arange(len(fields))[None] / 10 + state = np.zeros((n, 20), np.float64) + state[:, 0] = time + (1 if bad_state_time else 0) + # Poison unused arrays: preparation must not consume low-res RGB or tracks. + np.savez( + native / "aligned.npz", + pose=pose, + time=time, + source_row=np.arange(n), + state=state, + track=np.array([np.nan]), + rgb=np.array([np.nan]), + ) + native_manifest = save_json( + native / "manifest.json", + dict( + schema="libero-same-recorded-state-v1", + **episode["native_source"], + pose_fields=fields, + cameras=["agentview", "wrist"], + artifacts={"aligned.npz": sha256(native / "aligned.npz")}, + ), + ) + rgb_path = root / f"{name}.rgb.h5" + pixels = ( + (np.arange(n * 2 * 8 * 8 * 3) % 256).astype(np.uint8).reshape(n, 2, 8, 8, 3) + ) + with h5py.File(rgb_path, "w") as data: + data.attrs["schema"] = "libero-native-rgb-v1" + data.attrs["native_render"] = True + data["rgb"] = pixels + data["time"] = time + receipt = save_json( + root / f"{name}.rgb.json", + dict( + schema="libero-native-rgb-v1", + status="PASS", + native_render=True, + native_geometry_replay_exact=True, + frames=n, + height=8, + width=8, + cameras=["agentview", "wrist"], + source_manifest_sha256=sha256(native_manifest), + source_aligned_sha256=sha256(native / "aligned.npz"), + rgb_cache_sha256=sha256(rgb_path), + ), + ) + return ( + dict( + embodiment=episode["embodiment"], + episode_id=name, + native_cache=str(native), + rgb_cache=str(rgb_path), + rgb_receipt=str(receipt), + ), + pose, + pixels, + ) + + +def native_corpus(tmp_path, *, bad_state_time=False): + episodes, locations, poses, pixels = [], [], {}, {} + schemas = { + emb: {"dimension": dim, "fields": [f"coordinate_{i}" for i in range(dim)]} + for emb, dim in (("libero_gripper", 9), ("libero_sharpa", 29)) + } + for emb, schema in schemas.items(): + for split in ("train", "val"): + name = f"{emb}-{split}" + episode = dict( + embodiment=emb, + episode_id=name, + split=split, + frames=4, + path="MISSING_ENCODER_FEATURES.h5", + sha256="e" * 64, + native_source={ + "source_demo": f"/fixture/{emb}.hdf5", + "source_demo_sha256": "a" * 64, + "episode": split, + }, + ) + location, pose, rgb = native_episode( + tmp_path, + episode, + schema["fields"], + pose_offset=0 if split == "train" else 100, + bad_state_time=bad_state_time, + ) + episodes.append(episode) + locations.append(location) + poses[name], pixels[name] = pose, rgb + source = save_json( + tmp_path / "source.json", + dict(schema="fixture-source-metadata", episodes=episodes, pose_schemas=schemas), + ) + spec = save_json(tmp_path / "locations.json", {"episodes": locations}) + return source, spec, poses, pixels + + +def test_area_resize_is_exact_rgb_block_mean_without_crop_or_channel_reordering(): + image = np.arange(2 * 8 * 8 * 3, dtype=np.uint8).reshape(2, 8, 8, 3) + actual = resize_rgb_area(image, (4, 4)) + expected = np.empty((2, 4, 4, 3), np.uint8) + for camera in range(2): + for row in range(4): + for col in range(4): + expected[camera, row, col] = np.rint( + image[camera, row * 2 : row * 2 + 2, col * 2 : col * 2 + 2].mean( + (0, 1) + ) + ) + np.testing.assert_array_equal(actual, expected) + # Document the otherwise ambiguous uint8 half-way rounding. + half = np.array([0, 1, 1, 0], np.uint8).reshape(2, 2, 1).repeat(3, axis=-1) + assert np.all(resize_rgb_area(half, (1, 1)) == 0) + with pytest.raises(ValueError, match="integer downsampling"): + resize_rgb_area(image, (3, 3)) + + +def test_prepare_mixed_corpus_preserves_split_time_cameras_and_train_only_pose_statistics( + tmp_path, monkeypatch +): + source, spec, poses, pixels = native_corpus(tmp_path) + original = np.lib.npyio.NpzFile.__getitem__ + + def guard(data, key): + assert key in ("pose", "time", "source_row", "state") + return original(data, key) + + monkeypatch.setattr(np.lib.npyio.NpzFile, "__getitem__", guard) + output = tmp_path / "prepared" + result = prepare_pixels( + source, spec, output, image_size=[4, 4], patch_size=2, resize_batch_size=2 + ) + assert result["episodes"] == 4 and result["frames"] == 16 + manifest = json.loads((output / "manifest.json").read_text()) + stats = json.loads((output / manifest["normalization"]["path"]).read_text()) + assert "dino_identity" not in manifest and "action_schemas" not in manifest + assert set(stats["statistics"]) == {"pose/libero_gripper", "pose/libero_sharpa"} + assert manifest["split_provenance"]["counts"] == { + "libero_gripper": {"train": 1, "val": 1}, + "libero_sharpa": {"train": 1, "val": 1}, + } + assert [ + (e["embodiment"], e["episode_id"], e["split"]) for e in manifest["episodes"] + ] == [ + (e["embodiment"], e["episode_id"], e["split"]) + for e in json.loads(source.read_text())["episodes"] + ] + for ep in manifest["episodes"]: + with ( + h5py.File(output / ep["rgb_path"], "r") as rgb, + h5py.File(output / ep["path"], "r") as data, + ): + np.testing.assert_array_equal( + rgb["rgb"][:], resize_rgb_area(pixels[ep["episode_id"]], (4, 4)) + ) + np.testing.assert_array_equal(rgb["time"][:], data["time"][:]) + assert ( + rgb["cameras"].asstr()[:].tolist() + == ep["cameras"] + == ["agentview", "wrist"] + ) + assert rgb["rgb"].dtype == np.uint8 and "action" not in data + item = stats["statistics"][f"pose/{ep['embodiment']}"] + train_pose = poses[f"{ep['embodiment']}-train"] + np.testing.assert_allclose( + item["mean"], train_pose.astype(np.float32).mean(0), rtol=1e-6 + ) + assert item["count"] == 4 + dataset = PreparedPixelProprioDataset(output / "manifest.json", split="val") + assert dataset[0]["world"]["image"]["value"].shape[-1] == 12 + assert dataset[4]["world"]["proprio"]["value"].shape[-1] == 29 + assert ( + dataset[0]["world"]["proprio"]["value"].min() > 80 + ) # Validation poses did not fit stats. + before = sha256(output / "manifest.json") + with pytest.raises(FileExistsError, match="immutable"): + prepare_pixels(source, spec, output, image_size=[4, 4], patch_size=2) + assert sha256(output / "manifest.json") == before + + +@pytest.mark.parametrize( + "damage", ["state_time", "rgb_checksum", "pose_fields", "split"] +) +def test_prepare_rejects_alignment_corruption_or_split_changes_before_output( + tmp_path, damage +): + source, spec, _, _ = native_corpus(tmp_path, bad_state_time=damage == "state_time") + locations = json.loads(spec.read_text())["episodes"] + if damage == "rgb_checksum": + with h5py.File(locations[0]["rgb_cache"], "a") as data: + data["rgb"][0, 0, 0, 0, 0] += 1 + elif damage in ("pose_fields", "split"): + original = json.loads(source.read_text()) + if damage == "pose_fields": + original["pose_schemas"]["libero_gripper"]["fields"][0] = ( + "incorrect_coordinate" + ) + else: + original["episodes"][1]["split"] = "train" + save_json(source, original) + output = tmp_path / "prepared" + with pytest.raises(ValueError): + prepare_pixels(source, spec, output, image_size=[4, 4], patch_size=2) + assert not output.exists() + + +def test_optional_real_action_export_produces_loader_compatible_labels(tmp_path): + from tests.test_libero_action_sidecar_export import bundle + + source, raw = bundle(tmp_path) + manifest = json.loads(source.read_text()) + locations = [] + for episode in manifest["episodes"]: + location, _, _ = native_episode( + tmp_path, episode, manifest["pose_schemas"][episode["embodiment"]]["fields"] + ) + # The raw fixture stores timestamps starting at zero; use its actual + # rows for all native arrays, then rebind the receipt hashes. + native = Path(location["native_cache"]) + with np.load(native / "aligned.npz") as data: + arrays = {key: data[key] for key in data.files} + arrays["time"] = np.arange(episode["frames"], dtype=np.float64) * 0.05 + arrays["state"][:, 0] = arrays["time"] + np.savez(native / "aligned.npz", **arrays) + nmanifest = json.loads((native / "manifest.json").read_text()) + nmanifest["artifacts"]["aligned.npz"] = sha256(native / "aligned.npz") + save_json(native / "manifest.json", nmanifest) + with h5py.File(location["rgb_cache"], "a") as data: + data["time"][:] = arrays["time"] + receipt = json.loads(Path(location["rgb_receipt"]).read_text()) + receipt.update( + source_manifest_sha256=sha256(native / "manifest.json"), + source_aligned_sha256=sha256(native / "aligned.npz"), + rgb_cache_sha256=sha256(location["rgb_cache"]), + ) + save_json(Path(location["rgb_receipt"]), receipt) + locations.append(location) + spec = save_json(tmp_path / "native.json", {"episodes": locations}) + output = tmp_path / "prepared" + result = prepare_pixels( + source, spec, output, image_size=[4, 4], patch_size=2, raw_export_dir=raw + ) + assert result["action_episodes"] == 2 + dataset = PreparedPixelProprioDataset( + output / "manifest.json", split="val", task="direct_action" + ) + np.testing.assert_array_equal( + dataset[2]["world"]["action"]["value"].numpy(), + (np.arange(28).reshape(4, 7) / 30)[2:3].astype(np.float32), + ) + + +def test_gripper_selection_preserves_frozen_split_without_opening_sharpa(tmp_path): + source, spec, _, _ = native_corpus(tmp_path) + locations = json.loads(spec.read_text()) + for item in locations["episodes"]: + if item["embodiment"] == "libero_sharpa": + item.update( + native_cache="/nonexistent/sharpa", + rgb_cache="/nonexistent/sharpa.h5", + rgb_receipt="/nonexistent/sharpa.json", + ) + save_json(spec, locations) + output = tmp_path / "gripper" + result = prepare_pixels( + source, + spec, + output, + image_size=[4, 4], + patch_size=2, + embodiments=["libero_gripper"], + ) + assert result["episodes"] == 2 + prepared = json.loads((output / "manifest.json").read_text()) + assert prepared["selected_embodiments"] == ["libero_gripper"] + assert set(prepared["pose_schemas"]) == {"libero_gripper"} + assert [(ep["episode_id"], ep["split"]) for ep in prepared["episodes"]] == [ + ("libero_gripper-train", "train"), + ("libero_gripper-val", "val"), + ] + # Selection does not excuse a corrupted frozen source split in the excluded source. + original = json.loads(source.read_text()) + original["episodes"][-1]["split"] = "train" + save_json(source, original) + with pytest.raises(ValueError, match="nonempty train/val"): + prepare_pixels( + source, + spec, + tmp_path / "bad", + image_size=[4, 4], + patch_size=2, + embodiments=["libero_gripper"], + ) + + +def task_episodes(tasks=10, episodes_per_task=50): + return [ + dict( + embodiment="libero_gripper", + episode_id=f"task_{task}-demo_{i}", + frames=75, + split="val" if task < min(tasks, 5) and i == 0 else "train", + native_source=dict( + source_demo=f"/fixture/task_{task}_demo.hdf5", + source_demo_sha256=f"{task + 1:064x}", + episode=f"demo_{i}", + ), + ) + for task in range(tasks) + for i in range(episodes_per_task) + ] + + +def test_balanced_split_is_order_independent_preserves_old_holdout_and_covers_all_tasks(): + from types import SimpleNamespace + + from egomimic.rldb.pixel_proprio import EvenlySpacedEpisodeBatchSampler + + episodes = task_episodes() + unchanged = copy.deepcopy(episodes) + result, policy = derive_split(episodes, 5, 42) + reversed_result, reversed_policy = derive_split(list(reversed(episodes)), 5, 42) + assert episodes == unchanged + assert result == reversed_result and policy == reversed_policy + assert policy["counts"] == {"libero_gripper": {"train": 450, "val": 50}} + assert all( + count["train"] == 45 and count["val"] == 5 + for count in policy["task_counts"].values() + ) + assert ( + sum(count["retained_original_val"] for count in policy["task_counts"].values()) + == 5 + ) + assert all( + ep["split"] == "val" for ep in result.values() if ep["original_split"] == "val" + ) + assert policy["original_identities_sha256"] != policy["derived_identities_sha256"] + _, other_seed = derive_split(episodes, 5, 43) + assert ( + other_seed["derived_identities_sha256"] != policy["derived_identities_sha256"] + ) + heldout = [ep for ep in result.values() if ep["split"] == "val"] + dataset = SimpleNamespace( + split="val", episodes=heldout, counts=[ep["frames"] for ep in heldout] + ) + sampler = EvenlySpacedEpisodeBatchSampler(dataset, 2, 1) + assert len(sampler) == len(list(sampler)) == 100 + for i in range(50): + assert list(sampler)[2 * i : 2 * i + 2] == [[i * 75], [i * 75 + 74]] + + +@pytest.mark.parametrize( + "damage", ["task_name", "source_hash", "excess_old_val", "empty_train"] +) +def test_balanced_split_rejects_ambiguous_task_or_destructive_split(damage): + episodes = task_episodes(tasks=2, episodes_per_task=4) + count = 2 + if damage == "task_name": + episodes[0]["episode_id"] = "different-task-demo_0" + elif damage == "source_hash": + episodes[0]["native_source"]["source_demo_sha256"] = "e" * 64 + elif damage == "excess_old_val": + episodes[1]["split"] = episodes[2]["split"] = "val" + else: + count = 4 + with pytest.raises(ValueError): + derive_split(episodes, count, 42) + + +def balanced_native_corpus(tmp_path): + episodes = task_episodes(tasks=2, episodes_per_task=4) + schema = dict(dimension=9, fields=[f"p{i}" for i in range(9)]) + derived, _ = derive_split(episodes, 2, 42) + locations, poses = [], {} + for ep in episodes: + ep["frames"] = 4 + key = ep["embodiment"], ep["episode_id"] + # Every old/new validation episode is extreme, exposing any stats leak. + offset = 10000 if derived[key]["split"] == "val" else 0 + location, pose, _ = native_episode( + tmp_path, ep, schema["fields"], pose_offset=offset + ) + locations.append(location) + poses[ep["episode_id"]] = pose + source = save_json( + tmp_path / "source.json", + dict(episodes=episodes, pose_schemas={"libero_gripper": schema}), + ) + spec = save_json(tmp_path / "locations.json", dict(episodes=locations)) + return source, spec, poses + + +def test_derived_preparation_normalizes_only_new_training_and_records_original_assignments( + tmp_path, +): + source, spec, poses = balanced_native_corpus(tmp_path) + before = sha256(source) + output = tmp_path / "derived" + result = prepare_pixels( + source, + spec, + output, + image_size=[4, 4], + patch_size=2, + validation_episodes_per_task=2, + split_seed=42, + ) + assert sha256(source) == before + assert result["split_counts"] == {"libero_gripper": {"train": 4, "val": 4}} + manifest = json.loads((output / "manifest.json").read_text()) + saved = manifest["split_provenance"] + assert sha256(output / saved["path"]) == saved["sha256"] + policy = json.loads((output / saved["path"]).read_text()) + assert policy["kind"] == "balanced_native_libero_tasks_v1" and policy["seed"] == 42 + assert saved["original_identities_sha256"] == policy["original_identities_sha256"] + assert sum(ep["split"] != ep["original_split"] for ep in manifest["episodes"]) == 2 + norm = json.loads((output / "normalization.json").read_text()) + stats = norm["statistics"]["pose/libero_gripper"] + expected = np.concatenate( + [ + poses[ep["episode_id"]] + for ep in manifest["episodes"] + if ep["split"] == "train" + ] + ) + np.testing.assert_allclose(stats["mean"], expected.mean(0), atol=1e-6) + assert max(stats["mean"]) < 10 and stats["count"] == 16 + assert len(norm["train_episodes"]) == 4 + data = PreparedPixelProprioDataset(output / "manifest.json", split="val") + assert len(data.episodes) == 4 and data[0]["world"]["proprio"]["value"].min() > 1000 + + +def test_derived_split_still_rejects_original_raw_command_split_mismatch( + tmp_path, monkeypatch +): + from scripts.world_model import prepare_pixel_proprio as module + + source, spec, _ = balanced_native_corpus(tmp_path) + original = json.loads(source.read_text())["episodes"] + raw_values = {} + for ep in original: + raw = copy.deepcopy(ep) + raw["split"] = "train" if ep["split"] == "val" else "val" + raw_values[(ep["embodiment"], ep["episode_id"])] = ( + np.zeros((4, 7), np.float32), + 0.25 + np.arange(4) * 0.05, + raw, + ) + # The low-level binary bundle verifier has independent corruption tests; + # this boundary specifically must compare against ORIGINAL assignments. + monkeypatch.setattr(module, "_raw_actions", lambda *_: ({}, {}, raw_values)) + output = tmp_path / "bad" + with pytest.raises(ValueError, match="Raw action source/split/state time"): + prepare_pixels( + source, + spec, + output, + image_size=[4, 4], + patch_size=2, + raw_export_dir=tmp_path, + validation_episodes_per_task=2, + split_seed=42, + ) + assert not output.exists() diff --git a/tests/test_pixel_readout_graph.py b/tests/test_pixel_readout_graph.py new file mode 100644 index 000000000..ba242d9b2 --- /dev/null +++ b/tests/test_pixel_readout_graph.py @@ -0,0 +1,459 @@ +"""Explicit readout graph, strict initialization and frozen-image AdamW behavior.""" + +import copy +from pathlib import Path + +import pytest +import torch +from hydra import compose, initialize_config_dir +from hydra.utils import instantiate +from omegaconf import OmegaConf +from omegaconf.errors import InterpolationToMissingValueError + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, + PlainPixelMSE, +) +from egomimic.pipeline.algo import PipelineAlgo +from egomimic.pipeline.core import Pipeline, sum_losses +from egomimic.pipeline.stages_pixel_readout import ( + PixelReadoutDenoiserStage, + PixelReadoutImageDenoiserStage, + PixelReadoutImageLossStage, + PixelReadoutMSEStage, + PixelReadoutNoisingStage, + PixelStateActionReadoutStage, +) +from egomimic.pipeline.stages_pixel_video import ( + PixelVideoNoisingStage, + PixelVideoSamplerStage, +) +from scripts.world_model.initialize_pixel_readout import ( + initialize_readout, + sha256_file, + transfer_shared_weights, + validate_transition, +) +from tools.config_graph import build_graph + +IDS = { + "pixel_noising": 0, + "pixel_denoiser": 1, + "pixel_loss": 2, + "readout_noising": 3, + "readout_denoiser": 4, + "pixel_sampler": 5, + "state_action_readout": 6, + "measured_state_loss": 7, + "control_command_loss": 8, + "image_proprio": 1, +} +CONFIG_ROOT = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + + +@pytest.fixture(autouse=True) +def single_cpu_thread(): + before = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(before) + + +def arguments(mode="continuous"): + return dict( + objective="visual_state_action", + phase_a_condition_on_proprio=False, + pose_dims={"robot": 3}, + action_dims={"robot": 2}, + image_dim=12, + width=16, + depth=2, + heads=2, + mlp_ratio=2, + max_cameras=2, + max_tokens=16, + time_mode=mode, + loss_reduction="mean_per_example", + ) + + +def fixture_batch(): + image = { + "value": torch.linspace(-1, 1, 2 * 4 * 12).reshape(2, 4, 12), + "time": torch.tensor([[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0]]), + "index": torch.tensor([[0, 0, 1, 1], [0, 0, 1, 1]]), + "valid": torch.tensor([[True, True, True, True], [True, True, True, False]]), + "observed": torch.tensor( + [[True, True, False, False], [True, False, False, False]] + ), + "xy": torch.tensor([[[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [1.0, 0.0]]]).repeat( + 2, 1, 1 + ), + "camera": torch.tensor([[0, 1, 0, 1], [0, 1, 0, 1]]), + } + image["value"][1, -1] = float("nan") + pose = { + "value": torch.tensor( + [[[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]], [[-0.1, 0.2, 0.3], [0.3, -0.5, 0.6]]] + ), + "time": torch.tensor([[0.0, 1.0], [0.0, 1.0]]), + "index": torch.tensor([[0, 1], [0, 1]]), + "valid": torch.ones(2, 2, dtype=torch.bool), + "observed": torch.tensor([[True, False], [True, False]]), + } + action = { + **{k: v.clone() for k, v in pose.items() if k != "value"}, + "value": torch.tensor([[[0.2, -0.3], [-0.4, 0.5]], [[-0.2, 0.1], [0.3, -0.4]]]), + "observed": torch.zeros(2, 2, dtype=torch.bool), + } + return { + "world": {"image": image, "proprio": pose, "action": action}, + "embodiment": ["robot", "robot"], + } + + +def stages(model, training_mode="frozen_image"): + return [ + PixelVideoNoisingStage(LinearGaussianImageNoising()), + PixelReadoutDenoiserStage(model, training_mode), + PixelReadoutImageLossStage(PlainPixelMSE(model.loss_reduction)), + PixelReadoutNoisingStage(LinearGaussianImageNoising(0.2)), + PixelReadoutImageDenoiserStage(), + PixelVideoSamplerStage( + LinearDDIMImageSampler(3), "pixel_video/sampled_prediction" + ), + PixelStateActionReadoutStage(), + PixelReadoutMSEStage(PlainPixelMSE(model.loss_reduction), "proprio"), + PixelReadoutMSEStage(PlainPixelMSE(model.loss_reduction), "action"), + ] + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_joint_graph_preserves_legacy_three_losses_gradients_and_rng(mode): + model = ImageProprioModel(**arguments(mode)) + legacy = copy.deepcopy(model) + graph = Pipeline(stages(model, "joint"), stage_ids=IDS) + batch = fixture_batch() + torch.manual_seed(61) + expected = legacy.training_losses(batch["world"], "robot") + expected_rng = torch.get_rng_state().clone() + expected["total"].backward() + torch.manual_seed(61) + result = graph.execute(batch, mode="train") + assert set(k for k in result if k.startswith("loss/")) == { + "loss/pixel_video", + "loss/measured_state", + "loss/control_command", + } + assert "log/image_proprio_total" not in result + for key, legacy_key in ( + ("pixel_video", "image"), + ("measured_state", "pose"), + ("control_command", "action"), + ): + torch.testing.assert_close( + result[f"loss/{key}"], expected[legacy_key], rtol=0, atol=0 + ) + torch.testing.assert_close(sum_losses(result), expected["total"], rtol=0, atol=0) + assert torch.equal(torch.get_rng_state(), expected_rng) + assert result["pixel_readout/noise_levels"].max() <= 0.2 + sum_losses(result).backward() + for name, parameter in model.named_parameters(): + reference = dict(legacy.named_parameters())[name] + assert parameter.grad is not None and reference.grad is not None, name + torch.testing.assert_close(parameter.grad, reference.grad, rtol=0, atol=0) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_frozen_image_adamw_updates_only_new_branch_and_preserves_samples(mode): + target_args = arguments(mode) + source_args = { + **target_args, + "objective": "visual_only", + "pose_dims": None, + "action_dims": None, + } + source = ImageProprioModel(**source_args) + model = ImageProprioModel(**target_args) + transfer_shared_weights(source, model) + graph = Pipeline(stages(model), stage_ids=IDS) + assert all(key.startswith("stages.1.model.") for key in graph.state_dict()) + assert sum(p.numel() for p in graph.parameters()) == sum( + p.numel() for p in model.parameters() + ) + assert graph.stage_by_id("image_proprio") is graph.stage_by_id("pixel_denoiser") + before = {name: p.detach().clone() for name, p in model.named_parameters()} + batch = fixture_batch() + sampler = LinearDDIMImageSampler(3) + torch.manual_seed(78) + expected_images = sampler(model, batch["world"], "robot") + torch.manual_seed(78) + source_images = sampler(source, {"image": batch["world"]["image"]}, "robot") + torch.testing.assert_close(expected_images, source_images, rtol=0, atol=0) + # Include frozen tensors in AdamW too: no gradient means no decay/update. + optimizer = torch.optim.AdamW(graph.parameters(), lr=5e-5, weight_decay=0.01) + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + output = graph.execute(batch, mode="train") + assert not output["loss/pixel_video"].requires_grad + loss = sum_losses(output) + assert loss.requires_grad and loss.dtype == torch.float32 + loss.backward() + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + assert ( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + ), name + else: + assert parameter.grad is None, name + optimizer.step() + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + assert torch.equal(parameter, before[name]), name + else: + assert not torch.equal(parameter, before[name]), name + assert not model.phase.weight.requires_grad + assert torch.equal(model.phase.weight, before["phase.weight"]) + for name, tensor in source.state_dict().items(): + actual = model.state_dict()[name] + if name == "phase.weight": + actual = actual[:1] + assert torch.equal(tensor, actual), name + torch.manual_seed(78) + actual_images = sampler(model, batch["world"], "robot") + torch.testing.assert_close(actual_images, expected_images, rtol=0, atol=0) + # Restoring ordinary tensors into the same configured freeze policy is strict. + restored = Pipeline(stages(ImageProprioModel(**target_args)), stage_ids=IDS) + restored.load_state_dict(graph.state_dict(), strict=True) + assert [p.requires_grad for p in restored.parameters()] == [ + p.requires_grad for p in graph.parameters() + ] + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_inference_matches_legacy_sampling_and_never_consumes_query_labels(mode): + model = ImageProprioModel(**arguments(mode)) + graph = Pipeline(stages(model), stage_ids=IDS) + batch = fixture_batch() + torch.manual_seed(91) + expected = model.sample(batch["world"], "robot", steps=3) + changed = copy.deepcopy(batch) + changed["world"]["proprio"]["value"][~changed["world"]["proprio"]["observed"]] = ( + float("nan") + ) + changed["world"]["action"]["value"].fill_(float("nan")) + torch.manual_seed(91) + result = graph.execute(changed, mode="inference") + assert not any(key.startswith("loss/") for key in result) + assert result["pixel_video/denoiser"] is model + for key in ("image", "proprio", "action"): + torch.testing.assert_close( + result["image_proprio_prediction"][key], expected[key], rtol=0, atol=0 + ) + + +def test_measured_and_command_targets_are_separate_and_missing_labels_masked(): + graph = Pipeline(stages(ImageProprioModel(**arguments())), stage_ids=IDS) + batch = fixture_batch() + batch["world"]["proprio"]["target_valid"] = batch["world"]["proprio"][ + "observed" + ].clone() + batch["world"]["proprio"]["value"][:, 1] = float("nan") + result = graph.execute(batch, mode="train") + assert result["loss/measured_state"] == 0 + assert result["loss/control_command"] > 0 + assert torch.isfinite(sum_losses(result)) + # This graph has two experts, with separate state and command output heads. + model = graph.stage_by_id("pixel_denoiser").model + assert all(len(block.qkv) == 2 for block in model.blocks) + assert model.pose_out["robot"] is not model.action_out["robot"] + + +def recipe(mode, *, require_initialization=True): + overrides = [ + f"+experiment=pixel_proprio/readout_{mode}", + "hydra/launcher=basic", + "world_model_attention=native", + "data.manifest_path=/fixture/actions.json", + "paths.output_dir=/fixture/output", + "paths.work_dir=/fixture/source", + ] + if require_initialization: + overrides.extend( + [ + "pixel_readout.initial_weights=/fixture/readout.pt", + "pixel_readout.initial_weights_sha256=" + "a" * 64, + ] + ) + with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): + return compose(config_name="train_zarr_cartesian", overrides=overrides) + + +@pytest.mark.parametrize( + "mode,pose_dim,action_dim,capacity", + [("continuous", 9, 7, 25856), ("discrete", 36, 28, 3232)], +) +def test_actual_stage2_recipes_graph_and_counts( + tmp_path, mode, pose_dim, action_dim, capacity +): + cfg = recipe(mode) + plain = OmegaConf.to_container(cfg, resolve=True) + pipeline = plain["model"]["pipeline"] + expected_ids = { + "pixel_noising": 0, + "pixel_denoiser": 1, + "pixel_prediction": 2, + "pixel_loss": 3, + "readout_noising": 4, + "readout_denoiser": 5, + "pixel_sampler": 6, + "state_action_readout": 7, + "measured_state_loss": 8, + "control_command_loss": 9, + "image_proprio": 1, + } + assert pipeline["stage_ids"] == expected_ids + assert len(pipeline["stages"]) == 10 + model_cfg = pipeline["stages"][1]["model"] + assert model_cfg["pose_dims"] == {"libero_gripper": pose_dim} + assert model_cfg["action_dims"] == {"libero_gripper": action_dim} + assert model_cfg["max_tokens"] == capacity + assert not cfg.data.image_only and cfg.data.task == "visual_state_action" + assert cfg.evaluator.objective == "visual_state_action" + assert cfg.model.optimizer.lr == 5e-5 and cfg.model.optimizer.weight_decay == 0.01 + assert pipeline["stages"][1]["training_mode"] == "frozen_image" + assert pipeline["stages"][4]["noising"]["maximum"] == 0.2 + assert [pipeline["stages"][i]["weight"] for i in (3, 8, 9)] == [1.0, 1.0, 1.0] + assert pipeline["initialization"] == [ + { + "stage_id": "image_proprio", + "module_path": "model", + "source": "/fixture/readout.pt", + "sha256": "a" * 64, + } + ] + with torch.device("meta"): + stage = instantiate(pipeline["stages"][1]) + frozen = sum(p.numel() for p in stage.parameters() if not p.requires_grad) + assert frozen == 88_178_688 # Exactly the copied AdaLN/3-axis-RoPE image weights. + assert stage.model.noise_conditioning == "adaln_zero" + assert not hasattr(stage.model, "phase") + assert sum(p.numel() for p in stage.parameters()) > frozen + # Exact stage-one model declaration is compatible with the strict migrator. + with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): + source_cfg = compose( + config_name="train_zarr_cartesian", + overrides=[ + f"+experiment=pixel_proprio/video_{mode}", + "hydra/launcher=basic", + "world_model_attention=native", + ], + ) + validate_transition( + OmegaConf.to_container(source_cfg.model.pipeline.stages[1].model, resolve=True), + model_cfg, + ) + plain["data"]["train_datasets"] = { + "libero_gripper": {"batch_keys": ["world", "embodiment"]} + } + path = tmp_path / "graph.yaml" + OmegaConf.save(OmegaConf.create(plain), path) + with torch.device("meta"): + training = build_graph(path, mode="train") + inference = build_graph(path, mode="inference") + assert training["lint"] == inference["lint"] == [] + assert len(training["nodes"]) == 9 and len(inference["nodes"]) == 3 + assert [node["t"] for node in inference["nodes"]] == [ + "PixelReadoutModelStage", + "PixelReadoutSampledImageStage", + "PixelStateActionReadoutStage", + ] + + +def test_recipe_cannot_resolve_without_exact_initialization_artifact(): + cfg = recipe("continuous", require_initialization=False) + with pytest.raises(InterpolationToMissingValueError): + OmegaConf.to_container(cfg.model, resolve=True) + + +def test_strict_index_one_checkpoint_expansion_loads_into_frozen_graph(tmp_path): + target_args = arguments() + source_args = { + **target_args, + "objective": "visual_only", + "pose_dims": None, + "action_dims": None, + } + source = ImageProprioModel(**source_args) + checkpoint = tmp_path / "stage1.ckpt" + torch.save( + { + "state_dict": { + f"nets.pipeline.stages.1.model.{name}": value + for name, value in source.state_dict().items() + }, + "hyper_parameters": { + "config": { + "model": {"pipeline": {"stages": [{}, {"model": source_args}]}} + } + }, + }, + checkpoint, + ) + receipt = initialize_readout( + source_checkpoint=checkpoint, + source_sha256=sha256_file(checkpoint), + source_config=source_args, + target_config=target_args, + output_path=tmp_path / "expanded.pt", + receipt_path=tmp_path / "expanded.json", + seed=42, + state_dict_path=("state_dict",), + source_prefix="nets.pipeline.stages.1.model.", + source_config_path=( + "hyper_parameters", + "config", + "model", + "pipeline", + "stages", + 1, + "model", + ), + ) + model = ImageProprioModel(**target_args) + algo = PipelineAlgo( + stages(model), + device="cpu", + stage_ids=IDS, + initialization=receipt["initialization"], + ) + assert receipt["phase_a_preservation"]["bit_exact"] + assert algo.initialization_receipts[0]["sha256"] == receipt["output"]["sha256"] + state = torch.load(receipt["output"]["path"], weights_only=True) + assert all( + torch.equal(value, state[name]) for name, value in model.state_dict().items() + ) + assert ( + not model.image_in.weight.requires_grad + and model.pose_in["robot"].weight.requires_grad + ) + bad = copy.deepcopy(receipt["initialization"]) + bad[0]["sha256"] = "0" * 64 + with pytest.raises(ValueError, match="hash mismatch"): + PipelineAlgo( + stages(ImageProprioModel(**target_args)), + device="cpu", + stage_ids=IDS, + initialization=bad, + ) + + +def test_freeze_rejects_unknown_parameters_or_policy(): + model = ImageProprioModel(**arguments()) + model.unclassified = torch.nn.Parameter(torch.ones(1)) + with pytest.raises(ValueError, match="unknown or shared"): + PixelReadoutDenoiserStage(model) + with pytest.raises(ValueError, match="frozen_image or joint"): + PixelReadoutDenoiserStage(ImageProprioModel(**arguments()), "partial") diff --git a/tests/test_pixel_readout_initialization.py b/tests/test_pixel_readout_initialization.py new file mode 100644 index 000000000..27778c67c --- /dev/null +++ b/tests/test_pixel_readout_initialization.py @@ -0,0 +1,366 @@ +"""Offline readout expansion preserves phase A and loads through PipelineAlgo.""" + +import copy +import json +import pickle +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +from omegaconf import OmegaConf + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.pipeline.algo import PipelineAlgo +from egomimic.pipeline.stages_image_proprio import ImageProprioStage +from scripts.world_model.initialize_pixel_readout import ( + initialize_readout, + resolve_model_config, + sha256_file, + validate_transition, + verify_phase_a, +) + + +@pytest.fixture(autouse=True) +def single_cpu_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def configs(mode="continuous", pose_dim=9, action_dim=7): + source = { + "objective": "visual_only", + "pose_dims": {"libero_gripper": pose_dim}, + "image_dim": 12, + "width": 16, + "heads": 2, + "depth": 2, + "mlp_ratio": 2, + "max_cameras": 2, + "max_tokens": 12, + "time_mode": mode, + "use_camera_embedding": True, + "use_modality_embedding": False, + "use_embodiment_embedding": False, + } + target = copy.deepcopy(source) + target.update( + objective="visual_state_action", + action_dims={"libero_gripper": action_dim}, + max_tokens=14, + ) + return source, target + + +def write_source(tmp_path, source_config): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(123) + source = ImageProprioModel(**source_config) + path = tmp_path / "stage_one.pt" + torch.save(source.state_dict(), path) + return path, source + + +def run_initialization(tmp_path, source_config=None, target_config=None, **overrides): + if source_config is None: + source_config, target_config = configs() + path, source = write_source(tmp_path, source_config) + arguments = { + "source_checkpoint": path, + "source_sha256": sha256_file(path), + "source_config": source_config, + "target_config": target_config, + "output_path": tmp_path / "stage_two.pt", + "receipt_path": tmp_path / "receipt.json", + "seed": 42, + } + arguments.update(overrides) + receipt = initialize_readout(**arguments) + return source, receipt, arguments + + +@pytest.mark.parametrize( + "mode,pose_dim,action_dim", [("continuous", 9, 7), ("discrete", 36, 28)] +) +def test_preserves_phase_a_and_strict_pipeline_load( + tmp_path, mode, pose_dim, action_dim +): + source_config, target_config = configs(mode, pose_dim, action_dim) + rng_before = torch.random.get_rng_state().clone() + source, receipt, arguments = run_initialization( + tmp_path, source_config, target_config + ) + assert torch.equal(torch.random.get_rng_state(), rng_before) + assert receipt["phase_a_preservation"]["bit_exact"] + assert len(receipt["phase_a_preservation"]["checks"]) == 2 + assert not receipt["optimizer_state_transferred"] + assert not receipt["training_step_transferred"] + assert receipt == json.loads(arguments["receipt_path"].read_text()) + assert receipt["output"]["sha256"] == sha256_file(arguments["output_path"]) + state = torch.load(arguments["output_path"], weights_only=True) + assert all(isinstance(value, torch.Tensor) for value in state.values()) + assert "camera.weight" in state + assert "modality" not in state and "embodiment.weight" not in state + target = ImageProprioModel(**target_config) + pipeline = PipelineAlgo( + [ImageProprioStage(target)], + device="cpu", + stage_ids={"image_proprio": 0}, + initialization=receipt["initialization"], + ) + assert pipeline.initialization_receipts[0]["sha256"] == receipt["output"]["sha256"] + for key, value in state.items(): + assert torch.equal(target.state_dict()[key], value) + for name, value in source.state_dict().items(): + actual = target.state_dict()[name] + if name in ("phase.weight", "proprio_role.weight"): + actual = actual[:1] + assert torch.equal(value, actual) + assert verify_phase_a(source, target)["bit_exact"] + + +def test_new_parameters_and_rows_match_declared_seed_constructor(tmp_path): + _, receipt, arguments = run_initialization(tmp_path) + state = torch.load(arguments["output_path"], weights_only=True) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(arguments["seed"]) + fresh = ImageProprioModel(**receipt["target_model"]).state_dict() + for name in receipt["transfer"]["new_tensors"]: + assert torch.equal(state[name], fresh[name]) + for name in receipt["transfer"]["expanded_embeddings"]: + assert torch.equal(state[name][1:], fresh[name][1:]) + # A second invocation creates identical weights, not an optimizer continuation. + again = initialize_readout( + **{ + **arguments, + "output_path": tmp_path / "repeat.pt", + "receipt_path": tmp_path / "repeat.json", + } + ) + repeated = torch.load(again["output"]["path"], weights_only=True) + assert all(torch.equal(state[key], repeated[key]) for key in state) + changed = initialize_readout( + **{ + **arguments, + "seed": 43, + "output_path": tmp_path / "other_seed.pt", + "receipt_path": tmp_path / "other_seed.json", + } + ) + changed_state = torch.load(changed["output"]["path"], weights_only=True) + assert not torch.equal(state["action_query"], changed_state["action_query"]) + assert torch.equal(state["image_in.weight"], changed_state["image_in.weight"]) + + +@pytest.mark.parametrize( + "failure", ["hash", "missing", "unexpected", "shape", "dtype", "nan"] +) +def test_bad_source_never_creates_output(tmp_path, failure): + source_config, target_config = configs() + path, source = write_source(tmp_path, source_config) + state = source.state_dict() + if failure == "missing": + state.pop("image_in.bias") + elif failure == "unexpected": + state["surprise"] = torch.zeros(1) + elif failure == "shape": + state["image_in.bias"] = torch.zeros(99) + elif failure == "dtype": + state["image_in.bias"] = state["image_in.bias"].double() + elif failure == "nan": + state["image_in.bias"][0] = float("nan") + torch.save(state, path) + output, receipt = tmp_path / "output.pt", tmp_path / "receipt.json" + with pytest.raises(ValueError): + initialize_readout( + source_checkpoint=path, + source_sha256="0" * 64 if failure == "hash" else sha256_file(path), + source_config=source_config, + target_config=target_config, + output_path=output, + receipt_path=receipt, + seed=42, + ) + assert not output.exists() and not receipt.exists() + + +@pytest.mark.parametrize( + "field,value", + [ + ("time_mode", "discrete"), + ("heads", 4), + ("depth", 3), + ("pose_dims", {"libero_gripper": 36}), + ("use_camera_embedding", False), + ("use_modality_embedding", True), + ("use_embodiment_embedding", True), + ("pose_conditioning", "clean"), + ("max_tokens", 10), + ], +) +def test_rejects_architecture_or_conditioning_drift(field, value): + source, target = configs() + target[field] = value + with pytest.raises(ValueError): + validate_transition(source, target) + + +def test_config_order_and_targets_cannot_be_inferred(): + source, target = configs() + source["pose_dims"] = {"one": 9, "two": 9} + target["pose_dims"] = {"two": 9, "one": 9} + target["action_dims"] = {"two": 7, "one": 7} + with pytest.raises(ValueError, match="order"): + validate_transition(source, target) + with pytest.raises(ValueError, match="Only ImageProprioModel"): + resolve_model_config({**source, "_target_": "arbitrary.Code"}) + with pytest.raises(ValueError, match="resolved"): + resolve_model_config({**source, "time_mode": "${data.mode}"}) + with pytest.raises(TypeError): + resolve_model_config({**source, "unknown_parameter": True}) + + +def test_exact_nested_namespace_and_checkpoint_config_binding(tmp_path): + source_config, target_config = configs() + path, source = write_source(tmp_path, source_config) + prefix = "nets.pipeline.stages.0.model." + payload = { + "state_dict": { + prefix + key: value for key, value in source.state_dict().items() + }, + "model_config": source_config, + "optimizer_states": [{"do_not_transfer": 123}], + "global_step": 123, + } + payload["state_dict"]["another_stage.weight"] = torch.zeros(1) + torch.save(payload, path) + arguments = { + "source_checkpoint": path, + "source_sha256": sha256_file(path), + "source_config": source_config, + "target_config": target_config, + "output_path": tmp_path / "output.pt", + "receipt_path": tmp_path / "receipt.json", + "seed": 42, + "source_prefix": prefix, + "state_dict_path": ["state_dict"], + } + with pytest.raises(ValueError, match="requires source_config_path"): + initialize_readout(**arguments) + arguments["source_config_path"] = ["model_config"] + with pytest.raises(ValueError, match="namespace"): + initialize_readout(**{**arguments, "source_prefix": "wrong."}) + receipt = initialize_readout(**arguments) + assert receipt["source"]["config_binding"] == "checkpoint" + state = torch.load(arguments["output_path"], weights_only=True) + assert "optimizer_states" not in state and "global_step" not in state + # Head count and physical/discrete time are not recoverable from weight shapes. + payload["model_config"] = {**source_config, "heads": 4} + torch.save(payload, path) + with pytest.raises(ValueError, match="Checkpoint model configuration"): + initialize_readout( + **{ + **arguments, + "source_sha256": sha256_file(path), + "output_path": tmp_path / "bad.pt", + "receipt_path": tmp_path / "bad.json", + } + ) + + +def test_omegaconf_checkpoint_requires_explicit_trust_after_hash(tmp_path): + source_config, target_config = configs() + path, source = write_source(tmp_path, source_config) + torch.save( + {"state_dict": source.state_dict(), "config": OmegaConf.create(source_config)}, + path, + ) + arguments = dict( + source_checkpoint=path, + source_sha256=sha256_file(path), + source_config=source_config, + target_config=target_config, + output_path=tmp_path / "out.pt", + receipt_path=tmp_path / "receipt.json", + seed=42, + state_dict_path=["state_dict"], + source_config_path=["config"], + ) + with pytest.raises(pickle.UnpicklingError): + initialize_readout(**arguments) + with pytest.raises(ValueError, match="SHA256"): + initialize_readout( + **{ + **arguments, + "source_sha256": "0" * 64, + "trusted_lightning_checkpoint": True, + } + ) + result = initialize_readout(**arguments, trusted_lightning_checkpoint=True) + assert result["source"]["trusted_pickle"] + + +def test_refuses_overwrite_and_output_aliases(tmp_path): + _, _, arguments = run_initialization(tmp_path) + with pytest.raises(FileExistsError): + initialize_readout(**arguments) + with pytest.raises(ValueError, match="distinct"): + initialize_readout( + **{**arguments, "output_path": arguments["source_checkpoint"]} + ) + + +def test_phase_a_guard_catches_accidental_visual_weight_change(tmp_path): + source_config, target_config = configs() + source, receipt, arguments = run_initialization( + tmp_path, source_config, target_config + ) + target = ImageProprioModel(**target_config) + target.load_state_dict(torch.load(arguments["output_path"], weights_only=True)) + with torch.no_grad(): + target.image_out.bias.add_(0.1) + with pytest.raises(ValueError, match="Phase A preservation"): + verify_phase_a(source, target) + + +def test_standalone_cli_from_outside_checkout(tmp_path): + source_config, target_config = configs() + path, _ = write_source(tmp_path, source_config) + source_json, target_json = tmp_path / "source.json", tmp_path / "target.json" + source_json.write_text(json.dumps(source_config)) + target_json.write_text(json.dumps(target_config)) + script = ( + Path(__file__).resolve().parents[1] + / "scripts/world_model/initialize_pixel_readout.py" + ) + output, receipt = tmp_path / "output.pt", tmp_path / "receipt.json" + result = subprocess.run( + [ + sys.executable, + str(script), + "--source-checkpoint", + str(path), + "--source-sha256", + sha256_file(path), + "--source-config", + str(source_json), + "--target-config", + str(target_json), + "--output", + str(output), + "--receipt", + str(receipt), + "--seed", + "42", + ], + cwd=tmp_path, + check=True, + text=True, + capture_output=True, + ) + assert json.loads(result.stdout)["status"] == "PASS" + assert json.loads(receipt.read_text())["output"]["sha256"] == sha256_file(output) diff --git a/tests/test_pixel_readout_modularity.py b/tests/test_pixel_readout_modularity.py new file mode 100644 index 000000000..b7de1c98d --- /dev/null +++ b/tests/test_pixel_readout_modularity.py @@ -0,0 +1,218 @@ +"""Readout providers and optional image losses share one stable model owner.""" + +import copy + +import pytest +import torch + +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, + PlainPixelMSE, +) +from egomimic.pipeline.core import Pipeline, sum_losses +from egomimic.pipeline.stages_pixel_readout import ( + PixelReadoutImageDenoiserStage, + PixelReadoutImageLossStage, + PixelReadoutModelStage, + PixelReadoutMSEStage, + PixelReadoutNoisingStage, + PixelReadoutSampledImageStage, + PixelReadoutVideoDenoiserStage, + PixelStateActionReadoutStage, +) +from egomimic.pipeline.stages_pixel_video import PixelVideoNoisingStage +from tests.test_pixel_noise_conditioning_profiles import _trained_image_source +from tests.test_pixel_readout_graph import arguments, fixture_batch, stages + +MODEL_IDS = {"pixel_denoiser": 1, "image_proprio": 1} +IMAGE_PORT = "pixel_readout/image_prediction" + + +@pytest.fixture(autouse=True) +def single_cpu_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def model(mode="continuous", spatial=False): + kwargs = { + **arguments(mode), + "noise_conditioning": "adaln_zero", + "use_phase_embedding": False, + } + if spatial: + kwargs.update( + position_encoding="spatiotemporal_rope", + rope_axis_dims=[4, 2, 2], + spatial_grid_size=[2, 2], + ) + return _trained_image_source(kwargs) + + +def modular_stages(network, training_mode="frozen_image"): + return [ + PixelVideoNoisingStage(LinearGaussianImageNoising()), + PixelReadoutModelStage(network, training_mode), + PixelReadoutVideoDenoiserStage(), + PixelReadoutImageLossStage(PlainPixelMSE(network.loss_reduction)), + PixelReadoutNoisingStage(LinearGaussianImageNoising(0.2)), + PixelReadoutImageDenoiserStage(prediction_key=IMAGE_PORT), + PixelReadoutSampledImageStage(LinearDDIMImageSampler(3)), + PixelStateActionReadoutStage(image_key=IMAGE_PORT), + PixelReadoutMSEStage(PlainPixelMSE(network.loss_reduction), "proprio"), + PixelReadoutMSEStage(PlainPixelMSE(network.loss_reduction), "action"), + ] + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("policy", ["frozen_image", "joint"]) +@pytest.mark.parametrize("spatial", [False, True]) +def test_default_modular_graph_preserves_outputs_gradients_rng_and_checkpoints( + mode, policy, spatial +): + network = model(mode, spatial) + legacy = Pipeline(stages(copy.deepcopy(network), policy), stage_ids=MODEL_IDS) + graph = Pipeline(modular_stages(network, policy), stage_ids=MODEL_IDS) + assert set(graph.state_dict()) == set(legacy.state_dict()) + assert all(key.startswith("stages.1.model.") for key in graph.state_dict()) + graph.load_state_dict(legacy.state_dict(), strict=True) + assert sum(p.numel() for p in graph.parameters()) == sum( + p.numel() for p in network.parameters() + ) + batch = fixture_batch() + torch.manual_seed(171) + expected = legacy.execute(batch, mode="train") + expected_rng = torch.get_rng_state().clone() + sum_losses(expected).backward() + torch.manual_seed(171) + actual = graph.execute(batch, mode="train") + assert torch.equal(torch.get_rng_state(), expected_rng) + for key in ( + "loss/pixel_video", + "loss/measured_state", + "loss/control_command", + IMAGE_PORT, + ): + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + sum_losses(actual).backward() + for name, parameter in graph.named_parameters(): + reference = dict(legacy.named_parameters())[name] + if reference.grad is None: + assert parameter.grad is None, name + else: + torch.testing.assert_close(parameter.grad, reference.grad, rtol=0, atol=0) + torch.manual_seed(31) + expected = legacy.execute(batch, mode="inference")["image_proprio_prediction"] + expected_rng = torch.get_rng_state().clone() + torch.manual_seed(31) + actual = graph.execute(batch, mode="inference")["image_proprio_prediction"] + assert torch.equal(torch.get_rng_state(), expected_rng) + for key in ("image", "proprio", "action"): + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + + +def test_high_noise_branch_can_be_removed_without_losing_owner_or_readout_gradients(): + network = model(spatial=True) + full = modular_stages(network) + # Keep the owner at index1 while dropping all three high-noise branch nodes. + selected = [full[4], full[1], full[5], full[6], full[7], full[8], full[9]] + graph = Pipeline(selected, stage_ids=MODEL_IDS) + phases = [] + hook = network.register_forward_pre_hook( + lambda module, args, kwargs: phases.append(kwargs["phase"]), with_kwargs=True + ) + before = {key: value.clone() for key, value in network.state_dict().items()} + optimizer = torch.optim.AdamW(graph.parameters(), lr=5e-5) + result = graph.execute(fixture_batch(), mode="train") + hook.remove() + assert phases == ["A", "B"] + assert "loss/pixel_video" not in result and "pixel_video/noisy_image" not in result + assert graph.stage_by_id("image_proprio").model is network + assert all(key.startswith("stages.1.model.") for key in graph.state_dict()) + sum_losses(result).backward() + for name, parameter in network.named_parameters(): + if parameter.requires_grad: + assert ( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + ), name + else: + assert parameter.grad is None, name + optimizer.step() + for name in full[1].trainability["frozen_parameters"]: + assert torch.equal(network.state_dict()[name], before[name]), name + assert not torch.equal( + network.pose_out["robot"].weight, before["pose_out.robot.weight"] + ) + assert not torch.equal( + network.action_out["robot"].weight, before["action_out.robot.weight"] + ) + + +def test_sampling_provider_swaps_training_images_with_same_weights_and_tensor_port(): + network = model(spatial=True) + baseline = Pipeline(modular_stages(copy.deepcopy(network)), stage_ids=MODEL_IDS) + selected = modular_stages(network) + selected[5] = PixelReadoutSampledImageStage( + LinearDDIMImageSampler(3), + modes=["train"], + prediction_key=IMAGE_PORT, + gradient_policy="stop", + ) + del selected[4] # Sampling does not consume the low-noise corruption branch. + graph = Pipeline(selected, stage_ids=MODEL_IDS) + before = {key: value.clone() for key, value in graph.state_dict().items()} + batch = fixture_batch() + torch.manual_seed(41) + expected = baseline.execute(batch, mode="train") + torch.manual_seed(41) + actual = graph.execute(batch, mode="train") + unknown = batch["world"]["image"]["valid"] & ~batch["world"]["image"]["observed"] + assert not torch.equal(actual[IMAGE_PORT][unknown], expected[IMAGE_PORT][unknown]) + assert not actual[IMAGE_PORT].requires_grad + assert "pixel_readout/noisy_image" not in actual + reference = network( + batch["world"], + "robot", + actual[IMAGE_PORT], + torch.zeros_like(batch["world"]["image"]["time"]), + phase="B", + ) + for stream in ("proprio", "action"): + torch.testing.assert_close( + actual["pixel_readout/prediction"][stream], + reference[stream], + rtol=0, + atol=0, + ) + assert all( + torch.equal(value, before[key]) for key, value in graph.state_dict().items() + ) + sum_losses(actual).backward() + assert network.action_out["robot"].weight.grad is not None + + +def test_sampled_provider_declares_stopped_gradients_even_for_joint_model(): + network = model() + PixelReadoutModelStage(network, training_mode="joint") + provider = PixelReadoutSampledImageStage( + LinearDDIMImageSampler(3), modes=["train", "inference"], gradient_policy="stop" + ) + batch = {**fixture_batch(), "pixel_video/denoiser": network} + result = provider.execute(batch, mode="train") + assert not result[IMAGE_PORT].requires_grad + assert all(parameter.grad is None for parameter in network.parameters()) + assert not provider.train_only and not provider.inference_only + with pytest.raises(ValueError, match="gradient_policy=stop"): + PixelReadoutSampledImageStage( + LinearDDIMImageSampler(3), gradient_policy="through_sampler" + ) + with pytest.raises(ValueError, match="modes"): + PixelReadoutSampledImageStage(LinearDDIMImageSampler(3), modes="train") + owner = PixelReadoutModelStage(network) + assert owner.contract("train") == ((), ("pixel_video/denoiser",)) + assert owner.contract("inference") == owner.contract("train") + readout = PixelStateActionReadoutStage(image_key=IMAGE_PORT) + assert readout.contract("train")[0] == readout.contract("inference")[0] diff --git a/tests/test_pixel_rgb_frame_reads.py b/tests/test_pixel_rgb_frame_reads.py new file mode 100644 index 000000000..3eea094f9 --- /dev/null +++ b/tests/test_pixel_rgb_frame_reads.py @@ -0,0 +1,144 @@ +"""Selected HDF5 reads preserve complete samples and deterministic draw state.""" + +import json + +import h5py +import numpy as np +import pytest +import torch + +import egomimic.rldb.pixel_proprio as pixels +from tests.test_pixel_proprio_data import digest, pixel_manifest + + +@pytest.mark.parametrize("indices", [[1, 4, 9], [9, 1, 4, 1], [], [0]]) +def test_frame_chunk_reads_keep_exact_order_bytes_and_two_cameras( + tmp_path, monkeypatch, indices +): + rgb = np.random.default_rng(2).integers( + 0, 256, size=(10, 2, 160, 160, 3), dtype=np.uint8 + ) + path = tmp_path / "rgb.h5" + with h5py.File(path, "w") as handle: + handle.create_dataset( + "rgb", data=rgb, chunks=(1, 1, 160, 160, 3), compression="lzf" + ) + getitem = h5py.Dataset.__getitem__ + reads = [] + + def basic_reads_only(dataset, key, *args, **kwargs): + assert isinstance( + key, int + ), "RGB must not use HDF5 fancy selection or load all frames" + reads.append(key) + return getitem(dataset, key, *args, **kwargs) + + monkeypatch.setattr(h5py.Dataset, "__getitem__", basic_reads_only) + indices = np.asarray(indices, dtype=np.int64) + with h5py.File(path, "r") as handle: + actual = pixels.read_selected_rgb_frames(handle["rgb"], indices) + expected = rgb[indices] + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert actual.tobytes() == expected.tobytes() + assert reads == indices.tolist() + + +def test_npz_members_keep_numpy_indexing_behavior(tmp_path): + rgb = np.arange(12 * 2 * 4 * 4 * 3, dtype=np.uint8).reshape(12, 2, 4, 4, 3) + path = tmp_path / "rgb.npz" + np.savez_compressed(path, rgb=rgb) + with np.load(path, allow_pickle=False) as handle: + for indices in (np.array([9, 2, 9, -1]), np.array([], dtype=int)): + actual = pixels.read_selected_rgb_frames(handle["rgb"], indices) + assert actual.dtype == rgb.dtype + assert actual.tobytes() == rgb[indices].tobytes() + + +def _equal_tree(actual, expected): + if torch.is_tensor(expected): + assert actual.dtype == expected.dtype and torch.equal(actual, expected) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys() + for key in expected: + _equal_tree(actual[key], expected[key]) + elif isinstance(expected, list): + assert len(actual) == len(expected) + for a, e in zip(actual, expected): + _equal_tree(a, e) + else: + assert actual == expected + + +@pytest.mark.parametrize("storage", ["hdf5", "npz"]) +@pytest.mark.parametrize("split", ["train", "val"]) +@pytest.mark.parametrize( + "task,image_only,mode", + [ + ("visual_diffusion", True, "continuous"), + ("visual_diffusion", True, "discrete"), + ("visual_diffusion", False, "continuous"), + ("visual_diffusion", False, "discrete"), + ("visual_state_action", False, "continuous"), + ("visual_state_action", False, "discrete"), + ("direct_action", False, "continuous"), + ], +) +def test_full_samples_match_previous_fancy_read_and_rng( + tmp_path, monkeypatch, storage, mode, split, task, image_only +): + path = pixel_manifest(tmp_path, include_sharpa=False, frames=75) + manifest = json.loads(path.read_text()) + for episode in manifest["episodes"]: + original = path.parent / episode["rgb_path"] + with h5py.File(original, "r") as handle: + rgb, time = handle["rgb"][:], handle["time"][:] + if storage == "npz": + replacement = original.with_suffix(".npz") + np.savez_compressed(replacement, rgb=rgb, time=time) + else: + replacement = original + with h5py.File(replacement, "w") as handle: + handle.create_dataset( + "rgb", data=rgb, chunks=(1, 1, 4, 4, 3), compression="lzf" + ) + handle["time"] = time + episode["rgb_path"], episode["rgb_sha256"] = ( + replacement.name, + digest(replacement), + ) + path.write_text(json.dumps(manifest)) + options = dict( + split=split, + task=task, + image_only=image_only, + mode=mode, + min_frames=2, + max_frames=64, + discrete_frames=16, + stride=4, + sampling_strategy="discrete_window" + if mode == "discrete" + else "full_episode_uniform", + ) + actual, previous = ( + pixels.PreparedPixelProprioDataset(path, **options) for _ in range(2) + ) + for dataset in (actual, previous): + dataset.rng = np.random.default_rng(937) + dataset.worker_seed = torch.initial_seed() + # Both ordinary worker RNG draws and checkpointed (cursor, global slot) + # draws must preserve every image, label, time, camera, mask and metadata. + indices = [0, 37, 74] + if split == "train": + indices += [(0, 7, 0), (37, 7, 9), (74, 7, 31)] + actual_samples = [actual[index] for index in indices] + with monkeypatch.context() as context: + context.setattr(pixels, "read_selected_rgb_frames", lambda rgb, rows: rgb[rows]) + previous_samples = [previous[index] for index in indices] + _equal_tree(actual_samples, previous_samples) + for start in range(0, len(indices), 3): + _equal_tree( + pixels.collate_pixel_proprio(actual_samples[start : start + 3]), + pixels.collate_pixel_proprio(previous_samples[start : start + 3]), + ) + assert actual.rng.bit_generator.state == previous.rng.bit_generator.state diff --git a/tests/test_pixel_token_execution_profiles.py b/tests/test_pixel_token_execution_profiles.py new file mode 100644 index 000000000..8d78afc3f --- /dev/null +++ b/tests/test_pixel_token_execution_profiles.py @@ -0,0 +1,104 @@ +"""Packed execution is opt-in and remains bound across readout initialization.""" + +import copy +from pathlib import Path + +import pytest +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf + +from scripts.world_model.initialize_pixel_readout import ( + resolve_model_config, + validate_transition, +) + +CONFIG_ROOT = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + + +def configuration(stage, arm, profile=None): + overrides = [ + f"+experiment=pixel_proprio/{stage}_{arm}", + "hydra/launcher=basic", + "data.manifest_path=/fixture/manifest.json", + "paths.output_dir=/fixture/output", + "paths.work_dir=/fixture/source", + # Compare whole configs without wall-clock run-name differences between + # independent Hydra compositions. Keep all scientific fields asserted. + "logger.wandb.id=fixture-run", + "++logger.wandb.name=fixture-run", + "+pixel_execution=ddp4_prefetch", + ] + if profile is not None: + overrides.append(f"+world_model_execution={profile}") + if stage == "readout": + overrides.extend( + [ + "pixel_readout.initial_weights=/fixture/expanded.pt", + "pixel_readout.initial_weights_sha256=" + "a" * 64, + ] + ) + with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): + config = compose(config_name="train_zarr_cartesian", overrides=overrides) + return OmegaConf.to_container(config, resolve=True) + + +def denoiser(config): + pipeline = config["model"]["pipeline"] + return pipeline["stages"][pipeline["stage_ids"]["pixel_denoiser"]]["model"] + + +@pytest.mark.parametrize("stage", ["video", "readout"]) +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +@pytest.mark.parametrize("profile", [None, "batched_eager", "regional_compile"]) +def test_existing_profiles_keep_padded_token_execution(stage, arm, profile): + config = configuration(stage, arm, profile) + assert denoiser(config)["token_execution"] == "padded" + + +@pytest.mark.parametrize("stage", ["video", "readout"]) +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_packed_profile_changes_only_token_execution(stage, arm): + baseline = configuration(stage, arm, "regional_compile") + packed = configuration(stage, arm, "packed_regional_compile") + assert denoiser(packed)["token_execution"] == "packed_valid" + assert packed["world_model_execution"]["token_execution"] == "packed_valid" + assert packed["world_model_execution"]["compile_mode"] == "default" + assert packed["world_model_execution"]["compile_dynamic"] is True + assert packed["world_model_execution"]["loss_reduction"] == "mean_per_example" + normalized = copy.deepcopy(packed) + denoiser(normalized)["token_execution"] = "padded" + normalized["world_model_execution"].pop("token_execution") + assert normalized == baseline + + +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_old_checkpoint_config_omission_resolves_to_padded(arm): + source = denoiser(configuration("video", arm, "regional_compile")) + target = denoiser(configuration("readout", arm, "regional_compile")) + del source["token_execution"] + assert resolve_model_config(source)["token_execution"] == "padded" + resolved_source, resolved_target = validate_transition(source, target) + assert ( + resolved_source["token_execution"] + == resolved_target["token_execution"] + == "padded" + ) + + +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_readout_initialization_requires_matching_token_execution(arm): + source = denoiser(configuration("video", arm, "packed_regional_compile")) + target = denoiser(configuration("readout", arm, "packed_regional_compile")) + left, right = validate_transition(source, target) + assert left["token_execution"] == right["token_execution"] == "packed_valid" + for changed in (source, target): + changed["token_execution"] = "padded" + with pytest.raises(ValueError, match="configuration mismatch.*token_execution"): + validate_transition(source, target) + changed["token_execution"] = "packed_valid" + + +@pytest.mark.parametrize("arm", ["continuous", "discrete"]) +def test_deferred_readout_does_not_inherit_full_video_budget(arm): + assert configuration("video", arm)["trainer"]["max_steps"] == 300000 + assert configuration("readout", arm)["trainer"]["max_steps"] == 1000 diff --git a/tests/test_pixel_video_graph.py b/tests/test_pixel_video_graph.py new file mode 100644 index 000000000..60566dc06 --- /dev/null +++ b/tests/test_pixel_video_graph.py @@ -0,0 +1,246 @@ +"""Real diffusion graph contracts, shared weights, and legacy numerical parity.""" + +import copy +from pathlib import Path + +import pytest +import torch +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf + +from egomimic.models.image_proprio import ImageProprioModel +from egomimic.models.pixel_diffusion import ( + LinearDDIMImageSampler, + LinearGaussianImageNoising, + PlainPixelMSE, +) +from egomimic.pipeline.core import Pipeline, sum_losses +from egomimic.pipeline.stages_pixel_video import ( + PixelVideoDenoiserStage, + PixelVideoLossStage, + PixelVideoNoisingStage, + PixelVideoSamplerStage, +) +from tools.config_graph import build_graph + +IDS = { + "pixel_noising": 0, + "pixel_denoiser": 1, + "pixel_loss": 2, + "pixel_sampler": 3, + "image_proprio": 1, +} + + +@pytest.fixture(autouse=True) +def single_cpu_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def fixture_batch(): + image = { + "value": torch.linspace(-1, 1, 2 * 4 * 12).reshape(2, 4, 12), + "time": torch.tensor([[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 0.0]]), + "index": torch.tensor([[0, 0, 1, 1], [0, 0, 1, 0]]), + "valid": torch.tensor([[True, True, True, True], [True, True, True, False]]), + "observed": torch.tensor( + [[True, True, False, False], [True, False, False, False]] + ), + "xy": torch.tensor([[[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [1.0, 0.0]]]).repeat( + 2, 1, 1 + ), + "camera": torch.tensor([[0, 1, 0, 1], [0, 1, 0, 0]]), + } + image["value"][1, -1] = float("nan") + return {"world": {"image": image}, "embodiment": ["robot", "robot"]} + + +def small_model(mode="continuous", reduction="mean_per_example"): + return ImageProprioModel( + objective="visual_only", + phase_a_condition_on_proprio=False, + image_dim=12, + width=16, + depth=1, + heads=2, + mlp_ratio=2, + max_cameras=2, + max_tokens=32, + time_mode=mode, + loss_reduction=reduction, + ) + + +def pipeline(model, *, loss=None, steps=3, prediction_key="image_proprio_prediction"): + return Pipeline( + [ + PixelVideoNoisingStage(LinearGaussianImageNoising()), + PixelVideoDenoiserStage(model), + PixelVideoLossStage(loss or PlainPixelMSE(reduction=model.loss_reduction)), + PixelVideoSamplerStage(LinearDDIMImageSampler(steps), prediction_key), + ], + stage_ids=IDS, + ) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("reduction", ["token_mean", "mean_per_example"]) +def test_split_training_matches_legacy_loss_gradients_and_rng(mode, reduction): + model = small_model(mode, reduction) + legacy = copy.deepcopy(model) + graph = pipeline(model) + batch = fixture_batch() + torch.manual_seed(19) + expected = legacy.training_losses(batch["world"], "robot")["total"] + expected_rng = torch.get_rng_state().clone() + expected.backward() + torch.manual_seed(19) + output = graph.execute(batch, mode="train") + actual = sum_losses(output) + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert torch.equal(torch.get_rng_state(), expected_rng) + actual.backward() + for (name, actual_parameter), (expected_name, expected_parameter) in zip( + model.named_parameters(), legacy.named_parameters() + ): + assert name == expected_name + if expected_parameter.grad is None: + assert actual_parameter.grad is None + else: + torch.testing.assert_close( + actual_parameter.grad, expected_parameter.grad, rtol=0, atol=0 + ) + assert "optimizer_step" not in batch + assert "image_proprio_prediction" not in output + assert output["pixel_video/denoiser"] is model + assert set(batch) == {"world", "embodiment"} + assert not any(key.startswith("pixel_video/") for key in batch) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_sampler_matches_legacy_and_uses_one_registered_model(mode): + model = small_model(mode) + graph = pipeline(model) + batch = fixture_batch() + keys_before = tuple(graph.state_dict()) + assert keys_before and all(key.startswith("stages.1.model.") for key in keys_before) + assert sum(p.numel() for p in graph.parameters()) == sum( + p.numel() for p in model.parameters() + ) + assert graph.stage_by_id("image_proprio") is graph.stage_by_id("pixel_denoiser") + torch.manual_seed(23) + expected = model.sample(batch["world"], "robot", steps=3)["image"] + expected_rng = torch.get_rng_state().clone() + torch.manual_seed(23) + output = graph.execute(batch, mode="inference") + actual = output["image_proprio_prediction"]["image"] + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert torch.equal(torch.get_rng_state(), expected_rng) + image = batch["world"]["image"] + assert torch.equal(actual[image["observed"]], image["value"][image["observed"]]) + assert not actual[~image["valid"]].any() + assert tuple(graph.state_dict()) == keys_before + assert not any(key.startswith("loss/") for key in output) + restored = pipeline(small_model(mode)) + restored.load_state_dict(graph.state_dict(), strict=True) + for key, value in restored.state_dict().items(): + assert torch.equal(value, graph.state_dict()[key]) + + +def test_loss_is_replaceable_and_receives_noise_time_and_masks(): + class RecordingObjective: + def __call__(self, prediction, target, mask, **kwargs): + self.received = prediction, target, mask, kwargs + return PlainPixelMSE("mean_per_example")(prediction, target, mask, **kwargs) + + loss = RecordingObjective() + graph = pipeline(small_model(), loss=loss) + output = graph.execute(fixture_batch(), mode="train") + prediction, target, mask, kwargs = loss.received + assert prediction is output["pixel_video/prediction"] + assert target is output["pixel_video/target"] + assert mask is output["pixel_video/loss_mask"] + assert kwargs["noise_levels"] is output["pixel_video/noise_levels"] + metadata = kwargs["metadata"] + assert metadata is output["pixel_video/noise_metadata"] + assert set(metadata) == {"timestamps", "valid", "observed"} + assert torch.equal(mask, metadata["valid"] & ~metadata["observed"]) + assert not kwargs["noise_levels"][metadata["observed"]].any() + assert kwargs["check_target"] is False + sum_losses(output).backward() + + +def test_sampler_can_publish_intermediate_for_separate_readout(): + graph = pipeline(small_model(), prediction_key="pixel_video/sampled_prediction") + output = graph.execute(fixture_batch(), mode="inference") + assert "image_proprio_prediction" not in output + assert set(output["pixel_video/sampled_prediction"]) == {"image"} + assert graph.stage_by_id("pixel_sampler").contract("inference")[1] == ( + "pixel_video/sampled_prediction", + ) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_actual_recipe_fresh_canonical_graph_has_real_diffusion_edges(tmp_path, mode): + root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + with initialize_config_dir(config_dir=str(root), version_base=None): + cfg = compose( + config_name="train_zarr_cartesian", + overrides=[ + f"+experiment=pixel_proprio/video_{mode}", + "hydra/launcher=basic", + "data.manifest_path=/fixture/manifest.json", + "paths.output_dir=/fixture/output", + "paths.work_dir=/fixture/source", + "+pixel_execution=ddp4_prefetch", + "+world_model_execution=regional_compile", + ], + ) + plain = OmegaConf.to_container(cfg, resolve=True) + assert plain["model"]["pipeline"]["stage_ids"] == IDS + stages = plain["model"]["pipeline"]["stages"] + assert all( + "warmup_steps" not in stage and "pose_ramp_steps" not in stage + for stage in stages + ) + assert stages[2]["loss"]["reduction"] == "mean_per_example" + assert stages[3]["sampler"]["steps"] == 20 + assert stages[1]["model"]["max_tokens"] == (3200 if mode == "discrete" else 25600) + # The canonical builder's generic seed inventory needs only actual dataset + # outputs. Stage-one no longer reads injected optimizer-step metadata. + plain["data"]["train_datasets"] = { + "libero_gripper": {"batch_keys": ["world", "embodiment"]} + } + config_path = tmp_path / f"{mode}.yaml" + OmegaConf.save(OmegaConf.create(plain), config_path) + with torch.device("meta"): + training = build_graph(config_path, mode="train") + inference = build_graph(config_path, mode="inference") + assert training["lint"] == inference["lint"] == [] + assert [node["t"] for node in training["nodes"]] == [ + "PixelVideoNoisingStage", + "PixelVideoDenoiserStage", + "PixelVideoLossStage", + ] + assert [node["t"] for node in inference["nodes"]] == [ + "PixelVideoDenoiserStage", + "PixelVideoSamplerStage", + ] + assert {(edge["a"], edge["b"]) for edge in training["edges"]} == { + (0, 1), + (0, 2), + (1, 2), + } + assert inference["edges"] == [ + {"a": 0, "b": 1, "k": "pixel_video/denoiser", "s": "shared"} + ] + assert ( + set(training["seed_keys"]) + == set(inference["seed_keys"]) + == {"world", "embodiment"} + ) diff --git a/tests/test_train_hydra_torch_elastic.py b/tests/test_train_hydra_torch_elastic.py new file mode 100644 index 000000000..3588abfaf --- /dev/null +++ b/tests/test_train_hydra_torch_elastic.py @@ -0,0 +1,97 @@ +"""Explicit torchrun policy must supersede inherited one-task Slurm metadata.""" + +import pytest +from lightning.pytorch.plugins.environments import ( + SLURMEnvironment, + TorchElasticEnvironment, +) +from omegaconf import OmegaConf + +from egomimic.trainHydra import ( + _instantiate_slurm_callbacks, + _instantiate_trainer_plugins, +) + + +@pytest.fixture +def elastic(monkeypatch): + for key in ("ICE_REQUEUE_OWNER", "ICE_CHILD_REQUEUE_DISABLED"): + monkeypatch.delenv(key, raising=False) + for key, value in { + "TORCHELASTIC_RUN_ID": "separate-run-rendezvous", + "RANK": "2", "LOCAL_RANK": "2", "WORLD_SIZE": "4", + "LOCAL_WORLD_SIZE": "4", "GROUP_RANK": "0", + "MASTER_ADDR": "localhost", "MASTER_PORT": "29451", + "SLURM_JOB_ID": "1234", "SLURM_NTASKS": "1", + "SLURM_PROCID": "0", "SLURM_LOCALID": "0", + }.items(): + monkeypatch.setenv(key, value) + return OmegaConf.create({ + "mmap_checkpoint": False, + "runtime": {"cluster_environment": "torch_elastic", + "slurm_requeue_owner": "none"}, + }) + + +def test_torchrun_overrides_one_task_slurm_without_mutating_provenance(elastic): + import os + + before = dict(os.environ) + plugins = _instantiate_trainer_plugins(elastic) + assert len(plugins) == 1 + plugin = plugins[0] + assert isinstance(plugin, TorchElasticEnvironment) + assert not isinstance(plugin, SLURMEnvironment) + assert (plugin.world_size(), plugin.global_rank(), plugin.local_rank()) == (4, 2, 2) + assert plugin.main_port == 29451 + plugin.validate_settings(num_devices=4, num_nodes=1) + assert _instantiate_slurm_callbacks(elastic) == [] + assert dict(os.environ) == before + + +@pytest.mark.parametrize("key,value", [ + ("RANK", None), ("LOCAL_RANK", "four"), ("WORLD_SIZE", "0"), + ("LOCAL_WORLD_SIZE", "3"), ("RANK", "4"), ("RANK", "-1"), + ("LOCAL_RANK", "4"), ("GROUP_RANK", "1"), + ("MASTER_ADDR", ""), ("MASTER_PORT", "0"), + ("MASTER_PORT", "65536"), ("MASTER_PORT", "invalid"), +]) +def test_malformed_torchrun_environment_fails(elastic, monkeypatch, key, value): + if value is None: + monkeypatch.delenv(key, raising=False) + else: + monkeypatch.setenv(key, value) + with pytest.raises(ValueError, match="torch_elastic"): + _instantiate_trainer_plugins(elastic) + + +def test_explicit_elastic_requires_actual_torchrun(elastic, monkeypatch): + monkeypatch.delenv("TORCHELASTIC_RUN_ID") + with pytest.raises(RuntimeError, match="launched by torchrun"): + _instantiate_trainer_plugins(elastic) + + +@pytest.mark.parametrize("owner", ["lightning", "runner"]) +def test_elastic_requires_no_slurm_requeue_owner(elastic, owner): + elastic.runtime.slurm_requeue_owner = owner + with pytest.raises(ValueError, match="slurm_requeue_owner=none"): + _instantiate_trainer_plugins(elastic) + + +def test_elastic_rejects_inherited_runner_ownership(elastic, monkeypatch): + monkeypatch.setenv("ICE_REQUEUE_OWNER", "runner") + with pytest.raises(RuntimeError, match="inconsistent Slurm requeue ownership"): + _instantiate_trainer_plugins(elastic) + + +def test_unknown_environment_fails(elastic): + elastic.runtime.cluster_environment = "torchrun_typo" + with pytest.raises(ValueError, match="cluster_environment"): + _instantiate_trainer_plugins(elastic) + + +def test_default_still_uses_slurm(elastic): + del elastic.runtime.cluster_environment + plugins = _instantiate_trainer_plugins(elastic) + assert len(plugins) == 1 and isinstance(plugins[0], SLURMEnvironment) + assert plugins[0].world_size() == 1 diff --git a/tests/test_world_model_distributed_rng.py b/tests/test_world_model_distributed_rng.py new file mode 100644 index 000000000..ca4d6bba9 --- /dev/null +++ b/tests/test_world_model_distributed_rng.py @@ -0,0 +1,260 @@ +"""Rank-specific random streams survive a real distributed checkpoint/resume.""" + +import json +import os +import random +import subprocess +import sys +import textwrap +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import egomimic.pl_utils.world_model_behavior as behavior + + +@pytest.fixture(autouse=True) +def preserve_rng(): + state = behavior._capture_rng_state() + yield + behavior._restore_rng_state(state) + + +def test_single_process_checkpoint_keeps_legacy_schema_and_random_streams(): + checkpoint = {} + behavior.WorldModelTrainingBehavior().on_save_checkpoint(checkpoint) + saved = checkpoint["mot_rng_state"] + assert set(saved) == {"torch_cpu", "torch_cuda", "numpy", "python"} + expected = (torch.rand(3), np.random.rand(3), random.random()) + behavior._restore_rng_state(behavior._rank_rng_state(saved)) + assert torch.equal(torch.rand(3), expected[0]) + assert np.array_equal(np.random.rand(3), expected[1]) + assert random.random() == expected[2] + + +def test_distributed_resume_rejects_legacy_or_changed_world_size(monkeypatch): + local = behavior._capture_rng_state(current_cuda_only=True) + monkeypatch.setattr(behavior, "_distributed_world", lambda: (2, 1)) + with pytest.raises(ValueError, match="per-rank RNG"): + behavior._rank_rng_state(local) + saved = {"schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 4, "ranks": [local] * 4} + with pytest.raises(ValueError, match="saved RNG world size"): + behavior._rank_rng_state(saved) + + +@pytest.mark.parametrize("update,match", [ + ({"schema": "unknown"}, "Unknown"), + ({"ranks": []}, "one state per rank"), + ({"ranks": [{}, {}]}, "current CUDA device"), +]) +def test_corrupt_distributed_rng_fails_closed(monkeypatch, update, match): + local = behavior._capture_rng_state(current_cuda_only=True) + monkeypatch.setattr(behavior, "_distributed_world", lambda: (2, 1)) + saved = {"schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 2, "ranks": [local, local]} + saved.update(update) + with pytest.raises(ValueError, match=match): + behavior._rank_rng_state(saved) + + +def test_cuda_capture_and_restore_touch_only_current_device(monkeypatch): + state = torch.get_rng_state() + restored = [] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_rng_state", lambda: state) + monkeypatch.setattr(torch.cuda, "get_rng_state_all", lambda: pytest.fail("all devices")) + monkeypatch.setattr(torch.cuda, "set_rng_state", lambda value: restored.append(value)) + monkeypatch.setattr(torch.cuda, "set_rng_state_all", lambda _: pytest.fail("all devices")) + saved = behavior._capture_rng_state(current_cuda_only=True) + assert saved["cuda_scope"] == "current_device" + behavior._restore_rng_state(saved) + assert len(restored) == 1 and torch.equal(restored[0], state) + + +class _Context: + global_step = 0 + trainer = SimpleNamespace(datamodule=None) + + def _default_training_step(self, batch, batch_idx): + return (float(torch.rand(())), float(np.random.rand()), random.random()) + + +def _bound_behavior(seed): + context = _Context() + instance = behavior.WorldModelTrainingBehavior(distributed_noise_seed=seed) + instance.bind(context) + return instance, context # Retain the behavior's weakly referenced context. + + +def test_fresh_training_seeds_ranks_once_without_changing_initialization(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + first_draws = [] + for rank in range(4): + monkeypatch.setattr(behavior, "_distributed_world", lambda: (4, rank)) + initial_state = torch.get_rng_state().clone() + instance, context = _bound_behavior(42) + assert torch.equal(torch.get_rng_state(), initial_state) + generator = torch.Generator().manual_seed(42 + rank) + numpy_rng, python_rng = np.random.RandomState(42 + rank), random.Random(42 + rank) + for step in range(2): + expected = (float(torch.rand((), generator=generator)), + float(numpy_rng.rand()), python_rng.random()) + actual = instance.training_step({"probe": {}}, step) + assert actual == expected + if step == 0: + first_draws.append(actual) + del context + assert len(set(first_draws)) == 4 + + +def test_saved_rank_rng_takes_precedence_over_fresh_noise_seed(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(behavior, "_distributed_world", lambda: (2, 1)) + torch.manual_seed(987) + np.random.seed(987) + random.seed(987) + local = behavior._capture_rng_state(current_cuda_only=True) + expected = _Context()._default_training_step({}, 0) + instance, context = _bound_behavior(42) + instance.on_load_checkpoint({"mot_rng_state": { + "schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 2, "ranks": [local, local], + }}) + assert instance.training_step({"probe": {}}, 0) == expected + assert instance._pending_rng is None and instance._training_rng_initialized + del context + + +def test_fresh_seed_targets_current_cuda_device_only(monkeypatch): + seeds = [] + monkeypatch.setattr(behavior, "_distributed_world", lambda: (4, 2)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "manual_seed", seeds.append) + monkeypatch.setattr(torch.cuda, "manual_seed_all", lambda _: pytest.fail("all devices")) + instance, context = _bound_behavior(42) + instance.training_step({"probe": {}}, 0) + instance.training_step({"probe": {}}, 1) + assert seeds == [44] + del context + + +def test_default_does_not_reseed_legacy_training(monkeypatch): + monkeypatch.setattr(behavior, "_distributed_world", lambda: (4, 2)) + state = behavior._capture_rng_state() + expected = _Context()._default_training_step({}, 0) + behavior._restore_rng_state(state) + instance, context = _bound_behavior(None) + assert instance.training_step({"probe": {}}, 0) == expected + del context + + +@pytest.mark.parametrize("seed", [-1, 2**32, True, "42"]) +def test_distributed_noise_seed_rejects_invalid_values(seed): + with pytest.raises(ValueError, match="distributed_noise_seed"): + behavior.WorldModelTrainingBehavior(distributed_noise_seed=seed) + + +_WORKER = r''' +import json +import os +from pathlib import Path +import random +import sys +import numpy as np +import torch +from lightning.pytorch import LightningModule, Trainer +from lightning.pytorch.strategies import DDPStrategy +from torch.utils.data import DataLoader, TensorDataset +from omegaconf import OmegaConf +from egomimic.pl_utils.world_model_behavior import WorldModelTrainingBehavior +from egomimic.trainHydra import _instantiate_trainer_plugins + +class Probe(LightningModule): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(0.4)) + self.behavior = WorldModelTrainingBehavior(distributed_noise_seed=42) + self.behavior.bind(self) + self.draws = [] + def on_train_start(self): + # The ordinary entrypoint starts all ranks from the same global seed. + seed = 1234 + torch.manual_seed(seed) + np.random.seed(seed) + random.seed(seed) + def training_step(self, batch, idx): + return self.behavior.training_step({'probe': {'value': batch[0]}}, idx) + def _default_training_step(self, batch, idx): + draws = [float(torch.rand(())), float(np.random.rand()), random.random()] + self.draws.append({'step': self.global_step, 'draws': draws}) + return (self.weight - sum(draws)).square() + def configure_optimizers(self): + return torch.optim.AdamW(self.parameters(), lr=0.01) + def on_save_checkpoint(self, checkpoint): + self.behavior.on_save_checkpoint(checkpoint) + def on_load_checkpoint(self, checkpoint): + self.behavior.on_load_checkpoint(checkpoint) + +out, mode = Path(sys.argv[1]), sys.argv[2] +cfg = OmegaConf.create({'mmap_checkpoint': False, 'runtime': { + 'cluster_environment': 'torch_elastic', 'slurm_requeue_owner': 'none'}}) +model = Probe() +trainer = Trainer(accelerator='cpu', devices=2, num_nodes=1, + strategy=DDPStrategy(process_group_backend='gloo', find_unused_parameters=False), + plugins=_instantiate_trainer_plugins(cfg), logger=False, + enable_checkpointing=False, enable_progress_bar=False, enable_model_summary=False, + max_steps=2 if mode == 'partial' else 4, max_epochs=-1, + num_sanity_val_steps=0, limit_val_batches=0, default_root_dir=out) +loader = DataLoader(TensorDataset(torch.ones(8, 1)), batch_size=1, shuffle=False) +trainer.fit(model, train_dataloaders=loader, + ckpt_path=str(out / 'partial.ckpt') if mode == 'resume' else None, + weights_only=False) +trainer.save_checkpoint(out / (mode + '.ckpt')) +(out / f'{mode}-rank{trainer.global_rank}.json').write_text(json.dumps(model.draws)) +''' + + +def test_real_torchrun_checkpoint_resume_preserves_each_rank(tmp_path): + if not torch.distributed.is_available() or not torch.distributed.is_gloo_available(): + pytest.skip("CPU distributed backend unavailable") + worker = tmp_path / "worker.py" + worker.write_text(textwrap.dedent(_WORKER)) + env = dict(os.environ) + env.update({ + "PYTHONPATH": str(Path(__file__).resolve().parents[1]), + "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", "CUDA_VISIBLE_DEVICES": "", + # A one-task Slurm parent must not override torchrun's actual world. + "SLURM_JOB_ID": "999", "SLURM_NTASKS": "1", + "SLURM_PROCID": "0", "SLURM_LOCALID": "0", + }) + for key in ("ICE_REQUEUE_OWNER", "ICE_CHILD_REQUEUE_DISABLED"): + env.pop(key, None) + for mode in ("full", "partial", "resume"): + result = subprocess.run([ + sys.executable, "-m", "torch.distributed.run", "--standalone", + "--nnodes=1", "--nproc-per-node=2", str(worker), str(tmp_path), mode, + ], env=env, capture_output=True, text=True, timeout=90) + assert result.returncode == 0, result.stdout + result.stderr + first = [] + for rank in range(2): + full = json.loads((tmp_path / f"full-rank{rank}.json").read_text()) + partial = json.loads((tmp_path / f"partial-rank{rank}.json").read_text()) + resume = json.loads((tmp_path / f"resume-rank{rank}.json").read_text()) + assert full == partial + resume + first.append(full[0]["draws"]) + assert first[0] != first[1] + full = torch.load(tmp_path / "full.ckpt", weights_only=False) + resume = torch.load(tmp_path / "resume.ckpt", weights_only=False) + assert full["global_step"] == resume["global_step"] == 4 + assert torch.equal(full["state_dict"]["weight"], resume["state_dict"]["weight"]) + for name, value in full["optimizer_states"][0]["state"][0].items(): + assert torch.equal(value, resume["optimizer_states"][0]["state"][0][name]) + saved = resume["mot_rng_state"] + assert saved["schema"] == behavior._DISTRIBUTED_RNG_SCHEMA + assert saved["world_size"] == len(saved["ranks"]) == 2 + assert not torch.equal(saved["ranks"][0]["torch_cpu"], saved["ranks"][1]["torch_cpu"]) From 0e4baebd87051bd43692600d7b0eb0bbf0b3c1fc Mon Sep 17 00:00:00 2001 From: ElmoPA Date: Sun, 18 Oct 2026 14:10:54 -0400 Subject: [PATCH 2/4] Format pixel training changes with the pinned CI formatter --- egomimic/eval/image_proprio_eval.py | 43 +- egomimic/models/attention_execution.py | 214 +++-- egomimic/models/compilation.py | 4 +- egomimic/models/image_proprio.py | 183 ++-- egomimic/models/world_model.py | 144 +++- egomimic/pipeline/stages_image_proprio.py | 4 +- egomimic/pl_utils/world_model_behavior.py | 19 +- egomimic/trainHydra.py | 19 +- .../world_model/benchmark_pixel_execution.py | 804 +++++++++++++----- tests/test_attention_execution.py | 235 +++-- tests/test_image_proprio_adaln.py | 75 +- tests/test_image_proprio_attention.py | 49 +- tests/test_image_proprio_distributed.py | 41 +- tests/test_image_proprio_image_only.py | 51 +- ...image_proprio_image_only_initialization.py | 81 +- tests/test_image_proprio_integration.py | 4 +- tests/test_image_proprio_mixed_precision.py | 31 +- tests/test_image_proprio_no_phase.py | 47 +- tests/test_image_proprio_positions.py | 100 ++- tests/test_pixel_attention_profiles.py | 34 +- tests/test_pixel_compile_ddp.py | 36 +- tests/test_pixel_execution_benchmark.py | 209 +++-- tests/test_pixel_prediction_diagnostics.py | 9 +- tests/test_train_hydra_torch_elastic.py | 54 +- tests/test_world_model_distributed_rng.py | 124 ++- 25 files changed, 1916 insertions(+), 698 deletions(-) diff --git a/egomimic/eval/image_proprio_eval.py b/egomimic/eval/image_proprio_eval.py index a14904704..f2cddb195 100644 --- a/egomimic/eval/image_proprio_eval.py +++ b/egomimic/eval/image_proprio_eval.py @@ -82,13 +82,21 @@ def _sampling_index(batch, batch_idx): if all(value is None for value in indexed): return int(batch_idx) if any(value is None for value in indexed): - raise ValueError("Global validation batch identities must cover every source") + raise ValueError( + "Global validation batch identities must cover every source" + ) indices = [] for item, values in zip(batch.values(), indexed): - if not isinstance(values, (tuple, list)) or len(values) != len(item["task_kind"]): - raise ValueError("One global validation batch identity is required per sample") + if not isinstance(values, (tuple, list)) or len(values) != len( + item["task_kind"] + ): + raise ValueError( + "One global validation batch identity is required per sample" + ) if any(type(value) is not int or value < 0 for value in values): - raise ValueError("Global validation batch identities must be nonnegative integers") + raise ValueError( + "Global validation batch identities must be nonnegative integers" + ) indices.extend(values) if len(set(indices)) != 1: raise ValueError("A validation batch must share one global batch identity") @@ -100,7 +108,8 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0): ) with torch.random.fork_rng(devices=devices): torch.manual_seed( - self.sample_seed + self._sampling_index(batch, batch_idx) + self.sample_seed + + self._sampling_index(batch, batch_idx) + 100000 * dataloader_idx ) output = self.model.forward_eval(batch) @@ -162,12 +171,16 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0): def _reduce_validation_totals(self): """Sum target-weighted errors, including keys absent on some ranks.""" self.global_validation_batches = self._validation_batches - if not (torch.distributed.is_available() and torch.distributed.is_initialized()): + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): return self.totals keys_by_rank = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(keys_by_rank, sorted(self.totals)) keys = sorted({key for rank_keys in keys_by_rank for key in rank_keys}) - totals = torch.zeros(len(keys), 2, dtype=torch.float64, device=self.model.device) + totals = torch.zeros( + len(keys), 2, dtype=torch.float64, device=self.model.device + ) for index, key in enumerate(keys): if key in self.totals: numerator, denominator = self.totals[key] @@ -176,13 +189,19 @@ def _reduce_validation_totals(self): valid = (torch.isfinite(totals).all() & (totals >= 0).all()).to(torch.int64) torch.distributed.all_reduce(valid, op=torch.distributed.ReduceOp.MIN) if not valid.item(): - raise ValueError("Invalid accumulated validation metric on a distributed rank") + raise ValueError( + "Invalid accumulated validation metric on a distributed rank" + ) if keys: torch.distributed.all_reduce(totals, op=torch.distributed.ReduceOp.SUM) - batches = torch.tensor(self._validation_batches, dtype=torch.int64, device=self.model.device) + batches = torch.tensor( + self._validation_batches, dtype=torch.int64, device=self.model.device + ) torch.distributed.all_reduce(batches, op=torch.distributed.ReduceOp.SUM) self.global_validation_batches = int(batches.item()) - return {key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys)} + return { + key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys) + } def log_epoch_metrics(self): self._metrics_logged = False @@ -236,7 +255,9 @@ def on_validation_end(self): and any(denominator > 0 for _, denominator in self._global_totals.values()) ) if torch.distributed.is_available() and torch.distributed.is_initialized(): - success = torch.tensor(int(complete), dtype=torch.int64, device=self.model.device) + success = torch.tensor( + int(complete), dtype=torch.int64, device=self.model.device + ) torch.distributed.all_reduce(success, op=torch.distributed.ReduceOp.MIN) complete = bool(success.item()) if publication_error is not None: diff --git a/egomimic/models/attention_execution.py b/egomimic/models/attention_execution.py index eb056c6a2..596870947 100644 --- a/egomimic/models/attention_execution.py +++ b/egomimic/models/attention_execution.py @@ -30,7 +30,9 @@ class AttentionLayout: max_seqlen_k: int = 0 -def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_valid=None): +def prepare_attention_layout( + valid, *, query_valid=None, query_tokens=None, all_valid=None +): """Prepare dense or packed attention without assuming prefix-only padding. ``query_valid=None`` retains every query, including for query chunks or @@ -44,8 +46,12 @@ def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_ if batch < 1 or key_tokens < 1: raise ValueError("Attention requires nonempty batch and key dimensions") if query_valid is not None: - if (query_valid.ndim != 2 or query_valid.dtype != torch.bool - or query_valid.shape[0] != batch or query_valid.device != valid.device): + if ( + query_valid.ndim != 2 + or query_valid.dtype != torch.bool + or query_valid.shape[0] != batch + or query_valid.device != valid.device + ): raise ValueError("Query validity must be boolean [B,Q] on the key device") if query_tokens is not None and query_tokens != query_valid.shape[1]: raise ValueError("Query token count differs from query validity") @@ -65,7 +71,8 @@ def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_ return AttentionLayout("dense", batch, query_tokens, key_tokens, key_mask) key_counts = valid.sum(-1, dtype=torch.int32) query_counts = ( - query_valid.sum(-1, dtype=torch.int32) if query_valid is not None + query_valid.sum(-1, dtype=torch.int32) + if query_valid is not None else torch.full((batch,), query_tokens, device=valid.device, dtype=torch.int32) ) # This one eager transfer also verifies the optional kernel preconditions. @@ -73,16 +80,26 @@ def prepare_attention_layout(valid, *, query_valid=None, query_tokens=None, all_ if min(counts[1]) < 1: raise ValueError("Fused attention requires at least one valid key per example") if min(counts[0]) < 1: - raise ValueError("Packed attention requires at least one valid query per example") + raise ValueError( + "Packed attention requires at least one valid query per example" + ) zero = key_counts.new_zeros(1) return AttentionLayout( - "varlen", batch, query_tokens, key_tokens, key_mask, - query_indices=(query_valid.reshape(-1).nonzero(as_tuple=True)[0] - if query_valid is not None else None), + "varlen", + batch, + query_tokens, + key_tokens, + key_mask, + query_indices=( + query_valid.reshape(-1).nonzero(as_tuple=True)[0] + if query_valid is not None + else None + ), key_indices=valid.reshape(-1).nonzero(as_tuple=True)[0], cu_seqlens_q=torch.cat((zero, query_counts.cumsum(0, dtype=torch.int32))), cu_seqlens_k=torch.cat((zero, key_counts.cumsum(0, dtype=torch.int32))), - max_seqlen_q=max(counts[0]), max_seqlen_k=max(counts[1]), + max_seqlen_q=max(counts[0]), + max_seqlen_k=max(counts[1]), ) @@ -91,9 +108,11 @@ def prepare_temporal_rope(times, head_dim): if times.ndim != 2 or type(head_dim) is not int or head_dim < 2 or head_dim % 2: raise ValueError("RoPE requires [B,N] times and a positive even head width") half = head_dim // 2 - freq = torch.exp(-math.log(10000) * torch.arange( - half, device=times.device, dtype=torch.float32 - ) / half) + freq = torch.exp( + -math.log(10000) + * torch.arange(half, device=times.device, dtype=torch.float32) + / half + ) phase = times.float()[:, None, :, None] * freq return phase.cos(), phase.sin() @@ -101,22 +120,40 @@ def prepare_temporal_rope(times, head_dim): def resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size): """Validate time/x/y channel allocation and explicit image patch grid H/W.""" if type(head_dim) is not int or head_dim < 6 or head_dim % 2: - raise ValueError("Spatial-temporal RoPE requires an even head width of at least 6") + raise ValueError( + "Spatial-temporal RoPE requires an even head width of at least 6" + ) if axis_dims is None: if head_dim % 8: - raise ValueError("Automatic RoPE axis dimensions require head width divisible by 8") + raise ValueError( + "Automatic RoPE axis dimensions require head width divisible by 8" + ) axis_dims = (head_dim // 2, head_dim // 4, head_dim // 4) - if (not isinstance(axis_dims, Sequence) or isinstance(axis_dims, (str, bytes)) or len(axis_dims) != 3 - or any(type(value) is not int or value < 2 or value % 2 for value in axis_dims) - or sum(axis_dims) != head_dim): - raise ValueError("rope_axis_dims must be positive-even time/x/y dimensions summing to head width") - if (not isinstance(spatial_grid_size, Sequence) or isinstance(spatial_grid_size, (str, bytes)) or len(spatial_grid_size) != 2 - or any(type(value) is not int or value < 1 for value in spatial_grid_size)): - raise ValueError("spatial_grid_size must explicitly name positive integer patch-grid height/width") + if ( + not isinstance(axis_dims, Sequence) + or isinstance(axis_dims, (str, bytes)) + or len(axis_dims) != 3 + or any(type(value) is not int or value < 2 or value % 2 for value in axis_dims) + or sum(axis_dims) != head_dim + ): + raise ValueError( + "rope_axis_dims must be positive-even time/x/y dimensions summing to head width" + ) + if ( + not isinstance(spatial_grid_size, Sequence) + or isinstance(spatial_grid_size, (str, bytes)) + or len(spatial_grid_size) != 2 + or any(type(value) is not int or value < 1 for value in spatial_grid_size) + ): + raise ValueError( + "spatial_grid_size must explicitly name positive integer patch-grid height/width" + ) return tuple(axis_dims), tuple(spatial_grid_size) -def prepare_spatiotemporal_rope(times, xy, head_dim, *, axis_dims=None, spatial_grid_size=None): +def prepare_spatiotemporal_rope( + times, xy, head_dim, *, axis_dims=None, spatial_grid_size=None +): """FP32 axial rotations for time plus normalized image-patch center XY. Normalized centers scale by(grid_width/2,grid_height/2), producing centered @@ -129,21 +166,33 @@ def prepare_spatiotemporal_rope(times, xy, head_dim, *, axis_dims=None, spatial_ concatenates axis frequency vectors in that order. A64/32/32 allocation means32 time,16 x and16 y *pairs*, not a contiguous-slice rotation. """ - axes, grid = resolve_spatiotemporal_rope_geometry(head_dim, axis_dims, spatial_grid_size) + axes, grid = resolve_spatiotemporal_rope_geometry( + head_dim, axis_dims, spatial_grid_size + ) if times.ndim != 2 or times.dtype == torch.bool or times.is_complex(): raise ValueError("Spatial-temporal RoPE requires real-valued times[B,N]") - if xy is not None and (xy.shape != (*times.shape, 2) or xy.device != times.device - or not xy.is_floating_point()): - raise ValueError("Spatial RoPE coordinates must be floating-point[B,N,2] on the time device") + if xy is not None and ( + xy.shape != (*times.shape, 2) + or xy.device != times.device + or not xy.is_floating_point() + ): + raise ValueError( + "Spatial RoPE coordinates must be floating-point[B,N,2] on the time device" + ) positions = [times.float()] - positions.extend((torch.zeros_like(times, dtype=torch.float32),) * 2 if xy is None else - (xy[..., 0].float() * (grid[1] / 2), xy[..., 1].float() * (grid[0] / 2))) + positions.extend( + (torch.zeros_like(times, dtype=torch.float32),) * 2 + if xy is None + else (xy[..., 0].float() * (grid[1] / 2), xy[..., 1].float() * (grid[0] / 2)) + ) phases = [] for coordinate, dimension in zip(positions, axes): pairs = dimension // 2 - frequency = torch.exp(-math.log(10000) * torch.arange( - pairs, device=times.device, dtype=torch.float32 - ) / pairs) + frequency = torch.exp( + -math.log(10000) + * torch.arange(pairs, device=times.device, dtype=torch.float32) + / pairs + ) phases.append(coordinate[:, None, :, None] * frequency) phase = torch.cat(phases, dim=-1) return phase.cos(), phase.sin() @@ -152,7 +201,9 @@ def prepare_spatiotemporal_rope(times, xy, head_dim, *, axis_dims=None, spatial_ def apply_temporal_rope(x, rope): cos, sin = rope left, right = x.float().chunk(2, dim=-1) - return torch.cat((left * cos - right * sin, left * sin + right * cos), -1).to(x.dtype) + return torch.cat((left * cos - right * sin, left * sin + right * cos), -1).to( + x.dtype + ) @dataclass(frozen=True) @@ -174,10 +225,16 @@ def load_flash_attention_kernels(backend, device): raise RuntimeError(f"{backend} requires Hopper or newer CUDA hardware") if backend == "flash4": from packaging.version import Version + if Version(torch.__version__.split("+")[0]) < Version("2.9"): - raise RuntimeError("flash4 requires the separately verified Torch>=2.9 runtime") - module_names = (("flash_attn.cute",) if backend == "flash4" - else ("flash_attn_3.flash_attn_interface", "flash_attn_interface")) + raise RuntimeError( + "flash4 requires the separately verified Torch>=2.9 runtime" + ) + module_names = ( + ("flash_attn.cute",) + if backend == "flash4" + else ("flash_attn_3.flash_attn_interface", "flash_attn_interface") + ) try: for index, module_name in enumerate(module_names): try: @@ -187,8 +244,10 @@ def load_flash_attention_kernels(backend, device): # Older official FA3 builds expose a top-level alias. Only # missing module/package names permit that compatibility path; # an installed backend with a missing dependency must fail. - if (index == len(module_names) - 1 - or error.name not in {module_name, module_name.split(".")[0]}): + if index == len(module_names) - 1 or error.name not in { + module_name, + module_name.split(".")[0], + }: raise dense, varlen = module.flash_attn_func, module.flash_attn_varlen_func except (ImportError, AttributeError) as error: @@ -206,8 +265,12 @@ def _output_tensor(result, expected): if len(result) != 2: raise RuntimeError("Unsupported fused-attention result tuple") result = result[0] - if (not torch.is_tensor(result) or result.shape != expected.shape - or result.dtype != expected.dtype or result.device != expected.device): + if ( + not torch.is_tensor(result) + or result.shape != expected.shape + or result.dtype != expected.dtype + or result.device != expected.device + ): raise RuntimeError("Unsupported fused-attention output shape/dtype/device") return result @@ -217,10 +280,15 @@ def execute_flash_attention(query, key, value, layout, kernels): if any(x.dtype != torch.bfloat16 for x in (query, key, value)): raise ValueError("Fused attention requires bfloat16 Q/K/V") b, h, nq, d = query.shape - if (b != layout.batch_size or nq != layout.query_tokens - or key.shape != (b, h, layout.key_tokens, d) or value.shape != key.shape - or key.device != query.device or value.device != query.device - or layout.key_mask.device != query.device): + if ( + b != layout.batch_size + or nq != layout.query_tokens + or key.shape != (b, h, layout.key_tokens, d) + or value.shape != key.shape + or key.device != query.device + or value.device != query.device + or layout.key_mask.device != query.device + ): raise ValueError("Prepared attention layout differs from Q/K/V shape/device") q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value)) if layout.kind == "dense": @@ -233,11 +301,19 @@ def execute_flash_attention(query, key, value, layout, kernels): q = q.index_select(0, layout.query_indices) k = k.reshape(-1, h, d).index_select(0, layout.key_indices) v = v.reshape(-1, h, d).index_select(0, layout.key_indices) - out = _output_tensor(kernels.varlen( - q, k, v, cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, - max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, - causal=False, - ), q) + out = _output_tensor( + kernels.varlen( + q, + k, + v, + cu_seqlens_q=layout.cu_seqlens_q, + cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, + max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), + q, + ) if layout.query_indices is not None: out = out.new_zeros(b * nq, h, d).index_copy(0, layout.query_indices, out) return out.reshape(b, nq, h, d).transpose(1, 2) @@ -253,16 +329,32 @@ def execute_packed_flash_attention(query, key, value, layout, kernels): raise ValueError("Direct packed attention requires a varlen layout") if any(x.ndim != 3 or x.dtype != torch.bfloat16 for x in (query, key, value)): raise ValueError("Packed fused attention requires bfloat16 [tokens,H,D] Q/K/V") - query_count = (layout.query_indices.numel() if layout.query_indices is not None - else layout.batch_size * layout.query_tokens) - if (query.shape[0] != query_count - or key.shape != (layout.key_indices.numel(), *query.shape[1:]) - or value.shape != key.shape or key.device != query.device - or value.device != query.device or layout.key_mask.device != query.device): - raise ValueError("Prepared packed attention layout differs from Q/K/V shape/device") - return _output_tensor(kernels.varlen( - query.contiguous(), key.contiguous(), value.contiguous(), - cu_seqlens_q=layout.cu_seqlens_q, cu_seqlens_k=layout.cu_seqlens_k, - max_seqlen_q=layout.max_seqlen_q, max_seqlen_k=layout.max_seqlen_k, - causal=False, - ), query) + query_count = ( + layout.query_indices.numel() + if layout.query_indices is not None + else layout.batch_size * layout.query_tokens + ) + if ( + query.shape[0] != query_count + or key.shape != (layout.key_indices.numel(), *query.shape[1:]) + or value.shape != key.shape + or key.device != query.device + or value.device != query.device + or layout.key_mask.device != query.device + ): + raise ValueError( + "Prepared packed attention layout differs from Q/K/V shape/device" + ) + return _output_tensor( + kernels.varlen( + query.contiguous(), + key.contiguous(), + value.contiguous(), + cu_seqlens_q=layout.cu_seqlens_q, + cu_seqlens_k=layout.cu_seqlens_k, + max_seqlen_q=layout.max_seqlen_q, + max_seqlen_k=layout.max_seqlen_k, + causal=False, + ), + query, + ) diff --git a/egomimic/models/compilation.py b/egomimic/models/compilation.py index a7ff7d458..1b2c17194 100644 --- a/egomimic/models/compilation.py +++ b/egomimic/models/compilation.py @@ -16,7 +16,9 @@ def compile_region(function, *, mode="none", dynamic=False, backend="inductor"): if mode == "none": return function if torch._dynamo.config.suppress_errors: - raise RuntimeError("Compilation requires suppress_errors=False; no eager fallback") + raise RuntimeError( + "Compilation requires suppress_errors=False; no eager fallback" + ) compiled = torch.compile( function, backend=backend, mode=mode, dynamic=dynamic, fullgraph=True ) diff --git a/egomimic/models/image_proprio.py b/egomimic/models/image_proprio.py index 909aef38d..678f03b2b 100644 --- a/egomimic/models/image_proprio.py +++ b/egomimic/models/image_proprio.py @@ -207,17 +207,23 @@ def __init__( raise ValueError("Direct action BC has no diffusion noise conditioning") self.noise_conditioning = noise_conditioning if position_encoding not in ("linear_xy_temporal_rope", "spatiotemporal_rope"): - raise ValueError("position_encoding must be linear_xy_temporal_rope or spatiotemporal_rope") + raise ValueError( + "position_encoding must be linear_xy_temporal_rope or spatiotemporal_rope" + ) self.position_encoding = position_encoding if position_encoding == "spatiotemporal_rope": if width % heads: raise ValueError("Model width must be divisible by attention heads") - self.rope_axis_dims, self.spatial_grid_size = resolve_spatiotemporal_rope_geometry( - width // heads, rope_axis_dims, spatial_grid_size + self.rope_axis_dims, self.spatial_grid_size = ( + resolve_spatiotemporal_rope_geometry( + width // heads, rope_axis_dims, spatial_grid_size + ) ) else: if rope_axis_dims is not None or spatial_grid_size is not None: - raise ValueError("Spatial RoPE geometry requires position_encoding='spatiotemporal_rope'") + raise ValueError( + "Spatial RoPE geometry requires position_encoding='spatiotemporal_rope'" + ) self.rope_axis_dims = self.spatial_grid_size = None self.image_dim, self.width = image_dim, width self.allow_patch_compression = allow_patch_compression @@ -276,7 +282,8 @@ def __init__( ) self.noise_level = nn.Sequential( nn.Linear(256 if noise_conditioning == "adaln_zero" else 1, width), - nn.SiLU(), nn.Linear(width, width) + nn.SiLU(), + nn.Linear(width, width), ) if use_phase_embedding: self.phase = nn.Embedding(1 if objective == "visual_only" else 2, width) @@ -297,10 +304,13 @@ def __init__( ) for _ in range(depth) ) - self.norm = nn.ModuleList([ - nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) - if noise_conditioning == "adaln_zero" else nn.LayerNorm(width) - ]) + self.norm = nn.ModuleList( + [ + nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) + if noise_conditioning == "adaln_zero" + else nn.LayerNorm(width) + ] + ) if self.has_proprio_branch: self.norm.append( nn.Identity() if objective == "visual_only" else nn.LayerNorm(width) @@ -505,8 +515,11 @@ def forward( def _prepare_rope(self, times, xy=None): if self.position_encoding == "spatiotemporal_rope": return prepare_spatiotemporal_rope( - times, xy, self.blocks[0].head_dim, - axis_dims=self.rope_axis_dims, spatial_grid_size=self.spatial_grid_size, + times, + xy, + self.blocks[0].head_dim, + axis_dims=self.rope_axis_dims, + spatial_grid_size=self.spatial_grid_size, ) return prepare_temporal_rope(times, self.blocks[0].head_dim) @@ -524,10 +537,16 @@ def _prepare_execution(self, times, valid, image_xy=None): load_flash_attention_kernels(backend, joint_valid.device) layout = prepare_attention_layout(joint_valid, query_valid=joint_valid) if self.position_encoding == "spatiotemporal_rope" and image_xy is None: - raise ValueError("Image spatial-temporal attention requires patch coordinates") - safe_xy = None if image_xy is None else torch.where(valid[0][..., None], image_xy, 0) - rope = tuple(self._prepare_rope(t, safe_xy if index == 0 else None) - for index, t in enumerate(times)) + raise ValueError( + "Image spatial-temporal attention requires patch coordinates" + ) + safe_xy = ( + None if image_xy is None else torch.where(valid[0][..., None], image_xy, 0) + ) + rope = tuple( + self._prepare_rope(t, safe_xy if index == 0 else None) + for index, t in enumerate(times) + ) return layout, rope @staticmethod @@ -545,39 +564,65 @@ def _prepare_noise_frame_map(image): positions = valid[batch_index].nonzero(as_tuple=False).flatten() if positions.numel() == 0: continue - keys = torch.stack((image["time"][batch_index, positions].double(), - image["observed"][batch_index, positions].double()), dim=1) + keys = torch.stack( + ( + image["time"][batch_index, positions].double(), + image["observed"][batch_index, positions].double(), + ), + dim=1, + ) unique, inverse = torch.unique(keys, dim=0, return_inverse=True) - first = torch.full((unique.shape[0],), positions.numel(), device=positions.device, dtype=torch.long) - first.scatter_reduce_(0, inverse, torch.arange(positions.numel(), device=positions.device), reduce="amin") + first = torch.full( + (unique.shape[0],), + positions.numel(), + device=positions.device, + dtype=torch.long, + ) + first.scatter_reduce_( + 0, + inverse, + torch.arange(positions.numel(), device=positions.device), + reduce="amin", + ) representatives.append(batch_index * valid.shape[1] + positions[first]) groups[batch_index, positions] = inverse + offset offset += unique.shape[0] - indices = (torch.cat(representatives) if representatives else - torch.zeros(1, device=valid.device, dtype=torch.long)) + indices = ( + torch.cat(representatives) + if representatives + else torch.zeros(1, device=valid.device, dtype=torch.long) + ) return _NoiseFrameMap(indices, groups) def _embed_noise_frames(self, levels): # Embedding convention only: tau remains the original linear-corruption # strength in [0,1]. Multiplying by1000 does not discretize/change it. - frequency = torch.exp(-math.log(10000) * torch.arange( - 128, device=levels.device, dtype=torch.float32 - ) / 128) + frequency = torch.exp( + -math.log(10000) + * torch.arange(128, device=levels.device, dtype=torch.float32) + / 128 + ) angles = (1000 * levels.float())[:, None] * frequency[None] return self.noise_level(torch.cat((angles.cos(), angles.sin()), dim=-1)) def _image_noise_condition(self, image, levels, frame_map): effective = levels.masked_fill(image["observed"] | ~image["valid"], 0) grouped = effective.reshape(-1).index_select(0, frame_map.representatives) - reconstructed = grouped.index_select(0, frame_map.token_groups.reshape(-1)).reshape_as(effective) + reconstructed = grouped.index_select( + 0, frame_map.token_groups.reshape(-1) + ).reshape_as(effective) if ((effective != reconstructed) & image["valid"]).any(): - raise ValueError("Image noise levels must agree within each timestamp and observed/query group") + raise ValueError( + "Image noise levels must agree within each timestamp and observed/query group" + ) return self._embed_noise_frames_execute(grouped), frame_map.token_groups def _prepare_image_execution(self, image, valid): execution = self._prepare_image_attention_execution(image, valid) if self.noise_conditioning == "adaln_zero": - return _ConditionedImageExecution(execution, self._prepare_noise_frame_map(image)) + return _ConditionedImageExecution( + execution, self._prepare_noise_frame_map(image) + ) return execution def _prepare_image_attention_execution(self, image, valid): @@ -601,7 +646,11 @@ def _prepare_image_attention_execution(self, image, valid): if any(block.attention_query_chunk_size for block in self.blocks): raise ValueError("Packed token execution requires unchunked attention") indices = layout.key_indices - metadata_keys = ("xy", "observed", "camera") if self.use_camera_embedding else ("xy", "observed") + metadata_keys = ( + ("xy", "observed", "camera") + if self.use_camera_embedding + else ("xy", "observed") + ) packed_image = { key: image[key].flatten(0, 1).index_select(0, indices).unsqueeze(0) for key in metadata_keys @@ -609,17 +658,24 @@ def _prepare_image_attention_execution(self, image, valid): packed_times = times.reshape(-1).index_select(0, indices).unsqueeze(0) packed_valid = valid.new_ones((1, indices.numel())) return _PackedImageExecution( - layout, packed_image, packed_valid, + layout, + packed_image, + packed_valid, self._prepare_rope(packed_times, packed_image["xy"]), ) - def _forward_packed_image(self, values, levels, embodiment, execution, image_condition=None): + def _forward_packed_image( + self, values, levels, embodiment, execution, image_condition=None + ): indices = execution.layout.key_indices packed_values = values.flatten(0, 1).index_select(0, indices).unsqueeze(0) packed_levels = levels.reshape(-1).index_select(0, indices).unsqueeze(0) if image_condition is not None: condition, frame_groups = image_condition - image_condition = (condition, frame_groups.reshape(-1).index_select(0, indices).unsqueeze(0)) + image_condition = ( + condition, + frame_groups.reshape(-1).index_select(0, indices).unsqueeze(0), + ) stream = self._embed_image_execute( execution.image, packed_values, packed_levels, execution.valid, "A", False ) @@ -630,27 +686,38 @@ def _forward_packed_image(self, values, levels, embodiment, execution, image_con for block in self.blocks: if self.activation_checkpointing and torch.is_grad_enabled(): stream = checkpoint( - block.forward_packed, stream, execution.valid, - attention_layout=execution.layout, rope=execution.rope, + block.forward_packed, + stream, + execution.valid, + attention_layout=execution.layout, + rope=execution.rope, image_condition=image_condition, use_reentrant=False, ) else: stream = block.forward_packed( - stream, execution.valid, - attention_layout=execution.layout, rope=execution.rope, + stream, + execution.valid, + attention_layout=execution.layout, + rope=execution.rope, image_condition=image_condition, ) - predicted = self._decode_image_execute(stream, execution.valid, image_condition).squeeze(0) + predicted = self._decode_image_execute( + stream, execution.valid, image_condition + ).squeeze(0) batch, tokens = execution.layout.batch_size, execution.layout.key_tokens - return predicted.new_zeros(batch * tokens, self.image_dim).index_copy( - 0, indices, predicted - ).reshape(batch, tokens, self.image_dim) + return ( + predicted.new_zeros(batch * tokens, self.image_dim) + .index_copy(0, indices, predicted) + .reshape(batch, tokens, self.image_dim) + ) def _embed_image(self, image, values, levels, image_valid, phase, direct): image_x = self.image_in(torch.where(image_valid[..., None], values, 0)) if self.position_encoding == "linear_xy_temporal_rope": - image_x = image_x + self.xy(torch.where(image_valid[..., None], image["xy"], 0)) + image_x = image_x + self.xy( + torch.where(image_valid[..., None], image["xy"], 0) + ) if self.use_camera_embedding: image_x = image_x + self.camera( image["camera"].masked_fill(~image_valid, 0) @@ -669,7 +736,9 @@ def _decode_image(self, stream, valid, image_condition=None): if self.noise_conditioning == "adaln_zero": condition, frame_groups = image_condition shift, scale = self.final_adaln(condition).chunk(2, dim=-1) - normalized = normalized * (1 + MoTBlock._frame_values(scale, frame_groups)) + MoTBlock._frame_values(shift, frame_groups) + normalized = normalized * ( + 1 + MoTBlock._frame_values(scale, frame_groups) + ) + MoTBlock._frame_values(shift, frame_groups) return self.image_out(normalized).masked_fill(~valid[..., None], 0) def _forward_validated( @@ -701,7 +770,12 @@ def _forward_validated( image["observed"][..., None], image["value"], image_values ) frame_map = None - if self.noise_conditioning == "adaln_zero" and phase == "A" and not include_pose and execution is None: + if ( + self.noise_conditioning == "adaln_zero" + and phase == "A" + and not include_pose + and execution is None + ): execution = self._prepare_image_execution(image, image_valid) if isinstance(execution, _ConditionedImageExecution): frame_map = execution.noise_frames @@ -713,12 +787,17 @@ def _forward_validated( # Phase B retains the existing multi-expert route. Packing here changes # execution only; no token, camera, or per-example boundary is removed. if self.token_execution == "packed_valid" and phase == "A" and not include_pose: - execution = (self._prepare_image_execution(image, image_valid) - if execution is None else execution) + execution = ( + self._prepare_image_execution(image, image_valid) + if execution is None + else execution + ) if isinstance(execution, _PackedImageExecution): - return {"image": self._forward_packed_image( - values, levels, embodiment, execution, image_condition - )} + return { + "image": self._forward_packed_image( + values, levels, embodiment, execution, image_condition + ) + } image_x = self._embed_image_execute( image, values, levels, image_valid, phase, direct ) @@ -785,11 +864,17 @@ def _forward_validated( ) else: streams = block( - streams, times, valid, attention_layout=attention_layout, rope=rope, + streams, + times, + valid, + attention_layout=attention_layout, + rope=rope, image_condition=image_condition, ) image_result = ( - None if direct else self._decode_image_execute(streams[0], image_valid, image_condition) + None + if direct + else self._decode_image_execute(streams[0], image_valid, image_condition) ) if include_pose: streams[1] = self.norm[1](streams[1]) diff --git a/egomimic/models/world_model.py b/egomimic/models/world_model.py index 58679be61..1c5ca3c62 100644 --- a/egomimic/models/world_model.py +++ b/egomimic/models/world_model.py @@ -32,7 +32,10 @@ def masked_mse(pred, target, weight, *, reduction="token_mean", check_target=Tru if reduction not in ("token_mean", "mean_per_example"): raise ValueError("Unknown masked MSE reduction") active = weight > 0 - if check_target and not torch.isfinite(torch.where(active[..., None], target, 0)).all(): + if ( + check_target + and not torch.isfinite(torch.where(active[..., None], target, 0)).all() + ): raise ValueError("Non-finite supervised target") safe_target = torch.where(active[..., None], target, torch.zeros_like(target)) error = (pred.float() - safe_target.float()).square().mean(-1) @@ -88,15 +91,22 @@ def __init__( raise ValueError("RoPE requires an even attention head width") self.heads = heads self.head_dim = width // heads - if type(attention_query_chunk_size) is not int or attention_query_chunk_size < 0: + if ( + type(attention_query_chunk_size) is not int + or attention_query_chunk_size < 0 + ): raise ValueError("attention_query_chunk_size must be a nonnegative integer") self.attention_query_chunk_size = attention_query_chunk_size if attention_backend not in ("sdpa", "efficient", "flash3", "flash4"): - raise ValueError("attention_backend must be sdpa, efficient, flash3 or flash4") + raise ValueError( + "attention_backend must be sdpa, efficient, flash3 or flash4" + ) if attention_dtype not in ("native", "bfloat16"): raise ValueError("attention_dtype must be native or bfloat16") if attention_backend in ("flash3", "flash4") and attention_dtype != "bfloat16": - raise ValueError(f"{attention_backend} requires explicit bfloat16 attention") + raise ValueError( + f"{attention_backend} requires explicit bfloat16 attention" + ) self.attention_backend = attention_backend self.attention_dtype = attention_dtype if type(image_adaln) is not bool: @@ -104,14 +114,16 @@ def __init__( self.image_adaln = image_adaln self.norm1 = nn.ModuleList( nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) - if image_adaln and i == 0 else nn.LayerNorm(width) + if image_adaln and i == 0 + else nn.LayerNorm(width) for i, _ in enumerate(modalities) ) self.qkv = nn.ModuleList(nn.Linear(width, 3 * width) for _ in modalities) self.out = nn.ModuleList(nn.Linear(width, width) for _ in modalities) self.norm2 = nn.ModuleList( nn.LayerNorm(width, elementwise_affine=False, eps=1e-6) - if image_adaln and i == 0 else nn.LayerNorm(width) + if image_adaln and i == 0 + else nn.LayerNorm(width) for i, _ in enumerate(modalities) ) self.ffn = nn.ModuleList( @@ -133,9 +145,13 @@ def configure_compilation(self, mode="none", dynamic=False, backend="inductor"): if self.image_adaln: names += ("_image_modulation",) for name in names: - object.__setattr__(self, name + "_executor", compile_region( - getattr(self, name), mode=mode, dynamic=dynamic, backend=backend - )) + object.__setattr__( + self, + name + "_executor", + compile_region( + getattr(self, name), mode=mode, dynamic=dynamic, backend=backend + ), + ) def _image_modulation(self, condition): # Project once per physical-frame/role group, not once per patch. @@ -147,25 +163,47 @@ def _frame_values(values, frame_indices): *frame_indices.shape, values.shape[-1] ) - def _project_inputs(self, streams, times, rope, image_modulation=None, frame_indices=None): + def _project_inputs( + self, streams, times, rope, image_modulation=None, frame_indices=None + ): qs, ks, vs = [], [], [] for i, x in enumerate(streams): b, n, d = x.shape normalized = self.norm1[i](x) if i == 0 and self.image_adaln: - shift, scale = (self._frame_values(value, frame_indices) for value in image_modulation[:2]) + shift, scale = ( + self._frame_values(value, frame_indices) + for value in image_modulation[:2] + ) normalized = normalized * (1 + scale) + shift q, k, v = ( self.qkv[i](normalized) .reshape(b, n, 3, self.heads, d // self.heads) .permute(2, 0, 3, 1, 4) ) - qs.append(temporal_rope(q, times[i]) if rope is None else apply_temporal_rope(q, rope[i])) - ks.append(temporal_rope(k, times[i]) if rope is None else apply_temporal_rope(k, rope[i])) + qs.append( + temporal_rope(q, times[i]) + if rope is None + else apply_temporal_rope(q, rope[i]) + ) + ks.append( + temporal_rope(k, times[i]) + if rope is None + else apply_temporal_rope(k, rope[i]) + ) vs.append(v) return torch.cat(qs, 2), torch.cat(ks, 2), torch.cat(vs, 2) - def forward(self, streams, times, valid, *, attention_layout=None, rope=None, image_condition=None): + def forward( + self, + streams, + times, + valid, + *, + attention_layout=None, + rope=None, + image_condition=None, + ): sizes = [s.shape[1] for s in streams] modulation = None frame_indices = None @@ -174,11 +212,16 @@ def forward(self, streams, times, valid, *, attention_layout=None, rope=None, im raise ValueError("Image AdaLN requires frame conditioning") condition, frame_indices = image_condition modulation = self._image_modulation_executor(condition) - query, key, value = self._project_inputs_executor(streams, times, rope, modulation, frame_indices) + query, key, value = self._project_inputs_executor( + streams, times, rope, modulation, frame_indices + ) else: query, key, value = self._project_inputs_executor(streams, times, rope) - mask = (torch.cat(valid, 1)[:, None, None, :] if attention_layout is None - else attention_layout.key_mask) + mask = ( + torch.cat(valid, 1)[:, None, None, :] + if attention_layout is None + else attention_layout.key_mask + ) chunk = self.attention_query_chunk_size if chunk and query.shape[2] > chunk: # Split only queries. Every row retains all keys/values and the same @@ -187,29 +230,42 @@ def forward(self, streams, times, valid, *, attention_layout=None, rope=None, im # attention intermediates and losing the intended memory bound. parts = [] for q in query.split(chunk, dim=2): - if torch.is_grad_enabled() and any(t.requires_grad for t in (q, key, value)): - part = checkpoint(self._apply_attention, q, key, value, mask, use_reentrant=False) + if torch.is_grad_enabled() and any( + t.requires_grad for t in (q, key, value) + ): + part = checkpoint( + self._apply_attention, q, key, value, mask, use_reentrant=False + ) else: part = self._apply_attention(q, key, value, mask) parts.append(part) result = torch.cat(parts, dim=2) else: # Keep the historical four-argument call for native SDPA overrides. - result = (self._apply_attention(query, key, value, mask) - if attention_layout is None or self.attention_backend not in ("flash3", "flash4") - else self._apply_attention( - query, key, value, mask, attention_layout=attention_layout)) + result = ( + self._apply_attention(query, key, value, mask) + if attention_layout is None + or self.attention_backend not in ("flash3", "flash4") + else self._apply_attention( + query, key, value, mask, attention_layout=attention_layout + ) + ) if self.image_adaln: - return self._finish_outputs_executor(streams, result, valid, sizes, modulation, frame_indices) + return self._finish_outputs_executor( + streams, result, valid, sizes, modulation, frame_indices + ) return self._finish_outputs_executor(streams, result, valid, sizes) - def _finish_outputs(self, streams, result, valid, sizes, image_modulation=None, frame_indices=None): + def _finish_outputs( + self, streams, result, valid, sizes, image_modulation=None, frame_indices=None + ): result = result.transpose(1, 2).flatten(2).split(sizes, 1) outputs = [] for i, (x, y) in enumerate(zip(streams, result)): if i == 0 and self.image_adaln: gate_attention, shift, scale, gate_mlp = ( - self._frame_values(value, frame_indices) for value in image_modulation[2:] + self._frame_values(value, frame_indices) + for value in image_modulation[2:] ) x = x + gate_attention * self.out[i](y) x = x + gate_mlp * self.ffn[i](self.norm2[i](x) * (1 + scale) + shift) @@ -219,7 +275,9 @@ def _finish_outputs(self, streams, result, valid, sizes, image_modulation=None, outputs.append(x.masked_fill(~valid[i][..., None], 0)) return outputs - def forward_packed(self, stream, valid, *, attention_layout, rope, image_condition=None): + def forward_packed( + self, stream, valid, *, attention_layout, rope, image_condition=None + ): """Run expert zero on packed image tokens, preserving example boundaries. Pure tensor regions reuse the same modules and compilation policy as @@ -228,7 +286,9 @@ def forward_packed(self, stream, valid, *, attention_layout, rope, image_conditi if self.attention_backend not in ("flash3", "flash4"): raise ValueError("Packed token execution requires flash3 or flash4") if self.attention_dtype != "bfloat16": - raise ValueError("Packed token execution requires explicit bfloat16 attention") + raise ValueError( + "Packed token execution requires explicit bfloat16 attention" + ) if self.attention_query_chunk_size: raise ValueError("Packed token execution requires unchunked attention") if stream.ndim != 3 or stream.shape[0] != 1 or valid.shape != stream.shape[:2]: @@ -240,7 +300,9 @@ def forward_packed(self, stream, valid, *, attention_layout, rope, image_conditi raise ValueError("Image AdaLN requires frame conditioning") condition, frame_indices = image_condition modulation = self._image_modulation_executor(condition) - query, key, value = self._project_inputs_executor([stream], None, (rope,), modulation, frame_indices) + query, key, value = self._project_inputs_executor( + [stream], None, (rope,), modulation, frame_indices + ) else: query, key, value = self._project_inputs_executor([stream], None, (rope,)) original_dtype = query.dtype @@ -271,15 +333,20 @@ def _apply_attention(self, query, key, value, mask, *, attention_layout=None): if self.attention_dtype == "bfloat16": query, key, value = (x.to(torch.bfloat16) for x in (query, key, value)) if self.attention_backend in ("flash3", "flash4"): - kernels = load_flash_attention_kernels(self.attention_backend, query.device) + kernels = load_flash_attention_kernels( + self.attention_backend, query.device + ) layout = attention_layout or prepare_attention_layout( mask[:, 0, 0], query_tokens=query.shape[2] ) result = execute_flash_attention(query, key, value, layout, kernels) elif self.attention_backend == "efficient": if query.device.type != "cuda": - raise RuntimeError("efficient attention requires CUDA; no math fallback") + raise RuntimeError( + "efficient attention requires CUDA; no math fallback" + ) from torch.nn.attention import SDPBackend, sdpa_kernel + with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): result = self._attend(query, key, value, mask) else: @@ -342,7 +409,9 @@ def __init__( if type(strict_fp32) is not bool: raise ValueError("strict_fp32 must be a boolean") if strict_fp32 and attention_dtype != "native": - raise ValueError("Explicit BF16 attention conflicts with strict_fp32; select the mixed-attention execution profile") + raise ValueError( + "Explicit BF16 attention conflicts with strict_fp32; select the mixed-attention execution profile" + ) if image_prediction_type not in ("v", "x0"): raise ValueError("image_prediction_type must be v or x0") self.image_prediction_type = image_prediction_type @@ -387,9 +456,14 @@ def __init__( nn.Linear(2, width), nn.SiLU(), nn.Linear(width, width) ) self.blocks = nn.ModuleList( - MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size, - attention_backend=attention_backend, - attention_dtype=attention_dtype) + MoTBlock( + width, + heads, + mlp_ratio, + attention_query_chunk_size, + attention_backend=attention_backend, + attention_dtype=attention_dtype, + ) for _ in range(depth) ) # One private image adapter per embodiment per block. Zero disables the diff --git a/egomimic/pipeline/stages_image_proprio.py b/egomimic/pipeline/stages_image_proprio.py index 3d9fa1e5a..a2cc4b37b 100644 --- a/egomimic/pipeline/stages_image_proprio.py +++ b/egomimic/pipeline/stages_image_proprio.py @@ -73,8 +73,8 @@ def execute(self, batch, *, mode): image = batch["world"]["image"] batch[f"log/{self.log_prefix}_image_tokens"] = image["valid"].sum().float() batch[f"log/{self.log_prefix}_image_query_tokens"] = ( - image["valid"] & ~image["observed"] - ).sum().float() + (image["valid"] & ~image["observed"]).sum().float() + ) if self.model.objective in ("visual_proprio", "visual_state_action"): batch["log/pose_weight"] = losses["total"].new_tensor(scale) if self.model.objective == "visual_state_action": diff --git a/egomimic/pl_utils/world_model_behavior.py b/egomimic/pl_utils/world_model_behavior.py index 74a79090f..1b61a3f93 100644 --- a/egomimic/pl_utils/world_model_behavior.py +++ b/egomimic/pl_utils/world_model_behavior.py @@ -19,11 +19,16 @@ def _distributed_world(): def _capture_rng_state(*, current_cuda_only=False): cuda = [] if torch.cuda.is_available(): - cuda = ([torch.cuda.get_rng_state()] if current_cuda_only - else torch.cuda.get_rng_state_all()) + cuda = ( + [torch.cuda.get_rng_state()] + if current_cuda_only + else torch.cuda.get_rng_state_all() + ) state = { - "torch_cpu": torch.get_rng_state(), "torch_cuda": cuda, - "numpy": np.random.get_state(), "python": random.getstate(), + "torch_cpu": torch.get_rng_state(), + "torch_cuda": cuda, + "numpy": np.random.get_state(), + "python": random.getstate(), } if current_cuda_only: state["cuda_scope"] = "current_device" @@ -129,6 +134,8 @@ def on_save_checkpoint(self, checkpoint): def on_load_checkpoint(self, checkpoint): self._pending_rng = checkpoint.get("mot_rng_state") - if (getattr(self.context.trainer.datamodule, "checkpointable_sampling", False) - and self._pending_rng is None): + if ( + getattr(self.context.trainer.datamodule, "checkpointable_sampling", False) + and self._pending_rng is None + ): raise ValueError("Exact MoT resume requires saved model RNG state") diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 648fe201c..4768bf385 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -59,7 +59,9 @@ def _configure_strict_fp32(cfg: DictConfig) -> None: if enabled and str(precision) not in {"32", "32-true"}: raise ValueError("runtime.strict_fp32 requires trainer.precision=32-true") if os.environ.get("NVIDIA_TF32_OVERRIDE") != "0": - raise ValueError("FP32/disable_tf32 backend policy requires NVIDIA_TF32_OVERRIDE=0") + raise ValueError( + "FP32/disable_tf32 backend policy requires NVIDIA_TF32_OVERRIDE=0" + ) torch.set_float32_matmul_precision("highest") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False @@ -152,7 +154,8 @@ def _instantiate_slurm_callbacks(cfg: DictConfig) -> List[Callback]: checkpoint_dir=str(checkpoint_dir), save_signal=signal.SIGUSR2, train_batch_only=OmegaConf.select( - cfg, "runtime.slurm_signal_train_batch_only", default=False), + cfg, "runtime.slurm_signal_train_batch_only", default=False + ), ) ] @@ -195,10 +198,14 @@ def _torch_elastic_environment(cfg: DictConfig) -> TorchElasticEnvironment: except ValueError as exc: raise ValueError(f"torch_elastic requires integer {name}") from exc world, local_world = integers["WORLD_SIZE"], integers["LOCAL_WORLD_SIZE"] - if (world < 1 or local_world < 1 or world % local_world - or not 0 <= integers["RANK"] < world - or not 0 <= integers["LOCAL_RANK"] < local_world - or not 0 <= integers["GROUP_RANK"] < world // local_world): + if ( + world < 1 + or local_world < 1 + or world % local_world + or not 0 <= integers["RANK"] < world + or not 0 <= integers["LOCAL_RANK"] < local_world + or not 0 <= integers["GROUP_RANK"] < world // local_world + ): raise ValueError("torch_elastic has inconsistent world size or rank variables") if not os.environ.get("MASTER_ADDR"): raise ValueError("torch_elastic requires MASTER_ADDR") diff --git a/scripts/world_model/benchmark_pixel_execution.py b/scripts/world_model/benchmark_pixel_execution.py index d809df63a..faa25938f 100644 --- a/scripts/world_model/benchmark_pixel_execution.py +++ b/scripts/world_model/benchmark_pixel_execution.py @@ -53,7 +53,9 @@ def require(condition, message): def write_once(path, value): with Path(path).open("x") as stream: - stream.write(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") + stream.write( + json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n" + ) def digest(path): @@ -70,38 +72,78 @@ def frame_schedule(scenario, batch_size): if scenario != "mixed": raise ValueError(f"Unknown scenario: {scenario}") cycle = len(MIXED_FRAMES) // math.gcd(batch_size, len(MIXED_FRAMES)) - return [tuple(MIXED_FRAMES[(step * batch_size + row) % len(MIXED_FRAMES)] - for row in range(batch_size)) for step in range(cycle)] + return [ + tuple( + MIXED_FRAMES[(step * batch_size + row) % len(MIXED_FRAMES)] + for row in range(batch_size) + ) + for step in range(cycle) + ] def complete_cycles(steps, schedule): return math.ceil(steps / len(schedule)) * len(schedule) -def model_arguments(backend, scenario, *, compile_mode="none", compile_dynamic=False, - token_execution="padded"): - if backend not in BACKENDS or scenario not in SCENARIOS or token_execution not in TOKEN_EXECUTIONS: +def model_arguments( + backend, + scenario, + *, + compile_mode="none", + compile_dynamic=False, + token_execution="padded", +): + if ( + backend not in BACKENDS + or scenario not in SCENARIOS + or token_execution not in TOKEN_EXECUTIONS + ): raise ValueError("Unknown execution candidate") return dict( - objective="visual_only", phase_a_condition_on_proprio=False, - pose_dims=None, action_dims=None, image_dim=768, width=768, - depth=8, heads=6, mlp_ratio=4, max_cameras=2, + objective="visual_only", + phase_a_condition_on_proprio=False, + pose_dims=None, + action_dims=None, + image_dim=768, + width=768, + depth=8, + heads=6, + mlp_ratio=4, + max_cameras=2, max_tokens=3200 if scenario == "fixed16" else 25600, time_mode="discrete" if scenario == "fixed16" else "continuous", - use_camera_embedding=True, use_modality_embedding=False, - use_embodiment_embedding=False, allow_patch_compression=False, - use_phase_embedding=False, noise_conditioning="adaln_zero", - position_encoding="spatiotemporal_rope", rope_axis_dims=(64, 32, 32), + use_camera_embedding=True, + use_modality_embedding=False, + use_embodiment_embedding=False, + allow_patch_compression=False, + use_phase_embedding=False, + noise_conditioning="adaln_zero", + position_encoding="spatiotemporal_rope", + rope_axis_dims=(64, 32, 32), spatial_grid_size=(10, 10), - attention_backend=backend, attention_dtype="bfloat16", strict_fp32=False, - attention_query_chunk_size=0, activation_checkpointing=False, - loss_reduction="mean_per_example", compile_mode=compile_mode, - compile_dynamic=compile_dynamic, compile_backend="inductor", token_execution=token_execution, + attention_backend=backend, + attention_dtype="bfloat16", + strict_fp32=False, + attention_query_chunk_size=0, + activation_checkpointing=False, + loss_reduction="mean_per_example", + compile_mode=compile_mode, + compile_dynamic=compile_dynamic, + compile_backend="inductor", + token_execution=token_execution, ) -def make_batch(frames, *, seed, device, image_size=160, patch_size=16, cameras=2, - frame_seconds=0.05): +def make_batch( + frames, + *, + seed, + device, + image_size=160, + patch_size=16, + cameras=2, + frame_seconds=0.05, +): """True RGB patch geometry; only batch padding is invalid.""" if not frames or any(type(length) is not int or length < 2 for length in frames): raise ValueError("Every synthetic episode needs at least two frames") @@ -124,32 +166,49 @@ def make_batch(frames, *, seed, device, image_size=160, patch_size=16, cameras=2 size = length * per_frame value[row, :size] = torch.rand(size, dimension, generator=generator) * 2 - 1 valid[row, :size] = True - observed[row, :max(1, length // 2) * per_frame] = True + observed[row, : max(1, length // 2) * per_frame] = True indices = torch.arange(length, dtype=torch.float32).repeat_interleave(per_frame) index[row, :size] = indices clock[row, :size] = indices * frame_seconds - return {"image": {key: tensor.to(device) for key, tensor in dict( - value=value, valid=valid, observed=observed, time=clock, index=index, - camera=camera[None].expand(len(frames), -1).clone(), - xy=xy[None].expand(len(frames), -1, -1).clone(), - ).items()}} + return { + "image": { + key: tensor.to(device) + for key, tensor in dict( + value=value, + valid=valid, + observed=observed, + time=clock, + index=index, + camera=camera[None].expand(len(frames), -1).clone(), + xy=xy[None].expand(len(frames), -1, -1).clone(), + ).items() + } + } def single_example(batch, row): indices = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] - return {"image": {key: value[row:row + 1].index_select(1, indices) - for key, value in batch["image"].items()}} + return { + "image": { + key: value[row : row + 1].index_select(1, indices) + for key, value in batch["image"].items() + } + } def fixed_corruption(batch, seed): image = batch["image"] generator = torch.Generator(device=image["value"].device).manual_seed(seed) - noise = torch.randn(image["value"].shape, dtype=torch.float32, - device=image["value"].device, generator=generator) + noise = torch.randn( + image["value"].shape, + dtype=torch.float32, + device=image["value"].device, + generator=generator, + ) levels = (0.15 + 0.7 * torch.remainder(image["index"], 7) / 6).masked_fill( image["observed"] | ~image["valid"], 0 ) - noisy = ((1 - levels[..., None]) * image["value"] + levels[..., None] * noise) + noisy = (1 - levels[..., None]) * image["value"] + levels[..., None] * noise return noisy.masked_fill(~image["valid"][..., None], 0), levels @@ -169,6 +228,7 @@ def execution_context(backend, device): with torch.autocast(device.type, dtype=torch.bfloat16): if backend == "sdpa": from torch.nn.attention import SDPBackend, sdpa_kernel + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): yield else: @@ -177,9 +237,11 @@ def execution_context(backend, device): def relative_error(actual, expected): error = (actual.float() - expected.float()).double() - return dict(squared_error=float(error.square().sum()), - squared_reference=float(expected.double().square().sum()), - max_abs=float(error.abs().max())) + return dict( + squared_error=float(error.square().sum()), + squared_reference=float(expected.double().square().sum()), + max_abs=float(error.abs().max()), + ) @contextmanager @@ -193,12 +255,18 @@ def nontrivial_conditioning_probe(model, seed): if getattr(model, "noise_conditioning", "additive") != "adaln_zero": yield {"nonzero_adaln_probe": False} return - selected = {name: parameter for name, parameter in model.named_parameters() - if name.startswith(("image_out.", "final_adaln.")) - or ".adaln_image." in name} - require(selected and any("adaln_image" in name for name in selected), - "AdaLN correctness probe could not find modulation parameters") - original = {name: parameter.detach().clone() for name, parameter in selected.items()} + selected = { + name: parameter + for name, parameter in model.named_parameters() + if name.startswith(("image_out.", "final_adaln.")) or ".adaln_image." in name + } + require( + selected and any("adaln_image" in name for name in selected), + "AdaLN correctness probe could not find modulation parameters", + ) + original = { + name: parameter.detach().clone() for name, parameter in selected.items() + } device = next(model.parameters()).device devices = [device.index] if device.type == "cuda" else [] try: @@ -209,8 +277,11 @@ def nontrivial_conditioning_probe(model, seed): torch.cuda.default_generators[device.index].manual_seed(seed) for parameter in selected.values(): parameter.normal_(mean=0.0, std=0.02) - yield {"nonzero_adaln_probe": True, "seed": seed, - "restores_training_initialization": True} + yield { + "nonzero_adaln_probe": True, + "seed": seed, + "restores_training_initialization": True, + } finally: with torch.no_grad(): for name, parameter in selected.items(): @@ -226,13 +297,21 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): with execution_context(backend, device): prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] loss = example_mean_loss(prediction, batch) - require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), "Nonfinite/non-FP32 probe loss") + require( + loss.dtype == torch.float32 and torch.isfinite(loss).item(), + "Nonfinite/non-FP32 probe loss", + ) actual = prediction.detach() batched_gradients = None if gradients: loss.backward() - require(all(p.grad is not None for p in model.parameters()), "Unused trainable parameter") - batched_gradients = {name: p.grad.detach().cpu().clone() for name, p in model.named_parameters()} + require( + all(p.grad is not None for p in model.parameters()), + "Unused trainable parameter", + ) + batched_gradients = { + name: p.grad.detach().cpu().clone() for name, p in model.named_parameters() + } del prediction, loss model.zero_grad(set_to_none=True) errors = [] @@ -241,36 +320,59 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): selected = batch["image"]["valid"][row].nonzero(as_tuple=True)[0] independent = single_example(batch, row) with execution_context(backend, device): - predicted = model(independent, "libero_gripper", - noisy[row:row + 1].index_select(1, selected), - levels[row:row + 1].index_select(1, selected), phase="A")["image"] + predicted = model( + independent, + "libero_gripper", + noisy[row : row + 1].index_select(1, selected), + levels[row : row + 1].index_select(1, selected), + phase="A", + )["image"] individual_loss = example_mean_loss(predicted, independent) / batch_size - errors.append(relative_error(actual[row:row + 1].index_select(1, selected), predicted.detach())) + errors.append( + relative_error( + actual[row : row + 1].index_select(1, selected), predicted.detach() + ) + ) if gradients: individual_loss.backward() del predicted, individual_loss - output_relative = math.sqrt(sum(v["squared_error"] for v in errors) / - max(sum(v["squared_reference"] for v in errors), 1e-30)) - require(math.isfinite(output_relative) and output_relative <= tolerance, - f"Batched output differs from independent B1: relative L2={output_relative}") - report = {"examples": batch_size, "output_relative_l2": output_relative, - "output_max_abs": max(v["max_abs"] for v in errors)} + output_relative = math.sqrt( + sum(v["squared_error"] for v in errors) + / max(sum(v["squared_reference"] for v in errors), 1e-30) + ) + require( + math.isfinite(output_relative) and output_relative <= tolerance, + f"Batched output differs from independent B1: relative L2={output_relative}", + ) + report = { + "examples": batch_size, + "output_relative_l2": output_relative, + "output_max_abs": max(v["max_abs"] for v in errors), + } if gradients: error_sum, reference_sum = 0.0, 0.0 for name, parameter in model.named_parameters(): - require(parameter.grad is not None and parameter.grad.dtype == torch.float32, - "Missing/non-FP32 reference gradient") - measured = relative_error(batched_gradients[name], parameter.grad.detach().cpu()) + require( + parameter.grad is not None and parameter.grad.dtype == torch.float32, + "Missing/non-FP32 reference gradient", + ) + measured = relative_error( + batched_gradients[name], parameter.grad.detach().cpu() + ) error_sum += measured["squared_error"] reference_sum += measured["squared_reference"] gradient_relative = math.sqrt(error_sum / max(reference_sum, 1e-30)) - require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, - f"Batched gradient differs from mean independent B1: relative L2={gradient_relative}") + require( + math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Batched gradient differs from mean independent B1: relative L2={gradient_relative}", + ) report["gradient_relative_l2"] = gradient_relative model.zero_grad(set_to_none=True) # Perturb every other example, preserving row zero and its noise exactly. if batch_size > 1: - changed = {"image": {key: value.clone() for key, value in batch["image"].items()}} + changed = { + "image": {key: value.clone() for key, value in batch["image"].items()} + } changed["image"]["value"][1:] *= -0.7 perturbed = noisy.clone() perturbed[1:] *= -0.7 @@ -278,17 +380,29 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): # Inductor may select different kernels for grad-enabled and # no-grad execution. Compare intervention/control under the same # mode and graph; separately bound the mode's BF16 rounding drift. - control = model(batch, "libero_gripper", noisy, levels, phase="A")["image"][0].clone() - independent = model(changed, "libero_gripper", perturbed, levels, phase="A")["image"][0] + control = model(batch, "libero_gripper", noisy, levels, phase="A")["image"][ + 0 + ].clone() + independent = model( + changed, "libero_gripper", perturbed, levels, phase="A" + )["image"][0] mode_error = relative_error(control, actual[0]) - mode_relative = math.sqrt(mode_error["squared_error"] / max(mode_error["squared_reference"], 1e-30)) - require(math.isfinite(mode_relative) and mode_relative <= tolerance, - f"Grad/no-grad output differs beyond BF16 tolerance: relative L2={mode_relative}") + mode_relative = math.sqrt( + mode_error["squared_error"] / max(mode_error["squared_reference"], 1e-30) + ) + require( + math.isfinite(mode_relative) and mode_relative <= tolerance, + f"Grad/no-grad output differs beyond BF16 tolerance: relative L2={mode_relative}", + ) report["grad_vs_no_grad_relative_l2"] = mode_relative error = relative_error(independent, control) - leak = math.sqrt(error["squared_error"] / max(error["squared_reference"], 1e-30)) - require(math.isfinite(leak) and leak <= 1e-6, - f"Cross-example attention leakage detected: same-mode relative L2={leak}") + leak = math.sqrt( + error["squared_error"] / max(error["squared_reference"], 1e-30) + ) + require( + math.isfinite(leak) and leak <= 1e-6, + f"Cross-example attention leakage detected: same-mode relative L2={leak}", + ) report["cross_example_relative_l2"] = leak else: report["cross_example_check"] = "not_applicable_batch_one" @@ -297,55 +411,102 @@ def check_outputs(model, batch, *, backend, seed, tolerance, gradients=False): def check_eager_compiled(eager, compiled, batch, *, backend, seed, tolerance): """Require state-identical eager/compiled AMP forward and FP32 backward parity.""" - require(eager.token_execution == compiled.token_execution, - "Eager/compiled token execution differs") + require( + eager.token_execution == compiled.token_execution, + "Eager/compiled token execution differs", + ) reference_state, candidate_state = eager.state_dict(), compiled.state_dict() - require(reference_state.keys() == candidate_state.keys(), "Eager/compiled state keys differ") - require(all(torch.equal(value, candidate_state[name]) for name, value in reference_state.items()), - "Eager/compiled initial state differs") - require(all(p.dtype == torch.float32 for model in (eager, compiled) for p in model.parameters()), - "Eager/compiled stored parameters must be FP32") + require( + reference_state.keys() == candidate_state.keys(), + "Eager/compiled state keys differ", + ) + require( + all( + torch.equal(value, candidate_state[name]) + for name, value in reference_state.items() + ), + "Eager/compiled initial state differs", + ) + require( + all( + p.dtype == torch.float32 + for model in (eager, compiled) + for p in model.parameters() + ), + "Eager/compiled stored parameters must be FP32", + ) noisy, levels = fixed_corruption(batch, seed) device = batch["image"]["value"].device def forward_backward(model): model.zero_grad(set_to_none=True) with execution_context(backend, device): - prediction = model(batch, "libero_gripper", noisy, levels, phase="A")["image"] + prediction = model(batch, "libero_gripper", noisy, levels, phase="A")[ + "image" + ] loss = example_mean_loss(prediction, batch) - require(loss.dtype == torch.float32 and torch.isfinite(loss).item(), - "Eager/compiled loss must be finite FP32") + require( + loss.dtype == torch.float32 and torch.isfinite(loss).item(), + "Eager/compiled loss must be finite FP32", + ) # Match standard mixed precision: backward is outside forward autocast. loss.backward() - require(all(p.grad is not None and p.grad.dtype == torch.float32 - and torch.isfinite(p.grad).all() for p in model.parameters()), - "Eager/compiled gradient must be present and finite FP32") + require( + all( + p.grad is not None + and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() + for p in model.parameters() + ), + "Eager/compiled gradient must be present and finite FP32", + ) return prediction.detach().cpu(), float(loss.detach()) try: reference_output, reference_loss = forward_backward(eager) - reference_gradients = {name: p.grad.detach().cpu().clone() for name, p in eager.named_parameters()} + reference_gradients = { + name: p.grad.detach().cpu().clone() for name, p in eager.named_parameters() + } eager.zero_grad(set_to_none=True) output, loss = forward_backward(compiled) measured = relative_error(output, reference_output) - output_relative = math.sqrt(measured["squared_error"] / max(measured["squared_reference"], 1e-30)) + output_relative = math.sqrt( + measured["squared_error"] / max(measured["squared_reference"], 1e-30) + ) loss_relative = abs(loss - reference_loss) / max(abs(reference_loss), 1e-30) - require(math.isfinite(output_relative) and output_relative <= tolerance, - f"Compiled output differs from eager: relative L2={output_relative}") - require(math.isfinite(loss_relative) and loss_relative <= tolerance, - f"Compiled loss differs from eager: relative error={loss_relative}") - errors = [relative_error(p.grad.detach().cpu(), reference_gradients[name]) - for name, p in compiled.named_parameters()] - gradient_relative = math.sqrt(sum(value["squared_error"] for value in errors) / - max(sum(value["squared_reference"] for value in errors), 1e-30)) - require(math.isfinite(gradient_relative) and gradient_relative <= tolerance, - f"Compiled gradient differs from eager: relative L2={gradient_relative}") - return {"output_relative_l2": output_relative, "output_max_abs": measured["max_abs"], - "loss_relative_error": loss_relative, "eager_loss": reference_loss, - "compiled_loss": loss, "gradient_relative_l2": gradient_relative, - "parameters_and_gradients_dtype": "float32", "backward_autocast": "off", - "same_state_verified": True, "backend": backend, - "token_execution": compiled.token_execution} + require( + math.isfinite(output_relative) and output_relative <= tolerance, + f"Compiled output differs from eager: relative L2={output_relative}", + ) + require( + math.isfinite(loss_relative) and loss_relative <= tolerance, + f"Compiled loss differs from eager: relative error={loss_relative}", + ) + errors = [ + relative_error(p.grad.detach().cpu(), reference_gradients[name]) + for name, p in compiled.named_parameters() + ] + gradient_relative = math.sqrt( + sum(value["squared_error"] for value in errors) + / max(sum(value["squared_reference"] for value in errors), 1e-30) + ) + require( + math.isfinite(gradient_relative) and gradient_relative <= tolerance, + f"Compiled gradient differs from eager: relative L2={gradient_relative}", + ) + return { + "output_relative_l2": output_relative, + "output_max_abs": measured["max_abs"], + "loss_relative_error": loss_relative, + "eager_loss": reference_loss, + "compiled_loss": loss, + "gradient_relative_l2": gradient_relative, + "parameters_and_gradients_dtype": "float32", + "backward_autocast": "off", + "same_state_verified": True, + "backend": backend, + "token_execution": compiled.token_execution, + } finally: eager.zero_grad(set_to_none=True) compiled.zero_grad(set_to_none=True) @@ -353,7 +514,11 @@ def forward_backward(model): def compiler_counters(): from torch._dynamo.utils import counters - return {group: {str(k): int(v) for k, v in values.items()} for group, values in counters.items()} + + return { + group: {str(k): int(v) for k, v in values.items()} + for group, values in counters.items() + } def counter_total(counters, group, key=None): @@ -362,9 +527,11 @@ def counter_total(counters, group, key=None): def backend_inventory(backend): - module_names = {"flash4": ("flash_attn.cute.interface",), - "flash3": ("flash_attn_3.flash_attn_interface", "flash_attn_interface"), - "sdpa": ("torch",)}[backend] + module_names = { + "flash4": ("flash_attn.cute.interface",), + "flash3": ("flash_attn_3.flash_attn_interface", "flash_attn_interface"), + "sdpa": ("torch",), + }[backend] try: for index, module_name in enumerate(module_names): try: @@ -373,27 +540,45 @@ def backend_inventory(backend): except ModuleNotFoundError as error: # Match the model's official FA3 package/legacy alias support; # an installed backend with a missing dependency must fail. - if (index == len(module_names) - 1 - or error.name not in {module_name, module_name.split(".")[0]}): + if index == len(module_names) - 1 or error.name not in { + module_name, + module_name.split(".")[0], + }: raise except ImportError as error: - raise UnsupportedBackend(f"Required backend import failed: {module_name}: {error}") from error + raise UnsupportedBackend( + f"Required backend import failed: {module_name}: {error}" + ) from error module_path = Path(module.__file__) versions = {} - for name in ("torch", "flash-attn-4", "flash-attn-3", "flash-attn", "nvidia-cutlass-dsl"): + for name in ( + "torch", + "flash-attn-4", + "flash-attn-3", + "flash-attn", + "nvidia-cutlass-dsl", + ): try: versions[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: versions[name] = None - return {"module": module_name, "path": str(module_path), "sha256": digest(module_path), - "versions": versions, "sdpa_kernel": "FLASH_ATTENTION_ONLY" if backend == "sdpa" else None} + return { + "module": module_name, + "path": str(module_path), + "sha256": digest(module_path), + "versions": versions, + "sdpa_kernel": "FLASH_ATTENTION_ONLY" if backend == "sdpa" else None, + } def timed_step(model, batch, optimizer, backend, device): optimizer.zero_grad(set_to_none=True) with execution_context(backend, device): losses = model.training_losses(batch, "libero_gripper", pose_weight=0.0) - require(losses["total"].dtype == losses["image"].dtype == torch.float32, "MSE must remain FP32") + require( + losses["total"].dtype == losses["image"].dtype == torch.float32, + "MSE must remain FP32", + ) losses["total"].backward() optimizer.step() return losses["total"].detach() @@ -403,73 +588,131 @@ def run_candidate(request, destination): from egomimic.models.image_proprio import ImageProprioModel if not torch.cuda.is_available(): - raise UnsupportedBackend("A scheduled H100 CUDA device is required; no CPU benchmark fallback") + raise UnsupportedBackend( + "A scheduled H100 CUDA device is required; no CPU benchmark fallback" + ) device = torch.device("cuda", request["device"]) torch.cuda.set_device(device) properties = torch.cuda.get_device_properties(device) if "H100" not in properties.name: - raise UnsupportedBackend(f"This benchmark binds H100 hardware, got {properties.name}") + raise UnsupportedBackend( + f"This benchmark binds H100 hardware, got {properties.name}" + ) torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision("highest") torch.manual_seed(request["seed"]) inventory = backend_inventory(request["backend"]) - arguments = model_arguments(request["backend"], request["scenario"], - compile_mode=request["compile_mode"], compile_dynamic=request["compile_dynamic"], - token_execution=request.get("token_execution", "padded")) + arguments = model_arguments( + request["backend"], + request["scenario"], + compile_mode=request["compile_mode"], + compile_dynamic=request["compile_dynamic"], + token_execution=request.get("token_execution", "padded"), + ) model = ImageProprioModel(**arguments).to(device).train() - require(sum(p.numel() for p in model.parameters()) == PARAMETERS, "Exact architecture parameter count changed") - require(all(p.dtype == torch.float32 for p in model.parameters()), "Stored parameters must be FP32") + require( + sum(p.numel() for p in model.parameters()) == PARAMETERS, + "Exact architecture parameter count changed", + ) + require( + all(p.dtype == torch.float32 for p in model.parameters()), + "Stored parameters must be FP32", + ) before_counters = compiler_counters() schedule = frame_schedule(request["scenario"], request["batch_size"]) - batches = [make_batch(frames, seed=request["seed"] + index, device=device, - frame_seconds=0.2 if request["scenario"] == "fixed16" else 0.05) - for index, frames in enumerate(schedule)] + batches = [ + make_batch( + frames, + seed=request["seed"] + index, + device=device, + frame_seconds=0.2 if request["scenario"] == "fixed16" else 0.05, + ) + for index, frames in enumerate(schedule) + ] with nontrivial_conditioning_probe(model, request["seed"] + 1000) as probe_record: - probe = make_batch((2, 3) if request["batch_size"] > 1 else (2,), seed=request["seed"], device=device) - correctness = {"bounded_gradient_probe": check_outputs( - model, probe, backend=request["backend"], seed=request["seed"], - tolerance=request["tolerance"], gradients=True)} + probe = make_batch( + (2, 3) if request["batch_size"] > 1 else (2,), + seed=request["seed"], + device=device, + ) + correctness = { + "bounded_gradient_probe": check_outputs( + model, + probe, + backend=request["backend"], + seed=request["seed"], + tolerance=request["tolerance"], + gradients=True, + ) + } if request["compile_mode"] != "none": # This is a separate eager model, not another invocation of the same # compiled callable. Keep synthetic noise RNG unchanged by construction. with torch.random.fork_rng(devices=[device.index]): - eager = ImageProprioModel(**{**arguments, "compile_mode": "none"}).to(device).train() + eager = ( + ImageProprioModel(**{**arguments, "compile_mode": "none"}) + .to(device) + .train() + ) eager.load_state_dict(model.state_dict(), strict=True) correctness["compiled_vs_eager"] = check_eager_compiled( - eager, model, probe, backend=request["backend"], seed=request["seed"], - tolerance=request["tolerance"]) + eager, + model, + probe, + backend=request["backend"], + seed=request["seed"], + tolerance=request["tolerance"], + ) del eager del probe correctness["conditioning_probe"] = probe_record correctness["full_shape_outputs"] = [] with torch.no_grad(): for batch in batches: - correctness["full_shape_outputs"].append(check_outputs( - model, batch, backend=request["backend"], seed=request["seed"], - tolerance=request["tolerance"], gradients=False)) + correctness["full_shape_outputs"].append( + check_outputs( + model, + batch, + backend=request["backend"], + seed=request["seed"], + tolerance=request["tolerance"], + gradients=False, + ) + ) torch.cuda.empty_cache() # Release unused correctness-probe allocator blocks. optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01) warmup = complete_cycles(request["warmup"], schedule) steps = complete_cycles(request["steps"], schedule) warm_start = time.perf_counter() for step in range(warmup): - timed_step(model, batches[step % len(batches)], optimizer, request["backend"], device) + timed_step( + model, batches[step % len(batches)], optimizer, request["backend"], device + ) torch.cuda.synchronize(device) warmup_seconds = time.perf_counter() - warm_start warm_counters = compiler_counters() if request["compile_mode"] != "none": - require(counter_total(warm_counters, "stats", "unique_graphs") > counter_total(before_counters, "stats", "unique_graphs"), - "Requested compilation captured no tensor regions") - require(counter_total(warm_counters, "graph_break") == counter_total(before_counters, "graph_break"), - "A compiled tensor region broke its fullgraph contract") + require( + counter_total(warm_counters, "stats", "unique_graphs") + > counter_total(before_counters, "stats", "unique_graphs"), + "Requested compilation captured no tensor regions", + ) + require( + counter_total(warm_counters, "graph_break") + == counter_total(before_counters, "graph_break"), + "A compiled tensor region broke its fullgraph contract", + ) torch.cuda.reset_peak_memory_stats(device) walls, cuda_times, loss_values = [], [], [] valid_tokens = 0 for step in range(steps): batch = batches[step % len(batches)] torch.cuda.synchronize(device) - start_event, end_event = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start_event, end_event = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) start = time.perf_counter() start_event.record() loss = timed_step(model, batch, optimizer, request["backend"], device) @@ -480,41 +723,92 @@ def run_candidate(request, destination): loss_values.append(float(loss)) valid_tokens += sum(schedule[step % len(schedule)]) * 200 after_counters = compiler_counters() - require(all(math.isfinite(value) for value in loss_values), "Nonfinite measured training loss") - require(counter_total(after_counters, "stats", "unique_graphs") == counter_total(warm_counters, "stats", "unique_graphs"), - "New compiled graphs appeared during steady-state timing") - require(all(p.grad is not None and p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() - for p in model.parameters()), "Missing, nonfinite or non-FP32 measured gradient") - require(all(state[key].dtype == torch.float32 and torch.isfinite(state[key]).all() - for state in optimizer.state.values() for key in ("exp_avg", "exp_avg_sq")), - "Nonfinite or non-FP32 AdamW moments") - require(all(p.dtype == torch.float32 and torch.isfinite(p).all() for p in model.parameters()), - "Nonfinite or non-FP32 optimized parameters") + require( + all(math.isfinite(value) for value in loss_values), + "Nonfinite measured training loss", + ) + require( + counter_total(after_counters, "stats", "unique_graphs") + == counter_total(warm_counters, "stats", "unique_graphs"), + "New compiled graphs appeared during steady-state timing", + ) + require( + all( + p.grad is not None + and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() + for p in model.parameters() + ), + "Missing, nonfinite or non-FP32 measured gradient", + ) + require( + all( + state[key].dtype == torch.float32 and torch.isfinite(state[key]).all() + for state in optimizer.state.values() + for key in ("exp_avg", "exp_avg_sq") + ), + "Nonfinite or non-FP32 AdamW moments", + ) + require( + all( + p.dtype == torch.float32 and torch.isfinite(p).all() + for p in model.parameters() + ), + "Nonfinite or non-FP32 optimized parameters", + ) result = { - "status": "PASS", "candidate": request, "parameters_total_trainable": PARAMETERS, - "model_arguments": arguments, "backend_inventory": inventory, - "gpu": {"name": properties.name, "uuid": str(properties.uuid), "total_memory_bytes": properties.total_memory}, - "correctness": correctness, "frame_schedule": schedule, - "warmup_steps": warmup, "timed_steps": steps, "warmup_seconds": warmup_seconds, - "wall_seconds": walls, "cuda_event_seconds": cuda_times, - "median_wall_seconds": statistics.median(walls), "median_cuda_seconds": statistics.median(cuda_times), + "status": "PASS", + "candidate": request, + "parameters_total_trainable": PARAMETERS, + "model_arguments": arguments, + "backend_inventory": inventory, + "gpu": { + "name": properties.name, + "uuid": str(properties.uuid), + "total_memory_bytes": properties.total_memory, + }, + "correctness": correctness, + "frame_schedule": schedule, + "warmup_steps": warmup, + "timed_steps": steps, + "warmup_seconds": warmup_seconds, + "wall_seconds": walls, + "cuda_event_seconds": cuda_times, + "median_wall_seconds": statistics.median(walls), + "median_cuda_seconds": statistics.median(cuda_times), "examples_per_second": steps * request["batch_size"] / sum(walls), - "valid_tokens_per_second": valid_tokens / sum(walls), "valid_tokens_measured": valid_tokens, + "valid_tokens_per_second": valid_tokens / sum(walls), + "valid_tokens_measured": valid_tokens, "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), - "resident_input_bytes": sum(value.numel() * value.element_size() - for batch in batches for value in batch["image"].values()), - "losses": loss_values, "compiler_counters_before": before_counters, - "compiler_counters_after_warmup": warm_counters, "compiler_counters_after_timing": after_counters, + "resident_input_bytes": sum( + value.numel() * value.element_size() + for batch in batches + for value in batch["image"].values() + ), + "losses": loss_values, + "compiler_counters_before": before_counters, + "compiler_counters_after_warmup": warm_counters, + "compiler_counters_after_timing": after_counters, "compile_scope": "declared_fullgraph_tensor_regions; attention_adapter_and_validation_eager", - "compiled_backward_autocast": "off" if request["compile_mode"] != "none" else None, + "compiled_backward_autocast": "off" + if request["compile_mode"] != "none" + else None, "measurement_scope": "single_gpu_device_resident_synthetic_model; no_loader_H2D_NCCL_or_checkpoint_IO", } if request["profile"]: trace = destination / "trace.json" - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], - record_shapes=True, profile_memory=True) as profiler: - with torch.profiler.record_function("synthetic_device_resident_optimizer_step"): + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + profile_memory=True, + ) as profiler: + with torch.profiler.record_function( + "synthetic_device_resident_optimizer_step" + ): timed_step(model, batches[0], optimizer, request["backend"], device) torch.cuda.synchronize(device) profiler.export_chrome_trace(str(trace)) @@ -532,15 +826,29 @@ def worker(request_path): except UnsupportedBackend as error: result = {"status": "UNSUPPORTED", "error": str(error), "candidate": request} except CorrectnessError as error: - result = {"status": "CORRECTNESS_FAILED", "error": str(error), "candidate": request} + result = { + "status": "CORRECTNESS_FAILED", + "error": str(error), + "candidate": request, + } except Exception as error: # Backend/API/compiler failures are never retried with another kernel. - unsupported = (isinstance(error, RuntimeError) and any( - phrase in str(error) for phrase in ("No available kernel", "No viable backend") - )) or (isinstance(error, ValueError) and - str(error) == "Ragged packed_valid execution requires flash3 or flash4") - result = {"status": "UNSUPPORTED" if unsupported else "ERROR", "error": f"{type(error).__name__}: {error}", - "traceback": traceback.format_exc(), "candidate": request} + unsupported = ( + isinstance(error, RuntimeError) + and any( + phrase in str(error) + for phrase in ("No available kernel", "No viable backend") + ) + ) or ( + isinstance(error, ValueError) + and str(error) == "Ragged packed_valid execution requires flash3 or flash4" + ) + result = { + "status": "UNSUPPORTED" if unsupported else "ERROR", + "error": f"{type(error).__name__}: {error}", + "traceback": traceback.format_exc(), + "candidate": request, + } write_once(request_path.parent / "result.json", result) @@ -553,8 +861,13 @@ def csv_choices(value, choices): def run_worker_process(command, folder, timeout): """Bound compiler subprocesses as well as the GPU-owning worker on timeout.""" - with (folder / "stdout.log").open("x") as out, (folder / "stderr.log").open("x") as err: - process = subprocess.Popen(command, stdout=out, stderr=err, start_new_session=True) + with ( + (folder / "stdout.log").open("x") as out, + (folder / "stderr.log").open("x") as err, + ): + process = subprocess.Popen( + command, stdout=out, stderr=err, start_new_session=True + ) try: return process.wait(timeout=timeout), False except subprocess.TimeoutExpired: @@ -569,30 +882,57 @@ def run_worker_process(command, folder, timeout): def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path) - parser.add_argument("--backends", type=lambda s: csv_choices(s, BACKENDS), default=list(BACKENDS)) - parser.add_argument("--scenarios", type=lambda s: csv_choices(s, SCENARIOS), default=list(SCENARIOS)) - parser.add_argument("--batch-sizes", type=lambda s: [int(v) for v in csv_choices(s, ("1", "2", "4", "8"))], default=[1, 2, 4, 8]) + parser.add_argument( + "--backends", type=lambda s: csv_choices(s, BACKENDS), default=list(BACKENDS) + ) + parser.add_argument( + "--scenarios", type=lambda s: csv_choices(s, SCENARIOS), default=list(SCENARIOS) + ) + parser.add_argument( + "--batch-sizes", + type=lambda s: [int(v) for v in csv_choices(s, ("1", "2", "4", "8"))], + default=[1, 2, 4, 8], + ) parser.add_argument("--compile-mode", choices=COMPILE_MODES, default="none") parser.add_argument("--compile-dynamic", action="store_true") - parser.add_argument("--token-execution", choices=TOKEN_EXECUTIONS, default="padded", - help="Keep padded tensors, or pack valid image tokens across the image-only expert") + parser.add_argument( + "--token-execution", + choices=TOKEN_EXECUTIONS, + default="padded", + help="Keep padded tensors, or pack valid image tokens across the image-only expert", + ) parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--steps", type=int, default=10) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--device", type=int, default=0) parser.add_argument("--tolerance", type=float, default=0.03) parser.add_argument("--worker-timeout", type=int, default=1800) - parser.add_argument("--python-bootstrap", type=Path, - help="Run every child through this reviewed runtime bootstrap (.pth hooks included)") - parser.add_argument("--profile-candidate", help="Exact candidate ID such as flash4-mixed-b4-none") - parser.add_argument("--dry-run", action="store_true", help="Print the immutable candidate plan without CUDA work") + parser.add_argument( + "--python-bootstrap", + type=Path, + help="Run every child through this reviewed runtime bootstrap (.pth hooks included)", + ) + parser.add_argument( + "--profile-candidate", help="Exact candidate ID such as flash4-mixed-b4-none" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the immutable candidate plan without CUDA work", + ) parser.add_argument("--worker-request", type=Path, help=argparse.SUPPRESS) args = parser.parse_args() if args.worker_request: worker(args.worker_request) return - if min(args.warmup, args.steps, args.worker_timeout) < 1 or args.device < 0 or not 0 < args.tolerance < 1: - parser.error("Positive bounded warmup/steps/timeout/tolerance and nonnegative device required") + if ( + min(args.warmup, args.steps, args.worker_timeout) < 1 + or args.device < 0 + or not 0 < args.tolerance < 1 + ): + parser.error( + "Positive bounded warmup/steps/timeout/tolerance and nonnegative device required" + ) candidates = [] for backend in args.backends: for scenario in args.scenarios: @@ -600,16 +940,32 @@ def main(): identity = f"{backend}-{scenario}-b{batch_size}-{args.compile_mode}" if args.token_execution != "padded": identity += f"-{args.token_execution}" - candidates.append(dict(id=identity, backend=backend, scenario=scenario, - batch_size=batch_size, compile_mode=args.compile_mode, - compile_dynamic=args.compile_dynamic, token_execution=args.token_execution, - seed=args.seed, - device=args.device, warmup=args.warmup, steps=args.steps, - tolerance=args.tolerance, profile=args.profile_candidate == identity)) + candidates.append( + dict( + id=identity, + backend=backend, + scenario=scenario, + batch_size=batch_size, + compile_mode=args.compile_mode, + compile_dynamic=args.compile_dynamic, + token_execution=args.token_execution, + seed=args.seed, + device=args.device, + warmup=args.warmup, + steps=args.steps, + tolerance=args.tolerance, + profile=args.profile_candidate == identity, + ) + ) if args.profile_candidate and not any(c["profile"] for c in candidates): parser.error("--profile-candidate must identify an enabled candidate") if args.dry_run: - print(json.dumps({"parameters_total_trainable": PARAMETERS, "candidates": candidates}, indent=2)) + print( + json.dumps( + {"parameters_total_trainable": PARAMETERS, "candidates": candidates}, + indent=2, + ) + ) return if args.output is None: parser.error("--output is required") @@ -618,15 +974,31 @@ def main(): bootstrap_path = args.python_bootstrap.resolve(strict=True) bootstrap = {"path": str(bootstrap_path), "sha256": digest(bootstrap_path)} args.output.mkdir(parents=True, exist_ok=False) - source_files = [Path(__file__), ROOT / "egomimic/models/image_proprio.py", - ROOT / "egomimic/models/pixel_diffusion.py", - ROOT / "egomimic/models/world_model.py", ROOT / "egomimic/models/attention_execution.py", - ROOT / "egomimic/models/compilation.py"] - source = {str(path.relative_to(ROOT)): digest(path) for path in source_files if path.exists()} - write_once(args.output / "plan.json", {"schema": "pixel-execution-benchmark-v1", "candidates": candidates, - "source_sha256": source, "slurm_job_id": os.environ.get("SLURM_JOB_ID"), - "python_bootstrap": bootstrap, "python_executable": sys.executable, - "scope": "synthetic_single_gpu_execution_only; not_training_readiness_or_policy_quality"}) + source_files = [ + Path(__file__), + ROOT / "egomimic/models/image_proprio.py", + ROOT / "egomimic/models/pixel_diffusion.py", + ROOT / "egomimic/models/world_model.py", + ROOT / "egomimic/models/attention_execution.py", + ROOT / "egomimic/models/compilation.py", + ] + source = { + str(path.relative_to(ROOT)): digest(path) + for path in source_files + if path.exists() + } + write_once( + args.output / "plan.json", + { + "schema": "pixel-execution-benchmark-v1", + "candidates": candidates, + "source_sha256": source, + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "python_bootstrap": bootstrap, + "python_executable": sys.executable, + "scope": "synthetic_single_gpu_execution_only; not_training_readiness_or_policy_quality", + }, + ) results = [] for candidate in candidates: folder = args.output / candidate["id"] @@ -636,19 +1008,41 @@ def main(): command = [sys.executable] if args.python_bootstrap: command.append(str(args.python_bootstrap.resolve())) - command += [str(Path(__file__).resolve()), "--worker-request", str(request_path.resolve())] + command += [ + str(Path(__file__).resolve()), + "--worker-request", + str(request_path.resolve()), + ] returncode, timed_out = run_worker_process(command, folder, args.worker_timeout) if timed_out: - result = {"status": "TIMEOUT", "seconds": args.worker_timeout, "candidate": candidate} + result = { + "status": "TIMEOUT", + "seconds": args.worker_timeout, + "candidate": candidate, + } elif returncode != 0 or not (folder / "result.json").exists(): - result = {"status": "PROCESS_ERROR", "returncode": returncode, "candidate": candidate} + result = { + "status": "PROCESS_ERROR", + "returncode": returncode, + "candidate": candidate, + } else: result = json.loads((folder / "result.json").read_text()) results.append(result) - print(json.dumps({"candidate": candidate["id"], "status": result["status"]}), flush=True) - write_once(args.output / "report.json", {"schema": "pixel-execution-benchmark-v1", "source_sha256": source, - "status": "COMPLETED", "results": results, - "scope": "Only PASS candidates have measured and checked synthetic throughput; real-data DDP smoke still required"}) + print( + json.dumps({"candidate": candidate["id"], "status": result["status"]}), + flush=True, + ) + write_once( + args.output / "report.json", + { + "schema": "pixel-execution-benchmark-v1", + "source_sha256": source, + "status": "COMPLETED", + "results": results, + "scope": "Only PASS candidates have measured and checked synthetic throughput; real-data DDP smoke still required", + }, + ) if __name__ == "__main__": diff --git a/tests/test_attention_execution.py b/tests/test_attention_execution.py index f9a58931a..de1e95191 100644 --- a/tests/test_attention_execution.py +++ b/tests/test_attention_execution.py @@ -21,9 +21,15 @@ def _dense(q, k, v, *, causal): assert causal is False - return F.scaled_dot_product_attention( - q.float().transpose(1, 2), k.float().transpose(1, 2), v.float().transpose(1, 2) - ).transpose(1, 2).to(q.dtype) + return ( + F.scaled_dot_product_attention( + q.float().transpose(1, 2), + k.float().transpose(1, 2), + v.float().transpose(1, 2), + ) + .transpose(1, 2) + .to(q.dtype) + ) def _varlen(q, k, v, *, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal): @@ -31,10 +37,12 @@ def _varlen(q, k, v, *, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, qs, ks = cu_seqlens_q.tolist(), cu_seqlens_k.tolist() assert max_seqlen_q == max(b - a for a, b in zip(qs, qs[1:])) assert max_seqlen_k == max(b - a for a, b in zip(ks, ks[1:])) - return torch.cat([ - _dense(q[qa:qb][None], k[ka:kb][None], v[ka:kb][None], causal=causal)[0] - for qa, qb, ka, kb in zip(qs, qs[1:], ks, ks[1:]) - ]) + return torch.cat( + [ + _dense(q[qa:qb][None], k[ka:kb][None], v[ka:kb][None], causal=causal)[0] + for qa, qb, ka, kb in zip(qs, qs[1:], ks, ks[1:]) + ] + ) def _kernels(backend="flash4"): @@ -54,11 +62,14 @@ def varlen(*args, **kwargs): def test_batched_outputs_gradients_and_example_isolation(backend, packing): torch.manual_seed(43) b, h, nq, nk, d = 3, 2, 5, 7, 8 - q, k, v = [x.requires_grad_() for x in ( - torch.randn(b, h, nq, d).bfloat16(), - torch.randn(b, h, nk, d).bfloat16(), - torch.randn(b, h, nk, d).bfloat16(), - )] + q, k, v = [ + x.requires_grad_() + for x in ( + torch.randn(b, h, nq, d).bfloat16(), + torch.randn(b, h, nk, d).bfloat16(), + torch.randn(b, h, nk, d).bfloat16(), + ) + ] valid = torch.ones(b, nk, dtype=torch.bool) qvalid = torch.ones(b, nq, dtype=torch.bool) if packing != "dense": @@ -69,14 +80,20 @@ def test_batched_outputs_gradients_and_example_isolation(backend, packing): qvalid[2, 4] = False layout = prepare_attention_layout(valid, query_valid=qvalid, query_tokens=nq) out = execute_flash_attention(q, k, v, layout, _kernels(backend)) - expected = F.scaled_dot_product_attention( - q.float(), k.float(), v.float(), attn_mask=valid[:, None, None] - ).to(q.dtype).masked_fill(~qvalid[:, None, :, None], 0) - torch.testing.assert_close(out, expected, atol=.008, rtol=.02) - actual_grads = torch.autograd.grad(out.float().square().sum(), (q, k, v), retain_graph=True) + expected = ( + F.scaled_dot_product_attention( + q.float(), k.float(), v.float(), attn_mask=valid[:, None, None] + ) + .to(q.dtype) + .masked_fill(~qvalid[:, None, :, None], 0) + ) + torch.testing.assert_close(out, expected, atol=0.008, rtol=0.02) + actual_grads = torch.autograd.grad( + out.float().square().sum(), (q, k, v), retain_graph=True + ) reference_grads = torch.autograd.grad(expected.float().square().sum(), (q, k, v)) for actual, reference in zip(actual_grads, reference_grads): - torch.testing.assert_close(actual, reference, atol=.015, rtol=.03) + torch.testing.assert_close(actual, reference, atol=0.015, rtol=0.03) assert torch.isfinite(actual).all() assert actual_grads[1].transpose(1, 2)[~valid].count_nonzero() == 0 assert actual_grads[0].transpose(1, 2)[~qvalid].count_nonzero() == 0 @@ -94,14 +111,16 @@ def test_key_only_layout_retains_all_queries_with_different_lengths(): layout = prepare_attention_layout(valid, query_tokens=3) assert layout.query_indices is None out = execute_flash_attention(q, k, v, layout, _kernels()) - expected = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), - attn_mask=valid[:, None, None]).to(q.dtype) - torch.testing.assert_close(out, expected, atol=.008, rtol=.02) + expected = F.scaled_dot_product_attention( + q.float(), k.float(), v.float(), attn_mask=valid[:, None, None] + ).to(q.dtype) + torch.testing.assert_close(out, expected, atol=0.008, rtol=0.02) def test_verified_dense_layout_does_not_read_mask_values(monkeypatch): def forbidden(*args, **kwargs): raise AssertionError("Fully valid path must not inspect CUDA mask values") + valid = torch.ones(4, 11, dtype=torch.bool) monkeypatch.setattr(torch.Tensor, "all", forbidden) monkeypatch.setattr(torch.Tensor, "nonzero", forbidden) @@ -119,11 +138,18 @@ def test_layout_owns_mask_and_handles_interior_holes(): assert layout.cu_seqlens_k.tolist() == [0, 2, 4] -@pytest.mark.parametrize("valid,qvalid,match", [ - (torch.zeros(2, 3, dtype=torch.bool), None, "valid key"), - (torch.ones(2, 3, dtype=torch.bool), torch.zeros(2, 3, dtype=torch.bool), "valid query"), - (torch.ones(2, 3), None, "boolean"), -]) +@pytest.mark.parametrize( + "valid,qvalid,match", + [ + (torch.zeros(2, 3, dtype=torch.bool), None, "valid key"), + ( + torch.ones(2, 3, dtype=torch.bool), + torch.zeros(2, 3, dtype=torch.bool), + "valid query", + ), + (torch.ones(2, 3), None, "boolean"), + ], +) def test_bad_layout_fails_closed(valid, qvalid, match): with pytest.raises(ValueError, match=match): prepare_attention_layout(valid, query_valid=qvalid) @@ -136,28 +162,45 @@ def test_precomputed_rope_matches_historical_values_and_gradients(): assert all(value.dtype == torch.float32 for value in rope) actual, expected = apply_temporal_rope(x, rope), temporal_rope(x, times) assert torch.equal(actual, expected) - assert torch.equal(torch.autograd.grad(actual.square().sum(), x)[0], - torch.autograd.grad(expected.square().sum(), x)[0]) + assert torch.equal( + torch.autograd.grad(actual.square().sum(), x)[0], + torch.autograd.grad(expected.square().sum(), x)[0], + ) def test_mot_reuses_layout_and_rope_without_mask_discovery(monkeypatch): - block = MoTBlock(16, 2, modalities=("image", "proprio"), - attention_backend="flash4", attention_dtype="bfloat16") - streams = [torch.randn(2, 5, 16, requires_grad=True), torch.randn(2, 3, 16, requires_grad=True)] + block = MoTBlock( + 16, + 2, + modalities=("image", "proprio"), + attention_backend="flash4", + attention_dtype="bfloat16", + ) + streams = [ + torch.randn(2, 5, 16, requires_grad=True), + torch.randn(2, 3, 16, requires_grad=True), + ] times = [torch.randn(2, 5), torch.randn(2, 3)] - valid = [torch.tensor([[1, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.bool), - torch.tensor([[0, 1, 0], [1, 0, 1]], dtype=torch.bool)] + valid = [ + torch.tensor([[1, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.bool), + torch.tensor([[0, 1, 0], [1, 0, 1]], dtype=torch.bool), + ] joined = torch.cat(valid, 1) layout = prepare_attention_layout(joined, query_valid=joined) ropes = tuple(prepare_temporal_rope(t, 8) for t in times) reference = copy.deepcopy(block) reference.attention_backend = "sdpa" - monkeypatch.setattr("egomimic.models.world_model.load_flash_attention_kernels", lambda *args: _kernels()) - monkeypatch.setattr(torch.Tensor, "nonzero", lambda *args, **kwargs: pytest.fail("repeated packing")) + monkeypatch.setattr( + "egomimic.models.world_model.load_flash_attention_kernels", + lambda *args: _kernels(), + ) + monkeypatch.setattr( + torch.Tensor, "nonzero", lambda *args, **kwargs: pytest.fail("repeated packing") + ) actual = block(streams, times, valid, attention_layout=layout, rope=ropes) expected = reference(streams, times, valid) for a, e, mask in zip(actual, expected, valid): - torch.testing.assert_close(a, e, atol=.01, rtol=.02) + torch.testing.assert_close(a, e, atol=0.01, rtol=0.02) assert a[~mask].count_nonzero() == 0 sum(x.square().sum() for x in actual).backward() for x, mask in zip(streams, valid): @@ -166,21 +209,31 @@ def test_mot_reuses_layout_and_rope_without_mask_discovery(monkeypatch): def test_query_chunks_preserve_generic_batched_nonprefix_mask_semantics(monkeypatch): torch.manual_seed(53) - full = MoTBlock(16, 2, modalities=("image",), - attention_backend="flash4", attention_dtype="bfloat16") + full = MoTBlock( + 16, + 2, + modalities=("image",), + attention_backend="flash4", + attention_dtype="bfloat16", + ) chunked = copy.deepcopy(full) chunked.attention_query_chunk_size = 3 - monkeypatch.setattr("egomimic.models.world_model.load_flash_attention_kernels", lambda *args: _kernels()) + monkeypatch.setattr( + "egomimic.models.world_model.load_flash_attention_kernels", + lambda *args: _kernels(), + ) values = torch.randn(2, 7, 16, requires_grad=True) times = torch.randn(2, 7) - valid = torch.tensor([[1, 0, 1, 1, 0, 1, 1], [1, 1, 1, 0, 1, 1, 0]], dtype=torch.bool) + valid = torch.tensor( + [[1, 0, 1, 1, 0, 1, 1], [1, 1, 1, 0, 1, 1, 0]], dtype=torch.bool + ) layout = prepare_attention_layout(valid, query_valid=valid) - a, = full([values], [times], [valid], attention_layout=layout) - b, = chunked([values], [times], [valid], attention_layout=layout) - torch.testing.assert_close(a, b, rtol=.02, atol=.008) - ga, = torch.autograd.grad(a.square().sum(), values) - gb, = torch.autograd.grad(b.square().sum(), values) - torch.testing.assert_close(ga, gb, rtol=.03, atol=.015) + (a,) = full([values], [times], [valid], attention_layout=layout) + (b,) = chunked([values], [times], [valid], attention_layout=layout) + torch.testing.assert_close(a, b, rtol=0.02, atol=0.008) + (ga,) = torch.autograd.grad(a.square().sum(), values) + (gb,) = torch.autograd.grad(b.square().sum(), values) + torch.testing.assert_close(ga, gb, rtol=0.03, atol=0.015) @pytest.mark.parametrize("backend", ["flash3", "flash4"]) @@ -194,13 +247,17 @@ def test_flash3_import_uses_official_namespace_or_legacy_alias(monkeypatch, lega from types import SimpleNamespace calls = [] + def load(name): calls.append(name) if legacy and name == "flash_attn_3.flash_attn_interface": raise ModuleNotFoundError("legacy install", name="flash_attn_3") return SimpleNamespace(flash_attn_func=_dense, flash_attn_varlen_func=_varlen) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0)) - monkeypatch.setattr("egomimic.models.attention_execution.importlib.import_module", load) + monkeypatch.setattr( + "egomimic.models.attention_execution.importlib.import_module", load + ) load_flash_attention_kernels.cache_clear() try: result = load_flash_attention_kernels("flash3", torch.device("cuda:0")) @@ -213,11 +270,17 @@ def load(name): def test_installed_backend_missing_dependency_does_not_try_alias(monkeypatch): calls = [] + def load(name): calls.append(name) - raise ModuleNotFoundError("broken installed extension", name="flash_attn_3_cuda") + raise ModuleNotFoundError( + "broken installed extension", name="flash_attn_3_cuda" + ) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0)) - monkeypatch.setattr("egomimic.models.attention_execution.importlib.import_module", load) + monkeypatch.setattr( + "egomimic.models.attention_execution.importlib.import_module", load + ) load_flash_attention_kernels.cache_clear() try: with pytest.raises(RuntimeError, match="no backend fallback"): @@ -238,16 +301,19 @@ def test_regional_compile_preserves_checkpoint_keys_and_backward(): times = torch.randn(2, n) valid = torch.ones(2, n, dtype=torch.bool) rope = (prepare_temporal_rope(times, 8),) - a, = eager([x], [times], [valid], rope=rope) - b, = compiled([x], [times], [valid], rope=rope) + (a,) = eager([x], [times], [valid], rope=rope) + (b,) = compiled([x], [times], [valid], rope=rope) assert torch.equal(a, b) - assert torch.equal(torch.autograd.grad(a.square().sum(), x)[0], - torch.autograd.grad(b.square().sum(), x)[0]) + assert torch.equal( + torch.autograd.grad(a.square().sum(), x)[0], + torch.autograd.grad(b.square().sum(), x)[0], + ) compiled.load_state_dict(eager.state_dict(), strict=True) def test_compile_fails_instead_of_falling_back(monkeypatch): import torch._dynamo + with torch._dynamo.config.patch(suppress_errors=True): with pytest.raises(RuntimeError, match="no eager fallback"): compile_region(lambda x: x + 1, mode="default") @@ -255,7 +321,9 @@ def test_compile_fails_instead_of_falling_back(monkeypatch): compile_region(lambda x: x, mode="invalid") -def test_compiled_backward_autocast_policy_covers_lazy_invocation_and_restores(monkeypatch): +def test_compiled_backward_autocast_policy_covers_lazy_invocation_and_restores( + monkeypatch, +): from contextlib import contextmanager from types import SimpleNamespace @@ -286,6 +354,7 @@ def lazy_call(value): if value == "fail": raise RuntimeError("tracing failed") return function(value) + return lazy_call monkeypatch.setattr(torch, "compile", lazy_compile) @@ -298,15 +367,19 @@ def lazy_call(value): assert calls == ["off", "off"] -def test_missing_aot_autocast_policy_fails_closed_only_for_compiled_mixed_mode(monkeypatch): +def test_missing_aot_autocast_policy_fails_closed_only_for_compiled_mixed_mode( + monkeypatch, +): from types import SimpleNamespace import egomimic.models.compilation as compilation monkeypatch.setattr(compilation, "_aot_config", SimpleNamespace()) monkeypatch.setattr(torch, "compile", lambda function, **_: function) + def function(value): return value + 1 + execute = compile_region(function, mode="default") assert execute(torch.ones(1)).item() == 2 with torch.autocast("cpu", dtype=torch.bfloat16): @@ -318,6 +391,7 @@ def function(value): def test_aot_eager_bf16_forward_backward_matches_standard_amp(): from torch._functorch import config + if not hasattr(config, "backward_pass_autocast"): pytest.skip("This PyTorch lacks the required compiled AMP backward policy") torch.manual_seed(19) @@ -331,25 +405,27 @@ def test_aot_eager_bf16_forward_backward_matches_standard_amp(): copied = x.detach().clone().requires_grad_() before = config.backward_pass_autocast with torch.autocast("cpu", dtype=torch.bfloat16): - expected, = eager([x], [times], [valid], rope=rope) - actual, = compiled([copied], [times], [valid], rope=rope) + (expected,) = eager([x], [times], [valid], rope=rope) + (actual,) = compiled([copied], [times], [valid], rope=rope) expected_loss = expected.float().square().mean() actual_loss = actual.float().square().mean() assert config.backward_pass_autocast == before expected_loss.backward() actual_loss.backward() - torch.testing.assert_close(actual, expected, rtol=.02, atol=.01) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.01) for a, b in [(copied.grad, x.grad)] + [ (a.grad, b.grad) for a, b in zip(compiled.parameters(), eager.parameters()) ]: assert a.dtype == b.dtype == torch.float32 assert torch.isfinite(a).all() - torch.testing.assert_close(a, b, rtol=.02, atol=.002) - assert (a - b).norm() / b.norm().clamp_min(1e-8) < .02 + torch.testing.assert_close(a, b, rtol=0.02, atol=0.002) + assert (a - b).norm() / b.norm().clamp_min(1e-8) < 0.02 -@pytest.mark.skipif(os.environ.get("EGOMIMIC_RUN_FUSED_ATTENTION_TESTS") != "1", - reason="Explicit opt-in real CUDA attention proof") +@pytest.mark.skipif( + os.environ.get("EGOMIMIC_RUN_FUSED_ATTENTION_TESTS") != "1", + reason="Explicit opt-in real CUDA attention proof", +) @pytest.mark.parametrize("backend", ["flash3", "flash4"]) @pytest.mark.parametrize("dense", [True, False]) def test_real_cuda_batched_attention_matches_independent_gradients(backend, dense): @@ -358,27 +434,38 @@ def test_real_cuda_batched_attention_matches_independent_gradients(backend, dens device = torch.device("cuda", torch.cuda.current_device()) kernels = load_flash_attention_kernels(backend, device) torch.manual_seed(71) - q, k, v = [torch.randn(3, 6, 97, 128, device=device, dtype=torch.bfloat16, - requires_grad=True) for _ in range(3)] + q, k, v = [ + torch.randn( + 3, 6, 97, 128, device=device, dtype=torch.bfloat16, requires_grad=True + ) + for _ in range(3) + ] valid = torch.ones(3, 97, device=device, dtype=torch.bool) if not dense: valid[0, 1::3] = False valid[1, 1::2] = False layout = prepare_attention_layout(valid, query_valid=valid) actual = execute_flash_attention(q, k, v, layout, kernels) - reference = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), - attn_mask=valid[:, None, None]).masked_fill(~valid[:, None, :, None], 0) - torch.testing.assert_close(actual.float(), reference, rtol=.03, atol=.02) - ga = torch.autograd.grad(actual.float().square().sum(), (q, k, v), retain_graph=True) + reference = F.scaled_dot_product_attention( + q.float(), k.float(), v.float(), attn_mask=valid[:, None, None] + ).masked_fill(~valid[:, None, :, None], 0) + torch.testing.assert_close(actual.float(), reference, rtol=0.03, atol=0.02) + ga = torch.autograd.grad( + actual.float().square().sum(), (q, k, v), retain_graph=True + ) gr = torch.autograd.grad(reference.square().sum(), (q, k, v)) for a, r in zip(ga, gr): assert torch.isfinite(a).all() - assert (a.float() - r.float()).norm() / r.float().norm().clamp_min(1e-8) < .04 + assert (a.float() - r.float()).norm() / r.float().norm().clamp_min(1e-8) < 0.04 for i in range(3): - one = prepare_attention_layout(valid[i:i+1], query_valid=valid[i:i+1]) - independent = execute_flash_attention(q[i:i+1], k[i:i+1], v[i:i+1], one, kernels) - torch.testing.assert_close(actual[i:i+1], independent, rtol=.03, atol=.02) - independent_grads = torch.autograd.grad(independent.float().square().sum(), (q, k, v)) + one = prepare_attention_layout(valid[i : i + 1], query_valid=valid[i : i + 1]) + independent = execute_flash_attention( + q[i : i + 1], k[i : i + 1], v[i : i + 1], one, kernels + ) + torch.testing.assert_close(actual[i : i + 1], independent, rtol=0.03, atol=0.02) + independent_grads = torch.autograd.grad( + independent.float().square().sum(), (q, k, v) + ) for full, single in zip(ga, independent_grads): relative = (full[i].float() - single[i].float()).norm() - assert relative / full[i].float().norm().clamp_min(1e-8) < .04 + assert relative / full[i].float().norm().clamp_min(1e-8) < 0.04 diff --git a/tests/test_image_proprio_adaln.py b/tests/test_image_proprio_adaln.py index cbc47d282..97d3f7d4a 100644 --- a/tests/test_image_proprio_adaln.py +++ b/tests/test_image_proprio_adaln.py @@ -47,29 +47,47 @@ def test_additive_default_preserves_seeded_state_rng_and_outputs(): original = model(noise_conditioning="additive") rng = torch.get_rng_state().clone() torch.manual_seed(12) - explicit = packed_model(backend="sdpa", attention_dtype="native", use_phase_embedding=False) + explicit = packed_model( + backend="sdpa", attention_dtype="native", use_phase_embedding=False + ) assert torch.equal(rng, torch.get_rng_state()) - assert all(torch.equal(value, explicit.state_dict()[name]) for name, value in original.state_dict().items()) + assert all( + torch.equal(value, explicit.state_dict()[name]) + for name, value in original.state_dict().items() + ) world = batch(nonprefix=True) noisy, levels = fixed_corruption(world, 8) - assert torch.equal(original(world, "libero_gripper", noisy, levels)["image"], - explicit(world, "libero_gripper", noisy, levels)["image"]) + assert torch.equal( + original(world, "libero_gripper", noisy, levels)["image"], + explicit(world, "libero_gripper", noisy, levels)["image"], + ) def test_standard_zero_initialization_and_exact_parameter_count(): with torch.device("meta"): - full = ImageProprioModel(objective="visual_only", phase_a_condition_on_proprio=False, - max_cameras=2, use_phase_embedding=False, noise_conditioning="adaln_zero") + full = ImageProprioModel( + objective="visual_only", + phase_a_condition_on_proprio=False, + max_cameras=2, + use_phase_embedding=False, + noise_conditioning="adaln_zero", + ) assert sum(parameter.numel() for parameter in full.parameters()) == 88_180_992 network = model() assert not network.norm[0].elementwise_affine - assert all(not block.norm1[0].elementwise_affine and not block.norm2[0].elementwise_affine for block in network.blocks) + assert all( + not block.norm1[0].elementwise_affine and not block.norm2[0].elementwise_affine + for block in network.blocks + ) world = batch() noisy, levels = fixed_corruption(world, 7) prediction = network(world, "libero_gripper", noisy, levels)["image"] assert prediction.count_nonzero() == 0 network.training_losses(world, "libero_gripper")["total"].backward() - assert all(parameter.grad is not None and torch.isfinite(parameter.grad).all() for parameter in network.parameters()) + assert all( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + for parameter in network.parameters() + ) assert network.image_out.weight.grad.abs().sum() > 0 assert network.blocks[0].qkv[0].weight.grad.count_nonzero() == 0 @@ -80,10 +98,14 @@ def test_frame_modulation_is_shared_but_examples_and_levels_remain_distinct(): world = batch(nonprefix=True) noisy, levels = fixed_corruption(world, 11) frames = network._prepare_noise_frame_map(world["image"]) - assert frames.representatives.numel() == 5 #2 and3frames, each8camera/patchtokens. + assert frames.representatives.numel() == 5 # 2 and3frames, each8camera/patchtokens. shapes = [] - handles = [block.adaln_image[-1].register_forward_pre_hook(lambda _, args: shapes.append(args[0].shape)) - for block in network.blocks] + handles = [ + block.adaln_image[-1].register_forward_pre_hook( + lambda _, args: shapes.append(args[0].shape) + ) + for block in network.blocks + ] first = network(world, "libero_gripper", noisy, levels)["image"] for handle in handles: handle.remove() @@ -95,8 +117,10 @@ def test_frame_modulation_is_shared_but_examples_and_levels_remain_distinct(): assert not torch.allclose(first[1], second[1], atol=1e-7, rtol=1e-7) # Input embeddings no longer receive additive noise vectors. image = world["image"] - assert torch.equal(network._embed_image(image, noisy, levels, image["valid"], "A", False), - network._embed_image(image, noisy, changed, image["valid"], "A", False)) + assert torch.equal( + network._embed_image(image, noisy, levels, image["valid"], "A", False), + network._embed_image(image, noisy, changed, image["valid"], "A", False), + ) query = image["valid"] & ~image["observed"] row, token = query.nonzero()[0].tolist() invalid = levels.clone() @@ -124,7 +148,9 @@ def record_embed(levels): return original_embed(levels) object.__setattr__(network, "_embed_noise_frames_execute", record_embed) - output = network.sample(world, "libero_gripper", steps=4, generator=torch.Generator().manual_seed(4))["image"] + output = network.sample( + world, "libero_gripper", steps=4, generator=torch.Generator().manual_seed(4) + )["image"] assert len(calls) == 1 and len(levels_seen) == 4 assert [float(level.max()) for level in levels_seen] == [1.0, 0.75, 0.5, 0.25] assert all((level == 0).sum() == 2 for level in levels_seen) @@ -136,8 +162,12 @@ def record_embed(levels): @pytest.mark.parametrize("mixed", [False, True]) @pytest.mark.parametrize("time_mode", ["continuous", "discrete"]) @pytest.mark.usefixtures("reference_kernels") -def test_nonzero_packed_conditioning_matches_padded_outputs_and_gradients(mixed, time_mode): - arguments = dict(noise_conditioning="adaln_zero", use_phase_embedding=False, time_mode=time_mode) +def test_nonzero_packed_conditioning_matches_padded_outputs_and_gradients( + mixed, time_mode +): + arguments = dict( + noise_conditioning="adaln_zero", use_phase_embedding=False, time_mode=time_mode + ) padded = packed_model(**arguments) activate_conditioning(padded) packed = packed_model(token_execution="packed_valid", **arguments) @@ -168,7 +198,9 @@ def test_nonzero_packed_conditioning_matches_padded_outputs_and_gradients(mixed, def test_fullgraph_tensor_regions_match_eager_with_nonzero_conditioning(): eager = model() activate_conditioning(eager) - compiled = model(compile_mode="default", compile_backend="aot_eager", compile_dynamic=True) + compiled = model( + compile_mode="default", compile_backend="aot_eager", compile_dynamic=True + ) compiled.load_state_dict(eager.state_dict(), strict=True) world = batch(nonprefix=True) noisy, levels = fixed_corruption(world, 23) @@ -179,5 +211,10 @@ def test_fullgraph_tensor_regions_match_eager_with_nonzero_conditioning(): outputs.append(prediction.detach()) torch.testing.assert_close(outputs[0], outputs[1], rtol=1e-5, atol=1e-6) for name, parameter in eager.named_parameters(): - torch.testing.assert_close(parameter.grad, dict(compiled.named_parameters())[name].grad, - rtol=2e-5, atol=2e-6, msg=name) + torch.testing.assert_close( + parameter.grad, + dict(compiled.named_parameters())[name].grad, + rtol=2e-5, + atol=2e-6, + msg=name, + ) diff --git a/tests/test_image_proprio_attention.py b/tests/test_image_proprio_attention.py index 46ea095f8..8b2e5f47a 100644 --- a/tests/test_image_proprio_attention.py +++ b/tests/test_image_proprio_attention.py @@ -29,7 +29,7 @@ def example(): torch.manual_seed(31) image = { "value": torch.randn(1, 6, 8), - "time": torch.tensor([[-1., -1., 0., 0., 2., 2.]]), + "time": torch.tensor([[-1.0, -1.0, 0.0, 0.0, 2.0, 2.0]]), "valid": torch.tensor([[True, True, False, True, True, True]]), "observed": torch.tensor([[True, True, False, False, False, False]]), "xy": torch.randn(1, 6, 2), @@ -37,7 +37,7 @@ def example(): } proprio = { "value": torch.randn(1, 4, 3), - "time": torch.tensor([[-1., -.5, 0., 2.]]), + "time": torch.tensor([[-1.0, -0.5, 0.0, 2.0]]), "valid": torch.ones(1, 4, dtype=torch.bool), "observed": torch.tensor([[True, False, True, False]]), } @@ -46,15 +46,20 @@ def example(): def forward(network, batch): return network( - batch, "gripper", batch["image"]["value"], - torch.zeros_like(batch["image"]["time"]), phase="A", + batch, + "gripper", + batch["image"]["value"], + torch.zeros_like(batch["image"]["time"]), + phase="A", )["image"] def test_native_default_matches_historical_attention_and_strict_state_loading(): torch.manual_seed(8) default = model() - explicit = model(attention_backend="sdpa", attention_dtype="native", strict_fp32=True) + explicit = model( + attention_backend="sdpa", attention_dtype="native", strict_fp32=True + ) historical_attention = model() explicit.load_state_dict(default.state_dict(), strict=True) historical_attention.load_state_dict(default.state_dict(), strict=True) @@ -113,7 +118,9 @@ def test_two_stream_holes_have_no_label_leakage_or_invalid_value_gradients(dtype assert pose_grad[hidden].count_nonzero() == 0 assert pose_grad[~hidden].abs().sum() > 0 assert batch["image"]["value"].grad[~batch["image"]["valid"]].count_nonzero() == 0 - assert all(torch.isfinite(p.grad).all() for p in network.parameters() if p.grad is not None) + assert all( + torch.isfinite(p.grad).all() for p in network.parameters() if p.grad is not None + ) assert all(p.dtype == torch.float32 for p in network.parameters()) # Removing invalid keys (including interior holes) preserves valid outputs. @@ -124,8 +131,10 @@ def test_two_stream_holes_have_no_label_leakage_or_invalid_value_gradients(dtype ): compact[name] = {key: value[:, keep] for key, value in compact[name].items()} torch.testing.assert_close( - forward(network, compact), result[:, batch["image"]["valid"][0]], - rtol=1e-5, atol=1e-5, + forward(network, compact), + result[:, batch["image"]["valid"][0]], + rtol=1e-5, + atol=1e-5, ) @@ -144,25 +153,31 @@ def record(q, k, v, mask, original=original): result = forward(network, batch) assert result.dtype == torch.float32 result.square().mean().backward() - expected_keys = torch.cat((batch["image"]["valid"], batch["proprio"]["observed"]), 1) + expected_keys = torch.cat( + (batch["image"]["valid"], batch["proprio"]["observed"]), 1 + ) assert len(calls) == len(network.blocks) for q_dtype, k_dtype, v_dtype, mask in calls: assert q_dtype == k_dtype == v_dtype == torch.bfloat16 assert torch.equal(mask[:, 0, 0], expected_keys) assert all( p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() - for p in network.parameters() if p.grad is not None + for p in network.parameters() + if p.grad is not None ) def test_unchunked_noncheckpointed_model_does_not_recompute(monkeypatch): def forbidden(*args, **kwargs): - raise AssertionError("Checkpointing must be disabled by both execution controls") + raise AssertionError( + "Checkpointing must be disabled by both execution controls" + ) monkeypatch.setattr("egomimic.models.image_proprio.checkpoint", forbidden) monkeypatch.setattr("egomimic.models.world_model.checkpoint", forbidden) network = model( - attention_query_chunk_size=0, activation_checkpointing=False, + attention_query_chunk_size=0, + activation_checkpointing=False, attention_dtype="bfloat16", ) forward(network, example()).square().mean().backward() @@ -176,7 +191,11 @@ def test_flash4_model_does_not_silently_fallback_on_cpu(): def test_legacy_dino_wrapper_keeps_native_execution_defaults(): network = DinoProprioModel( - pose_dims={"gripper": 3}, image_dim=8, width=16, depth=1, heads=2, + pose_dims={"gripper": 3}, + image_dim=8, + width=16, + depth=1, + heads=2, ) assert network.attention_backend == "sdpa" assert network.attention_dtype == "native" @@ -188,9 +207,9 @@ def test_legacy_dino_wrapper_keeps_native_execution_defaults(): def test_generic_single_modality_block_has_valid_backward(): block = MoTBlock(24, 3, modalities=("image",), attention_dtype="bfloat16") values = torch.randn(1, 5, 24, requires_grad=True) - times = torch.tensor([[-1., 0., .5, 1., 2.]]) + times = torch.tensor([[-1.0, 0.0, 0.5, 1.0, 2.0]]) valid = torch.tensor([[True, False, True, True, True]]) - output, = block([values], [times], [valid]) + (output,) = block([values], [times], [valid]) assert len(block.qkv) == 1 assert output.shape == values.shape and output.dtype == torch.float32 assert output[~valid].count_nonzero() == 0 diff --git a/tests/test_image_proprio_distributed.py b/tests/test_image_proprio_distributed.py index 67b76a0aa..bd0d0bc13 100644 --- a/tests/test_image_proprio_distributed.py +++ b/tests/test_image_proprio_distributed.py @@ -17,13 +17,18 @@ def _distributed_metrics_worker(rank, rendezvous, output): torch.set_num_threads(1) torch.distributed.init_process_group( - "gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2, + "gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, timeout=timedelta(seconds=30), ) try: logged = {} model = SimpleNamespace( - device=torch.device("cpu"), global_step=4, global_rank=rank, + device=torch.device("cpu"), + global_step=4, + global_rank=rank, trainer=SimpleNamespace(sanity_checking=False), log=lambda key, value, **kwargs: logged.update({key: float(value)}), ) @@ -36,7 +41,10 @@ def prepare(): evaluator._validation_batches = (60, 40)[rank] # An unweighted average of rank means is wrong: (2 + 6)/2 != 5. evaluator.totals = { - "forecast": (torch.tensor((2.0, 18.0)[rank]), torch.tensor((1.0, 3.0)[rank])) + "forecast": ( + torch.tensor((2.0, 18.0)[rank]), + torch.tensor((1.0, 3.0)[rank]), + ) } if rank == 1: evaluator.totals["completion"] = (torch.tensor(10.0), torch.tensor(2.0)) @@ -80,12 +88,20 @@ def publish(_): torch.distributed.destroy_process_group() -@pytest.mark.skipif(not torch.distributed.is_gloo_available(), reason="Gloo unavailable") +@pytest.mark.skipif( + not torch.distributed.is_gloo_available(), reason="Gloo unavailable" +) def test_distributed_metrics_reduce_sums_and_global_completion(tmp_path): - mp.spawn(_distributed_metrics_worker, args=(str(tmp_path / "rendezvous"), str(tmp_path)), - nprocs=2, join=True) + mp.spawn( + _distributed_metrics_worker, + args=(str(tmp_path / "rendezvous"), str(tmp_path)), + nprocs=2, + join=True, + ) for rank in range(2): - assert json.loads((tmp_path / f"rank-{rank}.json").read_text()) == {"status": "PASS"} + assert json.loads((tmp_path / f"rank-{rank}.json").read_text()) == { + "status": "PASS" + } def test_nonzero_rank_does_not_publish_rgb_artifacts(tmp_path): @@ -93,8 +109,15 @@ def test_nonzero_rank_does_not_publish_rgb_artifacts(tmp_path): model, prediction, _ = fake_model(batch) model.global_rank = 1 diagnostics = PixelPredictionDiagnostics(artifact_dir=tmp_path, log_wandb=False) - diagnostics.record(model=model, source="robot", item=batch["robot"], - prediction=prediction, objective="visual_only", batch_idx=0, dataloader_idx=0) + diagnostics.record( + model=model, + source="robot", + item=batch["robot"], + prediction=prediction, + objective="visual_only", + batch_idx=0, + dataloader_idx=0, + ) assert diagnostics.written == [] and not list(tmp_path.iterdir()) diff --git a/tests/test_image_proprio_image_only.py b/tests/test_image_proprio_image_only.py index bc96d6d64..277af2fc7 100644 --- a/tests/test_image_proprio_image_only.py +++ b/tests/test_image_proprio_image_only.py @@ -11,8 +11,14 @@ def image_model(**kwargs): settings = dict( - objective="visual_only", phase_a_condition_on_proprio=False, - image_dim=8, width=24, heads=3, depth=2, max_cameras=2, max_tokens=32, + objective="visual_only", + phase_a_condition_on_proprio=False, + image_dim=8, + width=24, + heads=3, + depth=2, + max_cameras=2, + max_tokens=32, ) settings.update(kwargs) return ImageProprioModel(**settings) @@ -20,8 +26,10 @@ def image_model(**kwargs): def joint_model(**kwargs): return image_model( - objective="visual_state_action", pose_dims={"gripper": 3}, - action_dims={"gripper": 2}, **kwargs, + objective="visual_state_action", + pose_dims={"gripper": 3}, + action_dims={"gripper": 2}, + **kwargs, ) @@ -38,8 +46,11 @@ def joint_batch(): def predict(model, batch, phase="A"): return model( - batch, "gripper", batch["image"]["value"], - torch.zeros_like(batch["image"]["time"]), phase=phase, + batch, + "gripper", + batch["image"]["value"], + torch.zeros_like(batch["image"]["time"]), + phase=phase, ) @@ -52,7 +63,14 @@ def test_image_only_has_exact_parameters_and_no_proprio_modules(): assert not full.has_proprio_branch assert len(full.norm) == 1 assert all(len(block.qkv) == 1 for block in full.blocks) - for name in ("pose_in", "proprio_role", "pose_out", "action_out", "query", "action_query"): + for name in ( + "pose_in", + "proprio_role", + "pose_out", + "action_out", + "query", + "action_query", + ): assert not hasattr(full, name) assert not any("pose" in key or "proprio" in key for key in full.state_dict()) @@ -63,7 +81,9 @@ def test_image_only_trains_and_samples_without_state_data(): losses = model.training_losses(batch, "gripper") assert set(losses) == {"image", "total"} losses["total"].backward() - assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) + assert all( + p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() + ) query = ~batch["image"]["observed"] batch["image"]["value"][query] = float("nan") prediction = model.sample(batch, "gripper", steps=20)["image"] @@ -122,14 +142,21 @@ def test_joint_phase_b_uses_new_branch_but_withheld_labels_are_loss_only(): assert all(block.qkv[1].weight.grad.abs().sum() > 0 for block in model.blocks) -@pytest.mark.parametrize("readout_weight", [0., 1.]) -def test_joint_training_supports_zero_or_active_readout_after_image_only_phase_a(readout_weight): +@pytest.mark.parametrize("readout_weight", [0.0, 1.0]) +def test_joint_training_supports_zero_or_active_readout_after_image_only_phase_a( + readout_weight, +): model = joint_model() losses = model.training_losses( - joint_batch(), "gripper", pose_weight=readout_weight, action_weight=readout_weight, + joint_batch(), + "gripper", + pose_weight=readout_weight, + action_weight=readout_weight, ) assert set(losses) == {"image", "pose", "action", "total"} assert all(torch.isfinite(value) for value in losses.values()) losses["total"].backward() - assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None) + assert all( + torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None + ) assert (model.action_out["gripper"].weight.grad is None) == (readout_weight == 0) diff --git a/tests/test_image_proprio_image_only_initialization.py b/tests/test_image_proprio_image_only_initialization.py index 127ed243f..5ad3f98e9 100644 --- a/tests/test_image_proprio_image_only_initialization.py +++ b/tests/test_image_proprio_image_only_initialization.py @@ -16,14 +16,23 @@ def configs(mode="continuous", modality=False, backend="sdpa"): source = dict( - objective="visual_only", phase_a_condition_on_proprio=False, - image_dim=12, width=16, heads=2, depth=2, max_cameras=2, max_tokens=12, - time_mode=mode, use_modality_embedding=modality, + objective="visual_only", + phase_a_condition_on_proprio=False, + image_dim=12, + width=16, + heads=2, + depth=2, + max_cameras=2, + max_tokens=12, + time_mode=mode, + use_modality_embedding=modality, attention_backend=backend, attention_dtype="bfloat16" if backend == "flash4" else "native", ) target = dict( - source, objective="visual_state_action", max_tokens=16, + source, + objective="visual_state_action", + max_tokens=16, pose_dims={"gripper": 9 if mode == "continuous" else 36}, action_dims={"gripper": 7 if mode == "continuous" else 28}, ) @@ -32,10 +41,15 @@ def configs(mode="continuous", modality=False, backend="sdpa"): @pytest.mark.parametrize("mode", ["continuous", "discrete"]) @pytest.mark.parametrize("modality", [False, True]) -def test_image_weights_copy_exactly_and_added_branch_keeps_seeded_initialization(mode, modality): +def test_image_weights_copy_exactly_and_added_branch_keeps_seeded_initialization( + mode, modality +): source_config, target_config = configs(mode, modality) validate_transition(source_config, target_config) - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) initialized = {key: value.clone() for key, value in target.state_dict().items()} transfer = transfer_shared_weights(source, target) assert "proprio_role.weight" in transfer["new_tensors"] @@ -46,7 +60,7 @@ def test_image_weights_copy_exactly_and_added_branch_keeps_seeded_initialization for key, value in source.state_dict().items(): actual = target.state_dict()[key] if key in transfer["expanded_embeddings"]: - actual = actual[:value.shape[0]] + actual = actual[: value.shape[0]] assert torch.equal(value, actual) proof = verify_phase_a(source, target) assert proof["bit_exact"] and len(proof["checks"]) == 2 @@ -68,7 +82,10 @@ def test_full_twenty_step_generation_and_image_features_are_preserved(): source_config["image_dim"] = target_config["image_dim"] = 8 target_config["pose_dims"] = {"gripper": 3} target_config["action_dims"] = {"gripper": 2} - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) transfer_shared_weights(source, target) batch = joint_batch() image_features, target_features = [], [] @@ -82,11 +99,16 @@ def test_full_twenty_step_generation_and_image_features_are_preserved(): ] try: original = source.sample( - {"image": batch["image"]}, "gripper", steps=20, + {"image": batch["image"]}, + "gripper", + steps=20, generator=torch.Generator().manual_seed(17), ) migrated = target.sample( - batch, "gripper", steps=20, generator=torch.Generator().manual_seed(17), + batch, + "gripper", + steps=20, + generator=torch.Generator().manual_seed(17), ) finally: for handle in handles: @@ -98,7 +120,10 @@ def test_full_twenty_step_generation_and_image_features_are_preserved(): def test_fa4_cpu_probe_requires_declared_override_and_restores_model_execution(): source_config, target_config = configs(backend="flash4") - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) transfer_shared_weights(source, target) with pytest.raises(ValueError, match="explicit cpu_attention_override"): verify_phase_a(source, target) @@ -106,23 +131,36 @@ def test_fa4_cpu_probe_requires_declared_override_and_restores_model_execution() assert proof["bit_exact"] and proof["cpu_attention_override"] == "native" assert proof["declared_attention"] == ["flash4/bfloat16"] assert not proof["gpu_kernel_verified"] - assert all(block.attention_backend == "flash4" for model in (source, target) for block in model.blocks) + assert all( + block.attention_backend == "flash4" + for model in (source, target) + for block in model.blocks + ) with torch.no_grad(): target.image_out.bias.add_(1) with pytest.raises(ValueError, match="preservation failed"): verify_phase_a(source, target, cpu_attention_override="native") - assert all(block.attention_backend == "flash4" for model in (source, target) for block in model.blocks) + assert all( + block.attention_backend == "flash4" + for model in (source, target) + for block in model.blocks + ) -def test_offline_fa4_artifact_binds_original_configs_and_records_only_cpu_proof(tmp_path): +def test_offline_fa4_artifact_binds_original_configs_and_records_only_cpu_proof( + tmp_path, +): source_config, target_config = configs(backend="flash4") source = ImageProprioModel(**source_config) path = tmp_path / "source.pt" torch.save(source.state_dict(), path) kwargs = dict( - source_checkpoint=path, source_sha256=sha256_file(path), - source_config=source_config, target_config=target_config, - output_path=tmp_path / "target.pt", receipt_path=tmp_path / "receipt.json", + source_checkpoint=path, + source_sha256=sha256_file(path), + source_config=source_config, + target_config=target_config, + output_path=tmp_path / "target.pt", + receipt_path=tmp_path / "receipt.json", seed=42, ) with pytest.raises(ValueError, match="explicit cpu_attention_override"): @@ -136,7 +174,9 @@ def test_offline_fa4_artifact_binds_original_configs_and_records_only_cpu_proof( assert receipt["phase_a_preservation"]["cpu_attention_override"] == "native" assert not receipt["phase_a_preservation"]["gpu_kernel_verified"] target = ImageProprioModel(**target_config) - target.load_state_dict(torch.load(kwargs["output_path"], weights_only=True), strict=True) + target.load_state_dict( + torch.load(kwargs["output_path"], weights_only=True), strict=True + ) assert verify_phase_a(source, target, "native")["bit_exact"] assert not receipt["optimizer_state_transferred"] assert not receipt["training_step_transferred"] @@ -144,7 +184,10 @@ def test_offline_fa4_artifact_binds_original_configs_and_records_only_cpu_proof( def test_unknown_added_parameter_cannot_be_silently_initialized(): source_config, target_config = configs() - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) target.extra_parameter = torch.nn.Parameter(torch.ones(1)) with pytest.raises(ValueError, match="Unexpected parameter namespace"): transfer_shared_weights(source, target) diff --git a/tests/test_image_proprio_integration.py b/tests/test_image_proprio_integration.py index fc149a5ca..1c2c5f755 100644 --- a/tests/test_image_proprio_integration.py +++ b/tests/test_image_proprio_integration.py @@ -165,7 +165,9 @@ def test_pixel_pipeline_optimization_validation_and_checkpoint( if images_only: assert trainer.optimizers[0].param_groups[0]["lr"] == 5e-5 restored = ModelWrapper( - config_tree=cfg, evaluator=instantiate(cfg.evaluator), enable_grad_norm=False + config_tree=cfg, + evaluator=instantiate(cfg.evaluator), + enable_grad_norm=False, ) restored.load_state_dict(checkpoint["state_dict"], strict=True) for key, value in restored.state_dict().items(): diff --git a/tests/test_image_proprio_mixed_precision.py b/tests/test_image_proprio_mixed_precision.py index 5a8d5b899..67f0c27e2 100644 --- a/tests/test_image_proprio_mixed_precision.py +++ b/tests/test_image_proprio_mixed_precision.py @@ -11,7 +11,9 @@ def test_mixed_training_keeps_loss_timestamps_parameters_and_adamw_state_fp32(): model = image_model(attention_dtype="bfloat16") batch = {"image": example()["image"]} # Nearby physical timestamps must remain distinct under mixed precision. - batch["image"]["time"] = torch.tensor([[30.010, 30.010, 30.020, 30.020, 30.030, 30.030]]) + batch["image"]["time"] = torch.tensor( + [[30.010, 30.010, 30.020, 30.020, 30.030, 30.030]] + ) times = batch["image"]["time"].clone() optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) observed_dtypes = {} @@ -19,6 +21,7 @@ def test_mixed_training_keeps_loss_timestamps_parameters_and_adamw_state_fp32(): def capture(name): def record(module, inputs, output): observed_dtypes[name] = output.dtype + return record handles = [ @@ -30,11 +33,16 @@ def record(module, inputs, output): with torch.autocast("cpu", dtype=torch.bfloat16): losses = model.training_losses(batch, "gripper") assert set(observed_dtypes.values()) == {torch.bfloat16} - assert all(loss.dtype == torch.float32 and torch.isfinite(loss) for loss in losses.values()) + assert all( + loss.dtype == torch.float32 and torch.isfinite(loss) + for loss in losses.values() + ) losses["total"].backward() assert all( - p.dtype == torch.float32 and p.grad is not None - and p.grad.dtype == torch.float32 and torch.isfinite(p.grad).all() + p.dtype == torch.float32 + and p.grad is not None + and p.grad.dtype == torch.float32 + and torch.isfinite(p.grad).all() for p in model.parameters() ) before = model.image_out.weight.detach().clone() @@ -56,14 +64,18 @@ def record(module, inputs, output): @pytest.mark.parametrize("steps", [3, 20]) -def test_every_mixed_sampler_update_uses_fp32_arithmetic_before_clamping(monkeypatch, steps): +def test_every_mixed_sampler_update_uses_fp32_arithmetic_before_clamping( + monkeypatch, steps +): model = image_model(attention_dtype="bfloat16") batch = {"image": example()["image"]} image = batch["image"] valid, observed = image["valid"], image["observed"] - fixed_prediction = torch.linspace(-.73, .91, image["value"].numel()).reshape_as( - image["value"] - ).bfloat16() + fixed_prediction = ( + torch.linspace(-0.73, 0.91, image["value"].numel()) + .reshape_as(image["value"]) + .bfloat16() + ) actual_inputs, actual_levels = [], [] def predict(batch, embodiment, image_values, levels, phase, **execution): @@ -79,7 +91,8 @@ def predict(batch, embodiment, image_values, levels, phase, **execution): )["image"] anchor = torch.where(observed[..., None], image["value"], 0) current = torch.where( - observed[..., None], anchor, + observed[..., None], + anchor, torch.randn(image["value"].shape, generator=torch.Generator().manual_seed(101)), ).masked_fill(~valid[..., None], 0) prediction_fp32 = fixed_prediction.float() diff --git a/tests/test_image_proprio_no_phase.py b/tests/test_image_proprio_no_phase.py index 8fd0bc552..d54c13076 100644 --- a/tests/test_image_proprio_no_phase.py +++ b/tests/test_image_proprio_no_phase.py @@ -28,7 +28,10 @@ def test_legacy_default_has_identical_seeded_state_and_strict_reload(): torch.manual_seed(37) explicit = image_model(use_phase_embedding=True) assert original.use_phase_embedding and "phase.weight" in original.state_dict() - assert all(torch.equal(value, explicit.state_dict()[key]) for key, value in original.state_dict().items()) + assert all( + torch.equal(value, explicit.state_dict()[key]) + for key, value in original.state_dict().items() + ) explicit.load_state_dict(original.state_dict(), strict=True) with pytest.raises(ValueError, match="use_phase_embedding must be bool"): image_model(use_phase_embedding=0) @@ -36,14 +39,18 @@ def test_legacy_default_has_identical_seeded_state_and_strict_reload(): def test_no_phase_has_exact_parameter_delta_and_all_stage_one_gradients(): with torch.device("meta"): - model = image_model(image_dim=768, width=768, heads=6, depth=8, use_phase_embedding=False) + model = image_model( + image_dim=768, width=768, heads=6, depth=8, use_phase_embedding=False + ) assert sum(p.numel() for p in model.parameters()) == 58_481_664 assert not hasattr(model, "phase") assert not any("phase" in key for key in model.state_dict()) model = image_model(use_phase_embedding=False) batch = {"image": joint_batch()["image"]} model.training_losses(batch, "gripper")["total"].backward() - assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) + assert all( + p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() + ) with pytest.raises(ValueError, match="no proprio readout phase"): predict(model, batch, phase="B") @@ -55,7 +62,10 @@ def test_no_phase_readout_warm_start_preserves_images_and_routes_state_commands( config.update(use_phase_embedding=False, image_dim=8) target_config.update(pose_dims={"gripper": 3}, action_dims={"gripper": 2}) validate_transition(source_config, target_config) - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) transfer = transfer_shared_weights(source, target) assert "phase.weight" not in transfer["expanded_embeddings"] assert not hasattr(source, "phase") and not hasattr(target, "phase") @@ -63,15 +73,27 @@ def test_no_phase_readout_warm_start_preserves_images_and_routes_state_commands( clone = ImageProprioModel(**target_config) clone.load_state_dict(target.state_dict(), strict=True) policy = configure_readout_trainability(target, "frozen_image") - assert not policy["phase_embedding_present"] and not policy["phase_embedding_frozen"] - assert sum(p.numel() for p in target.parameters() if not p.requires_grad) == sum(p.numel() for p in source.parameters()) + assert ( + not policy["phase_embedding_present"] and not policy["phase_embedding_frozen"] + ) + assert sum(p.numel() for p in target.parameters() if not p.requires_grad) == sum( + p.numel() for p in source.parameters() + ) batch = joint_batch() if mode == "discrete": for stream in batch.values(): - stream["index"] = torch.arange(stream["valid"].shape[1]).expand_as(stream["valid"]) - original = source.sample({"image": batch["image"]}, "gripper", steps=20, - generator=torch.Generator().manual_seed(19)) - actual = target.sample(batch, "gripper", steps=20, generator=torch.Generator().manual_seed(19)) + stream["index"] = torch.arange(stream["valid"].shape[1]).expand_as( + stream["valid"] + ) + original = source.sample( + {"image": batch["image"]}, + "gripper", + steps=20, + generator=torch.Generator().manual_seed(19), + ) + actual = target.sample( + batch, "gripper", steps=20, generator=torch.Generator().manual_seed(19) + ) assert torch.equal(original["image"], actual["image"]) assert set(actual) == {"image", "proprio", "action"} assert all(torch.isfinite(value).all() for value in actual.values()) @@ -88,6 +110,9 @@ def test_initializer_refuses_to_drop_or_add_phase_parameters_implicitly(): target_config["use_phase_embedding"] = False with pytest.raises(ValueError, match="configuration mismatch.*use_phase_embedding"): validate_transition(source_config, target_config) - source, target = ImageProprioModel(**source_config), ImageProprioModel(**target_config) + source, target = ( + ImageProprioModel(**source_config), + ImageProprioModel(**target_config), + ) with pytest.raises(ValueError, match="use_phase_embedding must match"): transfer_shared_weights(source, target) diff --git a/tests/test_image_proprio_positions.py b/tests/test_image_proprio_positions.py index 3e966f0fb..811329c4b 100644 --- a/tests/test_image_proprio_positions.py +++ b/tests/test_image_proprio_positions.py @@ -31,9 +31,13 @@ def one_thread(): def spatial_model(**kwargs): - settings = dict(noise_conditioning="adaln_zero", use_phase_embedding=False, - position_encoding="spatiotemporal_rope", rope_axis_dims=[4, 2, 2], - spatial_grid_size=[2, 2]) + settings = dict( + noise_conditioning="adaln_zero", + use_phase_embedding=False, + position_encoding="spatiotemporal_rope", + rope_axis_dims=[4, 2, 2], + spatial_grid_size=[2, 2], + ) settings.update(kwargs) return model(**settings) @@ -41,7 +45,9 @@ def spatial_model(**kwargs): def test_axes_use_centered_patch_units_and_correct_half_split_pairs(): times = torch.tensor([[0.0, 1.0]]) xy = torch.tensor([[[-0.5, -0.5], [0.5, 0.5]]]) - rope = prepare_spatiotemporal_rope(times, xy, 8, axis_dims=[4, 2, 2], spatial_grid_size=[2, 2]) + rope = prepare_spatiotemporal_rope( + times, xy, 8, axis_dims=[4, 2, 2], spatial_grid_size=[2, 2] + ) expected_phase = torch.tensor([[[[0.0, 0.0, -0.5, -0.5], [1.0, 0.01, 0.5, 0.5]]]]) torch.testing.assert_close(rope[0], expected_phase.cos()) torch.testing.assert_close(rope[1], expected_phase.sin()) @@ -59,9 +65,14 @@ def test_joint_coordinate_translation_preserves_attention_dot_products(): times, xy = torch.randn(2, 5), torch.randn(2, 5, 2) arguments = dict(head_dim=16, axis_dims=[8, 4, 4], spatial_grid_size=[10, 10]) original = prepare_spatiotemporal_rope(times, xy, **arguments) - shifted = prepare_spatiotemporal_rope(times + 13, xy + torch.tensor([0.2, -0.3]), **arguments) + shifted = prepare_spatiotemporal_rope( + times + 13, xy + torch.tensor([0.2, -0.3]), **arguments + ) + def logits(table): - return apply_temporal_rope(q, table) @ apply_temporal_rope(k, table).transpose(-1, -2) + return apply_temporal_rope(q, table) @ apply_temporal_rope(k, table).transpose( + -1, -2 + ) torch.testing.assert_close(logits(original), logits(shifted), rtol=1e-5, atol=1e-5) with torch.autocast("cpu", dtype=torch.bfloat16): @@ -85,18 +96,32 @@ def test_hydra_geometry_and_discrete_integer_clock_match_float_tables(): assert resolve_spatiotemporal_rope_geometry(8, axes, grid) == ((4, 2, 2), (2, 2)) times = torch.tensor([[0, 1, 3]], dtype=torch.int64) xy = torch.zeros(1, 3, 2) - expected = prepare_spatiotemporal_rope(times.float(), xy, 8, axis_dims=axes, spatial_grid_size=grid) - actual = prepare_spatiotemporal_rope(times, xy, 8, axis_dims=axes, spatial_grid_size=grid) + expected = prepare_spatiotemporal_rope( + times.float(), xy, 8, axis_dims=axes, spatial_grid_size=grid + ) + actual = prepare_spatiotemporal_rope( + times, xy, 8, axis_dims=axes, spatial_grid_size=grid + ) assert all(torch.equal(a, b) for a, b in zip(expected, actual)) - network = spatial_model(backend="sdpa", attention_dtype="native", time_mode="discrete", - rope_axis_dims=axes, spatial_grid_size=grid) + network = spatial_model( + backend="sdpa", + attention_dtype="native", + time_mode="discrete", + rope_axis_dims=axes, + spatial_grid_size=grid, + ) world = batch() world["image"]["index"] = world["image"]["index"].long() noisy, levels = fixed_corruption(world, 8) - assert torch.isfinite(network(world, "libero_gripper", noisy, levels)["image"]).all() + assert torch.isfinite( + network(world, "libero_gripper", noisy, levels)["image"] + ).all() -@pytest.mark.parametrize("axes,grid", [([4, 3, 1], [2, 2]), ([4, 2, 4], [2, 2]), ([4, 2, 2], None), ([4, 2, 2], [0, 2])]) +@pytest.mark.parametrize( + "axes,grid", + [([4, 3, 1], [2, 2]), ([4, 2, 4], [2, 2]), ([4, 2, 2], None), ([4, 2, 2], [0, 2])], +) def test_explicit_spatial_geometry_is_validated(axes, grid): with pytest.raises(ValueError): resolve_spatiotemporal_rope_geometry(8, axes, grid) @@ -104,9 +129,16 @@ def test_explicit_spatial_geometry_is_validated(axes, grid): def test_new_parameter_count_and_legacy_rng_state_compatibility(): with torch.device("meta"): - full = ImageProprioModel(objective="visual_only", phase_a_condition_on_proprio=False, - noise_conditioning="adaln_zero", use_phase_embedding=False, max_cameras=2, - position_encoding="spatiotemporal_rope", rope_axis_dims=[64, 32, 32], spatial_grid_size=[10, 10]) + full = ImageProprioModel( + objective="visual_only", + phase_a_condition_on_proprio=False, + noise_conditioning="adaln_zero", + use_phase_embedding=False, + max_cameras=2, + position_encoding="spatiotemporal_rope", + rope_axis_dims=[64, 32, 32], + spatial_grid_size=[10, 10], + ) assert sum(parameter.numel() for parameter in full.parameters()) == 88_178_688 assert not hasattr(full, "xy") and hasattr(full, "camera") assert not any(name.startswith("xy.") for name in full.state_dict()) @@ -114,16 +146,25 @@ def test_new_parameter_count_and_legacy_rng_state_compatibility(): default = model() rng = torch.get_rng_state().clone() torch.manual_seed(29) - explicit = model(position_encoding="linear_xy_temporal_rope", rope_axis_dims=None, spatial_grid_size=None) + explicit = model( + position_encoding="linear_xy_temporal_rope", + rope_axis_dims=None, + spatial_grid_size=None, + ) assert torch.equal(rng, torch.get_rng_state()) - assert all(torch.equal(value, explicit.state_dict()[name]) for name, value in default.state_dict().items()) + assert all( + torch.equal(value, explicit.state_dict()[name]) + for name, value in default.state_dict().items() + ) assert resolve_spatiotemporal_rope_geometry(8, None, [2, 2])[0] == (4, 2, 2) @pytest.mark.parametrize("mixed", [False, True]) @pytest.mark.parametrize("time_mode", ["continuous", "discrete"]) @pytest.mark.usefixtures("reference_kernels") -def test_nonzero_spatial_packed_outputs_and_all_gradients_match_padded(mixed, time_mode): +def test_nonzero_spatial_packed_outputs_and_all_gradients_match_padded( + mixed, time_mode +): padded = spatial_model(time_mode=time_mode) activate_conditioning(padded) packed = spatial_model(time_mode=time_mode, token_execution="packed_valid") @@ -158,8 +199,9 @@ def prepare(times, xy=None): output = network.sample(world, "libero_gripper", steps=4)["image"] assert len(calls) == 1 and torch.isfinite(output).all() times = torch.tensor([[1.0, 2.0]]) - _, rope = network._prepare_execution([times, times], [torch.ones(1, 2, dtype=torch.bool)] * 2, - torch.ones(1, 2, 2)) + _, rope = network._prepare_execution( + [times, times], [torch.ones(1, 2, dtype=torch.bool)] * 2, torch.ones(1, 2, 2) + ) assert all(torch.equal(a[..., :2], b[..., :2]) for a, b in zip(*rope)) assert rope[1][1][..., 2:].count_nonzero() == 0 @@ -167,8 +209,13 @@ def prepare(times, xy=None): def test_spatial_fullgraph_tensor_regions_match_eager(): eager = spatial_model(backend="sdpa", attention_dtype="native") activate_conditioning(eager) - compiled = spatial_model(backend="sdpa", attention_dtype="native", compile_mode="default", - compile_backend="aot_eager", compile_dynamic=True) + compiled = spatial_model( + backend="sdpa", + attention_dtype="native", + compile_mode="default", + compile_backend="aot_eager", + compile_dynamic=True, + ) compiled.load_state_dict(eager.state_dict(), strict=True) world = batch(nonprefix=True) noisy, levels = fixed_corruption(world, 9) @@ -179,8 +226,13 @@ def test_spatial_fullgraph_tensor_regions_match_eager(): results.append(result.detach()) torch.testing.assert_close(*results, rtol=1e-5, atol=1e-6) for name, parameter in eager.named_parameters(): - torch.testing.assert_close(parameter.grad, dict(compiled.named_parameters())[name].grad, - rtol=2e-5, atol=2e-6, msg=name) + torch.testing.assert_close( + parameter.grad, + dict(compiled.named_parameters())[name].grad, + rtol=2e-5, + atol=2e-6, + msg=name, + ) def test_full_image_model_has_no_additive_xy_and_keeps_relative_time_origin(): diff --git a/tests/test_pixel_attention_profiles.py b/tests/test_pixel_attention_profiles.py index f8131bef4..69b484279 100644 --- a/tests/test_pixel_attention_profiles.py +++ b/tests/test_pixel_attention_profiles.py @@ -82,9 +82,9 @@ def test_pixel_profile_override_preserves_science( assert not model.activation_checkpointing assert scientific_stage(cfg) == scientific_stage(baseline) for section in ("data", "callbacks"): - assert OmegaConf.to_container(cfg[section], resolve=True) == OmegaConf.to_container( - baseline[section], resolve=True - ) + assert OmegaConf.to_container( + cfg[section], resolve=True + ) == OmegaConf.to_container(baseline[section], resolve=True) trainer = OmegaConf.to_container(cfg.trainer, resolve=True) baseline_trainer = OmegaConf.to_container(baseline.trainer, resolve=True) assert trainer.pop("precision") == trainer_precision @@ -106,8 +106,10 @@ def test_pixel_profile_override_preserves_science( assert not model.allow_patch_compression assert pixel_stage(cfg, "pixel_loss").weight == 1 assert pixel_stage(cfg, "pixel_sampler").sampler.steps == 20 - assert all("warmup_steps" not in stage and "pose_ramp_steps" not in stage - for stage in cfg.model.pipeline.stages) + assert all( + "warmup_steps" not in stage and "pose_ramp_steps" not in stage + for stage in cfg.model.pipeline.stages + ) assert cfg.data.mode == arm and model.time_mode == arm if arm == "discrete": assert cfg.data.sampling_strategy == "discrete_window" @@ -145,12 +147,18 @@ def test_pixel_native_precision_fallback_remains_cpu_compatible(arm): assert cfg.trainer.precision == "32-true" -@pytest.mark.parametrize("name", ["video_only", "video_discrete", "readout_continuous", "readout_discrete"]) +@pytest.mark.parametrize( + "name", ["video_only", "video_discrete", "readout_continuous", "readout_discrete"] +) def test_continuous_stage_one_cap_does_not_change_other_recipe_defaults(name): with initialize_config_dir(config_dir=str(CONFIG_ROOT), version_base=None): - cfg = compose(config_name="train_zarr_cartesian", overrides=[ - f"+experiment=pixel_proprio/{name}", "hydra/launcher=basic", - ]) + cfg = compose( + config_name="train_zarr_cartesian", + overrides=[ + f"+experiment=pixel_proprio/{name}", + "hydra/launcher=basic", + ], + ) assert cfg.data.min_frames == 2 and cfg.data.max_frames == 128 assert cfg.data.validation_sampling_strategy == "discrete_window" assert cfg.data.validation_discrete_frames == 16 @@ -178,7 +186,13 @@ def test_four_gpu_profile_changes_only_execution_and_global_batch(arm): assert actual_data == original_data actual_trainer = OmegaConf.to_container(cfg.trainer, resolve=True) original_trainer = OmegaConf.to_container(baseline.trainer, resolve=True) - for key in ("strategy", "devices", "num_nodes", "accumulate_grad_batches", "use_distributed_sampler"): + for key in ( + "strategy", + "devices", + "num_nodes", + "accumulate_grad_batches", + "use_distributed_sampler", + ): actual_trainer.pop(key, None) original_trainer.pop(key, None) assert actual_trainer == original_trainer diff --git a/tests/test_pixel_compile_ddp.py b/tests/test_pixel_compile_ddp.py index fdf605381..33cffb8a7 100644 --- a/tests/test_pixel_compile_ddp.py +++ b/tests/test_pixel_compile_ddp.py @@ -17,7 +17,9 @@ def __init__(self, *, compiled=False): self.norm = torch.nn.LayerNorm(4) self.linear = torch.nn.Linear(4, 3) self.execute = compile_region( - self.region, mode="default" if compiled else "none", dynamic=True, + self.region, + mode="default" if compiled else "none", + dynamic=True, backend="pixel_region_ddp_test", ) @@ -45,7 +47,10 @@ def checked_backend(graph, inputs): register_backend(checked_backend, name="pixel_region_ddp_test") dist.init_process_group( - "gloo", init_method=Path(directory, "rendezvous").as_uri(), rank=rank, world_size=2, + "gloo", + init_method=Path(directory, "rendezvous").as_uri(), + rank=rank, + world_size=2, ) try: torch.manual_seed(42) @@ -63,14 +68,22 @@ def checked_backend(graph, inputs): for actual, wanted in zip(network.parameters(), reference.parameters()): torch.testing.assert_close(actual.grad, wanted.grad, rtol=1e-5, atol=1e-6) torch.optim.SGD(network.parameters(), lr=0.01).step() - state = torch.cat([parameter.detach().reshape(-1) for parameter in network.parameters()]) + state = torch.cat( + [parameter.detach().reshape(-1) for parameter in network.parameters()] + ) ranks = [torch.empty_like(state) for _ in range(2)] dist.all_gather(ranks, state) assert torch.equal(ranks[0], ranks[1]) - Path(directory, f"rank{rank}.json").write_text(json.dumps({ - "rank": rank, "gradient_average_verified": True, - "compiled_regions": len(policies), "identical_updated_weights": True, - })) + Path(directory, f"rank{rank}.json").write_text( + json.dumps( + { + "rank": rank, + "gradient_average_verified": True, + "compiled_regions": len(policies), + "identical_updated_weights": True, + } + ) + ) finally: dist.destroy_process_group() @@ -78,9 +91,12 @@ def checked_backend(graph, inputs): @pytest.mark.skipif(not dist.is_gloo_available(), reason="Gloo is unavailable") def test_regional_compile_preserves_two_rank_ddp_gradient_average(tmp_path): torch.multiprocessing.spawn(_worker, args=(str(tmp_path),), nprocs=2, join=True) - assert all(json.loads((tmp_path / f"rank{rank}.json").read_text())[ - "gradient_average_verified" - ] for rank in range(2)) + assert all( + json.loads((tmp_path / f"rank{rank}.json").read_text())[ + "gradient_average_verified" + ] + for rank in range(2) + ) def test_region_restores_ddp_policy_after_backend_failure(monkeypatch): diff --git a/tests/test_pixel_execution_benchmark.py b/tests/test_pixel_execution_benchmark.py index 987877a0f..5c6ff1137 100644 --- a/tests/test_pixel_execution_benchmark.py +++ b/tests/test_pixel_execution_benchmark.py @@ -23,7 +23,9 @@ def test_mixed_schedule_matches_frame_histogram_across_batch_sizes(batch_size): assert benchmark.complete_cycles(10, schedule) >= 10 -@pytest.mark.parametrize("scenario,capacity", [("fixed16", 3200), ("mixed", 25600), ("worst128", 25600)]) +@pytest.mark.parametrize( + "scenario,capacity", [("fixed16", 3200), ("mixed", 25600), ("worst128", 25600)] +) def test_benchmark_binds_exact_architecture_and_explicit_loss(scenario, capacity): arguments = benchmark.model_arguments("flash4", scenario) with torch.device("meta"): @@ -42,14 +44,23 @@ def test_benchmark_binds_exact_architecture_and_explicit_loss(scenario, capacity def small_batch(): - return benchmark.make_batch((2, 3), seed=17, device=torch.device("cpu"), image_size=4, patch_size=2) + return benchmark.make_batch( + (2, 3), seed=17, device=torch.device("cpu"), image_size=4, patch_size=2 + ) def small_model(): return ImageProprioModel( - image_dim=12, width=16, depth=1, heads=2, max_cameras=2, max_tokens=32, - objective="visual_only", phase_a_condition_on_proprio=False, - attention_backend="sdpa", attention_dtype="bfloat16", + image_dim=12, + width=16, + depth=1, + heads=2, + max_cameras=2, + max_tokens=32, + objective="visual_only", + phase_a_condition_on_proprio=False, + attention_backend="sdpa", + attention_dtype="bfloat16", loss_reduction="mean_per_example", ) @@ -80,14 +91,22 @@ def test_mean_per_example_is_not_padding_or_token_weighted(): def test_batched_probe_checks_outputs_gradients_and_cross_example_isolation(): # Flash4's context means standard BF16 autocast. The tiny CPU model uses # native SDPA explicitly; hardware/kernel certification happens on CUDA. - result = benchmark.check_outputs(small_model(), small_batch(), backend="flash4", - seed=42, tolerance=0.03, gradients=True) + result = benchmark.check_outputs( + small_model(), + small_batch(), + backend="flash4", + seed=42, + tolerance=0.03, + gradients=True, + ) assert result["output_relative_l2"] <= 0.03 assert result["gradient_relative_l2"] <= 0.03 assert result["cross_example_relative_l2"] <= 1e-6 -def test_measured_training_loss_matches_explicit_mean_independent_example_contract(monkeypatch): +def test_measured_training_loss_matches_explicit_mean_independent_example_contract( + monkeypatch, +): model, batch = small_model(), small_batch() noisy, levels = benchmark.fixed_corruption(batch, 42) monkeypatch.setattr(model, "_noisy_images", lambda *args, **kwargs: (noisy, levels)) @@ -104,24 +123,38 @@ def test_eager_compiled_gate_accepts_equal_execution_and_rejects_bad_backward(): # CUDA Inductor. The scheduled harness supplies an actual compiled model. eager, candidate = small_model(), small_model() candidate.load_state_dict(eager.state_dict(), strict=True) - result = benchmark.check_eager_compiled(eager, candidate, small_batch(), - backend="flash4", seed=42, tolerance=0.03) + result = benchmark.check_eager_compiled( + eager, candidate, small_batch(), backend="flash4", seed=42, tolerance=0.03 + ) assert result["same_state_verified"] assert result["output_relative_l2"] == result["gradient_relative_l2"] == 0 assert result["loss_relative_error"] == 0 - hooks = [p.register_hook(lambda gradient: gradient * 2) for p in candidate.parameters()] - with pytest.raises(benchmark.CorrectnessError, match="Compiled gradient differs from eager"): - benchmark.check_eager_compiled(eager, candidate, small_batch(), - backend="flash4", seed=42, tolerance=0.03) - assert all(p.grad is None for model in (eager, candidate) for p in model.parameters()) + hooks = [ + p.register_hook(lambda gradient: gradient * 2) for p in candidate.parameters() + ] + with pytest.raises( + benchmark.CorrectnessError, match="Compiled gradient differs from eager" + ): + benchmark.check_eager_compiled( + eager, candidate, small_batch(), backend="flash4", seed=42, tolerance=0.03 + ) + assert all( + p.grad is None for model in (eager, candidate) for p in model.parameters() + ) for hook in hooks: hook.remove() def test_eager_compiled_gate_rejects_unmatched_weights(): with pytest.raises(benchmark.CorrectnessError, match="initial state differs"): - benchmark.check_eager_compiled(small_model(), small_model(), small_batch(), - backend="flash4", seed=42, tolerance=0.03) + benchmark.check_eager_compiled( + small_model(), + small_model(), + small_batch(), + backend="flash4", + seed=42, + tolerance=0.03, + ) def test_probe_rejects_cross_example_interference(): @@ -135,8 +168,14 @@ def leaking_forward(batch, *args, **kwargs): model.forward = leaking_forward with pytest.raises(benchmark.CorrectnessError, match="independent B1|leakage"): - benchmark.check_outputs(model, small_batch(), backend="flash4", seed=42, - tolerance=0.03, gradients=True) + benchmark.check_outputs( + model, + small_batch(), + backend="flash4", + seed=42, + tolerance=0.03, + gradients=True, + ) def test_leakage_gate_does_not_confuse_grad_mode_rounding_with_interference(): @@ -150,8 +189,9 @@ def mode_dependent_forward(*args, **kwargs): return result model.forward = mode_dependent_forward - result = benchmark.check_outputs(model, small_batch(), backend="flash4", seed=42, - tolerance=0.03, gradients=True) + result = benchmark.check_outputs( + model, small_batch(), backend="flash4", seed=42, tolerance=0.03, gradients=True + ) assert 1e-6 < result["grad_vs_no_grad_relative_l2"] < 0.03 assert result["cross_example_relative_l2"] == 0 @@ -169,7 +209,9 @@ def test_cpu_worker_is_explicitly_unsupported_without_fallback(tmp_path, monkeyp @pytest.mark.parametrize("legacy", [False, True]) -def test_fa3_inventory_matches_canonical_or_absent_package_alias(tmp_path, monkeypatch, legacy): +def test_fa3_inventory_matches_canonical_or_absent_package_alias( + tmp_path, monkeypatch, legacy +): module_file = tmp_path / "interface.py" module_file.write_text("# exact fake backend source\n") calls = [] @@ -182,7 +224,9 @@ def import_module(name): monkeypatch.setattr(benchmark.importlib, "import_module", import_module) inventory = benchmark.backend_inventory("flash3") - expected = ["flash_attn_3.flash_attn_interface"] + (["flash_attn_interface"] if legacy else []) + expected = ["flash_attn_3.flash_attn_interface"] + ( + ["flash_attn_interface"] if legacy else [] + ) assert calls == expected assert inventory["module"] == expected[-1] assert inventory["sha256"] == benchmark.digest(module_file) @@ -193,10 +237,14 @@ def test_fa3_installed_but_broken_dependency_never_uses_alias(monkeypatch): def import_module(name): calls.append(name) - raise ModuleNotFoundError("Missing installed backend dependency", name="flash_attn_3_cuda") + raise ModuleNotFoundError( + "Missing installed backend dependency", name="flash_attn_3_cuda" + ) monkeypatch.setattr(benchmark.importlib, "import_module", import_module) - with pytest.raises(benchmark.UnsupportedBackend, match="Required backend import failed"): + with pytest.raises( + benchmark.UnsupportedBackend, match="Required backend import failed" + ): benchmark.backend_inventory("flash3") assert calls == ["flash_attn_3.flash_attn_interface"] @@ -204,31 +252,56 @@ def import_module(name): def test_worker_subprocess_and_timeout_are_bounded(tmp_path): successful = tmp_path / "success" successful.mkdir() - code, timed_out = benchmark.run_worker_process([sys.executable, "-c", "print('done')"], successful, 10) + code, timed_out = benchmark.run_worker_process( + [sys.executable, "-c", "print('done')"], successful, 10 + ) assert code == 0 and not timed_out assert (successful / "stdout.log").read_text().strip() == "done" timeout = tmp_path / "timeout" timeout.mkdir() - code, timed_out = benchmark.run_worker_process([sys.executable, "-c", "import time; time.sleep(30)"], timeout, 0.1) + code, timed_out = benchmark.run_worker_process( + [sys.executable, "-c", "import time; time.sleep(30)"], timeout, 0.1 + ) assert timed_out and code != 0 -@pytest.mark.parametrize("token_execution,suffix", [("padded", ""), ("packed_valid", "-packed_valid")]) -def test_token_execution_cli_is_explicit_and_keeps_legacy_padded_ids(monkeypatch, capsys, token_execution, suffix): - monkeypatch.setattr(sys, "argv", [ - "benchmark_pixel_execution.py", "--dry-run", "--backends", "flash3", - "--scenarios", "mixed", "--batch-sizes", "2", "--compile-mode", "default", - "--compile-dynamic", "--token-execution", token_execution, - ]) +@pytest.mark.parametrize( + "token_execution,suffix", [("padded", ""), ("packed_valid", "-packed_valid")] +) +def test_token_execution_cli_is_explicit_and_keeps_legacy_padded_ids( + monkeypatch, capsys, token_execution, suffix +): + monkeypatch.setattr( + sys, + "argv", + [ + "benchmark_pixel_execution.py", + "--dry-run", + "--backends", + "flash3", + "--scenarios", + "mixed", + "--batch-sizes", + "2", + "--compile-mode", + "default", + "--compile-dynamic", + "--token-execution", + token_execution, + ], + ) benchmark.main() plan = json.loads(capsys.readouterr().out) - candidate, = plan["candidates"] + (candidate,) = plan["candidates"] assert candidate["id"] == "flash3-mixed-b2-default" + suffix assert candidate["token_execution"] == token_execution assert candidate["compile_mode"] == "default" and candidate["compile_dynamic"] arguments = benchmark.model_arguments( - candidate["backend"], candidate["scenario"], compile_mode=candidate["compile_mode"], - compile_dynamic=candidate["compile_dynamic"], token_execution=candidate["token_execution"], + candidate["backend"], + candidate["scenario"], + compile_mode=candidate["compile_mode"], + compile_dynamic=candidate["compile_dynamic"], + token_execution=candidate["token_execution"], ) assert arguments["token_execution"] == token_execution # The compiled/eager gate changes compilation only, not token layout. @@ -236,12 +309,22 @@ def test_token_execution_cli_is_explicit_and_keeps_legacy_padded_ids(monkeypatch def test_token_execution_cli_defaults_to_archived_padded_id(monkeypatch, capsys): - monkeypatch.setattr(sys, "argv", [ - "benchmark_pixel_execution.py", "--dry-run", "--backends", "flash4", - "--scenarios", "fixed16", "--batch-sizes", "1", - ]) + monkeypatch.setattr( + sys, + "argv", + [ + "benchmark_pixel_execution.py", + "--dry-run", + "--backends", + "flash4", + "--scenarios", + "fixed16", + "--batch-sizes", + "1", + ], + ) benchmark.main() - candidate, = json.loads(capsys.readouterr().out)["candidates"] + (candidate,) = json.loads(capsys.readouterr().out)["candidates"] assert candidate["id"] == "flash4-fixed16-b1-none" assert candidate["token_execution"] == "padded" assert benchmark.model_arguments("flash4", "fixed16")["token_execution"] == "padded" @@ -257,15 +340,22 @@ def test_compiled_reference_cannot_change_token_execution(): compiled.load_state_dict(eager.state_dict(), strict=True) compiled.token_execution = "packed_valid" with pytest.raises(benchmark.CorrectnessError, match="token execution differs"): - benchmark.check_eager_compiled(eager, compiled, small_batch(), backend="flash4", seed=42, tolerance=.03) + benchmark.check_eager_compiled( + eager, compiled, small_batch(), backend="flash4", seed=42, tolerance=0.03 + ) -def test_worker_reports_unsupported_ragged_backend_without_fallback(tmp_path, monkeypatch): +def test_worker_reports_unsupported_ragged_backend_without_fallback( + tmp_path, monkeypatch +): def unsupported(*args): raise ValueError("Ragged packed_valid execution requires flash3 or flash4") + monkeypatch.setattr(benchmark, "run_candidate", unsupported) request = tmp_path / "request.json" - benchmark.write_once(request, {"backend": "sdpa", "token_execution": "packed_valid"}) + benchmark.write_once( + request, {"backend": "sdpa", "token_execution": "packed_valid"} + ) benchmark.worker(request) result = json.loads((tmp_path / "result.json").read_text()) assert result["status"] == "UNSUPPORTED" @@ -274,10 +364,18 @@ def unsupported(*args): def test_adaln_correctness_probe_is_nontrivial_and_restores_state_and_rng(): model = ImageProprioModel( - image_dim=12, width=16, depth=1, heads=2, max_cameras=2, max_tokens=32, - objective="visual_only", phase_a_condition_on_proprio=False, - use_phase_embedding=False, noise_conditioning="adaln_zero", - attention_backend="sdpa", attention_dtype="bfloat16", + image_dim=12, + width=16, + depth=1, + heads=2, + max_cameras=2, + max_tokens=32, + objective="visual_only", + phase_a_condition_on_proprio=False, + use_phase_embedding=False, + noise_conditioning="adaln_zero", + attention_backend="sdpa", + attention_dtype="bfloat16", loss_reduction="mean_per_example", ) batch = small_batch() @@ -292,13 +390,18 @@ def test_adaln_correctness_probe_is_nontrivial_and_restores_state_and_rng(): second = model(batch, "libero_gripper", noisy, changed_levels)["image"] assert first.abs().sum() > 0 assert not torch.allclose(first, second) - result = benchmark.check_outputs(model, batch, backend="flash4", seed=42, - tolerance=0.03, gradients=True) + result = benchmark.check_outputs( + model, batch, backend="flash4", seed=42, tolerance=0.03, gradients=True + ) assert result["gradient_relative_l2"] < 0.03 assert torch.equal(rng, torch.get_rng_state()) - assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) + assert all( + torch.equal(value, model.state_dict()[name]) for name, value in state.items() + ) assert all(parameter.grad is None for parameter in model.parameters()) with pytest.raises(RuntimeError, match="intentional probe failure"): with benchmark.nontrivial_conditioning_probe(model, 20): raise RuntimeError("intentional probe failure") - assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) + assert all( + torch.equal(value, model.state_dict()[name]) for name, value in state.items() + ) diff --git a/tests/test_pixel_prediction_diagnostics.py b/tests/test_pixel_prediction_diagnostics.py index 44dc45f5c..c18838a0f 100644 --- a/tests/test_pixel_prediction_diagnostics.py +++ b/tests/test_pixel_prediction_diagnostics.py @@ -24,13 +24,16 @@ def test_wandb_media_uses_optimizer_axis_not_history_counter(tmp_path, monkeypat from lightning.pytorch.loggers import WandbLogger calls = [] - experiment = SimpleNamespace(log=lambda *args, **kwargs: calls.append((args, kwargs))) + experiment = SimpleNamespace( + log=lambda *args, **kwargs: calls.append((args, kwargs)) + ) logger = WandbLogger(experiment=experiment, offline=True) monkeypatch.setattr(wandb, "Image", lambda path: ("image", path)) diagnostic = PixelPredictionDiagnostics(tmp_path, log_wandb=True) diagnostic.written = [("robot", "forecast", "image", tmp_path / "example.png")] - model = SimpleNamespace(global_rank=0, global_step=250, - trainer=SimpleNamespace(loggers=[logger])) + model = SimpleNamespace( + global_rank=0, global_step=250, trainer=SimpleNamespace(loggers=[logger]) + ) diagnostic.log_to_wandb(model) assert len(calls) == 1 args, kwargs = calls[0] diff --git a/tests/test_train_hydra_torch_elastic.py b/tests/test_train_hydra_torch_elastic.py index 3588abfaf..efb8b1f00 100644 --- a/tests/test_train_hydra_torch_elastic.py +++ b/tests/test_train_hydra_torch_elastic.py @@ -19,18 +19,28 @@ def elastic(monkeypatch): monkeypatch.delenv(key, raising=False) for key, value in { "TORCHELASTIC_RUN_ID": "separate-run-rendezvous", - "RANK": "2", "LOCAL_RANK": "2", "WORLD_SIZE": "4", - "LOCAL_WORLD_SIZE": "4", "GROUP_RANK": "0", - "MASTER_ADDR": "localhost", "MASTER_PORT": "29451", - "SLURM_JOB_ID": "1234", "SLURM_NTASKS": "1", - "SLURM_PROCID": "0", "SLURM_LOCALID": "0", + "RANK": "2", + "LOCAL_RANK": "2", + "WORLD_SIZE": "4", + "LOCAL_WORLD_SIZE": "4", + "GROUP_RANK": "0", + "MASTER_ADDR": "localhost", + "MASTER_PORT": "29451", + "SLURM_JOB_ID": "1234", + "SLURM_NTASKS": "1", + "SLURM_PROCID": "0", + "SLURM_LOCALID": "0", }.items(): monkeypatch.setenv(key, value) - return OmegaConf.create({ - "mmap_checkpoint": False, - "runtime": {"cluster_environment": "torch_elastic", - "slurm_requeue_owner": "none"}, - }) + return OmegaConf.create( + { + "mmap_checkpoint": False, + "runtime": { + "cluster_environment": "torch_elastic", + "slurm_requeue_owner": "none", + }, + } + ) def test_torchrun_overrides_one_task_slurm_without_mutating_provenance(elastic): @@ -49,13 +59,23 @@ def test_torchrun_overrides_one_task_slurm_without_mutating_provenance(elastic): assert dict(os.environ) == before -@pytest.mark.parametrize("key,value", [ - ("RANK", None), ("LOCAL_RANK", "four"), ("WORLD_SIZE", "0"), - ("LOCAL_WORLD_SIZE", "3"), ("RANK", "4"), ("RANK", "-1"), - ("LOCAL_RANK", "4"), ("GROUP_RANK", "1"), - ("MASTER_ADDR", ""), ("MASTER_PORT", "0"), - ("MASTER_PORT", "65536"), ("MASTER_PORT", "invalid"), -]) +@pytest.mark.parametrize( + "key,value", + [ + ("RANK", None), + ("LOCAL_RANK", "four"), + ("WORLD_SIZE", "0"), + ("LOCAL_WORLD_SIZE", "3"), + ("RANK", "4"), + ("RANK", "-1"), + ("LOCAL_RANK", "4"), + ("GROUP_RANK", "1"), + ("MASTER_ADDR", ""), + ("MASTER_PORT", "0"), + ("MASTER_PORT", "65536"), + ("MASTER_PORT", "invalid"), + ], +) def test_malformed_torchrun_environment_fails(elastic, monkeypatch, key, value): if value is None: monkeypatch.delenv(key, raising=False) diff --git a/tests/test_world_model_distributed_rng.py b/tests/test_world_model_distributed_rng.py index ca4d6bba9..ab659dd09 100644 --- a/tests/test_world_model_distributed_rng.py +++ b/tests/test_world_model_distributed_rng.py @@ -40,22 +40,31 @@ def test_distributed_resume_rejects_legacy_or_changed_world_size(monkeypatch): monkeypatch.setattr(behavior, "_distributed_world", lambda: (2, 1)) with pytest.raises(ValueError, match="per-rank RNG"): behavior._rank_rng_state(local) - saved = {"schema": behavior._DISTRIBUTED_RNG_SCHEMA, - "world_size": 4, "ranks": [local] * 4} + saved = { + "schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 4, + "ranks": [local] * 4, + } with pytest.raises(ValueError, match="saved RNG world size"): behavior._rank_rng_state(saved) -@pytest.mark.parametrize("update,match", [ - ({"schema": "unknown"}, "Unknown"), - ({"ranks": []}, "one state per rank"), - ({"ranks": [{}, {}]}, "current CUDA device"), -]) +@pytest.mark.parametrize( + "update,match", + [ + ({"schema": "unknown"}, "Unknown"), + ({"ranks": []}, "one state per rank"), + ({"ranks": [{}, {}]}, "current CUDA device"), + ], +) def test_corrupt_distributed_rng_fails_closed(monkeypatch, update, match): local = behavior._capture_rng_state(current_cuda_only=True) monkeypatch.setattr(behavior, "_distributed_world", lambda: (2, 1)) - saved = {"schema": behavior._DISTRIBUTED_RNG_SCHEMA, - "world_size": 2, "ranks": [local, local]} + saved = { + "schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 2, + "ranks": [local, local], + } saved.update(update) with pytest.raises(ValueError, match=match): behavior._rank_rng_state(saved) @@ -66,9 +75,15 @@ def test_cuda_capture_and_restore_touch_only_current_device(monkeypatch): restored = [] monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_rng_state", lambda: state) - monkeypatch.setattr(torch.cuda, "get_rng_state_all", lambda: pytest.fail("all devices")) - monkeypatch.setattr(torch.cuda, "set_rng_state", lambda value: restored.append(value)) - monkeypatch.setattr(torch.cuda, "set_rng_state_all", lambda _: pytest.fail("all devices")) + monkeypatch.setattr( + torch.cuda, "get_rng_state_all", lambda: pytest.fail("all devices") + ) + monkeypatch.setattr( + torch.cuda, "set_rng_state", lambda value: restored.append(value) + ) + monkeypatch.setattr( + torch.cuda, "set_rng_state_all", lambda _: pytest.fail("all devices") + ) saved = behavior._capture_rng_state(current_cuda_only=True) assert saved["cuda_scope"] == "current_device" behavior._restore_rng_state(saved) @@ -99,10 +114,16 @@ def test_fresh_training_seeds_ranks_once_without_changing_initialization(monkeyp instance, context = _bound_behavior(42) assert torch.equal(torch.get_rng_state(), initial_state) generator = torch.Generator().manual_seed(42 + rank) - numpy_rng, python_rng = np.random.RandomState(42 + rank), random.Random(42 + rank) + numpy_rng, python_rng = ( + np.random.RandomState(42 + rank), + random.Random(42 + rank), + ) for step in range(2): - expected = (float(torch.rand((), generator=generator)), - float(numpy_rng.rand()), python_rng.random()) + expected = ( + float(torch.rand((), generator=generator)), + float(numpy_rng.rand()), + python_rng.random(), + ) actual = instance.training_step({"probe": {}}, step) assert actual == expected if step == 0: @@ -120,10 +141,15 @@ def test_saved_rank_rng_takes_precedence_over_fresh_noise_seed(monkeypatch): local = behavior._capture_rng_state(current_cuda_only=True) expected = _Context()._default_training_step({}, 0) instance, context = _bound_behavior(42) - instance.on_load_checkpoint({"mot_rng_state": { - "schema": behavior._DISTRIBUTED_RNG_SCHEMA, - "world_size": 2, "ranks": [local, local], - }}) + instance.on_load_checkpoint( + { + "mot_rng_state": { + "schema": behavior._DISTRIBUTED_RNG_SCHEMA, + "world_size": 2, + "ranks": [local, local], + } + } + ) assert instance.training_step({"probe": {}}, 0) == expected assert instance._pending_rng is None and instance._training_rng_initialized del context @@ -134,7 +160,9 @@ def test_fresh_seed_targets_current_cuda_device_only(monkeypatch): monkeypatch.setattr(behavior, "_distributed_world", lambda: (4, 2)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "manual_seed", seeds.append) - monkeypatch.setattr(torch.cuda, "manual_seed_all", lambda _: pytest.fail("all devices")) + monkeypatch.setattr( + torch.cuda, "manual_seed_all", lambda _: pytest.fail("all devices") + ) instance, context = _bound_behavior(42) instance.training_step({"probe": {}}, 0) instance.training_step({"probe": {}}, 1) @@ -158,7 +186,7 @@ def test_distributed_noise_seed_rejects_invalid_values(seed): behavior.WorldModelTrainingBehavior(distributed_noise_seed=seed) -_WORKER = r''' +_WORKER = r""" import json import os from pathlib import Path @@ -215,30 +243,52 @@ def on_load_checkpoint(self, checkpoint): weights_only=False) trainer.save_checkpoint(out / (mode + '.ckpt')) (out / f'{mode}-rank{trainer.global_rank}.json').write_text(json.dumps(model.draws)) -''' +""" def test_real_torchrun_checkpoint_resume_preserves_each_rank(tmp_path): - if not torch.distributed.is_available() or not torch.distributed.is_gloo_available(): + if ( + not torch.distributed.is_available() + or not torch.distributed.is_gloo_available() + ): pytest.skip("CPU distributed backend unavailable") worker = tmp_path / "worker.py" worker.write_text(textwrap.dedent(_WORKER)) env = dict(os.environ) - env.update({ - "PYTHONPATH": str(Path(__file__).resolve().parents[1]), - "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", - "MKL_NUM_THREADS": "1", "CUDA_VISIBLE_DEVICES": "", - # A one-task Slurm parent must not override torchrun's actual world. - "SLURM_JOB_ID": "999", "SLURM_NTASKS": "1", - "SLURM_PROCID": "0", "SLURM_LOCALID": "0", - }) + env.update( + { + "PYTHONPATH": str(Path(__file__).resolve().parents[1]), + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "CUDA_VISIBLE_DEVICES": "", + # A one-task Slurm parent must not override torchrun's actual world. + "SLURM_JOB_ID": "999", + "SLURM_NTASKS": "1", + "SLURM_PROCID": "0", + "SLURM_LOCALID": "0", + } + ) for key in ("ICE_REQUEUE_OWNER", "ICE_CHILD_REQUEUE_DISABLED"): env.pop(key, None) for mode in ("full", "partial", "resume"): - result = subprocess.run([ - sys.executable, "-m", "torch.distributed.run", "--standalone", - "--nnodes=1", "--nproc-per-node=2", str(worker), str(tmp_path), mode, - ], env=env, capture_output=True, text=True, timeout=90) + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc-per-node=2", + str(worker), + str(tmp_path), + mode, + ], + env=env, + capture_output=True, + text=True, + timeout=90, + ) assert result.returncode == 0, result.stdout + result.stderr first = [] for rank in range(2): @@ -257,4 +307,6 @@ def test_real_torchrun_checkpoint_resume_preserves_each_rank(tmp_path): saved = resume["mot_rng_state"] assert saved["schema"] == behavior._DISTRIBUTED_RNG_SCHEMA assert saved["world_size"] == len(saved["ranks"]) == 2 - assert not torch.equal(saved["ranks"][0]["torch_cpu"], saved["ranks"][1]["torch_cpu"]) + assert not torch.equal( + saved["ranks"][0]["torch_cpu"], saved["ranks"][1]["torch_cpu"] + ) From f791b17154d2f3a9417c49a4809dffd244c9448b Mon Sep 17 00:00:00 2001 From: ElmoPA Date: Sun, 18 Oct 2026 14:20:09 -0400 Subject: [PATCH 3/4] Audit prepared model configs without loading external artifacts --- egomimic/rldb/pixel_proprio_datamodule.py | 34 +++++ egomimic/rldb/world_model_datamodule.py | 108 ++++++++++--- scripts/audit_hydra_configs.py | 51 +++++++ tests/test_pixel_config_audit.py | 108 +++++++++++++ tests/test_prepared_component_audit.py | 177 ++++++++++++++++++++++ 5 files changed, 460 insertions(+), 18 deletions(-) create mode 100644 tests/test_pixel_config_audit.py create mode 100644 tests/test_prepared_component_audit.py diff --git a/egomimic/rldb/pixel_proprio_datamodule.py b/egomimic/rldb/pixel_proprio_datamodule.py index 838287deb..c606a3368 100644 --- a/egomimic/rldb/pixel_proprio_datamodule.py +++ b/egomimic/rldb/pixel_proprio_datamodule.py @@ -140,6 +140,40 @@ def __init__( if value is not None } + def preflight_configuration(self): + """Check both declared pixel samplers without reading a prepared cache.""" + result = super().preflight_configuration() + for split, settings in ( + ("train", self.settings), + ("val", self.settings | self.validation_settings), + ): + strategy = settings["sampling_strategy"] + if ( + (strategy, settings["mode"]) + not in ( + ("full_episode_uniform", "continuous"), + ("discrete_window", "continuous"), + ("discrete_window", "discrete"), + ) + or settings.get("max_past_seconds") is not None + or settings.get("max_future_seconds") is not None + ): + raise ValueError(f"Invalid {split} pixel sampling strategy/time bounds") + for key, minimum in (("discrete_frames", 2), ("stride", 1)): + value = settings[key] + if type(value) is not int or value < minimum: + raise ValueError(f"{split} {key} must be an integer >= {minimum}") + if ( + settings["task"] == "direct_action" + and strategy != "full_episode_uniform" + ): + raise ValueError("Direct action requires full_episode_uniform sampling") + return result | { + "validation_sampling": dict(self.validation_settings), + "distributed_world_size": self.distributed_world_size, + "validation_batch_size": self.validation_batch_size, + } + def setup(self, stage=None): if not self.datasets: verifier = ArtifactVerificationMemo() diff --git a/egomimic/rldb/world_model_datamodule.py b/egomimic/rldb/world_model_datamodule.py index e0a36065e..44da03e0f 100644 --- a/egomimic/rldb/world_model_datamodule.py +++ b/egomimic/rldb/world_model_datamodule.py @@ -1,6 +1,7 @@ """Prepared world-model caches through the existing trainHydra data boundary.""" import hashlib +import math from pathlib import Path from lightning import LightningDataModule @@ -10,6 +11,7 @@ from egomimic.rldb.world_model import ( HomogeneousWorldBatchSampler, PreparedWorldDataset, + _bounded_time_sampling, collate_world, ) from egomimic.rldb.world_model_preparation import ArtifactVerificationMemo @@ -39,7 +41,9 @@ def __init__( type(validation_batches_per_embodiment) is not int or validation_batches_per_embodiment < 1 ): - raise ValueError("validation_batches_per_embodiment must be a positive integer or None") + raise ValueError( + "validation_batches_per_embodiment must be a positive integer or None" + ) self.validation_batches_per_embodiment = validation_batches_per_embodiment self.batch_size, self.num_workers = batch_size, num_workers self.settings = dict( @@ -61,18 +65,74 @@ def __init__( self.datasets = {} self.checkpointable_sampling = bool(checkpointable_sampling) if self.checkpointable_sampling and self.num_workers != 0: - raise ValueError("Exact MoT sampling resume currently requires num_workers=0") + raise ValueError( + "Exact MoT sampling resume currently requires num_workers=0" + ) self.train_sampler = None self.consumed_batches = 0 + def preflight_configuration(self): + """Validate declared sampling/loader settings without opening artifacts. + + This checks configuration only. Manifest identities, episode lengths, + tensor schemas and normalization remain setup-time checks on real data. + """ + if ( + not isinstance(self.manifest_path, (str, Path)) + or not str(self.manifest_path).strip() + ): + raise ValueError("A nonempty prepared manifest path is required") + integers = { + "batch_size": (self.batch_size, 1), + "num_workers": (self.num_workers, 0), + "seed": (self.settings["seed"], 0), + "min_frames": (self.settings["min_frames"], 2), + "max_frames": (self.settings["max_frames"], 2), + "discrete_frames": (self.settings["discrete_frames"], 2), + "stride": (self.settings["stride"], 1), + } + for key, (value, minimum) in integers.items(): + if type(value) is not int or value < minimum: + raise ValueError(f"{key} must be an integer >= {minimum}") + if self.settings["mode"] not in ("continuous", "discrete"): + raise ValueError("Prepared sampling mode must be continuous or discrete") + if self.settings["max_frames"] < self.settings["min_frames"]: + raise ValueError("max_frames must be >= min_frames") + probability = self.settings["forecasting_probability"] + if ( + isinstance(probability, bool) + or not isinstance(probability, (int, float)) + or not math.isfinite(probability) + or not 0 <= probability <= 1 + ): + raise ValueError("forecasting_probability must be finite and in [0, 1]") + bounded = _bounded_time_sampling( + self.settings.get("max_past_seconds"), + self.settings.get("max_future_seconds"), + ) + if bounded and self.settings["mode"] != "continuous": + raise ValueError("Physical-time bounds apply only to continuous sampling") + return { + "scope": "configuration_only", + "artifacts_opened": False, + "sampling": dict(self.settings), + "batch_size": self.batch_size, + "num_workers": self.num_workers, + "checkpointable_sampling": self.checkpointable_sampling, + } + def setup(self, stage=None): if not self.datasets: # Reuse finalized fingerprints only for unchanged exact files; # otherwise hash once per setup and share that proof across splits. verification = ArtifactVerificationMemo() self.datasets = { - s: PreparedWorldDataset(self.manifest_path, split=s, - verification_memo=verification, **self.settings) + s: PreparedWorldDataset( + self.manifest_path, + split=s, + verification_memo=verification, + **self.settings, + ) for s in ("train", "val") } @@ -91,7 +151,10 @@ def prepare_context(self, *, mode, normalization, normalizer, restored_state=Non "sampling": self.settings, } if self.checkpointable_sampling: - state["sampling_resume"] = {"schema": "mot-sampling-state-v1", "batch_size": self.batch_size} + state["sampling_resume"] = { + "schema": "mot-sampling-state-v1", + "batch_size": self.batch_size, + } if self.validation_batches_per_embodiment is not None: state["validation_sampling"] = { "batches_per_embodiment": self.validation_batches_per_embodiment, @@ -129,15 +192,15 @@ def _loader(self, split): ) dataset = self.datasets[split] sampler = HomogeneousWorldBatchSampler( - dataset, - self.batch_size, - self.settings["seed"], - shuffle=split == "train", - checkpointable=self.checkpointable_sampling and split == "train", - validation_batches_per_embodiment=( - self.validation_batches_per_embodiment if split == "val" else None - ), - ) + dataset, + self.batch_size, + self.settings["seed"], + shuffle=split == "train", + checkpointable=self.checkpointable_sampling and split == "train", + validation_batches_per_embodiment=( + self.validation_batches_per_embodiment if split == "val" else None + ), + ) if split == "train": self.train_sampler = sampler sampler.consumed_batches = self.consumed_batches @@ -145,6 +208,7 @@ def _loader(self, split): generator = None if self.checkpointable_sampling: import torch + generator = torch.Generator().manual_seed(self.settings["seed"]) return DataLoader( dataset, @@ -166,14 +230,22 @@ def acknowledge_batch(self, cursor): def state_dict(self): if not self.checkpointable_sampling: return {} - return {"schema": "mot-sampling-state-v1", "consumed_batches": self.consumed_batches, - "manifest_sha256": hashlib.sha256(Path(self.manifest_path).read_bytes()).hexdigest(), - "settings": self.settings, "batch_size": self.batch_size} + return { + "schema": "mot-sampling-state-v1", + "consumed_batches": self.consumed_batches, + "manifest_sha256": hashlib.sha256( + Path(self.manifest_path).read_bytes() + ).hexdigest(), + "settings": self.settings, + "batch_size": self.batch_size, + } def load_state_dict(self, state_dict): if not self.checkpointable_sampling: if state_dict: - raise ValueError("Checkpointed sampling requires its original opt-in configuration") + raise ValueError( + "Checkpointed sampling requires its original opt-in configuration" + ) return expected = self.state_dict() for key in ("schema", "manifest_sha256", "settings", "batch_size"): diff --git a/scripts/audit_hydra_configs.py b/scripts/audit_hydra_configs.py index 1ffd6d616..5963d7435 100644 --- a/scripts/audit_hydra_configs.py +++ b/scripts/audit_hydra_configs.py @@ -19,6 +19,14 @@ CONFIGS = ROOT / "egomimic/hydra_configs" EXPERIMENT_CONTEXTS = { + "data/world_model/prepared": "world_model/continuous", + "model/world_model/base": "world_model/continuous", + "data/dino_proprio/prepared": "dino_proprio/visual_state", + "model/dino_proprio/base": "dino_proprio/visual_state", + "data/pixel_proprio/prepared": "pixel_proprio/video_continuous", + "model/pixel_proprio/base": "pixel_proprio/video_continuous", + "model/pixel_proprio/video": "pixel_proprio/video_continuous", + "model/pixel_proprio/readout": "pixel_proprio/readout_continuous", "data/e1/yam_human_keypoints_multiplier_proportional_val01": "e1/yam_human_matched_af64_multiplier", "data/e1/yam_human_matched_noaug_proportional_val01": "e1/yam_human_matched_dp_noaug", "data/pusht/planar_chain_manual4919_standard_retimed": "pusht/planar_usocket_chain_manual4919_standard_dp_retimed_masked", @@ -186,6 +194,49 @@ def _compose_for_audit(path): cfg.hydra.runtime.output_dir = "/tmp/egoverse-config-audit" cfg.hydra.job.num = 0 cfg.hydra.job.id = "config-audit" + # Prepared manifests and readout weights are launch-owned inputs. These + # synthetic paths resolve the offline graph only; no artifact is opened + # or certified here, and production YAML keeps its mandatory bindings. + data_target = OmegaConf.select(cfg, "data._target_") + prepared_manifests = { + "egomimic.rldb.world_model_datamodule.WorldModelDataModule": ( + "/tmp/egoverse-config-audit/world-manifest.json" + ), + "egomimic.rldb.dino_proprio_datamodule.DinoProprioDataModule": ( + "/tmp/egoverse-config-audit/dino-manifest.json" + ), + "egomimic.rldb.pixel_proprio_datamodule.PixelProprioDataModule": ( + "/tmp/egoverse-config-audit/pixel-manifest.json" + ), + } + if data_target in prepared_manifests: + if OmegaConf.is_missing(cfg.data, "manifest_path"): + cfg.data.manifest_path = prepared_manifests[data_target] + if ( + data_target + == "egomimic.rldb.dino_proprio_datamodule.DinoProprioDataModule" + and cfg.data.task == "direct_action" + and OmegaConf.select(cfg.data, "action_manifest_path") is None + ): + cfg.data.action_manifest_path = ( + "/tmp/egoverse-config-audit/dino-actions.json" + ) + for stage in OmegaConf.select(cfg, "model.pipeline.stages", default=[]): + model = stage.get("model") + if ( + OmegaConf.is_dict(model) + and model.get("_target_") == "egomimic.models.world_model.MoTWorldModel" + and OmegaConf.is_missing(model, "pose_dims") + ): + # Unverified structural width, not a prepared-data pose schema. + model.pose_dims = {"audit_measured_pose": 9} + if "pixel_readout" in cfg: + for key, value in { + "initial_weights": "/tmp/egoverse-config-audit/pixel-readout.pt", + "initial_weights_sha256": "0" * 64, + }.items(): + if OmegaConf.is_missing(cfg.pixel_readout, key): + cfg.pixel_readout[key] = value if "pi05" in cfg and OmegaConf.is_missing(cfg.pi05, "pretrained_weights"): cfg.pi05.pretrained_weights = "/tmp/egoverse-config-audit/pi-weights" if "e1" in cfg: diff --git a/tests/test_pixel_config_audit.py b/tests/test_pixel_config_audit.py new file mode 100644 index 000000000..a0cecba2f --- /dev/null +++ b/tests/test_pixel_config_audit.py @@ -0,0 +1,108 @@ +"""Offline composition supplies inputs without relaxing production contracts.""" + +import socket + +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.hydra_config import HydraConfig +from omegaconf import MissingMandatoryValue, OmegaConf + +from scripts.audit_hydra_configs import CONFIGS, compose_for_audit + +PREPARED_CONTEXTS = [ + "data/world_model/prepared.yaml", + "experiment/world_model/continuous.yaml", + "experiment/world_model/discrete.yaml", + "model/world_model/base.yaml", + "data/dino_proprio/prepared.yaml", + "data/pixel_proprio/prepared.yaml", + "experiment/dino_proprio/direct_action.yaml", + "experiment/dino_proprio/visual_state.yaml", + "experiment/pixel_proprio/video_continuous.yaml", + "experiment/pixel_proprio/video_discrete.yaml", + "experiment/pixel_proprio/video_only.yaml", + "experiment/pixel_proprio/readout_continuous.yaml", + "experiment/pixel_proprio/readout_discrete.yaml", + "model/dino_proprio/base.yaml", + "model/pixel_proprio/base.yaml", + "model/pixel_proprio/readout.yaml", + "model/pixel_proprio/video.yaml", +] + + +@pytest.mark.parametrize("name", PREPARED_CONTEXTS) +def test_prepared_recipes_and_fragments_resolve_without_external_artifacts( + name, monkeypatch +): + def no_network(*args, **kwargs): + pytest.fail("Offline configuration audit attempted network access") + + monkeypatch.setattr(socket.socket, "connect", no_network) + previous_hydra = HydraConfig.instance().cfg + with compose_for_audit(CONFIGS / name) as cfg: + # Exercise the same strict recursive resolution as the shipped auditor. + resolved = { + key: OmegaConf.to_container(cfg[key], resolve=True, throw_on_missing=True) + if OmegaConf.is_config(cfg[key]) + else cfg[key] + for key in cfg + if key != "hydra" + } + OmegaConf.to_container(cfg.hydra.launcher, resolve=True, throw_on_missing=True) + assert resolved["data"]["manifest_path"].startswith( + "/tmp/egoverse-config-audit/" + ) + assert resolved["data"]["mode"] in {"continuous", "discrete"} + if "pixel_readout" in resolved: + binding = resolved["model"]["pipeline"]["initialization"][0] + assert binding["source"] == resolved["pixel_readout"]["initial_weights"] + assert binding["sha256"] == "0" * 64 + assert resolved["trainer"]["max_steps"] == 1000 + if resolved["data"].get("task") == "direct_action": + assert resolved["data"]["action_manifest_path"].endswith( + "/dino-actions.json" + ) + assert HydraConfig.instance().cfg is previous_hydra + + +@pytest.mark.parametrize( + "experiment", + [ + "world_model/continuous", + "world_model/discrete", + "dino_proprio/direct_action", + "dino_proprio/visual_state", + "pixel_proprio/video_continuous", + "pixel_proprio/readout_continuous", + "pixel_proprio/readout_discrete", + ], +) +def test_audit_bindings_do_not_escape_into_fresh_production_composition(experiment): + with compose_for_audit(CONFIGS / f"experiment/{experiment}.yaml"): + pass + with initialize_config_dir(version_base=None, config_dir=str(CONFIGS)): + cfg = compose( + config_name="train_zarr_cartesian", + overrides=[f"+experiment={experiment}", "hydra/launcher=basic"], + ) + assert OmegaConf.is_missing(cfg.data, "manifest_path") + with pytest.raises(MissingMandatoryValue): + OmegaConf.to_container(cfg.data, resolve=True, throw_on_missing=True) + if "pixel_readout" in cfg: + assert OmegaConf.is_missing(cfg.pixel_readout, "initial_weights") + assert OmegaConf.is_missing(cfg.pixel_readout, "initial_weights_sha256") + if experiment == "dino_proprio/direct_action": + assert cfg.data.action_manifest_path is None + + +def test_world_model_pose_schema_remains_required_outside_audit(): + with compose_for_audit(CONFIGS / "model/world_model/base.yaml") as cfg: + assert cfg.model.pipeline.stages[0].model.pose_dims == { + "audit_measured_pose": 9 + } + with initialize_config_dir(version_base=None, config_dir=str(CONFIGS)): + cfg = compose( + config_name="train_zarr_cartesian", + overrides=["+experiment=world_model/continuous", "hydra/launcher=basic"], + ) + assert OmegaConf.is_missing(cfg.model.pipeline.stages[0].model, "pose_dims") diff --git a/tests/test_prepared_component_audit.py b/tests/test_prepared_component_audit.py new file mode 100644 index 000000000..6d85f554a --- /dev/null +++ b/tests/test_prepared_component_audit.py @@ -0,0 +1,177 @@ +"""Offline prepared-data preflight retains real graph construction and guards.""" + +from pathlib import Path + +import pytest +import torch + +from egomimic.pipeline.construction import restoring_parameters +from egomimic.rldb.dino_proprio_datamodule import DinoProprioDataModule +from egomimic.rldb.pixel_proprio_datamodule import PixelProprioDataModule +from egomimic.rldb.world_model_datamodule import WorldModelDataModule +from scripts import audit_components as components +from scripts.audit_hydra_configs import CONFIGS + + +@pytest.mark.parametrize( + "module, options", + [ + (WorldModelDataModule, {}), + (DinoProprioDataModule, {}), + ( + DinoProprioDataModule, + { + "task": "direct_action", + "action_manifest_path": "/not-opened/actions.json", + }, + ), + (PixelProprioDataModule, {"image_only": True}), + ( + PixelProprioDataModule, + { + "mode": "discrete", + "sampling_strategy": "discrete_window", + "discrete_frames": 16, + "stride": 4, + }, + ), + ( + PixelProprioDataModule, + { + "task": "visual_state_action", + "distributed_world_size": 4, + "checkpointable_sampling": True, + "validation_sampling_strategy": "discrete_window", + "validation_discrete_frames": 16, + "validation_stride": 4, + }, + ), + ], +) +def test_configuration_preflight_does_not_open_or_bind_data( + monkeypatch, module, options +): + def forbidden(*args, **kwargs): + pytest.fail( + "Configuration preflight must not access data or construct a loader" + ) + + monkeypatch.setattr(Path, "open", forbidden) + monkeypatch.setattr(module, "setup", forbidden) + dm = module("/not-opened/manifest.json", **options) + before = dict(dm.settings) + result = dm.preflight_configuration() + assert result["scope"] == "configuration_only" + assert result["artifacts_opened"] is False + assert result["sampling"] == before == dm.settings + assert dm.datasets == {} + assert dm.train_sampler is None + assert dm.consumed_batches == 0 + + +@pytest.mark.parametrize( + "options", + [ + {"batch_size": 0}, + {"num_workers": -1}, + {"seed": -1}, + {"min_frames": 1}, + {"min_frames": 8, "max_frames": 4}, + {"discrete_frames": True}, + {"stride": 0}, + {"mode": "unknown"}, + {"forecasting_probability": float("nan")}, + {"forecasting_probability": 1.1}, + {"max_past_seconds": 1.0}, + {"mode": "discrete", "max_past_seconds": 1.0, "max_future_seconds": 1.0}, + ], +) +def test_prepared_preflight_rejects_invalid_configuration_without_setup(options): + dm = WorldModelDataModule("/not-opened/manifest.json", **options) + with pytest.raises(ValueError): + dm.preflight_configuration() + assert dm.datasets == {} + + +@pytest.mark.parametrize( + "options", + [ + {"mode": "discrete", "sampling_strategy": "full_episode_uniform"}, + {"sampling_strategy": "unknown"}, + {"max_past_seconds": 1.0, "max_future_seconds": 1.0}, + {"validation_sampling_strategy": "unknown"}, + {"validation_discrete_frames": 1}, + {"validation_stride": 0}, + {"task": "direct_action", "validation_sampling_strategy": "discrete_window"}, + ], +) +def test_pixel_preflight_validates_validation_sampling_too(options): + dm = PixelProprioDataModule("/not-opened/manifest.json", **options) + with pytest.raises(ValueError): + dm.preflight_configuration() + assert dm.datasets == {} + + +def test_direct_action_still_requires_explicit_action_manifest(): + with pytest.raises(ValueError, match="verified action alignment manifest"): + DinoProprioDataModule("/not-opened/manifest.json", task="direct_action") + + +def test_family_component_audit_builds_real_meta_graphs_without_data_or_weights( + monkeypatch, +): + families = ("world_model", "dino_proprio", "pixel_proprio") + paths = [ + path + for group in ("data", "model", "experiment") + for family in families + for path in (CONFIGS / group / family).glob("*.yaml") + ] + assert len(paths) == 17 + original_open = Path.open + + def guarded_open(path, *args, **kwargs): + assert "/tmp/egoverse-config-audit/" not in str(path), path + return original_open(path, *args, **kwargs) + + def forbidden(*args, **kwargs): + pytest.fail("Offline component audit accessed data or external weights") + + for module in (WorldModelDataModule, DinoProprioDataModule, PixelProprioDataModule): + monkeypatch.setattr(module, "setup", forbidden) + monkeypatch.setattr(Path, "open", guarded_open) + monkeypatch.setattr(torch, "load", forbidden) + original_construct = components._instantiate_model_wrapper + constructed = [] + + def construct(cfg): + wrapper = original_construct(cfg) + parameters = list(wrapper.parameters()) + assert parameters and all(p.device.type == "meta" for p in parameters) + assert wrapper.model.initialization_receipts == [] + constructed.append( + { + "stages": len(wrapper.model.pipeline.stages), + "initialization": bool(cfg.model.pipeline.get("initialization")), + "frozen": any(not p.requires_grad for p in parameters), + "trainable": any(p.requires_grad for p in parameters), + } + ) + return wrapper + + monkeypatch.setattr(components, "_instantiate_model_wrapper", construct) + assert restoring_parameters() is False + records = components.audit_components(paths) + assert restoring_parameters() is False + assert not [row for row in records if row["status"] != "passed"], records + for record in records: + assert set(record["components"]) == {"data", "evaluator", "model"} + assert record["components"]["data"]["scope"] == "configuration_only" + assert record["components"]["model"]["optimizer_scheduler"] == "passed" + assert ( + record["components"]["model"]["late_bound_parameter_placeholder"] is False + ) + assert any(item["stages"] == 4 for item in constructed) + readouts = [item for item in constructed if item["initialization"]] + assert readouts and all(item["stages"] == 10 for item in readouts) + assert all(item["frozen"] and item["trainable"] for item in readouts) From 2039eadf0b481f3dcda4b097cea401e7bb41caeb Mon Sep 17 00:00:00 2001 From: ElmoPA Date: Sun, 18 Oct 2026 14:41:56 -0400 Subject: [PATCH 4/4] Validate and checkpoint pixel video runs every 20000 steps --- docs/PIXEL_VIDEO_PILOTS.md | 32 ++++++++++++------- .../pixel_proprio/readout_continuous.yaml | 5 +++ .../experiment/pixel_proprio/video_only.yaml | 6 ++-- tests/test_pixel_attention_profiles.py | 7 +++- tests/test_pixel_token_execution_profiles.py | 14 ++++++-- 5 files changed, 46 insertions(+), 18 deletions(-) diff --git a/docs/PIXEL_VIDEO_PILOTS.md b/docs/PIXEL_VIDEO_PILOTS.md index 7526828a0..5a252576a 100644 --- a/docs/PIXEL_VIDEO_PILOTS.md +++ b/docs/PIXEL_VIDEO_PILOTS.md @@ -28,26 +28,33 @@ explicit smoke step overrides, and deferred stage2 budgets remain unchanged. The allocation profile `+pixel_execution=ddp4` runs each pilot on four GPUs on one node. Its reusable default is batch 1 per GPU; the selected full runs use the explicit override `data.batch_size=8`, retaining global batch 32 with no -gradient accumulation or learning-rate scaling. Each arm uses its own four-GPU -process group; its scheduler allocation and node are recorded independently. -DDP uses one image expert with all trainable parameters used; automatic +gradient accumulation or learning-rate scaling. Two separate four-GPU runs share the allocated eight-GPU +node. DDP uses one image expert with all trainable parameters used; automatic sampler replacement is disabled because the pixel adapter owns sharding. The explicit TorchElastic environment reads ranks from each run's `torchrun` process group, while Slurm IDs remain available for artifact provenance. Validation covers two fixed anchors from each of the 50 held-out -episodes every 250 steps, with identical 16-timestamp evaluation grids in both arms. +episodes every 20,000 optimizer steps, with identical 16-timestamp evaluation grids in both arms. The 100 global validation batches are divided into four contiguous shards of 25 without padding or duplicated examples. Sampling noise uses global batch IDs so rank assignment does not change it. Metrics sum squared errors and valid-target counts across ranks before dividing; RGB diagnostics are written only by rank zero, whose shard contains both forecast and completion examples. -All scheduled checkpoints and `last` are retained: 1,200 scheduled checkpoints -per run at steps 250 through 300,000, including the 100,000-step review point. +Retain the 20,000-step checkpoints: a fresh run produces 15 scheduled +checkpoints through step 300,000, including the 100,000-step review point. +`save_last: link` makes `last.ckpt` a symbolic link to the latest saved +checkpoint on the local filesystem, including mounted NFS paths. Removing +older checkpoints outside these milestones is a separate operation after +verifying the resumed trainer and retained checkpoints; this configuration +does not delete existing files. Changing these defaults does not update an +active trainer. A resumed run must explicitly select the new validation and +checkpoint intervals and linked `last` while restoring its optimizer, step, +sampler, and per-rank RNG states. FP32 weights and two AdamW moment tensors alone use approximately 1.06 GB per -checkpoint for this model, or 1.27 TB per run and 2.54 TB for both runs at this -retention. These are decimal raw-tensor estimates; additional checkpoint -metadata, the separate `last` files, and diagnostic artifacts require more -space. Validation reports held-out +checkpoint for this model, or 15.9 GB per fresh run and 31.7 GB for both runs at +this cadence. These are decimal raw-tensor estimates; additional checkpoint +metadata and diagnostic artifacts require more space. A successful `last` +symlink adds no duplicate tensor payload. Validation reports held-out image MSE and RGB diagnostics; it is not a policy-success evaluation. The default execution profile is `world_model_attention=flash4_bf16_mixed`, @@ -72,7 +79,10 @@ later work; those fields are not stage1 inputs. Stage2 is deferred until review of the generated videos. Its separate model may introduce proprioception plus measured-state and control-command MSE heads, initialized from the audited stage1 image weights. Selecting these -stage1 recipes does not activate or schedule stage2 training. +stage1 recipes does not activate or schedule stage2 training. The deferred +readout recipes retain their 1,000-step budget, 250-step validation and +checkpoint intervals, and `save_last: true`. The shared 1,000-step stage1 pilot still validates at +the end through `val_at_end=true`. ## Source and publication provenance diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml index 22fe4780e..a87083315 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml @@ -23,6 +23,11 @@ data: max_frames: 128 # Preserve deferred stage2 sampling independently of stage1. trainer: max_steps: 1000 # Deferred readout keeps its original pilot budget. + val_check_interval: 250 +callbacks: + model_checkpoint: + every_n_train_steps: 250 + save_last: true evaluator: objective: visual_state_action logger: diff --git a/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml index 63141a7c9..0a60b565e 100644 --- a/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml +++ b/egomimic/hydra_configs/experiment/pixel_proprio/video_only.yaml @@ -37,7 +37,7 @@ trainer: max_epochs: -1 limit_train_batches: 1.0 limit_val_batches: 1.0 - val_check_interval: 250 + val_check_interval: 20000 check_val_every_n_epoch: null max_steps: 1000 num_sanity_val_steps: 0 @@ -46,10 +46,10 @@ callbacks: model_checkpoint: filename: "epoch_{epoch:04d}-step_{step:09d}" every_n_epochs: 0 - every_n_train_steps: 250 + every_n_train_steps: 20000 save_on_train_epoch_end: false save_top_k: -1 - save_last: true + save_last: link runtime: strict_fp32: ${world_model_attention.strict_fp32} disable_tf32: ${world_model_attention.disable_tf32} diff --git a/tests/test_pixel_attention_profiles.py b/tests/test_pixel_attention_profiles.py index 69b484279..3e123a33f 100644 --- a/tests/test_pixel_attention_profiles.py +++ b/tests/test_pixel_attention_profiles.py @@ -90,7 +90,12 @@ def test_pixel_profile_override_preserves_science( assert trainer.pop("precision") == trainer_precision baseline_trainer.pop("precision") assert trainer == baseline_trainer - assert cfg.trainer.max_steps == 300000 and cfg.trainer.val_check_interval == 250 + assert cfg.trainer.max_steps == 300000 and cfg.trainer.val_check_interval == 20000 + assert cfg.callbacks.model_checkpoint.every_n_train_steps == 20000 + assert cfg.callbacks.model_checkpoint.every_n_epochs == 0 + assert cfg.callbacks.model_checkpoint.save_top_k == -1 + assert cfg.callbacks.model_checkpoint.save_last == "link" + assert not cfg.callbacks.model_checkpoint.save_on_train_epoch_end assert cfg.model.optimizer == baseline.model.optimizer assert cfg.model.optimizer.lr == 5e-5 and cfg.model.optimizer.weight_decay == 0.01 assert cfg.model.scheduler is None and not cfg.model.enable_grad_norm diff --git a/tests/test_pixel_token_execution_profiles.py b/tests/test_pixel_token_execution_profiles.py index 8d78afc3f..d5b269f52 100644 --- a/tests/test_pixel_token_execution_profiles.py +++ b/tests/test_pixel_token_execution_profiles.py @@ -99,6 +99,14 @@ def test_readout_initialization_requires_matching_token_execution(arm): @pytest.mark.parametrize("arm", ["continuous", "discrete"]) -def test_deferred_readout_does_not_inherit_full_video_budget(arm): - assert configuration("video", arm)["trainer"]["max_steps"] == 300000 - assert configuration("readout", arm)["trainer"]["max_steps"] == 1000 +def test_deferred_readout_keeps_its_short_budget_and_validation_schedule(arm): + video = configuration("video", arm) + readout = configuration("readout", arm) + assert video["trainer"]["max_steps"] == 300000 + assert video["trainer"]["val_check_interval"] == 20000 + assert video["callbacks"]["model_checkpoint"]["every_n_train_steps"] == 20000 + assert video["callbacks"]["model_checkpoint"]["save_last"] == "link" + assert readout["trainer"]["max_steps"] == 1000 + assert readout["trainer"]["val_check_interval"] == 250 + assert readout["callbacks"]["model_checkpoint"]["every_n_train_steps"] == 250 + assert readout["callbacks"]["model_checkpoint"]["save_last"] is True