ASK in the Dark: Uncertainty-Gated LLM Assistance under Partial Observability

Abstract: Reinforcement learning agents operating under partial observability must act on incomplete information, making them natural candidates for guidance from small LLMs (SLMs) that carry broad reasoning priors. Yet integrating SLM guidance into this setting has proven difficult: across all test environments, vanilla uncertainty-gated approaches achieve an overwrite rate at or near zero, meaning the SLM almost never contributes an independent action. We trace this failure to the bare egocentric prompt, which provides insufficient context for genuine reasoning, and identify it as a context problem rather than a capacity problem. We propose ASK+, which supplies the SLM with trajectory-aware context (a partially revealed map, visited positions, and action history) and structured chain-of-thought reasoning, converting it from a passive redundancy check into a more informative consultant that occasionally corrects the policy. We further establish that the predictive entropy signal used for selective querying measures action uncertainty rather than state uncertainty and remains informative in POMDPs, making uncertainty-gated assistance viable beyond fully observable settings. The stateful prompt drives substantial gains: on DoorKey, where vanilla ASK matches PPO (both 89%), ASK+ reaches 93% success; on FourRooms, success climbs from 53% to 70%; on HigherLower, accuracy reaches 73.7%, matching the SLM-only upper bound. Across all environments, Qwen3.5-2B matches or exceeds Qwen3.5-4B, confirming that prompt design and selective gating dominate the impact of model scale, enabling guidance without large models.

• The paper presents ASK+, a framework that enriches prompts with episodic context to enable selective SLM interventions, significantly boosting RL performance in POMDPs.
• It employs Monte Carlo Dropout to gauge uncertainty and trigger corrective actions, achieving improvements such as a success rate increase to 93% in specific tasks.
• Experimental results indicate that precise prompt design, rather than larger model size, is key to effective SLM assistance in overcoming partial observability.

ASK+: Uncertainty-Gated SLM Assistance for RL under Partial Observability

Problem Setting and Motivation

Reinforcement learning (RL) in partially observable Markov decision processes (POMDPs) presents a longstanding challenge due to the inherent ambiguity agents face when only provided with localized, egocentric observations. The inability to reconstruct the complete underlying state solely from observations impedes reliable value estimations and hinders generalization, especially in distributionally shifted settings. Prior efforts have explored leveraging LMs as external consultants, exploiting their pre-trained world knowledge and reasoning capabilities. However, naïve integration—querying LMs based on policy uncertainty—frequently fails to yield performance gains in partially observable environments, as LMs lack sufficient context to reason meaningfully when provided only with local observations. Empirically, this leads to nearly zero overwrite rates, rendering LMs mere passive validators of RL policy choices rather than active, corrective agents.

The ASK+ Framework: Architecture and Methodology

The ASK+ approach advances the uncertainty-gated LM-RL hybrid paradigm by addressing the principal bottleneck: insufficient context provided to the SLM. Concretely, ASK+ augments the vanilla ASK framework by enriching prompts with trajectory-aware, episodic context. The system-level decision mechanism operates as follows: each time step, Monte Carlo Dropout is applied to the PPO actor-critic, producing a distribution over actions whose entropy $H(o)$ captures total predictive (aleatoric + epistemic) uncertainty. If $H(o)$ exceeds a learned threshold $\tau$, the SLM is queried. Crucially, the SLM receives not only the current observation and candidate action but also a history-augmented prompt containing:

• A partially revealed global map
• The agent's position and visitation history
• The chronological sequence of actions taken (episodic memory)
• Structured chain-of-thought, with rationale justification for proposal actions (constrained reasoning budget)

(Figure 1)

Figure 1: Threshold $\tau$ for reward and IR effects, with the selected $\tau$ for each model marked with a line.

This trajectory-informed prompt transforms the SLM's operational mode from reactive pattern matching to context-dependent deliberative reasoning. The SLM can thus resolve spatial ambiguity, track subgoal progress, and selectively intervene with corrective actions.

Experimental Protocol and Benchmarks

