Pruning with Scaled Policy Constraints for Light-weight Reinforcement Learning

Abstract

The increasing computational demands of Deep Reinforcement Learning (DRL) models, particularly for embedded systems in autonomous vehicles and drones, present significant challenges owing to their extensive neural network complexities. Previous DRL compression strategies predominantly focused on unstructured pruning, effective for reducing model size but requiring specialized hardware for computational acceleration. Conversely, DRL models with structured pruning applied can be accelerated on standard hardware, though they typically encounter performance issues at higher pruning rates due to structural constraints. In response to these challenges, this paper introduces an advanced structured pruning methodology, combined with scaled policy constraints (SPC) for DRL models. Our approach overcomes the performance limitations of conventional structured pruning, achieving high pruning rates while maintaining robust model performance. Enhanced performance restoration after pruning is achieved by fine-tuning with SPC and applying structural regularization, thus ensuring efficient decision-making with a minimal computational burden. Extensive evaluations on the D4RL benchmark and in a drone control simulation environment confirm the effectiveness of our method. Our approach maintains performance integrity even at high pruning rates, with less than a 2% decrease in normalized score at 90% pruning in D4RL and preserving cumulative reward at 87% pruning in drone control simulation. Significantly, our approach also enables considerable computational acceleration on standard hardware. We implemented our method on the NVIDIA Jetson Xavier NX board and achieved a 2.5-fold speed-up on devices with NVIDIA Volta GPUs and over double the speed-up on those with NVIDIA Carmel ARMv8.2 CPUs. These outcomes highlight our method’s suitability for real-time, resource-constrained applications, demonstrating its practicality and efficiency.