[Paper Review] Single-Shot Pruning for Offline Reinforcement Learning

Samin Yeasar Arnob, Riyasat Ohib+2 more|arXiv (Cornell University)|Dec 31, 2021

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes a single-shot pruning method applied before training in offline reinforcement learning (RL), leveraging fixed datasets to identify sparse, high-performing subnetworks. By applying SNIP and GraSP pruning at initialization, the approach achieves performance comparable to dense networks even at 95% sparsity, reducing memory usage by 4x without compression and enabling faster inference on resource-constrained devices.

ABSTRACT

Deep Reinforcement Learning (RL) is a powerful framework for solving complex real-world problems. Large neural networks employed in the framework are traditionally associated with better generalization capabilities, but their increased size entails the drawbacks of extensive training duration, substantial hardware resources, and longer inference times. One way to tackle this problem is to prune neural networks leaving only the necessary parameters. State-of-the-art concurrent pruning techniques for imposing sparsity perform demonstrably well in applications where data distributions are fixed. However, they have not yet been substantially explored in the context of RL. We close the gap between RL and single-shot pruning techniques and present a general pruning approach to the Offline RL. We leverage a fixed dataset to prune neural networks before the start of RL training. We then run experiments varying the network sparsity level and evaluating the validity of pruning at initialization techniques in continuous control tasks. Our results show that with 95% of the network weights pruned, Offline-RL algorithms can still retain performance in the majority of our experiments. To the best of our knowledge, no prior work utilizing pruning in RL retained performance at such high levels of sparsity. Moreover, pruning at initialization techniques can be easily integrated into any existing Offline-RL algorithms without changing the learning objective.

Motivation & Objective

• To address the high computational and memory costs of large neural networks in offline RL by enabling efficient model compression before training.
• To explore whether single-shot pruning techniques—originally designed for supervised learning—can be effectively applied to offline RL with fixed datasets.
• To evaluate the performance of pruned networks in continuous control tasks under extreme sparsity (up to 95%)
• To demonstrate that pruning at initialization can be seamlessly integrated into existing offline RL algorithms without modifying the learning objective.
• To enable faster training and inference in real-time, low-resource applications such as robotics and edge AI.

Proposed method

• Applying single-shot pruning methods—SNIP and GraSP—on randomly initialized neural networks before offline RL training.
• Using a fixed dataset from the D4RL benchmark to compute pruning criteria based on gradient sensitivity (GraSP) or connection importance (SNIP).
• Pruning the network once at initialization, without any retraining or iterative fine-tuning, to generate sparse subnetworks.
• Evaluating the pruned networks using standard offline RL algorithms (BCQ, BC) across continuous control environments (HalfCheetah, Hopper, Walker2D).
• Measuring performance, memory usage, and layer-wise sparsity to assess the impact of pruning on generalization and efficiency.
• Using PyTorch's sparse tensor functionality to store and compare memory requirements of dense vs. pruned models.

Experimental results

Research questions

• RQ1 Can single-shot pruning at initialization maintain performance in offline RL when applied before training?
• RQ2 What is the maximum level of sparsity (e.g., 95%) that can be achieved without degrading performance in offline RL?
• RQ3 How do different pruning criteria (SNIP vs. GraSP) affect the structure and performance of pruned networks in offline RL?
• RQ4 To what extent does pruning reduce memory usage and improve inference efficiency in offline RL models?
• RQ5 Can pruning at initialization be integrated into existing offline RL algorithms without altering their learning objectives?

Key findings

• The proposed method maintains performance comparable to dense networks even at 95% sparsity, with no significant drop in return across multiple environments.
• Pruning with GraSP preserves more weights in the final layer than SNIP, aligning with its objective of preserving gradient flow.
• Memory usage is reduced by a factor of 4x for 95%-sparse models compared to dense models, even without compression.
• The method is plug-and-play: it can be applied to any offline RL algorithm without modifying the learning objective or requiring additional training.
• Experiments show that varying batch size or pruning iterations does not improve performance over a single-shot pruning loop.
• The results demonstrate that fixed-dataset offline RL enables effective single-shot pruning, which is not feasible in online RL due to data distribution shifts.