Learning to Prompt in Unknown Environments: A POMDP Framework with Compositional Actions for Large Language Models

The most general approach treats the action as a monolithic unit:

Parameters: $O\left({\prod }_{i=1}^n|{𝒜}_i|\right)$

This approach can capture arbitrary correlations between components but suffers from the curse of dimensionality. For even modest component spaces (e.g., 10 choices per component with 5 components), this requires ${10}^5=100,000$ parameters per state.

At the opposite extreme, we assume complete independence:

Parameters: $O\left({\sum }_{i=1}^n|{𝒜}_i|\right)$

While tractable, this assumes that action components are selected independently, which may produce invalid or incoherent combinations.

We propose modeling the action distribution as a Bayesian network, where edges encode conditional dependencies:

where $\text{pa}(a_i)$ denotes the parents of component $a_i$ in the dependency graph.

Parameters: $O\left({\sum }_{i=1}^n|{𝒜}_i|\cdot {\prod }_{j\in \text{pa}(i)}|{𝒜}_j|\right)$

This framework allows us to:

• •

  Encode domain knowledge about component relationships

• •

  Dramatically reduce parameter complexity while preserving essential correlations

• •

  Maintain interpretability of the learned policy

• •

  Apply efficient inference algorithms from graphical model theory

A critical insight motivating our approach is that the LLM’s internal reasoning process constitutes an unknown environment whose dynamics must be learned through interaction. Unlike traditional MDP settings where transition probabilities $P(s^{\prime }|s,a)$ might be known or easily estimated, the LLM’s response to prompts involves:

• •

  Latent reasoning states that are fundamentally unobservable

• •

  Complex, non-stationary dynamics that vary across problem domains

• •

  Partial observability where we only see generated text, not internal computations

• •

  Stochastic transitions due to temperature sampling and inherent model uncertainty

This necessitates a POMDP formulation with model-free learning, where the controller must simultaneously:

1. 1.

   Learn the value of actions through trial and error (Q-learning)

2. 2.

   Maintain beliefs about the current reasoning state

3. 3.

   Adapt to the specific LLM’s behavior patterns without access to its weights

We model the prompting process as a POMDP defined by the tuple $(𝒮,𝒜,𝒪,T,\mathrm{\Omega },R,\gamma ,b_0)$:

• •

  Hidden States $s\in 𝒮$:
  The true state $s_t$ represents the LLM’s latent reasoning configuration, which is fundamentally unobservable. This includes:

  • –

    The LLM’s current attention patterns and activated knowledge

  • –

    Internal problem representation and reasoning paths being considered

  • –

    Confidence distributions over possible solutions

  • –

    Contextual priming effects from the dialogue history

  Since we cannot directly observe $s_t$, we work with observations and belief states.

• •

  Observations $o\in 𝒪$:
  We observe only the LLM’s generated text responses and can extract features from them:

  $$o_t=⟨e_c,f_{\text{struct}},f_{\text{prog}},f_{\text{hist}},f_{\text{meta}},f_{\text{trans}}⟩\in {ℝ}^{948}$$

  where each component provides complementary information:

  • –

    Semantic Embedding $e_c\in {ℝ}^{768}$: Dense vector representation of the current context from a pretrained encoder (e.g., Sentence-BERT), capturing semantic content and relationships.

  • –

    Structural Features $f_{\text{struct}}\in {ℝ}^{20}$:

    • *

      Lexical statistics: token count, average sentence length, vocabulary diversity

    • *

      Content indicators: presence of equations ($\mathrm{𝟙}[\text{has math}]$), code blocks ($\mathrm{𝟙}[\text{has code}]$), lists

    • *

      Question classification: $\text{softmax}(\text{type}\in \{\text{yes/no, MCQ, open, computational}\})$

    • *

      Complexity metrics: Flesch-Kincaid score, parse tree depth

  • –

    Progress Indicators $f_{\text{prog}}\in {ℝ}^{10}$:

    • *

      Normalized step count: $t/T_{\max }$

    • *

      Token utilization: $\text{tokens\_used}/\text{token\_budget}$

    • *

      Reasoning depth: number of completed compositional actions

    • *

      Solution proximity: estimated completion $\in [0,1]$

  • –

    Historical Context $f_{\text{hist}}\in {ℝ}^{100}$:

    • *

      Previous action encoding: $\text{embed}(a_{t-1})\in {ℝ}^{40}$

    • *

      Compressed previous embedding: $\text{PCA}(e_{c,t-1})\in {ℝ}^{50}$

    • *

      Action history summary: $\text{mean}({\{a_{t-k}\}}_{k=1}^K)\in {ℝ}^{10}$

  • –

    Meta-cognitive Features $f_{\text{meta}}\in {ℝ}^{20}$:

    • *

      Domain classification: $\text{softmax}(\text{domain}\in \{\text{math, code, reasoning, creative}\})$

    • *

      Confidence scores from auxiliary networks

    • *

      Coherence metric: $\text{coherence}(c_t,c_{t-1})$

    • *

      Perplexity delta: change in LLM perplexity

  • –

    Transition Features $f_{\text{trans}}\in {ℝ}^{30}$:

    • *

      Semantic shift: $\cos (e_c,e_{c,t-1})$

    • *

      Progress delta: $f_{\text{prog},t}-f_{\text{prog},t-1}$

    • *

      Complexity change: $\mathrm{\Delta }\text{complexity}$

    • *

      Information gain: KL divergence between consecutive states

  This observation vector of 948 dimensions provides partial information about the hidden state.

