flock_robotics_adapter.py

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)