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

arXiv: https://arxiv.org/abs/2601.22647 · HTML: https://mlsanity.com/paper/2601.22647/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- LLM agents and tool use / [Multimodal and embodied agents](https://mlsanity.com/agents/multimodal-and-embodied-agents/)
- Mixture-of-experts / [MoE in reinforcement learning and robotics](https://mlsanity.com/moe/moe-in-reinforcement-learning-and-robotics/)

## Builds on (most cited first, 25 of 25)

- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Do As I Can, Not As I Say: Grounding Language in Robotic Affordances](https://mlsanity.com/paper/2204.01691/) (2022)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [Language Models as Zero-Shot Planners: Extracting Actionable Knowledge for Embodied Agents](https://mlsanity.com/paper/2201.07207/) (2022)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [RLBench: The Robot Learning Benchmark & Learning Environment](https://mlsanity.com/paper/1909.12271/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge](https://mlsanity.com/paper/2206.08853/) (2022)
- [LLM-Planner: Few-Shot Grounded Planning for Embodied Agents with Large Language Models](https://mlsanity.com/paper/2212.04088/) (2022)
- [VirtualHome: Simulating Household Activities via Programs](https://mlsanity.com/paper/1806.07011/) (2018)
- [Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets](https://mlsanity.com/paper/2504.02792/) (2025)
- [MixLoRA: Enhancing Large Language Models Fine-Tuning with LoRA-based Mixture of Experts](https://mlsanity.com/paper/2404.15159/) (2024)
- [Efficient Test-Time Adaptation of Vision-Language Models](https://mlsanity.com/paper/2403.18293/) (2024)
- [Bootstrap Your Own Skills: Learning to Solve New Tasks with Large Language Model Guidance](https://mlsanity.com/paper/2310.10021/) (2023)
- [How Does BERT Answer Questions? A Layer-Wise Analysis of Transformer Representations](https://mlsanity.com/paper/1909.04925/) (2019)
- [See and Think: Embodied Agent in Virtual Environment](https://mlsanity.com/paper/2311.15209/) (2023)
- [Meta-DMoE: Adapting to Domain Shift by Meta-Distillation from Mixture-of-Experts](https://mlsanity.com/paper/2210.03885/) (2022)
- [SayCanPay: Heuristic Planning with Large Language Models using Learnable Domain Knowledge](https://mlsanity.com/paper/2308.12682/) (2023)
- [Few-Shot Class-Incremental Learning via Training-Free Prototype Calibration](https://mlsanity.com/paper/2312.05229/) (2023)
- [Online Continual Learning For Interactive Instruction Following Agents](https://mlsanity.com/paper/2403.07548/) (2024)
- [Multi-Modal Grounded Planning and Efficient Replanning For Learning Embodied Agents with A Few Examples](https://mlsanity.com/paper/2412.17288/) (2024)
- [Adapted-MoE: Mixture of Experts with Test-Time Adaption for Anomaly Detection](https://mlsanity.com/paper/2409.05611/) (2024)

## Built on by (most cited first, 2 of 2)

- [Why AI systems don't learn and what to do about it: Lessons on autonomous learning from cognitive science](https://mlsanity.com/paper/2603.15381/) (2026)
- [Human Cognition in Machines: A Unified Perspective of World Models](https://mlsanity.com/paper/2604.16592/) (2026)

## Abstract (arXiv)

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.