[Paper Review] rlpyt: A Research Code Base for Deep Reinforcement Learning in PyTorch

Adam Stooke, Pieter Abbeel|arXiv (Cornell University)|Sep 3, 2019

Reinforcement Learning in RoboticsComputer Science33 references52 citations

TL;DR

rlpyt provides a modular PyTorch-based codebase implementing all three major deep RL algorithm families (policy gradients, DQN variants, and Q-function policy gradients) with a shared high-throughput infrastructure and various sampling/optimization configurations. It emphasizes single-node parallelism, reproducibility, and practical tools for small- to medium-scale RL research.

ABSTRACT

Since the recent advent of deep reinforcement learning for game play and simulated robotic control, a multitude of new algorithms have flourished. Most are model-free algorithms which can be categorized into three families: deep Q-learning, policy gradients, and Q-value policy gradients. These have developed along separate lines of research, such that few, if any, code bases incorporate all three kinds. Yet these algorithms share a great depth of common deep reinforcement learning machinery. We are pleased to share rlpyt, which implements all three algorithm families on top of a shared, optimized infrastructure, in a single repository. It contains modular implementations of many common deep RL algorithms in Python using PyTorch, a leading deep learning library. rlpyt is designed as a high-throughput code base for small- to medium-scale research in deep RL. This white paper summarizes its features, algorithms implemented, and relation to prior work, and concludes with detailed implementation and usage notes. rlpyt is available at https://github.com/astooke/rlpyt.

Motivation & Objective

• Motivate shared, high-throughput infrastructure to unify three major families of deep RL algorithms.
• Provide modular, reusable implementations of common RL algorithms in PyTorch.
• Enable flexible experimentation through serial, parallel, and asynchronous sampling and optimization configurations.

Proposed method

• Implement three algorithm families on a common infrastructure: policy gradient (A2C, PPO), DQN and variants (Double, Dueling, Categorical, Rainbow, R2D2-like), and Q-function policy gradient (DDPG, TD3, SAC).
• Support replay buffers with n-step returns, sequence replay, prioritized replay, and frame-based buffering for memory efficiency.
• Offer multiple sampling configurations (Serial, Parallel-CPU, Parallel-GPU, Alternating-GPU) and both synchronous and asynchronous optimization using PyTorch DistributedDataParallel (NCCL/gloo).
• Introduce the namedarraytuple data structure to organize arrays with flexible leading dimensions and multi-modal data.
• Ensure OpenAI Gym compatibility with wrappers for env_info and spaces, plus launching utilities for running many experiments on local hardware.

Experimental results

Research questions

• RQ1 Can a single, modular codebase efficiently support multiple deep RL algorithm families with high throughput on single-node hardware?
• RQ2 How do different sampling and optimization configurations (serial, parallel CPU/GPU, asynchronous) affect throughput and learning performance in common RL benchmarks?
• RQ3 What practical data structures (like namedarraytuple) improve data organization and extensibility across diverse RL algorithms?

Key findings

• rlpyt reproduces competitive learning curves on Atari and MuJoCo environments using a single codebase.
• The R2D2-like recurrent replay setup is demonstrated in a non-distributed setting, achieving high sample throughput.
• Asynchronous and parallel sampling modes can improve hardware utilization and enable multi-GPU training within a single workstation.
• A new namedarraytuple data structure is introduced to manage multi-modal observations and batched data without flattening.
• The framework emphasizes serial mode for debugging and progressively adopts parallel configurations as needed.