Add the Cosmos3-Edge world-action model (LIBERO)

CHANGELOG.md

@@ -6,6 +6,14 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke
 
 ### Added
 
+- **New model: Cosmos3 Policy** (`lerobot/Cosmos3-Edge-Policy-LIBERO`, a world-action model: a 28-layer
+  Mixture-of-Transformers denoises 4 latent video frames and a 16-step action chunk together, 30 UniPC
+  steps). `src/models/cosmos3_policy.cpp` follows lerobot's `policies/cosmos3`: the Wan2.2 VAE encoder and the
+  two-stream backbone as ggml graphs, the UniPC solver on the host, the canvas, the prompt (JSON caption, chat
+  template, byte-level BPE from the GGUF) and the rot6d decode for `vla-cli --text ... --libero`.
+  `scripts/convert_cosmos3_policy_to_gguf.py` converts the checkpoint. Agrees with lerobot's policy in bf16 to
+  2e-3 to 2e-2 on the action chunk; the client preset `cosmos3_policy` and `eval/run_libero.sh -m cosmos3_policy`
+  serve it. CPU and CUDA only.
 - **Hexagon NPU on the Qualcomm Dragonwing IQ-10 (Linux).**
   `scripts/build_linux_hexagon.sh` cross-builds vla.cpp with `GGML_HEXAGON=ON`
   in Qualcomm's toolchain container, `vla-server` included. On the IQ-10 (v81,

CMakeLists.txt

