Add Elmo continuous objectives and verified interruption recovery

docs/integration/ELMO_SNAPSHOT_IMPORT_20261026.md

@@ -0,0 +1,39 @@
+# Elmo deployed source import (2026-10-26)
+
+This change imports the verified source snapshot used by the Elmo image,
+proprioception and action model into maintained Git history. It has the real
+parent `255fdaf5137f7043b35df7ef48a3c1b8adc70815`, published on
+`codex/joint-modality-ablation-20261019-v1` in `GaTech-RL2/EgoVerse-graph`.
+
+The deployment recorded source label
+`153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920`. Its original Git object and ancestry
+could not be recovered. **This import does not claim that label is a parent or
+ancestor, or reconstruct its missing author/commit history.** The new import
+receives its own Git identity. Existing runs and checkpoints retain their
+original source records.
+
+The independently retained inputs are:
+
+| Input | SHA-256 |
+| --- | --- |
+| `source-153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920.tar.gz` (2,061,765 bytes) | `5e3f11a43ae33411591d4149c77432ad355c0671f9764a68e18fdddbf180aa2f` |
+| `source-153d4a5484bb61eb113d9ac7a7b7a5cbabbbf920.files.sha256` | `208ba5fa51d0514fd27d00f0f202cbf73f6e31ab7f3537e26ca55004e4fe3f99` |
+
+All 1,030 manifest paths are reproduced byte-for-byte with archive executable
+bits: 988 already matched this parent, 25 are updated, and 17 are added. Every
+parent path outside the export is retained. This change imports the deployed
+baseline only; the larger-model configuration, fused objective and boundary
+sampling experiments belong in separate changes.
+
+Runtime packaging must use the exact 1,030-path snapshot allowlist, plus only
+explicitly reviewed experimental additions. Retaining older recipes and tests
+in Git does not select them for execution or certify their compatibility. The
+native model is `egomimic.models.joint_image_proprio.JointImageProprioModel`;
+the Elmo selectors live under `experiment/pixel_proprio/native_joint/`.
+Operational services remain separately hash-bound; do not accidentally use
+older operational modules retained from this parent.
+
+Byte equivalence and local CPU checks do not certify cluster readiness. A new
+immutable source export/readback, explicit runtime/data bindings and matching
+allocated optimization, validation and full-state/RNG smoke are required for
+new experimental runs. No historical run identity is relabeled by this import.

egomimic/eval/image_proprio_eval.py

@@ -9,11 +9,7 @@ class ImageProprioEval(Eval):
     prediction_key = "image_proprio_prediction"
 
     def __init__(
-        self,
-        objective="visual_proprio",
-        sample_seed=42,
-        pixel_diagnostics=None,
-        prediction_streams=None,
+        self, objective="visual_proprio", sample_seed=42, pixel_diagnostics=None, active_streams_from_data=False
     ):
         if objective not in (
             "visual_only",
@@ -23,16 +19,7 @@ def __init__(
         ):
             raise ValueError("Unknown image/proprio evaluation objective")
         self.objective, self.sample_seed = objective, int(sample_seed)
-        self.prediction_streams = (
-            None if prediction_streams is None else tuple(prediction_streams)
-        )
-        if self.prediction_streams is not None and (
-            not self.prediction_streams
-            or len(set(self.prediction_streams)) != len(self.prediction_streams)
-            or set(self.prediction_streams) - {"image", "proprio", "action"}
-            or (objective != "direct_action" and self.prediction_streams[0] != "image")
-        ):
-            raise ValueError("Declare unique supported prediction_streams")
+        self.active_streams_from_data = bool(active_streams_from_data)
         self.totals = {}
         self._completed_validation_step = None
         self._validation_step = None
@@ -96,21 +83,13 @@ def _sampling_index(batch, batch_idx):
         if all(value is None for value in indexed):
             return int(batch_idx)
         if any(value is None for value in indexed):
-            raise ValueError(
-                "Global validation batch identities must cover every source"
-            )
+            raise ValueError("Global validation batch identities must cover every source")
         indices = []
         for item, values in zip(batch.values(), indexed):
-            if not isinstance(values, (tuple, list)) or len(values) != len(
-                item["task_kind"]
-            ):
-                raise ValueError(
-                    "One global validation batch identity is required per sample"
-                )
+            if not isinstance(values, (tuple, list)) or len(values) != len(item["task_kind"]):
+                raise ValueError("One global validation batch identity is required per sample")
             if any(type(value) is not int or value < 0 for value in values):
-                raise ValueError(
-                    "Global validation batch identities must be nonnegative integers"
-                )
+                raise ValueError("Global validation batch identities must be nonnegative integers")
             indices.extend(values)
         if len(set(indices)) != 1:
             raise ValueError("A validation batch must share one global batch identity")
@@ -122,29 +101,26 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0):
         )
         with torch.random.fork_rng(devices=devices):
             torch.manual_seed(
-                self.sample_seed
-                + self._sampling_index(batch, batch_idx)
+                self.sample_seed + self._sampling_index(batch, batch_idx)
                 + 100000 * dataloader_idx
             )
             output = self.model.forward_eval(batch)
