World Model Implanting for Test-time Adaptation of Embodied Agents
Authors: Minjong Yoo
Jinwoo Jang
Sihyung Yoon
Honguk Woo
Venue: ICML 2025 poster (2025)
Keywords: Embodied AI
Model implanting
World models
Large language model
Primary area: reinforcement_learning->everything_else
TL;DR: World Model Implanting for Test-time Adaptation of Embodied Agents

Abstract: In embodied AI, a persistent challenge is enabling agents to robustly adapt to novel domains without requiring extensive data collection or retraining. To address this, we present a world model implanting framework (WorMI) that combines the reasoning capabilities of large language models (LLMs) with independently learned, domain-specific world models through test-time composition. By allowing seamless implantation and removal of the world models, the embodied agent's policy achieves and maintains cross-domain adaptability. In the WorMI framework, we employ a prototype-based world model retrieval approach, utilizing efficient trajectory-based abstract representation matching, to incorporate relevant models into test-time composition. We also develop a world-wise compound attention method that not only integrates the knowledge from the retrieved world models but also aligns their intermediate representations with the reasoning model's representation within the agent's policy. This framework design effectively fuses domain-specific knowledge from multiple world models, ensuring robust adaptation to unseen domains. We evaluate our WorMI on the VirtualHome and ALFWorld benchmarks, demonstrating superior zero-shot and few-shot performance compared to several LLM-based approaches across a range of unseen domains. These results highlight the framework’s potential for scalable, real-world deployment in embodied agent scenarios where adaptability and data efficiency are essential.

