rsj2003/StreamingWAM-DiT4DiT-LIBERO

StreamingWAM on DiT4DiT — LIBERO

Two checkpoints from a matched additional-training comparison on DiT4DiT, using the Cosmos-Predict2.5-2B video backbone.

Code and runnable commands · Project · Paper

Checkpoint	Additional updates	Global batch	H / K	LIBERO-Plus raw-prompt success
joint/pytorch_model.pt	10,000	64	8 / 4	70.88% (7,109 / 10,030)
streaming/pytorch_model.pt	10,000	64	8 / 4	76.37% (7,660 / 10,030)

Both start independently from the same official DiT4DiT-LIBERO initialization. The unmodified official weights are not duplicated in this repository. These are the 10K points of 20K cosine schedules; they are not checkpoints trained with a 10K schedule. Equal update and batch counts do not imply equal FLOPs.

What is included

Each directory contains a full inference state dict (about 21.3 GB), portable training and inference configurations, dataset normalization statistics and training metadata. No optimizer states or training logs are included. release.json records exact SHA256 identities.

The streaming variant conditions on five real visual observations and eight executed actions, refines an eight-action future buffer, and executes four actions per replan. It uses a four-step warm start followed by the first staircase update, then one action update for each steady replan. The ordinary Joint variant retains one current image and four action updates per replan. Both use paired primary and wrist camera views.

Download and run

git clone https://github.com/renshaojie233/StreamingWAM.git
cd StreamingWAM/extensions/dit4dit
# Follow README.md to create the separate DiT4DiT environment.
python scripts/download.py --method both --output checkpoints/dit4dit

hf download nvidia/Cosmos-Predict2.5-2B \
  --revision diffusers/base/post-trained --local-dir checkpoints/cosmos

CUDA_VISIBLE_DEVICES=0 python scripts/smoke_inference.py \
  --method streaming --checkpoint-dir checkpoints/dit4dit/streaming \
  --cosmos-path checkpoints/cosmos

Base-model access may require accepting NVIDIA's terms using your own Hugging Face account. The code README also provides ordinary Joint inference, LIBERO-Plus evaluation, data preparation and both continued-training commands. Use the bundled runtime and pinned Diffusers version; this is not a generic Transformers AutoModel checkpoint.

Release checks cover strict checkpoint loading, two replans, one simulator episode and one four-GPU optimizer update per variant. See validation details; the archived full-benchmark results were not rerun during packaging.

Evaluation notes

The reported comparison uses all 10,030 LIBERO-Plus entries with raw task.language, including benchmark suffixes, seed 7, initial-state index 0, 20 Hz control, five settling steps and a 700-step limit. It does not use the FastWAM prompt sanitizer. The two percentages should not be compared to clean-prompt FastWAM results as though their protocols were identical.

The streaming 10K/K4 point was selected using a 200-entry screen contained in the full manifest, so the full evaluation is not entirely held out. The continuation baseline is evaluated at the matched point. This tests the combined adaptation, not an isolated attribution to history or the streaming objective. The separate full-call latency measurements were 135.526 ms for ordinary DiT4DiT and 153.685 ms for streaming; we do not claim an end-to-end speedup for this adaptation.

Training uses the same four LIBERO no-noop suites for both variants (1,693 trajectories). This release covers LIBERO simulation, not a DiT4DiT real-robot or RoboTwin model.

Attribution and licenses

Built on NVIDIA Cosmos.

DiT4DiT is by Teli Ma, Jia Zheng, Zifan Wang, Chunli Jiang, Andy Cui, Junwei Liang and Shuo Yang. See DiT4DiT and its official checkpoints.

The upstream DiT4DiT model card specifies Apache-2.0; underlying Cosmos components are subject to the NVIDIA Open Model License. See LICENSE.md, NOTICE.txt and the included agreements. Source code retains the upstream MIT and Apache-2.0 notices as applicable.