Source-linked AI summary
Acme: A Research Framework for Distributed Reinforcement Learning
Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Danila Sinopalnikov, Piotr Stańczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Arnold, Manu Orsini, Alexis Jacq, Johan Ferret, Nino Vieillard, Seyed Kamyar Seyed Ghasemipour, Sertan Girgin, Olivier Pietquin, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, Sarah Henderson, Abe Friesen, Ruba Haroun, Alex Novikov, Sergio Gómez Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Andrew Cowie, Ziyu Wang, Bilal Piot, Nando de Freitas
TL;DR
Acme addresses the growing scale and complexity of deep RL, which hinder rapid prototyping and reproducibility, by providing modular components and reusable agent interfaces. The framework supports local and distributed execution, includes reference agents across online, offline, imitation, and demonstration settings, and reports baselines and state-of-the-art implementations across varied domains. The paper’s second version increases modularity, expands offline and imitation coverage, and adds more agents.
Problem
Increasing model scale and RL algorithm complexity make it harder to rapidly prototype ideas and reproduce published algorithms.
Method
Acme provides modular components and agent builders spanning policies, actors, replay, learners, training loops, evaluation, logging, and checkpointing, reusable from synchronous to distributed execution.
Results
Acme provides reference implementations and baselines with state-of-the-art performance across varied domains, including online, offline, imitation, and demonstration-based agents.
Takeaways & Limitations
Acme offers reusable building blocks intended to improve RL reproducibility and provide yardsticks for measuring progress.
Takeaways & Limitations
The paper gives a high-level overview of core interfaces, while precise APIs may change as the living framework evolves.
Abstract
from arXiv · showhide
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.
1. Introduction
Modern RL has gained performance through larger models, more data, and increasingly complex agent designs, but these trends make experimentation and reproduction harder. Acme addresses this tension with modular components that support both rapid iteration and scalable execution.
- Motivation: Increasingly integrationist agent designs combine many independent components, contributing to implementation complexity.Examples include intrinsic rewards, auxiliary tasks, and specialized neural network architectures.
- Motivation: Modern RL performance benefits from scaling function approximators and training data, including distributed interaction with parallel environments.This scaling introduces engineering, algorithmic, infrastructure, and reproducibility challenges.
- Acme’s approach: Acme provides modular components spanning networks, losses, policies, actors, learners, replay buffers, training loops, logging, and checkpointing.These components can be combined or scaled while preserving straightforward abstractions and simpler synchronous execution.
- Evaluation: The paper experiments with Acme agents across varied domains to demonstrate state-of-the-art performance.Section 5 evaluates whether the framework’s agents can achieve strong performance across multiple domains.
- Paper update: The paper’s second version increases modularity, emphasizes offline and imitation algorithms, adds learning-from-demonstrations coverage, and implements more core algorithms.The update also reflects broader use of JAX while retaining platform-agnostic base interfaces.
2. Modern Reinforcement Learning
RL studies agents that interact with environments by mapping observations to actions and optimizing future rewards. Acme supports online, offline, imitation, and demonstration-based formulations that differ in how data and rewards are obtained.
- Core formulation: An RL agent repeatedly selects actions from observations, while the environment returns rewards and subsequent observations.This interaction forms the standard agent-environment loop.
- Core formulation: A policy maps an agent’s experienced observation history to actions and may use recurrent state in partially observable environments.Acme discusses both feed-forward and recurrent policy representations.
- Core formulation: RL objectives commonly maximize expected discounted future rewards, with the discount factor prioritizing near-term rewards.Acme also supports agents optimizing other aggregate performance measures.
- Learning settings: Online RL learns through trial-and-error interaction and active exploration, whereas offline RL learns from an existing experience collection without environment interaction.Offline RL is relevant when environment interactions are costly or dangerous.
- Learning settings: Imitation learning uses demonstrations when a reward is unavailable or difficult to specify, while learning from demonstrations combines environmental rewards with demonstration data.Demonstrations can facilitate exploration in sparse-reward tasks.
3. Acme
Acme organizes RL agents into reusable components and shared execution interfaces, allowing the same implementation to run locally or in distributed settings. Its environment loop, actors, replay, learners, and builders separate interaction, storage, and updating responsibilities.
- Framework design: Acme reuses the same components between single-process implementations and large distributed agents, enabling algorithms to be implemented once across execution modes.This design targets readable and efficient RL implementations.
- Framework design: An Acme agent comprises an environment loop, actor, replay system, learner, builder, and experiment runners for local or distributed training.These components also support reuse in offline settings with limited or zero modification.
- Framework scope: Acme’s core interfaces are presented as a high-level overview because precise APIs may change as the framework evolves.The authors direct readers to documentation for up-to-date API details.
- Environment loop: The environment step produces a reward, environmental discount, new observation, and end-of-episode indicator.An environmental discount distinguishes termination from truncation; bootstrapping applies to truncation but not termination.
- Environment loop: The environment loop repeatedly coordinates actor-environment interaction for a specified number of transitions or complete episodes.The accompanying pseudocode resets the environment, selects actions, steps the environment, observes the result, and updates the actor.
- Actors: Actors separate action selection, observation handling, initial-state setup, and parameter updates from the environment loop.This delineation supports simpler distribution and parallelization by separating data generation from training.
- Actors: Actor updates commonly retrieve the learner’s latest policy parameters through a variable source in distributed agents.GenericActor provides reusable boilerplate for common distributed-agent logic, while agent-specific behavior is delegated to other components.
- Actors: Actor initialization can create recurrent state at episode beginnings, including the memory state used by recurrent agents such as R2D2.This state supports agents whose policies maintain memory across observations.
3.3. Experience replay and data storage
Acme separates acting, learning, and storage so experience can be collected, replayed, and consumed across online, off-policy, and offline settings. Reverb supplies distributed replay storage, sampling, prioritization, and rate control.
- Reverb provides a client-server storage system for high-throughput experience replay, with servers colocated with agents or deployed separately.
- Reverb supports prioritized replay by assigning scalar priorities when data is written or updating them later.
- Acme exposes replay data to learners through dataset iterators, while Adders convert action and timestep observations into algorithm-specific stored data.
- Reverb’s RateLimiter enforces a desired learning-to-acting rate while allowing actor and learner processes to run unblocked within tolerance.
- The learner interface can use a fixed offline dataset instead of experience replay, simplifying offline algorithm implementations.
3.5. Defining agents using Builders
Acme uses builders to encapsulate complete agent algorithms while keeping environment and network choices configurable. Experiment tooling then assembles these components, supports evaluation and logging, and can run the same interface at distributed scale.
- 3.5. Defining agents using Builders: The builder abstraction defines an agent algorithm with minimal assumptions about the learning environment.
- 3.5. Defining agents using Builders: Environment specifications provide input-output types and shapes, while separate network modules define the policy architecture and enable agent reuse.
- 3.5. Defining agents using Builders: Builder interfaces construct policies, replay tables, adders, actors, and learners from environment, network, replay, and dataset components.
- 3.6. Running an agent: An Acme experiment combines a builder, network definitions, and an environment through factories collected in ExperimentConfig.
- 3.6. Running an agent: Practitioners can adapt an algorithm to a new environment or network by supplying factories, without modifying the builder.
- 3.6. Running an agent: Acme evaluates agents in parallel with training using non-recording evaluation actors, configurable evaluation policies, and optional custom metrics.
3.7. Distributed agents
Acme scales reinforcement learning by separating data generation, learning, and storage into components that can run asynchronously across processes or machines. Its distributed experiment mechanism uses Launchpad to construct and launch this component topology, while alternative runners support fixed offline data.
- 3.7. Distributed agents: Distributed data generation addresses RL scaling by running multiple environment interactions asynchronously and in parallel with learning.
- 3.7. Distributed agents: Acme’s acting, learning, and storage split can be deployed across threads, processes, or machines using the same agent building blocks.
- 3.7. Distributed agents: Launchpad represents distributed computation as a directed graph whose service nodes communicate through remote procedure calls.
- 3.7. Distributed agents: Acme’s distributed experiment constructs Launchpad nodes for environment loops, learners, and the Reverb service.
- 3.7. Distributed agents: Separating execution from algorithm logic allows scaling and performance optimizations to be handled during distributed experiment construction.
- 3.7. Distributed agents: The offline experiment runner removes actor and adder components and trains from a fixed dataset, while evaluation may still use an environment or another mechanism.
4. Agent Algorithms
Acme organizes reinforcement-learning agents from data collection through policy learning, while supporting shared value-based principles and diverse online, recurrent, continuous-control, and offline algorithms.
- Agent structure: An Acme agent includes the apparatus for collecting data and continually learning improved policies, with primary algorithmic differences often concentrated in the learner.The learner consumes data and updates policy parameters.
- Shared principles: Acme supports value estimates ranging from direct Monte Carlo formulations to recursive Bellman formulations, including intermediate methods such as n-step TD errors.n-step returns replace a single reward with a discounted sum over multiple rewards and bootstrap from a later state-action pair.
- Online RL: R2D2 extends DQN with recurrent Q-functions and sequence replay to better address partially observable environments.Its implementation replaces feed-forward networks with recurrent networks and adapts learning to trajectories.
- Online RL: D4PG is an off-policy actor-critic algorithm for continuous actions that combines distributional critics, n-step transitions, and distributed experience generation.Acme experiments also use a single-actor, non-distributed D4PG variant for comparison with other continuous-control methods.
- Online RL: TD3 modifies DDPG with twin conservative Q estimates, delayed policy updates, and target-action noise to reduce overestimation and smooth critic targets.These changes occur in the learner’s loss computation.
- Offline RL: BVE learns behavior values from offline datasets rather than recovering the optimal policy, eliminating the target-network max operator during Bellman updates.For Atari, the paper uses SARSA tuples with bootstrapping and one-step policy improvement at evaluation.
5. Experiments
Acme evaluates modular agents across online, offline, imitation, demonstration, and distributed settings using standardized tasks and fixed interaction budgets. The experiments report broad algorithmic comparisons and show that scaling learner hardware can accelerate training without changing episode return.
- Experimental design: Acme evaluates agents across multiple settings, including online, offline, imitation learning, and learning from demonstrations.Agents can be reused across setups, such as TD3 online, TD3+BC offline, AIL with TD3, and TD3fD with demonstrations.
- Continuous Control: Continuous-control experiments evaluate PPO, TD3, SAC, D4PG, and MPO variants with shared three-layer, 256-unit network architectures.Gym environments run for 10 million steps, DM Control environments for 4 million, except humanoid for 25 million.
- Discrete Control: Discrete-control experiments evaluate DQN, R2D2, Impala, and Munchausen-DQN on five Atari games, with recurrent architectures still helping in the 200-million-frame regime.Results use five random seeds in a distributed setup and overall match reported literature results.
- Offline Reinforcement Learning: Offline results compare TD3, TD3+BC, CQL, CRR, and CRR+SARSA on D4RL locomotion datasets, while Atari experiments compare BVE, REM, DQN, and BC.On dense-reward Atari games BVE performs better than REM and DQN, whereas DQN and REM perform better on sparse-reward Gravitar.
- Imitation and Demonstrations: Performance of all imitation-learning methods improves with more demonstration trajectories, while demonstration data produces better learning-curve performance in sparse Adroit tasks.Imitation learning uses 1, 4, or 11 trajectories; learning-from-demonstrations experiments vary the replay-buffer ratio of demonstration to agent transitions.
- Speed: 1.7x and 5.6x learner speed-ups over GPU-V100 are obtained on TPU-v2 1x1 and TPU-v2 2x2 respectively, with identical episode return.The comparison uses distributed R2D2 on Pong with 256 actors and changes only learner hardware.
6. Related work
Acme is situated among deep RL frameworks that emphasize single-process algorithm implementations or distributed computation. Its design instead targets a balance between simplicity, modularity, and scale.
- Related frameworks: OpenAI Baselines and TF-Agents are TensorFlow 1.X frameworks that strive to express numerous algorithms in single-process format.
- Related frameworks: Dopamine focuses on single-process agents in the DQN family and distributional variants, while Fiber and Ray support distributed computations.
- Acme: Acme is designed to balance simplicity, modularity, and scale, supporting small-scale experimentation and high data-throughput execution.
7. Conclusion
Acme is presented as a modular, lightweight framework for scalable and fast RL research iteration, with expanded modularity, algorithm coverage, and emphasis on offline and imitation learning. Its tools aim to improve reproducibility and support creation of new RL agents.
- Acme supports scalable and fast iteration of RL research ideas through a modular, lightweight framework.
- The framework supports both single-actor and distributed training paradigms while providing agent baselines with state-of-the-art performance.
- This second paper version increases modularity, emphasizes offline and imitation algorithms, and expands the set of implemented core algorithms.
- Acme’s tools are intended to improve RL reproducibility and give researchers simple building blocks for creating new agents.
8. Author Contributions & Acknowledgments
The paper credits a broad group of contributors for Acme’s design, agents, infrastructure, experiments, documentation, and manuscript development. Acknowledgments also recognize external feedback on the paper and codebase.
- The initial Acme design and implementation involved Hoffman, Barth-Maron, Aslanides, and Shahriari, with contributions spanning actors, policies, environment loops, DQN, and JAX agents.
- Momchev, Sinopalnikov, and Stanczyk developed the Builder design and helped integrate modular and imitation-learning agents into Acme.
- Ramos, Raichuk, Hussenot, and Dadashi contributed Builder policy construction, Adders, efficiency improvements, SAC, agents, and experimental coordination.
- Additional contributors implemented or improved agents including GAIL, DQN, R2D2, PPO, and MPO, while others supported documentation, efficiency, integrations, and tuning.
- The authors thank contributors who provided manuscript, graphical, and earlier codebase feedback.
A. Algorithm Tag Descriptions
The algorithm tags describe the action space, training regime, approximated neural-network function, and value-learning logic associated with each agent.
- Discrete and Continuous Actions describe the nature of the action space supported by agents.
- On-Policy, Off-Policy, and Offline describe how training data is gathered or whether data gathering is absent.
- Q-Network, V-Network, and Policy-Network identify the functions approximated with neural networks.
- MC and Bootstrapping describe whether value learning uses full trajectories or the learner’s own estimate as the target.