• •

  Belief State $b_t\in \mathrm{\Delta }(𝒮)$:
  Since the true state is hidden, we maintain a belief distribution over possible states:

  $$b_t(s)=P(s_t=s|o_{1:t},a_{1:t-1})$$

  In practice, we approximate the belief state using a recurrent neural network that processes the observation history:

  $$b_t=\text{RNN}(o_t,a_{t-1},b_{t-1})\in {ℝ}^{256}$$

  This learned belief representation captures:

  • –

    Uncertainty about the LLM’s current reasoning state

  • –

    Historical context that influences future responses

  • –

    Patterns in how the LLM responds to different action sequences

• •

  Compositional Actions $a\in 𝒜$:
  Our action space leverages the compositional prompting framework (rlpilot package) which provides a rich, extensible set of prompting operations. Each action $a$ is a tuple of linguistic components:

  $$a=(\omega ,\varphi ,\sigma ,\kappa ,\tau ,\mu ,\rho ,\delta )\in \mathrm{\Omega }\times \mathrm{\Phi }\times \mathrm{\Sigma }\times K\times T\times M\times C\times D$$

  The core components from our implementation include:

  • –

    $\omega \in \mathrm{\Omega }$ - Cognitive Operations: {decompose, analyze, generate, verify, synthesize, abstract} plus extended operations

  • –

    $\varphi \in \mathrm{\Phi }$ - Focus Aspects: {structure, constraints, patterns, solution, correctness, efficiency, examples, relationships, alternatives, edge_cases}

  • –

    $\sigma \in \mathrm{\Sigma }$ - Reasoning Styles: {systematic, creative, critical, formal, intuitive}

  • –

    $\kappa \in K$ - Connection Types: {therefore, however, building_on, alternatively, verify, continue}

  • –

    $\tau \in T$ - Output Formats: {steps, list, mathematical, narrative, code, solution, explanation, table}

  • –

    $\mu \in M$ - Meta-Strategies: {forward_chaining, backward_chaining, means_ends_analysis, case_based, analogical, constraint_satisfaction}

  • –

    $\rho \in C$ - Confidence Levels: {very_low, low, medium, high, very_high, certain}

  • –

    $\delta \in D$ - Reasoning Depths: {surface, shallow, moderate, deep, profound}

  The compositional framework supports:

  • –

    Fluid API construction: Actions are built programmatically, enabling dynamic prompt generation

  • –

    LLM-augmented actions: Meta-reasoning capabilities for adaptive prompt refinement

  • –

    Parallel execution: Multiple prompting strategies can be explored simultaneously

  • –

    RAG integration: Retrieved examples can be dynamically incorporated

  This yields a potential action space of over 10 million combinations, but our Bayesian network factorization (described below) makes learning tractable.

  Retrieval and In-Context Learning Actions

  Beyond the compositional prompting dimensions, our action space includes sophisticated retrieval mechanisms for incorporating external knowledge and examples:

  1. Graph-Based Example Retrieval: Our framework models examples as nodes in a knowledge graph, enabling advanced retrieval strategies:

  • –

    Hub identification: Finding highly-connected examples that serve as conceptual anchors

  • –

    Bridge detection: Identifying examples that connect different knowledge communities

  • –

    Community clustering: Grouping related examples for coherent few-shot learning

  • –

    Diverse sampling: Selecting examples from different communities to maximize coverage

  Each retrieval action $a_{ret}$ includes:

  • –

    $n\in \{1,2,\mathrm{\dots },10\}$ - Number of examples for ICL

  • –

    $\text{strategy}\in$ {similar, diverse, hub_aware, community, bridge}

  • –

    $\text{threshold}\in [0,1]$ - Similarity threshold for relevance filtering

  2. Classical Knowledge Base Lookup: Fast, deterministic retrieval from structured sources:

  • –

    Wikipedia API: Factual information injection with latency 100ms

  • –

    Mathematical databases: Theorem and formula lookup

  • –

    Code repositories: Algorithm and implementation patterns

  • –

    Domain ontologies: Hierarchical concept relationships

  These lookup actions provide predictable sub-environments within the unknown LLM dynamics—we know exactly what will be retrieved, even if we don’t know how the LLM will use it.

  Bayesian Network Structure:
  Rather than treating this as a fully joint or fully independent distribution, we impose a Bayesian network structure based on linguistic and cognitive principles:

  $$\pi (a|s)=P(\omega ,\varphi ,\sigma |s)\cdot P(\kappa |s,\omega ,a_{t-1})\cdot P(\tau |s,\omega ,\varphi )$$

  This factorization reflects that:

  • –

    The cognitive operation ($\omega$), focus ($\varphi$), and style ($\sigma$) form a semantic unit that should be chosen jointly

  • –

    The connection type ($\kappa$) depends on the current operation and previous action for coherence

  • –

    The output format ($\tau$) is constrained by the operation and focus

  Complexity Reduction:
  This structure reduces parameters from:

  • –

    Fully joint: $30,720$ parameters per state (without retrieval actions)

  • –

    With retrieval: Over ${10}^7$ when including all retrieval strategies

  • –

    Our factorization: $(10\times 8\times 6)+(8\times 10)+(8\times 10\times 8)+250=480+80+640+250=1,450$ parameters

  Over 6,800× reduction from the full space while preserving essential correlations!

  Graph-Based Example Network

  The retrieval actions operate on a dynamically-maintained graph $G=(V,E)$ implemented through the complex-network-rag package. The graph construction proceeds as follows:

  Graph Construction:

  • –

    Vertices $V$: Examples with embeddings $e_i\in {ℝ}^d$ (typically $d=384$ for sentence transformers or $d=768$ for BERT-based models) and metadata $m_i$

  • –

    Edges $E$: Weighted connections where $w_{ij}=\cos (e_i,e_j)$ if $\cos (e_i,e_j)\ge {\tau }_{\text{min}}$ (typically ${\tau }_{\text{min}}=0.7$)

  • –

    Strong edges: Subset $E_s\subseteq E$ where $w_{ij}\ge {\tau }_{\text{strong}}$ (typically ${\tau }_{\text{strong}}=0.8$)

  Network Analysis Metrics:

  1. 1.

     Hub Centrality (Degree): $h(v)=\mathrm{deg}(v)=|\{u:(v,u)\in E\}|$ - High-degree nodes serve as conceptual anchors with broad applicability

  2. 2.

     Bridge Centrality (Betweenness): $b(v)={\sum }_{s\ne v\ne t}\frac{{\sigma }_{st}(v)}{{\sigma }_{st}}$ where ${\sigma }_{st}$ is the number of shortest paths from $s$ to $t$, and ${\sigma }_{st}(v)$ is the number passing through $v$. Bridge nodes connect disparate knowledge domains.

  3. 3.

     Community Structure: Louvain modularity maximization (Blondel et al., 2008) partitions $V$ into communities $C=\{c_1,\mathrm{\dots },c_k\}$ by optimizing:

     $$Q=\frac{1}{2m}\sum _{ij}\left[A_{ij}-\frac{k_i{k}_j}{2m}\right]\delta (c_i,c_j)$$

     where $m=|E|$, $A_{ij}$ is the adjacency matrix, $k_i=\mathrm{deg}(v_i)$, and $\delta (c_i,c_j)=1$ if nodes share a community.

  4. 4.

     k-Hop Neighborhoods: $N_k(v)=\{u:d(v,u)\le k\}$ provides local context around retrieved examples

  Retrieval Strategies: The Q-network learns to select among five retrieval strategies based on the current belief state:

  • –

    Similarity: Pure cosine similarity ranking - $\text{rank}(v)=\cos (e_q,e_v)$

  • –

    Community-aware: Find most similar node’s community, then rank within that community for coherent examples

  • –

    Bridge-focused: Prioritize high-betweenness nodes when crossing domains

  • –

    Hub-focused: Prioritize high-degree nodes for general problem understanding

  • –

    Hybrid: Combine similarity with network position - $\text{score}(v)=\alpha \cos (e_q,e_v)+\beta \frac{h(v)}{{\max }_{u}h(u)}+\gamma \frac{b(v)}{{\max }_{u}b(u)}$

  The policy learns which strategy to employ based on problem characteristics encoded in the belief state, discovering patterns such as using hub examples early in reasoning and bridge examples when integrating knowledge from multiple domains.

  Prompt Construction Function:
  Given an action tuple $a$ sampled from this distribution and current context $c$, a deterministic function $g:A\times C\to P$ constructs the actual prompt:

  $$p=g(a,c)=f_{\kappa }(c_{prev})\circ f_{\omega ,\varphi ,\sigma }(c_{current})\circ f_{\tau }$$

  where $f_{\kappa }$ generates the connection phrase, $f_{\omega ,\varphi ,\sigma }$ constructs the main directive, and $f_{\tau }$ specifies the output format.

