🤖 AI

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

This paper introduces ASK+, a framework that enhances reinforcement learning agents in partially observable environments by replacing ineffective bare prompts with trajectory-aware context and chain-of-thought reasoning, thereby enabling small language models to provide meaningful, uncertainty-gated guidance that significantly outperforms vanilla approaches and scales efficiently without requiring large models.

Original authors: Juarez Monteiro, Nathan Gavenski, Guilherme Lima, Francisco Galuppo, Odinaldo Rodrigues, Adriano Veloso

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Juarez Monteiro, Nathan Gavenski, Guilherme Lima, Francisco Galuppo, Odinaldo Rodrigues, Adriano Veloso

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are playing a video game where you can only see a small circle of light around your character. Everything outside that circle is pitch black. You don't know if there's a wall, a treasure chest, or a monster just a few steps away. This is what computer scientists call partial observability.

In this paper, the authors try to solve a problem: How do you help a computer "player" (an AI agent) make good decisions when it can't see the whole picture?

The Problem: The "Blind" Player and the "Amnesiac" Guide

The authors use two types of AI:

The Player (RL Agent): A robot trained to play the game. It's good at what it knows, but if the game changes slightly (like a door moving), it gets confused because it only remembers what it saw right now.
The Guide (Small Language Model - SLM): A smart, knowledgeable assistant (like a mini-ChatGPT) that knows general rules about the world. It can reason, "If I see a key, I probably need a door."

The Failed Attempt (Vanilla ASK):
Previously, researchers tried to let the Guide help the Player only when the Player was "uncertain." They built a system called ASK.

The Setup: The Player looks at the dark screen and says, "I'm not sure what to do."
The Mistake: The researchers only showed the Guide the exact same tiny, dark circle the Player saw.
The Result: The Guide was just as blind as the Player. It couldn't figure out the solution because it lacked context. It would just say, "Do what you think is best," effectively doing nothing. The Guide was like a tour guide standing in a dark room with you, unable to see the map either.

The Solution: ASK+ (The "Memory-Enhanced" Guide)

The authors realized the Guide wasn't "dumb"; it just wasn't given enough information. They created ASK+, which fixes the problem by changing what they show the Guide.

Instead of just showing the Guide the current dark screen, ASK+ gives it a Travel Log. This log includes:

The Map: A partially revealed map showing where the player has been.
The History: A list of moves the player has already made.
The Context: "You are in a room, you found a key, and you tried to open a locked door."

The Analogy:
Think of the Player as a hiker in a foggy forest.

Vanilla ASK: The hiker asks a friend, "What should I do?" The friend is standing right next to them in the fog, seeing nothing. The friend says, "I don't know, you decide."
ASK+: The hiker asks the friend, "What should I do?" But this time, the friend has a backpack full of notes. The notes say: "We started at the trailhead, walked north for 10 minutes, found a red rock, and now we are facing a cliff." With this history, the friend can say, "Ah, you're at the cliff! Turn left, there's a path you missed earlier."

How It Works (The "Uncertainty Gate")

The system uses a clever "gatekeeper" mechanism:

The Player tries to make a move.
The system checks: "Is the Player confused?" (It measures this using a mathematical signal called entropy).
If the Player is confident: It keeps going. No help needed.
If the Player is confused: It opens the gate and asks the Guide.
The Guide, now holding the Travel Log (the map and history), thinks carefully and says, "Actually, don't go that way. Go here instead."

The Results: Small Brains, Big Wins

The authors tested this on three different "games":

FourRooms: Navigating a maze.
DoorKey: Finding a key to unlock a door.
HigherLower: Guessing card ranks based on memory.

Key Findings:

Context is King: The "Travel Log" (the prompt) was the most important part. Without it, the Guide did nothing. With it, the Guide fixed the Player's mistakes.
Small is Beautiful: They tested a "small" Guide (2 billion parameters) and a "bigger" Guide (4 billion parameters). Surprisingly, the smaller Guide performed just as well as the bigger one. This proves that giving the AI the right information (the prompt) matters more than making the AI bigger.
Success Rates:
- In the maze game, success jumped from 53% to 70%.
- In the key/door game, success went from 89% to 93%.
- In the card game, the Guide matched the best possible performance.

The Bottom Line

The paper argues that we don't need massive, expensive AI models to help robots play games in the dark. We just need to stop treating the AI like a blindfolded person and start giving it a memory book. By showing the AI where it has been and what it has seen, even a small, smart assistant can guide a robot to victory.

The authors conclude that the failure of the old method wasn't because the AI wasn't smart enough; it was because the AI was asked to solve a puzzle while wearing a blindfold. Once they took the blindfold off (by adding context), the system worked perfectly.

Technical Summary: ASK in the Dark

Problem Statement

Reinforcement learning (RL) agents operating under partial observability (POMDPs) must act on incomplete information, often relying on egocentric observations that fail to distinguish between distinct underlying states. While small language models (SLMs) possess broad reasoning priors that could guide such agents, integrating them into partially observable settings has proven difficult. Previous work, specifically the ASK framework, utilized uncertainty-gated queries to consult SLMs only when an RL policy's predictive entropy exceeded a threshold. However, in partially observable environments, vanilla ASK failed to improve performance, achieving an overwrite rate near zero. The SLM was queried but almost never deviated from the policy's suggestion, acting as a passive validator rather than an active corrector. The authors hypothesize this failure stems from a context bottleneck: the "bare egocentric prompt" provided insufficient global context for the SLM to exercise its reasoning capabilities, rather than a limitation in the model's capacity.

Methodology: ASK+

The paper proposes ASK+, an extension of the ASK framework designed to resolve the context bottleneck in POMDPs. The core innovation lies in transforming the SLM prompt from a static observation into a stateful, trajectory-aware interface.

1. Stateful Prompting

Unlike vanilla ASK, which queries the SLM with only the current observation, ASK+ constructs a prompt ($Pr$) that includes:

Partially revealed map: Derived environmental features and local action previews to enable spatial inference.
Episodic memory: A history of visited positions and action history accumulated throughout the episode ( $\sigma$ ).
Structured reasoning: A constrained "chain-of-thought" mechanism where the SLM is allocated a budget (10 tokens) for deliberate computation before outputting an action.

2. Uncertainty-Gated Querying

The system retains the uncertainty-gating mechanism of ASK but applies it to the enriched context:

Uncertainty Signal: The RL policy (PPO) uses Monte Carlo (MC) Dropout to estimate total predictive entropy ( $H(o)$ ) via $N$ stochastic forward passes. This signal captures both aleatoric uncertainty (genuine observation ambiguity) and epistemic uncertainty (insufficient training coverage).
Gating Logic: If $H(o) \geq \tau$ (a threshold tuned via Bayesian optimization), the SLM is queried. If the SLM's response is valid, it overrides the policy's action ( $a'$ ); otherwise, the policy executes its default action ( $a$ ).

3. Architecture

RL Policy: A standard actor-critic PPO policy trained independently without SLM interaction. It remains small and focused on in-distribution states.
SLM Consultant: Uses Qwen3.5 models (2B and 4B parameters) without fine-tuning. The system relies on the SLM's pre-trained world knowledge to handle high-uncertainty states that the RL policy cannot resolve.

Key Contributions

Diagnosis of Context Bottleneck: The authors identify that the failure of uncertainty-gated assistance in POMDPs is due to insufficient prompting context, not model capacity. They demonstrate that providing trajectory-aware context converts the SLM from a passive redundancy check into an informative consultant.
ASK+ Framework: They introduce a stateful prompting strategy that supplies the SLM with maps, visited positions, and action history, enabling the model to reason over trajectories rather than isolated snapshots.
Validation of Uncertainty Signals: The paper establishes that total predictive entropy remains an informative signal for selective querying in POMDPs, effectively measuring action uncertainty even when state uncertainty is high.
Model Scale vs. Prompt Design: Empirical results show that prompt design and selective gating dominate performance over model scale, with the 2B parameter model matching or exceeding the 4B parameter model across all environments.

Experimental Results

The authors evaluated ASK+ across three complementary partially observable environments: FourRooms (spatial grounding), DoorKey (sequential subtasks), and HigherLower (temporal memory/card tracking).

Performance Gains:
- FourRooms: Success rate increased from 53% (PPO) and 55% (vanilla ASK) to 70% (ASK+).
- DoorKey: Success rate improved from 89% (PPO) and 89-92% (vanilla ASK) to 93% (ASK+).
- HigherLower: Accuracy reached 73.7%, matching the SLM-only upper bound and improving over PPO (72%).
Overwrite Rates: Vanilla ASK exhibited overwrite rates near zero (0.00–0.01), confirming it rarely corrected the policy. ASK+ achieved non-zero overwrite rates (0.01–0.06), indicating the SLM provided genuinely corrective guidance.
Model Efficiency: Qwen3.5-2B consistently matched or outperformed Qwen3.5-4B, confirming that scaling model parameters is less critical than optimizing the prompt structure and gating mechanism.
Robustness to Suboptimal Policies: ASK+ demonstrated the ability to compensate for weak PPO checkpoints. In scenarios where the PPO policy was poorly trained, the intervention rate (IR) increased, and ASK+ significantly recovered performance, suggesting the system functions as a general-purpose corrective layer.

Significance and Claims

The paper claims that uncertainty-gated SLM assistance is fundamentally viable under partial observability, provided the SLM is given sufficient context to reason. The primary contribution is not a new RL algorithm or a larger model, but a prompt engineering strategy that unlocks the latent reasoning capabilities of small models in ambiguous environments.

The authors assert that the "bare egocentric prompt" was the critical failure point of previous approaches. By closing this context gap, ASK+ allows small models to act as effective consultants, correcting policy errors in spatial and sequential tasks without requiring large-scale models or retraining the underlying policy. The work suggests that for partially observable tasks, reasoning capacity can be leveraged through prompt design and selective gating rather than by scaling up model size.

The evaluation is scoped to three specific environments to establish a proof of concept. The authors note that future work could explore richer POMDPs, the impact of stale history, and the potential for fine-tuning SLMs on egocentric views to further isolate the contributions of the gating mechanism.