Architectural Inductive Biases Can Be Enough for State Abstraction in Deep Reinforcement Learning

ICLR 2026 Conference Submission22516 Authors

Keywords: Deep Reinforcement Learning, State Abstraction, Representation Learning

Abstract: The ability to ignore irrelevant sensory information is central to intelligent behavior. In reinforcement learning (RL), existing methods typically rely on auxiliary objectives to achieve similar forms of abstraction. Such objectives tend to add significant complexity to the base RL algorithm. In this work, we take a step back and ask: can selective abstraction emerge naturally from reward optimization alone, without any additional objectives? Following prior work, we show that standard deep RL learns slowly or not at all in the presence of distracting, task-irrelevant state variables, failing to learn meaningful state abstractions. We then introduce a surprisingly simple neural network architecture change: a learnable, observation-independent attention mask applied to the inputs of the policy and value networks and trained end-to-end using only the RL objective. Despite its simplicity, this architectural modification consistently improves sample efficiency and learns to mask out distracting input variables across 12 continuous control tasks. We analyze the dynamics of gradient descent using this method on a linear regression task and demonstrate improved feature credit assignment. Finally, we conduct experiments on toy MDPs and show that the attention mask leads to accurate Q-value estimation and induces soft abstractions over a factored state space. Our findings challenge the need for complex auxiliary objectives to learn state abstractions in deep RL and suggest a simple baseline for future research.

Primary Area: reinforcement learning

Submission Number: 22516