Source-linked AI summary
AdaWorld: Learning Adaptable World Models with Latent Actions
Shenyuan Gao, Siyuan Zhou, Yilun Du, Jun Zhang, Chuang Gan
TL;DR
Existing world models depend on substantial action-labeled data and costly training, limiting adaptation to novel environments with heterogeneous actions. AdaWorld extracts latent actions from videos for action-aware pretraining and conditions an autoregressive world model on them. Across multiple environments, it demonstrates superior adaptability and promising simulation and visual-planning results, while remaining limited by non-real-time inference and long-term rollout challenges.
Problem
Existing world models require substantial action labels and costly training to acquire precise action controllability and adapt to new environments with varying action specifications.
Method
AdaWorld extracts compact latent actions from consecutive frames in unlabeled videos and uses them as conditions for autoregressive world-model pretraining.
Results
Experiments across multiple environments report promising results in action transfer, world-model adaptation, and visual planning.
Takeaways & Limitations
AdaWorld supports efficient transfer and learning of new actions with limited interactions and finetuning across various environments.
Takeaways & Limitations
AdaWorld does not operate at real-time frequency and struggles with novel content beyond the initial scene and extremely long-term rollouts.
Abstract
from arXiv · showhide
World models aim to learn action-controlled future prediction and have proven essential for the development of intelligent agents. However, most existing world models rely heavily on substantial action-labeled data and costly training, making it challenging to adapt to novel environments with heterogeneous actions through limited interactions. This limitation can hinder their applicability across broader domains. To overcome this limitation, we propose AdaWorld, an innovative world model learning approach that enables efficient adaptation. The key idea is to incorporate action information during the pretraining of world models. This is achieved by extracting latent actions from videos in a self-supervised manner, capturing the most critical transitions between frames. We then develop an autoregressive world model that conditions on these latent actions. This learning paradigm enables highly adaptable world models, facilitating efficient transfer and learning of new actions even with limited interactions and finetuning. Our comprehensive experiments across multiple environments demonstrate that AdaWorld achieves superior performance in both simulation quality and visual planning.
1. Introduction
AdaWorld addresses the costly, action-label-dependent adaptation of world models by pretraining with transferable latent actions extracted from videos. Its action-aware pretraining supports cross-context transfer and efficient specialization, with experiments reporting promising action transfer, adaptation, and visual planning.
- Existing world models require substantial action labels and costly training to control actions across environments with varying action specifications.
- A single action demonstration can support transfer to different contexts without further training, while raw-action models can be specialized through minor interactions and finetuning.
- AdaWorld pretrains an autoregressive world model with latent actions extracted from unlabeled videos, using an information-bottleneck autoencoder to represent critical frame transitions.
- AdaWorld is pretrained on a large-scale corpus spanning thousands of diverse environments and demonstrates generalization across domains.
- Experiments across multiple environments report promising results in action transfer, world model adaptation, and visual planning.
2. Method
AdaWorld extracts compact latent actions from unlabeled videos and uses them to condition an autoregressive world model. This action-aware pretraining supports transfer, composition, and adaptation across environments.
- Latent action autoencoder: An information bottleneck compresses critical transitions between consecutive frames into continuous latent actions.The compact representation is designed to capture variations relative to the former frame while reducing contextual information.
- Latent action autoencoder: The latent action autoencoder uses a β-VAE formulation to balance latent-action expressiveness against context disentangling.The adjustable β hyperparameter controls the information contained in latent actions and addresses difficulty expressing diverse transitions.
- Highly adaptable world models: Latent actions can be extracted from demonstrations and reused across contexts to transfer actions without additional training.The model applies the extracted latent-action sequence to an initial frame from a new context.
- Core architecture: AdaWorld uses a latent action autoencoder and an autoregressive world model conditioned on extracted latent actions.The encoder extracts actions from unlabeled videos, while the world model uses them as unified conditioning information.
- Action-aware pretraining: The world model predicts the next frame from historical frames and a latent action, then generates future sequences through autoregressive rollout.A short-term memory of K historical frames supports frame-level control and repeated next-frame prediction.
- Highly adaptable world models: Averaging latent actions enables action composition, producing new actions that merge the functions of the originals.The result supports action creation and suggests semantic continuity in the latent action space.
3. Experiments
AdaWorld is evaluated for action transfer, adaptation, simulation quality, and visual planning across unseen environments and tasks. Across these experiments, action-aware pretraining supports faster adaptation, stronger simulation fidelity, and more effective planning with limited interactions and finetuning.
- Action Transfer: AdaWorld transfers demonstrated action sequences across different contexts in qualitative and quantitative evaluations.The qualitative evaluation generates action sequences of length 20, while the quantitative set contains 1300 paired videos from LIBERO and Something-Something v2.
- World Model Adaptation: AdaWorld achieves the best action-controlled simulation fidelity after limited finetuning across four unseen environments.The evaluation covers Habitat, Minecraft, DMLab, and nuScenes, using 100 samples per discrete action or 100 continuous trajectory samples and testing after 800 finetuning steps.
- World Model Adaptation: AdaWorld adapts faster than conventional pretraining methods as sample numbers and finetuning steps vary.PSNR curves show stronger early performance and significantly faster improvement after a few finetuning steps, indicating a superior initialization for adaptation.
- Visual Planning in Games: AdaWorld substantially increases visual-planning success rates across video-game environments, whereas the action-agnostic baseline performs similarly to random planning.Planning uses sampling-based model predictive control with the Cross-Entropy Method on Procgen scenes.
- Visual Planning in Games: Without finetuning, AdaWorld still outperforms the finetuned action-agnostic baseline using averaged latent actions as scene-specific action embeddings.AdaWorld also significantly outperforms a Q-learning comparison using the same collected samples.
- Visual Planning in Robot Tasks: In robot-task planning, AdaWorld adapts more efficiently with limited finetuning steps and improves planning performance by a clear margin.Table 4 reports success rates on VP2 while omitting two RoboDesk tasks that yield meaningless scores under constrained adaptation.
- Method Generality: Applying action-aware pretraining to iVideoGPT significantly improves its adaptability, indicating that the approach generalizes across world-model architectures.The comparison evaluates action-controlled simulation quality on the BAIR robot-pushing dataset after finetuning.
4. Conclusion
AdaWorld facilitates adaptation across environments by transferring and learning new actions with limited interactions and finetuning. The paper reports superior adaptability while identifying real-time operation, novel-content generation, and extremely long-term rollouts as remaining challenges.
- AdaWorld facilitates efficient adaptation across various environments, including transferring and learning new actions with limited interactions and finetuning.
- AdaWorld does not operate at real-time frequency, limiting inference speed in its current form.
- Rollouts can struggle to create novel content beyond the initial scene and achieve extremely long-term predictions.
- The training data combines publicly accessible datasets and generated gaming transitions to support diverse action patterns and scenarios.
A.2. Data Mixture
The training mixture balances datasets by weighting subsets according to video counts, while also incorporating diverse generated environments and biased action sampling for longer-horizon exploration.
- Training subsets are weighted according to their number of videos because dataset size and diversity make perfect balancing challenging.Detailed training-data statistics are reported in Table 7.
- The data organization records sources, generation procedures, approximate frame counts, and mixture ratios.
- Biased action sampling enables agents to explore longer horizons than uniform action sampling.
- The collection includes 1000 Gym Retro environments generated from imported ROMs.
B. Implementation Details
AdaWorld uses a 500M-parameter latent action autoencoder and a 1.5B-parameter autoregressive world model, both operating at 256 × 256 resolution.
- The latent action autoencoder is a 500M-parameter Transformer with 16 encoder blocks, 16 decoder blocks, 1024 channels, and 16 attention heads.Its latent actions have dimension 32.
- The autoregressive world model is a 1.5B-parameter 3D UNet with memory length 6.The architecture follows SVD.
- Both models use a default input resolution of 256 × 256.
B.2. Training
Training uses regularized latent-action learning and large-scale world-model pretraining, with frame augmentation and sampling settings intended to improve generalization and generation quality. Visual planning then uses sampling-based model predictive control with iterative distribution updates toward high-reward action sequences.
- Training: The latent action autoencoder is trained for 200K steps with batch size 960, learning rate 2.5 × 10^-5, weight decay 0.01, and β = 2 × 10^-4.
- Training: The autoregressive world model is trained for 80K steps with batch size 64 and learning rate 5 × 10^-5 on 16 NVIDIA A100 GPUs.Training uses cosine scheduling with 10K warmup steps and exponential moving average.
- Training: Brightness jitter, center cropping, and approximately 10 Hz video downsampling are applied during training.
- Training: New frames are generated with 5 sampling steps and classifier-free guidance scale 1.05, while historical frames remain unnoised.The world model uses augmentation level 0.1 and timestep shifting.
- Visual Planning: Visual planning defines reward as RGB-space cosine similarity between the current observation and final-state image, maximizing reward along the trajectory.
- Visual Planning: Cross-Entropy Method planning samples action sequences, predicts trajectories with the world model, retains the highest-reward sequences, and updates their action distribution.
- Visual Planning: After optimization, the first T actions from the highest-probability sequence are executed until the goal is reached or the search limit is exceeded.
- Visual Planning: Planning uses i = 2 iterations, N = 100 sequences of length L = 15, K = 10 selected samples, and executes T = 5 actions.The search limit is 20 steps, with 3 denoising steps and classifier-free guidance disabled during planning.
B.5. Visual Planning on the VP2 Benchmark
The VP2 evaluation uses low-resolution models and standardized adaptation settings across Robosuite and RoboDesk. Success in Robosuite pushing is defined by a cost below 0.05.
- VP2 control-centric evaluation uses models trained at 64 × 64 resolution.
- Adaptation protocol: Robosuite adaptation uses 5K trajectories, while RoboDesk uses 35K scripted trajectories with perturbations.
- Adaptation protocol: Both models are finetuned for 1K steps with batch size 32 and learning rate 5 × 10−4.
- Success criterion: Robosuite tabletop pushing considers costs below 0.05 successful.
B.6. iVideoGPT Training Details
The iVideoGPT experiment compares action-aware and action-agnostic pretraining before robot-action finetuning on BAIR. Both models predict 15 future frames and are evaluated on 256 test videos.
- The experiment resumes from the official OpenX iVideoGPT checkpoint without finetuning its tokenizer.
- Each model predicts 15 future frames from an initial frame.
- Training protocol: After 27K extra OpenX pretraining steps, each model is finetuned for 1K steps with robot actions on BAIR robot pushing.
- Evaluation: The resulting models are tested on 256 BAIR test videos.
C. Additional Results
Additional results show AdaWorld transferring nuanced actions across contexts, acquiring precise controllability after adaptation, and supporting customizable control options. The paper also documents failures in physics understanding and long-term or large-view-shift rollouts.
- Action transfer: AdaWorld represents nuanced actions and transfers them strongly across contexts in qualitative comparisons.
- Action transfer: Action transfer visualizations include multiple environments, with each sample transferring a latent action sequence of length 20.
- Failure cases: AdaWorld does not always understand physics and dynamics, and struggles with high-quality content during long-term rollouts or dramatic view shifts.
- World model adaptation: AdaWorld acquires precise action controllability rapidly, whereas action-agnostic pretraining struggles with the same action inputs and adaptation steps.
- Customizable control: Latent action clustering creates a flexible number of control options by applying K-means to encoded Procgen and Gym Retro actions.
- Customizable control: Table 9 reports strong controllability when AdaWorld customizes different numbers of actions.
D. Something-Something v2 Categories for Action Transfer
The action-transfer evaluation uses the ten most frequent Something-Something v2 action categories. These categories cover object placement, movement, covering, taking, throwing, and related interactions.
- The evaluation selects the top-10 most frequently appearing categories from Something-Something v2.
- Included categories: The categories include placing, moving, pushing, covering, uncovering, taking, throwing, and putting objects into containers.
- Included categories: The list distinguishes directional pushing actions, including left-to-right and right-to-left movement.
E. Selected Scenes for Visual Planning
Visual planning is evaluated on 120 Procgen scenes selected for feasible completion through goal-image matching. Additional figures assess action transfer, Minecraft adaptation, planning scenes, and failure cases, including limitations in physics, dynamic agents, long rollouts, and viewpoint changes.
- Scene Selection: 120 Procgen scenes were selected by exhaustive ground-truth search for feasible completion within an acceptable number of steps using goal-image matching.Survival-oriented environments such as BigFish and BossFight were excluded from the evaluation.
- Action Transfer: AdaWorld accurately identifies demonstrated actions and transfers them to new contexts, whereas other baselines fall short.Additional figures report further action-transfer examples by AdaWorld.
- Action Transfer: After 800 finetuning steps with 100 samples per action, AdaWorld achieves precise action controllability, while action-agnostic pretraining fails to perform the correct actions.The result concerns Minecraft adaptation with minimal action-labeled data and finetuning.
- Failure Cases: AdaWorld remains imperfect at simulating real-world physics, dynamic agents, long-term rollouts, and substantial viewpoint changes.These are presented as examples of failure cases rather than quantified limitations.