diff --git a/README.md b/README.md
index 343c5ed7..0992d4d0 100644
--- a/README.md
+++ b/README.md
@@ -87,6 +87,7 @@ This repository is the home of the models: everything for exploring, running, an
| Explore runnable notebooks | Cookbooks | browse |
| Evaluate a model | Evaluation suites — PAIBench, Physics-IQ, VLMEvalKit | hours |
| Check latency & throughput | Benchmarks | browse |
+ | Run Cosmos3-Edge on device | Reasoner and policy with TensorRT-Edge-LLM | setup |
## Generate your first video
@@ -164,7 +165,7 @@ Training, optimization, and deployment tooling lives in [Cosmos Framework](https
| [Cosmos Curator](https://github.com/NVIDIA/cosmos-curator) | Distributed data curation: processing, annotation, filtering, dedup |
| [Cosmos Evaluator](https://github.com/NVIDIA/cosmos-evaluator) | Automated evaluation system for world generation & reasoning outputs |
-Cosmos 3 runs on Diffusers, Transformers, vLLM, vLLM-Omni, SGLang, TensorRT-LLM, and NIM — pick a backend in the [environment setup guide](cookbooks/cosmos3/README.md).
+Cosmos 3 runs on Diffusers, Transformers, vLLM, vLLM-Omni, SGLang, TensorRT-LLM, TensorRT-Edge-LLM, and NIM — pick a backend in the [environment setup guide](cookbooks/cosmos3/README.md).
## Community & contributing
diff --git a/cookbooks/cosmos3/README.md b/cookbooks/cosmos3/README.md
index 9dd21aa5..f53ad1ae 100644
--- a/cookbooks/cosmos3/README.md
+++ b/cookbooks/cosmos3/README.md
@@ -10,6 +10,7 @@ backend you want to run and follow that one section.
| [Diffusers](#diffusers) | Direct generation with `Cosmos3OmniPipeline` | Generator (Audiovisual) |
| [TensorRT-LLM Generator](#tensorrt-llm-generator) | OpenAI-compatible VisualGen server (image/video/audio/action/transfer generation) | Generator (Audiovisual, Action, **Transfer**) |
| [TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) | OpenAI-compatible image/video reasoning server | Reasoner |
+| [TensorRT-Edge-LLM](#tensorrt-edge-llm) | On-device TensorRT runtime for Cosmos3-Edge reasoning and DROID policy | Reasoner (Edge), Generator (Action / Policy) |
| [Transformers](#transformers) | Hugging Face Transformers inference | Reasoner |
| [vLLM](#vllm) | OpenAI-compatible reasoning server (image/video understanding) | Reasoner |
| [vLLM-Omni](#vllm-omni) | OpenAI-compatible generation server (image/video/audio/action/transfer) | Generator (Audiovisual, Action, **Transfer**) |
@@ -448,6 +449,42 @@ The server exposes `/health` and the OpenAI-compatible API at
[Reasoner TensorRT-LLM notebook](reasoner/run_with_tensorrt_llm.ipynb) for image
and video requests.
+## TensorRT-Edge-LLM
+
+On-device **TensorRT** path for **Cosmos3-Edge** multimodal reasoning and
+**Cosmos3-Edge-Policy-DROID** policy generation. This is a different stack from
+[TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) / [Generator](#tensorrt-llm-generator):
+export ONNX on an x86 host, build TensorRT engines on the target (Jetson
+Thor/Orin, DRIVE, DGX Spark, or an x86 developer build), then run the C++
+inference binaries.
+
+Use [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) **v0.11.0**.
+Complete the upstream
+[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)
+guide first (Python export package on the host; C++ runtime on the device).
+
+Policy examples also require building the C++ runtime with
+`-DBUILD_EXPERIMENTAL_MODELS=ON`.
+
+Set a workspace and the Edge checkpoints:
+
+```bash
+export EDGELLM_ROOT="${EDGELLM_ROOT:-$PWD/TensorRT-Edge-LLM}"
+export POLICY_CHECKPOINT=nvidia/Cosmos3-Edge-Policy-DROID
+export REASONING_CHECKPOINT=nvidia/Cosmos3-Edge
+export ONNX_DIR="${ONNX_DIR:-$HOME/tensorrt-edgellm-workspace/Cosmos3-Edge/onnx}"
+export ENGINE_DIR="${ENGINE_DIR:-$HOME/tensorrt-edgellm-workspace/Cosmos3-Edge/engines}"
+```
+
+Runnable cookbooks:
+
+- Reasoner: [`reasoner/run_with_trt_edge_llm.ipynb`](reasoner/run_with_trt_edge_llm.ipynb)
+- Policy: [`generator/action/run_policy_with_trt_edge_llm.ipynb`](generator/action/run_policy_with_trt_edge_llm.ipynb)
+
+Upstream reference:
+[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html)
+(export → build → run for reasoning and policy).
+
## Transformers
Local Python inference for the Cosmos3 Reasoner. This backend uses the
diff --git a/cookbooks/cosmos3/generator/action/README.md b/cookbooks/cosmos3/generator/action/README.md
index 8cd94264..63d8d41f 100644
--- a/cookbooks/cosmos3/generator/action/README.md
+++ b/cookbooks/cosmos3/generator/action/README.md
@@ -24,13 +24,15 @@ The rest of this doc shows how to run these modes on selected embodiments direct
- [Run with vLLM-Omni](#run-with-vllm-omni)
- [Quickstart](#quickstart-3)
- [Notebook walkthrough](#vllm-omni-notebook-walkthrough)
+- [Run with TensorRT-Edge-LLM](#run-with-tensorrt-edge-llm)
- [Post-Train for Cosmos3-Nano-Policy-DROID](#post-train-for-cosmos3-nano-policy-droid)
## Overview
-Examples are shown across native PyTorch (Cosmos Framework), Diffusers,
-TensorRT-LLM, and vLLM-Omni. Every backend uses the sample
-assets under [`assets/`](./assets) and covers three tasks:
+Examples for native PyTorch (Cosmos Framework), Diffusers, TensorRT-LLM,
+and vLLM-Omni use the sample assets under [`assets/`](./assets) and cover
+the three tasks above. TensorRT-Edge-LLM covers the Edge policy task only,
+with the same sample assets.
Environment setup for all backends is centralized in the shared
[Cosmos3 cookbooks environment setup](../../README.md) guide; each backend below
@@ -292,6 +294,29 @@ write outputs under `outputs/cosmos3_action_vllm/`:
- [`run_policy_with_vllm_omni.ipynb`](./run_policy_with_vllm_omni.ipynb) — policy
inference for DROID through the async video API.
+## Run with TensorRT-Edge-LLM
+
+### Quickstart
+
+On-device TensorRT path for **Cosmos3-Edge-Policy-DROID**. Set up the shared
+[TensorRT-Edge-LLM](../../README.md#tensorrt-edge-llm) environment (upstream
+install; C++ runtime with `-DBUILD_EXPERIMENTAL_MODELS=ON`). Then follow
+[`run_policy_with_trt_edge_llm.ipynb`](./run_policy_with_trt_edge_llm.ipynb):
+export → `cosmos3_policy_build` → `cosmos3_policy_inference`.
+
+This is separate from workstation [vLLM-Omni](#run-with-vllm-omni) policy
+serving. Multimodal reasoning on the same Edge-LLM runtime is documented in
+[`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb).
+
+### Notebook walkthrough
+
+[`run_policy_with_trt_edge_llm.ipynb`](./run_policy_with_trt_edge_llm.ipynb)
+builds the same 640×540 DROID `concat_view` as the vLLM-Omni / SGLang policy
+notebooks (from `assets/droid_lerobot_example`), runs
+`cosmos3_policy_inference` with `--viewPoint concat_view`, and writes `action.json`
+(`[batch, chunk, action_dimension]`). Optional `--steps` / `--seed` match the
+upstream VLA guide.
+
## Post-Train for Cosmos3-Nano-Policy-DROID
To reproduce our post-training recipe for [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID), use the
diff --git a/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb b/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb
new file mode 100644
index 00000000..ad66c454
--- /dev/null
+++ b/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb
@@ -0,0 +1,369 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Cosmos3-Edge Policy with TensorRT-Edge-LLM\n",
+ "\n",
+ "This notebook runs **Cosmos3-Edge-Policy-DROID** with\n",
+ "[TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) (v0.11.0):\n",
+ "export the policy checkpoint, build the experimental Cosmos3 policy engines,\n",
+ "then run `cosmos3_policy_inference`.\n",
+ "\n",
+ "Policy sits under the **Generator / Action** cookbook surface (action chunks),\n",
+ "not the Reasoner text path. For Edge multimodal reasoning on the same runtime,\n",
+ "see [`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb).\n",
+ "\n",
+ "1. Points at a TensorRT-Edge-LLM checkout built with\n",
+ " `-DBUILD_EXPERIMENTAL_MODELS=ON`.\n",
+ "2. Exports `nvidia/Cosmos3-Edge-Policy-DROID` with `--task policy`.\n",
+ "3. Builds policy engines.\n",
+ "4. Builds a 640×540 DROID `concat_view` from `assets/droid_lerobot_example`\n",
+ " (same layout as the vLLM-Omni / SGLang policy notebooks) and writes\n",
+ " `action.json`.\n",
+ "\n",
+ "Complete the upstream\n",
+ "[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)\n",
+ "first. Command reference:\n",
+ "[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html).\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Setup paths\n",
+ "\n",
+ "`EDGELLM_ROOT` must contain `build/experimental_models/cosmos3/examples/` after\n",
+ "a CMake build with `-DBUILD_EXPERIMENTAL_MODELS=ON`.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "66d95612",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from pathlib import Path\n",
+ "import os\n",
+ "\n",
+ "\n",
+ "def find_repo_root(start: Path) -> Path:\n",
+ " for candidate in [start, *start.parents]:\n",
+ " if (candidate / \"README.md\").exists() and (candidate / \"cookbooks\").exists():\n",
+ " return candidate\n",
+ " raise FileNotFoundError(\"Could not locate the cosmos repository root\")\n",
+ "\n",
+ "\n",
+ "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n",
+ "DROID_ASSET_ROOT = (\n",
+ " COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"generator\" / \"action\" / \"assets\" / \"droid_lerobot_example\"\n",
+ ").resolve()\n",
+ "assert DROID_ASSET_ROOT.exists(), DROID_ASSET_ROOT\n",
+ "\n",
+ "EDGELLM_ROOT = Path(\n",
+ " os.environ.get(\"EDGELLM_ROOT\", Path.home() / \"TensorRT-Edge-LLM\")\n",
+ ").expanduser().resolve()\n",
+ "ONNX_DIR = Path(\n",
+ " os.environ.get(\n",
+ " \"ONNX_DIR\",\n",
+ " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"onnx\",\n",
+ " )\n",
+ ").expanduser().resolve()\n",
+ "ENGINE_DIR = Path(\n",
+ " os.environ.get(\n",
+ " \"ENGINE_DIR\",\n",
+ " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"engines\",\n",
+ " )\n",
+ ").expanduser().resolve()\n",
+ "POLICY_CHECKPOINT = os.environ.get(\n",
+ " \"POLICY_CHECKPOINT\", \"nvidia/Cosmos3-Edge-Policy-DROID\"\n",
+ ")\n",
+ "POLICY_PROMPT = os.environ.get(\n",
+ " \"POLICY_PROMPT\",\n",
+ " \"Pick up the object and place it in the target container.\",\n",
+ ")\n",
+ "WORK_DIR = ENGINE_DIR / \"policy\" / \"cookbook_work\"\n",
+ "WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
+ "OUTPUT_JSON = WORK_DIR / \"action.json\"\n",
+ "\n",
+ "os.environ[\"COSMOS_ROOT\"] = str(COSMOS_ROOT)\n",
+ "os.environ[\"EDGELLM_ROOT\"] = str(EDGELLM_ROOT)\n",
+ "os.environ[\"ONNX_DIR\"] = str(ONNX_DIR)\n",
+ "os.environ[\"ENGINE_DIR\"] = str(ENGINE_DIR)\n",
+ "os.environ[\"POLICY_CHECKPOINT\"] = POLICY_CHECKPOINT\n",
+ "os.environ[\"DROID_ASSET_ROOT\"] = str(DROID_ASSET_ROOT)\n",
+ "os.environ[\"POLICY_PROMPT\"] = POLICY_PROMPT\n",
+ "os.environ[\"OUTPUT_JSON\"] = str(OUTPUT_JSON)\n",
+ "\n",
+ "print(\"cosmos root:\", COSMOS_ROOT)\n",
+ "print(\"EDGELLM_ROOT:\", EDGELLM_ROOT, \"(exists:\", EDGELLM_ROOT.exists(), \")\")\n",
+ "print(\"ONNX_DIR:\", ONNX_DIR)\n",
+ "print(\"ENGINE_DIR:\", ENGINE_DIR)\n",
+ "print(\"POLICY_CHECKPOINT:\", POLICY_CHECKPOINT)\n",
+ "print(\"DROID_ASSET_ROOT:\", DROID_ASSET_ROOT)\n",
+ "print(\"prompt:\", POLICY_PROMPT)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. Export the Policy on CPU\n",
+ "\n",
+ "Requires the TensorRT-Edge-LLM Python package (`tensorrt-edgellm-export`).\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${POLICY_CHECKPOINT:?run the setup cell first}\"\n",
+ ": \"${ONNX_DIR:?run the setup cell first}\"\n",
+ "\n",
+ "mkdir -p \"$ONNX_DIR\"\n",
+ "tensorrt-edgellm-export \\\n",
+ " \"$POLICY_CHECKPOINT\" \\\n",
+ " \"$ONNX_DIR\" \\\n",
+ " --task policy\n",
+ "\n",
+ "echo \"Exported policy ONNX under $ONNX_DIR\"\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 3. Build all policy engines\n",
+ "\n",
+ "The C++ runtime must have been configured with\n",
+ "`-DBUILD_EXPERIMENTAL_MODELS=ON`. Use `--maxBatchSize N` when the runtime must\n",
+ "accept more than one prompt.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n",
+ ": \"${ONNX_DIR:?run the setup cell first}\"\n",
+ ": \"${ENGINE_DIR:?run the setup cell first}\"\n",
+ "\n",
+ "cd \"$EDGELLM_ROOT\"\n",
+ "mkdir -p \"$ENGINE_DIR\"\n",
+ "\n",
+ "./build/experimental_models/cosmos3/examples/cosmos3_policy_build \\\n",
+ " --onnxDir \"$ONNX_DIR\" \\\n",
+ " --engineDir \"$ENGINE_DIR\"\n",
+ "\n",
+ "echo \"Built policy engines under $ENGINE_DIR\"\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "3b7a74e3",
+ "metadata": {},
+ "source": [
+ "## 4. Prepare a DROID observation\n",
+ "\n",
+ "This cell extracts the first frame from each checked-in DROID camera video and\n",
+ "builds the 640×540 `concat_view` used by the vLLM-Omni / SGLang policy\n",
+ "notebooks (wrist on top, two exterior views below). Pass that composite to\n",
+ "`--image`. The Edge-LLM runtime still resizes internally (736×544 in the\n",
+ "model contract); the **input** should be the DROID multiview, not a reasoner\n",
+ "caption asset.\n",
+ "\n",
+ "Override `POLICY_PROMPT` in the setup cell if your instruction differs. The\n",
+ "upstream VLA guide's banana/bowl string is only a CLI placeholder.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e14dba81",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import subprocess\n",
+ "\n",
+ "from PIL import Image\n",
+ "\n",
+ "try:\n",
+ " import imageio_ffmpeg\n",
+ "except ImportError as exc:\n",
+ " raise RuntimeError(\n",
+ " \"Install imageio-ffmpeg in this notebook kernel: pip install imageio-ffmpeg\"\n",
+ " ) from exc\n",
+ "\n",
+ "FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()\n",
+ "CAMERA_VIDEO_PATHS = {\n",
+ " \"observation/wrist_image_left\": DROID_ASSET_ROOT\n",
+ " / \"videos\"\n",
+ " / \"observation.image.wrist_image_left\"\n",
+ " / \"chunk-000\"\n",
+ " / \"file-000.mp4\",\n",
+ " \"observation/exterior_image_1_left\": DROID_ASSET_ROOT\n",
+ " / \"videos\"\n",
+ " / \"observation.image.exterior_image_1_left\"\n",
+ " / \"chunk-000\"\n",
+ " / \"file-000.mp4\",\n",
+ " \"observation/exterior_image_2_left\": DROID_ASSET_ROOT\n",
+ " / \"videos\"\n",
+ " / \"observation.image.exterior_image_2_left\"\n",
+ " / \"chunk-000\"\n",
+ " / \"file-000.mp4\",\n",
+ "}\n",
+ "for key, video_path in CAMERA_VIDEO_PATHS.items():\n",
+ " assert video_path.exists(), f\"missing {key}: {video_path}\"\n",
+ "\n",
+ "\n",
+ "def extract_first_frame(video_path: Path, out_path: Path) -> Path:\n",
+ " if not out_path.exists() or out_path.stat().st_mtime < video_path.stat().st_mtime:\n",
+ " subprocess.run(\n",
+ " [\n",
+ " FFMPEG,\n",
+ " \"-y\",\n",
+ " \"-loglevel\",\n",
+ " \"error\",\n",
+ " \"-i\",\n",
+ " str(video_path),\n",
+ " \"-frames:v\",\n",
+ " \"1\",\n",
+ " str(out_path),\n",
+ " ],\n",
+ " check=True,\n",
+ " )\n",
+ " return out_path\n",
+ "\n",
+ "\n",
+ "frame_paths = {\n",
+ " key: extract_first_frame(video_path, WORK_DIR / f\"policy_{key.split('/')[-1]}.png\")\n",
+ " for key, video_path in CAMERA_VIDEO_PATHS.items()\n",
+ "}\n",
+ "frames = {key: Image.open(path).convert(\"RGB\") for key, path in frame_paths.items()}\n",
+ "\n",
+ "# Same geometry as run_policy_with_vllm_omni.ipynb / run_policy_with_sglang.ipynb:\n",
+ "# native 640×360 wrist on top; exteriors resized to 320×180 on the bottom row.\n",
+ "wrist = frames[\"observation/wrist_image_left\"]\n",
+ "source_size = wrist.size\n",
+ "assert all(frame.size == source_size for frame in frames.values()), (\n",
+ " \"DROID camera frames must have matching dimensions\",\n",
+ " {key: frame.size for key, frame in frames.items()},\n",
+ ")\n",
+ "target_w = wrist.width\n",
+ "bottom_h, half_w = wrist.height // 2, wrist.width // 2\n",
+ "target_h = wrist.height + bottom_h\n",
+ "left = frames[\"observation/exterior_image_1_left\"].resize(\n",
+ " (half_w, bottom_h), Image.Resampling.BILINEAR\n",
+ ")\n",
+ "right = frames[\"observation/exterior_image_2_left\"].resize(\n",
+ " (half_w, bottom_h), Image.Resampling.BILINEAR\n",
+ ")\n",
+ "policy_image = Image.new(\"RGB\", (target_w, target_h))\n",
+ "policy_image.paste(wrist, (0, 0))\n",
+ "policy_image.paste(left, (0, wrist.height))\n",
+ "policy_image.paste(right, (half_w, wrist.height))\n",
+ "\n",
+ "OBSERVATION_IMAGE = WORK_DIR / \"droid_policy_first_frame.png\"\n",
+ "policy_image.save(OBSERVATION_IMAGE)\n",
+ "os.environ[\"OBSERVATION_IMAGE\"] = str(OBSERVATION_IMAGE)\n",
+ "\n",
+ "print(\"observation:\", OBSERVATION_IMAGE, policy_image.size)\n",
+ "print(\"prompt:\", POLICY_PROMPT)\n",
+ "try:\n",
+ " from IPython.display import display\n",
+ "\n",
+ " display(policy_image)\n",
+ "except Exception:\n",
+ " pass\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "9ca81f19",
+ "metadata": {},
+ "source": [
+ "## 5. Run the policy\n",
+ "\n",
+ "`--image` is the 640×540 DROID concat from the previous cell. Pass\n",
+ "`--viewPoint concat_view` so the policy prompt matches that multiview layout\n",
+ "(the CLI default is `ego_view`). For an ordered frame list, replace `--image`\n",
+ "with `--video frame_00.png,frame_01.png,frame_02.png`. The current policy\n",
+ "conditions on the most recent frame. Optional: `--steps` (denoise steps) and\n",
+ "`--seed`.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e466b7c4",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n",
+ ": \"${ENGINE_DIR:?run the setup cell first}\"\n",
+ ": \"${OBSERVATION_IMAGE:?run the DROID observation cell first}\"\n",
+ ": \"${POLICY_PROMPT:?run the setup cell first}\"\n",
+ ": \"${OUTPUT_JSON:?run the setup cell first}\"\n",
+ "\n",
+ "cd \"$EDGELLM_ROOT\"\n",
+ "./build/experimental_models/cosmos3/examples/cosmos3_policy_inference \\\n",
+ " --engineDir \"$ENGINE_DIR\" \\\n",
+ " --image \"$OBSERVATION_IMAGE\" \\\n",
+ " --prompt \"$POLICY_PROMPT\" \\\n",
+ " --viewPoint concat_view \\\n",
+ " --output \"$OUTPUT_JSON\"\n",
+ "\n",
+ "echo \"Wrote $OUTPUT_JSON\"\n",
+ "python3 - <<'PY'\n",
+ "import json\n",
+ "import os\n",
+ "from pathlib import Path\n",
+ "data = json.loads(Path(os.environ[\"OUTPUT_JSON\"]).read_text())\n",
+ "print(json.dumps(data, indent=2)[:4000])\n",
+ "PY\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e55845c9",
+ "metadata": {},
+ "source": [
+ "## 6. Next steps\n",
+ "\n",
+ "- Override `POLICY_PROMPT` in the setup cell, or rebuild the concat with your\n",
+ " own DROID camera videos.\n",
+ "- Multimodal reasoning on TensorRT-Edge-LLM:\n",
+ " [`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb).\n",
+ "- OpenAI-compatible policy serving on workstation GPUs:\n",
+ " [`run_policy_with_vllm_omni.ipynb`](./run_policy_with_vllm_omni.ipynb).\n",
+ "- Cosmos Framework policy server:\n",
+ " [`run_policy_with_cosmos_framework.md`](./run_policy_with_cosmos_framework.md).\n"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/cookbooks/cosmos3/reasoner/README.md b/cookbooks/cosmos3/reasoner/README.md
index 89a1cd1e..cbf6e95c 100644
--- a/cookbooks/cosmos3/reasoner/README.md
+++ b/cookbooks/cosmos3/reasoner/README.md
@@ -359,3 +359,24 @@ and walks through image and video examples covering detailed captioning,
temporal localization, embodied reasoning, common-sense reasoning, 2D
grounding, describe-anything, action CoT trajectories, driving scenes,
physical plausibility, and situation understanding.
+
+## Run with TensorRT-Edge-LLM
+
+### Quickstart
+
+Set up the environment with the shared
+[TensorRT-Edge-LLM setup](../README.md#tensorrt-edge-llm) (upstream install,
+`EDGELLM_ROOT`, ONNX/engine workspace). This path targets **Cosmos3-Edge** on
+Jetson Thor/Orin, DRIVE, DGX Spark, or an x86 developer build. It is separate
+from the datacenter [TensorRT-LLM](#run-with-tensorrt-llm) server.
+
+[`run_with_trt_edge_llm.ipynb`](./run_with_trt_edge_llm.ipynb) exports
+`nvidia/Cosmos3-Edge`, builds the LLM and visual engines, writes an image
+reasoning `input.json` for `assets/robot_153.jpg`, and runs `llm_inference`.
+
+### Notebook walkthrough
+
+[`run_with_trt_edge_llm.ipynb`](./run_with_trt_edge_llm.ipynb) is a
+command-driven walkthrough (export → engine build → image caption). For
+**Cosmos3-Edge-Policy-DROID** on the same runtime, see
+[`../generator/action/run_policy_with_trt_edge_llm.ipynb`](../generator/action/run_policy_with_trt_edge_llm.ipynb).
diff --git a/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb b/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb
new file mode 100644
index 00000000..3a235b34
--- /dev/null
+++ b/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb
@@ -0,0 +1,285 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "c0762016",
+ "metadata": {},
+ "source": [
+ "# Cosmos3-Edge Reasoner with TensorRT-Edge-LLM\n",
+ "\n",
+ "This notebook runs **Cosmos3-Edge** multimodal reasoning with\n",
+ "[TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) (v0.11.0):\n",
+ "export the checkpoint to ONNX, build TensorRT engines, then run the C++\n",
+ "`llm_inference` binary.\n",
+ "\n",
+ "It is **not** the datacenter [TensorRT-LLM](./run_with_tensorrt_llm.ipynb)\n",
+ "OpenAI-compatible server path.\n",
+ "\n",
+ "1. Points at a TensorRT-Edge-LLM checkout and a local workspace for ONNX/engines.\n",
+ "2. Exports `nvidia/Cosmos3-Edge` with `--task reasoning`.\n",
+ "3. Builds the LLM and visual engines.\n",
+ "4. Runs an image caption request against the cookbook asset `robot_153.jpg`.\n",
+ "\n",
+ "Complete the upstream\n",
+ "[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)\n",
+ "first (Python export on an x86 host; C++ runtime on Jetson Thor/Orin, DRIVE,\n",
+ "DGX Spark, or an x86 developer build). The authoritative command sequence is\n",
+ "the upstream\n",
+ "[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html).\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Setup paths\n",
+ "\n",
+ "Set `EDGELLM_ROOT` to your TensorRT-Edge-LLM checkout (the directory that\n",
+ "contains `build/` after you compile the runtime). Optional: override\n",
+ "`ONNX_DIR` / `ENGINE_DIR`.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from pathlib import Path\n",
+ "import os\n",
+ "import json\n",
+ "\n",
+ "\n",
+ "def find_repo_root(start: Path) -> Path:\n",
+ " for candidate in [start, *start.parents]:\n",
+ " if (candidate / \"README.md\").exists() and (candidate / \"cookbooks\").exists():\n",
+ " return candidate\n",
+ " raise FileNotFoundError(\"Could not locate the cosmos repository root\")\n",
+ "\n",
+ "\n",
+ "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n",
+ "REASONER_ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"reasoner\" / \"assets\"\n",
+ "assert REASONER_ASSETS.exists(), REASONER_ASSETS\n",
+ "\n",
+ "EDGELLM_ROOT = Path(\n",
+ " os.environ.get(\"EDGELLM_ROOT\", Path.home() / \"TensorRT-Edge-LLM\")\n",
+ ").expanduser().resolve()\n",
+ "ONNX_DIR = Path(\n",
+ " os.environ.get(\n",
+ " \"ONNX_DIR\",\n",
+ " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"onnx\",\n",
+ " )\n",
+ ").expanduser().resolve()\n",
+ "ENGINE_DIR = Path(\n",
+ " os.environ.get(\n",
+ " \"ENGINE_DIR\",\n",
+ " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"engines\",\n",
+ " )\n",
+ ").expanduser().resolve()\n",
+ "REASONING_CHECKPOINT = os.environ.get(\"REASONING_CHECKPOINT\", \"nvidia/Cosmos3-Edge\")\n",
+ "WORK_DIR = ENGINE_DIR / \"reasoning\" / \"cookbook_work\"\n",
+ "WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
+ "\n",
+ "os.environ[\"COSMOS_ROOT\"] = str(COSMOS_ROOT)\n",
+ "os.environ[\"REASONER_ASSETS\"] = str(REASONER_ASSETS)\n",
+ "os.environ[\"EDGELLM_ROOT\"] = str(EDGELLM_ROOT)\n",
+ "os.environ[\"ONNX_DIR\"] = str(ONNX_DIR)\n",
+ "os.environ[\"ENGINE_DIR\"] = str(ENGINE_DIR)\n",
+ "os.environ[\"REASONING_CHECKPOINT\"] = REASONING_CHECKPOINT\n",
+ "os.environ[\"WORK_DIR\"] = str(WORK_DIR)\n",
+ "\n",
+ "print(\"cosmos root:\", COSMOS_ROOT)\n",
+ "print(\"EDGELLM_ROOT:\", EDGELLM_ROOT, \"(exists:\", EDGELLM_ROOT.exists(), \")\")\n",
+ "print(\"ONNX_DIR:\", ONNX_DIR)\n",
+ "print(\"ENGINE_DIR:\", ENGINE_DIR)\n",
+ "print(\"REASONING_CHECKPOINT:\", REASONING_CHECKPOINT)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. Export the Reasoner on CPU\n",
+ "\n",
+ "Run this on the export host (after `pip install -e .` in the TensorRT-Edge-LLM\n",
+ "checkout). It downloads `nvidia/Cosmos3-Edge` on first use.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${REASONING_CHECKPOINT:?run the setup cell first}\"\n",
+ ": \"${ONNX_DIR:?run the setup cell first}\"\n",
+ "\n",
+ "mkdir -p \"$ONNX_DIR/reasoning\"\n",
+ "tensorrt-edgellm-export \\\n",
+ " \"$REASONING_CHECKPOINT\" \\\n",
+ " \"$ONNX_DIR/reasoning\" \\\n",
+ " --task reasoning\n",
+ "\n",
+ "echo \"Exported reasoning ONNX under $ONNX_DIR/reasoning\"\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 3. Build the LLM and visual engines\n",
+ "\n",
+ "Run these from the TensorRT-Edge-LLM tree after the C++ runtime is built\n",
+ "(`./build/...` binaries exist). Use the platform CMake flags from the upstream\n",
+ "installation guide; reasoning does **not** require\n",
+ "`-DBUILD_EXPERIMENTAL_MODELS=ON`.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n",
+ ": \"${ONNX_DIR:?run the setup cell first}\"\n",
+ ": \"${ENGINE_DIR:?run the setup cell first}\"\n",
+ "\n",
+ "cd \"$EDGELLM_ROOT\"\n",
+ "mkdir -p \"$ENGINE_DIR/reasoning\"\n",
+ "\n",
+ "./build/examples/llm/llm_build \\\n",
+ " --onnxDir \"$ONNX_DIR/reasoning/llm\" \\\n",
+ " --engineDir \"$ENGINE_DIR/reasoning\" \\\n",
+ " --maxInputLen 2048 \\\n",
+ " --maxKVCacheCapacity 4096\n",
+ "\n",
+ "./build/examples/multimodal/visual_build \\\n",
+ " --onnxDir \"$ONNX_DIR/reasoning/visual\" \\\n",
+ " --engineDir \"$ENGINE_DIR/reasoning\"\n",
+ "\n",
+ "echo \"Built reasoning engines under $ENGINE_DIR/reasoning\"\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 4. Write an image reasoning request\n",
+ "\n",
+ "The C++ `llm_inference` CLI accepts the standard\n",
+ "[image message format](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/format/input-format.html#message-fields).\n",
+ "This cell writes `input.json` pointing at the cookbook asset.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "image_path = (REASONER_ASSETS / \"robot_153.jpg\").resolve()\n",
+ "assert image_path.exists(), image_path\n",
+ "\n",
+ "input_path = WORK_DIR / \"input.json\"\n",
+ "output_path = WORK_DIR / \"output.json\"\n",
+ "\n",
+ "payload = {\n",
+ " \"batch_size\": 1,\n",
+ " \"max_generate_length\": 256,\n",
+ " \"requests\": [\n",
+ " {\n",
+ " \"messages\": [\n",
+ " {\n",
+ " \"role\": \"user\",\n",
+ " \"content\": [\n",
+ " {\"type\": \"image\", \"image\": str(image_path)},\n",
+ " {\"type\": \"text\", \"text\": \"Caption the image in detail.\"},\n",
+ " ],\n",
+ " }\n",
+ " ]\n",
+ " }\n",
+ " ],\n",
+ "}\n",
+ "input_path.write_text(json.dumps(payload, indent=2) + \"\\n\")\n",
+ "os.environ[\"INPUT_JSON\"] = str(input_path)\n",
+ "os.environ[\"OUTPUT_JSON\"] = str(output_path)\n",
+ "\n",
+ "print(\"wrote\", input_path)\n",
+ "print(\"image:\", image_path)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 5. Run reasoning\n",
+ "\n",
+ "Inference reads the engines under `ENGINE_DIR/reasoning` and writes\n",
+ "`output.json`.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%bash\n",
+ "set -euo pipefail\n",
+ ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n",
+ ": \"${ENGINE_DIR:?run the setup cell first}\"\n",
+ ": \"${INPUT_JSON:?run the input-json cell first}\"\n",
+ ": \"${OUTPUT_JSON:?run the input-json cell first}\"\n",
+ "\n",
+ "cd \"$EDGELLM_ROOT\"\n",
+ "./build/examples/llm/llm_inference \\\n",
+ " --engineDir \"$ENGINE_DIR/reasoning\" \\\n",
+ " --multimodalEngineDir \"$ENGINE_DIR/reasoning\" \\\n",
+ " --inputFile \"$INPUT_JSON\" \\\n",
+ " --outputFile \"$OUTPUT_JSON\"\n",
+ "\n",
+ "echo \"Wrote $OUTPUT_JSON\"\n",
+ "python3 - <<'PY'\n",
+ "import os\n",
+ "from pathlib import Path\n",
+ "out = Path(os.environ[\"OUTPUT_JSON\"])\n",
+ "print(out.read_text()[:4000])\n",
+ "PY\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 6. Next steps\n",
+ "\n",
+ "- Change the prompt or image path in the input-json cell and re-run inference.\n",
+ "- For **Cosmos3-Edge-Policy-DROID** on TensorRT-Edge-LLM, see\n",
+ " [`../generator/action/run_policy_with_trt_edge_llm.ipynb`](../generator/action/run_policy_with_trt_edge_llm.ipynb)\n",
+ " (build the runtime with `-DBUILD_EXPERIMENTAL_MODELS=ON`).\n",
+ "- Prefer an OpenAI-compatible server on workstation GPUs? Use\n",
+ " [`run_with_vllm.ipynb`](./run_with_vllm.ipynb) or\n",
+ " [`run_with_tensorrt_llm.ipynb`](./run_with_tensorrt_llm.ipynb) (Nano/Super).\n",
+ "- Python-first Edge Reasoner without TensorRT engines:\n",
+ " [`run_with_transformers.ipynb`](./run_with_transformers.ipynb).\n"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}