@@ -259,6 +259,7 @@ add_library(vla_core ${VLA_CORE_LIB_TYPE}
     src/models/turbovla.cpp
     src/models/act.cpp
     src/models/octo.cpp
+    src/models/cosmos3_policy.cpp
     src/tokenizer.cpp
 )
 target_include_directories(vla_core

CONTRIBUTING.md

@@ -29,7 +29,7 @@ case say so in the commit message and back it with a LIBERO sweep.
 
 `VLA_IMG_SIZE` must match the model or `predict` returns empty: 512 for
 SmolVLA, 448 for Evo-1, 256 for GR00T N1.7 and VLA-JEPA, 256 with 2 views for
-TurboVLA, 224 for the rest. Other knobs: `VLA_BENCH_ITERS` (timing),
+TurboVLA and Cosmos3 Policy, 224 for the rest. Other knobs: `VLA_BENCH_ITERS` (timing),
 `VLA_TIMING=phase`, `VLA_EXTRA_TOKEN` / `VLA_EXTRA_COUNT` (VLA-JEPA needs its
 `<embodied>` tokens; Octo needs exactly 16 tokens, so
 `VLA_EXTRA_TOKEN=0 VLA_EXTRA_COUNT=10`), `VLA_OCTO_UNNORM_DATASET` (Octo, e.g.

README.md

@@ -101,7 +101,8 @@ Use [ggml-compiler](https://github.com/khanhnd61-vr/ggml-compiler) and follow [t
 | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero)                | Y | Y | Y | Y | Y | Y | Y |
 | [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf)       | Y | Y | Y | Y | Y | - | Y |
 | [TurboVLA](https://hf.co/vrfai/turbovla-libero-gguf)           | Y | Y | Y | Y | Y | Y | Y |
-| [ACT](https://hf.co/ravediamond/act-aloha-sim-transfer-cube-gguf) | Y | Y | Y | ~ | Y | ~ | Y |
+| [ACT](https://hf.co/ravediamond/act-aloha-sim-transfer-cube-gguf) | Y | Y | Y | ~ | Y | ~ | ~ |
+| [Cosmos3 Policy](https://hf.co/vrfai/cosmos3-policy-libero-gguf)   | ~ | ~ | - | - | - | - | - |
 
 Legend: `Y` = supported (released and benchmarked), `~` = in progress, `-` = planned.
 

eval/client/run_sim_client_direct.py

@@ -95,6 +95,11 @@
              "it, gr00t_n1_7 returns the raw normalized [40, 132] chunk and gr00t_n1_6 "
              "returns the raw [50, 128] chunk (won't drive LIBERO).",
     )
+    parser.add_argument(
+        "--gguf", type=str, default=None,
+        help="[cosmos3_policy] the GGUF vla-server runs: the caption settings (view description, fps, canvas, ...) "
+             "are read from it, so the prompt is the one the server's model was converted for.",
+    )
     parser.add_argument(
         "--bitvla-unnorm-key", type=str, default=None,
         help="[bitvla/gr00t_n1_6/gr00t_n1_7] embodiment key inside dataset_statistics.json "
@@ -120,6 +125,7 @@
         n_action_steps=args.n_action_steps,
         stats_json=args.stats_json,
         bitvla_unnorm_key=args.bitvla_unnorm_key,
+        cosmos3_gguf=args.gguf,
     )
     if args.arch == "evo1":
         client = Evo1PipelineAdapter(client=client)

eval/client/vla_cpp_client.py

@@ -58,6 +58,12 @@
     },
     "gr00t_n1_7": {"image_size": 256, "tokenizer": "nvidia/Cosmos-Reason2-2B", "max_state_dim": 132},
 
+    # Cosmos3 Policy (LIBERO): two 256x256 views, no state, a 16x10 rot6d chunk that the client turns into 16 LIBERO
+    # commands; run it with --n-action-steps 16 (the policy executes the whole chunk). The tokenizer is the one
+    # the policy was trained with (config.json tokenizer_name); a local copy works too (`tokenizer/` of the
+    # checkpoint) with --tokenizer.
+    "cosmos3_policy": {"image_size": 256, "tokenizer": "nvidia/Cosmos3 Policy", "max_state_dim": 0, "max_length": 192},
+
     "gr00t_n1_5": {"image_size": 224, "tokenizer": "lerobot/eagle2hg-processor-groot-n1p5",
                    "max_state_dim": 64, "trust_remote_code": True},
 
@@ -91,6 +97,90 @@
     + "."
 )
 
+# Cosmos3 Policy: the prompt is lerobot's Cosmos3PromptProcessorStep caption in the chat template of
+# Cosmos3TokenizerProcessorStep. The caption's settings (view description, fps, chunk, canvas, aspect ratio,
+# use_state) are the ones the GGUF was converted with (`cosmos3_policy.caption.*`, what src/models/cosmos3_policy_prompt.cpp
+# reads for vla-cli --text), so the client reads them from the GGUF vla-server runs (--gguf). Without one it assumes
+# the values of the published LIBERO checkpoint, and says so.
+COSMOS3_POLICY_N_VIEWS  = 2
+COSMOS3_POLICY_CONCAT_VIEW = "This video contains concatenated views from multiple camera perspectives."
+COSMOS3_POLICY_START_OF_GENERATION = "<|vision_start|>"
+COSMOS3_POLICY_LIBERO_SETTINGS = {  # lerobot/Cosmos3-Edge-Policy-LIBERO, compiled for a 192x320 canvas
+    "view_description": "The left half shows the third-person view; the right half shows the wrist-mounted camera.",
+    "aspect_ratio": "16,9", "fps": 20.0, "use_state": False, "chunk_size": 16,
+    "canvas_h": 192, "canvas_w": 320, "prompt_len": 192,
+}
+
+
+def _cosmos3_policy_settings(gguf_path: str | Path) -> dict:
+    """The caption settings of a Cosmos3 Policy GGUF (written by scripts/convert_cosmos3_policy_to_gguf.py --policy-dir)."""
+    try:
+        import gguf  # llama.cpp's gguf-py
+    except ImportError as e:
+        raise ImportError("reading the GGUF needs the gguf package (pip install gguf)") from e
+    reader = gguf.GGUFReader(str(gguf_path))
+
+    def field(key: str, cast):
+        f = reader.fields.get(key)
+        if f is None:
+            raise KeyError(f"{gguf_path} has no {key}; convert it with --policy-dir")
+        v = f.parts[-1]
+        return bytes(v).decode() if cast is str else cast(v[0])
+
+    return {
+        "view_description": field("cosmos3_policy.caption.view_description", str),
+        "aspect_ratio": field("cosmos3_policy.caption.aspect_ratio", str),
+        "fps": field("cosmos3_policy.caption.fps", float),
+        "use_state": bool(field("cosmos3_policy.caption.use_state", int)),
+        "chunk_size": field("cosmos3_policy.chunk_size", int),
+        "canvas_h": field("cosmos3_policy.canvas_h", int),
+        "canvas_w": field("cosmos3_policy.canvas_w", int),
+        "prompt_len": field("cosmos3_policy.prompt_len", int),
+    }
+
+
+def _cosmos3_policy_caption(task: str, s: dict) -> str:
+    """Cosmos3PromptProcessorStep.caption with 0 idle frames for the settings `s`, as the C++ builds it."""
+    task = task.strip()
+    if not task:
+        return ""
+    if not task.endswith((".", "!", "?")):
+        task = f"{task}."
+    duration = (s["chunk_size"] + 1) / s["fps"]
+    rounded = round(duration)
+    framing = " ".join(t for t in (COSMOS3_POLICY_CONCAT_VIEW, s["view_description"].strip()) if t)
+    return json.dumps({
+        "cinematography": {"framing": framing},
+        "actions": [{"time": f"0:00-{rounded // 60}:{rounded % 60:02d}", "description": task,
+                     "idle_frame": f"0 out of {s['chunk_size'] + (1 if s['use_state'] else 0)}."}],
+        "duration": f"{int(duration)}s",
+        "fps": float(int(s["fps"])),
+        "resolution": {"H": s["canvas_h"], "W": s["canvas_w"]},
+        "aspect_ratio": s["aspect_ratio"],
+    })
+
+
+def _cosmos3_policy_libero(chunk: np.ndarray) -> np.ndarray:
+    """[16, 10] rot6d rows -> [16, 7] LIBERO commands, as lerobot's Cosmos3LiberoActionProcessorStep does:
+    the two columns and their cross product make a matrix, projected onto SO(3) by an SVD, then converted to an
+    axis-angle vector (all in float64); the gripper goes from 1 open / 0 closed to -1 open / +1 closed."""
+    a = np.asarray(chunk, dtype=np.float64)
+    first, second = a[:, 3:6], a[:, 6:9]
+    m = np.stack([first, second, np.cross(first, second)], axis=-1)  # columns
+    u, _, vh = np.linalg.svd(m)
+    reflection = np.linalg.det(u @ vh) < 0
+    u = np.where(reflection[:, None, None], u * np.array([1.0, 1.0, -1.0]), u)
+    r = u @ vh
+    skew = np.stack([r[:, 2, 1] - r[:, 1, 2], r[:, 0, 2] - r[:, 2, 0], r[:, 1, 0] - r[:, 0, 1]], -1)
+    sin_a = np.linalg.norm(skew, axis=-1) / 2
+    cos_a = (np.trace(r, axis1=-2, axis2=-1) - 1) / 2
+    angle = np.arctan2(sin_a, cos_a)
+    scale = np.where(sin_a > 1e-8, angle / np.maximum(sin_a, 1e-8), 1.0)
+    axis_angle = skew * scale[:, None] / 2
+    gripper = np.clip(1.0 - 2.0 * a[:, 9:], -1.0, 1.0)
+    return np.concatenate([a[:, :3], axis_angle, gripper], axis=-1).astype(np.float32)
+
+
 def _load_pb():
 
     if "vla_pb2" in sys.modules:
@@ -159,13 +249,25 @@ def __init__(
         stats_json: str | Path | None = None,
         bitvla_unnorm_key: str | None = None,
         rel_stats_json: str | Path | None = None,
+        cosmos3_gguf: str | Path | None = None,
     ):
         if arch not in ARCH_PRESETS:
             raise ValueError(f"unknown arch {arch!r}; expected one of {sorted(ARCH_PRESETS)}")
         preset = ARCH_PRESETS[arch]
         tokenizer_name = tokenizer_name if tokenizer_name is not None else preset["tokenizer"]
         image_size     = image_size     if image_size     is not None else preset["image_size"]
         max_state_dim  = max_state_dim  if max_state_dim  is not None else preset.get("max_state_dim", 32)
+        # cosmos3_policy: the caption settings and the prompt length come from the GGUF the server runs.
+        self._cosmos3 = None
+        if arch == "cosmos3_policy":
+            if cosmos3_gguf:
+                self._cosmos3 = _cosmos3_policy_settings(cosmos3_gguf)
+                if max_length is None:
+                    max_length = self._cosmos3["prompt_len"]
+            else:
+                self._cosmos3 = dict(COSMOS3_POLICY_LIBERO_SETTINGS)
+                print("vla-cpp-direct: no --gguf given; the caption is built with the settings of "
+                      "lerobot/Cosmos3-Edge-Policy-LIBERO (pass the server's GGUF to read them from it)", flush=True)
         max_length     = max_length     if max_length     is not None else preset.get("max_length", 48)
         if tokenizer_name is None:
             raise ValueError(
@@ -670,6 +772,8 @@ def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray:
 
         if self.arch == "turbovla":
             return self._predict_chunk_turbovla(observations)
+        if self.arch == "cosmos3_policy":
+            return self._predict_chunk_cosmos3_policy(observations)
         if self.arch == "evo1":
             return self._predict_chunk_evo1(observations)
         if self.arch == "gr00t_n1_7":
@@ -820,6 +924,63 @@ def _predict_chunk_turbovla(self, observations: dict[str, Any]) -> np.ndarray:
         )
         return self._turbovla_action_unnorm(chunk)
 
+    def _cosmos3_policy_prompt_ids(self, task: str) -> list[int]:
+        """Token ids of the prompt of `task`, as Cosmos3TokenizerProcessorStep.encode builds them: one user turn
+        in the chat template, the end-of-turn token and <|vision_start|>, no special tokens added."""
+        messages = [{"role": "user", "content": [{"type": "text", "text": _cosmos3_policy_caption(task, self._cosmos3)}]}]
+        text = self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        text += self.tok.eos_token + COSMOS3_POLICY_START_OF_GENERATION
+        return self.tok(text, add_special_tokens=False)["input_ids"]
+
+    def _predict_chunk_cosmos3_policy(self, observations: dict[str, Any]) -> np.ndarray:
+        """Two views, the instruction as tokens, no state; the server returns a [16, 10] rot6d chunk in dataset
+        units (the engine de-normalizes), decoded here to [16, 7] LIBERO commands."""
+        if len(self.image_keys) != COSMOS3_POLICY_N_VIEWS:
+            raise ValueError(f"Cosmos3 Policy takes {COSMOS3_POLICY_N_VIEWS} views [third person, wrist], "
+                             f"got image keys {self.image_keys}")
+        images_f32: list[np.ndarray] = []
+        for key in self.image_keys:
+            if key not in observations:
+                raise KeyError(f"Cosmos3 Policy image key {key!r} missing; got {list(observations)}")
+            image = observations[key]
+            if isinstance(image, torch.Tensor):
+                image = image.numpy()
+            image = np.asarray(image, dtype=np.float32)
+            if image.ndim != 3 or image.shape[0] != 3:
+                raise ValueError(f"{key}: expected CHW float32 [3,H,W] in [0,1], got {image.shape}")
+            images_f32.append(np.ascontiguousarray(np.transpose(image, (1, 2, 0)), dtype=np.float32))
+
+        task = observations.get("task", "")
+        if isinstance(task, bytes):
+            task = task.decode()
+        if not isinstance(task, str) or not task.strip():
+            raise ValueError("Cosmos3 Policy requires a non-empty task instruction")
+        lang = self._cosmos3_policy_prompt_ids(task)
+        if len(lang) > self.max_length:
+            raise ValueError(f"the prompt is {len(lang)} tokens; the model was compiled for {self.max_length}")
+
+        req = self.pb.PredictRequest()
+        req.request_id = self._step
+        self._step += 1
+        for image in images_f32:
+            proto_image = req.images.add()
+            proto_image.encoding = self.pb.Image.F32_RGB_01
+            proto_image.height, proto_image.width = image.shape[:2]
+            proto_image.data = image.tobytes()
+        req.lang_tokens.extend(lang)
+        # no state, and no noise: the server samples it
+
+        self.sock.send(req.SerializeToString())
+        resp = self.pb.PredictResponse()
+        resp.ParseFromString(self.sock.recv())
+        if resp.error:
+            raise RuntimeError(f"vla-server error: {resp.error}")
+        self._last_response = resp
+        chunk = np.array(resp.action_chunk, dtype=np.float32).reshape(resp.chunk_size, resp.action_dim)
+        if chunk.shape != (self._cosmos3["chunk_size"], 10):
+            raise RuntimeError(f"expected a {self._cosmos3['chunk_size']}x10 chunk from the Cosmos3 Policy server, got {chunk.shape}")
+        return _cosmos3_policy_libero(chunk)
+
     _VJ_PS, _VJ_TPS, _VJ_MERGE, _VJ_SIDE = 16, 2, 2, 256
 
     @classmethod

eval/run_libero.sh

@@ -38,7 +38,7 @@ Usage: $(basename "$0") -i <MODELS_ROOT> [-o <OUTPUT_ROOT>] [-n <N_EPISODES>] [-
   -m MODEL         which model to run: smol | pi0 | pi05 | bit | evo1 |
                                        vla_adapter | openvla_oft |
                                        gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 |
-                                       octo | turbovla | vla_jepa | all
+                                       octo | turbovla | vla_jepa | cosmos3_policy | all
                    (default: all)
   -h               show this help
 
@@ -49,6 +49,10 @@ Env overrides: BIND_ADDR, CLIENT_ADDR, BITVLA_TOKENIZER, GR00T_N1_6_TOKENIZER,
                SERVER_BIN (prebuilt vla-server; setting it skips the build),
                VLA_JEPA_PYTHON (python for the vla_jepa client; it needs LIBERO and
                                 transformers>=5.4, default is the LIBERO venv),
+               COSMOS3_POLICY_PYTHON (python for the cosmos3_policy client, which needs LIBERO and a
+                                transformers that reads the policy's tokenizer, default is the LIBERO venv),
+               COSMOS3_POLICY_TOKENIZER (the policy's tokenizer/ directory, default nvidia/Cosmos3 Policy
+                                from the Hub), COSMOS3_POLICY_RECV_TIMEOUT_MS (client wait per chunk, default 300000),
                VLA_FIXED_NOISE_SEED (client-side noise, for paired A/B runs)
 EOF
 }
@@ -74,9 +78,9 @@ done
 shift $((OPTIND - 1))
 
 case "${MODEL}" in
-    smol|pi0|pi05|bit|evo1|vla_adapter|openvla_oft|gr00t_n1_5|gr00t_n1_6|gr00t_n1_7|octo|turbovla|vla_jepa|all) ;;
+    smol|pi0|pi05|bit|evo1|vla_adapter|openvla_oft|gr00t_n1_5|gr00t_n1_6|gr00t_n1_7|octo|turbovla|vla_jepa|cosmos3_policy|all) ;;
     *)
-        echo "ERROR: -m must be one of: smol | pi0 | pi05 | bit | evo1 | vla_adapter | openvla_oft | gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 | octo | turbovla | vla_jepa | all (got '${MODEL}')" >&2
+        echo "ERROR: -m must be one of: smol | pi0 | pi05 | bit | evo1 | vla_adapter | openvla_oft | gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 | octo | turbovla | vla_jepa | cosmos3_policy | all (got '${MODEL}')" >&2
         exit 1
         ;;
 esac
@@ -113,6 +117,9 @@ TASK_IDS="${TASK_IDS:-0 1 2 3 4 5 6 7 8 9}"
 PALIGEMMA_TOKENIZER="${PALIGEMMA_TOKENIZER:-}"
 VENV_PY="${REPO_ROOT}/eval/sim/libero/libero_uv/.venv/bin/python"
 VLA_JEPA_PYTHON="${VLA_JEPA_PYTHON:-${VENV_PY}}"
+COSMOS3_POLICY_PYTHON="${COSMOS3_POLICY_PYTHON:-${VENV_PY}}"
+COSMOS3_POLICY_TOKENIZER="${COSMOS3_POLICY_TOKENIZER:-}"
+COSMOS3_POLICY_RECV_TIMEOUT_MS="${COSMOS3_POLICY_RECV_TIMEOUT_MS:-300000}"
 CLIENT="${REPO_ROOT}/eval/client/run_sim_client_direct.py"
 BIND_ADDR="${BIND_ADDR:-tcp://*:5555}"
 CLIENT_ADDR="${CLIENT_ADDR:-tcp://localhost:5555}"
@@ -147,6 +154,7 @@ N_ACTION_STEPS_GR00T_N1_7="${N_ACTION_STEPS_GR00T_N1_7:-16}" # N1.7 H4 closeout
 N_ACTION_STEPS_OCTO="${N_ACTION_STEPS_OCTO:-4}"
 N_ACTION_STEPS_TURBOVLA="${N_ACTION_STEPS_TURBOVLA:-12}"
 N_ACTION_STEPS_VLA_JEPA="${N_ACTION_STEPS_VLA_JEPA:-7}"
+N_ACTION_STEPS_COSMOS3_POLICY="${N_ACTION_STEPS_COSMOS3_POLICY:-16}" # Cosmos3 Policy executes the whole 16-action chunk
 
 mkdir -p "${OUTPUT_ROOT}"
 OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)"
