WorMI: World Model Implanting Framework

• WorMI is a test-time adaptation framework that integrates domain-specific world models with a fixed LLM, enabling cross-domain adaptability.
• It utilizes a prototype-based retrieval and compound attention mechanism to dynamically implant relevant models without retraining the core reasoning policy.
• Empirical evaluations on VirtualHome and ALFWorld demonstrate significant improvements in success rates and efficiency, highlighting its parameter‐efficient meta-learning approach.

World Model Implanting Framework (WorMI) is a test-time adaptation framework for embodied agents that combines the reasoning capabilities of LLMs with independently learned, domain-specific world models through test-time composition, allowing seamless implantation and removal of the world models so that an embodied agent’s policy achieves and maintains cross-domain adaptability (Yoo et al., 4 Sep 2025). Its central premise is that cross-domain adaptation can be organized as retrieval and composition of specialized world models around a fixed reasoning policy, rather than as repeated retraining of a monolithic planner. In that sense, WorMI belongs to a broader conception of world models as reusable system-level components for perception, interaction, reasoning, and memory, rather than as isolated injections of “world knowledge” into narrowly defined tasks (Yoo et al., 4 Sep 2025, Zeng et al., 2 Feb 2026).

1. Problem setting and conceptual scope

WorMI is formulated for embodied agents that must follow instructions and act in environments such as VirtualHome and ALFWorld, where novel tasks and scenes appear at test time (Yoo et al., 4 Sep 2025). The motivating problem is how an embodied agent can adapt to unseen domains with no or very little new data and without retraining large models. In the formulation used by WorMI, each domain $D_j$ has a dataset

$\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$

where $I$ is a natural language task instruction, $s_t$ is the environment state at time $t$, $a_t$ is the action, and $s_{t+1}$ is the next state (Yoo et al., 4 Sep 2025).

Within this setting, a world model $M_j$ is a domain-specific neural model trained on $\mathcal{D}_j$ and supporting dynamics prediction, affordance prediction, and behavior cloning. WorMI does not merge these domain models into one retrained backbone. Instead, it maintains a library of independent world models $\{M_1,\dots,M_N\}$ and dynamically implants the most relevant subset into a fixed LLM-based reasoning policy at inference time (Yoo et al., 4 Sep 2025).

The framework also defines an operational distinction between implantation and removal. Implanting means retrieving the top-$\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$0 relevant world models for the current state and integrating their intermediate representations into the reasoning model through a learnable composition mechanism. Removal means excluding a model from the retrieved set and recomposing with the remaining models, without retraining the reasoning model or the world models themselves (Yoo et al., 4 Sep 2025). This modularity is central to the framework’s claim of continual model implanting and unlearning behavior.

A recurrent misconception in the surrounding literature is that world-model work is equivalent to adding more world-aware priors to a task-specific model. A broader design-spec view argues instead that a world model should be a system-level, normative framework that maintains internal state, supports interaction, reasoning, memory, and multimodal generation (Zeng et al., 2 Feb 2026). WorMI is consistent with that distinction insofar as it treats world models as reusable components with explicit roles in dynamics, affordances, and policy composition, rather than as prompt-time factual augmentations (Yoo et al., 4 Sep 2025).

2. Formal architecture and implanted policy

WorMI has four primary components: a reasoning model $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$1, a set of domain-specific world models $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$2, a prototype-based retrieval module, and a composition module $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$3 implemented as world-wise compound attention (Yoo et al., 4 Sep 2025). In the reported instantiation, the reasoning model is a fixed Llama‑3.2‑3B, while each world model is a fixed Llama‑3.2‑1B trained with auxiliary tasks for dynamics, affordance, and behavior cloning (Yoo et al., 4 Sep 2025).

The framework defines the implanted policy as

$\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$4

where the retrieved subset of $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$5 world models is composed with the frozen reasoning model (Yoo et al., 4 Sep 2025). At the level of internal states, if $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$6 denotes a world model layer output and $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$7 denotes the reasoning-model hidden state, composition modifies the reasoning model through

$\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$8

The final action distribution is obtained by decoding the modified hidden states through the reasoning model’s action head (Yoo et al., 4 Sep 2025).

This design makes the composition module the only trainable part during meta-learning. The reasoning model $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$9 remains frozen, and the world models $I$0 remain frozen after their own pretraining (Yoo et al., 4 Sep 2025). The implication is that adaptation is concentrated in a compact composer rather than in the full planner. This suggests a parameter-efficient route to cross-domain transfer, especially in settings where the cost and instability of full LLM fine-tuning are prohibitive.

