Test-Time Mixture of World Models for Embodied Agents in Dynamic Environments

TL;DR

This work introduces Test-time Mixture of World Models (TMoW), an extension of Mixture-of-Experts for LM-based embodied agents that reconfigures world-model routing at test time to adapt to dynamic environments. By leveraging a multi-granular prototype-based router, test-time prototype refinement, and distilled mixture-based augmentation, TMoW enables rapid adaptation to unseen domains and continual expansion with few-shot demonstrations, without full retraining. The approach is instantiated with an adapter-based MoE framework atop a language model, incorporating a hierarchical graph representation that spans object to scene levels and uses a layer-wise routing mechanism. Empirical results across VirtualHome, ALFWorld, RLBench, and real-world scenarios show substantial improvements in zero-shot and few-shot adaptation, robustness to visual inputs, and scalable knowledge expansion, validating the practical impact for dynamic embodied reasoning. The findings indicate that test-time reconfiguration of world-model mixtures, guided by prototypes, can significantly enhance the flexibility and longevity of embodied agents in changing environments.

Abstract

Language model (LM)-based embodied agents are increasingly deployed in real-world settings. Yet, their adaptability remains limited in dynamic environments, where constructing accurate and flexible world models is crucial for effective reasoning and decision-making. To address this challenge, we extend the Mixture-of-Experts (MoE) paradigm to embodied agents. While conventional MoE architectures modularize knowledge into expert components with pre-trained routing, they remain rigid once deployed, making them less effective for adapting to unseen domains in dynamic environments. We therefore propose Test-time Mixture of World Models (TMoW), a framework that enhances adaptability to unseen and evolving domains. TMoW updates its routing function over world models at test time, unlike conventional MoE where the function remains fixed, enabling agents to recombine existing models and integrate new ones for continual adaptation. It achieves this through (i) multi-granular prototype-based routing, which adapts mixtures across object- to scene-level similarities, (ii) test-time refinement that aligns unseen domain features with prototypes during inference, and (iii) distilled mixture-based augmentation, which efficiently constructs new models from few-shot data and existing prototypes. We evaluate TMoW on VirtualHome, ALFWorld, and RLBench benchmarks, demonstrating strong performance in both zero-shot adaptation and few-shot expansion scenarios, and showing that it enables embodied agents to operate effectively in dynamic environments.

Figures (18)

• Figure 1: Existing Mixture-of-Experts employs a fixed router after training, making adaptation to unseen domains costly and requiring retraining (left). In contrast, Test-time Mixture of World Models (TMoW) overcomes this through prototype-based routing and test-time training of the routing function to reflect new domain characteristics without demonstrations (right).
• Figure 2: Overall framework of TMoW.
• Figure 3: Graph representation in multi-granular prototype. As layers deepen, node representations evolve from local object features to global scene structures.
• Figure 4: Comparison of the routing score heatmap before (left) and after (right) test-time prototype refinement for (a) seen and (b) unseen domains.
• Figure 5: Examples of Real-world environment.