• •

  Unknown Transition Function $T(s,a,s^{\prime })$:
  The transition dynamics $P(s^{\prime }|s,a)$ governing how the LLM’s hidden state evolves are unknown and must be learned through interaction. This is a fundamental departure from classical planning where models are given. The transitions depend on:

  • –

    The specific LLM’s training data and architecture (which we don’t have access to)

  • –

    Complex interactions between the prompt and the LLM’s learned representations

  • –

    Stochastic sampling procedures within the LLM

  Since we cannot model $T$ directly, we adopt a model-free Q-learning approach that learns the value of actions without explicitly modeling the transition dynamics.

• •

  Observation Function $\mathrm{\Omega }(o|s^{\prime },a)$:
  The observation function maps hidden states to observations (LLM outputs). We observe features extracted from the generated text:

  $$P(o|s^{\prime },a)=\text{LLM\_Generate}(s^{\prime },\text{prompt}(a))$$

  The observation includes both the raw text and extracted features. Importantly, this function is also unknown—we don’t know how the LLM will respond to a given prompt until we try it.

• •

  Reward $R(s,a,s^{\prime })$:
  Our reward structure primarily focuses on terminal rewards:

  • –

    Terminal Rewards: Significant reward upon reaching a state where the LLM produces an accurate solution, validated against ground-truth answers or evaluation metrics.

  • –

    Step Penalties: Small negative rewards for each action taken, encouraging efficiency in the reasoning process.

  Future Directions - Implicit Rewards:

  While our current implementation relies on sparse terminal rewards, we envision extending our framework with implicit rewards inspired by program synthesis systems like DreamCoder (Ellis et al., 2021). These could include:

  • –

    Compression Rewards: Incentives for representing information more concisely without loss of utility.

  • –

    Reuse Rewards: Bonuses for leveraging previously stored reasoning traces or code snippets from memory.

  • –

    Simplification Rewards: Rewards for reducing complex reasoning paths or code into more elegant, generalizable forms.

  • –

    Compositional Rewards: Encouraging the decomposition of complex functions into compositions of simpler, reusable components.

  Such intrinsic motivations could potentially lead to more efficient and generalizable policies, particularly for domains where explicit feedback is limited or costly to obtain. This aligns with the broader goal of developing systems that not only solve individual problems but also build increasingly powerful abstractions over time (Lake et al., 2017).

