Cosmos3-edge-generation-libero4in1

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

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

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