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 →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
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.
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
- 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
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)
- [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
- [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)
- [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.
- [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.
- [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.
- [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.
- [Section 4.1] Minor typo: 'SayCayPay' appears in the first paragraph; should be 'SayCanPay'.
- [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
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
free parameters (4)
- K (number of retrieved world models) =
3 (main experiments)
- k (number of prototypes per world model) =
15 (about 0.1% of 10,200 embeddings)
- N (number of world models in the pool) =
6 (VirtualHome), 4 scene types (ALFWorld)
- Intermediate connection layers for compound attention =
reasoning layers [13,27], world model layers [7,15]
assumptions (5)
- domain assumption Pre-trained world models M1..MN and per-domain datasets D1..DN are available and correctly labeled by domain.
- domain assumption The object-wise embedding model Φ_E produces a representation space where Euclidean/Wasserstein distance between prototypes reflects task-relevant domain similarity.
- standard math Triangle inequality holds for the chosen Wasserstein distance δ.
- standard math The k-center optimization in Eq. (3) gives a small radius ρ for every dataset.
- domain assumption Meta-learning (Reptile outer-loop update) on seen-domain subsets transfers to unseen domains and newly added world models.
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
Reference graph
Works this paper leans on
-
[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]
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
arXiv 2023
-
[3]
AI@Meta. Llama 3.2 model card. 2024. URL https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/MODEL_CARD.md
work page 2024
-
[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
2016
-
[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
work page Pith review arXiv 2024
-
[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
work page 2023
-
[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
2017
-
[8]
Ha, D., Dai, A., and Le, Q. V. Hypernetworks. arXiv preprint arXiv:1609.09106, 2016
arXiv 2016
Show all 28 references
-
[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
2024
-
[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
-
[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
2022
-
[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
2022 arXiv
-
[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
2023
-
[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
2024
-
[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
-
[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
2023
-
[17]
On first-order meta-learning algorithms
Nichol, A. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.02999, 2018
2018 arXiv
-
[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
2018
-
[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
2016
-
[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
2021
-
[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. ...
2023
-
[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
2023
-
[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
-
[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
2023 arXiv
-
[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
2024
-
[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
2022 arXiv
-
[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
2023
-
[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
2022
Reviewed August 5, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.