LaWAM β€” closed-loop staircase (variant A), LIBERO best checkpoint

Best-scoring LIBERO checkpoint from the LaWAM closed-loop "staircase" line of work: 0.973 mean success rate across the four LIBERO suites.

  • Run id: cl_A_gran10_bs32x2_fixall_bsz64rerun
  • Run dir: results/Checkpoints/libero/20260731_195651+cl_A_gran10_bs32x2_fixall_bsz64rerun
  • Step: 25 000 (final)
  • Seed: 2026

1. What the model is

Three parts, trained jointly, all inside the single .pt:

Component Params dtype Role
policy_backend.vlm β€” Qwen3-VL-2B 1.52 B bf16 language + image understanding, emits h_vlm and latent-action tokens z
policy_backend.lam β€” DINOv3 encoder + VQ + decoder 0.72 B fp32 latent world model: encodes an observation to h, decodes (h, z) to a sub-goal feature
policy_action_head.DiT β€” flow-matching action head 0.31 B fp32 16-layer VL-DiT, denoises a 50-step action chunk conditioned on h/z/h_vlm

Total 4.39 B parameters (policy_vlm_adapter and policy_action_head alias the same storages as policy_backend, so the on-disk file is 6.7 GB, not the naive sum).

The closed-loop staircase

A 50-action chunk is emitted as 5 groups of 10, and the K=5 flow-matching denoising steps are spread across execution rather than all run up front:

group j     0        1        2        3        4
depth       1        2        3        4        5
condition   o_0      o_10     o_20     o_30     o_40      <- geometric refresh at each boundary
executes    a0..9    a10..19  a20..29  a30..39  a40..49

At every boundary only the geometry is refreshed (a DINOv3 pass over the new observation, ~3 ms); z and h_vlm stay frozen at the draft frame. Training supervision equals inference execution: the strict executed loss only scores each group at its own execution depth under its own execution-time observation.

Trained with a shared tau (token_independent_noise=false) and a fully chained rollout (rollout_align_prob=1.0), plus all three fixes on:

Flag Value Effect
rollout_steer_grad true flow loss gradient flows into the LAM decoder and z
rollout_one_window_aux true perceptual/distill loss on a single boundary window (frame 0 β†’ frame 1)
rollout_nonexec_loss_weight 0.25 dense supervision at non-executed positions under strict loss

2. Results

LIBERO, variant A (num_steps=5, refresh=geo), 10 tasks Γ— 10 trials per suite (400 episodes per row).

step libero_10 goal object spatial mean
5 000 0.680 0.850 0.980 0.890 0.850
10 000 0.860 0.900 0.980 0.880 0.905
15 000 0.880 0.980 0.990 0.910 0.940
20 000 0.870 0.990 1.000 0.980 0.960
25 000 (this ckpt) 0.940 0.990 0.980 0.980 0.973

Read these numbers with the error bars. They are honest measurements, not tuned headline figures:

  • TRIALS=10 runs about 2 points optimistic. The same checkpoint measured at TRIALS=50 scores 0.953.
  • Seed variance at identical config spans ~3.5 points (0.973 / 0.938 for two seeds).
  • Per-suite noise floor: Β±5 points on libero_10, Β±2 on the others. Do not draw conclusions from differences under 3 points.

Ablations

Setting mean SR
A β€” geometric refresh at every boundary 0.973
open β€” no refresh at all (open loop) 0.105
A βˆ’ open 0.868

The model is extremely dependent on condition freshness. That is the intended property (refresh is what the staircase buys), but it also means this checkpoint degrades sharply if you feed it stale conditions β€” see Β§5.

Latency

Serving mode per-action policy block mean SR
chunk10 baseline 19.2 ms 0.693
closed-loop A, synchronous 6.2 ms 0.973
overlap, zero-staleness + full graph stack 1.52 ms (11.8Γ—) 0.950

All end-to-end measured, no arithmetic extrapolation.


3. Files

