# Optimizing Data Collection in Deep Reinforcement Learning

arXiv: https://arxiv.org/abs/2207.07736 · HTML: https://mlsanity.com/paper/2207.07736/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 23 references that are in the map.
- First posted on arXiv 15 Jul 2022.

## Builds on (most cited first, 6 of 6)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Large Batch Simulation for Deep Reinforcement Learning](https://mlsanity.com/paper/2103.07013/) (2021)
- [Accelerating Bandwidth-Bound Deep Learning Inference with Main-Memory Accelerators](https://mlsanity.com/paper/2012.00158/) (2020)
- [RL-Scope: Cross-Stack Profiling for Deep Reinforcement Learning Workloads](https://mlsanity.com/paper/2102.04285/) (2021)

## Abstract (arXiv)

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× 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× 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× and increase by up to 1024× 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.