Mind the Budget: Accelerating Deep Reinforcement Learning using Constrained Early Exit Neural Networks

Early exit neural networks, which adapt computation to input complexity, have proven effective in supervised learning but remain largely unexplored in deep reinforcement learning (DRL). In this paper, we propose Budgeted EXit Actor (BEXA), a novel actor-critic architecture that integrates early exit branches into the actor network. To ensure a mathematically principled trade-off between performance and inference expenditure, the exit decision is formulated as a constrained linear program during training, the solution of which is subsequently amortized to enable highly efficient runtime execution. BEXA is general, governed by an interpretable budget parameter, and compatible with any off-policy actor-critic method. We evaluate BEXA by integrating it with standard off-policy algorithms like SAC and TD3 on a suite of MuJoCo tasks. Our results demonstrate a substantial improvement in inference efficiency with minimal or no loss in performance. These findings highlight early exits as a promising direction for improving computational efficiency in DRL.