ASK+ is evaluated across three environments designed to test complementary aspects of partial observability:

• MiniGrid-FourRooms: Emphasizes spatial grounding and navigation from egocentric, local observations
• MiniGrid-DoorKey: Adds sequential task decomposition (key finding, unlocking, goal reaching)
• POPGym-HigherLower: Requires temporal memory for optimal play in a card prediction task

The base RL policy is a PPO agent; the SLM backbone uses Qwen3.5-2B and Qwen3.5-4B (unmodified, no finetuning). SLM queries are selectively triggered by uncertainty exceeding $\tau$, optimized per environment-model pair via Bayesian search. A suite of ablations systematically explores prompt design, threshold sensitivity, and dependence on PPO quality.

Empirical Results

ASK+ consistently surpasses both PPO-only and vanilla ASK (egocentric-only prompt) baselines in navigation-centric environments. Notably:

• DoorKey: ASK+ achieves success of 93%, compared to 89% for PPO and vanilla ASK—a statistically significant improvement.
• FourRooms: Success climbs from 53% (PPO/vanilla ASK) to 70% with ASK+ when using enriched stateful rationale prompts.
• HigherLower: ASK+ achieves 74% accuracy, matching the SLM-only regime (here, memory advantage is less pronounced compared to navigation tasks).

Importantly, vanilla ASK's overwrite rates are at or near zero, confirming that the SLM rarely corrects (or even disagrees with) the PPO policy in absence of sufficient global context. By contrast, ASK+ reports nonzero overwrite rates (2–6% depending on environment), directly tying these interventions to measurable gains. The results further demonstrate that the 2B SLM matches or exceeds the 4B SLM, underlining that performance is fundamentally bottlenecked by prompt engineering and information structuring rather than model scale.

Ablation studies reveal environment-specific dependencies for prompt design: rationale augmentation drives gains in settings demanding extended planning (FourRooms), while explicit state tracking is essential in sequential subtasks (DoorKey). Intervention rate (IR) and overwrite rate (OR) correlate with entropy threshold $\tau$, confirming the precision of uncertainty gating; IR is higher in tasks with greater observation ambiguity or policy uncertainty.

Implications, Limitations, and Future Research

ASK+ empirically establishes that SLM-guided RL—when equipped with a prompt interface capturing sufficient episodic and spatial context—can provide actionable, uncertainty-driven correction in POMDPs. Importantly, this can be achieved with small LMs, making the method viable for step-wise online inference settings. Key theoretical implications include:

• Decomposition of Generalization: The RL policy is permitted to specialize in in-distribution, low-uncertainty states, while high-uncertainty regions (due to state ambiguity or distributional shift) are delegated to the SLM, which leverages global knowledge embedded during pretraining.
• Separation of Gating and Correction: The combination of MC Dropout-derived entropy for selective intervention and prompt-context for actionable correction is necessary; either fails independently in POMDPs.
• Prompt/Interface as the Primary Bottleneck: Results support the hypothesis that the SLM's performance is fundamentally constrained by prompt informativeness rather than its scale or pretraining regime.

Practical implications suggest the feasibility of deploying uncertainty-gated SLM overlays atop suboptimal or memoryless RL policies, as ASK+ demonstrates graceful degradation and scalable intervention rates in response to policy uncertainty. Future research directions include extending to richer, more complex POMDPs with entropically misfiring gating signals, integrating trainable prompt augmentations, and jointly optimizing recurrent RL backbones with SLM overlays for domains where both trajectory memory and world knowledge are required.

Conclusion

ASK+ demonstrates that uncertainty-gated, trajectory-aware SLM assistance can substantively augment RL agent capabilities under partial observability, provided the prompt interface supplies sufficient episodic and environmental context. Strong numerical improvements in benchmark POMDPs, nonzero correction rates, and the dominance of prompt design over SLM scale collectively underscore the centrality of information structuring in SLM-assisted RL. This work opens further opportunities for robust, low-overhead hybrid intelligence systems in domains where full state observability is unattainable.