Pith. sign in

REVIEW 2 major objections 6 minor 28 references

World Model Implanting for Test-time Adaptation of Embodied Agents

T0 review · 2 major / 6 minor · reviewed 2026-08-05 · deepseek-v4-flash

Pith's one-line read By retrieving a few small domain-specific world models and fusing them into a frozen LLM at test time, WorMI achieves large zero-shot and few-shot gains in embodied task success, improving over SayCanPay by 20–27 percentage points.

desk verdict A practical, well-executed composition method for LLM-based embodied agents, hobbled by a bogus theory section that shouldn't mask the strong empirical work. read the letter →

arxiv 2509.03956 v1 pith:QL4L6H5A submitted 2025-09-04 cs.AI

classification cs.AI
keywords test-timeadaptationworldmodelsembodiedagentsLLMcompositionprototype-basedretrievalcompoundattentionzero-shotgeneralizationmeta-learning
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that an embodied agent can adapt to new, unseen environments without retraining by implanting a handful of domain-specific world models into a frozen large language model at test time. The central claim is that a prototype-based retrieval step—selecting the most relevant world models from the similarity of object-level state embeddings—combined with a learned 'compound attention' module that fuses and aligns those models' representations with the LLM's reasoning, yields large gains in zero-shot and few-shot task success. In VirtualHome, the framework (WorMI) improves average success rate over the SayCanPay baseline by 20.41 percentage points zero-shot and 26.58 points few-shot, with comparable gains in ALFWorld. If correct, this suggests a scalable recipe: keep the reasoning model fixed, acquire small per-domain models, and compose them on demand rather than fine-tuning the whole policy.

What carries the argument

The central mechanism is the 'world-wise compound attention' Cθ, a parameter-efficient module consisting of a linear projection, a world-level cross-attention layer, and a reasoning-level cross-attention layer. It takes the intermediate layer outputs of the K retrieved world models and the reasoning model's hidden state, and injects a fused, aligned representation into the next layer of the frozen LLM. The complementary machinery is the prototype-based retrieval: each domain's dataset is summarized by k cluster centers of object-wise embeddings, and at test time the agent retrieves the models whose prototype distance to the current observation is smallest under the Wasserstein metric. Togeth

What would settle it

Construct two domains with very similar object inventories but different transition dynamics (e.g., same objects, opposite affordances) and measure whether WorMI's prototype retrieval still selects the functionally correct world model. If prototype distance ranks the two domains as nearly identical, the zero-shot advantage over random retrieval should disappear; the paper's own adversarial-ratio experiment (Table A.9) already shows a collapse to 38% success at a 50% adversarial ratio.

Watch

Extended reading notes

Core claim

The discovery is that domain-specific world models—small models trained on transitions, affordances, and behavior cloning for one domain—can be treated as composable modules that plug into a frozen LLM-based policy. The framework's two-stage mechanism, world-to-world integration (via a world-level cross-attention that weighs the retrieved models) and world-to-reasoning alignment (via a reasoning-level cross-attention that conditions on the LLM's hidden states), is what lets the policy selectively use knowledge from multiple domains. The authors show that retrieval quality matters (random retrieval drops success from 66.1% to 62.0%, and using all models is worse than the retrieved three), and

Load-bearing premise

The retrieval step assumes that a small set of representative object-level prototypes of each domain preserves the true ordering of domain similarity, so that the most relevant world models are always selected; if the prototypes misorder domains, the framework's gains should vanish.

Editorial extensions