The belief state update follows Bayes’ rule:

where $\eta$ is a normalization constant. However, since both $T$ and $\mathrm{\Omega }$ are unknown, we approximate the belief update using a learned recurrent model:

where:

• •

  GRU is a Gated Recurrent Unit that learns to track belief dynamics

• •

  $\varphi (o_{t+1})$ is the observation feature vector (948 dimensions)

• •

  $\text{embed}(a_t)$ is a learned embedding of the compositional action

• •

  $b_t\in {ℝ}^{256}$ is the compressed belief representation

Our observation extraction preserves critical information about the hidden state:

1. 1.

   Multi-scale Feature Extraction: We extract features at multiple granularities from the LLM’s output:

   • •

     Semantic embeddings ($e_c$) capture high-level meaning

   • •

     Structural features ($f_{\text{struct}}$) identify problem characteristics

   • •

     Progress indicators ($f_{\text{prog}}$) track solution advancement

   • •

     Meta-cognitive features ($f_{\text{meta}}$) assess reasoning quality

2. 2.

   Temporal Context Integration: The belief state maintains a compressed history:

   $b_t=f(o_t,o_{t-1},\mathrm{\dots },o_1,a_{t-1},\mathrm{\dots },a_1)$

   (3)

   This allows the controller to recognize patterns in how the LLM responds to action sequences, even though the true state transitions are unknown.

