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.

Artificial intelligence (AI) systems can learn to make good decisions, for example when controlling a robot’s movements. These systems often use neural networks to choose their next actions. However, they usually spend the same amount of computation on every decision, even when some decisions are much easier than others. To make these systems more efficient, we introduce BEXA, a method that adds early exits to neural networks. These early exits allow the system to stop computation early when a simpler decision is enough. A key feature of BEXA is a user-defined computation budget. For example, when the budget is set to roughly half of the usual computation, BEXA tries to find the best trade-off between maintaining performance and staying within this limit. We show that BEXA can substantially reduce the computation needed for decision-making while keeping performance close to standard reinforcement learning methods. This kind of efficiency is especially important in settings where decisions must be made with limited computation, time or energy.