File size: 4,869 Bytes
0c3a82b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | --- 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__)
|