3. 3.

   Uncertainty Quantification: The belief state implicitly represents uncertainty about:

   • •

     Which reasoning paths the LLM is considering

   • •

     How close the LLM is to a solution

   • •

     What types of prompts will be most effective next

4. 4.

   Adaptive Learning: As the controller interacts with the LLM, the belief update model learns:

   • •

     Patterns in how different LLMs respond to compositional actions

   • •

     Domain-specific dynamics (e.g., mathematical vs. coding problems)

   • •

     Individual LLM quirks and biases

The POMDP formulation with unknown dynamics necessitates model-free reinforcement learning:

• •

  No Transition Model: We cannot predict how the LLM will respond to prompts

• •

  No Observation Model: We cannot predict what features we’ll observe

• •

  Learning Through Interaction: Q-values are learned purely from experience

This is why Q-learning is particularly suitable—it learns the value of actions without requiring a model:

where $b$ and $b^{\prime }$ are belief states rather than true states. The Q-function learns which compositional actions are valuable in different belief contexts, without ever needing to know the true state or dynamics.

The Q-learning update in our POMDP setting is:

Crucially, this update does not require knowledge of the transition or observation functions. The controller learns purely from:

• •

  The sequence of observations (LLM outputs)

• •

  The rewards received (task success/failure)

• •

  The actions taken (compositional prompts)

This model-free approach enables the controller to adapt to any LLM without access to its internals, learning its specific behavior patterns through interaction.

Our framework can integrate with MCTS in several ways, each with distinct advantages:

**Option 1: Policy-Prior MCTS** If we train our network to output action probabilities $\pi (a|s)$, we can use these as priors in MCTS:

**Option 2: Value-Function MCTS** Alternatively, we could train the network to directly estimate $V(s)$ or $Q(s,a)$, using these as initialization for MCTS node values. This provides better starting estimates but loses the explicit action guidance.

**Option 3: Hybrid Approach** We could learn both: $Q(s,a)$ for value estimation and $\pi (a|s)$ for action priors, though this requires more complex training.

The choice depends on what signal is easier to learn from limited data. $Q(s,a)$ is most universal since $V(s)={\max }_{a}Q(s,a)$ and $\pi (a|s)\propto \exp (Q(s,a)/\tau )$, but direct policy learning might be more sample-efficient for our structured action space.

**MCTS Policy vs. Learned Policy: A Fundamental Trade-off**

An important insight is that MCTS naturally produces policies through search without explicit learning:

where $N(s,a)$ is the visit count of action $a$ from state $s$ after running the full search tree. This creates a fundamental speed-accuracy trade-off:

• •

  MCTS Policy: Computationally expensive (full tree search per decision) but adapts precisely to each state through targeted exploration

• •

  Learned Policy: Fast inference (single forward pass) but relies on generalization from training data, potentially missing state-specific nuances

Interestingly, both approaches involve ”learning” - MCTS learns through visit count accumulation during search, while neural policies learn through parameter updates during training. The key difference is when the learning occurs: MCTS learns at inference time through search, while neural networks learn at training time through experience replay.

This suggests several hybrid strategies:

• •

  Policy Distillation: Train the learned policy to mimic MCTS policies by using ${\pi }_{\text{MCTS}}(a|s)$ as training targets

• •

  Adaptive Deployment: Use fast learned policy for routine problems, expensive MCTS for challenging ones

• •

  Search-Augmented Training: Use MCTS-discovered action sequences as high-quality training data

**Critical Challenge: Terminal Node Evaluation** A key challenge in applying MCTS to reasoning is recognizing and evaluating terminal nodes. Unlike game trees with clear win/loss outcomes, reasoning chains require:

1. **Terminal Detection**: When has the reasoning process concluded? This could be: - Fixed depth limit - LLM-based termination decision: ”Is this reasoning chain complete?” - Confidence threshold on solution quality

2. **Terminal Evaluation**: How good is the final reasoning chain? Options include: - Ground truth comparison (when available) - LLM-as-judge evaluation: ”How correct/convincing is this solution?” - Multi-metric evaluation (correctness, clarity, efficiency) - Human evaluation (expensive but high-quality)

Our compositional actions *may* help with terminal evaluation by making the reasoning process more interpretable, but this advantage requires empirical validation.

