flock_robotics_adapter.py

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

class SigLIPActionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk = int(config["chunk"])
        self.proprio = nn.Sequential(nn.Linear(32, 128), nn.LayerNorm(128), nn.SiLU())
        self.text = nn.Sequential(nn.Linear(128, 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(768 + 128 + 128 + 32 + 16, 512),
            nn.LayerNorm(512), nn.SiLU(), nn.Dropout(0.10),
            nn.Linear(512, 256), nn.SiLU(), nn.Linear(256, self.chunk * 7)
        )
    def forward(self, features, proprio, text, task_id, difficulty_id):
        fused = torch.cat([
            features.float(), 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("utf-8"), 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)])
    value = np.pad(value, (0, max(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, backbone, processor, head, config, device):
        self.backbone, self.processor = backbone, processor
        self.head, self.config, self.device = head, config, device
        self.chunk = int(config["chunk"])
        self.last_step, self.last_action = None, np.zeros(7, dtype=np.float32)
        self.chunks = {}
        self.task_to_id, self.difficulty_to_id = config["task_to_id"], config["difficulty_to_id"]
    @torch.inference_mode()
    def act(self, obs):
        image = np.asarray(obs.get("image", np.zeros((224, 224, 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()
        inputs = self.processor(images=image, return_tensors="pt")
        inputs = {key: value.to(self.device) for key, value in inputs.items()}
        result = self.backbone(**inputs)
        features = result.pooler_output if hasattr(result, "pooler_output") else result.last_hidden_state.mean(dim=1)
        task = str(obs.get("task", ""))
        difficulty = str(obs.get("difficulty", "") or "")
        text = f"task {task} difficulty {difficulty} instruction {obs.get('instruction', '')}"
        raw = self.head(
            features, torch.from_numpy(_proprio(obs, self.config)).unsqueeze(0).to(self.device),
            torch.from_numpy(_text_vector(text, 128)).unsqueeze(0).to(self.device),
            torch.tensor([self.task_to_id.get(task, len(self.task_to_id))], device=self.device),
            torch.tensor([self.difficulty_to_id.get(difficulty, len(self.difficulty_to_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):
    from transformers import AutoImageProcessor, AutoModel
    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")
    processor = AutoImageProcessor.from_pretrained(root / "backbone", local_files_only=True)
    backbone = AutoModel.from_pretrained(root / "backbone", local_files_only=True).to(selected).eval()
    head = SigLIPActionHead(config).to(selected)
    state = torch.load(root / "action_head.pt", map_location=selected, weights_only=True)
    head.load_state_dict(state["state_dict"], strict=True)
    head.eval()
    return Policy(backbone, processor, head, config, selected)