diff --git a/.gitignore b/.gitignore index 1d3f9d5..2be02dd 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ _deps/ # Models / data *.gguf +*.gguf.tensor_map.json *.bin *.safetensors *.pt diff --git a/CMakeLists.txt b/CMakeLists.txt index 01e34be..b0fc692 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -100,6 +100,29 @@ function(vla_exclude_fetched_targets dir) endfunction() vla_exclude_fetched_targets(${llama_SOURCE_DIR}) +# Octo is the one arch that tokenizes in-process, so it is also the one that +# needs SentencePiece and, through it, a protobuf. Optional so a build that does +# not want Octo does not inherit that dependency. +option(VLA_OCTO "Build the Octo arch (fetches SentencePiece, needs system protobuf)" ON) +if(VLA_OCTO) + # The system protobuf, not sentencepiece's vendored protobuf-lite: vla-server + # aborts at static-init if two protobuf runtimes reach the same binary. + set(SPM_ENABLE_SHARED OFF CACHE BOOL "" FORCE) + set(SPM_BUILD_TEST OFF CACHE BOOL "" FORCE) + set(SPM_ENABLE_TCMALLOC OFF CACHE BOOL "" FORCE) + set(SPM_PROTOBUF_PROVIDER "package" CACHE STRING "" FORCE) + FetchContent_Declare(sentencepiece + GIT_REPOSITORY https://github.com/google/sentencepiece + GIT_TAG v0.2.0 + GIT_SHALLOW TRUE + ) + # sentencepiece v0.2.0 still asks for cmake_minimum_required(VERSION 3.1), + # which CMake 4 refuses outright. + set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + FetchContent_MakeAvailable(sentencepiece) + unset(CMAKE_POLICY_VERSION_MINIMUM) +endif() + add_library(vla_core src/model.cpp src/loader.cpp @@ -129,6 +152,12 @@ target_include_directories(vla_core ) # The VLA archs call no llama_* API; only vlm_core needs llama. target_link_libraries(vla_core PUBLIC ggml) +if(VLA_OCTO) + target_sources(vla_core PRIVATE src/models/octo.cpp) + target_include_directories(vla_core PRIVATE ${sentencepiece_SOURCE_DIR}/src) + target_link_libraries(vla_core PRIVATE sentencepiece-static) + target_compile_definitions(vla_core PUBLIC VLA_USE_OCTO) +endif() if(GGML_CUDA) target_compile_definitions(vla_core PUBLIC GGML_USE_CUDA) diff --git a/README.md b/README.md index dfeae6d..7978201 100644 --- a/README.md +++ b/README.md @@ -325,6 +325,7 @@ supported (released and benchmarked), `~` = in progress, `-` = planned. | [VLA-Adapter](https://hf.co/vrfai/vla-adapter-libero-gguf) | Y | Y | ~ | Y | Y | - | | [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | - | Y | Y | - | | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | - | Y | Y | - | +| [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf) | Y | Y | Y | Y | - | - | --- @@ -362,6 +363,7 @@ Supported VLA models: - [OpenVLA-OFT](https://github.com/moojink/openvla-oft) - Moo Jin Kim et al. - [GR00T N1.x](https://github.com/NVIDIA/Isaac-GR00T) - NVIDIA Isaac. - [VLA-JEPA](https://github.com/ginwind/VLA-JEPA) - Jingwen Sun et al. +- [Octo](https://github.com/octo-models/octo) - Octo Model Team, UC Berkeley RAIL. Built on: diff --git a/eval/client/adapters.py b/eval/client/adapters.py index 10a203a..b9d3c8d 100644 --- a/eval/client/adapters.py +++ b/eval/client/adapters.py @@ -16,6 +16,7 @@ from typing import Any import numpy as np import torch +from PIL import Image from tree import map_structure from lerobot.envs.utils import preprocess_observation @@ -23,6 +24,17 @@ from lerobot.processor.pipeline import PolicyProcessorPipeline from lerobot.utils.constants import ACTION +def octo_preprocess_image(frame: np.ndarray, image_size: int = 256) -> np.ndarray: + # Rotate 180 like every other LIBERO adapter here (the off-screen render comes + # out upside-down), then resize. LIBERO renders at 256 and Octo's primary + # tokenizer wants 256, so the resize only bites if the camera is reconfigured. + rotated = np.ascontiguousarray(frame[::-1, ::-1]) + if rotated.shape[0] == image_size and rotated.shape[1] == image_size: + return rotated + resized = Image.fromarray(rotated).resize((image_size, image_size), resample=Image.LANCZOS) + return np.asarray(resized, dtype=np.uint8) + + class BasePipelineAdapter: def __init__(self, client: Any = None): self._client = client @@ -153,3 +165,21 @@ class Gr00tN15PipelineAdapter(Gr00tPipelineAdapter): def parse_action(self, action: np.ndarray) -> np.ndarray: return np.asarray(action[:7], dtype=np.float32).copy() + +# The LIBERO finetunes are single-camera: their image_obs_keys never held a wrist +# key, so sending the primary view alone is what the checkpoint trained on. Octo's +# observation tokenizers are image-only, so there is no state to send either. +class OctoPipelineAdapter(BasePipelineAdapter): + + def parse_observation(self, obs: dict[str, Any]) -> dict[str, Any]: + return { + "observation.images.image": octo_preprocess_image(obs["pixels"]["image"], image_size=256), + "task": obs.get("task_description", ""), + } + + def parse_action(self, action: np.ndarray) -> np.ndarray: + # The server already returned world units. Only the gripper needs Octo's + # +1=open/0=close turned into LIBERO's -1=open/+1=close. + action = np.asarray(action[:7], dtype=np.float32).copy() + action[6] = -1.0 if action[6] > 0.5 else 1.0 + return action diff --git a/eval/client/run_sim_client_direct.py b/eval/client/run_sim_client_direct.py index 1662b79..6d65a2c 100644 --- a/eval/client/run_sim_client_direct.py +++ b/eval/client/run_sim_client_direct.py @@ -29,6 +29,7 @@ Evo1PipelineAdapter, Gr00tPipelineAdapter, Gr00tN15PipelineAdapter, + OctoPipelineAdapter, ) ARCH_CHOICES = sorted(ARCH_PRESETS) @@ -127,6 +128,8 @@ elif args.arch in ("gr00t_n1_6", "gr00t_n1_7"): client = Gr00tPipelineAdapter(client=client) + elif args.arch == "octo": + client = OctoPipelineAdapter(client=client) else: client = LeRobotPipelineAdapter(client=client) diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 70f18da..0b2849e 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -55,6 +55,9 @@ "max_state_dim": 64, "trust_remote_code": True}, "gr00t_n1_6": {"image_size": 224, "tokenizer": None, "max_state_dim": 128, "trust_remote_code": True}, + + # Single-camera LIBERO finetune, image-only observations, so no state. + "octo": {"image_size": 256, "tokenizer": "t5-base", "max_state_dim": 0, "max_length": 16}, } BITVLA_N_PATCHES_PER_VIEW = 256 @@ -509,6 +512,8 @@ def get_action(self, observations: dict[str, Any]) -> np.ndarray: chunk = self._predict_chunk_vla_adapter(observations) elif self.arch == "openvla_oft": chunk = self._predict_chunk_openvla_oft(observations) + elif self.arch == "octo": + chunk = self._predict_chunk_octo(observations) else: chunk = self._predict_chunk(observations) for row in chunk[: self.n_action_steps, : self.real_action_dim]: @@ -846,6 +851,60 @@ def _predict_chunk_evo1(self, observations: dict[str, Any]) -> np.ndarray: return (np.array(resp.action_chunk, dtype=np.float32) .reshape(resp.chunk_size, resp.action_dim)) + def _predict_chunk_octo(self, observations: dict[str, Any]) -> np.ndarray: + # OctoPipelineAdapter has already rotated and resized these. A wrist view is + # sent when the observation carries one; a checkpoint without one just + # leaves that slot out of the sequence. + images_u8: list[np.ndarray] = [] + for key in self.image_keys[:2]: + if key not in observations: + continue + img = observations[key] + if isinstance(img, torch.Tensor): + img = img.numpy() + img = np.asarray(img, dtype=np.uint8) + if img.ndim != 3 or img.shape[2] != 3: + raise ValueError(f"octo: {key} expected HWC uint8 [H,W,3], got {img.shape}") + images_u8.append(np.ascontiguousarray(img, dtype=np.uint8)) + if not images_u8: + raise KeyError(f"octo: no image keys found in observations; got {list(observations.keys())}") + + task = observations.get("task", "") + if isinstance(task, bytes): + task = task.decode() + # The checkpoint's own text_processor: t5-base, max_length=16, + # padding="max_length", truncation=True. Octo's T5 encoder needs the real + # padding mask, so it is sent alongside the ids. + toks = self.tok(task, return_tensors="np", padding="max_length", + truncation=True, max_length=self.max_length) + input_ids = toks["input_ids"][0].astype(np.int32) + attn_mask = toks["attention_mask"][0].astype(np.int32) + + req = self.pb.PredictRequest() + req.request_id = self._step + self._step += 1 + for img in images_u8: + ip = req.images.add() + ip.encoding = self.pb.Image.RGB_U8 + ip.height = img.shape[0] + ip.width = img.shape[1] + ip.data = img.tobytes() + req.lang_tokens.extend(input_ids.tolist()) + req.attention_mask.extend(attn_mask.tolist()) + + self.sock.send(req.SerializeToString()) + body = self.sock.recv() + resp = self.pb.PredictResponse() + resp.ParseFromString(body) + if resp.error: + raise RuntimeError(f"vla-server error: {resp.error}") + self._last_response = resp + # World units already: Octo's dataset_statistics lives inside the checkpoint + # GGUF, so the server un-normalizes rather than handing the client a + # --stats-json it would have to extract from a multi-hundred-MB file. + return (np.array(resp.action_chunk, dtype=np.float32) + .reshape(resp.chunk_size, resp.action_dim)) + def _predict_chunk_bitvla(self, observations: dict[str, Any]) -> np.ndarray: images_u8: list[np.ndarray] = [] diff --git a/scripts/convert_octo_to_gguf.py b/scripts/convert_octo_to_gguf.py new file mode 100644 index 0000000..cf4fb90 --- /dev/null +++ b/scripts/convert_octo_to_gguf.py @@ -0,0 +1,543 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +import gguf + +ARCH = "octo" +MODEL_ID = "hf://rail-berkeley/octo-small-1.5" +# window_size for the rail-berkeley/octo-small-1.5 bridge pretrain checkpoint (the +# MODEL_ID default above). Used only as a last-resort fallback when converting that +# default checkpoint and its window_size can't be read back out of its own config +# (e.g. hf:// config.json fetch races) -- never used for a checkpoint passed via --ckpt. +DEFAULT_BRIDGE_WINDOW_SIZE = 2 + +OCTO_META: dict[str, Any] = { + "architecture": "octo-small-1.5", + "embedding_length": 384, + "block_count": 12, + "attention.head_count": 6, + "feed_forward_length": 1536, + "attention.layer_norm_eps": 1e-6, + # action.horizon / action.dim / action.head_type are set from the checkpoint's own + # config (model.config["model"]["heads"]["action"]) in main() -- they differ between + # the diffusion libero checkpoints (horizon=4) and L1 pytorch checkpoints (horizon=20). + "readout.count": 1, + "tokens.primary": 256, + "tokens.wrist": 64, + "tokens.language": 16, + "image.primary_size": 256, + "image.wrist_size": 128, + "diffusion.steps": 20, + "diffusion.beta_schedule": "cosine", + "diffusion.s": 0.008, + "diffusion.max_action": 5.0, + "diffusion.time_dim": 32, + "diffusion.hidden": 256, + "diffusion.num_blocks": 3, +} + +# octo.action.head_type values, keyed by the PyTorch action-head class name recorded in +# a checkpoint's own config.json (model.config["model"]["heads"]["action"]["name"]). +HEAD_TYPE_BY_CLASS: dict[str, str] = { + "L1ActionHeadPt": "l1", + "MSEActionHeadPt": "mse", + "DiffusionActionHeadPt": "diffusion", + "UNetDDPMActionHeadPt": "diffusion", +} + +IGNORED_PATTERNS = ( + # tied duplicate of hf_model.shared.weight (same underlying tensor, both names + # appear in state_dict()); only shared.weight is mapped to octo.t5.tok_embd.weight. + re.compile(r"module\.octo_transformer\.task_tokenizers\.language\.hf_model\.encoder\.embed_tokens\.weight"), +) + + +def _json_default(obj: Any) -> Any: + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.generic): + return obj.item() + if isinstance(obj, torch.Tensor): + return obj.detach().cpu().tolist() + raise TypeError(f"cannot JSON encode {type(obj).__name__}") + + +def _add_meta(writer: gguf.GGUFWriter, key: str, value: Any) -> None: + full = f"octo.{key}" + if isinstance(value, str): + writer.add_string(full, value) + elif isinstance(value, bool): + writer.add_bool(full, value) + elif isinstance(value, int): + writer.add_uint32(full, value) + elif isinstance(value, float): + writer.add_float32(full, value) + else: + raise TypeError(f"unsupported metadata {full}={value!r}") + + +def _f32(t: torch.Tensor) -> np.ndarray: + return t.detach().to(dtype=torch.float32, device="cpu").contiguous().numpy() + + +def _embed_tokenizer(writer: gguf.GGUFWriter, tokenizer_name: str = "t5-base") -> None: + """Embed the raw T5 SentencePiece unigram model (spiece.model) as a UINT8 GGUF + array (not a GGUF string: the serialized proto contains embedded NUL bytes, + which would truncate a null-terminated-string read). google-t5/t5-base's + tokenizer_name is "t5-base" (see octo-pytorch's octo_pretrain_config.py). + """ + from huggingface_hub import hf_hub_download + + spm_path = hf_hub_download(tokenizer_name, "spiece.model") + spm_bytes = Path(spm_path).read_bytes() + writer.add_array("octo.tokenizer.spm_model", spm_bytes) + # T5 unigram special tokens (fixed across all T5 SentencePiece vocabs): pad=0, eos==1. + writer.add_uint32("octo.tokenizer.eos_id", 1) + writer.add_uint32("octo.tokenizer.pad_id", 0) + + +def _strip_prefix(key: str) -> str: + return key.removeprefix("module.") + + +def map_key(pt_key: str) -> str | None: + k = _strip_prefix(pt_key) + + m = re.fullmatch(r"octo_transformer\.observation_tokenizers\.(primary|wrist)\.encoder_def\.layers\.(\d+)\.0\.(weight|bias)", k) + if m: + view, idx, leaf = m.groups() + return f"octo.obs.{view}.stem.{idx}.conv.{leaf}" + + m = re.fullmatch(r"octo_transformer\.observation_tokenizers\.(primary|wrist)\.encoder_def\.layers\.(\d+)\.1\.(weight|bias)", k) + if m: + view, idx, leaf = m.groups() + return f"octo.obs.{view}.stem.{idx}.gn.{leaf}" + + m = re.fullmatch(r"octo_transformer\.observation_tokenizers\.(primary|wrist)\.encoder_def\.embedding\.(weight|bias)", k) + if m: + view, leaf = m.groups() + return f"octo.obs.{view}.patch_embd.{leaf}" + + m = re.fullmatch(r"octo_transformer\.obs_projections\.obs_(primary|wrist|proprio)_projection\.(weight|bias)", k) + if m: + view, leaf = m.groups() + return f"octo.obs.{view}.proj.{leaf}" + + m = re.fullmatch(r"octo_transformer\.obs_(primary|wrist|proprio)_pos_embedding", k) + if m: + return f"octo.obs.{m.group(1)}.pos_embd" + + # LowdimObsTokenizerPt (proprio): fixed, non-trainable bin edges for the BinTokenizer + # quantization -- not a learned weight, but still needed by the engine to reproduce + # the same binning at inference time. + if k == "octo_transformer.observation_tokenizers.proprio.thresholds": + return "octo.obs.proprio.bin_thresholds" + + m = re.fullmatch(r"octo_transformer\.task_projections\.task_language_projection\.(weight|bias)", k) + if m: + return f"octo.task.language.proj.{m.group(1)}" + if k == "octo_transformer.task_language_pos_embedding": + return "octo.task.language.pos_embd" + if k == "octo_transformer.readout_action_pos_embedding": + return "octo.readout.action.pos_embd" + + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.layer_norm1\.(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.attn_norm.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.self_attention\.in_proj_(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.attn_qkv.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.self_attention\.out_proj\.(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.attn_o.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.layer_norm2\.(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.ffn_norm.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.mlp_block\.dense1\.(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.ffn_up.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.encoder_blocks\.(\d+)\.mlp_block\.dense2\.(weight|bias)", k) + if m: + return f"octo.blk.{m.group(1)}.ffn_down.{m.group(2)}" + m = re.fullmatch(r"octo_transformer\.block_transformer\.transformer\.layer_norm\.(weight|bias)", k) + if m: + return f"octo.output_norm.{m.group(1)}" + + p = "heads.action.map_head." + if k == p + "probe": + return "octo.head.l1.map.probe" + m = re.fullmatch(re.escape(p) + r"attention\.(in_proj_weight|in_proj_bias)", k) + if m: + leaf = "weight" if m.group(1) == "in_proj_weight" else "bias" + return f"octo.head.l1.map.attn_qkv.{leaf}" + m = re.fullmatch(re.escape(p) + r"attention\.out_proj\.(weight|bias)", k) + if m: + return f"octo.head.l1.map.attn_o.{m.group(1)}" + m = re.fullmatch(re.escape(p) + r"layer_norm\.(weight|bias)", k) + if m: + return f"octo.head.l1.map.norm.{m.group(1)}" + m = re.fullmatch(re.escape(p) + r"mlp_block\.dense1\.(weight|bias)", k) + if m: + return f"octo.head.l1.map.ffn_up.{m.group(1)}" + m = re.fullmatch(re.escape(p) + r"mlp_block\.dense2\.(weight|bias)", k) + if m: + return f"octo.head.l1.map.ffn_down.{m.group(1)}" + m = re.fullmatch(r"heads\.action\.mean_proj\.(weight|bias)", k) + if m: + return f"octo.head.l1.mean_proj.{m.group(1)}" + + p = "heads.action.diffusion_model." + if k == p + "time_preprocess.w": + return "octo.head.diffusion.time_fourier.weight" + m = re.fullmatch(re.escape(p) + r"cond_encoder\.layers\.(0|2)\.(weight|bias)", k) + if m: + idx = "0" if m.group(1) == "0" else "1" + return f"octo.head.diffusion.cond.{idx}.{m.group(2)}" + m = re.fullmatch(re.escape(p) + r"reverse_network\.linear1\.(weight|bias)", k) + if m: + return f"octo.head.diffusion.reverse.in.{m.group(1)}" + m = re.fullmatch(re.escape(p) + r"reverse_network\.blocks\.(\d+)\.layer_norm\.(weight|bias)", k) + if m: + return f"octo.head.diffusion.reverse.blk.{m.group(1)}.ln.{m.group(2)}" + m = re.fullmatch(re.escape(p) + r"reverse_network\.blocks\.(\d+)\.linear1\.(weight|bias)", k) + if m: + return f"octo.head.diffusion.reverse.blk.{m.group(1)}.fc1.{m.group(2)}" + m = re.fullmatch(re.escape(p) + r"reverse_network\.blocks\.(\d+)\.linear2\.(weight|bias)", k) + if m: + return f"octo.head.diffusion.reverse.blk.{m.group(1)}.fc2.{m.group(2)}" + m = re.fullmatch(re.escape(p) + r"reverse_network\.linear2\.(weight|bias)", k) + if m: + return f"octo.head.diffusion.reverse.out.{m.group(1)}" + + # T5-base encoder (google-t5/t5-base, frozen; module.*.hf_model.* was skipped at M0). + t5p = "octo_transformer.task_tokenizers.language.hf_model." + if k == t5p + "shared.weight": + return "octo.t5.tok_embd.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.(\d+)\.layer\.0\.layer_norm\.weight", k) + if m: + return f"octo.t5.blk.{m.group(1)}.attn_norm.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.(\d+)\.layer\.0\.SelfAttention\.(q|k|v|o)\.weight", k) + if m: + return f"octo.t5.blk.{m.group(1)}.attn_{m.group(2)}.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.0\.layer\.0\.SelfAttention\.relative_attention_bias\.weight", k) + if m: + return "octo.t5.blk.0.attn_rel_b.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.(\d+)\.layer\.1\.layer_norm\.weight", k) + if m: + return f"octo.t5.blk.{m.group(1)}.ffn_norm.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.(\d+)\.layer\.1\.DenseReluDense\.wi\.weight", k) + if m: + return f"octo.t5.blk.{m.group(1)}.ffn_up.weight" + m = re.fullmatch(re.escape(t5p) + r"encoder\.block\.(\d+)\.layer\.1\.DenseReluDense\.wo\.weight", k) + if m: + return f"octo.t5.blk.{m.group(1)}.ffn_down.weight" + if k == t5p + "encoder.final_layer_norm.weight": + return "octo.t5.output_norm.weight" + + return None + + +def _finetune_config_window_size(finetune_cfg: dict) -> int | None: + if "window_size" in finetune_cfg: + return int(finetune_cfg["window_size"]) + return ( + finetune_cfg.get("dataset_kwargs", {}) + .get("traj_transform_kwargs", {}) + .get("window_size") + ) + + +def _resolve_window_size(model: Any, ckpt_arg: str | None, ckpt_path: str, + override: int | None) -> int: + """window_size actually trained into `model`'s checkpoint -- NOT a fixed constant, + since it differs between the rail-berkeley bridge pretrain (2) and LIBERO + finetunes such as cyrusneary/octo-finetuned-libero (1). + + Priority: --window-size override > finetune_config.json next to the checkpoint + > model.config["finetune_metadata"]["effective_window_size"] > model.config["window_size"] + (all three read from the checkpoint's own config.json / finetune_config.json). + finetune_config.json wins when present: it is the fully-resolved per-run training + recipe (real dataset_dir, real dataset_kwargs_list, real save paths), whereas a + checkpoint's saved config.json can retain the base architecture's window_size (the + pos-embedding weight table's native shape, inherited unchanged from the + octo-small-1.5 pretrain) even when the finetune's data pipeline only ever fed it + fewer timesteps -- confirmed by hand for cyrusneary/octo-finetuned-libero/ + 2025-06-20_..._175739: config.json says window_size=2, but finetune_config.json + (matching every other fact about that run -- 4 LIBERO datasets, primary-only + image_obs_keys, 60000 steps) says window_size=1 both at top level and under + dataset_kwargs.traj_transform_kwargs. + + Native PyTorch checkpoints (OctoModelPt.load_pretrained, e.g. the aloha + jitter2525 open-loop adapt run) have no finetune_config.json file at all, but + carry the same kind of discrepancy inside their own config.json: top-level + window_size=2 (inherited from the octo-small-1.5 pretrain this run was adapted + from) vs. config["finetune_metadata"]["effective_window_size"]=1 (the actual + window size this specific adapt run trained/evaluated with, logged by the + training script's own flags snapshot). effective_window_size, when present, + is therefore preferred over the bare top-level window_size for exactly the same + reason finetune_config.json is preferred over it. + + If neither resolves AND a custom --ckpt was given, fail loudly rather than + silently guessing -- only the unmodified default MODEL_ID (rail-berkeley bridge, + which has no finetune_config.json or finetune_metadata) falls back to + model.config, then to the known bridge constant. + """ + if override is not None: + return override + + if ckpt_arg is not None: + finetune_cfg_path = Path(ckpt_path) / "finetune_config.json" + if finetune_cfg_path.exists(): + finetune_cfg = json.loads(finetune_cfg_path.read_text()) + ws = _finetune_config_window_size(finetune_cfg) + if ws is not None: + return int(ws) + + cfg = getattr(model, "config", None) + if isinstance(cfg, dict): + effective_ws = (cfg.get("finetune_metadata") or {}).get("effective_window_size") + if effective_ws is not None: + return int(effective_ws) + if "window_size" in cfg: + return int(cfg["window_size"]) + + if ckpt_arg is None: + return DEFAULT_BRIDGE_WINDOW_SIZE + + raise SystemExit( + f"cannot determine window_size for checkpoint {ckpt_path!r} from " + "finetune_config.json or model.config; pass --window-size explicitly" + ) + + +def _head_type_from_class(head_class_name: str) -> str: + try: + return HEAD_TYPE_BY_CLASS[head_class_name] + except KeyError: + raise SystemExit( + f"unrecognized action head class {head_class_name!r}; add it to " + "HEAD_TYPE_BY_CLASS with its octo.action.head_type value" + ) + + +def _detect_ckpt_format(ckpt_arg: str | None, step: int | None) -> str: + """Auto-detect checkpoint format for --ckpt-format=auto. + + PyTorch checkpoints saved via OctoModelPt.save_pretrained() lay out + /config.json, /dataset_statistics.json, //weights.pth. + JAX/Orbax checkpoints (the only kind load_pretrained_from_jax reads) never + have a weights.pth. hf:// ids and the default MODEL_ID (rail-berkeley bridge + pretrain) are always jax. + """ + if ckpt_arg is None or ckpt_arg.startswith("hf://"): + return "jax" + ckpt_path = Path(ckpt_arg) + if not ckpt_path.is_dir(): + return "jax" + if step is not None: + return "pytorch" if (ckpt_path / str(step) / "weights.pth").exists() else "jax" + for sub in ckpt_path.iterdir(): + if sub.is_dir() and sub.name.isdigit() and (sub / "weights.pth").exists(): + return "pytorch" + return "jax" + + +def _validate_required(mapped: dict[str, str], head_type: str, has_proprio: bool) -> list[str]: + required: list[str] = [] + for view in ("primary", "wrist"): + for i in range(4): + for leaf in ("weight", "bias"): + required.append(f"octo.obs.{view}.stem.{i}.conv.{leaf}") + required.append(f"octo.obs.{view}.stem.{i}.gn.{leaf}") + for leaf in ("weight", "bias"): + required.append(f"octo.obs.{view}.patch_embd.{leaf}") + required.append(f"octo.obs.{view}.proj.{leaf}") + required.append(f"octo.obs.{view}.pos_embd") + if has_proprio: + # LowdimObsTokenizerPt has no conv stem/patch_embd (it's a BinTokenizer, not an + # image encoder) -- only a projection, a pos embedding, and the bin thresholds. + for leaf in ("weight", "bias"): + required.append(f"octo.obs.proprio.proj.{leaf}") + required.append("octo.obs.proprio.pos_embd") + required.append("octo.obs.proprio.bin_thresholds") + required += ["octo.task.language.proj.weight", "octo.task.language.proj.bias", "octo.task.language.pos_embd", "octo.readout.action.pos_embd"] + for i in range(12): + for stem in ("attn_norm", "attn_qkv", "attn_o", "ffn_norm", "ffn_up", "ffn_down"): + for leaf in ("weight", "bias"): + required.append(f"octo.blk.{i}.{stem}.{leaf}") + if head_type == "diffusion": + required += ["octo.head.diffusion.time_fourier.weight"] + for i in range(2): + for leaf in ("weight", "bias"): + required.append(f"octo.head.diffusion.cond.{i}.{leaf}") + for leaf in ("weight", "bias"): + required.append(f"octo.head.diffusion.reverse.in.{leaf}") + required.append(f"octo.head.diffusion.reverse.out.{leaf}") + for i in range(3): + for sub in ("ln", "fc1", "fc2"): + for leaf in ("weight", "bias"): + required.append(f"octo.head.diffusion.reverse.blk.{i}.{sub}.{leaf}") + elif head_type == "l1": + required.append("octo.head.l1.map.probe") + for leaf in ("weight", "bias"): + required.append(f"octo.head.l1.map.attn_qkv.{leaf}") + required.append(f"octo.head.l1.map.attn_o.{leaf}") + required.append(f"octo.head.l1.map.norm.{leaf}") + required.append(f"octo.head.l1.map.ffn_up.{leaf}") + required.append(f"octo.head.l1.map.ffn_down.{leaf}") + required.append(f"octo.head.l1.mean_proj.{leaf}") + else: + raise SystemExit(f"no required-tensor list for head_type {head_type!r}") + required += ["octo.t5.tok_embd.weight", "octo.t5.blk.0.attn_rel_b.weight", "octo.t5.output_norm.weight"] + for i in range(12): + for stem in ("attn_norm", "attn_q", "attn_k", "attn_v", "attn_o", "ffn_norm", "ffn_up", "ffn_down"): + required.append(f"octo.t5.blk.{i}.{stem}.weight") + have = set(mapped.values()) + return [k for k in required if k not in have] + + +def main() -> int: + ap = argparse.ArgumentParser(description="Convert Octo PyTorch state_dict to F32 GGUF.") + ap.add_argument("--out", type=Path, default=Path("octo-small-1.5-f32.gguf")) + ap.add_argument("--octo-root", type=Path, default=Path(__file__).resolve().parents[1] / "octo-pytorch") + ap.add_argument("--allow-unmapped", action="store_true", help="write known mapped tensors and report unmapped keys instead of failing") + ap.add_argument("--ckpt", type=str, default=None, + help="path or HF id (hf://...) to a checkpoint dir to convert, e.g. an Octo " + "finetune experiment dir with config.json/dataset_statistics.json//. " + f"Default: {MODEL_ID!r} (the rail-berkeley bridge pretrain).") + ap.add_argument("--step", type=int, default=None, + help="checkpoint step to load from --ckpt (default: latest available step). " + "Ignored/invalid when --ckpt is an hf:// id.") + ap.add_argument("--window-size", type=int, default=None, + help="override window_size written to octo.window_size GGUF meta; only needed " + "if it can't be read from the checkpoint's config.json/finetune_config.json.") + ap.add_argument("--ckpt-format", choices=("auto", "jax", "pytorch"), default="auto", + help="checkpoint format to load --ckpt as: 'jax' (Orbax, via " + "OctoModelPt.load_pretrained_from_jax -- the original/default path) or " + "'pytorch' (native, via OctoModelPt.load_pretrained -- config.json + " + "dataset_statistics.json + /weights.pth). 'auto' (default) detects " + "pytorch by the presence of /weights.pth next to --ckpt; the default " + f"{MODEL_ID!r} (no --ckpt) always resolves to 'jax'.") + args = ap.parse_args() + + if args.octo_root.exists(): + sys.path.insert(0, str(args.octo_root)) + + from octo.model.octo_model_pt import OctoModelPt + + model_id = args.ckpt if args.ckpt is not None else MODEL_ID + ckpt_format = args.ckpt_format + if ckpt_format == "auto": + ckpt_format = _detect_ckpt_format(args.ckpt, args.step) + print(f"--ckpt-format auto detected: {ckpt_format}") + + if ckpt_format == "pytorch": + if args.ckpt is None: + raise SystemExit("--ckpt-format pytorch requires --ckpt (a local PyTorch checkpoint dir)") + print(f"loading {model_id} via OctoModelPt.load_pretrained (step={args.step}) ...") + loaded = OctoModelPt.load_pretrained(model_id, step=args.step) + else: + print(f"loading {model_id} via OctoModelPt.load_pretrained_from_jax (step={args.step}) ...") + loaded = OctoModelPt.load_pretrained_from_jax(model_id, step=args.step, skip_keys_regex=".*hf_model") + m = loaded["octo_model"] + sd = m.state_dict() + + window_size = _resolve_window_size(m, args.ckpt, model_id, args.window_size) + print(f"window_size = {window_size} (from " + f"{'--window-size override' if args.window_size is not None else 'checkpoint config'})") + OCTO_META["window_size"] = window_size + + head_cfg = m.config["model"]["heads"]["action"] + head_type = _head_type_from_class(head_cfg["name"]) + OCTO_META["action.head_type"] = head_type + OCTO_META["action.horizon"] = int(head_cfg["kwargs"]["action_horizon"]) + OCTO_META["action.dim"] = int(head_cfg["kwargs"]["action_dim"]) + print(f"action head: {head_cfg['name']} -> head_type={head_type} " + f"horizon={OCTO_META['action.horizon']} dim={OCTO_META['action.dim']}") + + has_proprio = "proprio" in m.config["model"]["observation_tokenizers"] + print(f"has_proprio = {has_proprio} (from checkpoint config observation_tokenizers)") + + print("state_dict keys and shapes:") + for key in sorted(sd): + print(f"{key}\t{tuple(sd[key].shape)}\t{sd[key].dtype}") + + mapped: dict[str, str] = {} + ignored: list[str] = [] + unmapped: list[str] = [] + for key, tensor in sd.items(): + if not tensor.is_floating_point(): + continue + if any(pattern.match(key) for pattern in IGNORED_PATTERNS): + ignored.append(key) + continue + dst = map_key(key) + if dst is None: + unmapped.append(key) + elif dst in mapped.values(): + raise SystemExit(f"duplicate GGUF destination {dst} from {key}") + else: + mapped[key] = dst + + missing = _validate_required(mapped, head_type, has_proprio) + if ignored: + print("IGNORED STATE_DICT KEYS:") + for key in sorted(ignored): + print(f" {key} {tuple(sd[key].shape)}") + + if unmapped or missing: + print("TENSOR MAP REPORT:") + if unmapped: + print("unmapped state_dict keys:") + for key in sorted(unmapped): + print(f" {key} {tuple(sd[key].shape)}") + if missing: + print("missing required GGUF tensors:") + for key in missing: + print(f" {key}") + if unmapped or missing: + if not args.allow_unmapped: + raise SystemExit("Octo tensor map is incomplete; re-run with --allow-unmapped only for investigation") + + args.out.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter(str(args.out), arch=ARCH) + for key, value in OCTO_META.items(): + _add_meta(writer, key, value) + writer.add_string("octo.dataset_statistics", json.dumps(m.dataset_statistics, default=_json_default, sort_keys=True)) + _embed_tokenizer(writer) + + rows = [] + for src, dst in sorted(mapped.items(), key=lambda kv: kv[1]): + tensor = sd[src] + writer.add_tensor(dst, _f32(tensor), raw_dtype=gguf.GGMLQuantizationType.F32) + rows.append({"state_dict": src, "gguf": dst, "shape": list(tensor.shape)}) + print(f"map {src} {tuple(tensor.shape)} -> {dst}") + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + report = args.out.with_suffix(args.out.suffix + ".tensor_map.json") + report.write_text(json.dumps({"mapped": rows, "ignored": ignored, "unmapped": unmapped, "missing_required": missing}, indent=2), encoding="utf-8") + print(f"done: {args.out} ({args.out.stat().st_size / (1024 * 1024):.1f} MiB)") + print(f"tensor map: {report}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/arch.h b/src/arch.h index 45db734..517f4cb 100644 --- a/src/arch.h +++ b/src/arch.h @@ -66,6 +66,7 @@ enum class Arch { GR00T_N1_5, // NVIDIA Isaac GR00T N1.5 (Eagle VLM + DiT action head). GR00T_N1_6, // NVIDIA Isaac GR00T N1.6 (Eagle Block-2A + DiT). GR00T_N1_7, // NVIDIA Isaac GR00T N1.7 (Qwen3 backbone + DiT). + OCTO, // UC Berkeley Octo small 1.5 (SmallStem16 + T5 + diffusion/L1 head). BITVLA, // Microsoft BitVLA (1.58-bit ternary LM/ViT). VLA_ADAPTER,// OpenHelix VLA-Adapter DINOv2 + SigLIP + Bridge-Attention. OPENVLA_OFT,// DINOv2-L/14-reg4 + SigLIP-so400m/14 +Llama-2-7B + MLPResNet. @@ -169,6 +170,17 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, const std::string& config_path, const Options& opts); +/** + * @brief Build an Octo model. Vision, T5 text encoder and tokenizer vocab are + * all baked into @p ckpt_path. Only compiled when VLA_OCTO is on. + * @copydetails smolvla_create + */ +#ifdef VLA_USE_OCTO +std::unique_ptr octo_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string& config_path); +#endif + /** * @brief Build a BitVLA model. Vision is baked into @p ckpt_path. * @copydetails smolvla_create diff --git a/src/model.cpp b/src/model.cpp index f379a9c..5242efb 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -72,6 +72,7 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { try_str("gr00t_n1_5.architecture", arch_str) || try_str("gr00t_n1_6.architecture", arch_str) || try_str("gr00t_n1_7.architecture", arch_str) || + try_str("octo.architecture", arch_str) || try_str("bitvla.architecture", arch_str) || try_str("openvla_oft.architecture", arch_str) || try_str("vla_jepa.architecture", arch_str) || @@ -104,6 +105,12 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { *out = Arch::GR00T_N1_7; ok = true; } +#ifdef VLA_USE_OCTO + else if (arch_str == "octo" || arch_str == "octo-small-1.5") { + *out = Arch::OCTO; + ok = true; + } +#endif else if (arch_str == "bitvla") { *out = Arch::BITVLA; ok = true; @@ -251,6 +258,12 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, std::printf("vla: arch = gr00t_n1_7\n"); impl = gr00t_n1_7_create(mmproj_path, ckpt_path, config_path, opts); break; +#ifdef VLA_USE_OCTO + case Arch::OCTO: + std::printf("vla: arch = octo\n"); + impl = octo_create(mmproj_path, ckpt_path, config_path); + break; +#endif case Arch::BITVLA: std::printf("vla: arch = bitvla\n"); impl = bitvla_create(mmproj_path, ckpt_path, config_path, opts); diff --git a/src/model.h b/src/model.h index dfdb170..a918831 100644 --- a/src/model.h +++ b/src/model.h @@ -145,7 +145,7 @@ struct Inputs { /// (pad @c real_state_dim..max with zeros). const float* noise; ///< Initial noise for the action expert. - /// Optional per-token language attention mask. Only Evo-1 honors this; the + /// Optional per-token language attention mask. Evo-1 and Octo honor it; the /// other architectures derive their own mask and ignore it. const int32_t* attention_mask = nullptr; int attention_mask_n = 0; ///< Length of @ref attention_mask. diff --git a/src/models/octo.cpp b/src/models/octo.cpp new file mode 100644 index 0000000..0bd5bc3 --- /dev/null +++ b/src/models/octo.cpp @@ -0,0 +1,1953 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "arch.h" +#include "backend.h" +#include "gguf_reader.h" +#include "loader.h" +#include "model.h" +#include "models/octo.h" +#include "modules/preprocess.h" +#include "scratch_ctx.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include "nlohmann/json.hpp" +#include "sentencepiece_processor.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vla { +namespace { + +constexpr int kHidden = 384; +constexpr int kPatchEmbed = 512; +constexpr int kTaskTokens = 16; + +// The obs/task position-embedding tables are converted as one max-horizon slab, +// the same size for every checkpoint; window_size indexes into it. +constexpr int kMaxHorizon = 10; + +// LowdimObsTokenizer emits one token per state dimension, and octo-small-1.5's +// proprio is 7-dim in every checkpoint that has one. +constexpr int kProprioTokens = 7; + +// A caller that supplies Inputs::noise is replaying a chunk, so the reverse +// process has to draw the same per-step noise too, not just the same sample. +constexpr uint32_t kReplaySeed = 20260921u; + +struct OctoRuntime { + ggml_backend_t backend = nullptr; + ggml_context * ctx_w = nullptr; + std::unordered_map by_name; + + // One cache per camera: the views differ in side and token count, so they + // cannot share a graph. + struct ObsKey { + int side = -1; + int n_tok = -1; + int steps = -1; + bool operator==(const ObsKey& o) const { + return side == o.side && n_tok == o.n_tok && steps == o.steps; + } + }; + struct ObsIO { + ggml_tensor * input = nullptr; + ggml_tensor * rows = nullptr; + ggml_tensor * pos = nullptr; + }; + graph_cache obs_primary, obs_wrist; + + struct ProprioKey { + int in_dim = -1; + int steps = -1; + bool operator==(const ProprioKey& o) const { + return in_dim == o.in_dim && steps == o.steps; + } + }; + struct ProprioIO { + ggml_tensor * tokens = nullptr; + ggml_tensor * rows = nullptr; + ggml_tensor * pos = nullptr; + }; + graph_cache proprio; + + struct T5Key { + int seq = -1; + bool operator==(const T5Key& o) const { + return seq == o.seq; + } + }; + struct T5IO { + ggml_tensor * ids = nullptr; + ggml_tensor * bucket = nullptr; + ggml_tensor * padmask = nullptr; + ggml_tensor * out = nullptr; + }; + graph_cache t5; + + struct LangKey { + int steps = -1; + bool operator==(const LangKey& o) const { + return steps == o.steps; + } + }; + struct LangIO { + ggml_tensor * in = nullptr; + ggml_tensor * pos = nullptr; + ggml_tensor * repeated = nullptr; + }; + graph_cache language; + + struct BtKey { + int seq = -1; + int n_readout = -1; + bool operator==(const BtKey& o) const { + return seq == o.seq && n_readout == o.n_readout; + } + }; + struct BtIO { + ggml_tensor * input = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * readout_idx = nullptr; + ggml_tensor * out = nullptr; + }; + graph_cache transformer; + + // The whole reverse process lives in one graph. The steps are sequentially + // dependent so they cannot be batched, but with the graph cached a frame + // costs one submission instead of twenty. + struct DiffKey { + int width = -1; + int action = -1; + int steps = -1; + bool operator==(const DiffKey& o) const { + return width == o.width && action == o.action && steps == o.steps; + } + }; + struct DiffIO { + ggml_tensor * obs = nullptr; + ggml_tensor * x0 = nullptr; + ggml_tensor * z = nullptr; + ggml_tensor * times = nullptr; + ggml_tensor * out = nullptr; + }; + graph_cache diffusion; + + struct L1Key { + int width = -1; + int action = -1; + bool operator==(const L1Key& o) const { + return width == o.width && action == o.action; + } + }; + struct L1IO { + ggml_tensor * readout = nullptr; + ggml_tensor * out = nullptr; + }; + graph_cache l1_head; + + // The T5 encoder and the projection after it depend only on the + // instruction, which holds for as long as the robot pursues one task. + std::vector lang_key; + std::vector lang_pos; + std::vector lang_repeated; + int lang_steps = -1; + + struct ActionStats { + std::vector mean, stdv; + std::vector mask; + }; + struct ProprioStats { + std::vector mean, stdv; + }; + std::string stats_key; + bool stats_loaded = false; + ActionStats action_stats; + ProprioStats proprio_stats; + bool has_proprio_stats = false; + + void init(ggml_backend_t b, ggml_context * w) { + backend = b; + ctx_w = w; + // ggml_get_tensor is a linear strcmp scan, and building the stage graphs + // looks up a few hundred weights by name. + by_name.clear(); + for (ggml_tensor * t = ggml_get_first_tensor(w); t; t = ggml_get_next_tensor(w, t)) + by_name.emplace(t->name, t); + } + + ggml_tensor * weight(const char * name) const { + auto it = by_name.find(name); + if (it == by_name.end()) { + std::fprintf(stderr, "vla(octo): missing resident weight %s\n", name); + return nullptr; + } + return it->second; + } + + /// Must run before the backend the caches allocated from is freed. + void reset() { + obs_primary.release(); + obs_wrist.release(); + proprio.release(); + t5.release(); + language.release(); + transformer.release(); + diffusion.release(); + l1_head.release(); + } + + ~OctoRuntime() { + reset(); + } +}; + +// Relabels duplicate node names for ggml-openvino (a no-op elsewhere). +bool octo_compute(OctoRuntime& rt, ggml_cgraph * gf, const char * what) { + graph_unique_names(gf); + const ggml_status st = ggml_backend_graph_compute(rt.backend, gf); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): %s graph compute failed (%d)\n", what, (int) st); + return false; + } + return true; +} + +struct OctoModelArch : public ModelArchBase { + OctoModelArch() : ModelArchBase(Arch::OCTO) {} + ~OctoModelArch() override { + rt.reset(); + if (weight_buf) ggml_backend_buffer_free(weight_buf); + if (ctx_weights) ggml_free(ctx_weights); + if (backend) ggml_backend_free(backend); + } + + std::string gguf_path; + ggml_backend_t backend = nullptr; + ggml_context * ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + int n_threads = default_cpu_threads(); + + int64_t hidden = 384; + int64_t blocks = 12; + int64_t heads = 6; + int64_t ffn = 1536; + int64_t window_size = 2; + int64_t action_horizon = 4; + int64_t action_dim = 7; + int64_t primary_tokens = 256; + int64_t wrist_tokens = 64; + int64_t primary_size = 256; + int64_t wrist_size = 128; + int64_t language_tokens = kTaskTokens; + int64_t diffusion_steps = 20; + float diffusion_s = 0.008f; + float max_action = 5.0f; + int32_t pad_id = 0; + + // "diffusion" for the published checkpoints, "l1" for the fine-tuned ones. + std::string head_type = "diffusion"; + // Proprio has no metadata key of its own; presence of the projection weight + // is what says the checkpoint has a LowdimObsTokenizer. Its in-dim is 1 + // (continuous) or 256 (bin one-hot). + bool has_proprio = false; + int64_t proprio_in_dim = 0; + + gguf_reader io{"octo"}; + OctoRuntime rt; + std::mt19937 rng{std::random_device{}()}; + + std::vector predict(const Inputs& in) override; +}; + +void detect_proprio(const gguf_reader& g, bool& has_proprio, int64_t& proprio_in_dim) { + const ggml_tensor * proj = g.meta("octo.obs.proprio.proj.weight"); + has_proprio = proj != nullptr; + proprio_in_dim = has_proprio ? proj->ne[0] : 0; +} + +// The converter writes a whole-numbered scalar as an integer, so the diffusion +// constants are not all one GGUF type. +float scalar_key(const gguf_reader& g, const char * key) { + const int64_t id = gguf_find_key(g.gctx, key); + if (id >= 0) { + switch (gguf_get_kv_type(g.gctx, id)) { + case GGUF_TYPE_FLOAT32: return gguf_get_val_f32(g.gctx, id); + case GGUF_TYPE_FLOAT64: return (float) gguf_get_val_f64(g.gctx, id); + case GGUF_TYPE_UINT32: return (float) gguf_get_val_u32(g.gctx, id); + case GGUF_TYPE_INT32: return (float) gguf_get_val_i32(g.gctx, id); + default: break; + } + } + std::fprintf(stderr, "vla(octo): %s is not a scalar number\n", key); + return 0.f; +} + +bool require_key(const gguf_reader& g, const char * key) { + if (!g.has(key)) { + std::fprintf(stderr, "vla(octo): missing metadata %s\n", key); + return false; + } + return true; +} + +bool load_config(const gguf_reader& g, OctoModelArch& m) { + const char * keys[] = { + "octo.architecture", + "octo.embedding_length", + "octo.block_count", + "octo.attention.head_count", + "octo.feed_forward_length", + "octo.attention.layer_norm_eps", + "octo.window_size", + "octo.action.horizon", + "octo.action.dim", + "octo.readout.count", + "octo.tokens.primary", + "octo.tokens.wrist", + "octo.tokens.language", + "octo.image.primary_size", + "octo.image.wrist_size", + "octo.diffusion.steps", + "octo.diffusion.beta_schedule", + "octo.diffusion.s", + "octo.diffusion.max_action", + "octo.dataset_statistics", + }; + for (const char * key : keys) { + if (!require_key(g, key)) + return false; + } + if (g.str("octo.architecture") != "octo-small-1.5") { + std::fprintf(stderr, "vla(octo): octo.architecture=%s, expected octo-small-1.5\n", + g.str("octo.architecture").c_str()); + return false; + } + if (g.str("octo.diffusion.beta_schedule") != "cosine") { + std::fprintf(stderr, "vla(octo): octo.diffusion.beta_schedule=%s, only cosine is implemented\n", + g.str("octo.diffusion.beta_schedule").c_str()); + return false; + } + + m.hidden = g.u32("octo.embedding_length"); + m.blocks = g.u32("octo.block_count"); + m.heads = g.u32("octo.attention.head_count"); + m.ffn = g.u32("octo.feed_forward_length"); + m.window_size = g.u32("octo.window_size"); + m.action_horizon = g.u32("octo.action.horizon"); + m.action_dim = g.u32("octo.action.dim"); + m.primary_tokens = g.u32("octo.tokens.primary"); + m.wrist_tokens = g.u32("octo.tokens.wrist"); + m.primary_size = g.u32("octo.image.primary_size"); + m.wrist_size = g.u32("octo.image.wrist_size"); + m.language_tokens = g.u32("octo.tokens.language"); + m.diffusion_steps = g.u32("octo.diffusion.steps"); + m.diffusion_s = scalar_key(g, "octo.diffusion.s"); + m.max_action = scalar_key(g, "octo.diffusion.max_action"); + m.pad_id = g.has("octo.tokenizer.pad_id") ? (int32_t) g.u32("octo.tokenizer.pad_id") : 0; + m.head_type = g.has("octo.action.head_type") ? g.str("octo.action.head_type") : "diffusion"; + detect_proprio(g, m.has_proprio, m.proprio_in_dim); + + if (m.has_proprio && m.proprio_in_dim != 1 && m.proprio_in_dim != 256) { + std::fprintf(stderr, "vla(octo): octo.obs.proprio.proj.weight in-dim=%lld, expected 1 or 256\n", + (long long) m.proprio_in_dim); + return false; + } + if (m.hidden != kHidden || m.blocks != 12 || m.heads != 6 || m.ffn != 1536) { + std::fprintf(stderr, "vla(octo): metadata does not match the octo-small-1.5 backbone\n"); + return false; + } + if (m.diffusion_steps < 1 || m.diffusion_s <= 0.f || m.max_action <= 0.f) { + std::fprintf(stderr, "vla(octo): diffusion steps=%lld s=%g max_action=%g must all be positive\n", + (long long) m.diffusion_steps, (double) m.diffusion_s, (double) m.max_action); + return false; + } + if (m.language_tokens != kTaskTokens) { + std::fprintf(stderr, "vla(octo): octo.tokens.language=%lld, the T5 encoder is built for %d\n", + (long long) m.language_tokens, kTaskTokens); + return false; + } + if (g.u32("octo.readout.count") != 1) { + std::fprintf(stderr, "vla(octo): octo.readout.count=%u, only the single action readout is implemented\n", + g.u32("octo.readout.count")); + return false; + } + // window_size is per-checkpoint (bridge pretrain 2, the LIBERO finetunes 1); + // any value the shared position-embedding slab covers is a legal slice. + if (m.window_size < 1 || m.window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%lld outside [1, %d]\n", + (long long) m.window_size, kMaxHorizon); + return false; + } + + const int64_t proprio_dim = m.has_proprio ? kProprioTokens : 0; + + m.cfg.n_img = m.primary_tokens+m.wrist_tokens; + m.cfg.n_lang = m.language_tokens; + m.cfg.n_state = proprio_dim; + m.cfg.n_prefix = m.language_tokens+m.window_size*(m.primary_tokens+m.wrist_tokens+m.language_tokens); + m.cfg.n_suffix = m.action_horizon; + m.cfg.n_full = m.cfg.n_prefix+m.window_size; + m.cfg.hidden = m.hidden; + m.cfg.expert_h = 256; + m.cfg.intermediate = m.ffn; + m.cfg.expert_inter = 256; + m.cfg.n_q_heads = m.heads; + m.cfg.n_kv_heads = m.heads; + m.cfg.head_dim = m.hidden/m.heads; + m.cfg.q_full_dim = m.hidden; + m.cfg.kv_full_dim = m.hidden; + m.cfg.n_layers = m.blocks; + m.cfg.self_attn_every_n = 1; + m.cfg.max_state_dim = proprio_dim; + m.cfg.max_action_dim = m.action_dim; + m.cfg.real_state_dim = proprio_dim; + m.cfg.real_action_dim = m.action_dim; + m.cfg.norm_eps = g.f32("octo.attention.layer_norm_eps"); + m.cfg.num_steps = (int) m.diffusion_steps; + return true; +} + +// SmallStem's StdConv standardizes its kernel per output channel on every +// forward pass. The kernel is frozen at inference, so the result is too: bake it +// in once instead of paying a round trip per camera per frame. +bool is_stem_conv_weight(const char * name) { + return std::strstr(name, "octo.obs.") == name && + std::strstr(name, ".stem.") != nullptr && + std::strstr(name, ".conv.weight") != nullptr; +} + +void standardize_conv_weight(float * w, int64_t oc, int64_t n) { + for (int64_t o=0; o row; + for (ggml_tensor * t = ggml_get_first_tensor(m.ctx_weights); t; t = ggml_get_next_tensor(m.ctx_weights, t)) { + if (!is_stem_conv_weight(ggml_get_name(t))) + continue; + row.resize((size_t) ggml_nelements(t)); + ggml_backend_tensor_get(t, row.data(), 0, ggml_nbytes(t)); + // ggml ne = [kw, kh, in, out]: one contiguous block per output channel. + standardize_conv_weight(row.data(), t->ne[3], t->ne[0]*t->ne[1]*t->ne[2]); + ggml_backend_tensor_set(t, row.data(), 0, ggml_nbytes(t)); + } + return true; +} + +// Host copy of a resident weight, for the few that a call transforms before they +// become graph operands. Goes through the backend: the buffer may be on device. +std::vector tensor_to_vec(const ggml_tensor * t) { + std::vector out((size_t) ggml_nelements(t)); + ggml_backend_tensor_get(t, out.data(), 0, ggml_nbytes(t)); + return out; +} + +// SmallStem16 for one camera view: four standardized-conv + GroupNorm + ReLU +// stages at stride 2, a 1x1 patch embedding, a projection to the model width, +// and the per-timestep position embedding. +// +// `obs` holds the normalized CHW frames for `steps`, `task` the single goal +// frame, which is concatenated onto every one of them as channels 3..5. +bool run_obs_tokenizer_graph(OctoRuntime& rt, + graph_cache& cache, + const char * view, + const std::vector& obs, + const std::vector& task, + int side, + int n_tok, + int steps, + const std::vector& pos_rows, + std::vector& pos) { + const size_t frame = (size_t) 3*side*side; + if (obs.size() != frame*(size_t) steps || task.size() != frame) { + std::fprintf(stderr, "vla(octo): unexpected input image shape for side=%d\n", side); + return false; + } + if ((int) pos_rows.size() != steps) + return false; + pos.resize((size_t) steps*n_tok*kHidden); + + const OctoRuntime::ObsKey key{side, n_tok, steps}; + const bool built = cache.ensure(rt.backend, key, (size_t) 32*1024*1024, + [&](ggml_context * C, OctoRuntime::ObsIO& io) -> ggml_cgraph * { + ggml_tensor * x = ggml_new_tensor_4d(C, GGML_TYPE_F32, side, side, 6, steps); + ggml_set_name(x, "octo.obs.input_norm"); + ggml_set_input(x); + io.input = x; + + char rname[160]; + auto weight = [&](const char * suffix, int li) -> ggml_tensor * { + if (li < 0) std::snprintf(rname, sizeof(rname), "octo.obs.%s.%s", view, suffix); + else std::snprintf(rname, sizeof(rname), "octo.obs.%s.stem.%d.%s", view, li, suffix); + return rt.weight(rname); + }; + // Biases and scales arrive as [ch] against a [w,h,ch,steps] feature map; + // the reshape is a view, and ggml_add/ggml_mul broadcast from there. + auto per_channel = [&](ggml_tensor * t) { + return t ? ggml_reshape_4d(C, t, 1, 1, t->ne[0], 1) : nullptr; + }; + + for (int li=0; li<4; ++li) { + ggml_tensor * cw = weight("conv.weight", li); + ggml_tensor * cb = per_channel(weight("conv.bias", li)); + ggml_tensor * gw = per_channel(weight("gn.weight", li)); + ggml_tensor * gb = per_channel(weight("gn.bias", li)); + if (!cw || !cb || !gw || !gb) + return nullptr; + + x = ggml_conv_2d(C, cw, x, 2, 2, 1, 1, 1, 1); + x = ggml_add(C, x, cb); + x = ggml_group_norm(C, x, 32, 1e-5f); + x = ggml_add(C, ggml_mul(C, x, gw), gb); + x = ggml_relu(C, x); + } + + ggml_tensor * pw = weight("patch_embd.weight", -1); + ggml_tensor * pb = per_channel(weight("patch_embd.bias", -1)); + ggml_tensor * jw = weight("proj.weight", -1); + ggml_tensor * jb = weight("proj.bias", -1); + ggml_tensor * pos_r = weight("pos_embd", -1); + if (!pw || !pb || !jw || !jb || !pos_r) + return nullptr; + + ggml_tensor * patch = ggml_add(C, ggml_conv_2d(C, pw, x, 1, 1, 0, 0, 1, 1), pb); + ggml_tensor * tok = ggml_cont(C, ggml_reshape_3d(C, + ggml_cont(C, ggml_permute(C, patch, 1, 2, 0, 3)), kPatchEmbed, n_tok, steps)); + + // The timesteps in the sequence need not be the leading ones, so the + // rows are gathered rather than sliced. See OctoSeqLayout. + ggml_tensor * rows = ggml_new_tensor_1d(C, GGML_TYPE_I32, steps); + ggml_set_name(rows, "octo.obs.pos_embd.rows"); + ggml_set_input(rows); + io.rows = rows; + ggml_tensor * pe = ggml_reshape_3d(C, + ggml_get_rows(C, ggml_reshape_2d(C, pos_r, kHidden*n_tok, pos_r->ne[2]), rows), + kHidden, n_tok, steps); + + ggml_tensor * out = ggml_add(C, ggml_add(C, ggml_mul_mat(C, jw, tok), jb), pe); + ggml_set_name(out, "obs.tokenizer.pos"); + ggml_set_output(out); + io.pos = out; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_build_forward_expand(gf, out); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): obs tokenizer graph build failed\n"); + return false; + } + + OctoRuntime::ObsIO& io = cache.io(); + std::vector input((size_t) steps*2*frame); + for (int t=0; t& proprio_norm, + int in_dim, + int steps, + const std::vector& pos_rows, + std::vector& pos_out) { + constexpr int n_dims = kProprioTokens; + if (proprio_norm.size() != (size_t) steps*n_dims) + return false; + if ((int) pos_rows.size() != steps) + return false; + if (in_dim != 1 && in_dim != 256) { + std::fprintf(stderr, "vla(octo): proprio tokenizer in_dim=%d, expected 1 or 256\n", in_dim); + return false; + } + pos_out.resize((size_t) steps*n_dims*kHidden); + + std::vector tokens_in((size_t) in_dim*n_dims*steps, 0.0f); + if (in_dim == 1) { + for (size_t i=0; i thresholds = tensor_to_vec(thresholds_r); + const int n_thresh = (int) thresholds.size(); + for (int t=0; t ggml_cgraph * { + ggml_tensor * proj_w = rt.weight("octo.obs.proprio.proj.weight"); + ggml_tensor * proj_b = rt.weight("octo.obs.proprio.proj.bias"); + ggml_tensor * pos_r = rt.weight("octo.obs.proprio.pos_embd"); + if (!proj_w || !proj_b || !pos_r) + return nullptr; + if (pos_r->ne[2] < steps) { + std::fprintf(stderr, "vla(octo): octo.obs.proprio.pos_embd too small for %d timesteps\n", steps); + return nullptr; + } + + ggml_tensor * x = ggml_new_tensor_3d(C, GGML_TYPE_F32, in_dim, n_dims, steps); + ggml_set_name(x, "octo.obs.proprio.tokens_in"); + ggml_set_input(x); + io.tokens = x; + + ggml_tensor * rows = ggml_new_tensor_1d(C, GGML_TYPE_I32, steps); + ggml_set_name(rows, "octo.obs.proprio.pos_embd.rows"); + ggml_set_input(rows); + io.rows = rows; + + ggml_tensor * pe = ggml_reshape_3d(C, + ggml_get_rows(C, ggml_reshape_2d(C, pos_r, kHidden*n_dims, pos_r->ne[2]), rows), + kHidden, n_dims, steps); + ggml_tensor * out = ggml_add(C, ggml_add(C, ggml_mul_mat(C, proj_w, x), proj_b), pe); + ggml_set_name(out, "obs.proprio.pos"); + ggml_set_output(out); + io.pos = out; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 256, false); + ggml_build_forward_expand(gf, out); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): proprio tokenizer graph build failed\n"); + return false; + } + + OctoRuntime::ProprioIO& io = rt.proprio.io(); + ggml_backend_tensor_set(io.tokens, tokens_in.data(), 0, ggml_nbytes(io.tokens)); + ggml_backend_tensor_set(io.rows, pos_rows.data(), 0, ggml_nbytes(io.rows)); + if (!octo_compute(rt, rt.proprio.graph(), "proprio tokenizer")) + return false; + ggml_backend_tensor_get(io.pos, pos_out.data(), 0, ggml_nbytes(io.pos)); + return true; +} + +// Language projection and position embedding, then repeat_task_tokens: the same +// 16 projected tokens are duplicated into every observation timestep. +bool run_language_graph(OctoRuntime& rt, + const std::vector& t5, + int steps, + std::vector& pos, + std::vector& repeated) { + if (t5.size() != (size_t) kTaskTokens*768) { + std::fprintf(stderr, "vla(octo): expected a T5 output of %dx768\n", kTaskTokens); + return false; + } + pos.resize((size_t) kTaskTokens*kHidden); + repeated.resize((size_t) steps*kTaskTokens*kHidden); + + const bool built = rt.language.ensure(rt.backend, OctoRuntime::LangKey{steps}, (size_t) 4*1024*1024, + [&](ggml_context * C, OctoRuntime::LangIO& io) -> ggml_cgraph * { + ggml_tensor * jw = rt.weight("octo.task.language.proj.weight"); + ggml_tensor * jb = rt.weight("octo.task.language.proj.bias"); + ggml_tensor * pe = rt.weight("octo.task.language.pos_embd"); + if (!jw || !jb || !pe) + return nullptr; + + ggml_tensor * in = ggml_new_tensor_3d(C, GGML_TYPE_F32, 768, kTaskTokens, 1); + ggml_set_name(in, "octo.task.language.t5_inject"); + ggml_set_input(in); + io.in = in; + + ggml_tensor * pos_t = ggml_add(C, ggml_add(C, ggml_mul_mat(C, jw, in), jb), pe); + ggml_set_name(pos_t, "task_language.pos"); + ggml_set_output(pos_t); + io.pos = pos_t; + + ggml_tensor * rep = ggml_repeat_4d(C, pos_t, kHidden, kTaskTokens, steps, 1); + ggml_set_name(rep, "obs_task_language.repeated"); + ggml_set_output(rep); + io.repeated = rep; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 1024, false); + ggml_build_forward_expand(gf, rep); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): language graph build failed\n"); + return false; + } + + OctoRuntime::LangIO& io = rt.language.io(); + ggml_backend_tensor_set(io.in, t5.data(), 0, ggml_nbytes(io.in)); + if (!octo_compute(rt, rt.language.graph(), "language")) + return false; + ggml_backend_tensor_get(io.pos, pos.data(), 0, ggml_nbytes(io.pos)); + ggml_backend_tensor_get(io.repeated, repeated.data(), 0, ggml_nbytes(io.repeated)); + return true; +} + +// relative_position = key_pos - query_pos, bidirectional, 32 buckets, max +// distance 128 -- HF T5Attention._relative_position_bucket. +int32_t t5_relative_position_bucket(int32_t query_pos, int32_t key_pos, int32_t n_buckets, int32_t max_distance) { + const int32_t nb = n_buckets/2; + const int32_t relative_position = key_pos-query_pos; + const int32_t rp = std::abs(relative_position); + const int32_t max_exact = nb/2; + + int32_t bucket = relative_position > 0 ? nb : 0; + if (rp < max_exact) { + bucket += rp; + } else { + const float v = (float) max_exact+std::log((float) rp/(float) max_exact)/ + std::log((float) max_distance/(float) max_exact)*(float) (nb-max_exact); + bucket += std::min((int32_t) std::floor(v), nb-1); + } + return bucket; +} + +// T5-base encoder over the instruction: 12 pre-norm blocks, bidirectional +// attention with the relative-position bias shared from block 0. The bucket +// table and pad mask are host-computed, so the graph depends only on the length. +bool run_t5_encoder_graph(OctoRuntime& rt, + const std::vector& input_ids, + const std::vector& attention_mask, + std::vector& t5_out) { + constexpr int hidden = 768; + constexpr int heads = 12; + constexpr int head_dim = 64; + constexpr int seq = kTaskTokens; + constexpr int n_buckets = 32; + constexpr int max_distance = 128; + constexpr float ln_eps = 1e-6f; + if (input_ids.size() != seq || attention_mask.size() != seq) { + std::fprintf(stderr, "vla(octo): T5 encoder expected %d input_ids/attention_mask\n", seq); + return false; + } + + std::vector bucket_idx((size_t) seq*seq); + std::vector padmask((size_t) seq*seq); + for (int q=0; q ggml_cgraph * { + ggml_tensor * tok_embd = rt.weight("octo.t5.tok_embd.weight"); + ggml_tensor * rel_b = rt.weight("octo.t5.blk.0.attn_rel_b.weight"); + ggml_tensor * outw = rt.weight("octo.t5.output_norm.weight"); + if (!tok_embd || !rel_b || !outw) + return nullptr; + + char rname[160]; + ggml_tensor * blk_w[12][8]; + const char * leaves[8] = {"attn_norm.weight", "attn_q.weight", "attn_k.weight", "attn_v.weight", + "attn_o.weight", "ffn_norm.weight", "ffn_up.weight", "ffn_down.weight"}; + for (int i=0; i<12; ++i) { + for (int j=0; j<8; ++j) { + std::snprintf(rname, sizeof(rname), "octo.t5.blk.%d.%s", i, leaves[j]); + blk_w[i][j] = rt.weight(rname); + if (!blk_w[i][j]) + return nullptr; + } + } + + ggml_tensor * ids = ggml_new_tensor_1d(C, GGML_TYPE_I32, seq); + ggml_set_name(ids, "octo.t5.input_ids"); + ggml_set_input(ids); + io.ids = ids; + + ggml_tensor * bucket = ggml_new_tensor_2d(C, GGML_TYPE_I32, seq, seq); + ggml_set_name(bucket, "octo.t5.pos_bucket"); + ggml_set_input(bucket); + io.bucket = bucket; + + ggml_tensor * padmask_t = ggml_new_tensor_2d(C, GGML_TYPE_F32, seq, seq); + ggml_set_name(padmask_t, "octo.t5.padmask"); + ggml_set_input(padmask_t); + io.padmask = padmask_t; + + ggml_tensor * x = ggml_get_rows(C, tok_embd, ids); + ggml_set_name(x, "octo.t5.input_embed"); + + ggml_tensor * pos_bias = ggml_get_rows(C, rel_b, ggml_reshape_1d(C, bucket, (int64_t) seq*seq)); + pos_bias = ggml_reshape_3d(C, pos_bias, heads, seq, seq); + pos_bias = ggml_cont(C, ggml_permute(C, pos_bias, 2, 0, 1, 3)); + ggml_tensor * mask = ggml_add(C, pos_bias, padmask_t); + + for (int i=0; i<12; ++i) { + ggml_tensor * n1 = ggml_mul(C, ggml_rms_norm(C, x, ln_eps), blk_w[i][0]); + ggml_tensor * Q = ggml_mul_mat(C, blk_w[i][1], n1); + ggml_tensor * K = ggml_mul_mat(C, blk_w[i][2], n1); + ggml_tensor * V = ggml_mul_mat(C, blk_w[i][3], n1); + ggml_tensor * Qh = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, Q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * Kh = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, K, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * Vh = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, V, head_dim, heads, seq), 1, 2, 0, 3)); + + ggml_tensor * scores = ggml_mul_mat(C, Kh, Qh); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + // T5 folds 1/sqrt(d_k) into the weights, so the scale here is 1. + ggml_tensor * probs = ggml_soft_max_ext(C, scores, mask, 1.0f, 0.0f); + ggml_tensor * attended = ggml_mul_mat(C, Vh, probs); + ggml_tensor * merged = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, attended, 0, 2, 1, 3)), hidden, seq); + x = ggml_add(C, x, ggml_mul_mat(C, blk_w[i][4], merged)); + + ggml_tensor * n2 = ggml_mul(C, ggml_rms_norm(C, x, ln_eps), blk_w[i][5]); + ggml_tensor * h = ggml_relu(C, ggml_mul_mat(C, blk_w[i][6], n2)); + x = ggml_add(C, x, ggml_mul_mat(C, blk_w[i][7], h)); + } + + ggml_tensor * out = ggml_mul(C, ggml_rms_norm(C, x, ln_eps), outw); + ggml_set_name(out, "t5.out"); + ggml_set_output(out); + io.out = out; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 4096, false); + ggml_build_forward_expand(gf, out); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): T5 encoder graph build failed\n"); + return false; + } + + OctoRuntime::T5IO& io = rt.t5.io(); + ggml_backend_tensor_set(io.ids, input_ids.data(), 0, ggml_nbytes(io.ids)); + ggml_backend_tensor_set(io.bucket, bucket_idx.data(), 0, ggml_nbytes(io.bucket)); + ggml_backend_tensor_set(io.padmask, padmask.data(), 0, ggml_nbytes(io.padmask)); + if (!octo_compute(rt, rt.t5.graph(), "T5 encoder")) + return false; + t5_out.resize((size_t) hidden*seq); + ggml_backend_tensor_get(io.out, t5_out.data(), 0, ggml_nbytes(io.out)); + return true; +} + +// Which tokenizer a run of sequence tokens came from. TASK is the once-only +// language prefix; the repeated task tokens count as observation, matching +// repeat_task_tokens. +enum class OctoGroup { TASK, PRIMARY, WRIST, PROPRIO, LANGUAGE, READOUT }; + +struct OctoSeqRun { + OctoGroup group; + int timestep; ///< -1 for the task prefix. + int n_tokens; + int src_step; ///< Index of this timestep inside that group's own buffer. + bool key_valid; ///< False => masked out as a key for every query. +}; + +// The sequence the block transformer actually sees. +// +// Observation tokens belonging to a padded timestep are left out rather than +// emitted and then masked. They cannot affect the result: the pad mask makes +// them invalid keys for every query, and only the readout rows are read back. +// Dropping them shortens a cold-start window_size=2 sequence from 690 tokens to +// 370 and lets the conv stem run over the frames that are actually live. +struct OctoSeqLayout { + std::vector runs; + int seq = 0; + std::vector primary_steps; ///< Original timestep of each emitted batch entry. + std::vector wrist_steps; + std::vector proprio_steps; + std::vector readout_seq_idx; ///< Sequence position of each timestep's readout token. +}; + +// Per-timestep order is primary -> wrist -> proprio -> repeated language -> +// readout, with the task prefix once at the front. +OctoSeqLayout build_seq_layout(bool task_valid, + const std::vector& primary_valid, + const std::vector& wrist_valid, + const std::vector& timestep_valid, + int window_size, + int n_primary, + int n_wrist, + int n_proprio) { + OctoSeqLayout L; + auto emit = [&](OctoGroup g, int t, int n, int src, bool valid) { + if (n <= 0) + return; + L.runs.push_back({g, t, n, src, valid}); + L.seq += n; + }; + + emit(OctoGroup::TASK, -1, kTaskTokens, 0, task_valid); + for (int t=0; t 0) { + emit(OctoGroup::PROPRIO, t, n_proprio, (int) L.proprio_steps.size(), true); + L.proprio_steps.push_back(t); + } + // The repeated task tokens and the readout query stay in for every + // timestep: both are valid keys for later ones regardless of the mask. + emit(OctoGroup::LANGUAGE, t, kTaskTokens, t, task_valid); + L.readout_seq_idx.push_back(L.seq); + emit(OctoGroup::READOUT, t, 1, t, true); + } + return L; +} + +// Each group's buffer is indexed by the run's src_step: the compacted batch +// index for the tokenized groups, the original timestep for the rest. +bool assemble_transformer_input(const OctoSeqLayout& layout, + const std::vector& task_language, + const std::vector& obs_primary, + const std::vector& obs_wrist, + const std::vector& obs_proprio, + const std::vector& repeated_language, + const std::vector& readout_pos, + std::vector& input) { + input.assign((size_t) layout.seq*kHidden, 0.0f); + + size_t dst = 0; + for (const OctoSeqRun& r : layout.runs) { + const std::vector * src = nullptr; + switch (r.group) { + case OctoGroup::TASK: src = &task_language; break; + case OctoGroup::PRIMARY: src = &obs_primary; break; + case OctoGroup::WRIST: src = &obs_wrist; break; + case OctoGroup::PROPRIO: src = &obs_proprio; break; + case OctoGroup::LANGUAGE: src = &repeated_language; break; + case OctoGroup::READOUT: src = &readout_pos; break; + } + const size_t n = (size_t) r.n_tokens*kHidden; + const size_t off = (size_t) r.src_step*n; + if (src->size() < off+n) { + std::fprintf(stderr, "vla(octo): invalid tensor size while assembling block transformer input\n"); + return false; + } + std::copy_n(src->begin()+(ptrdiff_t) off, n, input.begin()+(ptrdiff_t) dst); + dst += n; + } + return true; +} + +// Additive mask, 0 where attention is allowed and -FLT_MAX where it is blocked. +// A task token sees only task tokens; an observation token sees the task prefix +// plus every observation token at its own timestep or earlier; a readout token +// sees those plus the readouts up to its own timestep. A key the pad mask says +// is not real is blocked for everyone. Shape is [seq,seq] with no head axis -- +// ggml_soft_max_ext broadcasts a mask whose ne2 is 1 over all heads. +void build_transformer_mask(const OctoSeqLayout& layout, std::vector& mask) { + const int seq = layout.seq; + + std::vector group((size_t) seq); + std::vector timestep((size_t) seq); + std::vector key_valid((size_t) seq); + int i = 0; + for (const OctoSeqRun& r : layout.runs) { + for (int k=0; k& input, + const std::vector& blocked_mask, + std::vector& readout_action) { + constexpr int heads = 6; + constexpr int head_dim = 64; + constexpr float ln_eps = 1e-6f; + constexpr float attn_scale = 0.125f; + const int seq = layout.seq; + const int n_readout = (int) layout.readout_seq_idx.size(); + if (input.size() != (size_t) kHidden*seq || blocked_mask.size() != (size_t) seq*seq) + return false; + + const OctoRuntime::BtKey key{seq, n_readout}; + const bool built = rt.transformer.ensure(rt.backend, key, (size_t) 32*1024*1024, + [&](ggml_context * C, OctoRuntime::BtIO& io) -> ggml_cgraph * { + char rname[160]; + ggml_tensor * blk_w[12][12]; + const char * leaves[12] = {"attn_norm.weight", "attn_norm.bias", "attn_qkv.weight", "attn_qkv.bias", + "attn_o.weight", "attn_o.bias", "ffn_norm.weight", "ffn_norm.bias", + "ffn_up.weight", "ffn_up.bias", "ffn_down.weight", "ffn_down.bias"}; + for (int i=0; i<12; ++i) { + for (int j=0; j<12; ++j) { + std::snprintf(rname, sizeof(rname), "octo.blk.%d.%s", i, leaves[j]); + blk_w[i][j] = rt.weight(rname); + if (!blk_w[i][j]) + return nullptr; + } + } + ggml_tensor * out_w = rt.weight("octo.output_norm.weight"); + ggml_tensor * out_b = rt.weight("octo.output_norm.bias"); + if (!out_w || !out_b) + return nullptr; + + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, kHidden, seq); + ggml_set_name(x, "octo.block_transformer.input"); + ggml_set_input(x); + io.input = x; + + ggml_tensor * mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, seq, seq); + ggml_set_name(mask, "octo.block_transformer.additive_mask"); + ggml_set_input(mask); + io.mask = mask; + + ggml_tensor * readout_idx = ggml_new_tensor_1d(C, GGML_TYPE_I32, n_readout); + ggml_set_name(readout_idx, "octo.block_transformer.readout_idx"); + ggml_set_input(readout_idx); + io.readout_idx = readout_idx; + + for (int i=0; i<12; ++i) { + ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), blk_w[i][0]), blk_w[i][1]); + ggml_tensor * qkv = ggml_add(C, ggml_mul_mat(C, blk_w[i][2], n1), blk_w[i][3]); + ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, kHidden, seq, qkv->nb[1], 0)); + ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, kHidden, seq, qkv->nb[1], (size_t) kHidden*qkv->nb[0])); + ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, kHidden, seq, qkv->nb[1], (size_t) 2*kHidden*qkv->nb[0])); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, heads, seq), 1, 2, 0, 3)); + + ggml_tensor * scores = ggml_mul_mat(C, K, Q); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + ggml_tensor * probs = ggml_soft_max_ext(C, scores, mask, attn_scale, 0.0f); + ggml_tensor * attended = ggml_mul_mat(C, V, probs); + ggml_tensor * merged = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, attended, 0, 2, 1, 3)), kHidden, seq); + ggml_tensor * attn_out = ggml_add(C, ggml_mul_mat(C, blk_w[i][4], merged), blk_w[i][5]); + ggml_tensor * residual = ggml_add(C, x, attn_out); + + ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, residual, ln_eps), blk_w[i][6]), blk_w[i][7]); + ggml_tensor * mlp = ggml_add(C, ggml_mul_mat(C, blk_w[i][8], n2), blk_w[i][9]); + mlp = ggml_gelu_erf(C, mlp); + mlp = ggml_add(C, ggml_mul_mat(C, blk_w[i][10], mlp), blk_w[i][11]); + x = ggml_add(C, residual, mlp); + } + + ggml_tensor * output = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), out_w), out_b); + // The readouts are not evenly spaced once padded timesteps drop their + // observation groups, so they are gathered rather than strided. + ggml_tensor * readout = ggml_get_rows(C, output, readout_idx); + ggml_set_name(readout, "bt.readout_action"); + ggml_set_output(readout); + io.out = readout; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_build_forward_expand(gf, readout); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): block transformer graph build failed\n"); + return false; + } + + OctoRuntime::BtIO& io = rt.transformer.io(); + ggml_backend_tensor_set(io.input, input.data(), 0, ggml_nbytes(io.input)); + ggml_backend_tensor_set(io.mask, blocked_mask.data(), 0, ggml_nbytes(io.mask)); + ggml_backend_tensor_set(io.readout_idx, layout.readout_seq_idx.data(), 0, ggml_nbytes(io.readout_idx)); + if (!octo_compute(rt, rt.transformer.graph(), "block transformer")) + return false; + readout_action.resize((size_t) kHidden*n_readout); + ggml_backend_tensor_get(io.out, readout_action.data(), 0, ggml_nbytes(io.out)); + return true; +} + +struct OctoDiffusion { + int steps = 20; + float s = 0.008f; + float max_action = 5.0f; +}; + +struct OctoDiffusionSchedule { + std::vector betas; + std::vector alphas; + std::vector alpha_hats; +}; + +OctoDiffusionSchedule make_cosine_schedule(int steps, float s) { + constexpr double pi = 3.141592653589793238462643383279502884; + const double ds = (double) s; + + std::vector alpha_cum((size_t) steps+1); + for (int i=0; i<=steps; ++i) { + const double t = (double) i/(double) steps; + const double v = std::cos((t+ds)/(1.0+ds)*pi*0.5); + alpha_cum[(size_t) i] = v*v; + } + + OctoDiffusionSchedule sched; + sched.betas.resize((size_t) steps); + sched.alphas.resize((size_t) steps); + sched.alpha_hats.resize((size_t) steps); + + const double first = alpha_cum[0]; + float cum = 1.0f; + for (int i=0; i& readout_action, + int window_size, + int action_total, + const OctoDiffusion& diff, + std::vector& final_actions) { + const int width = window_size; + const int action = action_total; + const int steps = diff.steps; + if (readout_action.size() != (size_t) kHidden*width || width < 1 || action < 1 || steps < 1) + return false; + + const OctoDiffusionSchedule sched = make_cosine_schedule(steps, diff.s); + const size_t chunk = (size_t) width*action; + + // Drawn up front in the order and count a step-at-a-time loop would use, so + // one RNG state yields one trajectory. The last step adds no noise. + std::normal_distribution normal(0.0f, 1.0f); + std::vector initial_noise(chunk); + if (noise) { + for (int i=0; i z(chunk*(size_t) steps, 0.0f); + std::vector times((size_t) width*steps); + for (int step=0; step 0) { + for (size_t i=0; i ggml_cgraph * { + OctoScoreActorWeights w{}; + if (!resolve_score_actor_weights(rt, w)) + return nullptr; + if (w.routw->ne[1] != action) { + std::fprintf(stderr, "vla(octo): diffusion head emits %lld values, not horizon*dim=%d\n", + (long long) w.routw->ne[1], action); + return nullptr; + } + + ggml_tensor * obs = ggml_new_tensor_2d(C, GGML_TYPE_F32, kHidden, width); + ggml_set_name(obs, "action_head.readout_embedding"); + ggml_set_input(obs); + io.obs = obs; + + ggml_tensor * x0 = ggml_new_tensor_2d(C, GGML_TYPE_F32, action, width); + ggml_set_name(x0, "action_head.initial_noise"); + ggml_set_input(x0); + io.x0 = x0; + + ggml_tensor * z_all = ggml_new_tensor_3d(C, GGML_TYPE_F32, action, width, steps); + ggml_set_name(z_all, "action_head.step_noise"); + ggml_set_input(z_all); + io.z = z_all; + + ggml_tensor * times_all = ggml_new_tensor_3d(C, GGML_TYPE_F32, 1, width, steps); + ggml_set_name(times_all, "action_head.time"); + ggml_set_input(times_all); + io.times = times_all; + + ggml_tensor * x = x0; + for (int step=0; stepnb[1], + (size_t) step*times_all->nb[2]); + ggml_tensor * eps = build_score_actor(C, w, time, obs, x); + ggml_tensor * y = ggml_scale(C, ggml_add(C, x, ggml_scale(C, eps, -alpha_2)), alpha_1); + if (time_value > 0) { + ggml_tensor * zs = ggml_view_2d(C, z_all, action, width, z_all->nb[1], + (size_t) step*z_all->nb[2]); + y = ggml_add(C, y, ggml_scale(C, zs, std::sqrt(beta))); + } + x = ggml_clamp(C, y, -diff.max_action, diff.max_action); + } + + // sample_actions returns the last window timestep's chunk. + ggml_tensor * out = ggml_cont(C, ggml_view_1d(C, x, action, (size_t) (width-1)*x->nb[1])); + ggml_set_name(out, "action_head.final_actions"); + ggml_set_output(out); + io.out = out; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 4096, false); + ggml_build_forward_expand(gf, out); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): diffusion graph build failed\n"); + return false; + } + + OctoRuntime::DiffIO& io = rt.diffusion.io(); + ggml_backend_tensor_set(io.obs, readout_action.data(), 0, ggml_nbytes(io.obs)); + ggml_backend_tensor_set(io.x0, initial_noise.data(), 0, ggml_nbytes(io.x0)); + ggml_backend_tensor_set(io.z, z.data(), 0, ggml_nbytes(io.z)); + ggml_backend_tensor_set(io.times, times.data(), 0, ggml_nbytes(io.times)); + if (!octo_compute(rt, rt.diffusion.graph(), "diffusion")) + return false; + final_actions.resize((size_t) action); + ggml_backend_tensor_get(io.out, final_actions.data(), 0, ggml_nbytes(io.out)); + return true; +} + +// ContinuousActionHead: a MAP head over the readout rows, then the mean +// projection. Replaces the diffusion head entirely when head_type is "l1". +bool run_l1_action_head_graph(OctoRuntime& rt, + const std::vector& readout_action, + int window_size, + int action_total, + float max_action, + std::vector& final_actions) { + constexpr int map_heads = 8; // 8, not the block transformer's 6. + constexpr int map_head_dim = kHidden/map_heads; + constexpr float ln_eps = 1e-6f; + const int width = window_size; + if (readout_action.size() != (size_t) kHidden*width || action_total <= 0) + return false; + + const OctoRuntime::L1Key key{width, action_total}; + const bool built = rt.l1_head.ensure(rt.backend, key, (size_t) 8*1024*1024, + [&](ggml_context * C, OctoRuntime::L1IO& io) -> ggml_cgraph * { + ggml_tensor * probe = rt.weight("octo.head.l1.map.probe"); + ggml_tensor * qkv_w = rt.weight("octo.head.l1.map.attn_qkv.weight"); + ggml_tensor * qkv_b = rt.weight("octo.head.l1.map.attn_qkv.bias"); + ggml_tensor * o_w = rt.weight("octo.head.l1.map.attn_o.weight"); + ggml_tensor * o_b = rt.weight("octo.head.l1.map.attn_o.bias"); + ggml_tensor * norm_w = rt.weight("octo.head.l1.map.norm.weight"); + ggml_tensor * norm_b = rt.weight("octo.head.l1.map.norm.bias"); + ggml_tensor * ffn_up_w = rt.weight("octo.head.l1.map.ffn_up.weight"); + ggml_tensor * ffn_up_b = rt.weight("octo.head.l1.map.ffn_up.bias"); + ggml_tensor * ffn_down_w = rt.weight("octo.head.l1.map.ffn_down.weight"); + ggml_tensor * ffn_down_b = rt.weight("octo.head.l1.map.ffn_down.bias"); + ggml_tensor * mean_w = rt.weight("octo.head.l1.mean_proj.weight"); + ggml_tensor * mean_b = rt.weight("octo.head.l1.mean_proj.bias"); + if (!probe || !qkv_w || !qkv_b || !o_w || !o_b || !norm_w || !norm_b || + !ffn_up_w || !ffn_up_b || !ffn_down_w || !ffn_down_b || !mean_w || !mean_b) + return nullptr; + + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, kHidden, width); + ggml_set_name(x, "l1_head.readout_action"); + ggml_set_input(x); + io.readout = x; + + // Q comes from the probe and K/V from x through the SAME combined + // in_proj_weight, so each needs its own mul_mat and drops the slices it + // does not use -- nn.MultiheadAttention keeps the [Wq;Wk;Wv] row blocks + // whatever is fed through it. + ggml_tensor * qkv_probe = ggml_add(C, ggml_mul_mat(C, qkv_w, probe), qkv_b); + ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv_probe, kHidden, 1, qkv_probe->nb[1], 0)); + ggml_tensor * qkv_x = ggml_add(C, ggml_mul_mat(C, qkv_w, x), qkv_b); + ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv_x, kHidden, width, qkv_x->nb[1], (size_t) kHidden*qkv_x->nb[0])); + ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv_x, kHidden, width, qkv_x->nb[1], (size_t) 2*kHidden*qkv_x->nb[0])); + + // Heads on ne2 and window on ne3 are both batch axes mul_mat loops over, + // never cross-multiplied, which keeps each timestep independent. + ggml_tensor * Qh = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, q, map_head_dim, map_heads, 1, 1), 0, 2, 1, 3)); + ggml_tensor * Kh = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, k, map_head_dim, map_heads, 1, width), 0, 2, 1, 3)); + ggml_tensor * Vh = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, v, map_head_dim, map_heads, 1, width), 1, 2, 0, 3)); + + ggml_tensor * scores = ggml_mul_mat(C, Qh, Kh); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + // One readout token per timestep, so this softmax is over a single logit + // and always yields 1.0. Kept as the real op in case that changes. + ggml_tensor * probs = ggml_soft_max_ext(C, scores, nullptr, 1.0f/std::sqrt((float) map_head_dim), 0.0f); + ggml_tensor * attended = ggml_mul_mat(C, Vh, probs); + ggml_tensor * merged = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, attended, 0, 2, 1, 3)), kHidden, width); + + ggml_tensor * attn_out = ggml_add(C, ggml_mul_mat(C, o_w, merged), o_b); + ggml_tensor * y = ggml_add(C, ggml_mul(C, ggml_norm(C, attn_out, ln_eps), norm_w), norm_b); + ggml_tensor * h = ggml_gelu_erf(C, ggml_add(C, ggml_mul_mat(C, ffn_up_w, y), ffn_up_b)); + h = ggml_add(C, ggml_mul_mat(C, ffn_down_w, h), ffn_down_b); + // The residual is onto attn_out, before the norm, as in MAPHead. + ggml_tensor * emb = ggml_add(C, attn_out, h); + + ggml_tensor * mean_raw = ggml_add(C, ggml_mul_mat(C, mean_w, emb), mean_b); + ggml_tensor * out = ggml_scale(C, ggml_tanh(C, ggml_scale(C, mean_raw, 1.0f/max_action)), max_action); + ggml_set_name(out, "l1_head.mean_normalized"); + ggml_set_output(out); + if (ggml_nelements(out) != (int64_t) action_total*width) { + std::fprintf(stderr, "vla(octo): L1 mean_proj out-dim=%lld does not match horizon*dim=%d\n", + (long long) out->ne[0], action_total); + return nullptr; + } + io.out = out; + + ggml_cgraph * gf = ggml_new_graph_custom(C, 512, false); + ggml_build_forward_expand(gf, out); + return gf; + }); + if (!built) { + std::fprintf(stderr, "vla(octo): L1 head graph build failed\n"); + return false; + } + + OctoRuntime::L1IO& io = rt.l1_head.io(); + ggml_backend_tensor_set(io.readout, readout_action.data(), 0, ggml_nbytes(io.readout)); + if (!octo_compute(rt, rt.l1_head.graph(), "L1 head")) + return false; + + std::vector mean_normalized((size_t) action_total*width); + ggml_backend_tensor_get(io.out, mean_normalized.data(), 0, ggml_nbytes(io.out)); + // predict_action takes the last window timestep, as the diffusion head does. + final_actions.resize((size_t) action_total); + std::copy_n(mean_normalized.begin()+(ptrdiff_t) ((size_t) action_total*(width-1)), + (size_t) action_total, final_actions.begin()); + return true; +} + +bool read_kv_u8_array(const gguf_reader& g, const char * key, std::vector& out) { + const int64_t id = gguf_find_key(g.gctx, key); + if (id < 0) { + std::fprintf(stderr, "vla(octo): missing metadata %s\n", key); + return false; + } + if (gguf_get_kv_type(g.gctx, id) != GGUF_TYPE_ARRAY || gguf_get_arr_type(g.gctx, id) != GGUF_TYPE_UINT8) { + std::fprintf(stderr, "vla(octo): %s is not a UINT8 array\n", key); + return false; + } + const size_t n = gguf_get_arr_n(g.gctx, id); + const uint8_t * data = (const uint8_t *) gguf_get_arr_data(g.gctx, id); + out.assign(data, data+n); + return true; +} + +// Which top-level key of octo.dataset_statistics to un-normalize against, when +// the caller did not pin one down. In order: VLA_OCTO_UNNORM_DATASET, the sole +// key if there is only one, then bridge_dataset, which is what the pretrain +// checkpoint's ~25-key OXE mix was always read against. Several real candidates +// and no hint is a failure, not a guess. +bool resolve_unnorm_dataset_key(const nlohmann::json& stats, std::string& key) { + if (const char * env = std::getenv("VLA_OCTO_UNNORM_DATASET"); env && env[0] != '\0') { + key = env; + return true; + } + if (stats.is_object() && stats.size() == 1) { + key = stats.begin().key(); + return true; + } + if (stats.is_object() && stats.contains("bridge_dataset")) { + key = "bridge_dataset"; + return true; + } + std::fprintf(stderr, + "vla(octo): cannot resolve the unnorm dataset key (%zu candidates in " + "octo.dataset_statistics); set VLA_OCTO_UNNORM_DATASET\n", + stats.is_object() ? stats.size() : (size_t) 0); + return false; +} + +// Single-dataset checkpoints write the stats block flat, with no dataset-name +// wrapper. A member literally named "action" holding a "mean" is what tells the +// two shapes apart; no real dataset name collides with that. +bool resolve_stats_block(const nlohmann::json& j, const std::string& dataset_key_in, + const nlohmann::json ** out) { + if (j.is_object() && j.contains("action") && j["action"].is_object() && j["action"].contains("mean")) { + *out = &j; + return true; + } + std::string dataset_key = dataset_key_in; + if (dataset_key.empty() && !resolve_unnorm_dataset_key(j, dataset_key)) + return false; + if (!j.contains(dataset_key)) { + std::fprintf(stderr, "vla(octo): dataset_statistics has no key %s\n", dataset_key.c_str()); + return false; + } + *out = &j[dataset_key]; + return true; +} + +// Parses octo.dataset_statistics once and caches the blocks on `rt`: it is a +// JSON blob in the metadata, ~25 datasets wide for the pretrain checkpoint, and +// re-reading 21 floats out of it per request is pure overhead. +bool ensure_stats(OctoRuntime& rt, gguf_reader& g, const std::string& dataset_key_in, int64_t action_dim) { + if (rt.stats_loaded && rt.stats_key == dataset_key_in) + return true; + + const std::string stats_json = g.str("octo.dataset_statistics"); + if (stats_json.empty()) { + std::fprintf(stderr, "vla(octo): missing octo.dataset_statistics\n"); + return false; + } + nlohmann::json j = nlohmann::json::parse(stats_json, nullptr, false); + if (j.is_discarded()) { + std::fprintf(stderr, "vla(octo): octo.dataset_statistics is not valid JSON\n"); + return false; + } + const nlohmann::json * block = nullptr; + if (!resolve_stats_block(j, dataset_key_in, &block)) + return false; + if (!block->contains("action")) { + std::fprintf(stderr, "vla(octo): dataset_statistics block has no .action\n"); + return false; + } + + const auto& act = (*block)["action"]; + OctoRuntime::ActionStats a; + a.mean = act.at("mean").get>(); + a.stdv = act.at("std").get>(); + for (bool b : act.at("mask").get>()) + a.mask.push_back(b ? 1 : 0); + if (a.mask.size() != (size_t) action_dim || a.mean.size() != (size_t) action_dim || + a.stdv.size() != (size_t) action_dim) { + std::fprintf(stderr, "vla(octo): dataset_statistics/action is not %lld-dim\n", (long long) action_dim); + return false; + } + + OctoRuntime::ProprioStats pr; + bool has_pr = false; + if (block->contains("proprio")) { + const auto& p = (*block)["proprio"]; + pr.mean = p.at("mean").get>(); + pr.stdv = p.at("std").get>(); + if (pr.mean.size() != (size_t) kProprioTokens || pr.stdv.size() != (size_t) kProprioTokens) { + std::fprintf(stderr, "vla(octo): dataset_statistics/proprio is not %d-dim\n", kProprioTokens); + return false; + } + has_pr = true; + } + + rt.action_stats = std::move(a); + rt.proprio_stats = std::move(pr); + rt.has_proprio_stats = has_pr; + rt.stats_key = dataset_key_in; + rt.stats_loaded = true; + return true; +} + +// unnorm[d] = mask[d] ? norm[d]*std[d]+mean[d] : norm[d]. The masked-out dims +// (the gripper, in every LIBERO block) pass through untouched. +bool unnormalize_action(const OctoRuntime& rt, + const std::vector& normalized_flat, + std::vector& unnorm_flat) { + const OctoRuntime::ActionStats& a = rt.action_stats; + const size_t dim = a.mask.size(); + if (dim == 0 || normalized_flat.empty() || normalized_flat.size()%dim != 0) { + std::fprintf(stderr, "vla(octo): unexpected action shape for un-normalization\n"); + return false; + } + + const size_t horizon = normalized_flat.size()/dim; + unnorm_flat.resize(normalized_flat.size()); + for (size_t t=0; t& dst) { + dst.resize((size_t) 3*side*side); + for (int y=0; y& out) { + if (v.format != PixelFormat::U8 || !v.data) { + std::fprintf(stderr, "vla(octo): predict only supports PixelFormat::U8 images\n"); + return false; + } + if (v.w == side && v.h == side) + return preprocess_image_chw("octo", v, side, out); + + std::vector resized; + resize_hwc((const uint8_t *) v.data, v.w, v.h, side, resized); + const ImageView rv{resized.data(), side, side, PixelFormat::U8}; + return preprocess_image_chw("octo", rv, side, out); +} + +struct OctoFrame { + std::vector primary; ///< [3,primary_size,primary_size] normalized. + std::vector wrist; ///< Empty when the client sent no wrist view. + bool wrist_real = false; + std::vector input_ids; + std::vector attention_mask; + std::vector proprio_raw; ///< Original units; z-scored here, not by the caller. + const float * noise = nullptr; +}; + +bool run_pipeline(OctoModelArch& m, const OctoFrame& f, + std::vector& unnormalized_out, + float& ms_vision_out, float& ms_inference_out) { + using clock = std::chrono::steady_clock; + + OctoRuntime& rt = m.rt; + if (m.head_type != "diffusion" && m.head_type != "l1") { + std::fprintf(stderr, "vla(octo): head_type=%s is not implemented\n", m.head_type.c_str()); + return false; + } + + const int window_size = (int) m.window_size; + const int action_total = (int) (m.action_horizon*m.action_dim); + const int n_proprio = m.has_proprio ? kProprioTokens : 0; + if (!ensure_stats(rt, m.io, "", m.action_dim)) + return false; + + // Cold start: history is filled with copies of the one live frame and every + // slot but the last is marked padding, as HistoryWrapper.reset() does. + std::vector primary_valid ((size_t) window_size, 1); + std::vector wrist_valid ((size_t) window_size, f.wrist_real ? 1 : 0); + std::vector timestep_valid((size_t) window_size, 0); + timestep_valid[(size_t) window_size-1] = 1; + + const OctoSeqLayout layout = build_seq_layout(true, primary_valid, wrist_valid, timestep_valid, + window_size, (int) m.primary_tokens, + (int) m.wrist_tokens, n_proprio); + + const auto t_vision0 = clock::now(); + std::vector primary_pos, wrist_pos, proprio_pos; + // Every live window slot holds the same frame, so it is replicated rather + // than re-decoded. Language-only conditioning means no goal image: the task + // frame is a zero uint8 image, which normalizes to -1. + auto tokenize_view = [&](graph_cache& cache, + const char * view, const std::vector& frame, int side, int n_tok, + const std::vector& steps, std::vector& out) { + const size_t n = (size_t) 3*side*side; + if (frame.size() != n) { + std::fprintf(stderr, "vla(octo): %s frame is %zu floats, expected %zu\n", view, frame.size(), n); + return false; + } + std::vector obs(n*steps.size()); + for (size_t i=0; i task(n, -1.0f); + return run_obs_tokenizer_graph(rt, cache, view, obs, task, side, n_tok, (int) steps.size(), steps, out); + }; + + if (!layout.primary_steps.empty() && + !tokenize_view(rt.obs_primary, "primary", f.primary, (int) m.primary_size, + (int) m.primary_tokens, layout.primary_steps, primary_pos)) + return false; + if (!layout.wrist_steps.empty() && + !tokenize_view(rt.obs_wrist, "wrist", f.wrist, (int) m.wrist_size, + (int) m.wrist_tokens, layout.wrist_steps, wrist_pos)) + return false; + + if (!layout.proprio_steps.empty()) { + if (!rt.has_proprio_stats) { + std::fprintf(stderr, "vla(octo): proprio checkpoint but dataset_statistics has no .proprio block\n"); + return false; + } + const std::vector& mean = rt.proprio_stats.mean; + const std::vector& stdv = rt.proprio_stats.stdv; + const int n_steps = (int) layout.proprio_steps.size(); + + std::vector proprio_norm((size_t) n_steps*kProprioTokens); + for (int t=0; t(clock::now()-t_vision0).count(); + + const auto t_inference0 = clock::now(); + + std::vector lang_key = f.input_ids; + lang_key.insert(lang_key.end(), f.attention_mask.begin(), f.attention_mask.end()); + if (rt.lang_key != lang_key || rt.lang_steps != window_size) { + std::vector t5_out; + if (!run_t5_encoder_graph(rt, f.input_ids, f.attention_mask, t5_out)) + return false; + if (!run_language_graph(rt, t5_out, window_size, rt.lang_pos, rt.lang_repeated)) + return false; + rt.lang_key = std::move(lang_key); + rt.lang_steps = window_size; + } + + ggml_tensor * readout_pos_r = rt.weight("octo.readout.action.pos_embd"); + if (!readout_pos_r) + return false; + const std::vector readout_pos = tensor_to_vec(readout_pos_r); + + std::vector input, mask; + if (!assemble_transformer_input(layout, rt.lang_pos, primary_pos, wrist_pos, proprio_pos, + rt.lang_repeated, readout_pos, input)) + return false; + build_transformer_mask(layout, mask); + + std::vector readout_action; + if (!run_transformer_graph(rt, layout, input, mask, readout_action)) + return false; + + std::vector normalized; + if (m.head_type == "l1") { + if (!run_l1_action_head_graph(rt, readout_action, window_size, action_total, m.max_action, normalized)) + return false; + } else { + const OctoDiffusion diff{(int) m.diffusion_steps, m.diffusion_s, m.max_action}; + if (!run_diffusion(rt, m.rng, f.noise, readout_action, window_size, action_total, diff, normalized)) + return false; + } + ms_inference_out = std::chrono::duration(clock::now()-t_inference0).count(); + + return unnormalize_action(rt, normalized, unnormalized_out); +} + +} // namespace + +std::unique_ptr octo_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string&) { + if (!mmproj_path.empty()) + std::printf("vla(octo): note - mmproj '%s' is ignored (Octo ships one GGUF)\n", mmproj_path.c_str()); + + auto m = std::make_unique(); + m->gguf_path = ckpt_path; + + if (!m->io.open(ckpt_path)) + return nullptr; + if (!m->io.has("octo.architecture")) { + std::fprintf(stderr, "vla(octo): %s is not an Octo GGUF\n", ckpt_path.c_str()); + return nullptr; + } + if (!load_config(m->io, *m)) + return nullptr; + + const Backend b = backend_init("vla(octo)", m->n_threads); + if (!b.handle) + return nullptr; + m->backend = b.handle; + + if (!load_weights(*m, m->io)) + return nullptr; + m->rt.init(m->backend, m->ctx_weights); + + std::printf("vla(octo): weights resident in %.2f GiB (F32) - head=%s window=%lld horizon=%lld " + "action_dim=%lld proprio=%s\n", + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), + m->head_type.c_str(), (long long) m->window_size, (long long) m->action_horizon, + (long long) m->action_dim, m->has_proprio ? "yes" : "no"); + return m; +} + +bool octo_tokenize_text(const std::string& ckpt_path, + const std::string& text, + std::vector& input_ids, + std::vector& attention_mask) { + gguf_reader g{"octo"}; + if (!g.open(ckpt_path)) + return false; + + std::vector spm_bytes; + if (!read_kv_u8_array(g, "octo.tokenizer.spm_model", spm_bytes)) + return false; + const uint32_t eos_id = g.has("octo.tokenizer.eos_id") ? g.u32("octo.tokenizer.eos_id") : 1; + const uint32_t pad_id = g.has("octo.tokenizer.pad_id") ? g.u32("octo.tokenizer.pad_id") : 0; + const int64_t max_length = g.has("octo.tokens.language") ? g.u32("octo.tokens.language") : kTaskTokens; + + sentencepiece::SentencePieceProcessor sp; + const auto status = sp.LoadFromSerializedProto( + absl::string_view(reinterpret_cast(spm_bytes.data()), spm_bytes.size())); + if (!status.ok()) { + std::fprintf(stderr, "vla(octo): sentencepiece LoadFromSerializedProto failed: %s\n", + status.ToString().c_str()); + return false; + } + + std::vector ids = sp.EncodeAsIds(text); + if ((int64_t) ids.size() > max_length-1) + ids.resize((size_t) (max_length-1)); + input_ids.assign(ids.begin(), ids.end()); + input_ids.push_back((int32_t) eos_id); + attention_mask.assign(input_ids.size(), 1); + input_ids.resize((size_t) max_length, (int32_t) pad_id); + attention_mask.resize((size_t) max_length, 0); + return true; +} + +// Unlike the other archs, this returns the action in world units rather than the +// normalized one: Octo's dataset_statistics lives inside the multi-hundred-MB +// checkpoint, not in a small sibling stats.json a client could hold, so +// un-normalizing server-side is the only way a client avoids shipping the GGUF +// to read its metadata. The client still owns the gripper convention. +std::vector OctoModelArch::predict(const Inputs& in) { + const auto t_total0 = std::chrono::steady_clock::now(); + stats = Stats{}; + + if (in.n_images < 1 || !in.images) { + std::fprintf(stderr, "vla(octo): predict needs at least the primary image\n"); + return {}; + } + if (in.n_lang != (int) language_tokens) { + std::fprintf(stderr, "vla(octo): predict expects exactly %lld language tokens " + "(t5-base, padding=\"max_length\"), got %d\n", + (long long) language_tokens, in.n_lang); + return {}; + } + if (in.attention_mask && in.attention_mask_n != (int) language_tokens) { + std::fprintf(stderr, "vla(octo): attention_mask_n=%d does not match the %lld language tokens\n", + in.attention_mask_n, (long long) language_tokens); + return {}; + } + + OctoFrame f; + f.wrist_real = in.n_images >= 2; + if (!octo_prepare_view(in.images[0], (int) primary_size, f.primary)) + return {}; + // A wrist view we were not given is left out of the sequence entirely rather + // than zero-filled and masked. See build_seq_layout. + if (f.wrist_real && !octo_prepare_view(in.images[1], (int) wrist_size, f.wrist)) + return {}; + + f.input_ids.assign(in.lang_tokens, in.lang_tokens+in.n_lang); + if (in.attention_mask) { + f.attention_mask.assign(in.attention_mask, in.attention_mask+in.attention_mask_n); + } else { + // T5's encoder needs real padding information. With right-padded ids the + // pad token carries it, which reproduces the mask the client's own + // tokenizer would have sent. + f.attention_mask.resize(f.input_ids.size()); + for (size_t i=0; i unnormalized; + float ms_vision = 0.f; + float ms_inference = 0.f; + if (!run_pipeline(*this, f, unnormalized, ms_vision, ms_inference)) + return {}; + + stats.ms_vision = ms_vision; + stats.ms_inference = ms_inference; + stats.ms_total = std::chrono::duration( + std::chrono::steady_clock::now()-t_total0).count(); + return unnormalized; +} + +} // namespace vla diff --git a/src/models/octo.h b/src/models/octo.h new file mode 100644 index 0000000..0f253f3 --- /dev/null +++ b/src/models/octo.h @@ -0,0 +1,31 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace vla { + +// T5 SentencePiece-unigram tokenization, against the vocab the converter baked +// into the GGUF. Pads or truncates to octo.tokens.language and appends EOS, +// matching t5-base at max_length=16, padding="max_length", truncation=True. +bool octo_tokenize_text(const std::string& ckpt_path, + const std::string& text, + std::vector& input_ids, + std::vector& attention_mask); + +} // namespace vla diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index 67c9c23..7d31060 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -14,8 +14,11 @@ // One-shot action prediction from the command line. Loads a model, decodes an // image plus an instruction, runs one predict(), and prints the action chunk. -// No server, no simulator. There is no tokenizer in the C++ core, so --text -// shells out to scripts/tokenize_prompt.py; --tokens takes ids directly. +// No server, no simulator. Most archs have no tokenizer in the C++ core, so +// --text shells out to scripts/tokenize_prompt.py; --tokens takes ids directly. +// Octo is the exception: its T5 SentencePiece vocab is baked into the GGUF, so +// --text is tokenized in-process (no Python) and also yields the attention mask +// that Octo's predict() requires. // // vla-cli [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image img2.jpg] // (--text "pick up the bowl" | --tokens id,id,...) [--state f,f,...] [--pretty] @@ -23,6 +26,9 @@ #include "arch.h" #include "model.h" #include "serving/hf_fetch.h" +#ifdef VLA_USE_OCTO +#include "models/octo.h" +#endif #define STB_IMAGE_IMPLEMENTATION #define STB_IMAGE_STATIC @@ -123,10 +129,36 @@ const char * arch_slug(Arch a) { case Arch::VLA_ADAPTER: return "vla_adapter"; case Arch::OPENVLA_OFT: return "openvla_oft"; case Arch::VLA_JEPA: return "vla_jepa"; + case Arch::OCTO: return "octo"; } return ""; } +#ifdef VLA_USE_OCTO +bool octo_ckpt(const std::string & ckpt) { + Arch a; + return detect_arch_from_ckpt(ckpt, &a) && a == Arch::OCTO; +} + +// Octo's tokenizer ships inside the checkpoint, so --text needs no Python here +// and yields the attention mask its T5 encoder wants alongside the ids. +bool octo_tokens(const std::string & ckpt, const std::string & text, + std::vector & lang, std::vector & attn) { + if (octo_tokenize_text(ckpt, text, lang, attn)) + return true; + std::fprintf(stderr, "vla-cli: octo tokenization failed\n"); + return false; +} +#else +bool octo_ckpt(const std::string &) { + return false; +} + +bool octo_tokens(const std::string &, const std::string &, std::vector &, std::vector &) { + return false; +} +#endif + // The instruction reaches a shell command, so keep it to plain prose. bool text_ok(const std::string & s) { if (s.empty() || s.size() > 512) @@ -198,7 +230,8 @@ void usage(const char * prog) { " --ckpt model checkpoint GGUF\n" " -hf HuggingFace repo, user/repo[:file.gguf], cached under $VLA_CACHE\n" " --image image file, repeat for multi-view (decoded via stb_image)\n" - " --text instruction, tokenized by scripts/tokenize_prompt.py (needs transformers)\n" + " --text instruction; tokenized by scripts/tokenize_prompt.py (needs\n" + " transformers), or in-process for Octo, whose vocab is in the GGUF\n" " --tokens language token ids, comma-separated, if you tokenized already\n" " --state proprioception floats, comma-separated (default zeros)\n" " --pretty print one action row (max_action_dim values) per line\n", @@ -257,17 +290,25 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "vla-cli: pass --text or --tokens, not both\n"); return 1; } - if (!text_s.empty()) { - tokens_s = tokenize_text(ckpt, text_s); - if (tokens_s.empty()) - return 1; - std::fprintf(stderr, "vla-cli: --text tokenized to %s\n", tokens_s.c_str()); - } - // Validate the cheap args before loading the model. std::vector lang; + std::vector attn; // Octo only; empty leaves Inputs::attention_mask null. std::vector state; - if (!parse_ints(tokens_s, lang) || !parse_floats(state_s, state)) + + if (octo_ckpt(ckpt) && !text_s.empty()) { + if (!octo_tokens(ckpt, text_s, lang, attn)) + return 1; + } else { + if (!text_s.empty()) { + tokens_s = tokenize_text(ckpt, text_s); + if (tokens_s.empty()) + return 1; + std::fprintf(stderr, "vla-cli: --text tokenized to %s\n", tokens_s.c_str()); + } + if (!parse_ints(tokens_s, lang)) + return 1; + } + if (!parse_floats(state_s, state)) return 1; if (lang.empty()) { std::fprintf(stderr, "vla-cli: --tokens parsed to nothing\n"); @@ -302,6 +343,10 @@ int main(int argc, char ** argv) { in.n_images = (int) views.size(); in.lang_tokens = lang.data(); in.n_lang = (int) lang.size(); + if (!attn.empty()) { + in.attention_mask = attn.data(); + in.attention_mask_n = (int) attn.size(); + } in.state = state.data(); in.noise = nullptr; // predict() samples N(0,1) when omitted