[Paper Review] Reinforcement Learning through Asynchronous Advantage Actor-Critic on a GPU

Mohammad Babaeizadeh, Iuri Frosio+3 more|arXiv (Cornell University)|Nov 18, 2016

Reinforcement Learning in RoboticsComputer Science10 references174 citations

TL;DR

GA3C is a hybrid CPU/GPU implementation of A3C that centralizes the neural network on the GPU, uses prediction and training queues, and automatically tunes resource configuration to significantly speed up training compared to CPU-only A3C.

ABSTRACT

We introduce a hybrid CPU/GPU version of the Asynchronous Advantage Actor-Critic (A3C) algorithm, currently the state-of-the-art method in reinforcement learning for various gaming tasks. We analyze its computational traits and concentrate on aspects critical to leveraging the GPU's computational power. We introduce a system of queues and a dynamic scheduling strategy, potentially helpful for other asynchronous algorithms as well. Our hybrid CPU/GPU version of A3C, based on TensorFlow, achieves a significant speed up compared to a CPU implementation; we make it publicly available to other researchers at https://github.com/NVlabs/GA3C .

Motivation & Objective

• Motivate and analyze the computational bottlenecks of asynchronous deep RL methods when trained on GPUs.
• Propose GA3C architecture to better utilize GPU parallelism while maintaining convergence properties.
• Investigate dynamic auto-tuning of system knobs (predictors, trainers, agents) to maximize training throughput (TPS).
• Evaluate how GA3C scales with network size and hardware, and compare convergence speed against CPU-based A3C.

Proposed method

• Implement a CPU/GPU hybrid version of A3C (GA3C) in TensorFlow with a single GPU model instance.
• Introduce a prediction queue for agents to request policy in batches and a training queue for batches used to update the GPU model.
• Use one or more GPU predictors to batch-serve policy inferences and multiple trainers to batch-update the model on the GPU.
• Centralize DNN weights on a single parameter server-like structure to remove per-agent model copies and synchronization overhead.
• Incorporate an entropy regularization term in the policy loss to foster exploration.
• Add a small epsilon to log-probabilities and entropy computations to stabilize gradients when delays cause stale parameters (policy lag).
• Dynamically adjust N_P (predictors), N_T (trainers), and N_A (agents) every minute based on observed TPS to find near-optimal configurations.

Experimental results

Research questions

• RQ1 Can GA3C achieve substantial speedups over CPU-based A3C by centralizing the DNN on the GPU and batching predictions and training?
• RQ2 How do predictor/trainer/agent counts affect GPU utilization, training throughput (TPS), and convergence stability in GA3C?
• RQ3 Does dynamic auto-tuning of N_P, N_T, and N_A improve convergence speed and robustness across different games and hardware?
• RQ4 What is the impact of policy lag and training batch size on learning stability and convergence in GA3C?
• RQ5 How does GA3C scale with larger DNN architectures compared to the CPU implementation?

Key findings

• GA3C yields up to ~6x speedup for small DNNs and ~45x for larger DNNs over CPU A3C on comparable tasks.
• Optimal configurations typically use multiple agents (often many more than CPU A3C) with balanced predictor/trainer ratios (e.g., N_P ≈ N_T and N_A set to exploit CPU cores).
• Dynamic auto-tuning of N_P, N_T, and N_A can discover near-optimal configurations and adapt across games (e.g., Pong, Boxing) on the same hardware.
• Increasing training batch size (within a practical range, ~20–40) can improve learning stability and convergence speed by mitigating policy lag, though overly large batches reduce throughput.
• GA3C achieves faster wall-clock learning curves on several Atari games compared to A3C, with GA3C trained in ~1 day achieving competitive scores relative to A3C trained in ~4 days.
• MANUAL measurements indicate substantial GPU idle time in naive setups, underscoring the benefit of batching and centralized GPU usage.