WHALE: Towards Generalizable and Scalable World Models for Embodied Decision-making

Abstract: World models play a crucial role in decision-making within embodied environments, enabling cost-free explorations that would otherwise be expensive in the real world. To facilitate effective decision-making, world models must be equipped with strong generalizability to support faithful imagination in out-of-distribution (OOD) regions and provide reliable uncertainty estimation to assess the credibility of the simulated experiences, both of which present significant challenges for prior scalable approaches. This paper introduces WHALE, a framework for learning generalizable world models, consisting of two key techniques: behavior-conditioning and retracing-rollout. Behavior-conditioning addresses the policy distribution shift, one of the primary sources of the world model generalization error, while retracing-rollout enables efficient uncertainty estimation without the necessity of model ensembles. These techniques are universal and can be combined with any neural network architecture for world model learning. Incorporating these two techniques, we present Whale-ST, a scalable spatial-temporal transformer-based world model with enhanced generalizability. We demonstrate the superiority of Whale-ST in simulation tasks by evaluating both value estimation accuracy and video generation fidelity. Additionally, we examine the effectiveness of our uncertainty estimation technique, which enhances model-based policy optimization in fully offline scenarios. Furthermore, we propose Whale-X, a 414M parameter world model trained on 970K trajectories from Open X-Embodiment datasets. We show that Whale-X exhibits promising scalability and strong generalizability in real-world manipulation scenarios using minimal demonstrations.

• The paper presents the WHALE framework that uses behavior-conditioning to mitigate policy distribution shift, enhancing world model generalization.
• It introduces retracing-rollout to efficiently estimate uncertainty without expensive ensembles, improving simulated rollout fidelity.
• Experiments on Meta-World and ARX5 platforms demonstrate significant gains in video fidelity, value estimation accuracy, and out-of-distribution performance.

Scalable World Models in Embodied Decision-Making: An Insight into "WHALE"

The paper "WHALE: Towards Generalizable and Scalable World Models for Embodied Decision-making" (2411.05619) explores the development of scalable world models crucial for decision-making processes in embodied environments. The paper introduces a framework, WHALE, which aims to address the challenges of generalization and uncertainty estimation in world models—key hurdles in enhancing the fidelity and reliability of imagined experiences generated by such models.

Framework and Methodology

Key Techniques

WHALE is primarily defined by two pivotal techniques: behavior-conditioning and retracing-rollout.

1. Behavior-conditioning: This technique mitigates the policy distribution shift, a primary contributor to generalization error in world models. It enables the model to adaptively recognize and adjust to policy-induced distribution shifts by embedding behavioral patterns within the latent space.
2. Retracing-rollout: This is a novel approach to uncertainty estimation. It circumvents the need for computationally expensive model ensembles by using retracing actions that leverage the action space's semantic structure, significantly enhancing uncertainty quantification in simulated rollouts.

Implementation: Whale-ST and Whale-X

• Whale-ST: Built upon a spatial-temporal transformer (ST-transformer) architecture, this model ensures scalable and efficient training by reducing computational demands. Whale-ST is evaluated for its capability in maintaining high fidelity in video generation and accurate value estimation in simulated environments like Meta-World.

  Figure 1: Overall architecture of Whale-ST. The behavior-conditioning model encodes observation and action subsequences into behavior embedding $z_i$, then passed to the dynamics model with observation tokens to predict the next tokens $\hat x_{i+1}.$

• Whale-X: With 414M parameters, Whale-X is pre-trained on the Open X-Embodiment dataset, demonstrating substantial potential in scaling and generalizability. This model includes a more comprehensive architecture enabling finer adaptability and efficiency across diverse real-world tasks, showcasing notable OOD generalization ability with minimal demonstration data.

Experimental Analysis

Simulated and Real-World Evaluation

• Simulated Tasks: Extensive tests on the Meta-World benchmark reveal Whale-ST's superior performance in both video fidelity and value estimation metrics compared to baseline models like FitVid and DreamerV3. The model shows about 18% improvement in reducing value estimation errors, underscoring its enhanced generalization capabilities when behavior-conditioning is included.
• Real-World Tasks: Deployments on the ARX5 robotic platform for unseen tasks demonstrate Whale-X's viability in OOD generalization. Pre-trained on internet-scale data, it shows significant improvement in consistency rates, evidenced by a 63% increase over models lacking behavior-conditioning.

  Figure 2: Physical robot evaluation on unseen scenarios showing the consistency rate and the tasks used for testing. Whale-X exhibits strong generalization in unseen scenarios.

Scalability and Efficiency

The scalability of Whale-X is validated through experiments demonstrating the benefits of increasing both model capacity and pre-training data size. Larger models and data sets consistently lead to reduced training loss, marking a log-linear relationship that suggests clear guidance for designing larger scaled models.

Figure 3: Scaling experiment results of Whale-X. Shows significant relationship between training loss and parameters indicating effective scaling.

Implications and Future Directions

The WHALE framework, supported by its implementation in Whale-ST and Whale-X, provides significant advancements towards scalable and generalizable world models. The introduction of techniques like behavior-conditioning and retracing-rollout sets a foundation for future explorations in adaptive and robust decision-making systems in AI. Future work should aim to incorporate diverse data sources to further enhance generalization capabilities and address challenges in real-time uncertainty estimation. The integration of pre-existing contextual and domain-specific knowledge into data-driven models may open new avenues for research, promoting broader generalization and applicability across varied decision-making contexts.

In conclusion, WHALE represents a critical step in advancing the scalability and robustness of world models, potentially transforming their role in embodied AI systems navigating complex real-world environments.