evaluation/action_distribution/runtime_adapter.patch

--- archived/src/b1k/models/pi_behavior.py
+++ runtime/src/b1k/models/pi_behavior.py
@@ -19,7 +19,6 @@
 
 from openpi.models import model as _model
 from openpi.models import gemma as _gemma
-from b1k.models import spatial_da3 as _spatial_da3
 from openpi.models import siglip as _siglip
 from openpi.models.pi0 import make_attn_mask, posemb_sincos
 from openpi.shared import array_typing as at
@@ -215,6 +214,7 @@
         
         # Store config for later use
         self.config = config
+        assert config.da3 is None and not config.use_spatial_action_cross_attention and not config.use_skill_conditioning, "RGB-only evaluation adapter"
         
         paligemma_config = _gemma.get_config(config.paligemma_variant)
         action_expert_config = _gemma.get_config(config.action_expert_variant)
@@ -230,17 +230,6 @@
                 configs=[paligemma_config, action_expert_config],
                 embed_dtype=config.dtype,
                 adarms=True,
-                spatial_inject=spatial_inject,
-                num_spatial_layers=config.da3.num_inject_layers if spatial_inject else 6,
-                spatial_scale=config.da3.spatial_scale if spatial_inject else 2.0,
-                spatial_init_std=config.da3.spatial_init_std if spatial_inject else 0.0,
-                spatial_logit_gain=config.da3.attn_logit_gain if spatial_inject else False,
-                spatial_logit_gain_init=config.da3.attn_logit_gain_init if spatial_inject else 1.0,
-                spatial_logit_gain_max=config.da3.attn_logit_gain_max if spatial_inject else 0.0,
-                spatial_qk_norm=config.da3.qk_norm if spatial_inject else False,
-                spatial_s1=getattr(config.da3, "spatial_s1", False) if spatial_inject else False,
-                spatial_no_bias=getattr(config.da3, "spatial_no_bias", False) if spatial_inject else False,
-                spatial_inject_vlm=getattr(config.da3, "spatial_inject_vlm", False) if spatial_inject else False,
             )
         )
         llm.lazy_init(rngs=rngs, method="init", use_adarms=[False, True])
@@ -767,6 +756,10 @@
                 "Call load_correlation_matrix() first."
             )
         
+        if hasattr(self, '_eval_correction'):
+            assert O_indices.shape == (128,) and U_indices.shape == (832,)
+            return {'O_indices': O_indices, 'U_indices': U_indices, 'correction_matrix': self._eval_correction.value}
+
         L = self.action_correlation_cholesky.value
         Sigma = L @ L.T  # Full covariance matrix [hd, hd]
         
@@ -1124,7 +1117,6 @@
             [prefix_tokens, None],
             mask=prefix_attn_mask,
             positions=positions_prefix,
-            banks=_prefix_banks,
         )
 
         # 3. Predict stage from VLM output of base task token
@@ -1255,7 +1247,6 @@
                 positions=suffix_positions,
                 kv_cache=kv_cache_for_actions,  # Original cache, reused for all samples
                 adarms_cond=[None, adarms_cond],
-                banks=spatial_banks,
             )
             
             # Compute velocity and loss
@@ -1488,7 +1479,7 @@
         _prefix_banks = spatial_banks if (
             self.spatial_bank_builder is not None and getattr(self.da3_cfg, "spatial_inject_vlm", False)
         ) else None
-        (prefix_out, _), kv_cache = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions, banks=_prefix_banks)
+        (prefix_out, _), kv_cache = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions)
         
         # Predict stage from VLM output of base task token
         # Find base task token position (same logic as in compute_detailed_loss)
@@ -1535,7 +1526,6 @@
                 positions=positions,
                 kv_cache=kv_cache,
                 adarms_cond=[None, adarms_cond],
-                banks=spatial_banks,
             )
             assert prefix_out is None
             action_hidden = self.apply_spatial_action_conditioning(
--- upstream/src/b1k/policies/b1k_policy.py
+++ runtime/src/b1k/policies/b1k_policy.py
@@ -12,7 +12,7 @@
 
 from openpi import transforms
 from openpi.models import model as _model
-from omnigibson.learning.utils.eval_utils import PROPRIOCEPTION_INDICES
+from b1k_eval_contract import PROPRIOCEPTION_INDICES
 
 
 def make_b1k_example() -> dict:
--- upstream/src/b1k/shared/eval_b1k_wrapper.py
+++ runtime/src/b1k/shared/eval_b1k_wrapper.py
@@ -11,7 +11,7 @@
 from b1k.policies.b1k_policy import extract_state_from_proprio
 from b1k.models.pi_behavior_config import TASK_NUM_STAGES
 from b1k.shared.correction_rules import apply_correction_rules, check_gripper_variation
-from omnigibson.learning.utils.eval_utils import PROPRIOCEPTION_INDICES
+from b1k_eval_contract import PROPRIOCEPTION_INDICES
 
 logger = logging.getLogger(__name__)