World Model Implanting for Test-time Adaptation of Embodied Agents
Minjong Yoo Jinwoo Jang'
Sihyung Yoon'
Honguk Woo
Abstract
Zero-shot approach
Model integration approach
In embodied AL, persistent challenge is enabling agents to robustly adapt to novel domains with- out requiring extensive data collection or retrain- ing. To address this, we present a world model implanting framework (WorMI) that combines the reasoning capabilities of large language mod- els (LLMs) with independently learned, domain- specific world models through test-time com- position. By allowing scamless implantation and removal of the world models, the embod- ied agent's 's policy achieves and maintains cross- domain adaptability. In the WorMI framework, we employ a prototype-based world model re- trieval approach, utiliziig efficient trajectory- based abstract representation matching, incor- porate relevant models into test-time composition. We also develop world-wise compound attention method that not only integrates the knowledge from the retrieved world models but also aligns their intermediate representations with the rea- soning model's 's representation within the agent's policy. This framework design effectively fuses domain-specifte knowledge from multiple world models, ensuring robust adaptation to unseen do- mains. We evaluate our WorMI on the Virtual- Home and FLFWorl benchmarks, demonstrating superior zero-shot and few-shot performance com- pared to several L.L.-b approaches across a range of unseen domains. These results highlight the framework's 's potential for scalable, real-world deployment in embodied agent scenarios where adaptability and data efficiency are essential.
Action
Action
O
State
State
Theget
LLM
domain
Domain specific model
In-context approach
World model implanting (WerMI)
Action
Action
State
State
I.I.M
(
Compound attention
Retrieve
In-context
Data
Escemples
World models
Implam
Figure 1: Concept of WorMI. For sequential decision- making policies of embodied agents, our WorMI frame- work, world model implanting approach, is built upon an adaptive and composable policy structure that incorporates world-to-world integration and world-to-reasoning (world- to-LLM) alignment stages, enabling the flexible, test-time fusion of domain-specinf knowledge. This dual-stage de- sign not only combines the strengths of model integra- tion (which embeds a domain-specific model as part of the policy) and in-context adaptation (which incorporates ex- amples relevant to target domains) but also significantly enhances adaptability unseen domains.
AI, a field focused on creating intelligent agents capable of making sequential decisions and interacting with the physical environment through tasks such as navigation and manipulation (Huang al., 2022a;b; Yao et al., 2022; Wang et al., 2023). However, a significant challenge remains in enabling agents adapt effectively unseen domains with- out the need for extensive data collection or retraining. This adaptability is crucial for real-world applications, where environmental variations and diverse objectives often render rigid, domain-specific policies inadequate or ineffective.
1. Introduction
In recent years, policies powered by large language models (LLMs) have demonstrated remarkable success in embodied
Figure/ illustrates different approaches used for policy learning in the field of LLM-based embodied agents, in- cluding in-context adaptation (Song al., 2023) and model integration (Brohan et al., 2023; Hazra et al., 2024). The in-context adaptation approach retrieves relevant data from multiple domains to identify the most pertinent information for given situation. While offering flexibility to some extent, often suffers from inefficiencies due to the over-
'Department of Computer Science and Engineering, Sungkyunkwan University. Correspondence to: Honguk Woo <hwoo@ skku.edu>>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

World Model Implanting for Test-time Adaptation of Embodied Agents
head of searching, retrieving, and processing large volumes of information. On the other hand, the model integration approach combines two or more models with distinct prop- aspects (e.g.,
cross-domain embodied policy adaptation at test time, exploring the dual-stage model compositionality: world-to-world knowledge integration and world-to- reasoning alignment.
erties, explicitly separating domain-speciin affordance (Brohan et al., 2023), cost (Hazra et al., 2024)) from general knowledge. While this aims for efficient use of prior knowledge, it inherently lacks the flexibility to ex- pand knowledge beyond its learned domains. Our direction enhances these approaches by incorporating the test-time model composition into embodied policy adaptation.
We implement the prototype-based retrieval method to efficiently provide a robust adaptation process over multiple domain-specitee world models.
We develop the world-wise compound attention method, where the hierarchical cross-attention between the world models and the reasoning model is meta- learned to enable the dynamic integration of pertinent domain-specific knowledge.
Specifically, we present world model implanting frame- work, WorMI, which enables an embodied policy to adapt across diverse domains through test time, stal-stage com- position of domain-specific world models. This process encompasses both world-to-world integration and world-to- reasoning alignment, where LLMs' reasoning capabilities are leveraged. By seamlessly implanting and removing world models upon varied target domains, the agent's 's pol- icy tends to maintain grounded reasoning capabilities. It not only effectively leverages knowledge from multiple do- mains but also flexibly expands its knowledge, selectively utilizing only the most relevant information.
We demonstrate the superiority of WorMI through ex- periments in zero-shot and few-shot scenarios on the VirtualHome and ALFWorld benchmarks, highlight- ing its robustness across varied domains as well as its data-efficient ability to handle unseen domains.
2. Related Work
LLM-based Embodied Agent. LLM-based embodied agents, whose tasks are often referred to as embodied in- struction following, interact with the physical environment through activities such object manipulation and naviga- tion and require a high-level understanding of environmental dynamics and observations. Recent research has explored various techniques to enhance the reasoning and planning capabilities of these agents. They include implementing code-driven policies (Singh et al., 2023; Liang et al., 2023), generating reward funetions (Yu et al., 2023; Adeniji et al., 2023; Kim et al.), and integrating LLMs with additional domain-specific models to harness the affordances values of skills in the environment (Brohan et al., 2023; Hazra et al., 2024). Furthermore, in-context learning approaches, which incorporate previous demonstrations as part of the reasoning process, have also been introduced (Song et al., 2023).
Our WorMI framework integrates two key methods into adaptive, composable policy structure tailored for LLM- based embodied agents. (a) A prototype-based world model retrieval method selectively activates only a set of relevant world models. To determine relevance, each model's ' similarity to the current target domain is measured using object-wise state embeddings and clustering outcomes derived from trajectory- based prototypes. This ensures more robust and interpretable adaptation process across di- verse domains, particularly in zero-shot or few-shot scenar- ios. (b) A world-wise compound attention method effec- tively integrates the world models with the reasoning model by adaptively combining the pertinent knowledge from the retrieved model set. This facilitates effective and efficient policy adaptation during test-time execution. The interplay of these two methods enables the agent to dynamically com- pose and contextualize domain-specifer knowledge in its policy, through coherent integration and alignment of world models and a reasoning model.
While these approaches underscore the flexibility and ro- bustness of LLMs for embodied instruction following, they often rely on external and disconnected integration of en- virommental data or additional models, thus limiting their cohesion. Furthermore, they rarely address the challenge of effectively integrating multiple domain-specific models to facilitate adaptation to unseen domains. In contrast, our work introduces a scalable approach that implants multiple world models into LLM reasoning, enabling highly flex- ible and easily replaceable policy. This design facilitates strong generalization across diverse domains and rapid adap- tation, achieved through attention-driven, coherent knowl- edge integration from carefully selected multiple models.
Through experiments with VirtualHome (Puig et al., 2018), and ALFWorld (Shridhar et al., 2021), we demonstrate that the WorMI framework achieves competitive performance in both effectiveness and efficiency compared to several -tate-of-the-art LL---aseed embodied agents. For example, WorMI improves the average success rate in VirtualHome over SayCanPay (Hazra et al., 2024), a state-of-the- LLM- based embodied agent, by 20.41% % in zero-shot and 26.58% in few-shot scenarios. Our contributions are summarized as follows.
We present the novel WorMI framework to enable
Cross-domain Policy Adaptation. To tackle the challenge
2

---

World Model Implanting for Test-time Adaptation of Embodied Agents
of handling diverse environmental features and tasks, re- searchers have investigated a range of domain generaliza- tion methods (Zhou et al., 2022). Representative approaches include meta-learning, which learns how to learn multiple tasks (Finn et al., 2017; Nichol, 2018; Andrychowicz et al., 2016; Ha et al., 2016), hierarchical learning, which struc- tures knowledge across multiple levels by decomposing domain properties, and ensemble learning, which trains domain-specific neural networks for more robust perfor- mance. However, these methods often fall short of fully leveraging the specialized knowledge contained in domain- specific models. Our framework extends the benefits of these existing approaches, by implanting diverse, domain- specific world models into a single policy. This enables effective adaptation across a wide range of domains.
Di...., Dy from domains Dis..., Dx- A reasoning model nw. built upon LLM,, provides general reason- ing and decision-making capabilities. Given these indi- vidual models, we introduce a trainable composition mod- ule Co that selectively integrates only the relevant sub- set of pre-trained world models {M } with the fixed reasoning model TR. Our hierarchical approach first fuses multiple world models into a unified representation. which then aligns with nn to form the implanted policy Th Co((M). Mx); }, TR)-
As illustrated in Figure 2, to achieve this, our framework WorMI integrates two methods into an LLM--based policy structure. (a) Prototype-basee world model retrieval: at test time, only the most relevant world models for the cur- rent target domain are retrieved and integrated. To achieve this, we assess similarity using trajectory- based prototypes, which are derived from embedding and clustering outcomes of object-wise states. This ensures that the selected world models closely align with the agent's current environment and task requirements.
Model Merging. Researchers have explored variety of model merging methods that bring multiple models together in collaboration, often achieving performance levels that exceed those of any single model. These approaches, en- compassing model fusion (Wan et al.; Jiang et al., 2023), agentic workflows (Wu et al.; Hong et al.), and ensemble techniques (Shazeer et al., 2016), capitalize on the com- plementary strengths of individual component models to enhance overall performance. Recently, CALM (Bansal et al., 2024) introduced model composition framework that connects two models via a cross-attention layer for tasks like translation or multi-language coding, allowing the two models to combine their respective representations. Unlike CALM, we focus on embodied agents that operate in the physical environment, requiring adaptive decision- making across diverse and unseen domains. To address this, our WorMI framework employs a compound attention mechanism that dynamically combines and aligns multiple world models within the agent's policy. This enables the agent to flexibly fuse relevant knowledge at test time, thus ensuring robust reasoning and optimal decision-making. making.
(b) World-wise compound atten- tion: we train a world-wise compound attention module to integrate the intermediate representations of multiple world models and then align them with the reasoning model, en- abling effective and efficient knowledge integration. Meta- learning on various subsets of world models makes the com- pound attention flexibly integrate domain-specific knowl- edge from any combination of world models, even when new world models are added. With this hierarchical atten- tion mechanism, which aligns the integrated representations of domain-specific knowledge for reasoning, WorMI facili- tates zero-shot and few-shot adaptation at test time.
3.2. Prototype-based World Model Retrieval
World Models. Each domain-specific world model Mj is trained on its associated dataset D, {[(,, St, ae st+1)), where I denotes the embodied task instructions, s, is the agent's 's state at time t, and a is the executed action. Training comprises three auxiliary tasks: the dynamics task, predict- ing the next state st+1 given (st, at); the action affordance task, identifying feasible actions ag from si; and the behav- ior cloning task, learning a policy p(a, | 8t: 7). Collectively, these tasks equip each world model with domain-specifee knowledge of transitions, affordances, and decision-making
3. Approach
3.1. Overall Framework
A key challenge in LLM-based embodied agents is com- posing domain-specific knowledge into a general reasoming model, enabling agents to handle wide range of unseen do- mains. We present the WorMI framework which enables an agent to dynamically retrieve and compose relevant world models at test time, fusing them into its --aaeedd policy. This approach aims ensuring robust zero-shot and few- shot adaptation across diverse domains in ever-changing environments while avoiding the need for retraining or fine- tuning large models.
Prototype-basee Retrieval. For a given state st, we retrieve a set of world models Miet based on their relevance to 8t- This selection is made by measuring the embedding set distance between the set of embeddings E, which encodes object-wise states in the environment for each dataset D, and the set of embeddings & derived from st Accordingly, the retrieved set M.s can be formalized as
In WorMI, we assume the availability of pre-trained world models My...., My along with corresponding datasets
Ms & (1)8(((-( k)} (1)
M.a.
3

---

World Model Implanting for Test-time Adaptation of Embodied Agents
(a) Prototype-based world model retrieval
(b) World-wise compound attention
Embodied environments
Action grab applle
Aliging
Target domain
Domain
Domain 2
Domain N
Wedld wise contpound attemiion
Protolypes
S World model M2
Iniegnaed knowlelge
Prototype
Activate
Activate
Deactivate
World model
Similarity
Reaseming model
Instruction: Place the apple the fridge, Observations: apple bedroom.
model M;
World model M;
World model
Target domain
Figure 2: Overall procedure of WorMI. (a) By the prototype-based world model retrieval, relevant world models are selected using trajectory- based prototypes derived from obsect-wise state embeddings, given target domains. (b) By the world-witg compound attention, multiple world models' representations are integrated and then aligned with the fixed reasoning model for effective knowledge fusion.
where K is the number of selected world models and 8 (-:-) the Wasserstein distance. However, directly computing such set distance incurs high computational cost and tends to over-represent frequently occurring objects in the dataset. Therefore, we adopt a prototype-baser similarity to reduce computational overhead, especially at test time while maintaining representational diversity.
as bounded proxy for the distance between the under- lying datasets. Let p be the minimal maximum distance, optimized by (3), such that each point in &, is within p of prototype in Рş. Then, applying the triangle inequality under the Wasserstein distance yields the following.
5 8(P., I 8(E,, Eg) + 2p
(4)
To obtain such reliable prototypes, we construct the embed- ding set Ej in object-wise manner, using the object de- tection model dp that transforms input states object-wise ones and the embedding model dE that can be implemented using a language or vision embedding model. We have
3.3. World-wise Compound Attention
To integrate set of world models with the reasoning model in a policy, we develop the world-wise compound attention method, which utilizes hierarchical cross-attention to ef- fectively fuse domain-specifie knowledge and align it with LLM reasoning at test time.
E = (b)(() lo foi.....,, )n} = @D(s), SE D,}
(2)
where o denotes object-wise states for state s. We derive the prototypes by clustering &, using the k-center method, identifying the top-k embeddings that minimize the max- imum distance from any point to its nearest center. This approach both retains crucial object embeddings and re- duces the need for extensive distance computations over the entire dataset, formally described
Compound Attention. The compound attention Co maps the tth layer outputs of world models las, (a2:... IMx and 5th layer output of the reasoning model la to inputs of (j + 1)th layer of the reasoning model.
Inn + (MMM,, ..., Lax },
(5)
As illustrated in Figure 3, the compound attention Co con- sists of a linear projection layer Lg. world-level cross- attention layer Attn W and reasoning-level cross-attention layer Attn The linear projection layer matches the dimen- sions the world model's intermediate-layer representation with those of the reasoning model. The world-level cross- attention layer integrates these intermediate representations through weighted combination. The reasoning-leww cross- attention layer then aligns the integrated representation with the reasoning model based on its queries.
Pj =
argmin ICl=,
max min || ell [xEE aEC
(3)
These & centers serve as the prototypes py for each world model My. Using Py. which comprises only about 0.1% of the total embeddings in Ej. and p from the current observa- tion's 's embedding set ε, we replace ō(E,5E) E) with ((P,, p) in (1). In our implementation with VirtualHome, the number of embeddings in E, is 10,200, and K is 15.
Boundedness of Prototype-ae based Similarity, Here, we show that the distance between prototype sets can serve
Linear Projection. We first apply a linear projection Lg to each M,, matching its dimension to that of the reasoning

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Reusoning model Na
Reasoning-level Cross-attention. After obtaining the inte- grated representation of world models via the world-level cross-attention, align it to the reasoning model. In this step, the query is derived from L the keys are the sum of the world models' outputs iM, + + [Mx and the values are the aggregated representation from the world-wise at- tention Attng (Qw.Ww: Vw). ). Then, the query, key, and value of cross-attention Atin are defined as
World model M
World-wise compound attention Co
Reasoning-level cross-attention Attna
W
W
W
World model 2
QR WolTe KR Wi [Mm. + +.. + iMM]]
Sum
World-level cross-atiention Attri
VR = Wtttttng y (Qw. Kw, Vw).
(8)
W
Wo
3.4. Meta-Learning for Model Implanting
enable rapid adaptation to unseen domains, WorMI in- corporates a meta-learning approach (Nichol, 2018) for its compound attention Co, treating composer that dynamically aggregates and aligns domain- specific knowledge from multiple world models rather than merely linking existing models with the reasoning model.
World model Mg
Linear projection Lg
parameter-effict
Iaf 1
Learning Algorithm. At each inner-loop update, the pa- rameters 0, are initialized by meta- meta-paranne eeta-parameters 0. They are then adapted to such subset of world models M, S.U...UD, This en-
Figure Overall structure of the compound attention. Com- pound attention allows for the test-time composition of mul- tiple world models and the reasoning model. The world- level cross-attention first integrates the representation of multiple world models. The reasoning-level cross-attention then aligns the integrated representation to the reasoning model. Together, these hierarchical attention layers enable the policy to flexibly attend to relevant domain- specific knowledge and tightly couple it reasoning.
[M Mm. and datasets D, ables the compound attention Ce to learn how to weigh and align features from each M, for task planning. Formally, the loss function L is calculated as
Σ - log Te, (als)
(9)
(+,a)
for batch B sampled from Dj. the outer-loop update, we combine the adapted parameters back into the meta- parameters by 0 +- θ B m E 0) where 3 is the learning rate for the outer-loop update. This process encourages Co to learn a general integration and alignment strategy that can be specialized to an unseen domain with only a few gradient steps. Consequently, compound atten- tion acts as composer, fusing Nomain--specife knowledge in the reasoning process and enabling rapid adaptation to both new domains and newly introduced world models.
model's embedding space.
ÍA,, LI(,),) j (1,2....... K}
(6)
World-level Cross-attention. The world-level cross- attention Attn quantifies the relative importance of each world model's representation, conditioned on the current hidden state of the reasoning model. The query of the world- level cross-attention is derived from the reasoning model's layer output 4an- We treat each output of world model IM, "..,.... as key-value pairs. Then, the query, key, and value of cross-attention Allny are defined
4. Experiments
Environments and Datasets. We evaluate our approach in two embodied environments: VirtualHome (Puig et al., 2018). 3D simulation environment for household task exe- cution, and ALFWorld (Shridhar et al., 2021), text-based environment for indoor task simulation through language interaction. We treat tasks and scenes as separate domains for adaptation. For VirtualHome, we collect 1,023 episodes, covering 78 tasks (16 seen, 62 unseen) across 20 distinct scenes (6 seen, 14 unseen), each featuring unique room lay- outs and objects. For ALFWorld, we collect 3,554 episodes
Qw = Wo Kw = WK [lani... J.ss] Vi = [.
Here, Wo, WK are learned matrices in 0. The output of Attn W produces an integrated representation, determining the relevance each world model and assigning greater attention to the most pertinent ones for the given queries.
5

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table 1: Zero-shot performance in VirtualHome and ALFWorld. We use the 95% confidence interval, using 5 random seeds.
Seen Tasks & Seen Scenes
Seen Tasks & Unseen Scenes
Unseen Tasks & Unseen Scenes
Model
SR (1)
PS (U)
SR (1) PS (L)
SR (1) PS (1)
Evaluation VirtualHome
7SP
1...5%
29.02::0.06
8.95%::1.139
29.25:10.11
8.199%+00333%
29.36=:0556
LLM-FT
58.55%%+2..18%
17.37440.4
53.42%:11.799
17.70:10.23
42.82%%:10.66
20.79=0017
LLM-Pllaned
35.67%: 1.2%
27.15+0.13
28.55%:+0.42%
27.04::012
21.45%:21.
27.73=0005
SayCanPay
09.88%%:+2332%
14.53+00.55
64.74%%+4087%
15.64+0018
15.711%:22.0991
19.04=006
WorMI
85.78% 4..5%%
10.76=0.19
80.26 1.02%
12.42=0.09
66.12 % 0.80%
15.17:00.08
Evaluation in ALFWorld
ZSP
2.30%:+0.1%%
49.3441:066
0.00%
49.04:10.95
2.13%
49.68
LLM+FT
45.72%11..99
14.63
29. 10.72%
38.4912287
29.40%
16.8
LLM--Planner
17.73%.10..19
32.09
18.63% 1..82%
40.50
12.31%::10.00.
46.33:1
SayCanPay
10.67%.11.244%
18.3741.93
34.10%
20.8211.21
39.66% 1,36%
23.86:11990
WorMI
62.51 11.6%
9.96:1..29
52.67% 1..9%%
17.74=008
51.67 1.22233
10.18::0.3
across diverse scenes. Following CL-ALFRED benchmark settings (Kim al., 2024), the data is clustered into 4 scene types (3 seen, unseen) and 6 task types (4 seen, 2 unseen).
scenes). Specifically, WorMI outperforms the most com- petitive baseline SayCayPay in unseen tasks and scenes, achieving 20.41 % increase in SR and a 3.87 step reduction (20.32% improvement) in PS in VirtualHome, and 12.01% % increase in SR and a 3.68 reduction (18.23% improvement) in ALFWorld. These zero-shot results on robust generaliza- tion for unseen domains highlight the effectiveness and data efficiency of knowledge integration from multiple world models in WorMI. Moreover, the improvement in seen do- mains demonstrates the impact of the world-to-reasoning alignment which can enhance the agent's 's understanding of the environment.
Evaluation Metries. We employ two evaluation metrics: Success Rate (SR) and Pending Steps (PS). SR measures the proportion of tasks completed in VirtualHome and sub- goals completed in ALFWorld. PS represents the average number of timesteps required complete tasks, akin to cost-effectiveness in Hazra et al. (2024).
Baselines. For comparison, we evaluate four base- lines that span diverse cross-domain adaptation strategies. ZSP (Huang et al., 2022a) is a zero-shot approach that applies pre-trained model new domains without addi- tional adaptation. LLM+FT (Fine-Tuuned LLM) performs adaptation through fine-tuning limited domain-specific data.
Few-shot Adaptation. We further evaluate few-shot adap- tation scenarios, where each target domain (unseen dur- ing training) provides only a small number of demonstra- tions at test time. This setup allows to examine how effectively WorMI adapts with minimal additional data. Table 2 compares performance across different few-shot settings (1, 5, and 10 shots), illustrating how increasing the available demonstrations affects adaptation quality. As shown, WorMI surpasses all baselines, achieving on average 26.58 gain in SR and 4.98 step reduction in PS in Virtu- alHome, and 19.16% gain in SR and 7.36 step reduction in PS in ALFWorld, compared to SayCanPay. These results indicate that the lightweight compound attention module, trained via meta-learning, provides parameter-efficient up- dates while keeping the world models and the reasoning model frozen, thereby facilitating rapid, efficient adaptation even in minimal data conditions.
(Song et al., 2023) employs in-context learning to generate and refine high-level plans based on examples. Lastly. Say CanPay (Hazra et al., 2024) is a ttate-of-the--at LLM-based model integrated approach that incorporates heuristic cost minimization into LLM reason- ing and planning. These basclines collectively allow us to assess both the efficiency and robustness of our WorMI.
In our implementation, we use a fixed Llama-3.2- 3B (AI@ MMta, 2024) model for ZSP, ALM--llnnee the Say model in SayCanPay, and the reasoning model in WorMI. For LLM+FT, the Pay model in SayCanPay, and the world models in WorMI, we use a trainable L.ama-3.2-1B model.
4.1. Main Results
4.2. Analysis
Zero-shot Adaptation. We evaluate each method in a zero-shot scenario, relying solely on seen domain training data without any target domain data, As shown in Table 1, WorMI consistently outperforms all baselines across both seen and unseen domains, particularly demonstrating ro- bust generalization to unseen domains (unseen tasks and
We conduct several analyses on WorMI with unseen tasks and scenes in VirtualHome.
World-level Attention Map. In Figure 4, visualize the attention weights of the world-level cross-attention used in the compound attention of WorMI. We employ three
6

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table 2: Few-shot performance in VirtualHome and ALFWorld. We use the 95% confidence interval, using 5 random seeds.
1-Shot
5-Shot
10-Shot
Average
Model
(t)
PS (L)
SR (t)
(1)
(t)
PS (4)
SR (1)
(L)
Evaluation in VirtualHome
LLM-FT
12.35%==0.854
20.46+0.15
47.22%++0466
16.82=0.9
51.45%%=...o 1.06%
16.32=0.39
77.01%%40.99
18.03
L.MM-Planner
24.63%6=0.34
19.194
29. 29.80%=+0..55
26.98=008
26.60 06.60
29.3 4.44
26.59+ 18
PayCanPay
46.75%6211m 022
19.12:0.16
49.80%tt1.115
13.32=0.11
56.2 56.24%6:t077 79%
16.00=0.13
50.93%:+nn974
16.88
WorMI
1.57%
12.18:0031
78.04 1..34%
11.85
% 1.099%
11.6
77.51
0.98%
11.90:10.7
Evaluation in ALFWorld
LLM-FT
26. 36.78%%=1..0
46.51+0227
29.79%%+1..66
46,21=037
33.67%:=40998
43.70=0.83
30.05%++
45.47+0.49
LLM-Planner
18. 18.28%%=1100
43.94+287
40.50=3.14
17.76% 1.2%%
42.45=2264
17.61 17.61%++11
7616-....0 26%
42.30+2..6 88
SayCanPay
40.94%%21... 34%
21.58:+1.14
36.70%tt1.319
22. 22.15=3355
39.54%6::1 488%
24.74=3.48
39.06%6.: 1.38%
22.82+2772
WorMI
51.50%:12..14
21.46:33.6
58.69 11948
13.14::.00
64.46 0.88%
11.79::081
58.22 116888
15.46::192
lected subset of world models. In both WorMI and WorMI- R, selectively use three out of the six available world models. As shown in Table 3(a), WorMI outperforms WorMI-E and WorMI-R by 17.49% and 4.08% in SR, while reducing PS by 3.76 and 1.79 steps, respectively. These results confirm that the prototype-based retrieval effectively selects the most suitable world models test time.
Bedroom
Livingroom
Kitchen
20
25
Turn on.tv
Open cabinet
Put folder sofa
Place breadslice microwave
For the ablation study on the world-wise compound at- tention method, we compare WorMI against two variants: WorMI-CONCAT, which concatenates all world model rep- resentations, and WorMI-ADD, which instead sums these representations. As shown in Table 3(b), WorMI outper- forms WorMI-CONCAT and WorMI-ADD by 18.33 and 9. 9.65% in SR, while reducing PS by 4.86 and 1.79 steps, re- spectively. These improvements attributed to the world- level cross-ates ausss-attention's 's ability integrate relevant knowl- edge from the world models effectively, consistent with the domain-dependent attention outcomes shown in Figure 4.
Figure 4: Visualization of world-level attention map
world models, each derived from different rooms (e.g., bed- room, livingroom, and kitchen), to facilitate adaptation to unseen domains. We observe that WorMI dynamically shifts its focus among the three world models as the target do- main varies along with different tasks, assigning higher attention weights to the model most relevant to the current task. This adaptive prioritization of knowledge enhances context-aware reasoning, enabling more informed and effec- tive decision- mocision-makkng.
Table 4: Impact of LLMs on WorMI and baselines
Table 3: Ablation study on prototype-bet based world model retrieval and world-wa wise compound attention
L.M
Model
SR (t)
PS (L)
1..M-Planne SayCanPay
51.11%=027%
18.33+4008
(a) Prototype-based world model retricval
Llama-3.2-1 IB
49.28%=0020s
18.88:10.10
WorMI
71.18% 10.62%
13.15
Model
SR (1)
PS (L)
I.IMMPlanner
21.45%600424
27. 27.73:10.00
WorMI-E
18.63% 23%
18.93#4449
Llama-3.2-3R
SayCanPay
45.71%=00599
19.04:10.05
WorMI-R
62, 62.04% 00414
16.961020
WorMI
66.12%
15.17=0.00
WorMI
66.12% % 10.80%
15.17...08
LLM-Pllaned
9.93%.20.279
27.37=10.11
Llama-3.2-1B
SayCanPay
42.22% 3.3%
21.34:+:.35
(b) World-wise compound attention
WorMI
49.65 1-0.83%
18.99
Model
SR (1)
PS (L)
WorMI-CONCAT
47. 47.79%+2.99 39%
20.03
Impact of LLMs. Table shows that WorMI consistently outperforms the baseline methods across LLMs of various sizes, including 1B, 3B, and 11B parameters. As the LLM size increases, WorMI capitalizes on enhanced reasoning to attain progressively larger gains, ultimately surpassing LLM- -ll-nneeeer by 20.07% when using the larger 1BB model. In contrast, on the L.LM size, exhibiting substantial degradation with
WorMI-ADD
56.477
17.43
WorMI
66.12% 1.12%%1..00
15.17=0.08
Ablation Study. For the ablation study on the prototype- based world model retrieval, we compare WorMI against two variants: WorMI-E, which uses all world models for each inference, and WorMI-R, which uses a randomly se-
performance heavily depends

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table 5: Impact of world model scale (WMs)
steps. Notably, the substantial decrease in PS for the mul- tiple instructions highlights WorMI's effective reasoning to facilitate integrated task planning over set of related instructions.
Num. of WMs
SR (1)
PS (L)
12.75%:10.500
20.54:10.02
62. 22.944%:t00122
17.21=00200
66. 12%6:t0.59%
15.17=00.8
61. 01.96%2:002255
16.68.+0.34
SR (%)
73
48. 2.222%0::
18.93-+049
50
23
Table Performance for complex instruction scenarios
WM3
(a) Continual model implanting from WMI WM6
Model
SR (t)
GC (1)
PS (1)
100
Long horizon instructions
-34.92%
23.66%
-24.22%
33..00%
-24.3%
21,91%
0.65%.10.519
17.16%:10.58%
16.12::1.2
SayCanPay
6.55%%.10.51%
39.22%%10.666
57.87
60-
WorMI
19.61 % 1..88%%
49.35 1..55%
53.15:10.70
With full world models
Multiple instructions
Without pertinent world model
LLM- Planner
1.31%
29.02%=20.08%
17.17777.6
0
Domain
SayCanPay WorMI
6.54%
42.32%
57 97.94+092
20.26%=+1..4%%
60.62 9.89%
43.49:+1.23
(b) Effect of removing the pertinent world model
Figure 5: Performance for continual model implanting
the smaller 1B model. This highlights the limitation of in-context adaptation, whose benefits are maximized with larger, more capable LLMs. SayCanPay, on the other hand, maintains relatively consistent performance across differ- ent model sizes but does not fully exploit the capabilities of larger LLMs. We speculate that its Can (affordance) and Pay (cost) models contribute more SayCanPay's 'ss performance than the Say (reasoning) model, which relies on LLMs.
Continual Model Implanting. We consider a case of dy- namic model addition and removal, which referred to as continual model implanting. In Figure 5(a), we demonstrate how performance improves as we incrementally add new world models from WMI WM6. In this evaluation, we intentionally employ only the world models pertinent to the target domains, confirming that each domain-specific model provides complementary knowledge and leads to pro- gressively better overall results. Conversely, Figure 5(b) demonstrates how removing the world models can reduce performance. We observe machine unlearning effect by discarding specific world models longer available, con- firming the impact of lost domain-specific knowledge on overall performance. This shows WorMI's extensibility to various unlearning scenarios, where irrelevant or outdated knowledge can be seamlessly removed.
Scalability for World Models. Table 5 shows how varying the number of implanted world models affects WorMI's performance. While using between and world models yields higher SR and lower PS, employing only single model or as many as six leads marked performance drop. By selectively retrieving only the most relevant world models through compound attention, WorMI can expand Nomain-specific knowledge without exceeding its capacity, thereby balancing scalability and performance.
Complex Instructions. conduct two case studies Vir- tualHome verify the robustness of WorMI against com- plex task scenarios. Long-horizon instructions specify sequence of tasks (or sub-goals) that must be completed step by step. Multiple instructions define set of tasks to be executed concurrently, requiring integrated planning and execution for optimal completion. In both scenarios, we measure goal-conditioned success rate (GC) as the propor- tion of completed sub-goals alongside SR and PS, taking into account the extended, sequential nature of long-horizon instructions. As shown in Table 6, WorMI achieves both the highest SR and GC and the lowest PS for both scenarios. For long-horizon instructions, WorMI surpasses SayCanPay by 13.1% SR, 13.13% in GC, and reduces PS by 4.72 steps, while for multiple instructions, achieves gains of 13.72% in SR, 18.30% GC, and reduces PS by 14.45
5. Conclusion
We presented the WorMI framework which enables bodied policy dynamically compose multiple world mod- els with its fixed reasoning model at test time, thereby enhancing cross-domain adaptability. By combining the prototype-based world model retrieval with the world-wise compound attention, WorMI achieves dual-stage knowledge fusion, encompassing world-to- worldd knowledge integra- tion and borld-to-reasoning alignment. Evaluation results in VirtualHome and ALFWorld confirm that WorMI achieves robust adaptation to unseen domains in both zero-shot and few-shot scenarios, outperforming several LL-M-bseed base- lines.
Limitation. Despite its advantages, WorMI faces two key limitations. First, inferring over multiple domain-speciff.
8

---

World Model Implanting for Test-time Adaptation of Embodied Agents
world models increases computational overhead, potentially posing challenges in resource-constrainee setrings. Sec- ond, the framework's reasoning relies LLM, making its performance inherently dependent on the strengths and weaknesses of the underlying language model.
composition, 2024. abs/2401.02411
URL https://arxiv..o
Brohan, A., Chebotar, Y., Finn, C., Hausman, K., Herzog, A., Ho, D., Ibarz, J., Irpan, A., Jang, E., Julian, R., et al. Do as i can, not as i say: Grounding language in robotic affordances. In Conference on robot learning, pp. 287- 318. PMLR, 2023.
Acknowledgement
This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP)
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
grant funded by the Korea government (MSIT), (RS-2022-
11220043 (2022-0-00043), Adaptive Personality for Intel- ligent Agents, RS-20222 2--2--122221155 (2022-0-01045), Self- directed multi-modal Intelligence for solving unknown, open domain problems, RSS-2025-0221866 1768, Accelerated Insight Reasoning via Continual Learning, and RS-2019- 11190421, Artificial Intelligence Graduate School Program (Sungkyunkwan University)), the National Research Foun- dation of Korea (NRF) grant funded by the Korea govern- ment (MSIT) (No. RS-2023-00213 181),, 18), IITP-ITRC (Infor- mation Technology Research Center) grant funded by the Korea government (MIST) 2015-2222---2202- 2024-00437633, 10%). IITP-ICT Creative Consilience Program grant funded by the Korea government (MSIT) (ITT--025-RS-202)) I1201821, 10%), and by Samsung Electronics.
Ha, D., Dai, A., and Lee Q. V. Hypernetworks. arXiv preprint arXiv: 1601: 09906, 2016.
Hazra, R., Dos Martires, P. Z., and De Raedt, L. Say- canpay: Heuristic planning with large language models using learnable domain knowledge. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp.
2024.
Hong, S., Zhuge, M., Chen, J., Zheng, X., Cheng. Y., Wang, J., Zhang, C., Wang, Z., Yau, S. K. S., Lin, Z., et al. Metagpt: Meta programming for a multi-agent collabora- tive framework. In The Twelfth International Conference on Learning Representations.
Impact Statement
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Huang, W., Abbeel, P., Pathak, D., and Mordatch, Lan- guage models as zero-shot planners: Extracting action- able knowledge for embodied agents. In International conference machine learning, pp. 9118 9147. PMLR, 2022a.
References
Huang. W., Xia, F., Xiao, T., Chan, H., Liang, J., Florence, P., Zeng, A., Tompson, J., Mordatch, I., Chebotar, Y., et al. Inner monologue: Embodied reasoning through planning with language models. arXiv preprint arXiv:2207.05600 2022b.
Adeniji, A., Xie, A., Sferrazza, C., Seo, Y., James, S., and Abbeel, P. Language reward modulation for pretraining reinforcement learning. arXiv preprint arXiv:2308.12277 2023.
Jiang, D., Ren, X., and Lin, B. Y. LIm-blender: Ensembling large language models with pairwise ranking and genera- tive fusion. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 14165-144 11155-1418..111 2023.
AI Meta. Llama 3.2 model card. 2024. URL https://git com/ htt///mtaalllmaaa 11 11ama-modell selleesslllloo ob/maa///////// ieaaiin/modd s 111mmm33 2 2/MODEL_CARD mdd
Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and De Freitas, N. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Kim, B., Seo, M., and Choi, J. Online continual learning for interactive instruction following agents. In The Twelfth International Conference on Learning Representations, 2024.
Kim, W. K., Lee, Y., Kim, J., and Woo, H. Llm-based skill diffusion for zero-shot policy adaptation. In The ehirth-eighth Annual Conference on Neural Information Processing Systems.
Bansal, R., Samanta, B., Dalmia, S., Gupta, N., Vashishth, S., Ganapathy, S., Bapna, A., Jain, P., and Talukdar, Lim augmented lms: Expanding capabilities through

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Liang, J., Huang, W., Xia, F, Xu, P., Hausman, K., Ichter, B., Florence, P., and Zeng. A. Code policies: Language model programs for embodied control. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 9493-9500. IEFE, 2023.
Yu, W., Gileadi, N., Fu, C., Kirmani, S., Lee, K.-H., Are- nas, M. G., Chiang, H.-T. L.., Erez, T.,, Hasenclever, L.,, Humplik, J., et al. Language to rewards for robotic skill synthesis. In Conference on Robot Learning, pp. 374 404, PMLR, 2023.
Nichol, A. On first-order meta- leaannen algorithms. arXiv preprint 1XX3::8800009999 2018.
Zhou, K., Liu, Z.,, Qiao, Y., Xiang, T., and Loy, C. C. Do- main generalization: A survey, IEEE Transactions Pattern Analysis and Machine Intelligence, 45(4):4336- 4415, 2022.
Puig, X., Ra, K., Boben, M., Li, J., Wang, T., Fidler, S., and Torralba, A. Virtualhome: Simulating household activities via programs. In Proceedings of the IEEE con- ference computer vision and pattern recognition, pp. 1494-8502, 2018.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Lc, Q. Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gate mixture-of-expers layer. In International Conference on Learning Representations, 2016.
Shridhar, M., Yuan, X., Cote, M.-A., Bisk, Y., Trischler, A., and Hausknecht, M. Alfworld: Aligning text and embod- ied environments for interactive learning. In International Conference on Learning Representations, 2021.
Singh, I., Blukis, V., Mousavian, A., Goyal, A., Xu, D., Tremblay, J., Fox, D., Thomason, J., and Garg, A. Prog- prompt: Generating situated robot task plans using large language models. In 2023 IEEE International Conference Robotics and Automation (ICRA), pp. 11523-11530. IEEE, 2023.
Song, C. H., Wu, J., Washington, C., Sadler, B. M., Chao, W.-L., and Su, Y. Lim-planner: Few-shot grounded plan- ning for embodied agents with large language models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 2998-3009, 2023.
Wan, F., Huang, X., Cai, D., Quan, X., Bi, W., and Shi, S. Knowledge fusion of large language models. In The Twelfth International Conference Learning Represen- tations.
Wang, 7.,, Cai, S., Chen, G., Liu, A., Ma, X., and Liang. Y. Describe, explain, plan and select: Interactive planning with large language models enables world multi- task agents. arXiv preprint urXiv:2302.01560, 2023.
Wu, Q. Bansal, G., Zhang, J., Wu, Y., Li, B., Zhu, E., Jiang, L., Zhang, X., Zhang. S., Liu, J., al. Autogen: Enabling next-gen IIm applications via multi-agent conversation. In ICLR 2024 Workshop on Large Language Model (LLM) Agents.
Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., and Cao, Y. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629, 2022.
10

