Skip to content

cookbook: add Cosmos3-Edge TensorRT-Edge-LLM reasoner and policy notebooks #300

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ConstBob wants to merge 2 commits into
base: main
Choose a base branch
from
Open

cookbook: add Cosmos3-Edge TensorRT-Edge-LLM reasoner and policy notebooks #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9afe232
docs(cosmos3): add TensorRT-Edge-LLM reasoner and policy cookbooks
ConstBob Oct 21, 2026
f19c188
docs(cosmos3): build the Edge-LLM policy cookbook with the direct bui…
ConstBob Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions .agents/skills/cookbooks/cosmos3/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,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**), Simulation ([Transfer-Auto](simulation/auto/auto.md), [Sim-Auto](simulation/auto/run_sim_auto_with_vllm_omni.ipynb), [Sim-Transfer](simulation/transfer/transfer.md#vllm-omni)) |
Expand Down Expand Up @@ -434,6 +435,54 @@ 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).
Reasoning exports ONNX, then builds engines on the target (Jetson Thor/Orin,
DRIVE, DGX Spark, or an x86 developer build). Policy on v0.11.0 builds engines
with the direct builder.

Use [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) **v0.11.0**.

@pengcuo pengcuo Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Per the PR description, v0.11.0's cosmos3_policy_inference does not load engines produced by the v0.11.0 tensorrt-edgellm-export --task policy + cosmos3_policy_build path documented here, and the only end-to-end smoke was on v0.10.0. Pin what was actually run, or re-run on the version that is pinned, before this merges.

@ConstBob ConstBob Oct 23, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The policy cookbook now uses tensorrt-edgellm-build on the pinned v0.11.0 tag (f19c188), instead of tensorrt-edgellm-export --task policy followed by cosmos3_policy_build. On that tag the export writes the UND input attention_pos_id and VAE height / width only, while cosmos3_policy_inference binds attention_position_ids and reads builder_config.content_height / content_width. The direct builder writes that contract.

I re-ran the documented build and inference on v0.11.0 (95515c2) on an x86 RTX PRO 6000 (SM120): tensorrt-edgellm-build --components und-prefill,gen,vae-encoder, then cosmos3_policy_inference --viewPoint concat_view --state 0,0,0,0,0,0,0,0. action.json is [1, 32, 8], finite=true. The zero state is a runtime smoke, not a DROID quality check. The PR testing section now records this run.

Complete the upstream
[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)
guide first (Python package on the host; C++ runtime on the device).

Policy examples need `-DBUILD_EXPERIMENTAL_MODELS=ON` so
`cosmos3_policy_inference` exists, and `-DENABLE_CUTE_DSL=ALL` for
[`tensorrt-edgellm-build`](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/direct-engine-builder.html).
v0.11.0 `tensorrt-edgellm-export --task policy` writes `attention_pos_id` and
VAE `height` / `width` only. The v0.11.0 runtime binds `attention_position_ids`
and reads `builder_config.content_height` / `content_width`. The direct builder
writes that contract. `cosmos3_policy_inference` in this tag accepts
`--viewPoint concat_view`.

Set a workspace and the Edge checkpoints. `EDGELLM_ROOT` matches the notebook
default:

```bash
export EDGELLM_ROOT="${EDGELLM_ROOT:-$HOME/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)
(engines under `$ENGINE_DIR/reasoning`)
- Policy: [`generator/action/run_policy_with_trt_edge_llm.ipynb`](generator/action/run_policy_with_trt_edge_llm.ipynb)
(engines under `$ENGINE_DIR/policy`)

Upstream references:
[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html)
for reasoning, and the
[direct engine builder](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/direct-engine-builder.html)
for policy.

## Transformers

Local Python inference for the Cosmos3 Reasoner. This backend uses the
Expand Down
32 changes: 29 additions & 3 deletions .agents/skills/cookbooks/cosmos3/generator/action/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -292,6 +294,30 @@ 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` and
`-DENABLE_CUTE_DSL=ALL`). Then follow
[`run_policy_with_trt_edge_llm.ipynb`](./run_policy_with_trt_edge_llm.ipynb):
`tensorrt-edgellm-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 an 8-value
`--state`, 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
Expand Down
Loading
Loading