File size: 6,003 Bytes
109a8f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | 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)
|