-        names = (
-            self.prediction_streams
-            or {
-                "visual_only": ("image",),
-                "visual_proprio": ("image", "proprio"),
-                "visual_state_action": ("image", "proprio", "action"),
-                "direct_action": ("action",),
-            }[self.objective]
-        )
+        names = {
+            "visual_only": ("image",),
+            "visual_proprio": ("image", "proprio"),
+            "visual_state_action": ("image", "proprio", "action"),
+            "direct_action": ("action",),
+        }[self.objective]
         for source, item in batch.items():
-            if self.prediction_streams is not None and set(
-                output[source][self.prediction_key]
-            ) != set(names):
-                raise ValueError(
-                    "Evaluator prediction_streams differ from model output"
-                )
             kinds = item["task_kind"]
-            for name in names:
+            selected_names = names
+            if self.active_streams_from_data:
+                selected_names = tuple(item['world'])
+                if (not {'image','proprio','camera'} <= set(selected_names)
+                        or not set(selected_names) <= {'image','proprio','camera','action'}
+                        or set(output[source][self.prediction_key]) != set(selected_names)):
+                    raise ValueError('Actual multimodal prediction keys must match native world streams')
+            for name in selected_names:
                 stream = item["world"][name]
                 mask = (
                     stream["valid"]
@@ -194,16 +170,12 @@ def _evaluate_batch(self, batch, batch_idx, dataloader_idx=0):
     def _reduce_validation_totals(self):
         """Sum target-weighted errors, including keys absent on some ranks."""
         self.global_validation_batches = self._validation_batches
-        if not (
-            torch.distributed.is_available() and torch.distributed.is_initialized()
-        ):
+        if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
             return self.totals
         keys_by_rank = [None] * torch.distributed.get_world_size()
         torch.distributed.all_gather_object(keys_by_rank, sorted(self.totals))
         keys = sorted({key for rank_keys in keys_by_rank for key in rank_keys})
-        totals = torch.zeros(
-            len(keys), 2, dtype=torch.float64, device=self.model.device
-        )
+        totals = torch.zeros(len(keys), 2, dtype=torch.float64, device=self.model.device)
         for index, key in enumerate(keys):
             if key in self.totals:
                 numerator, denominator = self.totals[key]
@@ -212,19 +184,13 @@ def _reduce_validation_totals(self):
         valid = (torch.isfinite(totals).all() & (totals >= 0).all()).to(torch.int64)
         torch.distributed.all_reduce(valid, op=torch.distributed.ReduceOp.MIN)
         if not valid.item():
-            raise ValueError(
-                "Invalid accumulated validation metric on a distributed rank"
-            )
+            raise ValueError("Invalid accumulated validation metric on a distributed rank")
         if keys:
             torch.distributed.all_reduce(totals, op=torch.distributed.ReduceOp.SUM)
-        batches = torch.tensor(
-            self._validation_batches, dtype=torch.int64, device=self.model.device
-        )
+        batches = torch.tensor(self._validation_batches, dtype=torch.int64, device=self.model.device)
         torch.distributed.all_reduce(batches, op=torch.distributed.ReduceOp.SUM)
         self.global_validation_batches = int(batches.item())
-        return {
-            key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys)
-        }
+        return {key: (totals[index, 0], totals[index, 1]) for index, key in enumerate(keys)}
 
     def log_epoch_metrics(self):
         self._metrics_logged = False
