[Paper Review] QuaRL: Quantization for Sustainable Reinforcement Learning.

Srivatsan Krishnan, Maximilian Lam+5 more|arXiv (Cornell University)|Nov 29, 2021

Reinforcement Learning in Robotics35 references4 citations

TL;DR

QuaRL proposes ActorQ, a novel quantization method enabling 8-bit inference in reinforcement learning without reward degradation, achieving 1.5x–2.5x end-to-end training speedups on DeepMind Control Suite by combining full-precision learning with quantized actor inference. The method demonstrates that quantization’s computational gains outweigh accuracy costs in RL when properly applied.

ABSTRACT

Deep reinforcement learning has achieved significant milestones, however, the computational demands of reinforcement learning training and inference remain substantial. Quantization is an effective method to reduce the computational overheads of neural networks, though in the context of reinforcement learning, it is unknown whether quantization's computational benefits outweigh the accuracy costs introduced by the corresponding quantization error. To quantify this tradeoff we perform a broad study applying quantization to reinforcement learning. We apply standard quantization techniques such as post-training quantization (PTQ) and quantization aware training (QAT) to a comprehensive set of reinforcement learning tasks (Atari, Gym), algorithms (A2C, DDPG, DQN, D4PG, PPO), and models (MLPs, CNNs) and show that policies may be quantized to 8-bits without degrading reward, enabling significant inference speedups on resource-constrained edge devices. Motivated by the effectiveness of standard quantization techniques on reinforcement learning policies, we introduce a novel quantization algorithm, extit{ActorQ}, for quantized actor-learner distributed reinforcement learning training. By leveraging full precision optimization on the learner and quantized execution on the actors, extit{ActorQ} enables 8-bit inference while maintaining convergence. We develop a system for quantized reinforcement learning training around extit{ActorQ} and demonstrate end to end speedups of $>$ 1.5 $ imes$ - 2.5 $ imes$ over full precision training on a range of tasks (Deepmind Control Suite). Finally, we break down the various runtime costs of distributed reinforcement learning training (such as communication time, inference time, model load time, etc) and evaluate the effects of quantization on these system attributes.

Motivation & Objective

• To investigate whether quantization’s computational benefits outweigh accuracy costs in deep reinforcement learning.
• To evaluate standard quantization techniques (PTQ and QAT) across diverse RL tasks, algorithms, and architectures.
• To design a new quantization-aware training method, ActorQ, for distributed actor-learner RL that maintains convergence with 8-bit inference.
• To analyze and optimize system-level runtime costs (communication, inference, model loading) in distributed RL training under quantization.
• To demonstrate end-to-end speedups in training while preserving policy performance through quantization.

Proposed method

• Applying post-training quantization (PTQ) and quantization-aware training (QAT) to a broad set of RL environments (Atari, Gym), algorithms (A2C, DDPG, DQN, D4PG, PPO), and models (MLPs, CNNs).
• Introducing ActorQ, a novel training framework that uses full-precision optimization on the learner and quantized inference on the actors, enabling 8-bit policy execution.
• Designing a system around ActorQ to support end-to-end quantized RL training with minimal performance drop.
• Breaking down distributed RL training costs (communication, inference, model loading) and measuring quantization’s impact on each.
• Using standard quantization techniques with careful calibration to maintain policy performance during training and inference.
• Evaluating the tradeoff between computational efficiency and policy accuracy across multiple benchmarks and model configurations.

Experimental results

Research questions

• RQ1 Can standard quantization techniques like PTQ and QAT be effectively applied to deep reinforcement learning policies without significant reward degradation?
• RQ2 To what extent do 8-bit quantized policies maintain performance across diverse RL tasks and architectures?
• RQ3 Can a new training algorithm like ActorQ enable stable convergence during distributed RL training while using 8-bit inference?
• RQ4 How does quantization affect the various components of distributed RL training workloads (e.g., communication, inference, model loading)?
• RQ5 What end-to-end speedups can be achieved in RL training through quantization without sacrificing policy performance?

Key findings

• Policies can be quantized to 8 bits without degrading cumulative return on Atari and Gym environments, enabling significant inference speedups on edge devices.
• ActorQ enables stable training convergence in distributed actor-learner RL while maintaining 8-bit inference, preserving policy performance.
• End-to-end training speedups of 1.5× to 2.5× were achieved on DeepMind Control Suite using the proposed quantized training system.
• Quantization reduces inference and model loading times significantly, with minimal impact on communication overhead in distributed training.
• The study confirms that the computational benefits of quantization outweigh the accuracy costs when using appropriate techniques like QAT and ActorQ.
• System-level profiling shows that inference and model loading are major contributors to training latency, which are effectively reduced via quantization.