cosmos_edge_policy_so101

cosmos_edge_policy_so101

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

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 scores 3.1% on the same benchmark.

Runs on a Jetson Thor: see thor_deploy/ 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

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.

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.

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/ 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.

# 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, 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.