Sword: Forging Style-Robust World Models for Precise VLA Policy Training
Sword: Style-Robust World Models as Simulators via Dynamic Latent Bootstrapping for VLA Policy Post-Training
Sword is a robust World Model framework designed as a generative simulator for Vision-Language-Action (VLA) policy post-training. It introduces Structure-Guided Style Augmentation and Dynamic Latent Bootstrapping (DLB) to achieve SOTA performance on the LIBERO benchmark, significantly enhancing generalization and physical fidelity.
TL;DR
Sword (Style-Robust World Models) addresses the "hallucination" problem in generative robotic simulators. By combining Structure-Guided Style Augmentation to prevent visual overfitting and Dynamic Latent Bootstrapping (DLB) to eliminate the mismatch between training and inference, Sword provides a high-fidelity environment for VLA (Vision-Language-Action) models to practice in "imagination," leading to significantly higher real-world success rates.
Background: The Simulator Bottleneck
In the quest for General Purpose Robots, VLA models are traditionally trained via Imitation Learning. However, to surpass human-level performance, Reinforcement Learning (RL) is necessary. Since physical interaction is expensive, World Models act as generative simulators—"imagining" the results of a robot's actions.
The problem? Current World Models (like the baseline WoVR) are brittle. Change the lighting slightly or change the table color, and the model's physics collapse. Furthermore, due to Teacher Forcing during training, these models never learn how to recover from their own small errors during autonomous rollouts, leading to catastrophic visual "drift."
Methodology: The Two Pillars of Robustness
1. Structure-Guided Style Augmentation
To stop the model from simply memorizing the pixels of the LIBERO environment, Sword uses a style transfer pipeline based on Cosmos-Transfer 2.5. It randomizes textures, colors, and lighting.
To ensure the "physics" stay intact during these shifts, the authors use Geometric and Task Priors. By feeding depth maps and semantic segmentation masks back into the generation process, the model learns that while the table color might change, the distance to the handle and the shape of the gripper are invariant truths.
2. Dynamic Latent Bootstrapping (DLB)
Standard models are trained with "Teacher Forcing"—they always see perfect historical frames. During inference, they see their own imperfect previous predictions, creating a distribution shift known as Exposure Bias.
Sword introduces a Dynamic Latent Cache. Instead of storing raw pixels (which would take hundreds of GBs), it stores compressed VAE latents (reducing storage by 60x). During training, the model occasionally "reads" from its own previous predictions in the cache rather than the ground truth, forcing it to learn how to correct its own errors.
Experimental Performance: Beyond Realism to Reliability
Generalization under Shift
When tested on "LIBERO-Mixed" (out-of-distribution visual styles), the baseline WoVR completely loses structural integrity by frame 30. Sword, however, maintains sharp, physically consistent frames even under unseen illumination settings.
Quantitative Edge
The metrics tell a clear story:
- FVD (Video Quality): Sword achieves 35.61, nearly half of WoVR's 61.26 on the standard Libero set.
- Policy Success: When used to train a VLA model (OpenVLA-OFT) using the GRPO algorithm, Sword results in a 61.72% success rate, a massive leap over the 53.12% achieved by the previous baseline.
Deep Insight: Why It Matters
The core achievement of Sword isn't just "better-looking video." It is Action Controllability. In long-horizon tasks, if a gripper starts to blur or "hallucinate" its closing state, the RL agent receives false feedback. Sword's ability to maintain a clear gripper state and follow complex corrective actions (like re-attempting a failed grasp) makes it a viable replacement for traditional physics engines in the VLA era.
Limitations & Future Work
While Sword is a significant step forward, it still requires heavy computational lifting (~1.5 A100 GPU years for full evaluation). Future iterations may look to optimize the diffusion backbone for real-time inference to tighten the loop between imagination and execution.
Conclusion
Sword proves that for World Models to be useful simulators, they must look past the "style" of the world and master its "structure." By training on its own errors via DLB, Sword bridges the gap between static datasets and the dynamic, messy reality of robotic manipulation.