If this is right

  • Deployed agents can be given new capabilities by adding a small world model and retrieving it, without touching the base LLM's weights.
  • Because models can be removed as easily as added, the framework supports continual adaptation and machine unlearning of specific domains.
  • The number of simultaneously implanted models matters: the paper's experiments show peak performance at 2–4 models and a drop at 6, so selection is a real constraint, not a formality.
  • WorMI's advantage grows with the base LLM's size, indicating it leverages, rather than bypasses, the reasoning model.
  • The compound attention generalizes across model counts and even to a multimodal reasoning model with only a slight performance drop, per the appendix.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The retrieval's reliance on object-level embeddings suggests a testable boundary: two domains with similar objects but opposite physical dynamics might be misranked, so a dynamics-aware prototype (e.g., including transition embeddings) could be a stronger, more general selector.
  • A learned retriever trained end-to-end with the compound attention could replace the Wasserstein distance and remove the reliance on the quality of the offline clustering bound; this is not explored in the paper.
  • The paper's adversarial-model experiment implies that the compound attention filters out bad models up to a point; a safety-oriented deployment might add an explicit validity check that verifies a retrieved model's next-state predictions against the current observation before trusting it.
  • The k-center approximation bound in Eq. (4) is loose in the worst case, so the theoretical guarantee is weaker than the empirical evidence; a tighter, distribution-aware bound could make the framework's reliability claims more rigorous.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

2 major / 6 minor

Summary. The paper proposes WorMI, a framework for test-time adaptation of embodied agents. Domain-specific world models are first pre-trained on datasets from seen domains; at test time, a prototype-based retrieval step selects K relevant world models using Wasserstein distances between cluster prototypes of object-wise state embeddings. A 'world-wise compound attention' module then fuses the intermediate representations of the retrieved world models and aligns them with a frozen LLM-based reasoning model, with the composition module trained by REPTILE-style meta-learning. The method is evaluated on VirtualHome and ALFWorld in zero-shot and few-shot settings, reporting large improvements in success rate and pending steps over ZSP, LLM+FT, LLM-Planner, and SayCanPay, together with ablations, scalability studies, and attention-map analyses.

Significance. If the empirical findings hold, WorMI is a useful step toward modular, test-time composition of domain knowledge in embodied LLM agents. The paper reports 95% confidence intervals, includes ablations showing that both prototype-based retrieval and compound attention contribute to performance, and provides resource-usage comparisons. The gains over SayCanPay are substantial and consistent across two benchmarks and several analysis settings. The main weakness is that the paper's one formal justification for the prototype-based retrieval — the bound in Eq. (4) and Appendix A.1 — is mathematically unsound as written. Since the empirical ablations (Table A.11) show that prototype retrieval performs nearly as well as full retrieval, the central performance claim may still survive, but the theoretical claim needs correction or removal. Some implementation details that are central to reproducing the method are also missing.

major comments (2)
  1. [Section 3.2 / Eq. (4) / Appendix A.1] The claimed bound δ(p_i,p_j) ≤ δ(E_i,E_j) + 2ρ does not follow from the k-center condition. The proof assumes W(E_j,p_j) ≤ ρ, where p_j is the uniform measure on the k centers. But k-center clustering only guarantees that every point of E_j is within ρ of some center; it does not guarantee that the uniform measure on the centers is close in Wasserstein distance to the uniform empirical measure on E_j. For example, if E_j contains 99 copies of 0 and 1 copy of 100, with k=2, optimal centers are {0,100}, ρ=0, but W_1(E_j,p_j)=49 because the uniform center measure forces 0.49 mass to move from 0 to 100. Thus Eq. (4) is not a valid proxy bound. This is a load-bearing point because it is the stated theoretical justification for prototype-based retrieval. The issue is fixable: define prototypes as cluster-weighted measures (mass proportional to cluster size), in which case W(E_j,p_j) ≤ ρ follow
  2. [Section 3.2 / Appendix C.2 / Algorithm 1] The retrieval pipeline is under-specified for reproducibility. Eq. (2) and Algorithm 1 refer to an object detection model Φ_D and an embedding model Φ_E, but the concrete models are never named, and no details are given for how object-wise embeddings are extracted for either VirtualHome or ALFWorld. Also, Algorithm 1 line 26 computes the current prototype p from a single state s_t, while the text describes 'trajectory-based' prototypes; it is unclear whether p uses one state, a history of states, or a full episode. These details are necessary to reproduce the method and to interpret the Wasserstein-distance retrieval. Please specify the exact models, the prototype computation from current observations, and the distance metric used in practice.