The architecture also assigns distinct semantic roles to the components. The reasoning model provides generic high-level planning capacity. Each world model contributes domain-specific transition, affordance, and behavioral regularities. The compound attention module determines how those specialized regularities should be aligned with the reasoning state at the current step (Yoo et al., 4 Sep 2025).

3. Prototype-based world model retrieval

Retrieval determines which world models are implanted for the current state. WorMI constructs object-wise embeddings by first extracting object-wise states with an object detection or parsing model $I$1, then embedding each object with $I$2. For domain $I$3, the embedding set is

$I$4

Rather than comparing full embedding sets at test time, the framework compresses each domain into a prototype set using a $I$5-center objective (Yoo et al., 4 Sep 2025): $I$6

At inference time, the current state $I$7 is processed in the same way to obtain a prototype set $I$8. Retrieval then selects the top-$I$9 world models by Wasserstein distance: $s_t$0 The paper reports that $s_t$1 is about $s_t$2 of $s_t$3; for VirtualHome, $s_t$4 and $s_t$5 (Yoo et al., 4 Sep 2025).

The retrieval mechanism is justified by a boundedness argument. If $s_t$6 is the maximal distance from any element in $s_t$7 to its closest prototype in $s_t$8, then

$s_t$9

Accordingly, prototype similarity is a bounded proxy for full-set similarity (Yoo et al., 4 Sep 2025).

Empirically, retrieval is not a minor optimization but a substantive part of the framework. On VirtualHome unseen tasks and unseen scenes, using all world models every time yields $t$0 SR and $t$1 PS, random subset selection yields $t$2 SR and $t$3 PS, and prototype-based WorMI yields $t$4 SR and $t$5 PS (Yoo et al., 4 Sep 2025). The paper also reports robustness analyses in which prototype retrieval remains more stable than random selection as the world-model pool grows and as adversarial world models are introduced.

4. World-wise compound attention

The world-wise compound attention module $t$6 performs two distinct operations: world-level integration across retrieved world models, and reasoning-level alignment between the integrated world representation and the reasoning model’s hidden state (Yoo et al., 4 Sep 2025).

Because the world models and the reasoning model may have different hidden dimensions, each world-model representation is first projected: $t$7 World-level cross-attention $t$8 is then constructed with a query from the reasoning model and keys and values from the projected world-model outputs: $t$9

$a_t$0

This produces an integrated representation that weights the retrieved world models according to the current reasoning state (Yoo et al., 4 Sep 2025).

Reasoning-level cross-attention $a_t$1 then aligns that integrated representation back into the reasoning model: $a_t$2 The resulting aligned representation is the output of $a_t$3 that is added to the reasoning-model hidden state (Yoo et al., 4 Sep 2025).

The framework’s ablations indicate that this hierarchy matters. On VirtualHome unseen tasks and unseen scenes, WorMI‑CONCAT achieves $a_t$4 SR and $a_t$5 PS, WorMI‑ADD achieves $a_t$6 SR and $a_t$7 PS, and the full compound-attention WorMI achieves $a_t$8 SR and $a_t$9 PS (Yoo et al., 4 Sep 2025). The paper therefore attributes a material share of performance to the explicit weighting and alignment mechanism, rather than to the mere presence of multiple world-model streams.

The attention analysis is also presented as semantically meaningful. The reported visualizations show attention concentrating on the model that contains knowledge about a required object, shifting across models when object positions differ by domain, and dividing labor across subtasks in composite instructions such as “Place plum in cabinet” (Yoo et al., 4 Sep 2025). This suggests that the composition module is not merely blending features but performing task-dependent routing over domain knowledge.

5. Meta-learning, benchmarks, and empirical performance

Each world model $s_{t+1}$0 is trained separately on its domain with three auxiliary tasks: dynamics prediction $s_{t+1}$1, affordance prediction from $s_{t+1}$2, and behavior cloning $s_{t+1}$3 (Yoo et al., 4 Sep 2025). The compound attention module is then trained by a first-order meta-learning procedure similar to Reptile. For a sampled subset of world models $s_{t+1}$4 and the union of their datasets $s_{t+1}$5, the task-specific implanted policy is

$s_{t+1}$6

and the inner-loop behavior-cloning loss is

$s_{t+1}$7

After several inner-loop updates, the outer-loop meta-update is

$s_{t+1}$8

In zero-shot deployment, no parameters are updated. In few-shot deployment, only $s_{t+1}$9 is lightly fine-tuned; $M_j$0 and $M_j$1 remain fixed (Yoo et al., 4 Sep 2025).

