DiWINE-VLA

Dynamics-integrated World-model Informed Embodied Vision-Language-Action learning

DiWINE-VLA augments X-VLA with a stochastic latent dynamics model. Instead of using predicted future states only as an auxiliary training target, DiWINE-VLA converts its future belief into tokens and injects those tokens back into the vision-language encoder before action generation. The action policy therefore conditions on both the present observation and a learned distribution over plausible futures.

This repository is a research artifact and experiment tracker. Training and matched LIBERO evaluation are in progress; validated checkpoints and task-success numbers will be added only after evaluation. Source code lives at shahils01/DiWINE-VLA.

Core idea

flowchart LR
    A[Multi-view images] --> B[Florence-2 early layers]
    L[Language instruction] --> B
    P[Robot proprioception] --> C[State compressor]
    B --> C
    C --> D[RSSM stochastic rollout]
    D --> E[Mean + uncertainty projection]
    E --> F[Predicted future tokens]
    F --> G[Florence-2 remaining layers]
    B --> G
    G --> H[Flow-matching action transformer]
    H --> I[Continuous action chunks]
    J[Future observations<br/>training only] -. ELBO target .-> D

At training time, actual future observations are encoded without gradients to provide targets for the RSSM evidence lower bound. They are not passed to the action head. Both training and inference feed the RSSM's predicted future tokens into the remaining vision-language layers, preventing future-observation leakage while making the policy explicitly future-conditioned.

What is new

1. Future-token reinjection. The predicted latent future is returned to the multimodal token stream rather than being discarded after an auxiliary loss.
2. Stochastic predictive state. An RSSM represents multiple plausible futures and is trained with reconstruction likelihood and KL regularization.
3. Uncertainty-aware tokenization. Multiple latent samples are summarized through their mean and standard deviation before projection into future tokens.
4. Long-horizon conditioning. The recovered experiment predicts 30 future steps with four tokens per step and eight stochastic samples.
5. Optional GT-MHA action transformer. A Generative Transformer Multi-Head Attention variant tests whether task performance can be retained with a substantially smaller attention stack.

Architecture

DiWINE-VLA preserves X-VLA's Florence-2 multimodal backbone and flow-matching continuous-action decoder. The future module branches after the early encoder layers, combines compressed multimodal state with proprioception, rolls an RSSM forward, and inserts the resulting future tokens before the remaining encoder layers. A 24-layer action transformer then predicts action velocity fields.

The GT-MHA variant changes only self-attention in the action transformer. All data, future modeling, optimization, and evaluation settings are held fixed for a controlled comparison.

DiWINE-VLA
├── multimodal backbone (Florence-2)
│   ├── early encoder layers
│   ├── future-token insertion
│   └── remaining encoder layers
├── future world model
│   ├── current-state compressor
│   ├── RSSM stochastic rollout
│   └── uncertainty-aware token projector
└── flow-matching action head (24 layers)
    ├── MHA baseline
    └── residual GT-MHA (optional)

Controlled efficiency hypothesis

Model	Total parameters	Action head	Attention modules
MHA baseline	879,482,456	308,775,512	100,761,600
Residual GT-MHA	824,383,064	253,676,120	45,662,208
Reduction	55,099,392 (6.27%)	55,099,392 (17.84%)	55,099,392 (54.68%)

These are implementation-level parameter counts, not performance claims. The current experiment tests whether the reduction preserves action loss, throughput, memory use, and LIBERO task success.

Recovered training recipe

The configuration was recovered from the original Palmetto training launcher and is reproduced in configs/recovered_libero_mha.yaml. The GT-MHA run changes only the attention fields in configs/recovered_libero_gt_mha.yaml.

Setting	Value
Initialization	2toINF/X-VLA-Libero
Dataset	LIBERO-XVLA-format
GPUs	2 × A100 80 GB
Precision	BF16 DDP
Batch size	12/GPU, 24 global
Learning rate	1e-4
Future horizon	30 steps
Future tokens	4 per step
Future samples	8
RSSM latent / hidden	1024 / 1024
Checkpoint interval	5,000 steps

Evaluation plan

The matched study reports:

• action, future NLL, and future KL training curves;
• parameter count, peak GPU memory, and throughput;
• identical-checkpoint LIBERO Spatial, Object, Goal, and Long success rates;
• qualitative successes and a categorized failure analysis;
• at least one matched pilot before any claim about GT-MHA quality.

Current status

• Corrected LIBERO assets and metadata: validated locally on Palmetto.
• MHA recovered-configuration training: submitted.
• GT-MHA pretrained-weight conversion and finite-loss guards: unit tested.
• Matched GT-MHA training: submitted behind a successful smoke-test dependency.
• LIBERO checkpoint evaluation: queued.
• Public checkpoint weights and final metrics: pending validation.

Acknowledgements

DiWINE-VLA builds on X-VLA, the LIBERO benchmark, and the GT-MHA research implementation in Attention is Transformed.

Citation

@software{shaik2026diwinevla,
  author = {Shaik, Shahil},
  title = {DiWINE-VLA: Dynamics-integrated World-model Informed Embodied Vision-Language-Action Learning},
  year = {2026},
  url = {https://github.com/shahils01/DiWINE-VLA}
}