diff --git a/.gitignore b/.gitignore index 8894aa1..c9727da 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,7 @@ _deps/ # Models / data *.gguf +*.gguf.tensor_map.json *.bin *.safetensors *.pt @@ -54,3 +55,5 @@ eval/sim/libero/libero_uv/ eval/sim/simpler/SimplerEnv/ eval/sim/simpler/simpler_uv/ /models/ + +LOCAL_WSL_RUN_COMMANDS.md \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e757e1..b7f909f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -25,6 +25,22 @@ FetchContent_Declare(llama ) FetchContent_MakeAvailable(llama) +# T5 SentencePiece tokenizer (Octo language instructions). Static, self-contained +# -- no shared lib, no test binaries, no tcmalloc. Uses the SYSTEM protobuf +# (not sentencepiece's vendored protobuf-lite 3.14) so it shares one protobuf +# runtime with the server's own .pb.cc code (system protobuf) -- vla-server +# aborts at static-init if two protobuf versions end up in 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 +) +FetchContent_MakeAvailable(sentencepiece) + add_library(vla_core src/model.cpp src/models/smolvla.cpp @@ -34,6 +50,7 @@ add_library(vla_core src/models/gr00tn1d5.cpp src/models/gr00tn1d6.cpp src/models/gr00tn1d7.cpp + src/models/octo.cpp src/models/bitvla.cpp src/models/vla_adapter.cpp src/models/openvla_oft.cpp @@ -43,8 +60,10 @@ target_include_directories(vla_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src ${llama_SOURCE_DIR}/vendor + PRIVATE + ${sentencepiece_SOURCE_DIR}/src ) -target_link_libraries(vla_core PUBLIC llama ggml) +target_link_libraries(vla_core PUBLIC llama ggml PRIVATE sentencepiece) if(GGML_CUDA) target_compile_definitions(vla_core PUBLIC GGML_USE_CUDA) diff --git a/docs/assets/open_loop/vla_cpp_ep0raw_step4500_train_gt_vs_pred.png b/docs/assets/open_loop/vla_cpp_ep0raw_step4500_train_gt_vs_pred.png new file mode 100644 index 0000000..fff5bfd Binary files /dev/null and b/docs/assets/open_loop/vla_cpp_ep0raw_step4500_train_gt_vs_pred.png differ diff --git a/docs/octo_open_loop_vla_cpp_en.md b/docs/octo_open_loop_vla_cpp_en.md new file mode 100644 index 0000000..c63ad47 --- /dev/null +++ b/docs/octo_open_loop_vla_cpp_en.md @@ -0,0 +1,147 @@ +# Open-Loop Evaluation Report: Jitter-Adapted Octo (step 4500) via the vla.cpp C++ Engine + +Evaluation date: **2026-08-04** +Status: **complete** + +## 1. Objective + +This report presents an open-loop evaluation of the fine-tuned Octo L1/proprio +checkpoint `one_episode_ep0_jitter1cm_command_v4_adapt` (checkpoint step +`4500`), run through the **vla.cpp C++ inference engine** (GGUF checkpoint, +`head_type=l1` action head + proprio tokenizer, served over `vla-server`) +rather than the original PyTorch model. The objective is to determine whether +the C++ engine reproduces the checkpoint's published open-loop numbers +(computed with the original PyTorch model, `octo-pytorch-kamusarj`) closely +enough to trust it as a drop-in inference replacement. This is not a +simulation rollout and does not measure task success. + +## 2. Protocol + +The protocol was adapted from: + +- [Isaac-GR00T: Step 4 — Open Loop Evaluation](https://github.com/NVIDIA/Isaac-GR00T/blob/main/getting_started/finetune_new_embodiment.md#step-4-open-loop-evaluation) +- [NVIDIA `gr00t/eval/open_loop_eval.py`](https://github.com/NVIDIA/Isaac-GR00T/blob/main/gr00t/eval/open_loop_eval.py) + +At each inference point, the evaluator (`eval/client/run_open_loop_octo.py`): + +1. retrieves the ground-truth observation at that RLDS trajectory step (top + camera 256², wrist camera 128², raw 7-D proprioceptive state, language + instruction); +2. sends the observation to `vla-server`, which runs the head_type=l1 GGUF + checkpoint end-to-end (proprio z-score normalization, obs tokenizers, T5 + language encoder, block transformer, MAPHead action head) on the ggml CPU + backend; +3. receives a predicted `(20, 7)` action chunk, already un-normalized to + original units server-side; +4. takes the first `execution_horizon = 8` actions from the chunk; +5. concatenates the predicted chunks over time; +6. compares them with the corresponding ground-truth actions; +7. computes MAE, MSE, and RMSE over valid action values; +8. saves a GT-vs-prediction plot and a NumPy trace for auditing. + +The observations always come from the recorded dataset trajectory +(teacher-forced). Predicted actions are not fed back into a simulator or used +to generate subsequent observations. The results therefore measure action +imitation quality only, reproduced through a different (C++/ggml) inference +engine than the one used to produce the published numbers. + +The `original_units` metrics are computed after reversing the checkpoint's +normalization (z-score, `octo.dataset_statistics.action`). The first six +dimensions are joint actions, and the final dimension is the follower-gripper +command. The `normalized` metrics are retained for normalization debugging. + +## 3. Evaluation Configuration + +| Engine | Checkpoint | Evaluation dataset | Window | Action horizon | Execution horizon | +|---|---:|---|---:|---:|---:| +| vla.cpp (ggml, CPU) | 4500 | `aloha_carrot_easy_rlds` | 1 | 20 | 8 | + +The GGUF checkpoint was converted from the same PyTorch checkpoint +(`one_episode_ep0_jitter1cm_command_v4_adapt`, step 4500) used to produce the +published numbers, via `scripts/convert_octo_to_gguf.py`. Both the +vla.cpp run and the published PyTorch run use the same original RLDS +evaluation dataset (`aloha_carrot_easy_rlds`, episode 0) and the same +episode-0 normalization statistics baked into the checkpoint. + +| Engine | Split | Trajectory index | Source episode | Number of steps | Inference calls | +|---|---|---:|---:|---:|---:| +| vla.cpp | train | 0 | 0 | 149 | 19 | + +## 4. Main Results — vla.cpp vs. Pytorch + +### 4.1 Metrics in Original Action Units + +| Source | MAE | MSE | RMSE | Gripper accuracy | +|---|---:|---:|---:|---:| +| **vla.cpp** | 0.008952 | 0.000212 | 0.014549 | 99.33% | +| **Pytorch** | 0.008951 | 0.000212 | 0.014549 | 99.33% | + +### 4.2 MAE by Action Dimension (Original Units) + +| Source | j0 | j1 | j2 | j3 | j4 | j5 | gripper | +|---|---:|---:|---:|---:|---:|---:|---:| +| **vla.cpp** | 0.009523 | 0.012367 | 0.009324 | 0.003325 | 0.008640 | 0.008890 | 0.010593 | +| **Pytorch** | 0.009521 | 0.012370 | 0.009322 | 0.003325 | 0.008641 | 0.008889 | 0.010592 | + +### 4.3 Normalized Metrics for Auditing + +| Source | Normalized MAE | Normalized MSE | Normalized RMSE | +|---|---:|---:|---:| +| **vla.cpp** | 0.032435 | 0.003465 | 0.058863 | +| **Pytorch** | 0.032432 | 0.003465 | 0.058861 | + +## 5. Plots + +Blue lines show ground-truth actions, orange lines show predicted actions, and +dashed gray lines show the reference state. Purple circles and vertical dotted +lines mark inference points. + +**Jitter-adapted checkpoint at step 4500 — train trajectory 0 (vla.cpp engine)** + +![Jitter-adapted checkpoint at step 4500, train trajectory 0, vla.cpp engine: ground-truth actions versus predicted actions](assets/open_loop/vla_cpp_ep0raw_step4500_train_gt_vs_pred.png) + +## 6. Agreement with Published Results + +The vla.cpp C++ engine reproduces every published metric to within **~1e-5**: + +| Metric | vla.cpp | Pytorch | \|difference\| | +|---|---:|---:|---:| +| original MAE | 0.008952 | 0.008951 | 4.4e-7 | +| original MSE | 0.000212 | 0.000212 | 2.0e-8 | +| original RMSE | 0.014549 | 0.014549 | 6.8e-7 | +| gripper accuracy | 0.993289 | 0.993289 | 0.0 | +| normalized MAE | 0.032435 | 0.032432 | 2.7e-6 | +| normalized MSE | 0.003465 | 0.003465 | 2.3e-7 | +| normalized RMSE | 0.058863 | 0.058861 | 2.0e-6 | + +## 7. Reproduction + +Convert the PyTorch checkpoint to GGUF: + +```bash +python scripts/convert_octo_to_gguf.py \ + --ckpt-format pytorch \ + --ckpt ~/octo_ckpts/kamusarj_ep0raw_4500 \ + --step 4500 \ + --out ~/octo_gguf/octo-aloha-ep0raw-4500.gguf +``` + +Start `vla-server` on the GGUF checkpoint (clean shell, no conda env active): + +```bash +build_cpu/vla-server ~/octo_gguf/octo-aloha-ep0raw-4500.gguf +``` + +Run the open-loop evaluation client (separate shell, with `tensorflow_datasets` +available to read RLDS): + +```bash +python eval/client/run_open_loop_octo.py \ + --vla-addr tcp://localhost:5555 \ + --rlds-dir ~/aloha_carrot_ep0_rlds/aloha_carrot_easy_rlds/1.0.0 \ + --dataset-statistics ~/octo_ckpts/kamusarj_ep0raw_4500/dataset_statistics.json \ + --output-dir ~/outputs/open_loop_octo_4500 +``` + +`--execution-horizon 8` matches the published evaluation's execution horizon, +so both runs use identical inference points and are directly comparable. diff --git a/eval/client/adapters.py b/eval/client/adapters.py index 10a203a..350661c 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,29 @@ 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: + """Octo LIBERO preprocessing segment 1 (raw sim frame -> model_entry image). + + Two steps, matching OctoPt's get_libero_image / dlimp preprocessing: + 1. Rotate 180 (`[::-1, ::-1]`) -- the raw off-screen render comes out upside-down, + same convention already handled for evo1/gr00t above (Evo1PipelineAdapter, + Gr00tPipelineAdapter). + 2. Resize to `image_size`. LIBERO's sim renders at 256x256 and Octo's primary + tokenizer also expects 256x256, so with the sim kept at 256 (see TIP-CLIENT -- + do not change LIBERO's camera_heights/widths away from 256) this resize is an + identity; PIL LANCZOS only actually resamples if the render size ever differs + from `image_size`. + Verified against golden tier-A (raw.npy/model_entry.npy pairs): exact on + lossless-JPEG synthetic frames, within JPEG round-trip noise (~20 uint8 max_abs) on + photographic ones -- see TIP-P report. + """ + 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 +177,30 @@ class Gr00tN15PipelineAdapter(Gr00tPipelineAdapter): def parse_action(self, action: np.ndarray) -> np.ndarray: return np.asarray(action[:7], dtype=np.float32).copy() + +class OctoPipelineAdapter(BasePipelineAdapter): + """cyrusneary/octo-finetuned-libero (window=1). Primary camera only, matching this + checkpoint's own single-camera finetune (its finetune_config.json image_obs_keys + never included a wrist key) -- vla-server zero-fills+masks-invalid the wrist slot for + us (octo.cpp:predict(), TIP-CLIENT), exactly reproducing what the checkpoint actually + trained on. No proprio/state input: Octo's observation_tokenizers are image-only. + """ + + def __init__(self, client: Any = None): + super().__init__(client) + + def parse_observation(self, obs: dict[str, Any]) -> dict[str, Any]: + primary = octo_preprocess_image(obs["pixels"]["image"], image_size=256) + return { + "observation.images.image": primary, + "task": obs.get("task_description", ""), + } + + def parse_action(self, action: np.ndarray) -> np.ndarray: + # octo.cpp:predict() already un-normalized (world units, dims 0..5) -- only the + # gripper (dim 6, Octo's own +1=open/0=close convention, passthrough/un-masked by + # unnormalize_action) needs converting to LIBERO's -1=open/+1=close and binarizing. + # Same formula as Evo1PipelineAdapter/Gr00tPipelineAdapter above. + 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_open_loop_octo.py b/eval/client/run_open_loop_octo.py new file mode 100644 index 0000000..adb7d97 --- /dev/null +++ b/eval/client/run_open_loop_octo.py @@ -0,0 +1,440 @@ +# 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. + +"""TIP-07: teacher-forced open-loop evaluation of the Octo L1/proprio (head_type=l1) +checkpoint through vla-server, reusing VlaCppClient (not a new client). + +Protocol mirrors octo-pytorch-kamusarj's scripts/evaluate_octo_open_loop.py (TIP-02 +Task-4; formulas cross-checked against that script's _masked_metrics/_plot_trajectory/ +_aggregate, ~/work/octo-pytorch-kamusarj at commit 47f1a3e): read a real RLDS episode, +run inference every `execution_horizon` dataset steps with the recorded (teacher-forced) +observation, stitch the executed prefix of each predicted chunk into a full-length +prediction, and report normalized + original-unit MAE/MSE/RMSE, gripper accuracy, and a +GT-vs-pred plot. + +VlaCppClient's octo path (_predict_chunk_octo) predates this checkpoint's proprio input +(its ARCH_PRESETS["octo"] entry hardcodes max_state_dim=0 for the older, image-only +cyrusneary/octo-finetuned-libero checkpoint) -- it builds images + language but never +reads/sends Inputs::state. Per TIP-07 ("KHONG sua VlaCppClient cu, chi import"), +vla_cpp_client.py itself is untouched; OctoL1Client below subclasses it and overrides +just that one method to also send raw proprio via req.state (mirrors the parent's +image/tokenize/send logic, since Python has no clean way to inject one extra line into a +method without overriding the whole body). + +Usage: + python run_open_loop_octo.py \ + --vla-addr tcp://localhost:5555 \ + --rlds-dir ~/aloha_carrot_ep0_rlds/aloha_carrot_easy_rlds/1.0.0 \ + --dataset-statistics ~/octo_ckpts/kamusarj_jitter2525/dataset_statistics.json \ + --golden-t0-dir ~/octo_l1_golden/ep0_t0 \ + --output-dir outputs/open_loop_octo +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import sys +from pathlib import Path +from typing import Any + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +import matplotlib + +matplotlib.use("Agg") +from matplotlib import pyplot as plt +import numpy as np +import tensorflow_datasets as tfds + +from client.vla_cpp_client import VlaCppClient + +ACTION_NAMES = ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"] +PROTOCOL_SOURCE = ( + "https://github.com/NVIDIA/Isaac-GR00T/blob/main/" + "getting_started/finetune_new_embodiment.md#step-4-open-loop-evaluation" +) + + +class OctoL1Client(VlaCppClient): + """TIP-07: adds raw-proprio state to the octo predict request. See module docstring + for why this is an override (not an edit) of the parent's _predict_chunk_octo.""" + + def _predict_chunk_octo(self, observations: dict[str, Any]) -> np.ndarray: + images_u8: list[np.ndarray] = [] + for key in self.image_keys[:2]: + if key not in observations: + continue + img = np.asarray(observations[key], 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() + 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()) + + # TIP-07 addition: raw (un-normalized) proprio -- octo.cpp's proprio tokenizer + # z-scores it server-side (TIP-05 convention), matching Inputs::state's doc + # comment in model.h. self.max_state_dim must be overridden to 7 at construction + # (the "octo" ARCH_PRESETS default of 0 is for the older proprio-less checkpoint). + state = observations.get("state") + if state is not None: + state = np.asarray(state, dtype=np.float32).reshape(-1) + if state.shape[0] != self.max_state_dim: + raise ValueError( + f"octo: state has {state.shape[0]} dims, expected max_state_dim={self.max_state_dim}") + req.state.extend(float(x) for x in state) + + 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 + return (np.array(resp.action_chunk, dtype=np.float32) + .reshape(resp.chunk_size, resp.action_dim)) + + +def load_episode(rlds_dir: Path, traj_index: int) -> list[dict[str, Any]]: + builder = tfds.builder_from_directory(str(rlds_dir)) + ds = builder.as_dataset(split="train") + episode = None + for i, ep in enumerate(ds): + if i == traj_index: + episode = ep + break + if episode is None: + raise IndexError(f"{rlds_dir}: no trajectory index {traj_index}") + steps = [] + for s in episode["steps"]: + steps.append({ + "top": s["observation"]["top"].numpy(), + "wrist": s["observation"]["wrist"].numpy(), + "state": s["observation"]["state"].numpy().astype(np.float32), + "action": s["action"].numpy().astype(np.float64), + "instruction": s["language_instruction"].numpy().decode("utf-8"), + }) + return steps + + +def zscore_normalize(values: np.ndarray, stats: dict[str, Any]) -> np.ndarray: + mean = np.asarray(stats["mean"], dtype=np.float64) + std = np.asarray(stats["std"], dtype=np.float64) + mask = np.asarray(stats.get("mask", np.ones_like(mean, dtype=bool)), dtype=bool) + values = np.asarray(values, dtype=np.float64) + return np.where(mask, (values - mean) / std, values) + + +def masked_metrics(predicted: np.ndarray, target: np.ndarray, valid: np.ndarray) -> dict[str, Any]: + # TIP-02 formula, verbatim from evaluate_octo_open_loop.py::_masked_metrics. + if predicted.shape != target.shape or predicted.shape != valid.shape: + raise ValueError(f"metric shapes differ: predicted={predicted.shape} target={target.shape} valid={valid.shape}") + count = int(valid.sum()) + if count == 0: + raise ValueError("trajectory contains no valid action targets") + difference = predicted - target + abs_error = np.abs(difference) + sq_error = np.square(difference) + dim_count = valid.sum(axis=0) + per_dim_mae = np.divide((abs_error * valid).sum(axis=0), dim_count, + out=np.full(predicted.shape[1], np.nan, dtype=np.float64), where=dim_count > 0) + per_dim_mse = np.divide((sq_error * valid).sum(axis=0), dim_count, + out=np.full(predicted.shape[1], np.nan, dtype=np.float64), where=dim_count > 0) + return { + "valid_action_values": count, + "absolute_error_sum": float((abs_error * valid).sum()), + "squared_error_sum": float((sq_error * valid).sum()), + "mae": float((abs_error * valid).sum() / count), + "mse": float((sq_error * valid).sum() / count), + "rmse": float(math.sqrt((sq_error * valid).sum() / count)), + "per_dim_mae": per_dim_mae.tolist(), + "per_dim_mse": per_dim_mse.tolist(), + } + + +def aggregate(reports: list[dict[str, Any]]) -> dict[str, Any]: + # TIP-02 formula, verbatim from evaluate_octo_open_loop.py::_aggregate. + result: dict[str, Any] = {"trajectories": len(reports)} + for key in ("normalized", "original_units"): + count = sum(r[key]["valid_action_values"] for r in reports) + absolute_error_sum = sum(r[key]["absolute_error_sum"] for r in reports) + squared_error_sum = sum(r[key]["squared_error_sum"] for r in reports) + result[key] = { + "valid_action_values": count, + "mae": absolute_error_sum / count, + "mse": squared_error_sum / count, + "rmse": math.sqrt(squared_error_sum / count), + } + gripper_values = [r["gripper_accuracy"] for r in reports if r["gripper_accuracy"] is not None] + result["mean_gripper_accuracy"] = float(np.mean(gripper_values)) if gripper_values else None + return result + + +def plot_trajectory(*, predicted, target, state, valid, inference_points, title, output_path: Path) -> None: + # TIP-02 formula, verbatim from evaluate_octo_open_loop.py::_plot_trajectory. + action_dim = target.shape[1] + names = ACTION_NAMES[:action_dim] + [f"action_{i}" for i in range(len(ACTION_NAMES), action_dim)] + fig, axes = plt.subplots(action_dim, 1, figsize=(12, max(4, 2.6 * action_dim)), sharex=True, dpi=140) + if action_dim == 1: + axes = [axes] + timesteps = np.arange(len(target)) + for dim, ax in enumerate(axes): + target_values = np.where(valid[:, dim], target[:, dim], np.nan) + predicted_values = np.where(valid[:, dim], predicted[:, dim], np.nan) + if state is not None and state.shape == target.shape: + ax.plot(timesteps, state[:, dim], color="#666666", linestyle="--", linewidth=1.3, + alpha=0.8, label="state", zorder=1) + ax.plot(timesteps, target_values, color="#0072B2", linewidth=2.1, label="GT action", zorder=3) + ax.plot(timesteps, predicted_values, color="#D55E00", linewidth=1.9, label="pred action", zorder=2) + for point in inference_points: + ax.axvline(point, color="#CC79A7", linestyle=":", alpha=0.35, linewidth=1.0, zorder=0) + visible = [p for p in inference_points if 0 <= p < len(predicted_values) and np.isfinite(predicted_values[p])] + if visible: + ax.scatter(visible, predicted_values[visible], color="#CC79A7", edgecolors="white", + linewidths=0.5, s=24, label="inference point", zorder=4) + ax.set_ylabel(names[dim]) + ax.grid(True, color="#D0D0D0", alpha=0.45, linewidth=0.7) + if dim == 0: + ax.legend(loc="upper right", ncol=4, framealpha=0.95) + axes[-1].set_xlabel("dataset timestep") + fig.suptitle(title) + fig.tight_layout(rect=(0, 0, 1, 0.98)) + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path) + plt.close(fig) + + +def cross_check_preprocess(steps: list[dict[str, Any]], golden_dir: Path) -> bool: + """TIP-07 AC: model-input at t=0 (raw images + raw proprio the runner is about to + send) must match TIP-06's golden dump exactly -- same RLDS episode/step, proving the + runner reads the identical source data octo_l1_parity's golden was built from. Does + NOT re-check the server's internal normalize/tokenize forward pass -- that's already + verified bit-exact (T0-T5, <1e-4) by TIP-06's ctest, independent of this client.""" + ok = True + golden_top = np.load(golden_dir / "input.top_hwc_u8.npy") + golden_wrist = np.load(golden_dir / "input.wrist_hwc_u8.npy") + golden_proprio = np.load(golden_dir / "input.proprio_raw.npy") + top_diff = int(np.abs(steps[0]["top"].astype(np.int32) - golden_top.astype(np.int32)).max()) + wrist_diff = int(np.abs(steps[0]["wrist"].astype(np.int32) - golden_wrist.astype(np.int32)).max()) + proprio_diff = float(np.abs(steps[0]["state"] - golden_proprio).max()) + print(f"preprocess cross-check vs {golden_dir}:") + print(f" top image max|diff| = {top_diff} (expect 0)") + print(f" wrist image max|diff| = {wrist_diff} (expect 0)") + print(f" proprio_raw max|diff| = {proprio_diff:.8f} (expect < 1e-4)") + if top_diff != 0 or wrist_diff != 0 or proprio_diff >= 1e-4: + ok = False + print(f" PREPROCESS CROSS-CHECK: {'PASS' if ok else 'FAIL'}") + return ok + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--vla-addr", default="tcp://localhost:5555") + ap.add_argument("--rlds-dir", required=True, type=Path, + help="RLDS builder dir, e.g. ~/aloha_carrot_ep0_rlds/aloha_carrot_easy_rlds/1.0.0") + ap.add_argument("--dataset-statistics", required=True, type=Path, + help="jitter2525's dataset_statistics.json (action mean/std/mask/min/max)") + ap.add_argument("--golden-t0-dir", type=Path, default=None, + help="TIP-06 octo_l1_golden/ep0_t0 dir for the preprocess cross-check; skipped if omitted") + ap.add_argument("--traj-index", type=int, default=0) + ap.add_argument("--execution-horizon", type=int, default=8) + ap.add_argument("--steps", type=int, default=None, help="default: full episode length") + ap.add_argument("--output-dir", required=True, type=Path) + ap.add_argument("--recv-timeout-ms", type=int, default=120_000) + ap.add_argument("--model-label", default=None, + help="Plot-title label for the checkpoint under evaluation (TIP-12: was " + "hardcoded to a stale checkpoint name). Default: the " + "--dataset-statistics parent directory name, e.g. " + "'kamusarj_ep0raw_4500' for ~/octo_ckpts/kamusarj_ep0raw_4500/" + "dataset_statistics.json -- this client never loads the GGUF " + "itself (it only talks to an already-running vla-server), so the " + "checkpoint directory name is the most reliable identifier " + "available client-side.") + args = ap.parse_args() + + dataset_statistics_path = args.dataset_statistics.expanduser() + model_label = args.model_label or dataset_statistics_path.parent.name + dataset_statistics = json.loads(dataset_statistics_path.read_text()) + action_stats = dataset_statistics["action"] + action_dim = len(action_stats["mean"]) + + print(f"loading RLDS episode {args.traj_index} from {args.rlds_dir} ...") + rlds_steps = load_episode(args.rlds_dir.expanduser(), args.traj_index) + trajectory_length = len(rlds_steps) + print(f"episode has {trajectory_length} steps") + + preprocess_ok = None + if args.golden_t0_dir is not None: + preprocess_ok = cross_check_preprocess(rlds_steps, args.golden_t0_dir.expanduser()) + + client = OctoL1Client( + vla_addr=args.vla_addr, + arch="octo", + max_state_dim=action_dim, # override ARCH_PRESETS["octo"]'s stale max_state_dim=0 + real_action_dim=action_dim, + image_keys=["observation.images.image", "observation.images.image2"], + recv_timeout_ms=args.recv_timeout_ms, + ) + + actual_steps = min(args.steps, trajectory_length) if args.steps else trajectory_length + predicted = np.full((actual_steps, action_dim), np.nan, dtype=np.float64) + target = np.array([s["action"] for s in rlds_steps[:actual_steps]], dtype=np.float64) + state_units = np.array([s["state"] for s in rlds_steps[:actual_steps]], dtype=np.float64) + valid = np.zeros((actual_steps, action_dim), dtype=np.bool_) + inference_points: list[int] = [] + instruction = rlds_steps[0]["instruction"] + + for t in range(0, actual_steps, args.execution_horizon): + step = rlds_steps[t] + observations = { + "observation.images.image": step["top"], + "observation.images.image2": step["wrist"], + "task": instruction, + "state": step["state"], + } + chunk = client._predict_chunk_octo(observations) # (20,7), already unnormalized + take = min(args.execution_horizon, actual_steps - t, chunk.shape[0]) + predicted[t:t + take] = chunk[:take, :action_dim] + valid[t:t + take] = True + inference_points.append(t) + print(f"t={t:3d} take={take} chunk[0]={chunk[0, :action_dim].round(4).tolist()}") + + n_calls = len(inference_points) + print(f"inference_calls={n_calls} evaluated_steps={actual_steps}") + + predicted_norm = zscore_normalize(predicted, action_stats) + target_norm = zscore_normalize(target, action_stats) + normalized_metrics = masked_metrics(predicted_norm, target_norm, valid) + original_metrics = masked_metrics(predicted, target, valid) + + stats_min = np.asarray(action_stats["min"], dtype=np.float64) + stats_max = np.asarray(action_stats["max"], dtype=np.float64) + close_threshold = float((stats_min[-1] + stats_max[-1]) / 2.0) + gripper_valid = valid[:, -1] + predicted_closed = predicted[:, -1] <= close_threshold + target_closed = target[:, -1] <= close_threshold + gripper_accuracy = float(np.mean(predicted_closed[gripper_valid] == target_closed[gripper_valid])) \ + if gripper_valid.any() else None + + output_dir = args.output_dir.expanduser() + split_output = output_dir / "train" + split_output.mkdir(parents=True, exist_ok=True) + stem = "trajectory_000" + plot_path = split_output / f"{stem}_gt_vs_pred.png" + trace_path = split_output / f"{stem}_actions.npz" + + plot_trajectory( + predicted=predicted, target=target, state=state_units, valid=valid, + inference_points=inference_points, + title=f"{model_label} | train trajectory {args.traj_index} | execution horizon {args.execution_horizon}", + output_path=plot_path, + ) + np.savez_compressed( + trace_path, + predicted_actions=predicted, + ground_truth_actions=target, + predicted_actions_normalized=predicted_norm, + ground_truth_actions_normalized=target_norm, + valid_action_mask=valid, + inference_points=np.asarray(inference_points, dtype=np.int32), + state_actions_units=state_units, + ) + + report = { + "split": "train", + "trajectory_index": args.traj_index, + "source_episode_id": args.traj_index, + "trajectory_length": trajectory_length, + "evaluated_steps": actual_steps, + "inference_calls": n_calls, + "execution_horizon": args.execution_horizon, + "instruction": instruction, + "normalized": normalized_metrics, + "original_units": original_metrics, + "gripper_accuracy": gripper_accuracy, + "plot": str(plot_path.resolve()), + "action_trace": str(trace_path.resolve()), + } + reports = [report] + agg = aggregate(reports) + summary = { + "protocol": "teacher_forced_open_loop_action_chunk_stitching", + "protocol_source": PROTOCOL_SOURCE, + "vla_addr": args.vla_addr, + "rlds_dir": str(args.rlds_dir.resolve()), + "dataset_statistics": str(args.dataset_statistics.resolve()), + "traj_index": args.traj_index, + "action_horizon": 20, + "execution_horizon": args.execution_horizon, + "requested_steps": args.steps, + "preprocess_cross_check_pass": preprocess_ok, + "aggregate": agg, + "split_aggregates": {"train": agg}, + "trajectories": reports, + } + (output_dir / "summary.json").write_text(json.dumps(summary, indent=2)) + + csv_path = output_dir / "trajectory_metrics.csv" + with csv_path.open("w", newline="") as stream: + fieldnames = ["split", "trajectory_index", "source_episode_id", "trajectory_length", + "evaluated_steps", "inference_calls", "execution_horizon", + "normalized_mae", "normalized_mse", "original_mae", "original_mse", + "gripper_accuracy", "plot", "action_trace"] + writer = csv.DictWriter(stream, fieldnames=fieldnames) + writer.writeheader() + writer.writerow({ + "split": report["split"], "trajectory_index": report["trajectory_index"], + "source_episode_id": report["source_episode_id"], "trajectory_length": report["trajectory_length"], + "evaluated_steps": report["evaluated_steps"], "inference_calls": report["inference_calls"], + "execution_horizon": report["execution_horizon"], + "normalized_mae": report["normalized"]["mae"], "normalized_mse": report["normalized"]["mse"], + "original_mae": report["original_units"]["mae"], "original_mse": report["original_units"]["mse"], + "gripper_accuracy": report["gripper_accuracy"], + "plot": report["plot"], "action_trace": report["action_trace"], + }) + + print(f"summary={output_dir / 'summary.json'}") + print(f"metrics_csv={csv_path}") + print(f"original: mae={original_metrics['mae']:.6f} mse={original_metrics['mse']:.6f} rmse={original_metrics['rmse']:.6f}") + print(f"normalized: mae={normalized_metrics['mae']:.6f} mse={normalized_metrics['mse']:.6f} rmse={normalized_metrics['rmse']:.6f}") + print(f"per_dim_mae={[round(v, 4) for v in original_metrics['per_dim_mae']]}") + print(f"gripper_accuracy={gripper_accuracy}") + + +if __name__ == "__main__": + main() 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 a5db4d9..1ba2947 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -55,6 +55,13 @@ "max_state_dim": 64, "trust_remote_code": True}, "gr00t_n1_6": {"image_size": 224, "tokenizer": None, "max_state_dim": 128, "trust_remote_code": True}, + + # cyrusneary/octo-finetuned-libero (window=1) is genuinely single-camera -- its own + # finetune_config.json image_obs_keys={"primary": "image"} never fed a wrist view, so + # OctoPipelineAdapter sends primary only; the server zero-fills+masks-invalid the wrist + # slot (matches training distribution exactly). No proprio/state input (Octo's + # observation_tokenizers are image-only, see octo_pretrain_config.py) -> max_state_dim=0. + "octo": {"image_size": 256, "tokenizer": "t5-base", "max_state_dim": 0, "max_length": 16}, } BITVLA_N_PATCHES_PER_VIEW = 256 @@ -509,6 +516,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]: @@ -826,6 +835,66 @@ 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: + # observations come from OctoPipelineAdapter.parse_observation (adapters.py): + # already rotate180+resize256/128'd uint8 HWC images (TIP-P), primary-only for the + # genuinely single-camera cyrusneary checkpoint (image2 sent too if present, for a + # future two-camera Octo checkpoint -- vla-server zero-fills+masks-invalid whichever + # view it doesn't receive, octo.cpp:predict()). + 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() + # Same recipe as the checkpoint's own text_processor (octo_pretrain_config.py): + # t5-base, max_length=16, padding="max_length", truncation=True. Server's T5 encoder + # requires both input_ids and attention_mask at exactly this length (octo.cpp + # rejects anything else) -- unlike most archs, Octo needs the real mask, not a + # server-derived one. + 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 + # Already UN-normalized (world units) -- octo.cpp:predict() un-normalizes + # server-side (VLA_OCTO_UNNORM_DATASET / auto-resolve), unlike most archs which + # return normalized actions for the client to un-normalize via --stats-json. See + # TIP-CLIENT report for why (dataset_statistics is embedded in the multi-hundred-MB + # checkpoint GGUF, not a small sibling file a client can cheaply hold). + 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/compare_action_dist.py b/scripts/compare_action_dist.py new file mode 100644 index 0000000..190e9a7 --- /dev/null +++ b/scripts/compare_action_dist.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); + +"""TIP-009: compare OctoPt vs vla.cpp free-sampled normalized-action distributions. + +Per element (28 = 4 action_horizon x 7 action_dim): + - mean/SE check: |mean_cpp - mean_pt| <= 4*SE, where SE = sqrt(std_pt^2/N_pt + + std_cpp^2/N_cpp) -- except near-constant elements (std_pt < 0.05), which use a + fixed |dmean| <= 0.02 threshold instead (SE-based 4*SE is unstable/too strict + when variance is tiny -- e.g. the gripper dimension). + - std-ratio check: only when std_pt > 1e-3 (near-zero variance makes a ratio + meaningless); std_cpp/std_pt in [0.8, 1.25]. + - KS two-sample test: skipped for near-constant elements ("gan hang" -> KS is + unreliable on near-degenerate distributions). Applied to the rest; overall + budget of at most 1/28 rejections at Bonferroni alpha = 0.05/28. +Also reports noise ~ N(0,1) sanity for both sides' recorded initial DDPM noise. +""" + +from __future__ import annotations + +import argparse +import math +from pathlib import Path + +import numpy as np + +try: + from scipy.stats import ks_2samp as _scipy_ks_2samp + HAVE_SCIPY = True +except ImportError: + HAVE_SCIPY = False + + +def ks_2samp_manual(a: np.ndarray, b: np.ndarray) -> tuple[float, float]: + """Two-sample two-sided KS test with the asymptotic Kolmogorov p-value + approximation (same formula scipy uses for `method="asymp"`). Fallback only -- + used when scipy isn't installed. + """ + a = np.sort(a) + b = np.sort(b) + n1, n2 = len(a), len(b) + all_vals = np.concatenate([a, b]) + cdf_a = np.searchsorted(a, all_vals, side="right") / n1 + cdf_b = np.searchsorted(b, all_vals, side="right") / n2 + d = float(np.max(np.abs(cdf_a - cdf_b))) + ne = n1 * n2 / (n1 + n2) + lam = (math.sqrt(ne) + 0.12 + 0.11 / math.sqrt(ne)) * d + p = 0.0 + for k in range(1, 101): + term = 2.0 * ((-1) ** (k - 1)) * math.exp(-2.0 * k * k * lam * lam) + p += term + if abs(term) < 1e-12: + break + p = max(0.0, min(1.0, p)) + return d, p + + +def ks_test(a: np.ndarray, b: np.ndarray) -> tuple[float, float]: + if HAVE_SCIPY: + res = _scipy_ks_2samp(a, b) + return float(res.statistic), float(res.pvalue) + return ks_2samp_manual(a, b) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--pt-samples", type=Path, required=True) + ap.add_argument("--cpp-samples", type=Path, required=True) + ap.add_argument("--pt-noise", type=Path) + ap.add_argument("--cpp-noise", type=Path) + ap.add_argument("--alpha", type=float, default=0.05) + ap.add_argument("--near-constant-std", type=float, default=0.05) + ap.add_argument("--near-constant-dmean-tol", type=float, default=0.02) + ap.add_argument("--std-ratio-min-std", type=float, default=1e-3) + ap.add_argument("--std-ratio-bounds", type=float, nargs=2, default=(0.8, 1.25)) + ap.add_argument("--max-ks-rejects", type=int, default=1) + args = ap.parse_args() + + pt = np.load(args.pt_samples).reshape(-1, 28).astype(np.float64) + cpp = np.load(args.cpp_samples).reshape(-1, 28).astype(np.float64) + n_pt, n_cpp = pt.shape[0], cpp.shape[0] + + n_dims = 28 + bonferroni_alpha = args.alpha / n_dims + + print(f"N_pt={n_pt} N_cpp={n_cpp} scipy={'yes' if HAVE_SCIPY else 'no (manual KS fallback)'}") + print(f"bonferroni alpha = {args.alpha}/{n_dims} = {bonferroni_alpha:.6g}") + print() + + noise_ok = True + if args.pt_noise and args.cpp_noise: + pt_noise = np.load(args.pt_noise).reshape(-1) + cpp_noise = np.load(args.cpp_noise).reshape(-1) + print("NOISE SANITY (initial DDPM noise, expected ~ N(0,1)):") + for name, arr in (("pt_noise", pt_noise), ("cpp_noise", cpp_noise)): + mean, std = float(arr.mean()), float(arr.std()) + ok = (-0.05 <= mean <= 0.05) and (0.95 <= std <= 1.05) + noise_ok = noise_ok and ok + print(f" {name}: n={arr.size} mean={mean:+.6f} std={std:.6f} {'PASS' if ok else 'FAIL'}") + print() + + action_names = [f"t{t}d{d}" for t in range(4) for d in range(7)] + + header = (f"{'elem':>5} {'mean_pt':>11} {'mean_cpp':>11} {'|dmean|':>9} {'SE':>9} " + f"{'std_pt':>9} {'std_cpp':>9} {'ratio':>7} {'KS_p':>9} {'ks?':>4} {'status':>7}") + print(header) + print("-" * len(header)) + + fail_dims = [] + ks_rejects = [] + ks_tested = 0 + + for k in range(n_dims): + pt_k = pt[:, k] + cpp_k = cpp[:, k] + mean_pt, mean_cpp = float(pt_k.mean()), float(cpp_k.mean()) + std_pt, std_cpp = float(pt_k.std(ddof=1)), float(cpp_k.std(ddof=1)) + dmean = abs(mean_cpp - mean_pt) + se = math.sqrt(std_pt ** 2 / n_pt + std_cpp ** 2 / n_cpp) + std_ratio = std_cpp / std_pt if std_pt > 1e-12 else float("nan") + + near_constant = std_pt < args.near_constant_std + mean_ok = (dmean <= args.near_constant_dmean_tol) if near_constant else (dmean <= 4 * se) + + std_ok = True + if std_pt > args.std_ratio_min_std: + lo, hi = args.std_ratio_bounds + std_ok = lo <= std_ratio <= hi + + ks_p = float("nan") + ks_applied = not near_constant + ks_reject = False + if ks_applied: + ks_tested += 1 + _, ks_p = ks_test(pt_k, cpp_k) + ks_reject = ks_p < bonferroni_alpha + if ks_reject: + ks_rejects.append(action_names[k]) + + status = "PASS" if (mean_ok and std_ok) else "FAIL" + if status == "FAIL": + fail_dims.append(action_names[k]) + + ks_flag = ("R" if ks_reject else ".") if ks_applied else "skip" + ks_p_s = f"{ks_p:9.6f}" if ks_applied else f"{'--':>9}" + print(f"{action_names[k]:>5} {mean_pt:11.6f} {mean_cpp:11.6f} {dmean:9.6f} {se:9.6f} " + f"{std_pt:9.6f} {std_cpp:9.6f} {std_ratio:7.3f} {ks_p_s} {ks_flag:>4} {status:>7}") + + print() + ks_ok = len(ks_rejects) <= args.max_ks_rejects + mean_std_ok = len(fail_dims) == 0 + verdict = "PASS" if (mean_std_ok and ks_ok and noise_ok) else "FAIL" + + print(f"mean/std per-element: {'PASS' if mean_std_ok else f'FAIL ({len(fail_dims)}/{n_dims}: {fail_dims})'}") + print(f"KS (tested {ks_tested}/{n_dims}, skipped near-constant): " + f"{len(ks_rejects)} reject(s) {ks_rejects} -> " + f"{'PASS' if ks_ok else 'FAIL'} (budget <= {args.max_ks_rejects})") + print(f"noise sanity: {'PASS' if noise_ok else 'FAIL'}") + print(f"VERDICT: {verdict}") + return 0 if verdict == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/compare_open_loop_vs_golden.py b/scripts/compare_open_loop_vs_golden.py new file mode 100644 index 0000000..bf6bdde --- /dev/null +++ b/scripts/compare_open_loop_vs_golden.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""TIP-08: direct comparison of vla.cpp's open-loop output (TIP-07) against the local +OctoPt golden trace (TIP-02) for the SAME RLDS trajectory -- the most direct parity check +available (same model weights, same input, only the inference engine differs). + +Compares npz arrays directly (predicted_actions, predicted_actions_normalized, +ground_truth_actions, valid_action_mask), recomputes MAE/MSE/RMSE/gripper/per-dim/ +normalized metrics for both sides with the identical formula (mirrors +evaluate_octo_open_loop.py's _masked_metrics, already duplicated once in +run_open_loop_octo.py for TIP-07), and renders a 7-subplot overlay plot. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib +matplotlib.use("Agg") +from matplotlib import pyplot as plt +import numpy as np + +ACTION_NAMES = ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"] + + +def masked_metrics(predicted: np.ndarray, target: np.ndarray, valid: np.ndarray) -> dict[str, Any]: + # Verbatim formula, TIP-02's evaluate_octo_open_loop.py::_masked_metrics. + count = int(valid.sum()) + difference = predicted - target + abs_error = np.abs(difference) + sq_error = np.square(difference) + dim_count = valid.sum(axis=0) + per_dim_mae = np.divide((abs_error * valid).sum(axis=0), dim_count, + out=np.full(predicted.shape[1], np.nan), where=dim_count > 0) + per_dim_mse = np.divide((sq_error * valid).sum(axis=0), dim_count, + out=np.full(predicted.shape[1], np.nan), where=dim_count > 0) + return { + "valid_action_values": count, + "mae": float((abs_error * valid).sum() / count), + "mse": float((sq_error * valid).sum() / count), + "rmse": float(math.sqrt((sq_error * valid).sum() / count)), + "per_dim_mae": per_dim_mae.tolist(), + "per_dim_mse": per_dim_mse.tolist(), + } + + +def gripper_accuracy(predicted: np.ndarray, target: np.ndarray, valid: np.ndarray, threshold: float) -> float: + gripper_valid = valid[:, -1] + predicted_closed = predicted[:, -1] <= threshold + target_closed = target[:, -1] <= threshold + return float(np.mean(predicted_closed[gripper_valid] == target_closed[gripper_valid])) + + +def max_abs_diff(a: np.ndarray, b: np.ndarray) -> float: + return float(np.abs(a.astype(np.float64) - b.astype(np.float64)).max()) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--vla-npz", type=Path, required=True) + ap.add_argument("--golden-npz", type=Path, required=True) + ap.add_argument("--golden-metrics-json", type=Path, default=None, + help="golden's own trajectory_000_metrics.json (published target)") + ap.add_argument("--dataset-statistics", type=Path, required=True) + ap.add_argument("--output-plot", type=Path, required=True) + args = ap.parse_args() + + vla = np.load(args.vla_npz.expanduser()) + golden = np.load(args.golden_npz.expanduser()) + + action_stats = json.loads(args.dataset_statistics.expanduser().read_text())["action"] + stats_min = np.asarray(action_stats["min"], dtype=np.float64) + stats_max = np.asarray(action_stats["max"], dtype=np.float64) + gripper_threshold = float((stats_min[-1] + stats_max[-1]) / 2.0) + + print("=== 1. Direct array comparison (vla.cpp vs local golden) ===") + print(f"{'array':<32}{'max|diff|':>14}") + array_diffs = {} + for key in ("predicted_actions", "predicted_actions_normalized", "ground_truth_actions", "valid_action_mask"): + diff = max_abs_diff(vla[key], golden[key]) + array_diffs[key] = diff + print(f"{key:<32}{diff:>14.8f}") + + valid = golden["valid_action_mask"].astype(bool) + vla_metrics = masked_metrics(vla["predicted_actions"], vla["ground_truth_actions"], valid) + vla_metrics_norm = masked_metrics(vla["predicted_actions_normalized"], vla["ground_truth_actions_normalized"], valid) + vla_gripper = gripper_accuracy(vla["predicted_actions"], vla["ground_truth_actions"], valid, gripper_threshold) + + golden_metrics = masked_metrics(golden["predicted_actions"], golden["ground_truth_actions"], valid) + golden_metrics_norm = masked_metrics(golden["predicted_actions_normalized"], golden["ground_truth_actions_normalized"], valid) + golden_gripper = gripper_accuracy(golden["predicted_actions"], golden["ground_truth_actions"], valid, gripper_threshold) + + published = None + if args.golden_metrics_json is not None: + published = json.loads(args.golden_metrics_json.expanduser().read_text()) + + print() + print("=== 2. Metric comparison: vla.cpp vs local golden (recomputed) vs published target ===") + header = f"{'metric':<20}{'vla.cpp':>14}{'golden(npz)':>14}{'published':>14}{'|vla-golden|':>15}{'|vla-pub|':>13}" + print(header) + rows = [ + ("original.mae", vla_metrics["mae"], golden_metrics["mae"], + published["original_units"]["mae"] if published else None), + ("original.mse", vla_metrics["mse"], golden_metrics["mse"], + published["original_units"]["mse"] if published else None), + ("original.rmse", vla_metrics["rmse"], golden_metrics["rmse"], + published["original_units"]["rmse"] if published else None), + ("normalized.mae", vla_metrics_norm["mae"], golden_metrics_norm["mae"], + published["normalized"]["mae"] if published else None), + ("normalized.mse", vla_metrics_norm["mse"], golden_metrics_norm["mse"], + published["normalized"]["mse"] if published else None), + ("normalized.rmse", vla_metrics_norm["rmse"], golden_metrics_norm["rmse"], + published["normalized"]["rmse"] if published else None), + ("gripper_accuracy", vla_gripper, golden_gripper, + published["gripper_accuracy"] if published else None), + ] + for name, v, g, p in rows: + d_vg = abs(v - g) + d_vp = abs(v - p) if p is not None else float("nan") + p_str = f"{p:14.6f}" if p is not None else f"{'n/a':>14}" + print(f"{name:<20}{v:14.6f}{g:14.6f}{p_str}{d_vg:15.8f}{d_vp:13.8f}") + + print() + print("per_dim_mae (original units):") + print(f"{'dim':<10}{'vla.cpp':>12}{'golden':>12}{'|diff|':>12}") + for i, name in enumerate(ACTION_NAMES): + v = vla_metrics["per_dim_mae"][i] + g = golden_metrics["per_dim_mae"][i] + print(f"{name:<10}{v:12.6f}{g:12.6f}{abs(v - g):12.8f}") + + # 3. Overlay plot: GT, vla.cpp pred, golden pred per dim. + action_dim = vla["ground_truth_actions"].shape[1] + fig, axes = plt.subplots(action_dim, 1, figsize=(12, max(4, 2.6 * action_dim)), sharex=True, dpi=140) + timesteps = np.arange(vla["ground_truth_actions"].shape[0]) + for dim, ax in enumerate(axes): + gt = np.where(valid[:, dim], vla["ground_truth_actions"][:, dim], np.nan) + pred_vla = np.where(valid[:, dim], vla["predicted_actions"][:, dim], np.nan) + pred_golden = np.where(valid[:, dim], golden["predicted_actions"][:, dim], np.nan) + ax.plot(timesteps, gt, color="#0072B2", linewidth=2.1, label="GT action", zorder=3) + ax.plot(timesteps, pred_golden, color="#009E73", linewidth=1.6, label="pred (OctoPt golden)", + zorder=2, alpha=0.9) + ax.plot(timesteps, pred_vla, color="#D55E00", linewidth=1.4, linestyle="--", + label="pred (vla.cpp)", zorder=4, alpha=0.9) + ax.set_ylabel(ACTION_NAMES[dim]) + ax.grid(True, color="#D0D0D0", alpha=0.45, linewidth=0.7) + if dim == 0: + ax.legend(loc="upper right", ncol=3, framealpha=0.95) + axes[-1].set_xlabel("dataset timestep") + fig.suptitle("vla.cpp vs OctoPt golden -- open-loop prediction overlay (TIP-08)") + fig.tight_layout(rect=(0, 0, 1, 0.98)) + args.output_plot.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.output_plot) + plt.close(fig) + print() + print(f"overlay plot: {args.output_plot}") + + # Gate check. + ok = True + if array_diffs["predicted_actions"] >= 1e-4: + ok = False + if array_diffs["ground_truth_actions"] != 0.0 or array_diffs["valid_action_mask"] != 0.0: + ok = False + print() + print(f"G8 array-diff gate: {'PASS' if ok else 'FAIL'}") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/convert_octo_to_gguf.py b/scripts/convert_octo_to_gguf.py new file mode 100644 index 0000000..11edae2 --- /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, + "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/scripts/patch_libero_golden_for_harness.py b/scripts/patch_libero_golden_for_harness.py new file mode 100644 index 0000000..6fcefe6 --- /dev/null +++ b/scripts/patch_libero_golden_for_harness.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Add zero-filled wrist placeholder tensors to single-camera LIBERO golden cases. + +octo_parity_dump's loader (octo_dump_tokenizer_case in src/models/octo.cpp) hard-requires +tensors/input.observation.image_wrist.npy and .../pad_mask_dict.image_wrist.npy to exist in +any golden case directory it reads (unlike the task-side wrist tensor, which has a zero-fill +fallback keyed off the observation tensor's own shape). The cyrusneary LIBERO checkpoint +(TIP-GOLD) is genuinely single-camera -- its golden traces never ran a wrist observation +through the model at all, so these files don't exist there. + +This script adds all-zero / all-False placeholders so the C++ loader can read the case +directory without crashing. It does not touch any tensor that was actually produced by the +model (Tier B trace values are untouched). The wrist-derived boundaries this unblocks +(obs.wrist.*, bt.obs_wrist, bt.input, bt.mask, bt.output) are excluded from parity comparison +by verify_octo_parity.py's --exclude flag for LIBERO cases -- see TIP-HARNESS report for why. + +Idempotent: skips files that already exist. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +WRIST_IMAGE_SIZE = 128 # matches the bridge golden's own input.observation.image_wrist.npy + + +def patch_case(case_dir: Path) -> list[str]: + tensors_dir = case_dir / "tensors" + primary_path = tensors_dir / "input.observation.image_primary.npy" + if not primary_path.exists(): + raise SystemExit(f"missing {primary_path}; not a valid golden case dir") + primary = np.load(primary_path) + if primary.ndim != 5: + raise SystemExit(f"expected 5D image_primary [B,T,C,H,W], got {primary.shape}") + batch, window = primary.shape[0], primary.shape[1] + + written = [] + wrist_img_path = tensors_dir / "input.observation.image_wrist.npy" + if not wrist_img_path.exists(): + wrist = np.zeros((batch, window, 3, WRIST_IMAGE_SIZE, WRIST_IMAGE_SIZE), dtype=np.uint8) + np.save(wrist_img_path, wrist) + written.append(str(wrist_img_path)) + + wrist_mask_path = tensors_dir / "input.observation.pad_mask_dict.image_wrist.npy" + if not wrist_mask_path.exists(): + mask = np.zeros((batch, window), dtype=bool) + np.save(wrist_mask_path, mask) + written.append(str(wrist_mask_path)) + + manifest_path = case_dir / "manifest.json" + if manifest_path.exists() and written: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + note = manifest.setdefault("metadata", {}).setdefault("harness_patch", {}) + note["patched_wrist_placeholder"] = True + note["reason"] = ( + "checkpoint is single-camera; placeholder added only so octo_parity_dump's " + "loader can read the case dir, excluded from comparison via --exclude" + ) + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8") + + return written + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("golden_dir", type=Path, help="LIBERO golden root (contains / subdirs)") + ap.add_argument("--case", action="append", default=[], help="case name(s); default: all subdirs") + args = ap.parse_args() + + cases = args.case or [p.name for p in args.golden_dir.iterdir() if (p / "tensors").is_dir()] + for case in cases: + written = patch_case(args.golden_dir / case) + if written: + print(f"{case}: wrote {written}") + else: + print(f"{case}: already patched (nothing to do)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/sample_action_dist_pt.py b/scripts/sample_action_dist_pt.py new file mode 100644 index 0000000..4c6b07b --- /dev/null +++ b/scripts/sample_action_dist_pt.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); + +"""TIP-009: OctoPt reference sampler for statistical action-distribution parity. + +Loads a golden case's OWN saved observation/task tensors directly (dataset-agnostic: +whatever input.observation.*/input.task.* keys exist in the case's manifest.json are +loaded, none are assumed present -- e.g. bridge_debug has no +input.task.image_primary/wrist at all, since it's a real dataset with a language-only +task space; those keys are simply absent from `tasks`, matching what OctoPt's own +dump_octo_golden_trace_bridge_debug_pt.py passes to sample_actions() for that case). + +Calls model.sample_actions(...) n_samples times with fresh per-sample noise (a +torch.Generator seeded seed+i each call, matching vla.cpp's std::mt19937(seed+i)), +collecting the normalized [N,4,7] action and the [N,2,28] initial DDPM noise actually +consumed each call (captured via a minimal trace shim -- DiffusionActionHeadPt. +predict_action already calls trace.write("action_head.predict_action.initial_noise", +current_x) unconditionally when a trace object is passed). +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Dict, Optional + +import numpy as np +import torch + + +def _load_tensor(case_dir: Path, manifest: dict, key: str) -> torch.Tensor: + info = manifest["tensors"][key] + arr = np.load(case_dir / info["file"]) + return torch.from_numpy(arr) + + +def _build_tree(case_dir: Path, manifest: dict, prefix: str) -> Dict[str, Any]: + tree: Dict[str, Any] = {} + full_prefix = prefix + "." + for key in manifest["tensors"]: + if not key.startswith(full_prefix): + continue + rel = key[len(full_prefix):] + parts = rel.split(".") + node = tree + for p in parts[:-1]: + node = node.setdefault(p, {}) + node[parts[-1]] = _load_tensor(case_dir, manifest, key) + return tree + + +def _to_device(tree, device: str): + if isinstance(tree, dict): + return {k: _to_device(v, device) for k, v in tree.items()} + return tree.to(device) + + +class NoiseCapture: + """Duck-typed trace shim. DiffusionActionHeadPt.predict_action calls + trace.write(name, value) for many boundaries during sampling; we keep only the + one we need and no-op everything else (including write_tree, called from the + transformer forward pass for unrelated boundaries). + """ + + def __init__(self) -> None: + self.initial_noise: Optional[torch.Tensor] = None + + def write(self, name: str, value) -> None: + if name == "action_head.predict_action.initial_noise": + self.initial_noise = value.detach().cpu().clone() + + def write_tree(self, *_args, **_kwargs) -> None: + pass + + +def main() -> int: + ap = argparse.ArgumentParser(description="OctoPt free-sample reference for TIP-009 distribution parity.") + ap.add_argument("--case", required=True, type=Path, help="golden case directory (contains manifest.json + tensors/)") + ap.add_argument("--checkpoint", default="hf://rail-berkeley/octo-small-1.5") + ap.add_argument("--octo-root", type=Path, default=Path(__file__).resolve().parents[1] / "octo-pytorch") + ap.add_argument("--n-samples", type=int, default=50) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--device", default="auto") + ap.add_argument("--out-samples", type=Path, default=Path("pt_samples.npy")) + ap.add_argument("--out-noise", type=Path, default=Path("pt_noise.npy")) + 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 # noqa: E402 (needs sys.path insert first) + + if args.device == "auto": + args.device = "cuda:0" if torch.cuda.is_available() else "cpu" + + manifest = json.loads((args.case / "manifest.json").read_text(encoding="utf-8")) + + observations = _to_device(_build_tree(args.case, manifest, "input.observation"), args.device) + tasks = _to_device(_build_tree(args.case, manifest, "input.task"), args.device) + print(f"observation keys: {sorted(observations.keys())}") + print(f"task keys: {sorted(tasks.keys())}") + if "image_primary" not in tasks: + print("note: no input.task.image_primary in this case (language-only task space) " + "-- passing tasks dict as-is, matching how this case's own golden trace was generated") + + timestep_pad_mask = observations["timestep_pad_mask"] + + print(f"loading {args.checkpoint} on {args.device} ...") + loaded = OctoModelPt.load_pretrained_from_jax(args.checkpoint, skip_keys_regex=".*hf_model") + model = loaded["octo_model"].to(args.device).eval() + + samples = np.zeros((args.n_samples, 4, 7), dtype=np.float32) + noise = np.zeros((args.n_samples, 2, 28), dtype=np.float32) + + with torch.no_grad(): + for i in range(args.n_samples): + torch.manual_seed(args.seed + i) # belt-and-suspenders; sampling itself is generator-scoped + cap = NoiseCapture() + gen = torch.Generator(device=args.device).manual_seed(args.seed + i) + action = model.sample_actions( + observations, + tasks, + timestep_pad_mask=timestep_pad_mask, + train=False, + generator=gen, + trace=cap, + ) + samples[i] = action.detach().cpu().numpy().reshape(4, 7) + if cap.initial_noise is None: + raise RuntimeError("trace shim did not capture initial_noise -- predict_action's trace contract changed") + noise[i] = cap.initial_noise.numpy().reshape(2, 28) + if (i + 1) % 10 == 0 or i == args.n_samples - 1: + print(f" sample {i + 1}/{args.n_samples}") + + args.out_samples.parent.mkdir(parents=True, exist_ok=True) + np.save(args.out_samples, samples) + np.save(args.out_noise, noise) + print(f"wrote {args.out_samples} {samples.shape}, {args.out_noise} {noise.shape}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/verify_octo_l1_parity.py b/scripts/verify_octo_l1_parity.py new file mode 100644 index 0000000..6eaad8b --- /dev/null +++ b/scripts/verify_octo_l1_parity.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""TIP-06: compare octo_l1_parity_dump's T0-T5 C++ activations against the OctoPt stagewise +golden dump (octo-pytorch-kamusarj's scripts/dump_l1_stagewise_golden.py). Mirrors +verify_octo_parity.py's .npy/.f32 parsing (no numpy dependency -- ctest's Python3_EXECUTABLE +may be a bare interpreter) but is deliberately standalone/simpler: this golden format has no +dataset-name wrapper or boundary-renaming map, since both sides were designed together and +use the SAME stage names (t0.proprio_normalized, t0.proprio_tokens, t1t2.readout_action, +t3.map_attn_out, t3.map_emb, t4.mean_normalized, t5.action_unnormalized). +""" + +from __future__ import annotations + +import argparse +import ast +import array +import math +from pathlib import Path + +STAGES = ( + "t0.proprio_normalized", + "t0.proprio_tokens", + "t1t2.readout_action", + "t3.map_attn_out", + "t3.map_emb", + "t4.mean_normalized", + "t5.action_unnormalized", +) + + +def load_dump_manifest(path: Path) -> dict[str, tuple[Path, str, tuple[int, ...]]]: + rows: dict[str, tuple[Path, str, tuple[int, ...]]] = {} + for raw in path.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if not line: + continue + parts = line.split() + if len(parts) < 4 or parts[2] != "float32": + raise ValueError(f"bad dump manifest line: {raw}") + rows[parts[0]] = (Path(parts[1]), parts[2], tuple(int(x) for x in parts[3:])) + return rows + + +def load_npy_float32(path: Path) -> tuple[tuple[int, ...], array.array]: + raw = path.read_bytes() + if len(raw) < 16 or raw[:6] != b"\x93NUMPY": + raise ValueError(f"not a .npy file: {path}") + major = raw[6] + pos = 8 + if major == 1: + hlen = int.from_bytes(raw[pos : pos + 2], "little") + pos += 2 + elif major in (2, 3): + hlen = int.from_bytes(raw[pos : pos + 4], "little") + pos += 4 + else: + raise ValueError(f"unsupported .npy version {major}: {path}") + header = raw[pos : pos + hlen].decode("latin1").strip() + pos += hlen + meta = ast.literal_eval(header) + if meta.get("descr") not in (" array.array: + data = array.array("f") + with path.open("rb") as f: + data.fromfile(f, n) + if len(data) != n: + raise ValueError(f"truncated dump payload: {path}") + return data + + +def max_abs_diff(a: array.array, b: array.array) -> float: + return max((abs(float(x) - float(y)) for x, y in zip(a, b)), default=0.0) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--golden", required=True, type=Path, help="ep0_t dir from dump_l1_stagewise_golden.py") + ap.add_argument("--dump", required=True, type=Path, help="dump dir containing manifest.txt (octo_l1_parity_dump --out)") + ap.add_argument("--tol", type=float, default=1e-4) + args = ap.parse_args() + + dump = load_dump_manifest(args.dump / "manifest.txt") + + print("stage\tgolden_shape\tdump_shape\tmax_abs_diff\tstatus") + ok = True + for name in STAGES: + golden_path = args.golden / f"{name}.npy" + if not golden_path.exists(): + raise SystemExit(f"missing golden file: {golden_path}") + if name not in dump: + raise SystemExit(f"missing dump boundary: {name} (dump manifest={sorted(dump)})") + golden_shape, golden_data = load_npy_float32(golden_path) + dump_path, _, dump_shape = dump[name] + dump_n = math.prod(dump_shape) if dump_shape else 1 + golden_n = math.prod(golden_shape) if golden_shape else 1 + if dump_n != golden_n: + raise SystemExit( + f"{name}: element count mismatch dump={dump_shape}({dump_n}) golden={golden_shape}({golden_n})") + # manifest.txt already stores a directly-usable path (write_f32_dump wrote + # "/.f32" verbatim) -- use it as-is, don't re-prefix with --dump. + dump_data = load_raw_float32(dump_path, dump_n) + diff = max_abs_diff(dump_data, golden_data) + status = "PASS" if diff < args.tol else "FAIL" + if status == "FAIL": + ok = False + print(f"{name}\t{golden_shape}\t{dump_shape}\t{diff:.8f}\t{status}") + + if not ok: + print(f"FAILED: one or more stages exceeded tol={args.tol}") + return 1 + print(f"PASS: all {len(STAGES)} stages < tol={args.tol}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/verify_octo_parity.py b/scripts/verify_octo_parity.py new file mode 100644 index 0000000..b65c3f1 --- /dev/null +++ b/scripts/verify_octo_parity.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +"""Compare Octo tokenizer boundary dumps against npy+manifest golden traces.""" + +from __future__ import annotations + +import argparse +import ast +import array +import json +import math +from pathlib import Path + + +BOUNDARY_MAP = { + "obs.primary.tok": "octo_transformer.obs_primary.tokens_after_tokenizer", + "obs.primary.proj": "octo_transformer.obs_primary.tokens_after_projection", + "obs.primary.pos": "octo_transformer.obs_primary.tokens_after_pos_embedding", + "obs.wrist.tok": "octo_transformer.obs_wrist.tokens_after_tokenizer", + "obs.wrist.proj": "octo_transformer.obs_wrist.tokens_after_projection", + "obs.wrist.pos": "octo_transformer.obs_wrist.tokens_after_pos_embedding", +} + +T5_BOUNDARY_MAP = { + "t5.out": "octo_transformer.task_language.tokens_after_tokenizer", +} + +LANGUAGE_BOUNDARY_MAP = { + "lang.proj": "octo_transformer.task_language.tokens_after_projection", + "lang.pos": "octo_transformer.task_language.tokens_after_pos_embedding", + "repeated_language": "octo_transformer.obs_task_language.tokens_repeated_task", +} + +TRANSFORMER_BOUNDARY_MAP = { + "bt.input": "block_transformer.input_tokens", + "bt.mask": "block_transformer.attention_mask", + "bt.output": "block_transformer.output_tokens", + "bt.task_language": "block_transformer.prefix_output.task_language.tokens", + "bt.obs_primary": "block_transformer.timestep_output.obs_primary.tokens", + "bt.obs_wrist": "block_transformer.timestep_output.obs_wrist.tokens", + "bt.obs_task_language": "block_transformer.timestep_output.obs_task_language.tokens", + "bt.readout_action": "block_transformer.timestep_output.readout_action.tokens", +} + +DIFFUSION_BOUNDARY_MAP = { + "diff.initial_noise": "action_head.predict_action.initial_noise", + "diff.action_mask": "action_head.predict_action.action_mask", + "diff.flat_action_mask": "action_head.predict_action.flat_action_mask", + "diff.actions_all_timesteps": "action_head.predict_action.actions_all_timesteps", + "sample_actions.final_action_normalized": "sample_actions.final_action_normalized", + "action_final": "final_action", +} + +# The golden unnormalized-action tensor is named differently per case (tier1 uses +# "final_action_unnormalized"; tier2/bridge_debug uses +# "final_action_unnormalized_bridge_debug") -- resolved dynamically per golden +# manifest rather than hardcoded, since the tensor name isn't a fixed contract. +UNNORM_BOUNDARY_CANDIDATES = ("final_action_unnormalized", "final_action_unnormalized_bridge_debug") + +for _step in range(20): + _time = 19 - _step + _prefix = f"action_head.predict_action.step_{_step:02d}.t_{_time:02d}" + DIFFUSION_BOUNDARY_MAP.update({ + f"diff.step{_step:02d}.current_x_before": f"{_prefix}.current_x_before", + f"diff.step{_step:02d}.pred_eps": f"{_prefix}.pred_eps", + f"diff.step{_step:02d}.z": f"{_prefix}.z", + f"diff.step{_step:02d}.x_after_denoise": f"{_prefix}.current_x_after_denoise", + f"diff.step{_step:02d}.x_after_noise_add": f"{_prefix}.current_x_after_noise_add", + f"diff.step{_step:02d}.x_after_clip": f"{_prefix}.current_x_after_clip", + f"diff.step{_step:02d}.x_after_mask": f"{_prefix}.current_x_after_mask", + }) + + +def load_dump_manifest(path: Path) -> dict[str, tuple[Path, str, tuple[int, ...]]]: + rows: dict[str, tuple[Path, str, tuple[int, ...]]] = {} + for raw in path.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if not line: + continue + parts = line.split() + if len(parts) < 4 or parts[2] not in ("float32", "bool"): + raise ValueError(f"bad dump manifest line: {raw}") + rows[parts[0]] = (Path(parts[1]), parts[2], tuple(int(x) for x in parts[3:])) + return rows + + +def load_npy_float32(path: Path) -> tuple[tuple[int, ...], array.array]: + raw = path.read_bytes() + if len(raw) < 16 or raw[:6] != b"\x93NUMPY": + raise ValueError(f"not a .npy file: {path}") + major = raw[6] + pos = 8 + if major == 1: + hlen = int.from_bytes(raw[pos : pos + 2], "little") + pos += 2 + elif major in (2, 3): + hlen = int.from_bytes(raw[pos : pos + 4], "little") + pos += 4 + else: + raise ValueError(f"unsupported .npy version {major}: {path}") + header = raw[pos : pos + hlen].decode("latin1").strip() + pos += hlen + meta = ast.literal_eval(header) + if meta.get("descr") not in (" array.array: + n = math.prod(shape) + data = array.array("f") + with path.open("rb") as f: + data.fromfile(f, n) + if len(data) != n: + raise ValueError(f"truncated dump payload: {path}") + return data + + +def load_npy_bool(path: Path) -> tuple[tuple[int, ...], array.array]: + raw = path.read_bytes() + if len(raw) < 16 or raw[:6] != b"\x93NUMPY": + raise ValueError(f"not a .npy file: {path}") + major = raw[6] + pos = 8 + if major == 1: + hlen = int.from_bytes(raw[pos : pos + 2], "little") + pos += 2 + elif major in (2, 3): + hlen = int.from_bytes(raw[pos : pos + 4], "little") + pos += 4 + else: + raise ValueError(f"unsupported .npy version {major}: {path}") + meta = ast.literal_eval(raw[pos : pos + hlen].decode("latin1").strip()) + pos += hlen + if meta.get("descr") != "|b1" or meta.get("fortran_order"): + raise ValueError(f"expected C-order bool .npy: {path}") + shape = tuple(int(x) for x in meta["shape"]) + n = math.prod(shape) + data = array.array("B", raw[pos : pos + n]) + if len(data) != n: + raise ValueError(f"truncated .npy payload: {path}") + return shape, data + + +def load_raw_bool(path: Path, shape: tuple[int, ...]) -> array.array: + n = math.prod(shape) + data = array.array("B") + with path.open("rb") as f: + data.fromfile(f, n) + if len(data) != n: + raise ValueError(f"truncated dump payload: {path}") + return data + + +def stats(actual: array.array, golden: array.array) -> tuple[float, float, float]: + max_abs = 0.0 + ref_max = 0.0 + dot = 0.0 + aa = 0.0 + bb = 0.0 + for a, b in zip(actual, golden): + af = float(a) + bf = float(b) + diff = abs(af - bf) + max_abs = max(max_abs, diff) + ref_max = max(ref_max, abs(bf)) + dot += af * bf + aa += af * af + bb += bf * bf + max_rel = max_abs / max(ref_max, 1e-12) + cos = dot / (math.sqrt(aa) * math.sqrt(bb)) if aa and bb else 1.0 + return max_abs, max_rel, cos + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--golden", required=True, type=Path, help="golden case directory containing manifest.json") + ap.add_argument("--dump", required=True, type=Path, help="VLA_OCTO_DUMP directory containing manifest.txt") + ap.add_argument("--oracle-dump", type=Path, help="optional prior oracle dump directory containing manifest.txt") + ap.add_argument("--tol", type=float, default=1e-3) + ap.add_argument("--transformer-tol", type=float, default=2e-3) + ap.add_argument("--oracle-tol", type=float, default=1e-4) + ap.add_argument("--report", type=Path) + ap.add_argument( + "--known-mismatch", action="append", default=[], metavar="BOUNDARY:REASON", + help="boundary that is allowed to FAIL without failing the overall run (printed and " + "recorded as FAIL_KNOWN in the report, never silently dropped); repeatable") + ap.add_argument( + "--exclude", action="append", default=[], metavar="TOKEN[,TOKEN...]", + help="boundary(es) to drop entirely before comparison (not printed, not counted as " + "FAIL) -- for cases whose golden legitimately doesn't have a matching tensor, e.g. " + "single-camera LIBERO golden omitting wrist. A boundary is excluded if TOKEN is an " + "exact match, or a substring of the boundary name (checked with both '_' and '.' " + "forms of TOKEN, so 'obs_wrist' matches 'obs.wrist.tok' and 'bt.obs_wrist' alike). " + "Repeatable and/or comma-separated. Unlike --known-mismatch, an excluded boundary " + "is never looked up in dump/golden at all, so it can't raise a missing-boundary " + "SystemExit either.") + args = ap.parse_args() + known_mismatch = {} + for spec in args.known_mismatch: + boundary, _, reason = spec.partition(":") + known_mismatch[boundary] = reason or "(no reason given)" + exclude_tokens: list[str] = [] + for spec in args.exclude: + exclude_tokens.extend(t.strip() for t in spec.split(",") if t.strip()) + + def is_excluded(name: str) -> bool: + return any( + name == token or token in name or token.replace("_", ".") in name + for token in exclude_tokens + ) + + golden_manifest = json.loads((args.golden / "manifest.json").read_text(encoding="utf-8")) + if golden_manifest.get("format") != "npy+manifest.v1": + raise SystemExit(f"unsupported golden format: {golden_manifest.get('format')}") + if golden_manifest.get("layout") != "row-major": + raise SystemExit(f"unsupported golden layout: {golden_manifest.get('layout')}") + tensors = golden_manifest["tensors"] + dump = load_dump_manifest(args.dump / "manifest.txt") + oracle = load_dump_manifest(args.oracle_dump / "manifest.txt") if args.oracle_dump else None + + rows = [] + ok = True + print("boundary\tgolden\tshape\tmax_abs_err\tmax_rel_err\tcosine\tstatus\toracle_max_rel_err\toracle_status") + boundary_map = dict(BOUNDARY_MAP) + if any(name in dump for name in T5_BOUNDARY_MAP): + boundary_map.update(T5_BOUNDARY_MAP) + if any(name in dump for name in LANGUAGE_BOUNDARY_MAP): + boundary_map.update(LANGUAGE_BOUNDARY_MAP) + if any(name in dump for name in TRANSFORMER_BOUNDARY_MAP): + boundary_map.update(TRANSFORMER_BOUNDARY_MAP) + if any(name in dump for name in DIFFUSION_BOUNDARY_MAP): + boundary_map.update(DIFFUSION_BOUNDARY_MAP) + if "action_final_unnormalized" in dump: + matches = [name for name in UNNORM_BOUNDARY_CANDIDATES if name in tensors] + if not matches: + raise SystemExit( + f"action_final_unnormalized dumped but golden has none of {UNNORM_BOUNDARY_CANDIDATES}") + boundary_map["action_final_unnormalized"] = matches[0] + + excluded = [name for name in boundary_map if is_excluded(name)] + for name in excluded: + del boundary_map[name] + for name in excluded: + print(f"{name}\t\t\t\t\t\tEXCLUDED\t\tSKIP") + rows.append({"boundary": name, "golden": None, "status": "EXCLUDED"}) + + for dump_name, golden_name in boundary_map.items(): + if dump_name not in dump: + raise SystemExit(f"missing dump boundary: {dump_name}") + if golden_name not in tensors: + raise SystemExit(f"missing golden boundary: {golden_name}") + dump_path, dump_dtype, dump_shape = dump[dump_name] + ginfo = tensors[golden_name] + if dump_dtype == "bool": + golden_shape, golden = load_npy_bool(args.golden / ginfo["file"]) + actual = load_raw_bool(dump_path, dump_shape) + else: + golden_shape, golden = load_npy_float32(args.golden / ginfo["file"]) + actual = load_raw_float32(dump_path, dump_shape) + if golden_shape != dump_shape: + raise SystemExit(f"shape mismatch {dump_name}: dump={dump_shape} golden={golden_shape}") + max_abs, max_rel, cos = stats(actual, golden) + is_transformer_tol = (dump_name.startswith("bt.") and dump_name not in ("bt.input", "bt.mask")) or dump_name == "t5.out" + boundary_tol = args.transformer_tol if is_transformer_tol else args.tol + if dump_name.startswith("diff.") or dump_name in ("action_final", "action_final_unnormalized") or dump_name.startswith("sample_actions."): + status = "PASS" if (max_abs == 0.0 if dump_dtype == "bool" else max_abs <= boundary_tol) else "FAIL" + else: + status = "PASS" if (max_abs == 0.0 if dump_dtype == "bool" else max_rel <= boundary_tol) else "FAIL" + oracle_max_rel = None + oracle_status = "SKIP" + if oracle is not None: + if dump_name not in oracle: + raise SystemExit(f"missing oracle boundary: {dump_name}") + oracle_path, oracle_dtype, oracle_shape = oracle[dump_name] + if oracle_shape != dump_shape: + raise SystemExit(f"oracle shape mismatch {dump_name}: dump={dump_shape} oracle={oracle_shape}") + if oracle_dtype != dump_dtype: + raise SystemExit(f"oracle dtype mismatch {dump_name}: dump={dump_dtype} oracle={oracle_dtype}") + oracle_data = load_raw_bool(oracle_path, oracle_shape) if dump_dtype == "bool" else load_raw_float32(oracle_path, oracle_shape) + _, oracle_max_rel, _ = stats(actual, oracle_data) + oracle_status = "PASS" if oracle_max_rel <= args.oracle_tol else "FAIL" + known_reason = None + if status == "FAIL" and dump_name in known_mismatch: + known_reason = known_mismatch[dump_name] + status = "FAIL_KNOWN" + ok = ok and status in ("PASS", "FAIL_KNOWN") and oracle_status != "FAIL" + row = { + "boundary": dump_name, + "golden": golden_name, + "shape": list(dump_shape), + "max_abs_err": max_abs, + "max_rel_err": max_rel, + "cosine": cos, + "status": status, + "tolerance": 0.0 if dump_dtype == "bool" else boundary_tol, + "oracle_max_rel_err": oracle_max_rel, + "oracle_status": oracle_status, + } + if known_reason is not None: + row["known_mismatch_reason"] = known_reason + rows.append(row) + oracle_rel_s = "" if oracle_max_rel is None else f"{oracle_max_rel:.6g}" + status_s = status if known_reason is None else f"{status} ({known_reason})" + print(f"{dump_name}\t{golden_name}\t{list(dump_shape)}\t{max_abs:.6g}\t{max_rel:.6g}\t{cos:.9f}\t{status_s}\t{oracle_rel_s}\t{oracle_status}") + + for i in range(12): + dump_name = f"bt.blk{i}.out" + if dump_name in dump: + _, _, dump_shape = dump[dump_name] + print(f"{dump_name}\t\t{list(dump_shape)}\t\t\t\tUNVERIFIED\t\tSKIP") + rows.append({ + "boundary": dump_name, + "golden": None, + "shape": list(dump_shape), + "status": "UNVERIFIED_NO_GOLDEN", + }) + + if args.report: + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps({"ok": ok, "tol": args.tol, "transformer_tol": args.transformer_tol, "oracle_tol": args.oracle_tol, "rows": rows}, indent=2), encoding="utf-8") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/arch.h b/src/arch.h index 633e8ff..1f0c19e 100644 --- a/src/arch.h +++ b/src/arch.h @@ -55,6 +55,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 (M0 GGUF load only). 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. @@ -151,6 +152,15 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, const std::string& ckpt_path, const std::string& config_path); +/** + * @brief Build an Octo model. TIP-001 M0 supports GGUF load/shape validation, + * not inference. + * @copydetails smolvla_create + */ +std::unique_ptr octo_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string& config_path); + /** * @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 b99ba3c..6fc179e 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -66,6 +66,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) || @@ -77,6 +78,8 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { else if (arch_str == "gr00t_n1_5") { *out = Arch::GR00T_N1_5; ok = true; } else if (arch_str == "gr00t_n1_6") { *out = Arch::GR00T_N1_6; ok = true; } else if (arch_str == "gr00t_n1_7") { *out = Arch::GR00T_N1_7; ok = true; } + else if (arch_str == "octo" || + arch_str == "octo-small-1.5") { *out = Arch::OCTO; ok = true; } else if (arch_str == "bitvla") { *out = Arch::BITVLA; ok = true; } else if (arch_str == "vla_adapter"){ *out = Arch::VLA_ADAPTER;ok = true; } else if (arch_str == "openvla_oft"){ *out = Arch::OPENVLA_OFT;ok = true; } @@ -165,6 +168,10 @@ 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); break; + case Arch::OCTO: + std::printf("vla: arch = octo\n"); + impl = octo_create(mmproj_path, ckpt_path, config_path); + break; case Arch::BITVLA: std::printf("vla: arch = bitvla\n"); impl = bitvla_create(mmproj_path, ckpt_path, config_path); diff --git a/src/models/octo.cpp b/src/models/octo.cpp new file mode 100644 index 0000000..b90f593 --- /dev/null +++ b/src/models/octo.cpp @@ -0,0 +1,3056 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +#include "arch.h" +#include "model.h" +#include "models/gguf_reader.h" +#include "models/octo.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" +#ifdef GGML_USE_CUDA +#include "ggml-cuda.h" +#endif +#ifdef GGML_USE_METAL +#include "ggml-metal.h" +#endif +#include "gguf.h" + +#include "nlohmann/json.hpp" +#include "sentencepiece_processor.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vla { +namespace { + +// Per-timestep token count in the block-transformer sequence: primary(256) + wrist(64) +// + repeated-language(16) + readout(1) = 337. Constant across window_size -- only the +// number of timesteps (window_size) and the derived total seq length change. +// TIP-05: an L1/proprio checkpoint inserts a 3rd obs group (7 proprio tokens) right after +// wrist, per-timestep: primary(256) + wrist(64) + proprio(0 or 7) + repeated-language(16) +// + readout(1). kStepTokens is the historical no-proprio constant (still what every +// diffusion checkpoint uses, n_proprio=0); octo_step_tokens/octo_seq_len generalize it. +constexpr int kStepTokens = 337; +// task_language prefix tokens (once, not repeated per timestep). +constexpr int kTaskTokens = 16; +// Shared max-horizon slab size of the obs/task pos-embedding tables written by +// scripts/convert_octo_to_gguf.py (same for every checkpoint regardless of the +// window_size actually trained/used) -- see TIP-C1. window_size must fit within it. +constexpr int kMaxHorizon = 10; +// LowdimObsTokenizerPt(obs_keys=["proprio"]) always emits exactly 1 token per state +// dimension (discretize=False or True doesn't change the token *count*, only how each +// dimension is encoded before the shared Linear projection) -- octo-small-1.5's proprio +// is 7-dim, so this is fixed regardless of checkpoint. +constexpr int kProprioTokens = 7; + +inline int octo_step_tokens(int n_proprio_tokens) { + return kStepTokens + n_proprio_tokens; +} + +// seq = kTaskTokens + window_size * octo_step_tokens(n_proprio_tokens). n_proprio_tokens +// defaults to 0 (every pre-TIP-05 call site: diffusion checkpoints have no proprio group) +// so this stays byte-identical to the old window_size-only signature for them. +inline int octo_seq_len(int64_t window_size, int n_proprio_tokens = 0) { + return kTaskTokens + (int) window_size * octo_step_tokens(n_proprio_tokens); +} + +struct OctoModelArch : public ModelArchBase { + OctoModelArch() : ModelArchBase(Arch::OCTO) {} + ~OctoModelArch() override { + 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; + ggml_type matmul_type = GGML_TYPE_F32; + int n_threads = default_cpu_threads(); + + // TIP-BUILD-OCTO-GPU-B: predict()'s 5 stages now compute on `backend` directly (CUDA on + // GPU-capable builds, CPU otherwise -- see octo_create) and reference weight tensors + // resident on `weight_buf`/`ctx_weights` (populated once by load_all_tensors, same call + // as before TIP-A/B). TIP-A's separate CPU-only residency buffer + // (weight_cpu_backend/ctx_weights_cpu/weight_buf_cpu/tensors_cpu/load_all_tensors_cpu) + // has been removed: it's no longer needed now that compute and residency both key off + // the same real `backend` -- one residency path, matching the model's actual backend, + // exactly like every sibling. + + 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; + // "diffusion" (default, backward-compat for GGUFs converted before this key existed) + // or "l1" (aloha jitter-adapted L1 head; forward wired in TIP-05). + std::string head_type = "diffusion"; + // TIP-05: proprio tokenizer is detected from tensor presence (no dedicated GGUF KV + // key), independent of head_type -- see detect_proprio(). proprio_in_dim is + // octo.obs.proprio.proj.weight's in-dim: 1 (LowdimObsTokenizerPt discretize=False, + // the expected/documented case) or 256 (discretize=True, bin one-hot). + bool has_proprio = false; + int64_t proprio_in_dim = 0; + int64_t primary_tokens = 256; + int64_t wrist_tokens = 64; + int64_t language_tokens = 16; + int64_t diffusion_steps = 20; + + gguf_reader io{"octo"}; + std::vector tensors; + + std::vector predict(const Inputs& in) override; +}; + +// TIP-05: proprio has no dedicated GGUF metadata key (scripts/convert_octo_to_gguf.py only +// writes the octo.obs.proprio.* tensors when the source checkpoint's config has a "proprio" +// observation_tokenizer) -- detect it from tensor presence instead, same way every other +// per-checkpoint-optional group in this file is handled. Called from both load_config (the +// resident server/dump/free-sample paths, which already have a gguf_reader) and +// octo_predict_from_images (the CLI path, which builds its OctoModelArch by hand rather than +// through load_config). +static 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; +} + +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.diffusion.time_dim", + "octo.diffusion.hidden", + "octo.diffusion.num_blocks", + "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; + } + + 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"); + // octo.action.head_type is optional: GGUFs converted before this key existed (all + // pre-TIP-03 diffusion checkpoints) default to "diffusion", their sole prior behavior. + m.head_type = g.has("octo.action.head_type") ? g.str("octo.action.head_type") : "diffusion"; + m.primary_tokens = g.u32("octo.tokens.primary"); + m.wrist_tokens = g.u32("octo.tokens.wrist"); + m.language_tokens = g.u32("octo.tokens.language"); + m.diffusion_steps = g.u32("octo.diffusion.steps"); + 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): unexpected octo.obs.proprio.proj.weight in-dim=%lld (expected 1 or 256)\n", + (long long) m.proprio_in_dim); + return false; + } + + // action_horizon is per-checkpoint (diffusion libero=4, L1 aloha jitter-adapted=20); + // action_dim stays fixed at 7 (the octo-small-1.5 backbone's action-dim constant, + // shared by every head type) along with the other M0 backbone consts below. + if (m.hidden != 384 || m.blocks != 12 || m.heads != 6 || m.ffn != 1536 || m.action_dim != 7) { + std::fprintf(stderr, "vla(octo): metadata does not match octo-small-1.5 M0 constants\n"); + return false; + } + // window_size is per-checkpoint (bridge pretrain=2, LIBERO finetunes such as + // cyrusneary/octo-finetuned-libero=1); the pos-embedding table is a shared + // max_horizon=10 slab (see scripts/convert_octo_to_gguf.py), so any value in + // [1, kMaxHorizon] is a legal slice of it. + if (m.window_size < 1 || m.window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%lld out of supported range [1, %d]\n", + (long long) m.window_size, kMaxHorizon); + return false; + } + + // TIP-05: n_state/max_state_dim/real_state_dim were 0 (Octo had no proprio input) + // before this TIP; a proprio-tokenizer checkpoint has exactly 7 state dims (one + // token/dim -- see kProprioTokens), reported here so callers (server.cpp) know to + // supply Inputs::state. + 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; +} + +bool is_matmul_tensor(const char * name) { + return std::strstr(name, ".weight") != nullptr && + std::strstr(name, ".gn.") == nullptr && + std::strstr(name, "_norm.") == nullptr && + std::strstr(name, ".ln.") == nullptr && + std::strstr(name, "pos_embd") == nullptr && + std::strstr(name, "time_fourier") == nullptr; +} + +bool load_all_tensors(OctoModelArch& m, gguf_reader& g) { + ggml_init_params wp = { (size_t) 16 * 1024 * 1024, nullptr, true }; + m.ctx_weights = ggml_init(wp); + if (!m.ctx_weights) { + std::fprintf(stderr, "vla(octo): ggml_init(ctx_weights) failed\n"); + return false; + } + + const int64_t n = gguf_get_n_tensors(g.gctx); + m.tensors.reserve((size_t) n); + ggml_context * W = m.ctx_weights; + auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { + const ggml_tensor * gt = g.meta(name); + if (!gt) { + std::fprintf(stderr, "vla(octo): missing tensor %s\n", name); + return nullptr; + } + if (gt->type != GGML_TYPE_F32) { + std::fprintf(stderr, "vla(octo): tensor %s type=%d, expected F32 for M0\n", name, (int) gt->type); + return nullptr; + } + ggml_tensor * t = ggml_new_tensor(W, type, ggml_n_dims(gt), gt->ne); + ggml_set_name(t, name); + return t; + }; + auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; + auto mk_mm = [&](const char * name) { return mk(name, m.matmul_type); }; + + for (int64_t i = 0; i < n; ++i) { + const char * name = gguf_get_tensor_name(g.gctx, i); + ggml_tensor * t = is_matmul_tensor(name) ? mk_mm(name) : mk_f32(name); + if (!t) return false; + m.tensors.push_back(t); + } + + m.weight_buf = ggml_backend_alloc_ctx_tensors(m.ctx_weights, m.backend); + if (!m.weight_buf) { + std::fprintf(stderr, "vla(octo): ggml_backend_alloc_ctx_tensors failed\n"); + return false; + } + + for (ggml_tensor * t : m.tensors) { + std::vector bytes = g.read_convert(t->name, t->type); + if (bytes.empty()) return false; + ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + } + return true; +} + +// TIP-ND1-B: shared backend-selection logic (CUDA -> Metal -> CPU fallback), factored out of +// octo_create so every caller that needs its own resident model load (octo_create itself, +// octo_dump_tokenizer_case_resident, octo_predict_from_images, octo_free_sample_case) shares +// the identical selection order instead of re-duplicating it. `log_prefix` distinguishes the +// log lines of the non-octo_create callers (e.g. "[dump] ") from the live server's. +static ggml_backend_t octo_select_backend(int n_threads, const char * log_prefix) { + ggml_backend_t backend = nullptr; +#ifdef GGML_USE_CUDA + backend = ggml_backend_cuda_init(0); + if (backend) std::printf("vla(octo): %sbackend = CUDA (device 0)\n", log_prefix); + else std::fprintf(stderr, "vla(octo): %sggml_backend_cuda_init failed; falling back to CPU\n", log_prefix); +#elif defined(GGML_USE_METAL) + backend = ggml_backend_metal_init(); + if (backend) std::printf("vla(octo): %sbackend = Metal\n", log_prefix); + else std::fprintf(stderr, "vla(octo): %sggml_backend_metal_init failed; falling back to CPU\n", log_prefix); +#endif + if (!backend) { + backend = ggml_backend_cpu_init(); + if (!backend) { + std::fprintf(stderr, "vla(octo): %sggml_backend_cpu_init failed\n", log_prefix); + return nullptr; + } + ggml_backend_cpu_set_n_threads(backend, n_threads); + std::printf("vla(octo): %sbackend = CPU (%d threads)\n", log_prefix, n_threads); + } + return backend; +} + +// TIP-BUILD-OCTO-GPU-A: looks up an already-resident weight tensor by its GGUF name (same +// name strings the *_resident graph-building functions below already used as ggml_set_name +// literals, so this is a drop-in replacement for the old "build a fresh tensor + upload from +// a freshly-disk-read host vector" pattern). +static ggml_tensor * wt(ggml_context * ctx_w, const char * name) { + ggml_tensor * t = ggml_get_tensor(ctx_w, name); + if (!t) std::fprintf(stderr, "vla(octo): missing resident weight %s\n", name); + return t; +} + +// Host-side copy of a resident weight tensor's data, for the handful of weights that need a +// per-call CPU-side transform (standardize_conv_weight, expand_1d, pos-embed slicing) before +// they become graph operands -- those transforms are unchanged, just fed from the resident +// tensor's bytes instead of a freshly-read host vector. TIP-BUILD-OCTO-GPU-B: uses +// ggml_backend_tensor_get (not a raw ->data dereference) so this stays correct now that the +// resident tensor may live on a CUDA device buffer, not just host memory. +static 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; +} + +struct NpyU8 { + std::vector shape; + std::vector data; +}; + +struct NpyF32 { + std::vector shape; + std::vector data; +}; + +struct NpyBool { + std::vector shape; + std::vector data; +}; + +struct NpyI32 { + std::vector shape; + std::vector data; +}; + +static bool read_file_all(const std::string& path, std::vector& out) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "vla(octo): cannot open %s\n", path.c_str()); + return false; + } + f.seekg(0, std::ios::end); + const std::streamoff n = f.tellg(); + f.seekg(0, std::ios::beg); + out.resize((size_t) n); + return n == 0 || (bool) f.read(reinterpret_cast(out.data()), n); +} + +static bool parse_npy_u8(const std::string& path, NpyU8& out) { + std::vector bytes; + if (!read_file_all(path, bytes)) return false; + if (bytes.size() < 16 || std::memcmp(bytes.data(), "\x93NUMPY", 6) != 0) { + std::fprintf(stderr, "vla(octo): %s is not a .npy file\n", path.c_str()); + return false; + } + const int major = bytes[6]; + size_t pos = 8; + uint32_t hlen = 0; + if (major == 1) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8); + pos += 2; + } else if (major == 2 || major == 3) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8) | + ((uint32_t) bytes[pos + 2] << 16) | ((uint32_t) bytes[pos + 3] << 24); + pos += 4; + } else { + std::fprintf(stderr, "vla(octo): unsupported .npy version %d in %s\n", major, path.c_str()); + return false; + } + if (pos + hlen > bytes.size()) return false; + const std::string header(reinterpret_cast(bytes.data() + pos), hlen); + pos += hlen; + if (header.find("'descr': '|u1'") == std::string::npos && + header.find("\"descr\": \"|u1\"") == std::string::npos) { + std::fprintf(stderr, "vla(octo): %s expected uint8 .npy\n", path.c_str()); + return false; + } + if (header.find("'fortran_order': False") == std::string::npos && + header.find("\"fortran_order\": False") == std::string::npos) { + std::fprintf(stderr, "vla(octo): %s expected C-order .npy\n", path.c_str()); + return false; + } + const size_t l = header.find('('); + const size_t r = header.find(')', l == std::string::npos ? 0 : l); + if (l == std::string::npos || r == std::string::npos) return false; + out.shape.clear(); + size_t s = l + 1; + while (s < r) { + while (s < r && (header[s] == ' ' || header[s] == ',')) ++s; + size_t e = s; + while (e < r && header[e] >= '0' && header[e] <= '9') ++e; + if (e > s) out.shape.push_back(std::strtoll(header.substr(s, e - s).c_str(), nullptr, 10)); + s = e + 1; + } + int64_t ne = 1; + for (int64_t d : out.shape) ne *= d; + if (pos + (size_t) ne > bytes.size()) { + std::fprintf(stderr, "vla(octo): %s truncated .npy payload\n", path.c_str()); + return false; + } + out.data.assign(bytes.begin() + (ptrdiff_t) pos, bytes.begin() + (ptrdiff_t) pos + ne); + return true; +} + +static bool parse_npy_f32(const std::string& path, NpyF32& out) { + std::vector bytes; + if (!read_file_all(path, bytes)) return false; + if (bytes.size() < 16 || std::memcmp(bytes.data(), "\x93NUMPY", 6) != 0) { + std::fprintf(stderr, "vla(octo): %s is not a .npy file\n", path.c_str()); + return false; + } + const int major = bytes[6]; + size_t pos = 8; + uint32_t hlen = 0; + if (major == 1) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8); + pos += 2; + } else if (major == 2 || major == 3) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8) | + ((uint32_t) bytes[pos + 2] << 16) | ((uint32_t) bytes[pos + 3] << 24); + pos += 4; + } else { + std::fprintf(stderr, "vla(octo): unsupported .npy version %d in %s\n", major, path.c_str()); + return false; + } + if (pos + hlen > bytes.size()) return false; + const std::string header(reinterpret_cast(bytes.data() + pos), hlen); + pos += hlen; + if (header.find("'descr': '= '0' && header[e] <= '9') ++e; + if (e > s) out.shape.push_back(std::strtoll(header.substr(s, e - s).c_str(), nullptr, 10)); + s = e + 1; + } + int64_t ne = 1; + for (int64_t d : out.shape) ne *= d; + if (pos + (size_t) ne * sizeof(float) > bytes.size()) { + std::fprintf(stderr, "vla(octo): %s truncated .npy payload\n", path.c_str()); + return false; + } + out.data.resize((size_t) ne); + std::memcpy(out.data.data(), bytes.data() + pos, (size_t) ne * sizeof(float)); + return true; +} + +static bool parse_npy_bool(const std::string& path, NpyBool& out) { + std::vector bytes; + if (!read_file_all(path, bytes)) return false; + if (bytes.size() < 16 || std::memcmp(bytes.data(), "\x93NUMPY", 6) != 0) { + std::fprintf(stderr, "vla(octo): %s is not a .npy file\n", path.c_str()); + return false; + } + const int major = bytes[6]; + size_t pos = 8; + uint32_t hlen = 0; + if (major == 1) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8); + pos += 2; + } else if (major == 2 || major == 3) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8) | + ((uint32_t) bytes[pos + 2] << 16) | ((uint32_t) bytes[pos + 3] << 24); + pos += 4; + } else { + std::fprintf(stderr, "vla(octo): unsupported .npy version %d in %s\n", major, path.c_str()); + return false; + } + if (pos + hlen > bytes.size()) return false; + const std::string header(reinterpret_cast(bytes.data() + pos), hlen); + pos += hlen; + if (header.find("'descr': '|b1'") == std::string::npos && + header.find("\"descr\": \"|b1\"") == std::string::npos) { + std::fprintf(stderr, "vla(octo): %s expected bool .npy\n", path.c_str()); + return false; + } + if (header.find("'fortran_order': False") == std::string::npos && + header.find("\"fortran_order\": False") == std::string::npos) { + std::fprintf(stderr, "vla(octo): %s expected C-order .npy\n", path.c_str()); + return false; + } + const size_t l = header.find('('); + const size_t r = header.find(')', l == std::string::npos ? 0 : l); + if (l == std::string::npos || r == std::string::npos) return false; + out.shape.clear(); + size_t s = l + 1; + while (s < r) { + while (s < r && (header[s] == ' ' || header[s] == ',')) ++s; + size_t e = s; + while (e < r && header[e] >= '0' && header[e] <= '9') ++e; + if (e > s) out.shape.push_back(std::strtoll(header.substr(s, e - s).c_str(), nullptr, 10)); + s = e + 1; + } + int64_t ne = 1; + for (int64_t d : out.shape) ne *= d; + if (pos + (size_t) ne > bytes.size()) { + std::fprintf(stderr, "vla(octo): %s truncated .npy payload\n", path.c_str()); + return false; + } + out.data.assign(bytes.begin() + (ptrdiff_t) pos, bytes.begin() + (ptrdiff_t) pos + ne); + return true; +} + +// Some golden cases (e.g. tier2/bridge_debug, real-robot language-only conditioning) +// never recorded a task-image tensor at all: no goal image was ever provided for +// that dump, matching OctoModelPt.create_tasks(texts=...)'s own zero-fill for +// absent task modalities (octo-pytorch/octo/model/octo_model_pt.py:139-152). Parse +// the file if present (tier1's synthetic zero-content task image); zero-fill to +// match obs's spatial shape if the file is simply absent. +static bool parse_npy_u8_or_zero_task(const std::string& path, const NpyU8& obs, NpyU8& task) { + std::error_code ec; + if (std::filesystem::exists(path, ec)) return parse_npy_u8(path, task); + if (obs.shape.size() != 5) { + std::fprintf(stderr, "vla(octo): cannot infer zero-fill task shape from obs\n"); + return false; + } + task.shape = {1, obs.shape[2], obs.shape[3], obs.shape[4]}; + task.data.assign((size_t) obs.shape[2] * obs.shape[3] * obs.shape[4], 0); + return true; +} + +// Observation-side counterpart to parse_npy_u8_or_zero_task: some golden cases (the +// cyrusneary LIBERO checkpoint, TIP-GOLD) come from a genuinely single-camera +// finetune -- its example_batch/finetune_config never fed a wrist observation at all, +// so tensors/input.observation.image_wrist.npy simply doesn't exist. Zero-fill to +// {1, window_size, 3, side, side} if absent (matches +// scripts/patch_libero_golden_for_harness.py's WRIST_IMAGE_SIZE=128 placeholder, now +// built in so that script is no longer required). `used_real` reports which happened, +// so the caller can mark the corresponding wrist_valid pad-mask entries false instead +// of pretending a placeholder camera is a real, valid observation. +static bool parse_npy_u8_or_zero_obs(const std::string& path, int window_size, int side, + NpyU8& obs, bool& used_real) { + std::error_code ec; + if (std::filesystem::exists(path, ec)) { + used_real = true; + return parse_npy_u8(path, obs); + } + used_real = false; + obs.shape = {1, window_size, 3, side, side}; + obs.data.assign((size_t) window_size * 3 * side * side, 0); + return true; +} + +// Bool-array counterpart: zero-fill (all-False = "not valid") {1, window_size} if the +// golden case has no wrist pad-mask file at all (same single-camera scenario above). +static bool parse_npy_bool_or_zero(const std::string& path, int window_size, NpyBool& out) { + std::error_code ec; + if (std::filesystem::exists(path, ec)) return parse_npy_bool(path, out); + out.shape = {1, window_size}; + out.data.assign((size_t) window_size, 0); + return true; +} + +static bool parse_npy_i32(const std::string& path, NpyI32& out) { + std::vector bytes; + if (!read_file_all(path, bytes)) return false; + if (bytes.size() < 16 || std::memcmp(bytes.data(), "\x93NUMPY", 6) != 0) { + std::fprintf(stderr, "vla(octo): %s is not a .npy file\n", path.c_str()); + return false; + } + const int major = bytes[6]; + size_t pos = 8; + uint32_t hlen = 0; + if (major == 1) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8); + pos += 2; + } else if (major == 2 || major == 3) { + hlen = (uint32_t) bytes[pos] | ((uint32_t) bytes[pos + 1] << 8) | + ((uint32_t) bytes[pos + 2] << 16) | ((uint32_t) bytes[pos + 3] << 24); + pos += 4; + } else { + std::fprintf(stderr, "vla(octo): unsupported .npy version %d in %s\n", major, path.c_str()); + return false; + } + if (pos + hlen > bytes.size()) return false; + const std::string header(reinterpret_cast(bytes.data() + pos), hlen); + pos += hlen; + // Golden dump scripts aren't consistent about token-id width (tier1 uses int32, + // tier2/bridge_debug uses torch's default int64) -- accept either and downcast; + // token ids are always well within int32 range. + bool is_i32 = header.find("'descr': '= '0' && header[e] <= '9') ++e; + if (e > s) out.shape.push_back(std::strtoll(header.substr(s, e - s).c_str(), nullptr, 10)); + s = e + 1; + } + int64_t ne = 1; + for (int64_t d : out.shape) ne *= d; + const size_t elsz = is_i64 ? sizeof(int64_t) : sizeof(int32_t); + if (pos + (size_t) ne * elsz > bytes.size()) { + std::fprintf(stderr, "vla(octo): %s truncated .npy payload\n", path.c_str()); + return false; + } + out.data.resize((size_t) ne); + if (is_i64) { + std::vector tmp((size_t) ne); + std::memcpy(tmp.data(), bytes.data() + pos, (size_t) ne * elsz); + for (size_t i = 0; i < tmp.size(); ++i) out.data[i] = (int32_t) tmp[i]; + } else { + std::memcpy(out.data.data(), bytes.data() + pos, (size_t) ne * elsz); + } + return true; +} + +static bool write_f32_dump(const std::string& dir, + const char * name, + const std::vector& data, + const std::vector& shape, + std::ofstream& manifest) { + static std::string made; + if (made != dir) { + std::error_code ec; + std::filesystem::create_directories(dir, ec); + if (ec) { + std::fprintf(stderr, "vla(octo): cannot create %s: %s\n", dir.c_str(), ec.message().c_str()); + return false; + } + made = dir; + } + const std::string path = dir + "/" + name + ".f32"; + std::ofstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "vla(octo): cannot write %s\n", path.c_str()); + return false; + } + f.write(reinterpret_cast(data.data()), (std::streamsize) (data.size() * sizeof(float))); + manifest << name << " " << path << " float32"; + for (int64_t d : shape) manifest << " " << d; + manifest << "\n"; + return (bool) f; +} + +static bool write_bool_dump(const std::string& dir, + const char * name, + const std::vector& data, + const std::vector& shape, + std::ofstream& manifest) { + const std::string path = dir + "/" + name + ".bool"; + std::ofstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "vla(octo): cannot write %s\n", path.c_str()); + return false; + } + f.write(reinterpret_cast(data.data()), (std::streamsize) data.size()); + manifest << name << " " << path << " bool"; + for (int64_t d : shape) manifest << " " << d; + manifest << "\n"; + return (bool) f; +} + +static void standardize_conv_weight(const std::vector& src, + int oc, int ic, int kh, int kw, + std::vector& dst) { + dst.resize(src.size()); + const int n = ic * kh * kw; + for (int o = 0; o < oc; ++o) { + double mean = 0.0, var = 0.0; + const size_t base = (size_t) o * n; + for (int i = 0; i < n; ++i) mean += src[base + i]; + mean /= n; + for (int i = 0; i < n; ++i) { + const double d = (double) src[base + i] - mean; + var += d * d; + } + const float inv = 1.0f / std::sqrt((float) (var / n) + 1e-10f); + for (int i = 0; i < n; ++i) dst[base + i] = ((float) src[base + i] - (float) mean) * inv; + } +} + +static ggml_tensor * make_4d(ggml_context * ctx, const char * name, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + ggml_tensor * t = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne0, ne1, ne2, ne3); + ggml_set_name(t, name); + return t; +} + +// TIP-ND1-B: sole obs-tokenizer implementation (was run_one_obs_tokenizer_graph / +// run_one_obs_tokenizer_graph_resident, unified -- TIP-ND1-A proved resident-on-CPU is +// bit-exact with the deleted disk-read/cpu_init original, 20/20 golden cases). Weights come +// from the resident context (`wt()` + `tensor_to_vec()` for the handful needing a host-side +// transform); compute runs on the model's real `backend` (CUDA when available, CPU +// otherwise). `backend` is NOT owned by this function -- never freed here. +static bool run_one_obs_tokenizer_graph_resident(ggml_context * ctx_w, + ggml_backend_t backend, + const char * view, + const NpyU8& obs, + const NpyU8& task, + int side, + int n_tok, + int window_size, + std::vector& tok, + std::vector& proj, + std::vector& pos) { + if (obs.shape.size() != 5 || task.shape.size() != 4 || + obs.shape[0] != 1 || obs.shape[1] != window_size || obs.shape[2] != 3 || + task.shape[0] != 1 || task.shape[1] != 3 || + obs.shape[3] != side || obs.shape[4] != side || + task.shape[2] != side || task.shape[3] != side) { + std::fprintf(stderr, "vla(octo): unexpected input image shape for side=%d\n", side); + return false; + } + tok.assign((size_t) 1 * window_size * n_tok * 512, 0.0f); + proj.assign((size_t) 1 * window_size * n_tok * 384, 0.0f); + pos.assign((size_t) 1 * window_size * n_tok * 384, 0.0f); + + const int stem_oc[4] = {32, 96, 192, 384}; + const int stem_ic[4] = {6, 32, 96, 192}; + std::vector input((size_t) side * side * 6 * window_size, 0.0f); + for (int t = 0; t < window_size; ++t) { + for (int c = 0; c < 3; ++c) { + for (int yy = 0; yy < side; ++yy) { + for (int xx = 0; xx < side; ++xx) { + const size_t oi = ((((size_t) t * 3 + c) * side + yy) * side + xx); + const size_t ti = (((size_t) c * side + yy) * side + xx); + input[(((size_t) t * 6 + c) * side + yy) * side + xx] = (float) obs.data[oi] / 127.5f - 1.0f; + input[(((size_t) t * 6 + c + 3) * side + yy) * side + xx] = (float) task.data[ti] / 127.5f - 1.0f; + } + } + } + } + + ggml_init_params gp = {(size_t) 96 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) { + std::fprintf(stderr, "vla(octo): ggml_init(tokenizer graph ctx) failed\n"); + return false; + } + + std::vector tensors; + std::vector> payloads; + auto add_payload = [&](ggml_tensor * t, std::vector data) { + tensors.push_back(t); + payloads.push_back(std::move(data)); + }; + auto expand_1d = [](const std::vector& src, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + std::vector out((size_t) ne0 * ne1 * ne2 * ne3, 0.0f); + for (int64_t i3 = 0; i3 < ne3; ++i3) + for (int64_t i2 = 0; i2 < ne2; ++i2) + for (int64_t i1 = 0; i1 < ne1; ++i1) + for (int64_t i0 = 0; i0 < ne0; ++i0) + out[((size_t) i3 * ne2 * ne1 * ne0) + (size_t) i2 * ne1 * ne0 + (size_t) i1 * ne0 + i0] = src[(size_t) i2]; + return out; + }; + + ggml_tensor * x = make_4d(ctx, "octo.obs.input_norm", side, side, 6, window_size); + add_payload(x, std::move(input)); + + char rname[160]; + bool ok = true; + for (int li = 0; li < 4 && ok; ++li) { + std::snprintf(rname, sizeof(rname), "octo.obs.%s.stem.%d.conv.weight", view, li); + ggml_tensor * conv_w_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.stem.%d.conv.bias", view, li); + ggml_tensor * conv_b_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.stem.%d.gn.weight", view, li); + ggml_tensor * gn_w_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.stem.%d.gn.bias", view, li); + ggml_tensor * gn_b_r = wt(ctx_w, rname); + if (!conv_w_r || !conv_b_r || !gn_w_r || !gn_b_r) { ok = false; break; } + + std::vector ws; + standardize_conv_weight(tensor_to_vec(conv_w_r), stem_oc[li], stem_ic[li], 3, 3, ws); + char name[64]; + std::snprintf(name, sizeof(name), "octo.obs.stem.%d.conv.weight_std", li); + ggml_tensor * cw = make_4d(ctx, name, 3, 3, stem_ic[li], stem_oc[li]); + add_payload(cw, std::move(ws)); + std::snprintf(name, sizeof(name), "octo.obs.stem.%d.conv.bias", li); + ggml_tensor * cb = make_4d(ctx, name, 1, 1, stem_oc[li], 1); + add_payload(cb, expand_1d(tensor_to_vec(conv_b_r), 1, 1, stem_oc[li], 1)); + std::snprintf(name, sizeof(name), "octo.obs.stem.%d.gn.weight", li); + ggml_tensor * gw = make_4d(ctx, name, 1, 1, stem_oc[li], 1); + add_payload(gw, expand_1d(tensor_to_vec(gn_w_r), 1, 1, stem_oc[li], 1)); + std::snprintf(name, sizeof(name), "octo.obs.stem.%d.gn.bias", li); + ggml_tensor * gb = make_4d(ctx, name, 1, 1, stem_oc[li], 1); + add_payload(gb, expand_1d(tensor_to_vec(gn_b_r), 1, 1, stem_oc[li], 1)); + + x = ggml_conv_2d(ctx, cw, x, 2, 2, 1, 1, 1, 1); + x = ggml_add(ctx, x, cb); + x = ggml_group_norm(ctx, x, 32, 1e-5f); + x = ggml_add(ctx, ggml_mul(ctx, x, gw), gb); + x = ggml_relu(ctx, x); + } + if (!ok) { + ggml_free(ctx); + return false; + } + + std::snprintf(rname, sizeof(rname), "octo.obs.%s.patch_embd.weight", view); + ggml_tensor * patch_w_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.patch_embd.bias", view); + ggml_tensor * patch_b_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.proj.weight", view); + ggml_tensor * proj_w_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.proj.bias", view); + ggml_tensor * proj_b_r = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.obs.%s.pos_embd", view); + ggml_tensor * pos_r = wt(ctx_w, rname); + if (!patch_w_r || !patch_b_r || !proj_w_r || !proj_b_r || !pos_r) { + ggml_free(ctx); + return false; + } + + ggml_tensor * pw = make_4d(ctx, "octo.obs.patch_embd.weight", 1, 1, 384, 512); + add_payload(pw, tensor_to_vec(patch_w_r)); + ggml_tensor * pb = make_4d(ctx, "octo.obs.patch_embd.bias", 1, 1, 512, 1); + add_payload(pb, expand_1d(tensor_to_vec(patch_b_r), 1, 1, 512, 1)); + ggml_tensor * jw = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, 384); + ggml_set_name(jw, "octo.obs.proj.weight"); + add_payload(jw, tensor_to_vec(proj_w_r)); + ggml_tensor * jb = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 384, 1, 1); + ggml_set_name(jb, "octo.obs.proj.bias"); + add_payload(jb, tensor_to_vec(proj_b_r)); + const std::vector pos_full = tensor_to_vec(pos_r); + std::vector pos2(pos_full.begin(), pos_full.begin() + (size_t) window_size * n_tok * 384); + ggml_tensor * pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 384, n_tok, window_size); + ggml_set_name(pe, "octo.obs.pos_embd.window"); + add_payload(pe, std::move(pos2)); + + ggml_tensor * patch = ggml_conv_2d(ctx, pw, x, 1, 1, 0, 0, 1, 1); + patch = ggml_add(ctx, patch, pb); + ggml_tensor * tok_t = ggml_cont(ctx, ggml_reshape_3d(ctx, ggml_cont(ctx, ggml_permute(ctx, patch, 1, 2, 0, 3)), 512, n_tok, window_size)); + ggml_set_name(tok_t, "obs.tokenizer.tok"); + ggml_set_output(tok_t); + + ggml_tensor * proj_t = ggml_add(ctx, ggml_mul_mat(ctx, jw, tok_t), jb); + ggml_set_name(proj_t, "obs.tokenizer.proj"); + ggml_set_output(proj_t); + ggml_tensor * pos_t = ggml_add(ctx, proj_t, pe); + ggml_set_name(pos_t, "obs.tokenizer.pos"); + ggml_set_output(pos_t); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 8192, false); + ggml_build_forward_expand(graph, pos_t); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): tokenizer ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + for (size_t i = 0; i < tensors.size(); ++i) { + ggml_backend_tensor_set(tensors[i], payloads[i].data(), 0, ggml_nbytes(tensors[i])); + } + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): tokenizer ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_get(tok_t, tok.data(), 0, ggml_nbytes(tok_t)); + ggml_backend_tensor_get(proj_t, proj.data(), 0, ggml_nbytes(proj_t)); + ggml_backend_tensor_get(pos_t, pos.data(), 0, ggml_nbytes(pos_t)); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +// TIP-05 Part A: LowdimObsTokenizerPt(obs_keys=["proprio"]) forward. p_norm (z-scored by +// proprio-stats, done by the caller before this graph -- see load_proprio_stats) is turned +// into one token per state dim (7 total): continuous (in_dim=1) feeds p_norm straight into +// the shared Linear(1,384) projection; discretize (in_dim=256) bins p_norm through +// octo.obs.proprio.bin_thresholds into a one-hot(256) vector first. Either way the +// projection's weight/bias is the SAME [in_dim,384] matrix applied independently per +// dim-token (matches BinTokenizerPt/LowdimObsTokenizerPt's nn.Linear applied on the last +// axis of a (...,7,in_dim) tensor -- one shared projection, not 7 per-dim ones), so both +// branches converge on an identical graph shape ([in_dim,7,window_size] -> [384,7,window_size]) +// once the host-side `tokens_in` buffer is built. +static bool run_proprio_tokenizer_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + const std::vector& proprio_norm, // [window_size,7], z-scored + int in_dim, // 1 or 256 + int window_size, + std::vector& pos_out) { // [window_size,7,384] + constexpr int n_dims = kProprioTokens; + if (proprio_norm.size() != (size_t) window_size * n_dims) return false; + if (in_dim != 1 && in_dim != 256) { + std::fprintf(stderr, "vla(octo): proprio tokenizer in_dim=%d unsupported (expected 1 or 256)\n", in_dim); + return false; + } + pos_out.assign((size_t) window_size * n_dims * 384, 0.0f); + + ggml_tensor * proj_w_r = wt(ctx_w, "octo.obs.proprio.proj.weight"); + ggml_tensor * proj_b_r = wt(ctx_w, "octo.obs.proprio.proj.bias"); + ggml_tensor * pos_r = wt(ctx_w, "octo.obs.proprio.pos_embd"); + if (!proj_w_r || !proj_b_r || !pos_r) return false; + + std::vector tokens_in((size_t) in_dim * n_dims * window_size, 0.0f); + if (in_dim == 1) { + // Continuous: token = p_norm.unsqueeze(-1) -- the raw z-scored scalar itself. + for (size_t i = 0; i < proprio_norm.size(); ++i) tokens_in[i] = proprio_norm[i]; + } else { + // Discretize: BinTokenizerPt buckets each dim's p_norm against the fixed (n_bins-1) + // thresholds octo.obs.proprio.bin_thresholds and one-hot encodes the bucket index -- + // torch.bucketize(x, boundaries) semantics (index = count of boundaries <= x, i.e. + // TIP-05's "argmax_j (p_norm >= thr[j] & p_norm < thr[j+1])" with thr padded by + // implicit +/-inf at the ends). NOT exercised by any GGUF available in this sandbox + // (the only observed checkpoint config is discretize=False/in_dim=1) -- implemented + // per spec for completeness; TIP-06 parity should confirm bucket-boundary behavior + // against a real discretize=True golden trace if one ever exists. + ggml_tensor * thresholds_r = wt(ctx_w, "octo.obs.proprio.bin_thresholds"); + if (!thresholds_r) return false; + const std::vector thresholds = tensor_to_vec(thresholds_r); + const int n_thresh = (int) thresholds.size(); + for (int t = 0; t < window_size; ++t) { + for (int d = 0; d < n_dims; ++d) { + const float v = proprio_norm[(size_t) t * n_dims + d]; + int bucket = 0; + while (bucket < n_thresh && thresholds[(size_t) bucket] <= v) ++bucket; + bucket = std::min(bucket, in_dim - 1); + tokens_in[((size_t) t * n_dims + d) * in_dim + bucket] = 1.0f; + } + } + } + + ggml_init_params gp = {(size_t) 8 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) return false; + + ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, n_dims, window_size); + ggml_set_name(x, "octo.obs.proprio.tokens_in"); + + ggml_tensor * pb = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 384, 1, 1); + ggml_set_name(pb, "octo.obs.proprio.proj.bias3d"); + + const std::vector pos_full = tensor_to_vec(pos_r); + if (pos_full.size() < (size_t) window_size * n_dims * 384) { + std::fprintf(stderr, "vla(octo): octo.obs.proprio.pos_embd too small for window_size=%d\n", window_size); + ggml_free(ctx); + return false; + } + std::vector pos_slice(pos_full.begin(), pos_full.begin() + (size_t) window_size * n_dims * 384); + ggml_tensor * pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 384, n_dims, window_size); + ggml_set_name(pe, "octo.obs.proprio.pos_embd.window"); + + ggml_tensor * proj_t = ggml_add(ctx, ggml_mul_mat(ctx, proj_w_r, x), pb); + ggml_tensor * pos_t = ggml_add(ctx, proj_t, pe); + ggml_set_name(pos_t, "obs.proprio.pos"); + ggml_set_output(pos_t); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 256, false); + ggml_build_forward_expand(graph, pos_t); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): proprio tokenizer ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(x, tokens_in.data(), 0, ggml_nbytes(x)); + ggml_backend_tensor_set(pb, tensor_to_vec(proj_b_r).data(), 0, ggml_nbytes(pb)); + ggml_backend_tensor_set(pe, pos_slice.data(), 0, ggml_nbytes(pe)); + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): proprio tokenizer ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_get(pos_t, pos_out.data(), 0, ggml_nbytes(pos_t)); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +// TIP-ND1-B: sole language-stage implementation (was run_language_graph / +// run_language_graph_resident, unified). proj_w/proj_b/pos are referenced straight from the +// resident context (no copy, no per-call upload); `inp` (the T5 encoder's per-call output) is +// genuinely new data each call and still gets uploaded fresh. +static bool run_language_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + const NpyF32& t5, + int window_size, + std::vector& proj, + std::vector& pos, + std::vector& repeated) { + if (t5.shape.size() != 3 || t5.shape[0] != 1 || t5.shape[1] != 16 || t5.shape[2] != 768) { + std::fprintf(stderr, "vla(octo): expected T5 inject shape [1,16,768]\n"); + return false; + } + proj.assign((size_t) 1 * 16 * 384, 0.0f); + pos.assign((size_t) 1 * 16 * 384, 0.0f); + repeated.assign((size_t) 1 * window_size * 16 * 384, 0.0f); + + ggml_tensor * jw = wt(ctx_w, "octo.task.language.proj.weight"); + ggml_tensor * jb = wt(ctx_w, "octo.task.language.proj.bias"); + ggml_tensor * pe = wt(ctx_w, "octo.task.language.pos_embd"); + if (!jw || !jb || !pe) return false; + + ggml_init_params gp = {(size_t) 8 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) { + std::fprintf(stderr, "vla(octo): ggml_init(language graph ctx) failed\n"); + return false; + } + + ggml_tensor * inp = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 768, 16, 1); + ggml_set_name(inp, "octo.task.language.t5_inject"); + + ggml_tensor * proj_t = ggml_add(ctx, ggml_mul_mat(ctx, jw, inp), jb); + ggml_set_name(proj_t, "task_language.proj"); + ggml_set_output(proj_t); + ggml_tensor * pos_t = ggml_add(ctx, proj_t, pe); + ggml_set_name(pos_t, "task_language.pos"); + ggml_set_output(pos_t); + ggml_tensor * repeated_t = ggml_repeat_4d(ctx, pos_t, 384, 16, window_size, 1); + ggml_set_name(repeated_t, "obs_task_language.repeated"); + ggml_set_output(repeated_t); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1024, false); + ggml_build_forward_expand(graph, repeated_t); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): language ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(inp, t5.data.data(), 0, ggml_nbytes(inp)); + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): language ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_get(proj_t, proj.data(), 0, ggml_nbytes(proj_t)); + ggml_backend_tensor_get(pos_t, pos.data(), 0, ggml_nbytes(pos_t)); + ggml_backend_tensor_get(repeated_t, repeated.data(), 0, ggml_nbytes(repeated_t)); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +// T5 relative-position bucket, encoder self-attention (bidirectional=true). +// Matches HF T5Attention._relative_position_bucket / llama.cpp's llama_relative_position_bucket: +// relative_position = key_pos - query_pos, 32 buckets, max_distance 128. +static 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; + int32_t bucket = (relative_position > 0) ? nb : 0; + const int32_t rp = std::abs(relative_position); + const int32_t max_exact = nb / 2; + 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); + int32_t rp_large = (int32_t) std::floor(v); + rp_large = std::min(rp_large, nb - 1); + bucket += rp_large; + } + return bucket; +} + +// T5-base encoder-only forward as a ggml graph: embedding lookup (host-side row fetch) -> 12x +// [T5LayerNorm(RMS) -> self-attn (shared relative-position bias + padding mask, no query scaling) +// -> residual -> T5LayerNorm -> DenseReluDense(ReLU) -> residual] -> final T5LayerNorm. +// TIP-ND1-B: sole T5-encoder implementation (was run_t5_encoder_graph / +// run_t5_encoder_graph_resident, unified). The 12 blocks' attn/ffn weights, attn_rel_b, and +// output_norm are referenced directly from the resident context; compute runs on the model's +// real `backend` (CUDA when available, CPU otherwise). The embedding lookup is an in-graph +// ggml_get_rows against the resident octo.t5.tok_embd.weight tensor (same op already used for +// the relative-position-bias gather two lines below), so on a CUDA build the gather itself +// happens on-device, no host round-trip. `backend` is NOT owned by this function -- never +// freed here. +static bool run_t5_encoder_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + 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 = 16; + 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; + } + + ggml_tensor * tok_embd_r = wt(ctx_w, "octo.t5.tok_embd.weight"); + ggml_tensor * rel_b_r = wt(ctx_w, "octo.t5.blk.0.attn_rel_b.weight"); + ggml_tensor * outw_r = wt(ctx_w, "octo.t5.output_norm.weight"); + if (!tok_embd_r || !rel_b_r || !outw_r) return false; + char rname[160]; + ggml_tensor * blk_w[12][8]; // attn_norm, q, k, v, o, ffn_norm, ffn_up, ffn_down + 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] = wt(ctx_w, rname); + if (!blk_w[i][j]) return false; + } + } + + std::vector bucket_idx((size_t) seq * seq); + std::vector padmask((size_t) seq * seq); + for (int j = 0; j < seq; ++j) { // query + for (int i = 0; i < seq; ++i) { // key + bucket_idx[(size_t) j * seq + i] = t5_relative_position_bucket(j, i, n_buckets, max_distance); + padmask[(size_t) j * seq + i] = attention_mask[(size_t) i] != 0 ? 0.0f : -FLT_MAX; + } + } + + ggml_init_params gp = {(size_t) 32 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) return false; + + std::vector tensors; + std::vector> payloads_f32; + std::vector> payloads_i32; + auto add_f32 = [&](ggml_tensor * t, std::vector data) { + tensors.push_back(t); + payloads_f32.push_back(std::move(data)); + return t; + }; + auto add_i32 = [&](ggml_tensor * t, std::vector data) { + tensors.push_back(t); + payloads_i32.push_back(std::move(data)); + return t; + }; + + ggml_tensor * input_ids_t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, seq); + ggml_set_name(input_ids_t, "octo.t5.input_ids"); + add_i32(input_ids_t, input_ids); + ggml_tensor * x = ggml_get_rows(ctx, tok_embd_r, input_ids_t); + ggml_set_name(x, "octo.t5.input_embed"); + + ggml_tensor * bucket = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, seq, seq); + ggml_set_name(bucket, "octo.t5.pos_bucket"); + add_i32(bucket, bucket_idx); + ggml_tensor * rel_b = rel_b_r; + ggml_tensor * padmask_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, seq, seq); + ggml_set_name(padmask_t, "octo.t5.padmask"); + add_f32(padmask_t, padmask); + + ggml_tensor * pos_bucket_1d = ggml_reshape_1d(ctx, bucket, (int64_t) seq * seq); + ggml_tensor * pos_bias = ggml_get_rows(ctx, rel_b, pos_bucket_1d); + pos_bias = ggml_reshape_3d(ctx, pos_bias, heads, seq, seq); + pos_bias = ggml_cont(ctx, ggml_permute(ctx, pos_bias, 2, 0, 1, 3)); + ggml_tensor * mask = ggml_add(ctx, pos_bias, padmask_t); + + for (int i = 0; i < 12; ++i) { + ggml_tensor * n1w = blk_w[i][0]; + ggml_tensor * Wq = blk_w[i][1]; + ggml_tensor * Wk = blk_w[i][2]; + ggml_tensor * Wv = blk_w[i][3]; + ggml_tensor * Wo = blk_w[i][4]; + ggml_tensor * n2w = blk_w[i][5]; + ggml_tensor * Wup = blk_w[i][6]; + ggml_tensor * Wdown = blk_w[i][7]; + + ggml_tensor * n1 = ggml_mul(ctx, ggml_rms_norm(ctx, x, ln_eps), n1w); + ggml_tensor * Q = ggml_mul_mat(ctx, Wq, n1); + ggml_tensor * K = ggml_mul_mat(ctx, Wk, n1); + ggml_tensor * V = ggml_mul_mat(ctx, Wv, n1); + ggml_tensor * Qh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, Q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * Kh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, K, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * Vh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, V, head_dim, heads, seq), 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(ctx, Kh, Qh); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + ggml_tensor * probs = ggml_soft_max_ext(ctx, scores, mask, 1.0f, 0.0f); // T5: no 1/sqrt(d_k) scaling + ggml_tensor * attended = ggml_mul_mat(ctx, Vh, probs); + ggml_tensor * merged = ggml_reshape_2d(ctx, ggml_cont(ctx, ggml_permute(ctx, attended, 0, 2, 1, 3)), hidden, seq); + ggml_tensor * attn_out = ggml_mul_mat(ctx, Wo, merged); + x = ggml_add(ctx, x, attn_out); + + ggml_tensor * n2 = ggml_mul(ctx, ggml_rms_norm(ctx, x, ln_eps), n2w); + ggml_tensor * h = ggml_relu(ctx, ggml_mul_mat(ctx, Wup, n2)); + ggml_tensor * ffn_out = ggml_mul_mat(ctx, Wdown, h); + x = ggml_add(ctx, x, ffn_out); + } + + ggml_tensor * out = ggml_mul(ctx, ggml_rms_norm(ctx, x, ln_eps), outw_r); + ggml_set_name(out, "t5.out"); + ggml_set_output(out); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 4096, false); + ggml_build_forward_expand(graph, out); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): T5 encoder ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + size_t fi = 0, ii = 0; + for (ggml_tensor * t : tensors) { + if (t->type == GGML_TYPE_I32) { + ggml_backend_tensor_set(t, payloads_i32[ii].data(), 0, ggml_nbytes(t)); + ++ii; + } else { + ggml_backend_tensor_set(t, payloads_f32[fi].data(), 0, ggml_nbytes(t)); + ++fi; + } + } + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): T5 encoder ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + t5_out.resize((size_t) hidden * seq); + ggml_backend_tensor_get(out, t5_out.data(), 0, ggml_nbytes(out)); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +enum class OctoTokenKind { TASK, OBS, READOUT }; + +struct OctoTokenMetadata { + OctoTokenKind kind; + int timestep; +}; + +struct OctoTransformerResult { + std::vector input; + std::vector blocked_mask; + std::array, 12> block_outputs; + std::vector output; + std::vector task_language; + std::vector obs_primary; + std::vector obs_wrist; + std::vector obs_proprio; // TIP-05: empty unless n_proprio_tokens > 0. + std::vector obs_task_language; + std::vector readout_action; +}; + +struct OctoDiffusionSchedule { + std::array betas{}; + std::array alphas{}; + std::array alpha_hats{}; +}; + +struct OctoDiffusionResult { + std::vector initial_noise; + std::vector action_mask; + std::vector flat_action_mask; + std::array, 20> current_x_before; + std::array, 20> pred_eps; + std::array, 20> z; + std::array, 20> after_denoise; + std::array, 20> after_noise_add; + std::array, 20> after_clip; + std::array, 20> after_mask; + std::vector actions_all_timesteps; + std::vector final_actions; +}; + +// TIP-05: obs_proprio/n_proprio_tokens are additive -- n_proprio_tokens=0 (obs_proprio +// ignored/may be empty) reproduces the pre-TIP-05 layout byte-for-byte (every existing call +// site keeps doing exactly that). When n_proprio_tokens=7, the proprio group is inserted +// right after wrist and before the repeated-language group, per timestep: primary(256) + +// wrist(64) + proprio(7) + repeated-language(16) + readout(1) -- see kStepTokens's TIP-05 +// comment. (Order chosen from the "primary -> wrist -> proprio" obs-group ordering TIP-05 +// specifies, placed before the task-derived repeated-language/readout entries exactly like +// primary/wrist already are; TIP-06 should confirm this ordering against a golden L1 dump.) +static bool assemble_transformer_input(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, + int window_size, + int n_proprio_tokens, + std::vector& input) { + constexpr int hidden = 384; + const int seq = octo_seq_len(window_size, n_proprio_tokens); + const int step_tokens = octo_step_tokens(n_proprio_tokens); + if (task_language.size() != (size_t) 16 * hidden || + obs_primary.size() != (size_t) window_size * 256 * hidden || + obs_wrist.size() != (size_t) window_size * 64 * hidden || + (n_proprio_tokens > 0 && obs_proprio.size() != (size_t) window_size * n_proprio_tokens * hidden) || + repeated_language.size() != (size_t) window_size * 16 * hidden || + readout_pos.size() < (size_t) window_size * hidden) { + std::fprintf(stderr, "vla(octo): invalid tensor size while assembling block transformer input\n"); + return false; + } + input.assign((size_t) seq * hidden, 0.0f); + std::copy(task_language.begin(), task_language.end(), input.begin()); + for (int t = 0; t < window_size; ++t) { + const size_t dst = (size_t) (16 + t * step_tokens) * hidden; + std::copy_n(obs_primary.begin() + (size_t) t * 256 * hidden, (size_t) 256 * hidden, input.begin() + dst); + std::copy_n(obs_wrist.begin() + (size_t) t * 64 * hidden, (size_t) 64 * hidden, input.begin() + dst + (size_t) 256 * hidden); + size_t off = (size_t) 320 * hidden; + if (n_proprio_tokens > 0) { + std::copy_n(obs_proprio.begin() + (size_t) t * n_proprio_tokens * hidden, (size_t) n_proprio_tokens * hidden, + input.begin() + dst + off); + off += (size_t) n_proprio_tokens * hidden; + } + std::copy_n(repeated_language.begin() + (size_t) t * 16 * hidden, (size_t) 16 * hidden, + input.begin() + dst + off); + off += (size_t) 16 * hidden; + std::copy_n(readout_pos.begin() + (size_t) t * hidden, hidden, + input.begin() + dst + off); + } + return true; +} + +static OctoDiffusionSchedule make_cosine_schedule() { + OctoDiffusionSchedule s; + constexpr int steps = 20; + constexpr double ds = 0.008; + constexpr double pi = 3.141592653589793238462643383279502884; + std::array alpha_cum{}; + 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; + } + const double first = alpha_cum[0]; + float cum = 1.0f; + for (int i = 0; i < steps; ++i) { + const double a0 = alpha_cum[(size_t) i] / first; + const double a1 = alpha_cum[(size_t) i + 1] / first; + const float beta = (float) std::min(std::max(1.0 - a1 / a0, 0.0), 0.999); + s.betas[(size_t) i] = beta; + s.alphas[(size_t) i] = 1.0f - beta; + cum *= s.alphas[(size_t) i]; + s.alpha_hats[(size_t) i] = cum; + } + return s; +} + +// TIP-ND1-B: sole score-actor implementation (was run_score_actor_graph / +// run_score_actor_graph_resident, unified) -- called once per denoising step (20x per +// predict()), so this is the highest call-frequency of the 5 stages. All 9 fixed weights + the +// 3 residual blocks' 6 weights each are referenced directly from the resident context; only +// the 3 genuinely-per-call input tensors (time, readout embedding, noisy action) get built + +// uploaded fresh each call. +static bool run_score_actor_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + const std::vector& readout_action, + const std::vector& noisy_action, + int time_value, + int window_size, + int action_total, + std::vector& pred_eps) { + constexpr int hidden = 384; + const int action = action_total; + const int width = window_size; + constexpr float ln_eps = 1e-6f; + constexpr float two_pi = 6.2831853071795864769f; + if (readout_action.size() != (size_t) hidden * width || noisy_action.size() != (size_t) action * width) return false; + + ggml_tensor * time_w = wt(ctx_w, "octo.head.diffusion.time_fourier.weight"); + ggml_tensor * c0w = wt(ctx_w, "octo.head.diffusion.cond.0.weight"); + ggml_tensor * c0b = wt(ctx_w, "octo.head.diffusion.cond.0.bias"); + ggml_tensor * c1w = wt(ctx_w, "octo.head.diffusion.cond.1.weight"); + ggml_tensor * c1b = wt(ctx_w, "octo.head.diffusion.cond.1.bias"); + ggml_tensor * rinw = wt(ctx_w, "octo.head.diffusion.reverse.in.weight"); + ggml_tensor * rinb = wt(ctx_w, "octo.head.diffusion.reverse.in.bias"); + ggml_tensor * routw = wt(ctx_w, "octo.head.diffusion.reverse.out.weight"); + ggml_tensor * routb = wt(ctx_w, "octo.head.diffusion.reverse.out.bias"); + if (!time_w || !c0w || !c0b || !c1w || !c1b || !rinw || !rinb || !routw || !routb) return false; + + char rname[160]; + ggml_tensor * blk_w[3][6]; // ln.weight, ln.bias, fc1.weight, fc1.bias, fc2.weight, fc2.bias + const char * leaves[6] = {"ln.weight", "ln.bias", "fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"}; + for (int i = 0; i < 3; ++i) { + for (int j = 0; j < 6; ++j) { + std::snprintf(rname, sizeof(rname), "octo.head.diffusion.reverse.blk.%d.%s", i, leaves[j]); + blk_w[i][j] = wt(ctx_w, rname); + if (!blk_w[i][j]) return false; + } + } + + ggml_init_params gp = {(size_t) 16 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) return false; + + std::vector in_tensors; + std::vector> in_payloads; + auto add_input = [&](ggml_tensor * t, const std::vector& data) { + in_tensors.push_back(t); + in_payloads.push_back(data); + return t; + }; + + std::vector time_data(width, (float) time_value); + ggml_tensor * time = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, width); + ggml_set_name(time, "action_head.time"); + add_input(time, time_data); + ggml_tensor * obs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, width); + ggml_set_name(obs, "action_head.readout_embedding"); + add_input(obs, readout_action); + ggml_tensor * actions = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, action, width); + ggml_set_name(actions, "action_head.noisy_action"); + add_input(actions, noisy_action); + + ggml_tensor * f = ggml_scale(ctx, ggml_mul_mat(ctx, time_w, time), two_pi); + ggml_tensor * time_ff = ggml_concat(ctx, ggml_cos(ctx, f), ggml_sin(ctx, f), 0); + ggml_tensor * cond = ggml_silu(ctx, ggml_add(ctx, ggml_mul_mat(ctx, c0w, time_ff), c0b)); + cond = ggml_add(ctx, ggml_mul_mat(ctx, c1w, cond), c1b); + ggml_tensor * reverse_input = ggml_concat(ctx, ggml_concat(ctx, cond, obs, 0), actions, 0); + ggml_tensor * x = ggml_add(ctx, ggml_mul_mat(ctx, rinw, reverse_input), rinb); + for (int i = 0; i < 3; ++i) { + ggml_tensor * lnw = blk_w[i][0]; + ggml_tensor * lnb = blk_w[i][1]; + ggml_tensor * fc1w = blk_w[i][2]; + ggml_tensor * fc1b = blk_w[i][3]; + ggml_tensor * fc2w = blk_w[i][4]; + ggml_tensor * fc2b = blk_w[i][5]; + ggml_tensor * residual = x; + ggml_tensor * h = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, ln_eps), lnw), lnb); + h = ggml_silu(ctx, ggml_add(ctx, ggml_mul_mat(ctx, fc1w, h), fc1b)); + h = ggml_add(ctx, ggml_mul_mat(ctx, fc2w, h), fc2b); + x = ggml_add(ctx, residual, h); + } + ggml_tensor * out = ggml_add(ctx, ggml_mul_mat(ctx, routw, ggml_silu(ctx, x)), routb); + ggml_set_name(out, "action_head.pred_eps"); + ggml_set_output(out); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 2048, false); + ggml_build_forward_expand(graph, out); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): score actor ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + for (size_t i = 0; i < in_tensors.size(); ++i) { + ggml_backend_tensor_set(in_tensors[i], in_payloads[i].data(), 0, ggml_nbytes(in_tensors[i])); + } + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): score actor ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + pred_eps.resize((size_t) action * width); + ggml_backend_tensor_get(out, pred_eps.data(), 0, ggml_nbytes(out)); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +static bool load_f32_shape(const std::string& path, const std::vector& shape, std::vector& dst) { + NpyF32 npy; + if (!parse_npy_f32(path, npy)) return false; + if (npy.shape != shape) { + std::fprintf(stderr, "vla(octo): %s unexpected shape\n", path.c_str()); + return false; + } + dst = std::move(npy.data); + return true; +} + +// TIP-ND1-B: which noise the diffusion reverse process consumes. REPLAY reads golden +// z.npy/initial_noise.npy from a case dir (ctest parity -- deterministic, must match exactly). +// RANDOM samples N(0,1) from a caller-owned RNG (live predict()/free-sample -- stochastic by +// design, not compared bit-exact anywhere). CLIENT_NOISE is reserved for a future round +// (client-supplied deterministic noise, e.g. a PredictRequest.noise field) -- NOT implemented +// here; the slot exists so this signature doesn't need to change again when that lands. +enum class OctoNoiseSourceKind { + REPLAY, + RANDOM, + // CLIENT_NOISE, // TODO(in.noise): caller-supplied noise buffer, not implemented yet. +}; + +struct OctoNoiseSource { + OctoNoiseSourceKind kind; + std::string case_dir; // REPLAY only: golden case directory to read z.npy/initial_noise.npy from. + std::mt19937 * rng = nullptr; // RANDOM only: not owned by this struct. +}; + +// TIP-ND1-B: sole diffusion implementation (was run_diffusion_replay / _resident / +// run_diffusion_live / _resident, unified -- those 4 differed along two orthogonal axes: +// compute-source [gone now, resident-only] and noise-source [now `noise`]). REPLAY behavior is +// byte-for-byte the same as the deleted run_diffusion_replay_resident (same npy reads, same +// unconditional per-step z.npy read, same masked-noise-substitution using flat_action_mask). +// RANDOM behavior matches the deleted run_diffusion_live_resident exactly: z is only drawn +// from `noise.rng` when time_value > 0 (same RNG draw count/order as before -- the old code +// never drew a step-20 noise sample either, since DDPM's reverse process adds no noise at +// t=0), and flat_action_mask is all-valid so the masked-noise-substitution step below is +// structurally a no-op for it, matching the old live path that omitted that step outright. +static bool run_diffusion_resident(ggml_context * ctx_w, ggml_backend_t backend, + const OctoNoiseSource& noise, + const std::vector& readout_action, + int window_size, + int action_total, + OctoDiffusionResult& result) { + constexpr int steps = 20; + const int width = window_size; + const int action = action_total; + constexpr float max_action = 5.0f; + if (readout_action.size() != (size_t) 384 * width) return false; + + std::normal_distribution normal(0.0f, 1.0f); + std::string prefix; + if (noise.kind == OctoNoiseSourceKind::REPLAY) { + prefix = noise.case_dir + "/tensors/action_head.predict_action."; + if (!load_f32_shape(prefix + "initial_noise.npy", {1, window_size, action}, result.initial_noise)) return false; + NpyBool action_mask; + if (!parse_npy_bool(prefix + "action_mask.npy", action_mask) || action_mask.shape != std::vector({1, window_size, 4, 7})) return false; + result.action_mask = std::move(action_mask.data); + NpyBool flat_mask; + if (!parse_npy_bool(prefix + "flat_action_mask.npy", flat_mask) || flat_mask.shape != std::vector({1, window_size, action})) return false; + result.flat_action_mask = std::move(flat_mask.data); + } else { + result.initial_noise.resize((size_t) width * action); + for (float& v : result.initial_noise) v = normal(*noise.rng); + // All-true: octo-small-1.5 has real_action_dim == max_action_dim == 7 (full + // action_horizon), so this makes the masked-noise-substitution step below a + // structural no-op, matching the old live path that omitted it outright. + result.action_mask.assign((size_t) width * 4 * 7, 1); + result.flat_action_mask.assign((size_t) width * action, 1); + } + + OctoDiffusionSchedule sched = make_cosine_schedule(); + std::vector x = result.initial_noise; + for (int step = 0; step < steps; ++step) { + const int time_value = steps - 1 - step; + result.current_x_before[(size_t) step] = x; + if (!run_score_actor_graph_resident(ctx_w, backend, readout_action, x, time_value, window_size, action_total, result.pred_eps[(size_t) step])) return false; + + if (noise.kind == OctoNoiseSourceKind::REPLAY) { + char stem[96]; + std::snprintf(stem, sizeof(stem), "step_%02d.t_%02d.", step, time_value); + if (!load_f32_shape(prefix + stem + "z.npy", {1, window_size, action}, result.z[(size_t) step])) return false; + } else if (time_value > 0) { + result.z[(size_t) step].resize((size_t) width * action); + for (float& v : result.z[(size_t) step]) v = normal(*noise.rng); + } else { + result.z[(size_t) step].assign((size_t) width * action, 0.0f); // never read below (time_value == 0) + } + + std::vector y((size_t) width * action); + const float alpha = sched.alphas[(size_t) time_value]; + const float beta = sched.betas[(size_t) time_value]; + const float alpha_hat = sched.alpha_hats[(size_t) time_value]; + const float alpha_1 = 1.0f / std::sqrt(alpha); + const float alpha_2 = (1.0f - alpha) / std::sqrt(1.0f - alpha_hat); + for (size_t i = 0; i < y.size(); ++i) y[i] = alpha_1 * (x[i] - alpha_2 * result.pred_eps[(size_t) step][i]); + result.after_denoise[(size_t) step] = y; + if (time_value > 0) { + const float sigma = std::sqrt(beta); + for (size_t i = 0; i < y.size(); ++i) y[i] += sigma * result.z[(size_t) step][i]; + } + result.after_noise_add[(size_t) step] = y; + for (float& v : y) v = std::min(std::max(v, -max_action), max_action); + result.after_clip[(size_t) step] = y; + const float masked_noise_scale = std::sqrt(1.0f - alpha_hat); + for (size_t i = 0; i < y.size(); ++i) { + if (!result.flat_action_mask[i]) y[i] = masked_noise_scale * result.z[(size_t) step][i]; + } + result.after_mask[(size_t) step] = y; + x = std::move(y); + } + + result.actions_all_timesteps = x; + // Final action chunk = the LAST timestep's readout (index window_size-1), matching + // OctoPt's sample_actions() which returns actions[:, -1] after the block-transformer + // diffusion head runs over the full observation window. For window_size=2 this is + // exactly the old hardcoded [action, 2*action) slice; for window_size=1 it is the + // sole timestep's block [0, action). + const size_t final_begin = (size_t) (window_size - 1) * action; + const size_t final_end = (size_t) window_size * action; + if (final_end > x.size()) { + std::fprintf(stderr, "vla(octo): final action slice [%zu,%zu) out of bounds (x.size()=%zu)\n", + final_begin, final_end, x.size()); + return false; + } + result.final_actions.assign(x.begin() + (ptrdiff_t) final_begin, x.begin() + (ptrdiff_t) final_end); + return true; +} + +// TIP-05 Part B: L1 action head forward -- ContinuousActionHeadPt.forward/predict_action, +// replacing the diffusion head entirely when head_type=="l1" (routed in +// octo_run_pipeline_resident). Dump points for TIP-06 parity: map_attn_out (post out_proj, +// pre-LN/MLP), map_emb (post residual+MLP -- MAPHeadPt's final per-timestep embedding), +// mean_normalized (post mean_proj+tanh*max_action, ALL window timesteps), final_actions +// (mean_normalized's last window timestep only, what predict_action returns). +struct OctoL1HeadResult { + std::vector map_attn_out; // [384,window_size] + std::vector map_emb; // [384,window_size] + std::vector mean_normalized; // [action_total,window_size] + std::vector final_actions; // [action_total] -- mean_normalized[:, window_size-1] +}; + +// readout_action: [384,window_size] (the block-transformer's readout_action output, one +// 384-dim token per timestep -- readouts.action=1 so there is exactly one token/timestep +// already, matching MAPHeadPt's (b,w,1,384) input shape with the "1" dim implicit). No +// attention happens ACROSS window_size or across batch: MAPHeadPt's nn.MultiheadAttention +// runs independently per (batch,timestep) pair (query=probe, key=value=that timestep's +// single token) -- window_size is therefore modeled as a batch axis (ggml ne3) here, heads +// as a second batch axis (ne2), never mixed via the ne0/ne1 axes mul_mat actually contracts +// over/compares. See TIP-05 Completion Report for the full ggml shape derivation. +static bool run_l1_action_head_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + const std::vector& readout_action, + int window_size, + int action_total, + OctoL1HeadResult& result) { + constexpr int hidden = 384; + constexpr int map_heads = 8; // TIP-05: 8, NOT the block-transformer's 6. + constexpr int map_head_dim = hidden / map_heads; // 48 + constexpr float ln_eps = 1e-6f; + constexpr float max_action = 5.0f; + const int width = window_size; + if (readout_action.size() != (size_t) hidden * width || action_total <= 0) return false; + + ggml_tensor * probe_r = wt(ctx_w, "octo.head.l1.map.probe"); + ggml_tensor * qkv_w_r = wt(ctx_w, "octo.head.l1.map.attn_qkv.weight"); + ggml_tensor * qkv_b_r = wt(ctx_w, "octo.head.l1.map.attn_qkv.bias"); + ggml_tensor * o_w_r = wt(ctx_w, "octo.head.l1.map.attn_o.weight"); + ggml_tensor * o_b_r = wt(ctx_w, "octo.head.l1.map.attn_o.bias"); + ggml_tensor * norm_w_r = wt(ctx_w, "octo.head.l1.map.norm.weight"); + ggml_tensor * norm_b_r = wt(ctx_w, "octo.head.l1.map.norm.bias"); + ggml_tensor * ffn_up_w_r = wt(ctx_w, "octo.head.l1.map.ffn_up.weight"); + ggml_tensor * ffn_up_b_r = wt(ctx_w, "octo.head.l1.map.ffn_up.bias"); + ggml_tensor * ffn_down_w_r = wt(ctx_w, "octo.head.l1.map.ffn_down.weight"); + ggml_tensor * ffn_down_b_r = wt(ctx_w, "octo.head.l1.map.ffn_down.bias"); + ggml_tensor * mean_w_r = wt(ctx_w, "octo.head.l1.mean_proj.weight"); + ggml_tensor * mean_b_r = wt(ctx_w, "octo.head.l1.mean_proj.bias"); + if (!probe_r || !qkv_w_r || !qkv_b_r || !o_w_r || !o_b_r || !norm_w_r || !norm_b_r || + !ffn_up_w_r || !ffn_up_b_r || !ffn_down_w_r || !ffn_down_b_r || !mean_w_r || !mean_b_r) { + return false; + } + + ggml_init_params gp = {(size_t) 16 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) return false; + + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, width); + ggml_set_name(x, "l1_head.readout_action"); + + // probe is resident weight octo.head.l1.map.probe, PyTorch shape (1,1,384) -> GGUF + // drops the leading size-1 dims at load (ggml_n_dims), so it's already a plain + // ne=[384,1,1,1] resident tensor -- exactly the "[384,1] single query token" shape + // needed below, used directly with no reshape (same convention as every other + // resident weight in this file: mul_mat/add operands reference wt()'s result as-is + // unless a host-side transform is needed, per house style -- see conv weights below + // needing tensor_to_vec vs. attention weights that don't). + // + // Q comes from probe (query), K/V come from x (key=value=this timestep's readout + // token) -- two DIFFERENT inputs through the SAME combined in_proj_weight, so (unlike + // the block-transformer's self-attention, where q/k/v all come from one input and can + // share one mul_mat) each needs its own mul_mat against the full [384,1152] qkv + // weight; the slice not needed from each (k/v from the probe pass, q from the x pass) + // is simply left unused, same in_proj_weight both passes (matches + // nn.MultiheadAttention where in_proj_weight's 3 row-blocks are always [Wq;Wk;Wv] + // regardless of what query/key/value tensors get fed through it). + ggml_tensor * qkv_probe = ggml_add(ctx, ggml_mul_mat(ctx, qkv_w_r, probe_r), qkv_b_r); + ggml_tensor * q = ggml_cont(ctx, ggml_view_2d(ctx, qkv_probe, hidden, 1, qkv_probe->nb[1], 0)); + ggml_tensor * qkv_x = ggml_add(ctx, ggml_mul_mat(ctx, qkv_w_r, x), qkv_b_r); + ggml_tensor * k = ggml_cont(ctx, ggml_view_2d(ctx, qkv_x, hidden, width, qkv_x->nb[1], (size_t) hidden * qkv_x->nb[0])); + ggml_tensor * v = ggml_cont(ctx, ggml_view_2d(ctx, qkv_x, hidden, width, qkv_x->nb[1], (size_t) 2 * hidden * qkv_x->nb[0])); + + // Split into heads with heads on ne2 (a batch axis mul_mat loops over, never + // cross-multiplied) and width on ne3 (a SECOND batch axis) -- this is what keeps each + // timestep's attention independent (no cross-timestep mixing) while still doing + // per-head dot products correctly. Qh's ne3=1 broadcasts into Kh/Vh's ne3=width per + // ggml_can_mul_mat (t1->ne3 % t0->ne3 == 0, t0=the smaller/A operand). + ggml_tensor * Qh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, q, map_head_dim, map_heads, 1, 1), 0, 2, 1, 3)); + ggml_tensor * Kh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, k, map_head_dim, map_heads, 1, width), 0, 2, 1, 3)); + ggml_tensor * Vh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, v, map_head_dim, map_heads, 1, width), 1, 2, 0, 3)); + + ggml_tensor * scores = ggml_mul_mat(ctx, Qh, Kh); // [1(seq_q),1(seq_k),heads,width] + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + // seq_k==1 always (readouts.action==1 token/timestep) so this softmax is over a + // single logit -- always exactly 1.0 -- computed via the real op (not hand-simplified + // to "skip attention") so this stays correct if that ever changes, and so TIP-06 can + // dump a genuine attn-probability boundary. + ggml_tensor * probs = ggml_soft_max_ext(ctx, scores, nullptr, 1.0f / std::sqrt((float) map_head_dim), 0.0f); + ggml_tensor * attended = ggml_mul_mat(ctx, Vh, probs); // [head_dim,1,heads,width] + ggml_tensor * merged = ggml_reshape_2d(ctx, ggml_cont(ctx, ggml_permute(ctx, attended, 0, 2, 1, 3)), hidden, width); + + ggml_tensor * attn_out = ggml_add(ctx, ggml_mul_mat(ctx, o_w_r, merged), o_b_r); + ggml_set_name(attn_out, "l1_head.map.attn_out"); + ggml_set_output(attn_out); + + ggml_tensor * y = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, attn_out, ln_eps), norm_w_r), norm_b_r); + ggml_tensor * h = ggml_gelu_erf(ctx, ggml_add(ctx, ggml_mul_mat(ctx, ffn_up_w_r, y), ffn_up_b_r)); + h = ggml_add(ctx, ggml_mul_mat(ctx, ffn_down_w_r, h), ffn_down_b_r); + // Residual is onto attn_out (pre-LN), not onto y -- per TIP-05 spec, matches + // MlpBlockPt's usage inside MAPHeadPt (out = attn_out + MlpBlock(LayerNorm(attn_out))). + ggml_tensor * emb = ggml_add(ctx, attn_out, h); + ggml_set_name(emb, "l1_head.map.emb"); + ggml_set_output(emb); + + ggml_tensor * mean_raw = ggml_add(ctx, ggml_mul_mat(ctx, mean_w_r, emb), mean_b_r); + ggml_tensor * mean_final = ggml_scale(ctx, ggml_tanh(ctx, ggml_scale(ctx, mean_raw, 1.0f / max_action)), max_action); + ggml_set_name(mean_final, "l1_head.mean_normalized"); + ggml_set_output(mean_final); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 512, false); + ggml_build_forward_expand(graph, attn_out); + ggml_build_forward_expand(graph, emb); + ggml_build_forward_expand(graph, mean_final); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): L1 head ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(x, readout_action.data(), 0, ggml_nbytes(x)); + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): L1 head ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + + result.map_attn_out.resize((size_t) hidden * width); + ggml_backend_tensor_get(attn_out, result.map_attn_out.data(), 0, ggml_nbytes(attn_out)); + result.map_emb.resize((size_t) hidden * width); + ggml_backend_tensor_get(emb, result.map_emb.data(), 0, ggml_nbytes(emb)); + result.mean_normalized.resize((size_t) action_total * width); + if ((size_t) ggml_nelements(mean_final) != result.mean_normalized.size()) { + std::fprintf(stderr, "vla(octo): L1 head mean_proj out-dim=%lld does not match action_horizon*action_dim=%d\n", + (long long) mean_final->ne[0], action_total); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_get(mean_final, result.mean_normalized.data(), 0, ggml_nbytes(mean_final)); + + // predict_action: last window timestep only (matches diffusion's own final-action-slice + // convention -- see run_diffusion_resident's "final action slice" comment/the + // octo_action_slice tripwire test). + result.final_actions.assign(result.mean_normalized.end() - action_total, result.mean_normalized.end()); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +static 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; +} + +// Nearest-neighbor resize from interleaved HWC RGB8 (arbitrary size) to a +// dside x dside planar CHW RGB8 buffer, matching the [C,H,W] layout +// run_one_obs_tokenizer_graph expects for its obs/task tensors. +static void resize_to_planar_chw(const uint8_t * src_hwc, int sw, int sh, int dside, std::vector& dst_chw) { + dst_chw.assign((size_t) 3 * dside * dside, 0); + for (int yy = 0; yy < dside; ++yy) { + const int sy = std::min(sh - 1, (int) ((int64_t) yy * sh / dside)); + for (int xx = 0; xx < dside; ++xx) { + const int sx = std::min(sw - 1, (int) ((int64_t) xx * sw / dside)); + const uint8_t * px = src_hwc + ((size_t) sy * sw + sx) * 3; + for (int c = 0; c < 3; ++c) { + dst_chw[((size_t) c * dside + yy) * dside + xx] = px[c]; + } + } + } +} + +// Resolves which top-level key of octo.dataset_statistics to un-normalize against when +// the caller didn't pin one down explicitly (dataset_key_in empty). Priority: +// 1. VLA_OCTO_UNNORM_DATASET env var -- the per-checkpoint-config idiom this codebase +// already uses for things a loaded GGUF can't self-describe (c.f. gr00t's +// VLA_GR00T_EMBODIMENT); the only channel available to server predict(), which has +// no per-call CLI flag. +// 2. dataset_statistics has exactly one top-level key -> unambiguous, use it. +// 3. "bridge_dataset" if present -- preserves the historical hardcoded default for the +// rail-berkeley bridge pretrain checkpoint (whose dataset_statistics has ~25 OXE-mix +// keys, so rule 2 doesn't apply to it). +// Otherwise: fail loudly rather than silently guessing among several real candidates +// (e.g. cyrusneary's libero_object/libero_spatial/libero_goal/liber_o10). +static 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 auto-resolve unnorm dataset key (%zu candidate keys in " + "octo.dataset_statistics); set VLA_OCTO_UNNORM_DATASET or pass an explicit key " + "(e.g. --unnorm-dataset libero_object)\n", + stats.is_object() ? stats.size() : (size_t) 0); + return false; +} + +// TIP-05V: octo.dataset_statistics comes in two shapes depending on how many datasets the +// checkpoint's OctoModelPt.dataset_statistics covers -- nested-by-dataset-key (bridge/libero: +// {"bridge_dataset": {"action":..., "proprio":...}, "": {...}, ...}, needing +// resolve_unnorm_dataset_key to pick one) vs. FLAT single-dataset (the real +// octo-aloha-jitter2525.gguf checkpoint: {"action":..., "proprio":..., "num_transitions":..., +// "num_trajectories":...} directly, no dataset-name wrapper at all -- confirmed by inspecting +// the GGUF's raw KV bytes; resolve_unnorm_dataset_key previously misread this flat object's 4 +// members as 4 candidate dataset NAMES and failed loudly, since none of them is a +// single-key/bridge_dataset/env-var match). Detected by whether the object itself already has +// an "action" member shaped like a stats block (has "mean") -- real dataset names never +// collide with that. Every existing (nested) GGUF this codebase ships/tests against is +// unaffected: their top level never has a member literally named "action". +static 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 missing key %s\n", dataset_key.c_str()); + return false; + } + *out = &j[dataset_key]; + return true; +} + +// TIP-05 Part A step 1: proprio z-score stats -- octo.dataset_statistics[dataset_key].proprio +// (NOT .action; a sibling stat block in the same per-dataset JSON object). Convention: the +// caller (harness/TIP-07) sends RAW proprio (RLDS state, original units) via Inputs::state; +// this engine normalizes it here, once, before the proprio tokenizer runs -- mirrors +// unnormalize_action's dataset-key resolution (same resolve_stats_block, so both stat blocks +// always come from the same resolved dataset), but is otherwise a separate, self-contained +// JSON read (no mask field -- TIP-05 doesn't mention one for proprio, unlike action's +// gripper-dim mask). +static bool load_proprio_stats(gguf_reader& g, const std::string& dataset_key_in, + std::vector& mean, std::vector& stdv) { + 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("proprio")) { + std::fprintf(stderr, "vla(octo): dataset_statistics stats block missing .proprio\n"); + return false; + } + const auto& p = (*block)["proprio"]; + mean = p.at("mean").get>(); + stdv = p.at("std").get>(); + if (mean.size() != (size_t) kProprioTokens || stdv.size() != (size_t) kProprioTokens) { + std::fprintf(stderr, "vla(octo): unexpected dataset_statistics/proprio shape\n"); + return false; + } + return true; +} + +// Un-normalizes a [horizon,7] flattened action (horizon = normalized_flat.size()/7 -- +// action_dim is fixed at 7 across every head type, only action_horizon varies: 4 for +// diffusion libero, 20 for the L1 aloha jitter-adapted head) using octo.dataset_statistics's +// per-dataset action mean/std/mask: unnorm[d] = mask[d] ? norm[d]*std[d]+mean[d] +// : norm[d] (dims with mask=false, e.g. bridge_dataset's/libero_object's gripper dim 6, +// are left as-is -- confirmed against golden: final_action_unnormalized[...,6] == +// sample_actions.final_action_normalized[...,6] exactly, while masked-in dims match +// a*std+mean to ~1e-5). dataset_key_in must match golden's metadata.unnormalization.dataset +// when comparing against a golden trace (all shipped bridge golden cases use +// "bridge_dataset"); pass "" to auto-resolve via resolve_stats_block (live/serving paths, +// where there's no golden to match against; also transparently handles a flat/single-dataset +// stats blob with no dataset-name wrapper, e.g. octo-aloha-jitter2525.gguf -- TIP-05V). +static bool unnormalize_action(gguf_reader& g, const std::string& dataset_key_in, + const std::vector& normalized_flat, std::vector& unnorm_flat) { + 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 stats block missing .action\n"); + return false; + } + const auto& act = (*block)["action"]; + std::vector mean = act.at("mean").get>(); + std::vector stdv = act.at("std").get>(); + std::vector mask = act.at("mask").get>(); + const size_t dim = mask.size(); + if (mean.size() != dim || stdv.size() != dim || dim != 7 || + normalized_flat.empty() || normalized_flat.size() % dim != 0) { + std::fprintf(stderr, "vla(octo): unexpected dataset_statistics/action shape\n"); + return false; + } + const size_t horizon = normalized_flat.size() / dim; + unnorm_flat.resize(normalized_flat.size()); + for (size_t t = 0; t < horizon; ++t) { + for (size_t d = 0; d < dim; ++d) { + const float norm = normalized_flat[t * dim + d]; + unnorm_flat[t * dim + d] = mask[d] ? (norm * stdv[d] + mean[d]) : norm; + } + } + return true; +} + +// TIP-05: n_proprio_tokens=0 (the default every pre-TIP-05 call site passes) reproduces the +// old metadata/key_valid sequence byte-for-byte. n_proprio_tokens=7 inserts a proprio OBS +// group right after wrist (same position assemble_transformer_input uses -- the two must +// stay in lock-step, both keyed off the same n_proprio_tokens). Proprio validity is gated +// solely by timestep_valid (CAUSAL, same rule as primary/wrist/readout): unlike a camera, +// proprioception has no "was this sensor even connected" ambiguity to model with a separate +// per-token pad mask, so there's no proprio_valid input here (matches TIP-05 Part A step 4: +// "mask = timestep_pad_mask"). +static bool build_transformer_mask(const NpyBool& task_valid, + const NpyBool& primary_valid, + const NpyBool& wrist_valid, + const NpyBool& timestep_valid, + int window_size, + int n_proprio_tokens, + std::vector& blocked, + std::vector& blocked_f32) { + const int seq = octo_seq_len(window_size, n_proprio_tokens); + constexpr int heads = 6; + const int step_tokens = octo_step_tokens(n_proprio_tokens); + if (task_valid.shape != std::vector{1} || + primary_valid.shape != std::vector({1, window_size}) || + wrist_valid.shape != std::vector({1, window_size}) || + timestep_valid.shape != std::vector({1, window_size})) { + std::fprintf(stderr, "vla(octo): unexpected input pad-mask shape\n"); + return false; + } + + std::vector metadata; + std::vector key_valid; + metadata.reserve(seq); + key_valid.reserve(seq); + for (int i = 0; i < 16; ++i) { + metadata.push_back({OctoTokenKind::TASK, -1}); + key_valid.push_back(task_valid.data[0] != 0); + } + for (int t = 0; t < window_size; ++t) { + const uint8_t timestep_ok = timestep_valid.data[(size_t) t] != 0; + for (int i = 0; i < 256; ++i) { + metadata.push_back({OctoTokenKind::OBS, t}); + key_valid.push_back(timestep_ok && primary_valid.data[(size_t) t]); + } + for (int i = 0; i < 64; ++i) { + metadata.push_back({OctoTokenKind::OBS, t}); + key_valid.push_back(timestep_ok && wrist_valid.data[(size_t) t]); + } + for (int i = 0; i < n_proprio_tokens; ++i) { + metadata.push_back({OctoTokenKind::OBS, t}); + key_valid.push_back(timestep_ok); + } + for (int i = 0; i < 16; ++i) { + metadata.push_back({OctoTokenKind::OBS, t}); + key_valid.push_back(task_valid.data[0] != 0); + } + metadata.push_back({OctoTokenKind::READOUT, t}); + key_valid.push_back(1); + } + if (metadata.size() != (size_t) seq || key_valid.size() != (size_t) seq || + 16 + window_size * step_tokens != seq) return false; + + const size_t plane = (size_t) seq * seq; + std::vector blocked_one_head(plane, 0); + blocked_f32.assign(plane, 0.0f); + for (int q = 0; q < seq; ++q) { + const OctoTokenMetadata qm = metadata[(size_t) q]; + for (int k = 0; k < seq; ++k) { + const OctoTokenMetadata km = metadata[(size_t) k]; + bool allowed = false; + if (qm.kind == OctoTokenKind::TASK) { + allowed = km.kind == OctoTokenKind::TASK; + } else if (qm.kind == OctoTokenKind::OBS) { + allowed = km.kind == OctoTokenKind::TASK || + (km.kind == OctoTokenKind::OBS && km.timestep <= qm.timestep); + } else { + allowed = km.kind == OctoTokenKind::TASK || + (km.kind == OctoTokenKind::OBS && km.timestep <= qm.timestep) || + (km.kind == OctoTokenKind::READOUT && km.timestep <= qm.timestep); + } + const size_t index = (size_t) q * seq + k; + const bool is_blocked = !allowed || !key_valid[(size_t) k]; + blocked_one_head[index] = is_blocked ? 1 : 0; + blocked_f32[index] = is_blocked ? 1.0f : 0.0f; + } + } + blocked.resize((size_t) heads * plane); + for (int h = 0; h < heads; ++h) { + std::copy(blocked_one_head.begin(), blocked_one_head.end(), blocked.begin() + (size_t) h * plane); + } + return true; +} + +// TIP-ND1-B: sole block-transformer implementation (was run_transformer_graph / +// run_transformer_graph_resident, unified). All 12 blocks' attn/ffn weights + the final +// output_norm are referenced directly from the resident context; `input`/`blocked_mask` (this +// call's assembled sequence + mask) are genuinely new per-call data and still get uploaded +// fresh. Graph topology and op order are unchanged. +static bool run_transformer_graph_resident(ggml_context * ctx_w, ggml_backend_t backend, + const std::vector& input, + const std::vector& blocked_mask, + int window_size, + int n_proprio_tokens, + OctoTransformerResult& result) { + constexpr int hidden = 384; + constexpr int heads = 6; + constexpr int head_dim = 64; + const int seq = octo_seq_len(window_size, n_proprio_tokens); + constexpr float ln_eps = 1e-6f; + constexpr float attn_scale = 0.125f; + if (input.size() != (size_t) hidden * seq || blocked_mask.size() != (size_t) seq * seq) return false; + + char rname[160]; + ggml_tensor * blk_w[12][10]; // attn_norm.{w,b}, attn_qkv.{w,b}, attn_o.{w,b}, ffn_norm.{w,b}, ffn_up.{w,b}, ffn_down.{w,b} + const char * leaves[10] = {"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"}; + for (int i = 0; i < 12; ++i) { + for (int j = 0; j < 10; ++j) { + std::snprintf(rname, sizeof(rname), "octo.blk.%d.%s", i, leaves[j]); + blk_w[i][j] = wt(ctx_w, rname); + if (!blk_w[i][j]) return false; + } + } + // ffn_down.{weight,bias} didn't fit the 10-wide table above cleanly (name pattern differs + // only in the leaf, kept separate to avoid a confusing 12-wide array); fetched per-block below. + ggml_tensor * ffn_down_w[12]; + ggml_tensor * ffn_down_b[12]; + for (int i = 0; i < 12; ++i) { + std::snprintf(rname, sizeof(rname), "octo.blk.%d.ffn_down.weight", i); + ffn_down_w[i] = wt(ctx_w, rname); + std::snprintf(rname, sizeof(rname), "octo.blk.%d.ffn_down.bias", i); + ffn_down_b[i] = wt(ctx_w, rname); + if (!ffn_down_w[i] || !ffn_down_b[i]) return false; + } + ggml_tensor * out_w_r = wt(ctx_w, "octo.output_norm.weight"); + ggml_tensor * out_b_r = wt(ctx_w, "octo.output_norm.bias"); + if (!out_w_r || !out_b_r) return false; + + ggml_init_params gp = {(size_t) 32 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(gp); + if (!ctx) return false; + + std::vector tensors; + std::vector> payloads; + auto add_payload = [&](ggml_tensor * t, std::vector data) { + tensors.push_back(t); + payloads.push_back(std::move(data)); + return t; + }; + + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, seq); + ggml_set_name(x, "octo.block_transformer.input"); + add_payload(x, input); + ggml_tensor * mask_blocked = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, seq, seq); + ggml_set_name(mask_blocked, "octo.block_transformer.blocked_mask"); + add_payload(mask_blocked, blocked_mask); + ggml_tensor * mask = ggml_scale(ctx, ggml_repeat_4d(ctx, mask_blocked, seq, seq, heads, 1), -FLT_MAX); + ggml_set_name(mask, "octo.block_transformer.additive_mask"); + + std::array block_out{}; + for (int i = 0; i < 12; ++i) { + ggml_tensor * n1w = blk_w[i][0]; + ggml_tensor * n1b = blk_w[i][1]; + ggml_tensor * Wqkv = blk_w[i][2]; + ggml_tensor * bqkv = blk_w[i][3]; + ggml_tensor * Wo = blk_w[i][4]; + ggml_tensor * bo = blk_w[i][5]; + ggml_tensor * n2w = blk_w[i][6]; + ggml_tensor * n2b = blk_w[i][7]; + ggml_tensor * Wup = blk_w[i][8]; + ggml_tensor * bup = blk_w[i][9]; + ggml_tensor * Wdown = ffn_down_w[i]; + ggml_tensor * bdown = ffn_down_b[i]; + + ggml_tensor * n1 = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, ln_eps), n1w), n1b); + ggml_tensor * qkv = ggml_add(ctx, ggml_mul_mat(ctx, Wqkv, n1), bqkv); + ggml_tensor * q = ggml_cont(ctx, ggml_view_2d(ctx, qkv, hidden, seq, qkv->nb[1], 0)); + ggml_tensor * k = ggml_cont(ctx, ggml_view_2d(ctx, qkv, hidden, seq, qkv->nb[1], (size_t) hidden * qkv->nb[0])); + ggml_tensor * v = ggml_cont(ctx, ggml_view_2d(ctx, qkv, hidden, seq, qkv->nb[1], (size_t) 2 * hidden * qkv->nb[0])); + ggml_tensor * Q = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, q, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, k, head_dim, heads, seq), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, v, head_dim, heads, seq), 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(ctx, K, Q); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + ggml_tensor * probs = ggml_soft_max_ext(ctx, scores, mask, attn_scale, 0.0f); + ggml_tensor * attended = ggml_mul_mat(ctx, V, probs); + ggml_tensor * merged = ggml_reshape_2d(ctx, ggml_cont(ctx, ggml_permute(ctx, attended, 0, 2, 1, 3)), hidden, seq); + ggml_tensor * attn_out = ggml_add(ctx, ggml_mul_mat(ctx, Wo, merged), bo); + ggml_tensor * residual = ggml_add(ctx, x, attn_out); + ggml_tensor * n2 = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, residual, ln_eps), n2w), n2b); + ggml_tensor * mlp = ggml_add(ctx, ggml_mul_mat(ctx, Wup, n2), bup); + mlp = ggml_gelu_erf(ctx, mlp); + mlp = ggml_add(ctx, ggml_mul_mat(ctx, Wdown, mlp), bdown); + x = ggml_add(ctx, residual, mlp); + char name[64]; + std::snprintf(name, sizeof(name), "bt.blk%d.out", i); + ggml_set_name(x, name); + ggml_set_output(x); + block_out[(size_t) i] = x; + } + + ggml_tensor * output = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, ln_eps), out_w_r), out_b_r); + ggml_set_name(output, "bt.output"); + ggml_set_output(output); + const int step_tokens = octo_step_tokens(n_proprio_tokens); + // TIP-05: proprio (when present) sits right after wrist, before repeated-language -- + // same offsets assemble_transformer_input used to place it. n_proprio_tokens=0 + // collapses off_language/off_readout back to the original 336/352 constants. + constexpr size_t off_primary = 16; + constexpr size_t off_wrist = off_primary + 256; + constexpr size_t off_proprio = off_wrist + 64; + const size_t off_language = off_proprio + (size_t) n_proprio_tokens; + const size_t off_readout = off_language + 16; + ggml_tensor * split_task = ggml_cont(ctx, ggml_view_2d(ctx, output, hidden, 16, output->nb[1], 0)); + ggml_tensor * split_primary = ggml_cont(ctx, ggml_view_3d(ctx, output, hidden, 256, window_size, output->nb[1], (size_t) step_tokens * output->nb[1], off_primary * output->nb[1])); + ggml_tensor * split_wrist = ggml_cont(ctx, ggml_view_3d(ctx, output, hidden, 64, window_size, output->nb[1], (size_t) step_tokens * output->nb[1], off_wrist * output->nb[1])); + ggml_tensor * split_proprio = nullptr; + if (n_proprio_tokens > 0) { + split_proprio = ggml_cont(ctx, ggml_view_3d(ctx, output, hidden, n_proprio_tokens, window_size, output->nb[1], (size_t) step_tokens * output->nb[1], off_proprio * output->nb[1])); + } + ggml_tensor * split_language = ggml_cont(ctx, ggml_view_3d(ctx, output, hidden, 16, window_size, output->nb[1], (size_t) step_tokens * output->nb[1], off_language * output->nb[1])); + ggml_tensor * split_readout = ggml_cont(ctx, ggml_view_3d(ctx, output, hidden, 1, window_size, output->nb[1], (size_t) step_tokens * output->nb[1], off_readout * output->nb[1])); + for (ggml_tensor * t : {split_task, split_primary, split_wrist, split_language, split_readout}) ggml_set_output(t); + if (split_proprio) ggml_set_output(split_proprio); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 8192, false); + ggml_build_forward_expand(graph, split_task); + ggml_build_forward_expand(graph, split_primary); + ggml_build_forward_expand(graph, split_wrist); + if (split_proprio) ggml_build_forward_expand(graph, split_proprio); + ggml_build_forward_expand(graph, split_language); + ggml_build_forward_expand(graph, split_readout); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_alloc_graph(gallocr, graph)) { + std::fprintf(stderr, "vla(octo): transformer ggml_gallocr_alloc_graph failed\n"); + if (gallocr) ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + for (size_t i = 0; i < tensors.size(); ++i) { + ggml_backend_tensor_set(tensors[i], payloads[i].data(), 0, ggml_nbytes(tensors[i])); + } + const ggml_status st = ggml_backend_graph_compute(backend, graph); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(octo): transformer ggml_backend_graph_compute failed (%d)\n", (int) st); + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return false; + } + + for (int i = 0; i < 12; ++i) { + result.block_outputs[(size_t) i].resize((size_t) hidden * seq); + ggml_backend_tensor_get(block_out[(size_t) i], result.block_outputs[(size_t) i].data(), 0, ggml_nbytes(block_out[(size_t) i])); + } + auto get = [&](ggml_tensor * t, std::vector& dst) { + dst.resize(ggml_nelements(t)); + ggml_backend_tensor_get(t, dst.data(), 0, ggml_nbytes(t)); + }; + get(output, result.output); + get(split_task, result.task_language); + get(split_primary, result.obs_primary); + get(split_wrist, result.obs_wrist); + if (split_proprio) get(split_proprio, result.obs_proprio); + else result.obs_proprio.clear(); + get(split_language, result.obs_task_language); + get(split_readout, result.readout_action); + + ggml_gallocr_free(gallocr); + ggml_free(ctx); + return true; +} + +} // 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 M0 uses one GGUF)\n", mmproj_path.c_str()); + } + + auto m = std::make_unique(); + m->gguf_path = ckpt_path; + m->matmul_type = GGML_TYPE_F32; + + 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; + + m->backend = octo_select_backend(m->n_threads, ""); + if (!m->backend) return nullptr; + + // TIP-BUILD-OCTO-GPU-B: predict()'s 5 stages now consume these same weights (resident on + // `m->backend` -- CUDA when available, CPU otherwise) directly; no separate CPU-only copy. + if (!load_all_tensors(*m, m->io)) return nullptr; + std::printf("vla(octo): loaded %lld F32 tensors, hidden=%lld blocks=%lld heads=%lld horizon=%lld " + "action_dim=%lld head_type=%s window_size=%lld has_proprio=%s proprio_in_dim=%lld " + "proprio_tokens=%d\n", + (long long) m->tensors.size(), (long long) m->hidden, (long long) m->blocks, + (long long) m->heads, (long long) m->action_horizon, (long long) m->action_dim, + m->head_type.c_str(), (long long) m->window_size, m->has_proprio ? "true" : "false", + (long long) m->proprio_in_dim, m->has_proprio ? kProprioTokens : 0); + return m; +} + +bool octo_dump_gguf_inventory(const std::string& ckpt_path) { + gguf_init_params p{}; + p.no_alloc = true; + ggml_context * meta = nullptr; + p.ctx = &meta; + gguf_context * gctx = gguf_init_from_file(ckpt_path.c_str(), p); + if (!gctx) { + std::fprintf(stderr, "vla(octo): gguf_init_from_file failed for %s\n", ckpt_path.c_str()); + return false; + } + + const int64_t n_kv = gguf_get_n_kv(gctx); + const int64_t n_tensors = gguf_get_n_tensors(gctx); + std::printf("gguf_kv=%lld\n", (long long) n_kv); + std::printf("gguf_tensors=%lld\n", (long long) n_tensors); + for (int64_t i = 0; i < n_kv; ++i) { + const char * key = gguf_get_key(gctx, i); + std::printf("kv\t%s\ttype=%s\n", key, gguf_type_name(gguf_get_kv_type(gctx, i))); + } + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(gctx, i); + const ggml_tensor * t = ggml_get_tensor(meta, name); + std::printf("tensor\t%s\ttype=%s\tshape=[", name, ggml_type_name(gguf_get_tensor_type(gctx, i))); + const int nd = t ? ggml_n_dims(t) : 0; + for (int d = nd - 1; d >= 0; --d) { + std::printf("%lld%s", (long long) t->ne[d], d == 0 ? "" : ","); + } + std::printf("]\tbytes=%zu\n", gguf_get_tensor_size(gctx, i)); + } + gguf_free(gctx); + if (meta) ggml_free(meta); + return true; +} + +// TIP-ND1-B: sole dump implementation (was octo_dump_tokenizer_case / _resident, unified). +// Dumps every M0-M8 boundary needed by verify_octo_parity.py, computed via the *_resident +// stage functions (on `backend`, CUDA when available) using the SAME weight residency setup +// octo_create() builds for the live server (load_all_tensors onto a backend selected the +// identical way) -- this exercises the literal code OctoModelArch::predict() runs, not a +// separate parallel implementation. +bool octo_dump_tokenizer_case_resident(const std::string& ckpt_path, + const std::string& case_dir, + const std::string& dump_dir, + const std::string& t5_inject_path, + const std::string& unnorm_dataset) { + gguf_reader g{"octo"}; + if (!g.open(ckpt_path)) return false; + const int window_size = (int) g.u32("octo.window_size"); + if (window_size < 1 || window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%d out of supported range [1, %d]\n", + window_size, kMaxHorizon); + return false; + } + const int seq = octo_seq_len(window_size); + + // Backend selection + weight residency: byte-for-byte the same sequence octo_create() + // uses for the live server, so this dump path is built exactly like a real model load. + OctoModelArch m; + m.matmul_type = GGML_TYPE_F32; + m.action_horizon = g.u32("octo.action.horizon"); + m.action_dim = g.u32("octo.action.dim"); + const int action_total = (int) (m.action_horizon * m.action_dim); + m.backend = octo_select_backend(default_cpu_threads(), "[dump] "); + if (!m.backend) return false; + if (!load_all_tensors(m, g)) return false; + ggml_context * ctx_w = m.ctx_weights; + ggml_backend_t backend = m.backend; + + NpyU8 primary_obs, wrist_obs, primary_task, wrist_task; + bool wrist_obs_real = true; + if (!parse_npy_u8(case_dir + "/tensors/input.observation.image_primary.npy", primary_obs)) return false; + if (!parse_npy_u8_or_zero_obs(case_dir + "/tensors/input.observation.image_wrist.npy", + window_size, 128, wrist_obs, wrist_obs_real)) return false; + if (!parse_npy_u8_or_zero_task(case_dir + "/tensors/input.task.image_primary.npy", primary_obs, primary_task)) return false; + if (!parse_npy_u8_or_zero_task(case_dir + "/tensors/input.task.image_wrist.npy", wrist_obs, wrist_task)) return false; + + std::ofstream mf(dump_dir + "/manifest.txt"); + if (!mf) { + std::error_code ec; + std::filesystem::create_directories(dump_dir, ec); + if (ec) { + std::fprintf(stderr, "vla(octo): cannot create %s: %s\n", dump_dir.c_str(), ec.message().c_str()); + return false; + } + mf.open(dump_dir + "/manifest.txt"); + } + if (!mf) { + std::fprintf(stderr, "vla(octo): cannot write %s/manifest.txt\n", dump_dir.c_str()); + return false; + } + + std::vector tok, proj, pos, primary_pos, wrist_pos; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "primary", primary_obs, primary_task, 256, 256, window_size, tok, proj, pos)) return false; + if (!write_f32_dump(dump_dir, "obs.primary.tok", tok, {1, window_size, 256, 512}, mf)) return false; + if (!write_f32_dump(dump_dir, "obs.primary.proj", proj, {1, window_size, 256, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "obs.primary.pos", pos, {1, window_size, 256, 384}, mf)) return false; + primary_pos = pos; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "wrist", wrist_obs, wrist_task, 128, 64, window_size, tok, proj, pos)) return false; + if (!write_f32_dump(dump_dir, "obs.wrist.tok", tok, {1, window_size, 64, 512}, mf)) return false; + if (!write_f32_dump(dump_dir, "obs.wrist.proj", proj, {1, window_size, 64, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "obs.wrist.pos", pos, {1, window_size, 64, 384}, mf)) return false; + wrist_pos = pos; + + NpyI32 input_ids, attn_mask; + if (!parse_npy_i32(case_dir + "/tensors/input.task.language_instruction.input_ids.npy", input_ids)) return false; + if (!parse_npy_i32(case_dir + "/tensors/input.task.language_instruction.attention_mask.npy", attn_mask)) return false; + if (input_ids.shape != std::vector({1, 16}) || attn_mask.shape != std::vector({1, 16})) { + std::fprintf(stderr, "vla(octo): unexpected input_ids/attention_mask shape\n"); + return false; + } + std::vector t5_native_out; + if (!run_t5_encoder_graph_resident(ctx_w, backend, input_ids.data, attn_mask.data, t5_native_out)) return false; + if (!write_f32_dump(dump_dir, "t5.out", t5_native_out, {1, 16, 768}, mf)) return false; + + NpyF32 t5; + if (!t5_inject_path.empty()) { + if (!parse_npy_f32(t5_inject_path, t5)) return false; + } else { + t5.shape = {1, 16, 768}; + t5.data = t5_native_out; + } + std::vector lang_proj, lang_pos, repeated; + if (!run_language_graph_resident(ctx_w, backend, t5, window_size, lang_proj, lang_pos, repeated)) return false; + if (!write_f32_dump(dump_dir, "lang.proj", lang_proj, {1, 16, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "lang.pos", lang_pos, {1, 16, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "repeated_language", repeated, {1, window_size, 16, 384}, mf)) return false; + + ggml_tensor * readout_pos_r = wt(ctx_w, "octo.readout.action.pos_embd"); + if (!readout_pos_r) return false; + const std::vector readout_pos = tensor_to_vec(readout_pos_r); + + NpyBool task_valid, primary_valid, wrist_valid, timestep_valid; + if (!parse_npy_bool(case_dir + "/tensors/input.task.pad_mask_dict.language_instruction.npy", task_valid) || + !parse_npy_bool(case_dir + "/tensors/input.observation.pad_mask_dict.image_primary.npy", primary_valid) || + !parse_npy_bool_or_zero(case_dir + "/tensors/input.observation.pad_mask_dict.image_wrist.npy", window_size, wrist_valid) || + !parse_npy_bool(case_dir + "/tensors/input.observation.timestep_pad_mask.npy", timestep_valid)) return false; + if (!wrist_obs_real) { + std::fill(wrist_valid.data.begin(), wrist_valid.data.end(), (uint8_t) 0); + } + + OctoTransformerResult bt; + if (!assemble_transformer_input(lang_pos, primary_pos, wrist_pos, {}, repeated, readout_pos, window_size, /*n_proprio_tokens=*/0, bt.input)) return false; + std::vector blocked_mask; + if (!build_transformer_mask(task_valid, primary_valid, wrist_valid, timestep_valid, window_size, /*n_proprio_tokens=*/0, bt.blocked_mask, blocked_mask)) return false; + if (!write_f32_dump(dump_dir, "bt.input", bt.input, {1, seq, 384}, mf)) return false; + if (!write_bool_dump(dump_dir, "bt.mask", bt.blocked_mask, {6, seq, seq}, mf)) return false; + if (!run_transformer_graph_resident(ctx_w, backend, bt.input, blocked_mask, window_size, /*n_proprio_tokens=*/0, bt)) return false; + for (int i = 0; i < 12; ++i) { + char boundary[32]; + std::snprintf(boundary, sizeof(boundary), "bt.blk%d.out", i); + if (!write_f32_dump(dump_dir, boundary, bt.block_outputs[(size_t) i], {1, seq, 384}, mf)) return false; + } + if (!write_f32_dump(dump_dir, "bt.output", bt.output, {1, seq, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "bt.task_language", bt.task_language, {1, 16, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "bt.obs_primary", bt.obs_primary, {1, window_size, 256, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "bt.obs_wrist", bt.obs_wrist, {1, window_size, 64, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "bt.obs_task_language", bt.obs_task_language, {1, window_size, 16, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "bt.readout_action", bt.readout_action, {1, window_size, 1, 384}, mf)) return false; + + OctoDiffusionResult diff; + const OctoNoiseSource replay_noise{OctoNoiseSourceKind::REPLAY, case_dir, nullptr}; + if (!run_diffusion_resident(ctx_w, backend, replay_noise, bt.readout_action, window_size, action_total, diff)) return false; + if (!write_f32_dump(dump_dir, "diff.initial_noise", diff.initial_noise, {1, window_size, action_total}, mf)) return false; + if (!write_bool_dump(dump_dir, "diff.action_mask", diff.action_mask, {1, window_size, 4, 7}, mf)) return false; + if (!write_bool_dump(dump_dir, "diff.flat_action_mask", diff.flat_action_mask, {1, window_size, action_total}, mf)) return false; + for (int step = 0; step < 20; ++step) { + char boundary[64]; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.current_x_before", step); + if (!write_f32_dump(dump_dir, boundary, diff.current_x_before[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.pred_eps", step); + if (!write_f32_dump(dump_dir, boundary, diff.pred_eps[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.z", step); + if (!write_f32_dump(dump_dir, boundary, diff.z[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.x_after_denoise", step); + if (!write_f32_dump(dump_dir, boundary, diff.after_denoise[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.x_after_noise_add", step); + if (!write_f32_dump(dump_dir, boundary, diff.after_noise_add[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.x_after_clip", step); + if (!write_f32_dump(dump_dir, boundary, diff.after_clip[(size_t) step], {1, window_size, action_total}, mf)) return false; + std::snprintf(boundary, sizeof(boundary), "diff.step%02d.x_after_mask", step); + if (!write_f32_dump(dump_dir, boundary, diff.after_mask[(size_t) step], {1, window_size, action_total}, mf)) return false; + } + if (!write_f32_dump(dump_dir, "diff.actions_all_timesteps", diff.actions_all_timesteps, {1, window_size, 4, 7}, mf)) return false; + if (!write_f32_dump(dump_dir, "sample_actions.final_action_normalized", diff.final_actions, {1, 4, 7}, mf)) return false; + if (!write_f32_dump(dump_dir, "action_final", diff.final_actions, {1, 4, 7}, mf)) return false; + + std::vector unnorm_actions; + if (!unnormalize_action(g, unnorm_dataset, diff.final_actions, unnorm_actions)) return false; + if (!write_f32_dump(dump_dir, "action_final_unnormalized", unnorm_actions, {1, 4, 7}, mf)) return false; + return true; +} + +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") : 16; + + 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)); // reserve 1 slot for EOS + 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; +} + +// TIP-ND1-B: sole pipeline implementation (was octo_run_pipeline / _resident, unified). Shared +// tail of the live/cold-start prediction path: given already-assembled per-view +// observation+task images (obs.shape={1,window_size,3,side,side}, task.shape={1,3,side,side}) +// and already-tokenized language, runs obs tokenizer x2 -> T5 encoder -> language -> +// assemble -> causal mask -> block transformer -> diffusion head -> unnormalize, entirely via +// the resident stage functions on `backend` (the model's real backend: CUDA when available, +// CPU otherwise), referencing weights resident on `ctx_w` (populated once by load_all_tensors). +// `io` is the caller's already-open gguf_reader -- unnormalize_action only reads the small +// in-memory "octo.dataset_statistics" metadata string via it (metadata is parsed once at +// gguf_reader::open() time, so this is not a disk touch as long as `io` is already open). Used +// by both octo_predict_from_images (CLI: tokenizes its own instruction via SentencePiece) and +// OctoModelArch::predict (server: receives already-tokenized Inputs::lang_tokens from the +// client, matching every other arch's client-tokenizes/server-embeds convention). +// wrist_real=false marks the wrist observation as a zero-fill placeholder (single-camera +// checkpoint or a caller that supplied no wrist view) -- its causal-mask key_valid entries +// are forced false so the transformer treats it as absent padding, not a real observation. +// +// TIP-05: has_proprio/proprio_in_dim/proprio_state_raw are additive -- has_proprio=false +// (every pre-TIP-05 checkpoint) skips Part A entirely (n_proprio=0, byte-identical sequence +// geometry to before). head_type routes Part B: "l1" runs the L1 action head instead of +// diffusion (bypassing the whole noise/denoising-loop pipeline per TIP-05's routing spec); +// "diffusion" keeps the untouched pre-TIP-05 path. The two are independent switches (a +// checkpoint could in principle have proprio without an L1 head or vice versa; every actual +// checkpoint pairs them, but nothing here assumes that). +static bool octo_run_pipeline_resident(ggml_context * ctx_w, ggml_backend_t backend, gguf_reader& io, int window_size, + int action_total, const std::string& head_type, + bool has_proprio, int proprio_in_dim, + const NpyU8& primary_obs, const NpyU8& primary_task, + const NpyU8& wrist_obs, const NpyU8& wrist_task, bool wrist_real, + const std::vector& input_ids, + const std::vector& attention_mask, + const std::vector& proprio_state_raw, + const std::string& unnorm_dataset, + std::vector& normalized_out, + std::vector& unnormalized_out, + float& ms_vision_out, + float& ms_inference_out) { + using clock = std::chrono::steady_clock; + + if (head_type != "diffusion" && head_type != "l1") { + std::fprintf(stderr, "vla(octo): head_type=%s forward not implemented\n", head_type.c_str()); + return false; + } + const int n_proprio = has_proprio ? kProprioTokens : 0; + + const auto t_vision0 = clock::now(); + std::vector tok, primary_proj, primary_pos, wrist_proj, wrist_pos; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "primary", primary_obs, primary_task, 256, 256, window_size, tok, primary_proj, primary_pos)) return false; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "wrist", wrist_obs, wrist_task, 128, 64, window_size, tok, wrist_proj, wrist_pos)) return false; + std::vector proprio_pos; + if (has_proprio) { + std::vector proprio_mean, proprio_std; + if (!load_proprio_stats(io, unnorm_dataset, proprio_mean, proprio_std)) return false; + std::vector proprio_norm((size_t) window_size * kProprioTokens); + for (int t = 0; t < window_size; ++t) { + for (int d = 0; d < kProprioTokens; ++d) { + const float raw = (d < (int) proprio_state_raw.size()) ? proprio_state_raw[(size_t) d] : 0.0f; + proprio_norm[(size_t) t * kProprioTokens + d] = (raw - proprio_mean[(size_t) d]) / proprio_std[(size_t) d]; + } + } + if (!run_proprio_tokenizer_graph_resident(ctx_w, backend, proprio_norm, proprio_in_dim, window_size, proprio_pos)) return false; + } + ms_vision_out = std::chrono::duration(clock::now() - t_vision0).count(); + + const auto t_inference0 = clock::now(); + std::vector t5_out; + if (!run_t5_encoder_graph_resident(ctx_w, backend, input_ids, attention_mask, t5_out)) return false; + + NpyF32 t5; + t5.shape = {1, 16, 768}; + t5.data = t5_out; + std::vector lang_proj, lang_pos, repeated; + if (!run_language_graph_resident(ctx_w, backend, t5, window_size, lang_proj, lang_pos, repeated)) return false; + + ggml_tensor * readout_pos_r = wt(ctx_w, "octo.readout.action.pos_embd"); + if (!readout_pos_r) return false; + const std::vector readout_pos = tensor_to_vec(readout_pos_r); + + NpyBool task_valid, primary_valid, wrist_valid, timestep_valid; + task_valid.shape = {1}; + task_valid.data = {1}; + primary_valid.shape = {1, window_size}; + primary_valid.data.assign((size_t) window_size, 1); + wrist_valid.shape = {1, window_size}; + wrist_valid.data.assign((size_t) window_size, wrist_real ? 1 : 0); + timestep_valid.shape = {1, window_size}; + timestep_valid.data.assign((size_t) window_size, 0); + timestep_valid.data.back() = 1; // cold start: only the last slot is the live frame (matches HistoryWrapper.reset) + + OctoTransformerResult bt; + if (!assemble_transformer_input(lang_pos, primary_pos, wrist_pos, proprio_pos, repeated, readout_pos, window_size, n_proprio, bt.input)) return false; + std::vector blocked_mask; + if (!build_transformer_mask(task_valid, primary_valid, wrist_valid, timestep_valid, window_size, n_proprio, bt.blocked_mask, blocked_mask)) return false; + if (!run_transformer_graph_resident(ctx_w, backend, bt.input, blocked_mask, window_size, n_proprio, bt)) return false; + + if (head_type == "l1") { + OctoL1HeadResult l1; + if (!run_l1_action_head_graph_resident(ctx_w, backend, bt.readout_action, window_size, action_total, l1)) return false; + normalized_out = std::move(l1.final_actions); + } else { + std::random_device rd; + std::mt19937 rng(rd()); + OctoDiffusionResult diff; + const OctoNoiseSource live_noise{OctoNoiseSourceKind::RANDOM, "", &rng}; + if (!run_diffusion_resident(ctx_w, backend, live_noise, bt.readout_action, window_size, action_total, diff)) return false; + normalized_out = std::move(diff.final_actions); + } + ms_inference_out = std::chrono::duration(clock::now() - t_inference0).count(); + + return unnormalize_action(io, unnorm_dataset, normalized_out, unnormalized_out); +} + +// Cold start: only ONE live frame is available (a single vla-cli snapshot or server +// request), but the model expects `window_size` observation timesteps. Repeat the live +// frame into every slot; the caller's timestep_valid marks all-but-the-last as padding +// (see octo_run_pipeline_resident) so the causal mask treats only the last slot as "real" -- matches +// OctoPt's HistoryWrapper.reset() cold start for window_size=2, and generalizes to +// window_size=1 (no padding slot at all, the single slot IS the live frame). No goal image +// (language-only conditioning): task is zero-filled, matching OctoModelPt.create_tasks(texts=...). +static void octo_build_cold_start_obs_task(const uint8_t* rgb, int sw, int sh, int side, + int window_size, NpyU8& obs, NpyU8& task) { + std::vector frame; + resize_to_planar_chw(rgb, sw, sh, side, frame); + obs.shape = {1, window_size, 3, side, side}; + obs.data.resize((size_t) window_size * 3 * side * side); + for (int t = 0; t < window_size; ++t) { + std::copy(frame.begin(), frame.end(), obs.data.begin() + (ptrdiff_t) t * (ptrdiff_t) frame.size()); + } + task.shape = {1, 3, side, side}; + task.data.assign((size_t) 3 * side * side, 0); +} + +bool octo_predict_from_images(const std::string& ckpt_path, + const uint8_t* primary_rgb, int primary_w, int primary_h, + const uint8_t* wrist_rgb, int wrist_w, int wrist_h, + const std::string& instruction, + OctoCliAction& out, + const std::string& unnorm_dataset) { + gguf_reader g{"octo"}; + if (!g.open(ckpt_path)) return false; + const int window_size = (int) g.u32("octo.window_size"); + if (window_size < 1 || window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%d out of supported range [1, %d]\n", + window_size, kMaxHorizon); + return false; + } + + // TIP-ND1-B: CLI one-shot resident load -- same backend-selection + load_all_tensors + // sequence octo_create()/octo_dump_tokenizer_case_resident() use, scoped to this single + // call (no persistent OctoModelArch survives past this function, unlike the live server). + OctoModelArch m; + m.matmul_type = GGML_TYPE_F32; + m.action_horizon = g.u32("octo.action.horizon"); + m.action_dim = g.u32("octo.action.dim"); + 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); + m.backend = octo_select_backend(default_cpu_threads(), "[cli] "); + if (!m.backend) return false; + if (!load_all_tensors(m, g)) return false; + + std::vector input_ids, attention_mask; + if (!octo_tokenize_text(ckpt_path, instruction, input_ids, attention_mask)) return false; + + NpyU8 primary_obs, primary_task, wrist_obs, wrist_task; + octo_build_cold_start_obs_task(primary_rgb, primary_w, primary_h, 256, window_size, primary_obs, primary_task); + octo_build_cold_start_obs_task(wrist_rgb, wrist_w, wrist_h, 128, window_size, wrist_obs, wrist_task); + + // TIP-05: octo_predict_from_images (the images-only CLI entry point, octo.h) has no + // proprio parameter -- an L1/proprio checkpoint driven from the CLI gets an all-zero + // proprio reading (z-scored to (0-mean)/std, NOT world-unit zero) rather than a real + // robot state. Fine for the CLI's smoke-test purpose; the server path (predict() below) + // is what actually threads Inputs::state through. + const std::vector proprio_state_raw; + + float ms_vision = 0.f, ms_inference = 0.f; + const int action_total = (int) (m.action_horizon * m.action_dim); + return octo_run_pipeline_resident(m.ctx_weights, m.backend, g, window_size, action_total, m.head_type, + m.has_proprio, (int) m.proprio_in_dim, + primary_obs, primary_task, wrist_obs, wrist_task, + /*wrist_real=*/true, input_ids, attention_mask, proprio_state_raw, unnorm_dataset, + out.normalized, out.unnormalized, ms_vision, ms_inference); +} + +// TIP-06: sole L1/proprio stagewise dump implementation. Mirrors +// octo_dump_tokenizer_case_resident's structure/boilerplate (backend selection, +// load_all_tensors, manifest handling) but drives the head_type=l1 resident stage functions +// (TIP-05) instead of the diffusion ones, and reads its input from a case_dir written by the +// Part A python dump script rather than a JAX/OctoPt golden trace directory. +bool octo_dump_l1_stagewise_case_resident(const std::string& ckpt_path, + const std::string& case_dir, + const std::string& dump_dir, + const std::string& unnorm_dataset) { + gguf_reader g{"octo"}; + if (!g.open(ckpt_path)) return false; + const int window_size = (int) g.u32("octo.window_size"); + if (window_size < 1 || window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%d out of supported range [1, %d]\n", + window_size, kMaxHorizon); + return false; + } + + OctoModelArch m; + m.matmul_type = GGML_TYPE_F32; + m.action_horizon = g.u32("octo.action.horizon"); + m.action_dim = g.u32("octo.action.dim"); + 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.head_type != "l1" || !m.has_proprio) { + std::fprintf(stderr, "vla(octo): octo_dump_l1_stagewise_case_resident requires head_type=l1 " + "and a proprio tokenizer (got head_type=%s has_proprio=%s)\n", + m.head_type.c_str(), m.has_proprio ? "true" : "false"); + return false; + } + const int action_total = (int) (m.action_horizon * m.action_dim); + m.backend = octo_select_backend(default_cpu_threads(), "[l1-dump] "); + if (!m.backend) return false; + if (!load_all_tensors(m, g)) return false; + ggml_context * ctx_w = m.ctx_weights; + ggml_backend_t backend = m.backend; + + NpyU8 top_hwc, wrist_hwc; + if (!parse_npy_u8(case_dir + "/input.top_hwc_u8.npy", top_hwc)) return false; + if (!parse_npy_u8(case_dir + "/input.wrist_hwc_u8.npy", wrist_hwc)) return false; + if (top_hwc.shape != std::vector{256, 256, 3} || wrist_hwc.shape != std::vector{128, 128, 3}) { + std::fprintf(stderr, "vla(octo): unexpected l1-dump input image shape\n"); + return false; + } + NpyF32 proprio_raw_npy; + if (!parse_npy_f32(case_dir + "/input.proprio_raw.npy", proprio_raw_npy) || + proprio_raw_npy.data.size() != (size_t) kProprioTokens) { + std::fprintf(stderr, "vla(octo): unexpected l1-dump input.proprio_raw.npy shape\n"); + return false; + } + std::ifstream instr_f(case_dir + "/input.instruction.txt"); + if (!instr_f) { + std::fprintf(stderr, "vla(octo): cannot open %s/input.instruction.txt\n", case_dir.c_str()); + return false; + } + const std::string instruction((std::istreambuf_iterator(instr_f)), std::istreambuf_iterator()); + + std::vector input_ids, attention_mask; + if (!octo_tokenize_text(ckpt_path, instruction, input_ids, attention_mask)) return false; + + NpyU8 primary_obs, primary_task, wrist_obs, wrist_task; + octo_build_cold_start_obs_task(top_hwc.data.data(), 256, 256, 256, window_size, primary_obs, primary_task); + octo_build_cold_start_obs_task(wrist_hwc.data.data(), 128, 128, 128, window_size, wrist_obs, wrist_task); + + std::ofstream mf(dump_dir + "/manifest.txt"); + if (!mf) { + std::error_code ec; + std::filesystem::create_directories(dump_dir, ec); + if (ec) { + std::fprintf(stderr, "vla(octo): cannot create %s: %s\n", dump_dir.c_str(), ec.message().c_str()); + return false; + } + mf.open(dump_dir + "/manifest.txt"); + } + if (!mf) { + std::fprintf(stderr, "vla(octo): cannot write %s/manifest.txt\n", dump_dir.c_str()); + return false; + } + + // T0: proprio z-score normalize (same formula/stats source as octo_run_pipeline_resident) + // + proprio tokenizer forward (proj + pos_embd, PRE-transformer). + std::vector proprio_mean, proprio_std; + if (!load_proprio_stats(g, unnorm_dataset, proprio_mean, proprio_std)) return false; + std::vector proprio_norm((size_t) kProprioTokens); + for (int d = 0; d < kProprioTokens; ++d) { + proprio_norm[(size_t) d] = (proprio_raw_npy.data[(size_t) d] - proprio_mean[(size_t) d]) / proprio_std[(size_t) d]; + } + if (!write_f32_dump(dump_dir, "t0.proprio_normalized", proprio_norm, {kProprioTokens}, mf)) return false; + + std::vector proprio_tokens; + if (!run_proprio_tokenizer_graph_resident(ctx_w, backend, proprio_norm, (int) m.proprio_in_dim, window_size, proprio_tokens)) return false; + if (!write_f32_dump(dump_dir, "t0.proprio_tokens", proprio_tokens, {1, window_size, kProprioTokens, 384}, mf)) return false; + + std::vector tok, primary_proj, primary_pos, wrist_proj, wrist_pos; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "primary", primary_obs, primary_task, 256, 256, window_size, tok, primary_proj, primary_pos)) return false; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "wrist", wrist_obs, wrist_task, 128, 64, window_size, tok, wrist_proj, wrist_pos)) return false; + + std::vector t5_out; + if (!run_t5_encoder_graph_resident(ctx_w, backend, input_ids, attention_mask, t5_out)) return false; + NpyF32 t5; + t5.shape = {1, 16, 768}; + t5.data = t5_out; + std::vector lang_proj, lang_pos, repeated; + if (!run_language_graph_resident(ctx_w, backend, t5, window_size, lang_proj, lang_pos, repeated)) return false; + + ggml_tensor * readout_pos_r = wt(ctx_w, "octo.readout.action.pos_embd"); + if (!readout_pos_r) return false; + const std::vector readout_pos = tensor_to_vec(readout_pos_r); + + NpyBool task_valid, primary_valid, wrist_valid, timestep_valid; + task_valid.shape = {1}; + task_valid.data = {1}; + primary_valid.shape = {1, window_size}; + primary_valid.data.assign((size_t) window_size, 1); + wrist_valid.shape = {1, window_size}; + wrist_valid.data.assign((size_t) window_size, 1); // real wrist image, always valid + timestep_valid.shape = {1, window_size}; + timestep_valid.data.assign((size_t) window_size, 0); + timestep_valid.data.back() = 1; // cold start: only the last slot is the live frame + + const int n_proprio = kProprioTokens; + OctoTransformerResult bt; + if (!assemble_transformer_input(lang_pos, primary_pos, wrist_pos, proprio_tokens, repeated, readout_pos, window_size, n_proprio, bt.input)) return false; + std::vector blocked_mask; + if (!build_transformer_mask(task_valid, primary_valid, wrist_valid, timestep_valid, window_size, n_proprio, bt.blocked_mask, blocked_mask)) return false; + if (!run_transformer_graph_resident(ctx_w, backend, bt.input, blocked_mask, window_size, n_proprio, bt)) return false; + + // T1/T2: readout_action tokens -- transformer OUTPUT, also the L1 head's INPUT. + if (!write_f32_dump(dump_dir, "t1t2.readout_action", bt.readout_action, {1, window_size, 1, 384}, mf)) return false; + + OctoL1HeadResult l1; + if (!run_l1_action_head_graph_resident(ctx_w, backend, bt.readout_action, window_size, action_total, l1)) return false; + // T3: MAPHead internals (post out_proj pre-LN/MLP; post residual+MLP). + if (!write_f32_dump(dump_dir, "t3.map_attn_out", l1.map_attn_out, {1, window_size, 384}, mf)) return false; + if (!write_f32_dump(dump_dir, "t3.map_emb", l1.map_emb, {1, window_size, 1, 384}, mf)) return false; + // T4: mean_proj + tanh*max_action, ALL window timesteps. + if (!write_f32_dump(dump_dir, "t4.mean_normalized", l1.mean_normalized, + {1, window_size, (int) m.action_horizon, (int) m.action_dim}, mf)) return false; + + // T5: final unnormalized action (last window timestep only). + std::vector unnorm_actions; + if (!unnormalize_action(g, unnorm_dataset, l1.final_actions, unnorm_actions)) return false; + if (!write_f32_dump(dump_dir, "t5.action_unnormalized", unnorm_actions, + {1, (int) m.action_horizon, (int) m.action_dim}, mf)) return false; + return true; +} + +// Server-facing entry point (vla-server, TIP-CLIENT): in.images[0] is the primary view, +// in.images[1] the wrist view if the client sent one (single-camera checkpoints/clients +// omit it -- zero-filled and masked invalid via octo_run_pipeline_resident's wrist_real=false, same +// fallback as the golden-case loaders above). Unlike most other archs' predict(), this +// returns the UN-normalized (world-unit) action rather than the normalized one: Octo's +// dataset_statistics lives embedded in the multi-hundred-MB checkpoint GGUF itself (not a +// small sibling stats.json a client can cheaply hold locally the way gr00t's +// --stats-json works), so un-normalizing server-side and reusing the already +// golden-verified unnormalize_action() is the only practical option without shipping the +// whole GGUF to the client just to read its JSON metadata. See TIP-CLIENT report for the +// full rationale; the client (adapters.py) only needs to invert+binarize the gripper dim, +// not repeat the mean/std un-normalization. +std::vector OctoModelArch::predict(const Inputs& in) { + const auto t_total0 = std::chrono::steady_clock::now(); + if (in.n_images < 1 || !in.images) { + std::fprintf(stderr, "vla(octo): predict needs at least 1 image (primary)\n"); + return {}; + } + if (in.images[0].format != PixelFormat::U8) { + std::fprintf(stderr, "vla(octo): predict only supports PixelFormat::U8 images " + "(client must send RGB_U8, already rotated+resized -- TIP-P)\n"); + return {}; + } + const bool wrist_real = in.n_images >= 2; + if (wrist_real && in.images[1].format != PixelFormat::U8) { + std::fprintf(stderr, "vla(octo): predict only supports PixelFormat::U8 images\n"); + return {}; + } + if (in.n_lang != (int) language_tokens || in.attention_mask_n != (int) language_tokens || !in.attention_mask) { + std::fprintf(stderr, + "vla(octo): predict expects lang_tokens AND attention_mask of exactly %lld " + "entries each (client tokenizes with t5-base, max_length=%lld, " + "padding=\"max_length\" -- Octo's T5 encoder needs real padding info, unlike " + "archs that derive their own mask); got n_lang=%d attention_mask=%s attention_mask_n=%d\n", + (long long) language_tokens, (long long) language_tokens, in.n_lang, + in.attention_mask ? "set" : "null", in.attention_mask_n); + return {}; + } + + // TIP-BUILD-OCTO-GPU-B: no per-call GGUF reopen -- weights are resident on `backend` + // (loaded once in octo_create via load_all_tensors), and the 5 stages now compute on + // that same real backend (CUDA when available, CPU otherwise) instead of a fresh + // per-call CPU backend. `io` (opened once in octo_create too) covers the one remaining + // in-memory metadata read (octo.dataset_statistics, inside unnormalize_action). + + NpyU8 primary_obs, primary_task, wrist_obs, wrist_task; + octo_build_cold_start_obs_task((const uint8_t*) in.images[0].data, in.images[0].w, in.images[0].h, + 256, (int) window_size, primary_obs, primary_task); + if (wrist_real) { + octo_build_cold_start_obs_task((const uint8_t*) in.images[1].data, in.images[1].w, in.images[1].h, + 128, (int) window_size, wrist_obs, wrist_task); + } else { + wrist_obs.shape = {1, window_size, 3, 128, 128}; + wrist_obs.data.assign((size_t) window_size * 3 * 128 * 128, 0); + wrist_task.shape = {1, 3, 128, 128}; + wrist_task.data.assign((size_t) 3 * 128 * 128, 0); + } + + const std::vector input_ids(in.lang_tokens, in.lang_tokens + in.n_lang); + const std::vector attention_mask(in.attention_mask, in.attention_mask + in.attention_mask_n); + + // TIP-05: Inputs::state is RAW proprio (RLDS state, original units) per the doc comment + // in model.h -- octo_run_pipeline_resident z-scores it internally via load_proprio_stats + // (harness/TIP-07 must NOT pre-normalize, to avoid double-normalizing). Missing + // Inputs::state on a proprio checkpoint zero-fills rather than hard-failing (matches + // pi0.cpp's `in.state ? in.state[i] : 0.f` convention for the same field), but is + // surfaced loudly since a silently-zeroed proprio reading will silently skew every + // predicted action. + std::vector proprio_state_raw; + if (has_proprio) { + proprio_state_raw.assign((size_t) kProprioTokens, 0.0f); + if (in.state) { + for (int64_t d = 0; d < kProprioTokens; ++d) proprio_state_raw[(size_t) d] = in.state[d]; + } else { + std::fprintf(stderr, "vla(octo): predict: proprio checkpoint but Inputs::state is null -- " + "using all-zero proprio (z-scored, not world-unit zero)\n"); + } + } + + std::vector normalized, unnormalized; + float ms_vision = 0.f, ms_inference = 0.f; + const int action_total = (int) (action_horizon * action_dim); + if (!octo_run_pipeline_resident(ctx_weights, backend, io, (int) window_size, action_total, head_type, + has_proprio, (int) proprio_in_dim, + primary_obs, primary_task, wrist_obs, wrist_task, + wrist_real, input_ids, attention_mask, proprio_state_raw, /*unnorm_dataset=*/"", + normalized, 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; +} + +bool octo_free_sample_case(const std::string& ckpt_path, + const std::string& case_dir, + int n_samples, + uint32_t seed, + std::vector& samples_out, + std::vector& noise_out) { + if (n_samples <= 0) { + std::fprintf(stderr, "vla(octo): free-sample n_samples must be > 0\n"); + return false; + } + gguf_reader g{"octo"}; + if (!g.open(ckpt_path)) return false; + const int window_size = (int) g.u32("octo.window_size"); + if (window_size < 1 || window_size > kMaxHorizon) { + std::fprintf(stderr, "vla(octo): octo.window_size=%d out of supported range [1, %d]\n", + window_size, kMaxHorizon); + return false; + } + + // Observation: load once from the golden case, exactly as octo_dump_tokenizer_case + // does (same helpers, same zero-fill fallback for cases with no task image or, + // for a single-camera checkpoint, no wrist observation at all). + NpyU8 primary_obs, wrist_obs, primary_task, wrist_task; + bool wrist_obs_real = true; + if (!parse_npy_u8(case_dir + "/tensors/input.observation.image_primary.npy", primary_obs)) return false; + if (!parse_npy_u8_or_zero_obs(case_dir + "/tensors/input.observation.image_wrist.npy", + window_size, 128, wrist_obs, wrist_obs_real)) return false; + if (!parse_npy_u8_or_zero_task(case_dir + "/tensors/input.task.image_primary.npy", primary_obs, primary_task)) return false; + if (!parse_npy_u8_or_zero_task(case_dir + "/tensors/input.task.image_wrist.npy", wrist_obs, wrist_task)) return false; + NpyI32 input_ids, attention_mask; + if (!parse_npy_i32(case_dir + "/tensors/input.task.language_instruction.input_ids.npy", input_ids)) return false; + if (!parse_npy_i32(case_dir + "/tensors/input.task.language_instruction.attention_mask.npy", attention_mask)) return false; + NpyBool task_valid, primary_valid, wrist_valid, timestep_valid; + if (!parse_npy_bool(case_dir + "/tensors/input.task.pad_mask_dict.language_instruction.npy", task_valid) || + !parse_npy_bool(case_dir + "/tensors/input.observation.pad_mask_dict.image_primary.npy", primary_valid) || + !parse_npy_bool_or_zero(case_dir + "/tensors/input.observation.pad_mask_dict.image_wrist.npy", window_size, wrist_valid) || + !parse_npy_bool(case_dir + "/tensors/input.observation.timestep_pad_mask.npy", timestep_valid)) return false; + if (!wrist_obs_real) { + std::fill(wrist_valid.data.begin(), wrist_valid.data.end(), (uint8_t) 0); + } + std::vector blocked_bytes; + std::vector blocked_mask; + if (!build_transformer_mask(task_valid, primary_valid, wrist_valid, timestep_valid, window_size, /*n_proprio_tokens=*/0, blocked_bytes, blocked_mask)) return false; + + // TIP-ND1-B: weights resident, loaded once (same load_all_tensors() sequence every other + // resident caller uses) -- matches how OctoPt holds a loaded model in memory across + // repeated sample_actions() calls. Unlike the deleted disk-read path, the resident stage + // functions only ever READ tensor bytes via wt()/tensor_to_vec() -- nothing is + // moved/consumed out of ctx_w, so every one of the N passes below sees the identical, + // uncorrupted weight set with no per-iteration "copy the master" step needed (that copy + // existed only because the deleted run_t5_encoder_graph/run_transformer_graph MOVED their + // weight-struct argument's vectors into the ggml payload, consuming them after one use -- + // that move-semantics hazard no longer exists once the weight source is a read-only + // resident lookup). + OctoModelArch m; + m.matmul_type = GGML_TYPE_F32; + m.action_horizon = g.u32("octo.action.horizon"); + m.action_dim = g.u32("octo.action.dim"); + m.backend = octo_select_backend(default_cpu_threads(), "[free-sample] "); + if (!m.backend) return false; + if (!load_all_tensors(m, g)) return false; + ggml_context * ctx_w = m.ctx_weights; + ggml_backend_t backend = m.backend; + + ggml_tensor * readout_pos_r = wt(ctx_w, "octo.readout.action.pos_embd"); + if (!readout_pos_r) return false; + const std::vector readout_pos = tensor_to_vec(readout_pos_r); + + const int action = (int) (m.action_horizon * m.action_dim); + samples_out.assign((size_t) n_samples * (size_t) action, 0.0f); + noise_out.assign((size_t) n_samples * (size_t) window_size * action, 0.0f); + + for (int i = 0; i < n_samples; ++i) { + std::vector tok, primary_proj, primary_pos, wrist_proj, wrist_pos; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "primary", primary_obs, primary_task, 256, 256, window_size, tok, primary_proj, primary_pos)) return false; + if (!run_one_obs_tokenizer_graph_resident(ctx_w, backend, "wrist", wrist_obs, wrist_task, 128, 64, window_size, tok, wrist_proj, wrist_pos)) return false; + + std::vector t5_out; + if (!run_t5_encoder_graph_resident(ctx_w, backend, input_ids.data, attention_mask.data, t5_out)) return false; + + NpyF32 t5; + t5.shape = {1, 16, 768}; + t5.data = t5_out; + std::vector lang_proj, lang_pos, repeated; + if (!run_language_graph_resident(ctx_w, backend, t5, window_size, lang_proj, lang_pos, repeated)) return false; + + OctoTransformerResult bt; + if (!assemble_transformer_input(lang_pos, primary_pos, wrist_pos, {}, repeated, readout_pos, window_size, /*n_proprio_tokens=*/0, bt.input)) return false; + if (!run_transformer_graph_resident(ctx_w, backend, bt.input, blocked_mask, window_size, /*n_proprio_tokens=*/0, bt)) return false; + + std::mt19937 rng(seed + (uint32_t) i); + OctoDiffusionResult diff; + const OctoNoiseSource sample_noise{OctoNoiseSourceKind::RANDOM, "", &rng}; + if (!run_diffusion_resident(ctx_w, backend, sample_noise, bt.readout_action, window_size, action, diff)) return false; + std::copy(diff.final_actions.begin(), diff.final_actions.end(), samples_out.begin() + (size_t) i * (size_t) action); + std::copy(diff.initial_noise.begin(), diff.initial_noise.end(), noise_out.begin() + (size_t) i * (size_t) window_size * action); + } + return true; +} + +} // namespace vla diff --git a/src/models/octo.h b/src/models/octo.h new file mode 100644 index 0000000..a7b89c3 --- /dev/null +++ b/src/models/octo.h @@ -0,0 +1,96 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +#pragma once + +#include +#include +#include + +namespace vla { + +bool octo_dump_gguf_inventory(const std::string& ckpt_path); + +// TIP-ND1-B: sole tokenizer-case dump path (was octo_dump_tokenizer_case / _resident before +// TIP-ND1-A proved the resident/GPU-wired path is bit-exact with the deleted disk-read +// original -- 20/20 golden cases). Dumps every M0-M8 boundary needed by +// verify_octo_parity.py, computed via the resident stage functions on the model's real +// backend (CUDA when available, CPU otherwise) -- the literal code OctoModelArch::predict() +// runs. +bool octo_dump_tokenizer_case_resident(const std::string& ckpt_path, + const std::string& case_dir, + const std::string& dump_dir, + const std::string& t5_inject_path = "", + const std::string& unnorm_dataset = "bridge_dataset"); + +// T5 SentencePiece-unigram tokenization (vocab embedded in the GGUF at convert +// time). Pads/truncates to octo.tokens.language (16), appends EOS, matching +// HFTokenizer(t5-base, 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); + +// Action chunk produced by one live end-to-end Octo forward pass. +struct OctoCliAction { + std::vector normalized; ///< [4,7] normalized action; verified vs golden (M1-M5 parity). + std::vector unnormalized; ///< [4,7] world-unit action via octo.dataset_statistics + ///< ("bridge_dataset"). The unnormalize formula itself is + ///< verified vs golden (M8, see action_final_unnormalized in + ///< the ctest harness); this specific *live* CLI call is not, + ///< since it samples fresh diffusion noise each run (see + ///< TIP-007's Completion Report for why an exact-match golden + ///< comparison isn't meaningful for the live/stochastic path). +}; + +// Runs the full Octo pipeline (SmallStem x2 -> T5 encoder -> block transformer +// -> diffusion action head) from a live image pair + raw instruction text. +// primary/wrist images are single current frames (interleaved RGB8, arbitrary +// size); resized internally to 256x256 / 128x128 and duplicated across the +// window with timestep_pad_mask=[0,1] (matches OctoPt's HistoryWrapper cold +// start: history filled with the first frame, only the latest slot valid). No +// goal image (language-only conditioning): task image is zero-filled, matching +// OctoModelPt.create_tasks(texts=...). +// unnorm_dataset: octo.dataset_statistics key to un-normalize against; "" (default) +// auto-resolves (VLA_OCTO_UNNORM_DATASET env var, else the sole key if unambiguous, +// else "bridge_dataset" if present) -- see resolve_unnorm_dataset_key in octo.cpp. +bool octo_predict_from_images(const std::string& ckpt_path, + const uint8_t* primary_rgb, int primary_w, int primary_h, + const uint8_t* wrist_rgb, int wrist_w, int wrist_h, + const std::string& instruction, + OctoCliAction& out, + const std::string& unnorm_dataset = ""); + +// TIP-009: statistical action-distribution parity. Loads the golden case's own +// observation (SmallStem images, task-language input_ids/attention_mask, all +// four pad masks) once, then runs the full pipeline (tokenizer -> T5 -> block +// transformer -> diffusion) n_samples times end-to-end, each with fresh +// N(0,1) diffusion noise seeded from std::mt19937(seed + i). Does not touch +// any already-verified graph function's behavior; reuses them unchanged. +// samples_out: [n_samples,4,7] normalized action, row-major, sample-major. +// noise_out: [n_samples,2,28] initial DDPM noise actually consumed, same layout. +bool octo_free_sample_case(const std::string& ckpt_path, + const std::string& case_dir, + int n_samples, + uint32_t seed, + std::vector& samples_out, + std::vector& noise_out); + +// TIP-06: stagewise (T0-T5) golden-parity dump for the L1/proprio forward path (TIP-05), +// mirroring octo_dump_tokenizer_case_resident's role for the diffusion path. ckpt_path must +// be a head_type=l1 + proprio GGUF. case_dir must contain (written by +// octo-pytorch-kamusarj's scripts/dump_l1_stagewise_golden.py): +// input.top_hwc_u8.npy (256,256,3) uint8, input.wrist_hwc_u8.npy (128,128,3) uint8, +// input.proprio_raw.npy (7,) float32, input.instruction.txt (plain text). +// Runs octo.cpp's own real T5/SentencePiece tokenizer (octo_tokenize_text) and the exact +// head_type=l1 resident stage functions (run_proprio_tokenizer_graph_resident, +// run_l1_action_head_graph_resident, ...) -- the same code OctoModelArch::predict() runs for +// an L1 checkpoint, not a separate reimplementation. Writes T0-T5 to dump_dir as .f32 (same +// write_f32_dump format octo_dump_tokenizer_case_resident already uses) + manifest.txt. +bool octo_dump_l1_stagewise_case_resident(const std::string& ckpt_path, + const std::string& case_dir, + const std::string& dump_dir, + const std::string& unnorm_dataset = ""); + +} // namespace vla diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index 597a010..cbe5b03 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -19,8 +19,17 @@ // // vla-cli [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image img2.jpg] // --tokens id,id,... [--state f,f,...] [--pretty] +// +// --model octo is the exception: Octo tokenizes its own instruction text +// in-process (TIP-007/M6), so that path takes raw images + text instead: +// +// vla-cli --model octo --ckpt octo-small-1.5-f32.gguf +// --image-primary p.png --image-wrist w.png --instruction "..." +// [--normalized] [--pretty] +#include "arch.h" #include "model.h" +#include "models/octo.h" #define STB_IMAGE_IMPLEMENTATION #define STB_IMAGE_STATIC @@ -101,8 +110,18 @@ void usage(const char * prog) { " --image image file, repeat for multi-view (decoded via stb_image)\n" " --tokens language token ids, comma-separated (tokenize in the client)\n" " --state proprioception floats, comma-separated (default zeros)\n" - " --pretty print one action row (max_action_dim values) per line\n", - prog); + " --pretty print one action row (max_action_dim values) per line\n" + "\n" + "octo mode (in-process SentencePiece tokenization, no --tokens needed):\n" + " %s --model octo --ckpt octo-small-1.5-f32.gguf\n" + " --image-primary p.png --image-wrist w.png --instruction \"...\" [--normalized] [--pretty]\n" + " --model octo (auto-detected from --ckpt if omitted)\n" + " --image-primary primary camera view, single current frame\n" + " --image-wrist wrist camera view, single current frame\n" + " --instruction raw language instruction text\n" + " --normalized print the verified normalized action instead of the\n" + " un-normalized (world-unit, NOT verified vs golden -- M8) default\n", + prog, prog); } } // namespace @@ -111,6 +130,8 @@ int main(int argc, char ** argv) { std::string mmproj, ckpt, tokens_s, state_s; std::vector image_paths; bool pretty = false; + std::string model_flag, image_primary, image_wrist, instruction; + bool normalized = false; for (int i = 1; i < argc; ++i) { const std::string a = argv[i]; @@ -124,10 +145,58 @@ int main(int argc, char ** argv) { else if (a == "--tokens") tokens_s = need("--tokens"); else if (a == "--state") state_s = need("--state"); else if (a == "--pretty") pretty = true; + else if (a == "--model") model_flag = need("--model"); + else if (a == "--image-primary") image_primary = need("--image-primary"); + else if (a == "--image-wrist") image_wrist = need("--image-wrist"); + else if (a == "--instruction") instruction = need("--instruction"); + else if (a == "--normalized") normalized = true; else if (a == "-h" || a == "--help") { usage(argv[0]); return 0; } else { std::fprintf(stderr, "vla-cli: unknown argument %s\n", a.c_str()); usage(argv[0]); return 1; } } - if (ckpt.empty() || image_paths.empty() || tokens_s.empty()) { usage(argv[0]); return 1; } + if (ckpt.empty()) { usage(argv[0]); return 1; } + + bool octo_mode = (model_flag == "octo"); + if (!octo_mode && model_flag.empty() && (!image_primary.empty() || !image_wrist.empty() || !instruction.empty())) { + Arch arch; + if (detect_arch_from_ckpt(ckpt, &arch) && arch == Arch::OCTO) octo_mode = true; + } + + if (octo_mode) { + if (image_primary.empty() || image_wrist.empty() || instruction.empty()) { + std::fprintf(stderr, "vla-cli: --model octo needs --image-primary, --image-wrist, and --instruction\n"); + usage(argv[0]); + return 1; + } + std::vector pbuf, wbuf; + int pw = 0, ph = 0, ww = 0, wh = 0; + if (!load_image(image_primary.c_str(), pbuf, pw, ph)) return 1; + if (!load_image(image_wrist.c_str(), wbuf, ww, wh)) return 1; + + OctoCliAction result; + if (!octo_predict_from_images(ckpt, pbuf.data(), pw, ph, wbuf.data(), ww, wh, instruction, result)) { + std::fprintf(stderr, "vla-cli: octo predict failed\n"); + return 2; + } + const std::vector & act = normalized ? result.normalized : result.unnormalized; + if (!normalized) { + std::fprintf(stderr, + "vla-cli: printing UN-normalized action (world units, via octo.dataset_statistics " + "bridge_dataset); this path is NOT verified against golden yet (M8). " + "Pass --normalized for the verified normalized action.\n"); + } + constexpr int64_t cols = 7; + if (pretty) { + for (size_t i = 0; i < act.size(); ++i) + std::printf("%.6g%c", act[i], ((int64_t) (i + 1) % cols == 0) ? '\n' : ' '); + } else { + std::printf("action_len=%zu\n", act.size()); + for (float x : act) std::printf("%.9g\n", x); + } + std::fflush(stdout); + return 0; + } + + if (image_paths.empty() || tokens_s.empty()) { usage(argv[0]); return 1; } // Validate the cheap args before loading the model. std::vector lang; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3e20804..73dfb70 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1,11 +1,203 @@ # vla.cpp tests, enabled with -DVLA_BUILD_TESTS=ON. predict_check needs a real # GGUF so it is built but not run by ctest; the unit tests below are pure. +find_package(Python3 COMPONENTS Interpreter) + add_executable(vla_predict_check predict_check.cpp) target_link_libraries(vla_predict_check PRIVATE vla_core) target_compile_options(vla_predict_check PRIVATE -Wall -Wextra) +add_executable(octo_bench octo_bench.cpp) +target_include_directories(octo_bench PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(octo_bench PRIVATE vla_core) +target_compile_options(octo_bench PRIVATE -Wall -Wextra) + +add_executable(octo_load_check octo_load_check.cpp) +target_include_directories(octo_load_check PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(octo_load_check PRIVATE vla_core) +target_compile_options(octo_load_check PRIVATE -Wall -Wextra) + +add_executable(octo_parity_dump octo_parity_dump.cpp) +target_include_directories(octo_parity_dump PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(octo_parity_dump PRIVATE vla_core) +target_compile_options(octo_parity_dump PRIVATE -Wall -Wextra) +# All shipped golden traces (tier1 + tier2/bridge_debug) were generated against +# "bridge_dataset" statistics (manifest metadata.unnormalization.dataset / +# metadata.dataset, confirmed for every case below) -- passed explicitly here +# rather than left to octo_parity_dump's default, so the key is visible and easy +# to audit if a future golden case uses a different dataset. +set(OCTO_UNNORM_DATASET bridge_dataset) + +# TIP-ND1-B: registers one ctest case running octo_parity_dump (now resident-only -- the +# gốc/"_resident" split this helper used to register in parallel, TIP-ND1-A, was collapsed +# back to one test per case once TIP-ND1-A proved resident-on-CPU is bit-exact with the +# now-deleted gốc path and octo.cpp was unified onto it, TIP-ND1-B). Test names are unchanged +# from before TIP-ND1-A. +function(add_octo_parity_test) + set(oneValueArgs NAME CKPT CASE_DIR T5_INJECT TRANSFORMER_TOL UNNORM_DATASET KNOWN_MISMATCH EXCLUDE DUMP_DIR) + cmake_parse_arguments(APT "" "${oneValueArgs}" "" ${ARGN}) + add_test( + NAME ${APT_NAME} + COMMAND ${CMAKE_COMMAND} + -DCKPT=${APT_CKPT} + -DCASE_DIR=${APT_CASE_DIR} + -DT5_INJECT=${APT_T5_INJECT} + -DTRANSFORMER_TOL=${APT_TRANSFORMER_TOL} + -DUNNORM_DATASET=${APT_UNNORM_DATASET} + -DKNOWN_MISMATCH=${APT_KNOWN_MISMATCH} + -DEXCLUDE=${APT_EXCLUDE} + -DDUMP_DIR=${APT_DUMP_DIR} + -DDUMPER=$ + -DVERIFY=${CMAKE_SOURCE_DIR}/scripts/verify_octo_parity.py + -DPYTHON_EXECUTABLE=${Python3_EXECUTABLE} + -P ${CMAKE_SOURCE_DIR}/tests/run_octo_parity.cmake) +endfunction() + +if(Python3_FOUND AND DEFINED ENV{VLA_OCTO_GOLDEN_DIR}) + foreach(OCTO_PARITY_CASE example_batch synthetic_zero synthetic_random second_timestep_padded wrist_valid_synthetic synthetic_ramp) + if(EXISTS "$ENV{VLA_OCTO_GOLDEN_DIR}/${OCTO_PARITY_CASE}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy") + # Oracle mode: --t5-inject feeds the golden tokens downstream; native T5 output is still + # dumped as "t5.out" and checked against golden independently (M5, TIP-006). + add_octo_parity_test( + NAME octo_parity_${OCTO_PARITY_CASE} + CKPT ${CMAKE_SOURCE_DIR}/octo-small-1.5-f32.gguf + CASE_DIR $ENV{VLA_OCTO_GOLDEN_DIR}/${OCTO_PARITY_CASE} + T5_INJECT $ENV{VLA_OCTO_GOLDEN_DIR}/${OCTO_PARITY_CASE}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_UNNORM_DATASET} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity/${OCTO_PARITY_CASE}) + # Native mode: no --t5-inject, T5-base encoder output drives the whole pipeline end-to-end. + add_octo_parity_test( + NAME octo_parity_native_${OCTO_PARITY_CASE} + CKPT ${CMAKE_SOURCE_DIR}/octo-small-1.5-f32.gguf + CASE_DIR $ENV{VLA_OCTO_GOLDEN_DIR}/${OCTO_PARITY_CASE} + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_UNNORM_DATASET} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity_native/${OCTO_PARITY_CASE}) + endif() + endforeach() +endif() + +# tier2/bridge_debug: real (non-synthetic) images + label evaluation, exercised the +# same way as tier1 (oracle + native), one directory level up from VLA_OCTO_GOLDEN_DIR +# (.../golden_traces/octo_small_tier1 -> .../golden_traces/octo_small_tier2/bridge_debug), +# gated on VLA_OCTO_GOLDEN_DIR_TIER2 so it stays opt-in when that data isn't present. +if(Python3_FOUND AND DEFINED ENV{VLA_OCTO_GOLDEN_DIR_TIER2}) + set(OCTO_BRIDGE_DEBUG_DIR "$ENV{VLA_OCTO_GOLDEN_DIR_TIER2}/bridge_debug") + # bridge_debug's own dataset_statistics.action.{mean,std} (recorded per-case in its + # golden manifest, sourced from its debug-only "./tests/debug_dataset") do NOT match + # the octo-small-1.5 checkpoint's "bridge_dataset" stats embedded in the GGUF -- both + # are labeled "bridge_dataset" but are numerically different snapshots (verified: our + # unnormalize formula reproduces bridge_debug's golden action_final_unnormalized to + # ~2e-9 when fed bridge_debug's OWN recorded mean/std, so the formula is correct; the + # mismatch is golden-data provenance, not a vla.cpp defect -- see TIP-008 Completion + # Report. All 6 tier1 cases' own recorded stats match the GGUF exactly, so this is + # specific to the bridge_debug trace, not our GGUF conversion. + set(OCTO_BRIDGE_DEBUG_KNOWN_MISMATCH + "action_final_unnormalized:bridge_debug's golden dataset_statistics differs numerically from the checkpoint's embedded bridge_dataset stats (see TIP-008 report); unnormalize formula itself is verified correct") + if(EXISTS "${OCTO_BRIDGE_DEBUG_DIR}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy") + add_octo_parity_test( + NAME octo_parity_bridge_debug + CKPT ${CMAKE_SOURCE_DIR}/octo-small-1.5-f32.gguf + CASE_DIR ${OCTO_BRIDGE_DEBUG_DIR} + T5_INJECT ${OCTO_BRIDGE_DEBUG_DIR}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_UNNORM_DATASET} + KNOWN_MISMATCH ${OCTO_BRIDGE_DEBUG_KNOWN_MISMATCH} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity/bridge_debug) + add_octo_parity_test( + NAME octo_parity_native_bridge_debug + CKPT ${CMAKE_SOURCE_DIR}/octo-small-1.5-f32.gguf + CASE_DIR ${OCTO_BRIDGE_DEBUG_DIR} + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_UNNORM_DATASET} + KNOWN_MISMATCH ${OCTO_BRIDGE_DEBUG_KNOWN_MISMATCH} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity_native/bridge_debug) + endif() +endif() + +# LIBERO (window=1): golden traces from the cyrusneary octo-finetuned-libero checkpoint +# (TIP-GOLD), a genuinely single-camera model -- its example_batch/finetune_config never fed +# a wrist observation at all, so the golden has no octo_transformer.obs_wrist.* / +# block_transformer...obs_wrist tensors, and its block_transformer.{input_tokens,attention_mask, +# output_tokens} are shorter than the port's (vla.cpp always includes masked-but-present wrist +# tokens in the full sequence -- numerically equivalent for downstream outputs, since a masked +# token contributes nothing, but shape-incompatible for those specific full-sequence/wrist +# boundaries). OCTO_LIBERO_EXCLUDE drops exactly those from comparison; every other boundary +# (obs_primary, task_language, readout, diff.*, final_action(+unnormalized)) is still compared +# at the same tolerance as bridge. See TIP-HARNESS report for the full rationale. +# Gated on two env vars since neither the golden traces nor the LIBERO GGUF are committed: +# VLA_OCTO_LIBERO_GOLDEN_DIR -- root containing /tensors/ + manifest.json per case +# VLA_OCTO_LIBERO_CKPT -- path to the window=1 GGUF (scripts/convert_octo_to_gguf.py +# --ckpt --step 60000) +set(OCTO_LIBERO_UNNORM_DATASET libero_object) +set(OCTO_LIBERO_EXCLUDE "bt.input,bt.mask,bt.output,obs_wrist") + +if(Python3_FOUND AND DEFINED ENV{VLA_OCTO_LIBERO_GOLDEN_DIR} AND DEFINED ENV{VLA_OCTO_LIBERO_CKPT}) + foreach(OCTO_LIBERO_CASE example_frame synthetic_zero synthetic_ramp) + if(EXISTS "$ENV{VLA_OCTO_LIBERO_GOLDEN_DIR}/${OCTO_LIBERO_CASE}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy") + # Oracle mode: --t5-inject feeds the golden tokens downstream; native T5 output is still + # dumped as "t5.out" and checked against golden independently. + add_octo_parity_test( + NAME octo_parity_libero_${OCTO_LIBERO_CASE} + CKPT $ENV{VLA_OCTO_LIBERO_CKPT} + CASE_DIR $ENV{VLA_OCTO_LIBERO_GOLDEN_DIR}/${OCTO_LIBERO_CASE} + T5_INJECT $ENV{VLA_OCTO_LIBERO_GOLDEN_DIR}/${OCTO_LIBERO_CASE}/tensors/octo_transformer.task_language.tokens_after_tokenizer.npy + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_LIBERO_UNNORM_DATASET} + EXCLUDE ${OCTO_LIBERO_EXCLUDE} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity_libero/${OCTO_LIBERO_CASE}) + # Native mode: no --t5-inject, T5-base encoder output drives the whole pipeline end-to-end. + add_octo_parity_test( + NAME octo_parity_libero_native_${OCTO_LIBERO_CASE} + CKPT $ENV{VLA_OCTO_LIBERO_CKPT} + CASE_DIR $ENV{VLA_OCTO_LIBERO_GOLDEN_DIR}/${OCTO_LIBERO_CASE} + TRANSFORMER_TOL 2e-3 + UNNORM_DATASET ${OCTO_LIBERO_UNNORM_DATASET} + EXCLUDE ${OCTO_LIBERO_EXCLUDE} + DUMP_DIR ${CMAKE_BINARY_DIR}/octo_parity_libero_native/${OCTO_LIBERO_CASE}) + endif() + endforeach() +endif() + +# TIP-06: octo_l1_parity_dump + verify_octo_l1_parity.py, mirroring add_octo_parity_test's +# shape but for the head_type=l1 forward path (TIP-05) against a stagewise OctoPt golden +# dump (octo-pytorch-kamusarj's scripts/dump_l1_stagewise_golden.py, one case dir per +# ep0 timestep -- ep0_t0/ep0_t8/ep0_t16). Gated on two env vars since neither the L1 GGUF +# nor the golden dumps are committed: +# VLA_OCTO_L1_GOLDEN_DIR -- root containing ep0_t/{input.*,t0..t5.*.npy,manifest.json} +# VLA_OCTO_L1_CKPT -- path to the head_type=l1 GGUF (octo-aloha-jitter2525.gguf) +add_executable(octo_l1_parity_dump octo_l1_parity_dump.cpp) +target_include_directories(octo_l1_parity_dump PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(octo_l1_parity_dump PRIVATE vla_core) +target_compile_options(octo_l1_parity_dump PRIVATE -Wall -Wextra) + +if(Python3_FOUND AND DEFINED ENV{VLA_OCTO_L1_GOLDEN_DIR} AND DEFINED ENV{VLA_OCTO_L1_CKPT}) + foreach(OCTO_L1_CASE ep0_t0 ep0_t8 ep0_t16) + if(EXISTS "$ENV{VLA_OCTO_L1_GOLDEN_DIR}/${OCTO_L1_CASE}/manifest.json") + add_test( + NAME octo_l1_parity_${OCTO_L1_CASE} + COMMAND ${CMAKE_COMMAND} + -DCKPT=$ENV{VLA_OCTO_L1_CKPT} + -DCASE_DIR=$ENV{VLA_OCTO_L1_GOLDEN_DIR}/${OCTO_L1_CASE} + -DUNNORM_DATASET= + -DTOL=1e-4 + -DDUMP_DIR=${CMAKE_BINARY_DIR}/octo_l1_parity/${OCTO_L1_CASE} + -DDUMPER=$ + -DVERIFY=${CMAKE_SOURCE_DIR}/scripts/verify_octo_l1_parity.py + -DPYTHON_EXECUTABLE=${Python3_EXECUTABLE} + -P ${CMAKE_SOURCE_DIR}/tests/run_octo_l1_parity.cmake) + endif() + endforeach() +endif() + add_executable(test_vision_common test_vision_common.cpp) target_include_directories(test_vision_common PRIVATE ${CMAKE_SOURCE_DIR}/src) target_compile_options(test_vision_common PRIVATE -Wall -Wextra) add_test(NAME vision_common COMMAND test_vision_common) + +# Model-free tripwire for the final-action-slice formula (Blueprint risk item 1445/1492). +# Always registered (no golden dir / GGUF needed), runs in well under a second. +add_executable(octo_action_slice_test octo_action_slice_test.cpp) +target_compile_options(octo_action_slice_test PRIVATE -Wall -Wextra) +add_test(NAME octo_action_slice COMMAND octo_action_slice_test) diff --git a/tests/octo_action_slice_test.cpp b/tests/octo_action_slice_test.cpp new file mode 100644 index 0000000..96671b0 --- /dev/null +++ b/tests/octo_action_slice_test.cpp @@ -0,0 +1,59 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +// Permanent tripwire for Octo's final-action-slice formula, the single riskiest line in the +// window-size port (Blueprint risk item 1445/1492): OctoPt's sample_actions() returns +// actions[:, -1] -- the LAST timestep of the observation window -- after the diffusion head +// runs over the full window. src/models/octo.cpp (run_diffusion_resident, around the "final +// action slice" comment) implements this as: +// +// final_begin = (window_size - 1) * action; +// final_end = window_size * action; +// +// where `action` is the flattened per-timestep action chunk (action_horizon * action_dim, e.g. +// 4 * 7 = 28 for octo-small). This test does NOT call into octo.cpp -- that formula lives in a +// static function with no public entry point, and TIP-HARNESS's scope explicitly excludes +// touching src/models/octo.cpp to expose one. Instead it re-asserts the formula's two concrete, +// specifically-called-out cases (window=1 -> [0,28), window=2 -> [28,56)) as an independent, +// model-free, every-build check. If this test and octo.cpp's formula ever diverge, a human needs +// to reconcile them -- that divergence is exactly the failure mode this test exists to catch. + +#include +#include +#include + +namespace { + +struct Slice { + size_t begin; + size_t end; +}; + +// Mirrors src/models/octo.cpp's final_begin/final_end computation verbatim. +Slice final_action_slice(int window_size, size_t action) { + return { + (size_t) (window_size - 1) * action, + (size_t) window_size * action, + }; +} + +bool check(int window_size, size_t action, size_t expect_begin, size_t expect_end) { + const Slice s = final_action_slice(window_size, action); + const bool ok = (s.begin == expect_begin) && (s.end == expect_end); + std::printf("window_size=%d action=%zu -> [%zu,%zu) expected [%zu,%zu) %s\n", + window_size, action, s.begin, s.end, expect_begin, expect_end, + ok ? "PASS" : "FAIL"); + return ok; +} + +} // namespace + +int main() { + // action_horizon=4, action_dim=7 -> action=28, matching octo-small's DiffusionActionHead. + const size_t action = 28; + bool ok = true; + ok &= check(/*window_size=*/1, action, /*expect_begin=*/0, /*expect_end=*/28); + ok &= check(/*window_size=*/2, action, /*expect_begin=*/28, /*expect_end=*/56); + return ok ? 0 : 1; +} diff --git a/tests/octo_bench.cpp b/tests/octo_bench.cpp new file mode 100644 index 0000000..d555d43 --- /dev/null +++ b/tests/octo_bench.cpp @@ -0,0 +1,110 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +// TIP-BUILD-OCTO-GPU-B AC3: per-call latency of the live vla::predict() path for Octo. +// predict_check.cpp (the generic multi-arch harness) hardcodes a 6-token, no-mask language +// input; Octo's predict() rejects anything but exactly 16 lang tokens + a 16-entry attention +// mask (see the "expects lang_tokens AND attention_mask of exactly 16 entries" check in +// octo.cpp), so it can't be reused here without changing behavior for every other arch that +// binary also exercises. This is a small Octo-only sibling: same model.h API, same timing +// loop shape, real 16-token/mask input built via octo_tokenize_text (the tokenizer already +// verified against golden in the M-series parity tests), so it hits the exact +// OctoModelArch::predict() -> octo_run_pipeline_resident() -> 5 *_resident stage functions +// live path with no separate reimplementation. +// +// octo_bench [iters] +// env: VLA_BENCH_ITERS overrides the CLI iters arg if set. + +#include "model.h" +#include "models/octo.h" + +#include +#include +#include +#include +#include +#include + +using namespace vla; + +int main(int argc, char** argv) { + if (argc < 2) { + std::fprintf(stderr, "usage: %s [iters]\n", argv[0]); + return 1; + } + const std::string ckpt = argv[1]; + int iters = argc > 2 ? std::atoi(argv[2]) : 20; + if (const char* env = std::getenv("VLA_BENCH_ITERS")) iters = std::atoi(env); + + Model* m = model_load("", ckpt, ""); + if (!m) { + std::fprintf(stderr, "model_load failed\n"); + return 1; + } + const Config& cfg = model_config(m); + + std::vector input_ids, attention_mask; + if (!octo_tokenize_text(ckpt, "pick up the block and place it on the plate", input_ids, attention_mask)) { + std::fprintf(stderr, "octo_tokenize_text failed\n"); + return 1; + } + std::fprintf(stderr, "tokenized: n_lang=%zu n_mask=%zu\n", input_ids.size(), attention_mask.size()); + + const int primary_w = 256, primary_h = 256; + const int wrist_w = 128, wrist_h = 128; + std::vector primary_buf((size_t) 3 * primary_w * primary_h); + std::vector wrist_buf((size_t) 3 * wrist_w * wrist_h); + for (int y = 0; y < primary_h; ++y) + for (int x = 0; x < primary_w; ++x) + for (int c = 0; c < 3; ++c) + primary_buf[((size_t) y * primary_w + x) * 3 + c] = (uint8_t) ((x + 2 * y + 40 * c) & 0xFF); + for (int y = 0; y < wrist_h; ++y) + for (int x = 0; x < wrist_w; ++x) + for (int c = 0; c < 3; ++c) + wrist_buf[((size_t) y * wrist_w + x) * 3 + c] = (uint8_t) ((x + 3 * y + 60 * c) & 0xFF); + + ImageView views[2]; + views[0] = ImageView{primary_buf.data(), primary_w, primary_h, PixelFormat::U8}; + views[1] = ImageView{wrist_buf.data(), wrist_w, wrist_h, PixelFormat::U8}; + + std::vector state((size_t) cfg.max_state_dim, 0.0f); + const size_t noise_n = (size_t) cfg.max_action_dim * (size_t) cfg.n_suffix; + std::vector noise(noise_n); + for (size_t i = 0; i < noise_n; ++i) noise[i] = 0.001f * (float) ((i * 2654435761u) % 1000) - 0.5f; + + Inputs in{}; + in.images = views; + in.n_images = 2; + in.lang_tokens = input_ids.data(); + in.n_lang = (int) input_ids.size(); + in.attention_mask = attention_mask.data(); + in.attention_mask_n = (int) attention_mask.size(); + in.state = state.data(); + in.noise = noise.data(); + in.timing_detail = TimingDetail::NONE; + + std::vector act = predict(m, in); + std::printf("action_len=%zu\n", act.size()); + if (act.empty()) { + std::fprintf(stderr, "predict() returned empty action -- see stderr above for the reason\n"); + model_free(m); + return 2; + } + + for (int w = 0; w < 3; ++w) (void) predict(m, in); + double best = 1e30, sum = 0.0; + for (int i = 0; i < iters; ++i) { + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + std::vector a = predict(m, in); + clock_gettime(CLOCK_MONOTONIC, &t1); + const double ms = (t1.tv_sec - t0.tv_sec) * 1e3 + (t1.tv_nsec - t0.tv_nsec) / 1e6; + best = ms < best ? ms : best; + sum += ms; + } + std::fprintf(stderr, "octo predict() over %d iters: min=%.3f ms avg=%.3f ms\n", iters, best, sum / iters); + + model_free(m); + return 0; +} diff --git a/tests/octo_l1_parity_dump.cpp b/tests/octo_l1_parity_dump.cpp new file mode 100644 index 0000000..ad24928 --- /dev/null +++ b/tests/octo_l1_parity_dump.cpp @@ -0,0 +1,64 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +// TIP-06: CLI wrapper around vla::octo_dump_l1_stagewise_case_resident, mirroring +// octo_parity_dump.cpp's role for the diffusion path -- one small binary, all the real work +// lives in octo.cpp so this stays a thin driver ctest can shell out to. +// +// octo_l1_parity_dump --ckpt --case --out [--unnorm-dataset ] + +#include "models/octo.h" + +#include +#include +#include + +int main(int argc, char ** argv) { + std::string ckpt; + std::string case_dir; + std::string dump_dir; + std::string unnorm_dataset; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto need = [&](const char * opt) -> const char * { + if (i + 1 >= argc) { + std::fprintf(stderr, "%s requires a value\n", opt); + return nullptr; + } + return argv[++i]; + }; + if (a == "--ckpt") { + const char * v = need("--ckpt"); + if (!v) return 1; + ckpt = v; + } else if (a == "--case") { + const char * v = need("--case"); + if (!v) return 1; + case_dir = v; + } else if (a == "--out") { + const char * v = need("--out"); + if (!v) return 1; + dump_dir = v; + } else if (a == "--unnorm-dataset") { + const char * v = need("--unnorm-dataset"); + if (!v) return 1; + unnorm_dataset = v; + } else { + std::fprintf(stderr, "unknown arg: %s\n", a.c_str()); + return 1; + } + } + if (dump_dir.empty()) { + if (const char * env = std::getenv("VLA_OCTO_L1_DUMP")) dump_dir = env; + } + if (ckpt.empty() || case_dir.empty() || dump_dir.empty()) { + std::fprintf(stderr, + "usage: %s --ckpt --case [--out |VLA_OCTO_L1_DUMP=]\n" + " [--unnorm-dataset ]\n", argv[0]); + return 1; + } + if (!vla::octo_dump_l1_stagewise_case_resident(ckpt, case_dir, dump_dir, unnorm_dataset)) return 2; + std::printf("octo_l1_parity_dump wrote %s\n", dump_dir.c_str()); + return 0; +} diff --git a/tests/octo_load_check.cpp b/tests/octo_load_check.cpp new file mode 100644 index 0000000..55ef394 --- /dev/null +++ b/tests/octo_load_check.cpp @@ -0,0 +1,30 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +#include "model.h" +#include "models/octo.h" + +#include +#include + +int main(int argc, char ** argv) { + if (argc != 3 || std::string(argv[1]) != "--ckpt") { + std::fprintf(stderr, "usage: %s --ckpt octo-small-1.5-f32.gguf\n", argv[0]); + return 1; + } + const std::string ckpt = argv[2]; + if (!vla::octo_dump_gguf_inventory(ckpt)) return 2; + + vla::Model * m = vla::model_load("", ckpt, ""); + if (!m) { + std::fprintf(stderr, "octo model_load failed\n"); + return 3; + } + const vla::Config& cfg = vla::model_config(m); + std::printf("loaded_octo hidden=%lld layers=%lld action_horizon=%lld action_dim=%lld\n", + (long long) cfg.hidden, (long long) cfg.n_layers, + (long long) cfg.n_suffix, (long long) cfg.real_action_dim); + vla::model_free(m); + return 0; +} diff --git a/tests/octo_parity_dump.cpp b/tests/octo_parity_dump.cpp new file mode 100644 index 0000000..99a3b2f --- /dev/null +++ b/tests/octo_parity_dump.cpp @@ -0,0 +1,141 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); + +#include "models/octo.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +// Minimal NPY v1.0 writer for a C-contiguous float32 array (TIP-009: cpp_samples.npy +// / cpp_noise.npy need to be loadable by numpy in scripts/compare_action_dist.py). +bool write_npy_f32(const std::string& path, const std::vector& data, const std::vector& shape) { + std::string shape_str = "("; + for (size_t i = 0; i < shape.size(); ++i) { + shape_str += std::to_string(shape[i]); + if (shape.size() == 1 || i + 1 < shape.size()) shape_str += ", "; + } + shape_str += ")"; + std::string header = "{'descr': '(&hlen), 2); + f.write(header.data(), (std::streamsize) header.size()); + f.write(reinterpret_cast(data.data()), (std::streamsize) (data.size() * sizeof(float))); + return (bool) f; +} + +} // namespace + +int main(int argc, char ** argv) { + std::string ckpt; + std::string case_dir; + std::string dump_dir; + std::string t5_inject_path; + std::string unnorm_dataset = "bridge_dataset"; + int free_sample_n = 0; + uint32_t seed = 0; + std::string out_samples_path = "cpp_samples.npy"; + std::string out_noise_path = "cpp_noise.npy"; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto need = [&](const char * opt) -> const char * { + if (i + 1 >= argc) { + std::fprintf(stderr, "%s requires a value\n", opt); + return nullptr; + } + return argv[++i]; + }; + if (a == "--ckpt") { + const char * v = need("--ckpt"); + if (!v) return 1; + ckpt = v; + } else if (a == "--case") { + const char * v = need("--case"); + if (!v) return 1; + case_dir = v; + } else if (a == "--out") { + const char * v = need("--out"); + if (!v) return 1; + dump_dir = v; + } else if (a == "--t5-inject") { + const char * v = need("--t5-inject"); + if (!v) return 1; + t5_inject_path = v; + } else if (a == "--unnorm-dataset") { + const char * v = need("--unnorm-dataset"); + if (!v) return 1; + unnorm_dataset = v; + } else if (a == "--free-sample") { + const char * v = need("--free-sample"); + if (!v) return 1; + free_sample_n = std::atoi(v); + } else if (a == "--seed") { + const char * v = need("--seed"); + if (!v) return 1; + seed = (uint32_t) std::strtoul(v, nullptr, 10); + } else if (a == "--out-samples") { + const char * v = need("--out-samples"); + if (!v) return 1; + out_samples_path = v; + } else if (a == "--out-noise") { + const char * v = need("--out-noise"); + if (!v) return 1; + out_noise_path = v; + } else { + std::fprintf(stderr, "unknown arg: %s\n", a.c_str()); + return 1; + } + } + + if (free_sample_n > 0) { + // TIP-009: statistical action-distribution parity mode. Runs the full + // pipeline free_sample_n times end-to-end on --case's own observation. + if (ckpt.empty() || case_dir.empty()) { + std::fprintf(stderr, + "usage: %s --ckpt octo-small-1.5-f32.gguf --case --free-sample N\n" + " [--seed S] [--out-samples cpp_samples.npy] [--out-noise cpp_noise.npy]\n", argv[0]); + return 1; + } + std::vector samples, noise; + if (!vla::octo_free_sample_case(ckpt, case_dir, free_sample_n, seed, samples, noise)) return 2; + if (!write_npy_f32(out_samples_path, samples, {free_sample_n, 4, 7})) return 2; + if (!write_npy_f32(out_noise_path, noise, {free_sample_n, 2, 28})) return 2; + std::printf("octo_parity_dump free-sample: N=%d seed=%u -> %s, %s\n", + free_sample_n, seed, out_samples_path.c_str(), out_noise_path.c_str()); + return 0; + } + + if (dump_dir.empty()) { + if (const char * env = std::getenv("VLA_OCTO_DUMP")) dump_dir = env; + } + if (ckpt.empty() || case_dir.empty() || dump_dir.empty()) { + std::fprintf(stderr, + "usage: %s --ckpt octo-small-1.5-f32.gguf --case [--t5-inject ]\n" + " [--unnorm-dataset ] [--out |VLA_OCTO_DUMP=]\n" + " %s --ckpt --case --free-sample N [--seed S]\n" + " [--out-samples ] [--out-noise ]\n", argv[0], argv[0]); + return 1; + } + if (!vla::octo_dump_tokenizer_case_resident(ckpt, case_dir, dump_dir, t5_inject_path, unnorm_dataset)) return 2; + std::printf("octo_parity_dump wrote %s\n", dump_dir.c_str()); + return 0; +} diff --git a/tests/py/test_octo_preprocessing.py b/tests/py/test_octo_preprocessing.py new file mode 100644 index 0000000..55dae26 --- /dev/null +++ b/tests/py/test_octo_preprocessing.py @@ -0,0 +1,143 @@ +# 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. + +"""Tier-A parity tests for Octo's client-side image preprocessing (rotate180 + resize +to 256) -- eval/client/adapters.py:octo_preprocess_image. Ref: TIP-P. + +octo_preprocess_image() is pure numpy/PIL, but the module it lives in +(eval/client/adapters.py) unconditionally imports torch/tree/lerobot at module scope +for LeRobotPipelineAdapter. Those aren't needed here and may not be installed outside +the LIBERO eval venv, so they're stubbed before import (same technique as +test_converters.py's `gguf` stub) to keep this test runnable with just numpy+PIL. + +Golden tier-A cases (raw.npy/model_entry.npy pairs) are not committed to the repo, so +the parity tests are gated on VLA_OCTO_LIBERO_GOLDEN_DIR (same env var CMake's +octo_parity_libero_* ctest cases use) and skipped/no-op if it isn't set. The rotate +recovery test needs no golden data and always runs. +""" + +import importlib.util +import os +import pathlib +import sys +import types + +import numpy as np + + +def _load_adapters(): + torch_stub = sys.modules.setdefault("torch", types.ModuleType("torch")) + if not hasattr(torch_stub, "from_numpy"): + torch_stub.from_numpy = object() + tree_stub = sys.modules.setdefault("tree", types.ModuleType("tree")) + if not hasattr(tree_stub, "map_structure"): + tree_stub.map_structure = object() + lerobot = sys.modules.setdefault("lerobot", types.ModuleType("lerobot")) + for sub, attrs in { + "lerobot.envs.utils": ["preprocess_observation"], + "lerobot.processor.env_processor": ["LiberoProcessorStep"], + "lerobot.processor.pipeline": ["PolicyProcessorPipeline"], + "lerobot.utils.constants": ["ACTION"], + }.items(): + mod = sys.modules.setdefault(sub, types.ModuleType(sub)) + for attr in attrs: + if not hasattr(mod, attr): + setattr(mod, attr, object()) + del lerobot + + path = pathlib.Path(__file__).resolve().parents[2] / "eval" / "client" / "adapters.py" + spec = importlib.util.spec_from_file_location("octo_tipp_adapters", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +# Per-case tier-A tolerance (uint8 max_abs vs golden model_entry.npy). synthetic_zero is +# an all-black frame -> JPEG is lossless there, so rotate+resize must be bit-exact +# (max_abs=0). The other two carry real JPEG round-trip noise in golden itself +# (max_abs_diff_vs_source_frame = 2 / 20 respectively, per +# libero_golden_summary.json) -- <=25 is soft slack for that, not a license to get the +# rotation/resize wrong: a wrong rotation or an axis-swapped resize produces max_abs in +# the 127-255 range (near-uncorrelated pixels), which this tolerance would still catch. +TIER_A_TOLERANCE = { + "synthetic_zero": 0, + "synthetic_ramp": 25, + "example_frame": 25, +} + + +def test_octo_rotate_recovers_frame(): + """Step 3: preprocess(frame[::-1,::-1]) must recover frame exactly (no JPEG in the + loop here, so this isolates just the rotate+identity-resize logic).""" + adapters = _load_adapters() + rng = np.random.default_rng(0) + frame = rng.integers(0, 256, size=(256, 256, 3), dtype=np.uint8) + + upside_down = frame[::-1, ::-1] + recovered = adapters.octo_preprocess_image(upside_down, image_size=256) + + assert recovered.shape == frame.shape + assert recovered.dtype == np.uint8 + diff = np.abs(recovered.astype(np.int32) - frame.astype(np.int32)) + assert diff.max() == 0, f"rotate180 did not recover the original frame: max_abs={diff.max()}" + + +def _tier_a_case(adapters, golden_dir: pathlib.Path, case: str) -> tuple[int, float]: + case_dir = golden_dir / case / "preprocessing" + raw = np.load(case_dir / "raw.npy") + model_entry = np.load(case_dir / "model_entry.npy") + + got = adapters.octo_preprocess_image(raw, image_size=256) + assert got.shape == model_entry.shape, f"{case}: shape {got.shape} != golden {model_entry.shape}" + + diff = np.abs(got.astype(np.int32) - model_entry.astype(np.int32)) + return int(diff.max()), float(diff.mean()) + + +def _run_tier_a(case: str): + golden_dir_env = os.environ.get("VLA_OCTO_LIBERO_GOLDEN_DIR") + if not golden_dir_env: + return # gated: no golden dir configured, nothing to verify (matches ctest gating) + golden_dir = pathlib.Path(golden_dir_env) + case_dir = golden_dir / case / "preprocessing" + if not (case_dir / "raw.npy").exists(): + return + + adapters = _load_adapters() + max_abs, mean_abs = _tier_a_case(adapters, golden_dir, case) + tol = TIER_A_TOLERANCE[case] + assert max_abs <= tol, ( + f"{case}: max_abs={max_abs} exceeds tolerance {tol} vs golden model_entry.npy " + f"(mean_abs={mean_abs:.4f}) -- check rotation direction/axis if this is large " + f"(>~30), not just JPEG noise" + ) + print(f"tier_a[{case}]: max_abs={max_abs} mean_abs={mean_abs:.4f} tol={tol} PASS") + + +def test_octo_tier_a_synthetic_zero(): + _run_tier_a("synthetic_zero") + + +def test_octo_tier_a_synthetic_ramp(): + _run_tier_a("synthetic_ramp") + + +def test_octo_tier_a_example_frame(): + _run_tier_a("example_frame") + + +if __name__ == "__main__": + test_octo_rotate_recovers_frame() + print("test_octo_rotate_recovers_frame: OK") + for case in TIER_A_TOLERANCE: + _run_tier_a(case) + print("test_octo_preprocessing: done") diff --git a/tests/run_octo_l1_parity.cmake b/tests/run_octo_l1_parity.cmake new file mode 100644 index 0000000..942286e --- /dev/null +++ b/tests/run_octo_l1_parity.cmake @@ -0,0 +1,31 @@ +if(NOT EXISTS "${CKPT}") + message(FATAL_ERROR "Octo L1 GGUF not found: ${CKPT}") +endif() + +get_filename_component(DUMP_PARENT "${DUMP_DIR}" DIRECTORY) +file(MAKE_DIRECTORY "${DUMP_PARENT}") + +set(UNNORM_DATASET_ARGS "") +if(DEFINED UNNORM_DATASET AND NOT "${UNNORM_DATASET}" STREQUAL "") + list(APPEND UNNORM_DATASET_ARGS --unnorm-dataset "${UNNORM_DATASET}") +endif() + +# TIP-06: octo_l1_parity_dump runs octo_dump_l1_stagewise_case_resident -- the head_type=l1 +# resident stage functions (proprio tokenizer, L1 MAPHead), NOT the diffusion path. See +# octo.cpp's own head_type=="l1" gate in that function for the call-site proof this +# grep-verified at TIP-06 report time. +execute_process( + COMMAND "${DUMPER}" --ckpt "${CKPT}" --case "${CASE_DIR}" ${UNNORM_DATASET_ARGS} --out "${DUMP_DIR}" + RESULT_VARIABLE dump_rc +) +if(NOT dump_rc EQUAL 0) + message(FATAL_ERROR "octo_l1_parity_dump failed: ${dump_rc}") +endif() + +execute_process( + COMMAND "${PYTHON_EXECUTABLE}" "${VERIFY}" --golden "${CASE_DIR}" --dump "${DUMP_DIR}" --tol "${TOL}" + RESULT_VARIABLE verify_rc +) +if(NOT verify_rc EQUAL 0) + message(FATAL_ERROR "verify_octo_l1_parity.py failed: ${verify_rc}") +endif() diff --git a/tests/run_octo_parity.cmake b/tests/run_octo_parity.cmake new file mode 100644 index 0000000..7b23945 --- /dev/null +++ b/tests/run_octo_parity.cmake @@ -0,0 +1,50 @@ +if(NOT EXISTS "${CKPT}") + message(FATAL_ERROR "Octo GGUF not found: ${CKPT}") +endif() + +get_filename_component(DUMP_PARENT "${DUMP_DIR}" DIRECTORY) +file(MAKE_DIRECTORY "${DUMP_PARENT}") + +set(T5_ARGS "") +if(DEFINED T5_INJECT AND NOT "${T5_INJECT}" STREQUAL "") + list(APPEND T5_ARGS --t5-inject "${T5_INJECT}") +endif() + +set(TRANSFORMER_TOL_ARGS "") +if(DEFINED TRANSFORMER_TOL AND NOT "${TRANSFORMER_TOL}" STREQUAL "") + list(APPEND TRANSFORMER_TOL_ARGS --transformer-tol "${TRANSFORMER_TOL}") +endif() + +set(UNNORM_DATASET_ARGS "") +if(DEFINED UNNORM_DATASET AND NOT "${UNNORM_DATASET}" STREQUAL "") + list(APPEND UNNORM_DATASET_ARGS --unnorm-dataset "${UNNORM_DATASET}") +endif() + +set(KNOWN_MISMATCH_ARGS "") +if(DEFINED KNOWN_MISMATCH AND NOT "${KNOWN_MISMATCH}" STREQUAL "") + list(APPEND KNOWN_MISMATCH_ARGS --known-mismatch "${KNOWN_MISMATCH}") +endif() + +set(EXCLUDE_ARGS "") +if(DEFINED EXCLUDE AND NOT "${EXCLUDE}" STREQUAL "") + list(APPEND EXCLUDE_ARGS --exclude "${EXCLUDE}") +endif() + +# TIP-ND1-B: octo_parity_dump always runs the resident (m->backend) path now -- the gốc +# disk-read dump path it used to alternate with (RESIDENT=1 toggle, TIP-ND1-A) was deleted +# once TIP-ND1-A proved resident-on-CPU is bit-exact with it (20/20 golden cases). +execute_process( + COMMAND "${DUMPER}" --ckpt "${CKPT}" --case "${CASE_DIR}" ${T5_ARGS} ${UNNORM_DATASET_ARGS} --out "${DUMP_DIR}" + RESULT_VARIABLE dump_rc +) +if(NOT dump_rc EQUAL 0) + message(FATAL_ERROR "octo_parity_dump failed: ${dump_rc}") +endif() + +execute_process( + COMMAND "${PYTHON_EXECUTABLE}" "${VERIFY}" --golden "${CASE_DIR}" --dump "${DUMP_DIR}" ${TRANSFORMER_TOL_ARGS} ${KNOWN_MISMATCH_ARGS} ${EXCLUDE_ARGS} --report "${DUMP_DIR}/parity_report.json" + RESULT_VARIABLE verify_rc +) +if(NOT verify_rc EQUAL 0) + message(FATAL_ERROR "verify_octo_parity.py failed: ${verify_rc}") +endif()