Our factored action space *potentially* enables more structured search strategies, though their effectiveness compared to well-designed atomic action spaces remains an empirical question:

1. Component-wise Beam Search: Instead of expanding the full action space, we could beam search through components sequentially:

• •

  Beam search over $\omega$ (cognitive operations)

• •

  For each ${\omega }^{\ast }$, beam search over $\varphi$ given $P(\varphi |s,{\omega }^{\ast })$

• •

  Continue through $\sigma ,\kappa ,\tau$ in dependency order

This could reduce search complexity from $O({|𝒜|}^d)$ to $O({\sum }_i{|{𝒜}_i|}^d)$ where $d$ is search depth, but only if the factorization assumptions hold and the sequential dependencies are meaningful.

2. Hierarchical Reasoning Trees: We can structure search hierarchically: - **Strategic level**: Search over high-level reasoning approaches ($\omega ,\varphi$) - **Tactical level**: For each strategy, search over implementation details ($\sigma ,\kappa ,\tau$)

This mirrors human reasoning: first deciding ”how to approach the problem” then ”how to execute that approach.”

3. Multi-Path Consistency Checking: Our compositional actions enable exploring multiple reasoning paths simultaneously and checking for consistency:

“‘ Path 1: (decompose, structure, systematic, …) Path 2: (analyze, constraints, formal, …) Path 3: (generate, solution, direct, …) “‘

If multiple paths converge to the same answer with high confidence, we gain additional validation.

The finite context window $W$ creates natural budget constraints for search. We can adaptively allocate search budget based on:

• •

  **Problem complexity**: Estimated from $f_{\text{meta}}$ features

• •

  **Current uncertainty**: Entropy of policy distribution $H(\pi (\cdot |s))$

• •

  **Progress indicators**: Change in $f_{\text{prog}}$ features

• •

  **Remaining budget**: Available tokens in context window

A particularly promising direction is using LLMs not just as reasoning engines, but as meta-reasoners about their own prompting strategies. Our compositional action space enables several forms of LLM-guided enhancement:

1. Context-Adaptive Action Modification: Instead of using fixed compositional templates, we can ask the LLM to adapt actions to specific contexts:

”Given this mathematical problem: [context]
Here are potential approaches: (decompose, structure, systematic, therefore, steps)
How should this be modified for this specific problem type?”

This allows dynamic prompt engineering based on problem characteristics.

2. LLM-Based Action Evaluation: Rather than relying solely on visit counts, we can use LLM judgment to evaluate action quality:

where $\text{LLMScore}(s,a)$ asks the LLM: ”Rate how appropriate this prompting strategy is for the current reasoning state (1-10).”

3. Compositional Strategy Suggestion: The LLM can propose entirely new action combinations:

”Given the current context, what values of $(\omega ,\phi ,\sigma ,\kappa ,\tau )$ would be most effective? Explain your reasoning.”

4. Three-Layer Reasoning Architecture: This suggests a hybrid system with complementary strengths:

• •

  MCTS Layer: Strategic exploration through principled search, handles long-term planning

• •

  Neural Network Layer: Fast pattern recognition and learned priors, provides efficiency

• •

  LLM Meta-Reasoning Layer: Context-specific adaptation and strategy refinement, adds flexibility

The compositional action space serves as a common vocabulary enabling all three layers to communicate. The neural network provides fast priors, MCTS provides systematic exploration, and the LLM provides contextual adaptation.

Potential Implementation: 1. NN suggests initial action probabilities based on state features 2. LLM refines these suggestions based on full context understanding 3. MCTS uses both as priors for systematic search 4. Visit counts and outcomes update both NN and LLM guidance strategies

This architecture could combine the best aspects of each approach: NN speed, MCTS systematicity, and LLM contextual intelligence.

This search integration connects our work to recent advances in search-based reasoning:

Tree of Thoughts (Yao et al., 2023): Our approach generalizes ToT by: - Using learned policies instead of hand-crafted heuristics - Providing finer-grained action spaces beyond ”continue/backtrack” - Enabling compositional combination of reasoning strategies

ReAct (Yao et al., 2023): Our $\omega$ component can include ”tool_use” operations similar to ReAct’s action space, but within a more general compositional framework that handles reasoning + acting as special cases.

Graph of Thoughts (Besta et al., 2024): Our compositional actions can construct arbitrary reasoning graphs, not just trees, by allowing actions that reference and combine previous reasoning steps.

