Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LLM-Guided Probabilistic Program Induction for POMDP Model Estimation

Aidan Curtis, Hao Tang, Thiago Veloso, Kevin Ellis, Joshua Tenenbaum, Tomás Lozano-Pérez, Leslie Pack Kaelbling

TL;DR

This work addresses learning interpretable, low-complexity POMDP models by representing the environment components—initial state, transitions, observations, and rewards—as short probabilistic programs. An LLM provides an informative prior to generate candidate programs, which are evaluated against observed data and refined via feedback, enabling data-efficient model discovery. The authors introduce POMDP Coder, which combines this LLM-guided probabilistic-program induction with an online, determinized belief-space planner to act under partial observability. Across classical POMDP benchmarks, MiniGrid simulations, and real-robot experiments, the approach achieves superior sample efficiency and predictive accuracy compared with tabular learning, behavior cloning, and direct LLM planning, demonstrating the practical impact of integrating probabilistic programming, LLM priors, and online planning for robust world-model learning in uncertain domains.

Abstract

Partially Observable Markov Decision Processes (POMDPs) model decision making under uncertainty. While there are many approaches to approximately solving POMDPs, we aim to address the problem of learning such models. In particular, we are interested in a subclass of POMDPs wherein the components of the model, including the observation function, reward function, transition function, and initial state distribution function, can be modeled as low-complexity probabilistic graphical models in the form of a short probabilistic program. Our strategy to learn these programs uses an LLM as a prior, generating candidate probabilistic programs that are then tested against the empirical distribution and adjusted through feedback. We experiment on a number of classical toy POMDP problems, simulated MiniGrid domains, and two real mobile-base robotics search domains involving partial observability. Our results show that using an LLM to guide in the construction of a low-complexity POMDP model can be more effective than tabular POMDP learning, behavior cloning, or direct LLM planning.

LLM-Guided Probabilistic Program Induction for POMDP Model Estimation

TL;DR

This work addresses learning interpretable, low-complexity POMDP models by representing the environment components—initial state, transitions, observations, and rewards—as short probabilistic programs. An LLM provides an informative prior to generate candidate programs, which are evaluated against observed data and refined via feedback, enabling data-efficient model discovery. The authors introduce POMDP Coder, which combines this LLM-guided probabilistic-program induction with an online, determinized belief-space planner to act under partial observability. Across classical POMDP benchmarks, MiniGrid simulations, and real-robot experiments, the approach achieves superior sample efficiency and predictive accuracy compared with tabular learning, behavior cloning, and direct LLM planning, demonstrating the practical impact of integrating probabilistic programming, LLM priors, and online planning for robust world-model learning in uncertain domains.

Abstract

Partially Observable Markov Decision Processes (POMDPs) model decision making under uncertainty. While there are many approaches to approximately solving POMDPs, we aim to address the problem of learning such models. In particular, we are interested in a subclass of POMDPs wherein the components of the model, including the observation function, reward function, transition function, and initial state distribution function, can be modeled as low-complexity probabilistic graphical models in the form of a short probabilistic program. Our strategy to learn these programs uses an LLM as a prior, generating candidate probabilistic programs that are then tested against the empirical distribution and adjusted through feedback. We experiment on a number of classical toy POMDP problems, simulated MiniGrid domains, and two real mobile-base robotics search domains involving partial observability. Our results show that using an LLM to guide in the construction of a low-complexity POMDP model can be more effective than tabular POMDP learning, behavior cloning, or direct LLM planning.
Paper Structure (25 sections, 3 equations, 4 figures, 3 tables, 3 algorithms)

This paper contains 25 sections, 3 equations, 4 figures, 3 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Partially Observable Markov Decision Processes
  5. Probabilistic Programs
  6. POMDP Solvers
  7. POMDP Coder
  8. Learning Models
  9. Experiments
  10. Simulated Experiments
  11. Baselines
  12. Simulation Results
  13. Real Robot Experiments
  14. Discussion
  15. Limitations
  16. ...and 10 more sections

Figures (4)

  • Figure 1: An architecture diagram for our POMDP coder method
  • Figure 2: A visualization of the final belief state for each of the MiniGrid tasks. The green square is the goal, the red triangle is the agent, and the blue squares are places that the agent has not viewed.
  • Figure 3: Experimental results for the MiniGrid and Classical POMDP domains. We show the expected discounted returns ($\gamma=0.98$) of each method across five learning seeds with ten episodes per seed. The error bars show standard error across all episodes. We normalize the expected discounted returns by the performance of Oracle.
  • Figure 4: The two real-world experimental setups wherein a robot is searching for an apple in a partially observable world. The blue cells represent the robot's belief about where the apple could be in the world. In the uniform initial belief, the robot thinks the apple could be anywhere it has not looked yet. The learned initial belief found by POMDP Coder has a narrower initial belief leading to more efficient exploration.