checkpoints/steps_25000_pytorch_model.pt   6.7 GB β€” the model (self-contained state dict)
config.json                                resolved training config for the run
config.yaml                                same, YAML form
dataset_statistics.json                    action/state normalization statistics β€” REQUIRED at inference
summary.jsonl                              which steps were checkpointed
logs/train.log                             full training log
logs/train_libero_spec50.yaml              dataset/spec config used
train_scripts/                             the exact launch script + shared helpers
wandb/                                     offline wandb run

dataset_statistics.json is not optional β€” the action head predicts normalized actions and this file is how they get mapped back.

Only step 25 000 is published here. Steps 5 000–20 000 exist in the original run directory if intermediate checkpoints are needed.


4. Usage

The weights alone are not runnable β€” they need the LaWAM code and two config-only assets.

git clone https://github.com/s1ghhh/lawam_s.git
cd lawam_s
git checkout closedloop-staircase

You additionally need:

Asset Why
Qwen3-VL base dir (results/Checkpoints/qwen3_weights) tokenizer / processor / model config. The weights are already in the .pt, but the loader still reads the config dir.
LAM yaml (latent_action_model/logs/dino_large_vae/lam_release/dino_large_vae.yaml) LAM architecture spec. Again weights come from the .pt.
LIBERO dataset in LeRobot v3.0 form only if you want to reproduce eval or resume training

Evaluate:

# variant must match the checkpoint β€” this is an A checkpoint
bash closedloop/eval_closedloop.sh <path>/steps_25000_pytorch_model.pt A
# with ablations
bash closedloop/eval_closedloop.sh <path>/steps_25000_pytorch_model.pt A open chunk10

Reproduce training (2 GPUs, global batch 64, 25 000 steps):

GPUS=0,1 PER_DEVICE_BS=32 GRAD_ACCUM=1 STEPS=25000 \
  bash closedloop/train_A_gran10_bs32_fixall_bsz64.sh

Key hyperparameters: lr 1e-4 (base and action model), 1500 warmup steps, action_horizon=50, num_frames=6, sec_chunk=2.5, horizon_sec=2.5, num_inference_steps=10, repeated_diffusion_steps=2, perceptual_weight=0.1, lam_encoder_distill_weight=0.1.

The learning rate is on a sharp peak. 5e-5 β†’ 0.925 and 5e-4 β†’ 0.920, i.e. about 5 points lost in either direction from 1e-4. Likely because rollout_steer_grad couples the world model and the action head to one learning rate.


5. Known caveats

  1. Trained with an off-by-one frame grid. This run predates the frame-grid fix. The K+1 boundary frames were laid over [0, N-1] instead of [0, N], so the sampled video deltas were [0, 10, 20, 29, 39, 49] where the execution grid the evaluator actually uses is [0, 10, 20, 30, 40, 50]. Boundaries 3 and 4 were therefore trained one frame early. The 0.973 was measured with that mismatch present in the weights and absent from the evaluator. A retrained, grid-corrected run scores 0.9475 β€” inside the same-config seed band, so the fix has not been shown to change quality either way. If you are building on this line of work, prefer the corrected training code; if you are reproducing 0.973 exactly, you need the pre-fix dataloader.
  2. Do not enable vlm_prefetch at inference β€” measured βˆ’25 points.
  3. Off-distribution in chunk mode. chunk10 evaluation of this checkpoint is a coarse reference only; the model was never trained for truncated chunk execution.
  4. Stale conditions are fatal, not merely costly (open = 0.105). Any async or pipelined serving scheme must keep the boundary condition fresh, or the model must first be fine-tuned on stale conditions.
  5. Training used teacher forcing: the "new observation" at each boundary is the demo frame, which assumes model actions β‰ˆ demo actions. On-policy residual is not addressed by this recipe.
  6. LIBERO-only. This checkpoint scores 0.606 on LIBERO-Plus, 15–20 points below models trained on Plus.

6. Citation

Part of the LaWAM closed-loop streaming work (Sony proposal line, 2026). Code: https://github.com/s1ghhh/lawam_s, branch closedloop-staircase.

Downloads last month
-
Video Preview
loading