Optimizing Data Collection in Deep Reinforcement Learning

TL;DR

This paper introduces GPU-based optimizations, including vectorization and kernel fusion, to significantly accelerate data collection in deep reinforcement learning, reducing training time and improving efficiency.

Contribution

It presents novel GPU-centric techniques for simulation acceleration in RL, achieving up to 1024x speedup and demonstrating their effectiveness over traditional CPU-based methods.

Findings

GPU vectorization achieves up to 1024x speedup.

ML compiler implementations outperform DNN frameworks by 13.4x.

Kernel fusion provides up to 1024x speedup, especially with complex simulators.

Abstract

Reinforcement learning (RL) workloads take a notoriously long time to train due to the large number of samples collected at run-time from simulators. Unfortunately, cluster scale-up approaches remain expensive, and commonly used CPU implementations of simulators induce high overhead when switching back and forth between GPU computations. We explore two optimizations that increase RL data collection efficiency by increasing GPU utilization: (1) GPU vectorization: parallelizing simulation on the GPU for increased hardware parallelism, and (2) simulator kernel fusion: fusing multiple simulation steps to run in a single GPU kernel launch to reduce global memory bandwidth requirements. We find that GPU vectorization can achieve up to $1024\times$ speedup over commonly used CPU simulators. We profile the performance of different implementations and show that for a simple simulator, ML compiler implementations (XLA) of GPU vectorization outperform a DNN framework (PyTorch) by $13.4\times$ by reducing CPU overhead from repeated Python to DL backend API calls. We show that simulator kernel fusion speedups with a simple simulator are $11.3\times$ and increase by up to $1024\times$ as simulator complexity increases in terms of memory bandwidth requirements. We show that the speedups from simulator kernel fusion are orthogonal and combinable with GPU vectorization, leading to a multiplicative speedup.