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

AI summary

Overview

• Research area: Artificial Intelligence / embodied agents — specifically Mixture-of-Experts (MoE) architectures, world models, and test-time adaptation for language-model-based embodied agents operating in dynamic environments.
• Technical level: Advanced. The paper assumes familiarity with MoE routing, parameter-efficient adapters (LoRA), message-passing neural networks, and test-time adaptation.
• Scope: The paper proposes Test-time Mixture of World Models (TMoW), a framework that lets embodied agents reconfigure their mixture of world models at test time and expand with new models from few-shot data, evaluated on VirtualHome, ALFWorld, RLBench, and a real-world robot scenario.

What This Paper Is About

LM-based embodied agents freeze their capabilities at training time, so adapting to new households, tasks, or scenes usually requires costly retraining or domain-specific prompting with inflated context windows. Conventional Mixture-of-Experts offers modularity through expert modules, but its routing function is fixed after training, so it cannot flexibly recombine or extend experts when the environment changes. TMoW addresses this by making the routing function itself trainable at test time, so an agent can recombine existing world models or build new ones for unseen domains.

Key Contributions

1. Test-time mixture framework (TMoW). A novel extension of MoE that supports test-time reconfiguration of expert mixtures, allowing embodied agents to adapt to unseen domains without costly retraining.
2. Multi-granular prototype-based routing. A routing mechanism that compares input observations with learned prototypes across spatial abstraction levels ranging from local objects to global scenes, enabling layer-wise model mixture and partial domain similarity sharing.
3. Distilled mixture-based model augmentation. A data-efficient strategy that distills knowledge from a weighted mixture of existing world models to construct a new world model from few-shot demonstrations, then integrates it directly into the router.
4. Empirical validation. Evaluation on VirtualHome, ALFWorld, RLBench, and real-world robotic scenarios across zero-shot adaptation, few-shot expansion, and continuous expansion settings.

Main Findings

• Zero-shot adaptation, unseen domains: TMoW achieves an average improvement of 27.21% in Success Rate (SR) and a reduction of 3.45 steps (14.81% improvement) in Pending Steps (PS) across environments, compared with baselines. The abstract attributes the 27.21% figure to comparison with SayCanPay.
• Zero-shot adaptation, seen domains: TMoW achieves an average improvement of 14.61% in SR and a 4.42 reduction (35.31% improvement) in PS across environments.
• Absolute unseen-domain numbers (Table 1): VirtualHome 80.16% SR / 13.20 PS, ALFWorld 68.83% SR / 37.44 PS, RLBench 62.75% SR / 8.95 PS. Seen-domain numbers: VirtualHome 83.61% SR / 11.07 PS, ALFWorld 72.05% SR / 6.94 PS, RLBench 71.89% SR / 6.30 PS.
• Few-shot expansion: In VirtualHome, TMoW gains on average 25.66% in SR and reduces PS by 4.73 steps versus baselines, reaching 82.59% average SR across 1-shot (81.56%) and 5-shot (83.61%) settings, compared with SayCanPay's 56.93% average SR.
• Real-world scenario: Using a Franka Research 3 robot arm, TMoW improves 34.36% in SR and reduces PS by 2.77 (39.57% improvement) for seen domains, and improves 38.60% in SR and reduces PS by 2.98 (37.86% improvement) for unseen domains compared with the best baseline, FLARE.
• Multi-granular routing matters: TMoW outperforms TMoW-Object (object features only) by 15.49% in SR and 3.60 in PS, and TMoW-Scene (scene features only) by 72.00% in SR and 14.26% in PS. TMoW-Scene alone reaches only 8.74% SR / 27.38 PS.
• Test-time refinement matters: TMoW-NoRefine reaches 73.30% SR / 14.85 PS versus TMoW's 80.74% SR / 13.12 PS, a 7.65% SR improvement from dynamic refinement.
• Distillation beats scratch training: TMoW reaches 81.56% SR / 13.20 PS versus TMoW-Scratch's 59.84% SR / 18.16 PS — 18.39% higher SR while using 40% less data.
• Top-K routing: K=3 is optimal (80.16% SR / 13.20 PS). Top-1 gives 65.43% SR / 17.34 PS, Top-5 gives 77.52% SR / 15.19 PS, and Top-7 drops to 66.01% SR / 17.67 PS, indicating excessive expert activation adds noise.
• Entropy analysis: Routing score entropy is high in early (object-level) layers and lower in later (scene-level) layers, and prototype refinement increases average entropy across all layers, indicating broader use of previously underutilized world models.
• Continuous expansion: Adding domains in new phases improves performance on new domains and existing ones, with no forgetting on previously encountered domains.

