2505.02216v2

# LLM-Guided Probabilistic Program Induction for POMDP Model Estimation ## 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. Keywords: POMDPs, LLMs, Probabilistic Program, Induction ## 1 Introduction Decision making under uncertainty is a central challenge in robotics, autonomous systems, and artificial intelligence more broadly. Partially Observable Markov Decision Processes (POMDPs) provide a principled framework for modeling and solving such problems by explicitly representing uncertainty in state perception, transitions, and rewards. Many prior works have used the POMDP formulation to solve real-world problems such as intention-aware decision-making for autonomous vehicles [1], collaborative control of smart assistive wheelchairs [2], robotic manipulation in cluttered environments [3], and generalized object search [4]. Despite their conceptual clarity, practical application of POMDPs is bottlenecked by difficulties in specifying accurate models of environments, which requires careful engineering and a thorough understanding of the theory and available solvers. In this work, we address the critical challenge of learning interpretable, low-complexity POMDP models directly from data. We specifically target a class of POMDPs whose components, namely the observation function, reward function, transition dynamics, and initial state distribution, can be succinctly represented as probabilistic graphical models encoded by short probabilistic programs. To efficiently identify these programs, we leverage recent advancements in Large Language Models (LLMs) to serve as informative priors, generating candidate probabilistic programs. These candidates are evaluated against empirical observations and iteratively refined through LLM-generated feedback. We evaluate our approach across several domains including classical POMDP problems, minigrid navigation and manipulation problems, and a real-world robotics setting involving a mobile-base robot searching for a target object. Our experimental results demonstrate that guiding model construction with an LLM significantly enhances sample efficiency compared to traditional tabular model or behavior learning methods and achieves greater accuracy than directly querying an LLM. ## 2 Related Work Learning world models for partially observable decision-making is a form of model-based reinforcement learning [5], for which there are many methods. In this section, we focus on methods that emphasize extreme data efficiency through the use of explicit models and representations, such as probabilistic programs, that facilitate efficient learning. Additionally, we discuss methods that use LLMs to synthesize and refine these models, enabling the integration of human priors and domain-specific constraints to create world models for downstream solvers. POMDP Model Learning A substantial body of research has addressed the challenge of learning Partially Observable Markov Decision Process (POMDP) models from experience. For example, Mossel and Roch [6] provide an average-case complexity result showing that estimating the parameters of certain Hidden Markov Models, an essential subproblem in POMDP learning, is computationally intractable in general. Despite these challenges, several approaches have made significant progress by introducing tractability under specific assumptions. Bayesian methods, such as Bayes-Adaptive POMDPs [7], incorporate model uncertainty directly into decision-making by unifying model learning, information gathering, and exploitation. This, however, increases the overall complexity of the POMDP. In contrast, spectral techniques like Predictive State Representations (PSRs) [8] offer computational efficiency by bypassing full Bayesian inference, although they require strong structural assumptions and extensive exploratory data. Recent work on optimism-based exploration algorithms has yielded theoretical guarantees for efficient learning in specific POMDP subclasses, even though these methods can be challenging to apply directly to real-world, complex domains [9]. Additionally, apprenticeship learning approaches estimate POMDP parameters by leveraging expert demonstrations, assuming that expert behavior encapsulates informative state-transition dynamics [10]. Such methods help reduce the burden of exploration but are sensitive to the quality of the expert demonstrations. In dialogue systems, for instance, these techniques have successfully learned user models without relying on manual annotations [11]. Probabilistic Program Induction. Probabilistic programming provides a powerful framework for modeling complex systems using concise, symbolic representations [12, 13, 14]. Several works in probabilistic program induction have demonstrated that such representations can lead to markedly improved data efficiency and enable few-shot learning, in stark contrast to more data-intensive conventional methods [15, 16]. Nonetheless, the vast search space inherent to probabilistic programs can be a major computational bottleneck. Recent advances have attempted to address this issue by integrating language models with probabilistic programming, thereby infusing human priors into the model discovery process [17, 18, 19]. LLM Model Learning for Decision-Making. The application of large language models (LLMs) to world-modeling for decision-making is an emerging and rapidly evolving area. Prior work has primarily focused on fully observable settings, where transition dynamics and reward structures are represented using frameworks such as PDDL or code-based models [20, 21]. Other approaches have utilized code-based representations as constraints or as part of optimization frameworks [22, 23, 24]. To our knowledge, our work is the first to extend these techniques to the POMDP setting, thereby addressing the additional complexities introduced by partial observability. ## 3 Background ### 3.1 Partially Observable Markov Decision Processes Partially Observable Markov Decision Processes (POMDPs) provide a principled framework for sequential decision making under uncertainty. A POMDP is defined by the tuple $(\mathcal{S},\mathcal{A},\mathcal{O},\mathcal{T},\mathcal{Z},\mathcal{R},\gamma)$ , where $\mathcal{S}$ , $\mathcal{A}$ , and $\mathcal{O}$ denote the state, action, and observation spaces, respectively. In this work we assume all these spaces are discrete, but the POMDP formulation supports continuous spaces in general. $\mathcal{T}(s_{t+1}\mid s_{t},a_{t})$ indicates the distribution over next states expected when taking action $a_{t}$ in state $s_{t}$ . The observation model $\mathcal{Z}(o_{t}\mid s_{t+1},a_{t})$ represents the distribution over observations expected when executing action $a_{t}$ resulting in a subsequent state $s_{t+1}$ . $\mathcal{R}(s_{t},a_{t},s_{t+1})$ is the reward function. Lastly, $\gamma\in(0,1)$ is the discount factor which weighs the value of current rewards over future ones. Since the agent does not have direct access to the true state $s_{t}$ , it maintains a belief $b_{t}(s)$ , or a probability distribution over possible states, which must be updated for replanning after each action is taken and new observation received. Given an action and observation, a belief can be updated via particle filtering [25]. We use particle filtering as our belief-updating mechanism across all domains. <details> <summary>extracted/6430121/figures/arch.png Details</summary>  ### Visual Description ## System Diagram: Reinforcement Learning with Model Feedback ### Overview The image is a system diagram illustrating a reinforcement learning process, incorporating model templates, code APIs, and feedback loops. It shows the flow of information and interactions between different components, including an experiences database, coverage analysis, a model, a POMDP solver, and the environment. ### Components/Axes * **Experiences Database:** A cylindrical shape labeled "Experiences Database" stores data. * **Model Templates and Code API:** A rectangular box at the top labeled "Model Templates and Code API". * **Coverage Analysis:** A box containing two bell curves labeled "Empirical" and "Model" above "Coverage Analysis". * **Model:** A green circle containing a white logo, presumably representing the model. * **Z(o | s):** A light blue rounded rectangle containing the expression "Z(o | s)". * **T(s_t+1 | s, a):** A light pink rounded rectangle containing the expression "T(s_t+1 | s, a)". * **R(s, a):** A light green rounded rectangle containing the expression "R(s, a)". * **b₀(s):** A light purple rounded rectangle containing the expression "b₀(s)". * **POMDP Solver:** A rectangular box labeled "POMDP Solver". * **Actions:** A label indicating the output of the POMDP Solver. * **Environment:** A rounded rectangle containing a globe image and labeled "Environment". * **Experiences Feedback:** A label indicating the feedback loop from the environment to the experiences database. ### Detailed Analysis or ### Content Details 1. **Data Flow:** * Data flows from the "Experiences Database" to the "Coverage Analysis" module. * The "Coverage Analysis" module compares "Empirical" data with the "Model". * The output of the "Coverage Analysis" module feeds into the green circle model. * The "Model Templates and Code API" also feeds into the green circle model. * The model outputs to "Z(o | s)", "T(s_t+1 | s, a)", "R(s, a)", and "b₀(s)". * These outputs feed into the "POMDP Solver". * The "POMDP Solver" outputs "Actions" to the "Environment". * The "Environment" provides "Experiences Feedback" back to the "Experiences Database", closing the loop. 2. **Mathematical Expressions:** * Z(o | s): Represents the probability of observation 'o' given state 's'. * T(s_t+1 | s, a): Represents the transition probability from state 's' to state 's_t+1' given action 'a'. * R(s, a): Represents the reward function for taking action 'a' in state 's'. * b₀(s): Represents the initial belief state over states 's'. ### Key Observations * The diagram illustrates a closed-loop reinforcement learning system. * The "Coverage Analysis" module compares empirical data with the model's predictions. * The model receives input from both the "Coverage Analysis" and "Model Templates and Code API". * The POMDP solver uses the model's outputs to determine actions in the environment. * The environment provides feedback to the experiences database, allowing the system to learn and improve over time. ### Interpretation The diagram represents a reinforcement learning system that incorporates model-based learning and feedback mechanisms. The "Coverage Analysis" module likely serves to identify discrepancies between the model's predictions and real-world experiences, allowing the model to be refined and improved. The inclusion of "Model Templates and Code API" suggests that the model can be customized and extended using external code and templates. The use of a POMDP solver indicates that the environment is partially observable, requiring the agent to maintain a belief state over possible states. The feedback loop from the environment to the experiences database enables the system to learn from its interactions with the environment and adapt its behavior accordingly. </details> Figure 1: An architecture diagram for our POMDP coder method The objective of a POMDP is to find a policy $\pi$ that maximizes the expected discounted reward: $$ \max_{\pi}\;\mathbb{E}_{b_{0},\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}\sum_{s \in\mathcal{S}}b_{t}(s)\sum_{s^{\prime}\in\mathcal{S}}\mathcal{T}(s^{\prime} \mid s,a_{t})\,\mathcal{R}(s,a_{t},s^{\prime})\right] $$ To illustrate the POMDP formulation and serve as a running example, consider the classic Tiger problem [26]. An agent faces two doors: one hides a tiger, the other a treasure. The true state $s$ consisting of the the tiger’s location is hidden. The agent can listen to receive a noisy observation $o$ of which door the tiger is behind, or open a door to gain a reward or incur a penalty. An ideal policy is to wait long enough to be confident in the tiger’s location before opening the door with the treasure. ### 3.2 Probabilistic Programs Probabilistic programming offers an expressive and concise way to represent complex probabilistic models as executable code. In our work, each component of the POMDP including the initial state model, transition dynamics, observation function, and reward structure is encoded as a short probabilistic program. We leverage Pyro [12], a flexible probabilistic programming framework built on Python, to specify these models. Pyro enables us to define generative models with inherent stochastic behavior. In our tiger problem example, the below implementation would be a correct probabilistic program for the observation model. [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python def tiger_observation_func(state: TigerState, act: TigerActions): if act != TigerActions.LISTEN: return TigerObservation(NONE) correct = bool(pyro.sample(”listen_correct”, Bernoulli(torch.tensor(0.85)))) tiger_left = state.tiger_location == 0 hear_left = (correct and tiger_left) or (not correct and not tiger_left) return TigerObservation(HEAR_LEFT if hear_left else HEAR_RIGHT) ### 3.3 POMDP Solvers While traditional offline solvers aim to compute a global policy over the entire belief space [27, 28, 29], these methods often face scalability challenges in high-dimensional or continuous environments. In contrast, online solvers focus on finding a solution from a specific initial belief and replan after every step as new observations become available. Online approaches such as Partially Observable Monte Carlo Planning (POMCP) [30, 31] and Partially Observable Upper Confidence Trees (POUCT) [32] combine Monte Carlo sampling with tree search to approximate optimal policies in real time. Additionally, determinized belief space planners simplify the stochastic nature of the problem by converting it into a deterministic surrogate, thereby enabling rapid replanning [33, 34, 35]. In our work, we adopt the determinized belief space planning approach due to its computational efficiency in larger problems. Please refer to Appendix C for specifics on our solver implementation. ## 4 POMDP Coder In our approach, which we call POMDP Coder, we decompose the problem into two major components: learning the probabilistic models that define the POMDP and using these learned models for online planning. The first part involves leveraging a Large Language Model (LLM) to generate, refine, and validate candidate probabilistic programs that represent the initial state, transition, observation, and reward functions. The second component uses these models within an online POMDP solver along with the current belief to find optimal actions to take in the environment. We assume the agent is provided access to an initial set of ten human-generated demonstrations $\mathcal{D}$ , where each demonstration consists of $(s_{t},a_{t},o_{t},s_{t+1})$ transitions. Since we are learning models instead of policies, there are no assumptions made about the quality of these demonstrations. However, we do make an assumption of post-hoc full observability [36]. That is, we assume the agent gets access to the intermediate states after an episode has terminated (see Section 7 for details). Additionally POMDP Coder is provided a code-based API defining the structure of the state, action, and observation space. Below is an example for the Tiger domain. [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python class TigerActions(enum.IntEnum): OPEN_LEFT = 0, OPEN_RIGHT = 1, LISTEN = 2 class TigerObservation(Observation): obs: int # 0 = hear left, 1 = hear right, 2 = none class TigerState(State): tiger_location: int # 0 = left, 1 = right 1: Input: A demo dataset $\mathcal{D}$ , max episodes $E$ , num particles $N$ , empty models $\theta=\emptyset$ 2: for $\text{episode}=1$ to $E$ do 3: $\theta=(\theta_{\text{trans}},\theta_{\text{rew}},\theta_{\text{obs}},\theta_{ \text{init}})\leftarrow\texttt{LearnModels}(\mathcal{D},\theta)$ $\triangleright$ Update all models using $\mathcal{D}$ 4: $b\leftarrow N\text{ samples from }\theta_{\text{init}}$ 5: while episode not terminated do 6: $a\leftarrow\texttt{POMDPSolver}(b,\theta)$ $\triangleright$ Plan best next action, see Appendix C 7: Execute $a$ in the world, observe $o$ 8: $b\leftarrow\texttt{ParticleFilter}(b,a,o,\theta)$ $\triangleright$ Update belief 9: Append $(a,o)$ to trajectory $\tau$ 10: end while 11: $\mathcal{D}\leftarrow\mathcal{D}\cup\tau$ $\triangleright$ Update demo data with new trajectory 12: end for 13: return $\theta$ Algorithm 1 POMDP Coder Given these inputs, POMDP Coder proceeds as outlined in Algorithm 1. It proposes an initial set of models that comprise the POMDP problem (Line 3), using a learning procedure detailed in Section 4.1. After an initial set of models is decided on, these models are passed to a POMDP solver along with the initial belief (Line 6) to find an optimal first action to take in the environment (Line 7). After an action is taken and an observation received, the belief is updated using particle filtering to form a new belief (Line 8). This process continues until the episode terminates or times out. Lastly, at the end of each episode, the trajectory is added to the dataset (Line 8) and the learning process repairs any inaccuracies that the previous model may have had under the new data (Line 3). ### 4.1 Learning Models We aim to learn four core components of a POMDP: the initial state distribution $P(s_{0})$ , the transition model $P(s_{t+1}\mid s_{t},a_{t})$ , the observation model $P(o_{t}\mid s_{t+1},a_{t})$ , and the reward model $R(s_{t},a_{t},s_{t+1})$ . Each of these components is expressed as a short probabilistic program. The objective of our learning procedure is to maximize a dataset coverage metric, which we define to be the proportion of data in $\mathcal{D}$ that has support under the model as follows: $$ \text{coverage}(P_{\theta},\mathcal{D})\;=\;\frac{1}{|\mathcal{D}|}\sum_{i=1}^ {|\mathcal{D}|}\mathbf{1}\!\Bigl{[}\,P_{\theta}\bigl{(}y_{i}\mid x_{i}\bigr{)} >0\,\Bigr{]}. $$ Although we experimented with other metrics such distributional distance metrics, we found that those were more prone to overfitting and less interpretable to an LLM than binary coverage feedback. Still, the coverage metric has its own limitations, which we discuss in Section 7. Our approach to learning these models builds on strategies previously developed for reward model learning [37], but extends them to the more general setting of POMDP model learning across multiple stochastic components (transition, observation, initial state, and reward), replaces accuracy with coverage, and introduces the notion of a testing and training split to avoid overfitting. At its core, our model learning strategy uses two operations: (1) LLM program proposal given a model function template and a set of examples from the database and and (2) LLM program repair given a previous model and set of examples that the previous model failed to cover. We run a stochastic procedure for sampling which program to repair next, which is biased toward repairing programs that have high coverage. The pseudocode and additional details can be found in the Appendix B. ## 5 Experiments ### 5.1 Simulated Experiments Simulated experiments were conducted on two categories of problems: classical POMDP problems from the literature and MiniGrid tasks. The classical POMDP problems serve as simplified benchmarks that capture the core challenges of decision making under uncertainty while keeping the problem domains small and tractable. In addition to the tiger problem, we evaluate on the rock sample problem [38]. <details> <summary>x1.png Details</summary>  ### Visual Description ## Diagram: Game Level Layout Examples ### Overview The image presents five example layouts for a game level, each demonstrating different environmental features or challenges. The layouts are displayed as grids with various colored blocks representing different elements. ### Components/Axes * **Grid Structure:** Each layout is presented on a grid, likely representing the game's tile-based system. * **Colors:** * Black: Represents solid, impassable blocks. * Gray: Represents passable floor. * Blue with diagonal lines: Represents walls. * Orange with wavy lines: Represents lava. * Green: Represents the player's starting position. * Red Triangle: Represents the player. * **Labels (Bottom of each layout):** * Empty * Corners * Lava * Unlock * Rooms ### Detailed Analysis or ### Content Details 1. **Empty:** * Grid size: 5x5 * Top-left corner: Blue diagonal lines. * Top-right and right side: Gray. * Center: Black. * Bottom-left: Black. * Bottom-right: Green square with a red triangle pointing down. 2. **Corners:** * Grid size: 7x7 * Top and left: Blue diagonal lines. * Center: Black. * Top-left: Green square with a red triangle pointing up. * Bottom-right: Black. 3. **Lava:** * Grid size: 7x7 * Top and right: Blue diagonal lines. * Center: Black. * Vertical strip of orange with wavy lines in the center. * Bottom-right: Green square with a red triangle pointing down. 4. **Unlock:** * Grid size: 7x7 * Top and left: Gray. * Bottom: Blue diagonal lines. * Center: Black. * Top-left: Gray. * Center: Green square with a red triangle pointing right. 5. **Rooms:** * Grid size: 9x9 * Background: Blue diagonal lines. * Rooms: Black and gray. * Center-right: Green square with a red triangle pointing up. ### Key Observations * Each layout demonstrates a different level design concept. * The "Empty" layout is the simplest, while "Rooms" is the most complex. * The red triangle consistently represents the player's position, and the green square indicates the starting point. * The blue diagonal lines represent walls, and the orange wavy lines represent lava. ### Interpretation The image serves as a visual guide to different level design elements that can be incorporated into a game. The layouts showcase how different combinations of walls, floors, obstacles (lava), and room structures can create varied gameplay experiences. The "Empty" layout might represent a tutorial level, while "Rooms" could be a more advanced stage. The "Corners" and "Unlock" layouts demonstrate specific challenges or mechanics that the player might encounter. The "Lava" layout introduces a hazard that the player must avoid. </details> 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. MiniGrid is a set of minimalistic gridworld tasks for testing navigation and planning under partial observability originally designed for reinforcement learning [39]. We evaluate on five of these environments. In the empty environment, the agent is deterministically placed in the top left cell of a 5×5 grid and must navigate to the green square in the bottom right. In the corners environment, the agent is randomly positioned with an arbitrary orientation in a 10×10 grid while the green square appears in a randomly selected corner. In the lava environment, the agent starts in the upper left corner of a 10×10 grid that features a randomly positioned column of lava with a gap forming a narrow passage to the green square. In the unlock environment, the agent is randomly placed in the left room of a two-room layout, must collect a randomly placed key to open a locked door, and then proceeds to the green square fixed at the center of the right room. In the rooms environment, the agent is initialized in the upper right-hand corner of a multi-room setting and must traverse through the rooms to reach the green square randomly located in the bottom right room. Each MiniGrid environment is modified from the original implementation [39]. This modification demonstrates the ability of our method to generalize to new environments not seen during LLM pretraining. ### 5.2 Baselines In our experiments, we evaluate our method against several diverse baselines to comprehensively assess its performance in partially observable environments. One baseline, termed the oracle, uses POMDP models that are hardcoded to exactly match the true dynamics of the environment, thereby serving as an upper-bound on achievable performance. In contrast, the random baseline takes actions arbitrarily at every step, establishing a lower-bound benchmark for comparison. Another baseline, referred to as direct LLM, involves querying a large language model for the next action at each decision point. In this setup, the LLM is provided with all the same information provided to POMDP Coder during model learning. The exact prompt template used for this method is detailed in Appendix E.4. Next, our evaluation includes a behavior cloning baseline, where a policy is constructed by mapping states to actions using a dictionary learned from the demonstration dataset. In addition, we consider a tabular baseline in which the POMDP models are learned as conditional probability tables derived from counts in the demonstration dataset. Lastly, we test against two ablations of POMDP Coder. The first is the offline only ablation which only makes use of the human demonstrations and does not update the model with its own experiences. Conversely, the online only ablation does not make use of the expert demonstrations, learning only from its own experiences. ### 5.3 Simulation Results We evaluate various methods using expected discounted reward defined in Section 3.1, measuring both total cumulative reward and efficiency. The results of our evaluations can be seen in Figure 3. We see POMDP Coder match or outperform all baseline methods across all domains. We observe that the behavior cloning and tabular baselines were fundamentally limited in their ability to generalize. This is because the set of possible initial states for many tasks was orders of magnitude larger than the training set. In contrast, the probabilistic programs written by POMDP Coder use symbolic abstraction to cover large portions of the state and observation spaces, allowing them to generalize to new situations. While the direct LLM approach was sometimes effective, such as in the rock sample domain, it frequently got stuck in infinite loops, failing to understand constraints such as obstacle obstruction despite the examples of collision it had access to in the dataset. POMDP Coder outperformed both the online-only and offline-only ablations across most environments. A common failure mode of the offline-only ablation was missing transitions outside expert demonstrations. For example, it will run into lava without knowing it causes death. In contrast, the online-only ablation struggled to discover informative actions due to inefficient random exploration. For instance, in the unlock environment, if the agent never used the key, the model failed to learn that behavior, and random actions rarely uncovered it. In addition to success metrics, we record some additional runtime statistics such as the number of candidate programs generated during offline and online learning in Table 2 as well as the training and testing coverage scores after offline model learning in Table 3. <details> <summary>x2.png Details</summary>  ### Visual Description ## Bar Chart: Episode Reward Across Environments & Approaches ### Overview The image is a bar chart comparing the performance of different reinforcement learning approaches across various environments. The y-axis represents the normalized discounted reward, ranging from 0.0 to 1.0. The x-axis represents different environments: Tiger, RockSample, Empty, Corners, Lava, Rooms, and Unlock. Each environment has a set of bars representing different approaches: Oracle, Ours, Ours (Offline), Ours (Online), Direct LLM, Behavior Cloning, Tabular, and Random. Error bars are included on each bar to indicate the variance in performance. ### Components/Axes * **Title:** Episode Reward Across Environments & Approaches * **Y-axis:** Normalised Discounted Reward, ranging from 0.0 to 1.0 in increments of 0.2. * **X-axis:** Environments: Tiger, RockSample, Empty, Corners, Lava, Rooms, Unlock. * **Legend:** Located at the bottom of the chart. * Gray: Oracle * Dark Blue: Ours * Purple: Ours (Offline) * Pinkish-Purple: Ours (Online) * Pink: Direct LLM * Orange: Behavior Cloning * Yellow-Orange: Tabular * Yellow: Random ### Detailed Analysis **Tiger Environment:** * Oracle (Gray): ~0.99 * Ours (Dark Blue): ~0.95 * Ours (Offline) (Purple): ~0.90 * Ours (Online) (Pinkish-Purple): ~0.80 * Direct LLM (Pink): ~0.70 * Behavior Cloning (Orange): ~0.60 * Tabular (Yellow-Orange): ~0.50 * Random (Yellow): ~0.45 **RockSample Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~0.75 * Ours (Offline) (Purple): ~0.68 * Ours (Online) (Pinkish-Purple): ~0.65 * Direct LLM (Pink): ~0.15 * Behavior Cloning (Orange): ~0.30 * Tabular (Yellow-Orange): ~0.25 * Random (Yellow): ~0.10 **Empty Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~1.0 * Ours (Offline) (Purple): ~0.98 * Ours (Online) (Pinkish-Purple): ~0.95 * Direct LLM (Pink): ~0.90 * Behavior Cloning (Orange): ~0.85 * Tabular (Yellow-Orange): ~0.30 * Random (Yellow): ~0.10 **Corners Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~0.75 * Ours (Offline) (Purple): ~0.70 * Ours (Online) (Pinkish-Purple): ~0.75 * Direct LLM (Pink): ~0.75 * Behavior Cloning (Orange): ~0.10 * Tabular (Yellow-Orange): ~0.05 * Random (Yellow): ~0.05 **Lava Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~0.35 * Ours (Offline) (Purple): ~0.20 * Ours (Online) (Pinkish-Purple): ~0.10 * Direct LLM (Pink): ~0.05 * Behavior Cloning (Orange): ~0.02 * Tabular (Yellow-Orange): ~0.01 * Random (Yellow): ~0.01 **Rooms Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~0.70 * Ours (Offline) (Purple): ~0.20 * Ours (Online) (Pinkish-Purple): ~0.30 * Direct LLM (Pink): ~0.10 * Behavior Cloning (Orange): ~0.05 * Tabular (Yellow-Orange): ~0.02 * Random (Yellow): ~0.01 **Unlock Environment:** * Oracle (Gray): ~1.0 * Ours (Dark Blue): ~0.75 * Ours (Offline) (Purple): ~0.10 * Ours (Online) (Pinkish-Purple): ~0.05 * Direct LLM (Pink): ~0.02 * Behavior Cloning (Orange): ~0.01 * Tabular (Yellow-Orange): ~0.01 * Random (Yellow): ~0.01 ### Key Observations * The "Oracle" approach consistently achieves the highest normalized discounted reward across all environments, often reaching a value of 1.0. * The "Ours" approach (dark blue) generally performs well, but its performance varies depending on the environment. * The "Random" approach (yellow) consistently performs the worst across all environments. * The performance of "Ours (Offline)", "Ours (Online)", "Direct LLM", "Behavior Cloning", and "Tabular" varies significantly depending on the environment. ### Interpretation The bar chart provides a comparative analysis of different reinforcement learning approaches across a range of environments. The "Oracle" approach serves as a benchmark, representing the optimal performance achievable in each environment. The performance of the other approaches relative to the "Oracle" indicates their effectiveness in solving the specific challenges posed by each environment. The significant performance differences across environments highlight the importance of choosing an appropriate learning approach based on the characteristics of the environment. The error bars indicate the variability in performance, which can be attributed to factors such as the stochastic nature of the learning process and the complexity of the environment. The data suggests that the "Ours" approach is a viable alternative to the "Oracle" approach in some environments, while other approaches may be more suitable for specific tasks. The consistently poor performance of the "Random" approach underscores the importance of using informed learning strategies. </details> 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. ### 5.4 Real Robot Experiments Our real robot experiments are conducted using a Boston Dynamics Spot robot. The robot carries an in-hand camera mounted on a 6-DoF arm at its back. April tags are distributed throughout the area, enabling precise localization. The goal is for the robot to find an pick up an apple placed within the scene. Before any demonstrations are gathered, we construct a map of the empty room by scanning it with PolyCam [40]. This scan is used to construct a scene representation that includes an object-centric scene graph, encoding “on” relationships derived from geometric cues, and an occupancy grid delineating forbidden zones corresponding to physical obstacles. For each real-world task, ten demonstrations are collected by commanding the robot via keyboard. The agent’s action space is discretized into fixed theta rotations to the left and right and movements in the four cardinal directions. Objects are detected online using Grounding SAM, an open vocabulary object detector [41, 42, 43]. We test our method in two distinct spaces. The first is a small, closed-off room, as shown in the top row of Figure 4, which contains a few tables, chairs, and drawer cabinets. The second is a large, open lobby area depicted in the bottom row of Figure 4, furnished with more than twenty pieces of furniture. Within these environments, task distributions are defined by varying the location of the apple. In the Small-Cabinets configuration, which takes place in the small room, the apple is consistently placed on top of one of the three drawer cabinets for each demonstration. In the Large-Tables setting within the large room, the apple is positioned on one of the five round tables. Our approach is compared against a subset of baselines evaluated in Section 5.2. The evaluation includes behavior cloning, direct LLM execution, and a hardcoded uniform baseline that assumes the object is placed uniformly throughout the search space. Although the uniform baseline requires additional task-specific human input and is not strictly an apples-to-apples comparison, it demonstrates that our method can outperform a naively designed initial state distribution. The results shown in Table 1 demonstrate that POMDP Coder achieves more efficient and accurate exploration by understanding and generalizing trends in initial state distribution seen in the training data. Specifically, our approach learns that objects are always on top of objects of a particular class, and constructs an initial state distribution that captures that without overfitting to specific initial states. <details> <summary>x3.png Details</summary>  ### Visual Description ## Comparison of Robot Localization Methods ### Overview The image presents a comparison of robot localization methods in two different environments: one with small cabinets and another with large tables. For each environment, the image shows the initial state, initial observation, uniform initial belief, and learned initial belief. The robot's position is represented in yellow, and the environment is represented in blue in the belief states. ### Components/Axes * **Rows:** * Small-Cabinets * Large-Tables * **Columns:** * Initial State: A photograph of the environment. * Initial Observation: A photograph from the robot's perspective. * Uniform Initial Belief: A 3D representation of the environment with a uniform probability distribution (blue) and the robot's location (yellow). * Learned Initial Belief: A 3D representation of the environment with a learned probability distribution (blue) and the robot's location (yellow). ### Detailed Analysis or ### Content Details **Row 1: Small-Cabinets** * **Initial State:** A room with two cabinets, a chair, and two robots (one yellow, one black). There are also several markers on the wall. * **Initial Observation:** A view of a cabinet with a marker on the wall. * **Uniform Initial Belief:** A 3D representation of the room with a uniform probability distribution (blue) and the robot's location (yellow). The blue dots cover the entire floor. * **Learned Initial Belief:** A 3D representation of the room with a learned probability distribution (blue) and the robot's location (yellow). The blue dots are concentrated around the objects in the room. **Row 2: Large-Tables** * **Initial State:** A room with a large table, a blue sofa, and several chairs. A yellow robot is visible. * **Initial Observation:** A view of a black chair. * **Uniform Initial Belief:** A 3D representation of the room with a uniform probability distribution (blue) and the robot's location (yellow). The blue dots cover the entire floor. * **Learned Initial Belief:** A 3D representation of the room with a learned probability distribution (blue) and the robot's location (yellow). The blue dots are concentrated around the objects in the room. ### Key Observations * The "Uniform Initial Belief" shows a uniform distribution of possible locations for the robot, indicated by the blue dots covering the entire floor area. * The "Learned Initial Belief" shows a more refined probability distribution, with the blue dots concentrated around the actual objects in the environment. This indicates that the robot has learned to localize itself more accurately. * The robot's position is indicated by the yellow figure in the "Uniform Initial Belief" and "Learned Initial Belief" columns. ### Interpretation The image demonstrates the process of robot localization using different methods. The "Uniform Initial Belief" represents a state where the robot has no prior knowledge of its location. The "Learned Initial Belief" represents a state where the robot has learned to localize itself more accurately by observing its environment. The comparison between the two belief states highlights the effectiveness of the learning process in improving robot localization. The concentration of the blue dots in the "Learned Initial Belief" around the objects in the environment indicates that the robot has successfully learned to associate its observations with the objects in the environment. </details> 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. $$ \textbf{0.89}\pm{\textbf{0.09}}\;(10) 0.68\pm{0.21}\;(10) 0.25\pm{0.42}\;(3) 0.28\pm{0.39}\;(4) 0.35\pm{0.41}\;(5) \textbf{0.73}\pm{\textbf{0.08}}\;(10) 0.40\pm{0.26}\;(8) 0.15\pm{0.32}\;(2) 0.13\pm{0.28}\;(2) 0.13\pm{0.22}\;(3) \tag{10} $$ Table 1: Real-world experiment results on the small room domain in the top row of Figure 4 and the large room domain in the bottom row of Figure 4. The table shows the mean and standard deviation of expected discounted reward (with $\gamma=0.98$ ) under the ground-truth reward model (1 if the agent is holding the apple and 0 otherwise) along with the number of successes over ten runs in parenthesis. ## 6 Discussion In this work, we have presented a novel approach to learning interpretable, low-complexity POMDP models by integrating LLM-guided probabilistic program induction with online planning. Our method leverages large language models to generate and iteratively refine candidate probabilistic programs that capture the dynamics, observation functions, initial state distributions, and reward structures of complex environments. Experimental results on simulated MiniGrid domains and real-world robotics scenarios demonstrate that our approach can significantly enhance sample efficiency and predictive accuracy compared to traditional tabular learning methods, behavior cloning, or direct LLM planning. Our findings further suggest that environments represented with structured scene graphs and other rich input representations can be better modeled by learning a world model within which a reasoning agent can operate, rather than by directly learning a policy that maps observation histories to actions or by attempting to apply a language model in a zero-shot setting. This is particularly evident in large, partially observable worlds where the belief space is considerably more complex and challenging to cover with training examples than the space of states itself. The use of code to represent these models is especially advantageous, as language models are adept at generating concise, executable snippets that can be interpreted, debugged, and evaluated post-hoc, thus providing an additional layer of transparency and robustness in model evaluation. ## 7 Limitations Despite these promising results, several limitations remain. Our approach currently relies on human expertise to design the underlying representation over which the world model is learned, which may constrain its applicability to domains where such structured representations are not readily available. Additionally, due to our post-hoc observability assumption, collecting datasets outside of a simulator requires one of the following: human state annotation, complete robot exploration after the episode, or third-party perspectives such as externally mounted cameras. In our real robot experiments, this was not a challenge because the only state variability was in the position of the goal object, which is fully determined upon completion of the task. More complex problems with multiple dimensions of both task-relevant and task-irrelevant uncertainty would require more than just the agent’s perspective. Alleviating this assumption may require jointly reasoning about the interrelated structure of the constituent models, and is a valuable direction for future work. Moreover, the particle filter employed in our current implementation does not scale well to arbitrarily large state spaces. Future work may address this limitation by incorporating more advanced inference techniques, such as factored particle filters or other scalable methods, to improve performance in high-dimensional settings. Another area for improvement is in the sometimes overly broad distributions proposed by the LLM due to the coverage metric indirectly rewarding broader distributions. While this doesn’t make the problem infeasible, it can lead to less efficient behavior. A direction for future work could be to use inference methods on the hidden variables of the proposed probabilistic program to strike a balance between the empirical distribution and the overly broad model distribution. Lastly, our study focuses exclusively on discrete state and action spaces, despite robotics tasks requiring search over continuous spaces such as grasps and poses. Extending our learning strategy and adopting continuous-space POMDP solvers would broaden our framework to these domains, enabling more complex manipulation and navigation tasks. Ultimately, our work opens up exciting avenues for combining the strengths of probabilistic programming and large language models to construct robust, interpretable models for decision-making under uncertainty. ## References - Song et al. [2016] W. Song, G. Xiong, and H. Chen. Intention-aware autonomous driving decision-making in an uncontrolled intersection. Mathematical Problems in Engineering, 2016:1–15, 01 2016. doi: 10.1155/2016/1025349. - Ghorbel et al. [2018] M. Ghorbel, J. Pineau, R. Gourdeau, S. Javdani, and S. Srinivasa. A decision-theoretic approach for the collaborative control of a smart wheelchair. International Journal of Social Robotics, 10, 01 2018. doi: 10.1007/s12369-017-0434-7. - Pajarinen and Kyrki [2017] J. Pajarinen and V. Kyrki. Robotic manipulation of multiple objects as a pomdp. Artificial Intelligence, 247:213–228, 2017. ISSN 0004-3702. doi: https://doi.org/10.1016/j.artint.2015.04.001. URL https://www.sciencedirect.com/science/article/pii/S0004370215000570. Special Issue on AI and Robotics. - Zheng et al. [2023] K. Zheng, A. Paul, and S. Tellex. A system for generalized 3d multi-object search. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 1638–1644, 2023. doi: 10.1109/ICRA48891.2023.10161387. - Moerland et al. [2020] T. M. Moerland, J. Broekens, and C. M. Jonker. Model-based reinforcement learning: A survey. CoRR, abs/2006.16712, 2020. URL https://arxiv.org/abs/2006.16712. - Mossel and Roch [2005] E. Mossel and S. Roch. Learning nonsingular phylogenies and hidden markov models. CoRR, abs/cs/0502076, 2005. URL http://arxiv.org/abs/cs/0502076. - Ross et al. [2007] S. Ross, B. Chaib-draa, and J. Pineau. Bayes-adaptive pomdps. In J. Platt, D. Koller, Y. Singer, and S. Roweis, editors, Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc., 2007. URL https://proceedings.neurips.cc/paper_files/paper/2007/file/3b3dbaf68507998acd6a5a5254ab2d76-Paper.pdf. - Boots et al. [2009] B. Boots, S. M. Siddiqi, and G. J. Gordon. Closing the learning-planning loop with predictive state representations, 2009. URL https://arxiv.org/abs/0912.2385. - Jin et al. [2020] C. Jin, S. M. Kakade, A. Krishnamurthy, and Q. Liu. Sample-efficient reinforcement learning of undercomplete pomdps. CoRR, abs/2006.12484, 2020. URL https://arxiv.org/abs/2006.12484. - Makino and Takeuchi [2012] T. Makino and J. Takeuchi. Apprenticeship learning for model parameters of partially observable environments, 2012. URL https://arxiv.org/abs/1206.6484. - Thomson et al. [2010] B. Thomson, F. Jurčíček, M. Gašić, S. Keizer, F. Mairesse, K. Yu, and S. Young. Parameter learning for pomdp spoken dialogue models. In 2010 IEEE Spoken Language Technology Workshop, pages 271–276, 2010. doi: 10.1109/SLT.2010.5700863. - Bingham et al. [2018] E. Bingham, J. P. Chen, M. Jankowiak, F. Obermeyer, N. Pradhan, T. Karaletsos, R. Singh, P. A. Szerlip, P. Horsfall, and N. D. Goodman. Pyro: Deep universal probabilistic programming. CoRR, abs/1810.09538, 2018. URL http://arxiv.org/abs/1810.09538. - Cusumano-Towner et al. [2019] M. F. Cusumano-Towner, F. A. Saad, A. K. Lew, and V. K. Mansinghka. Gen: a general-purpose probabilistic programming system with programmable inference. In Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2019, page 221–236, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367127. doi: 10.1145/3314221.3314642. URL https://doi.org/10.1145/3314221.3314642. - Goodman et al. [2012] N. D. Goodman, V. Mansinghka, D. M. Roy, K. A. Bonawitz, and J. B. Tenenbaum. Church: a language for generative models. CoRR, abs/1206.3255, 2012. URL http://arxiv.org/abs/1206.3255. - Ellis et al. [2020] K. Ellis, C. Wong, M. I. Nye, M. Sablé-Meyer, L. Cary, L. Morales, L. B. Hewitt, A. Solar-Lezama, and J. B. Tenenbaum. Dreamcoder: Growing generalizable, interpretable knowledge with wake-sleep bayesian program learning. CoRR, abs/2006.08381, 2020. URL https://arxiv.org/abs/2006.08381. - Lake et al. [2015] B. M. Lake, R. Salakhutdinov, and J. B. Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332–1338, 2015. doi: 10.1126/science.aab3050. URL https://doi.org/10.1126/science.aab3050. - Li et al. [2024] M. Y. Li, E. B. Fox, and N. D. Goodman. Automated statistical model discovery with language models, 2024. URL https://arxiv.org/abs/2402.17879. - Wong et al. [2023] L. Wong, G. Grand, A. K. Lew, N. D. Goodman, V. K. Mansinghka, J. Andreas, and J. B. Tenenbaum. From word models to world models: Translating from natural language to the probabilistic language of thought, 2023. URL https://arxiv.org/abs/2306.12672. - Grand et al. [2024] G. Grand, V. Pepe, J. Andreas, and J. B. Tenenbaum. Loose lips sink ships: Asking questions in battleship with language-informed program sampling, 2024. URL https://arxiv.org/abs/2402.19471. - Liang et al. [2025] Y. Liang, N. Kumar, H. Tang, A. Weller, J. B. Tenenbaum, T. Silver, J. F. Henriques, and K. Ellis. Visualpredicator: Learning abstract world models with neuro-symbolic predicates for robot planning, 2025. URL https://arxiv.org/abs/2410.23156. - Tang et al. [2024] H. Tang, D. Key, and K. Ellis. Worldcoder, a model-based llm agent: Building world models by writing code and interacting with the environment, 2024. URL https://arxiv.org/abs/2402.12275. - Curtis et al. [2024] A. Curtis, N. Kumar, J. Cao, T. Lozano-Pérez, and L. P. Kaelbling. Trust the proc3s: Solving long-horizon robotics problems with llms and constraint satisfaction, 2024. URL https://arxiv.org/abs/2406.05572. - Hao et al. [2024] Y. Hao, Y. Chen, Y. Zhang, and C. Fan. Large language models can plan your travels rigorously with formal verification tools. In arxiv preprint, 2024. URL https://arxiv.org/abs/2404.11891. - Ye et al. [2024] X. Ye, Q. Chen, I. Dillig, and G. Durrett. Satlm: Satisfiability-aided language models using declarative prompting. Advances in Neural Information Processing Systems (NeurIPS), 2024. URL https://arxiv.org/pdf/2305.09656. - Thrun et al. [2005] S. Thrun, W. Burgard, and D. Fox. Probabilistic Robotics (Intelligent Robotics and Autonomous Agents). The MIT Press, 2005. ISBN 0262201623. - Kaelbling et al. [1998] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101(1):99–134, 1998. ISSN 0004-3702. doi: https://doi.org/10.1016/S0004-3702(98)00023-X. URL https://www.sciencedirect.com/science/article/pii/S000437029800023X. - Mundhenk et al. [1997] M. Mundhenk, J. Goldsmith, and E. Allender. The complexity of policy evaluation for finite-horizon partially-observable markov decision processes. In I. Prívara and P. Ružička, editors, Mathematical Foundations of Computer Science 1997, pages 129–138, Berlin, Heidelberg, 1997. Springer Berlin Heidelberg. ISBN 978-3-540-69547-9. - Littman [1995] M. Littman. The witness algorithm: Solving partially observable markov decision processes. 02 1995. - Cassandra et al. [2013] A. R. Cassandra, M. L. Littman, and N. L. Zhang. Incremental pruning: A simple, fast, exact method for partially observable markov decision processes. CoRR, abs/1302.1525, 2013. URL http://arxiv.org/abs/1302.1525. - Silver and Veness [2010] D. Silver and J. Veness. Monte-Carlo planning in large POMDPs. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010. - Curtis et al. [2022] A. Curtis, L. Kaelbling, and S. Jain. Task-directed exploration in continuous pomdps for robotic manipulation of articulated objects, 2022. URL https://arxiv.org/abs/2212.04554. - Sunberg and Kochenderfer [2017] Z. Sunberg and M. J. Kochenderfer. POMCPOW: an online algorithm for pomdps with continuous state, action, and observation spaces. CoRR, abs/1709.06196, 2017. URL http://arxiv.org/abs/1709.06196. - Yoon et al. [2007] S. W. Yoon, A. Fern, and R. Givan. Ff-replan: A baseline for probabilistic planning. In International Conference on Automated Planning and Scheduling, 2007. URL https://api.semanticscholar.org/CorpusID:15013602. - Kaelbling and Lozano-Perez [2013] L. Kaelbling and T. Lozano-Perez. Integrated task and motion planning in belief space. The International Journal of Robotics Research, 32:1194–1227, 08 2013. doi: 10.1177/0278364913484072. - Curtis et al. [2024] A. Curtis, G. Matheos, N. Gothoskar, V. Mansinghka, J. Tenenbaum, T. Lozano-Pérez, and L. P. Kaelbling. Partially observable task and motion planning with uncertainty and risk awareness, 2024. URL https://arxiv.org/abs/2403.10454. - Pinto et al. [2017] L. Pinto, M. Andrychowicz, P. Welinder, W. Zaremba, and P. Abbeel. Asymmetric actor critic for image-based robot learning. CoRR, abs/1710.06542, 2017. URL http://arxiv.org/abs/1710.06542. - Tang et al. [2024] H. Tang, K. Hu, J. P. Zhou, S. Zhong, W.-L. Zheng, X. Si, and K. Ellis. Code repair with llms gives an exploration-exploitation tradeoff. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 117954–117996. Curran Associates, Inc., 2024. URL https://proceedings.neurips.cc/paper_files/paper/2024/file/d5c56ec4f69c9a473089b16000d3f8cd-Paper-Conference.pdf. - Smith and Simmons [2012] T. Smith and R. G. Simmons. Heuristic search value iteration for pomdps. CoRR, abs/1207.4166, 2012. URL http://arxiv.org/abs/1207.4166. - Chevalier-Boisvert et al. [2023] M. Chevalier-Boisvert, B. Dai, M. Towers, R. de Lazcano, L. Willems, S. Lahlou, S. Pal, P. S. Castro, and J. Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.13831, 2023. - Polycam [2025] Polycam. Polycam – capture and share 3d models, 2025. URL https://poly.cam/. Accessed: 2025-04-30. - Liu et al. [2023] S. Liu, Z. Zeng, T. Ren, F. Li, H. Zhang, J. Yang, C. Li, J. Yang, H. Su, J. Zhu, et al. Grounding dino: Marrying dino with grounded pre-training for open-set object detection. arXiv preprint arXiv:2303.05499, 2023. - Ren et al. [2024] T. Ren, S. Liu, A. Zeng, J. Lin, K. Li, H. Cao, J. Chen, X. Huang, Y. Chen, F. Yan, Z. Zeng, H. Zhang, F. Li, J. Yang, H. Li, Q. Jiang, and L. Zhang. Grounded sam: Assembling open-world models for diverse visual tasks, 2024. - Smith and Simmons [2012] T. Smith and R. G. Simmons. Heuristic search value iteration for pomdps. CoRR, abs/1207.4166, 2012. URL http://arxiv.org/abs/1207.4166. ## Appendix A Code Release The full implementation, including all code, trained models, and experiment configurations, will be released publicly upon publication to ensure reproducibility and facilitate future research. ## Appendix B Model Learning The learning algorithm follows the pseudocode in Algorithm 2. Firstly, in order to mitigate overfitting to spurious patterns, we split the demonstration dataset into separate training and test sets (Line 3). Given the training set and a code-based interface that specifies how states, actions, and observations are represented, we query an LLM to propose an initial code snippet for a given component (Line 4, see Appendix E.2 for prompt). 1: Input: demonstration dataset $\mathcal{D}$ , initial model $\theta_{\text{prev}}$ , budget $N$ 2: Output: Learned model $\theta_{\text{new}}$ 3: Split $\mathcal{D}$ into $\mathcal{D}_{\text{train}},\mathcal{D}_{\text{test}}$ 4: if $\theta=\emptyset$ then $\theta_{\text{prev}}\leftarrow\text{LLM Init using }\mathcal{D}_{\text{train}}$ $\triangleright$ Appendix E.2 5: $\text{coverage},\mathcal{D}_{\text{errors}}\leftarrow\texttt{Eval}(\theta_{ \text{prev}},\mathcal{D}_{\text{test}},\mathcal{D}_{\text{train}})$ $\triangleright$ Discrepancy between model & empirical 6: $\text{beta}\leftarrow\text{Beta}(1+C\cdot\text{coverage},1+C\cdot(1-\text{ coverage}))$ 7: $\mathcal{M}\leftarrow\{(\theta_{\text{prev}},\text{beta},\mathcal{D}_{\text{ errors}})\}$ 8: while $\text{coverage}<1.0$ and $\text{iterations}<M$ do 9: $(\theta_{\text{new}},\text{Beta}(\alpha,\beta),\mathcal{D}_{\text{errors}}) \leftarrow\text{argmax}_{\mathcal{M}}(p\sim\text{beta})$ $\triangleright$ Thompson sampling 10: $\theta_{\text{new}}^{\prime}\leftarrow$ LLM refinement using $\theta_{\text{new}}$ and $\mathcal{D}_{\text{errors}}$ $\triangleright$ Appendix E.3 11: $\text{coverage}^{\prime},\mathcal{D}_{\text{errors}}^{\prime}\leftarrow\texttt {Eval}(\theta_{\text{new}},\mathcal{D}_{\text{test}},\mathcal{D}_{\text{train}})$ 12: $\text{beta}^{\prime}\leftarrow\text{Beta}(\alpha+C\cdot\text{coverage}^{\prime },\beta+C\cdot(1-\text{coverage}^{\prime}))$ 13: insert $(\theta_{\text{new}}^{\prime},\text{beta}^{\prime},\text{coverage}^{\prime})$ into $\mathcal{M}$ 14: end while 15: return $\text{argmax}_{\mathcal{M}}(\text{coverage})[0]$ $\triangleright$ Return the model with the best overall coverage Algorithm 2 LearnModel Following the initial candidate model proposal, we evaluate them against the empirical conditional probability distributions observed in the demonstration data (Line 5). Determining if a particular outcome is possible under an arbitrary code model is not analytically possible in most probabilistic programming languages, including Pyro, so we use Monte Carlo approximation of the model density. Any empirical sample that is never produced is treated as a failure and recorded in an error set $\mathcal{D}_{\text{errors}}$ . During evaluation, we estimate the model’s coverage on a combination of the training and testing sets, evaluating models based on their ability to generalize beyond the training examples (Line 5). We proceed with an iterative learning procedure that uses a Thompson sampling exploration strategy to build out a tree of candidate models. In addition to a candidate code block, each node contains a Beta distribution capturing uncertainty over its true coverage performance. Initially, the root node contains the LLM’s first code proposal evaluated against the data (Line 7). At each iteration, we select a node to expand using Thompson sampling: we sample from each node’s Beta distribution and pick the node with the highest sampled value (Line 9). The selected node is refined by prompting the LLM with its associated training set coverage mismatch errors $\mathcal{D}_{\text{errors}}$ , encouraging the LLM to generate a corrected or improved version of the model (Line 10, see Appendix E.3 for prompt). This produces a child node with updated code, which is re-evaluated to obtain new train and test coverage statistics (Line 11). The parent’s Beta distribution is updated using a smoothing constant $C$ to encourage stable learning from finite samples, and the child node is added to the tree (Line 13). This iterative process continues until the overall coverage across the empirical distribution is sufficiently high, or until a pre-defined iteration budget is exhausted. Ultimately, we return the candidate model with the highest empirical coverage among all nodes (Line 26). ## Appendix C Belief-Space Planner 1: Input: initial belief $b_{0}$ , models $\mathcal{T}$ , $\mathcal{O}$ , $\mathcal{R}$ , horizon $H$ , hyper-parameters $\lambda,\alpha$ 2: Output: best first action $a^{\ast}$ 3: $\text{Open}\leftarrow\{(b_{0},g{=}0)\}$ , Closed $\leftarrow\emptyset$ , Cost $\leftarrow\{b_{0}:0\}$ 4: while Open $\neq\emptyset$ and iterations $<H$ do 5: $(b,g)\leftarrow\text{pop\_lowest}(\text{Open})$ 6: if $b$ is terminal then 7: continue 8: end if 9: if $b\in$ Closed then 10: continue 11: end if 12: add $b$ to Closed 13: for each action $a\in\mathcal{A}$ do 14: Draw $n$ state particles $s^{\prime}\sim\mathcal{T}(s,a)$ for $s\sim b$ 15: Draw observations $o\sim\mathcal{O}(s^{\prime},a)$ 16: Form child belief $b^{\prime}=\texttt{Branch}(b,a,o)$ 17: $\hat{r}\leftarrow\mathbb{E}_{s,s^{\prime}}[\mathcal{R}(s,a,s^{\prime})]$ 18: $\hat{p}\leftarrow\Pr[o\mid b,a]$ , $\hat{h}\leftarrow H(b^{\prime})$ 19: $g^{\prime}\leftarrow g-\hat{r}-\lambda\log\hat{p}+\alpha\hat{h}+\text{cost}(a)$ 20: if $b^{\prime}\notin$ Cost or $g^{\prime}<\text{Cost}[b^{\prime}]$ then 21: $\text{Cost}[b^{\prime}]\leftarrow g^{\prime}$ 22: insert $(b^{\prime},g^{\prime})$ into Open 23: end if 24: end for 25: end while 26: return first action in the path to the node in Open $\cup$ Closed with minimal $g$ Algorithm 3 BeliefPlanner Our online planning routine conducts forward search directly in belief space, determinizing the stochastic dynamics to enable an A*-style expansion strategy that balances exploitation (reward) and exploration (information gain). Algorithm 3 provides the complete procedure. We begin by inserting the initial belief $b_{0}$ into an open priority queue with zero cost-to-come and initializing an empty closed set (Line 3). Each queue element stores the belief, its cumulative cost $g$ , and bookkeeping metadata such as depth. During each iteration (Line 4), we pop the node with the lowest priority value; if it is terminal or has already been expanded (i.e., in the closed set), we skip further expansion (Lines 6 – 9). Otherwise, for every action $a$ (Line 13), we draw next-state particles from the transition model $\mathcal{T}$ (Line 14) and sample the corresponding observations through the observation model $\mathcal{O}$ (Line 15). Conditioning on the sampled observation yields a child belief $b^{\prime}$ (Line 16). We estimate the expected reward $\hat{r}$ under $\mathcal{R}$ , the likelihood $\hat{p}$ of the observation, and the entropy $\hat{h}$ of $b^{\prime}$ (Lines 17 – 18). These statistics define the child’s cost $$ g^{\prime}\;=\;g\;-\;\hat{r}\;-\;\lambda\log\hat{p}\;+\;\alpha\hat{h}\;+\; \text{cost}(a), $$ where $\lambda$ and $\alpha$ trade off risk sensitivity and information gain (Line 19). The child node is inserted into the queue only if it is not yet discovered or has a lower cumulative cost than a previously seen version (Line 20). The process continues until the queue is empty or a computational budget is exhausted. Finally, we return the first action in the path from $b_{0}$ to the node with the minimum accumulated cost (Line 26), thereby maximizing the composite objective of long-term reward, low risk, and maximal information gathering. While this planner is similar in many ways to POUCT [32], it has the additional feature that enables graph-based search rather than strictly tree-based search, which proved computationally necessary for many of our larger MiniGrid problems. It is important to note that this planner is not optimal, but it suffices for all of the problems we tested, and can easily be substituted for other planning methods in the POMDP Coder framework. ## Appendix D Hyperparameters Table D shows the hyperparameters used per domain for both planning and learning. The planning hyperparameters (in grey) were tuned to work best for the ground truth models used in oracle. With the exception of Thompson smoothing coefficient which was selected based on [37], the other hyperparameters were selected to be as large as possible under computational and budgetary constraints. | Hyperparameter | Classical | MiniGrid | Spot Robot | | --- | --- | --- | --- | | gray!10 Action cost penalty | 0.01 | 0.01 | 0.01 | | gray!10 $\alpha$ (Entropy coefficient) | 0.0 | 0.0 | 1.0 | | gray!10 $\lambda$ (log-prob reward shaping) | 0.1 | 0.1 | 0.1 | | gray!10 Rollouts per stochastic model query | 5 | 1 | 1 | | gray!10 $H$ (Max expansions) | 50 | 5000 | 5000 | | $N$ (Num initial particles), | 50 | 10 | 10 | | Max particle rejuvenations, | 2,500,00 | 500,000 | 25,000 | | $M$ Max refinements | 25 | 25 | 25 | | $C$ Thompson smoothing | 25 | 25 | 25 | | $ND$ (# datapoints shown - initial, E.2) | 5 | 5 | 5 | | $NC$ (# conditions shown - refinement, E.3) | 5 | 5 | 5 | | $NS$ (# samples per condition - refinement, E.3) | 5 | 5 | 5 | ## Appendix E Prompts ### E.1 Function Templates [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python def initial_func(empty_state:MiniGridState): ””” Input: empty_state (MiniGridState): An empty state with only the walls filled into the grid Returns: state (MiniGridState): the initial state of the environment ””” raise NotImplementedError [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python def observation_func(state, action, empty_obs): ””” Args: state (MiniGridState): the state of the environment action (int): the previous action that was taken empty_obs (MiniGridObservation): an empty observation that needs to be filled and returned Returns: obs (MiniGridObservation): observation of the agent ””” raise NotImplementedError [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python def reward_func(state, action, next_state): ””” Args: state (MiniGridState): the state of the environment action (int): the action to be executed next_state (MiniGridState): the next state of the environment Returns: reward (float): the reward of that state done (bool): whether the episode is done ””” raise NotImplementedError [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]python def transition_func(state, action): ””” Args: state (MiniGridState): the state of the environment action (int): action to be taken in state ‘state‘ Returns: new_state (MiniGridState): the new state of the environment ””” raise NotImplementedError ### E.2 Initial Prompt [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]text You are a robot exploring its environment. Environment Description: env_description Goal Description: goal_description Your goal is to model the what_to_model. You need to implement the python code to model the world, as seen in the provided experiences. Please follow the template to implement the code. The code needs to be directly runnable model_input and return model_output. Below are a few samples from the environment distribution. These are only samples from a larger distribution that your should model. exp Here is the template for the model_name function. Please implement the reward function following the template. The code needs to be directly runnable. “‘ code_api code_template “‘ Explain what you believe is the what_to_model in english. Additionally, please implement code to model the logic of the world. Please implement the code following the template. Only output the definition for ‘ model_name ‘. You must implement the ‘ model_name ‘ function. Create any helper function inside the scope of ‘ model_name ‘. Do not create any helper function outside the scope of ‘ model_name ‘. Do not output examples usage. Do not create any new classes. Do not rewrite existing classes. Do not import any new modules from anywhere. Do not overfit to the specific samples. Put the ‘ model_name ‘ function in a python code block. Implement any randomness with ‘pyro.sample‘ ### E.3 Refinement Prompt [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]text You are a robot exploring its environment. env_description Your goal is to model what_to_model of the world in python. You have tried it before and came up with one partially correct solution, but it is not perfect. The observed distribution disagrees with the generated model in several cases. You need to improve your code to come closer to the true distribution. Environment Description: env_description Goal Description: goal_description Here is a solution you came up with before. “‘ code_api code “‘ experiences Explain what you believe is what_to_model in english, then improve your code to better model the true distribution. Please implement the code for the following the template. You must implement the ‘ model_name ‘ function. The code needs to be directly runnable model_input and return model_output. Do not output examples. Do not create any new classes. Do not rewrite existing classes. Do not import any new modules from anywhere. Do not list out specific indices that overfit to the examples, but include ranges. Put the ‘ model_name ‘ function in a python code block. Implement any randomness with ‘pyro.sample‘ The experiences in the refinement prompt are structured as follows. First, we sample conditions for which there is coverage less than 1. For example, in the transition model $P(s_{t+1}|s_{t},a_{t})$ , we search through the set of $(s_{t},a_{t})$ tuples in the database and we find the set of those tuples where the distribution over $(s_{t+1})$ contains at least 1 element that can not be achieved by the LLM-generated model. We select out $NC$ of those conditions. Then, given those conditions, we select $NS$ samples that were not covered by the model to show as examples to the LLM. An example template for a single condition and an $NS=3$ is shown below. The values we used for $NC$ , $NS$ are in Table D. [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]text Here are some samples from the real world that were impossible under your model condition -¿ dataset_outcome condition -¿ dataset_outcome condition -¿ dataset_outcome And here are some samples from your code under the same conditions condition -¿ model_outcome condition -¿ model_outcome condition -¿ model_outcome ### E.4 Direct LLM Baseline Prompt [fontsize=, breaklines, tabsize=4, bgcolor=LightGray]text You are a robot exploring its environment. env_description Your goal is to predict the next best action to take to reach the goal and maximize reward. Here is the template for the reward function. Please implement the reward function following the template. The code needs to be directly runnable on the inputs of (state) and return (reward) in python. “‘ code_api “‘ Here are some example rollouts from the environment exp Here is the current episode history for the task that you are doing right now current_episode Output the next aciton in the form where you fill in ¡action-here¿ with the action that is best for reaching the goal and maximizing reward. For example, your code will look like this: “‘ next_action:int = 0 “‘ The action should be an integer with no additional code. Explan your reasoning in one sentence. ### E.5 Runtime Statistics | Offline Transition Online Transition Offline Reward | $2.00\,\pm\,0.00$ $0.00\,\pm\,0.00$ $2.40\,\pm\,0.24$ | $2.60\,\pm\,0.24$ $0.06\,\pm\,0.06$ $11.80\,\pm\,4.33$ | $2.20\,\pm\,0.20$ $0.00\,\pm\,0.00$ $2.00\,\pm\,0.00$ | $2.00\,\pm\,0.00$ $0.00\,\pm\,0.00$ $2.00\,\pm\,0.00$ | $2.00\,\pm\,0.00$ $0.00\,\pm\,0.00$ $2.20\,\pm\,0.20$ | $2.60\,\pm\,0.24$ $0.00\,\pm\,0.00$ $2.00\,\pm\,0.00$ | $4.60\,\pm\,0.81$ $0.40\,\pm\,0.15$ $2.20\,\pm\,0.20$ | | --- | --- | --- | --- | --- | --- | --- | --- | | Online Reward | $0.05\,\pm\,0.05$ | $0.05\,\pm\,0.05$ | $0.00\,\pm\,0.00$ | $0.00\,\pm\,0.00$ | $0.19\,\pm\,0.09$ | $0.00\,\pm\,0.00$ | $0.00\,\pm\,0.00$ | | Offline Observation | $2.20\,\pm\,0.20$ | $5.80\,\pm\,1.46$ | $10.40\,\pm\,4.01$ | $7.80\,\pm\,4.61$ | $2.40\,\pm\,0.24$ | $17.25\,\pm\,5.71$ | $15.00\,\pm\,3.63$ | | Online Observation | $0.00\,\pm\,0.00$ | $0.00\,\pm\,0.00$ | $0.05\,\pm\,0.05$ | $0.08\,\pm\,0.06$ | $0.00\,\pm\,0.00$ | $0.28\,\pm\,0.17$ | $0.05\,\pm\,0.05$ | | Offline Initial | $2.00\,\pm\,0.00$ | $2.60\,\pm\,0.24$ | $2.00\,\pm\,0.00$ | $2.80\,\pm\,0.37$ | $22.60\,\pm\,3.40$ | $1.75\,\pm\,0.25$ | $18.80\,\pm\,4.59$ | | Online Initial | $0.00\,\pm\,0.00$ | $0.02\,\pm\,0.02$ | $0.00\,\pm\,0.00$ | $0.17\,\pm\,0.09$ | $0.86\,\pm\,0.21$ | $0.06\,\pm\,0.06$ | $0.49\,\pm\,0.15$ | Table 2: The average and standard deviation of the number of nodes, or LLM-generated candidate programs, sampled during the online and offline phases of model learning. | Transition Train Transition Test Reward Train | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $0.99\,\pm\,0.01$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ $1.00\,\pm\,0.00$ | | --- | --- | --- | --- | --- | --- | --- | --- | | Reward Test | $1.00\,\pm\,0.00$ | $0.99\,\pm\,0.01$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | | Observation Train | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $0.92\,\pm\,0.08$ | $0.81\,\pm\,0.19$ | $1.00\,\pm\,0.00$ | $0.94\,\pm\,0.06$ | $0.94\,\pm\,0.06$ | | Observation Test | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $0.92\,\pm\,0.08$ | $0.88\,\pm\,0.12$ | $1.00\,\pm\,0.00$ | $0.93\,\pm\,0.06$ | $0.94\,\pm\,0.06$ | | Initial Train | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $0.68\,\pm\,0.16$ | $0.75\,\pm\,0.25$ | $0.88\,\pm\,0.12$ | | Initial Test | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $1.00\,\pm\,0.00$ | $0.56\,\pm\,0.18$ | $0.75\,\pm\,0.25$ | $0.24\,\pm\,0.19$ | Table 3: The average and standard deviation of coverages achieved after the offline model learning step has completed running split into training and testing coverages.