diff --git a/docs/PIXEL_VIDEO_PILOTS.md b/docs/PIXEL_VIDEO_PILOTS.md new file mode 100644 index 000000000..5a252576a --- /dev/null +++ b/docs/PIXEL_VIDEO_PILOTS.md @@ -0,0 +1,107 @@ +# 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. 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 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. +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 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`, +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. 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 + +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..f2cddb195 --- /dev/null +++ b/egomimic/eval/image_proprio_eval.py @@ -0,0 +1,267 @@ +"""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..a87083315 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pixel_proprio/readout_continuous.yaml @@ -0,0 +1,36 @@ +# @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. + val_check_interval: 250 +callbacks: + model_checkpoint: + every_n_train_steps: 250 + save_last: true +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..0a60b565e --- /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: 20000 + 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: 20000 + save_on_train_epoch_end: false + save_top_k: -1 + save_last: link +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..596870947 --- /dev/null +++ b/egomimic/models/attention_execution.py @@ -0,0 +1,360 @@ +"""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..1b2c17194 --- /dev/null +++ b/egomimic/models/compilation.py @@ -0,0 +1,50 @@ +"""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..678f03b2b --- /dev/null +++ b/egomimic/models/image_proprio.py @@ -0,0 +1,1037 @@ +"""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..1c5ca3c62 100644 --- a/egomimic/models/world_model.py +++ b/egomimic/models/world_model.py @@ -16,15 +16,35 @@ 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 +75,153 @@ 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 - if type(attention_query_chunk_size) is not int or attention_query_chunk_size < 0: + 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 forward(self, streams, times, valid): - sizes = [s.shape[1] for s in streams] + 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) + + @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 @@ -108,23 +230,101 @@ def forward(self, streams, times, valid): # 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: - 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,12 +332,21 @@ 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") + 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: @@ -146,40 +355,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): @@ -231,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 @@ -276,8 +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_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_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..a2cc4b37b --- /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..1b61a3f93 100644 --- a/egomimic/pl_utils/world_model_behavior.py +++ b/egomimic/pl_utils/world_model_behavior.py @@ -7,8 +7,80 @@ 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 +88,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,14 +116,26 @@ 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): 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/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..c606a3368 --- /dev/null +++ b/egomimic/rldb/pixel_proprio_datamodule.py @@ -0,0 +1,397 @@ +"""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 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() + 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/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/egomimic/trainHydra.py b/egomimic/trainHydra.py index b1e605c32..4768bf385 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 @@ -56,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 @@ -149,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 + ), ) ] @@ -160,7 +166,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 +182,42 @@ 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/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/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..faa25938f --- /dev/null +++ b/scripts/world_model/benchmark_pixel_execution.py @@ -0,0 +1,1049 @@ +"""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..de1e95191 --- /dev/null +++ b/tests/test_attention_execution.py @@ -0,0 +1,471 @@ +"""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=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=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 + 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=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) + 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=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): + 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=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"]) +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=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=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.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=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) < 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=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) < 0.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..97d3f7d4a --- /dev/null +++ b/tests/test_image_proprio_adaln.py @@ -0,0 +1,220 @@ +"""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..8b2e5f47a --- /dev/null +++ b/tests/test_image_proprio_attention.py @@ -0,0 +1,218 @@ +"""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.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), + "camera": torch.tensor([[0, 1, 0, 1, 0, 1]]), + } + proprio = { + "value": torch.randn(1, 4, 3), + "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]]), + } + 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, 0.0, 0.5, 1.0, 2.0]]) + 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..bd0d0bc13 --- /dev/null +++ b/tests/test_image_proprio_distributed.py @@ -0,0 +1,153 @@ +"""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..277af2fc7 --- /dev/null +++ b/tests/test_image_proprio_image_only.py @@ -0,0 +1,162 @@ +"""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.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, + ) + 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..5ad3f98e9 --- /dev/null +++ b/tests/test_image_proprio_image_only_initialization.py @@ -0,0 +1,193 @@ +"""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..1c2c5f755 --- /dev/null +++ b/tests/test_image_proprio_integration.py @@ -0,0 +1,189 @@ +"""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..67f0c27e2 --- /dev/null +++ b/tests/test_image_proprio_mixed_precision.py @@ -0,0 +1,119 @@ +"""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(-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): + 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..d54c13076 --- /dev/null +++ b/tests/test_image_proprio_no_phase.py @@ -0,0 +1,118 @@ +"""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..811329c4b --- /dev/null +++ b/tests/test_image_proprio_positions.py @@ -0,0 +1,249 @@ +"""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..3e123a33f --- /dev/null +++ b/tests/test_pixel_attention_profiles.py @@ -0,0 +1,205 @@ +"""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 == 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 + 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..33cffb8a7 --- /dev/null +++ b/tests/test_pixel_compile_ddp.py @@ -0,0 +1,112 @@ +"""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_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_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..5c6ff1137 --- /dev/null +++ b/tests/test_pixel_execution_benchmark.py @@ -0,0 +1,407 @@ +"""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=0.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..c18838a0f --- /dev/null +++ b/tests/test_pixel_prediction_diagnostics.py @@ -0,0 +1,231 @@ +"""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..d5b269f52 --- /dev/null +++ b/tests/test_pixel_token_execution_profiles.py @@ -0,0 +1,112 @@ +"""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_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 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_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) diff --git a/tests/test_train_hydra_torch_elastic.py b/tests/test_train_hydra_torch_elastic.py new file mode 100644 index 000000000..efb8b1f00 --- /dev/null +++ b/tests/test_train_hydra_torch_elastic.py @@ -0,0 +1,117 @@ +"""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..ab659dd09 --- /dev/null +++ b/tests/test_world_model_distributed_rng.py @@ -0,0 +1,312 @@ +"""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"] + )