flock_robotics_adapter.py

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)