Methodology in Plain English

The researchers start from a pre-trained language-model base and attach one lightweight adapter per domain (for example, LoRA modules), where each adapter acts as a "world model" that captures that domain's environment dynamics and policy. Adapters are trained on instruction–trajectory pairs from their own domain, and at inference the base model is combined with a weighted mixture of adapters.

To choose the mixture, they build multi-granular prototypes: each observation is converted into a graph, and a message-passing network with instruction-aware edge weighting produces layer-by-layer embeddings that progress from local object features to global scene structure. The per-domain average of these embeddings at each layer becomes that world model's prototype. At test time, the router compares the current input's layer-wise embedding against each prototype using cosine similarity, keeps only the top-K scores, normalizes them with softmax and a temperature, and uses them to weight the adapters at each layer.

When the agent meets an unseen domain, test-time prototype refinement updates each prototype by interpolating it with a similarity-weighted combination of other prototypes, controlled by a refinement rate, without any demonstrations. When a domain is very different, distilled mixture-based model augmentation initializes a new world model as a router-weighted mixture of existing ones and fine-tunes it on a few-shot demonstration set with a teacher-forcing loss on next actions and observations; its prototype joins the router automatically. Experiments use Success Rate and Pending Steps as metrics, report 95% confidence intervals over 5 random seeds, and use Llama-3.2-3B for ZSP, LLM-Planner, FLARE, and the Say model in SayCanPay, and trainable Llama-3.2-1B for LLM+FT, the Pay model in SayCanPay, and TMoW.

Why This Matters

• Research impact: The work reframes MoE from a static, pre-trained routing architecture into something that can be reconfigured and grown after deployment, connecting MoE, world models, and test-time adaptation for embodied agents.
• Real-world applications (as reported or directly implied by the paper's settings):
  • Household robots operating in new home layouts and object arrangements.
  • Factory and industrial automation environments with changing task configurations.
  • Tabletop robotic manipulation, evaluated via RLBench and a Franka Research 3 arm.
  • Simulated or virtual game environments with temporally and spatially changing conditions.
• Industry relevance: Avoiding full retraining and large-scale data collection lowers computational and data-collection burdens for deployment, which the paper identifies as a key obstacle to real-world deployment of embodied agents. Data-efficient few-shot expansion (40% less data than scratch training in the reported ablation) is directly relevant to cost-sensitive deployments.

Future Directions

1. Safety and interpretability of routing decisions, which the authors list as planned future work.
2. Extension to multi-agent systems, since the paper notes that world-model approaches may struggle in highly non-stationary settings where other agents continuously change environment dynamics.
3. Overcoming the ceiling imposed by the underlying LLM, since the paper states performance is inherently bounded by the planning LLM's capabilities.
4. Broader validation of the modular expansion mechanism, as the paper's conclusion references evaluation on VirtualHome and ALFWorld while the abstract and results sections also report RLBench and real-world robotic experiments.

Target Audience

Researchers and practitioners working on embodied AI, language-model agents, Mixture-of-Experts architectures, world models, and test-time adaptation. It is most useful for readers already comfortable with MoE routing, parameter-efficient fine-tuning, and graph-based or message-passing representations, and for engineers evaluating modular alternatives to retraining for robot or simulation deployment.

Authors’ 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.