Acme: A Research Framework for Distributed Reinforcement Learning

Abstract: Deep reinforcement learning (RL) has led to many recent and groundbreaking advances. However, these advances have often come at the cost of both increased scale in the underlying architectures being trained as well as increased complexity of the RL algorithms used to train them. These increases have in turn made it more difficult for researchers to rapidly prototype new ideas or reproduce published RL algorithms. To address these concerns this work describes Acme, a framework for constructing novel RL algorithms that is specifically designed to enable agents that are built using simple, modular components that can be used at various scales of execution. While the primary goal of Acme is to provide a framework for algorithm development, a secondary goal is to provide simple reference implementations of important or state-of-the-art algorithms. These implementations serve both as a validation of our design decisions as well as an important contribution to reproducibility in RL research. In this work we describe the major design decisions made within Acme and give further details as to how its components can be used to implement various algorithms. Our experiments provide baselines for a number of common and state-of-the-art algorithms as well as showing how these algorithms can be scaled up for much larger and more complex environments. This highlights one of the primary advantages of Acme, namely that it can be used to implement large, distributed RL algorithms that can run at massive scales while still maintaining the inherent readability of that implementation. This work presents a second version of the paper which coincides with an increase in modularity, additional emphasis on offline, imitation and learning from demonstrations algorithms, as well as various new agents implemented as part of Acme.

• The paper details how Acme simplifies distributed reinforcement learning research by modularizing agents into scalable components like actors, learners, and replay systems.
• It demonstrates the integration of various RL algorithms including DQN, SAC, and TD3, enabling efficient experimentation across discrete and continuous control tasks.
• The framework’s support for both online and offline learning highlights its potential to accelerate advancements in robotics, autonomous systems, and complex decision-making applications.

Acme: A Distributed Reinforcement Learning Framework

The paper "Acme: A Research Framework for Distributed Reinforcement Learning" details the development and functionalities of Acme, a framework explicitly designed to facilitate the construction and experimentation of distributed reinforcement learning (RL) algorithms. The primary motivation behind the framework is to address the increasing complexity and computational demands encountered in modern RL research, which often involve large-scale architectures and intricate algorithms.

Core Features of Acme

Acme distinguishes itself through its modular and scalable architecture, allowing researchers to easily prototype and test new ideas. This is achieved by dividing RL agents into well-defined components, including actors, learners, and replay systems, which can be composed and scaled across different computation settings, from local to distributed environments.

• Actors are responsible for interacting with an environment to generate experience data. They are designed to evaluate policies and collect observations seamlessly, supporting both synchronous and asynchronous execution modes.
• Replay Systems are implemented via Reverb, providing a robust and high-throughput data storage mechanism to manage and sample experience data efficiently. This allows for various data sampling strategies, supporting off-policy, on-policy, and mixed approaches.
• Learners update agent parameters based on samples drawn from the replay system. The learner architecture is flexible, enabling the use of various algorithms, whether they rely on bootstrapping or Monte Carlo methods for value estimation.

The framework also includes comprehensive support for offline reinforcement learning, allowing for the direct use of static datasets in experiments. This is particularly beneficial when online data collection is costly or impractical.

Significant Algorithms and Extensions

Acme implements a wide range of reinforcement learning algorithms, both classic and contemporary, effectively offering state-of-the-art reference implementations. These include:

• DQN Variants (including Double DQN and Dueling DQN): These cater to discrete action spaces and emphasize enhancements like distributional value functions and prioritized experience replay.
• SAC and TD3: Built for continuous control tasks, these algorithms introduce specific optimizer adaptations and noise strategies to stabilize learning in continuous settings.
• MPO and Distributional Variants: These offer a perspective grounded in the reinforcement learning as inference paradigm, emphasizing policy improvements through proximal optimization techniques.

Furthermore, through its support for distributed systems, Acme can harness computational resources to scale up experiments considerably, allowing for parallel environment interactions and accelerated learning processes.

Practical Implications and Future Directions

The advent of Acme has significant implications for both academic research and industry applications. By reducing the complexity barrier associated with implementing sophisticated RL systems, it democratizes access to cutting-edge algorithms, facilitating reproducibility and progress in computational sciences.

In practical terms, Acme's composability and scalability enable large-scale RL experiments that were previously infeasible, potentially leading to advancements in various domains such as robotics, autonomous systems, and complex decision-making applications.

Looking ahead, Acme's integration with evolving machine learning technologies like JAX and future support for even broader algorithmic extensions could further cement its role as a pivotal tool in the reinforcement learning community. As RL continues to push boundaries in performance and applicability, frameworks like Acme will be crucial in bridging theoretical innovation with experimental proficiency.