---

World Model Implanting for Test-time Adaptation of Embodied Agents
A. Supplementary Proofs
A.I. Boundedness of Prototype-basee Similarity
By leveraging prototypes, the similarity between datasets can be bounded by the similarity between these prototypes. Based on the triangle inequality for the Wasserstein distance, we have
(p: Pý) (E,, &g) 8(p;, Ei) + k(p,; Eg).
(A.1)
By optimizing (3), we obtain minimal maximum distance such that
Vo E E 3c E C: | - cll S
(A.2)
Since any point in & and & can be matched to a prototype within a distance p. the distance between the prototype sets can be bounded
8(p, P2) 8(E,, Ej) 2p.
(A.3)
B. Environments
B.1. VirtualHome
VirtualHome (Puig et al., 2018) simulation environment developed with Unity, designed to let embodied agents interact with everyday household items in order to complete given instructions. In VirtualHome, we use 20 varied house configurations each with unique room layouts and object placements, creating a complex testbed that closely mirrors real-world domestic scenarios. Within this embodied space, agents can perform wide spectrum of actions from picking up and moving objects switching appliances and off, opening and closing doors and drawers.
Figure A.1: Visualization of VirtualHome
(faucet, inside, bathroom), (stall, inside, bathroom), (kiichencoumterdrawwr,, inside, kitchen), (mousemat, inside, livingroom), (bookshelf, inside, livingroom), (character, close, plum), (bedroom, adjacent, kitchen), (stove, inside, kitchen), (peach, inside, bedroom), (character, inside, bedroom), (tv, inside, livingroom), (character, hold, none).
Figure A.2: An example of an observation in VirtualHome
Observations and Actions. We 6 action types (e.g., walk, grab, open, put, putin, and switchon) in VirtualHome, which update the graph according to the rules of VirtualHome Environments. We represent VirtualHome's graph list of triples, where each triple describes object node and its relationship to another node. The details of VirtualHome actions provided in Table A.1. Action examples are from VirtualHome Documents.
Instructions. There are 78 instructions consisting of types of tasks. These tasks instruct the agents on what they should do in situation. The goal of an instruction to reach a successful state, that represented by triples environment graph. Details are Table A.2.
B.2, ALFWorld
ALFWorld is text-driven simulation for household tasks, created to help train and evaluate agents on formulating high-level strategies from natural language instructions in an embodied context. Drawing on series tasks originally introduced by

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table A.I: Actions in VirtualHome
Types
Example
Walk [Object] or Walk(room]
walk kitchen grab apple
Grab [Grabbable Object]
Open [Object]
open fridge
Put [Grabbable object] [Object]
put apple table
Putin [Grabbable object] [Containers]
putin apple fridge
SwitchOn [Switchable object]
switchon stove
Table A.2: Instructions in VirtualHome
Type
Amount
Example
Goal State
Turn on tv
[(tv, is, on)]
Turnn
Turn on radio
[(radio, is, on)]
Turn on microwave
[(microwave, is, on)]
Open cabinet
[(cabinet, is open)]
Open
7
Open dishwasher
[(dishwasher, is, open)]
Open microwave
[(microwave, is, open)]
Put apple desk
[(apple, desk)]
PutOn
30
Put clock sofa
[(clock, on, sofa)]
Put bananas on microwave
[(bananas, on, sofa)]
Place towel closet Place paper bookshelf
[(towel, inside, closet)) [(paper, inside, bookshelf)) [(plum, inside, fridge)]
PlaceIn
32
Place plum in fridge
the ALFRED benchmark, it maintains long task sequences and unidirectional state changes key elements that give the environment high degree of authenticity and practicality.
Observations and Actions. ALFWorld is combination of text and visual simulation. When a scenario begins, the agent's s current observation and the assigned task are displayed to the user in text form. The goal of the game is to reach a certain successful state by entering the appropriate text commands. A description of the game's 's observation-action sequence provided in Figure A.3, and the list of possible actions can be found in Table A.3.
Welcome to TextWorld, ALFRED! =-
Observation: You are in the middle of room. Looking quickly around you, you see a armchair 1, coffeetable 1, garbagecan 1, shelf 14, shelf 13, a shelf 12, a shelf 11, a shelf 10, a shelf 9, shelf 8, a shelf 7, a shelf 6, a shelf 5, shelf 4, a shelf 3, a shelf 2, shelf I, a sofa 1, a tvstand 2, and a tvstand 1. Your task is to: put a pillow armchair.
Action: look
Observation: You
the middle of room. Looking quickly around you, you nothing.
Action:
go to sofa 1
Observation: You arrive at sofa 1. On the sofa 1, you see box 2, creditcard 3, a pillow 2, pillow 1, and a remotecontrol 2.
Figure A.3: example of trajectory in ALFWorld
Instructions. There are main instruction types in ALFWorld, which imply the overall content that an agent must perform in scenario. By combining numerous objects, receptacles, and rooms with one instruction, 3,554 scenarios are created. The details are explained Table A.4.
12

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table A.3: Actions in ALFWorld.
Type
Example
Goto [Receptacle Ohject]
Goto dresser
Open [Receptacle Object]
Open microwave
Close [Receptacle Object]
Close microwave
Pickup [Object] [Receptacle Object]
Take alarmelock from dresser
Put [Object] |Receptacle Object]
Put book in/on dresser
Heat [Ohject] |Receptacle Object]
Hcat bread with microwave Cool cup with fridge
Cool [Object] |Receptacle Object]
Clean [Object] [Receptacle Object]
Clean kettle with sinkbassin
Slice [Object] [Instrument Object]
Slice bread with knife
Examine [Object]
Examine clock
Examine |Receptacle Object]
Examine countertop
Table A.4: Instructions in ALFWorld
Type
Example
Pick & Place
Place box with keys in it on the coffee table.
Pick Two & Place
Place two wine bottles in bin.
Clean Place
Put clean of soap in the drawer.
Heat & Place
Put a warmed slice of tomato on the table.
Cool & Place
Place chilled potato slice inside the microwave.
Examine & in Light
Inspect racket with lamp desk.
C. Implementation Details
C.1. Baselines
ZSP (Huang et al., 2022a) is a zero-shot policy approach that directly applies pre-trained model to new domains without any adaptation procedure. This baseline serves as a reference to assess the improvements achieved through our WorMI. uses a single LLM (Llama-3.2-3B -BB--tstuuc,,, and takes an observation from the environment, then injects it into the prompt described Figure A.4 and A.5. So. it takes the prompt and generates an action step by step. For implementation, we refer to the opensource
LLM+FT is representative baseline that demonstrates adaptation through fine-tuning limited domain-specins data. By comparing against this approach, we evaluate the efficiency and effectiveness of WorMI relative to fine-tuning in cross-domain adaptation. In the few-shot adaptation scenario, we additionally train the fine-tuned LLM with the few-shot data sampled from the target domain.
L.MM-lla (Song et al., 2023) is an embodied planner that leverages L.MSs for reasoning and task execution within the environment. It employs in-context learning to generate high-level plans and adapt them based on current observations. This baseline allows us evaluate the gains achieved by WorMI over the in-context learning method. For evaluation, we employ the
sentence embedding model for the retriever, which uses the same dataset employed in the training of other approaches requiring additional fine-tuning (e.g., LLM+FT, SayCanPay, WorMI). Furthermore, for the few-shot scenarios, we add few-shot data examples into the retriever. For implementation, we refer to the opensource 2
SayCanPay (Hazra et al., 2024) is a state-of-the-aa reinforcement lcarning- planning approach that incorporates pre-trained skills to assess feasibility and employs heuristic cost minimization. By comparing against SayCanPay, we aim to demonstrate the planning efficiency of WorMI. uses three models under the same prompts: (1) We use the Llama-3.2-3B model Say model, (2) We use optimal affordance from the environment for Can model, and (3) We train the Pay model based on the Llama-3.2-1BB For implementation, we refer to the opensource 3, The hyperparameter settings of SayCayPay
are in Table A.5.
planner
htnps//gittta.cmmmSSSSAttiCCaonr Pllannne
13

---

World Model Implanting for Test-time Adaptation of Embodied Agents
The hyperparameter settings of baselines are summarized in Table A.5. The prompts we described in Figure A.4 and
A.5.
Table A.5: Hyperparameter settings and configurations of baselines
Hyperparameter
Value
Trainable model (LLM+FT, and Pay model in SayCayPay)
Llama-3.2-1B
Reasoning model (ZSP, LLM--Plnner, and Say model in SayCayPay)
Llama-3.2-3B
Batch size
4
Gradient steps
200
Learning rate scheduler
cosine
Initial learning rate
x 10
Learning rate (for few-shot learning)
1 x 10-6
Temperature (both of A.lama-3.2-1BB and Llama-3.2-3B)
1.0
[System]
You are a home robot agent. You can use skills, (walk [object or room], grab [object], switch [object]. open [object], putin [target object), put [target object]). You should return only a skill after "Action.". Room: livingroom, bathroom, kitchen, bedroom.
[User]
Instruction: {instruction)
Observation: {observation)
Action:
Figure A.4: System prompt in VirtualHome
[System]
You are a home robot agent. You can use 10 skills, (go to [object], take [object] from [object], put [object] on [object], open [object], close [object], toggle [object], heat [object] with [object], cool [object] with [object]. clean [object] with [object], look). You should return only skill after "Action:". Room: livingroom, bathroom, kitchen, bedroom.
[User]
Instruction: {instruction]
Observation: {observationg)
Action: {actiona}
Observation: {observation:
Action: {action)}
Action:
Figure A.5: System prompt in ALFWorld
C.2. WorMI (Ours)
Our WorMI framework integrates two key methods into an adaptive, composable policy structure tailored for LLM-based embodied agents, facilitating the test-time, dual-stage composition of world models. (a) A prototype-based world model retrieval method selectively activates only a set relevant world models. To determine relevance, each model's similarity to the current target domain measured using object-wist state embeddings and clustering outcomes derived from trajectory- based prototypes. This ensures a more robust and interpretable adaptation process across diverse domains, particularly in
14

---

World Model Implanting for Test-time Adaptation of Embodied Agents
zero-shot few-shot settings. (b) world-wise compound attention method effectively integrates the world models with the reasoning model by selectively combining the most pertinent knowledge from the set. This facilitates effective and efficient policy adaptation during test-time execution. The interplay of these two methods enables the agent to dynamically compose and contextualize domain-specific knowledge in its policy across domains, through coherent integration and alignment of world models and a reasoning model. The hyperparameter settings of WorMi are in Table A.6. Algorithm shows the entire procedure of the world model implanting framework.
Algorithm WorMI Framework
World models (M1,.... Mv), Reasoning model mn, Datasets D1,..., DN. learning rate o,
2: World-wise compound attention Ce for policy we
3: // Training world-wise compound attrition via meta-learning
4: Initialize parameters 0 randomly
5: Sample the subset of world models M1, C {M1,..,, ... 3v and associated datasets
... ... (D1, Dx}
6: for meta update steps 1, 2,... AM do
7: for each M, and D, do
8: 0g- <- 0
Implant the world model me Ca, (Mj, RR) 10: for iteration=1, 2....,,, X do
9:
11:
Sample batch B from Dataset D,
12:
0 x 8j Vo, [(0) B) in (9)
13:
end for
14:
end for
15:
+ 0 B- M1 1 (0, - 0)
end for
17:
18: // Test-time adaptation with the prototype-based world model retrieval and world-wise compound attention
19: Object detection model dp, Embedding model 4p, Environment enc
20: for = 1, 2,.,,.., NN do
21: = ( foi,-.. ON) = (
22: Optimize prototype Pj by (3)
end for
24:
0, 8 res.reset()
25: loop
26:
E (ff( (oi...., ON) = (s),
27: Optimize prototype p by (3)
28: Myet = {My j jE Topk(
p) K)}
Th C9 RR), a: (s:))
30:
8++1 env.step(a:), t ++
31:
end loop
15

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table A.6: Hyperparameter settings and configurations of WorMI
Parameter
Value
Learning world models Mi,... MN
Base model
Llama-3.2-IB
Batch Size
4
Gradient steps
2000
Learning Rate Scheduler
cosine
Learning Rate
3 x 10-6
Temperature Intermediate connection layer
1.0
[13, 27]
Learning compound attention
Reasoning model IR
Llama-3.2-3B
Batch Size
4
Meta update steps Am
8
Inner-loop gradient steps À Learning Rate Scheduler
30
cosine
Learning Rate a Meta learning Rate
1 x 10-5
1 x 10-1
Temperature
1.0
Learning Rate (for few-shot learning)
1 x 10
[13, 27] for reasoning model
Intermediate connection layer
[7.15] for world models
Prototype-based world model retrieval
The number of embeddings in prototype k
15
The number world models N
The number retrieved world models K 3
D. Additional Analysis
D.1. Analysis world model usage
We categorize different cases to analyze the world-level attention map.
(a) Object existence
(b) Object position
(c) Task decomposition
Figure A.6: Analysis for the world-level attention value in various cases
Case 1. Object existence. First, examine how the attention changes when particular object exists in only world model. For example, Figure A.6(a) illustrates the scenario where the object "pie" appears only in domain A3, and not in domains Al or A2, under the instruction "Place pie in microwave." The attention value remains high in world models from domain A3 throughout all steps of the episodes, demonstrating that the world-wise cross-attention effectively focuses on the world model containing the relevant information.
Case 2. Object position. We examined how the world model operates when the positions of objects differ. For instance, in the task "Put mug on coffeetable," the target domain places the mug in the livingroom and the coffectable in the bedroom. In domains B1 and B2, the mug is located in the kitchen while the coffeetable remains in the bedroom. In contrast, in domain B3, both the mug and the coffeetable are situated in the livingroom. As shown in Figure A.6(b), initially, high attention values are assigned to B3. Over time, attention values for world models from domains B1 and B2 increase. This behavior demonstrates that world-level cross-attention selectively utilizes the world models containing the relevant information needed to successfully execute the task.
Case 3. Task decomposition. We analyzed scenarios where tasks are partially present in each world model. For example,
16

---

World Model Implanting for Test-time Adaptation of Embodied Agents
when performing the task "Place plum in cabinet" in the target domain, the world model from domain CI is trained with data related cabinet tasks (e.g.. "Place mug in cabinet"), while the world model from domain C3 is trained with data related to plum tasks (e.g.. "Place plum in fridge"). The word model from domain C2 is less relevant. As shown in Figure A.6, during the process of picking up the plum, the world model from domain C3 exhibits a high attention value. When placing the plum in the cabinet, the world model from domain C1 shows a high attention value. This demonstrates that through knowledge integration and alignment in the world-wise compound attention, WorMI sequentially utilizes partially learned task information, achieving high performance even on unseen tasks.
D.2. Multi-modal agents
We consider framework is designed to support this by allowing the reasoning model and individual world models to come from different modalities.
is critical for real robotic systems, which often rely on various sensor inputs. Our WorMI
The table below shows the performance of multi-modal WorMI, which employs VLM as its reasoning model, using both text and image states in VirtualHome. Multi-modal WorMI exhibits only slight performance drop compared to WorMI, demonstrating the applicability for multi-modal experiment setups. Additionally, there is certainly room for improvement of Multi-modal WorMI, as we do not have enough time to optimize the hyperparameters. Even so, our approach still demonstrates superior performance compared to the baselines. We will include these experimental results in the final version.
Table A.7: Performance comparisons for Multi-modaa WorMI and WorMI
Model
SR (t)
PS (L)
Multi-modal WorMI WorMI
57.65%
17.21
66.12%
15.17
D.3. Analysis of resource usage
Below table shows the inference times and memory usage among the baselines and our WorMI. We use LLaMA-3.2-111 as the reasoning model and LLaMA-3.2-1 -BB2-1B for the world models. The same model configurations are also used for all baselines. L.LM-Planner requires a longer prompt for in-context examples, leading to increased inference time. SayCanPay utilizes three separate models and repetitive infers for the action log probabilities, which need further inference time. WorMI uses relatively smaller world models and selects only the K most relevant ones rather than using all them. Additionally, WorMI keeps each domain-specift world model much smaller than the main reasoning model, making it easier to scale up the number of world models.
Table A.8: Inference time and memory for comparisons and WorMI
Inference Time
Memory
LLM+FT, ZSP
298
21877 MiB
LLM-Plannes
401 ms
21877 MiB
SayCanPay
609 ms
46230 MiB
WorMi (K=2, N=4)
339 ms
30020 MiB
WorMI (K=2, N=6)
348 ms
33445 MiB
WorMI (K=3, N=6)
385
33445 MiB
D.4. Robustness and scalable analysis for prototype-based retrieval
In Table 3(a), we compare a random selection strategy (WorMI-R) to our prototype-based retrieval approach, illustrating that incorrect retrieval slightly degrades performance. To investigate this more thoroughly, conducted two additional experiments. First, in Table A.9 we increased the proportion of adversarial world models by replacing some with untrained Llama-3.2- B. models. With prototypes unchanged, the retrieval could not distinguish adversarial from valid models. At lower proportions, performance stayed relatively stable, but it declined sharply once the adversarial ratio exceeded a certain threshold. This indicates that our compound attention helps filter out adversarial models, maintaining robustness unless a critical mass of them is adversarial.
17

---

World Model Implanting for Test-time Adaptation of Embodied Agents
Table A.9: Robustness of prototype-based retrieval under adversarial world model settings
Ady. ratio
SR (1)
PS (J)
66.12%
15.17
16%
66.67%
14.96
33%
58.58%
16.00
50%
38.03%
21.42
We also scaled the number of world models from N=6 N=12. As shown in Table A.10, WorMI-R suffers more from incorrect retrieval N grows, whereas prototype-based retrieval remains relatively resilient.
Table A.10: Scalability of prototype-ber baseee retrieval across varying world model pool sizes.
Model
SR (1)
PS (1)
WorMI-R (N=6)
62.04%
16.96
WorMI (N=6)
66.12%
15.17
WorMI-R (N=12) WorMI (N=12)
51.17%
19.22
66.51 %
14.90
Overall, these results confirm that suboptimal retrieval or misleading world models can lower performance, but WorMI's world-to-world knowledge integration in compound attention provides robust outcomes. Although adversarial models do pose a risk, our retrieval method and compound attention mitigate their impact, unless the majority of models are adversarial.
D.5. Analysis prototypes world model retrieval
The Table A. shows the results comparing prototype retrieval and full retrieval (WorMI-P). The performance difference between the two minimal, yet prototype retrieval significantly reduces inference time.
Table A.11: Effect of prototype on world model retrieval
Model
SR (t)
PS ())
Inference Time
WorMI-P
66.54%
15.04
81
WorMI
66.12%
15.17
385ms
D.6. Analysis of multiple world models
The Table A.12 shows that performance comparison between WorMI and the variant that uses only the single most relevant world model without fusion (WorMI-F). For seen tasks and scenes, a single world model suffices as it captures the domain knowledge. In contrast, for unseen tasks and scenes, combining multiple world models via world-to-woff integration advantageous since single model fully represents the unseen domain.
Table A. 12: Effect of multiple world models in model implanting
Model
SR (1)
PS (L)
Unseen tasks & scenes
WorMI-F 42.75%
20.54
WorMI
62.12%
15.17
Seen tasks & scenes
WorMI-F 82.72%
11.16
WorMI 85.78%
10.76
18


Reviews:
[{"review_id": "TosfGiQM7k", "reviewer": "Reviewer_MES3", "summary": "This paper presents WorMI, a framework enabling embodied agents to adapt to new domains at test time by combining large language models with domain-specific world models. It introduces prototype-based world model retrieval and world-wise compound attention to effectively integrate knowledge from multiple models. Experiments show WorMI outperforms existing methods in zero-shot and few-shot scenarios, demonstrating robust adaptation to unseen domains. The framework's design allows for scalable, efficient deployment in real-world settings where adaptability and data efficiency are crucial.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your detailed comments. We will include the following experimental results and clarifications in the final version.\n\n> Q1. While the combination of methods is novel, the individual components (prototype-based retrieval, attention mechanisms, meta-learning) are not entirely new. Some readers might argue that the innovation is incremental rather than groundbreaking. \n\nTo the best of our knowledge, our model implanting framework is the first to enable selective addition and removal of domain-specific world models in an agent’s policy at test time. \nThe novelty lies in how we orchestrate them to enable a plug-and-play composition of world models.\nRather than training a single monolithic model or simply combining multiple models externally, we propose a framework where each world model is fully implanted in the reasoning model only when relevant to the current domain.\n\nOur approach is novel in its three-fold strategy. First, prototype-based retrieval efficiently selects only the most relevant world models for an unseen domain, reducing overhead. \nSecond, our compound attention mechanism fuses domain-specific knowledge and aligns it with LLM reasoning at test time. Finally, meta-learning enables seamless adaptation to new models or domains without retraining components.\nBy combining these three ideas, we achieve a flexible test-time architecture that is easily extensible and maintainable, providing a clear step forward over existing methods that rely solely on in-context learning or static model ensembles.\n\n> Q2. The framework's performance is inherently dependent on the underlying language model, which could be seen as a limitation since it inherits any weaknesses of the LLM.\n\nWorMI leverages domain-specific world models to reduce hallucination and bias instead of relying solely on the LLM’s own knowledge. As shown in Table 4, it remains robust even with a smaller LLM while keeping expansion costs low. However, completely eliminating LLM weaknesses remains challenging. These limitations are explicitly addressed in the main text.\n\n> Q3. What specific challenges have you identified in deploying WorMI in real robotic systems with sensorimotor embodiments?\n\nWe consider multi-modality is critical for real robotic systems, which often rely on various sensor inputs. Our WorMI framework is designed to support this by allowing the reasoning model and individual world models to come from different modalities. \n\nThe table below shows the performance of multi-modal WorMI, which employs a VLM as its reasoning model, using both text and image states in VirtualHome.\nMulti-modal WorMI exhibits only a slight performance drop compared to WorMI, demonstrating the applicability for multi-modal experiment setups.\nAdditionally, there is certainly room for improvement of Multi-modal WorMI, as we do not have enough time to optimize the hyperparameters.\nEven so, our approach still demonstrates superior performance compared to the baselines.\nWe will include these experimental results in the final version.\n\n| Model             | SR (↑) | PS (↓) |\n|------------------|--------|--------|\n| Multi-modal WorMI| 57.65% | 17.21  |\n| WorMI            | 66.12% | 15.17  |", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "The paper introduces WorMI, a framework that advances embodied AI by enabling test-time adaptation through dynamic composition of domain-specific world models with LLMs, building upon and extending prior work in model composition, knowledge retrieval, and meta-learning. The novel prototype-based retrieval and compound attention mechanisms in WorMI efficiently select and integrate relevant models, addressing limitations in computational efficiency and adaptability found in previous approaches, and demonstrate superior performance in zero-shot and few-shot scenarios across established benchmarks.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n1) The paper presents a novel framework (WorMI) that creatively combines prototype-based retrieval with a world-wise compound attention mechanism for embodied AI adaptation. This approach addresses a significant challenge in the field by enabling dynamic composition of world models at test time, representing a meaningful advancement over previous methods.\n2)The work tackles a crucial problem for real-world embodied AI applications—adaptation to new domains without extensive retraining. The demonstrated performance improvements, particularly in zero-shot and few-shot scenarios, suggest practical impact and scalability.\n3)The paper is well-structured.\n\n\nWeaknesses:\n1)While the combination of methods is novel, the individual components (prototype-based retrieval, attention mechanisms, meta-learning) are not entirely new. Some readers might argue that the innovation is incremental rather than groundbreaking.\n2)The framework's performance is inherently dependent on the underlying language model, which could be seen as a limitation since it inherits any weaknesses of the LLM.\n3) What specific challenges have you identified in deploying WorMI in real robotic systems with sensorimotor embodiments?", "other_comments_or_suggestions": "See the weaknesses", "questions_for_authors": "See the weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "waPmd7VLtb", "reviewer": "Reviewer_zfbw", "summary": "This paper presents World Model Implanting (WorMI), a framework to improve the test-time adaptation of embodied AI agents. This work assumes access to a set of world models which are pre-trained on a set of datasets. During the adaptation phases, they select a subset of models that are most relevant to the input state (done via prototype-based retrieval), and then they perform \"world-wise\" hierarchical cross attention that is the input to the reasoning model.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your detailed comments. We will include the following experimental results and clarifications in the final version.\n> Q1. How are the \"world models\" trained?\n\nFollowing the reviewer’s comment, we clarify that the world models imitate the environment by capturing dynamics and affordances. \nEach world model is fine-tuned with LLaMA-3.2-1B using text-based states $s_t$, actions $a_t$, and instructions $I$. Specifically, we use three auxiliary tasks: (1) predict $s_{t+1}$ from $(s_t, a_t)$ to learn transition dynamics; (2) identify feasible actions $a_t$ from $s_t$ to capture affordances; (3) predict $a_t$ from $(s_t, I)$ to account for instruction conditioning.\nThe dataset is collected only from seen scenes, and instructions are sampled solely from seen tasks. Training prompts for each world model build on the environment prompts in the appendix, with added instructions for each auxiliary task.\n\n> Q2. What is $I$ in $D_j=(I,s_t,a_t,s_{t+1})$ denote?\n\n$I$ is the embodied task instructions. We will add the descriptions in the main text Line 146.\n\n> Q3. Comparison for the inference time and memory efficiency.\n\nThe table below shows the inference times and memory usages among the baselines and our WorMI. \nWe use LLaMA-3.2-11B for reasoning and LLaMA-3.2-1B for world models, with the same configurations for all baselines. LLM-Planner’s long in-context prompts increase inference time, and SayCanPay’s use of three models with repetitive action log probability computations further slows performance.\nWorMI uses smaller world models and selects only the K world models, with each domain-specific world model kept much smaller than the reasoning model, easing scalability.\n\n|Model|Inference Time|Memory|\n|-|-|-|\n|LLM+FT, ZSP|298ms|21877MiB|\n|LLM-Planner|401ms|21877MiB|\n|SayCanPay|609ms|46230MiB|\n|WorMI(K=2,N=4)|339ms|30020MiB|\n|WorMI(K=2,N=6)|348ms|33445MiB|\n|WorMI(K=3,N=6)|385ms|33445MiB|\n\n> Q4. For figure 5(a) what are WM {1-6} trained on?\n\nWM1 to WM6 represent six pre-trained world models, with each model trained on a dataset from a distinct domain as demonstrated in Figure 5(b) showing Domain 1 through Domain 6.\n\n> Q5. Exploration capability for WorMI\n\nWorMI is capable of exploration even if the object is located outside its seen domain. \nFor instance, if the TV is usually in the living room but is found instead in the bedroom or kitchen, it checks that likely location, and if the TV is missing there, it searches other rooms. \nThis works because, beyond each world model’s domain-specific knowledge, the world-to-reasoning alignment also leverages the reasoning model's general knowledge, letting the agent systematically explore instead of relying exclusively on its seen domains. \n\n> Q6. Clarification for object detection module $\\phi_D$.\n\nWe use the environment’s built-in object detection, which is error-free. \nTherefore, errors due to missed or incorrect object recognition do not affect our method or any of the baselines.\nThis setup aligns with text-based simulation environments such as VirtualHome and ALFWorld, where object descriptions are provided without recognition errors, and is consistently applied in other comparison studies as well [1, 2, 3, 4, 5].\n\nIn addition, for real-world scenarios requiring visual input, our framework can integrate an object detection module $\\phi_D$ by using a vision-language model as its reasoning component, extending the approach beyond purely text-based domains.\n\nThe table below shows the performance of multi-modal WorMI, which employs a VLM as its reasoning model, using both text and image states in VirtualHome.\nMulti-modal WorMI exhibits only a slight performance drop compared to WorMI, demonstrating the applicability for  multi-modal experiment setups.\nAdditionally, there is certainly room for improvement of Multi-modal WorMI, as we do not have enough time to optimize the hyperparameters.\nEven so, our approach still demonstrates superior performance compared to the baselines.\n\n|Model|SR (↑)|PS (↓)|\n|-|-|-|\n|Multi-modal WorMI|57.65%|17.21|\n|WorMI|66.12%|15.17|\n\n[1] Huang et al. \"Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.\" ICML 2022.\n\n[2] Song et al. \"Llm-planner: Few-shot grounded planning for embodied agents with large language models.\" ICCV 2023.\n\n[3] Hazra et al. \"Saycanpay: Heuristic planning with large language models using learnable domain knowledge.\" AAAI 2024.\n\n[4] Singh  et al. \"Progprompt: Generating situated robot task plans using large language models.\", ICRA 2023.\n\n[5] Yoo et al. \"Exploratory retrieval-augmented planning for continual embodied instruction following.\", NeurIPS 2024.\n\n> Q7. Clarification for the configuration in Table 3(a)\n\nYes. All variants in Table 3(a) use the same world-wise compound attention during training. They only differ in which world models are selected at test time. WorMI-E uses all models, while WorMI-R randomly picks a subset, and WorMI does prototype-based retrieval.", "claims_and_evidence": "Yes. All the claims have experiments to support them.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No significant theory (mostly an empirical work). One exception is the bound for prototype-based similarity that is discussed in the appendix. I have not given much importance to that part of the paper.", "experimental_designs_or_analyses": "Yes, I have checked the soundness and validity of the experiments. No critical issues in experimentation.", "relation_to_broader_scientific_literature": "This work is very much related to the field of Embodied AI, and agents being able to adapt in their environment is a very critical problem to address since they are bound to fail to some new unseen environments/tasks.", "essential_references_not_discussed": "Not to the best of my knowledge.", "other_strengths_and_weaknesses": "**Strengths**:\n\n1. The paper shows really strong few-shot results in unseen environments as well as on unseen tasks in Alfred and VirtualHome embodied environments.\n\n**Weaknesses**: \n\n1. I find the notion of the world model in the paper to be misleading as the term is predominantly used to mean a transition dynamics model with an optional reward model. However, there are no details regarding the World model training or its description apart from one small paragraph at the beginning of Sec 3.2. \n\n2. The writing of the paper needs to be improved. Lots of missing definitions and missing details (mentioned above as well as in other comments) make me question the reproducibility of this work.", "other_comments_or_suggestions": "No specific typo that I could find.", "questions_for_authors": "1. How are the \"world models\" trained?\n\n2. What is $I$ in $D_j = {(I, s_t, a_t, s_{t+1})}$ denote?\n\n3. Is it possible to compare the inference time for WorMI and the baselines? I'm curious to see how much time does the selection of relevant $K$ models takes.\n\n4. For figure 5(a) what are WM {1-6} trained on? \n\n5. I am curious if WorMI can adapt to finding objects that are not typically in the \"desired\" location. For instance for the task of \"Place breadslice in microwave,\" the agent's attention is over the kitchen as it is the most likely place to be -- however if the microwave (for whatever) reason is in the living room/bedroom - does the agent have the capability to explore?\n\n6. What object detection module $\\phi_D$ is used? If it's a learned model -- how many of the errors in the final performance are due to either misrecognition or missing the recognitions?\n\n7. [Clarification question]: For the results in Table 3 (a), are the models still trained with World-wise compound attention and only differ in the number of models used to perform the retrieval?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vP62BJOPxZ", "reviewer": "Reviewer_SQ6Q", "summary": "The paper introduces WorMI, a framework designed to improve the adaptability of embodied AI agents across diverse and unseen domains at test time, without requiring extensive retraining or additional data collection. Experiments on the VirtualHome and ALFWorld benchmarks demonstrate that WorMI outperforms state-of-the-art LLM-based embodied agents in zero-shot and few-shot adaptation. The framework also supports continual model implanting, allowing new world models to be added or removed flexibly.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your detailed comments. We will include references, the following experimental results, and additional clarifications in the final version.\n> Q1. Experimental results for inference time and memory usage\n\nBelow shows the inference times and memory usages among the baselines and WorMI.\nA detailed explanation is provided in response to Reviewer zfbw, Q3.\n| |Inference Time|Memory|\n|-|-|-|\n|LLM+FT, ZSP|298ms|21877MiB|\n|LLM-Planner|401ms|21877MiB|\n|SayCanPay|609ms|46230MiB|\n|WorMI(K=2,N=4)|339ms|30020MiB|\n|WorMI(K=2,N=6)|348ms|33445MiB|\n|WorMI(K=3,N=6)|385ms|33445MiB|\n> Q2. Clarification on the limitations due to utilizing LLMs\n\nWorMI leverages domain-specific world models to reduce hallucination and bias instead of relying solely on the LLM’s own knowledge. As shown in Table 4, it remains robust even with a smaller LLM while keeping expansion costs low. However, completely eliminating LLM weaknesses remains challenging. These limitations are explicitly addressed in the main text.\n> Q3. Experimental results for real-time environmental change\n\nWorMI composes pre-trained models at test time, forming a composite world model that adapts to unseen or dynamic environments.\nBelow are the results under real-time environment changes, such as shifting object locations or state changes over time.\n|Model|SR (↑)|PS (↓)|\n|-|-|-|\n|LLM-Planner|45.35%|20.16|\n|SayCanPay|41.33%|21.91|\n|WorMI|58.36%|19.33|\n> Q4. Robustness and scalable tests for prototype-based retrieval\n\nIn Table 3(a), we compare a random selection strategy (WorMI-R) to our prototype-based retrieval approach, illustrating that incorrect retrieval slightly degrades performance. To investigate this more thoroughly, we conducted two additional experiments.\n\nFirst, we increased the proportion of adversarial world models by replacing some with untrained Llama-3.2-1B models. With prototypes unchanged, the retrieval could not distinguish adversarial from valid models. At moderate proportions of adversarial world models, performance remains stable, but it drops sharply when their proportion becomes too high.\n|Adv. ratio|SR(↑)|PS(↓)|\n|-|-|-|\n|0%|66.12%|15.17|\n|16%|66.67%|14.96|\n|33%|58.58%|16.00|\n|50%|38.03%|21.42|\n\nWe also scaled the number of world models from N=6 to N=12. As shown in the table below, WorMI-R suffers more from incorrect retrieval as N grows, whereas our prototype-based retrieval remains relatively resilient.\n|Model|SR(↑)|PS(↓)|\n|-|-|-|\n|WorMI-R(N=6)|62.04%|16.96|\n|WorMI(N=6)|66.12%|15.17|\n|WorMI-R(N=12)|51.17%|19.22|\n|WorMI(N=12) |66.51%|14.90|\n\nOverall, these results show that while suboptimal retrieval or misleading world models can reduce performance, WorMI’s compound attention integrating world-to-world knowledge remains robust, unless most models are adversarial. We will include these results, along with an experiment using more world models, in the final version.\n> Q5. Experimental results comparing prototype retrieval vs. full retrieval\n\nThe table below shows the results comparing prototype retrieval and full retrieval (WorMI-P). There is almost no performance difference, but prototype retrieval significantly reduces inference time.\n|Model|SR(↑)|PS(↓)|Inference Time|\n|-|-|-|-|\n|WorMI-P| 66.54%|15.04|811ms|\n|WorMI|66.12%|15.17|385ms|\n> Q6. Qualitative analysis\n\nOne notable failure case arises when there is no relevant world model for the unseen domain. \nFor instance, if the instruction is “Place the slipper in the closet\", but the agent has never encountered a slipper in any seen domain, the success rate drops and the agent might attempt actions for a similar object (e.g., “towel”).\nIt is because world-to-reasoning alignment training encourages reliance on the knowledge derived from the world model, the agent struggles in scenarios where it must rely solely on the reasoning model’s common knowledge. \nIn future work, we plan to explore more flexible architectures and training methods for controlling the degree of utilizing world models.\n> Q7. Comparison with using most relevant world models\n\nThe table below shows that a performance comparison between WorMI and the variant that uses only the single most relevant world model without fusion (WorMI-F).\nFor seen tasks and scenes, a single world model suffices as it captures the domain knowledge. In contrast, for unseen tasks and scenes, combining multiple world models via world-to-world integration is advantageous since no single model fully represents the unseen domain.\n|Model|SR(↑)|PS(↓)|\n|-|-|-|\n|Unseen tasks & scenes|\n|WorMI-F|42.75%|20.54|\n|WorMI|66.12%|15.17|\n|Seen tasks & scenes|\n|WorMI-F|82.72%|11.16|\n|WorMI|85.78%|10.76|\n> Q8. Challenges for real world settings\n\nWe consider multi-modality is critical for real robotic systems, which often rely on various sensor inputs. Our WorMI framework is designed to support this by allowing the reasoning model and individual world models to come from different modalities. We add a multi-modal experiment in response to Reviewer MES3, Q3.", "claims_and_evidence": "WorMI’s effectiveness in zero-shot and few-shot adaptation is strongly supported. In zero-shot tests, WorMI outperforms all baselines on unseen tasks and scenes.\n\nThe paper also claims advantages of the prototype-based world model retrieval mechanism, which selects relevant domain-specific models at test time. This claim is backed by an ablation study: using WorMI’s prototype retrieval to pick a few pertinent world models yields better results than naive alternatives.\n\nThe authors further claim that WorMI’s design generalizes across unseen domains. This is reflected in the experimental setup: in VirtualHome and ALFWorld, “unseen” test scenarios involve new tasks and/or environments not encountered in training.", "methods_and_evaluation_criteria": "The WorMI framework logically addresses the problem of test-time adaptation for embodied agents. It does so by allowing the agent to dynamically retrieve and integrate domain-specific knowledge at inference time, rather than requiring retraining.\n\nHowever, there are several limitations of the proposed method.\n\nIt introduces computational overhead due to retrieving and integrating multiple world models at test time, which may not scale well in resource-constrained settings. \n\nThe framework heavily depends on large language models, making its performance sensitive to the LLM’s limitations, such as biases and hallucinations. \n\nAdaptation to dynamic environments is limited, as the approach primarily focuses on structured, predefined world models rather than handling real-time environmental changes.\n\nThe lack of explicit robustness tests makes it unclear how well WorMI handles incorrect retrieval or misleading world models.", "theoretical_claims": "The paper includes a key theoretical component related to prototype-based world model retrieval. It provides a proof that the distance between prototype sets can serve as a bounded proxy for the distance between full datasets. While the proof shows a bound, it does not quantify how much information is lost by using prototypes instead of full embeddings. The paper does not provide experimental results comparing prototype retrieval vs. full retrieval to validate this bound in practice.", "experimental_designs_or_analyses": "The experimental setup appears mostly sound and well-motivated for the problem of test-time adaptation in embodied AI. The study evaluates WorMI on VirtualHome (a 3D household task simulation) and ALFWorld (a text-based embodied environment). These benchmarks are reasonable choices because they provide structured tasks, unseen environments, and diverse domain shifts. However, both are simulations, meaning results may not fully translate to real-world robotics applications.\n\nThe paper compares WorMI against four baselines, covering a range of adaptation approaches. These baselines are appropriate and diverse, making the comparisons meaningful.\n\nThe experiments are well-described and appear repeatable, as dataset splits, baselines, and architectures are clearly outlined.", "relation_to_broader_scientific_literature": "The concept of world models has been widely explored in robotics and reinforcement learning. Prior work, such as World Models by Ha & Schmidhuber (2018), introduced learned world models to enable agents to simulate future states and make more efficient decisions.\nThe WorMI framework extends world models by allowing multiple world models to be retrieved and fused dynamically at test time. Unlike prior methods where a single learned world model guides decision-making, WorMI selects the most relevant world models per task, combining them with an LLM for reasoning.\n\nTest-time adaptation aims to help models generalize to new environments without retraining. WorMI introduces a retrieval-based approach to test-time adaptation. Instead of updating a model’s parameters during test time, it retrieves and fuses world models dynamically, reducing the need for computationally expensive online updates.\n\nWorMI follows a modular learning paradigm by treating each world model as a domain-specific knowledge module that can be combined dynamically. (Modular Multitask Reinforcement Learning with Policy Sketches)", "essential_references_not_discussed": "Some test-time adaptation papers should be added to the related work section for a more comprehensive discussion:\n\nTest-time adaptation: Tent: Fully Test-Time Adaptation by Entropy Minimization\nTest-Time Classifier Adjustment Module for Model-Agnostic Domain Generalization\nTest-Time Training with Self-Supervision for Generalization under Distribution Shifts", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper presents an innovative approach to test-time adaptation by dynamically retrieving and fusing world models to enhance embodied agent reasoning. \n\nThe empirical results demonstrate strong improvements over baselines in both zero-shot and few-shot adaptation scenarios.\n\nWeaknesses:\n\nComputational or memory efficiency is not analyzed, even though WorMI requires additional computation for retrieval and fusion.\n\nThe study does not provide qualitative insights into failure cases, making it unclear why adaptation fails in some scenarios.\n\nThe paper claims that WorMI is scalable, but it only evaluates retrieval from a small number of world models. It does not test how well the approach scales when hundreds of world models must be retrieved dynamically.\n\nThere is no explicit analysis of what happens when the retrieval mechanism selects irrelevant world models, which could be a key failure mode in complex environments.\n\nThe ablations confirm that both prototype retrieval and compound attention contribute positively, but they do not explore simpler alternatives like selecting the most relevant world model without fusion, which could reveal whether the full attention mechanism is necessary.", "other_comments_or_suggestions": "The paper is well-written, with clear explanations and a logical structure. The writing is concise and easy to follow, making the technical contributions accessible. \n\nThe figures are well-designed and effectively illustrate key concepts, particularly the retrieval mechanism and world-wise compound attention.", "questions_for_authors": "How well does WorMI scale when retrieving from a large number of world models (e.g., 50–100 instead of ≤6)? Have you tested how retrieval accuracy or computational cost changes as the number of world models increases?\n\nWhat happens when the retrieval mechanism selects an irrelevant or suboptimal world model? Is there a mechanism for detecting and correcting retrieval errors at test time?\n\nHave you considered evaluating WorMI in a real-world robotic setting instead of simulation? What are the main challenges in transferring the approach from VirtualHome and ALFWorld to physical environments?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3hFmnakVoa", "reviewer": "Reviewer_sgRi", "summary": "This paper proposes a world model implanting framework to augment the LLM-based agents. The world models are learned using domain/task-specific data to capture the in-domain characteristics, serving as domain experts. With these world models, this paper introduces a prototype-based retrieval method together with an attention-based knowledge integration to allow test-time adaptation. This framework also applies a meta-learning objective to improve the adaptability to unseen tasks/domains. Experiments are conducted on Virtual-Home and ALFWorld benchmarks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your detailed comments. We will include the following experimental results and clarifications in the final version.\n\n> Q1. How does this method guarantee that the knowledge of the seen world models will benefit the target task? If the target domain/task is highly out-of-the-training distributions, how does this paper handle this case?\n\nOur method guarantees that the knowledge of seen world models benefits the target task by combining two key components. \nFirst, the prototype-based retrieval selects only those world models that were trained on data with distributions similar to the target environment, ensuring that the most relevant domain-specific information is leveraged. \nSecond, our compound attention model integrates the selected world models with the reasoning model on two levels, world-level integration and reasoning-level alignment, while meta-learning is used to adapt effectively to diverse world model combinations. This process not only creates a robust, domain-specific composite model but also aligns the integrated knowledge with the reasoning model to enhance decision-making.\nThis approach is validated by our results in Table 1, where our method outperforms baselines even in unseen domains, confirming that the aligned knowledge is indeed beneficial for the target task.\n\nMoreover, if the environment is completely unrelated to any available world model, the agent can still utilize the reasoning model’s common knowledge via world-to-reasoning alignment, maintaining at least a baseline level of performance. Please refer to our response to Reviewer SQ6Q, Q4 for details on the adversarial world models experiments, which involve models that are highly out-of-training distribution.\nOf course, in highly out-of-distribution scenarios that lie well beyond our assumptions, the success rate may drop significantly.\n\n> Q2. It might be risky to discard crucial visual/structural information about the environment. In addition, without being grounded in the visual environment, how does this paper ensure the LLM provides accurate actions without potential hallucination issues?\n\nIn VirtualHome and ALFWorld, the environment provides text-based observations that include visual and structural information for performing tasks. \nThese text-based outputs have been widely used in prior work [1, 2, 3, 4, 5] and the agent can select suitable actions without discarding crucial spatial or visual context. \nMoreover, our WorMI framework is designed to be modular with respect to input modalities: the reasoning model and individual world models can each handle different data forms.\n\nThe table below shows the performance of multi-modal WorMI, which employs a VLM as its reasoning model, using both text and image states in VirtualHome.\nMulti-modal WorMI exhibits only a slight performance drop compared to WorMI, demonstrating the applicability for multi-modal experiment setups.\nAdditionally, there is certainly room for improvement of Multi-modal WorMI, as we do not have enough time to optimize the hyperparameters.\nEven so, our approach still demonstrates superior performance compared to the baselines.\nWe will include these experimental results in the final version.\n\n| Model             | SR (↑) | PS (↓) |\n|------------------|--------|--------|\n| Multi-modal WorMI| 57.65% | 17.21  |\n| WorMI            | 66.12% | 15.17  |\n\n[1] Huang, Wenlong, et al. \"Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.\" ICML 2022.\n\n[2] Song, Chan Hee, et al. \"Llm-planner: Few-shot grounded planning for embodied agents with large language models.\" ICCV 2023.\n\n[3] Hazra, Rishi, Pedro Zuidberg Dos Martires, and Luc De Raedt. \"Saycanpay: Heuristic planning with large language models using learnable domain knowledge.\" AAAI 2024.\n\n[4] Singh, Ishika, et al. \"Progprompt: Generating situated robot task plans using large language models.\", ICRA 2023.\n\n[5] Yoo, Minjong, et al. \"Exploratory retrieval-augmented planning for continual embodied instruction following.\", NeurIPS 2024.\n\n> Q3.  For the meta-learning, this paper does not provide an ablation study or analysis to assess the effectiveness of this objective. The improvement of generalizability contributed by meta-learning remains unclear.\n\nThe table below shows an ablation study on meta-learning for unseen tasks and scenes. WorMI-M is a variant that learns world model composition sequentially instead of using meta-learning. \nCompared to the WorMI, WorMI-M exhibits lower performance. This indicates that meta-learning equips our framework to better handle world model combinations it may not have encountered during training, thereby improving generalizability.\nWe will include these results in the final version.\n\n| Model    | SR (↑) | PS (↓) |\n|----------|--------|--------|\n| WorMI-M  | 53.31% | 18.23  |\n| WorMI    | 66.12% | 15.17  |", "claims_and_evidence": "This paper claims that the proposed WorMI framework is able to generalize to unseen domains in zero-shot or few-shot manners by leveraging the most relevant information from the seen world models. However, how does this method guarantee that the knowledge of the seen world models will benefit the target task? If the target domain/task is highly out-of-the-training distributions, how does this paper handle this case?", "methods_and_evaluation_criteria": "The proposed method applies off-the-shelf models to extract abstract state features. However, it might be risky to discard crucial visual/structural information about the environment. In addition, without being grounded in the visual environment, how does this paper ensure the LLM provides accurate actions without potential hallucination issues?", "theoretical_claims": "No theoretical claims are provided in the main paper.", "experimental_designs_or_analyses": "I have checked the Experiments section. For the meta-learning, this paper does not provide an ablation study or analysis to assess the effectiveness of this objective. The improvement of generalizability contributed by meta-learning remains unclear.", "relation_to_broader_scientific_literature": "The main contribution compared with previous LLM-based embodied agents is the integration of the domain-specific models, which provide domain knowledge to augment the LLM.", "essential_references_not_discussed": "To the best of my knowledge, the references are sufficiently covered.", "other_strengths_and_weaknesses": "Overall, this paper is well-written and easy to follow. For the weaknesses, please refer to the issues raised above.", "other_comments_or_suggestions": "Please refer to the issues raised above.", "questions_for_authors": "Please refer to the issues raised above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this work, the authors propose WorMI, a world model implanting framework that can leverage a set of pre-trained domain-specific (neural) world models during test-time by fusing their representations with reasoning LLMs. The main contributions include a prototype-based retrieval mechanism, which selects relevant world models from a set of candidate based on the specific test domain; a compound attention mechanism that integrates the selected world models with the reasoning model during test-time. 

All reviewers agree that the WorMI framework is novel; the idea of composing pretrianed world models and fusing with reasoning LLMs is interesting; the prototype-based retrieval mechanism and the compound attention mechanism make sense and are well designed. 

The reviewers' concerns were mostly regarding the insufficient experiments and analysis. The authors did a good job addressing the concerns by providing a set of additional experiments, new results and some new qualitative analysis have indeed strengthened the work. 

---

Additional comments from AC:

1. I agree with Reviewer sgRi that the work could benefit from more insights/discussions on leveraging (highly) out-of-distribution world models, for example, can agents use a set of VirtualHome world models in a ALFWorld environment, or can agents use a set of world models pretrained on a wide spectrum of embodied tasks in an unseen benchmark? I acknowledge that there are already "ood" test sets in benchmarks like ALFWorld, but I believe a more realistic setting would be trying to leverage pre-trained world models on completely unseen scenarios.
2. Related to the above point, it may be worth investigating how the diversity of pre-trained world models affects WorMI's performance, especially on ood settings. I believe to fully unleash the potentials of test-time world model fusing, one might need to have a large and at the same time very diverse set of world models. 

---

AC Recommendation:

The initial scores were 3/2/3/2 (avg 2.5); after the rebuttal, one reviewer increased their score, resulting 3/2/3/3 (avg 2.75). It's worth noting that the reviewer who gave the remaining score of 2 only acknowledged but did not respond to the authors, which may suggest they do not think the rebuttal is convincing enough. 
To me this is a good example where the reviewers helped the authors to improve their work by suggesting reasonable additional experiments, the authors also did well on addressing these requests.
Decision: Accept (poster)