Upload folder using huggingface_hub

README.md

@@ -0,0 +1,2 @@
+Self-contained Robotics VLA policy for FLock task 23.
+The adapter only consumes the documented validator observation dictionary.

flock_robotics_adapter.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+import hashlib
+import json
+import math
+import re
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+VARIANTS = {
+    "small_cnn": {"width": 32, "depth": 3, "chunk": 1, "dropout": 0.05},
+    "base_cnn": {"width": 64, "depth": 4, "chunk": 1, "dropout": 0.10},
+    "chunk_cnn": {"width": 48, "depth": 4, "chunk": 8, "dropout": 0.10},
+    "ensemble_cnn": {"width": 64, "depth": 5, "chunk": 8, "dropout": 0.15},
+}
+
+class VisionBackbone(nn.Module):
+    def __init__(self, width, depth):
+        super().__init__()
+        layers = [nn.Conv2d(3, width, 5, 2, 2), nn.GroupNorm(max(1, width // 8), width), nn.SiLU()]
+        channels = width
+        for index in range(depth - 1):
+            next_channels = min(width * (2 ** min(index + 1, 2)), 256)
+            layers += [nn.Conv2d(channels, next_channels, 3, 2, 1),
+                       nn.GroupNorm(max(1, next_channels // 8), next_channels), nn.SiLU()]
+            channels = next_channels
+        layers += [nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(channels, 384), nn.SiLU()]
+        self.net = nn.Sequential(*layers)
+    def forward(self, images):
+        return self.net(images)
+
+class PolicyNet(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        variant = config["variant"]
+        cfg = VARIANTS[variant]
+        self.chunk = int(cfg["chunk"])
+        self.vision = VisionBackbone(int(cfg["width"]), int(cfg["depth"]))
+        self.proprio = nn.Sequential(nn.Linear(int(config["proprio_dim"]), 128), nn.LayerNorm(128), nn.SiLU())
+        self.text = nn.Sequential(nn.Linear(int(config["text_dim"]), 128), nn.LayerNorm(128), nn.SiLU())
+        self.task = nn.Embedding(int(config["task_count"]), 32)
+        self.difficulty = nn.Embedding(int(config["difficulty_count"]), 16)
+        self.head = nn.Sequential(
+            nn.Linear(384 + 128 + 128 + 32 + 16, 512), nn.LayerNorm(512), nn.SiLU(),
+            nn.Dropout(float(cfg["dropout"])), nn.Linear(512, 256), nn.SiLU(),
+            nn.Linear(256, self.chunk * 7))
+    def forward(self, images, proprio, text, task_id, difficulty_id):
+        if images.ndim != 4:
+            raise ValueError("image batch must have four dimensions")
+        if images.shape[-1] == 3:
+            images = images.permute(0, 3, 1, 2)
+        images = images.float()
+        if images.max() > 2:
+            images = images / 255.0
+        images = F.interpolate(images, size=(96, 96), mode="bilinear", align_corners=False)
+        fused = torch.cat([
+            self.vision(images), self.proprio(proprio.float()), self.text(text.float()),
+            self.task(task_id.clamp(0, self.task.num_embeddings - 1)),
+            self.difficulty(difficulty_id.clamp(0, self.difficulty.num_embeddings - 1)),
+        ], dim=-1)
+        return torch.tanh(self.head(fused)).reshape(-1, self.chunk, 7)
+
+def _text_vector(text, dim):
+    result = np.zeros(dim, dtype=np.float32)
+    for token in re.findall(r"[a-z0-9_]+", str(text).lower()):
+        value = int.from_bytes(hashlib.blake2b(token.encode(), digest_size=8).digest(), "little")
+        result[value % dim] += 1.0 if value & 1 else -1.0
+    norm = float(np.linalg.norm(result))
+    return result / norm if norm else result
+
+def _proprio(obs, config):
+    value = np.asarray(obs.get("proprio", np.zeros(25, dtype=np.float32)), dtype=np.float32).reshape(-1)
+    if value.size == 21:
+        value = np.pad(value, (0, 4))
+    if value.size != 25:
+        raise ValueError(f"proprio must have 25 values, got {value.size}")
+    step = float(obs.get("step", 0))
+    horizon = float(obs.get("horizon", 320) or 320)
+    value = np.concatenate([value, np.asarray([step / max(horizon, 1.0)], dtype=np.float32)])
+    if value.size < int(config["proprio_dim"]):
+        value = np.pad(value, (0, int(config["proprio_dim"]) - value.size))
+    value = value[: int(config["proprio_dim"])]
+    mean = np.asarray(config["proprio_mean"], dtype=np.float32)
+    std = np.asarray(config["proprio_std"], dtype=np.float32)
+    return ((value - mean) / np.maximum(std, 1e-4)).astype(np.float32)
+
+class Policy:
+    def __init__(self, model, config, device):
+        self.model, self.config, self.device = model, config, device
+        self.chunk = int(config["chunk"])
+        self.last_step = None
+        self.chunks = {}
+        self.last_action = np.zeros(7, dtype=np.float32)
+        self.task_to_id = config["task_to_id"]
+        self.difficulty_to_id = config["difficulty_to_id"]
+    @torch.inference_mode()
+    def act(self, obs):
+        image = np.asarray(obs.get("image", np.zeros((96, 96, 3), dtype=np.uint8)))
+        if image.ndim == 2:
+            image = np.repeat(image[..., None], 3, axis=-1)
+        if image.shape[-1] == 4:
+            image = image[..., :3]
+        if image.shape[-1] != 3:
+            raise ValueError("image must have 3 or 4 channels")
+        step = int(obs.get("step", 0))
+        if self.last_step is None or step == 0 or step != self.last_step + 1:
+            self.chunks = {}
+        if step == self.last_step:
+            return self.last_action.copy()
+        task = str(obs.get("task", ""))
+        difficulty = str(obs.get("difficulty", "") or "")
+        text = f"task {task} difficulty {difficulty} instruction {obs.get('instruction', '')}"
+        text_feature = _text_vector(text, int(self.config["text_dim"]))
+        proprio = _proprio(obs, self.config)
+        task_id = self.task_to_id.get(task, len(self.task_to_id))
+        difficulty_id = self.difficulty_to_id.get(difficulty, len(self.difficulty_to_id))
+        raw = self.model(
+            torch.from_numpy(image).unsqueeze(0).to(self.device),
+            torch.from_numpy(proprio).unsqueeze(0).to(self.device),
+            torch.from_numpy(text_feature).unsqueeze(0).to(self.device),
+            torch.tensor([task_id], device=self.device),
+            torch.tensor([difficulty_id], device=self.device),
+        )[0].float().cpu().numpy()
+        self.chunks[step] = raw
+        candidates, weights = [], []
+        for start, chunk in list(self.chunks.items()):
+            offset = step - start
+            if 0 <= offset < self.chunk:
+                candidates.append(chunk[offset])
+                weights.append(math.exp(-0.55 * offset))
+            else:
+                del self.chunks[start]
+        action = np.average(np.asarray(candidates), axis=0, weights=np.asarray(weights))
+        action = np.clip(action, -1.0, 1.0).astype(np.float32)
+        if action[6] > 0.12:
+            action[6] = 1.0
+        elif action[6] < -0.12:
+            action[6] = -1.0
+        self.last_step, self.last_action = step, action.copy()
+        return action
+
+def load_policy(model_dir: str, device: str, dtype: str):
+    root = Path(model_dir)
+    config = json.loads((root / "vla_config.json").read_text())
+    selected = torch.device(device if str(device).startswith("cuda") and torch.cuda.is_available() else "cpu")
+    model = PolicyNet(config).to(selected)
+    checkpoint = torch.load(root / "model.pt", map_location=selected, weights_only=True)
+    state = checkpoint["state_dict"] if "state_dict" in checkpoint else checkpoint
+    model.load_state_dict(state, strict=True)
+    model.eval()
+    return Policy(model, config, selected)

model.pt

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

vla_config.json

@@ -0,0 +1,116 @@
+{
+  "schema_version": "task23_self_contained_policy_v1",
+  "model_type": "cnn_conditioned_action_chunk",
+  "variant": "small_cnn",
+  "chunk": 1,
+  "text_dim": 128,
+  "proprio_dim": 32,
+  "task_count": 14,
+  "difficulty_count": 5,
+  "proprio_mean": [
+    -0.013340266421437263,
+    0.7383145093917847,
+    0.011660933494567871,
+    -1.9436376094818115,
+    0.019467052072286606,
+    2.826350212097168,
+    0.7630648612976074,
+    0.0225653275847435,
+    0.052658289670944214,
+    0.009153264574706554,
+    0.07257223129272461,
+    -0.0217689611017704,
+    -0.014813058078289032,
+    0.045801397413015366,
+    0.07692108303308487,
+    -0.06413360685110092,
+    0.9298723936080933,
+    0.9961475133895874,
+    -0.004522569477558136,
+    0.07633855938911438,
+    -0.005218956153839827,
+    0.03139699995517731,
+    -0.030553586781024933,
+    0.0001240656420122832,
+    5.485760084411595e-06,
+    0.49797627329826355,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "proprio_std": [
+    0.18346060812473297,
+    0.2751201093196869,
+    0.08362191170454025,
+    0.4013335406780243,
+    0.1880636066198349,
+    0.22030314803123474,
+    0.36815357208251953,
+    0.09236042201519012,
+    0.24793101847171783,
+    0.07140849530696869,
+    0.27341264486312866,
+    0.12188053876161575,
+    0.17399784922599792,
+    0.20990528166294098,
+    0.09051761776208878,
+    0.1660534143447876,
+    0.07342937588691711,
+    0.0047762468457221985,
+    0.021420147269964218,
+    0.03227067366242409,
+    0.016867103055119514,
+    0.008550216443836689,
+    0.009234834462404251,
+    0.015089103952050209,
+    0.012418420054018497,
+    0.2886728048324585,
+    1.0,
+    1.0,
+    1.0,
+    1.0,
+    1.0,
+    1.0
+  ],
+  "task_to_id": {
+    "lift_cube": 0,
+    "pick_place_can": 1,
+    "pick_place_milk": 2,
+    "pick_place_bread": 3,
+    "pick_place_cereal": 4,
+    "pick_place_clutter": 5,
+    "stack_blocks": 6,
+    "open_door": 7,
+    "nut_assembly": 8,
+    "nut_assembly_square": 9,
+    "nut_assembly_round": 10,
+    "wipe_table": 11,
+    "tool_hang": 12
+  },
+  "difficulty_to_id": {
+    "low": 0,
+    "medium": 1,
+    "hard": 2,
+    "very_high": 3
+  },
+  "adapter_inputs": [
+    "image",
+    "instruction",
+    "proprio",
+    "task",
+    "step",
+    "difficulty",
+    "horizon"
+  ],
+  "forbidden_adapter_inputs": [
+    "raw_obs",
+    "object_state",
+    "sim_state",
+    "mujoco_data"
+  ],
+  "training_data": "data/hf",
+  "trainval": false
+}