The framework is evaluated on VirtualHome and ALFWorld. VirtualHome includes $M_j$2 house configurations and $M_j$3 instructions, with $M_j$4 episodes partitioned into seen and unseen tasks and scenes. ALFWorld includes $M_j$5 episodes and domain splits based on CL-ALFRED clustering over scene and task types (Yoo et al., 4 Sep 2025). Performance is reported using Success Rate (SR), Pending Steps (PS), and, for complex instructions, Goal-conditioned success (GC).

The principal zero-shot and few-shot comparisons against SayCanPay are as follows (Yoo et al., 4 Sep 2025):

Setting	SayCanPay	WorMI
VirtualHome, zero-shot, unseen tasks & unseen scenes	SR $M_j$6, PS $M_j$7	SR $M_j$8, PS $M_j$9
ALFWorld, zero-shot, unseen tasks & unseen scenes	SR $\mathcal{D}_j$0, PS $\mathcal{D}_j$1	SR $\mathcal{D}_j$2, PS $\mathcal{D}_j$3
VirtualHome, few-shot average over $\mathcal{D}_j$4 shot	SR $\mathcal{D}_j$5, PS $\mathcal{D}_j$6	SR $\mathcal{D}_j$7, PS $\mathcal{D}_j$8
ALFWorld, few-shot average	SR $\mathcal{D}_j$9, PS $\{M_1,\dots,M_N\}$0	SR $\{M_1,\dots,M_N\}$1, PS $\{M_1,\dots,M_N\}$2

The framework also reports performance under increased instruction complexity. For long-horizon instructions in VirtualHome, LLM-Planner achieves SR $\{M_1,\dots,M_N\}$3, GC $\{M_1,\dots,M_N\}$4, PS $\{M_1,\dots,M_N\}$5; SayCanPay achieves SR $\{M_1,\dots,M_N\}$6, GC $\{M_1,\dots,M_N\}$7, PS $\{M_1,\dots,M_N\}$8; and WorMI achieves SR $\{M_1,\dots,M_N\}$9, GC $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$00, PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$01 (Yoo et al., 4 Sep 2025). For multiple instructions, LLM-Planner achieves SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$02, GC $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$03, PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$04; SayCanPay achieves SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$05, GC $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$06, PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$07; and WorMI achieves SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$08, GC $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$09, PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$10 (Yoo et al., 4 Sep 2025). The paper interprets these results as evidence that fused world-model knowledge improves planning over extended and multi-goal horizons.

6. Modularity, interpretation, and limitations

WorMI presents world models as removable and composable modules rather than permanent parameters of a single policy. This is reflected in the reported continual model implanting and unlearning behavior: adding relevant world models incrementally improves performance, while removing a relevant world model causes performance drops “as if unlearning” the corresponding domain (Yoo et al., 4 Sep 2025). The framework therefore treats domain knowledge as an attachable library rather than as a globally entangled internal state of the planner.

At the same time, the framework identifies a practical upper bound on simultaneous fusion. On VirtualHome unseen tasks and unseen scenes, using one world model yields SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$11 and PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$12; two models yield SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$13 and PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$14; three models yield SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$15 and PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$16; four models yield SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$17 and PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$18; and six models yield SR $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$19 and PS $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$20 (Yoo et al., 4 Sep 2025). The paper accordingly argues that there is an optimal range, here $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$21–$\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$22 models, beyond which conflicting or noisy knowledge becomes detrimental even with prototype retrieval and compound attention.

The principal limitations identified by the authors are computational overhead, dependence on the reasoning LLM, and dependence on world-model coverage and retrieval quality (Yoo et al., 4 Sep 2025). Running several $\mathcal{D}_j = \{(I, s_t, a_t, s_{t+1})\},$23B world models plus an LLM at each step is heavier than a single-model policy. Performance inherits the strengths and weaknesses of the underlying LLM. If the world-model library does not include domains similar to the target, or if retrieval is misled, performance drops. The paper therefore suggests better world-model libraries, on-the-fly world-model construction, tighter LLM–world-model alignment, and more efficient retrieval and sparse fusion as future directions (Yoo et al., 4 Sep 2025).

In the broader world-model literature, these limitations connect to a continuing debate over what a world model should be. A normative view holds that a robust world model should not be a loose collection of capabilities but a framework that integrates interaction, perception, symbolic reasoning, and spatial representation (Zeng et al., 2 Feb 2026). WorMI does not claim to solve that full unification problem. Its contribution is narrower and more operational: it shows that independently trained, domain-specific world models can be retrieved and composed into a frozen reasoning policy at test time, and that this yields improved zero-shot and few-shot embodied adaptation across unseen domains (Yoo et al., 4 Sep 2025).