TLDR: This research introduces OFEXiNet, a method for simultaneously training and pruning deep neural networks in reinforcement learning, particularly those enhanced by Online Feature Extractor Networks (OFENets). By integrating a sparsity-promoting regularization scheme, the approach dynamically prunes less useful network units during training. Evaluations on continuous control tasks show that OFEXiNet can significantly reduce network size (up to 40%) with minimal performance loss, demonstrating that pruning large networks during training yields more efficient and higher-performing RL agents than training smaller networks from scratch.
Deep reinforcement learning (RL) has achieved remarkable success in various domains, often by employing increasingly large and complex neural networks. While these expansive networks can boost performance, they come with a significant drawback: increased computational and memory demands. This challenge has been successfully tackled in supervised learning through neural network pruning, but its application in the realm of reinforcement learning has remained largely unexplored.
A new research paper, titled “Online Training and Pruning of Deep Reinforcement Learning Networks,” by Valentin Frank Ingmar Guenter and Athanasios Sideris, proposes a novel approach to integrate simultaneous training and pruning directly within advanced RL methods. Their work focuses on RL algorithms that are enhanced by the Online Feature Extractor Network (OFENet), a component designed to extract rich features from observations and actions, allowing RL agents to leverage larger deep neural networks effectively.
The core of their innovation lies in what they call XiNet, a modification of neural networks where each unit is scaled by a 0/1 Random Variable (RV) called ξ. These RVs are governed by variational Bernoulli distributions, and the networks are trained to solve stochastic optimization problems over both the network weights and the parameters of these distributions. This unique formulation introduces regularization terms that encourage the variational parameters to converge to zero when a unit contributes minimally to performance. When this happens, the corresponding network structure is permanently deactivated and pruned.
OFENets, which are implemented as DenseNets, are crucial for scaling RL networks. However, their dense architecture also adds to the complexity. The researchers’ method, dubbed OFEXiNet, integrates the pruning mechanism directly into the OFENet structure. They introduce a cost-aware, sparsity-promoting regularization scheme specifically tailored to the DenseNet architecture of OFENets. This scheme automatically selects the numerous hyperparameters associated with the regularization terms, effectively combining the RL objectives with network compression goals.
The method was rigorously evaluated on continuous control benchmarks, specifically MuJoCo environments, using the popular Soft Actor-Critic (SAC) RL agent. The results are compelling: OFENets can be pruned considerably, reducing their initial size by up to 40%, with only minimal loss in performance. Furthermore, the study confirms a significant finding from supervised learning that also holds true for RL: starting with larger networks and pruning them during training produces more efficient and higher-performing RL agents compared to training smaller networks from scratch.
Also Read:
- Metropolis-Hastings Sampling Unlocks Efficient AI Control
- Advancing Policy Optimization: A New Algorithm for Efficient Reinforcement Learning
This research offers a promising direction for making advanced deep reinforcement learning more computationally efficient and deployable, especially on resource-constrained devices. By dynamically optimizing network size during the training process, OFEXiNet paves the way for powerful yet compact RL agents. You can read the full paper here.