minor comments (6)
  1. [Section 1 / Tables 1–2] The claims '20.41%' and '26.58%' improvements are percentage-point differences (e.g., 45.71 to 66.12 in VirtualHome zero-shot). Please state 'percentage points' or give relative improvements, since the two readings differ substantially.
  2. [Table 5] The caption 'Num. of WMs' is ambiguous: it could mean the number of selected models K or the total pool size N. Since the main setup uses K=3 from N=6, please clarify whether Table 5 varies K while holding N fixed or varies N.
  3. [Appendix D.2 / Table A.7] The table is labeled 'Multi-modal WorMI' and the text says 'We will include these experimental results in the final version.' A submitted paper should not contain placeholder text or incomplete experimental sections. Please finalize the results or remove them.
  4. [Figure 5] The text of Figure 5 is garbled with '/uni00000031/uni00000052/...' tokens in the provided manuscript. Please check the source and regenerate the figure.
  5. [Section 4.1] Minor typo: 'SayCayPay' appears in the first paragraph; should be 'SayCanPay'.
  6. [Section 3.2] The notation δ(E_j,E) treats Ej and E as objects with a Wasserstein distance, but Wasserstein distance is defined on probability measures. Please state explicitly that each embedding set is viewed as a uniform empirical measure (or, if the prototype definition is changed, as the corresponding weighted measure).

Circularity Check

0 steps flagged · score 1.0 of 10

No significant circularity: reported gains are held-out measurements; the only overlap-author citation is non-load-bearing.

full rationale

I walked the paper's derivation chain: world models are trained on per-domain datasets; the compound attention is meta-learned on training-domain subsets; and evaluation is on held-out tasks/scenes (Tables 1, 2). The SR/PS numbers are empirical measurements, not outputs that have been redefined as inputs. The prototype-retrieval step is a selection mechanism: p_j are k-center summaries of the training embedding sets and p comes from the current observation, so δ(p_j,p) is not the same quantity as the training objective or the reported success rate. Ablations (WorMI-E, WorMI-R; Tables 3a and A.9–A.11) show that retrieval quality changes performance, so the component has independent content. The only author-overlap citation is Kim et al. (with Woo) in Related Work, listed alongside other works on reward-function generation; it is not used to justify any load-bearing premise, uniqueness claim, or modeling ansatz. The Appendix A.1 boundedness argument is questionable as mathematics (W(E_j,p_j)≤ρ does not follow from the k-center condition for uniform empirical measures in imbalanced clusters), but that is a correctness gap, not circularity: Eq. (4) is an attempted bound, not a fitted parameter renamed as a prediction or a definitional equivalence. Appendix D.2's note 'We will include these experimental results in the final version' indicates an unfinished ablation, but that is a completeness issue, not circularity. I therefore find no circular step.

Assumptions & free parameters 4 free parameters · 5 assumptions · 0 invented entities

The central empirical claims rest on a set of hand-chosen architectural constants (K, k, N, layer indices) and assumptions about the availability and representativeness of per-domain datasets and the transferability of meta-learned attention. The embedding-space relevance assumption is the most consequential: the paper's ablations show that random retrieval degrades performance, so the method's gains hinge on it.

free parameters (4)
  • K (number of retrieved world models) = 3 (main experiments)
    Chosen by hand; Table 5 shows performance peaks near K=3 and drops with more or fewer models, indicating the result is sensitive to this choice.
  • k (number of prototypes per world model) = 15 (about 0.1% of 10,200 embeddings)
    Chosen to reduce distance computation; no sensitivity analysis is provided and the bound in Eq. (4) depends on this value.
  • N (number of world models in the pool) = 6 (VirtualHome), 4 scene types (ALFWorld)
    Arbitrary pool size; Table 5 shows N=6 underperforms N=3, suggesting the retrieval and attention are sensitive to pool size.
  • Intermediate connection layers for compound attention = reasoning layers [13,27], world model layers [7,15]
    Hand-picked; no ablation or justification is given for why these layers are used for cross-attention.
