Add BridgeAttention adapter + config + policy definition

adapter.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c15b7174ab82dc4590b3e0a35a63c2d833353496a7239ef54feaada727e00e1
+size 55237642

adapter.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3749947322390daf4ca9093633518845885f2bd289bea321ff4ab8c691fa47af
+size 55220344

config.json

@@ -0,0 +1,27 @@
+{
+  "vision_model_id": "google/siglip-base-patch16-224",
+  "text_model_id": "Qwen/Qwen2.5-0.5B-Instruct",
+  "image_size": 224,
+  "num_action_queries": 64,
+  "policy_dim": 512,
+  "policy_layers": 4,
+  "n_heads": 8,
+  "dropout": 0.1,
+  "action_dim": 43,
+  "state_dim": 43,
+  "per_device_batch": 16,
+  "max_steps": 400,
+  "lr": 0.0003,
+  "warmup_ratio": 0.03,
+  "weight_decay": 0.01,
+  "log_every": 20,
+  "hf_repo": "nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1",
+  "split": "train[:2%]",
+  "num_workers": 4,
+  "enable_vla_download": false,
+  "vla_patterns": [
+    "g1-pick-apple/**",
+    "g1-pick-pear/**"
+  ],
+  "local_data_dir": "/kaggle/working/gr00t_g1_subset"
+}

policy_definition.py

@@ -0,0 +1,37 @@
+
+import math, torch
+import torch.nn as nn
+from einops import repeat
+
+class BridgeAttentionPolicy(nn.Module):
+    def __init__(self, v_hidden, t_hidden, state_dim, policy_dim, n_heads, n_layers, n_queries, action_dim, dropout=0.1):
+        super().__init__()
+        self.n_queries = n_queries
+        self.query = nn.Parameter(torch.randn(n_queries, policy_dim) / math.sqrt(policy_dim))
+        self.v_proj = nn.Linear(v_hidden, policy_dim)
+        self.t_proj = nn.Linear(t_hidden, policy_dim)
+        self.s_proj = nn.Linear(state_dim, policy_dim)
+        self.alpha_v = nn.Parameter(torch.tensor(0.7))
+        self.alpha_t = nn.Parameter(torch.tensor(0.7))
+        self.alpha_s = nn.Parameter(torch.tensor(0.7))
+        enc = nn.TransformerEncoderLayer(d_model=policy_dim, nhead=n_heads, dim_feedforward=policy_dim*4,
+                                         dropout=dropout, activation="gelu", batch_first=True, norm_first=True)
+        self.blocks = nn.TransformerEncoder(enc, num_layers=n_layers)
+        self.norm = nn.LayerNorm(policy_dim)
+        self.head = nn.Sequential(nn.Linear(policy_dim, policy_dim), nn.GELU(), nn.Linear(policy_dim, action_dim))
+
+    def forward(self, v_feats_layers, t_feats_layers, state_vec):
+        B = state_vec.size(0)
+        v_cat = torch.cat(v_feats_layers, dim=1) if v_feats_layers else None
+        t_cat = torch.cat(t_feats_layers, dim=1)
+        s_tok = self.s_proj(state_vec).unsqueeze(1)
+        toks = [s_tok]
+        if v_cat is not None:
+            toks.append(self.v_proj(v_cat) * torch.sigmoid(self.alpha_v))
+        toks.append(self.t_proj(t_cat) * torch.sigmoid(self.alpha_t))
+        ctx = torch.cat(toks, dim=1)
+        q = repeat(self.query, 'Q D -> B Q D', B=B)
+        tokens = torch.cat([q, ctx], dim=1)
+        tokens = self.blocks(tokens)
+        pooled = self.norm(tokens[:, :self.n_queries].mean(dim=1))
+        return self.head(pooled)