Cosmos3-Edge Action Policy — LIBERO 4in1 SFT

Fine-tuned action-policy checkpoints on LIBERO 4in1 benchmark. Backbone: Cosmos3-Edge Policy (OmniMoT + Wan2.2-VAE + Nemotron-3 3B reasoner).

Project

Action SFT (supervised fine-tuning) on top of the Cosmos3-Edge Policy DROID backbone, targeting the four LIBERO suites (libero_spatial, libero_object, libero_goal, libero_10). The repo ships Megatron-Core Distributed Checkpoint (DCP) snapshots at multiple training iterations, plus the resolved training config and the latest loss-curve PNG.

The four suites are sampled proportionally and concatenated into a single mixed dataset; the SFT objective is rectified-flow action generation (action loss weight = 10.0, loss scale = 10.0), with an effective batch size of 2048 (16 grad-accum × 128 max-samples-per-batch).

Architecture

Component	Implementation
Backbone	`OmniMoTModel` (`cosmos_framework.model.generator.omni_mot_model`)
Latent tokenizer	`Wan2pt2VAEInterface` — Wan2.2 VAE (`Wan2.2_VAE.pth`)
Reasoning LLM	`nvidia/Cosmos3-Edge-Reasoner` (Nemotron-3 Dense VL MoT, 3B backbone + lm_head ≈ 3.4B total)
Action head	Custom action modality embed + `action2llm` / `llm2action` MLPs (lr × 5)
Diffusion expert	Rectified-flow, logitnormal timestep, train-time shift table
Precision	bfloat16
Distributed	FSDP (data-parallel sharded)

Training Configuration

Data: LIBERO 4in1 (4 suites × 10 tasks ≈ 1700 episodes), pre-encoded with Wan2.2-VAE (see Related Resources).
Frame: 256×256, 20 FPS, windowed encoding, latent cache root /disk/rl/data/LIBERO_LeRobot_v3_cosmos_exact_window_shared_vae_v1/.
Action representation: frame_wise_relative, 6d rotation, quantile_rot normalization, dim = 10.
Optimizer: FusedAdam (betas=0.9/0.99, eps=1e-8, wd=0.05), lr=5e-5 with 5× multiplier for action2llm / llm2action / action_modality_embed.
Scheduler: LambdaLinear, 200 warm-up steps, 16000-step cosine cycle, f_max=1.0 → f_min=0.0.
Effective batch size: 2048 (16 grad-accum × 128 max-samples).
Max iter: 5000; grad-clip = 1.0; AMP off (bfloat16).
Init: resume from iter_000000450.
Save iter: 50.
Seeds: training=42, dataloader=42, episode-shuffle=42.

The full Hydra / OmegaConf-style config is committed at config.yaml.

Training Loss

Latest loss curve (flow_matching_loss_action + flow_matching_loss_vision):

Checkpoints

Iter	Components	Approx. size
`iter_000002000`	`model/` `optim/` `scheduler/` `trainer/` (DCP shards)	~16 GB
`iter_000002800`	`model/` `optim/` `scheduler/` `trainer/` (DCP shards)	~16 GB

Each iter_xxxxx/ directory contains the DCP shards (__0_0.distcp, plus .metadata). Load with torch.distributed.checkpoint / Megatron-Core DistributedCheckpointer (checkpoint.type = cosmos_framework.checkpoint.dcp.DistributedCheckpointer).

Loading example (single-rank)

import torch
from cosmos_framework.checkpoint.dcp import DistributedCheckpointer

ckpt = DistributedCheckpointer(config=cfg.checkpoint)
state = ckpt.load(
    checkpoint_path='/path/to/iter_000002800',
    model=model,
    optimizer=optimizer,
    scheduler=scheduler,
)

For the inference-only path (scripts/action_policy_server_libero.py), use --no-use-ema-weights and pass the iter directory via --checkpoint-path.

Repo Layout

MangoGoes/Cosmos3-edge-generation-libero4in1/
├── README.md
├── config.yaml                              # resolved training config
├── libero_4in1_sft_loss_latest.png          # latest loss-curve snapshot
├── iter_000002000/
│   ├── model/{__0_0.distcp, .metadata}
│   ├── optim/{__0_0.distcp, .metadata}
│   ├── scheduler/{__0_0.distcp, .metadata}
│   └── trainer/{__0_0.distcp, .metadata}
└── iter_000002800/
    ├── model/{__0_0.distcp, .metadata}
    ├── optim/{__0_0.distcp, .metadata}
    ├── scheduler/{__0_0.distcp, .metadata}
    └── trainer/{__0_0.distcp, .metadata}

Related Resources

Pre-encoded latent cache (Wan2.2-VAE): MangoGoes/libero4in1_wan2.2vae_latent_dataset (53.19 GB, 1717 episodes).
Original LIBERO dataset (un-encoded): Lifelong-Robot-Learning/LIBERO.
Cosmos3-Edge Policy backbone (pre-trained SFT base): nvidia/Cosmos-Reason1-Edge.
Wan2.2 VAE weights: bundled inside Cosmos3 project (examples/checkpoints/wan22_vae/Wan2.2_VAE.pth).

Reproducing the Run

# 1. Ensure latent cache is on disk (or download from the dataset repo above)
ls /disk/rl/data/LIBERO_LeRobot_v3_cosmos_exact_window_shared_vae_v1/{libero_spatial,libero_object,libero_goal,libero_10}

# 2. Launch training with the committed config
cd /disk/rl/psm_wma/cosmos-framework
python -m cosmos_framework.scripts.train \
  --config-path .. \
  --config-name config \
  --config_file=/path/to/this/repo/config.yaml \
  checkpoint.load_path=/disk/rl/psm_wma/cosmos-framework/outputs/train/cosmos3_action_libero/action_sft/edge_libero_4in1/checkpoints/iter_000000450

# 3. Launch inference server against the latest iter
CHECKPOINT_PATH=/path/to/iter_000002800 \
NUM_STEPS=30 GUIDANCE=1.0 \
  bash examples/launch_action_server_libero_edge_all.sh

Downloads last month: 7

Video Preview

Robotics