assumptions (5)
  • domain assumption Pre-trained world models M1..MN and per-domain datasets D1..DN are available and correctly labeled by domain.
    Section 3.1 states this availability as a premise of the framework.
  • domain assumption The object-wise embedding model Φ_E produces a representation space where Euclidean/Wasserstein distance between prototypes reflects task-relevant domain similarity.
    Section 3.2 builds retrieval on this premise; the paper only provides a heuristic bound (Eq. 4) and offers no validation that the embedding space captures what matters for action selection.
  • standard math Triangle inequality holds for the chosen Wasserstein distance δ.
    Used in Appendix A.1 to derive Eq. (4); standard for metric Wasserstein distances.
  • standard math The k-center optimization in Eq. (3) gives a small radius ρ for every dataset.
    The k-center problem is NP-hard, so an approximate solution is used; the paper does not report the achieved ρ values.
  • domain assumption Meta-learning (Reptile outer-loop update) on seen-domain subsets transfers to unseen domains and newly added world models.
    Section 3.4 claims this; no theory or analysis is given for this transfer.

how reviews work

0 comments
Cite this review

Pith. "Pith review of World Model Implanting for Test-time Adaptation of Embodied Agents." pith.science (2026). https://pith.science/paper/QL4L6H5A

@misc{pith2026250903956,
  author       = {Pith},
  title        = {Pith review of: World Model Implanting for Test-time Adaptation of Embodied Agents},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/QL4L6H5A}},
  note         = {Machine review of arXiv:2509.03956}
}
read the original 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 frameworks potential for scalable, real-world deployment in embodied agent scenarios where adaptability and data efficiency are essential.

Figures

Figures reproduced from arXiv: 2509.03956 by the authors.