@@ -175,7 +183,12 @@ if [[ ! -x "${SERVER_BIN}" ]]; then
     echo "ERROR: ${SERVER_BIN} not found after build." >&2
     exit 1
 fi
-if [[ ! -x "${VENV_PY}" ]]; then
+# The cosmos3_policy client may run in another python that has LIBERO (COSMOS3_POLICY_PYTHON), as vla_jepa's may.
+need_libero_venv=1
+if [[ "${MODEL}" == cosmos3_policy && "${COSMOS3_POLICY_PYTHON}" != "${VENV_PY}" && -x "${COSMOS3_POLICY_PYTHON}" ]]; then
+    need_libero_venv=0
+fi
+if (( need_libero_venv )) && [[ ! -x "${VENV_PY}" ]]; then
     echo "ERROR: LIBERO venv not found at ${VENV_PY}." >&2
     echo "       Run: bash eval/sim/libero/setup_libero.sh" >&2
     exit 1
@@ -316,6 +329,17 @@ run_model() {
     if [[ "${arch}" == vla_jepa ]]; then
         client_py="${VLA_JEPA_PYTHON}"
     fi
+    # cosmos3_policy: the tokenizer is the policy's (nvidia/Cosmos3 Policy on the Hub unless a local copy is given), the
+    # caption settings are read from the GGUF the server runs, and a chunk takes seconds on a GPU but about a minute
+    # on a CPU, so the client waits longer than its 30 s default.
+    if [[ "${arch}" == cosmos3_policy ]]; then
+        client_py="${COSMOS3_POLICY_PYTHON}"
+        client_extra+=(--gguf "${server_args[0]}")
+        if [[ -n "${COSMOS3_POLICY_TOKENIZER}" ]]; then
+            client_extra+=(--tokenizer "${COSMOS3_POLICY_TOKENIZER}")
+        fi
+        client_extra+=(--recv-timeout-ms "${COSMOS3_POLICY_RECV_TIMEOUT_MS}")
+    fi
 
     # bitvla auto-loads tokenizer + dataset_statistics.json from the GGUF repo on
     # the Hub; only pass --tokenizer when BITVLA_TOKENIZER overrides with a local dir.
@@ -552,5 +576,20 @@ if should_run vla_jepa; then
     fi
 fi
 
+# cosmos3_policy: the GGUF of scripts/convert_cosmos3_policy_to_gguf.py, not
+# published yet, so a missing one is skipped instead of failing a sweep over all models.
+if should_run cosmos3_policy; then
+    cosmos3_gguf="${MODELS_ROOT}/cosmos3-policy-libero-gguf/cosmos3-policy-libero-bf16.gguf"
+    if [[ ! -f "${cosmos3_gguf}" ]]; then
+        echo "[skip] cosmos3_policy: ${cosmos3_gguf} not found"
+    else
+        run_model cosmos3_policy \
+            "${MODELS_ROOT}/cosmos3-policy-libero-gguf" \
+            "${N_ACTION_STEPS_COSMOS3_POLICY}" \
+            "" \
+            "${cosmos3_gguf}"
+    fi
+fi
+
 echo "===================="
 echo "Done. Results under ${OUTPUT_ROOT}"