Reward-Conditioned Attention: How Reward Design Shapes What Autonomous Driving Agents See

Technical Summary: Reward-Conditioned Attention in Autonomous Driving

Problem Statement
While Reinforcement Learning (RL) has achieved expert-level performance in autonomous driving (AD) simulations, the internal mechanisms driving these policies remain opaque. A critical gap exists in understanding how the reward function—the primary signal guiding RL agents—shapes the agent's internal attention patterns. Although the validity of attention weights as post-hoc explanations for model decisions is debated, this paper posits a distinct question: does the reward design predictably determine what the encoder attends to? Understanding this relationship is essential for verifying that reward functions teach agents to monitor the correct aspects of the driving scene and for advancing principled reward engineering in safety-critical systems.

Methodology
The study utilizes the V-Max framework, employing three Soft Actor-Critic (SAC) agents trained on the Waymo Open Motion Dataset (WOMD). All three agents share an identical Perceiver-based Latent-Query (LQ) encoder architecture, which processes 280 input tokens (representing the ego vehicle, other agents, road geometry, traffic lights, and GPS waypoints) into a compact latent representation via cross-attention.

The core experimental design isolates the effect of reward design by comparing three agents trained on the same data with identical initialization seeds but differing reward configurations:

1. Basic: Penalizes only safety violations (collision, off-road, red-light).
2. Minimal: Adds navigation incentives (route adherence and forward progress).
3. Complete: Adds continuous proximity penalties based on Time-to-Collision (TTC), speed limits, and comfort.

Key Methodological Innovation: Within-Episode Correlation
A central methodological contribution is the rejection of "naïve pooling" (aggregating all timesteps across all episodes) for analyzing attention-risk relationships. The authors demonstrate that pooling conflates within-episode dynamics with between-episode heterogeneity (e.g., urban vs. suburban risk profiles), leading to a significant underestimation of the attention-risk link.

Instead, the study employs within-episode Spearman rank correlation ($\rho_i$) between collision risk ($R$) and agent-directed attention for each scenario individually. These correlations are aggregated across scenarios using the Fisher z-transform. This approach correctly handles the variance inherent in heterogeneous RL episodes and reveals robust statistical signals that pooled metrics obscure.

Key Results
The analysis, conducted across 50 real-world scenarios (3,676 timesteps for the complete model), yields three primary findings:

1. Validation of Within-Episode Analysis:
   For the complete-reward model, the within-episode correlation between collision risk and attention to "Other Agents" is robustly positive ($\bar{\rho} = +0.291$, 95% CI $[+0.125, +0.442]$). In the 16 most risk-reactive scenarios, this correlation strengthens significantly ($\bar{\rho} = +0.522$), with agents increasing attention to dynamic threats by 77% under threat. In contrast, naïve pooling yields a correlation of only $+0.088$, a 3.3× underestimate.

2. Reward Content Shapes Attention Baselines:
   The specific components of the reward function directly dictate the allocation of attention resources. Agents trained with navigation rewards (Minimal and Complete) allocate significantly more attention to GPS-path tokens than those trained with basic safety penalties.

   • Minimal vs. Complete: The Minimal model allocates 33.5% attention to GPS paths, compared to 16.4% for the Complete model (a 2.0× difference).
   • Basic vs. Minimal: The range spans from 7.1% (Basic) to 33.5% (Minimal), a 4.7× difference.
     This confirms that reward content determines which scene elements the encoder prioritizes.

3. Continuous Safety Rewards Create a Vigilance Prior:
   The inclusion of continuous TTC-based proximity penalties in the Complete model establishes a "learned vigilance prior." Even during collision-free phases, the Complete model maintains elevated "resting" surveillance of other agents compared to the Minimal model.

   • In 16 of 26 qualifying scenarios, the Complete model maintains a mean of +151% more resting agent surveillance than the Minimal model.
   • This elevated baseline is present from the first timestep, suggesting the agent learns a state of constant readiness rather than reacting only when risk spikes.

4. Qualitative Reversal of Strategy:
   In several scenarios, the Complete and Minimal models exhibit opposite directions of attention-risk correlation. This demonstrates that reward design can qualitatively reverse an agent's attentional strategy, not merely modulate the magnitude of its response.

Significance and Claims
The paper claims that attention analysis serves as a practical diagnostic tool for verifying that a reward function produces the intended representational behavior in safety-critical RL systems. By establishing that reward design directly shapes internal attention patterns, the authors argue that:

• Practitioners can use attention metrics to verify if an agent is monitoring the correct scene elements.
• The "within-episode correlation with Fisher z-transform" is the appropriate statistical standard for attention analysis in heterogeneous RL environments, as aggregate metrics conceal critical behavioral heterogeneity.
• Reward engineering must account for how reward components alter not just policy actions, but the fundamental representational focus of the agent.

The study concludes that while attention weights may not always explain why a decision was made, they reliably reveal what the agent was trained to look at, making them a vital component of reward design validation.