diff --git a/docs/WORLD_MODEL.md b/docs/WORLD_MODEL.md new file mode 100644 index 0000000000..5516702caa --- /dev/null +++ b/docs/WORLD_MODEL.md @@ -0,0 +1,429 @@ +# Continuous-time MoT world model + +This is an independent research model, not an Action Flow variant. The implementation +is on `codex/continuous-time-mot-world-model`, based on the verified common release +7043c025. It uses the existing PipelineAlgo, ModelWrapper, DataContext and Eval +interfaces. No existing model, launcher or run is changed. + +## Confirmed design + +- Human, YAM, LIBERO Sharpa and LIBERO gripper co-training. +- Three modality experts: DINO image features, XYZ tracks and proprioception. + Each expert has its own Q/K/V/output projections, normalization and FFN. + All valid tokens within one example participate in joint attention. Experts + are shared across embodiments; pose input/output adapters are private. + After each block, image tokens also pass through the selected embodiment's + private residual adapter. Each block has separate adapters for each embodiment; + the image output head remains shared. +- Frozen DINO targets; denoising objectives on both image features and tracks. +- Pose readout sees observed context and predicted visual/track features. Pose + gradients propagate through those predictions to the visual generator. +- Visual/track warm-up followed by a pose-loss ramp. +- Poses and tracks expressed in one fixed reference frame per sequence. +- Regular image-grid tracking, persistent point IDs and confidence/visibility. + Do not substitute interpolated endpoints for actual tracks. +- Source sampling proportional to dataset size. This implementation defines + size as eligible temporal anchors/windows, and returns homogeneous source + batches. PipelineAlgo sees one source per update; it must not apply a second + source-balancing factor. Exact source probabilities depend on the corpus and + temporal mode and must be reported for a comparison. +- Continuous: 2–16 random recorded timestamps (user-selected maximum), each within + -3 to +3 physical seconds of the current/anchor frame (user-selected range). + Bounds clip to episode edges; the anchor is included. Draw without replacement + from eligible recorded times, not uniformly over continuous seconds. Anchors + cover the episode; individual examples span at most six seconds. Isolated + anchors with fewer than two eligible times are excluded, never replaced with + out-of-range frames. This replaces the earlier unbounded frame-index heuristic + in the configured continuous recipe. Direct calls without time bounds retain + the historical sampler for checkpoint-bound compatibility. + Discrete: four fixed-stride frames (provisional), integer temporal positions. + This compares complete recipes, not time encoding alone. + +## Proposed engineering defaults, not measured training settings + +Width384, depth8, heads6, FFN expansion4. There are 43,767,555 shared trainable +parameters excluding embodiment pose input/output adapters and the frozen DINO +backbone. Pose adapter count is sum_m(769*D_m + 384). Pose widths deliberately remain +required config values until measured-pose schemas are verified; historical +7D gripper or 69D Sharpa policy command widths are not measured-pose schemas. + +The image residual adapters add 1,218,048 trainable parameters for four +embodiments: 304,512 per embodiment across eight blocks. The configurable +bottleneck is 48, giving 384 → 48 → 384 per adapter: + + image_next = image_shared + Up_e(GELU(Down_e(LayerNorm_e(image_shared)))) + +Only the current example's embodiment adapter is used. The up projection starts +at zero, so initialization exactly preserves the shared model's forward output. +On the first update the up projection learns; the down projection and adapter +normalization begin receiving nonzero gradients once the up weights change. +Adapters are reused in both visual denoising and pose-readout passes. Padding is +masked again after the adapter. The total excluding pose adapters and frozen +DINO is now 44,985,603 parameters; 43,767,555 of those remain shared. + +Set image_adapter_dim=0 explicitly for an adapter-free ablation or the original +adapter-free checkpoint architecture. Strict loading rejects a missing adapter +state; enabling adapters is an architecture change, not transparent resume. + +The maintained upstream feature extractor is +`scripts/embedding_process/dinov3_embedding.py`, with a DINOv3 ViT-B/16 default. +It supports pinned model/processor revisions and explicit full-FOV letterboxing; +patch counts are checked against the actual processed pixel grid and model output. +Preprocessing geometry, normalization, precision and revision are saved with each +output. The correspondence diagnostic in `dinov3_correspondence.py` uses hashed +native frames and independent visible-point labels, excludes padding, and saves +feature caches. Its tests and a tiny untrained Transformers DINOv3 runtime check +verify geometry, not pretrained accuracy. Official pretrained access has now +been verified in the task's FP32 DINO pilot. Whole-hand correspondence uses +estimated projected Aria references with unknown occlusion; it does not measure +independently verified visible-joint accuracy or establish a final resolution. +The proposal explicitly favors high-resolution +features for small gripper detail. The earlier224x224/196-patch estimate is only +a debugging budget, NOT an approved scientific preprocessing choice. The actual +processor/checkpoint revision, resolution and produced grid must be pinned and +checked during preprocessing. The prepared loader rejects declared patch grids +that disagree with the cache. + +The isolated `scripts/embedding_process/dinov3_pose_probe.py` diagnostic reuses +SHA-verified image features and trains only a linear patch-to-landmark heatmap +head. The world-model pose readout cannot express this control because it uses +MoT context and track predictions. The probe belongs with the maintained DINO +diagnostics, accepts paths and settings as arguments, and leaves entire episodes +out. It compares training and held-out image-only predictions with a training-only +mean-pose baseline and predictions from shuffled held-out images. It never selects +a checkpoint on the test set. Projected references retain their uncertainty; +this is current-image 2D landmark recovery, not future 3D pose forecasting. + +A proposed1120-point28x40 grid is selected in ONE reference camera, with its +3D positions expressed in the common reference frame. It is not1120 tracks +per camera. Four timestamps, three cameras and196 image patches produce +4*(3*196 +1120 +1)=6836 tokens. The initial hard limit is8192 tokens. This is a +budget guard, not evidence the full-size model fits a particular GPU. Joint +attention scales quadratically in this token count; memory/time need a scheduled +GPU measurement before training. No silent point subsampling or pooling occurs. +For comparison,448x448/784 patches at four timestamps gives13892 tokens and +exceeds this guard; two timestamps gives6946. Resolution and temporal/track +coverage must be budgeted together, without silently reducing the requested +image detail. Neither example is an approved full training resolution. + +The requested continuous maximum is now16 timestamps, including observed and +target times together. With640x640 DINO inputs, three YAM cameras,1120 tracks +and one pose token per time,16 times require94736 tokens. The8192-token guard +remains active: this configuration is not ready for a full-resolution launch. +Sixteen versus four full times gives4x as many tokens and16x the dense attention +pair count. A materialized FP32 six-head attention-score tensor alone would +occupy about200.6GiB per layer at batch1; that is a backend-dependent warning, +not a measured peak or a claim that memory-efficient attention materializes it. +Do not silently raise the guard, reduce resolution, pool features, subsample +tracks, or switch precision to make it fit. GPU feasibility and any redesigned +token budget remain separate work. Historical direct sampler defaults stay4. + +Memory controls are available without changing parameters or token coverage: +\`attention_query_chunk_size\` splits attention queries while retaining all keys +and values; non-reentrant checkpointing recomputes each query chunk in backward. +\`activation_checkpointing\` recomputes each MoT block instead of retaining its +internal activations. Both default off for historical behavior. FP32 fixture +tests compare outputs, losses, all gradients and one AdamW update, including +padding and randomly placed observations. These controls change execution and +roundoff, not the attention graph; they do not remove quadratic compute cost. +The maintained \`scripts/benchmarks/profile_world_model.py\` measures synthetic +full-size FP32 forward/backward peak memory on one allocated GPU. Its explicit +diagnostic pose widths and token guard override are not training schema or +launch approvals. It performs no optimizer updates; real-data optimizer, +validation and reload checks remain required separately. + +Warm-up10000/ramp10000, AdamW1e-4/weight-decay0.01, loss weights1/1/1 and20 sampler +steps are provisional configurable defaults. No CFG, Action Flow Jacobian loss, +velocity augmentation or historical policy-specific optimizer is inherited. +Actual training duration, split/data identities and physical temporal span are +not selected by the config. `max_steps`, `manifest_path` and `pose_dims` are +mandatory before use. Both variants share the same parent recipe. + +## Training computation + +Physical time tau and diffusion noise s are separate. Temporal RoPE uses actual +relative seconds in continuous mode, grid indices in discrete mode. A common +random time-origin shift preserves all temporal differences. Image tokens also +carry patch XY and camera ID; tracks carry reference XY, camera and point ID. +Observed/query status and s have their own embedding. + +For each withheld visual frame and each modality, independently sample s in +(0,1) and Gaussian epsilon. All image patches/cameras at that timestamp use the +same s. With alpha=cos(pi*s/2), sigma=sin(pi*s/2): + + x_s = alpha*x_0 + sigma*epsilon + v_target = alpha*epsilon - sigma*x_0 + x0_estimate = alpha*x_s - sigma*v_prediction + +This is the default `image_prediction_type: v`. The opt-in `x0` setting instead +interprets the image head as a clean-feature estimate and applies image MSE +directly against x_0. Track targets remain v in both modes. The pose readout +receives the corresponding differentiable clean estimate. The setting is part +of checkpoint/config identity: matching weight shapes do not make v and x0 +checkpoints interchangeable. + +A linear image output head narrower than the DINO feature dimension cannot +represent arbitrary full-dimensional noise in v. The x0 alternative lets the +head model structured clean features while the sampler carries noise outside +the head analytically. Its DDIM update uses +`alpha_next*x0 + (sigma_next/sigma)*(x_s-alpha*x0)` with zero-level anchors +handled separately; this avoids overflow at tiny positive refinement strengths. +Clean-feature MSE corresponds to sigma-squared-weighted v MSE, so this is an +explicit objective ablation, not a numerically equivalent optimization. It +preserves all image channels but does not guarantee improved generalization. + +Observed tokens remain clean. The image and track losses are weighted mean MSE +on supervised withheld tokens. Track confidence is a loss weight; unknown-future +visibility is NOT an attention mask or input signal. Padding is separately +masked. Inputs never contain withheld clean poses. + +Half of sampled masks are past-only forecasting by default. The other half can +condition/query before, between or after observations, with at least one observed +and one withheld image frame. Neither endpoint is forced in this random branch. +The50/50 mixture is provisional. Fractional confidence weights are normalized +by their actual sum, including sums below one; empty objectives return zero. + +### Missing-track policy + +Missing coordinates must never become clean zero coordinates before partial +noising. A track is unavailable when its target-valid mask is false or confidence +is zero. The default `missing_track_policy: error` rejects an incomplete training +batch with a clear error. It also rejects invalid/zero-confidence observed visual +values and unsupported missing image features. Padding is not a missing target. +Geometric target availability is distinct from camera visibility: an occluded +simulator point can still have an exact measured XYZ target. Do not automatically +set its confidence to zero solely because it is occluded; confidence describes +coordinate reliability. Future visibility must not become an input feature. + +For explicitly approved incomplete-data experiments, the existing stage accepts +`missing_track_policy: conditional_sample` and `imputation_steps` (default20). +This is approximate conditional data completion, not measured geometry: + +1. Build a private auxiliary batch with available recorded image/track targets + clamped as clean anchors. Keep only the originally observed poses; sanitize + all withheld poses and missing coordinate placeholders. +2. Use the existing sampler under `no_grad` and detach its missing-track outputs. +3. Insert only those missing coordinates before ordinary per-frame forward + noising. Their direct track-loss weights remain zero. Do not promote them to + supervised or observed values in the main pass. + +Target availability selects anchors in the private auxiliary sampler; it is +not added as a flag, attention mask or noise-level change in the main denoising +pass. The main pose loss still differentiates through the denoised predictions, +but never through the auxiliary completion sampler. The current model supplies +the completions, so early/untrained completions can be inaccurate. This is not +an exact marginal-likelihood objective or proof of correct occluded geometry. + +Complete-data batches preserve their previous losses, gradients and RNG state +bit-exact in tested CPU FP32/BF16 cases and incur no extra sampling. An incomplete +batch with20 completion steps costs21 extra backbone forwards (the shared sampler +also computes a final pose readout). `world_imputed_track_tokens` and +`world_imputation_forwards` report this cost. The default remains strict to avoid +silently introducing that compute and approximation into a real run. Choose and +measure a completion policy on actual data before full training. + +Regression tests cover strict rejection, zero-confidence finite placeholders, +NaN/large-value placeholder invariance, withheld-pose isolation, no batch mutation, +detached completion, unchanged main frame-noise levels, finite CPU BF16 gradients, +pipeline controls/cost logging and complete-data compatibility. + +After warm-up a second forward through the SAME backbone uses differentiable +same-trajectory x0 estimates plus empty pose queries. Its regression loss updates +both visual output heads and their upstream shared parameters. It must never +pair an independently sampled alternative future with the recorded pose target. +There are two network forwards per training batch after warm-up, one before; +there is no unrolled differentiable sampler or higher-order/Jacobian objective. + +Inference uses deterministic DDIM updates from Gaussian initialization, clamps +observed visual values at every step, then runs one pose readout. All unknown +input values can be NaN placeholders without affecting output. The sampler never +reads training confidence/target-valid fields. Image outputs are DINO FEATURES, +not RGB video; no image decoder is included. + +### Selective prediction refinement + +The same sampler optionally accepts `world["refinement"]`. Each image/track +entry has `strength[B,N]` and, when retaining or revising existing predictions, +`prediction[B,N,C]`. The ordinary pipeline stage passes these controls through +inside its existing world input; there is no separate inference launcher. + +| Strength | Behavior | +| --- | --- | +| 0 | Keep the supplied clean prediction bit-exact as an anchor throughout sampling | +| Between 0 and 1 | Add the corresponding noise to a supplied clean prediction, then refine | +| 1 | Generate from fresh Gaussian noise, ignoring any supplied previous value | + +For a clean previous prediction y and strength s, initialization is +`cos(pi*s/2)*y + sin(pi*s/2)*epsilon`. Each query then descends from its own s +to zero over the configured number of DDIM steps. Strength is a noise level, +not a percentage change or a guarantee of geometric closeness. All image patches +and cameras at one unknown timestamp must agree on strength; track points at +one timestamp must also agree, independently of image strength. + +Example for a caller that has already aligned previous outputs to the requested +token layout: + +```python +world["refinement"] = { + "image": {"prediction": previous["image"], "strength": image_levels}, + "track": {"prediction": previous["track"], "strength": track_levels}, +} +# e.g. frame levels: [0.0, 0.0, 0.4, 1.0], repeated over each frame's tokens +result = model.sample(world, embodiment, steps=20, generator=generator) +``` + +Actual observations must have strength zero and always use their measured +values, even if the supplied previous prediction disagrees. An unknown frame +at strength below one requires an explicit finite prediction; the sampler never +falls back to withheld values in the batch. Fresh-frame prediction placeholders +may be NaN. Prediction arrays must match the value shape/device/dtype, normalized +feature/track representation, timestamps, camera IDs and point ordering. +The caller owns this alignment; shape checks cannot prove semantic alignment. +Omitted modality controls retain ordinary fresh generation for that modality. + +Fixed generated anchors are marked clean only in a private sampling copy; +caller inputs and observation masks are unchanged. Pose queries are regressed +again from the refined visuals; pose diffusion/refinement is not introduced. +Keeping a previously generated pose fixed requires explicitly treating it as +conditioning rather than passing it to the visual refinement controls. + +Tests cover mixed keep/refine/generate requests in continuous and discrete time, +CPU BF16 mixed, exact fixed anchors at every denoising call, repeated calls with +changed query times, ordinary pipeline inference, and rejection before RNG use +of invalid strengths, missing predictions or attempts to alter observations. +Automatic timestamp insertion/reordering and observation-update orchestration +remain separate future work. Learned refinement quality needs real-data tests. + +Open scientific issues: the one-pass training x0 distribution differs from +multi-step sampled features, direct context can provide a readout shortcut, +frozen features may omit contact/gripper detail, track estimates may be biased, +and deterministic pose regression can average ambiguous motions. Parameter +sharing does not establish transfer. Evaluate forecasting and interpolation +separately and compare matching physical query times. Do not label normalized +feature/pose MSE as robot task success or action-distribution fidelity. + +## Prepared data interface + +`rldb/world_model.py` owns the typed adapter and timestamp sampling. It consumes +an immutable JSON manifest plus per-episode NPZ or HDF5 arrays; it does not modify source +Zarr datasets. The existing DINO extraction and raw dataset readers remain the +upstream authorities. This new cache contract is necessary because the old policy +loaders do not supply synchronized irregular image/track/pose targets. + +Manifest schema `mot-world-prepared-v1` requires: +- `dino_identity`: checkpoint/revision, feature layer, preprocessing/grid identity. +- `normalization`: train-only provenance and statistics SHA; arrays already normalized. +- `pose_schemas`: per-embodiment dimension and coordinate/rotation/field semantics. +- `track_provenance`: tracker/simulator identity, calibration and frame binding; + `interpolated_endpoints` must be false. +- `episodes`: embodiment, episode_id, relative path, SHA256, frames, explicit + train/val assignment, frame=`fixed_reference`, pose_kind=`measured_pose`. + Duplicate IDs and resolved physical paths across splits are rejected. + +Per-episode arrays: + +| Array | Shape / semantics | +| --- | --- | +| time | [T], increasing physical seconds; subtract origin before casting float32 | +| image | [T,V,P,D], frozen normalized DINO patch features | +| image_xy | [P,2], patch positions | +| track | [T,Q,3], normalized common-reference positions of persistent points | +| track_xy | [Q,2], regular-grid reference image positions | +| track_camera | [Q], integer reference camera identifiers | +| track_confidence | [T,Q], [0,1], unavailable=0 | +| pose | [T,Dm], normalized measured poses in same fixed reference frame | +| pose_valid | [T], measured target availability | + +The current adapter assumes synchronized arrays on a verified physical timeline. +Do not fabricate time from arbitrary indices or interpolate unseen trajectories. +Do not copy policy action arrays into pose. Use original measured state and +calibrated transforms. `points_to_reference` validates proper rigid transforms +for XYZ; orientation conversions should use existing pose utilities. Sampling +and padded collation preserve actual times and masks. Different query schedules +per modality are supported by the model token contract; the initial prepared +loader samples the same timestamp set for all three modalities. + +NPZ remains a correctness/prototype interface: selecting a few frames reads +the entire episode field before slicing. Dense caches should use the supported +HDF5 backend with one-frame image chunks; the reader selects frame indices on +disk before materializing image/track/pose arrays. Both backends retain the same +SHA-bound manifest, sampling, normalization and masks. Regression tests require +bit-exact samples in both temporal modes and intercept image reads to reject +whole-episode access. The finalizer hashes each new immutable payload and embeds +its SHA-bound filesystem fingerprint in the trusted prepared manifest. Startup +reuses that verification only when the resolved path, device, inode, byte size, +mtime and ctime all match exactly. Copies, relocations, mutations, malformed +fingerprints and historical manifests fall back to a full content hash. This +proof is authored only if ctime was at least one second old when hashing began; +recent files are hashed without creating a reusable proof, guarding against +same-clock-tick rewrites. Aging later cannot revive a recent hash's proof. This +fast path relies on a trusted immutable manifest and filesystem metadata; it is +not a substitute for hashing new artifacts or an independent silent-bit-rot +check. Startup also reads the small timestamp vectors. Metadata assertions are +not proof of actual calibration or leakage-free normalization; preprocessing +receipts must establish those properties. Exact data-RNG continuation across +checkpoint resume is not implemented or claimed. + +Single-device loading only: distributed training fails explicitly until source +sampling and data sharding are implemented and tested. Training source choice +is proportional to eligible windows, with replacement; validation enumerates +anchors once with deterministic timestamp/mask choices. This is not a claim of +exhaustive evaluation over every possible temporal subset. + +## What has been verified + +Tests cover real tiny optimizer updates for both embodiment heads, strict model +and optimizer reload, shared Lightning fit/validation/resume, both Hydra recipes, +BF16 mixed finite gradients on CPU, pose gradients into both visual heads, +no withheld-target leakage at inference, clean observation clamping, padding +invariance, temporal shift invariance, discrete indices, irregular sampling, +proper reference-frame transforms and cache split/provenance rejection. +Adapter tests cover all four embodiment routes, pose-loss gradients, unchanged +unused adapters, zero-initialization equivalence, embodiment-specific effects, +invalid bottleneck rejection and strict checkpoint restoration. + +These are synthetic fixtures. They do not establish real-data performance, +CUDA numerical equivalence, full-size memory fit or GPU training readiness. +Regression tests also cover backward/interpolated image queries, fractional +confidence normalization, empty objectives and declared-versus-actual patch grids. +Ordered hierarchical plan tokens are intentionally deferred by the original +proposal's first-version decision, rather than an accidentally missing component. + +## Data inspection and remaining work + +Training uses Lightning `32-true` (FP32), explicitly selected by the user. +The discrete recipe inherits this setting from the continuous parent. Historical +BF16 fixture tests remain compatibility checks, not the selected training mode. + +Live shell inspection confirmed both500-episode LIBERO Spatial Zarr roots exist. +Sharpa data has61535 rows, EEF position3/quaternion, hand state22, two RGB cameras, +and action/observation host-nanosecond timestamps. Original gripper data has62250 +rows, EEF position3/quaternion, gripper state2, joint state and two RGB cameras. +Neither listed processed replay contains scene-wide tracks. A bounded audit +recovered original HDF5 simulator states and XML for one episode of each LIBERO +embodiment; both physical clocks advance in 0.05-second steps. Exact state/image +alignment, asset restoration and persistent-point extraction remain unverified. +Use simulator time, not Sharpa host wall-clock timestamps. + +The human sample has real device timestamps, including dropped-frame gaps, and +moving camera poses. Its 21 hand landmarks follow raw Aria-supported ordering: +wrist index5 and fingertip indices0–4, not canonical MANO. Aria lacks MANO's thumb +CMC joint, so conversion is not a full permutation. Prepared pose schemas must +bind the original producer, input/output topology and any fitted conversion; +63 coordinates alone are not topology evidence. Preserve original data. + +The inspected YAM sample has measured EEF/gripper/joint states, all three camera +intrinsics and front-camera/inter-arm calibration, but no acquisition timestamps. +Nominal30fps is not a verified physical clock. Wrist-camera metric extrinsics are +needed if using those views for 3D geometry; front-reference tracks do not require +them for image-only wrist-camera feature tokens. No inspected embodiment has a +verified scene-grid track cache. Occlusion and coordinate availability differ: +an occluded simulated point can still have valid XYZ. Keep visibility separate +from coordinate confidence and never expose unknown future visibility as input. + +Before any real training: resolve human/YAM camera poses and calibration, exact +DINO runtime/weights, source simulator assets/state replay, 3D track extraction +for each domain, synchronized physical clocks, measured-pose schemas, train-only +normalization, dataset splits and complete preprocessing receipts. Then verify +an actual batch and a scheduled full-size GPU optimizer/validation/reload check. +Do not treat this code milestone as a prepared four-domain training dataset. diff --git a/docs/cluster-pipeline.md b/docs/cluster-pipeline.md index 39b8fa40f7..5fa4bb7848 100644 --- a/docs/cluster-pipeline.md +++ b/docs/cluster-pipeline.md @@ -131,6 +131,45 @@ status and signal names and never requeues a cancellation or ordinary child failure. A failed upstream dependency automatically retires the dependent request; it does not permit bypassing its smoke gate. +For a single-process trainer sharing the runner's Linux container, the generic +`ice_requeue_runner.py` also supports opt-in `--checkpoint-forwarding trainer-pid`. +Run the runner from the exact source root, and pass its own Python executable +as the direct child command (no shell or `srun` wrapper). The existing +`runtime.slurm_requeue_owner=runner` callback installs SIGUSR2 before publishing +readiness. The runner binds that proof to its nonce, job/restart, source hashes, +UID, node, boot ID, PID namespace, cgroup, executable, working directory and +process start time; `pidfd` prevents signaling a reused PID. It signals only its +known child and never invokes `scancel`, including on failure. This route requires +one task/node/process and a working `scontrol` client inside the same container; +verify the site's client/config/Munge-socket bindings with a read-only query and +a dedicated requeue test before training. Do not expose or copy a Munge key. + +SIGUSR2 only latches a request; the callback saves at a completed batch hook, +publishes an immutable full-state checkpoint and receipt, and the runner strictly +validates a newer optimizer step before asking Slurm to requeue. Work after that +checkpoint can be replayed after restart. Keep periodic checkpoint/validation +cadences unchanged; allocation-boundary recovery saves are additional. Pass +`--target-global-step` with a completion sentinel to reject an early zero exit or +an older COMPLETE record. Boundary and validator time must fit the signal lead; +restart limits must cover the measured run budget. When the runner lives inside +an `srun` container, a batch-only Slurm signal may not reach it. Opt in to +`--allocation-boundary-lead-seconds` to arm a relative kernel timer from the +native `SLURM_JOB_END_TIME` epoch before discovery. Missing/invalid/past end times +fail closed; already inside the lead window uses the existing early-boundary +path. The timer only latches the same safe save request and cannot override +TERM/INT. It is disarmed/restored on every runner exit and is not inherited by +the trainer. This route requires a fixed allocation end time for the attempt; +do not shorten its walltime externally without coordinating the boundary. +Prove one actual deadline-triggered requeue separately from an injected signal. +When mid-validation continuation is unverified, set +`runtime.slurm_signal_train_batch_only=true`. Requests then remain latched through +validation and epoch/fit-end hooks until the next completed optimizer update. +This avoids conflicting with a scheduled checkpoint at the same global step. +The allocation lead and signal grace must exceed the measured **complete** +validation pass, longest training step, save and strict checkpoint verification; +a small diagnostic validation cap is not a production timing bound. +Defaults and the existing `slurm-steps` transport remain unchanged. + In this Action Flow branch, `egomimic/eval/artifact_paths.py` supplies the same namespace rule to EnergyScore, UNITE diagnostics, and Action Flow diagnostics. The Action Flow smoke verifier selects a matching active job/restart namespace diff --git a/egomimic/benchmarks/libero/replay.py b/egomimic/benchmarks/libero/replay.py index d31aec7d45..910713703a 100644 --- a/egomimic/benchmarks/libero/replay.py +++ b/egomimic/benchmarks/libero/replay.py @@ -223,7 +223,7 @@ def fetch(row): return raw / suite -def rebase_demo_xml(xml, robosuite_root, assets_root): +def rebase_demo_xml(xml, robosuite_root, assets_root, asset_remaps=()): """Relocate saved asset paths, preserving every model and physics parameter.""" tree = ET.fromstring(xml) for element in tree.iter(): @@ -231,7 +231,14 @@ def rebase_demo_xml(xml, robosuite_root, assets_root): if old is None: continue parts = Path(old).parts - if "robosuite" in parts: + matches = [(Path(prefix), Path(target).resolve()) for prefix, target in asset_remaps + if Path(old).is_relative_to(Path(prefix))] + if len(matches) > 1: + raise ValueError(f"Ambiguous demonstration asset remap: {old}") + if matches: + prefix, target = matches[0] + new = target / Path(old).relative_to(prefix) + elif "robosuite" in parts: index = max(i for i, part in enumerate(parts) if part == "robosuite") new = Path(robosuite_root).joinpath(*parts[index + 1 :]) elif "assets" in parts: diff --git a/egomimic/benchmarks/libero/state_tracks.py b/egomimic/benchmarks/libero/state_tracks.py new file mode 100644 index 0000000000..1f7bbfc48c --- /dev/null +++ b/egomimic/benchmarks/libero/state_tracks.py @@ -0,0 +1,737 @@ +"""Aligned native-state RGB, measured-pose and persistent-track preparation. + +This deliberately rerenders observations; original cached RGB/poses are never +paired with saved states by row or shifted heuristically. Outputs are raw +physical diagnostic data, not a normalized/DINO training cache. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import multiprocessing +import time +import xml.etree.ElementTree as ET +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +import h5py +import mujoco +import numpy as np + +from egomimic.benchmarks.libero.replay import rebase_demo_xml + +CONTRACT = "libero-same-recorded-state-v1" +RGB_POLICIES = ("exact", "uint8-roundoff-v1", "uint8-roundoff-v2") + + +def compare_rgb(original, repeated, policy="exact"): + """Separate bounded GPU raster rounding from exact physical-state replay. + + The opt-in policy was measured on repeated native GPU renders. It never + changes cached pixels, geometry tolerances, or the historical exact default. + """ + if policy not in RGB_POLICIES: + raise ValueError("Unknown RGB repeatability policy") + if original.shape != repeated.shape or original.dtype != np.uint8 or repeated.dtype != np.uint8: + raise ValueError("RGB replay requires matching uint8 arrays") + delta = np.abs(original.astype(np.int16) - repeated.astype(np.int16)) + result = {"policy": policy, "exact": bool(not np.any(delta)), + "max_abs": int(delta.max(initial=0)), + "changed_channels": int(np.count_nonzero(delta)), + "total_channels": int(delta.size), + "changed_fraction": float(np.count_nonzero(delta) / max(1, delta.size))} + if policy == "exact": + np.testing.assert_array_equal(original, repeated, err_msg="rgb") + elif result["max_abs"] > 2 or result["changed_fraction"] > ( + 0.00001 if policy == "uint8-roundoff-v1" else 0.0001 + ): + raise ValueError(f"RGB replay exceeds measured uint8 rounding policy: {result}") + if policy != "exact": + result["max_abs_limit"] = 2 + result["changed_fraction_limit"] = 0.00001 if policy == "uint8-roundoff-v1" else 0.0001 + return result + + +def digest(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def restore_state(model, data, state): + if model.na or model.nmocap: + raise ValueError("Recorded state omits activation/mocap state; unsupported") + if state.shape != (1 + model.nq + model.nv,) or not np.isfinite(state).all(): + raise ValueError("Invalid recorded state") + mujoco.mj_resetData(model, data) + data.time = state[0] + data.qpos[:] = state[1 : 1 + model.nq] + data.qvel[:] = state[1 + model.nq :] + mujoco.mj_forward(model, data) + np.testing.assert_array_equal(data.qpos, state[1 : 1 + model.nq]) + + +def camera_geometry(model, data, camera, height, width): + focal = height / (2 * np.tan(np.deg2rad(model.cam_fovy[camera]) / 2)) + K = np.array( + [[focal, 0, (width - 1) / 2], [0, focal, (height - 1) / 2], [0, 0, 1.0]] + ) + transform = np.eye(4) + # Camera convention: +X right, +Y down, +Z forward. + transform[:3, :3] = data.cam_xmat[camera].reshape(3, 3) @ np.diag([1.0, -1.0, -1.0]) + transform[:3, 3] = data.cam_xpos[camera] + return K, transform + + +def pose_joint_addresses(model, names): + """Resolve explicitly named measured scalar joints; never infer policy actions.""" + if len(set(names)) != len(names): + raise ValueError("Duplicate pose joint") + result = [] + for name in names: + joint = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) + if joint < 0 or model.jnt_type[joint] not in ( + mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE + ): + raise ValueError(f"Pose joint must be a known scalar joint: {name}") + result.append(int(model.jnt_qposadr[joint])) + return result + + +def build_arrays(model, states, cameras, pose_site, height, width, grid, pose_joints=()): + if states.ndim != 2 or len(states) < 2 or not np.all(np.diff(states[:, 0]) > 0): + raise ValueError( + "Require at least two strictly increasing physical state times" + ) + if min(height, width, *grid) < 2: + raise ValueError("Invalid image/grid dimensions") + data = mujoco.MjData(model) + pose_addresses = pose_joint_addresses(model, pose_joints) + camids = [mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, c) for c in cameras] + site = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, pose_site) + if not cameras or len(set(cameras)) != len(cameras) or min([site, *camids]) < 0: + raise ValueError("Unknown or duplicate camera / unknown pose site") + option = mujoco.MjvOption() + option.geomgroup[0] = 0 + hit = np.zeros(1, np.int32) + + def ray(origin, direction): + distance = mujoco.mj_ray( + model, data, origin, direction, option.geomgroup, 1, -1, hit + ) + return distance, int(hit[0]) + + restore_state(model, data, states[0]) + K, transform = camera_geometry(model, data, camids[0], height, width) + ys, xs = np.meshgrid( + np.linspace(0, height - 1, grid[0]), + np.linspace(0, width - 1, grid[1]), + indexing="ij", + ) + pixels = np.c_[xs.ravel(), ys.ravel()] + local, bodies, geoms, valid_indices = [], [], [], [] + for index, pixel in enumerate(pixels): + direction = transform[:3, :3] @ np.linalg.solve(K, [*pixel, 1.0]) + direction /= np.linalg.norm(direction) + distance, geom = ray(transform[:3, 3], direction) + if distance < 0 or geom < 0: + continue + body = int(model.geom_bodyid[geom]) + world = transform[:3, 3] + distance * direction + local.append(data.xmat[body].reshape(3, 3).T @ (world - data.xpos[body])) + bodies.append(body) + geoms.append(geom) + valid_indices.append(index) + if not local: + raise ValueError("No reference ray hits") + local = np.asarray(local) + bodies = np.asarray(bodies, dtype=np.int64) + rows = { + key: [] + for key in ( + "rgb", + "pose", + "xyz", + "uv", + "visibility", + "in_view", + "world_from_camera", + "body_position", + "body_rotation", + ) + } + renderer = mujoco.Renderer(model, height=height, width=width) + try: + for state in states: + restore_state(model, data, state) + rotations = data.xmat[bodies].reshape(-1, 3, 3) + positions = data.xpos[bodies].copy() + world = np.einsum("nij,nj->ni", rotations, local) + positions + images, transforms, projections, in_views, visibles = [], [], [], [], [] + for camera in camids: + K, transform = camera_geometry(model, data, camera, height, width) + camera_xyz = (world - transform[:3, 3]) @ transform[:3, :3] + z = camera_xyz[:, 2] + uv = camera_xyz[:, :2] / np.where(z == 0, np.nan, z)[:, None] + uv = uv * np.diag(K)[:2] + K[:2, 2] + inside = (z > 0) & (uv[:, 0] >= 0) & (uv[:, 0] < width) + inside &= (uv[:, 1] >= 0) & (uv[:, 1] < height) + visible = np.zeros(len(local), bool) + for point in np.flatnonzero(inside): + delta = world[point] - transform[:3, 3] + distance = np.linalg.norm(delta) + hit_distance, geom = ray(transform[:3, 3], delta / distance) + visible[point] = ( + geom >= 0 + and model.geom_bodyid[geom] == bodies[point] + and abs(hit_distance - distance) < 1e-4 + ) + renderer.update_scene(data, camera=camera, scene_option=option) + images.append(renderer.render().copy()) + transforms.append(transform) + projections.append(uv) + in_views.append(inside) + visibles.append(visible) + quaternion = np.zeros(4) + mujoco.mju_mat2Quat(quaternion, data.site_xmat[site]) + rows["pose"].append(np.r_[data.site_xpos[site], quaternion, data.qpos[pose_addresses]]) + for key, value in ( + ("rgb", images), + ("xyz", world), + ("uv", projections), + ("world_from_camera", transforms), + ("visibility", visibles), + ("in_view", in_views), + ("body_position", positions), + ("body_rotation", rotations.copy()), + ): + rows[key].append(value) + finally: + renderer.close() + arrays = {key: np.asarray(value) for key, value in rows.items()} + arrays.update( + time=states[:, 0], + state=states, + source_row=np.arange(len(states)), + point_body_id=bodies, + point_geom_id=np.asarray(geoms), + body_local_xyz=local, + reference_grid_xy=pixels, + valid_grid_indices=np.asarray(valid_indices), + coordinate_available=np.ones((len(states), len(local)), bool), + camera_K=np.asarray( + [camera_geometry(model, data, c, height, width)[0] for c in camids] + ), + ) + validate_arrays(arrays) + return arrays + + +def validate_arrays(a): + """Reject broken row identity, geometry, timestamps and visibility contracts.""" + count = len(a["time"]) + if not np.array_equal(a["source_row"], np.arange(count)): + raise ValueError("Source row permutation/shift") + if not np.array_equal(a["time"], a["state"][:, 0]) or not np.all( + np.diff(a["time"]) > 0 + ): + raise ValueError("State/time mismatch") + for key in ( + "rgb", + "pose", + "xyz", + "uv", + "visibility", + "in_view", + "world_from_camera", + "body_position", + "body_rotation", + "coordinate_available", + ): + if len(a[key]) != count: + raise ValueError(f"Misaligned channel: {key}") + if ( + a["rgb"].dtype != np.uint8 + or a["rgb"].ndim != 5 + or a["pose"].ndim != 2 + or a["pose"].shape[1] < 7 + ): + raise ValueError("Invalid RGB/pose schema") + if not np.isfinite(a["xyz"]).all() or not np.isfinite(a["pose"]).all(): + raise ValueError("Nonfinite geometry") + if not np.all(a["coordinate_available"]): + raise ValueError("Exact simulator coordinates must survive occlusion") + if np.any(a["visibility"] & ~a["in_view"]): + raise ValueError("Visible point outside camera") + expected = ( + np.einsum("tpij,pj->tpi", a["body_rotation"], a["body_local_xyz"]) + + a["body_position"] + ) + np.testing.assert_allclose(a["xyz"], expected, rtol=0, atol=1e-10) + np.testing.assert_allclose(np.linalg.norm(a["pose"][:, 3:7], axis=-1), 1, atol=1e-10) + np.testing.assert_allclose( + a["uv"][0, 0], a["reference_grid_xy"][a["valid_grid_indices"]], atol=1e-8 + ) + + +def exclude_visualization_mocap(xml, names): + """Explicitly remove nonphysical marker bodies whose poses were not saved. + + Refuse collision geometry, child bodies, joints or surviving references. + Physical joint-state ordering must remain unchanged after removal. + """ + if not names: + return xml + if len(names) != len(set(names)): + raise ValueError("Duplicate marker body names") + original = mujoco.MjModel.from_xml_string(xml) + tree = ET.fromstring(xml) + removed_names = set() + for name in names: + body_id = mujoco.mj_name2id(original, mujoco.mjtObj.mjOBJ_BODY, name) + matches = [(p, e) for p in tree.iter() for e in p + if e.tag == "body" and e.get("name") == name] + if body_id < 0 or original.body_mocapid[body_id] < 0 or len(matches) != 1: + raise ValueError(f"Not a unique mocap visualization body: {name}") + parent, element = matches[0] + geoms = np.flatnonzero(original.geom_bodyid == body_id) + if (any(e.tag in ("body", "joint", "freejoint") for e in list(element.iter())[1:]) + or np.any(original.geom_contype[geoms]) + or np.any(original.geom_conaffinity[geoms])): + raise ValueError(f"Marker body has physical geometry or children: {name}") + removed_names.update(e.get("name") for e in element.iter() if e.get("name")) + parent.remove(element) + if any(value in removed_names for e in tree.iter() for key, value in e.attrib.items() + if key != "name"): + raise ValueError("Excluded marker still referenced by model") + result = ET.tostring(tree, encoding="unicode") + filtered = mujoco.MjModel.from_xml_string(result) + if (original.nq, original.nv, original.nu, original.na) != ( + filtered.nq, filtered.nv, filtered.nu, filtered.na): + raise ValueError("Marker exclusion changed physical state dimensions") + original_joints = [mujoco.mj_id2name(original, mujoco.mjtObj.mjOBJ_JOINT, i) + for i in range(original.njnt)] + filtered_joints = [mujoco.mj_id2name(filtered, mujoco.mjtObj.mjOBJ_JOINT, i) + for i in range(filtered.njnt)] + if original_joints != filtered_joints: + raise ValueError("Marker exclusion changed joint-state ordering") + return result + + +def prepare(args): + output = Path(args.output) + if output.exists(): + raise FileExistsError(f"Refusing to overwrite cache: {output}") + input_schema = getattr(args, "input_schema", "libero") + with h5py.File(args.demo, "r") as handle: + if input_schema == "quest": + states = handle["state"][:] + raw_xml = handle["destination_model_xml"][()] + original_pose = handle["numeric_observations/robot0_eef_pos"][:] + elif input_schema == "libero": + demo = handle[f"data/{args.episode}"] + states = demo["states"][:] + raw_xml = demo.attrs["model_file"] + original_pose = demo["obs/ee_pos"][:] if "obs/ee_pos" in demo else None + else: + raise ValueError("Unknown recorded input schema") + remaps = getattr(args, "asset_remap", None) or () + excluded = getattr(args, "exclude_visualization_body", None) or () + xml = rebase_demo_xml( + raw_xml, Path(args.robot_assets).resolve(), Path(args.scene_assets).resolve(), remaps + ) + xml = exclude_visualization_mocap(xml, excluded) + model = mujoco.MjModel.from_xml_string(xml) + pose_joints = getattr(args, "pose_joints", None) or [] + arrays = build_arrays( + model, states, args.cameras, args.pose_site, args.height, args.width, args.grid, pose_joints + ) + assets = { + str(Path(e.get("file")).resolve()): digest(e.get("file")) + for e in ET.fromstring(xml).iter() + if e.get("file") + } + output.mkdir(parents=True) + np.savez_compressed(output / "aligned.npz", **arrays) + (output / "model.xml").write_text(xml) + mismatch = ( + None + if original_pose is None + else float(np.linalg.norm(arrays["pose"][:, :3] - original_pose, axis=-1).max()) + ) + manifest = dict( + schema=CONTRACT, + source_demo=str(Path(args.demo).resolve()), + source_demo_sha256=digest(args.demo), + episode=args.episode, + input_schema=input_schema, + asset_remaps=list(remaps), + excluded_visualization_bodies=list(excluded), + alignment="Every RGB, measured pose and track row is recomputed after restoring the exact same source state. No original cached observation is consumed as a target.", + original_cache_max_eef_error_m=mismatch, + original_cache_repaired=False, + cameras=args.cameras, + pose_site=args.pose_site, + pose_joints=pose_joints, + height=args.height, + width=args.width, + grid=args.grid, + pose_fields=["x_m", "y_m", "z_m", "qw", "qx", "qy", "qz"] + [f"qpos:{name}" for name in pose_joints], + reference_frame="recorded simulator world", + time_unit="physical simulator seconds", + rgb_origin="top_left", + camera_axes="right_down_forward", + visibility="ray first-hit surface at current state; separate from coordinate availability", + track_method="reference visual-geometry ray-hit attached to persistent body-local XYZ", + normalized=False, + dino_features=False, + training_ready=False, + mujoco_version=mujoco.__version__, + preparation_source_sha256=digest(__file__), + asset_hashes=assets, + artifacts={ + name: digest(output / name) for name in ("aligned.npz", "model.xml") + }, + ) + (output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + return validate_cache(output, replay=True) + + +def validate_cache(output, replay=False, rgb_policy="exact"): + output = Path(output) + manifest = json.loads((output / "manifest.json").read_text()) + if manifest["schema"] != CONTRACT: + raise ValueError("Unsupported alignment contract") + for name, expected in manifest["artifacts"].items(): + if digest(output / name) != expected: + raise ValueError(f"Cache hash mismatch: {name}") + with np.load(output / "aligned.npz", allow_pickle=False) as archive: + arrays = dict(archive) + validate_arrays(arrays) + if arrays["pose"].shape[1] != len(manifest["pose_fields"]): + raise ValueError("Pose schema width disagrees with arrays") + if replay: + for name, expected in manifest["asset_hashes"].items(): + if digest(name) != expected: + raise ValueError(f"Asset changed: {name}") + model = mujoco.MjModel.from_xml_path(str((output / "model.xml").resolve())) + repeated = build_arrays( + model, + arrays["state"], + manifest["cameras"], + manifest["pose_site"], + manifest["height"], + manifest["width"], + manifest["grid"], + manifest.get("pose_joints", ()), + ) + # Geometry always remains bit-exact, independent of RGB rounding policy. + for key in [key for key in arrays if key != "rgb"]: + np.testing.assert_array_equal(arrays[key], repeated[key], err_msg=key) + rgb_result = compare_rgb(arrays["rgb"], repeated["rgb"], rgb_policy) + result = { + "status": "PASS", + "contract": CONTRACT, + "frames": len(arrays["time"]), + "points": len(arrays["point_body_id"]), + "all_channels_native_replay_exact": bool(replay and rgb_result["exact"]), + "native_geometry_replay_exact": bool(replay), + "training_ready": False, + } + if replay: + result["rgb_replay"] = rgb_result + return result + + +def export_rgb(args): + """Stream true high-resolution native RGB from verified saved states. + + Low-resolution persistent tracks are preserved; this does not regenerate + the track grid or substitute upsampled images for high-resolution renders. + """ + source = Path(args.output).resolve() + target = Path(args.rgb_cache).resolve() + receipt_path = target.with_suffix(target.suffix + ".json") + if target.exists() or receipt_path.exists(): + raise FileExistsError("Preserve existing native RGB output and receipt") + if min(args.height, args.width) < 2: + raise ValueError("Invalid native RGB resolution") + proof = validate_cache(source, replay=True, rgb_policy=args.rgb_policy) + manifest = json.loads((source / "manifest.json").read_text()) + with np.load(source / "aligned.npz", allow_pickle=False) as archive: + states = archive["state"] + times = archive["time"] + original_transforms = archive["world_from_camera"] + model = mujoco.MjModel.from_xml_path(str(source / "model.xml")) + # Increasing framebuffer capacity changes neither state nor visual settings. + model.vis.global_.offheight = max(model.vis.global_.offheight, args.height) + model.vis.global_.offwidth = max(model.vis.global_.offwidth, args.width) + data = mujoco.MjData(model) + cameras = [mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, name) + for name in manifest["cameras"]] + if min(cameras) < 0: + raise ValueError("Missing recorded native camera") + option = mujoco.MjvOption() + option.geomgroup[0] = 0 + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(target.suffix + ".partial") + if temporary.exists(): + raise FileExistsError("Preserve partial native RGB output for diagnosis") + renderer = mujoco.Renderer(model, height=args.height, width=args.width) + try: + with h5py.File(temporary, "x") as cache: + cache.attrs.update(schema="libero-native-rgb-v1", native_render=True, + height=args.height, width=args.width) + cache["time"] = times + rgb = cache.create_dataset("rgb", (len(times), len(cameras), args.height, + args.width, 3), dtype="u1", + chunks=(1, 1, args.height, args.width, 3), + compression="lzf") + transforms = cache.create_dataset("camera_to_world", (len(times), len(cameras), 4, 4), dtype="f8") + intrinsics = [] + for row, state in enumerate(states): + restore_state(model, data, state) + for view, camera in enumerate(cameras): + K, transform = camera_geometry(model, data, camera, args.height, args.width) + np.testing.assert_array_equal(transform, original_transforms[row, view], + err_msg="High-resolution camera/state alignment") + renderer.update_scene(data, camera=camera, scene_option=option) + rgb[row, view] = renderer.render() + transforms[row, view] = transform + if row == 0: + intrinsics.append(K) + cache["camera_K"] = np.asarray(intrinsics) + cache.flush() + finally: + renderer.close() + temporary.replace(target) + receipt = {"schema": "libero-native-rgb-v1", "status": "PASS", + "source_manifest_sha256": digest(source / "manifest.json"), + "source_aligned_sha256": digest(source / "aligned.npz"), + "rgb_cache_sha256": digest(target), "rgb_cache": str(target), + "native_render": True, "height": args.height, "width": args.width, + "cameras": manifest["cameras"], "frames": len(times), + "native_geometry_replay_exact": True, "low_resolution_replay": proof, + "source": str(source), "preparation_source_sha256": digest(__file__), + "mujoco_version": mujoco.__version__, "training_ready": False} + receipt_path.write_text(json.dumps(receipt, indent=2) + "\n") + return receipt + + +def _export_rgb_episode(options): + target = Path(options["rgb_cache"]) + receipt_path = target.with_suffix(target.suffix + ".json") + if target.exists() and receipt_path.exists(): + receipt = json.loads(receipt_path.read_text()) + source = Path(options["output"]).resolve() + if (receipt.get("status") != "PASS" or + receipt.get("preparation_source_sha256") != digest(__file__) or + receipt.get("source") != str(source) or + receipt.get("source_manifest_sha256") != digest(source / "manifest.json") or + receipt.get("source_aligned_sha256") != digest(source / "aligned.npz") or + receipt.get("rgb_cache_sha256") != digest(target) or + receipt.get("height") != options["height"] or + receipt.get("width") != options["width"] or + receipt.get("low_resolution_replay", {}).get("rgb_replay", {}).get("policy") != options["rgb_policy"] or + receipt.get("native_geometry_replay_exact") is not True): + raise ValueError("Native RGB reuse identity/proof mismatch") + receipt["reused"] = True + else: + receipt = export_rgb(argparse.Namespace(**options)) + return {key: receipt[key] for key in ("status", "native_geometry_replay_exact", "frames", + "height", "width", "rgb_cache_sha256", "low_resolution_replay")} + + +def export_rgb_corpus(args): + """Parallel native RGB export over an explicit inventory, with resumable proof.""" + spec_path = Path(args.spec).resolve() + spec = json.loads(spec_path.read_text()) + status_path = Path(args.status).resolve() + if args.workers < 1 or not spec.get("episodes"): + raise ValueError("Need positive workers and explicit native cache inventory") + if status_path.exists() and json.loads(status_path.read_text())["spec_sha256"] != digest(spec_path): + raise ValueError("Native RGB inventory changed; preserve previous status") + targets, identities, options_list = set(), set(), [] + for episode in spec["episodes"]: + key = (episode["embodiment"], episode["episode_id"]) + options = {**spec.get("common", {}), "output": episode["output"], "rgb_cache": episode["rgb_cache"]} + target = str(Path(options["rgb_cache"]).resolve()) + if target in targets or key in identities or Path(options["output"]).resolve() == Path(target): + raise ValueError("Duplicate native RGB identity/output") + targets.add(target); identities.add(key); options_list.append((key, options)) + status_path.parent.mkdir(parents=True, exist_ok=True) + status = {"spec_sha256": digest(spec_path), "status": "RUNNING", "expected": len(options_list), + "episodes": [], "training_ready": False, "preparation_source_sha256": digest(__file__)} + + def save(): + temporary = status_path.with_suffix(".tmp") + temporary.write_text(json.dumps(status, indent=2) + "\n") + temporary.replace(status_path) + + save() + with ProcessPoolExecutor(max_workers=args.workers, + mp_context=multiprocessing.get_context("spawn")) as pool: + futures = {pool.submit(_export_rgb_episode, options): key for key, options in options_list} + for future in as_completed(futures): + key = futures[future] + row = {"embodiment": key[0], "episode_id": key[1]} + try: + row.update(future.result()) + except Exception as error: + row.update(status="FAILED", error=f"{type(error).__name__}: {error}") + status["episodes"].append(row); save() + print(json.dumps({"completed": len(status["episodes"]), **row}), flush=True) + status["status"] = "PASS" if all(row["status"] == "PASS" for row in status["episodes"]) else "FAILED" + save() + if status["status"] != "PASS": + raise RuntimeError("Native RGB export failed episodes; preserve status and outputs") + return {key: value for key, value in status.items() if key != "episodes"} + + +def _prepare_corpus_episode(options): + """Reuse verified completed episodes; never accept a changed request as a hit.""" + deadline = time.monotonic() + options.get("wait_input_seconds", 0) + while not Path(options["demo"]).is_file(): + if time.monotonic() >= deadline: + raise FileNotFoundError(options["demo"]) + time.sleep(min(10, max(0, deadline - time.monotonic()))) + if options.get("expected_source_sha256") and digest(options["demo"]) != options["expected_source_sha256"]: + raise ValueError("Source demonstration hash differs from corpus inventory") + output = Path(options["output"]) + if output.exists(): + manifest = json.loads((output / "manifest.json").read_text()) + if manifest["preparation_source_sha256"] != digest(__file__): + raise ValueError("Existing cache preparation source mismatch") + for path, expected in manifest["asset_hashes"].items(): + if digest(path) != expected: + raise ValueError(f"Existing cache asset changed: {path}") + for key in ("episode", "input_schema", "cameras", "pose_site", "height", "width", "grid", "pose_joints"): + if manifest.get(key) != options.get(key): + raise ValueError(f"Existing cache request mismatch: {key}") + for recorded, requested in (("asset_remaps", "asset_remap"), + ("excluded_visualization_bodies", "exclude_visualization_body")): + if manifest.get(recorded, []) != options.get(requested, []): + raise ValueError(f"Existing cache request mismatch: {recorded}") + if manifest["source_demo"] != str(Path(options["demo"]).resolve()) or manifest["source_demo_sha256"] != digest(options["demo"]): + raise ValueError("Existing cache source mismatch") + # A manifest is written before native replay finishes. Its presence and + # hashes therefore cannot prove that an earlier preparation passed. + # Until replay success has a hash-bound receipt, every reuse must replay. + result = validate_cache(output, replay=True) + result["reused"] = True + return result + return prepare(argparse.Namespace(**options)) + + +def prepare_corpus(args): + """Run an explicit corpus inventory through the maintained native preparer.""" + spec_path = Path(args.spec).resolve() + spec = json.loads(spec_path.read_text()) + output = Path(args.output).resolve() + if args.workers < 1 or not spec.get("episodes"): + raise ValueError("Need positive workers and an explicit episode inventory") + options_list, identities, sources = [], set(), set() + for episode in spec["episodes"]: + embodiment, identity = episode["embodiment"], episode["episode_id"] + if any(not name or Path(name).name != name or name in (".", "..") for name in (embodiment, identity)): + raise ValueError("Unsafe corpus identity") + key = (embodiment, identity) + options = {**spec.get("common", {}), **spec["embodiments"][embodiment], + "demo": episode["demo"], "episode": episode["episode"], + "output": str(output / embodiment / identity), + "expected_source_sha256": episode.get("sha256"), + "wait_input_seconds": getattr(args, "wait_input_seconds", 0)} + physical = (str(Path(options["demo"]).resolve()), options["episode"]) + if key in identities or physical in sources: + raise ValueError("Duplicate corpus identity or physical episode") + identities.add(key); sources.add(physical) + options_list.append((key, options)) + output.mkdir(parents=True, exist_ok=True) + status_path = output / "CORPUS_STATUS.json" + if status_path.exists() and json.loads(status_path.read_text())["spec_sha256"] != digest(spec_path): + raise ValueError("Corpus specification changed; preserve the original output") + status = {"spec_sha256": digest(spec_path), "spec": str(spec_path), + "status": "RUNNING", "expected": len(options_list), "episodes": [], + "training_ready": False} + + def save(): + temporary = status_path.with_suffix(".tmp") + temporary.write_text(json.dumps(status, indent=2) + "\n") + temporary.replace(status_path) + + save() + with ProcessPoolExecutor(max_workers=args.workers, + mp_context=multiprocessing.get_context("spawn")) as pool: + futures = {pool.submit(_prepare_corpus_episode, options): key for key, options in options_list} + for future in as_completed(futures): + key = futures[future] + row = {"embodiment": key[0], "episode_id": key[1]} + try: + row.update(future.result()) + except Exception as error: + row.update(status="FAILED", error=f"{type(error).__name__}: {error}") + status["episodes"].append(row) + save() + print(json.dumps({"completed": len(status["episodes"]), **row}), flush=True) + status["status"] = "PASS" if all(r["status"] == "PASS" for r in status["episodes"]) else "FAILED" + save() + if status["status"] != "PASS": + raise RuntimeError("Native corpus preparation has failed episodes; see CORPUS_STATUS.json") + return {key: value for key, value in status.items() if key != "episodes"} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + prepare_parser = commands.add_parser("prepare") + for name in ("demo", "robot-assets", "scene-assets", "output", "pose-site"): + prepare_parser.add_argument(f"--{name}", required=True) + prepare_parser.add_argument("--episode", default="demo_0") + prepare_parser.add_argument("--input-schema", choices=["libero", "quest"], default="libero") + prepare_parser.add_argument("--asset-remap", nargs=2, action="append", metavar=("OLD_PREFIX", "NEW_ROOT")) + prepare_parser.add_argument("--exclude-visualization-body", action="append") + prepare_parser.add_argument("--cameras", nargs="+", required=True) + prepare_parser.add_argument("--pose-joints", nargs="+", help="Explicit measured scalar joint names appended after EEF pose") + prepare_parser.add_argument("--height", type=int, required=True) + prepare_parser.add_argument("--width", type=int, required=True) + prepare_parser.add_argument("--grid", nargs=2, type=int, required=True) + check = commands.add_parser("validate") + check.add_argument("--output", required=True) + check.add_argument("--replay", action="store_true") + check.add_argument("--rgb-policy", choices=RGB_POLICIES, default="exact") + rgb = commands.add_parser("export-rgb") + rgb.add_argument("--output", required=True, help="Existing native aligned episode") + rgb.add_argument("--rgb-cache", required=True) + rgb.add_argument("--height", type=int, required=True) + rgb.add_argument("--width", type=int, required=True) + rgb.add_argument("--rgb-policy", choices=RGB_POLICIES, default="exact") + rgb_corpus = commands.add_parser("export-rgb-corpus") + rgb_corpus.add_argument("--spec", required=True) + rgb_corpus.add_argument("--status", required=True) + rgb_corpus.add_argument("--workers", type=int, default=1) + corpus = commands.add_parser("prepare-corpus") + corpus.add_argument("--spec", required=True) + corpus.add_argument("--output", required=True) + corpus.add_argument("--workers", type=int, default=1) + corpus.add_argument("--wait-input-seconds", type=int, default=0, + help="Allow atomic dataset transfers to finish; never read temporary transfer files") + args = parser.parse_args() + print( + json.dumps( + prepare(args) + if args.command == "prepare" + else prepare_corpus(args) if args.command == "prepare-corpus" + else export_rgb(args) if args.command == "export-rgb" + else export_rgb_corpus(args) if args.command == "export-rgb-corpus" + else validate_cache(args.output, args.replay, args.rgb_policy), + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/eval/world_model_eval.py b/egomimic/eval/world_model_eval.py new file mode 100644 index 0000000000..0e08442829 --- /dev/null +++ b/egomimic/eval/world_model_eval.py @@ -0,0 +1,69 @@ +"""Normalized prediction metrics for the offline world-model family.""" + +import torch + +from egomimic.eval.eval import Eval, EvaluationDataRequirements +from egomimic.models.world_model import MODALITIES + + +class WorldModelEval(Eval): + def __init__(self, sample_seed=42): + self.sample_seed = int(sample_seed) + self.totals = {} + + def bind_data_context(self, *, normalizer): + if normalizer is not None: + raise ValueError("World-model evaluator expects prepared normalized caches") + + def data_requirements(self): + return EvaluationDataRequirements(required_keys=("world", "embodiment")) + + def trainer_overrides(self): + return {} # Respect configured step-based validation; no hidden override. + + def on_validation_start(self): + self.totals = {} + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + if self.model.device.type == "cuda": + devices = [self.model.device.index or 0] + else: + devices = [] + # Repeatable validation without advancing the training sampler RNG. + with torch.random.fork_rng(devices=devices): + torch.manual_seed( + self.sample_seed + int(batch_idx) + 100000 * dataloader_idx + ) + output = self.model.forward_eval(batch) + for source, item in batch.items(): + for name in MODALITIES: + s = item["world"][name] + weight = ( + s.get("target_valid", s["valid"]) & s["valid"] & ~s["observed"] + ).float() + if "confidence" in s: + weight *= s["confidence"] + active = weight > 0 + target = torch.where(active[..., None], s["value"], 0) + prediction = output[source]["world_prediction"][name] + error = (prediction.float() - target.float()).square().mean(-1) + if not torch.isfinite(error[active]).all(): + raise ValueError( + "Non-finite world-model validation prediction/target" + ) + key = f"val/{source}/{name}_normalized_mse" + numerator = (error * weight).sum().double() + denominator = weight.sum().double() + old = self.totals.get( + key, (numerator.new_zeros(()), denominator.new_zeros(())) + ) + self.totals[key] = (old[0] + numerator, old[1] + denominator) + return output + + def log_epoch_metrics(self): + for key, (numerator, denominator) in self.totals.items(): + if denominator > 0: + self.model.log(key, (numerator / denominator).float(), sync_dist=False) + + def on_validation_end(self): + return None diff --git a/egomimic/hydra_configs/data/world_model/prepared.yaml b/egomimic/hydra_configs/data/world_model/prepared.yaml new file mode 100644 index 0000000000..46b65527c7 --- /dev/null +++ b/egomimic/hydra_configs/data/world_model/prepared.yaml @@ -0,0 +1,16 @@ +_target_: egomimic.rldb.world_model_datamodule.WorldModelDataModule +manifest_path: ??? +batch_size: 1 +num_workers: 0 +mode: continuous +seed: 42 +min_frames: 2 +# Requested continuous count range; the independent model token guard still applies. +# This is not evidence that 16 full image/track timestamps fit GPU memory. +max_frames: 16 +discrete_frames: 4 +stride: 1 +forecasting_probability: 0.5 +# Physical offsets from the current/anchor frame, clipped to recorded times. +max_past_seconds: 3.0 +max_future_seconds: 3.0 diff --git a/egomimic/hydra_configs/experiment/world_model/continuous.yaml b/egomimic/hydra_configs/experiment/world_model/continuous.yaml new file mode 100644 index 0000000000..f023004cae --- /dev/null +++ b/egomimic/hydra_configs/experiment/world_model/continuous.yaml @@ -0,0 +1,38 @@ +# @package _global_ +# Offline world-model recipe. Requires a separately verified prepared-data manifest. +defaults: + - override /model: world_model/base + - override /data: world_model/prepared + - override /trainer: default + +name: mot_world_model +norm_stats: null +normalizer: null +data: + # Smoke may cap each embodiment; full validation remains exhaustive. + validation_batches_per_embodiment: null +inference_config: + enabled: false # This is not the generic action-policy robot deployment adapter. +evaluator: + _target_: egomimic.eval.world_model_eval.WorldModelEval + sample_seed: 42 +trainer: + accelerator: gpu + precision: 32-true # User-selected FP32; inherited by the discrete comparison. + devices: 1 + min_epochs: 0 + max_epochs: -1 + limit_train_batches: 1.0 + limit_val_batches: 1.0 + val_check_interval: 15000 + check_val_every_n_epoch: null + max_steps: 150000 # User-selected budget for the five MoT experiments. + num_sanity_val_steps: 0 +callbacks: + model_checkpoint: + filename: "epoch_{epoch:04d}-step_{step:09d}" + every_n_epochs: 0 + every_n_train_steps: 15000 + save_on_train_epoch_end: false + save_top_k: -1 + save_last: true diff --git a/egomimic/hydra_configs/experiment/world_model/discrete.yaml b/egomimic/hydra_configs/experiment/world_model/discrete.yaml new file mode 100644 index 0000000000..b6c3080b40 --- /dev/null +++ b/egomimic/hydra_configs/experiment/world_model/discrete.yaml @@ -0,0 +1,8 @@ +# @package _global_ +defaults: + - /experiment/world_model/continuous + - _self_ +data: + mode: discrete + max_past_seconds: null + max_future_seconds: null diff --git a/egomimic/hydra_configs/model/world_model/base.yaml b/egomimic/hydra_configs/model/world_model/base.yaml new file mode 100644 index 0000000000..846c0508dc --- /dev/null +++ b/egomimic/hydra_configs/model/world_model/base.yaml @@ -0,0 +1,44 @@ +_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 +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stage_ids: + world_model: 0 + stages: + - _target_: egomimic.pipeline.stages_world_model.WorldModelStage + warmup_steps: 10000 + pose_ramp_steps: 10000 + image_weight: 1.0 + track_weight: 1.0 + pose_weight: 1.0 + sampler_steps: 20 + # Never silently train on zero-filled missing geometry. Conditional sample + # completion is an explicit approximation with extra denoising cost. + missing_track_policy: error + imputation_steps: 20 + model: + _target_: egomimic.models.world_model.MoTWorldModel + # REQUIRED from prepared measured-pose schemas; never use policy action widths. + pose_dims: ??? + image_dim: 768 + # Opt-in x0 diagnostic changes image target/noise weighting; tracks stay v. + image_prediction_type: v + width: 384 + depth: 8 + heads: 6 + mlp_ratio: 4 + image_adapter_dim: 48 + max_cameras: 3 + max_tracks: 1120 + max_tokens: 8192 + # Exact full attention; opt-in memory controls pending full-size GPU proof. + attention_query_chunk_size: 0 + activation_checkpointing: false + time_mode: ${data.mode} diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index a5a429540e..80ddea791f 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -40,11 +40,17 @@ launch_params: nodes: 1 runtime: + # Opt in to disabling TF32 as well as autocast; requires precision=32-true + # and NVIDIA_TF32_OVERRIDE=0 in the launch environment. + strict_fp32: false # Lightning owns SIGUSR1/scontrol by default. An external ICE runner must # opt in explicitly and prove that child/framework requeue is disabled. slurm_requeue_owner: lightning # Runner-owned mode forwards SIGUSR2 only to request a full-state save. slurm_save_signal: SIGUSR2 + # Opt in when mid-validation checkpoints cannot be resumed exactly. + # Timer lead must cover complete validation plus the next optimizer step. + slurm_signal_train_batch_only: false slurm_signal_checkpoint_dir: ${paths.output_dir}/checkpoints seed: 42 diff --git a/egomimic/models/world_model.py b/egomimic/models/world_model.py new file mode 100644 index 0000000000..35f73aa351 --- /dev/null +++ b/egomimic/models/world_model.py @@ -0,0 +1,655 @@ +"""Three-modality continuous/discrete-time world model (independent of Action Flow). + +Inputs are normalized prepared tokens. Every stream has value[B,N,C], +time[B,N], valid[B,N] (padding only), observed[B,N]. Image/track streams +also have reference xy[B,N,2], camera[B,N]; tracks have persistent point_id. +Targets/confidence are used by the objective, never the backbone. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping + +import torch +from torch import nn +from torch.nn import functional as F +from torch.utils.checkpoint import checkpoint + +MODALITIES = ("image", "track", "proprio") + + +def masked_mse(pred, target, weight): + active = weight > 0 + if not torch.isfinite(target[active]).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) + 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. + denominator = torch.where(total_weight > 0, total_weight, 1.0) + return (error * weight.float()).sum() / denominator + + +def noise_coefficients(level): + angle = level.float() * (math.pi / 2) + return angle.cos()[..., None], angle.sin()[..., None] + + +def temporal_rope(x, times): + # x[B,H,N,D]; rotation evaluated in FP32 even under BF16 autocast. + half = x.shape[-1] // 2 + freq = torch.exp(-math.log(10000) * torch.arange(half, device=x.device) / half) + phase = times.float()[:, None, :, None] * freq + left, right = x.float().chunk(2, dim=-1) + return torch.cat( + ( + left * phase.cos() - right * phase.sin(), + left * phase.sin() + right * phase.cos(), + ), + -1, + ).to(x.dtype) + + +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): + 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: + raise ValueError("attention_query_chunk_size must be a nonnegative integer") + self.attention_query_chunk_size = attention_query_chunk_size + 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) + self.ffn = nn.ModuleList( + nn.Sequential( + nn.Linear(width, width * mlp_ratio), + nn.GELU(), + nn.Linear(width * mlp_ratio, width), + ) + for _ in MODALITIES + ) + + def forward(self, streams, times, valid): + sizes = [s.shape[1] for s in streams] + qs, ks, vs = [], [], [] + for i, x in enumerate(streams): + b, n, d = x.shape + q, k, v = ( + self.qkv[i](self.norm1[i](x)) + .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])) + 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) + 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 + # softmax denominator; this is full attention, not a local window. + # Recompute each chunk in backward rather than retaining all chunks' + # 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._attend, q, key, value, mask, use_reentrant=False) + else: + part = self._attend(q, key, value, mask) + parts.append(part) + result = torch.cat(parts, dim=2) + else: + result = self._attend(query, key, value, mask) + 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)) + outputs.append(x.masked_fill(~valid[i][..., None], 0)) + return outputs + + @staticmethod + def _attend(query, key, value, mask): + return F.scaled_dot_product_attention( + query, key, value, attn_mask=mask, dropout_p=0.0 + ) + + +class ImageResidualAdapter(nn.Module): + """Embodiment-specific image residual; starts as an identity operation.""" + + def __init__(self, width, bottleneck): + super().__init__() + self.norm = nn.LayerNorm(width) + self.down = nn.Linear(width, bottleneck) + self.up = nn.Linear(bottleneck, width) + nn.init.zeros_(self.up.weight) + nn.init.zeros_(self.up.bias) + + def forward(self, x): + return x + self.up(F.gelu(self.down(self.norm(x)))) + + +class MoTWorldModel(nn.Module): + def __init__( + self, + pose_dims: Mapping[str, int], + image_dim=768, + width=384, + depth=8, + heads=6, + mlp_ratio=4, + max_cameras=3, + max_tracks=1120, + max_tokens=8192, + time_mode="continuous", + image_adapter_dim=48, + image_prediction_type="v", + attention_query_chunk_size=0, + activation_checkpointing=False, + ): + super().__init__() + if type(activation_checkpointing) is not bool: + raise ValueError("activation_checkpointing must be bool") + self.activation_checkpointing = activation_checkpointing + 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 + if time_mode not in ("continuous", "discrete"): + raise ValueError("Unknown time mode") + if not pose_dims or any(int(d) < 1 for d in pose_dims.values()): + raise ValueError( + "Explicit positive per-embodiment pose dimensions required" + ) + self.pose_dims = {str(k): int(v) for k, v in pose_dims.items()} + if ( + isinstance(image_adapter_dim, bool) + or not isinstance(image_adapter_dim, int) + or image_adapter_dim < 0 + ): + raise ValueError("image_adapter_dim must be a nonnegative integer") + self.image_adapter_dim = image_adapter_dim + self.image_dim, self.max_tokens, self.time_mode = ( + image_dim, + max_tokens, + time_mode, + ) + self.image_in = nn.Linear(image_dim, width) + self.track_in = nn.Linear(3, width) + self.pose_in = nn.ModuleDict( + {k: nn.Linear(d, width) for k, d in self.pose_dims.items()} + ) + self.pose_out = nn.ModuleDict( + {k: nn.Linear(width, d) for k, d in self.pose_dims.items()} + ) + self.image_out = nn.Linear(width, image_dim) + self.track_out = nn.Linear(width, 3) + self.xy = nn.Linear(2, width) + self.camera = nn.Embedding(max_cameras, width) + self.point = nn.Embedding(max_tracks, width) + self.modality = nn.Parameter(torch.randn(3, width) * 0.02) + self.query = nn.Parameter(torch.randn(width) * 0.02) + self.embodiment = nn.Embedding(len(pose_dims), width) + self.embodiment_ids = {k: i for i, k in enumerate(self.pose_dims)} + # Noise level and observed/query status are distinct from physical time. + self.noise_status = nn.Sequential( + nn.Linear(2, width), nn.SiLU(), nn.Linear(width, width) + ) + self.blocks = nn.ModuleList( + MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size) + for _ in range(depth) + ) + # One private image adapter per embodiment per block. Zero disables the + # extension explicitly for adapter-free ablations/historical checkpoints. + self.image_adapters = ( + nn.ModuleList( + nn.ModuleDict( + { + k: ImageResidualAdapter(width, image_adapter_dim) + for k in self.pose_dims + } + ) + for _ in range(depth) + ) + if image_adapter_dim + else nn.ModuleList() + ) + self.norm = nn.ModuleList(nn.LayerNorm(width) for _ in MODALITIES) + + def validate(self, batch, embodiment): + if embodiment not in self.pose_dims: + raise ValueError(f"Unknown embodiment {embodiment!r}") + total = 0 + b = None + for name, dim in zip( + MODALITIES, (self.image_dim, 3, self.pose_dims[embodiment]) + ): + s = batch[name] + value = s["value"] + if value.ndim != 3 or value.shape[-1] != dim: + raise ValueError(f"{name}: invalid value shape") + if b is not None and b != value.shape[0]: + raise ValueError("Mixed batch sizes") + b = value.shape[0] + total += value.shape[1] + for key in ("time", "valid", "observed"): + if s[key].shape != value.shape[:2]: + raise ValueError(f"{name}: {key} shape mismatch") + if s["valid"].dtype != torch.bool or s["observed"].dtype != torch.bool: + raise ValueError("Masks must be boolean") + if "target_valid" in s and ( + s["target_valid"].shape != s["valid"].shape + or s["target_valid"].dtype != torch.bool + ): + raise ValueError("Target validity must be a boolean token mask") + if name != "proprio": + if s["xy"].shape != (*value.shape[:2], 2): + raise ValueError("Reference image coordinates must match tokens") + if not torch.isfinite(s["xy"][s["valid"]]).all(): + raise ValueError("Non-finite reference image coordinates") + for key, limit in (("camera", self.camera.num_embeddings),): + if s[key].shape != s["valid"].shape or s[key].dtype != torch.long: + raise ValueError(f"{key} must contain integer token IDs") + active_ids = s[key][s["valid"]] + if ((active_ids < 0) | (active_ids >= limit)).any(): + raise ValueError(f"{key} ID out of range") + if name == "track": + ids = s["point_id"] + if ids.shape != s["valid"].shape or ids.dtype != torch.long: + raise ValueError("Persistent point IDs must match tokens") + if ( + (ids[s["valid"]] < 0) + | (ids[s["valid"]] >= self.point.num_embeddings) + ).any(): + raise ValueError("Persistent point ID out of range") + if (s["observed"] & ~s["valid"]).any(): + raise ValueError("Padding cannot be observed") + if not torch.isfinite(s["time"][s["valid"]]).all(): + raise ValueError("Non-finite timestamp") + if not torch.isfinite(value[s["observed"]]).all(): + raise ValueError("Non-finite observed value") + if self.time_mode == "discrete": + if "index" not in s or s["index"].shape != s["time"].shape: + raise ValueError("Discrete mode requires explicit grid indices") + if not torch.isfinite(s["index"][s["valid"]]).all() or not torch.equal( + s["index"][s["valid"]], s["index"][s["valid"]].round() + ): + raise ValueError("Discrete indices must be finite integers") + if total > self.max_tokens: + raise ValueError(f"Token budget exceeded: {total} > {self.max_tokens}") + valid = torch.cat([batch[n]["valid"] for n in MODALITIES], 1) + if not valid.any(1).all(): + raise ValueError("Every sample needs at least one valid token") + + def forward(self, batch, embodiment, visual_values, levels): + self.validate(batch, embodiment) + embeds, times, valid = [], [], [] + for i, name in enumerate(MODALITIES): + s = batch[name] + observed = s["observed"] & s["valid"] + if name == "proprio": + # Never embed unknown target poses, even during teacher-forced training. + value = torch.where(observed[..., None], s["value"], 0) + x = self.pose_in[embodiment](value) + x = torch.where(observed[..., None], x, self.query) + level = torch.zeros_like(s["time"]) + else: + value = torch.where(s["valid"][..., None], visual_values[name], 0) + x = (self.image_in if name == "image" else self.track_in)(value) + x = x + self.xy(torch.where(s["valid"][..., None], s["xy"], 0)) + x = x + self.camera(s["camera"].masked_fill(~s["valid"], 0)) + if name == "track": + x = x + self.point(s["point_id"].masked_fill(~s["valid"], 0)) + level = levels[name] + condition = torch.stack((level, observed.to(level.dtype)), -1) + x = x + self.noise_status(condition.to(x.dtype)) + self.modality[i] + x = x + self.embodiment.weight[self.embodiment_ids[embodiment]] + embeds.append(x.masked_fill(~s["valid"][..., None], 0)) + clock = s["time"] if self.time_mode == "continuous" else s["index"] + times.append(clock.masked_fill(~s["valid"], 0)) + valid.append(s["valid"]) + for index, block in enumerate(self.blocks): + if self.activation_checkpointing and torch.is_grad_enabled(): + embeds = checkpoint(block, embeds, times, valid, use_reentrant=False) + else: + embeds = block(embeds, times, valid) + if self.image_adapter_dim: + embeds[0] = self.image_adapters[index][embodiment]( + embeds[0] + ).masked_fill(~valid[0][..., None], 0) + embeds = [norm(x) for norm, x in zip(self.norm, embeds)] + return { + "image": self.image_out(embeds[0]), + "track": self.track_out(embeds[1]), + "proprio": self.pose_out[embodiment](embeds[2]), + } + + @staticmethod + def frame_noise_levels(stream): + # One noise level per timestamp per example; all patches/cameras agree. + result = torch.zeros_like(stream["time"]) + for b in range(result.shape[0]): + _, inverse = torch.unique( + stream["time"][b], sorted=True, return_inverse=True + ) + draws = torch.rand( + int(inverse.max()) + 1, device=result.device, dtype=result.dtype + ).clamp(0.001, 0.999) + result[b] = draws[inverse] + return result.masked_fill(stream["observed"] | ~stream["valid"], 0) + + def _complete_training_tracks(self, batch, embodiment, policy, steps): + """Explicit approximate completion; never invent zero clean geometry. + + Available recorded visuals condition a detached auxiliary sample. Only + missing tracks are retained, without becoming supervised targets. The + main denoising pass receives neither completion nor availability flags. + """ + if policy not in ("error", "conditional_sample"): + raise ValueError("Unknown missing_track_policy") + if type(steps) is not int or steps < 1: + raise ValueError("imputation_steps must be a positive integer") + available = {} + for name in ("image", "track"): + s = batch[name] + confidence = s.get("confidence", torch.ones_like(s["time"])) + if confidence.shape != s["time"].shape or ( + not torch.isfinite(confidence).all() + or (confidence < 0).any() + or (confidence > 1).any() + ): + raise ValueError("Confidence must match tokens and be finite in [0,1]") + known = s["valid"] & s.get("target_valid", s["valid"]) & (confidence > 0) + if (s["observed"] & ~known).any(): + raise ValueError( + "Observed visual values require valid positive-confidence targets" + ) + if not torch.isfinite(s["value"][known]).all(): + raise ValueError("Non-finite visual target") + available[name] = known + if (batch["image"]["valid"] & ~available["image"]).any(): + raise ValueError( + "Missing image targets require explicit preprocessing; unsupported" + ) + track = batch["track"] + missing = track["valid"] & ~available["track"] + if not missing.any(): + # No forward or random draw: retain exact complete-data behavior. + return track["value"], missing + if policy == "error": + raise ValueError( + "Missing track targets: zero-fill is forbidden. Supply valid tracks " + "or explicitly select missing_track_policy='conditional_sample'." + ) + teacher = {n: dict(batch[n]) for n in MODALITIES} + for name in ("image", "track"): + teacher[name]["observed"] = available[name] + teacher[name]["value"] = torch.where( + available[name][..., None], batch[name]["value"], 0 + ) + # Even an auxiliary sampler must never receive withheld recorded poses. + teacher["proprio"]["value"] = torch.where( + batch["proprio"]["observed"][..., None], batch["proprio"]["value"], 0 + ) + with torch.no_grad(): + completion = self.sample(teacher, embodiment, steps=steps)["track"].detach() + if not torch.isfinite(completion[missing]).all(): + raise ValueError("Non-finite conditional track completion") + return torch.where(missing[..., None], completion, track["value"]), missing + + def training_losses( + self, + batch, + embodiment, + pose_weight=1.0, + image_weight=1.0, + track_weight=1.0, + missing_track_policy="error", + imputation_steps=20, + ): + self.validate(batch, embodiment) + track_values, imputed = self._complete_training_tracks( + batch, embodiment, missing_track_policy, imputation_steps + ) + noisy, levels, targets, weights = {}, {}, {}, {} + for name in ("image", "track"): + s = batch[name] + confidence = s.get("confidence", torch.ones_like(s["time"])) + if confidence.shape != s["time"].shape: + raise ValueError("Confidence must match token shape") + if ( + not torch.isfinite(confidence).all() + or (confidence < 0).any() + or (confidence > 1).any() + ): + raise ValueError("Confidence must be finite and between zero and one") + supervised = ( + s.get("target_valid", s["valid"]) & s["valid"] & (confidence > 0) + ) + if not torch.isfinite(s["value"][supervised | s["observed"]]).all(): + raise ValueError("Non-finite visual target") + if name == "track": + value = torch.where(s["valid"][..., None], track_values, 0) + else: + value = torch.where( + (supervised | s["observed"])[..., None], s["value"], 0 + ) + levels[name] = self.frame_noise_levels(s) + alpha, sigma = noise_coefficients(levels[name]) + eps = torch.randn_like(value) + noisy[name] = alpha * value + sigma * eps + targets[name] = ( + value + if name == "image" and self.image_prediction_type == "x0" + else alpha * eps - sigma * value + ) + weights[name] = (supervised & ~s["observed"]).float() * confidence + prediction = self(batch, embodiment, noisy, levels) + losses = { + "image": masked_mse( + prediction["image"], targets["image"], weights["image"] + ), + "track": masked_mse( + prediction["track"], targets["track"], weights["track"] + ), + } + # Same-trajectory x0 estimates. No detach: pose loss reaches both visual heads. + clean = {} + for name in ("image", "track"): + alpha, sigma = noise_coefficients(levels[name]) + x0 = ( + prediction[name] + if name == "image" and self.image_prediction_type == "x0" + else alpha * noisy[name] - sigma * prediction[name] + ) + clean[name] = torch.where( + batch[name]["observed"][..., None], batch[name]["value"], x0 + ) + pose = batch["proprio"] + if pose_weight > 0: + readout = self( + batch, + embodiment, + clean, + {k: torch.zeros_like(v) for k, v in levels.items()}, + ) + pose_mask = ( + pose.get("target_valid", pose["valid"]) + & pose["valid"] + & ~pose["observed"] + ) + losses["pose"] = masked_mse( + readout["proprio"], pose["value"], pose_mask.float() + ) + else: + losses["pose"] = prediction["proprio"].sum() * 0 + losses["total"] = ( + image_weight * losses["image"] + + track_weight * losses["track"] + + pose_weight * losses["pose"] + ) + losses["imputed_track_tokens"] = imputed.sum().to(losses["total"].dtype) + losses["imputation_forwards"] = losses["total"].new_tensor( + imputation_steps + 1 if imputed.any() else 0 + ) + return losses + + def _refinement_state(self, batch, generator): + """Validate explicit clean predictions and build a private sampling batch. + + Optional batch["refinement"][image|track] contains strength[B,N] and + prediction[B,N,C]. Strength is a diffusion noise level, not a linear + blend percentage. No withheld batch value supplies a starting estimate. + """ + controls = batch.get("refinement", {}) + if not isinstance(controls, Mapping) or set(controls) - {"image", "track"}: + raise ValueError("Refinement controls support only image and track") + sampling_batch = {n: dict(batch[n]) for n in MODALITIES} + starts, anchors = {}, {} + for name in ("image", "track"): + s = batch[name] + valid, observed = s["valid"], s["observed"] + control = controls.get(name) + strength = torch.ones_like(s["time"]).masked_fill(observed | ~valid, 0) + anchor = torch.where(observed[..., None], s["value"], 0) + if control is not None: + if ( + not isinstance(control, Mapping) + or "strength" not in control + or set(control) - {"strength", "prediction"} + ): + raise ValueError( + "Refinement requires strength and optional prediction" + ) + supplied = control["strength"] + if ( + not isinstance(supplied, torch.Tensor) + or supplied.shape != valid.shape + or supplied.device != s["value"].device + or not supplied.is_floating_point() + ): + raise ValueError( + "Refinement strength must be a floating token-shaped tensor" + ) + if ( + not torch.isfinite(supplied[valid]).all() + or (supplied[valid] < 0).any() + or (supplied[valid] > 1).any() + ): + raise ValueError("Refinement strength must be finite in [0,1]") + if (supplied[observed] != 0).any(): + raise ValueError("Observed values cannot be revised") + # Match training's one noise level per modality/timestamp. + for b in range(valid.shape[0]): + for time in torch.unique(s["time"][b][valid[b] & ~observed[b]]): + selected = valid[b] & ~observed[b] & (s["time"][b] == time) + if torch.unique(supplied[b][selected]).numel() != 1: + raise ValueError( + "Refinement strength must agree within each frame" + ) + strength = supplied.float().masked_fill(~valid, 0) + needs_prediction = valid & ~observed & (strength < 1) + previous = control.get("prediction") + if previous is not None: + if ( + not isinstance(previous, torch.Tensor) + or previous.shape != s["value"].shape + or previous.device != s["value"].device + or previous.dtype != s["value"].dtype + ): + raise ValueError( + "Refinement prediction must match value shape/device/dtype" + ) + if not torch.isfinite(previous[needs_prediction]).all(): + raise ValueError("Non-finite retained/refined prediction") + anchor = torch.where(needs_prediction[..., None], previous, anchor) + elif needs_prediction.any(): + raise ValueError( + "Retaining or refining requires an explicit prediction" + ) + fixed = valid & (strength == 0) + sampling_batch[name]["observed"] = fixed + # Fixed generated predictions act as clean anchors only in this call. + sampling_batch[name]["value"] = anchor + starts[name], anchors[name] = strength, anchor + current = {} + # Validate all controls before consuming the caller's random generator. + for name in ("image", "track"): + s = batch[name] + noise = torch.randn( + s["value"].shape, + device=s["value"].device, + dtype=s["value"].dtype, + generator=generator, + ) + alpha, sigma = noise_coefficients(starts[name]) + initial = alpha * anchors[name] + sigma * noise + # Explicit endpoints preserve old fresh-noise sampling and anchors exactly. + current[name] = torch.where( + (starts[name] == 1)[..., None], + noise, + torch.where((starts[name] == 0)[..., None], anchors[name], initial), + ).masked_fill(~s["valid"][..., None], 0) + return sampling_batch, current, starts, anchors + + @torch.no_grad() + def sample(self, batch, embodiment, steps=20, generator=None): + """Deterministic DDIM; randomness only in initial Gaussian draws. + + Unknown values/target masks/confidence are never read by the sampler. + Optional batch["refinement"] supplies explicit clean previous predictions + and per-frame start levels. Each level descends start -> 0 together. + Physical query times stay fixed. Omitted controls retain fresh sampling. + """ + self.validate(batch, embodiment) + if steps < 1: + raise ValueError("At least one denoising step is required") + sampling_batch, current, starts, anchors = self._refinement_state( + batch, generator + ) + for j in range(steps, 0, -1): + levels = {n: start * (j / steps) for n, start in starts.items()} + prediction = self(sampling_batch, embodiment, current, levels) + for name in current: + alpha, sigma = noise_coefficients(levels[name]) + next_level = starts[name] * ((j - 1) / steps) + an, sn = noise_coefficients(next_level) + if name == "image" and self.image_prediction_type == "x0": + # Predict structured clean features directly. The bounded + # ratio avoids forming epsilon at tiny positive strengths: + # division first can overflow, followed by zero * infinity. + x0 = prediction[name] + denominator = torch.where( + levels[name][..., None] > 0, sigma, torch.ones_like(sigma) + ) + candidate = an * x0 + (sn / denominator) * ( + current[name] - alpha * x0 + ) + candidate = candidate.masked_fill( + ~sampling_batch[name]["valid"][..., None], 0 + ) + else: + x0 = alpha * current[name] - sigma * prediction[name] + eps = sigma * current[name] + alpha * prediction[name] + candidate = an * x0 + sn * eps + current[name] = torch.where( + sampling_batch[name]["observed"][..., None], + anchors[name], + candidate, + ) + zeros = {n: torch.zeros_like(batch[n]["time"]) for n in current} + pose = self(sampling_batch, embodiment, current, zeros)["proprio"] + current["proprio"] = torch.where( + batch["proprio"]["observed"][..., None], batch["proprio"]["value"], pose + ) + return { + n: x.masked_fill(~batch[n]["valid"][..., None], 0) + for n, x in current.items() + } diff --git a/egomimic/pipeline/stages_world_model.py b/egomimic/pipeline/stages_world_model.py new file mode 100644 index 0000000000..3dd0161479 --- /dev/null +++ b/egomimic/pipeline/stages_world_model.py @@ -0,0 +1,77 @@ +"""Pipeline adapter for the independent image/track/proprio world model.""" + +from egomimic.pipeline.core import Stage, resolve_homogeneous_scalar + + +class WorldModelStage(Stage): + reads_by_mode = { + "train": ("world", "embodiment", "optimizer_step"), + "inference": ("world", "embodiment"), + } + writes_by_mode = {"train": ("loss/world",), "inference": ("world_prediction",)} + + def __init__( + self, + model, + warmup_steps=10000, + pose_ramp_steps=10000, + image_weight=1.0, + track_weight=1.0, + pose_weight=1.0, + sampler_steps=20, + missing_track_policy="error", + imputation_steps=20, + ): + super().__init__() + if ( + min(warmup_steps, pose_ramp_steps) < 0 + or min(image_weight, track_weight, pose_weight) < 0 + ): + raise ValueError("Schedules and weights must be nonnegative") + self.model = model + self.warmup_steps, self.pose_ramp_steps = warmup_steps, pose_ramp_steps + self.image_weight, self.track_weight, self.pose_weight = ( + image_weight, + track_weight, + pose_weight, + ) + self.sampler_steps = sampler_steps + if missing_track_policy not in ("error", "conditional_sample"): + raise ValueError("Unknown missing_track_policy") + if type(imputation_steps) is not int or imputation_steps < 1: + raise ValueError("imputation_steps must be a positive integer") + self.missing_track_policy = missing_track_policy + self.imputation_steps = imputation_steps + + def pose_scale(self, step): + if step < self.warmup_steps: + return 0.0 + if self.pose_ramp_steps == 0: + return self.pose_weight + return self.pose_weight * 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"])) + losses = self.model.training_losses( + batch["world"], + embodiment, + scale, + self.image_weight, + self.track_weight, + missing_track_policy=self.missing_track_policy, + imputation_steps=self.imputation_steps, + ) + batch["loss/world"] = losses["total"] + batch.update({f"log/world_{k}": v.detach() for k, v in losses.items()}) + batch["log/pose_weight"] = losses["total"].new_tensor(scale) + else: + batch["world_prediction"] = self.model.sample( + batch["world"], embodiment, steps=self.sampler_steps + ) + return batch diff --git a/egomimic/pl_utils/world_model_behavior.py b/egomimic/pl_utils/world_model_behavior.py new file mode 100644 index 0000000000..d1c36f783e --- /dev/null +++ b/egomimic/pl_utils/world_model_behavior.py @@ -0,0 +1,50 @@ +"""Use the shared Lightning wrapper; schedule follows actual optimizer steps.""" + +import random + +import numpy as np +import torch + +from egomimic.pl_utils.training_behavior import TrainingBehavior + + +class WorldModelTrainingBehavior(TrainingBehavior): + def on_validation_epoch_end(self): + if self.context.evaluator is not None: + self.context.evaluator.log_epoch_metrics() + + 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"]) + self._pending_rng = None + data = self.context.trainer.datamodule + if getattr(data, "checkpointable_sampling", False): + cursors = {values["sampling_cursor"] for values in batch.values()} + if len(cursors) != 1: + raise ValueError("Expected one consumed MoT batch cursor") + data.acknowledge_batch(cursors.pop()) + batch = { + name: {**values, "optimizer_step": int(self.context.global_step)} + for name, values in batch.items() + } + return self.context._default_training_step(batch, batch_idx) + + def on_save_checkpoint(self, checkpoint): + 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(), + } + + 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): + raise ValueError("Exact MoT resume requires saved model RNG state") diff --git a/egomimic/rldb/world_model.py b/egomimic/rldb/world_model.py new file mode 100644 index 0000000000..22bb10c4c8 --- /dev/null +++ b/egomimic/rldb/world_model.py @@ -0,0 +1,559 @@ +"""Prepared world-model data contract and temporal sampling. + +Owns no simulator or tracker: upstream extraction must supply real tracks and +fixed-frame measured poses. Never reinterpret policy commands as pose targets. +""" + +from __future__ import annotations + +import json +import re +from contextlib import contextmanager +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from egomimic.rldb.world_model_preparation import ( + ArtifactVerificationMemo, + load_normalization, +) + + +@contextmanager +def open_prepared_episode(path): + """Keep historical NPZ support; HDF5 slices only selected frame chunks.""" + path = Path(path) + if path.suffix in (".h5", ".hdf5"): + import h5py + + with h5py.File(path, "r") as data: + yield data + elif path.suffix == ".npz": + with np.load(path, allow_pickle=False) as data: + yield data + else: + raise ValueError("Prepared episodes must use NPZ or HDF5") + + +def points_to_reference(points, reference_from_source): + """Transform [..., P, 3] points with calibrated [..., 4, 4] rigid transforms.""" + if points.shape[-1] != 3 or reference_from_source.shape[-2:] != (4, 4): + raise ValueError("Expected XYZ and homogeneous rigid transforms") + if not np.isfinite(reference_from_source).all(): + raise ValueError("Non-finite calibration") + rotation = reference_from_source[..., :3, :3] + if not np.allclose(rotation @ np.swapaxes(rotation, -1, -2), np.eye(3), atol=1e-4): + raise ValueError("Calibration rotation is not orthonormal") + if not np.allclose(np.linalg.det(rotation), 1, atol=1e-4): + raise ValueError("Calibration must preserve handedness") + if not np.allclose(reference_from_source[..., 3, :], [0, 0, 0, 1], atol=1e-5): + raise ValueError("Invalid homogeneous transform") + return ( + points @ np.swapaxes(rotation, -1, -2) + reference_from_source[..., None, :3, 3] + ) + + +def _bounded_time_sampling(max_past_seconds, max_future_seconds): + if max_past_seconds is None and max_future_seconds is None: + return False + if any( + isinstance(value, (bool, np.bool_)) + or not isinstance(value, (int, float, np.integer, np.floating)) + or not np.isfinite(value) + or value < 0 + for value in (max_past_seconds, max_future_seconds) + ) or max_past_seconds + max_future_seconds <= 0: + raise ValueError("Time bounds require two finite nonnegative seconds values with positive span") + return True + + +def _physical_times(timestamps, length): + times = np.asarray(timestamps, dtype=np.float64) + if ( + times.shape != (length,) + or not np.isfinite(times).all() + or not (np.diff(times) > 0).all() + ): + raise ValueError("Physical timestamps must match episode length and strictly increase") + return times - times[0] + + +def sample_times( + length, + rng, + mode="continuous", + min_frames=2, + max_frames=4, + discrete_frames=4, + stride=1, + anchor=None, + timestamps=None, + max_past_seconds=None, + max_future_seconds=None, +): + """Select recorded times; optional bounds are relative to the anchor time. + + Bounded continuous mode draws uniformly without replacement among recorded + timestamps inside [anchor-past, anchor+future]. This is not a uniform draw + over continuous seconds. Unbounded calls retain the historical index-based + multiscale sampler and its RNG behavior for checkpoint-bound callers. + """ + if ( + mode not in ("continuous", "discrete") + or min_frames < 2 + or max_frames < min_frames + or stride < 1 + ): + raise ValueError("Invalid time sampler") + bounded = _bounded_time_sampling(max_past_seconds, max_future_seconds) + if bounded and mode != "continuous": + raise ValueError("Physical-time bounds apply only to continuous sampling") + if mode == "discrete": + available = length - (discrete_frames - 1) * stride + if discrete_frames < 2 or available <= 0: + raise ValueError("Episode too short for complete discrete window") + start = int(rng.integers(available)) if anchor is None else int(anchor) + if not 0 <= start < available: + raise ValueError("Invalid discrete anchor") + return start + np.arange(discrete_frames) * stride + if length < min_frames: + raise ValueError("Episode too short") + if bounded: + times = _physical_times(timestamps, length) + start = int(rng.integers(length)) if anchor is None else int(anchor) + if not 0 <= start < length: + raise ValueError("Invalid anchor") + lo = np.searchsorted(times, times[start] - max_past_seconds, side="left") + hi = np.searchsorted(times, times[start] + max_future_seconds, side="right") + eligible = np.arange(lo, hi, dtype=np.int64) + if len(eligible) < min_frames: + raise ValueError("Too few recorded timestamps inside anchor time bounds") + count = int(rng.integers(min_frames, min(max_frames, len(eligible)) + 1)) + others = rng.choice(eligible[eligible != start], size=count - 1, replace=False) + return np.sort(np.r_[start, others]).astype(np.int64) + count = int(rng.integers(min_frames, min(max_frames, length) + 1)) + start = int(rng.integers(length)) if anchor is None else int(anchor) + if not 0 <= start < length: + raise ValueError("Invalid anchor") + selected = {start} + # Cycle nearby / intermediate / episode-wide support; randomized first scale. + scales = rng.permutation([1, max(2, length // 8), length]) + for i in range(count - 1): + radius = int(scales[i % 3]) + candidates = np.array( + [ + j + for j in range(max(0, start - radius), min(length, start + radius + 1)) + if j not in selected + ] + ) + if not len(candidates): + candidates = np.array([j for j in range(length) if j not in selected]) + selected.add(int(rng.choice(candidates))) + return np.array(sorted(selected), dtype=np.int64) + + +def collate_world(samples): + """Homogeneous embodiment batches; pad streams independently without leaking masks.""" + embodiments = {s["embodiment"] for s in samples} + if len(embodiments) != 1: + raise ValueError("Batch must be homogeneous in embodiment") + embodiment = samples[0]["embodiment"] + result = {} + for name in ("image", "track", "proprio"): + size = max(s["world"][name]["value"].shape[0] for s in samples) + stream = {} + for key in samples[0]["world"][name]: + arrays = [] + for sample in samples: + value = sample["world"][name][key] + target = torch.zeros((size, *value.shape[1:]), dtype=value.dtype) + target[: len(value)] = value + arrays.append(target) + stream[key] = torch.stack(arrays) + result[name] = stream + item = {"embodiment": embodiment, "world": result} + if "sampling_cursor" in samples[0]: + cursors = {sample["sampling_cursor"] for sample in samples} + if len(cursors) != 1: + raise ValueError("A resumed batch must have one source-draw cursor") + item["sampling_cursor"] = cursors.pop() + return {embodiment: item} + + +class PreparedWorldDataset(Dataset): + """Read per-episode NPZ or sliced HDF5 caches with SHA-bound manifests. + + A flat global anchor index samples proportional to eligible frames/windows, + NOT equally per embodiment. Use HomogeneousWorldBatchSampler for batches. + """ + + def __init__( + self, + manifest_path, + split="train", + 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, + ): + self.path = Path(manifest_path) + self.manifest = json.loads(self.path.read_text()) + m = self.manifest + if m.get("schema") != "mot-world-prepared-v1": + raise ValueError("Unsupported prepared-data schema") + for key in ( + "dino_identity", + "normalization", + "pose_schemas", + "track_provenance", + ): + if not m.get(key): + raise ValueError(f"Missing {key}") + for key in ("model", "revision", "image_size", "patch_grid", "feature_layer"): + if not m["dino_identity"].get(key): + raise ValueError(f"Missing DINO identity field {key}") + for key in ("image_size", "patch_grid"): + shape = m["dino_identity"][key] + if ( + not isinstance(shape, list) + or len(shape) != 2 + or any(type(x) is not int or x < 1 for x in shape) + ): + raise ValueError(f"DINO {key} must contain two positive integers") + if ( + m["normalization"].get("scope") != "train_only" + or not re.fullmatch(r"[0-9a-f]{64}", m["normalization"].get("sha256", "")) + or type(m["normalization"].get("arrays_normalized")) is not bool + ): + raise ValueError("Require train-only normalization provenance") + for schema in m["pose_schemas"].values(): + if ( + schema.get("dimension", 0) < 1 + or len(schema.get("fields", [])) != schema["dimension"] + ): + raise ValueError( + "Pose schemas require explicit field names for every coordinate" + ) + for key in ("method", "reference_frame", "calibration_sha256"): + if not m["track_provenance"].get(key): + raise ValueError(f"Missing track provenance {key}") + if m["track_provenance"].get("interpolated_endpoints", True): + raise ValueError("Endpoint interpolation is not actual tracked supervision") + verifier = verification_memo if verification_memo is not None else ArtifactVerificationMemo() + self.normalization = load_normalization(m, self.path.parent, verifier) + ids, paths = set(), set() + for ep in m["episodes"]: + identity = (ep["embodiment"], ep["episode_id"]) + physical = (self.path.parent / ep["path"]).resolve() + if identity in ids or physical in paths: + raise ValueError( + "Duplicate episode identity or physical path across splits" + ) + ids.add(identity) + paths.add(physical) + if ( + ep.get("frame") != "fixed_reference" + or ep.get("pose_kind") != "measured_pose" + ): + raise ValueError( + "Require measured poses and tracks in a fixed reference frame" + ) + if ep["split"] not in ("train", "val"): + raise ValueError("Explicit train/val episode assignment required") + verifier.verify(physical, ep["sha256"], + trusted_fingerprint=ep.get("verified_fingerprint")) + self.episodes = [ep for ep in m["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 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 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") + # An isolated frame cannot yield min_frames without leaving the + # requested time range. Exclude that anchor instead of escaping. + 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)) + 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 = None + self.worker_seed = None + + 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) + # Training RNG advances on repeated visits; validation choices stay reproducible. + 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])) + episode_index = int(np.searchsorted(self.ends, index, side="right")) + ep = self.episodes[episode_index] + offset = 0 if episode_index == 0 else int(self.ends[episode_index - 1]) + anchor = int(self.anchor_indices[episode_index][index - offset]) + with open_prepared_episode(self.path.parent / ep["path"]) as data: + required = ( + "time", + "image", + "image_xy", + "track", + "track_xy", + "track_camera", + "track_confidence", + "pose", + "pose_valid", + ) + if any(k not in data for k in required): + raise ValueError("Prepared cache missing required arrays") + full_time = data["time"][:] + if ( + len(full_time) != ep["frames"] + or not np.isfinite(full_time).all() + or not (np.diff(full_time) > 0).all() + ): + raise ValueError( + "Timestamps must be finite, strictly increasing physical seconds" + ) + indices = sample_times( + ep["frames"], rng, anchor=anchor, timestamps=full_time, **self.sampling + ) + times = full_time[indices].astype(np.float64) + # Subtract origin before float32 conversion (important for epoch timestamps). + times = (times - times[int(rng.integers(len(times)))]).astype(np.float32) + image, track, pose = ( + data[k][indices].astype(np.float32) for k in ("image", "track", "pose") + ) + if self.normalization is not None: + values = [] + for name, value in (("image", image), ("track", track), (f"pose/{ep['embodiment']}", pose)): + mean, scale = self.normalization[name] + if value.shape[-1] != len(mean): + raise ValueError(f"Normalization schema width mismatch: {name}") + values.append((value - mean) / scale) + image, track, pose = values + if ( + image.ndim != 4 + or track.ndim != 3 + or track.shape[-1] != 3 + or pose.ndim != 2 + ): + raise ValueError("Expected image[T,V,P,D], track[T,Q,3], pose[T,Dm]") + if ( + pose.shape[-1] + != self.manifest["pose_schemas"][ep["embodiment"]]["dimension"] + ): + raise ValueError("Pose schema width mismatch") + f, v, p, _ = image.shape + grid_h, grid_w = self.manifest["dino_identity"]["patch_grid"] + if p != grid_h * grid_w: + raise ValueError("DINO patch grid differs from cached image tokens") + q = track.shape[1] + if ( + data["image_xy"].shape != (p, 2) + or data["track_xy"].shape != (q, 2) + or data["track_camera"].shape != (q,) + or data["track_confidence"].shape != (ep["frames"], q) + or data["pose_valid"].shape != (ep["frames"],) + ): + raise ValueError( + "Prepared geometry/mask shapes differ from token arrays" + ) + if not np.isfinite(image).all(): + raise ValueError("Non-finite DINO features") + if rng.random() < self.forecasting_probability: + cutoff = int(rng.integers(1, f)) + known = np.arange(f) < cutoff + observed = {k: known.copy() for k in ("image", "track", "proprio")} + # Independently omit observed proprio/track frames; never reveal future ones. + for k in ("track", "proprio"): + observed[k] &= rng.random(f) < 0.8 + else: + observed = { + k: rng.random(f) < 0.5 for k in ("image", "track", "proprio") + } + # Ensure context and a target without privileging either edge: + # random conditioning must include backcasting/interpolation. + if not observed["image"].any(): + observed["image"][rng.integers(f)] = True + if observed["image"].all(): + observed["image"][rng.integers(f)] = False + + def base(values, repeats, name): + n = f * repeats + return { + "value": values.reshape(n, -1), + "time": np.repeat(times, 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[name], repeats), + } + + streams = { + "image": base(image, v * p, "image"), + "track": base(track, q, "track"), + "proprio": base(pose, 1, "proprio"), + } + streams["image"].update( + xy=np.tile(data["image_xy"][:].astype(np.float32), (f * v, 1)), + camera=np.tile(np.repeat(np.arange(v), p), f), + ) + confidence = ( + data["track_confidence"][indices].reshape(-1).astype(np.float32) + ) + streams["track"].update( + xy=np.tile(data["track_xy"][:].astype(np.float32), (f, 1)), + camera=np.tile(data["track_camera"][:].astype(np.int64), f), + point_id=np.tile(np.arange(q), f), + confidence=confidence, + ) + streams["track"]["target_valid"] = (confidence > 0) & np.isfinite( + track.reshape(-1, 3) + ).all(-1) + streams["track"]["observed"] &= streams["track"]["target_valid"] + streams["proprio"]["target_valid"] = data["pose_valid"][indices].astype( + bool + ) + streams["proprio"]["observed"] &= streams["proprio"]["target_valid"] + sample = { + "embodiment": ep["embodiment"], + "world": { + name: { + key: torch.from_numpy(np.array(value, copy=True)) + for key, value in s.items() + } + for name, s in streams.items() + }, + } + if draw is not None: + sample["sampling_cursor"] = draw[0] + 1 + return sample + + +class HomogeneousWorldBatchSampler: + """Choose an embodiment proportional to eligible anchors, then draw its batch.""" + + def __init__(self, dataset, batch_size=1, seed=42, shuffle=True, checkpointable=False, + validation_batches_per_embodiment=None): + if batch_size < 1: + raise ValueError("Positive batch size required") + if validation_batches_per_embodiment is not None: + if 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") + if shuffle or getattr(dataset, "split", None) != "val": + raise ValueError("Per-embodiment batch caps require unshuffled validation data") + self.validation_batches_per_embodiment = validation_batches_per_embodiment + self.dataset, self.batch_size, self.seed, self.shuffle = ( + dataset, + batch_size, + seed, + shuffle, + ) + self.epoch = 0 + self.checkpointable = checkpointable + self.consumed_batches = 0 + self.ranges = {} + start = 0 + for ep, count in zip(dataset.episodes, dataset.counts): + self.ranges.setdefault(ep["embodiment"], []).append( + np.arange(start, start + count) + ) + start += count + self.groups = {k: np.concatenate(v) for k, v in self.ranges.items()} + + def set_epoch(self, epoch): + self.epoch = int(epoch) + + def __len__(self): + if self.shuffle: + return (len(self.dataset) + self.batch_size - 1) // self.batch_size + return sum( + min((len(v) + self.batch_size - 1) // self.batch_size, + self.validation_batches_per_embodiment) + if self.validation_batches_per_embodiment is not None + else (len(v) + self.batch_size - 1) // self.batch_size + for v in self.groups.values() + ) + + def __iter__(self): + if not self.shuffle: + for group in self.groups.values(): + end = len(group) + if self.validation_batches_per_embodiment is not None: + end = min(end, self.validation_batches_per_embodiment * self.batch_size) + for i in range(0, end, self.batch_size): + yield group[i : i + self.batch_size].tolist() + return + if self.checkpointable: + groups = list(self.groups.values()) + probabilities = np.array([len(v) for v in groups], float) + probabilities /= probabilities.sum() + start = self.consumed_batches + end = (start // len(self) + 1) * len(self) + for cursor in range(start, end): + rng = np.random.default_rng(np.random.SeedSequence([self.seed, cursor])) + group = groups[int(rng.choice(len(groups), p=probabilities))] + chosen = rng.choice(group, self.batch_size, replace=True) + yield [(int(index), cursor, slot) for slot, index in enumerate(chosen)] + return + rng = np.random.default_rng(self.seed + self.epoch) + groups = list(self.groups.values()) + probabilities = np.array([len(v) for v in groups], float) + probabilities /= probabilities.sum() + for _ in range(len(self)): + group = groups[int(rng.choice(len(groups), p=probabilities))] + yield rng.choice(group, self.batch_size, replace=True).tolist() + self.epoch += 1 diff --git a/egomimic/rldb/world_model_datamodule.py b/egomimic/rldb/world_model_datamodule.py new file mode 100644 index 0000000000..e0a36065e5 --- /dev/null +++ b/egomimic/rldb/world_model_datamodule.py @@ -0,0 +1,193 @@ +"""Prepared world-model caches through the existing trainHydra data boundary.""" + +import hashlib +from pathlib import Path + +from lightning import LightningDataModule +from torch.utils.data import DataLoader + +from egomimic.pl_utils.data_context import DataContext +from egomimic.rldb.world_model import ( + HomogeneousWorldBatchSampler, + PreparedWorldDataset, + collate_world, +) +from egomimic.rldb.world_model_preparation import ArtifactVerificationMemo + + +class WorldModelDataModule(LightningDataModule): + def __init__( + self, + manifest_path, + 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, + ): + super().__init__() + self.manifest_path = manifest_path + if validation_batches_per_embodiment is not None and ( + 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") + self.validation_batches_per_embodiment = validation_batches_per_embodiment + self.batch_size, self.num_workers = batch_size, num_workers + self.settings = dict( + mode=mode, + seed=seed, + min_frames=min_frames, + max_frames=max_frames, + discrete_frames=discrete_frames, + stride=stride, + forecasting_probability=forecasting_probability, + ) + # Preserve historical unbounded checkpoint data contexts. Configured + # bounds become part of strict context identity for new bounded runs. + if max_past_seconds is not None or max_future_seconds is not None: + self.settings.update( + max_past_seconds=max_past_seconds, + max_future_seconds=max_future_seconds, + ) + 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") + self.train_sampler = None + self.consumed_batches = 0 + + 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) + for s 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( + "World-model caches already carry train-only normalization" + ) + self.setup() + dataset = self.datasets["train"] + state = { + "kind": "mot-world-prepared-v1", + "manifest": dataset.manifest, + "manifest_sha256": hashlib.sha256(dataset.path.read_bytes()).hexdigest(), + "sampling": self.settings, + } + 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 restored_state is not None and restored_state != state: + raise ValueError("World-model checkpoint/data context mismatch") + return DataContext( + normalizer=None, + sample_schema={"world": "three_token_streams"}, + validation_groups=tuple( + sorted({e["embodiment"] for e in dataset.episodes}) + ), + state=state, + ) + + def configure_evaluation(self, requirements): + if requirements.complete_episodes or set(requirements.required_keys) - { + "world", + "embodiment", + }: + raise ValueError("Unsupported world-model evaluation requirements") + + def frame_counts(self): + return [ + (s, ep["embodiment"], ep["frames"]) + for s, dataset in self.datasets.items() + for ep in dataset.episodes + ] + + 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" + ) + 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 + ), + ) + if split == "train": + self.train_sampler = sampler + sampler.consumed_batches = self.consumed_batches + # Iterator construction must not advance the model's global Torch RNG. + generator = None + if self.checkpointable_sampling: + import torch + generator = torch.Generator().manual_seed(self.settings["seed"]) + return DataLoader( + dataset, + batch_sampler=sampler, + collate_fn=collate_world, + num_workers=self.num_workers, + generator=generator, + ) + + def acknowledge_batch(self, cursor): + if not self.checkpointable_sampling: + return + if cursor != self.consumed_batches + 1: + raise ValueError("Nonconsecutive world-model source draw") + self.consumed_batches = cursor + if self.train_sampler is not None: + self.train_sampler.consumed_batches = 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} + + 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") + return + expected = self.state_dict() + for key in ("schema", "manifest_sha256", "settings", "batch_size"): + if state_dict.get(key) != expected[key]: + raise ValueError(f"World-model sampling restore mismatch: {key}") + cursor = state_dict.get("consumed_batches") + if type(cursor) is not int or cursor < 0: + raise ValueError("Invalid consumed source-draw cursor") + self.consumed_batches = cursor + if self.train_sampler is not None: + self.train_sampler.consumed_batches = cursor + + def train_dataloader(self): + return self._loader("train") + + def val_dataloader(self): + return self._loader("val") diff --git a/egomimic/rldb/world_model_preparation.py b/egomimic/rldb/world_model_preparation.py new file mode 100644 index 0000000000..d91f3b041c --- /dev/null +++ b/egomimic/rldb/world_model_preparation.py @@ -0,0 +1,305 @@ +"""Finalize immutable feature receipts into train-only prepared world-model data. + +Raw feature caches stay unchanged. Statistics are fitted by streaming training +frames, then applied to selected frame slices by PreparedWorldDataset. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import time +from pathlib import Path + +import numpy as np + +from egomimic.rldb.zarr.episode_split import split_dataset_names + + +class ArtifactVerificationMemo: + """Reuse exact verified files, including fingerprints in a trusted manifest. + + Persistent fingerprints are created only after a successful content hash. + They are a fast path for immutable files on the same filesystem, not an + independent authenticity proof: the prepared manifest must remain trusted. + Copies, relocations, metadata changes, and old manifests fall back to SHA. + """ + + _STAT_FIELDS = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + _FINGERPRINT_SCHEMA = "mot-verified-artifact-stat-v1" + _MINIMUM_CTIME_AGE_NS = 1_000_000_000 + + def __init__(self): + self.verified = {} + self.fingerprints = {} + + @staticmethod + def _now_ns(): + return time.time_ns() + + @staticmethod + def identity(path): + stat = path.stat() + return (stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns) + + @classmethod + def _fingerprint(cls, path, expected, identity, verification_started_ns): + return {"schema": cls._FINGERPRINT_SCHEMA, "path": str(path), + "sha256": expected, "verification_started_ns": verification_started_ns, + **dict(zip(cls._STAT_FIELDS, identity))} + + @classmethod + def _matches_fingerprint(cls, fingerprint, path, expected, identity): + if not isinstance(fingerprint, dict): + return False + # Reject incomplete/extra fields and bool-as-int metadata. A malformed + # cache hint must never silently turn off content verification. + if any(type(fingerprint.get(key)) is not int for key in cls._STAT_FIELDS): + return False + started = fingerprint.get("verification_started_ns") + if (type(started) is not int or started > cls._now_ns() + or started - identity[-1] < cls._MINIMUM_CTIME_AGE_NS): + return False + return fingerprint == cls._fingerprint(path, expected, identity, started) + + def verify(self, path, expected, *, trusted_fingerprint=None): + path = Path(path).resolve(strict=True) + before = self.identity(path) + started = self._now_ns() + if self.verified.get((path, expected)) == before: + return self.fingerprints[(path, expected)].copy() + if self._matches_fingerprint(trusted_fingerprint, path, expected, before): + self.verified[(path, expected)] = before + self.fingerprints[(path, expected)] = trusted_fingerprint.copy() + return trusted_fingerprint.copy() + with path.open("rb") as stream: + actual = hashlib.file_digest(stream, "sha256").hexdigest() + if before != self.identity(path): + raise ValueError(f"Artifact changed while hashing: {path}") + if actual != expected: + raise ValueError(f"Artifact hash mismatch: {path}") + # A same-clock-tick rewrite can preserve ctime as well as restored mtime. + # Only author reusable proof when ctime was already old BEFORE hashing. + # Merely waiting after a recent hash must never turn it into a proof. + if started - before[-1] < self._MINIMUM_CTIME_AGE_NS: + self.verified.pop((path, expected), None) + self.fingerprints.pop((path, expected), None) + return None + fingerprint = self._fingerprint(path, expected, before, started) + self.verified[(path, expected)] = before + self.fingerprints[(path, expected)] = fingerprint + return fingerprint.copy() + + +def training_identities(episodes): + return sorted( + [ep["embodiment"], ep["episode_id"], ep["sha256"]] + for ep in episodes if ep["split"] == "train" + ) + + +def load_normalization(manifest, root, verifier): + contract = manifest["normalization"] + if contract.get("arrays_normalized") is True and not contract.get("path"): + return None # Historical pre-normalized cache contract. + if not contract.get("path"): + raise ValueError("Raw caches require a bound normalization statistics file") + path = root / contract["path"] + verifier.verify(path, contract["sha256"]) + stats = json.loads(path.read_text()) + if (stats.get("schema") != "mot-world-normalization-v1" + or stats.get("scope") != "train_only" + or stats.get("train_episodes") != training_identities(manifest["episodes"])): + raise ValueError("Normalization statistics do not match training episode identities") + keys = {"image", "track"} | {f"pose/{name}" for name in manifest["pose_schemas"]} + if set(stats.get("statistics", {})) != keys: + raise ValueError("Normalization statistics have incomplete modality/pose schemas") + arrays = {} + for key, item in stats["statistics"].items(): + mean, scale = np.asarray(item["mean"], np.float32), np.asarray(item["std"], np.float32) + if (mean.ndim != 1 or not len(mean) or mean.shape != scale.shape + or not np.isfinite(mean).all() or not np.isfinite(scale).all() + or np.any(scale <= 0) or item.get("count", 0) < 1): + raise ValueError(f"Invalid normalization statistics: {key}") + expected = (3 if key == "track" else + manifest["pose_schemas"][key[5:]]["dimension"] if key.startswith("pose/") else None) + if expected is not None and len(mean) != expected: + raise ValueError(f"Normalization schema width mismatch: {key}") + arrays[key] = (mean, scale) + return arrays if contract["arrays_normalized"] is False else None + + +class StreamingMoments: + def __init__(self): + self.count, self.mean, self.m2 = 0, None, None + + def update(self, values): + values = np.asarray(values, np.float64) + values = values.reshape(-1, values.shape[-1]) + if not len(values): + return + if not np.isfinite(values).all(): + raise ValueError("Non-finite supervised values while fitting normalization") + mean = values.mean(0) + m2 = ((values - mean) ** 2).sum(0) + self.merge({"count": len(values), "mean": mean, "m2": m2}) + + def merge(self, summary): + count = summary["count"] + mean, m2 = np.asarray(summary["mean"], np.float64), np.asarray(summary["m2"], np.float64) + if (type(count) is not int or count < 1 or mean.ndim != 1 or not len(mean) + or mean.shape != m2.shape or not np.isfinite(mean).all() + or not np.isfinite(m2).all() or np.any(m2 < 0)): + raise ValueError("Invalid per-episode normalization moments") + if self.count == 0: + self.count, self.mean, self.m2 = count, mean.copy(), m2.copy() + return + if self.mean.shape != mean.shape: + raise ValueError("Inconsistent feature dimensions in normalization corpus") + total = self.count + count + delta = mean - self.mean + self.m2 += m2 + delta ** 2 * (self.count * count / total) + self.mean += delta * (count / total) + self.count = total + + def result(self, floor): + if not self.count: + raise ValueError("No supervised training values for normalization") + return {"count": self.count, "mean": self.mean.tolist(), + "std": np.maximum(np.sqrt(self.m2 / self.count), floor).tolist()} + + +def _write_immutable(path, value): + payload = (json.dumps(value, indent=2, sort_keys=True) + "\n").encode() + if path.exists(): + if path.read_bytes() != payload: + raise ValueError(f"Refusing to replace a different finalized artifact: {path}") + else: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + f".{os.getpid()}.tmp") + temporary.write_bytes(payload) + temporary.replace(path) + return hashlib.sha256(payload).hexdigest() + + +def finalize(inventory_path, output, *, valid_ratio=0.01, seed=42, std_floor=1e-5, + verification_memo=None): + """Finalize an explicit inventory, optionally sharing verified immutable files. + + Reusing a memo across overlapping corpora avoids a second content scan; + each corpus still fits its own train-only statistics and episode split. + The verifier invalidates changed file identities or expected digests. + """ + from egomimic.rldb.world_model import open_prepared_episode + + inventory_path, output = Path(inventory_path).resolve(), Path(output).resolve() + inventory = json.loads(inventory_path.read_text()) + if not 0 < valid_ratio < 1 or not np.isfinite(std_floor) or std_floor <= 0: + raise ValueError("Require a nonempty train/val split and positive finite std floor") + episodes, identities, paths, dino, schemas = [], set(), set(), None, inventory["pose_schemas"] + source_episodes = set() + episode_moments = {} + verifier = verification_memo if verification_memo is not None else ArtifactVerificationMemo() + for receipt_path in inventory["receipts"]: + receipt_path = (inventory_path.parent / receipt_path).resolve() + receipt = json.loads(receipt_path.read_text()) + if (receipt.get("schema") != "mot-world-feature-cache-v1" + or receipt.get("status") != "PASS" or receipt.get("arrays_normalized") is not False): + raise ValueError(f"Feature cache is not a verified raw episode: {receipt_path}") + # The producer places its receipt beside , named .json. + artifact = Path(str(receipt_path)[:-5]).resolve() + identity = (receipt["embodiment"], receipt["episode_id"]) + if identity in identities or artifact in paths: + raise ValueError("Duplicate episode identity or physical path in corpus") + identities.add(identity) + paths.add(artifact) + native = receipt["native_source"] + native_identity = (native["source_demo_sha256"], native["episode"]) + if not all(native_identity) or native_identity in source_episodes: + raise ValueError("Duplicate or missing physical source episode identity") + source_episodes.add(native_identity) + if "moments" in receipt: + if set(receipt["moments"]) != {"image", "track", "pose"}: + raise ValueError("Incomplete per-episode normalization moments") + episode_moments[identity] = receipt["moments"] + schema = schemas[identity[0]] + if schema["fields"] != receipt["pose_fields"] or schema["dimension"] != len(receipt["pose_fields"]): + raise ValueError("Measured pose receipt disagrees with configured schema") + if dino is not None and dino != receipt["dino_identity"]: + raise ValueError("Mixed DINO extraction identities in corpus") + dino = receipt["dino_identity"] + fingerprint = verifier.verify(artifact, receipt["output_sha256"]) + episodes.append(dict(embodiment=identity[0], episode_id=identity[1], + path=os.path.relpath(artifact, output.parent), sha256=receipt["output_sha256"], + frames=receipt["frames"], frame="fixed_reference", pose_kind="measured_pose", + verified_fingerprint=fingerprint, + native_source=native, + receipt_sha256=hashlib.sha256(receipt_path.read_bytes()).hexdigest())) + if not episodes or set(schemas) != {e["embodiment"] for e in episodes}: + raise ValueError("Corpus and pose schema embodiment sets must match and be nonempty") + split_evidence = {} + for name in sorted(schemas): + ids = {e["episode_id"] for e in episodes if e["embodiment"] == name} + train, val = split_dataset_names(ids, valid_ratio=valid_ratio, seed=seed) + if not train or not val or train & val or train | val != ids: + raise ValueError(f"Incomplete episode split: {name}") + split_evidence[name] = {"train": sorted(train), "val": sorted(val)} + for ep in episodes: + if ep["embodiment"] == name: + ep["split"] = "train" if ep["episode_id"] in train else "val" + episodes.sort(key=lambda e: (e["embodiment"], e["episode_id"])) + moments = {key: StreamingMoments() for key in ["image", "track", *(f"pose/{n}" for n in schemas)]} + for ep in episodes: + if ep["split"] != "train": + continue + summaries = episode_moments.get((ep["embodiment"], ep["episode_id"])) + if summaries is not None: + for key in ("image", "track", "pose"): + destination = f"pose/{ep['embodiment']}" if key == "pose" else key + moments[destination].merge(summaries[key]) + continue + with open_prepared_episode(output.parent / ep["path"]) as cache: + for frame in range(ep["frames"]): + moments["image"].update(cache["image"][frame]) + track = cache["track"][frame] + confidence = cache["track_confidence"][frame] + if not np.isfinite(confidence).all() or np.any((confidence < 0) | (confidence > 1)): + raise ValueError("Track confidence must lie in [0,1]") + moments["track"].update(track[confidence > 0]) + if cache["pose_valid"][frame]: + moments[f"pose/{ep['embodiment']}"].update(cache["pose"][frame]) + stats = {"schema": "mot-world-normalization-v1", "scope": "train_only", + "method": "coordinate_population_mean_std", "std_floor": std_floor, + "train_episodes": training_identities(episodes), + "statistics": {key: value.result(std_floor) for key, value in moments.items()}} + stats_path = output.with_name(output.stem + "-normalization.json") + stats_sha = _write_immutable(stats_path, stats) + manifest = {"schema": "mot-world-prepared-v1", "dino_identity": dino, + "pose_schemas": schemas, "track_provenance": inventory["track_provenance"], + "episodes": episodes, + "normalization": {"scope": "train_only", "sha256": stats_sha, + "path": stats_path.name, "arrays_normalized": False}, + "split_provenance": {"seed": seed, "validation_ratio": valid_ratio, + "episode_ids": split_evidence, "id_overlap": 0, "physical_path_overlap": 0, + "complete_corpus_coverage": True}, + "inventory_sha256": hashlib.sha256(inventory_path.read_bytes()).hexdigest()} + _write_immutable(output, manifest) + return manifest + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--inventory", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--valid-ratio", type=float, default=0.01) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--std-floor", type=float, default=1e-5) + args = parser.parse_args() + result = finalize(args.inventory, args.output, valid_ratio=args.valid_ratio, seed=args.seed, std_floor=args.std_floor) + print(json.dumps({"status": "PASS", "manifest": args.output, "episodes": len(result["episodes"])})) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/dinov3_correspondence.py b/egomimic/scripts/embedding_process/dinov3_correspondence.py new file mode 100644 index 0000000000..8a5e4d6df7 --- /dev/null +++ b/egomimic/scripts/embedding_process/dinov3_correspondence.py @@ -0,0 +1,398 @@ +"""Screen frozen DINO resolutions using independent visible-point annotations. + +This is descriptor matching, not a learned tracker. Manifest schema: +{model, revision, frames: [{id, path, sha256, episode, +points: {left_index: {xy: [x,y], visible: true, annotation_source: "..."}}}], +pairs: [[source_id, target_id], ...]}. Coordinates are original-image pixels. +Run with --manifest FILE --output DIR; no dataset writes or model substitution. + +Optional reference_mode="projected_landmark" requires projection_provenance +and per-point reference_valid instead of visible. This reports agreement with +estimated landmarks, not visible-point accuracy; unknown occlusion is retained. +Use --feature-results for SHA-verified extraction reuse after changing labels. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from egomimic.scripts.embedding_process.dinov3_embedding import DINOv3ImageEmbedding + + +def patch_geometry(metadata): + """Native-image patch centers and a mask excluding all padded/mixed patches.""" + gh, gw = metadata["patch_grid"] + ph, pw = metadata["patch_hw"] + geometry = metadata["frames"][0] + h, w = geometry["original_hw"] + rh, rw = geometry["resized_hw"] + top, left = geometry["padding_top_left"] + yy, xx = np.meshgrid(np.arange(gh), np.arange(gw), indexing="ij") + valid = ( + (xx * pw >= left) + & ((xx + 1) * pw <= left + rw) + & (yy * ph >= top) + & ((yy + 1) * ph <= top + rh) + ) + xy = np.stack( + [ + ((xx + 0.5) * pw - left) * w / rw - 0.5, + ((yy + 0.5) * ph - top) * h / rh - 0.5, + ], + -1, + ) + if not valid.any(): + raise ValueError("No fully unpadded patches at this resolution") + return xy.reshape(-1, 2), valid.reshape(-1) + + +def match_point(source, target, source_meta, target_meta, xy): + """Bilinear source descriptor; cosine nearest neighbor among valid target patches.""" + g = source_meta["frames"][0] + h, w = g["original_hw"] + rh, rw = g["resized_hw"] + top, left = g["padding_top_left"] + processed_h, processed_w = source_meta["processed_hw"] + qx = 2 * (((xy[0] + 0.5) * rw / w + left) / processed_w) - 1 + qy = 2 * (((xy[1] + 0.5) * rh / h + top) / processed_h) - 1 + gh, gw = source_meta["patch_grid"] + _, source_valid = patch_geometry(source_meta) + gx, gy = (qx + 1) * gw / 2 - 0.5, (qy + 1) * gh / 2 - 0.5 + x0, y0 = int(np.floor(gx)), int(np.floor(gy)) + for iy in (y0, y0 + 1): + for ix in (x0, x0 + 1): + weight = max(0, 1 - abs(gx - ix)) * max(0, 1 - abs(gy - iy)) + index = np.clip(iy, 0, gh - 1) * gw + np.clip(ix, 0, gw - 1) + if weight > 1e-6 and not source_valid[index]: + raise ValueError( + "Source descriptor touches a padded or mixed patch; exclude this label at every resolution" + ) + feature_grid = torch.as_tensor(source).reshape(gh, gw, -1).permute(2, 0, 1)[None] + query = F.grid_sample( + feature_grid, + torch.tensor([[[[qx, qy]]]], dtype=torch.float32), + mode="bilinear", + padding_mode="border", + align_corners=False, + ).flatten() + descriptors = F.normalize(torch.as_tensor(target), dim=-1) + scores = descriptors @ F.normalize(query, dim=0) + centers, valid = patch_geometry(target_meta) + scores[~torch.from_numpy(valid)] = -torch.inf + index = int(scores.argmax()) + return centers[index], float(scores[index]) + + +def load_manifest(path, features_only=False): + manifest = json.loads(Path(path).read_text()) + mode = manifest.get("reference_mode", "visible_manual") + if mode not in {"visible_manual", "projected_landmark"}: + raise ValueError("Unknown reference_mode") + if mode == "projected_landmark" and not manifest.get("projection_provenance"): + raise ValueError("Projected references require projection_provenance") + if not manifest.get("model") or not manifest.get("revision"): + raise ValueError("An explicit model and immutable revision are required") + if len(manifest["revision"]) != 40 or any( + c not in "0123456789abcdef" for c in manifest["revision"] + ): + raise ValueError("revision must be a full lowercase Git commit") + frames = {} + for frame in manifest["frames"]: + if frame["id"] in frames: + raise ValueError("Duplicate frame id") + file = Path(frame["path"]) + if not file.is_absolute(): + file = Path(path).parent / file + if hashlib.sha256(file.read_bytes()).hexdigest() != frame["sha256"]: + raise ValueError("Frame SHA mismatch") + image = np.asarray(Image.open(file).convert("RGB")) + for point in frame["points"].values(): + if not point.get("annotation_source"): + raise ValueError("Independent annotation_source is required") + if ( + mode == "visible_manual" + and point.get("visible") + and point.get("visibility_status", "").startswith("UNKNOWN") + ): + raise ValueError("Unknown visibility is not a visible manual label") + if eligible(point, mode): + x, y = point["xy"] + if not ( + np.isfinite([x, y]).all() + and 0 <= x < image.shape[1] + and 0 <= y < image.shape[0] + ): + raise ValueError( + "Visible annotation must be inside the native image" + ) + frames[frame["id"]] = (frame, image) + for a, b in [] if features_only else manifest["pairs"]: + if a == b or frames[a][0]["episode"] != frames[b][0]["episode"]: + raise ValueError("Pairs must be distinct frames from the same episode") + if not any( + eligible(p, mode) and eligible(frames[b][0]["points"].get(k, {}), mode) + for k, p in frames[a][0]["points"].items() + ): + raise ValueError( + "Each pair needs at least one eligible correspondence for its reference mode" + ) + if not features_only and not manifest["pairs"]: + raise ValueError("At least one pair is required") + if not frames: + raise ValueError("At least one original frame is required") + return manifest, frames + + +def eligible(point, mode): + """Projection validity is not image visibility; never conflate the two.""" + key = "reference_valid" if mode == "projected_landmark" else "visible" + return point.get(key) is True + + +def cached_features(result_paths, model, revision): + cache = {} + for result_path in result_paths: + result_path = Path(result_path) + result = json.loads(result_path.read_text()) + if (result["model"], result["revision"], result["precision"]) != ( + model, + revision, + "float32", + ): + raise ValueError("Feature cache model/revision/precision mismatch") + for resolution in result["resolutions"]: + for entry in resolution["feature_caches"].values(): + path = (result_path.parent / entry["path"]).resolve() + if hashlib.sha256(path.read_bytes()).hexdigest() != entry["sha256"]: + raise ValueError("Feature cache SHA mismatch") + with np.load(path, allow_pickle=False) as saved: + metadata = json.loads(str(saved["provenance_json"])) + features = saved["features"] + prep = metadata["preprocessing"] + if ( + prep["model"] != model + or prep["resolved_revision"] != revision + or prep["policy"] != "full_fov_letterbox_bicubic" + or prep["precision"] != "torch.float32" + or prep["processed_hw"] != [resolution["side"]] * 2 + or features.dtype != np.float32 + or not np.isfinite(features).all() + or features.shape[0] != np.prod(prep["patch_grid"]) + ): + raise ValueError("Feature cache preprocessing/array mismatch") + cache[(resolution["side"], metadata["frame_sha256"])] = ( + features, + prep, + {"path": str(path), "sha256": entry["sha256"]}, + ) + return cache + + +def synchronize(device): + if device.type == "cuda": + torch.cuda.synchronize(device) + elif device.type == "mps": + torch.mps.synchronize() + + +def run(manifest_path, output, sides, device, feature_results=(), features_only=False): + manifest, frames = load_manifest(manifest_path, features_only=features_only) + mode = manifest.get("reference_mode", "visible_manual") + reuse = cached_features(feature_results, manifest["model"], manifest["revision"]) + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + if (output / "results.json").exists(): + raise FileExistsError( + "Preserve existing results; choose a fresh output directory" + ) + extractor = DINOv3ImageEmbedding( + ["image"], + model_name=manifest["model"], + revision=manifest["revision"], + image_side=sides[0], + device=device, + dtype=torch.float32, + batch_size=1, + ) + extractor.setup() # Access failure occurs before any model substitution or matching. + results = { + "model": manifest["model"], + "revision": manifest["revision"], + "manifest_sha256": hashlib.sha256(Path(manifest_path).read_bytes()).hexdigest(), + "precision": "float32", + "device": str(device), + "resolutions": [], + "method": "feature_extraction_only" + if features_only + else "cosine_dense_descriptor_matching_not_learned_tracking", + "reference_mode": mode, + "metric_interpretation": "No correspondence metrics; unlabelled feature extraction only" + if features_only + else "Agreement with projected estimated 3D landmarks; unknown occlusion, NOT verified visible-point accuracy" + if mode == "projected_landmark" + else "Error against visible manual annotations", + "manifest_limitations": manifest.get("limitations", []), + } + for side in sides: + extractor.image_side = side + first = next(iter(frames.values()))[1] + for _ in range(2): + extractor._embed_batch([first]) + synchronize(extractor.device) + if extractor.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(extractor.device) + cache, latencies, cache_files = {}, [], {} + cache_dir = output / f"features-{side}" + cache_dir.mkdir(exist_ok=True) + for key, (_, image) in frames.items(): + previous = reuse.get((side, frames[key][0]["sha256"])) + if previous is not None: + features, metadata, entry = previous + cache[key] = (features, metadata) + cache_files[key] = entry + continue + synchronize(extractor.device) + start = time.perf_counter() + features = extractor._embed_batch([image])[0] + synchronize(extractor.device) + latencies.append(time.perf_counter() - start) + cache[key] = (features, extractor.last_preprocessing) + cache_path = cache_dir / ( + hashlib.sha256(key.encode()).hexdigest()[:16] + ".npz" + ) + if cache_path.exists(): + raise FileExistsError(f"Preserve existing feature cache: {cache_path}") + provenance = { + "frame_id": key, + "frame_sha256": frames[key][0]["sha256"], + "preprocessing": extractor.last_preprocessing, + } + np.savez_compressed( + cache_path, features=features, provenance_json=json.dumps(provenance) + ) + cache_files[key] = { + "path": str(cache_path.relative_to(output)), + "sha256": hashlib.sha256(cache_path.read_bytes()).hexdigest(), + } + matches, excluded = [], 0 + for a, b in [] if features_only else manifest["pairs"]: + source, target = frames[a][0], frames[b][0] + sf, sm = cache[a] + tf, tm = cache[b] + for name, p in source["points"].items(): + q = target["points"].get(name) + if not eligible(p, mode) or q is None or not eligible(q, mode): + excluded += 1 + continue + pred, cosine = match_point(sf, tf, sm, tm, p["xy"]) + back, _ = match_point(tf, sf, tm, sm, pred) + candidates = { + k: np.linalg.norm(pred - v["xy"]) + for k, v in target["points"].items() + if eligible(v, mode) + } + closest = min(candidates, key=candidates.get) + matches.append( + { + "source": a, + "target": b, + "point": name, + "hand": p.get("hand"), + "region": p.get("region"), + "visibility_status": p.get("visibility_status"), + "prediction_xy": pred.tolist(), + "target_xy": q["xy"], + "error_native_px": float(np.linalg.norm(pred - q["xy"])), + "cycle_error_native_px": float(np.linalg.norm(back - p["xy"])), + "cosine": cosine, + "nearest_labeled_point": closest, + "wrong_labeled_point": closest != name + if len(candidates) > 1 + else None, + } + ) + errors = np.asarray([m["error_native_px"] for m in matches]) + metadata = next(iter(cache.values()))[1] + result = { + "side": side, + "patch_tokens": int(np.prod(metadata["patch_grid"])), + "image_only_attention_pairs_vs_first_resolution": (side / sides[0]) ** 4, + "extract_median_seconds_including_transfer": float(np.median(latencies)) + if latencies + else None, + "reused_frames": len(frames) - len(latencies), + "latencies_seconds": latencies, + "matched_points": len(matches), + "excluded_occluded_or_unlabeled": excluded, + "mean_error_native_px": float(errors.mean()) if len(errors) else None, + "pck": {str(t): float((errors <= t).mean()) for t in [8, 16, 32]} + if len(errors) + else None, + "wrong_labeled_point_fraction": ( + float( + np.mean( + [ + m["wrong_labeled_point"] + for m in matches + if m["wrong_labeled_point"] is not None + ] + ) + ) + if any(m["wrong_labeled_point"] is not None for m in matches) + else None + ), + "wrong_finger_metric_limit": "Nearest annotated point only; null when fewer than two visible identities are labeled, not proof of correct finger identity", + "preprocessing": {k: v[1] for k, v in cache.items()}, + "feature_caches": cache_files, + "matches": matches, + } + if extractor.device.type == "cuda": + result["peak_allocated_bytes"] = torch.cuda.max_memory_allocated( + extractor.device + ) + results["resolutions"].append(result) + (output / "results.json").write_text(json.dumps(results, indent=2) + "\n") + return results + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--sides", type=int, nargs="+", default=[224, 336, 448, 640]) + parser.add_argument("--device", default="cpu") + parser.add_argument( + "--features-only", + action="store_true", + help="Extract verified image features without running point correspondence; preserves labels for downstream probes", + ) + parser.add_argument( + "--feature-results", + nargs="*", + default=[], + help="SHA-verified existing result JSONs to reuse features", + ) + args = parser.parse_args() + if not args.sides or any(s < 16 or s % 16 for s in args.sides): + parser.error("sides must be positive multiples of 16") + run( + args.manifest, + args.output, + args.sides, + args.device, + args.feature_results, + args.features_only, + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/dinov3_embedding.py b/egomimic/scripts/embedding_process/dinov3_embedding.py index de59e6822d..6fcb743070 100644 --- a/egomimic/scripts/embedding_process/dinov3_embedding.py +++ b/egomimic/scripts/embedding_process/dinov3_embedding.py @@ -10,12 +10,14 @@ from __future__ import annotations +import copy import logging import numpy as np import simplejpeg import torch import zarr +from PIL import Image from egomimic.scripts.embedding_process.zarr_key_transform import ZarrKeyTransform @@ -44,6 +46,8 @@ def __init__( device: str | torch.device = "cuda", chunk_timesteps: int = 100, dtype: torch.dtype = torch.float16, + revision: str | None = None, + image_side: int | None = None, ): if output_keys is not None and len(input_keys) != len(output_keys): raise ValueError( @@ -60,6 +64,16 @@ def __init__( ) self.model_name = model_name self.dtype = dtype + if image_side is not None and ( + isinstance(image_side, bool) + or not isinstance(image_side, int) + or image_side < 1 + ): + raise ValueError("image_side must be a positive integer") + self.revision = revision + self.image_side = image_side + self.last_preprocessing = None + self.output_metadata = {} self.processor = None self.model = None @@ -67,39 +81,114 @@ def setup(self) -> None: from transformers import AutoImageProcessor, AutoModel logger.info("Loading DINOv3 model: %s", self.model_name) - self.processor = AutoImageProcessor.from_pretrained(self.model_name) - self.model = AutoModel.from_pretrained(self.model_name, torch_dtype=self.dtype) + self.processor = AutoImageProcessor.from_pretrained( + self.model_name, revision=self.revision + ) + self.model = AutoModel.from_pretrained( + self.model_name, revision=self.revision, torch_dtype=self.dtype + ) self.model.to(self.device).eval() - - cfg = self.model.config - side = int(cfg.image_size) // int(cfg.patch_size) - self.num_patches = side * side - logger.info( - "DINOv3 patch grid: %d x %d = %d tokens (hidden_dim=%d)", - side, - side, - self.num_patches, - int(cfg.hidden_size), + self.model.requires_grad_(False) + + def _prepare(self, images): + """Explicit full-FOV letterboxing; default processor behavior stays opt-in legacy.""" + geometry = [] + prepared = [] + for image in images: + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[2] != 3: + raise ValueError("DINO images must be uint8 HWC RGB") + h, w = image.shape[:2] + if self.image_side is None: + prepared.append(image) + geometry.append({"original_hw": [h, w]}) + continue + side = self.image_side + scale = side / max(h, w) + rh, rw = max(1, round(h * scale)), max(1, round(w * scale)) + top, left = (side - rh) // 2, (side - rw) // 2 + mean = np.asarray(self.processor.image_mean) + canvas = np.empty((side, side, 3), dtype=np.uint8) + canvas[:] = np.rint(mean * 255).astype(np.uint8) + canvas[top : top + rh, left : left + rw] = np.asarray( + Image.fromarray(image).resize((rw, rh), Image.Resampling.BICUBIC) + ) + prepared.append(canvas) + geometry.append( + { + "original_hw": [h, w], + "resized_hw": [rh, rw], + "padding_top_left": [top, left], + } + ) + kwargs = ( + {} + if self.image_side is None + else {"do_resize": False, "do_center_crop": False} ) + inputs = self.processor(images=prepared, return_tensors="pt", **kwargs) + if self.image_side is not None and tuple(inputs["pixel_values"].shape[-2:]) != ( + self.image_side, + self.image_side, + ): + raise ValueError("DINO processor changed explicit letterbox resolution") + return inputs, geometry @torch.no_grad() def _embed_batch(self, images: list[np.ndarray]) -> np.ndarray: - inputs = self.processor(images=images, return_tensors="pt") + inputs, geometry = self._prepare(images) pixel_values = inputs["pixel_values"].to(self.device, dtype=self.dtype) + h, w = pixel_values.shape[-2:] + patch = self.model.config.patch_size + ph, pw = (patch, patch) if isinstance(patch, int) else patch + if h % ph or w % pw: + raise ValueError( + "Processed image dimensions must be divisible by DINO patch size" + ) + grid = [h // ph, w // pw] + num_patches = grid[0] * grid[1] outputs = self.model(pixel_values=pixel_values) # last_hidden_state ends with the patch tokens; CLS + any register # tokens sit at the start. Slice the trailing num_patches tokens to # get just the patch grid in row-major order. - patches = outputs.last_hidden_state[:, -self.num_patches :] + prefix = 1 + int(getattr(self.model.config, "num_register_tokens", 0)) + if outputs.last_hidden_state.shape[1] != prefix + num_patches: + raise ValueError( + "DINO token count disagrees with actual processed patch grid" + ) + patches = outputs.last_hidden_state[:, prefix:] + self.last_preprocessing = { + "model": self.model_name, + "requested_revision": self.revision, + "resolved_revision": getattr(self.model.config, "_commit_hash", None), + "processed_hw": [h, w], + "patch_grid": grid, + "patch_hw": [ph, pw], + "feature_layer": "last_hidden_state_patch_tokens", + "ordering": "row_major", + "precision": str(self.dtype), + "policy": "full_fov_letterbox_bicubic" + if self.image_side + else "processor_default", + "image_mean": list(self.processor.image_mean), + "image_std": list(self.processor.image_std), + "frames": geometry, + } return patches.float().cpu().numpy() + def _write_numeric(self, store, key, arr): + super()._write_numeric(store, key, arr) + store[key].attrs["dinov3_preprocessing"] = self.output_metadata[key] + def compute(self, store: zarr.Group) -> dict[str, np.ndarray]: outputs: dict[str, np.ndarray] = {} + self.output_metadata = {} for in_key, out_key in zip(self.input_keys, self.output_keys): arr = store[in_key] n = arr.shape[0] total_frames = int(store.attrs.get("total_frames", n)) n = min(n, total_frames) + if n < 1: + raise ValueError("DINO embedding requires at least one frame") chunks: list[np.ndarray] = [] for start in range(0, n, self.batch_size): @@ -110,5 +199,20 @@ def compute(self, store: zarr.Group) -> dict[str, np.ndarray]: for b in jpeg_batch ] chunks.append(self._embed_batch(images)) + current = copy.deepcopy(self.last_preprocessing) + frames = current.pop("frames") + if any(frame != frames[0] for frame in frames): + raise ValueError( + "DINO frame geometry must be constant within an episode camera" + ) + current["geometry"] = frames[0] + if ( + out_key in self.output_metadata + and self.output_metadata[out_key] != current + ): + raise ValueError( + "DINO preprocessing geometry changed within an output key" + ) + self.output_metadata[out_key] = current outputs[out_key] = np.concatenate(chunks, axis=0).astype(np.float32) return outputs diff --git a/egomimic/scripts/embedding_process/dinov3_pose_probe.py b/egomimic/scripts/embedding_process/dinov3_pose_probe.py new file mode 100644 index 0000000000..0576c7577e --- /dev/null +++ b/egomimic/scripts/embedding_process/dinov3_pose_probe.py @@ -0,0 +1,379 @@ +"""Frozen-DINO-only linear landmark heatmap probe on cached image features. + +Leave one whole episode out; fixed-budget training without test-based selection. +No track, pose, timestamp, camera-pose or target coordinates enter the readout. +Projected references remain estimated targets, not visible-joint ground truth. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import platform +import shutil +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from egomimic.scripts.embedding_process.dinov3_correspondence import ( + cached_features, + eligible, + load_manifest, + patch_geometry, +) + + +class PoseHeatmapProbe(nn.Module): + def __init__(self, feature_dim, landmarks): + super().__init__() + self.head = nn.Linear(feature_dim, landmarks, bias=False) + + def forward(self, features, valid_patches): + features = F.layer_norm(features, (features.shape[-1],)) + logits = self.head(features).transpose(1, 2) + return logits.masked_fill(~valid_patches[None, None], -1e9) + + +def heatmap_targets(xy, centers, valid_patches, sigma): + distances = (xy[:, :, None] - centers[None, None]).square().sum(-1) + logits = -distances / (2 * sigma**2) + return logits.masked_fill(~valid_patches[None, None], -1e9).softmax(-1) + + +def objective(logits, targets, valid_labels): + per_point = -(targets * logits.log_softmax(-1)).sum(-1) + return (per_point * valid_labels).sum() / valid_labels.sum().clamp_min(1) + + +def mean_pose_baseline(xy, valid): + count = valid.sum(0) + if (count == 0).any(): + raise ValueError("Every landmark needs at least one training label") + return (xy * valid[..., None]).sum(0) / count[:, None] + + +def episode_split(metadata, heldout): + train = [i for i, row in enumerate(metadata) if row["episode"] != heldout] + test = [i for i, row in enumerate(metadata) if row["episode"] == heldout] + if not train or not test: + raise ValueError("Nonempty train and held-out episodes required") + if {metadata[i]["image_sha256"] for i in train} & { + metadata[i]["image_sha256"] for i in test + }: + raise ValueError("Identical image bytes occur in both train and test") + return train, test + + +def prepare(args): + manifest, frames = load_manifest(args.manifest) + cached = cached_features( + args.feature_results, manifest["model"], manifest["revision"] + ) + names = sorted(next(iter(frames.values()))[0]["points"]) + features, targets, valid, metadata = [], [], [], [] + centers = patch_valid = None + for key, (frame, _) in frames.items(): + if sorted(frame["points"]) != names: + raise ValueError("All frames require identical landmark schema") + feature, prep, entry = cached[(args.side, frame["sha256"])] + coordinates, mask = patch_geometry(prep) + if centers is not None and ( + not np.array_equal(centers, coordinates) + or not np.array_equal(patch_valid, mask) + ): + raise ValueError("Probe currently requires identical camera/image geometry") + centers, patch_valid = coordinates, mask + features.append(feature) + label_valid = [ + eligible( + frame["points"][n], manifest.get("reference_mode", "visible_manual") + ) + for n in names + ] + target = np.asarray([frame["points"][n]["xy"] for n in names], np.float32) + target[~np.asarray(label_valid)] = 0 + targets.append(target) + valid.append(label_valid) + metadata.append( + { + "id": key, + "episode": frame["episode"], + "image_sha256": frame["sha256"], + "feature": entry, + } + ) + episodes = sorted({r["episode"] for r in metadata}) + if len(episodes) < 3: + raise ValueError("At least three episodes required for this pilot") + for episode in episodes: + episode_split(metadata, episode) + if not np.isfinite(np.stack(targets)).all(): + raise ValueError("Valid references must have finite coordinates") + return ( + manifest, + names, + metadata, + episodes, + torch.from_numpy(np.stack(features)).float(), + torch.from_numpy(np.stack(targets)).float(), + torch.tensor(valid), + torch.tensor(centers, dtype=torch.float32), + torch.tensor(patch_valid), + ) + + +def parameter_count(manifest): + from transformers import AutoConfig, AutoModel + + config = AutoConfig.from_pretrained( + manifest["model"], revision=manifest["revision"], local_files_only=True + ) + with torch.device("meta"): + backbone = AutoModel.from_config(config) + return sum(p.numel() for p in backbone.parameters()) + + +def run(args): + if args.steps < 1 or args.save_every < 1 or args.lr <= 0 or args.sigma <= 0: + raise ValueError("Positive training settings required") + manifest, names, meta, episodes, x, xy, label_valid, centers, patch_valid = prepare( + args + ) + head = PoseHeatmapProbe(x.shape[-1], len(names)) + trainable = sum(p.numel() for p in head.parameters()) + frozen = parameter_count(manifest) + out = Path(args.output) + out.mkdir(parents=True, exist_ok=True) + if list(out.glob("fold-*")) or (out / "COMPLETE.json").exists(): + raise FileExistsError( + "Preserve previous probe checkpoints; use a new experiment output" + ) + checkpoint_count = len(episodes) * math.ceil(args.steps / args.save_every) + reserve = checkpoint_count * (trainable * 12 + 65536) + free = shutil.disk_usage(out).free + if free < reserve + 100_000_000: + raise RuntimeError("Insufficient immutable checkpoint reserve") + config = { + "model": manifest["model"], + "revision": manifest["revision"], + "precision": "float32", + "device": args.device, + "side": args.side, + "frozen_parameters": frozen, + "trainable_parameters": trainable, + "total_parameters": frozen + trainable, + "feature_dim": x.shape[-1], + "patch_tokens": x.shape[1], + "landmarks": len(names), + "output_coordinates": len(names) * 2, + "denoising_dim": None, + "lr": args.lr, + "weight_decay": args.weight_decay, + "scheduler": None, + "steps_per_fold": args.steps, + "seed": args.seed, + "sigma_native_pixels": args.sigma, + "save_every": args.save_every, + "checkpoint_count": checkpoint_count, + "checkpoint_reserve_bytes": reserve, + "free_bytes": free, + "initialization": "fresh linear head per fold; frozen cached official DINOv3", + "input_contract": "image features only; per-token nonlearned LayerNorm; no poses, tracks or timestamps", + "baseline": "training-only mean pose per landmark", + "reference_mode": manifest.get("reference_mode", "visible_manual"), + "manifest_sha256": hashlib.sha256(Path(args.manifest).read_bytes()).hexdigest(), + "folds": [ + { + "heldout_episode": e, + "train_frames": [r["id"] for r in meta if r["episode"] != e], + "test_frames": [r["id"] for r in meta if r["episode"] == e], + } + for e in episodes + ], + } + config.update( + { + "optimizer": "AdamW", + "betas": [0.9, 0.999], + "epsilon": 1e-8, + "batch_size_per_fold": [len(episode_split(meta, e)[0]) for e in episodes], + "gradient_accumulation": 1, + "runtime": { + "torch": torch.__version__, + "python": platform.python_version(), + }, + "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "features": meta, + "evaluation_cadence": "terminal only; no test-based model selection", + } + ) + (out / "PREFLIGHT.json").write_text(json.dumps(config, indent=2) + "\n") + print(json.dumps(config, indent=2), flush=True) + if args.preflight_only: + return + torch.set_num_threads(args.threads) + x = x.to(args.device) + xy = xy.to(args.device) + label_valid = label_valid.to(args.device) + centers = centers.to(args.device) + patch_valid = patch_valid.to(args.device) + target = heatmap_targets(xy, centers, patch_valid, args.sigma) + predictions = [] + receipts = [] + for fold, e in enumerate(episodes): + torch.manual_seed(args.seed) + model = PoseHeatmapProbe(x.shape[-1], len(names)).float().to(args.device) + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=args.weight_decay + ) + train_ids, test_ids = episode_split(meta, e) + train = torch.tensor(train_ids, device=args.device) + test = torch.tensor(test_ids, device=args.device) + baseline = mean_pose_baseline(xy[train], label_valid[train]) + fold_dir = out / f"fold-{fold}" + fold_dir.mkdir() + (fold_dir / "split.json").write_text( + json.dumps(config["folds"][fold], indent=2) + "\n" + ) + with (fold_dir / "train.jsonl").open("w") as log: + for step in range(1, args.steps + 1): + optimizer.zero_grad(set_to_none=True) + logits = model(x[train], patch_valid) + loss = objective(logits, target[train], label_valid[train]) + if not torch.isfinite(loss): + raise RuntimeError("Nonfinite probe loss") + loss.backward() + if any( + p.grad is None or not torch.isfinite(p.grad).all() + for p in model.parameters() + ): + raise RuntimeError("Invalid probe gradient") + optimizer.step() + if step == 1 or step % 100 == 0: + row = {"step": step, "loss": float(loss.detach().cpu())} + log.write(json.dumps(row) + "\n") + log.flush() + print(f'fold={fold} step={step} loss={row["loss"]:.4f}', flush=True) + if step % args.save_every == 0 or step == args.steps: + checkpoint = fold_dir / f"epoch_{step:04d}-step_{step:09d}.pt" + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "global_step": step, + "epoch": step, + "config": config, + "heldout_episode": e, + "torch_rng": torch.get_rng_state(), + }, + checkpoint, + ) + receipts.append( + { + "path": str(checkpoint.resolve()), + "global_step": step, + "epoch": step, + "bytes": checkpoint.stat().st_size, + "sha256": hashlib.sha256( + checkpoint.read_bytes() + ).hexdigest(), + } + ) + with torch.no_grad(): + logits = model(x, patch_valid) + pred = centers[logits.argmax(-1)] + shuffled = pred[test.roll(1)] + for split, indices in [("train", train), ("test", test)]: + for k, i in enumerate(indices.tolist()): + predictions.append( + { + "fold": fold, + "split": split, + "frame_id": meta[i]["id"], + "episode": meta[i]["episode"], + "landmarks": names, + "predicted_xy": pred[i].cpu().tolist(), + "target_xy": xy[i].cpu().tolist(), + "valid": label_valid[i].cpu().tolist(), + "baseline_xy": baseline.cpu().tolist(), + "shuffled_xy": shuffled[k].cpu().tolist() + if split == "test" + else None, + } + ) + # Verify strict reload of actual terminal head and optimizer state. + saved = torch.load(checkpoint, map_location=args.device, weights_only=False) + restored = PoseHeatmapProbe(x.shape[-1], len(names)).float().to(args.device) + restored.load_state_dict(saved["model"], strict=True) + restored_opt = torch.optim.AdamW( + restored.parameters(), lr=args.lr, weight_decay=args.weight_decay + ) + restored_opt.load_state_dict(saved["optimizer"]) + with torch.no_grad(): + if not torch.equal( + restored(x[test], patch_valid), model(x[test], patch_valid) + ): + raise RuntimeError("Strict reload output changed") + (out / "predictions.json").write_text(json.dumps(predictions) + "\n") + summary = {} + for split in ["train", "test"]: + rows = [v for v in predictions if v["split"] == split] + metrics = {} + for key in ["predicted_xy", "baseline_xy"] + ( + ["shuffled_xy"] if split == "test" else [] + ): + errors = np.concatenate( + [ + np.linalg.norm( + np.array(v[key]) - np.array(v["target_xy"]), axis=-1 + )[v["valid"]] + for v in rows + ] + ) + metrics[key] = { + "n": len(errors), + "mean_px": float(errors.mean()), + "median_px": float(np.median(errors)), + "pck16": float((errors <= 16).mean()), + } + summary[split] = metrics + result = { + "status": "COMPLETE", + "scope": "Small leave-one-episode-out pilot; estimated projected references, unknown occlusion; not verified visible-joint or future-pose accuracy", + "summary": summary, + "config": config, + "checkpoints": receipts, + "strict_reload": "PASS all folds", + "predictions_sha256": hashlib.sha256( + (out / "predictions.json").read_bytes() + ).hexdigest(), + } + (out / "COMPLETE.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(summary, indent=2), flush=True) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--manifest", required=True) + p.add_argument("--feature-results", nargs="+", required=True) + p.add_argument("--output", required=True) + p.add_argument("--side", type=int, default=768) + p.add_argument("--steps", type=int, default=1000) + p.add_argument("--save-every", type=int, default=250) + p.add_argument("--lr", type=float, default=0.001) + p.add_argument("--weight-decay", type=float, default=0.0001) + p.add_argument("--sigma", type=float, default=8.0) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--device", default="mps") + p.add_argument("--threads", type=int, default=1) + p.add_argument("--preflight-only", action="store_true") + run(p.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/dinov3_segmentation_probe.py b/egomimic/scripts/embedding_process/dinov3_segmentation_probe.py new file mode 100644 index 0000000000..fa985e248e --- /dev/null +++ b/egomimic/scripts/embedding_process/dinov3_segmentation_probe.py @@ -0,0 +1,387 @@ +"""FP32 frozen-feature binary segmentation probe with episode-held-out evaluation. + +Annotation manifests bind original RGB hashes to visible-hand polygons, holes, +and uncertain regions. Labels must be visually reviewed before training. This +tool never makes masks from landmarks or treats model-generated masks as truth. +Uses the existing DINO cache verifier and landmark probe's split/count helpers. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image, ImageDraw +from scipy.ndimage import binary_dilation, binary_erosion, distance_transform_edt +from torch import nn + +from egomimic.scripts.embedding_process.dinov3_correspondence import cached_features +from egomimic.scripts.embedding_process.dinov3_pose_probe import episode_split, parameter_count + + +def sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def save_json(path, value): + Path(path).write_text(json.dumps(value, indent=2) + "\n") + + +def rasterize(frame, hw, band): + h, w = hw + mask = Image.new("L", (w, h), 0) + draw = ImageDraw.Draw(mask) + for polygon in frame["hand_polygons"]: + draw.polygon([tuple(p) for p in polygon], fill=1) + for polygon in frame.get("background_holes", []): + draw.polygon([tuple(p) for p in polygon], fill=0) + mask = np.asarray(mask).astype(bool) + uncertain = Image.new("L", (w, h), 0) + for polygon in frame.get("ignore_polygons", []): + ImageDraw.Draw(uncertain).polygon([tuple(p) for p in polygon], fill=1) + certain = ~np.asarray(uncertain).astype(bool) + edge_band = binary_dilation(mask, iterations=band) ^ binary_erosion(mask, iterations=band) + return mask, certain, certain & ~edge_band + + +class SegmentationProbe(nn.Module): + def __init__(self, dim): + super().__init__() + self.head = nn.Linear(dim, 2) + + def forward(self, features): + return self.head(F.layer_norm(features, (features.shape[-1],))) + + +def patch_targets(mask, valid, prep): + """Pixel-area targets; ignore padding and uncertain labels before reduction.""" + geometry = prep["frames"][0] + rh, rw = geometry["resized_hw"] + top, left = geometry["padding_top_left"] + ph, pw = prep["processed_hw"] + gh, gw = prep["patch_grid"] + planes = torch.tensor(np.stack([(~mask) & valid, mask & valid]), dtype=torch.float32)[None] + # Area weights retain thin structures and ignore coverage, unlike nearest labels. + resized = F.interpolate(planes, size=(rh, rw), mode="area") + padded = F.pad(resized, (left, pw - rw - left, top, ph - rh - top)) + counts = F.adaptive_avg_pool2d(padded, (gh, gw))[0].flatten(1).T + if counts.shape != (gh * gw, 2) or not torch.isfinite(counts).all(): + raise ValueError("Invalid label/grid geometry") + return counts + + +def objective(logits, counts, class_weights): + weights = counts * class_weights + denom = weights.sum() + if denom <= 0: + raise ValueError("No labeled pixels") + return -(weights * logits.log_softmax(-1)).sum() / denom + + +def native_probability(logits, prep): + gh, gw = prep["patch_grid"] + geometry = prep["frames"][0] + top, left = geometry["padding_top_left"] + rh, rw = geometry["resized_hw"] + # Interpolate logits, crop letterbox, then restore original image coordinates. + grid = logits.T.reshape(1, 2, gh, gw) + grid = F.interpolate(grid, size=prep["processed_hw"], mode="bilinear", align_corners=False) + grid = grid[:, :, top:top + rh, left:left + rw] + native = F.interpolate(grid, size=geometry["original_hw"], mode="bilinear", align_corners=False) + return native.softmax(1)[0, 1].numpy() + + +def overlap(pred, target, valid): + tp = int((pred & target & valid).sum()) + fp = int((pred & ~target & valid).sum()) + fn = int((~pred & target & valid).sum()) + return {"iou": tp / max(1, tp + fp + fn), "dice": 2 * tp / max(1, 2 * tp + fp + fn), + "tp": tp, "fp": fp, "fn": fn} + + +def boundary_f1(pred, target, certain, tolerance): + # Ignore clipped image edges and artificial cut edges around ambiguous wrists. + valid = binary_erosion(certain, iterations=tolerance + 1, border_value=0) + p = (pred ^ binary_erosion(pred)) & valid + t = (target ^ binary_erosion(target)) & valid + if not p.any() or not t.any(): + return 0.0 + precision = float((distance_transform_edt(~t)[p] <= tolerance).mean()) + recall = float((distance_transform_edt(~p)[t] <= tolerance).mean()) + return 2 * precision * recall / max(1e-12, precision + recall) + + +def load_labels(path): + manifest = json.loads(Path(path).read_text()) + if not manifest.get("annotation_source") or manifest["boundary_ignore_px"] < 1: + raise ValueError("Explicit label origin and positive uncertainty band required") + data, meta = [], [] + for frame in manifest["frames"]: + if sha(frame["path"]) != frame["sha256"]: + raise ValueError("Image SHA mismatch") + image = np.asarray(Image.open(frame["path"]).convert("RGB")) + mask, certain, valid = rasterize(frame, image.shape[:2], manifest["boundary_ignore_px"]) + if not (mask & valid).any() or not (~mask & valid).any(): + raise ValueError("Each annotation needs certain foreground and background") + data.append((image, mask, certain, valid)) + meta.append({"episode": frame["episode"], "image_sha256": frame["sha256"]}) + if len({m["episode"] for m in meta}) < 3: + raise ValueError("At least three distinct episodes required") + for episode in sorted({m["episode"] for m in meta}): + episode_split(meta, episode) + return manifest, data, meta + + +def review(args): + import matplotlib.pyplot as plt + manifest, data, _ = load_labels(args.annotations) + fig, axes = plt.subplots(len(data), 2, figsize=(12, 3.6 * len(data)), squeeze=False) + for i, (image, mask, certain, valid) in enumerate(data): + axes[i, 0].imshow(image) + axes[i, 1].imshow(image) + overlay = np.zeros((*mask.shape, 4)) + overlay[mask & valid] = [0, 1, 0.5, .45] + overlay[~valid] = [1, .65, 0, .65] + axes[i, 1].imshow(overlay) + axes[i, 1].contour(mask, levels=[.5], colors="white", linewidths=.6) + axes[i, 0].set_title(manifest["frames"][i]["id"]) + axes[i, 1].set_title("Green: hand | orange: ignored | white: polygon") + for ax in axes[i]: + ax.axis("off") + fig.tight_layout() + fig.savefig(args.output, dpi=100) + plt.close(fig) + + +def prepare(args): + manifest, data, meta = load_labels(args.annotations) + if manifest.get("review_status") != "visually_reviewed_before_training": + raise ValueError("Mask visual review must precede training") + cache = cached_features(args.feature_results, manifest["model"], manifest["revision"]) + packed = {} + for side in args.sides: + xs, targets, preps, receipts = [], [], [], [] + for frame, (_, mask, _, valid) in zip(manifest["frames"], data): + features, prep, receipt = cache[(side, frame["sha256"])] + if tuple(prep["frames"][0]["original_hw"]) != mask.shape: + raise ValueError("Annotation and feature geometry disagree") + xs.append(torch.from_numpy(features)) + targets.append(patch_targets(mask, valid, prep)) + preps.append(prep) + receipts.append(receipt) + packed[side] = (torch.stack(xs), torch.stack(targets), preps, receipts) + return manifest, data, meta, packed + + +def preflight(args, manifest, meta, packed): + if min(args.steps, args.save_every, args.threads) < 1 or args.lr <= 0: + raise ValueError("Positive training settings required") + out = Path(args.output) + out.mkdir(parents=True, exist_ok=True) + if list(out.glob("side-*")) or (out / "COMPLETE.json").exists(): + raise FileExistsError("Preserve previous probe outputs; no overwrite or pruning") + dim = next(iter(packed.values()))[0].shape[-1] + head = SegmentationProbe(dim) + trainable = sum(p.numel() for p in head.parameters()) + frozen = parameter_count(manifest) + episodes = sorted({r["episode"] for r in meta}) + count = len(args.sides) * len(episodes) * ((args.steps + args.save_every - 1) // args.save_every) + estimated_bytes = trainable * 12 + 65536 + free = shutil.disk_usage(out).free + if free < count * estimated_bytes + 100_000_000: + raise RuntimeError("Insufficient storage for all immutable checkpoints") + config = {"status": "PASS", "total_parameters": frozen + trainable, "frozen_parameters": frozen, + "trainable_parameters": trainable, "feature_dim": dim, "output_classes": 2, + "action_dim": None, "denoising_dim": None, "latent_bottleneck": None, + "sides": args.sides, "tokens": {s: int(packed[s][0].shape[1]) for s in args.sides}, + "steps_per_fold": args.steps, "save_every": args.save_every, "seed": args.seed, + "optimizer": "AdamW", "lr": args.lr, "weight_decay": args.weight_decay, + "scheduler": None, "precision": "float32", "device": "cpu", "cpu_threads": args.threads, + "batch": "all training images per update", "folds": len(episodes), + "splits": {ep: dict(zip(("train", "test"), episode_split(meta, ep))) for ep in episodes}, + "checkpoint_count": count, "estimated_bytes_each": estimated_bytes, "free_bytes": free, + "naming": "side-S/fold-F/epoch_EEEE-step_SSSSSSSSS.pt", + "triggers": "fixed step cadence plus terminal only if off cadence; retain all; no signal saves", + "output": str(out.resolve()), "annotation_sha256": sha(args.annotations), + "source_sha256": sha(__file__), "features": {s: packed[s][3] for s in args.sides}, + "fairness": "FAIR WITH CAVEATS: paired seed, full-batch order, steps and labels; intended resolution/token difference; equal updates not equal compute; nine model-annotated images only", + "selection": "terminal checkpoint fixed before training; no selection on test scores", + "threshold": .5, "postprocessing": None, + "annotation_source": manifest["annotation_source"]} + save_json(out / "PREFLIGHT.json", config) + return config + + +def run(args): + torch.set_num_threads(args.threads) + manifest, data, meta, packed = prepare(args) + config = preflight(args, manifest, meta, packed) + print(json.dumps(config, indent=2), flush=True) + if args.preflight_only: + return + out = Path(args.output) + episodes = sorted({r["episode"] for r in meta}) + rows, checkpoints, init_hashes = [], [], [] + for side, (x, targets, preps, _) in packed.items(): + for fold, episode in enumerate(episodes): + train, test = episode_split(meta, episode) + folder = out / f"side-{side}" / f"fold-{fold}" + folder.mkdir(parents=True, exist_ok=False) + torch.manual_seed(args.seed) + model = SegmentationProbe(x.shape[-1]) + init_hashes.append(sha_tensor_state(model.state_dict())) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + counts = targets[train] + # Native labeled pixel counts keep class balancing exactly paired + # across resolutions, independent of grid/resize rounding. + frequencies = torch.tensor([ + sum(int(((~data[i][1]) & data[i][3]).sum()) for i in train), + sum(int((data[i][1] & data[i][3]).sum()) for i in train), + ], dtype=torch.float32) + if (frequencies <= 0).any(): + raise ValueError("Missing training class") + weights = frequencies.sum() / (2 * frequencies) + losses = [] + for step in range(1, args.steps + 1): + optimizer.zero_grad(set_to_none=True) + loss = objective(model(x[train]), counts, weights) + loss.backward() + gradients = torch.cat([p.grad.flatten() for p in model.parameters()]) + if not torch.isfinite(loss) or not torch.isfinite(gradients).all() or gradients.norm() == 0: + raise RuntimeError("Nonfinite loss/gradients or zero gradient") + optimizer.step() + losses.append(float(loss.detach())) + if step % args.save_every == 0 or step == args.steps: + path = folder / f"epoch_{step:04d}-step_{step:09d}.pt" + torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), + "epoch": step, "step": step, "config": config, + "side": side, "heldout_episode": episode, "class_weights": weights}, path) + saved = torch.load(path, weights_only=False, map_location="cpu") + restored = SegmentationProbe(x.shape[-1]) + restored.load_state_dict(saved["model"], strict=True) + restored_opt = torch.optim.AdamW(restored.parameters(), lr=args.lr) + restored_opt.load_state_dict(saved["optimizer"]) + with torch.no_grad(): + if not torch.equal(restored(x[:1]), model(x[:1])): + raise RuntimeError("Strict model reload changed predictions") + for key in optimizer.state_dict()["state"]: + for name, value in optimizer.state_dict()["state"][key].items(): + if not torch.equal(value, restored_opt.state_dict()["state"][key][name]): + raise RuntimeError("Optimizer reload mismatch") + checkpoints.append({"path": str(path.resolve()), "size_bytes": path.stat().st_size, + "sha256": sha(path), "global_step": step, + "strict_model_optimizer_reload": True}) + save_json(folder / "losses.json", losses) + with torch.no_grad(): + logits = model(x) + prior_numerator = np.stack([data[i][1] & data[i][3] for i in train]).sum(0) + prior_denominator = np.stack([data[i][3] for i in train]).sum(0) + prior = prior_numerator / np.maximum(1, prior_denominator) >= .5 + for index in train + test: + probability = native_probability(logits[index], preps[index]) + prediction = probability >= .5 + _, mask, certain, valid = data[index] + row = {"side": side, "fold": fold, "split": "test" if index in test else "train", + "id": manifest["frames"][index]["id"], "episode": meta[index]["episode"], + "full_polygon": overlap(prediction, mask, certain), + "uncertainty_excluded": overlap(prediction, mask, valid), + "training_spatial_prior": overlap(prior, mask, valid), + "boundary_f1_5px": boundary_f1(prediction, mask, certain, manifest["boundary_tolerance_px"])} + if index in test: + dest = folder / (row["id"] + ".npz") + np.savez_compressed(dest, probability=probability, prediction=prediction, + polygon_mask=mask, certain=certain, scored=valid) + row["prediction"] = str(dest.resolve()) + row["prediction_sha256"] = sha(dest) + rows.append(row) + print(f"side={side} fold={fold} loss={losses[-1]:.6f} complete", flush=True) + if len(set(init_hashes)) != 1: + raise RuntimeError("Paired head initialization mismatch") + summary = {} + for side in args.sides: + tests = [r for r in rows if r["side"] == side and r["split"] == "test"] + summary[side] = {"test_images": len(tests), "episodes": len(episodes)} + for mode in ("full_polygon", "uncertainty_excluded", "training_spatial_prior"): + summary[side][mode] = {name: float(np.mean([r[mode][name] for r in tests])) for name in ("iou", "dice")} + summary[side]["boundary_f1_5px"] = float(np.mean([r["boundary_f1_5px"] for r in tests])) + save_json(out / "predictions.json", rows) + save_json(out / "SUMMARY.json", summary) + save_json(out / "COMPLETE.json", {"status": "COMPLETE_EXPLORATORY_AGENT_ANNOTATED_MASKS", + "config": config, "summary": summary, "checkpoints": checkpoints, + "paired_initialization_sha256": init_hashes[0], "predictions_sha256": sha(out / "predictions.json")}) + + +def sha_tensor_state(state): + return hashlib.sha256(b"".join(t.detach().numpy().tobytes() for t in state.values())).hexdigest() + + +def plot_results(args): + import matplotlib.pyplot as plt + manifest, data, _ = load_labels(args.annotations) + out = Path(args.output) + rows = json.loads((out / "predictions.json").read_text()) + rows = {(r["side"], r["id"]): r for r in rows if r["split"] == "test"} + # Predeclared middle frame per episode, not selected by error. + for label, indices in (("heldout-examples", [1, 4, 7]), ("all-heldout", list(range(len(data))))): + fig, axes = plt.subplots(len(indices), 4, figsize=(16, 3.7 * len(indices) + 1), squeeze=False) + fig.suptitle("Frozen DINOv3 hand segmentation | FP32 | held-out episodes\n" + "Exploratory Codex polygon labels; not independently human-annotated ground truth", fontsize=15) + for y, i in enumerate(indices): + image, mask, _, valid = data[i] + frame = manifest["frames"][i] + for ax in axes[y]: + ax.imshow(image) + ax.axis("off") + axes[y, 0].set_title(f"Original • frame {frame['frame']}") + overlay = np.zeros((*mask.shape, 4)) + overlay[mask & valid] = [0, 1, .5, .5] + overlay[~valid] = [1, .65, 0, .55] + axes[y, 1].imshow(overlay) + axes[y, 1].set_title("Reference mask • orange ignored") + for col, side in enumerate(args.sides, 2): + row = rows[(side, frame["id"])] + if sha(row["prediction"]) != row["prediction_sha256"]: + raise ValueError("Prediction SHA mismatch") + pred = np.load(row["prediction"], allow_pickle=False)["prediction"] + overlay[:] = 0 + overlay[pred] = [.05, .65, 1., .55] + axes[y, col].imshow(overlay) + axes[y, col].contour(mask, levels=[.5], colors="lime", linewidths=.5) + axes[y, col].set_title(f"{side} prediction • IoU {row['uncertainty_excluded']['iou']:.1%}") + fig.subplots_adjust(left=.01, right=.99, bottom=.02, + top=.90 if len(indices) == 3 else .96, + hspace=.24, wspace=.04) + fig.savefig(out / f"{label}.png", dpi=140) + plt.close(fig) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--annotations", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--review-only", action="store_true") + parser.add_argument("--plot-only", action="store_true") + parser.add_argument("--preflight-only", action="store_true") + parser.add_argument("--feature-results", nargs="+", default=[]) + parser.add_argument("--sides", nargs="+", type=int, default=[448, 640]) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument("--save-every", type=int, default=250) + parser.add_argument("--lr", type=float, default=.001) + parser.add_argument("--weight-decay", type=float, default=.0001) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--threads", type=int, default=2) + args = parser.parse_args() + if args.review_only: + review(args) + elif args.plot_only: + plot_results(args) + else: + run(args) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/world_model_cache.py b/egomimic/scripts/embedding_process/world_model_cache.py new file mode 100644 index 0000000000..05413496dd --- /dev/null +++ b/egomimic/scripts/embedding_process/world_model_cache.py @@ -0,0 +1,262 @@ +"""Stream verified native RGB through the maintained DINO forward into raw HDF5. + +These caches deliberately remain unnormalized. Dataset splitting, train-only +statistics and the training manifest belong to world_model_preparation. +""" +from __future__ import annotations + +import argparse +import copy +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path + +import h5py +import numpy as np +import torch + +from egomimic.rldb.world_model import open_prepared_episode +from egomimic.rldb.world_model_preparation import StreamingMoments + + +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 _require_hash(path, expected): + if sha256(path) != expected: + raise ValueError(f"Artifact hash mismatch: {path}") + + +def _safe_name(name): + if not isinstance(name, str) or not name or Path(name).name != name or name in (".", ".."): + raise ValueError("Unsafe episode/embodiment identity") + return name + + +def _source_identity(native_cache, rgb_cache, rgb_receipt): + native_cache, rgb_cache, rgb_receipt = map(Path, (native_cache, rgb_cache, rgb_receipt)) + manifest_path = native_cache / "manifest.json" + manifest, receipt = _json(manifest_path), _json(rgb_receipt) + if manifest.get("schema") != "libero-same-recorded-state-v1": + raise ValueError("Unsupported native cache contract") + 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): + raise ValueError("Require a successful native RGB and exact geometry receipt") + _require_hash(manifest_path, receipt["source_manifest_sha256"]) + _require_hash(native_cache / "aligned.npz", receipt["source_aligned_sha256"]) + _require_hash(native_cache / "aligned.npz", manifest["artifacts"]["aligned.npz"]) + _require_hash(rgb_cache, receipt["rgb_cache_sha256"]) + if receipt["cameras"] != manifest["cameras"]: + raise ValueError("Native and RGB camera ordering differs") + return manifest, receipt, { + "native_manifest_sha256": sha256(manifest_path), + "native_aligned_sha256": receipt["source_aligned_sha256"], + "rgb_cache_sha256": receipt["rgb_cache_sha256"], + "rgb_receipt_sha256": sha256(rgb_receipt), + } + + +def _geometry(native, manifest, rgb, receipt, side): + times = native["time"] + shape = rgb["rgb"].shape + if (len(shape) != 5 or shape[-1] != 3 or rgb["rgb"].dtype != np.uint8 + or shape[:2] != (len(times), len(manifest["cameras"])) + or shape[2:4] != (receipt["height"], receipt["width"]) + or len(times) != receipt["frames"]): + raise ValueError("RGB dimensions disagree with native episode/receipt") + # No low-resolution source may be silently upscaled to the chosen DINO side. + if min(shape[2:4]) < side: + raise ValueError("Native RGB resolution is below requested DINO resolution") + if shape[2] != shape[3] or manifest["height"] != manifest["width"]: + raise ValueError("Native simulator cache requires square camera views; do not infer letterbox track coordinates") + if not np.isfinite(times).all() or not (np.diff(times) > 0).all(): + raise ValueError("Require finite increasing physical timestamps") + np.testing.assert_array_equal(times, rgb["time"][:], err_msg="Native/RGB time alignment") + np.testing.assert_array_equal(native["source_row"], np.arange(len(times))) + pose, track = native["pose"], native["xyz"] + if (pose.shape != (len(times), len(manifest["pose_fields"])) + or track.ndim != 3 or track.shape[0] != len(times) or track.shape[2] != 3 + or not np.isfinite(pose).all() or not np.isfinite(track).all()): + raise ValueError("Invalid measured pose/track arrays") + available = native["coordinate_available"] + if available.dtype != np.bool_ or available.shape != track.shape[:2] or not available.all(): + raise ValueError("Native simulator coordinates must remain available under occlusion") + xy = native["reference_grid_xy"][native["valid_grid_indices"]].astype(np.float32) + if xy.shape != (track.shape[1], 2): + raise ValueError("Track reference grid mismatch") + xy = xy / np.array([manifest["width"] - 1, manifest["height"] - 1], np.float32) * 2 - 1 + return { + "time": times.astype(np.float64), "track": track.astype(np.float32), + "track_xy": xy, "track_camera": np.zeros(track.shape[1], np.int64), + "track_confidence": available.astype(np.float32), "pose": pose.astype(np.float32), + "pose_valid": np.ones(len(times), bool), "visibility": native["visibility"], + } + + +def extract_episode(*, native_cache, rgb_cache, rgb_receipt, output, embodiment, + episode_id, embedder, model_identity, batch_size=4): + """Write one immutable, raw FP32 cache; no corpus-sized in-memory arrays.""" + if batch_size < 1 or embedder.dtype != torch.float32 or not embedder.image_side: + raise ValueError("Require positive batch size and explicit FP32 DINO resolution") + if any(p.dtype != torch.float32 for p in embedder.model.parameters()): + raise ValueError("DINO weights must be FP32") + if embedder.model.training or any(p.requires_grad for p in embedder.model.parameters()): + raise ValueError("DINO must be frozen in evaluation mode") + if embedder.device.type == "cuda" and (torch.backends.cuda.matmul.allow_tf32 or torch.backends.cudnn.allow_tf32): + raise ValueError("FP32 feature extraction requires TF32 disabled") + for key in ("model", "revision", "weights_sha256", "config_sha256", "processor_sha256"): + if not model_identity.get(key): + raise ValueError(f"Missing pinned DINO {key}") + _safe_name(embodiment); _safe_name(episode_id) + output = Path(output).resolve() + output.parent.mkdir(parents=True, exist_ok=True) + manifest, receipt, inputs = _source_identity(native_cache, rgb_cache, rgb_receipt) + request = {"inputs": inputs, "model": model_identity, "image_side": embedder.image_side, + "batch_size": batch_size, "embodiment": embodiment, "episode_id": episode_id, + "producer_sha256": sha256(__file__)} + saved_receipt = output.with_suffix(output.suffix + ".json") + if output.exists(): + prior = _json(saved_receipt) + if prior.get("status") != "PASS" or prior.get("request") != request: + raise ValueError("Existing feature cache request differs or is incomplete") + _require_hash(output, prior["output_sha256"]) + return {**prior, "reused": True} + lock = output.with_suffix(output.suffix + ".lock") + temporary = output.with_suffix(output.suffix + ".partial") + with lock.open("x"): + pass + try: + if temporary.exists(): + raise FileExistsError(f"Prior partial cache requires inspection: {temporary}") + metadata = None + moments = {name: StreamingMoments() for name in ("image", "track", "pose")} + with open_prepared_episode(Path(native_cache) / "aligned.npz") as native, h5py.File(rgb_cache, "r") as rgb: + arrays = _geometry(native, manifest, rgb, receipt, embedder.image_side) + moments["track"].update(arrays["track"][arrays["track_confidence"] > 0]) + moments["pose"].update(arrays["pose"][arrays["pose_valid"]]) + n, cameras = rgb["rgb"].shape[:2] + with h5py.File(temporary, "x") as target: + target.attrs["arrays_normalized"] = False + target.attrs["reference_frame"] = manifest["reference_frame"] + for name, value in arrays.items(): + target.create_dataset(name, data=value, compression="lzf") + image_dataset = None + for camera in range(cameras): + for start in range(0, n, batch_size): + end = min(n, start + batch_size) + images = list(rgb["rgb"][start:end, camera]) + with torch.autocast(device_type=embedder.device.type, enabled=False): + features = embedder._embed_batch(images) + current = copy.deepcopy(embedder.last_preprocessing) + frames = current.pop("frames") + if not frames or any(item != frames[0] for item in frames): + raise ValueError("Inconsistent source image geometry") + current["geometry"] = frames[0] + if metadata is not None and current != metadata: + raise ValueError("DINO preprocessing changed within episode") + metadata = current + grid = metadata["patch_grid"] + if (features.dtype != np.float32 or not np.isfinite(features).all() + or features.shape != (end - start, int(np.prod(grid)), 768) + or metadata["processed_hw"] != [embedder.image_side] * 2 + or metadata["precision"] != "torch.float32"): + raise ValueError("Invalid DINO FP32 patch features/metadata") + if image_dataset is None: + image_dataset = target.create_dataset( + "image", shape=(n, cameras, features.shape[1], features.shape[2]), + dtype="f4", chunks=(1, 1, features.shape[1], features.shape[2])) + yy, xx = np.meshgrid((np.arange(grid[0]) + .5) / grid[0] * 2 - 1, + (np.arange(grid[1]) + .5) / grid[1] * 2 - 1, + indexing="ij") + target.create_dataset("image_xy", data=np.stack([xx, yy], -1).reshape(-1, 2).astype(np.float32)) + image_dataset[start:end, camera] = features + moments["image"].update(features) + target.attrs["dino_preprocessing"] = json.dumps(metadata) + result = { + "schema": "mot-world-feature-cache-v1", "status": "PASS", "request": request, + "arrays_normalized": False, "training_ready": False, "embodiment": embodiment, + "episode_id": episode_id, "frames": n, "pose_fields": manifest["pose_fields"], + "cameras": manifest["cameras"], "reference_frame": manifest["reference_frame"], + "native_cache": str(Path(native_cache).resolve()), + "native_source": {key: manifest.get(key) for key in ("source_demo", "source_demo_sha256", "episode")}, + "xy_convention": "minus_one_to_one; image patch centers; camera0 reference track pixels", + "track_confidence_policy": "exact simulator coordinate availability, independent of visibility", + "dino_identity": {**model_identity, "image_size": metadata["processed_hw"], + "patch_grid": metadata["patch_grid"], "feature_layer": metadata["feature_layer"], + "precision": "float32"}, + "dino_preprocessing": metadata, "native_rgb_receipt": receipt, + "moments": {name: {"count": value.count, "mean": value.mean.tolist(), + "m2": value.m2.tolist()} for name, value in moments.items()}, + "output": str(output), "output_sha256": sha256(temporary), + } + # Hard-link publication refuses to replace an existing immutable artifact. + os.link(temporary, output) + temporary.unlink() + with saved_receipt.open("x") as stream: + json.dump(result, stream, indent=2) + stream.write("\n") + return result + finally: + lock.unlink() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--spec", required=True, help="Explicit episodes and pinned local DINO model identity") + parser.add_argument("--output", required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument("--shard-count", type=int, default=1) + args = parser.parse_args() + if not 0 <= args.shard_index < args.shard_count: + raise ValueError("Invalid shard") + spec = _json(args.spec) + versions = {name: importlib.metadata.version(name) for name in ("transformers", "tokenizers", "torch", "h5py", "numpy")} + for name, expected in spec.get("runtime_versions", {}).items(): + if versions.get(name) != expected: + raise ValueError(f"Runtime {name} differs from requested version {expected}") + if args.device.startswith("cuda") and (not torch.cuda.is_available() or torch.cuda.device_count() != 1): + raise ValueError("DINO cache extraction requires exactly one visible CUDA GPU") + print(json.dumps({"status": "DINO_RUNTIME_READY", "versions": versions, "device": args.device}), flush=True) + identities, physical = set(), set() + if not spec.get("episodes"): + raise ValueError("Explicit nonempty episode inventory required") + for episode in spec["episodes"]: + key = (_safe_name(episode["embodiment"]), _safe_name(episode["episode_id"])) + native = str(Path(episode["native_cache"]).resolve()) + if key in identities or native in physical: + raise ValueError("Duplicate feature-cache identity or native source") + identities.add(key); physical.add(native) + local_model = Path(spec["model"]["path"]) + _require_hash(local_model / "model.safetensors", spec["model"]["weights_sha256"]) + _require_hash(local_model / "config.json", spec["model"]["config_sha256"]) + _require_hash(local_model / "preprocessor_config.json", spec["model"]["processor_sha256"]) + from egomimic.scripts.embedding_process.dinov3_embedding import DINOv3ImageEmbedding + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + embedder = DINOv3ImageEmbedding(input_keys=["rgb"], model_name=str(local_model), + revision=spec["model"]["revision"], dtype=torch.float32, + image_side=spec["image_side"], device=args.device) + embedder.setup() + identity = {key: spec["model"][key] for key in ("model", "revision", "weights_sha256", "config_sha256", "processor_sha256")} + for index, episode in enumerate(spec["episodes"]): + key = (_safe_name(episode["embodiment"]), _safe_name(episode["episode_id"])) + if index % args.shard_count != args.shard_index: + continue + result = extract_episode(**episode, output=Path(args.output) / key[0] / (key[1] + ".h5"), + embedder=embedder, model_identity=identity, + batch_size=spec.get("batch_size", 4)) + print(json.dumps({key: result[key] for key in ("status", "episode_id", "frames", "output", "output_sha256")}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/world_model_pose_probe.py b/egomimic/scripts/embedding_process/world_model_pose_probe.py new file mode 100644 index 0000000000..1a4306d0e3 --- /dev/null +++ b/egomimic/scripts/embedding_process/world_model_pose_probe.py @@ -0,0 +1,774 @@ +"""Native MoT clean-input pose-readout component diagnostic, not full training. + +Real image features and measured hand coordinates; no invented scene tracks. +Oracle future images isolate readout quality. Whole-episode holdout is mandatory. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import shutil +from pathlib import Path + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from egomimic.models import world_model as world_model_module +from egomimic.models.world_model import MODALITIES, MoTWorldModel +from egomimic.scripts.embedding_process.dinov3_correspondence import ( + cached_features, + patch_geometry, +) +from egomimic.utils.pose_utils import ee_pose_to_cam_frame, wxyz_to_xyzw + + +def sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def select_contiguous_anchors( + timestamps_ns, + count=4, + horizons=(0.1, 0.5, 1.0), + max_gap=0.05, + tolerance=0.02, + separation=3.0, +): + """Select deterministic, spaced anchors using actual physical timestamps.""" + timestamps = np.asarray(timestamps_ns, np.int64) + if timestamps.ndim != 1 or len(timestamps) < 2 or np.any(np.diff(timestamps) <= 0): + raise ValueError("Physical timestamps must strictly increase") + gaps = np.diff(timestamps) * 1e-9 + eligible = [] + for a in range(len(timestamps)): + targets = [] + for horizon in horizons: + wanted = timestamps[a] + round(horizon * 1e9) + position = int(np.searchsorted(timestamps, wanted)) + candidates = [ + b for b in [position - 1, position] if a < b < len(timestamps) + ] + if not candidates: + break + b = min(candidates, key=lambda i: abs(int(timestamps[i]) - wanted)) + if abs( + (timestamps[b] - timestamps[a]) * 1e-9 - horizon + ) > tolerance or np.any(gaps[a:b] > max_gap): + break + targets.append(b) + if len(targets) == len(horizons): + eligible.append((a, targets)) + chosen = [] + for quantile in np.linspace(0, 1, count + 2)[1:-1]: + candidates = [ + entry + for entry in eligible + if all( + abs(int(timestamps[entry[0]]) - int(timestamps[other[0]])) + > separation * 1e9 + for other in chosen + ) + ] + if not candidates: + raise ValueError("Insufficient distinct contiguous anchors") + wanted = timestamps[0] + quantile * (timestamps[-1] - timestamps[0]) + chosen.append( + min(candidates, key=lambda entry: abs(timestamps[entry[0]] - wanted)) + ) + return chosen + + +def validate_raw_aria_wrist(points, wrist_xyz, tolerance=1e-7): + """Reject mixing fitted/MANO-like point arrays with raw Aria-supported order.""" + points, wrists = np.asarray(points), np.asarray(wrist_xyz) + if points.shape[-2:] != (21, 3) or wrists.shape != points.shape[:-2] + (3,): + raise ValueError("Raw-Aria wrist topology shape mismatch") + if not np.isfinite(points).all() or not np.isfinite(wrists).all(): + raise ValueError("Nonfinite recorded human poses") + error = float(np.max(np.abs(points[..., 5, :] - wrists))) + if error > tolerance: + raise ValueError(f"Raw-Aria wrist index5 mismatch: {error} m") + return error + + +def _reference_pairs(records, explicit=None): + """Resolve explicit native-frame pairs without inventing adjacent targets.""" + keys = [(r["episode"], r["frame"]) for r in records] + if len(keys) != len(set(keys)): + raise ValueError("Duplicate episode/frame reference") + lookup = {key: i for i, key in enumerate(keys)} + if explicit is None: + return [ + (a, a + 1) + for a in range(len(records) - 1) + if records[a]["episode"] == records[a + 1]["episode"] + ] + result = [] + for pair in explicit: + episode = pair["episode"] + a = lookup[episode, pair["source_frame"]] + b = lookup[episode, pair["target_frame"]] + if ( + a == b + or records[b]["obs_rgb_timestamps_ns"] + <= records[a]["obs_rgb_timestamps_ns"] + ): + raise ValueError("Explicit pairs require positive physical time") + result.append((a, b)) + if not result or len(result) != len(set(result)): + raise ValueError("Explicit pairs must be nonempty and unique") + return result + + +def prepare_inputs(reference_path, feature_results, model, revision, side, output): + """Shared generator/readout input contract; no training and no DINO rerun.""" + references = json.loads(Path(reference_path).read_text()) + records = sorted( + references["records"], key=lambda x: (x["episode"], x["obs_rgb_timestamps_ns"]) + ) + cache = cached_features(feature_results, model, revision) + features, world, transforms, metadata = [], [], [], [] + xy = valid = None + for r in records: + if sha(r["path"]) != r["sha256"] or not 0 <= r["frame"] < r["total_frames"]: + raise ValueError("Original frame hash or logical frame bound mismatch") + feature, prep, entry = cache[side, r["sha256"]] + centers, mask = patch_geometry(prep) + h, w = prep["frames"][0]["original_hw"] + normalized = (centers + 0.5) / [w, h] * 2 - 1 + if xy is not None and ( + not np.allclose(xy, normalized) or not np.array_equal(valid, mask) + ): + raise ValueError("Component probe requires matching camera geometry") + xy, valid = normalized.astype(np.float32), mask + pose = np.concatenate( + [r["left.obs_keypoints"], r["right.obs_keypoints"]] + ).reshape(42, 3) + if not np.isfinite(pose).all() or np.abs(pose).max() > 100: + raise ValueError("Invalid measured world XYZ") + head = np.asarray(r["obs_head_pose"]) + if not np.isfinite(head).all() or abs(np.linalg.norm(head[3:]) - 1) > 0.01: + raise ValueError("Invalid reference camera pose") + transform = np.eye(4) + transform[:3, :3] = Rotation.from_quat(wxyz_to_xyzw(head[3:])).as_matrix() + transform[:3, 3] = head[:3] + features.append(feature) + world.append(pose) + transforms.append(transform) + metadata.append( + { + "id": Path(r["path"]).stem, + "episode": r["episode"], + "frame": r["frame"], + "timestamp_ns": r["obs_rgb_timestamps_ns"], + "image_sha256": r["sha256"], + "feature": entry, + } + ) + pairs, pair_pose, dt = [], [], [] + for a, b in _reference_pairs(records, references.get("pairs")): + delta = ( + records[b]["obs_rgb_timestamps_ns"] - records[a]["obs_rgb_timestamps_ns"] + ) * 1e-9 + if not np.isfinite(delta) or delta <= 0: + raise ValueError("Physical timestamps must strictly increase") + pairs.append([a, b]) + dt.append(delta) + pair_pose.append( + np.stack( + [ + ee_pose_to_cam_frame(world[i], transforms[a]).reshape(-1) + for i in [a, b] + ] + ) + ) + episodes = sorted({r["episode"] for r in metadata}) + if len(episodes) < 3: + raise ValueError("At least three independent episodes required") + for e in episodes: + train_sha = {m["image_sha256"] for m in metadata if m["episode"] != e} + test_sha = {m["image_sha256"] for m in metadata if m["episode"] == e} + if train_sha & test_sha: + raise ValueError("Image bytes overlap across heldout episodes") + out = Path(output) + out.mkdir(parents=True, exist_ok=True) + if (out / "INPUTS.npz").exists(): + raise FileExistsError("Preserve prepared component inputs") + np.savez_compressed( + out / "INPUTS.npz", + features=np.stack(features), + patch_xy=xy, + patch_valid=valid, + pairs=np.asarray(pairs, np.int64), + poses=np.asarray(pair_pose, np.float32), + dt=np.asarray(dt, np.float32), + reference_world_from_camera=np.asarray(transforms), + timestamps_ns=np.asarray([m["timestamp_ns"] for m in metadata], np.int64), + ) + info = { + "frames": metadata, + "episodes": episodes, + "pose_schema": "42 recorded raw-Aria-supported world XYZ points; not canonical MANO;126 coordinates", + "pose_reference": "Both poses in each pair transformed with inverse first-frame world_from_camera; future camera pose is not model input", + "scene_tracks": "unavailable; empty track stream, never hand-joints-as-tracks", + "model": model, + "revision": revision, + "side": side, + "precision": "float32", + "source_reference_sha256": sha(reference_path), + "pair_selection": "explicit" + if "pairs" in references + else "adjacent_selected_frames", + "input_sha256": sha(out / "INPUTS.npz"), + } + (out / "INPUTS.json").write_text(json.dumps(info, indent=2) + "\n") + return info + + +def load_inputs(directory): + directory = Path(directory) + meta = json.loads((directory / "INPUTS.json").read_text()) + if sha(directory / "INPUTS.npz") != meta["input_sha256"]: + raise ValueError("Prepared component input SHA mismatch") + with np.load(directory / "INPUTS.npz", allow_pickle=False) as z: + data = {k: z[k].copy() for k in z.files} + return data, meta + + +def fold_data(data, meta, heldout): + train_frames = np.asarray( + [i for i, r in enumerate(meta["frames"]) if r["episode"] != heldout] + ) + train = np.asarray( + [ + i + for i, (a, b) in enumerate(data["pairs"]) + if meta["frames"][a]["episode"] != heldout + ] + ) + test = np.asarray( + [ + i + for i, (a, b) in enumerate(data["pairs"]) + if meta["frames"][a]["episode"] == heldout + ] + ) + if len(train) == 0 or len(test) == 0: + raise ValueError("Nonempty train and heldout pairs required") + return train, test, fit_normalization(data, train_frames, train) + + +def fit_normalization(data, train_frames, train): + """Fit only on explicitly selected training frames and target pairs.""" + f = data["features"][train_frames][:, data["patch_valid"]].reshape( + -1, data["features"].shape[-1] + ) + p = data["poses"][train].reshape(-1, data["poses"].shape[-1]) + norm = { + "image_mean": f.mean(0), + "image_std": f.std(0).clip(1e-5), + "pose_mean": p.mean(0), + "pose_std": p.std(0).clip(0.05), + } + return norm + + +def fixed_partition(data, meta, split): + """Separate validation selection from a final untouched episode test.""" + names = ("train", "validation", "test") + groups = {name: set(split[name]) for name in names} + episodes = {r["episode"] for r in meta["frames"]} + if any(not groups[n] or len(groups[n]) != len(split[n]) for n in names): + raise ValueError("Split groups must be nonempty without duplicate episodes") + if any( + groups[a] & groups[b] + for a, b in (("train", "validation"), ("train", "test"), ("validation", "test")) + ): + raise ValueError("Episode split overlap") + if set.union(*groups.values()) != episodes: + raise ValueError("Split must cover exactly the prepared episodes") + # Episode names alone cannot detect copied frames under different aliases. + image_groups = { + n: {r["image_sha256"] for r in meta["frames"] if r["episode"] in groups[n]} + for n in names + } + if any( + image_groups[a] & image_groups[b] + for a, b in (("train", "validation"), ("train", "test"), ("validation", "test")) + ): + raise ValueError("Image content overlaps across split groups") + indices = {name: [] for name in names} + for i, (a, b) in enumerate(data["pairs"]): + episode = meta["frames"][a]["episode"] + if episode != meta["frames"][b]["episode"]: + raise ValueError("A prediction pair crosses episodes") + for name in names: + if episode in groups[name]: + indices[name].append(i) + if any(not indices[n] for n in names): + raise ValueError("Each split requires prediction pairs") + indices = {k: np.asarray(v, dtype=np.int64) for k, v in indices.items()} + train_frames = np.asarray( + [i for i, r in enumerate(meta["frames"]) if r["episode"] in groups["train"]] + ) + return indices, fit_normalization(data, train_frames, indices["train"]) + + +def model_for_probe(width=128, depth=2): + model = MoTWorldModel( + {"human": 126}, width=width, depth=depth, heads=4, image_adapter_dim=16 + ).float() + # These outputs/track-private parameters are not used by a pose-only loss. + track_index = MODALITIES.index("track") + unused = [ + model.image_out, + model.track_in, + model.track_out, + model.point, + model.norm[track_index], + ] + for block in model.blocks: + unused.extend( + [ + block.norm1[track_index], + block.qkv[track_index], + block.out[track_index], + block.norm2[track_index], + block.ffn[track_index], + ] + ) + for module in unused: + module.requires_grad_(False) + return model + + +def make_batch(data, norm, indices, device="cpu", visual=True): + indices = np.asarray(indices) + pairs = data["pairs"][indices] + b = len(indices) + n = data["features"].shape[1] + features = (data["features"][pairs] - norm["image_mean"]) / norm["image_std"] + pose = (data["poses"][indices] - norm["pose_mean"]) / norm["pose_std"] + target = torch.tensor(pose[:, 1], device=device) + pose[:, 1] = 0 # Query targets never enter the forward batch. + times = np.stack([np.zeros(b), data["dt"][indices]], axis=1).astype(np.float32) + + def tensor(x): + return torch.as_tensor(x, device=device) + + image_valid = np.tile(data["patch_valid"], (b, 2)) & visual + image = { + "value": tensor(features.reshape(b, 2 * n, -1)), + "time": tensor(np.repeat(times, n, axis=1)), + "valid": tensor(image_valid), + "observed": tensor(image_valid.copy()), + "xy": tensor(np.tile(data["patch_xy"], (b, 2, 1))), + "camera": torch.zeros(b, 2 * n, dtype=torch.long, device=device), + } + empty = { + "value": torch.zeros(b, 0, 3, device=device), + "time": torch.zeros(b, 0, device=device), + "valid": torch.zeros(b, 0, dtype=torch.bool, device=device), + "observed": torch.zeros(b, 0, dtype=torch.bool, device=device), + "xy": torch.zeros(b, 0, 2, device=device), + "camera": torch.zeros(b, 0, dtype=torch.long, device=device), + "point_id": torch.zeros(b, 0, dtype=torch.long, device=device), + } + proprio = { + "value": tensor(pose), + "time": tensor(times), + "valid": torch.ones(b, 2, dtype=torch.bool, device=device), + "observed": tensor(np.tile([True, False], (b, 1))), + } + batch = {"image": image, "track": empty, "proprio": proprio} + return batch, target + + +def predict(model, batch): + return model( + batch, + "human", + {k: batch[k]["value"] for k in ("image", "track")}, + {k: torch.zeros_like(batch[k]["time"]) for k in ("image", "track")}, + )["proprio"][:, 1] + + +def evaluate(model, data, meta, norm, indices, device, visual): + rows = [] + model.eval() + with torch.no_grad(): + for i in indices: + batch, _ = make_batch(data, norm, [i], device, visual) + estimate = ( + predict(model, batch).cpu().numpy()[0] * norm["pose_std"] + + norm["pose_mean"] + ) + truth = data["poses"][i, 1] + hold = data["poses"][i, 0] + a, b = data["pairs"][i] + rows.append( + { + "source": meta["frames"][a]["id"], + "target": meta["frames"][b]["id"], + "dt_seconds": float(data["dt"][i]), + "prediction_xyz": estimate.reshape(42, 3).tolist(), + "target_xyz": truth.reshape(42, 3).tolist(), + "mpjpe_m": float( + np.linalg.norm( + (estimate - truth).reshape(42, 3), axis=-1 + ).mean() + ), + "hold_pose_mpjpe_m": float( + np.linalg.norm((hold - truth).reshape(42, 3), axis=-1).mean() + ), + } + ) + return rows + + +def run(args): + torch.set_num_threads(args.threads) + data, meta = load_inputs(args.inputs) + out = Path(args.output) + out.mkdir(parents=True, exist_ok=True) + if list(out.glob("fold-*")): + raise FileExistsError("Preserve existing diagnostic checkpoints") + model = model_for_probe(args.width, args.depth) + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + split_path = getattr(args, "split_manifest", None) + if split_path: + partition = json.loads(Path(split_path).read_text()) + indices, normalization = fixed_partition(data, meta, partition) + plans = [ + ( + "fixed_split", + indices["train"], + indices["test"], + normalization, + indices["validation"], + ) + ] + else: + plans = [ + (e, *fold_data(data, meta, e), np.asarray([], dtype=np.int64)) + for e in meta["episodes"] + ] + count = len(plans) * 2 * math.ceil(args.steps / args.save_every) + reserve = count * (total * 4 + trainable * 8 + 65536) + free = shutil.disk_usage(out).free + if free < reserve + 100_000_000: + raise RuntimeError("Insufficient immutable checkpoint reserve") + config = { + "input_sha256": meta["input_sha256"], + "width": args.width, + "depth": args.depth, + "heads": 4, + "image_adapter_dim": 16, + "parameters_total": total, + "parameters_requires_grad": trainable, + "denoising_dim": None, + "hidden_dim": args.width, + "output_pose_dim": 126, + "action_dim": None, + "input_image_dim": data["features"].shape[-1], + "image_tokens_per_sequence": 2 * data["features"].shape[1], + "pose_tokens": 2, + "track_tokens": 0, + "optimizer": "AdamW", + "lr": args.lr, + "weight_decay": 0.01, + "betas": [0.9, 0.999], + "epsilon": 1e-8, + "scheduler": None, + "batch": 2, + "steps_per_condition_per_fold": args.steps, + "conditions": ["observed_pose_only", "oracle_images_and_observed_pose"], + "seed": 42, + "precision": "FP32", + "device": args.device, + "cpu_threads": args.threads, + "save_every": args.save_every, + "checkpoint_count": count, + "checkpoint_reserve_bytes": reserve, + "free_bytes": free, + "folds": [p[0] for p in plans], + "split_manifest_sha256": sha(split_path) if split_path else None, + "source_sha256": sha(__file__), + "model_source_sha256": sha(world_model_module.__file__), + "evaluation": "validation at each saved checkpoint; select minimum validation MPJPE then evaluate test once" + if split_path + else "terminal only, all heldout pairs; no test-based selection", + "full_model_training": False, + } + (out / "PREFLIGHT.json").write_text(json.dumps(config, indent=2) + "\n") + source_snapshot = out / "source" + source_snapshot.mkdir(exist_ok=True) + for filename in (__file__, world_model_module.__file__): + shutil.copy2(filename, source_snapshot / Path(filename).name) + # Empty modality streams can pass MPS forward but fail its backward kernel. + # Probe the exact device, batch shape and both conditions before any update. + guard = {} + _, train, _, norm, _ = plans[0] + for visual in (False, True): + torch.manual_seed(42) + probe = model_for_probe(args.width, args.depth).to(args.device) + batch, target = make_batch(data, norm, train[:2], args.device, visual) + loss = (predict(probe, batch) - target).square().mean() + loss.backward() + if not torch.isfinite(loss) or not all( + p.grad is None or torch.isfinite(p.grad).all() for p in probe.parameters() + ): + raise RuntimeError("Exact-runtime backward preflight failed") + guard[str(visual)] = {"loss": float(loss.detach().cpu()), "finite": True} + (out / "RUNTIME_BACKWARD_GUARD.json").write_text( + json.dumps( + {"device": args.device, "conditions": guard, "optimizer_updates": 0}, + indent=2, + ) + + "\n" + ) + print(json.dumps(config, indent=2), flush=True) + if args.preflight_only: + return + receipts = [] + all_rows = [] + for fold, (heldout, train, test, norm, validation) in enumerate(plans): + for visual in (False, True): + condition = ( + "oracle_images_and_observed_pose" if visual else "observed_pose_only" + ) + directory = out / f"fold-{fold}" / condition + directory.mkdir(parents=True) + torch.manual_seed(42) + model = model_for_probe(args.width, args.depth).to(args.device) + initial_sha = hashlib.sha256( + b"".join(p.detach().cpu().numpy().tobytes() for p in model.parameters()) + ).hexdigest() + optimizer = torch.optim.AdamW( + [p for p in model.parameters() if p.requires_grad], + lr=args.lr, + weight_decay=0.01, + ) + rng = np.random.default_rng(42) + order = np.asarray([], dtype=np.int64) + epoch = 0 + orders = [] + split = { + "heldout_episode": heldout, + "train_pair_indices": train.tolist(), + "test_pair_indices": test.tolist(), + "validation_pair_indices": validation.tolist(), + "normalization": {k: v.tolist() for k, v in norm.items()}, + } + (directory / "split.json").write_text(json.dumps(split, indent=2) + "\n") + best_validation = float("inf") + selected_path = None + validation_history = [] + with (directory / "train.jsonl").open("w") as log: + for step in range(1, args.steps + 1): + if len(order) == 0: + order = rng.permutation(train) + epoch += 1 + ids, order = order[:2], order[2:] + orders.append(ids.tolist()) + batch, target = make_batch(data, norm, ids, args.device, visual) + model.train() + optimizer.zero_grad(set_to_none=True) + loss = (predict(model, batch) - target).square().mean() + if not torch.isfinite(loss): + raise RuntimeError("Nonfinite pose loss") + loss.backward() + active = ( + sum( + p.numel() + for p in model.parameters() + if p.grad is not None and torch.count_nonzero(p.grad) > 0 + ) + if step == 1 + else None + ) + if not torch.stack( + [ + torch.isfinite(p.grad).all() + for p in model.parameters() + if p.grad is not None + ] + ).all(): + raise RuntimeError("Nonfinite pose gradient") + optimizer.step() + if step == 1 or step % 25 == 0: + row = { + "step": step, + "epoch": epoch, + "loss": float(loss.detach().cpu()), + "nonzero_gradient_parameters": active, + } + log.write(json.dumps(row) + "\n") + log.flush() + print(f"fold{fold} {condition} {row}", flush=True) + if step % args.save_every == 0 or step == args.steps: + path = directory / f"epoch_{epoch:04d}-step_{step:09d}.pt" + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "step": step, + "epoch": epoch, + "normalization": norm, + "config": config, + "condition": condition, + "heldout_episode": heldout, + }, + path, + ) + receipts.append( + { + "path": str(path.resolve()), + "sha256": sha(path), + "bytes": path.stat().st_size, + } + ) + if len(validation): + val_rows = evaluate( + model, data, meta, norm, validation, args.device, visual + ) + score = float(np.mean([r["mpjpe_m"] for r in val_rows])) + if not math.isfinite(score): + raise RuntimeError("Nonfinite validation score") + validation_history.append( + { + "step": step, + "mpjpe_m": score, + "checkpoint": str(path.resolve()), + } + ) + if score < best_validation: + best_validation, selected_path = score, path + else: + selected_path = path + # Restore the validation-selected state before ever evaluating test. + selected = torch.load( + selected_path, map_location=args.device, weights_only=False + ) + model.load_state_dict(selected["model"], strict=True) + optimizer.load_state_dict(selected["optimizer"]) + (directory / "SELECTION.json").write_text( + json.dumps( + { + "checkpoint": str(selected_path.resolve()), + "selected_step": selected["step"], + "selected_by": "validation_mpjpe" + if len(validation) + else "terminal_budget", + "validation_history": validation_history, + "test_used_for_selection": False, + }, + indent=2, + ) + + "\n" + ) + before = evaluate(model, data, meta, norm, test, args.device, visual) + (directory / "RUN_IDENTITY.json").write_text( + json.dumps( + { + "initial_state_sha256": initial_sha, + "sample_order_sha256": hashlib.sha256( + json.dumps(orders).encode() + ).hexdigest(), + "sample_order": orders, + }, + indent=2, + ) + + "\n" + ) + clone = model_for_probe(args.width, args.depth).to(args.device) + loaded = torch.load( + selected_path, map_location=args.device, weights_only=False + ) + clone.load_state_dict(loaded["model"], strict=True) + restored_optimizer = torch.optim.AdamW( + [p for p in clone.parameters() if p.requires_grad], + lr=args.lr, + weight_decay=0.01, + ) + restored_optimizer.load_state_dict(loaded["optimizer"]) + if len(restored_optimizer.state) != len(optimizer.state): + raise RuntimeError("Optimizer reload state count differs") + for original, restored in zip( + optimizer.state.values(), restored_optimizer.state.values() + ): + for key, value in original.items(): + if torch.is_tensor(value) and not torch.equal( + value.to("cpu"), restored[key].to("cpu") + ): + raise RuntimeError("Optimizer reload tensor differs") + after = evaluate(clone, data, meta, norm, test, args.device, visual) + if before != after: + raise RuntimeError("Strict reload predictions differ") + for row in before: + row.update(fold=fold, condition=condition, split="test") + all_rows.extend(before) + train_rows = evaluate(model, data, meta, norm, train, args.device, visual) + for row in train_rows: + row.update(fold=fold, condition=condition, split="train") + all_rows.extend(train_rows) + (directory / "predictions.json").write_text( + json.dumps(before + train_rows, indent=2) + "\n" + ) + summary = {} + for condition in config["conditions"]: + summary[condition] = { + split: { + key: float( + np.mean( + [ + r[key] + for r in all_rows + if r["condition"] == condition and r["split"] == split + ] + ) + ) + for key in ["mpjpe_m", "hold_pose_mpjpe_m"] + } + for split in ["train", "test"] + } + complete = { + "status": "COMPLETE", + "config": config, + "summary": summary, + "checkpoints": receipts, + "strict_reload": "PASS", + "limits": "Small real recorded-coordinate readout diagnostic, not full model, not scene-track test, not visiblehandtruth; clean future images are oracle context.", + } + (out / "COMPLETE.json").write_text(json.dumps(complete, indent=2) + "\n") + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--inputs", required=True) + p.add_argument("--output", required=True) + p.add_argument("--width", type=int, default=128) + p.add_argument("--depth", type=int, default=2) + p.add_argument("--steps", type=int, default=300) + p.add_argument("--save-every", type=int, default=100) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--device", default="mps") + p.add_argument("--threads", type=int, default=2) + p.add_argument( + "--split-manifest", + help="Explicit train/validation/test episode lists; default retains historical leave-one-episode-out mode", + ) + p.add_argument("--preflight-only", action="store_true") + args = p.parse_args() + if args.steps < 1 or args.save_every < 1 or args.lr <= 0 or args.threads < 1: + p.error("Positive training settings required") + run(args) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/embedding_process/zarr_embedding.py b/egomimic/scripts/embedding_process/zarr_embedding.py index 15b005c32b..130283826a 100644 --- a/egomimic/scripts/embedding_process/zarr_embedding.py +++ b/egomimic/scripts/embedding_process/zarr_embedding.py @@ -58,6 +58,9 @@ def _build_transform(args: argparse.Namespace) -> ZarrKeyTransform: kwargs["output_keys"] = args.output_keys if args.model_name is not None: kwargs["model_name"] = args.model_name + if args.transform == "dinov3": + kwargs["revision"] = args.revision + kwargs["image_side"] = args.image_side return cls(**kwargs) @@ -109,12 +112,25 @@ def main(): ) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--model-name", default=None) + parser.add_argument( + "--revision", default=None, help="DINO model and processor revision" + ) + parser.add_argument( + "--image-side", + type=int, + default=None, + help="DINO full-FOV letterbox square side; must be patch divisible", + ) parser.add_argument( "--device", default="cuda" if torch.cuda.is_available() else "cpu", ) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() + if args.transform != "dinov3" and ( + args.revision is not None or args.image_side is not None + ): + parser.error("--revision and --image-side apply only to dinov3") if args.output_keys is not None and len(args.input_keys) != len(args.output_keys): parser.error("--input-keys and --output-keys must have the same length") diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index fa14c663f1..1eb69691ef 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -42,6 +42,29 @@ _SHA256_RE = re.compile(r"[0-9a-f]{64}") +def _configure_strict_fp32(cfg: DictConfig) -> None: + """Opt in to full FP32 math, including CUDA/cuDNN backend precision.""" + enabled = OmegaConf.select(cfg, "runtime.strict_fp32", default=False) + if not isinstance(enabled, bool): + raise ValueError("runtime.strict_fp32 must be a boolean") + if not enabled: + return + precision = OmegaConf.select(cfg, "trainer.precision") + if 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("runtime.strict_fp32 requires NVIDIA_TF32_OVERRIDE=0") + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + if ( + torch.get_float32_matmul_precision() != "highest" + or torch.backends.cuda.matmul.allow_tf32 + or torch.backends.cudnn.allow_tf32 + ): + raise RuntimeError("strict FP32 backend configuration did not take effect") + + def _slurm_auto_requeue(cfg: DictConfig) -> bool: """Resolve and cross-check the configured Slurm requeue owner.""" @@ -122,6 +145,8 @@ def _instantiate_slurm_callbacks(cfg: DictConfig) -> List[Callback]: SaveOnlySignalCheckpoint( checkpoint_dir=str(checkpoint_dir), save_signal=signal.SIGUSR2, + train_batch_only=OmegaConf.select( + cfg, "runtime.slurm_signal_train_batch_only", default=False), ) ] @@ -496,6 +521,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: :return: A tuple with metrics and dict with all instantiated objects. """ mode = _validate_run_config(cfg) + _configure_strict_fp32(cfg) if mode == "train": configure_runner_wandb(cfg) diff --git a/egomimic/utils/slurm_process_identity.py b/egomimic/utils/slurm_process_identity.py new file mode 100644 index 0000000000..acae0a0387 --- /dev/null +++ b/egomimic/utils/slurm_process_identity.py @@ -0,0 +1,84 @@ +"""Authenticate one direct-exec trainer in the runner's Linux PID namespace. + +This stdlib-only boundary is shared by the generic runner and the Lightning +save callback. It never discovers PIDs, reads process environments, or signals +Slurm job steps. Different namespaces/nodes/cgroups fail closed. +""" + +import hashlib +import json +import os +from pathlib import Path +import signal +import socket + + +def source_identity(): + root = Path(__file__).resolve().parents[2] + files = ("egomimic/trainHydra.py", "egomimic/utils/slurm_requeue.py", + "egomimic/utils/slurm_process_identity.py", + "scripts/ice/ice_requeue_runner.py") + return {"root": str(root), "sha256": { + name: hashlib.sha256((root / name).read_bytes()).hexdigest() for name in files}} + + +def process_identity(pid): + if type(pid) is not int or pid < 1: + raise ValueError("Invalid trainer PID") + proc = Path("/proc") / str(pid) + status = (proc / "status").read_text().splitlines() + uid = next(line.split()[1:] for line in status if line.startswith("Uid:")) + # /proc//stat comm may contain spaces and parentheses. + fields = (proc / "stat").read_text().rsplit(")", 1)[1].split() + return {"pid": pid, "uid": [int(value) for value in uid], + "starttime": int(fields[19]), "node": socket.gethostname(), + "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip(), + "pid_namespace": os.readlink(proc / "ns/pid"), + "cgroup": (proc / "cgroup").read_text(), + "exe": str((proc / "exe").resolve(strict=True)), + "cwd": str((proc / "cwd").resolve(strict=True))} + + +def ready_record(nonce, job_id, restart): + return {"schema": "egoverse-trainer-pid-ready-v1", "nonce": nonce, + "job_id": job_id, "restart_count": restart, + "handler_installed": True, "process": process_identity(os.getpid()), + "source": source_identity()} + + +def signal_ready_trainer(path, *, pid, nonce, job_id, restart, parent_identity, + expected_source, expected_process, sig=signal.SIGUSR2): + """Authenticate the known child and signal through a PID-reuse-safe fd.""" + if sig != signal.SIGUSR2: + raise ValueError("Trainer PID transport only permits save-only SIGUSR2") + path = Path(path) + if path.is_symlink() or path.stat().st_uid != os.getuid(): + raise ValueError("Untrusted trainer readiness file ownership") + record = json.loads(path.read_text()) + if (record.get("schema") != "egoverse-trainer-pid-ready-v1" + or record.get("nonce") != nonce or record.get("job_id") != job_id + or record.get("restart_count") != restart + or record.get("handler_installed") is not True + or record.get("source") != expected_source): + raise ValueError("Trainer readiness authority mismatch") + if not hasattr(os, "pidfd_open") or not hasattr(signal, "pidfd_send_signal"): + raise RuntimeError("Linux pidfd support is required; no unsafe PID fallback") + descriptor = os.pidfd_open(pid) + try: + actual = process_identity(pid) + if actual != record.get("process") or actual != expected_process: + raise ValueError("Trainer process identity changed or was forged") + for key in ("uid", "node", "boot_id", "pid_namespace", "cgroup", "exe", "cwd"): + if actual[key] != parent_identity[key]: + raise ValueError(f"Trainer differs from runner {key}") + if actual["uid"] != [os.getuid()] * 4: + raise ValueError("Trainer UID differs from current user") + if actual["cwd"] != expected_source["root"]: + raise ValueError("Trainer working directory differs from pinned source") + if source_identity() != expected_source: + raise ValueError("Signal-delivery source changed") + signal.pidfd_send_signal(descriptor, sig) + return {"transport": "trainer-pid", "pid": pid, + "starttime": actual["starttime"], "signal": sig.name} + finally: + os.close(descriptor) diff --git a/egomimic/utils/slurm_requeue.py b/egomimic/utils/slurm_requeue.py index 0a4f7ea989..62007f84f3 100644 --- a/egomimic/utils/slurm_requeue.py +++ b/egomimic/utils/slurm_requeue.py @@ -98,6 +98,7 @@ def __init__( self, checkpoint_dir: str | Path, save_signal: signal.Signals = signal.SIGUSR2, + train_batch_only: bool = False, ) -> None: super().__init__() checkpoint_dir = Path(checkpoint_dir).expanduser() @@ -107,6 +108,10 @@ def __init__( raise ValueError("external save-only checkpoint signal must be SIGUSR2") self.checkpoint_dir = checkpoint_dir.resolve() self.save_signal = save_signal + if type(train_batch_only) is not bool: + raise TypeError("train_batch_only must be a boolean") + self.train_batch_only = train_batch_only + self._train_batch_start_step: int | None = None self._request_generation = 0 self._handled_generation = 0 self._save_index = 0 @@ -138,6 +143,22 @@ def setup( self._previous_handler = previous self._installed_handler = self._request_checkpoint signal.signal(self.save_signal, self._installed_handler) + if os.environ.get("ICE_CHECKPOINT_FORWARDING") == "trainer-pid": + from egomimic.utils.slurm_process_identity import ready_record + if _trainer.world_size != 1 or not _trainer.is_global_zero: + raise RuntimeError("trainer-pid requires a single-process trainer") + job_id = os.environ.get("SLURM_JOB_ID", "") + nonce = os.environ.get("ICE_TRAINER_NONCE", "") + ready = Path(os.environ.get("ICE_TRAINER_READY_FILE", "")) + attempt = Path(os.environ.get("ICE_ATTEMPT_DIR", "")) + if (_SLURM_JOB_ID_RE.fullmatch(job_id) is None + or re.fullmatch(r"[0-9a-f]{64}", nonce) is None + or not ready.is_absolute() or not attempt.is_absolute() + or ready.parent.resolve() != attempt.resolve()): + raise RuntimeError("invalid trainer-pid readiness authority") + _atomic_json_no_overwrite(ready, ready_record( + nonce, job_id, int(os.environ.get("SLURM_RESTART_COUNT", "0")))) + def teardown( self, @@ -173,6 +194,15 @@ def _checkpoint_path(self, trainer: Trainer) -> tuple[Path, str, int]: return path, job_id, restart_count def _service(self, trainer: Trainer, hook: str) -> None: + # Preserve the request until an optimizer update completes. A save + # during validation can fork a scheduled checkpoint at the same step, + # and does not prove exact mid-validation RNG/loop continuation. + if self.train_batch_only: + if hook != "on_train_batch_end": + return + if (self._train_batch_start_step is None + or int(trainer.global_step) <= self._train_batch_start_step): + return observed_generation = self._request_generation local_requested = observed_generation > self._handled_generation if not local_requested: @@ -254,6 +284,7 @@ def _service(self, trainer: Trainer, hook: str) -> None: "epoch": int(trainer.current_epoch), "global_step": int(trainer.global_step), "hook": hook, + "train_batch_only": self.train_batch_only, "job_id": job_id, "restart_count": restart_count, "save_signal": self.save_signal.name, @@ -297,6 +328,12 @@ def _service(self, trainer: Trainer, hook: str) -> None: self._save_index += 1 + def on_train_batch_start( + self, trainer: Trainer, _pl_module: LightningModule, + _batch: Any, _batch_idx: int, + ) -> None: + self._train_batch_start_step = int(trainer.global_step) + def on_train_batch_end( self, trainer: Trainer, diff --git a/scripts/benchmarks/profile_world_model.py b/scripts/benchmarks/profile_world_model.py new file mode 100644 index 0000000000..5f4fa87b4b --- /dev/null +++ b/scripts/benchmarks/profile_world_model.py @@ -0,0 +1,126 @@ +"""Synthetic FP32 forward/backward memory probe; never a real-data launch gate. + +Run in a scheduled GPU allocation. No optimizer updates, dataset access, or +checkpoints. Pose widths are explicit diagnostic inputs, not schema approval. +""" +import argparse +import hashlib +import inspect +import json +import time +from pathlib import Path + +import torch + +from egomimic.models.world_model import MoTWorldModel + + +def synthetic_batch(frames, cameras, patches, tracks, pose_dim, device): + result = {} + for name, count, dim in ( + ("image", cameras * patches, 768), ("track", tracks, 3), + ("proprio", 1, pose_dim), + ): + clock = torch.linspace(-3, 3, frames, device=device).repeat_interleave(count)[None] + n = frames * count + result[name] = dict( + value=torch.randn(1, n, dim, device=device), time=clock, + valid=torch.ones(1, n, dtype=torch.bool, device=device), + target_valid=torch.ones(1, n, dtype=torch.bool, device=device), + observed=clock < 0, + ) + if name != "proprio": + result[name].update( + xy=torch.rand(1, n, 2, device=device), + camera=(torch.arange(cameras, device=device).repeat_interleave(patches) + .repeat(frames)[None] if name == "image" + else torch.zeros(1, n, dtype=torch.long, device=device)), + ) + if name == "track": + result[name].update( + point_id=torch.arange(tracks, device=device).repeat(frames)[None], + confidence=torch.ones(1, n, device=device), + ) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--frames", type=int, nargs="+", required=True) + parser.add_argument("--cameras", type=int, required=True) + parser.add_argument("--side", type=int, required=True) + parser.add_argument("--tracks", type=int, required=True) + parser.add_argument("--pose-dims", type=json.loads, required=True) + parser.add_argument("--embodiment", required=True) + parser.add_argument("--query-chunk", type=int, default=512) + parser.add_argument("--width", type=int, default=384) + parser.add_argument("--depth", type=int, default=8) + parser.add_argument("--heads", type=int, default=6) + parser.add_argument("--device", choices=["cuda", "cpu"], default="cuda") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if min(args.frames) < 2 or min(args.cameras, args.tracks, args.side) < 1 or args.side % 16: + parser.error("Need >=2 frames, positive counts and side divisible by16") + if args.embodiment not in args.pose_dims: + parser.error("Active embodiment missing from diagnostic pose widths") + if args.device == "cuda" and torch.cuda.device_count() != 1: + parser.error("Require exactly one visible allocated GPU") + args.output.parent.mkdir(parents=True, exist_ok=True) + if args.output.exists(): + raise FileExistsError(args.output) + torch.set_num_threads(4) + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cuda = args.device == "cuda" + report = { + "scope": "Synthetic memory/backward diagnostic; not actual data, optimizer or reload proof", + "torch": torch.__version__, "device": torch.cuda.get_device_name(0) if cuda else "cpu", + "precision": "FP32; autocast off; TF32 off", "results": [], + "arguments": {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}, + "model_sha256": hashlib.sha256(Path(inspect.getfile(MoTWorldModel)).read_bytes()).hexdigest(), + } + for frames in args.frames: + torch.manual_seed(42) + patches = (args.side // 16) ** 2 + tokens = frames * (args.cameras * patches + args.tracks + 1) + if cuda: + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + row = {"frames": frames, "tokens": tokens, "status": "RUNNING"} + report["results"].append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + try: + model = MoTWorldModel( + args.pose_dims, width=args.width, depth=args.depth, heads=args.heads, + max_cameras=args.cameras, max_tracks=args.tracks, max_tokens=tokens, + attention_query_chunk_size=args.query_chunk, activation_checkpointing=True, + ).float().to(args.device) + batch = synthetic_batch(frames, args.cameras, patches, args.tracks, + args.pose_dims[args.embodiment], args.device) + loss = model.training_losses(batch, args.embodiment) + loss["total"].backward() + if cuda: + torch.cuda.synchronize() + if not all(torch.isfinite(v).all() for v in loss.values()): + raise ValueError("Nonfinite loss") + grads = [p.grad for p in model.parameters() if p.grad is not None] + if not grads or not all(torch.isfinite(g).all() for g in grads): + raise ValueError("Nonfinite or absent gradients") + row.update(status="PASS", parameters=sum(p.numel() for p in model.parameters()), + loss={k: float(v.detach()) for k, v in loss.items()}, + elapsed_seconds=time.perf_counter()-started, + peak_allocated_bytes=torch.cuda.max_memory_allocated() if cuda else None, + peak_reserved_bytes=torch.cuda.max_memory_reserved() if cuda else None) + del model, batch, loss, grads + except Exception as error: + row.update(status="FAILED", error=str(error), elapsed_seconds=time.perf_counter()-started) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(row), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/ice/ice_requeue_runner.py b/scripts/ice/ice_requeue_runner.py index dc1789bb1d..ce7b82c18b 100644 --- a/scripts/ice/ice_requeue_runner.py +++ b/scripts/ice/ice_requeue_runner.py @@ -19,12 +19,14 @@ import json import os import re +import secrets import shlex import shutil import signal import subprocess import sys import time +from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path from typing import Any, Sequence @@ -645,11 +647,13 @@ def parser() -> argparse.ArgumentParser: ) result.add_argument( "--checkpoint-forwarding", - choices=("slurm-steps",), + choices=("slurm-steps", "trainer-pid"), required=True, - help="explicit checkpoint signal transport; slurm-steps excludes the batch shell", + help="explicit save-only transport; trainer-pid authenticates the direct child", ) result.add_argument("--checkpoint-grace-seconds", type=int, default=240) + result.add_argument("--allocation-boundary-lead-seconds", type=int, + help="opt-in relative timer from native SLURM_JOB_END_TIME; no step-wide signal") result.add_argument("--poll-seconds", type=float, default=2.0) result.add_argument("--max-restarts", type=int, default=32) result.add_argument( @@ -666,6 +670,7 @@ def parser() -> argparse.ArgumentParser: result.add_argument("--scontrol", default="scontrol") result.add_argument("--scancel", default="scancel") result.add_argument("--completion-sentinel", type=Path) + result.add_argument("--target-global-step", type=int, help="refuse completion before this exact optimizer step") result.add_argument("--dry-run", action="store_true") result.add_argument("command", nargs=argparse.REMAINDER) return result @@ -874,6 +879,11 @@ def validated_completion_sentinel(path: Path) -> dict[str, Any]: def main(argv: Sequence[str] | None = None) -> int: + with ExitStack() as cleanup: + return _main(argv, cleanup) + + +def _main(argv: Sequence[str] | None, cleanup: ExitStack) -> int: args = parser().parse_args(argv) if args.command and args.command[0] == "--": args.command = args.command[1:] @@ -884,6 +894,8 @@ def main(argv: Sequence[str] | None = None) -> int: # the authenticated no-work requeue path below; TERM/INT take priority. boundary_requested_at: float | None = None boundary_processed = False + boundary_trigger = None + allocation_boundary = None termination_signal: signal.Signals | None = None handled_signals = {signal.SIGUSR1, signal.SIGTERM, signal.SIGINT} launch_gate_write_fd: int | None = None @@ -900,15 +912,16 @@ def close_launch_gate() -> None: pass def handler(received: int, _frame: object) -> None: - nonlocal boundary_requested_at, boundary_processed, termination_signal + nonlocal boundary_requested_at, boundary_processed, termination_signal, boundary_trigger if not workload_launch_committed: # Closing the one-byte launch gate makes a signal racing with # Popen fail before the wrapper can exec the real workload. close_launch_gate() received_signal = signal.Signals(received) - if received_signal == signal.SIGUSR1: + if received_signal in (signal.SIGUSR1, signal.SIGALRM): if boundary_requested_at is None and not boundary_processed: boundary_requested_at = time.time() + boundary_trigger = "allocation-timer" if received_signal == signal.SIGALRM else "SIGUSR1" else: termination_signal = received_signal @@ -928,6 +941,35 @@ def handler(received: int, _frame: object) -> None: ready_path, {"pid": os.getpid(), "signal_handlers_ready": True} ) + if args.allocation_boundary_lead_seconds is not None: + lead = args.allocation_boundary_lead_seconds + end_text = os.environ.get("SLURM_JOB_END_TIME", "") + now = time.time() + if lead < 1 or re.fullmatch(r"[0-9]+", end_text) is None: + raise SystemExit("allocation timer requires positive lead and native SLURM_JOB_END_TIME") + end_time = int(end_text) + if end_time <= now: + raise SystemExit("SLURM_JOB_END_TIME is stale; refusing allocation timer") + delay = max(0., end_time - lead - now) + allocation_boundary = {"source": "SLURM_JOB_END_TIME", "end_unix": end_time, + "lead_seconds": lead, "relative_delay_seconds": delay} + if not args.dry_run: + if (signal.getitimer(signal.ITIMER_REAL) != (0., 0.) + or signal.SIGALRM in signal.sigpending()): + raise SystemExit("allocation timer refuses an existing process timer or pending alarm") + old_mask = signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGALRM}) + cleanup.callback(signal.pthread_sigmask, signal.SIG_SETMASK, old_mask) + previous_alarm = signal.signal(signal.SIGALRM, handler) + cleanup.callback(signal.signal, signal.SIGALRM, previous_alarm) + cleanup.callback(signal.setitimer, signal.ITIMER_REAL, 0.) + if delay == 0: + handler(signal.SIGALRM, None) + else: + # Relative kernel timer keeps counting during slow discovery. + # Child processes do not inherit setitimer timers. + signal.setitimer(signal.ITIMER_REAL, delay) + signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGALRM}) + if not args.command: raise SystemExit("a child command is required after --") if args.checkpoint_grace_seconds < 1 or args.poll_seconds <= 0: @@ -945,6 +987,23 @@ def handler(received: int, _frame: object) -> None: raise SystemExit( "slurm-step checkpoint forwarding requires --checkpoint-signal USR2" ) + pid_transport = args.checkpoint_forwarding == "trainer-pid" + pid_api = None + if pid_transport: + # Import only after the early signal latch. This module is stdlib-only. + from egomimic.utils import slurm_process_identity as pid_api + if args.requeue_owner != "runner" or args.checkpoint_signal != signal.SIGUSR2: + raise SystemExit("trainer-pid requires runner ownership and SIGUSR2") + if Path(args.command[0]).resolve() != Path(sys.executable).resolve(): + raise SystemExit("trainer-pid requires a direct child using the runner Python executable") + if Path.cwd().resolve() != Path(pid_api.source_identity()["root"]): + raise SystemExit("trainer-pid requires the pinned source working directory") + for key in ("SLURM_NTASKS", "SLURM_JOB_NUM_NODES", "WORLD_SIZE"): + if int(os.environ.get(key, "1")) != 1: + raise SystemExit("trainer-pid supports exactly one task/node/process") + if args.target_global_step is not None: + if args.target_global_step < 1 or args.completion_sentinel is None: + raise SystemExit("target-global-step requires a positive step and completion-sentinel") if args.require_initial_checkpoint and args.initial_checkpoint is None: raise SystemExit( "--require-initial-checkpoint requires an exact --initial-checkpoint" @@ -991,7 +1050,10 @@ def handler(received: int, _frame: object) -> None: ) if completion_sentinel is not None and completion_sentinel.exists(): - validated_completion_sentinel(completion_sentinel) + completed = validated_completion_sentinel(completion_sentinel) + if (args.target_global_step is not None + and completed["checkpoint"]["global_step"] != args.target_global_step): + raise SystemExit("completion sentinel does not match target-global-step") print( json.dumps( {"status": "ALREADY_COMPLETE", "sentinel": str(completion_sentinel)} @@ -999,6 +1061,7 @@ def handler(received: int, _frame: object) -> None: ) return 0 + child = None lock_stream = (state_dir / "runner.lock").open("a+") try: try: @@ -1009,7 +1072,7 @@ def handler(received: int, _frame: object) -> None: ) from exc scancel = None - if not args.dry_run: + if not args.dry_run and not pid_transport: scancel = resolve_executable(args.scancel) scontrol = None if args.requeue_owner == "runner" and not args.dry_run: @@ -1134,6 +1197,8 @@ def handler(received: int, _frame: object) -> None: "started_at_unix": time.time(), "status": "DRY_RUN" if args.dry_run else "STARTING", } + if allocation_boundary is not None: + record["allocation_boundary"] = allocation_boundary atomic_json(attempt_dir / "attempt.json", record) metadata_json = ( @@ -1158,6 +1223,14 @@ def handler(received: int, _frame: object) -> None: else "0", } ) + if pid_transport: + trainer_ready = attempt_dir / "trainer-ready.json" + trainer_nonce = secrets.token_hex(32) + expected_source = pid_api.source_identity() + parent_identity = pid_api.process_identity(os.getpid()) + child_env.update({"ICE_CHECKPOINT_FORWARDING": "trainer-pid", + "ICE_TRAINER_READY_FILE": str(trainer_ready), + "ICE_TRAINER_NONCE": trainer_nonce}) if completion_sentinel is not None: child_env["ICE_COMPLETION_SENTINEL"] = str(completion_sentinel) if args.dry_run: @@ -1195,6 +1268,8 @@ def handler(received: int, _frame: object) -> None: finally: os.close(gate_read_fd) + if pid_transport: + expected_process = pid_api.process_identity(child.pid) if termination_signal is None and boundary_requested_at is None: try: written = os.write(gate_write_fd, b"G") @@ -1421,7 +1496,8 @@ def forward_termination(received: signal.Signals) -> None: event_path, { "event": "signal", - "signal": "SIGUSR1", + "signal": "SIGALRM" if boundary_trigger == "allocation-timer" else "SIGUSR1", + "boundary_trigger": boundary_trigger, "checkpoint_signal": args.checkpoint_signal.name, "baseline": baseline.as_dict() if baseline else None, "at_unix": signal_at, @@ -1429,51 +1505,73 @@ def forward_termination(received: signal.Signals) -> None: ) if termination_signal is not None or child.poll() is not None: continue - assert scancel is not None - try: - signal_result = forward_checkpoint_to_slurm_steps( - scancel, - job_id, - args.checkpoint_signal, - ) - except OSError as exc: - record.update( - { - "status": "CHECKPOINT_SIGNAL_FORWARD_FAILED", - "checkpoint_signal_argv": [ - scancel, - "--signal=USR2", - job_id, - ], - "checkpoint_signal_error": str(exc), - } - ) - atomic_json(attempt_dir / "attempt.json", record) - return 75 - if signal_result.returncode != 0: - record.update( + if pid_transport: + ready_deadline = time.monotonic() + args.checkpoint_grace_seconds + while not trainer_ready.exists() and time.monotonic() < ready_deadline: + if termination_signal is not None or child.poll() is not None: + break + time.sleep(args.poll_seconds) + if termination_signal is not None: + continue + try: + delivery = pid_api.signal_ready_trainer( + trainer_ready, pid=child.pid, nonce=trainer_nonce, + job_id=job_id, restart=restart_count, + parent_identity=parent_identity, expected_source=expected_source, + expected_process=expected_process, sig=args.checkpoint_signal) + except (OSError, ValueError, RuntimeError) as exc: + record.update({"status": "CHECKPOINT_SIGNAL_FORWARD_FAILED", + "checkpoint_signal_error": str(exc)}) + atomic_json(attempt_dir / "attempt.json", record) + return 75 + append_event(event_path, {"event": "checkpoint_signal_forwarded", + **delivery, "at_unix": time.time()}) + else: + assert scancel is not None + try: + signal_result = forward_checkpoint_to_slurm_steps( + scancel, + job_id, + args.checkpoint_signal, + ) + except OSError as exc: + record.update( + { + "status": "CHECKPOINT_SIGNAL_FORWARD_FAILED", + "checkpoint_signal_argv": [ + scancel, + "--signal=USR2", + job_id, + ], + "checkpoint_signal_error": str(exc), + } + ) + atomic_json(attempt_dir / "attempt.json", record) + return 75 + if signal_result.returncode != 0: + record.update( + { + "status": "CHECKPOINT_SIGNAL_FORWARD_FAILED", + "checkpoint_signal_argv": [ + scancel, + "--signal=USR2", + job_id, + ], + "checkpoint_signal_returncode": signal_result.returncode, + "checkpoint_signal_stdout": signal_result.stdout, + "checkpoint_signal_stderr": signal_result.stderr, + } + ) + atomic_json(attempt_dir / "attempt.json", record) + return signal_result.returncode or 75 + append_event( + event_path, { - "status": "CHECKPOINT_SIGNAL_FORWARD_FAILED", - "checkpoint_signal_argv": [ - scancel, - "--signal=USR2", - job_id, - ], - "checkpoint_signal_returncode": signal_result.returncode, - "checkpoint_signal_stdout": signal_result.stdout, - "checkpoint_signal_stderr": signal_result.stderr, - } + "event": "checkpoint_signal_forwarded", + "argv": [scancel, "--signal=USR2", job_id], + "at_unix": time.time(), + }, ) - atomic_json(attempt_dir / "attempt.json", record) - return signal_result.returncode or 75 - append_event( - event_path, - { - "event": "checkpoint_signal_forwarded", - "argv": [scancel, "--signal=USR2", job_id], - "at_unix": time.time(), - }, - ) # Give an early queued boundary a complete save window after # forwarding; validator startup time must not consume grace. deadline = time.time() + args.checkpoint_grace_seconds @@ -1599,6 +1697,13 @@ def forward_termination(received: signal.Signals) -> None: ) atomic_json(attempt_dir / "attempt.json", record) return 73 + if (args.target_global_step is not None + and terminal.global_step != args.target_global_step): + record.update({"status": "COMPLETE_REFUSED_TARGET_STEP", + "target_global_step": args.target_global_step, + "checkpoint": terminal.as_dict()}) + atomic_json(attempt_dir / "attempt.json", record) + return 78 high_water = commit_checkpoint_state( high_water_path=high_water_path, observation_index_path=observation_index_path, @@ -1636,6 +1741,17 @@ def forward_termination(received: signal.Signals) -> None: return 128 - code if code < 0 else code time.sleep(args.poll_seconds) finally: + # The PID backend owns this direct child's new session. Never invoke + # scancel or discover other processes, even on a validation failure. + if pid_transport and child is not None and child.poll() is None: + try: + os.killpg(child.pid, signal.SIGTERM) + child.wait(timeout=5) + except ProcessLookupError: + pass + except subprocess.TimeoutExpired: + os.killpg(child.pid, signal.SIGKILL) + child.wait() lock_stream.close() diff --git a/tests/test_dinov3_pose_probe.py b/tests/test_dinov3_pose_probe.py new file mode 100644 index 0000000000..b09526f992 --- /dev/null +++ b/tests/test_dinov3_pose_probe.py @@ -0,0 +1,74 @@ +"""Probe isolation, padding and leakage controls; no pretrained-quality claims.""" + +import pytest +import torch + +from egomimic.scripts.embedding_process.dinov3_pose_probe import ( + PoseHeatmapProbe, + episode_split, + heatmap_targets, + mean_pose_baseline, + objective, +) + + +def test_image_only_probe_finds_informative_patch_and_excludes_padding(): + model = PoseHeatmapProbe(3, 1) + with torch.no_grad(): + model.head.weight.copy_(torch.tensor([[1.0, -1.0, 0.0]])) + features = torch.tensor([[[2.0, -2.0, 0.0], [0.0, 2.0, -2.0], [1.0, -1.0, 0.0]]]) + valid = torch.tensor([False, True, True]) + logits = model(features, valid) + assert logits.argmax(-1).item() == 2 + assert logits.dtype == torch.float32 + assert sum(p.numel() for p in PoseHeatmapProbe(768, 42).parameters()) == 32256 + assert not features.requires_grad + + +def test_targets_normalize_only_over_real_image_and_ignore_unlabeled_points(): + centers = torch.tensor([[0.0, 0.0], [10.0, 0.0], [20.0, 0.0]]) + target = heatmap_targets( + torch.tensor([[[10.0, 0.0], [0.0, 0.0]]]), + centers, + torch.tensor([False, True, True]), + 2.0, + ) + assert torch.equal(target[..., 0], torch.zeros(1, 2)) + torch.testing.assert_close(target.sum(-1), torch.ones(1, 2)) + logits = torch.randn(1, 2, 3, requires_grad=True) + loss = objective(logits, target, torch.tensor([[True, False]])) + loss.backward() + assert logits.grad[0, 0].abs().sum() > 0 + assert torch.equal(logits.grad[0, 1], torch.zeros(3)) + + +def test_empty_labels_produce_zero_loss_and_gradients(): + logits = torch.randn(1, 2, 3, requires_grad=True) + loss = objective( + logits, + torch.softmax(torch.randn_like(logits), -1), + torch.zeros(1, 2, dtype=torch.bool), + ) + loss.backward() + assert loss.item() == 0 + assert torch.equal(logits.grad, torch.zeros_like(logits)) + + +def test_episode_split_and_training_only_baseline_reject_image_leakage(): + metadata = [ + {"episode": e, "image_sha256": str(i)} + for i, e in enumerate(["a", "a", "b", "c"]) + ] + train, test = episode_split(metadata, "b") + assert train == [0, 1, 3] and test == [2] + xy = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]], [[1000.0, 2000.0]], [[5.0, 6.0]]]) + baseline = mean_pose_baseline(xy[train], torch.ones(3, 1, dtype=torch.bool)) + torch.testing.assert_close(baseline, torch.tensor([[3.0, 4.0]])) + metadata[2]["image_sha256"] = "0" + with pytest.raises(ValueError, match="Identical image bytes"): + episode_split(metadata, "b") + + +def test_missing_training_landmark_is_rejected(): + with pytest.raises(ValueError, match="training label"): + mean_pose_baseline(torch.zeros(2, 1, 2), torch.zeros(2, 1, dtype=torch.bool)) diff --git a/tests/test_dinov3_preprocessing.py b/tests/test_dinov3_preprocessing.py new file mode 100644 index 0000000000..3098611ac5 --- /dev/null +++ b/tests/test_dinov3_preprocessing.py @@ -0,0 +1,248 @@ +"""Preprocessing/correspondence contracts; fake features are not DINO quality tests.""" + +import hashlib +import json +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +from egomimic.scripts.embedding_process.dinov3_correspondence import ( + cached_features, + eligible, + load_manifest, + match_point, + patch_geometry, + run, +) +from egomimic.scripts.embedding_process.dinov3_embedding import DINOv3ImageEmbedding + + +def test_projected_reference_validity_is_not_visibility(tmp_path): + point = {"reference_valid": True, "visible": False} + assert eligible(point, "projected_landmark") + assert not eligible(point, "visible_manual") + path = tmp_path / "manifest.json" + path.write_text(json.dumps({"reference_mode": "projected_landmark"})) + with pytest.raises(ValueError, match="projection_provenance"): + load_manifest(path) + + +def test_feature_reuse_rejects_changed_identity(tmp_path): + path = tmp_path / "results.json" + path.write_text( + json.dumps({"model": "wrong", "revision": "a" * 40, "precision": "float32"}) + ) + with pytest.raises(ValueError, match="model/revision/precision"): + cached_features([path], "wanted", "a" * 40) + + +class Processor: + image_mean = [0.5, 0.5, 0.5] + image_std = [0.25, 0.25, 0.25] + + def __call__(self, images, return_tensors, **kwargs): + assert kwargs == {"do_resize": False, "do_center_crop": False} + pixels = torch.from_numpy(np.stack(images)).permute(0, 3, 1, 2).float() + return {"pixel_values": (pixels / 255 - 0.5) / 0.25} + + +class Model: + config = SimpleNamespace( + image_size=224, patch_size=16, num_register_tokens=4, _commit_hash="test" + ) + + def __call__(self, pixel_values): + n = pixel_values.shape[-1] * pixel_values.shape[-2] // 256 + tokens = torch.arange(n + 5).float().reshape(1, -1, 1) + return SimpleNamespace( + last_hidden_state=tokens.expand(pixel_values.shape[0], -1, 2) + ) + + +def extractor(side): + x = DINOv3ImageEmbedding( + ["image"], image_side=side, device="cpu", dtype=torch.float32 + ) + x.processor, x.model = Processor(), Model() + return x + + +@pytest.mark.parametrize( + "side,patches", [(224, 196), (336, 441), (448, 784), (640, 1600)] +) +def test_actual_processed_grid_not_static_model_config(side, patches): + x = extractor(side) + image = np.zeros((480, 640, 3), np.uint8) + features = x._embed_batch([image]) + assert features.shape == (1, patches, 2) + assert features[0, 0, 0] == 5 + assert features[0, -1, 0] == patches + 4 + assert x.last_preprocessing["patch_grid"] == [side // 16] * 2 + assert x.last_preprocessing["frames"][0]["original_hw"] == [480, 640] + + +def test_same_fov_padding_and_native_coordinate_roundtrip(): + x = extractor(640) + image = np.full((480, 640, 3), 255, np.uint8) + inputs, geometry = x._prepare([image]) + assert geometry[0] == { + "original_hw": [480, 640], + "resized_hw": [480, 640], + "padding_top_left": [80, 0], + } + assert torch.all(inputs["pixel_values"][0, :, 80:560] == 2) + x._embed_batch([image]) + centers, valid = patch_geometry(x.last_preprocessing) + assert valid.sum() == 30 * 40 + assert np.array_equal(centers[5 * 40], [7.5, 7.5]) + assert not valid[: 5 * 40].any() + + +def test_malformed_output_and_fractional_patch_rejected(): + x = extractor(225) + with pytest.raises(ValueError, match="divisible"): + x._embed_batch([np.zeros((32, 32, 3), np.uint8)]) + x = extractor(224) + x.model = lambda **kwargs: SimpleNamespace(last_hidden_state=torch.zeros(1, 196, 2)) + x.model.config = Model.config + with pytest.raises(ValueError, match="token count"): + x._embed_batch([np.zeros((32, 32, 3), np.uint8)]) + + +def test_identical_descriptor_grid_matches_same_native_point(): + x = extractor(64) + x._embed_batch([np.zeros((64, 64, 3), np.uint8)]) + features = np.eye(16, dtype=np.float32) + pred, cosine = match_point( + features, features, x.last_preprocessing, x.last_preprocessing, [23.5, 39.5] + ) + assert np.array_equal(pred, [23.5, 39.5]) + assert cosine == 1 + + +def test_padding_cannot_win_match_even_with_larger_cosine(): + x = extractor(64) + x._embed_batch([np.zeros((32, 64, 3), np.uint8)]) + source = np.ones((16, 2), np.float32) + target = np.tile([1.0, 0.0], (16, 1)).astype(np.float32) + target[:4] = [1.0, 1.0] + pred, _ = match_point( + source, target, x.last_preprocessing, x.last_preprocessing, [23.5, 7.5] + ) + assert 0 <= pred[1] < 32 + + +def test_setup_pins_processor_and_model_to_same_revision(monkeypatch): + calls = [] + + class Factory: + @staticmethod + def from_pretrained(name, **kwargs): + calls.append((name, kwargs)) + return SimpleNamespace( + to=lambda device: SimpleNamespace(eval=lambda: None), + requires_grad_=lambda value: None, + ) + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(AutoModel=Factory, AutoImageProcessor=Factory), + ) + x = DINOv3ImageEmbedding(["image"], device="cpu", revision="a" * 40) + x.setup() + assert len(calls) == 2 + assert all(kwargs["revision"] == "a" * 40 for _, kwargs in calls) + + +def test_source_descriptor_touching_padding_is_rejected(): + x = extractor(64) + x._embed_batch([np.zeros((32, 64, 3), np.uint8)]) + features = np.ones((16, 2), np.float32) + with pytest.raises(ValueError, match="padded or mixed"): + match_point( + features, features, x.last_preprocessing, x.last_preprocessing, [23.5, 0] + ) + + +def test_boolean_resolution_rejected(): + with pytest.raises(ValueError, match="positive integer"): + extractor(True) + + +@pytest.mark.parametrize("features_only", [False, True]) +@pytest.mark.parametrize("reference_mode", ["visible_manual", "projected_landmark"]) +def test_screen_writes_reusable_features_and_does_not_fake_wrong_finger_metric( + tmp_path, monkeypatch, reference_mode, features_only +): + def fake_setup(self): + self.processor, self.model = Processor(), Model() + self.model.config = SimpleNamespace(**vars(Model.config)) + self.model.config._commit_hash = self.revision + + monkeypatch.setattr(DINOv3ImageEmbedding, "setup", fake_setup) + frames = [] + for key in ["a", "b"]: + path = tmp_path / f"{key}.png" + Image.fromarray(np.zeros((32, 32, 3), np.uint8)).save(path) + frames.append( + { + "id": key, + "path": str(path), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "episode": "synthetic", + "points": { + "index": { + "xy": [15.5, 15.5], + "visible": reference_mode == "visible_manual", + "reference_valid": True, + "annotation_source": "synthetictest", + } + }, + } + ) + path = tmp_path / "manifest.json" + path.write_text( + json.dumps( + { + "model": "fake", + "reference_mode": reference_mode, + "projection_provenance": {"synthetic": True}, + "revision": "a" * 40, + "frames": frames, + "pairs": [["a", "b"]], + } + ) + ) + if features_only: + def reject_matching(*args, **kwargs): + raise AssertionError("Feature extraction must not require point matching") + monkeypatch.setattr( + "egomimic.scripts.embedding_process.dinov3_correspondence.match_point", + reject_matching, + ) + result = run(path, tmp_path / "output", [32, 64], "cpu", features_only=features_only) + if features_only: + assert result["method"] == "feature_extraction_only" + assert all(r["mean_error_native_px"] is None and r["pck"] is None for r in result["resolutions"]) + assert len(result["resolutions"]) == 2 + for r in result["resolutions"]: + assert r["wrong_labeled_point_fraction"] is None + for artifact in r["feature_caches"].values(): + file = tmp_path / "output" / artifact["path"] + assert hashlib.sha256(file.read_bytes()).hexdigest() == artifact["sha256"] + with np.load(file, allow_pickle=False) as cache: + assert cache["features"].shape[0] == r["patch_tokens"] + with pytest.raises(FileExistsError, match="Preserve"): + run(path, tmp_path / "output", [32], "cpu") + reused = run( + path, tmp_path / "reused", [32], "cpu", [tmp_path / "output/results.json"], features_only=features_only + ) + assert reused["reference_mode"] == reference_mode + assert reused["resolutions"][0]["reused_frames"] == 2 + assert reused["resolutions"][0]["extract_median_seconds_including_transfer"] is None + assert reused["resolutions"][0]["matches"] == result["resolutions"][0]["matches"] diff --git a/tests/test_dinov3_segmentation_probe.py b/tests/test_dinov3_segmentation_probe.py new file mode 100644 index 0000000000..0b35e26559 --- /dev/null +++ b/tests/test_dinov3_segmentation_probe.py @@ -0,0 +1,48 @@ +import numpy as np +import pytest +import torch + +from egomimic.scripts.embedding_process.dinov3_segmentation_probe import ( + boundary_f1, native_probability, objective, overlap, patch_targets, rasterize, +) + + +def test_letterbox_targets_ignore_padding_and_unknown(): + prep = {"frames": [{"resized_hw": [16, 32], "padding_top_left": [8, 0], + "original_hw": [16, 32]}], + "processed_hw": [32, 32], "patch_grid": [2, 2]} + mask = np.zeros((16, 32), bool) + mask[:, :16] = True + valid = np.ones_like(mask) + valid[:, 16:] = False + targets = patch_targets(mask, valid, prep) + assert torch.equal(targets[:, 0], torch.zeros(4)) + assert torch.equal(targets[:, 1], torch.tensor([.5, 0, .5, 0])) + # Constant log-odds must survive both letterbox removal and native resize. + logits = torch.tensor([[0., 2.]]).repeat(4, 1) + probability = native_probability(logits, prep) + assert probability.shape == (16, 32) + assert np.allclose(probability, torch.sigmoid(torch.tensor(2.)).item()) + + +def test_mask_holes_uncertainty_and_metrics(): + frame = {"hand_polygons": [[[8, 8], [30, 8], [30, 30], [8, 30]]], + "background_holes": [[[15, 15], [20, 15], [20, 20], [15, 20]]], + "ignore_polygons": [[[0, 0], [3, 0], [3, 39], [0, 39]]]} + mask, certain, valid = rasterize(frame, (40, 40), 2) + assert mask[12, 12] and not mask[17, 17] + assert not certain[20, 2] and not valid[8, 8] + assert overlap(mask, mask, valid)["iou"] == 1 + assert boundary_f1(mask, mask, certain, 2) == 1 + assert overlap(np.zeros_like(mask), mask, valid)["dice"] == 0 + + +def test_unknown_labels_have_zero_gradient_and_reject_empty(): + logits = torch.zeros(1, 2, 2, requires_grad=True) + targets = torch.tensor([[[1., 0.], [0., 0.]]]) + loss = objective(logits, targets, torch.ones(2)) + loss.backward() + assert torch.count_nonzero(logits.grad[0, 1]) == 0 + assert torch.count_nonzero(logits.grad[0, 0]) == 2 + with pytest.raises(ValueError, match="No labeled"): + objective(logits, torch.zeros_like(targets), torch.ones(2)) diff --git a/tests/test_ice_requeue_runner.py b/tests/test_ice_requeue_runner.py index c652e216f3..aec55ad96c 100644 --- a/tests/test_ice_requeue_runner.py +++ b/tests/test_ice_requeue_runner.py @@ -14,6 +14,7 @@ import tempfile import time import unittest +import pytest from pathlib import Path from unittest import mock @@ -241,6 +242,12 @@ def test_checkpoint_forwarding_is_exact_scancel_and_never_killpg(self): class RunnerSafetyTest(unittest.TestCase): + def setUp(self): + # Scheduled tests must not inherit the enclosing allocation identity. + patcher = mock.patch.dict(os.environ, {"SLURM_JOB_ID": ""}) + patcher.start() + self.addCleanup(patcher.stop) + def test_signal_handlers_are_installed_before_child_spawn(self): with tempfile.TemporaryDirectory() as raw: root = Path(raw) @@ -1474,3 +1481,108 @@ def test_completion_sentinel_contains_terminal_checkpoint(self): if __name__ == "__main__": unittest.main() + + +@pytest.mark.parametrize("deadline", [False, True]) +def test_trainer_pid_backend_requeue_and_target_guard(tmp_path, deadline): + """Real child signal transport; fake scontrol never mutates the allocation.""" + if not hasattr(os, "pidfd_open"): + import pytest + pytest.skip("Linux pidfd required") + source = MODULE_PATH.parents[2] + state = tmp_path / "state" + checkpoints = tmp_path / "checkpoints" + checkpoints.mkdir() + validator = make_validator(tmp_path) + control = tmp_path / "scontrol" + control.write_text("#!/bin/sh\nprintf '%s\\n' \"$*\" > \"$REQUEUE_LOG\"\n") + control.chmod(0o700) + child = tmp_path / "child.py" + child.write_text('''import json, os, signal, time +from pathlib import Path +from egomimic.utils.slurm_process_identity import ready_record +root = Path(os.environ["TEST_CHECKPOINTS"]) +if os.environ["ICE_RESTART_COUNT"] == "1": + assert os.environ["ICE_RESUME_GLOBAL_STEP"] == os.environ.get("EXPECTED_STEP", "2") + (root / "step-3.ckpt").write_text(json.dumps({"global_step":3})) +else: + def save(*_): + (root / "step-2.ckpt").write_text(json.dumps({"global_step":2})) + signal.signal(signal.SIGUSR2, save) + Path(os.environ["ICE_TRAINER_READY_FILE"]).write_text(json.dumps(ready_record( + os.environ["ICE_TRAINER_NONCE"], os.environ["SLURM_JOB_ID"], 0))) + if not os.environ.get("TEST_DEADLINE"): + os.kill(os.getppid(), signal.SIGUSR1) + time.sleep(30) +''') + env = os.environ.copy() + env.update(SLURM_JOB_ID="123456", SLURM_RESTART_COUNT="0", SLURM_NTASKS="1", + SLURM_JOB_NUM_NODES="1", WORLD_SIZE="1", TEST_CHECKPOINTS=str(checkpoints), + REQUEUE_LOG=str(tmp_path / "requeue.log"), + PATH=str(Path(sys.executable).parent)+os.pathsep+env["PATH"]) + args = [sys.executable, str(MODULE_PATH), "--state-dir", str(state), + "--checkpoint-glob", str(checkpoints / "*.ckpt"), + "--checkpoint-validator", str(validator), "--checkpoint-signal", "USR2", + "--checkpoint-forwarding", "trainer-pid", "--checkpoint-grace-seconds", "10", + "--poll-seconds", "0.02", "--requeue-owner", "runner", + "--confirm-child-requeue-disabled", "--scontrol", str(control), + "--scancel", "/must-never-resolve-scancel", "--completion-sentinel", str(tmp_path / "COMPLETE.json"), + "--target-global-step", "3", "--", sys.executable, str(child)] + if deadline: + env["TEST_DEADLINE"] = "1" + env["SLURM_JOB_END_TIME"] = str(int(time.time()) + 5) + args[args.index("--"):args.index("--")] = ["--allocation-boundary-lead-seconds", "2"] + first = subprocess.run(args, cwd=source, env=env, capture_output=True, text=True, timeout=40) + assert first.returncode == 0, first.stderr + assert (tmp_path / "requeue.log").read_text().strip() == "requeue 123456" + event = (state / "requeue/job-123456-restart-000/events.jsonl").read_text() + assert '"transport": "trainer-pid"' in event + if deadline: + assert '"boundary_trigger": "allocation-timer"' in event + assert not (tmp_path / "COMPLETE.json").exists() + env["SLURM_RESTART_COUNT"] = "1" + if deadline: + env["SLURM_JOB_END_TIME"] = str(int(time.time()) + 30) + second = subprocess.run(args, cwd=source, env=env, capture_output=True, text=True, timeout=40) + assert second.returncode == 0, second.stderr + assert json.loads((tmp_path / "COMPLETE.json").read_text())["checkpoint"]["global_step"] == 3 + # A previously complete shorter run cannot satisfy a larger target. + args[args.index("--target-global-step")+1] = "4" + refused = subprocess.run(args, cwd=source, env=env, capture_output=True, text=True, timeout=10) + assert refused.returncode != 0 and "target-global-step" in refused.stderr + (tmp_path / "COMPLETE.json").unlink() + env["EXPECTED_STEP"] = "3" + refused = subprocess.run(args, cwd=source, env=env, capture_output=True, text=True, timeout=40) + assert refused.returncode == 78 and not (tmp_path / "COMPLETE.json").exists() + + +@pytest.mark.parametrize("end", [None, "", "invalid", "0", "1"]) +def test_allocation_timer_rejects_missing_invalid_stale_end(tmp_path, end): + env = os.environ.copy() + if end is None: + env.pop("SLURM_JOB_END_TIME", None) + else: + env["SLURM_JOB_END_TIME"] = end + args = [sys.executable, str(MODULE_PATH), "--state-dir", str(tmp_path / "state"), + "--checkpoint-glob", str(tmp_path / "*.ckpt"), "--allow-unvalidated-checkpoints", + "--checkpoint-signal", "USR2", "--checkpoint-forwarding", "slurm-steps", + "--requeue-owner", "child", "--allocation-boundary-lead-seconds", "2", + "--dry-run", "--", "/bin/true"] + result = subprocess.run(args, env=env, capture_output=True, text=True, timeout=10) + assert result.returncode != 0 and "SLURM_JOB_END_TIME" in result.stderr + assert not (tmp_path / "state").exists() + + +def test_allocation_timer_cleanup_on_early_failure(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_END_TIME", str(int(time.time()) + 120)) + old_alarm = signal.getsignal(signal.SIGALRM) + old_mask = signal.pthread_sigmask(signal.SIG_BLOCK, set()) + with pytest.raises(SystemExit, match="checkpoint-validator"): + MODULE.main(["--state-dir", str(tmp_path / "state"), + "--checkpoint-glob", str(tmp_path / "*.ckpt"), + "--checkpoint-signal", "USR2", "--checkpoint-forwarding", "slurm-steps", + "--requeue-owner", "child", "--allocation-boundary-lead-seconds", "10", + "--", "/bin/true"]) + assert signal.getitimer(signal.ITIMER_REAL) == (0., 0.) + assert signal.getsignal(signal.SIGALRM) == old_alarm + assert signal.pthread_sigmask(signal.SIG_BLOCK, set()) == old_mask diff --git a/tests/test_libero_state_tracks.py b/tests/test_libero_state_tracks.py new file mode 100644 index 0000000000..3a2458f974 --- /dev/null +++ b/tests/test_libero_state_tracks.py @@ -0,0 +1,268 @@ +"""Native geometry tests, independent of original dataset availability.""" + +import copy +import json +from argparse import Namespace + +import h5py +import mujoco +import numpy as np +import pytest + +from egomimic.benchmarks.libero.state_tracks import ( + build_arrays, + digest, + prepare, + restore_state, + validate_arrays, + validate_cache, + exclude_visualization_mocap, + _prepare_corpus_episode, + prepare_corpus, + compare_rgb, + export_rgb, + _export_rgb_episode, + export_rgb_corpus, +) + +XML = """ + + + + +""" + + +def fixture(): + model = mujoco.MjModel.from_xml_string(XML) + states = np.array([[0.0, 0.0, 0.0], [0.1, 0.2, 0.0], [0.2, 0.4, 0.0]]) + return model, states + + +def test_actual_state_alignment_and_persistent_motion(): + model, states = fixture() + arrays = build_arrays(model, states, ["view", "second"], "eef", 32, 32, [16, 16]) + assert arrays["rgb"].shape == (3, 2, 32, 32, 3) + np.testing.assert_allclose(arrays["pose"][:, 0], states[:, 1]) + moving = arrays["point_body_id"] != 0 + assert moving.any() + np.testing.assert_allclose( + arrays["xyz"][2, moving, 0] - arrays["xyz"][0, moving, 0], 0.4 + ) + damaged = copy.deepcopy(arrays) + damaged["source_row"] = np.roll(damaged["source_row"], 1) + with pytest.raises(ValueError, match="row"): + validate_arrays(damaged) + damaged = copy.deepcopy(arrays) + damaged["xyz"][1] = arrays["xyz"][0] + with pytest.raises(AssertionError): + validate_arrays(damaged) + arrays["visibility"][:] = False + validate_arrays(arrays) # occlusion must not erase exact coordinates + arrays["coordinate_available"][0, 0] = False + with pytest.raises(ValueError, match="survive occlusion"): + validate_arrays(arrays) + + +def test_reject_missing_state_and_time_contracts(): + model, states = fixture() + with pytest.raises(ValueError, match="strictly increasing"): + build_arrays(model, states[[1, 0, 2]], ["view"], "eef", 32, 32, [8, 8]) + with pytest.raises(ValueError, match="Invalid recorded state"): + restore_state(model, mujoco.MjData(model), np.zeros(2)) + model = mujoco.MjModel.from_xml_string( + '' + ) + with pytest.raises(ValueError, match="mocap"): + restore_state(model, mujoco.MjData(model), np.zeros(1)) + + +@pytest.mark.parametrize("pose_joints", [[], ["slide"]]) +def test_cache_rejects_original_observation_substitution_and_tampering(tmp_path, pose_joints): + _, states = fixture() + original = tmp_path / "demo.hdf5" + with h5py.File(original, "w") as f: + d = f.create_group("data/demo_0") + d.attrs["model_file"] = XML + d["states"] = states + d["obs/ee_pos"] = np.full((3, 3), 1000.0) # deliberately wrong originals + d["obs/agentview_rgb"] = np.full((3, 32, 32, 3), 255, np.uint8) + args = Namespace( + demo=str(original), + episode="demo_0", + robot_assets=tmp_path, + scene_assets=tmp_path, + output=tmp_path / "aligned", + pose_site="eef", + pose_joints=pose_joints, + cameras=["view"], + height=32, + width=32, + grid=[8, 8], + ) + result = prepare(args) + assert result["all_channels_native_replay_exact"] + options = {**vars(args), "input_schema": "libero", "asset_remap": [], + "exclude_visualization_body": []} + reused = _prepare_corpus_episode(options) + assert reused["reused"] + assert reused["all_channels_native_replay_exact"] + with pytest.raises(ValueError, match="request mismatch"): + _prepare_corpus_episode({**options, "pose_site": "changed"}) + manifest = json.loads((args.output / "manifest.json").read_text()) + assert manifest["original_cache_max_eef_error_m"] > 100 + assert manifest["original_cache_repaired"] is False + assert manifest["training_ready"] is False + highres = args.output / "rgb64.h5" + receipt = export_rgb(Namespace(output=args.output, rgb_cache=highres, + height=64, width=64, rgb_policy="exact")) + assert receipt["native_render"] and receipt["native_geometry_replay_exact"] + assert receipt["rgb_cache_sha256"] == digest(highres) + with h5py.File(highres) as cache: + assert cache["rgb"].shape == (3, 1, 64, 64, 3) + np.testing.assert_array_equal(cache["time"][:], states[:, 0]) + assert cache["camera_K"][0, 0, 0] > 0 + reused_rgb = _export_rgb_episode(dict(output=str(args.output), rgb_cache=str(highres), + height=64, width=64, rgb_policy="exact")) + assert reused_rgb["status"] == "PASS" + with pytest.raises(ValueError, match="mismatch"): + _export_rgb_episode(dict(output=str(args.output), rgb_cache=str(highres), + height=128, width=128, rgb_policy="exact")) + with pytest.raises(FileExistsError): + export_rgb(Namespace(output=args.output, rgb_cache=highres, + height=64, width=64, rgb_policy="exact")) + with pytest.raises(FileExistsError): + prepare(args) + with np.load(args.output / "aligned.npz") as z: + arrays = dict(z) + assert arrays["pose"].shape[1] == 7 + len(pose_joints) + if pose_joints: + np.testing.assert_array_equal(arrays["pose"][:, 7], states[:, 1]) + arrays["rgb"] = np.roll(arrays["rgb"], 1, axis=0) + np.savez_compressed(args.output / "aligned.npz", **arrays) + with pytest.raises(ValueError, match="Cache hash mismatch"): + validate_cache(args.output) + + # Even a re-hashed accidental row shift is rejected by actual native replay. + manifest["artifacts"]["aligned.npz"] = digest(args.output / "aligned.npz") + (args.output / "manifest.json").write_text(json.dumps(manifest)) + with pytest.raises(AssertionError): + validate_cache(args.output, replay=True) + + # Failed replay leaves a hash-valid manifest and arrays on disk. A corpus + # restart must not convert this failed native replay into a successful hit. + assert validate_cache(args.output)["status"] == "PASS" + with pytest.raises(AssertionError): + _prepare_corpus_episode(options) + + # If RGB and geometry both differ, the geometry failure must be exposed + # first; an RGB-only assertion should imply all non-RGB channels matched. + arrays["pose"][0, 0] += 0.01 + np.savez_compressed(args.output / "aligned.npz", **arrays) + manifest["artifacts"]["aligned.npz"] = digest(args.output / "aligned.npz") + (args.output / "manifest.json").write_text(json.dumps(manifest)) + with pytest.raises(AssertionError, match="pose"): + validate_cache(args.output, replay=True) + with pytest.raises(AssertionError, match="pose"): + validate_cache(args.output, replay=True, rgb_policy="uint8-roundoff-v1") + + +def test_rgb_roundoff_policy_is_explicit_bounded_and_never_modifies_pixels(): + original = np.full((100000, 3), 100, dtype=np.uint8) + repeated = original.copy() + repeated[0, 0] += 2 + before = repeated.copy() + with pytest.raises(AssertionError): + compare_rgb(original, repeated) + proof = compare_rgb(original, repeated, "uint8-roundoff-v1") + assert proof["max_abs"] == 2 and not proof["exact"] + np.testing.assert_array_equal(repeated, before) + repeated[0, 0] = 103 + with pytest.raises(ValueError, match="exceeds"): + compare_rgb(original, repeated, "uint8-roundoff-v1") + with pytest.raises(ValueError, match="exceeds"): + compare_rgb(original, repeated, "uint8-roundoff-v2") + repeated[0, 0] = 101 + repeated[1:5, 0] += 1 + with pytest.raises(ValueError, match="exceeds"): + compare_rgb(original, repeated, "uint8-roundoff-v1") + accepted = compare_rgb(original, repeated, "uint8-roundoff-v2") + assert accepted["changed_fraction_limit"] == 0.0001 + assert not accepted["exact"] + repeated[5:35, 0] += 1 + with pytest.raises(ValueError, match="exceeds"): + compare_rgb(original, repeated, "uint8-roundoff-v2") + with pytest.raises(ValueError, match="policy"): + compare_rgb(original, repeated, "anything") + + +def test_rgb_corpus_rejects_duplicate_output_before_processing(tmp_path): + spec = {"common": {"height": 64, "width": 64, "rgb_policy": "exact"}, "episodes": [ + {"embodiment": "one", "episode_id": "first", "output": "source", "rgb_cache": str(tmp_path / "rgb.h5")}, + {"embodiment": "two", "episode_id": "second", "output": "source2", "rgb_cache": str(tmp_path / "rgb.h5")}]} + path = tmp_path / "spec.json" + path.write_text(json.dumps(spec)) + with pytest.raises(ValueError, match="Duplicate"): + export_rgb_corpus(Namespace(spec=path, status=tmp_path / "status.json", workers=2)) + + +def test_reject_unknown_or_duplicate_pose_joint(): + model, states = fixture() + for names in (["missing"], ["slide", "slide"]): + with pytest.raises(ValueError, match="joint"): + build_arrays(model, states, ["view"], "eef", 32, 32, [8, 8], names) + + +def test_corpus_rejects_duplicate_source_and_unsafe_identity(tmp_path): + spec = {"common": {}, "embodiments": {"gripper": {}}, "episodes": [ + {"embodiment": "gripper", "episode_id": "one", "episode": "demo_0", "demo": str(tmp_path / "source.h5")}, + {"embodiment": "gripper", "episode_id": "two", "episode": "demo_0", "demo": str(tmp_path / "source.h5")}]} + path = tmp_path / "spec.json" + path.write_text(json.dumps(spec)) + args = Namespace(spec=path, output=tmp_path / "output", workers=1) + with pytest.raises(ValueError, match="Duplicate"): + prepare_corpus(args) + spec["episodes"] = [spec["episodes"][0]] + spec["episodes"][0]["episode_id"] = "../outside" + path.write_text(json.dumps(spec)) + with pytest.raises(ValueError, match="Unsafe"): + prepare_corpus(args) + + +def test_quest_recording_and_explicit_nonphysical_marker_exclusion(tmp_path): + _, states = fixture() + marker = '' + xml = XML.replace('', marker + '') + source = tmp_path / 'quest.hdf5' + with h5py.File(source, 'w') as f: + f['state'] = states + f['destination_model_xml'] = xml + f['numeric_observations/robot0_eef_pos'] = np.zeros((3, 3)) + args = Namespace(demo=str(source), episode='root', input_schema='quest', + robot_assets=tmp_path, scene_assets=tmp_path, output=tmp_path/'quest', + pose_site='eef', cameras=['view'], height=32, width=32, grid=[8, 8], + exclude_visualization_body=['marker'], asset_remap=[]) + assert prepare(args)['all_channels_native_replay_exact'] + manifest = json.loads((args.output/'manifest.json').read_text()) + assert manifest['excluded_visualization_bodies'] == ['marker'] + assert manifest['input_schema'] == 'quest' + assert mujoco.MjModel.from_xml_string(exclude_visualization_mocap(xml, ['marker'])).nmocap == 0 + with pytest.raises(ValueError, match='physical'): + exclude_visualization_mocap(xml.replace('contype="0"', 'contype="1"'), ['marker']) + with pytest.raises(ValueError, match='referenced'): + exclude_visualization_mocap(xml.replace('', ''), ['marker']) + + +def test_explicit_asset_prefix_remap_rejects_ambiguity(tmp_path, monkeypatch): + from egomimic.benchmarks.libero.replay import rebase_demo_xml + (tmp_path/'hand.stl').write_bytes(b'fixture') + xml = '' + result = rebase_demo_xml(xml, tmp_path, tmp_path, [('/recorded/hand', tmp_path)]) + assert str(tmp_path/'hand.stl') in result + monkeypatch.chdir(tmp_path.parent) + relative = rebase_demo_xml(xml, tmp_path, tmp_path, + [('/recorded/hand', tmp_path.name)]) + assert relative == result + with pytest.raises(ValueError, match='Ambiguous'): + rebase_demo_xml(xml, tmp_path, tmp_path, + [('/recorded/hand', tmp_path), ('/recorded', tmp_path)]) diff --git a/tests/test_slurm_process_identity.py b/tests/test_slurm_process_identity.py new file mode 100644 index 0000000000..c71e38a534 --- /dev/null +++ b/tests/test_slurm_process_identity.py @@ -0,0 +1,62 @@ +"""Real Linux child delivery and fail-closed trainer identity regressions.""" +import copy +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time +from unittest import mock + +import pytest +from egomimic.utils import slurm_process_identity as identity + +pytestmark = pytest.mark.skipif(not hasattr(os, "pidfd_open"), reason="Linux pidfd required") + + +def test_real_child_pidfd_delivery(tmp_path): + ready = tmp_path / "ready.json" + hit = tmp_path / "received" + child_code = """ +import json, os, signal, sys, time +from pathlib import Path +from egomimic.utils.slurm_process_identity import ready_record +signal.signal(signal.SIGUSR2, lambda *_: Path(sys.argv[2]).write_text('USR2')) +Path(sys.argv[1]).write_text(json.dumps(ready_record('nonce', '123', 0))) +time.sleep(30) +""" + source = identity.source_identity() + child = subprocess.Popen([sys.executable, "-c", child_code, str(ready), str(hit)], cwd=source["root"]) + try: + deadline = time.monotonic() + 15 + while not ready.exists() and child.poll() is None and time.monotonic() < deadline: + time.sleep(.05) + assert ready.exists() + original = json.loads(ready.read_text()) + arguments = dict(pid=child.pid, nonce="nonce", job_id="123", restart=0, + parent_identity=identity.process_identity(os.getpid()), + expected_source=source, expected_process=identity.process_identity(child.pid)) + for field, bad in (("nonce", "forged"), ("job_id", "124"), ("restart_count", 1), + ("source", {}), ("handler_installed", False)): + changed = copy.deepcopy(original); changed[field] = bad + ready.write_text(json.dumps(changed)) + with pytest.raises(ValueError): identity.signal_ready_trainer(ready, **arguments) + for field, bad in (("uid", [999]*4), ("pid", child.pid+1), ("starttime", 0), + ("pid_namespace", "forged"), ("cgroup", "forged"), + ("cwd", "/tmp"), ("exe", "/bin/false"), ("boot_id", "forged")): + changed = copy.deepcopy(original); changed["process"][field] = bad + ready.write_text(json.dumps(changed)) + with pytest.raises(ValueError): identity.signal_ready_trainer(ready, **arguments) + ready.write_text(json.dumps(original)) + with mock.patch.object(identity, "source_identity", return_value={}): + with pytest.raises(ValueError, match="source changed"): + identity.signal_ready_trainer(ready, **arguments) + assert not hit.exists(), "forged records must never deliver the signal" + receipt = identity.signal_ready_trainer(ready, **arguments) + deadline = time.monotonic() + 5 + while not hit.exists() and time.monotonic() < deadline: time.sleep(.05) + assert hit.read_text() == "USR2" + assert receipt["pid"] == child.pid + finally: + child.terminate(); child.wait(timeout=10) diff --git a/tests/test_slurm_save_only_checkpoint.py b/tests/test_slurm_save_only_checkpoint.py index c2d543952f..dff7ec4fd2 100644 --- a/tests/test_slurm_save_only_checkpoint.py +++ b/tests/test_slurm_save_only_checkpoint.py @@ -269,3 +269,30 @@ def broadcast(value, src=0): assert existing.read_bytes() == b"immutable" assert not existing.with_name(f".{existing.name}.partial").exists() assert not existing.with_name(f".{existing.name}.reserve").exists() + + +def test_train_only_deferral_waits_for_completed_optimizer_update(tmp_path, monkeypatch): + callback = SaveOnlySignalCheckpoint(tmp_path, train_batch_only=True) + trainer = SimpleNamespace(global_step=15) + callback._request_checkpoint(callback.save_signal.value, None) + # Validation/epoch/fit hooks and an accumulation-only microbatch preserve + # the pending generation without touching checkpoint storage. + for hook in ("on_validation_batch_end", "on_validation_epoch_end", + "on_train_epoch_end", "on_fit_end", "on_train_batch_end"): + callback._service(trainer, hook) + callback.on_train_batch_start(trainer, None, None, 0) + callback.on_train_batch_end(trainer, None, None, None, 0) + assert callback._handled_generation == 0 + assert not list(tmp_path.iterdir()) + def reached_save(_trainer): + raise RuntimeError("completed update reaches checkpoint publication") + monkeypatch.setattr(callback, "_checkpoint_path", reached_save) + trainer.global_step = 16 + with pytest.raises(RuntimeError, match="completed update reaches"): + callback.on_train_batch_end(trainer, None, None, None, 1) + assert callback._handled_generation == 1 + + +def test_train_only_deferral_rejects_string_boolean(tmp_path): + with pytest.raises(TypeError, match="must be a boolean"): + SaveOnlySignalCheckpoint(tmp_path, train_batch_only="false") diff --git a/tests/test_train_hydra_slurm_environment.py b/tests/test_train_hydra_slurm_environment.py index 69ec177e5d..5a700abefb 100644 --- a/tests/test_train_hydra_slurm_environment.py +++ b/tests/test_train_hydra_slurm_environment.py @@ -19,6 +19,54 @@ ) +@pytest.fixture +def preserve_precision_flags(): + precision = torch.get_float32_matmul_precision() + matmul = torch.backends.cuda.matmul.allow_tf32 + cudnn = torch.backends.cudnn.allow_tf32 + yield + torch.set_float32_matmul_precision(precision) + torch.backends.cuda.matmul.allow_tf32 = matmul + torch.backends.cudnn.allow_tf32 = cudnn + + +def test_strict_fp32_opt_out_preserves_backend_flags(preserve_precision_flags): + torch.set_float32_matmul_precision("high") + torch.backends.cudnn.allow_tf32 = True + before = (torch.get_float32_matmul_precision(), + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + train_hydra._configure_strict_fp32(OmegaConf.create({})) + assert before == (torch.get_float32_matmul_precision(), + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + + +@pytest.mark.parametrize("precision", [32, "32-true"]) +def test_strict_fp32_disables_both_backends(monkeypatch, preserve_precision_flags, precision): + monkeypatch.setenv("NVIDIA_TF32_OVERRIDE", "0") + torch.set_float32_matmul_precision("high") + torch.backends.cudnn.allow_tf32 = True + train_hydra._configure_strict_fp32(OmegaConf.create({ + "runtime": {"strict_fp32": True}, "trainer": {"precision": precision}})) + assert torch.get_float32_matmul_precision() == "highest" + assert not torch.backends.cuda.matmul.allow_tf32 + assert not torch.backends.cudnn.allow_tf32 + + +@pytest.mark.parametrize("precision,override", [("bf16-mixed", "0"), ("32-true", None), ("32-true", "1")]) +def test_strict_fp32_rejects_incompatible_runtime(monkeypatch, precision, override): + monkeypatch.delenv("NVIDIA_TF32_OVERRIDE", raising=False) + if override is not None: + monkeypatch.setenv("NVIDIA_TF32_OVERRIDE", override) + with pytest.raises(ValueError): + train_hydra._configure_strict_fp32(OmegaConf.create({ + "runtime": {"strict_fp32": True}, "trainer": {"precision": precision}})) + + +def test_strict_fp32_rejects_string_boolean(): + with pytest.raises(ValueError, match="must be a boolean"): + train_hydra._configure_strict_fp32(OmegaConf.create({"runtime": {"strict_fp32": "false"}})) + + def _clear_ownership(monkeypatch): for name in _OWNERSHIP_VARIABLES: monkeypatch.delenv(name, raising=False) @@ -56,6 +104,7 @@ def test_base_config_records_lightning_as_default_requeue_owner(): assert cfg.runtime.slurm_requeue_owner == "lightning" assert cfg.runtime.slurm_save_signal == "SIGUSR2" + assert cfg.runtime.slurm_signal_train_batch_only is False assert ( cfg.runtime.slurm_signal_checkpoint_dir == "/tmp/egomimic-test-output/checkpoints" @@ -98,6 +147,9 @@ def test_runner_owned_job_adds_one_save_only_sigusr2_callback(monkeypatch, tmp_p assert len(callbacks) == 1 assert isinstance(callbacks[0], SaveOnlySignalCheckpoint) + assert callbacks[0].train_batch_only is False + cfg.runtime.slurm_signal_train_batch_only = True + assert train_hydra._instantiate_slurm_callbacks(cfg)[0].train_batch_only is True assert callbacks[0].save_signal == signal.SIGUSR2 assert callbacks[0].checkpoint_dir == (tmp_path / "checkpoints").resolve() diff --git a/tests/test_world_model.py b/tests/test_world_model.py new file mode 100644 index 0000000000..4684c767df --- /dev/null +++ b/tests/test_world_model.py @@ -0,0 +1,1216 @@ +import copy +import hashlib +import json + +import numpy as np +import pytest +import torch + +from egomimic.models.world_model import MoTWorldModel +from egomimic.pipeline.algo import PipelineAlgo +from egomimic.pipeline.stages_world_model import WorldModelStage +from egomimic.rldb.world_model import ( + HomogeneousWorldBatchSampler, + PreparedWorldDataset, + collate_world, + points_to_reference, + sample_times, +) + + +def make_batch(pose_dim=7, batch_size=2): + result = {} + for name, dim, n in (("image", 8, 6), ("track", 3, 4), ("proprio", pose_dim, 2)): + time = ( + torch.arange(2) + .repeat_interleave(n // 2) + .float() + .expand(batch_size, -1) + .clone() + ) + result[name] = dict( + value=torch.randn(batch_size, n, dim), + time=time, + index=time.clone(), + valid=torch.ones(batch_size, n, dtype=torch.bool), + observed=time == 0, + target_valid=torch.ones(batch_size, n, dtype=torch.bool), + ) + if name != "proprio": + result[name].update( + xy=torch.rand(batch_size, n, 2), + camera=torch.zeros(batch_size, n, dtype=torch.long), + ) + if name == "track": + result[name].update( + point_id=torch.tensor([0, 1, 0, 1]).expand(batch_size, -1), + confidence=torch.ones(batch_size, n), + ) + return result + + +def make_model(**kwargs): + return MoTWorldModel( + {"robot": 7, "human": 11}, + image_dim=8, + width=24, + heads=3, + depth=2, + max_tracks=8, + max_tokens=100, + **kwargs, + ) + + +@pytest.mark.parametrize("chunk,rematerialize", [(3, False), (0, True), (3, True)]) +def test_memory_controls_preserve_outputs_losses_gradients_and_update(chunk, rematerialize): + torch.manual_seed(48) + baseline = make_model().float() + optimized = make_model( + attention_query_chunk_size=chunk, activation_checkpointing=rematerialize + ).float() + optimized.load_state_dict(baseline.state_dict(), strict=True) + batch = make_batch() + # Include padding in one example and arbitrary observed/target ordering. + for name in ("image", "track"): + batch[name]["valid"][0, -1] = False + batch[name]["target_valid"][0, -1] = False + batch[name]["observed"][0, -1] = False + batch[name]["observed"][1] = batch[name]["observed"][1].flip(0) + levels = {n: torch.zeros_like(batch[n]["time"]) for n in ("image", "track")} + values = {n: batch[n]["value"] for n in ("image", "track")} + with torch.no_grad(): + expected = baseline(batch, "robot", values, levels) + actual = optimized(batch, "robot", values, levels) + for name in expected: + torch.testing.assert_close(actual[name], expected[name], atol=2e-6, rtol=2e-5) + losses = [] + for model in (baseline, optimized): + torch.manual_seed(99) + loss = model.training_losses(batch, "robot") + loss["total"].backward() + losses.append(loss) + for name in losses[0]: + torch.testing.assert_close(losses[0][name], losses[1][name], atol=2e-6, rtol=2e-5) + for (name, a), (other, b) in zip(baseline.named_parameters(), optimized.named_parameters()): + assert name == other + assert (a.grad is None) == (b.grad is None) + if a.grad is not None: + torch.testing.assert_close(a.grad, b.grad, atol=3e-6, rtol=1e-4, msg=name) + for model in (baseline, optimized): + torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01).step() + for name, a in baseline.state_dict().items(): + torch.testing.assert_close(a, optimized.state_dict()[name], atol=2e-6, rtol=2e-5, msg=name) + + +def test_chunked_attention_keeps_all_keys(monkeypatch): + from egomimic.models.world_model import MoTBlock + calls = [] + original = MoTBlock._attend + def observe(q, k, v, mask): + calls.append((q.shape[2], k.shape[2], v.shape[2])) + return original(q, k, v, mask) + monkeypatch.setattr(MoTBlock, "_attend", staticmethod(observe)) + model, batch = make_model(attention_query_chunk_size=3), make_batch() + with torch.no_grad(): + model(batch, "robot", {n: batch[n]["value"] for n in ("image", "track")}, + {n: torch.zeros_like(batch[n]["time"]) for n in ("image", "track")}) + assert calls and all(q <= 3 and k == v == 12 for q, k, v in calls) + + +def test_pose_gradient_reaches_image_and_track_generators(): + torch.manual_seed(7) + model = make_model() + losses = model.training_losses(make_batch(), "robot") + losses["pose"].backward() + for module in (model.image_out, model.track_out, model.image_in, model.track_in): + assert module.weight.grad is not None + assert module.weight.grad.norm() > 0 + for block in model.blocks: + assert all(m.weight.grad.norm() > 0 for m in block.qkv) + assert model.pose_out["human"].weight.grad is None + + +@pytest.mark.parametrize("active", ["image", "track", "proprio"]) +def test_each_modality_alone_excludes_other_streams_and_their_gradients(active): + """Isolated expert execution is a code contract, not learned real-data quality.""" + torch.manual_seed(42) + model, batch = make_model().float(), make_batch() + names = ("image", "track", "proprio") + for name in names: + if name != active: + for key in ("valid", "observed", "target_valid"): + batch[name][key].zero_() + levels = {n: torch.zeros_like(batch[n]["time"]) for n in names[:2]} + values = {n: batch[n]["value"] for n in names[:2]} + original = model(batch, "robot", values, levels)[active] + changed = copy.deepcopy(batch) + for name in names: + if name != active: + changed[name]["value"].fill_(float("nan")) + mutated = model( + changed, "robot", {n: changed[n]["value"] for n in names[:2]}, levels + )[active] + torch.testing.assert_close(original, mutated, rtol=0, atol=0) + original.square().mean().backward() + for block in model.blocks: + for i, name in enumerate(names): + grad = block.qkv[i].weight.grad + if name == active: + assert ( + grad is not None and torch.isfinite(grad).all() and grad.norm() > 0 + ) + else: + assert grad is None or torch.equal(grad, torch.zeros_like(grad)) + + +def test_oracle_visual_pose_readout_without_generator_uses_both_modalities(): + """Clean visual inputs isolate the readout; hidden pose targets cannot leak.""" + torch.manual_seed(42) + model, batch = make_model().float(), make_batch() + visual = {n: batch[n]["value"].clone().requires_grad_() for n in ("image", "track")} + batch["proprio"]["value"].requires_grad_() + levels = {n: torch.zeros_like(batch[n]["time"]) for n in visual} + prediction = model(batch, "robot", visual, levels)["proprio"] + query = ~batch["proprio"]["observed"] + prediction[query].square().mean().backward() + for value in visual.values(): + assert value.grad is not None and torch.isfinite(value.grad).all() + assert value.grad.norm() > 0 + pose_grad = batch["proprio"]["value"].grad + assert pose_grad[batch["proprio"]["observed"]].norm() > 0 + assert torch.equal(pose_grad[query], torch.zeros_like(pose_grad[query])) + assert model.image_out.weight.grad is None + assert model.track_out.weight.grad is None + + +def test_visual_velocity_target_inverts_to_clean_image_and_track_independently(): + from egomimic.models.world_model import noise_coefficients + + torch.manual_seed(42) + for width in (768, 3): + clean, noise = torch.randn(2, 4, width), torch.randn(2, 4, width) + level = torch.tensor([[0.0, 0.1, 0.9, 1.0]]).expand(2, -1) + alpha, sigma = noise_coefficients(level) + noisy = alpha * clean + sigma * noise + velocity = alpha * noise - sigma * clean + restored = alpha * noisy - sigma * velocity + torch.testing.assert_close(restored, clean, rtol=1e-5, atol=1e-6) + + +def test_zero_initialized_adapters_preserve_adapter_free_predictions(): + model = make_model(image_adapter_dim=4).eval() + baseline = make_model(image_adapter_dim=0).eval() + baseline.load_state_dict( + { + k: v + for k, v in model.state_dict().items() + if not k.startswith("image_adapters.") + }, + strict=True, + ) + batch = make_batch() + a = model.sample(batch, "robot", 2, torch.Generator().manual_seed(5)) + b = baseline.sample(batch, "robot", 2, torch.Generator().manual_seed(5)) + for name in a: + torch.testing.assert_close(a[name], b[name], rtol=0, atol=0) + # Old checkpoints must not silently skip the new trainable parameters. + with pytest.raises(RuntimeError, match="Missing key"): + model.load_state_dict(baseline.state_dict(), strict=True) + + +@pytest.mark.parametrize( + "selected", ["yam", "human", "libero_sharpa", "libero_gripper"] +) +def test_image_adapters_route_and_learn_for_all_four_embodiments(selected): + torch.manual_seed(14) + dims = dict(yam=7, human=11, libero_sharpa=9, libero_gripper=5) + model = MoTWorldModel( + dims, + image_dim=8, + width=24, + heads=3, + depth=2, + max_tracks=8, + max_tokens=100, + image_adapter_dim=4, + ) + untouched = { + name: p.detach().clone() + for name, p in model.named_parameters() + if name.startswith("image_adapters.") and f".{selected}." not in name + } + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + for step in range(2): + optimizer.zero_grad(set_to_none=True) + # Pose loss alone must reach adapters through generated visual features. + model.training_losses(make_batch(dims[selected]), selected)["pose"].backward() + for block in model.image_adapters: + adapter = block[selected] + assert adapter.up.weight.grad.norm() > 0 + if step: + assert adapter.down.weight.grad.norm() > 0 + assert adapter.norm.weight.grad.norm() > 0 + for other in dims.keys() - {selected}: + assert all(p.grad is None for p in block[other].parameters()) + assert model.blocks[0].qkv[0].weight.grad.norm() > 0 + optimizer.step() + for name, p in model.named_parameters(): + if name in untouched: + torch.testing.assert_close(p, untouched[name], rtol=0, atol=0) + + +def test_adapter_effect_is_embodiment_specific_and_checkpointed(): + model = make_model(image_adapter_dim=4).eval() + batches = {"robot": make_batch(7), "human": make_batch(11)} + + def predict(embodiment): + return model.sample( + batches[embodiment], embodiment, 2, torch.Generator().manual_seed(9) + )["image"] + + before = {e: predict(e) for e in batches} + with torch.no_grad(): + for block in model.image_adapters: + block["robot"].up.weight.normal_(std=0.1) + assert not torch.equal(before["robot"], predict("robot")) + torch.testing.assert_close(before["human"], predict("human"), rtol=0, atol=0) + restored = make_model(image_adapter_dim=4).eval() + restored.load_state_dict(model.state_dict(), strict=True) + torch.testing.assert_close( + predict("robot"), + restored.sample(batches["robot"], "robot", 2, torch.Generator().manual_seed(9))[ + "image" + ], + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize("dimension", [-1, 1.5, True]) +def test_invalid_image_adapter_width_is_rejected(dimension): + with pytest.raises(ValueError, match="image_adapter_dim"): + make_model(image_adapter_dim=dimension) + + +@pytest.mark.parametrize("embodiment,dimension", [("robot", 7), ("human", 11)]) +def test_actual_optimizer_and_strict_reload(embodiment, dimension): + torch.manual_seed(10) + model = make_model() + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + before = model.track_out.weight.detach().clone() + for _ in range(2): + optimizer.zero_grad() + loss = model.training_losses(make_batch(dimension), embodiment)["total"] + assert torch.isfinite(loss) + loss.backward() + optimizer.step() + assert not torch.equal(before, model.track_out.weight) + replacement = make_model() + replacement.load_state_dict(model.state_dict(), strict=True) + new_optimizer = torch.optim.AdamW(replacement.parameters(), lr=1e-3) + new_optimizer.load_state_dict(optimizer.state_dict()) + assert len(new_optimizer.state) == len(optimizer.state) + batch = make_batch(dimension) + a = model.sample(batch, embodiment, 3, torch.Generator().manual_seed(8)) + b = replacement.sample(batch, embodiment, 3, torch.Generator().manual_seed(8)) + for key in a: + torch.testing.assert_close(a[key], b[key], rtol=0, atol=0) + + +def test_sampler_does_not_read_withheld_ground_truth(): + model = make_model().eval() + batch = make_batch() + changed = copy.deepcopy(batch) + for stream in changed.values(): + stream["value"][~stream["observed"]] = float("nan") + stream["target_valid"][:] = False + if "confidence" in stream: + stream["confidence"][:] = 0 + a = model.sample(batch, "robot", 2, torch.Generator().manual_seed(4)) + b = model.sample(changed, "robot", 2, torch.Generator().manual_seed(4)) + for name in a: + torch.testing.assert_close(a[name], b[name], rtol=0, atol=0) + torch.testing.assert_close( + a[name][batch[name]["observed"]], + batch[name]["value"][batch[name]["observed"]], + ) + + +def test_hidden_pose_does_not_condition_first_pass(): + model = make_model() + batch = make_batch() + other = copy.deepcopy(batch) + other["proprio"]["value"][~other["proprio"]["observed"]] += 900 + values = {n: batch[n]["value"] for n in ("image", "track")} + levels = {n: torch.zeros_like(batch[n]["time"]) for n in values} + a, b = model(batch, "robot", values, levels), model(other, "robot", values, levels) + for name in a: + torch.testing.assert_close(a[name], b[name], rtol=0, atol=0) + + +def refinement_batch(): + batch = make_batch(batch_size=1) + for stream in batch.values(): + for key, value in list(stream.items()): + stream[key] = torch.cat((value, value.clone()), dim=1) + time = ( + torch.arange(4) + .repeat_interleave(stream["value"].shape[1] // 4)[None] + .float() + ) + stream["time"], stream["index"] = time, time.clone() + stream["observed"] = time == 0 + batch["refinement"] = {} + for name in ("image", "track"): + s = batch[name] + strength = torch.tensor([0.0, 0.0, 0.4, 1.0])[s["time"].long()] + previous = torch.randn_like(s["value"]) + previous[strength == 1] = float("nan") # Fresh queries need no previous value. + batch["refinement"][name] = dict(strength=strength, prediction=previous) + s["value"][~s["observed"]] = float("nan") # Never use held-out ground truth. + return batch + + +@pytest.mark.parametrize("mode,mixed", [("continuous", False), ("discrete", True)]) +def test_selective_refinement_schedule_and_clean_anchor_clamping(mode, mixed): + from egomimic.models.world_model import noise_coefficients + + model = make_model(time_mode=mode).eval() + batch = refinement_batch() + original = copy.deepcopy(batch) + calls = [] + handle = model.register_forward_pre_hook( + lambda module, args: calls.append( + ( + {k: v.clone() for k, v in args[2].items()}, + {k: v.clone() for k, v in args[3].items()}, + ) + ) + ) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=mixed): + output = model.sample(batch, "robot", 4, torch.Generator().manual_seed(29)) + handle.remove() + assert len(calls) == 5 # Four joint denoising passes, then pose readout. + rng = torch.Generator().manual_seed(29) + for name in ("image", "track"): + s, control = batch[name], batch["refinement"][name] + strength, previous = control["strength"], control["prediction"] + noise = torch.randn(s["value"].shape, generator=rng) + partial = strength == 0.4 + alpha, sigma = noise_coefficients(strength) + expected = alpha * previous + sigma * noise + torch.testing.assert_close(calls[0][0][name][partial], expected[partial]) + fresh = strength == 1 + torch.testing.assert_close( + calls[0][0][name][fresh], noise[fresh], rtol=0, atol=0 + ) + retained = (strength == 0) & ~s["observed"] + for values, levels in calls: + torch.testing.assert_close( + values[name][retained], previous[retained], rtol=0, atol=0 + ) + torch.testing.assert_close( + values[name][s["observed"]], s["value"][s["observed"]], rtol=0, atol=0 + ) + for j, (_, levels) in enumerate(calls[:4]): + torch.testing.assert_close( + levels[name], strength * ((4 - j) / 4), rtol=0, atol=0 + ) + torch.testing.assert_close( + output[name][retained], previous[retained], rtol=0, atol=0 + ) + assert torch.isfinite(output[name]).all() + assert not torch.equal(output[name][partial], previous[partial]) + for key in s: + torch.testing.assert_close( + s[key], original[name][key], rtol=0, atol=0, equal_nan=True + ) + assert torch.isfinite(output["proprio"]).all() + + +def test_refinement_fresh_controls_equal_default_and_ignore_seed_values(): + model = make_model().eval() + batch = make_batch() + baseline = model.sample(batch, "robot", 3, torch.Generator().manual_seed(31)) + for name in ("image", "track"): + s = batch[name] + batch.setdefault("refinement", {})[name] = dict( + strength=(~s["observed"]).float(), + prediction=torch.full_like(s["value"], float("nan")), + ) + controlled = model.sample(batch, "robot", 3, torch.Generator().manual_seed(31)) + for name in baseline: + torch.testing.assert_close(baseline[name], controlled[name], rtol=0, atol=0) + + +def test_repeat_refinement_and_new_timestamps_through_pipeline(): + model = make_model().eval() + batch = refinement_batch() + first = model.sample(batch, "robot", 2, torch.Generator().manual_seed(3)) + # Preserve the generated frame at t=2, revise t=1, generate t=3 afresh. + for name in ("image", "track"): + batch["refinement"][name] = dict( + strength=torch.tensor([0.0, 0.6, 0.0, 1.0])[batch[name]["time"].long()], + prediction=first[name], + ) + # Change the fresh query's physical time, preserving other token identities. + batch[name]["time"][batch[name]["time"] == 3] = 4.5 + algo = PipelineAlgo([WorldModelStage(model, sampler_steps=2)], device="cpu") + raw = {"robot": {"world": batch, "embodiment": "robot"}} + second = algo.forward_eval(raw)["robot"]["world_prediction"] + for name in ("image", "track"): + retained = batch[name]["time"] == 2 + torch.testing.assert_close( + second[name][retained], first[name][retained], rtol=0, atol=0 + ) + assert torch.isfinite(second[name]).all() + + +@pytest.mark.parametrize( + "error", + [ + "range", + "observed", + "missing", + "nan", + "frame", + "shape", + "proprio", + ], +) +def test_refinement_invalid_requests_fail_before_rng_use(error): + model, batch = make_model(), refinement_batch() + c = batch["refinement"]["image"] + if error == "range": + c["strength"][0, -1] = 1.1 + elif error == "observed": + c["strength"][0, 0] = 0.1 + elif error == "missing": + del c["prediction"] + elif error == "nan": + c["prediction"][0, 3] = float("nan") + elif error == "frame": + c["strength"][0, 3] = 0.2 + elif error == "shape": + c["strength"] = c["strength"][:, :-1] + else: + batch["refinement"]["proprio"] = {} + generator = torch.Generator().manual_seed(19) + state = generator.get_state().clone() + with pytest.raises(ValueError): + model.sample(batch, "robot", 2, generator) + assert torch.equal(state, generator.get_state()) + + +def test_padding_invariance(): + model = make_model().eval() + batch = make_batch() + padded = copy.deepcopy(batch) + for stream in padded.values(): + for key, x in list(stream.items()): + extra = torch.zeros_like(x[:, :1]) + if key == "value": + extra[:] = float("nan") + stream[key] = torch.cat((x, extra), 1) + levels = {n: torch.zeros_like(batch[n]["time"]) for n in ("image", "track")} + values = {n: batch[n]["value"] for n in levels} + a = model(batch, "robot", values, levels) + b = model( + padded, + "robot", + {n: padded[n]["value"] for n in levels}, + {n: torch.zeros_like(padded[n]["time"]) for n in levels}, + ) + for name in a: + torch.testing.assert_close(a[name], b[name][:, :-1], atol=2e-6, rtol=2e-5) + + +def test_common_time_shift_invariance_and_discrete_mode(): + model = make_model().eval() + batch, shifted = make_batch(), None + shifted = copy.deepcopy(batch) + for s in shifted.values(): + s["time"] += 3.5 + levels = {n: torch.zeros_like(batch[n]["time"]) for n in ("image", "track")} + values = {n: batch[n]["value"] for n in levels} + a, b = ( + model(batch, "robot", values, levels), + model(shifted, "robot", values, levels), + ) + for name in a: + torch.testing.assert_close(a[name], b[name], atol=2e-6, rtol=2e-5) + discrete = make_model(time_mode="discrete") + discrete.load_state_dict(model.state_dict()) + for name in a: + torch.testing.assert_close( + a[name], discrete(batch, "robot", values, levels)[name] + ) + del shifted["image"]["index"] + with pytest.raises(ValueError, match="indices"): + discrete(shifted, "robot", values, levels) + + +def test_budget_guard_and_warmup_in_pipeline(): + model = make_model() + stage = WorldModelStage(model, warmup_steps=3, pose_ramp_steps=2) + assert [stage.pose_scale(i) for i in range(6)] == [0, 0, 0, 0, 0.5, 1] + algo = PipelineAlgo([stage], device="cpu") + raw = {"robot": {"world": make_batch(), "embodiment": "robot", "optimizer_step": 0}} + predictions = algo.forward_training(raw) + loss = algo.compute_losses(predictions, raw)["loss"] + loss.backward() + assert torch.isfinite(loss) + assert predictions["robot"]["log/world_pose"] == 0 + assert "optimizer_step" not in raw["robot"]["world"] + model.max_tokens = 2 + with pytest.raises(ValueError, match="budget"): + model.sample(make_batch(), "robot") + + +def test_temporal_sampling_and_rigid_frame(): + counts = set() + rng = np.random.default_rng(2) + for _ in range(50): + selected = sample_times(100, rng) + counts.add(len(selected)) + assert (np.diff(selected) > 0).all() + assert counts == {2, 3, 4} + selected = sample_times(100, rng, mode="discrete", stride=3) + assert (np.diff(selected) == 3).all() + transform = np.eye(4) + transform[:3, 3] = [1, 2, 3] + np.testing.assert_allclose( + points_to_reference(np.zeros((2, 3)), transform), [[1, 2, 3]] * 2 + ) + transform[0, 0] = -1 + with pytest.raises(ValueError, match="handedness"): + points_to_reference(np.zeros((2, 3)), transform) + + +def test_sixteen_timestamps_preserve_physical_window(): + times = np.linspace(0, 12, 361, dtype=np.float64) + counts = set() + for seed in range(200): + selected = sample_times( + len(times), np.random.default_rng(seed), anchor=180, max_frames=16, + timestamps=times, max_past_seconds=3, max_future_seconds=3, + ) + counts.add(len(selected)) + assert 180 in selected + assert (np.diff(selected) > 0).all() + assert (np.abs(times[selected] - times[180]) <= 3).all() + assert counts == set(range(2, 17)) + + +def test_physical_time_bounds_use_seconds_and_cover_both_sides(): + times = np.array([0, 2, 5.9, 6, 7, 9, 10, 12, 12.1, 40], dtype=np.float64) + counts, seen = set(), set() + for seed in range(100): + selected = sample_times( + len(times), np.random.default_rng(seed), anchor=5, timestamps=times, + max_past_seconds=3, max_future_seconds=3, + ) + assert 5 in selected + assert (np.diff(selected) > 0).all() + assert ((times[selected] - times[5] >= -3) & (times[selected] - times[5] <= 3)).all() + counts.add(len(selected)) + seen.update(selected) + assert counts == {2, 3, 4} + assert seen == {3, 4, 5, 6, 7} + # The exact recorded endpoints are included; no nearest frame may cross them. + selected = sample_times( + 5, np.random.default_rng(1), anchor=2, min_frames=3, max_frames=3, + timestamps=np.array([6.9, 7, 10, 13, 13.1]), + max_past_seconds=3, max_future_seconds=3, + ) + np.testing.assert_array_equal(selected, [1, 2, 3]) + + +@pytest.mark.parametrize("anchor,expected", [(0, [0, 1, 2, 3]), (7, [4, 5, 6, 7])]) +def test_physical_sampling_clips_episode_edges_without_gap_escape(anchor, expected): + times = np.array([0, 1, 2, 3, 100, 101, 102, 103], dtype=np.float64) + selected = sample_times( + len(times), np.random.default_rng(9), anchor=anchor, + min_frames=4, max_frames=4, timestamps=times, + max_past_seconds=3, max_future_seconds=3, + ) + np.testing.assert_array_equal(selected, expected) + with pytest.raises(ValueError, match="Too few recorded"): + sample_times( + 3, np.random.default_rng(1), anchor=1, timestamps=[0, 10, 20], + max_past_seconds=3, max_future_seconds=3, + ) + + +@pytest.mark.parametrize("past,future", [(None, 3), (3, None), (-1, 3), (3, float("nan")), + (True, 3), (0, 0), (3, float("inf"))]) +def test_physical_sampling_rejects_invalid_bounds(past, future): + with pytest.raises(ValueError, match="Time bounds"): + sample_times(4, np.random.default_rng(1), timestamps=[0, 1, 2, 3], + max_past_seconds=past, max_future_seconds=future) + + +def test_physical_sampling_requires_valid_times_and_rejects_discrete_bounds(): + for times in (None, [0, 1], [0, 1, 1, 3], [0, 1, np.nan, 3]): + with pytest.raises(ValueError, match="Physical timestamps"): + sample_times(4, np.random.default_rng(1), timestamps=times, + max_past_seconds=3, max_future_seconds=3) + with pytest.raises(ValueError, match="only to continuous"): + sample_times(4, np.random.default_rng(1), mode="discrete", + max_past_seconds=3, max_future_seconds=3) + + +def cache_manifest(tmp_path): + entries = [] + for embodiment, d in (("robot", 7), ("human", 11)): + for split in ("train", "val"): + path = tmp_path / f"{embodiment}-{split}.npz" + rng = np.random.default_rng(1) + np.savez( + path, + time=np.arange(8, dtype=np.float64) * 0.1 + 1e9, + image=rng.normal(size=(8, 2, 3, 8)), + image_xy=rng.random((3, 2)), + track=rng.normal(size=(8, 2, 3)), + track_xy=rng.random((2, 2)), + track_camera=np.zeros(2, dtype=np.int64), + track_confidence=np.ones((8, 2)), + pose=rng.normal(size=(8, d)), + pose_valid=np.ones(8, dtype=bool), + ) + entries.append( + dict( + embodiment=embodiment, + episode_id=path.stem, + path=path.name, + frames=8, + split=split, + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + frame="fixed_reference", + pose_kind="measured_pose", + ) + ) + manifest = dict( + schema="mot-world-prepared-v1", + dino_identity={ + "model": "test-fixture", + "revision": "fixture", + "image_size": [2, 6], + "patch_grid": [1, 3], + "feature_layer": "last", + }, + normalization={ + "scope": "train_only", + "sha256": "0" * 64, + "arrays_normalized": True, + }, + pose_schemas={ + name: {"dimension": dim, "fields": [f"coordinate_{i}" for i in range(dim)]} + for name, dim in (("robot", 7), ("human", 11)) + }, + track_provenance={ + "interpolated_endpoints": False, + "method": "synthetic_fixture", + "reference_frame": "world", + "calibration_sha256": "0" * 64, + }, + episodes=entries, + ) + path = tmp_path / "manifest.json" + path.write_text(json.dumps(manifest)) + return path + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_hdf5_prepared_cache_reads_selected_frames_and_matches_npz(tmp_path, monkeypatch, mode): + import h5py + + path = cache_manifest(tmp_path) + kwargs = dict(split="val", mode=mode) + if mode == "continuous": + kwargs.update(max_past_seconds=3, max_future_seconds=3) + original = PreparedWorldDataset(path, **kwargs) + expected = [original[i] for i in range(len(original))] + manifest = json.loads(path.read_text()) + for ep in manifest["episodes"]: + target = (tmp_path / ep["path"]).with_suffix(".h5") + with np.load(tmp_path / ep["path"]) as data, h5py.File(target, "w") as out: + for name in data.files: + value = data[name] + out.create_dataset(name, data=value, chunks=(1, *value.shape[1:])) + ep.update(path=target.name, sha256=hashlib.sha256(target.read_bytes()).hexdigest()) + path.write_text(json.dumps(manifest)) + reads = [] + native = h5py.Dataset.__getitem__ + + def sliced_read(dataset, key): + if dataset.name == "/image": + assert isinstance(key, np.ndarray) + assert len(key) <= 4 + reads.append(key.copy()) + return native(dataset, key) + + monkeypatch.setattr(h5py.Dataset, "__getitem__", sliced_read) + actual = PreparedWorldDataset(path, **kwargs) + for i, reference in enumerate(expected): + result = actual[i] + assert result["embodiment"] == reference["embodiment"] + for stream, fields in reference["world"].items(): + for name, value in fields.items(): + assert torch.equal(result["world"][stream][name], value) + assert len(reads) == len(expected) + + +def test_prepared_cache_pipeline_and_reject_overlap(tmp_path): + path = cache_manifest(tmp_path) + dataset = PreparedWorldDataset(path) + sampler = HomogeneousWorldBatchSampler(dataset, batch_size=2, shuffle=False) + for indices in sampler: + batch = collate_world([dataset[i] for i in indices]) + embodiment, values = next(iter(batch.items())) + result = make_model().training_losses(values["world"], embodiment) + assert torch.isfinite(result["total"]) + manifest = json.loads(path.read_text()) + manifest["episodes"].append(manifest["episodes"][0]) + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="Duplicate"): + PreparedWorldDataset(path) + + +def test_validation_batch_cap_covers_both_embodiments_and_preserves_training(tmp_path): + from egomimic.rldb.world_model_datamodule import WorldModelDataModule + path = cache_manifest(tmp_path) + full = WorldModelDataModule(path, batch_size=2) + capped = WorldModelDataModule(path, batch_size=2, validation_batches_per_embodiment=2) + full.setup() + capped.setup() + validation = capped.val_dataloader() + batches = list(validation) + labels = [next(iter(batch)) for batch in batches] + assert len(validation) == len(batches) == 4 + assert {label: labels.count(label) for label in labels} == {"robot": 2, "human": 2} + assert len(full.val_dataloader()) > len(validation) + assert list(full.train_dataloader().batch_sampler) == list(capped.train_dataloader().batch_sampler) + default = HomogeneousWorldBatchSampler(full.datasets["val"], batch_size=2, shuffle=False) + explicit = HomogeneousWorldBatchSampler(full.datasets["val"], batch_size=2, shuffle=False, + validation_batches_per_embodiment=None) + assert len(default) == len(explicit) + assert list(default) == list(explicit) + + +def test_validation_batch_cap_keeps_short_final_batch(): + from types import SimpleNamespace + dataset = SimpleNamespace(split="val", episodes=[{"embodiment": "gripper"}, + {"embodiment": "sharpa"}], counts=[5, 1]) + sampler = HomogeneousWorldBatchSampler(dataset, batch_size=2, shuffle=False, + validation_batches_per_embodiment=2) + assert len(sampler) == 3 + assert list(sampler) == [[0, 1], [2, 3], [5]] + + +@pytest.mark.parametrize("cap", [0, -1, True, 1.5, "2"]) +def test_validation_batch_cap_rejects_invalid_value(cap): + from types import SimpleNamespace + from egomimic.rldb.world_model_datamodule import WorldModelDataModule + with pytest.raises(ValueError, match="positive integer"): + HomogeneousWorldBatchSampler(SimpleNamespace(split="val"), shuffle=False, + validation_batches_per_embodiment=cap) + with pytest.raises(ValueError, match="positive integer"): + WorldModelDataModule("unused.json", validation_batches_per_embodiment=cap) + + +@pytest.mark.parametrize("split,shuffle", [("train", False), ("train", True), ("val", True)]) +def test_validation_batch_cap_rejects_training_or_shuffle(split, shuffle): + from types import SimpleNamespace + with pytest.raises(ValueError, match="unshuffled validation"): + HomogeneousWorldBatchSampler(SimpleNamespace(split=split), shuffle=shuffle, + validation_batches_per_embodiment=2) + + +def test_bounded_dataset_filters_isolated_anchors_and_restores_context(tmp_path): + from egomimic.rldb.world_model_datamodule import WorldModelDataModule + path = cache_manifest(tmp_path) + manifest = json.loads(path.read_text()) + for episode in manifest["episodes"]: + file = path.parent / episode["path"] + with np.load(file, allow_pickle=False) as saved: + arrays = {k: saved[k] for k in saved.files} + arrays["time"] = np.array([0, 10, 20, 30, 100, 100.1, 100.2, 100.3]) + np.savez(file, **arrays) + episode["sha256"] = hashlib.sha256(file.read_bytes()).hexdigest() + path.write_text(json.dumps(manifest)) + dataset = PreparedWorldDataset(path, split="val", max_past_seconds=3, max_future_seconds=3) + assert dataset.counts == [4, 4] + for anchors in dataset.anchor_indices: + np.testing.assert_array_equal(anchors, [4, 5, 6, 7]) + for index in range(len(dataset)): + sample = dataset[index] + episode = dataset.episodes[index // 4] + with np.load(path.parent / episode["path"], allow_pickle=False) as saved: + expected_anchor_pose = saved["pose"][4 + index % 4].astype(np.float32) + poses = sample["world"]["proprio"]["value"].numpy() + assert any(np.array_equal(p, expected_anchor_pose) for p in poses) + assert sample["world"]["proprio"]["time"].abs().max() < .31 + data = WorldModelDataModule(str(path), max_past_seconds=3, max_future_seconds=3) + context = data.prepare_context(mode="train", normalization=None, normalizer=None) + assert context.state["sampling"]["max_past_seconds"] == 3 + assert context.state["sampling"]["max_future_seconds"] == 3 + changed = WorldModelDataModule(str(path), max_past_seconds=2, max_future_seconds=3) + with pytest.raises(ValueError, match="checkpoint/data context mismatch"): + changed.prepare_context(mode="train", normalization=None, normalizer=None, + restored_state=context.state) + + +def test_reject_interpolated_track_targets(tmp_path): + path = cache_manifest(tmp_path) + manifest = json.loads(path.read_text()) + manifest["track_provenance"]["interpolated_endpoints"] = True + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="interpolation"): + PreparedWorldDataset(path) + + +def test_bf16_mixed_finite_gradients(): + model = make_model() + with torch.autocast("cpu", dtype=torch.bfloat16): + loss = model.training_losses(make_batch(), "robot")["total"] + loss.backward() + assert torch.isfinite(loss) + assert all( + torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None + ) + + +def test_hydra_both_recipes_and_required_schema(): + from pathlib import Path + + from hydra import compose, initialize_config_dir + from omegaconf import OmegaConf + + root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + for mode in ("continuous", "discrete"): + with initialize_config_dir(config_dir=str(root), version_base=None): + cfg = compose( + config_name="train_zarr_cartesian", + overrides=[f"+experiment=world_model/{mode}", "hydra/launcher=basic"], + ) + assert cfg.data.mode == mode + assert cfg.data.min_frames == 2 + assert cfg.data.max_frames == 16 + assert cfg.data.discrete_frames == 4 + assert cfg.model.pipeline.stages[0].model.max_tokens == 8192 + assert cfg.model.pipeline.stages[0].model.time_mode == mode + assert cfg.trainer.check_val_every_n_epoch is None + assert cfg.callbacks.model_checkpoint.every_n_train_steps == 15000 + expected_bound = 3.0 if mode == "continuous" else None + assert cfg.data.max_past_seconds == expected_bound + assert cfg.data.max_future_seconds == expected_bound + assert OmegaConf.is_missing(cfg.model.pipeline.stages[0].model, "pose_dims") + assert cfg.evaluator._target_.endswith("WorldModelEval") + + +def test_shared_lightning_fit_validate_and_resume(tmp_path): + import inspect + from lightning import Trainer + from omegaconf import OmegaConf + + from egomimic.eval.world_model_eval import WorldModelEval + from egomimic.pl_utils.pl_model import ModelWrapper + from egomimic.rldb.world_model_datamodule import WorldModelDataModule + + path = cache_manifest(tmp_path) + data = WorldModelDataModule(str(path), batch_size=2, max_frames=3) + context = data.prepare_context(mode="train", normalization=None, normalizer=None) + with pytest.raises(ValueError, match="mismatch"): + data.prepare_context( + mode="train", normalization=None, normalizer=None, restored_state={} + ) + cfg = OmegaConf.create( + { + "model": { + "pipeline": { + "_target_": "egomimic.pipeline.algo.PipelineAlgo", + "device": "cpu", + "stages": [ + { + "_target_": "egomimic.pipeline.stages_world_model.WorldModelStage", + "warmup_steps": 0, + "pose_ramp_steps": 0, + "sampler_steps": 2, + "model": { + "_target_": "egomimic.models.world_model.MoTWorldModel", + "pose_dims": {"robot": 7, "human": 11}, + "image_dim": 8, + "width": 24, + "heads": 3, + "depth": 1, + "max_tracks": 8, + }, + } + ], + }, + "training_behavior": { + "_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior" + }, + "optimizer": { + "_target_": "torch.optim.AdamW", + "_partial_": True, + "lr": 0.001, + }, + } + } + ) + wrapper = ModelWrapper(config_tree=cfg, evaluator=WorldModelEval()) + wrapper.data_context = context + context.bind(wrapper.model, wrapper.evaluator) + trainer = Trainer( + accelerator="cpu", + devices=1, + max_steps=2, + max_epochs=1, + limit_val_batches=1, + val_check_interval=2, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + num_sanity_val_steps=0, + ) + trainer.fit(wrapper, datamodule=data) + assert trainer.global_step == 2 + assert any(k.startswith("val/") for k in trainer.callback_metrics) + checkpoint = tmp_path / "model.ckpt" + trainer.save_checkpoint(checkpoint) + restored = ModelWrapper(config_tree=cfg, evaluator=WorldModelEval()) + restored.data_context = context + trainer2 = Trainer( + accelerator="cpu", + devices=1, + max_steps=3, + max_epochs=2, + limit_val_batches=0, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + num_sanity_val_steps=0, + ) + # This is our own freshly saved full-state checkpoint. Match trainHydra's + # explicit restore mode on Lightning versions exposing weights_only. + load_options = {"weights_only": False} if "weights_only" in inspect.signature(trainer2.fit).parameters else {} + trainer2.fit(restored, datamodule=data, ckpt_path=str(checkpoint), **load_options) + assert trainer2.global_step == 3 + + +def test_invalid_unobserved_track_targets_are_masked(): + model = make_model() + batch = make_batch() + batch["track"]["value"][:, -1] = float("nan") + batch["track"]["target_valid"][:, -1] = False + batch["track"]["confidence"][:, -1] = 0 + with pytest.raises(ValueError, match="zero-fill is forbidden"): + model.training_losses(batch, "robot") + loss = model.training_losses( + batch, "robot", missing_track_policy="conditional_sample", imputation_steps=2 + )["total"] + assert torch.isfinite(loss) + loss.backward() + assert all( + torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None + ) + + +def test_noise_levels_are_per_frame_and_observed_are_clean(): + stream = make_batch()["image"] + levels = MoTWorldModel.frame_noise_levels(stream) + assert torch.equal(levels[:, :3], torch.zeros_like(levels[:, :3])) + assert torch.equal(levels[:, 3], levels[:, 4]) + assert torch.equal(levels[:, 4], levels[:, 5]) + assert ((levels[:, 3:] > 0) & (levels[:, 3:] < 1)).all() + + +def test_arbitrary_conditioning_includes_backcasting_and_endpoint_interpolation( + tmp_path, +): + torch.manual_seed(42) + dataset = PreparedWorldDataset( + cache_manifest(tmp_path), + min_frames=4, + max_frames=4, + forecasting_probability=0, + ) + backcast = interpolation = False + for i in range(100): + stream = dataset[i % len(dataset)]["world"]["image"] + times, mask = stream["time"], stream["observed"] + unique = torch.unique(times, sorted=True) + observed = torch.stack([mask[times == t].all() for t in unique]) + assert observed.any() and (~observed).any() + backcast |= bool(~observed[0] & observed[-1]) + interpolation |= bool(observed[0] & observed[-1] & (~observed[1:-1]).any()) + assert backcast, "No image query before the earliest image observation" + assert interpolation, "No withheld image between observed endpoint images" + + +def test_fractional_confidence_mse_is_a_weighted_mean(): + from egomimic.models.world_model import masked_mse + + pred = torch.tensor([[[2.0], [4.0]]], requires_grad=True) + target = torch.zeros_like(pred) + weight = torch.tensor([[0.1, 0.2]]) + loss = masked_mse(pred, target, weight) + torch.testing.assert_close(loss, torch.tensor(12.0)) + torch.testing.assert_close(loss, masked_mse(pred, target, weight * 0.1)) + loss.backward() + assert torch.isfinite(pred.grad).all() + pred.grad = None + empty = masked_mse(pred, torch.full_like(target, float("nan")), weight * 0) + assert empty.item() == 0 + empty.backward() + assert torch.equal(pred.grad, torch.zeros_like(pred)) + + +def test_declared_dino_patch_grid_must_match_cache(tmp_path): + path = cache_manifest(tmp_path) + manifest = json.loads(path.read_text()) + manifest["dino_identity"]["patch_grid"] = [2, 3] + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="patch grid"): + PreparedWorldDataset(path)[0] + + +@pytest.mark.parametrize("shape", [[0, 3], [True, 3], [3], "1x3", [1.0, 3]]) +def test_declared_dino_grid_requires_positive_integer_shape(tmp_path, shape): + path = cache_manifest(tmp_path) + manifest = json.loads(path.read_text()) + manifest["dino_identity"]["patch_grid"] = shape + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="two positive integers"): + PreparedWorldDataset(path) + + +def test_complete_tracks_bypass_imputation_and_preserve_rng(monkeypatch): + model, batch = make_model(), make_batch() + + def forbidden(*args, **kwargs): + raise AssertionError("Complete tracks must not invoke a teacher") + + monkeypatch.setattr(model, "sample", forbidden) + torch.manual_seed(19) + a = model.training_losses(batch, "robot") + rng = torch.get_rng_state().clone() + torch.manual_seed(19) + b = model.training_losses(batch, "robot", missing_track_policy="conditional_sample") + assert torch.equal(rng, torch.get_rng_state()) + for key in a: + torch.testing.assert_close(a[key], b[key], rtol=0, atol=0) + assert a["imputed_track_tokens"] == a["imputation_forwards"] == 0 + + +def test_missing_tracks_use_detached_completion_not_zero_geometry(monkeypatch): + from egomimic.models.world_model import noise_coefficients + + model, batch = make_model(), make_batch() + batch["track"]["value"][:, -1] = float("nan") + batch["track"]["target_valid"][:, -1] = False + batch["track"]["confidence"][:, -1] = 0 + before = copy.deepcopy(batch) + completed = torch.full_like(batch["track"]["value"], 7.0, requires_grad=True) + + def teacher(teacher_batch, embodiment, steps): + assert not torch.is_grad_enabled() + assert steps == 2 + for name in ("image", "track"): + assert torch.isfinite(teacher_batch[name]["value"]).all() + pose = teacher_batch["proprio"] + assert torch.equal(pose["observed"], batch["proprio"]["observed"]) + assert (pose["value"][~pose["observed"]] == 0).all() + assert not teacher_batch["track"]["observed"][:, -1].any() + return {"track": completed} + + monkeypatch.setattr(model, "sample", teacher) + monkeypatch.setattr(torch, "randn_like", lambda value: torch.zeros_like(value)) + calls = [] + original = model.forward + + def capture(b, e, values, levels): + calls.append((values, levels)) + return original(b, e, values, levels) + + monkeypatch.setattr(model, "forward", capture) + loss = model.training_losses( + batch, "robot", missing_track_policy="conditional_sample", imputation_steps=2 + ) + values, levels = calls[0] + alpha, _ = noise_coefficients(levels["track"]) + torch.testing.assert_close(values["track"][:, -1], 7 * alpha[:, -1].expand(-1, 3)) + assert torch.equal(levels["track"][:, -1], levels["track"][:, -2]) + assert loss["imputed_track_tokens"] == 2 + assert loss["imputation_forwards"] == 3 + loss["total"].backward() + assert completed.grad is None + assert model.track_out.weight.grad.norm() > 0 + for name in batch: + for key in batch[name]: + torch.testing.assert_close( + batch[name][key], before[name][key], equal_nan=True + ) + + +@pytest.mark.parametrize("bf16", [False, True]) +def test_conditional_completion_ignores_placeholders_and_withheld_poses(bf16): + model, batch = make_model(), make_batch() + # Even a finite coordinate with zero confidence must not become a target. + batch["track"]["confidence"][:, -1] = 0 + batch["track"]["value"][:, -1] = float("nan") + other = copy.deepcopy(batch) + other["track"]["value"][:, -1] = 1e9 + other["proprio"]["value"][~other["proprio"]["observed"]] = float("nan") + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=bf16): + torch.manual_seed(77) + a = model.training_losses( + batch, + "robot", + pose_weight=0, + missing_track_policy="conditional_sample", + imputation_steps=2, + ) + torch.manual_seed(77) + b = model.training_losses( + other, + "robot", + pose_weight=0, + missing_track_policy="conditional_sample", + imputation_steps=2, + ) + for key in a: + torch.testing.assert_close(a[key], b[key], rtol=0, atol=0) + a["total"].backward() + assert all( + torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None + ) + + +def test_training_rejects_unavailable_observations_and_image_targets(): + model, batch = make_model(), make_batch() + batch["track"]["confidence"][:, 0] = 0 + with pytest.raises(ValueError, match="Observed visual"): + model.training_losses(batch, "robot", missing_track_policy="conditional_sample") + batch = make_batch() + batch["image"]["target_valid"][:, -1] = False + with pytest.raises(ValueError, match="Missing image targets"): + model.training_losses(batch, "robot", missing_track_policy="conditional_sample") + + +def test_pipeline_exposes_missing_track_policy_and_cost(): + model, world = make_model(), make_batch() + world["track"]["target_valid"][:, -1] = False + world["track"]["value"][:, -1] = float("nan") + stage = WorldModelStage( + model, missing_track_policy="conditional_sample", imputation_steps=2 + ) + result = stage.execute( + dict(world=world, embodiment="robot", optimizer_step=0), mode="train" + ) + assert torch.isfinite(result["loss/world"]) + assert result["log/world_imputed_track_tokens"] == 2 + assert result["log/world_imputation_forwards"] == 3 diff --git a/tests/test_world_model_cache.py b/tests/test_world_model_cache.py new file mode 100644 index 0000000000..c24d690da8 --- /dev/null +++ b/tests/test_world_model_cache.py @@ -0,0 +1,154 @@ +"""Streaming producer contracts; fake frozen features, no pretrained download.""" +import json +from pathlib import Path + +import h5py +import numpy as np +import pytest +import torch + +from egomimic.rldb.world_model import open_prepared_episode +from egomimic.scripts.embedding_process.world_model_cache import extract_episode, sha256 + + +class FakeDino: + dtype = torch.float32 + device = torch.device("cpu") + image_side = 32 + + def __init__(self): + self.model = torch.nn.Linear(1, 1).eval().requires_grad_(False) + self.calls = [] + + def _embed_batch(self, images): + self.calls.append(len(images)) + assert all(image.shape == (32, 32, 3) for image in images) + self.last_preprocessing = { + "processed_hw": [32, 32], "patch_grid": [2, 2], + "precision": "torch.float32", "feature_layer": "last_hidden_state_patch_tokens", + "frames": [{"original_hw": [32, 32]} for _ in images], + } + return np.stack([np.full((4, 768), image[0, 0, 0], np.float32) for image in images]) + + +def fixture(tmp_path, side=32): + native = tmp_path / "native" + native.mkdir() + times = np.array([0., .05, .1], np.float64) + arrays = { + "time": times, "source_row": np.arange(3), + "pose": np.zeros((3, 9)), "xyz": np.zeros((3, 2, 3)), + "coordinate_available": np.ones((3, 2), bool), + "reference_grid_xy": np.array([[0, 0], [127, 127]]), + "valid_grid_indices": np.arange(2), "visibility": np.zeros((3, 2, 2), bool), + } + arrays["pose"][:, 3] = 1 + np.savez(native / "aligned.npz", **arrays) + manifest = {"schema": "libero-same-recorded-state-v1", "cameras": ["front", "wrist"], + "pose_fields": [f"pose{i}" for i in range(9)], "reference_frame": "recorded simulator world", + "height": 128, "width": 128, + "artifacts": {"aligned.npz": sha256(native / "aligned.npz")}} + (native / "manifest.json").write_text(json.dumps(manifest)) + rgb = tmp_path / "rgb.h5" + with h5py.File(rgb, "w") as stream: + pixels = np.zeros((3, 2, side, side, 3), np.uint8) + for t in range(3): + for camera in range(2): + pixels[t, camera] = t * 10 + camera + stream["rgb"] = pixels + stream["time"] = times + receipt = {"schema": "libero-native-rgb-v1", "status": "PASS", "native_render": True, + "native_geometry_replay_exact": True, + "source_manifest_sha256": sha256(native / "manifest.json"), + "source_aligned_sha256": sha256(native / "aligned.npz"), + "rgb_cache_sha256": sha256(rgb), "height": side, "width": side, + "cameras": manifest["cameras"], "frames": 3} + receipt_path = tmp_path / "rgb.json" + receipt_path.write_text(json.dumps(receipt)) + return dict(native_cache=native, rgb_cache=rgb, rgb_receipt=receipt_path, + output=tmp_path / "features.h5", embodiment="gripper", episode_id="ep1", + embedder=FakeDino(), batch_size=2, + model_identity={"model": "fake-test-only", "revision": "abc", "weights_sha256": "a" * 64, + "config_sha256": "b" * 64, "processor_sha256": "c" * 64}) + + +def test_streaming_fp32_cache_geometry_identity_and_reuse(tmp_path): + args = fixture(tmp_path) + result = extract_episode(**args) + assert result["status"] == "PASS" and not result["training_ready"] + assert result["moments"]["image"]["count"] == 24 + assert result["moments"]["track"]["count"] == 6 + assert result["moments"]["pose"]["count"] == 3 + assert args["embedder"].calls == [2, 1, 2, 1] + with open_prepared_episode(args["output"]) as saved: + assert saved["image"].shape == (3, 2, 4, 768) + assert saved["image"].dtype == np.float32 and saved["image"].chunks[0] == 1 + np.testing.assert_array_equal(saved["image"][:, :, 0, 0], [[0, 1], [10, 11], [20, 21]]) + np.testing.assert_array_equal(saved["track_confidence"][:], np.ones((3, 2))) + assert not saved["visibility"][:].any() # occluded does not mean unavailable + np.testing.assert_array_equal(saved["track_xy"][:], [[-1, -1], [1, 1]]) + np.testing.assert_array_equal(saved["image_xy"][:], [[-.5, -.5], [.5, -.5], [-.5, .5], [.5, .5]]) + assert not saved.attrs["arrays_normalized"] + values = saved["image"][:].reshape(-1, 768).astype(np.float64) + np.testing.assert_allclose(result["moments"]["image"]["mean"], values.mean(0)) + np.testing.assert_allclose(result["moments"]["image"]["m2"], ((values - values.mean(0)) ** 2).sum(0)) + args["embedder"].calls.clear() + assert extract_episode(**args)["reused"] + assert args["embedder"].calls == [] + changed = {**args, "model_identity": {**args["model_identity"], "processor_sha256": "d" * 64}} + with pytest.raises(ValueError, match="request differs"): + extract_episode(**changed) + with h5py.File(args["output"], "a") as saved: + saved["image"][0, 0, 0, 0] += 1 + with pytest.raises(ValueError, match="hash mismatch"): + extract_episode(**args) + + +def test_reject_upscaled_native_image(tmp_path): + args = fixture(tmp_path, side=16) + with pytest.raises(ValueError, match="Native RGB resolution"): + extract_episode(**args) + assert args["embedder"].calls == [] + assert not args["output"].exists() + + +def test_reject_missing_replay_proof(tmp_path): + args = fixture(tmp_path) + receipt = json.loads(args["rgb_receipt"].read_text()) + receipt["native_geometry_replay_exact"] = False + args["rgb_receipt"].write_text(json.dumps(receipt)) + with pytest.raises(ValueError, match="exact geometry receipt"): + extract_episode(**args) + + +def test_reject_rgb_time_mismatch(tmp_path): + args = fixture(tmp_path) + with h5py.File(args["rgb_cache"], "a") as saved: + saved["time"][1] += .01 + receipt = json.loads(args["rgb_receipt"].read_text()) + receipt["rgb_cache_sha256"] = sha256(args["rgb_cache"]) + args["rgb_receipt"].write_text(json.dumps(receipt)) + with pytest.raises(AssertionError, match="time alignment"): + extract_episode(**args) + + +@pytest.mark.parametrize("failure", ["nan", "float16", "grid"]) +def test_invalid_feature_batch_never_publishes(tmp_path, failure): + args = fixture(tmp_path) + original = args["embedder"]._embed_batch + + def damaged(images): + values = original(images) + if failure == "nan": + values[0, 0, 0] = np.nan + elif failure == "float16": + values = values.astype(np.float16) + else: + args["embedder"].last_preprocessing["patch_grid"] = [3, 3] + return values + + args["embedder"]._embed_batch = damaged + with pytest.raises(ValueError, match="Invalid DINO"): + extract_episode(**args) + assert not args["output"].exists() + assert not Path(str(args["output"]) + ".json").exists() diff --git a/tests/test_world_model_cohort_inputs.py b/tests/test_world_model_cohort_inputs.py new file mode 100644 index 0000000000..3c23e75396 --- /dev/null +++ b/tests/test_world_model_cohort_inputs.py @@ -0,0 +1,69 @@ +import hashlib +import json + +import pytest +from PIL import Image + +from egomimic.scripts.embedding_process.dinov3_correspondence import load_manifest +from egomimic.scripts.embedding_process.world_model_pose_probe import _reference_pairs + + +def test_explicit_pairs_preserve_selected_horizons(): + records = [ + dict(episode="e", frame=i, obs_rgb_timestamps_ns=i * 100) for i in [0, 1, 5, 10] + ] + explicit = [dict(episode="e", source_frame=0, target_frame=i) for i in [1, 5, 10]] + assert _reference_pairs(records, explicit) == [(0, 1), (0, 2), (0, 3)] + with pytest.raises(ValueError, match="positive physical time"): + _reference_pairs(records, [dict(episode="e", source_frame=5, target_frame=1)]) + with pytest.raises(ValueError, match="unique"): + _reference_pairs(records, explicit + explicit[:1]) + + +def test_unlabelled_manifest_only_for_features(tmp_path): + path = tmp_path / "frame.png" + Image.new("RGB", (32, 32)).save(path) + manifest = dict( + model="official", + revision="a" * 40, + frames=[ + dict( + id="a", + path=str(path), + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + points={}, + episode="e", + ) + ], + pairs=[], + ) + f = tmp_path / "manifest.json" + f.write_text(json.dumps(manifest)) + assert len(load_manifest(f, features_only=True)[1]) == 1 + with pytest.raises(ValueError, match="At least one pair"): + load_manifest(f) + + +def test_timestamp_selection_rejects_filtered_gaps(): + import numpy as np + + from egomimic.scripts.embedding_process.world_model_pose_probe import ( + select_contiguous_anchors, + validate_raw_aria_wrist, + ) + + timestamps = np.arange(1200, dtype=np.int64) * 33333333 + timestamps[500:] += 8_000_000_000 + anchors = select_contiguous_anchors(timestamps) + assert len(anchors) == 4 + for a, targets in anchors: + assert not (a < 500 <= targets[-1]) + assert np.allclose( + (timestamps[targets] - timestamps[a]) * 1e-9, [0.1, 0.5, 1], atol=0.02 + ) + points = np.zeros((2, 21, 3)) + wrists = np.zeros((2, 3)) + assert validate_raw_aria_wrist(points, wrists) == 0 + points[:, 5, 0] = 0.1 + with pytest.raises(ValueError, match="index5 mismatch"): + validate_raw_aria_wrist(points, wrists) diff --git a/tests/test_world_model_image_parameterization.py b/tests/test_world_model_image_parameterization.py new file mode 100644 index 0000000000..6d508b14cc --- /dev/null +++ b/tests/test_world_model_image_parameterization.py @@ -0,0 +1,133 @@ +"""Image x0 ablation: identical architecture, explicit target/solver semantics.""" +import copy + +import pytest +import torch + +from egomimic.models.world_model import masked_mse +from tests.test_world_model import make_batch, make_model + + +@pytest.mark.parametrize("bf16", [False, True]) +def test_explicit_v_preserves_implicit_default_losses_gradients_rng(bf16): + torch.manual_seed(42) + old = make_model() + new = make_model(image_prediction_type="v") + new.load_state_dict(old.state_dict(), strict=True) + batch = make_batch() + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=bf16): + torch.manual_seed(17) + expected = old.training_losses(batch, "robot") + state = torch.get_rng_state() + torch.manual_seed(17) + actual = new.training_losses(batch, "robot") + assert torch.equal(state, torch.get_rng_state()) + for key in expected: + torch.testing.assert_close(expected[key], actual[key], rtol=0, atol=0) + expected["total"].backward() + actual["total"].backward() + for a, b in zip(old.parameters(), new.parameters()): + assert (a.grad is None) == (b.grad is None) + if a.grad is not None: + torch.testing.assert_close(a.grad, b.grad, rtol=0, atol=0) + + +@pytest.mark.parametrize("steps", [1, 5, 20]) +def test_x0_oracle_cancels_unknown_noise_and_preserves_observations(monkeypatch, steps): + model = make_model(image_prediction_type="x0") + batch = make_batch() + target = batch["image"]["value"].clone() + batch["image"]["value"][~batch["image"]["observed"]] = float("nan") + checks = [] + + def oracle(b, embodiment, values, levels): + observed = b["image"]["observed"] + checks.append(torch.equal(values["image"][observed], target[observed])) + assert torch.isfinite(values["image"]).all() + return {"image": target, "track": torch.zeros_like(values["track"]), + "proprio": torch.zeros_like(b["proprio"]["value"])} + + monkeypatch.setattr(model, "forward", oracle) + result = model.sample(batch, "robot", steps=steps) + torch.testing.assert_close(result["image"], target, rtol=0, atol=0) + assert all(checks) + + +def test_x0_objective_is_clean_mse_and_track_remains_v(monkeypatch): + model = make_model(image_prediction_type="x0") + batch = make_batch() + original = model.forward + + def zero_image(*args): + output = original(*args) + output["image"] = torch.zeros_like(output["image"]) + return output + + monkeypatch.setattr(model, "forward", zero_image) + torch.manual_seed(7) + losses = model.training_losses(batch, "robot", pose_weight=0) + mask = batch["image"]["valid"] & ~batch["image"]["observed"] + expected = masked_mse(torch.zeros_like(batch["image"]["value"]), + batch["image"]["value"], mask.float()) + torch.testing.assert_close(losses["image"], expected, rtol=0, atol=0) + model.image_prediction_type = "v" + torch.manual_seed(7) + v_loss = model.training_losses(batch, "robot", pose_weight=0) + torch.testing.assert_close(losses["track"], v_loss["track"], rtol=0, atol=0) + + +@pytest.mark.parametrize("bf16", [False, True]) +def test_x0_pose_loss_reaches_both_generators(bf16): + model = make_model(image_prediction_type="x0") + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=bf16): + losses = model.training_losses(make_batch(), "robot") + losses["pose"].backward() + for module in (model.image_out, model.track_out): + assert module.weight.grad.norm() > 0 + assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None) + + +def test_x0_sampler_ignores_withheld_values_and_validity(): + model = make_model(image_prediction_type="x0") + batch = make_batch() + changed = copy.deepcopy(batch) + for stream in changed.values(): + stream["value"][~stream["observed"]] = float("nan") + stream["target_valid"].zero_() + if "confidence" in stream: + stream["confidence"].zero_() + torch.manual_seed(55) + expected = model.sample(batch, "robot", steps=3) + torch.manual_seed(55) + actual = model.sample(changed, "robot", steps=3) + for key in expected: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + + +def test_x0_missing_track_completion_still_uses_explicit_policy(): + model, batch = make_model(image_prediction_type="x0"), make_batch() + batch["track"]["target_valid"][:, -1] = False + batch["track"]["value"][:, -1] = float("nan") + with pytest.raises(ValueError, match="zero-fill"): + model.training_losses(batch, "robot") + loss = model.training_losses(batch, "robot", missing_track_policy="conditional_sample", imputation_steps=2) + assert torch.isfinite(loss["total"]) + loss["total"].backward() + assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None) + + +def test_invalid_image_parameterization_rejected(): + with pytest.raises(ValueError, match="image_prediction_type"): + make_model(image_prediction_type="epsilon") + + +@pytest.mark.parametrize("strength", [1e-40, 1e-20, 0.4]) +def test_x0_tiny_refinement_strength_does_not_overflow(strength): + model, batch = make_model(image_prediction_type="x0"), make_batch() + stream = batch["image"] + levels = torch.full_like(stream["time"], strength).masked_fill(stream["observed"], 0) + batch["refinement"] = {"image": {"strength": levels, "prediction": stream["value"].clone()}} + original = stream["value"].clone() + result = model.sample(batch, "robot", steps=3) + assert all(torch.isfinite(v).all() for v in result.values()) + torch.testing.assert_close(result["image"][stream["observed"]], original[stream["observed"]], rtol=0, atol=0) diff --git a/tests/test_world_model_pose_probe.py b/tests/test_world_model_pose_probe.py new file mode 100644 index 0000000000..37880b41f2 --- /dev/null +++ b/tests/test_world_model_pose_probe.py @@ -0,0 +1,221 @@ +"""Contracts for native MoT readout diagnostics; no quality claims from fixtures.""" + +import json +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from egomimic.models.world_model import MODALITIES +from egomimic.scripts.embedding_process.world_model_pose_probe import ( + fixed_partition, + fold_data, + make_batch, + model_for_probe, + predict, + run, + sha, +) + + +def inputs(): + rng = np.random.default_rng(42) + data = { + "features": rng.normal(size=(6, 4, 768)).astype(np.float32), + "patch_xy": np.zeros((4, 2), np.float32), + "patch_valid": np.ones(4, bool), + "pairs": np.array([[0, 1], [2, 3], [4, 5]]), + "poses": rng.normal(size=(3, 2, 126)).astype(np.float32), + "dt": np.array([0.5, 1, 2], np.float32), + } + meta = {"frames": [{"episode": str(i // 2)} for i in range(6)]} + return data, meta + + +def test_freeze_uses_named_track_branch_not_proprio(): + model = model_for_probe(width=16, depth=1) + track = MODALITIES.index("track") + pose = MODALITIES.index("proprio") + for block in model.blocks: + assert all(not p.requires_grad for p in block.qkv[track].parameters()) + assert all(p.requires_grad for p in block.qkv[pose].parameters()) + assert model.pose_out["human"].weight.requires_grad + + +def test_normalization_cannot_see_heldout_frames_or_poses(): + data, meta = inputs() + _, _, before = fold_data(data, meta, "2") + data["features"][4:] += 1000 + data["poses"][2] += 1000 + _, _, after = fold_data(data, meta, "2") + for key in before: + assert np.array_equal(before[key], after[key]) + + +def test_fixed_split_excludes_validation_and_test_from_normalization(): + data, meta = inputs() + for i, row in enumerate(meta["frames"]): + row["image_sha256"] = str(i) + split = {"train": ["0"], "validation": ["1"], "test": ["2"]} + indices, before = fixed_partition(data, meta, split) + assert {k: v.tolist() for k, v in indices.items()} == { + "train": [0], + "validation": [1], + "test": [2], + } + data["features"][2:] += 1000 + data["poses"][1:] += 1000 + _, after = fixed_partition(data, meta, split) + for key in before: + assert np.array_equal(before[key], after[key]) + + +def test_fixed_split_rejects_copied_images_and_cross_episode_pairs(): + data, meta = inputs() + for i, row in enumerate(meta["frames"]): + row["image_sha256"] = str(i) + split = {"train": ["0"], "validation": ["1"], "test": ["2"]} + meta["frames"][2]["image_sha256"] = "0" + with pytest.raises(ValueError, match="Image content overlaps"): + fixed_partition(data, meta, split) + meta["frames"][2]["image_sha256"] = "2" + data["pairs"][0] = [0, 2] + with pytest.raises(ValueError, match="crosses episodes"): + fixed_partition(data, meta, split) + + +def test_future_pose_is_loss_only_and_no_fake_tracks(): + data, meta = inputs() + _, _, norm = fold_data(data, meta, "2") + batch, target = make_batch(data, norm, [0]) + model = model_for_probe(width=16, depth=1).eval() + assert batch["track"]["value"].shape == (1, 0, 3) + assert torch.count_nonzero(batch["proprio"]["value"][:, 1]) == 0 + before = predict(model, batch) + data["poses"][0, 1] += 100 + changed, target2 = make_batch(data, norm, [0]) + assert not torch.equal(target, target2) + assert torch.equal(before, predict(model, changed)) + + +def test_observed_only_is_image_invariant_but_proprio_gradients_flow(): + data, meta = inputs() + _, _, norm = fold_data(data, meta, "2") + batch, target = make_batch(data, norm, [0], visual=False) + model = model_for_probe(width=16, depth=1) + before = predict(model, batch) + batch["image"]["value"] += 1000 + assert torch.equal(before, predict(model, batch)) + (before - target).square().mean().backward() + assert model.blocks[0].qkv[MODALITIES.index("proprio")].weight.grad.norm() > 0 + assert model.blocks[0].qkv[MODALITIES.index("track")].weight.grad is None + + +@pytest.mark.parametrize("explicit_split", [False, True]) +def test_preflight_exercises_both_backward_paths_without_optimizer_update( + tmp_path, monkeypatch, explicit_split +): + data, meta = inputs() + meta["episodes"] = ["0", "1", "2"] + for i, row in enumerate(meta["frames"]): + row["image_sha256"] = str(i) + split_file = tmp_path / "SPLIT.json" + split_file.write_text( + json.dumps({"train": ["0"], "validation": ["1"], "test": ["2"]}) + ) + np.savez(tmp_path / "INPUTS.npz", **data) + meta["input_sha256"] = sha(tmp_path / "INPUTS.npz") + (tmp_path / "INPUTS.json").write_text(json.dumps(meta)) + count = [] + original = torch.autograd.backward + + def backward(*args, **kwargs): + count.append(1) + return original(*args, **kwargs) + + def forbidden(*args, **kwargs): + raise AssertionError("Preflight must not update parameters") + + monkeypatch.setattr(torch.autograd, "backward", backward) + monkeypatch.setattr(torch.optim.AdamW, "step", forbidden) + run( + SimpleNamespace( + inputs=tmp_path, + output=tmp_path / "output", + threads=1, + steps=1, + save_every=1, + width=16, + depth=1, + device="cpu", + lr=1e-4, + preflight_only=True, + split_manifest=split_file if explicit_split else None, + ) + ) + assert len(count) == 2 + assert not list((tmp_path / "output").glob("fold-*")) + assert ( + json.loads((tmp_path / "output/RUNTIME_BACKWARD_GUARD.json").read_text())[ + "optimizer_updates" + ] + == 0 + ) + + +def test_selection_uses_validation_and_seals_test_until_selection( + tmp_path, monkeypatch +): + """Exercise orchestration without optimizer updates or a quality claim.""" + from egomimic.scripts.embedding_process import world_model_pose_probe as probe + + data, meta = inputs() + meta["episodes"] = ["0", "1", "2"] + for i, row in enumerate(meta["frames"]): + row["image_sha256"] = str(i) + np.savez(tmp_path / "INPUTS.npz", **data) + meta["input_sha256"] = sha(tmp_path / "INPUTS.npz") + (tmp_path / "INPUTS.json").write_text(json.dumps(meta)) + split_file = tmp_path / "SPLIT.json" + split_file.write_text( + json.dumps({"train": ["0"], "validation": ["1"], "test": ["2"]}) + ) + output = tmp_path / "output" + val_calls = {False: 0, True: 0} + + def evaluate(model, data, meta, norm, indices, device, visual): + if indices.tolist() == [1]: + val_calls[visual] += 1 + score = float(val_calls[visual]) # First checkpoint is best. + else: + condition = ( + "oracle_images_and_observed_pose" if visual else "observed_pose_only" + ) + selection = json.loads( + (output / "fold-0" / condition / "SELECTION.json").read_text() + ) + assert selection["selected_step"] == 1 + assert selection["test_used_for_selection"] is False + score = 3.0 + return [{"mpjpe_m": score, "hold_pose_mpjpe_m": 4.0}] + + monkeypatch.setattr(probe, "evaluate", evaluate) + monkeypatch.setattr(torch.optim.AdamW, "step", lambda *a, **k: None) + run( + SimpleNamespace( + inputs=tmp_path, + output=output, + threads=1, + steps=2, + save_every=1, + width=16, + depth=1, + device="cpu", + lr=1e-4, + preflight_only=False, + split_manifest=split_file, + ) + ) + assert val_calls == {False: 2, True: 2} + assert len(list(output.glob("fold-*/*/epoch*.pt"))) == 4 diff --git a/tests/test_world_model_preparation.py b/tests/test_world_model_preparation.py new file mode 100644 index 0000000000..80630f0a36 --- /dev/null +++ b/tests/test_world_model_preparation.py @@ -0,0 +1,490 @@ +"""Real sliced caches, train-only statistics, and one-setup integrity reuse.""" + +import hashlib +import json +import os +from pathlib import Path +import shutil +import time + +import h5py +import numpy as np +import pytest +import torch + +from egomimic.rldb.world_model import PreparedWorldDataset +from egomimic.rldb.world_model_datamodule import WorldModelDataModule +from egomimic.rldb.world_model_preparation import ( + ArtifactVerificationMemo, finalize, StreamingMoments, +) +from egomimic.rldb.zarr.episode_split import split_dataset_names + + +def aged_verification_clock(monkeypatch): + """Simulate already-immutable files without sleeping in the test suite.""" + now = time.time_ns() + 2_000_000_000 + monkeypatch.setattr(ArtifactVerificationMemo, "_now_ns", staticmethod(lambda: now)) + + +def make_inventory(tmp_path): + paths, schemas = [], {} + for name, dimension in (("robot", 3), ("human", 4)): + fields = [f"pose_{i}" for i in range(dimension)] + schemas[name] = {"dimension": dimension, "fields": fields} + _, validation = split_dataset_names(["a", "b", "c"], .01, 42) + for episode in ("a", "b", "c"): + file = tmp_path / f"{name}-{episode}.h5" + # Huge held-out values make accidental leakage into moments obvious. + offset = 10000 if episode in validation else (1 if episode == "a" else 3) + with h5py.File(file, "w") as cache: + cache["time"] = np.arange(4, dtype=np.float64) * .1 + cache.create_dataset("image", data=np.ones((4, 1, 2, 2), np.float32) * offset, + chunks=(1, 1, 2, 2)) + cache["image_xy"] = np.zeros((2, 2), np.float32) + cache["track"] = np.ones((4, 2, 3), np.float32) * offset + cache["track_xy"] = np.zeros((2, 2), np.float32) + cache["track_camera"] = np.zeros(2, np.int64) + cache["track_confidence"] = np.ones((4, 2), np.float32) + cache["pose"] = np.ones((4, dimension), np.float32) * offset + cache["pose_valid"] = np.ones(4, bool) + receipt = Path(str(file) + ".json") + receipt.write_text(json.dumps({ + "schema": "mot-world-feature-cache-v1", "status": "PASS", "arrays_normalized": False, + "output_sha256": hashlib.sha256(file.read_bytes()).hexdigest(), "embodiment": name, + "episode_id": episode, "frames": 4, "pose_fields": fields, "cameras": ["front"], + "native_source": {"source_demo_sha256": hashlib.sha256(f"{name}-{episode}".encode()).hexdigest(), + "source_demo": str(file), "episode": "demo_0"}, + "dino_identity": {"model": "fixture", "revision": "fixture", "image_size": [1, 2], + "patch_grid": [1, 2], "feature_layer": "last"}})) + paths.append(receipt.name) + inventory = tmp_path / "inventory.json" + inventory.write_text(json.dumps({"receipts": paths, "pose_schemas": schemas, + "track_provenance": {"method": "fixture", "reference_frame": "world", + "calibration_sha256": "0" * 64, + "interpolated_endpoints": False}})) + return inventory + + +def test_streaming_moments_match_full_population(): + values = np.random.default_rng(4).normal(size=(31, 3)) + moments = StreamingMoments() + for batch in np.array_split(values, 7): + moments.update(batch) + result = moments.result(1e-5) + np.testing.assert_allclose(result["mean"], values.mean(0), atol=1e-15) + np.testing.assert_allclose(result["std"], values.std(0), atol=1e-15) + + +def test_finalizer_merges_bound_train_moments_without_rereading_images(tmp_path, monkeypatch): + inventory = make_inventory(tmp_path) + reference_path = tmp_path / "reference.json" + reference = finalize(inventory, reference_path) + expected = json.loads((tmp_path / reference["normalization"]["path"]).read_text()) + for receipt_name in json.loads(inventory.read_text())["receipts"]: + receipt_path = tmp_path / receipt_name + receipt = json.loads(receipt_path.read_text()) + with h5py.File(str(receipt_path)[:-5]) as cache: + receipt["moments"] = {} + for key in ("image", "track", "pose"): + moments = StreamingMoments() + moments.update(cache[key][:]) + receipt["moments"][key] = {"count": moments.count, "mean": moments.mean.tolist(), + "m2": moments.m2.tolist()} + receipt_path.write_text(json.dumps(receipt)) + original = h5py.Dataset.__getitem__ + + def read(dataset, index): + assert dataset.name != "/image", "Bound producer moments should avoid rereading image payloads" + return original(dataset, index) + + monkeypatch.setattr(h5py.Dataset, "__getitem__", read) + output = tmp_path / "merged.json" + actual = finalize(inventory, output) + observed = json.loads((tmp_path / actual["normalization"]["path"]).read_text()) + assert {k: v for k, v in observed.items() if k != "statistics"} == { + k: v for k, v in expected.items() if k != "statistics"} + for key, actual_stats in observed["statistics"].items(): + expected_stats = expected["statistics"][key] + assert actual_stats["count"] == expected_stats["count"] + for field in ("mean", "std"): + np.testing.assert_allclose(actual_stats[field], expected_stats[field], rtol=1e-14, atol=1e-14) + + +def test_finalizer_has_no_validation_leak_and_normalizes_selected_frames(tmp_path, monkeypatch): + inventory = make_inventory(tmp_path) + output = tmp_path / "prepared.json" + manifest = finalize(inventory, output) + assert manifest["split_provenance"]["complete_corpus_coverage"] + stats = json.loads((tmp_path / manifest["normalization"]["path"]).read_text()) + assert all(np.max(np.abs(s["mean"])) < 10 for s in stats["statistics"].values()) + # Same inventory finalization is immutable and deterministic. + assert finalize(inventory, output) == manifest + calls = [] + original = h5py.Dataset.__getitem__ + + def read(dataset, index): + if dataset.name == "/image": + calls.append(index) + return original(dataset, index) + + monkeypatch.setattr(h5py.Dataset, "__getitem__", read) + dataset = PreparedWorldDataset(output, split="val", min_frames=2, max_frames=2) + sample = dataset[0] + image = sample["world"]["image"]["value"] + assert torch.isfinite(image).all() and image.min() > 100 + assert len(calls) == 1 and isinstance(calls[0], np.ndarray) and len(calls[0]) == 2 + + +def test_setup_hashes_payloads_once_and_rejects_changed_memo_artifact(tmp_path, monkeypatch): + inventory = make_inventory(tmp_path) + output = tmp_path / "prepared.json" + manifest = finalize(inventory, output) + # Historical manifests retain their full-hash behavior. + for episode in manifest["episodes"]: + episode.pop("verified_fingerprint") + output.write_text(json.dumps(manifest)) + aged_verification_clock(monkeypatch) + count = [] + original = hashlib.file_digest + + def digest(stream, *args, **kwargs): + count.append(stream.name) + return original(stream, *args, **kwargs) + + monkeypatch.setattr(hashlib, "file_digest", digest) + data = WorldModelDataModule(output, min_frames=2, max_frames=2) + data.setup() + assert len(count) == len(manifest["episodes"]) + 1 + memo = ArtifactVerificationMemo() + artifact = tmp_path / manifest["episodes"][0]["path"] + sha = manifest["episodes"][0]["sha256"] + memo.verify(artifact, sha) + with artifact.open("ab") as stream: + stream.write(b"changed") + with pytest.raises(ValueError, match="hash mismatch"): + memo.verify(artifact, sha) + + +def test_finalized_fingerprints_skip_payload_reads_across_fresh_setups(tmp_path, monkeypatch): + inventory = make_inventory(tmp_path) + aged_verification_clock(monkeypatch) + output = tmp_path / "prepared.json" + calls, original = [], hashlib.file_digest + + def digest(stream, *args, **kwargs): + calls.append(Path(stream.name).resolve()) + return original(stream, *args, **kwargs) + + monkeypatch.setattr(hashlib, "file_digest", digest) + manifest = finalize(inventory, output) + assert len(calls) == len(manifest["episodes"]) + calls.clear() + for _ in range(2): + WorldModelDataModule(output, min_frames=2, max_frames=2).setup() + # Each independent setup still verifies the small train-only statistics. + assert calls == [(tmp_path / manifest["normalization"]["path"]).resolve()] * 2 + for episode in manifest["episodes"]: + fingerprint = episode["verified_fingerprint"] + assert fingerprint["sha256"] == episode["sha256"] + assert fingerprint["path"] == str((tmp_path / episode["path"]).resolve()) + + +@pytest.mark.parametrize("change", ["missing", "extra", "schema", "sha", "path", "bool_stat"]) +def test_malformed_fingerprint_falls_back_to_content_hash(tmp_path, monkeypatch, change): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + aged_verification_clock(monkeypatch) + fingerprint = ArtifactVerificationMemo().verify(path, expected) + if change == "missing": + fingerprint.pop("st_ctime_ns") + elif change == "extra": + fingerprint["unchecked"] = True + elif change == "schema": + fingerprint["schema"] = "unknown" + elif change == "sha": + fingerprint["sha256"] = "0" * 64 + elif change == "path": + fingerprint["path"] = str(tmp_path / "elsewhere") + else: + fingerprint["st_dev"] = True + calls, original = [], hashlib.file_digest + + def digest(stream, *args, **kwargs): + calls.append(stream.name) + return original(stream, *args, **kwargs) + + monkeypatch.setattr(hashlib, "file_digest", digest) + ArtifactVerificationMemo().verify(path, expected, trusted_fingerprint=fingerprint) + assert len(calls) == 1 + + +def test_changed_expected_hash_is_not_trusted(tmp_path, monkeypatch): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + aged_verification_clock(monkeypatch) + fingerprint = ArtifactVerificationMemo().verify(path, expected) + with pytest.raises(ValueError, match="hash mismatch"): + ArtifactVerificationMemo().verify(path, "0" * 64, trusted_fingerprint=fingerprint) + + +def test_same_size_restored_mtime_mutation_fails_persistent_verification(tmp_path, monkeypatch): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + aged_verification_clock(monkeypatch) + fingerprint = ArtifactVerificationMemo().verify(path, expected) + before = path.stat() + path.write_bytes(b"modified") + os.utime(path, ns=(before.st_atime_ns, before.st_mtime_ns)) + assert path.stat().st_size == before.st_size + assert path.stat().st_mtime_ns == before.st_mtime_ns + # Explicitly exercise the changed-ctime branch even on a coarse-clock tmpfs. + native_identity = ArtifactVerificationMemo.identity + monkeypatch.setattr(ArtifactVerificationMemo, "identity", staticmethod( + lambda candidate: (*native_identity(candidate)[:-1], before.st_ctime_ns + 1))) + with pytest.raises(ValueError, match="hash mismatch"): + ArtifactVerificationMemo().verify(path, expected, trusted_fingerprint=fingerprint) + + +@pytest.mark.parametrize("change", ["replace_inode", "relocate"]) +def test_copy_or_relocation_rehashes_even_identical_payload(tmp_path, monkeypatch, change): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + aged_verification_clock(monkeypatch) + fingerprint = ArtifactVerificationMemo().verify(path, expected) + destination = tmp_path / "copy" + if change == "replace_inode": + shutil.copy2(path, destination) + destination.replace(path) + assert path.stat().st_ino != fingerprint["st_ino"] + else: + path.rename(destination) + path = destination + calls, original = [], hashlib.file_digest + + def digest(stream, *args, **kwargs): + calls.append(stream.name) + return original(stream, *args, **kwargs) + + monkeypatch.setattr(hashlib, "file_digest", digest) + ArtifactVerificationMemo().verify(path, expected, trusted_fingerprint=fingerprint) + assert len(calls) == 1 + + +def test_recent_hash_cannot_become_trusted_after_same_tick_mutation(tmp_path, monkeypatch): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + identity = ArtifactVerificationMemo.identity(path) + # Simulate the observed coarse-clock case: the file keeps every stat field. + monkeypatch.setattr(ArtifactVerificationMemo, "identity", staticmethod(lambda _: identity)) + monkeypatch.setattr(ArtifactVerificationMemo, "_now_ns", staticmethod(lambda: identity[-1])) + memo = ArtifactVerificationMemo() + fingerprint = memo.verify(path, expected) + assert fingerprint is None and not memo.verified + path.write_bytes(b"modified") + monkeypatch.setattr(ArtifactVerificationMemo, "_now_ns", + staticmethod(lambda: identity[-1] + 2_000_000_000)) + for verifier in (memo, ArtifactVerificationMemo()): + with pytest.raises(ValueError, match="hash mismatch"): + verifier.verify(path, expected, trusted_fingerprint=fingerprint) + + +def test_recent_verification_timestamp_in_fingerprint_is_rejected(tmp_path, monkeypatch): + path = tmp_path / "artifact" + path.write_bytes(b"original") + expected = hashlib.sha256(path.read_bytes()).hexdigest() + aged_verification_clock(monkeypatch) + fingerprint = ArtifactVerificationMemo().verify(path, expected) + fingerprint["verification_started_ns"] = fingerprint["st_ctime_ns"] + # Even after aging, this malformed proof must trigger a new content hash. + path.write_bytes(b"modified") + identity = tuple(fingerprint[key] for key in ArtifactVerificationMemo._STAT_FIELDS) + monkeypatch.setattr(ArtifactVerificationMemo, "identity", staticmethod(lambda _: identity)) + with pytest.raises(ValueError, match="hash mismatch"): + ArtifactVerificationMemo().verify(path, expected, trusted_fingerprint=fingerprint) + + +@pytest.mark.parametrize("corruption", ["stats_hash", "train_identity", "missing_stats", "raw_without_path"]) +def test_raw_normalization_requires_bound_train_statistics(tmp_path, corruption): + inventory = make_inventory(tmp_path) + output = tmp_path / "prepared.json" + manifest = finalize(inventory, output) + stats_path = tmp_path / manifest["normalization"]["path"] + if corruption == "stats_hash": + stats_path.write_text("{}") + elif corruption == "train_identity": + stats = json.loads(stats_path.read_text()) + stats["train_episodes"] = [] + stats_path.write_text(json.dumps(stats)) + manifest["normalization"]["sha256"] = hashlib.sha256(stats_path.read_bytes()).hexdigest() + elif corruption == "missing_stats": + stats_path.unlink() + else: + del manifest["normalization"]["path"] + output.write_text(json.dumps(manifest)) + with pytest.raises((ValueError, FileNotFoundError)): + PreparedWorldDataset(output) + + +def test_finalizer_rejects_duplicate_physical_cache(tmp_path): + inventory = make_inventory(tmp_path) + spec = json.loads(inventory.read_text()) + spec["receipts"].append(spec["receipts"][0]) + inventory.write_text(json.dumps(spec)) + with pytest.raises(ValueError, match="Duplicate"): + finalize(inventory, tmp_path / "prepared.json") + + +@pytest.mark.parametrize("mutation", ["artifact", "expected_digest"]) +def test_overlapping_finalizers_share_only_verified_unchanged_payloads(tmp_path, monkeypatch, mutation): + inventory = make_inventory(tmp_path) + aged_verification_clock(monkeypatch) + specification = json.loads(inventory.read_text()) + subset = {**specification, + "receipts": [name for name in specification["receipts"] if name.startswith("robot-")], + "pose_schemas": {"robot": specification["pose_schemas"]["robot"]}} + subset_path = tmp_path / "subset.json" + subset_path.write_text(json.dumps(subset)) + calls, original = [], hashlib.file_digest + + def digest(stream, *args, **kwargs): + calls.append(Path(stream.name).resolve()) + return original(stream, *args, **kwargs) + + monkeypatch.setattr(hashlib, "file_digest", digest) + memo = ArtifactVerificationMemo() + combined = finalize(inventory, tmp_path / "combined.json", verification_memo=memo) + assert len(calls) == 6 # A fresh memo never skips the initial full hashes. + single = finalize(subset_path, tmp_path / "single.json", verification_memo=memo) + assert len(calls) == 6 + assert single["split_provenance"]["episode_ids"]["robot"] == combined["split_provenance"]["episode_ids"]["robot"] + single_stats = json.loads((tmp_path / single["normalization"]["path"]).read_text()) + assert {row[0] for row in single_stats["train_episodes"]} == {"robot"} + assert set(single_stats["statistics"]) == {"image", "track", "pose/robot"} + finalize(subset_path, tmp_path / "fresh.json") + assert len(calls) == 9 # Default behavior remains independently verified. + receipt_path = tmp_path / subset["receipts"][0] + if mutation == "artifact": + with Path(str(receipt_path)[:-5]).open("ab") as stream: + stream.write(b"changed") + else: + receipt = json.loads(receipt_path.read_text()) + receipt["output_sha256"] = "0" * 64 + receipt_path.write_text(json.dumps(receipt)) + with pytest.raises(ValueError, match="hash mismatch"): + finalize(subset_path, tmp_path / "invalid.json", verification_memo=memo) + assert len(calls) == 10 # Both identity changes force a content hash and reject. + + +def test_sampling_cursor_ignores_prefetched_unconsumed_batches(tmp_path): + inventory = make_inventory(tmp_path) + path = tmp_path / "prepared.json" + finalize(inventory, path) + data = WorldModelDataModule(path, max_frames=3, checkpointable_sampling=True) + loader = data.train_dataloader() + iterator = iter(loader) + first = next(iterator) + data.acknowledge_batch(next(iter(first.values()))["sampling_cursor"]) + expected = next(iterator) # Prefetch without training must not advance saved cursor. + saved = data.state_dict() + restored = WorldModelDataModule(path, max_frames=3, checkpointable_sampling=True) + restored.load_state_dict(saved) + actual = next(iter(restored.train_dataloader())) + assert next(iter(actual.values()))["sampling_cursor"] == 2 + for name, stream in next(iter(actual.values()))["world"].items(): + for key, value in stream.items(): + torch.testing.assert_close(value, next(iter(expected.values()))["world"][name][key], rtol=0, atol=0) + changed = WorldModelDataModule(path, max_frames=3, batch_size=2, checkpointable_sampling=True) + with pytest.raises(ValueError, match="batch_size"): + changed.load_state_dict(saved) + with pytest.raises(ValueError, match="num_workers=0"): + WorldModelDataModule(path, num_workers=1, checkpointable_sampling=True) + + +@pytest.mark.parametrize("signal_boundary", [False, True, "validation"]) +def test_interrupted_world_training_matches_uninterrupted_updates(tmp_path, monkeypatch, signal_boundary): + import inspect + from lightning import Trainer, seed_everything + from omegaconf import OmegaConf + from egomimic.pl_utils.pl_model import ModelWrapper + path = tmp_path / "prepared.json" + finalize(make_inventory(tmp_path), path) + cfg = OmegaConf.create({"model": { + "_target_": "egomimic.pl_utils.pl_model.ModelWrapper", + "pipeline": {"_target_": "egomimic.pipeline.algo.PipelineAlgo", "device": "cpu", "stages": [{ + "_target_": "egomimic.pipeline.stages_world_model.WorldModelStage", + "warmup_steps": 0, "pose_ramp_steps": 0, "sampler_steps": 2, + "model": {"_target_": "egomimic.models.world_model.MoTWorldModel", "pose_dims": {"robot": 3, "human": 4}, + "image_dim": 2, "width": 24, "heads": 3, "depth": 1, "max_tracks": 8}}]}, + "training_behavior": {"_target_": "egomimic.pl_utils.world_model_behavior.WorldModelTrainingBehavior"}, + "optimizer": {"_target_": "torch.optim.AdamW", "_partial_": True, "lr": .001}}}) + + def make(steps, checkpoint=None, request_signal=False): + callbacks = [] + if request_signal: + import signal + from lightning import Callback + from egomimic.utils.slurm_requeue import SaveOnlySignalCheckpoint + class RequestDuringBackward(Callback): + def on_before_backward(self, trainer, module, loss): + if trainer.global_step == 0: + os.kill(os.getpid(), signal.SIGUSR2) + assert not list((tmp_path / "signals").glob("*.ckpt")) + class RequestDuringValidation(Callback): + def on_validation_batch_end(self, trainer, module, outputs, batch, index, dataloader_idx=0): + if trainer.global_step == 2: + os.kill(os.getpid(), signal.SIGUSR2) + assert not list((tmp_path / "signals").glob("*.ckpt")) + def on_validation_epoch_end(self, trainer, module): + if trainer.global_step == 2: + assert not list((tmp_path / "signals").glob("*.ckpt")) + validation = signal_boundary == "validation" + callbacks = [SaveOnlySignalCheckpoint(tmp_path / "signals", train_batch_only=validation), + RequestDuringValidation() if validation else RequestDuringBackward()] + seed_everything(123) + data = WorldModelDataModule(path, max_frames=3, checkpointable_sampling=True) + context = data.prepare_context(mode="train", normalization=None, normalizer=None) + model = ModelWrapper(config_tree=cfg) + model.data_context = context + context.bind(model.model, None) + validation_options = ({"limit_val_batches": 1, "val_check_interval": 2, + "check_val_every_n_epoch": None} if signal_boundary == "validation" + else {"limit_val_batches": 0}) + trainer = Trainer(accelerator="cpu", devices=1, max_steps=steps, max_epochs=-1, + **validation_options, logger=False, enable_checkpointing=False, + enable_progress_bar=False, num_sanity_val_steps=0, callbacks=callbacks) + kwargs = {"weights_only": False} if "weights_only" in inspect.signature(trainer.fit).parameters else {} + trainer.fit(model, datamodule=data, ckpt_path=checkpoint, **kwargs) + return model, trainer, data + + reference, uninterrupted, _ = make(4) + monkeypatch.setenv("SLURM_JOB_ID", "123456") + monkeypatch.setenv("SLURM_RESTART_COUNT", "0") + monkeypatch.delenv("ICE_CHECKPOINT_FORWARDING", raising=False) + split_step = 3 if signal_boundary == "validation" else 2 + _, first, data = make(split_step, request_signal=signal_boundary) + assert data.consumed_batches == split_step + if signal_boundary: + checkpoint, = (tmp_path / "signals").glob("*.ckpt") + payload = torch.load(checkpoint, map_location="cpu", weights_only=False) + boundary_step = 3 if signal_boundary == "validation" else 1 + assert payload["global_step"] == boundary_step + from scripts.ice.validate_lightning_checkpoint import validate_checkpoint + proof = validate_checkpoint(checkpoint, environment={}) + assert proof["strict_model_reload"] and proof["global_step"] == boundary_step + else: + checkpoint = tmp_path / "resume.ckpt" + first.save_checkpoint(checkpoint) + resumed, continued, restored_data = make(4, str(checkpoint)) + assert continued.global_step == 4 and restored_data.consumed_batches == 4 + for name, expected in reference.state_dict().items(): + torch.testing.assert_close(resumed.state_dict()[name], expected, rtol=0, atol=0) + expected_states = uninterrupted.optimizers[0].state_dict()["state"] + actual_states = continued.optimizers[0].state_dict()["state"] + for parameter, expected in expected_states.items(): + for key, value in expected.items(): + torch.testing.assert_close(actual_states[parameter][key], value, rtol=0, atol=0)