Figure 1
Figure 1. Concept of WorMI. For sequential decision￾making policies of embodied agents, our WorMI frame￾work, a 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-specific knowledge. This dual-stage de￾sign not only combines the strengths… view at source ↗
Figure 2
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 object-wise state embeddings, given target domains. (b) By the world-wise 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 select… view at source ↗
Figure 3
Figure 3. Overall structure of the compound attention. Com [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figures from the paper (2 more)
Figure 4
Figure 4. Figure 4: Visualization of world-level attention map [PITH_FULL_IMAGE:figures/full_fig_p007_4.png]
Figure 5
Figure 5. Figure 5: Performance for continual model implanting [PITH_FULL_IMAGE:figures/full_fig_p008_5.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

28 extracted references · 17 canonical work pages

  1. [1]

    write newline

    " write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION format.date year duplicate empty "emp...

  2. [2]

    Language reward modulation for pretraining reinforcement learning

    Adeniji, A., Xie, A., Sferrazza, C., Seo, Y., James, S., and Abbeel, P. Language reward modulation for pretraining reinforcement learning. arXiv preprint arXiv:2308.12270, 2023

  3. [3]

    Llama 3.2 model card

    AI@Meta. Llama 3.2 model card. 2024. URL https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/MODEL_CARD.md

  4. [4]

    W., Pfau, D., Schaul, T., Shillingford, B., and De Freitas, N

    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

  5. [5]

    LLM Augmented LLMs: Expanding Capabilities through Composition

    Bansal, R., Samanta, B., Dalmia, S., Gupta, N., Vashishth, S., Ganapathy, S., Bapna, A., Jain, P., and Talukdar, P. Llm augmented llms: Expanding capabilities through composition, 2024. URL https://arxiv.org/abs/2401.02412

  6. [6]

    Do as i can, not as i say: Grounding language in robotic affordances

    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

  7. [7]

    Model-agnostic meta-learning for fast adaptation of deep networks

    Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp.\ 1126--1135. PMLR, 2017

  8. [8]

    Ha, D., Dai, A., and Le, Q. V. Hypernetworks. arXiv preprint arXiv:1609.09106, 2016

Show all 28 references
  1. [9]

    Z., and De Raedt, L

    Hazra, R., Dos Martires, P. Z., and De Raedt, L. Saycanpay: Heuristic planning with large language models using learnable domain knowledge. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp.\ 20123--20133, 2024

  2. [10]

    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 collaborative framework. In The Twelfth International Conference on Learning Representations

  3. [11]

    Language models as zero-shot planners: Extracting actionable knowledge for embodied agents

    Huang, W., Abbeel, P., Pathak, D., and Mordatch, I. Language models as zero-shot planners: Extracting actionable knowledge for embodied agents. In International conference on machine learning, pp.\ 9118--9147. PMLR, 2022 a

  4. [12]

    Inner monologue: Embodied reasoning through planning with language models

    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.05608, 2022 b

  5. [13]

    Jiang, D., Ren, X., and Lin, B. Y. Llm-blender: Ensembling large language models with pairwise ranking and generative fusion. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp.\ 14165--14178, 2023

  6. [14]

    Online continual learning for interactive instruction following agents

    Kim, B., Seo, M., and Choi, J. Online continual learning for interactive instruction following agents. In The Twelfth International Conference on Learning Representations, 2024

  7. [15]

    K., Lee, Y., Kim, J., and Woo, H

    Kim, W. K., Lee, Y., Kim, J., and Woo, H. Llm-based skill diffusion for zero-shot policy adaptation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems

  8. [16]

    Code as policies: Language model programs for embodied control

    Liang, J., Huang, W., Xia, F., Xu, P., Hausman, K., Ichter, B., Florence, P., and Zeng, A. Code as policies: Language model programs for embodied control. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp.\ 9493--9500. IEEE, 2023

  9. [17]

    On first-order meta-learning algorithms

    Nichol, A. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.02999, 2018

  10. [18]

    Virtualhome: Simulating household activities via programs

    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 conference on computer vision and pattern recognition, pp.\ 8494--8502, 2018

  11. [19]

    Outrageously large neural networks: The sparsely-gated mixture-of-experts layer

    Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2016

  12. [20]

    Alfworld: Aligning text and embodied environments for interactive learning

    Shridhar, M., Yuan, X., Cote, M.-A., Bisk, Y., Trischler, A., and Hausknecht, M. Alfworld: Aligning text and embodied environments for interactive learning. In International Conference on Learning Representations, 2021

  13. [21]

    Progprompt: Generating situated robot task plans using large language models

    Singh, I., Blukis, V., Mousavian, A., Goyal, A., Xu, D., Tremblay, J., Fox, D., Thomason, J., and Garg, A. Progprompt: Generating situated robot task plans using large language models. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp.\ 11523--11530. ...

  14. [22]

    H., Wu, J., Washington, C., Sadler, B

    Song, C. H., Wu, J., Washington, C., Sadler, B. M., Chao, W.-L., and Su, Y. Llm-planner: Few-shot grounded planning for embodied agents with large language models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.\ 2998--3009, 2023

  15. [23]

    Knowledge fusion of large language models

    Wan, F., Huang, X., Cai, D., Quan, X., Bi, W., and Shi, S. Knowledge fusion of large language models. In The Twelfth International Conference on Learning Representations

  16. [24]

    Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents

    Wang, Z., Cai, S., Chen, G., Liu, A., Ma, X., and Liang, Y. Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents. arXiv preprint arXiv:2302.01560, 2023

  17. [25]

    Autogen: Enabling next-gen llm applications via multi-agent conversation

    Wu, Q., Bansal, G., Zhang, J., Wu, Y., Li, B., Zhu, E., Jiang, L., Zhang, X., Zhang, S., Liu, J., et al. Autogen: Enabling next-gen llm applications via multi-agent conversation. In ICLR 2024 Workshop on Large Language Model (LLM) Agents

  18. [26]

    React: Synergizing reasoning and acting in language models

    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

  19. [27]

    G., Chiang, H.-T

    Yu, W., Gileadi, N., Fu, C., Kirmani, S., Lee, K.-H., Arenas, 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

  20. [28]

    Zhou, K., Liu, Z., Qiao, Y., Xiang, T., and Loy, C. C. Domain generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45 0 (4): 0 4396--4415, 2022

Pith tools

Reviewed August 5, 2026 · model on record in the stance chip above.