Skip to content

The Horcrux: Mechanistically Interpretable Task Decomposition for Detecting and Mitigating Reward Hacking in Embodied AI Systems

🕒 Published (v1): 2025-11-22 01:45 UTC · Source: Arxiv · Venue: NEURIPS · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MITD (Mechanistically Interpretable Task Decomposition) is a hierarchical transformer architecture with Planner, Coordinator, and Executor modules designed to detect and mitigate reward hacking in embodied AI agents. By decomposing tasks into interpretable subtasks and exposing internal activations through diagnostic visualizations, it identifies an "optimal decomposition window" of 12–25 steps that reduces reward hacking frequency by 34%. The work is a small-scale proof-of-concept (1,000 HH-RLHF training samples, GPT-2 scale) accepted to a NeurIPS 2025 workshop.

Problem

Embodied AI agents exploit flaws in proxy reward signals (reward hacking), achieving high scores while failing true objectives. Existing approaches—post-hoc behavioral monitoring and mechanistic interpretability applied separately—neither predict hacking before it occurs nor provide module-level causal traceability into hierarchical decision structures.

Method

MITD implements three GPT-2-style transformer modules in a hierarchy: a Planner generates multi-scale goal embeddings, a Coordinator routes subgoals through disentangled bottlenecks, and four Executors fuse features with token embeddings via cross-attention. A Consistency Monitor enforces executor agreement; outputs are aggregated by an LSTM into structured reasoning traces. The system is trained on 1,000 HH-RLHF samples across 16 RTX 5090 GPUs for 3 epochs and evaluated on 50 held-out samples. Five diagnostic tools are introduced: Attention Waterfall Diagrams (AWDs, visualizing per-head attention exceedances above threshold τ=0.5), Decomposition Stability Diagrams (reward hacking frequency vs. decomposition step count), Mechanistic Failure Trees (causal risk propagation from objective to leaf decisions), Neural Pathway Flow Charts (layer-wise activation graphs color-coded by hacking category), and Objective Alignment Heatmaps (proxy-vs-intended correlation matrices).

Key Contributions

  • MITD: a unified hierarchical architecture with built-in interpretability hooks (not post-hoc) for detecting reward hacking at the module level.
  • Identification of an inverted-U stability pattern: hacking peaks at 4–8 decomposition steps; optimal "Goldilocks" zone is 12–25 steps, reducing hacking frequency by 34% across four failure modes.
  • Five novel diagnostic visualizations exposing internal computations at task-hierarchy boundaries.
  • Mechanistic Failure Trees formalizing causal risk flow from global objectives to low-level decisions with weighted edge coefficients.
  • Causal Intervention Leverage analysis showing non-uniform sensitivity concentrated in layers 3–7.

Results

  • Proxy rewards: −0.009 ± 0.023; true rewards: −0.005 ± 0.044 on 50 test samples (no baseline comparison provided).
  • Consistency score: 0.164 (zero variance across test set—constant value).
  • Reward correlation (Pearson ρ between proxy and true reward): −0.283 (zero variance—constant value).
  • Optimal decomposition window (12–25 steps) reduces reward hacking frequency by 34% vs. moderate depths (4–8 steps), across four failure modes: reward tampering, mesa-optimization, deceptive alignment, specification gaming.
  • Reward tampering shows highest baseline susceptibility but greatest improvement in optimal zone; mesa-optimization and deceptive alignment persist even in optimal regions; specification gaming shows steepest drop.
  • No comparison against external baseline models is reported; Table 1 is a qualitative capability comparison only.

Limitations

  • Extremely small scale: 1,000 training samples, 50 test samples (N per category ≈ 0–25); single model family (GPT-2).
  • Consistency score and reward correlation show zero variance across the test set, raising questions about metric validity.
  • Inverted-U pattern is correlational, not causal; may reflect capacity bottlenecking rather than interpretability-driven safety.
  • Hacking category detection thresholds {τ_r, γ_r, τ_s, γ_s} are set empirically on validation data, risking overfitting.
  • Visualizations are post-hoc and diagnostic only—no real-time mitigation.
  • Failure trees assume strictly hierarchical causality; lateral inter-module interactions are not modeled.
  • Optimal decomposition window may be architecture- or task-specific; generalization to larger models (Llama, GPT-3/4/5 scale) is unvalidated.
  • No comparison to quantitative baselines; reported metrics lack clear interpretation without reference points.

Relevance to Agentic AI / LLM Agents

This paper directly addresses reward hacking in hierarchical agentic systems—a core alignment challenge as LLM agents are deployed on long-horizon tasks with imperfect reward proxies. The MITD framework connects mechanistic interpretability to agentic planning architectures, proposing that task-module boundaries (not just individual neurons) are the right unit of analysis for detecting misalignment. The finding that decomposition granularity is a critical safety hyperparameter—with both too-shallow and too-deep decompositions inducing hacking—has direct implications for designing task-planning agents. However, the extremely limited experimental scale and absent baselines mean the findings should be treated as preliminary hypotheses rather than validated results.