An Acceleration Framework for Deep Reinforcement Learning Using Heterogeneous Systems

Abstract

Deep Reinforcement Learning (DRL) is vital in various AI applications. DRL algorithms comprise diverse compute primitives, which may not be simultaneously optimized using a homogeneous architecture. However, even with available heterogeneous architectures, optimizing DRL performance remains a challenge due to the complexity of design space in parallelizing DRL primitives and the variety of hardware employed in modern data centers. To address this, we introduce a framework for composing parallel DRL systems on heterogeneous platforms consisting of general-purpose processors (CPUs) and accelerators (GPUs, FPGAs). Our innovations include: 1. A general training protocol agnostic of the underlying hardware, enabling portable implementations across various processors and accelerators. 2. Efficient design exploration and automatic task placement enabling parallelization of tasks within each DRL primitive over one or multiple heterogeneous devices. 3. Incorporation of DRL-specific optimizations on runtime scheduling and resource allocation, facilitating parallelized training and enhancing the overall system performance. 4. High-level API for productive development using the framework. We showcase our framework through experimentation with three widely used DRL algorithms, DQN, DDPG, and SAC, on three heterogeneous platforms with diverse hardware characteristics and interconnections. The generated implementations outperform state-of-the-art libraries for CPU-GPU platforms by throughput improvements of up to 2×, and $1.7\times$ higher performance portability across platforms.