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)