End-to-End Differentiable Feature Selection for Deep Reinforcement Learning

Lightweight reinforcement learning (RL) agents must remove unnecessary raw observation features without sacrificing control performance. Differentiable feature gates are a natural approach, but their standard theory assumes a fixed dataset, as in supervised learning. In RL that assumption breaks down: the gate reshapes the policy, the policy reshapes its trajectories, and those trajectories drive gate updates. We systematically study end-to-end differentiable feature selection under this policy-dependent feedback, comparing four gate mechanisms on a live commercial racing game with a $4,244$-dimensional observation across two key design choices: stochastic versus deterministic gating, and expected-$L_0$ versus KL-to-prior regularization. Sweeping the number of selected features, the best mechanism, Hard Concrete, removes $70.1\%$ of the inputs while cutting completion time across twelve maps from $37.4$ to $34.0$ seconds, a $9.1\%$ improvement. Gate masks, however, reveal which features are selected, not how strongly the trained policy relies on them. We introduce GateKL, which toggles one gate at inference—closing an open gate or opening a closed one while keeping the policy parameters and all other learned gates fixed—and measures the resulting change in the action distribution using KL divergence. Across all $4,244$ gates and twelve maps, gate values learned by Hard Concrete and Stochastic Gates closely reflect policy reliance, whereas those learned by Concrete Dropout and Interpretable Feature Extractor carry little such information. These findings distinguish sparse selection from policy reliance and show why gate masks alone are insufficient to assess learned feature selection.