Importantly, search and policy learning create a virtuous cycle:

• •

  **Search → Better Training Data**: MCTS exploration discovers high-reward action sequences that provide better training targets for the policy

• •

  **Better Policy → More Efficient Search**: Improved policy provides better priors, reducing search budget needed for good performance

• •

  **Compositional Transfer**: Learning about component effectiveness in one search tree transfers to other trees with shared components

This suggests a training curriculum: 1. **Bootstrap Phase**: Learn basic policy from simple, direct reasoning 2. **Search-Augmented Phase**: Use policy-guided search to discover complex reasoning chains 3. **Distillation Phase**: Train policy to directly predict successful search outcomes 4. **Deployment Phase**: Use lightweight policy for easy problems, full search for hard problems

The compositional structure makes this particularly effective because component-level learning transfers across all phases, unlike end-to-end approaches that must relearn everything for each level of sophistication.

Our external controller cannot access the LLM’s internal representations or gradients, making traditional pre-training approaches like minimizing perplexity infeasible. Instead, we must evaluate complete reasoning chains based on their outcomes. This constraint, while limiting, has an important benefit: it forces us to optimize for what actually matters - correctness and efficiency - rather than proxy metrics like token likelihood.

This approach resonates with Sutton’s ”Bitter Lesson” (Sutton, 2019): rather than encoding specific reasoning strategies, we let the controller discover effective prompting patterns through trial and error, guided only by reward signals.

The key to successful training lies in rich, multi-scale reward signals that address the credit assignment problem inherent in sequential decision-making:

1. Terminal Rewards (Correctness Evaluation):

Evaluation methods depend on the domain:

• •

  Mathematics: Symbolic math solvers (SymPy) for verification

• •

  Code: Unit test execution and output matching

• •

  Reasoning: LLM-as-judge with calibrated confidence scores

• •

  Factual QA: Exact match or semantic similarity metrics

2. Intermediate Progress Rewards (and the Goodhart’s Law Dilemma):

To address the credit assignment problem, we provide dense rewards throughout the reasoning chain. However, this introduces a fundamental tension: we need intermediate signals to guide learning, but any proxy metric we optimize can be gamed (Goodhart’s Law). Our approach is to use a diverse portfolio of weakly correlated signals:

where components include:

• •

  Semantic coherence: $r_{\text{coh}}=\cos (e_{c,t},e_{c,t+1})$ - Maintains logical flow

• •

  Information gain: $r_{\text{info}}=H(s_t)-H(s_{t+1})$ - Rewards uncertainty reduction

• •

  Complexity reduction: $r_{\text{simp}}=\text{complex}(s_t)-\text{complex}(s_{t+1})$

• •

  Goal proximity: $r_{\text{goal}}=\text{dist}(s_{t+1},s_{\text{goal}})-\text{dist}(s_t,s_{\text{goal}})$

Critically, the weights $w_i(t)$ decrease over training time, gradually shifting emphasis from these proxies to terminal correctness. This curriculum prevents the model from overfitting to intermediate metrics while still providing early training signal.

3. Efficiency Penalties:

These penalties encourage concise, efficient reasoning paths.

Building reliable evaluation infrastructure is crucial for generating consistent reward signals:

1. Synthetic Data Generation:

• •

  Generate problems with known solutions using rule-based systems

• •

  Create compositional tasks by combining simpler verified components

• •

  Use data augmentation to increase diversity while maintaining correctness guarantees

2. LLM-as-Judge Framework:

For tasks without deterministic evaluation, we employ carefully calibrated LLM judges:

where we aggregate scores from multiple models to reduce bias.

3. Verification Pipelines:

• •

  Mathematical: SymPy for symbolic verification, numerical solvers for approximation

• •

  Code: Sandboxed execution with test suites

• •

  Logic: SAT/SMT solvers for formal verification

• •

  Factual: Knowledge base lookups and cross-referencing

Our training employs a modified Q-learning algorithm with prioritized experience replay:

The credit assignment problem - determining which actions contributed to eventual success or failure - is particularly acute in our setting where reasoning chains can span many steps. We address this through:

1. Reward Shaping: Carefully designed intermediate rewards that align with progress toward the goal while avoiding reward hacking. However, we must remain cognizant of Goodhart’s Law (Goodhart, 1975): ”When a measure becomes a target, it ceases to be a good measure.” For instance, optimizing solely for semantic coherence might lead to verbose but ultimately incorrect reasoning chains. To mitigate this:

• •

  Use multiple uncorrelated reward signals that are difficult to game simultaneously

