update

.gitattributes

@@ -0,0 +1 @@
+model.pt filter=lfs diff=lfs merge=lfs -text

flock_robotics_adapter.py

@@ -0,0 +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
+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(),
+        )
+        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,
+        )
+        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:8433d4893118995f0a3538134fe5faceb2ba2e8a9f5d0fa35df1b47581aba757
+size 97821951

training_report.json

@@ -0,0 +1,167 @@
+{
+  "samples": {
+    "trajectory_count": 103,
+    "loaded_frame_count": 25328,
+    "step_stride": 1,
+    "source": "random-sequence/flock-robotics-vla-training-v2"
+  },
+  "epochs": 25,
+  "batch_size": 64,
+  "learning_rate": 0.0003,
+  "weight_decay": 0.0001,
+  "device": "cuda",
+  "encoder_id": "facebook/dinov2-small",
+  "best_val_mse": 0.005434711876076663,
+  "history": [
+    {
+      "epoch": 1.0,
+      "train_mse": 0.052450973142346244,
+      "val_mse": 0.02464906360194864,
+      "lr": 0.0002988172051971717
+    },
+    {
+      "epoch": 2.0,
+      "train_mse": 0.015353698004698999,
+      "val_mse": 0.013575026146055304,
+      "lr": 0.0002952874741692947
+    },
+    {
+      "epoch": 3.0,
+      "train_mse": 0.012783010591259027,
+      "val_mse": 0.01241856964531621,
+      "lr": 0.0002894664728832377
+    },
+    {
+      "epoch": 4.0,
+      "train_mse": 0.010736383704168101,
+      "val_mse": 0.011335676267454738,
+      "lr": 0.00028144600200657956
+    },
+    {
+      "epoch": 5.0,
+      "train_mse": 0.0100194234388602,
+      "val_mse": 0.011257341734058625,
+      "lr": 0.00027135254915624217
+    },
+    {
+      "epoch": 6.0,
+      "train_mse": 0.009152803605718239,
+      "val_mse": 0.009612169949849847,
+      "lr": 0.0002593452941132118
+    },
+    {
+      "epoch": 7.0,
+      "train_mse": 0.00936236106123713,
+      "val_mse": 0.007432314051582962,
+      "lr": 0.00024561359846230347
+    },
+    {
+      "epoch": 8.0,
+      "train_mse": 0.008174229212043036,
+      "val_mse": 0.009825789436586142,
+      "lr": 0.0002303740192468495
+    },
+    {
+      "epoch": 9.0,
+      "train_mse": 0.008620002260551634,
+      "val_mse": 0.01041258188070491,
+      "lr": 0.0002138668937347609
+    },
+    {
+      "epoch": 10.0,
+      "train_mse": 0.007946362478426172,
+      "val_mse": 0.008026734568284587,
+      "lr": 0.0001963525491562421
+    },
+    {
+      "epoch": 11.0,
+      "train_mse": 0.007336850237303667,
+      "val_mse": 0.007538307731085432,
+      "lr": 0.00017810719718785873
+    },
+    {
+      "epoch": 12.0,
+      "train_mse": 0.00767612234281653,
+      "val_mse": 0.006661466446978822,
+      "lr": 0.00015941857792939703
+    },
+    {
+      "epoch": 13.0,
+      "train_mse": 0.00704610332234791,
+      "val_mse": 0.006911170980812155,
+      "lr": 0.00014058142207060302
+    },
+    {
+      "epoch": 14.0,
+      "train_mse": 0.006659559791542447,
+      "val_mse": 0.009048803709079564,
+      "lr": 0.00012189280281214133
+    },
+    {
+      "epoch": 15.0,
+      "train_mse": 0.006679477127921169,
+      "val_mse": 0.007556101494861581,
+      "lr": 0.00010364745084375795
+    },
+    {
+      "epoch": 16.0,
+      "train_mse": 0.006114454294127817,
+      "val_mse": 0.007176910378450843,
+      "lr": 8.61331062652391e-05
+    },
+    {
+      "epoch": 17.0,
+      "train_mse": 0.0062116264081053425,
+      "val_mse": 0.005434711876076663,
+      "lr": 6.962598075315047e-05
+    },
+    {
+      "epoch": 18.0,
+      "train_mse": 0.005705748372087895,
+      "val_mse": 0.006514828746489545,
+      "lr": 5.4386401537696536e-05
+    },
+    {
+      "epoch": 19.0,
+      "train_mse": 0.005474576261368675,
+      "val_mse": 0.005767398440657186,
+      "lr": 4.06547058867883e-05
+    },
+    {
+      "epoch": 20.0,
+      "train_mse": 0.0052186791698284486,
+      "val_mse": 0.006887162315010389,
+      "lr": 2.86474508437579e-05
+    },
+    {
+      "epoch": 21.0,
+      "train_mse": 0.004928951738128989,
+      "val_mse": 0.00600228262520844,
+      "lr": 1.8553997993420465e-05
+    },
+    {
+      "epoch": 22.0,
+      "train_mse": 0.004528539363558345,
+      "val_mse": 0.005473934499087845,
+      "lr": 1.05335271167623e-05
+    },
+    {
+      "epoch": 23.0,
+      "train_mse": 0.0043255689688378035,
+      "val_mse": 0.006198732651723034,
+      "lr": 4.712525830705356e-06
+    },
+    {
+      "epoch": 24.0,
+      "train_mse": 0.0040297096492367315,
+      "val_mse": 0.005554187416847895,
+      "lr": 1.1827948028283355e-06
+    },
+    {
+      "epoch": 25.0,
+      "train_mse": 0.0038712321390164763,
+      "val_mse": 0.005647414989951898,
+      "lr": 0.0
+    }
+  ]
+}

vla_config.json

@@ -0,0 +1,95 @@
+{
+  "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"
+}