Add Jetson Thor hardware-in-the-loop deployment (thor_deploy/)

28ff0ee verified about 1 month ago

6.21 kB

---
license: other
license_name: openmdw1.1-license
license_link: https://openmdw.ai/license/1-1/
base_model: nvidia/Cosmos3-Edge-Policy-DROID
library_name: cosmos
pipeline_tag: robotics
tags:
  - cosmos
  - cosmos3
  - cosmos3-edge
  - robotics
  - action-policy
  - so101
  - lerobot
  - isaac-lab
---
# cosmos_edge_policy_so101
A 6-DOF **SO-101** bin-placement action policy, post-trained from
[`nvidia/Cosmos3-Edge-Policy-DROID`](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID).
Each folder holds one training checkpoint with its LoRA adapters merged into the weights, exported
as consolidated safetensors.
| Folder | Iteration | Epoch (approx.) | Benchmark |
| --- | ---: | ---: | --- |
| `iter_6000/` | 6000 | 3.5 | not evaluated yet |
| `iter_6500/` | 6500 | 3.8 | **5 / 98 (5.1%)**, the best checkpoint |
| `iter_7000/` | 7000 (final) | 4.0 | 1 / 100 (1.0%) |

Scores are on the benchmark below with a 25 s episode limit. Earlier checkpoints of the same run
scored 0 / 50 (iteration 1500) and 0 / 100 (iteration 2500). For comparison, the
[Cosmos3-Nano SO-101 policy](https://huggingface.co/kabilanKB/cosmos_nano_policy_so101) scores
3.1% on the same benchmark.

**Runs on a Jetson Thor:** see [`thor_deploy/`](thor_deploy/README.md) for serving `iter_6500`
natively on a Thor and driving the Isaac Lab benchmark from a PC through a hardware-in-the-loop web UI.

## Training

| Setting | Value |
| --- | --- |
| Base model | `nvidia/Cosmos3-Edge-Policy-DROID` |
| Action-head warm start | `action2llm` / `llm2action` keep a separate weight row per embodiment. The DROID row (domain 8) was copied into the SO-101 row (domain 22) before training, so SO-101 starts from DROID's trained action mapping instead of random init. |
| Experiment | `action_policy_so101_edge_focus5_multi` (cosmos-framework `5e67049` plus local SO-101 support) |
| Data | `so101_bench_sim_6`: 5 single-object "Place the X in the plastic bin" instructions (97 episodes) plus the multi-object "Place each object in the plastic bin", capped at 45 episodes. 142 episodes, 10% held out. |
| Action space | Absolute `joint_pos`, 6-D (5 arm joints plus gripper, LeRobot `.pos` units). Row 0 is the current state. Chunk of 32 steps at 30 fps. |
| Normalization | minmax against the calibration bounds: joints [-100, 100], gripper [0, 100] (`so101_lerobot_stats.json`) |
| Video | `concat_view`: front wrist camera stacked on top of the overhead camera, 480p |
| Method | LoRA rank 64 / alpha 128 on `q/k/v/o_proj_moe_gen`, plus full training of the action heads |
| Schedule | Global batch 32, learning rate 1e-4, 200 warm-up steps, linear decay over 7000 iterations |
| Hardware | Iterations 0–5500 on 1× RTX PRO 6000 (batch 16 × 2 accumulation steps); 5500–7000 resumed on 2× RTX PRO 6000 (batch 16 per GPU) |
| Embodiment domain id | 22 (`so101`) |

## Benchmark setup

`so101_bench` Isaac Lab digital twin, `So101Bench-Bin-v0`: 100 single-object episodes
(`tasks/focus5.jsonl`), 25 s per episode, 32 actions executed per inference call.

## Serving

Served with `cosmos_framework.scripts.action_policy_server_robolab`, an openpi websocket server.
Every SO-101 flag below is required, and a missing one fails silently: the server starts but
returns wrong actions.

```shell
huggingface-cli download kabilanKB/cosmos_edge_policy_so101 --include "iter_7000/*" "so101_lerobot_stats.json" \
    --local-dir cosmos_edge_policy_so101
python -m cosmos_framework.scripts.action_policy_server_robolab \
    --checkpoint-path cosmos_edge_policy_so101/iter_7000 \
    --port 8000 \
    --domain-name so101 \
    --action-dim 6 \
    --arm-joint-dim 5 \
    --action-space joint_pos \
    --conditioning-fps 30 \
    --no-flip-gripper \
    --action-normalization minmax \
    --normalizer-stats-path cosmos_edge_policy_so101/so101_lerobot_stats.json \
    --view-description 'The top half is from the front-facing wrist camera. The bottom half is from the fixed overhead camera.' \
    --no-guardrails
```

These flags depend on SO-101 support in the policy server (`--arm-joint-dim`, `--no-flip-gripper`,
`--action-normalization`, `--view-description`), which is not in upstream cosmos-framework `5e67049`.
The pipeline for training, merge, export, serving and evaluation is in
[kabilankb/so101-cosmos-nano-policy](https://github.com/kabilankb/so101-cosmos-nano-policy).

**Request format:** `prompt` (the instruction), `observation/image` (the concatenated view),
`observation/joint_position` (5 values), `observation/gripper_position` (1 value).

**Response:** `action`, 32 × 6 absolute joint targets in raw `.pos` units.

## Run on a Jetson Thor

The policy server runs natively on an NVIDIA Jetson Thor (JetPack 7, CUDA 13); the ~9 GB of BF16
weights fit in the Thor's memory. [`thor_deploy/`](thor_deploy/README.md) has one-command setup and
serve scripts for the Thor, a setup script for the PC that runs the `so101_bench` simulation
(ported to Isaac Lab 3.0 / Isaac Sim 6.0), and a hardware-in-the-loop web UI that takes the Thor's IP
address, launches the policy server on the Thor over SSH, runs the benchmark and shows each camera
frame and action chunk.

```bash
# Thor
uvx hf@latest download kabilanKB/cosmos_edge_policy_so101 --include "thor_deploy/*" --local-dir ~/so101_deploy
bash ~/so101_deploy/thor_deploy/thor/setup_thor.sh && bash ~/so101-edge-thor/serve_thor.sh
```

On the Thor a 32-step chunk takes ~2.1 s with cuDNN attention and `torch.compile` (details and the
two fixes that make this possible are in the deploy README), versus 0.42 s on an RTX PRO 6000.

## Limitations

- Trained and evaluated only in simulation. It has not been tested on a physical SO-101.
- Covers five objects plus one multi-object instruction, with 18–45 demonstrations per instruction.
- Sampling is stochastic: the server draws a new seed for every request.

## License and attribution

This model is a derivative of **NVIDIA Cosmos3-Edge-Policy-DROID**, released under the
[OpenMDW License 1.1](https://openmdw.ai/license/1-1/), and is distributed under the same license.
The bundled vision encoder, processor and tokenizer files come from `nvidia/Cosmos3-Edge`.

Built with NVIDIA Cosmos. Post-trained by Kabilan KB.

	---
	license: other
	license_name: openmdw1.1-license
	license_link: https://openmdw.ai/license/1-1/
	base_model: nvidia/Cosmos3-Edge-Policy-DROID
	library_name: cosmos
	pipeline_tag: robotics
	tags:
	- cosmos
	- cosmos3
	- cosmos3-edge
	- robotics
	- action-policy
	- so101
	- lerobot
	- isaac-lab
	---

	# cosmos_edge_policy_so101

	A 6-DOF SO-101 bin-placement action policy, post-trained from
	[`nvidia/Cosmos3-Edge-Policy-DROID`](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID).
	Each folder holds one training checkpoint with its LoRA adapters merged into the weights, exported
	as consolidated safetensors.

	\| Folder \| Iteration \| Epoch (approx.) \| Benchmark \|
	\| --- \| ---: \| ---: \| --- \|
	\| `iter_6000/` \| 6000 \| 3.5 \| not evaluated yet \|
	\| `iter_6500/` \| 6500 \| 3.8 \| 5 / 98 (5.1%), the best checkpoint \|
	\| `iter_7000/` \| 7000 (final) \| 4.0 \| 1 / 100 (1.0%) \|

	Scores are on the benchmark below with a 25 s episode limit. Earlier checkpoints of the same run
	scored 0 / 50 (iteration 1500) and 0 / 100 (iteration 2500). For comparison, the
	[Cosmos3-Nano SO-101 policy](https://huggingface.co/kabilanKB/cosmos_nano_policy_so101) scores
	3.1% on the same benchmark.

	Runs on a Jetson Thor: see [`thor_deploy/`](thor_deploy/README.md) for serving `iter_6500`
	natively on a Thor and driving the Isaac Lab benchmark from a PC through a hardware-in-the-loop web UI.

	## Training

	\| Setting \| Value \|
	\| --- \| --- \|
	\| Base model \| `nvidia/Cosmos3-Edge-Policy-DROID` \|
	\| Action-head warm start \| `action2llm` / `llm2action` keep a separate weight row per embodiment. The DROID row (domain 8) was copied into the SO-101 row (domain 22) before training, so SO-101 starts from DROID's trained action mapping instead of random init. \|
	\| Experiment \| `action_policy_so101_edge_focus5_multi` (cosmos-framework `5e67049` plus local SO-101 support) \|
	\| Data \| `so101_bench_sim_6`: 5 single-object "Place the X in the plastic bin" instructions (97 episodes) plus the multi-object "Place each object in the plastic bin", capped at 45 episodes. 142 episodes, 10% held out. \|
	\| Action space \| Absolute `joint_pos`, 6-D (5 arm joints plus gripper, LeRobot `.pos` units). Row 0 is the current state. Chunk of 32 steps at 30 fps. \|
	\| Normalization \| minmax against the calibration bounds: joints [-100, 100], gripper [0, 100] (`so101_lerobot_stats.json`) \|
	\| Video \| `concat_view`: front wrist camera stacked on top of the overhead camera, 480p \|
	\| Method \| LoRA rank 64 / alpha 128 on `q/k/v/o_proj_moe_gen`, plus full training of the action heads \|
	\| Schedule \| Global batch 32, learning rate 1e-4, 200 warm-up steps, linear decay over 7000 iterations \|
	\| Hardware \| Iterations 0–5500 on 1× RTX PRO 6000 (batch 16 × 2 accumulation steps); 5500–7000 resumed on 2× RTX PRO 6000 (batch 16 per GPU) \|
	\| Embodiment domain id \| 22 (`so101`) \|

	## Benchmark setup

	`so101_bench` Isaac Lab digital twin, `So101Bench-Bin-v0`: 100 single-object episodes
	(`tasks/focus5.jsonl`), 25 s per episode, 32 actions executed per inference call.

	## Serving

	Served with `cosmos_framework.scripts.action_policy_server_robolab`, an openpi websocket server.
	Every SO-101 flag below is required, and a missing one fails silently: the server starts but
	returns wrong actions.

	```shell
	huggingface-cli download kabilanKB/cosmos_edge_policy_so101 --include "iter_7000/*" "so101_lerobot_stats.json" \
	--local-dir cosmos_edge_policy_so101

	python -m cosmos_framework.scripts.action_policy_server_robolab \
	--checkpoint-path cosmos_edge_policy_so101/iter_7000 \
	--port 8000 \
	--domain-name so101 \
	--action-dim 6 \
	--arm-joint-dim 5 \
	--action-space joint_pos \
	--conditioning-fps 30 \
	--no-flip-gripper \
	--action-normalization minmax \
	--normalizer-stats-path cosmos_edge_policy_so101/so101_lerobot_stats.json \
	--view-description 'The top half is from the front-facing wrist camera. The bottom half is from the fixed overhead camera.' \
	--no-guardrails
	```

	These flags depend on SO-101 support in the policy server (`--arm-joint-dim`, `--no-flip-gripper`,
	`--action-normalization`, `--view-description`), which is not in upstream cosmos-framework `5e67049`.
	The pipeline for training, merge, export, serving and evaluation is in
	[kabilankb/so101-cosmos-nano-policy](https://github.com/kabilankb/so101-cosmos-nano-policy).

	Request format: `prompt` (the instruction), `observation/image` (the concatenated view),
	`observation/joint_position` (5 values), `observation/gripper_position` (1 value).

	Response: `action`, 32 × 6 absolute joint targets in raw `.pos` units.

	## Run on a Jetson Thor

	The policy server runs natively on an NVIDIA Jetson Thor (JetPack 7, CUDA 13); the ~9 GB of BF16
	weights fit in the Thor's memory. [`thor_deploy/`](thor_deploy/README.md) has one-command setup and
	serve scripts for the Thor, a setup script for the PC that runs the `so101_bench` simulation
	(ported to Isaac Lab 3.0 / Isaac Sim 6.0), and a hardware-in-the-loop web UI that takes the Thor's IP
	address, launches the policy server on the Thor over SSH, runs the benchmark and shows each camera
	frame and action chunk.

	```bash
	# Thor
	uvx hf@latest download kabilanKB/cosmos_edge_policy_so101 --include "thor_deploy/*" --local-dir ~/so101_deploy
	bash ~/so101_deploy/thor_deploy/thor/setup_thor.sh && bash ~/so101-edge-thor/serve_thor.sh
	```

	On the Thor a 32-step chunk takes ~2.1 s with cuDNN attention and `torch.compile` (details and the
	two fixes that make this possible are in the deploy README), versus 0.42 s on an RTX PRO 6000.

	## Limitations

	- Trained and evaluated only in simulation. It has not been tested on a physical SO-101.
	- Covers five objects plus one multi-object instruction, with 18–45 demonstrations per instruction.
	- Sampling is stochastic: the server draws a new seed for every request.

	## License and attribution

	This model is a derivative of NVIDIA Cosmos3-Edge-Policy-DROID, released under the
	[OpenMDW License 1.1](https://openmdw.ai/license/1-1/), and is distributed under the same license.
	The bundled vision encoder, processor and tokenizer files come from `nvidia/Cosmos3-Edge`.

	Built with NVIDIA Cosmos. Post-trained by Kabilan KB.