• •

  Maintain a large weight on terminal correctness rewards relative to intermediate signals

• •

  Periodically audit learned behaviors for degenerate strategies

• •

  Employ reward curriculum that shifts emphasis from intermediate to terminal rewards over training

2. Hindsight Experience Replay: When a trajectory fails, we can still learn by relabeling it with alternative goals it did achieve.

3. Monte Carlo Returns: For stable tasks, we can use full episode returns:

4. Temporal Difference Learning: For faster learning, we bootstrap from value estimates:

Solomonoff’s theory of universal induction establishes that optimal prediction requires finding the shortest program that generates observed data. In our context, effective prompting means finding minimal representations that preserve task-relevant information:

where $K(\pi )$ is the Kolmogorov complexity of the policy and the second term measures prediction accuracy. This formulation reveals why our compositional action space is theoretically motivated: by factorizing actions into reusable components, we achieve compression through modularity.

The connection to the Minimum Description Length (MDL) principle (Rissanen, 1978) suggests that our policy should learn to:

• •

  Identify regularities in reasoning patterns

• •

  Compress these patterns into reusable action components

• •

  Apply compressed representations to novel problems

Given finite context windows, optimal prompting requires principled compression that preserves task-relevant information while discarding redundancy. We formalize this as:

where $I(\cdot ;\cdot )$ denotes mutual information with the target output. This constraint ensures compression doesn’t lose critical information.

In practice, our context management operator $\psi$ approximates this ideal through learned compression strategies:

Each strategy represents a different compression philosophy:

• •

  Abstractive compression: Identifies high-level patterns, discarding implementation details

• •

  Extractive compression: Selects most informative segments

• •

  Lossy compression: Accepts controlled information loss for greater compression

We modify our Q-learning objective to explicitly reward compression:

This augmented Q-value rewards actions that maximize information gain per unit of complexity. The hyperparameter $\lambda$ controls the trade-off between task performance and compression efficiency.

During training, we estimate $K(a)$ using practical proxies:

• •

  Token count: $K_{\text{tokens}}(a)=|\text{tokens}(a)|$

• •

  Template complexity: $K_{\text{template}}(a)=-\log P(a|𝒜)$

• •

  Semantic novelty: $K_{\text{semantic}}(a)=\text{dist}(a,\overline{a})$

Our compositional action space enables information-theoretic action selection that maximizes expected information gain:

where $H(Y|s)$ is the conditional entropy of the target given state $s$. This formulation naturally balances exploration (high uncertainty reduction) with exploitation (moving toward low-entropy states).

The Bayesian network structure of our action space provides computational advantages:

This decomposition allows us to estimate information gain for each component independently, dramatically reducing computational complexity.

We can view our entire framework through the lens of Solomonoff’s universal prior, which assigns higher probability to simpler (more compressible) hypotheses:

This prior naturally induces a hierarchy of prompting strategies:

1. 1.

   Level 0: Direct answering (minimal complexity)

2. 2.

   Level 1: Single-step decomposition

3. 3.

   Level 2: Multi-step reasoning with tool use

4. 4.

   Level 3: Recursive decomposition with backtracking

Our policy learns to select the minimum complexity level sufficient for each problem, embodying Occam’s Razor in automated reasoning.

The compression perspective reveals several critical training challenges:

1. Exploration-Compression Trade-off: Random exploration in compositional spaces is inefficient due to the $|\mathrm{\Omega }|\times |\mathrm{\Phi }|\times |\mathrm{\Sigma }|\times |K|\times |T|$ action space. Instead, we employ compression-guided exploration:

where $\beta$ penalizes complex actions during exploration.

2. Meta-Learning Compression Operators: The optimal compression strategy is task-dependent. We employ meta-learning to discover compression operators:

where $\tau$ represents task distributions and $ℒ$ measures compression-performance trade-off.

3. Curriculum Learning via Complexity: We structure training curricula by Kolmogorov complexity:

• •

  Begin with problems requiring minimal compression

• •

  Gradually increase complexity, forcing sophisticated compression

• •

  Transfer learned compression strategies across domains

This approach aligns with the ”Bitter Lesson” while maintaining theoretical grounding: we’re not hand-coding compression strategies but learning them through complexity-aware optimization.

We employ an enriched 948-dimensional state representation:

This captures semantic content (768D embedding), problem structure, reasoning progress, interaction history, metadata, and transition features.

The neural network learns Q-values for compositional actions using the Bayesian network factorization: