Sparse Feature Policy Unlearning Mitigates State Hallucination in Vision-Language-Action Models

AI summary

Overview

Research area: Robotics / embodied AI — specifically vision-language-action (VLA) models, mechanistic interpretability with sparse autoencoders, and machine unlearning. Technical level: Intermediate to Advanced (assumes familiarity with transformer policies, LoRA fine-tuning, and sparse autoencoders). Scope: The paper diagnoses a recurring VLA failure mode called "state hallucination," shows it corresponds to identifiable sparse features and attention patterns, and proposes a feature-level unlearning method (SOUL) that suppresses those features to reduce hallucinated failures while preserving successful manipulation.

What This Paper Is About

VLA models sometimes keep executing a manipulation sequence as if a robot-object state had already been achieved when it has not — for example, closing the gripper before reaching the object, then carrying on with transport and place motions as though the object were grasped. The authors call this state hallucination, and it is dangerous because a small local error propagates through later stages of a task. The goal is to determine whether the internal policy knowledge responsible for this behavior can be selectively removed, rather than patched at inference time, without destroying the manipulation skills that already work.

Key Contributions

1. A definition and taxonomy of state hallucination as a distinct recurring failure pattern in VLAs, broken down into grasp, transport, and place hallucinations, with explicit labeling rules (grasp failures involving object contact, such as inaccurate localization or slippage, are explicitly excluded).
2. An interpretability analysis of the failure: hallucinated rollouts show weakened attention to task-relevant regions (gripper, target object) and increased attention to background, and sparse autoencoders reveal features whose activation is selectively associated with hallucinated versus successful executions, quantified by a new feature association score.
3. SOUL (Sparse feature pOlicy UnLearning), which uses those identified sparse features as explicit forgetting targets (suppressed) and retention targets (enhanced) in a feature-guided unlearning objective applied to the policy parameters via LoRA.
4. Validation across two VLA architectures, two simulators, and a real robot, reporting large reductions in hallucination and gains in task success against the original policy and a gradient-ascent unlearning baseline.

Main Findings

• Hallucination is common in the original policies: across the benchmark settings the original pretrained policies show an average hallucination rate of 55.0% and an average hallucination failure rate of 48.6%.
• Attention differs systematically: clean-success rollouts concentrate attention on the gripper and target object, whereas hallucination-failure rollouts show weaker attention to those regions and more spatially dispersed attention toward the background.
• Sparse features separate the outcomes: hallucination features stay weak in clean successes but spike around hallucination events in failing rollouts (temporally localized, recurring across tasks and scenes), while success features are strong during successful subtask execution or task completion and weak in hallucination failures.
• SOUL cuts hallucination substantially: on LIBERO-Plus with OpenVLA, hallucination falls from 70.0% to 44.0%, hallucination failure from 60.0% to 38.0%, and clean success rises from 16.0% to 52.0%; overall success rises from 26.0% to 58.0%.
• The gains hold on a second architecture: on RoboCasa with π0.5, hallucination falls from 40.0% to 22.9%, hallucination failure from 37.1% to 21.0%, overall success rises from 41.9% to 54.3%, and clean success from 39.0% to 52.4%.
• Stage-specific rates improve too: grasp, transport, and place hallucination rates all drop (for example, LIBERO-Plus grasp hallucination 70.0% to 44.0%; RoboCasa transport hallucination 7.6% to 3.8%), so the effect is not confined to one manipulation stage.
• A naive unlearning baseline fails: gradient ascent, adapted by maximizing action prediction loss on hallucination-failure samples, drives overall success to 0.0% on LIBERO-Plus and 4.8% on RoboCasa, showing that directly forgetting failure samples destroys useful capability.
• The trade-off is favorable: SOUL preserves 89% of originally successful task executions and non-hallucinated events on average, converts 40% of originally failed executions into successes, and corrects 78% of hallucinated events.
• Region grounding matters: region-guided feature selection (hallucination 33.5%, clean success 52.2%) outperforms tokenized selection (41.5%, 41.3%) and global pooling (45.8%, 38.3%).
• Both score components and the retention term are needed: removing the retention objective raises hallucination to 36.8% and drops clean success to 48.3%; removing the aggregated score term raises hallucination to 39.9% and drops clean success to 47.7%, versus 33.5% and 52.2% for the full method.
• It transfers to hardware: on a Franka Research 3 with 50 rollouts across both models, SOUL raises clean success by an average of 20%p and reduces hallucination failure by 22%p.
• Headline summary numbers: the paper reports an average 21.7%p reduction in hallucination rate and a 22.1%p improvement in task success while preserving 89% of previously successful behaviors.

Methodology in Plain English

The authors start by watching the policy fail and sorting every rollout into one of three buckets: clean success (task completed, no hallucination), hallucination failure (task failed and at least one hallucination event occurred), and non-hallucination failure. Labels come from ground-truth robot-object states in simulation and observable interaction states in the real world.

They then look inside the model in two ways. First, they map the attention from action tokens onto visual tokens back onto the input image to see where the policy is "looking." Second, they train a sparse autoencoder — specifically a Top-k SAE with a latent dimension of 4096, k = 100 active features, and k_aux = 512 — on visual-token activations from the residual streams of selected transformer layers. The SAE is frozen after training and used purely as a readout tool.

To decide which features matter, they average feature activations within semantic scene regions (from simulator masks; hand-annotated for real-world images), take the maximum over time within each rollout, and average across rollouts per outcome. Two scores are defined: a hallucination association score comparing hallucination failures to the mean of clean successes and non-hallucination failures, and a success association score computed the other way around. Forgetting targets are chosen using a combined score that rewards separation from non-hallucination failures and penalizes association with success; retention targets are the features with high success association. The top five region-feature pairs from each side are used.

Finally, they update the policy itself. With the SAE frozen, they add LoRA adapters to the pretrained model and optimize a loss that squares and penalizes activations of forget features on hallucination-failure samples, while a weighted negative term rewards activations of retain features on clean-success samples. The result is baked into the weights, so no intervention is needed at inference time — a key contrast with prior work that steers VLA representations at test time.

Why This Matters

This paper argues that interpretable feature analysis can do more than diagnose failures: it can supply concrete optimization targets, turning a vague "the robot is unreliable" problem into a specific, editable set of internal representations. It also introduces an unlearning framing to robot policies, a setting where the authors note no prior unlearning baseline exists — they had to adapt gradient ascent themselves.

Real-world applications:

• Warehouse and logistics pick-and-place, where a gripper that closes on empty space and then carries nothing to a destination silently corrupts downstream sorting.
• Home and service robotics, where open-ended kitchen tasks (the RoboCasa setting) involve long sequences in which an early grasp error otherwise propagates.
• Manufacturing and assembly cells, where a policy acting on an assumed-but-unachieved state can drive a tool into the wrong location.
• Safety monitoring and deployment triage, since the identified hallucination features give an inspectable signal that a policy is operating on a false belief.

Industry relevance: because the fix is stored in the model weights via LoRA rather than applied as a per-step inference intervention, it is cheap to deploy on existing robot policies and compatible with the common practice of task-specific fine-tuning. The method is demonstrated on two widely used open VLA families (OpenVLA and π0.5), which lowers the barrier to adoption.

Future Directions

• Beyond pick-and-place: the authors deliberately restrict evaluation to pick-and-place with a structured grasp-transport-place sequence to control for planning confounds; whether hallucination features generalize to longer, more compositional, or contact-rich tasks is untested.
• Beyond two architectures and one robot: the study covers OpenVLA, π0.5, and a single 7-DoF Franka Research 3; scaling to other embodiments and action heads remains open.
• Reducing manual effort: real-world semantic regions are hand-annotated; automating that step would make the region-guided selection pipeline practical at scale.
• Relationship to inference-time steering: concurrent work steers VLA behavior by intervening on sparse features at test time, while SOUL edits weights permanently; combining or comparing the two, and understanding when each is preferable, is a natural next question. The paper's ablation also leaves the weighting of the retention term and the number of selected features as tunable choices whose sensitivity is only partially explored.

Target Audience

Robotics and embodied-AI researchers working on VLA reliability and failure analysis; mechanistic interpretability researchers interested in applying sparse autoencoders to control policies rather than language models; machine unlearning researchers looking for a domain where forgetting targets are behavioral rather than factual; and practitioners deploying foundation-model policies on real hardware who need persistent, low-overhead behavior correction.

Authors’ abstract

Vision-Language-Action (VLA) models have shown strong generalization in robotic manipulation by leveraging rich representations from pretrained vision-language models. However, their deployment in real-world environments remains limited by recurring unreliable behaviors. In this work, we study state hallucination, a recurring failure pattern in which a VLA continues acting as if an unrealized robot-object state had been achieved. Our analyses find that state hallucination coincides with weakened attention to task-relevant visual regions, and a mechanistic interpretation via sparse autoencoders reveals that hallucination-associated sparse features are activated when these failures occur. Based on this analysis, we propose SOUL (Sparse feature pOlicy UnLearning), which selectively unlearns policy knowledge associated with state hallucination behaviors, where sparse features identified from hallucination failures and successful behaviors serve as explicit forgetting and retention targets, respectively. Experiments across VLA architectures in simulated and real-world environments show that our method substantially reduces hallucinated failures and improves task success without substantially compromising the existing manipulation capabilities. These results suggest that interpretable feature analysis provides a practical basis for selectively modifying undesirable knowledge in robot policies.