uniwam-libero-checkpoints

UniWAM — LIBERO-10 checkpoints (Cosmos3-Edge fine-tunes)

Three fine-tunes of Cosmos3-Edge (3.37 B parameters, 1.42 B trainable; the vision–language tower stays frozen) on LIBERO-10, differing only in which tokens are given to the model as clean conditions during training. Every intermediate checkpoint is included: 1,000 · 2,000 · … · 9,000 iterations, so the whole training trajectory is available, not just the endpoint.

What a training sample looks like

A 17-frame clip at 20 fps. Frame 0 is the observation (two 256×256 cameras side by side) and frames 1–16 are the future. The action lane holds 16 tokens, one per future frame, each a 10-D end-effector delta: [Δpos(3), Δrotation as 6-D(6), gripper(1)]. A per-sample prefix length k decides how many action tokens are handed over clean:

Training recipe (identical across the three)

9,000 iterations × global batch 512 · bfloat16 · FusedAdam, lr 5e-5, weight decay 0.05, warm-up 500, linear decay · rectified-flow objective on video and action tokens with a shared noise level · full activation checkpointing · EMA off · 4 × H100.

Results

Read the data budgets before comparing the first two. uniwam_thirds_100pct trains on all 375 demonstrations; wam_k0_25pct trains on a task-stratified 25 % subset (94 demonstrations). They are not a controlled pair — the matching 100 %-data WAM run is not part of this release. A controlled comparison at 25 % data is UniWAM 76.2 % against WAM 77.4 %.

Layout

<run>/config.yaml                                   # required to load the checkpoint
<run>/checkpoints/iter_0000XX000/model/*.distcp     # torch DCP shards

Optimizer state is not included: these are for inference and evaluation, not for resuming training. Each model/ directory is ≈9.8 GB.

Loading

Checkpoints are PyTorch Distributed Checkpoint directories for cosmos-framework. The framework loads a DCP directory only together with the training config it was produced by, so pass both:

hf download jiachenjiang/uniwam-libero-checkpoints \
  --include "uniwam_thirds_100pct/*" --local-dir ./uniwam

python action_policy_server.py \
  --checkpoint-path ./uniwam/uniwam_thirds_100pct/checkpoints/iter_000009000 \
  --config-file    ./uniwam/uniwam_thirds_100pct/config.yaml

At inference the prefix length is a serving choice: give the model no actions and it samples video and actions jointly; give it all 16 and it predicts only the video that those actions produce.