Elamur Research⌕
☰
home / papers / 2024 / WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment

WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment

Paper · 2024 · Hao Tang, Darren Key and Kevin Ellis

Summary

The paper asks how an AI agent could learn how an unfamiliar environment works from only a few interactions, then reuse that knowledge for new goals. The authors present WorldCoder, a model-based reinforcement learning agent. It uses an LLM to write and debug a Python program that serves as an explicit world model, with separate transition and reward functions. The LLM does not simulate the world or plan. It builds the simulator, and a conventional planner chooses actions. The program must reproduce every logged transition and reward. It must also be optimistic: given a natural-language goal, a planner must be able to reach positive reward under the model. This constraint pushes the agent to invent plausible intermediate goals, which drives exploration even when rewards are sparse. In experiments on gridworlds and robot-style task planning, the agent needed far fewer environment interactions than deep RL and far fewer LLM calls than ReAct-style agents. It also transferred knowledge across environments by editing existing code. The work shows that general-purpose code can act as an interpretable, reusable world model.

Key points

  • Formal setup: The setting is a Contextual MDP (Hallak et al.) where goals are the contexts, dynamics are deterministic and shared across goals, and states are sets of objects (an Object Oriented MDP).
  • Two learning constraints: ϕ1 requires the model to fit all replayed data. ϕ2 is an optimism constraint requiring that some action sequence reaches the goal with positive reward. A world model satisfying both is found by LLM program synthesis, with back-prompting on any mismatched transitions.
  • Theoretical guarantee: The analysis is non-probabilistic and combinatorial. It shows that a good-enough world model is found in polynomial time, using new notions of "logical dimensionality" of the model space and MDP diameter.
  • Efficiency claims: Once a world model is learned, a new task needs at most one LLM query (O(1)) to update the reward function. ReAct-style agents need about O(T) calls per task, where T is the episode length. The optimism constraint is reported to cut the interactions needed to obtain reward by orders of magnitude.
  • Limitations: The method assumes deterministic dynamics and symbolic, discrete states. Its knowledge-transfer mechanism is basic, and the authors suggest library learning as future work.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 2The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 works

Cite

Cite as Hao Tang, Darren Key and Kevin Ellis. 2024. WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment. arXiv:2402.12275.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2402.12275 · arXiv 2402.12275

@misc{tang2024worldcoder,
  title = {{WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment}},
  author = {Hao Tang, Darren Key and Kevin Ellis},
  year = {2024},
  doi = {10.48550/arxiv.2402.12275},
  eprint = {2402.12275},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.12275}
}

References

