RNM-TD3: N:M Semi-structured Sparse Reinforcement Learning From Scratch

TL;DR

This paper introduces RNM-TD3, a reinforcement learning method utilizing N:M structured sparsity to compress neural networks, improve performance, and enable hardware acceleration, demonstrating superior results on continuous-control benchmarks.

Contribution

First study applying N:M structured sparsity in reinforcement learning, balancing compression, performance, and hardware efficiency during training.

Findings

Outperforms dense networks at 50-75% sparsity

Achieves up to 14% performance increase at 2:4 sparsity

Remains competitive at 87.5% sparsity, enabling training speedups

Abstract

Sparsity is a well-studied technique for compressing deep neural networks (DNNs) without compromising performance. In deep reinforcement learning (DRL), neural networks with up to 5% of their original weights can still be trained with minimal performance loss compared to their dense counterparts. However, most existing methods rely on unstructured fine-grained sparsity, which limits hardware acceleration opportunities due to irregular computation patterns. Structured coarse-grained sparsity enables hardware acceleration, yet typically degrades performance and increases pruning complexity. In this work, we present, to the best of our knowledge, the first study on N:M structured sparsity in RL, which balances compression, performance, and hardware efficiency. Our framework enforces row-wise N:M sparsity throughout training for all networks in off-policy RL (TD3), maintaining compatibility with accelerators that support N:M sparse matrix operations. Experiments on continuous-control benchmarks show that RNM-TD3, our N:M sparse agent, outperforms its dense counterpart at 50%-75% sparsity (e.g., 2:4 and 1:4), achieving up to a 14% increase in performance at 2:4 sparsity on the Ant environment. RNM-TD3 remains competitive even at 87.5% sparsity (1:8), while enabling potential training speedups.