from __future__ import annotations import json import math import sys from collections import deque from pathlib import Path import numpy as np import torch from safetensors.torch import load_file _MODEL_DIR = Path(__file__).resolve().parent if str(_MODEL_DIR) not in sys.path: sys.path.insert(0, str(_MODEL_DIR)) from competitive_vla_model import phase_vector, text_vector # noqa: E402 from multibackbone_vla_model import MultiBackboneRoutedVLAModel # noqa: E402 def _resolve_device(requested: str) -> torch.device: if requested.startswith("cuda") and torch.cuda.is_available(): return torch.device(requested) if ( requested == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): return torch.device("mps") return torch.device("cpu") def _resolve_dtype(name: str, device: torch.device) -> torch.dtype: if device.type != "cuda": return torch.float32 normalized = str(name).lower().replace("torch.", "") if normalized in {"float16", "fp16", "half"}: return torch.float16 if normalized in {"float32", "fp32"}: return torch.float32 return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 class MultiBackboneRoutedVLAPolicy: def __init__(self, model_dir: str, device: str, dtype: str) -> None: self.model_dir = Path(model_dir) self.config = json.loads((self.model_dir / "vla_config.json").read_text()) self.device = _resolve_device(device) self.dtype = _resolve_dtype(dtype, self.device) self.model = MultiBackboneRoutedVLAModel(self.config) state = load_file(str(self.model_dir / "model.safetensors"), device="cpu") missing, unexpected = self.model.load_state_dict(state, strict=True) if missing or unexpected: raise RuntimeError( f"invalid checkpoint; missing={missing}, unexpected={unexpected}" ) self.model.to(device=self.device, dtype=self.dtype).eval() self.base_router_config = self.config["base_router_config"] self.large_config = self.config["large_config"] self.condition_routes = { str(key): str(value) for key, value in self.config.get("condition_routes", {}).items() } self.condition_inference_overrides = { str(key): dict(value) for key, value in self.config.get( "condition_inference_overrides", {} ).items() } self.task_routes = { str(key): str(value) for key, value in self.config.get("fallback_task_routes", {}).items() } self.default_route = str(self.config.get("default_route", "general")) reference = self.base_router_config["route_configs"][self.default_route] self.history = int(reference["history"]) self.action_chunk = int(reference["action_chunk"]) self.visual_history: deque[torch.Tensor] = deque(maxlen=self.history) self.rgb_history: deque[torch.Tensor] = deque(maxlen=self.history) self.proprio_history: deque[torch.Tensor] = deque(maxlen=self.history) self.phase_history: deque[torch.Tensor] = deque(maxlen=self.history) self.prediction_chunks: dict[int, np.ndarray] = {} self.last_step = -1 self.episode_key: tuple[str, str, str] | None = None self.gripper_state = -1.0 def reset(self) -> None: self.visual_history.clear() self.rgb_history.clear() self.proprio_history.clear() self.phase_history.clear() self.prediction_chunks.clear() self.last_step = -1 self.gripper_state = -1.0 def _maybe_reset(self, obs: dict) -> None: step = int(obs.get("step", 0)) key = ( str(obs.get("task", "")), str(obs.get("difficulty", "") or ""), str(obs.get("instruction", "")), ) if step == 0 or step <= self.last_step or self.episode_key != key: self.reset() self.episode_key = key def _stack_history(self, values: deque[torch.Tensor]) -> torch.Tensor: if not values: raise RuntimeError("history is unexpectedly empty") padded = [values[0]] * (self.history - len(values)) + list(values) return torch.stack(padded, dim=0).unsqueeze(0) def _route(self, task: str, difficulty: str) -> str: condition = f"{task}::{difficulty}" return self.condition_routes.get( condition, self.task_routes.get(task, self.default_route) ) @torch.inference_mode() def act(self, obs: dict) -> np.ndarray: self._maybe_reset(obs) step = int(obs.get("step", 0)) horizon = int(obs.get("horizon", 320) or 320) task = str(obs.get("task", "")) difficulty = str(obs.get("difficulty", "") or "") route = self._route(task, difficulty) inference_override = self.condition_inference_overrides.get( f"{task}::{difficulty}", {} ) if route == "large": route_config = self.large_config active_model = self.model.large_model else: route_config = self.base_router_config["route_configs"][route] active_model = self.model.base_model 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.ndim != 3: raise ValueError(f"expected HWC image, got {image.shape}") image_tensor = torch.as_tensor( image[..., :3].copy(), device=self.device ).unsqueeze(0) visual, rgb = active_model.encode_images(image_tensor) self.visual_history.append(visual.squeeze(0)) self.rgb_history.append(rgb.squeeze(0)) proprio_dim = int(route_config["proprio_dim"]) proprio = np.asarray( obs.get("proprio", np.zeros(proprio_dim, dtype=np.float32)), dtype=np.float32, ).reshape(-1) if proprio.size < proprio_dim: proprio = np.pad(proprio, (0, proprio_dim - proprio.size)) self.proprio_history.append( torch.as_tensor(proprio[:proprio_dim], device=self.device, dtype=self.dtype) ) self.phase_history.append( torch.as_tensor( phase_vector(step, horizon), device=self.device, dtype=self.dtype ) ) instruction = str(obs.get("instruction", "")) condition_text = ( f"task {task} difficulty {difficulty} instruction {instruction}" ) text_features = torch.as_tensor( text_vector(condition_text, int(route_config["text_dim"])), device=self.device, dtype=self.dtype, ).unsqueeze(0) task_map = route_config["task_to_id"] difficulty_map = route_config["difficulty_to_id"] model_inputs = ( self._stack_history(self.visual_history), self._stack_history(self.rgb_history), self._stack_history(self.proprio_history), self._stack_history(self.phase_history), text_features, torch.tensor( [int(task_map.get(task, len(task_map)))], device=self.device, dtype=torch.long, ), torch.tensor( [int(difficulty_map.get(difficulty, len(difficulty_map)))], device=self.device, dtype=torch.long, ), ) if route == "large": raw = self.model.large_model.forward_cached(*model_inputs) else: raw = self.model.base_model.forward_cached(route, *model_inputs) chunk = torch.tanh(raw).mean(dim=0).squeeze(0).float().cpu().numpy() self.prediction_chunks[step] = chunk candidates: list[np.ndarray] = [] weights: list[float] = [] decay = float( inference_override.get( "temporal_ensemble_decay", route_config.get("temporal_ensemble_decay", 0.55), ) ) for start, predicted in list(self.prediction_chunks.items()): offset = step - start if offset < 0 or offset >= self.action_chunk: self.prediction_chunks.pop(start, None) continue candidates.append(predicted[offset]) weights.append(math.exp(-decay * float(offset))) if not candidates: raise RuntimeError("temporal ensemble has no candidate action") action = np.average( np.stack(candidates), axis=0, weights=np.asarray(weights, dtype=np.float32), ).astype(np.float32) action[:3] *= float( inference_override.get( "translation_scale", route_config.get("translation_scale", 1.0) ) ) hysteresis = float( inference_override.get( "gripper_hysteresis", route_config.get("gripper_hysteresis", 0.12) ) ) gripper_score = float(action[6]) if gripper_score > hysteresis: self.gripper_state = 1.0 elif gripper_score < -hysteresis: self.gripper_state = -1.0 action[6] = self.gripper_state if bool(route_config.get("fixed_rotation", False)): action[3:6] = 0.0 self.last_step = step return np.clip(action, -1.0, 1.0).astype(np.float32) def load_policy( model_dir: str, device: str, dtype: str ) -> MultiBackboneRoutedVLAPolicy: return MultiBackboneRoutedVLAPolicy(model_dir, device, dtype)