Table of Contents
Fetching ...

SpikeRL: A Scalable and Energy-efficient Framework for Deep Spiking Reinforcement Learning

Tokey Tahmid, Mark Gates, Piotr Luszczek, Catherine D. Schuman

TL;DR

SpikeRL addresses the pressing need for scalable and energy-efficient training of DeepRL-enabled Spiking Neural Networks for continuous control. It introduces a three-component architecture that combines population-encoded SNN observations, a TD3-based actor-critic, and distributed training via MPI and NCCL, augmented by mixed-precision BFLOAT16 optimization. The framework achieves 4.26X faster and 2.25X more energy-efficient performance than state-of-the-art DeepRL-SNN methods, and delivers notable gains over PopSAN while dramatically reducing energy consumption. Experiments across Ant-v4, Hopper-v4, and HalfCheetah-v4 demonstrate both higher rewards and substantial energy savings, highlighting SpikeRL's potential for sustainable, scalable real-world robotics and control applications.

Abstract

In this era of AI revolution, massive investments in large-scale data-driven AI systems demand high-performance computing, consuming tremendous energy and resources. This trend raises new challenges in optimizing sustainability without sacrificing scalability or performance. Among the energy-efficient alternatives of the traditional Von Neumann architecture, neuromorphic computing and its Spiking Neural Networks (SNNs) are a promising choice due to their inherent energy efficiency. However, in some real-world application scenarios such as complex continuous control tasks, SNNs often lack the performance optimizations that traditional artificial neural networks have. Researchers have addressed this by combining SNNs with Deep Reinforcement Learning (DeepRL), yet scalability remains unexplored. In this paper, we extend our previous work on SpikeRL, which is a scalable and energy efficient framework for DeepRL-based SNNs for continuous control. In our initial implementation of SpikeRL framework, we depended on the population encoding from the Population-coded Spiking Actor Network (PopSAN) method for our SNN model and implemented distributed training with Message Passing Interface (MPI) through mpi4py. Also, further optimizing our model training by using mixed-precision for parameter updates. In our new SpikeRL framework, we have implemented our own DeepRL-SNN component with population encoding, and distributed training with PyTorch Distributed package with NCCL backend while still optimizing with mixed precision training. Our new SpikeRL implementation is 4.26X faster and 2.25X more energy efficient than state-of-the-art DeepRL-SNN methods. Our proposed SpikeRL framework demonstrates a truly scalable and sustainable solution for complex continuous control tasks in real-world applications.

SpikeRL: A Scalable and Energy-efficient Framework for Deep Spiking Reinforcement Learning

TL;DR

SpikeRL addresses the pressing need for scalable and energy-efficient training of DeepRL-enabled Spiking Neural Networks for continuous control. It introduces a three-component architecture that combines population-encoded SNN observations, a TD3-based actor-critic, and distributed training via MPI and NCCL, augmented by mixed-precision BFLOAT16 optimization. The framework achieves 4.26X faster and 2.25X more energy-efficient performance than state-of-the-art DeepRL-SNN methods, and delivers notable gains over PopSAN while dramatically reducing energy consumption. Experiments across Ant-v4, Hopper-v4, and HalfCheetah-v4 demonstrate both higher rewards and substantial energy savings, highlighting SpikeRL's potential for sustainable, scalable real-world robotics and control applications.

Abstract

In this era of AI revolution, massive investments in large-scale data-driven AI systems demand high-performance computing, consuming tremendous energy and resources. This trend raises new challenges in optimizing sustainability without sacrificing scalability or performance. Among the energy-efficient alternatives of the traditional Von Neumann architecture, neuromorphic computing and its Spiking Neural Networks (SNNs) are a promising choice due to their inherent energy efficiency. However, in some real-world application scenarios such as complex continuous control tasks, SNNs often lack the performance optimizations that traditional artificial neural networks have. Researchers have addressed this by combining SNNs with Deep Reinforcement Learning (DeepRL), yet scalability remains unexplored. In this paper, we extend our previous work on SpikeRL, which is a scalable and energy efficient framework for DeepRL-based SNNs for continuous control. In our initial implementation of SpikeRL framework, we depended on the population encoding from the Population-coded Spiking Actor Network (PopSAN) method for our SNN model and implemented distributed training with Message Passing Interface (MPI) through mpi4py. Also, further optimizing our model training by using mixed-precision for parameter updates. In our new SpikeRL framework, we have implemented our own DeepRL-SNN component with population encoding, and distributed training with PyTorch Distributed package with NCCL backend while still optimizing with mixed precision training. Our new SpikeRL implementation is 4.26X faster and 2.25X more energy efficient than state-of-the-art DeepRL-SNN methods. Our proposed SpikeRL framework demonstrates a truly scalable and sustainable solution for complex continuous control tasks in real-world applications.

Paper Structure

This paper contains 13 sections, 1 equation, 13 figures, 3 tables.

Figures (13)

  • Figure 1: SpikeRL System Architecture.
  • Figure 2: Distributed Training with MPI using mpi4py.
  • Figure 3: Distributed Training with MPI/NCCL Backend.
  • Figure 4: Mixed Precision Training with BFLOAT16.
  • Figure 5: Ant-v4 Performance Benchmark of SpikeRL over PopSAN and DSRL.
  • ...and 8 more figures