Skip to content

Update policy example with vllm_omni #306

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
MaciejBalaNV wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Update policy example with vllm_omni #306

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
beeae11
Update policy example with vllm_omni
MaciejBalaNV Aug 7, 2026
b1ce090
Small setup improvements to action notebooks for vllm and sglang
MaciejBalaNV Aug 10, 2026
3f97f11
Updated COSMOS3_REPO documentation
MaciejBalaNV Aug 10, 2026
36256c4
Improve the lerobot vllm-omni example
MaciejBalaNV Aug 11, 2026
675292d
Improve fd vllm-omni example
MaciejBalaNV Aug 11, 2026
f86c8fd
Fixed installs
MaciejBalaNV Aug 11, 2026
2afff58
Fix video preview
MaciejBalaNV Aug 11, 2026
434b465
Improved install setup
MaciejBalaNV Aug 11, 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
38 changes: 34 additions & 4 deletions cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@
"docker run -d --init --name cosmos3-sglang-notebook \\\n",
" --runtime nvidia --gpus '\"device=0\"' \\\n",
" -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n",
" -v ~/.cache/huggingface:/root/.cache/huggingface \\\n",
" -v \"$PWD:/workspace\" \\\n",
" -p 30000:30000 --ipc=host \\\n",
" lmsysorg/sglang:dev \\\n",
Expand All @@ -45,14 +44,43 @@
"```\n"
]
},
{
"cell_type": "markdown",
"id": "fdvl-install-md",
"metadata": {},
"source": [
"## Install Notebook Dependencies\n",
"\n",
"Run this cell in the notebook kernel before the data-loading and visualization cells.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fdvl-install-code",
"metadata": {},
"outputs": [],
"source": [
"%pip install matplotlib pandas pyarrow \"lerobot[dataset]\"\n",
"%pip install av==18.0.0\n"
]
},
{
"cell_type": "markdown",
"id": "fdvl-vars-md",
"metadata": {},
"source": [
"## Configure Notebook Variables\n",
"\n",
"Run this cell after the SGLang server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_sglang/` by default.\n"
"Run this cell after the SGLang server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_sglang/` by default.\n",
"\n",
"Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n",
"\n",
"```bash\n",
"export COSMOS3_REPO=/path/to/cosmos-framework\n",
"```\n",
"\n",
"This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running this cell.\n"
]
},
{
Expand Down Expand Up @@ -462,7 +490,7 @@
"def make_preview(src: Path, crf: int = 28) -> Path:\n",
" \"\"\"Re-encode `src` to a compact, browser-friendly mp4 (cached).\"\"\"\n",
" preview = src.with_name(f\"{src.stem}_preview.mp4\")\n",
" if not preview.exists():\n",
" if not preview.exists() or preview.stat().st_mtime < src.stat().st_mtime:\n",
" subprocess.run(\n",
" [FFMPEG, \"-y\", \"-loglevel\", \"error\", \"-i\", str(src),\n",
" \"-c:v\", \"libx264\", \"-crf\", str(crf),\n",
Expand All @@ -488,7 +516,9 @@
"source": [
"## Robotics\n",
"\n",
"In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n"
"In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n",
"\n",
"> Note: Before running this LeRobot section, PyAV may need to be reinstalled as `av==15.1.0`.\n"
]
},
{
Expand Down
130 changes: 118 additions & 12 deletions cookbooks/cosmos3/generator/action/run_fd_with_vllm_omni.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,27 @@
"Forward dynamics predicts future visual observations from an initial image and an action trajectory. This notebook contains separate AV, robotics, UMI, and human hand-pose sections that each build their own input spec, run inference, and visualize generated videos.\n"
]
},
{
"cell_type": "markdown",
"id": "fdvl-install-md",
"metadata": {},
"source": [
"## Install Notebook Dependencies\n",
"\n",
"Run this cell in the notebook kernel before the data-loading and visualization cells.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fdvl-install-code",
"metadata": {},
"outputs": [],
"source": [
"%pip install matplotlib pandas pyarrow \"lerobot[dataset]\"\n",
"%pip install av==18.0.0\n"
]
},
{
"cell_type": "markdown",
"id": "4100dead",
Expand Down Expand Up @@ -96,7 +117,20 @@
"source": [
"## Configure Notebook Variables\n",
"\n",
"Run this cell after the vLLM-Omni server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_vllm/` by default.\n"
"Run this cell after the vLLM-Omni server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_vllm/` by default.\n",
"\n",
"The configuration supports these environment variables:\n",
"\n",
"- `COSMOS3_VLLM_OUTPUT_ROOT`: optional override for the output directory. The notebook resolves this into the `COSMOS3_OUTPUT_ROOT` Python variable, which contains generated inputs, chunk outputs, and stitched videos.\n",
"- `COSMOS3_DROID_ROOT`: optional override for the DROID LeRobot dataset. Set it to a version-named dataset root such as `/path/to/droid_plus_lerobot_640x360_20260412`, containing the `success/` split. If unset, the notebook creates a versioned alias for the bundled `assets/droid_lerobot_example` sample.\n",
"\n",
"Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n",
"\n",
"```bash\n",
"export COSMOS3_REPO=/path/to/cosmos-framework\n",
"```\n",
"\n",
"This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running this cell.\n"
]
},
{
Expand All @@ -108,6 +142,7 @@
"source": [
"from pathlib import Path\n",
"import os\n",
"import json\n",
"\n",
"\n",
"def find_repo_root(start: Path) -> Path:\n",
Expand Down Expand Up @@ -140,7 +175,42 @@
"print(\"COSMOS3_REPO:\", COSMOS3_REPO)\n",
"print(\"COSMOS3_INPUT_DIR:\", COSMOS3_INPUT_DIR)\n",
"print(\"COSMOS3_OUTPUT_ROOT:\", COSMOS3_OUTPUT_ROOT)\n",
"print(\"COSMOS3_VLLM_BASE_URL:\", VLLM_BASE_URL)\n"
"print(\"COSMOS3_VLLM_BASE_URL:\", VLLM_BASE_URL)\n",
"\n",
"\n",
"_ACTION_VIEWPOINT_FRAMING = {\n",
" \"ego_view\": \"This video is captured from a first-person perspective looking at the scene.\",\n",
" \"third_person_view\": \"This video is captured from a third-person perspective looking towards the agent from the front.\",\n",
" \"wrist_view\": \"This video is captured from a wrist-mounted camera.\",\n",
" \"concat_view\": \"This video contains concatenated views from multiple camera perspectives.\",\n",
"}\n",
"_ACTION_ASPECT_RATIOS = {\"1,1\": 1.0, \"4,3\": 4 / 3, \"3,4\": 3 / 4, \"16,9\": 16 / 9, \"9,16\": 9 / 16}\n",
"\n",
"\n",
"def build_action_prompt(record: dict, *, width: int, height: int, default_prompt: str) -> str:\n",
" description = str(record.get(\"prompt\") or \"\").strip() or default_prompt\n",
" if not description.endswith((\".\", \"!\", \"?\")):\n",
" description += \".\"\n",
"\n",
" num_frames = int(record[\"action_chunk_size\"]) + 1\n",
" fps = int(record[\"fps\"])\n",
" duration_seconds = num_frames / fps\n",
" aspect_ratio = min(_ACTION_ASPECT_RATIOS, key=lambda key: abs(width / height - _ACTION_ASPECT_RATIOS[key]))\n",
" minutes, seconds = divmod(round(duration_seconds), 60)\n",
" action_prompt = {}\n",
" framing = _ACTION_VIEWPOINT_FRAMING.get(record.get(\"view_point\"))\n",
" if framing is not None:\n",
" action_prompt[\"cinematography\"] = {\"framing\": framing}\n",
" action_prompt.update(\n",
" {\n",
" \"actions\": [{\"time\": f\"0:00-{minutes}:{seconds:02d}\", \"description\": description}],\n",
" \"duration\": f\"{int(duration_seconds)}s\",\n",
" \"fps\": float(fps),\n",
" \"resolution\": {\"H\": height, \"W\": width},\n",
" \"aspect_ratio\": aspect_ratio,\n",
" }\n",
" )\n",
" return json.dumps(action_prompt)\n"
]
},
{
Expand Down Expand Up @@ -412,11 +482,12 @@
" \"domain_name\": record[\"domain_name\"],\n",
" \"action_chunk_size\": record[\"action_chunk_size\"],\n",
" \"image_size\": record[\"image_size\"],\n",
" \"view_point\": record[\"view_point\"],\n",
" \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n",
" \"guardrails\": False,\n",
" }\n",
" prompt = str(record.get(\"prompt\") or \"\").strip() or \"A robot manipulates an object.\"\n",
" prompt = build_action_prompt(\n",
" record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n",
" )\n",
" form = {\n",
" \"prompt\": prompt,\n",
" \"num_frames\": record[\"action_chunk_size\"] + 1,\n",
Expand Down Expand Up @@ -509,7 +580,7 @@
"def make_preview(src: Path, crf: int = 28) -> Path:\n",
" \"\"\"Re-encode `src` to a compact, browser-friendly mp4 (cached).\"\"\"\n",
" preview = src.with_name(f\"{src.stem}_preview.mp4\")\n",
" if not preview.exists():\n",
" if not preview.exists() or preview.stat().st_mtime < src.stat().st_mtime:\n",
" subprocess.run(\n",
" [FFMPEG, \"-y\", \"-loglevel\", \"error\", \"-i\", str(src),\n",
" \"-c:v\", \"libx264\", \"-crf\", str(crf),\n",
Expand All @@ -535,7 +606,11 @@
"source": [
"## Robotics\n",
"\n",
"In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n"
"In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n",
"\n",
"> Note: Before running this LeRobot section, PyAV may need to be reinstalled as `av==15.1.0`.\n",
"\n",
"`DROIDLeRobotDataset` resolves its feature configuration from the dataset root name. The bundled sample uses a friendly asset name, so the setup cell exposes it through the matching `droid_plus_lerobot_640x360_20260412/success` layout. To use another dataset, set `COSMOS3_DROID_ROOT` to its version-named root directory.\n"
]
},
{
Expand Down Expand Up @@ -569,10 +644,39 @@
"\n",
"from cosmos_framework.data.generator.action.datasets import DROIDLeRobotDataset\n",
"\n",
"robotics_dataset_root = resolve_input(\"cookbooks/cosmos3/generator/action/assets/droid_lerobot_example\")\n",
"robotics_dataset = DROIDLeRobotDataset(root=robotics_dataset_root)\n",
"robotics_num_chunks = 5\n",
"robotics_chunk_length = 16\n",
"\n",
"# `DROIDLeRobotDataset` selects its feature configuration from the root directory\n",
"# name. The bundled sample comes from the 640x360 release but is stored under a\n",
"# friendly asset name, so expose it through the release layout expected by the loader.\n",
"droid_root_override = os.environ.get(\"COSMOS3_DROID_ROOT\")\n",
"if droid_root_override:\n",
" robotics_dataset_root = str(Path(droid_root_override).expanduser().resolve())\n",
"else:\n",
" bundled_droid_root = Path(resolve_input(\"cookbooks/cosmos3/generator/action/assets/droid_lerobot_example\"))\n",
" versioned_droid_root = (\n",
" COSMOS3_OUTPUT_ROOT\n",
" / \"datasets\"\n",
" / \"droid_plus_lerobot_640x360_20260412\"\n",
" )\n",
" success_root = versioned_droid_root / \"success\"\n",
" versioned_droid_root.mkdir(parents=True, exist_ok=True)\n",
" if success_root.is_symlink():\n",
" if success_root.resolve() != bundled_droid_root:\n",
" raise RuntimeError(f\"{success_root} points to {success_root.resolve()}, expected {bundled_droid_root}\")\n",
" elif success_root.exists():\n",
" if success_root.resolve() != bundled_droid_root:\n",
" raise RuntimeError(f\"{success_root} already exists and is not the bundled DROID sample\")\n",
" else:\n",
" success_root.symlink_to(bundled_droid_root, target_is_directory=True)\n",
" robotics_dataset_root = str(versioned_droid_root)\n",
"\n",
"robotics_dataset = DROIDLeRobotDataset(\n",
" root=robotics_dataset_root,\n",
" chunk_length=robotics_chunk_length,\n",
" use_success_only=True,\n",
")\n",
"robotics_chunk_starts = [chunk_idx * robotics_chunk_length for chunk_idx in range(robotics_num_chunks)]\n",
"assert robotics_chunk_starts[-1] < len(robotics_dataset)\n",
"\n",
Expand Down Expand Up @@ -693,14 +797,15 @@
" \"domain_name\": record[\"domain_name\"],\n",
" \"action_chunk_size\": record[\"action_chunk_size\"],\n",
" \"image_size\": record[\"image_size\"],\n",
" \"view_point\": record[\"view_point\"],\n",
" \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n",
" \"guardrails\": False,\n",
" }\n",
" if disable_guardrails:\n",
" extra_params[\"guardrails\"] = False\n",
"\n",
" prompt = str(record.get(\"prompt\") or \"\").strip() or \" \"\n",
" prompt = build_action_prompt(\n",
" record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n",
" )\n",
" form = {\n",
" \"prompt\": prompt,\n",
" \"num_frames\": record[\"action_chunk_size\"] + 1,\n",
Expand Down Expand Up @@ -1101,11 +1206,12 @@
" \"domain_name\": record[\"domain_name\"],\n",
" \"action_chunk_size\": record[\"action_chunk_size\"],\n",
" \"image_size\": record[\"image_size\"],\n",
" \"view_point\": record[\"view_point\"],\n",
" \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n",
" \"guardrails\": False,\n",
" }\n",
" prompt = str(record.get(\"prompt\") or \"\").strip() or \"A robot manipulates an object.\"\n",
" prompt = build_action_prompt(\n",
" record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n",
" )\n",
" form = {\n",
" \"prompt\": prompt,\n",
" \"num_frames\": record[\"action_chunk_size\"] + 1,\n",
Expand Down
17 changes: 16 additions & 1 deletion cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@
"docker run -d --init --name cosmos3-sglang-notebook \\\n",
" --runtime nvidia --gpus '\"device=0\"' \\\n",
" -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n",
" -v ~/.cache/huggingface:/root/.cache/huggingface \\\n",
" -v \"$PWD:/workspace\" \\\n",
" -p 30000:30000 --ipc=host \\\n",
" lmsysorg/sglang:dev \\\n",
Expand All @@ -51,6 +50,22 @@
"```\n"
]
},
{
"cell_type": "markdown",
"id": "idvl-vars-md",
"metadata": {},
"source": [
"## Configure Notebook Variables\n",
"\n",
"Before running the variables cell, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n",
"\n",
"```bash\n",
"export COSMOS3_REPO=/path/to/cosmos-framework\n",
"```\n",
"\n",
"This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running the variables cell.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand Down
16 changes: 16 additions & 0 deletions cookbooks/cosmos3/generator/action/run_id_with_vllm_omni.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,22 @@
"The request cells below use whichever Nano or Edge server is active at the configured endpoint. To keep outputs from separate model runs, set `COSMOS3_VLLM_OUTPUT_ROOT` to a different directory before running the configuration cell.\n"
]
},
{
"cell_type": "markdown",
"id": "idvl-vars-md",
"metadata": {},
"source": [
"## Configure Notebook Variables\n",
"\n",
"Before running the variables cell, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n",
"\n",
"```bash\n",
"export COSMOS3_REPO=/path/to/cosmos-framework\n",
"```\n",
"\n",
"This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running the variables cell.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand Down
11 changes: 8 additions & 3 deletions cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,13 @@
"source": [
"## Start SGLang Policy Server\n",
"\n",
"Start the server in a terminal from the `cosmos` repo root.\n",
"Start the server in a terminal from the `cosmos` repo root. Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n",
"\n",
"```bash\n",
"export COSMOS3_REPO=/path/to/cosmos-framework\n",
"```\n",
"\n",
"This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The server container mounts this checkout at `/workspace/cosmos-framework`.\n",
"\n",
"```bash\n",
"docker rm -f cosmos3-sglang-policy-notebook 2>/dev/null || true\n",
Expand All @@ -43,8 +49,7 @@
" --runtime nvidia --gpus '\"device=0\"' \\\n",
" -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n",
" -e PYTHONPATH=/workspace/cosmos-framework \\\n",
" -v ~/.cache/huggingface:/root/.cache/huggingface \\\n",
" -v \"$PWD:/workspace\" \\\n",
" -v \"$COSMOS3_REPO:/workspace/cosmos-framework\" \\\n",
" -p 30000:30000 --ipc=host \\\n",
" lmsysorg/sglang:dev \\\n",
" sglang serve \\\n",
Expand Down
Loading