--- 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__)