Interpretable by Design: Query-Specific Neural Modules for Explainable Reinforcement Learning

TL;DR

This work reframes reinforcement learning from purely action-driven optimization to queryable inference in deterministic environments. By introducing Query-Conditioned Deterministic Inference Networks (QDIN) with specialized heads and a query-conditioned fusion mechanism, the approach exposes diverse world-knowledge queries (reachability, paths, comparisons, policy) while maintaining competitive control. A key finding is the inference-control decoupling: inference accuracy can approach near-perfect levels (IoU ≈ 0.99) even when control performance remains suboptimal (e.g., episode return ≈ 0.31), implying distinct representations suffice for world knowledge versus action optimization. The results demonstrate architectural specialization, curriculum-based multi-query training, and calibrated confidence for selective answering, supporting a research agenda toward interpretable, verifiable, and collaborative RL systems that function as knowledge bases as well as controllers.

Abstract

Reinforcement learning has traditionally focused on a singular objective: learning policies that select actions to maximize reward. We challenge this paradigm by asking: what if we explicitly architected RL systems as inference engines that can answer diverse queries about their environment? In deterministic settings, trained agents implicitly encode rich knowledge about reachability, distances, values, and dynamics - yet current architectures are not designed to expose this information efficiently. We introduce Query Conditioned Deterministic Inference Networks (QDIN), a unified architecture that treats different types of queries (policy, reachability, paths, comparisons) as first-class citizens, with specialized neural modules optimized for each inference pattern. Our key empirical finding reveals a fundamental decoupling: inference accuracy can reach near-perfect levels (99% reachability IoU) even when control performance remains suboptimal (31% return), suggesting that the representations needed for accurate world knowledge differ from those required for optimal control. Experiments demonstrate that query specialized architectures outperform both unified models and post-hoc extraction methods, while maintaining competitive control performance. This work establishes a research agenda for RL systems designed from inception as queryable knowledge bases, with implications for interpretability, verification, and human-AI collaboration.

Figures (8)

• Figure 1: (a) QDIN architecture with query-conditioned processing and specialized heads. (b) Pareto frontier showing that mixed training achieves better trade-offs between inference accuracy and control performance than single-objective approaches. The shaded region indicates configurations dominated by our method.
• Figure 2: Detailed component architecture of QDIN. The state encoder produces hierarchical representations ($h_1, h_2, h_3$) while the query encoder creates a query embedding ($h_q$). Cross-attention fuses these representations based on query relevance. Specialized heads then process the fused representation according to the query type, with each head optimized for its specific inference pattern. The dashed line shows skip connections that preserve spatial information for the reachability head.
• Figure 3: The fundamental decoupling between inference accuracy and control performance. This surprising result suggests that the representations needed for accurate world knowledge differ from those required for optimal action selection.
• Figure 4: Query-specific performance across training modes and horizons. Mixed training (green) consistently outperforms single-objective approaches on inference tasks while maintaining competitive control.
• Figure 5: (a) QDIN generalizes to composite queries never seen during training, suggesting learned compositional representations. (b) Specialized architectures scale more efficiently than monolithic models as complexity increases.