@@ -278,9 +244,7 @@ def on_validation_end(self):
             and any(denominator > 0 for _, denominator in self._global_totals.values())
         )
         if torch.distributed.is_available() and torch.distributed.is_initialized():
-            success = torch.tensor(
-                int(complete), dtype=torch.int64, device=self.model.device
-            )
+            success = torch.tensor(int(complete), dtype=torch.int64, device=self.model.device)
             torch.distributed.all_reduce(success, op=torch.distributed.ReduceOp.MIN)
             complete = bool(success.item())
         if publication_error is not None:

egomimic/hydra_configs/experiment/pixel_proprio/native_joint/common.yaml

@@ -0,0 +1,98 @@
+# @package _global_
+defaults:
+- /experiment/pixel_proprio/video_continuous
+- /model/pixel_proprio/video@native_joint.reference
+- override /model: pixel_proprio/native_joint
+- _self_
+# Reference the historical image constructor; do not reconstruct its defaults.
+joint:
+  pose_dims: {libero_gripper: 9}
+  max_tokens: 12864
+  compile_mode: default
+world_model_execution:
+  compile_mode: ${joint.compile_mode}
+  compile_dynamic: true
+  compile_backend: inductor
+  token_execution: packed_valid
+  loss_reduction: mean_per_example
+pixel_video:
+  max_tokens: ${joint.max_tokens}
+modality_ablation:
+  active_streams: null
+  nodes:
+    joint_noising:
+      _target_: egomimic.pipeline.stages_joint_state_action.JointNoisingStage
+      noising:
+        _target_: egomimic.models.joint_image_proprio.CoupledImageProprioNoising
+        maximum: 1.0
+        active_streams: ${modality_ablation.active_streams}
+    joint_denoiser:
+      _target_: egomimic.pipeline.stages_joint_state_action.JointDenoiserStage
+      model:
+        _target_: egomimic.models.joint_image_proprio.JointImageProprioModel
+        image_model: ${native_joint.reference.pipeline.stages.1.model}
+        pose_dims: ${joint.pose_dims}
+        proprio_width: 256
+        mlp_ratio: 4
+        active_streams: ${modality_ablation.active_streams}
+    image_loss:
+      _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage
+      stream: image
+      weight: 1.0
+      loss:
+        _target_: egomimic.models.pixel_diffusion.PlainPixelMSE
+        reduction: mean_per_example
+    proprio_loss:
+      _target_: egomimic.pipeline.stages_joint_state_action.JointLossStage
+      stream: proprio
+      weight: 1.0
+      loss:
+        _target_: egomimic.models.pixel_diffusion.PlainPixelMSE
+        reduction: mean_per_example
+    joint_sampler:
+      _target_: egomimic.pipeline.stages_joint_state_action.JointSamplerStage
+      sampler:
+        _target_: egomimic.models.joint_image_proprio.JointImageProprioSampler
+        state_steps: 20
+  stage_ids: null
+  graph: null
+  initialization: Native image construction first; vector initialization preserves CPU RNG
+  training_noise: Native image draws unchanged; optional tied proprio draws preserve image RNG
+  screening_budget_updates: 300000
+data:
+  image_only: false
+  task: visual_diffusion
+  couple_state_observations: true
+  max_frames: 64
+  batch_size: 8
+  distributed_world_size: 4
+  validation_batch_size: 1
+  validation_anchors_per_episode: 2
+  num_workers: 2
+  pin_memory: true
+  persistent_workers: true
+  prefetch_factor: 2
+  balance_training_work: true
+trainer:
+  strategy: ddp
+  devices: 4
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  use_distributed_sampler: false
+  max_steps: 300000
+  val_check_interval: 20000
+callbacks:
+  model_checkpoint:
+    every_n_train_steps: 20000
+    save_top_k: -1
+  learning_rate_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+    logging_interval: step
+runtime:
+  cluster_environment: torch_elastic
+  slurm_requeue_owner: none
+logger:
+  wandb:
+    group: native-image-joint-state
+    resume: never
+    tags: [historical-image-core, native-compiled, no-action, paired-initialization]

egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_only.yaml

@@ -0,0 +1,22 @@
+# @package _global_
+defaults:
+- /experiment/pixel_proprio/native_joint/common
+- _self_
+name: native_mot_continuous_image_only_s42_300k
+description: Historical image core with minimal joint state extension, historical compilation preserved
+modality_ablation:
+  active_streams: [image]
+  stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, joint_sampler: 3}
+  graph:
+  - ${modality_ablation.nodes.joint_noising}
+  - ${modality_ablation.nodes.joint_denoiser}
+  - ${modality_ablation.nodes.image_loss}
+  - ${modality_ablation.nodes.joint_sampler}
+evaluator:
+  objective: visual_only
+
+joint:
+  max_tokens: 12800
+data:
+  image_only: true
+  couple_state_observations: false

egomimic/hydra_configs/experiment/pixel_proprio/native_joint/continuous_image_proprio.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+defaults:
+- /experiment/pixel_proprio/native_joint/common
+- _self_
+name: native_mot_continuous_image_proprio_s42_300k
+description: Historical image core with minimal joint state extension, historical compilation preserved
+modality_ablation:
+  active_streams: [image, proprio]
+  stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, joint_sampler: 4}
+  graph:
+  - ${modality_ablation.nodes.joint_noising}
+  - ${modality_ablation.nodes.joint_denoiser}
+  - ${modality_ablation.nodes.image_loss}
+  - ${modality_ablation.nodes.proprio_loss}
+  - ${modality_ablation.nodes.joint_sampler}
+evaluator:
+  objective: visual_proprio

egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_only.yaml

@@ -0,0 +1,30 @@
+# @package _global_
+defaults:
+- /experiment/pixel_proprio/native_joint/common
+- _self_
+name: native_mot_discrete_image_only_s42_300k
+description: Historical image core with minimal joint state extension, historical compilation preserved
+modality_ablation:
+  active_streams: [image]
+  stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, joint_sampler: 3}
+  graph:
+  - ${modality_ablation.nodes.joint_noising}
+  - ${modality_ablation.nodes.joint_denoiser}
+  - ${modality_ablation.nodes.image_loss}
+  - ${modality_ablation.nodes.joint_sampler}
+evaluator:
+  objective: visual_only
+joint:
+  pose_dims: {libero_gripper: 36}
+  max_tokens: 3200
+data:
+  image_only: true
+  couple_state_observations: false
+  mode: discrete
+  sampling_strategy: discrete_window
+  discrete_frames: 16
+  stride: 4
+  batch_size: 16
+  distributed_world_size: 2
+trainer:
+  devices: 2

egomimic/hydra_configs/experiment/pixel_proprio/native_joint/discrete_image_proprio.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+- /experiment/pixel_proprio/native_joint/common
+- _self_
+name: native_mot_discrete_image_proprio_s42_300k
+description: Historical image core with minimal joint state extension, historical compilation preserved
+modality_ablation:
+  active_streams: [image, proprio]
+  stage_ids: {joint_noising: 0, joint_denoiser: 1, image_loss: 2, proprio_loss: 3, joint_sampler: 4}
+  graph:
+  - ${modality_ablation.nodes.joint_noising}
+  - ${modality_ablation.nodes.joint_denoiser}
+  - ${modality_ablation.nodes.image_loss}
+  - ${modality_ablation.nodes.proprio_loss}
+  - ${modality_ablation.nodes.joint_sampler}
+evaluator:
+  objective: visual_proprio
+joint:
+  pose_dims: {libero_gripper: 36}
+  max_tokens: 3216
+data:
+  mode: discrete
+  sampling_strategy: discrete_window
+  discrete_frames: 16
+  stride: 4
+  batch_size: 16
+  distributed_world_size: 2
+trainer:
+  devices: 2