Bilinear Convolution Decomposition for Causal RL Interpretability
This work addresses interpretability in RL for researchers and practitioners, offering a method to move beyond correlational insights, though it is incremental as it builds on existing bilinear techniques.
The paper tackles the problem of interpreting reinforcement learning models by proposing bilinear convolutional layers that enable weight-based decomposition and causal validation of concept-based probes, achieving comparable performance to standard models on ProcGen environments.
Efforts to interpret reinforcement learning (RL) models often rely on high-level techniques such as attribution or probing, which provide only correlational insights and coarse causal control. This work proposes replacing nonlinearities in convolutional neural networks (ConvNets) with bilinear variants, to produce a class of models for which these limitations can be addressed. We show bilinear model variants perform comparably in model-free reinforcement learning settings, and give a side by side comparison on ProcGen environments. Bilinear layers' analytic structure enables weight-based decomposition. Previous work has shown bilinearity enables quantifying functional importance through eigendecomposition, to identify interpretable low rank structure. We show how to adapt the decomposition to convolution layers by applying singular value decomposition to vectors of interest, to separate the channel and spatial dimensions. Finally, we propose a methodology for causally validating concept-based probes, and illustrate its utility by studying a maze-solving agent's ability to track a cheese object.