update

flock_robotics_adapter.py

@@ -1,117 +1,195 @@
 from __future__ import annotations
 
+import hashlib
 import json
+import re
 from pathlib import Path
 
 import numpy as np
 import torch
 import torch.nn as nn
-from PIL import Image
-from safetensors.torch import load_file
-from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
-
-ACTION_DIM = 7
-
-
-def format_user_text(obs: dict) -> str:
-    """Build the text the VLM sees — must match training."""
-    instruction = str(obs.get("instruction", "")).strip()
-    task = str(obs.get("task", "")).strip()
-    difficulty = str(obs.get("difficulty", "")).strip()
-    step = obs.get("step", obs.get("step_index", ""))
-    horizon = obs.get("horizon", "")
-    parts = [instruction] if instruction else []
-    meta = []
-    if task:
-        meta.append(f"task={task}")
-    if difficulty:
-        meta.append(f"difficulty={difficulty}")
-    if step != "" and step is not None:
-        meta.append(f"step={step}" + (f"/{horizon}" if horizon not in ("", None) else ""))
-    if meta:
-        parts.append(" ".join(meta))
-    proprio = obs.get("proprio")
-    if proprio is not None:
-        arr = np.asarray(proprio, dtype=np.float32).reshape(-1)
-        # Compact fixed-width floats keep tokenizer length stable
-        props = " ".join(f"{x:.3f}" for x in arr.tolist())
-        parts.append(f"proprio: {props}")
-    return "\n".join(parts) if parts else "act"
-
-
-def build_action_head(hidden_size, head_hidden):
-    return nn.Sequential(
-        nn.Linear(hidden_size, head_hidden),
-        nn.GELU(),
-        nn.Dropout(0.05),
-        nn.Linear(head_hidden, head_hidden),
-        nn.GELU(),
-        nn.Linear(head_hidden, ACTION_DIM),
-    )
-
-
-def last_token_hidden(hidden, attention_mask):
-    # hidden: [B, T, H], mask: [B, T]
-    idx = attention_mask.long().sum(dim=1).clamp(min=1) - 1
-    b = torch.arange(hidden.size(0), device=hidden.device)
-    return hidden[b, idx]
-
-
-class QwenVLAPolicy:
-    def __init__(self, model, processor, head, device):
-        self.model = model
-        self.processor = processor
-        self.head = head
-        self.device = device
-        self.model.eval()
-        self.head.eval()
-
-    @torch.no_grad()
-    def act(self, obs):
-        image = np.asarray(obs["image"], dtype=np.uint8)
-        pil = Image.fromarray(image)
-        user_text = format_user_text(obs)
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": user_text},
-                ],
-            }
-        ]
-        text = self.processor.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
+import torch.nn.functional as F
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+class EncoderVLAPolicyNet(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        text_dim: int = 256,
+        proprio_dim: int = 32,
+        action_dim: int = 7,
+        image_size: int = 224,
+        vision_dim: int = 384,
+        modality_emb_dim: int = 256,
+        freeze_encoder: bool = True,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.image_size = int(image_size)
+        self.freeze_encoder = bool(freeze_encoder)
+        if self.freeze_encoder:
+            for p in self.encoder.parameters():
+                p.requires_grad = False
+        self.vision_proj = nn.Sequential(
+            nn.Linear(vision_dim, 512),
+            nn.SiLU(),
+            nn.LayerNorm(512),
+            nn.Linear(512, 512),
+            nn.SiLU(),
+        )
+        self.proprio_encoder = nn.Sequential(
+            nn.Linear(proprio_dim, modality_emb_dim),
+            nn.SiLU(),
+            nn.LayerNorm(modality_emb_dim),
+            nn.Linear(modality_emb_dim, modality_emb_dim),
+            nn.SiLU(),
         )
-        inputs = self.processor(
-            text=[text], images=[pil], return_tensors="pt", padding=True
+        self.text_encoder = nn.Sequential(
+            nn.Linear(text_dim, modality_emb_dim),
+            nn.SiLU(),
+            nn.LayerNorm(modality_emb_dim),
+            nn.Linear(modality_emb_dim, modality_emb_dim),
+            nn.SiLU(),
+        )
+        fused = 512 + 2 * modality_emb_dim
+        self.action_head = nn.Sequential(
+            nn.Linear(fused, 1024),
+            nn.SiLU(),
+            nn.LayerNorm(1024),
+            nn.Dropout(0.05),
+            nn.Linear(1024, 512),
+            nn.SiLU(),
+            nn.LayerNorm(512),
+            nn.Linear(512, 256),
+            nn.SiLU(),
+            nn.Linear(256, action_dim),
+            nn.Tanh(),
+        )
+        self.register_buffer(
+            "img_mean",
+            torch.tensor(IMAGENET_MEAN, dtype=torch.float32).view(1, 3, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "img_std",
+            torch.tensor(IMAGENET_STD, dtype=torch.float32).view(1, 3, 1, 1),
+            persistent=False,
+        )
+
+    def _encode_images(self, images: torch.Tensor) -> torch.Tensor:
+        if images.ndim != 4:
+            raise ValueError(f"expected image batch [B,H,W,3] or [B,3,H,W], got {tuple(images.shape)}")
+        if images.shape[-1] == 3:
+            images = images.permute(0, 3, 1, 2)
+        images = images.float()
+        if images.max() > 2.0:
+            images = images / 255.0
+        target = (self.image_size, self.image_size)
+        if images.shape[-2:] != target:
+            images = F.interpolate(images, size=target, mode="bilinear", align_corners=False)
+        images = (images - self.img_mean) / self.img_std
+        if self.freeze_encoder:
+            with torch.no_grad():
+                out = self.encoder(pixel_values=images)
+        else:
+            out = self.encoder(pixel_values=images)
+        # CLS token
+        return out.last_hidden_state[:, 0]
+
+    def forward(self, images, proprio, text_features):
+        vision = self.vision_proj(self._encode_images(images))
+        proprio_emb = self.proprio_encoder(proprio.float())
+        text_emb = self.text_encoder(text_features.float())
+        return self.action_head(torch.cat([vision, proprio_emb, text_emb], dim=-1))
+
+
+def _text_vector(text: str, dim: int) -> np.ndarray:
+    vec = np.zeros(dim, dtype=np.float32)
+    tokens = re.findall(r"[a-z0-9_]+", text.lower())
+    for token in tokens:
+        digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
+        value = int.from_bytes(digest, byteorder="little", signed=False)
+        vec[value % dim] += 1.0 if value & 1 else -1.0
+    norm = float(np.linalg.norm(vec))
+    if norm > 0:
+        vec /= norm
+    return vec
+
+
+def _proprio_vector(obs: dict, dim: int) -> np.ndarray:
+    raw = np.asarray(obs.get("proprio", np.zeros(25, dtype=np.float32)), dtype=np.float32).reshape(-1)
+    step = float(obs.get("step", 0))
+    horizon = float(obs.get("horizon", 320) or 320)
+    step_feature = np.asarray([step / max(horizon, 1.0)], dtype=np.float32)
+    combined = np.concatenate([raw, step_feature], axis=0)
+    if combined.size < dim:
+        combined = np.pad(combined, (0, dim - combined.size))
+    return combined[:dim].astype(np.float32)
+
+
+class EncoderVLAPolicy:
+    def __init__(self, model_dir: str, device: str, dtype: str):
+        from transformers import Dinov2Config, Dinov2Model
+
+        self.model_dir = Path(model_dir)
+        self.config = json.loads((self.model_dir / "vla_config.json").read_text())
+        if device == "cuda" and torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif device == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        else:
+            self.device = torch.device("cpu")
+
+        enc_cfg = Dinov2Config.from_dict(self.config["encoder_config"])
+        encoder = Dinov2Model(enc_cfg)
+        image_size = self.config.get("image_size", [224, 224])
+        if isinstance(image_size, list):
+            image_size = int(image_size[0])
+        self.model = EncoderVLAPolicyNet(
+            encoder=encoder,
+            text_dim=int(self.config.get("text_dim", 256)),
+            proprio_dim=int(self.config.get("proprio_dim", 32)),
+            action_dim=int(self.config.get("action_dim", 7)),
+            image_size=int(image_size),
+            vision_dim=int(self.config.get("vision_dim", enc_cfg.hidden_size)),
+            modality_emb_dim=int(self.config.get("modality_emb_dim", 256)),
+            freeze_encoder=True,
+        ).to(self.device)
+        checkpoint = torch.load(self.model_dir / "model.pt", map_location=self.device, weights_only=True)
+        state_dict = checkpoint.get("state_dict", checkpoint)
+        self.model.load_state_dict(state_dict, strict=True)
+        self.model.eval()
+        for p in self.model.parameters():
+            p.requires_grad = False
+
+    def act(self, obs: dict) -> np.ndarray:
+        image_size = int(self.config.get("image_size", [224, 224])[0])
+        image = np.asarray(
+            obs.get("image", np.zeros((image_size, image_size, 3), dtype=np.uint8)),
+            dtype=np.uint8,
         )
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
-        out = self.model(**inputs, output_hidden_states=True, use_cache=False)
-        hidden = out.hidden_states[-1]
-        pooled = last_token_hidden(hidden, inputs["attention_mask"])
-        action = torch.tanh(self.head(pooled.float())).squeeze(0).cpu().numpy()
-        return np.clip(action, -1.0, 1.0).astype(np.float32)
-
-
-def load_policy(model_dir, device, dtype):
-    model_root = Path(model_dir)
-    cfg = json.loads((model_root / "model_config.json").read_text())
-    use_cuda = device == "cuda" and torch.cuda.is_available()
-    torch_device = torch.device("cuda" if use_cuda else "cpu")
-    torch_dtype = (
-        torch.bfloat16
-        if (use_cuda and str(dtype) == "bfloat16")
-        else torch.float32
-    )
-    model = Qwen3VLForConditionalGeneration.from_pretrained(
-        model_root, dtype=torch_dtype
-    ).to(torch_device)
-    processor = AutoProcessor.from_pretrained(model_root)
-    head = build_action_head(int(cfg["hidden_size"]), int(cfg["head_hidden"])).to(
-        torch_device
-    )
-    head.load_state_dict(
-        load_file(str(model_root / "action_head.safetensors"), device=str(torch_device))
-    )
-    return QwenVLAPolicy(model, processor, head, torch_device)
+        if image.ndim == 2:
+            image = np.repeat(image[..., None], 3, axis=-1)
+        if image.shape[-1] > 3:
+            image = image[..., :3]
+        task = str(obs.get("task", ""))
+        difficulty = str(obs.get("difficulty", ""))
+        instruction = str(obs.get("instruction", ""))
+        text = f"task {task} difficulty {difficulty} instruction {instruction}"
+        text_features = _text_vector(text, int(self.config.get("text_dim", 256)))
+        proprio = _proprio_vector(obs, int(self.config.get("proprio_dim", 32)))
+        with torch.no_grad():
+            action = self.model(
+                torch.from_numpy(image).unsqueeze(0).to(self.device),
+                torch.from_numpy(proprio).unsqueeze(0).to(self.device),
+                torch.from_numpy(text_features).unsqueeze(0).to(self.device),
+            )
+        return np.clip(action.squeeze(0).detach().cpu().numpy(), -1.0, 1.0).astype(np.float32)
+
+
+def load_policy(model_dir: str, device: str, dtype: str):
+    return EncoderVLAPolicy(model_dir=model_dir, device=device, dtype=dtype)

model.pt

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

vla_config.json

@@ -0,0 +1,96 @@
+{
+  "schema_version": "robotics_vla_encoder_policy_v1",
+  "model_type": "dinov2_small_text_proprio_bc",
+  "encoder_id": "facebook/dinov2-small",
+  "encoder_config": {
+    "return_dict": true,
+    "output_hidden_states": false,
+    "dtype": "float32",
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "architectures": [
+      "Dinov2Model"
+    ],
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "problem_type": null,
+    "_name_or_path": "",
+    "transformers_version": "5.3.0",
+    "model_type": "dinov2",
+    "hidden_size": 384,
+    "num_hidden_layers": 12,
+    "num_attention_heads": 6,
+    "mlp_ratio": 4,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "attention_probs_dropout_prob": 0.0,
+    "initializer_range": 0.02,
+    "layer_norm_eps": 1e-06,
+    "image_size": 518,
+    "patch_size": 14,
+    "num_channels": 3,
+    "qkv_bias": true,
+    "layerscale_value": 1.0,
+    "drop_path_rate": 0.0,
+    "use_swiglu_ffn": false,
+    "stage_names": [
+      "stem",
+      "stage1",
+      "stage2",
+      "stage3",
+      "stage4",
+      "stage5",
+      "stage6",
+      "stage7",
+      "stage8",
+      "stage9",
+      "stage10",
+      "stage11",
+      "stage12"
+    ],
+    "apply_layernorm": true,
+    "reshape_hidden_states": true,
+    "use_mask_token": true,
+    "output_attentions": false,
+    "out_features": [
+      "stage12"
+    ],
+    "out_indices": [
+      12
+    ]
+  },
+  "text_dim": 256,
+  "proprio_dim": 32,
+  "action_dim": 7,
+  "image_size": [
+    224,
+    224
+  ],
+  "vision_dim": 384,
+  "modality_emb_dim": 256,
+  "parameter_count": 24434823,
+  "freeze_encoder": true,
+  "adapter_inputs": [
+    "image",
+    "proprio",
+    "instruction",
+    "task",
+    "difficulty",
+    "step",
+    "horizon"
+  ],
+  "forbidden_adapter_inputs": [
+    "raw_obs",
+    "object_state",
+    "sim_state",
+    "mujoco_data"
+  ],
+  "training_data": "hf:random-sequence/flock-robotics-vla-training-v2",
+  "run_tag": "dino_v10_baseline"
+}