80 references, 13 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey, Sally Jesmonth, Nikhil… — Do As I Can, Not As I Say: Grounding Language in Robotic Affordances
  2. Peter Auer. Using confidence bounds for exploitation-exploration trade-offs. Journal of Machine Learning Research, 3(Nov):397–422, 2002. find
  3. Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International conference on machine learning, pages 263–272. PMLR, 2017. find
  4. Matthew Bowers, Theo X. Olausson, Lionel Wong, Gabriel Grand, Joshua B. Tenenbaum, Kevin Ellis, and Armando Solar-Lezama. Top-down synthesis for library learning. Proc. ACM Program. Lang., 7(POPL), jan 2023. doi: 10.1145/3571234. URL https://doi.org/10.1145/3571234. ↗
  5. Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter I Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games, 4(1):1–43, 2012… find
  6. Swarat Chaudhuri, Kevin Ellis, Oleksandr Polozov, Rishabh Singh, Armando Solar-Lezama, Yisong Yue, et al. Neurosymbolic programming. Foundations and Trends® in Programming Languages, 7(3), 2021. find
  7. Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128, 2023. find
  8. Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym. https://github.com/maximecb/gym-minigrid, 2018. ↗
  9. Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.13831, 2023.
  10. Leah Chrestien, Stefan Edelkamp, Antonin Komenda, and Tomáš Pevný. Optimize planning heuristics to rank, not to estimate cost-to-goal. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=Mgy6sgslPY. ↗
  11. Stephen Chung, Ivan Anokhin, and David Krueger. Thinker: Learning to plan and act. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=mumEBl0arj. ↗
  12. Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pages 72–83. Springer, 2006. find
  13. Ria Das, Joshua B. Tenenbaum, Armando Solar-Lezama, and Zenna Tavares. Combining functional and automata synthesis to discover causal reactive programs. Proc. ACM Program. Lang., 7(POPL), jan 2023. doi: 10.1145/3571249. URL https://doi.org/10.1145/3571249. — [doi:10.1145/3571249](https://doi.org/10.1145/3571249) ↗
  14. Luc De Raedt, Angelika Kimmig, and Hannu Toivonen. Problog: a probabilistic prolog and its application in link discovery. In Proceedings of the 20th International Joint Conference on Artifical Intelligence, IJCAI’07, page 2468–2473, San Francisco, CA, USA, 2007. Morgan Kaufmann Publishers Inc. find
  15. Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforcement learning. In Proceedings of the 25th international conference on Machine learning, pages 240–247, 2008. find
  16. Kevin Ellis, Lionel Wong, Maxwell Nye, Mathias Sable-Meyer, Luc Cary, Lore Anaya Pozo, Luke Hewitt, Armando Solar-Lezama, and Joshua B Tenenbaum. Dreamcoder: growing generalizable, interpretable knowledge with wake–sleep bayesian program learning. Philosophical Transactions of the Royal Society A, 3… find
  17. Richard Evans, Matko Bošnjak, Lars Buesing, Kevin Ellis, David Pfau, Pushmeet Kohli, and Marek Sergot. Making sense of raw input. Artificial Intelligence, 299:103521, 2021. find
  18. Dieqiao Feng, Carla P Gomes, and Bart Selman. A novel automated curriculum strategy to solve hard sokoban planning instances. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3141–3152. Curran Associates,… find
  19. Dieqiao Feng, Carla Gomes, and Bart Selman. Left heavy tails and the effectiveness of the policy and value networks in dnn-based best-first search for sokoban planning, 2022. find
  20. Caelan Reed Garrett, Rohan Chitnis, Rachel Holladay, Beomjoon Kim, Tom Silver, Leslie Pack Kaelbling, and Tomás Lozano-Pérez. Integrated task and motion planning, 2020. find
  21. Noah D. Goodman, Vikash K. Mansinghka, Daniel Roy, Keith Bonawitz, and Joshua B. Tenenbaum. Church: a language for generative models. In Proceedings of the Twenty-Fourth Conference on Uncertainty in Artificial Intelligence, UAI’08, page 220–229, Arlington, Virginia, USA, 2008. AUAI Press. ISBN 09749… find
  22. Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael Mozer, and Yoshua Bengio. Neural production systems: Learning rule-governed visual dynamics. arXiv preprint arXiv:2103.01937, 2021. — [arXiv:2103.01937](https://arxiv.org/abs/2103.01937) find
  23. Gabriel Grand, Lionel Wong, Matthew Bowers, Theo X. Olausson, Muxin Liu, Joshua B. Tenenbaum, and Jacob Andreas. Lilo: Learning interpretable libraries by compressing and documenting code, 2023. find
  24. Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging pre-trained large language models to construct and utilize world models for model-based task planning. NeurIPS, 2023.
  25. David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems 31, pages 2451–2463. Curran Associates, Inc., 2018. https://worldmodels.github.io. ↗
  26. Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ICLR, 2021. find
  27. Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2023. find
  28. Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.02259, 2015. — [arXiv:1502.02259](https://arxiv.org/abs/1502.02259) find
  29. Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. arXiv preprint arXiv:2305.14992, 2023. find
  30. Yingdong Hu, Fanqi Lin, Tong Zhang, Li Yi, and Yang Gao. Look before you leap: Unveiling the power of gpt-4v in robotic vision-language planning. arXiv preprint arXiv:2311.17842, 2023. — [arXiv:2311.17842](https://arxiv.org/abs/2311.17842) find
  31. Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on learning theory, pages 2137– 2143. PMLR, 2020. find
  32. Leslie Pack Kaelbling and Tomás Lozano-Pérez. Hierarchical task and motion planning in the now. In 2011 IEEE International Conference on Robotics and Automation, pages 1470–1477, 2011. doi: 10.1109/ICRA.2011.5980391. — [doi:10.1109/icra.2011.5980391](https://doi.org/10.1109/icra.2011.5980391) find
  33. Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04906, 2020.
  34. Mei Kobayashi and Koichi Takeda. Information retrieval on the web. ACM computing surveys (CSUR), 32(2):144–173, 2000. find
  35. Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference on machine learning, pages 282–293. Springer, 2006. find
  36. George Konidaris, Leslie Pack Kaelbling, and Tomas Lozano-Perez. From skills to symbols: Learning symbolic representations for abstract high-level planning. Journal of Artificial Intelligence Research, 61:215–289, 2018. find
  37. Kalle Kujanpää, Joni Pajarinen, and Alexander Ilin. Hybrid search for efficient planning with completeness guarantees. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=bY0c46ZtXa. ↗
  38. Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring a sequence model trained on a synthetic task. ICLR, 2023. find
  39. Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. In arXiv preprint arXiv:2209.07753, 2022. find
  40. Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle: A generic model-based algorithm for partially observable sequential decision making. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing, pages 363–376, 2023. find
  41. Zhihan Liu, Hao Hu, Shenao Zhang, Hongyi Guo, Shuqi Ke, Boyi Liu, and Zhaoran Wang. Reason for future, act for now: A principled framework for autonomous llm agents with provable sample efficiency. arXiv preprint arXiv:2309.17382, 2023.
  42. Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayaraman, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Eureka: Human-level reward design via coding large language models. arXiv preprint arXiv: Arxiv-2310.12931, 2023.
  43. Christopher D Manning, Prabhakar Raghavan, and Hinrich Schütze. Introduction to information retrieval. Cambridge university press, 2008. find
  44. Jiayuan Mao, Tomás Lozano-Pérez, Josh Tenenbaum, and Leslie Kaelbling. Pdsketch: Integrated domain programming, learning, and planning. Advances in Neural Information Processing Systems, 35:36972–36984, 2022. find
  45. Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=vhFu1Acb0xb. ↗
  46. Kolby Nottingham, Prithviraj Ammanabrolu, Alane Suhr, Yejin Choi, Hannaneh Hajishirzi, Sameer Singh, and Roy Fox. Do Embodied Agents Dream of Pixelated Sheep: Embodied Decision Making using Language Guided World Modelling. In International Conference on Machine Learning (ICML), 2023. find
  47. Theo X. Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-Lezama. Is self-repair a silver bullet for code generation?, 2023. Rl baselines3
  48. Antonin Raffin. rl-baselines3-zoo, 2020. find
  49. Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1–8, 2021. URL http://jmlr.org/papers/v22/20-1364.html. ↗
  50. Stephen Robertson, Hugo Zaragoza, et al. The probabilistic relevance framework: Bm25 and beyond. Foundations and Trends® in Information Retrieval, 3(4):333–389, 2009. find
  51. Rafael Rodriguez-Sanchez, Benjamin A. Spiegel, Jennifer Wang, Roma Patel, Stefanie Tellex, and George Konidaris. Rlang: A declarative language for describing partial world knowledge to reinforcement learning agents, 2023. find
  52. Stuart J Russell and Peter Norvig. Artificial intelligence: a modern approach. Pearson, 2016. find
  53. Feras Saad. Scalable Structure Learning, Inference, and Analysis with Probabilistic Programs. PhD thesis, Massachusetts Institute of Technology, 2022. find
  54. Max-Philipp B. Schrader. gym-sokoban, 2018. gym-sokoban. find
  55. Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604–609, 2020. find
  56. Atharva Sehgal, Arya Grayeli, Jennifer J. Sun, and Swarat Chaudhuri. Neurosymbolic grounding for compositional world models, 2023. find
  57. Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Côté, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. ALFWorld: Aligning Text and Embodied Environments for Interactive Learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2021. URL https://arxiv.org/abs/2010.… ↗
  58. David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484–489, 2016. find
  59. Tom Silver and Rohan Chitnis. Pddlgym: Gym environments from pddl problems. In International Conference on Automated Planning and Scheduling (ICAPS) PRL Workshop, 2020. URL https://github.com/tomsilver/pddlgym. ↗
  60. Sean R Sinclair, Siddhartha Banerjee, and Christina Lee Yu. Adaptive discretization for episodic reinforcement learning in metric spaces. Proceedings of the ACM on Measurement and Analysis of Computing Systems, 3(3):1–44, 2019. https://github.com/DLR-RM/zoo. https://github.com/mpSchrader/ ↗
  61. Haotian Sun, Yuchen Zhuang, Lingkai Kong, Bo Dai, and Chao Zhang. Adaplanner: Adaptive planning from feedback with language models. NeurIPS, 2023.
  62. Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pages 216–224. Elsevier, 1990. find
  63. Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018. find
  64. Hao Tang and Kevin Ellis. From perception to programs: Regularize, overparameterize, and amortize. International Conference on Machine Learning (ICML), 2023. find
  65. Hao Tang, Keya Hu, Jin Peng Zhou, Sicheng Zhong, Wei-Long Zheng, Xujie Si, and Kevin Ellis. Code repair with llms gives an exploration-exploitation tradeoff. arXiv, 2024. find
  66. Edward Chace Tolman and Charles H Honzik. Introduction and removal of reward, and maze performance in rats. University of California publications in psychology, 1930. find
  67. Pedro A Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel J Gershman, and Joshua B Tenenbaum. Human-level reinforcement learning through theory-based modeling, exploration, and planning. arXiv preprint arXiv:2107.12544, 2021. — [arXiv:2107.12544](https://arxiv.org/abs/2107.12544) find
  68. Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models - a critical investigation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL — On the Planning Abilities of Large Language Models : A Critical Investigation https://openreview.net/forum?id=X6dEqXIsEW. ↗
  69. Abhinav Verma, Vijayaraghavan Murali, Rishabh Singh, Pushmeet Kohli, and Swarat Chaudhuri. Programmatically interpretable reinforcement learning. In International Conference on Machine Learning, pages 5045–5054. PMLR, 2018. find
  70. Abhinav Verma, Hoang Le, Yisong Yue, and Swarat Chaudhuri. Imitation-projected programmatic reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019. find
  71. Lev Semenovich Vygotsky. Mind in society: The development of higher psychological processes. Harvard university press, 1980. find
  72. Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv: Arxiv-2305.16291, 2023. find
  73. Huaxiaoyue Wang, Gonzalo Gonzalez-Pumariega, Yash Sharma, and Sanjiban Choudhury. Demo2code: From summarizing demonstrations to synthesizing code via extended chain-of-thought. NeurIPS, 2023. find
  74. Lionel Wong, Jiayuan Mao, Pratyusha Sharma, Zachary S. Siegel, Jiahai Feng, Noa Korneev, Joshua B. Tenenbaum, and Jacob Andreas. Learning adaptive planning representations with natural language guidance, 2023. find
  75. Chenyang Wu, Tianci Li, Zongzhang Zhang, and Yang Yu. Bayesian optimistic optimization: Optimistic exploration for model-based reinforcement learning. Advances in neural information processing systems, 35:14210–14223, 2022. find
  76. Jiannan Xiang, Tianhua Tao, Yi Gu, Tianmin Shu, Zirui Wang, Zichao Yang, and Zhiting Hu. Language models meet world models: Embodied experiences enhance language models. NeurIPS, 2023. find
  77. Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629, 2022.
  78. Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models as commonsense knowledge for large-scale task planning. NeurIPS, 2023.
  79. Feiyu Zhu and Reid Simmons. Bootstrapping cognitive agents with a large language model. In AAAI, 2024. find
  80. Yutao Zhu, Huaying Yuan, Shuting Wang, Jiongnan Liu, Wenhan Liu, Chenlong Deng, Zhicheng Dou, and Ji-Rong Wen. Large language models for information retrieval: A survey. arXiv preprint arXiv:2308.07107, 2023.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/tang-2024-worldcoder-model-based-llm/