Intro to DeepMind’s Reinforcement Learning Framework “Acme” | by Andreas Stöffelbauer | Jun, 2021 | Towards Data Science | Towards Data Science

Acme is a Python-based research framework for reinforcement learning, open sourced by Google’s DeepMind in 2020. It was designed to simplify the development of novel RL agents and accelerate RL research. According to their own statement, Acme is used on a daily basis at DeepMind, which is spearheading research in reinforcement learning and artificial intelligence.
For a recent university project, I decided to learn about Acme and use it to implement different RL algorithms. I found it to be brilliant and I really enjoyed working with it.
Getting started with Acme is relatively easy too. This is because there are several points of entry, at different levels of complexity. In other words, the framework is not only meant for advanced researchers but also allows beginners to implement even simple algorithms — somewhat similar to how TensorFlow and PyTorch are used by both beginners and experts.
On the downside, however, as the framework is still quite new, there isn’t really a thorough documentation available, and there aren’t any great tutorials out there either.
This blog post can hopefully be a step in the right direction. It is not meant to be or replace a thorough documentation, but rather a concise, practical introduction to Acme. Above all, it should give you an idea of the design choices underlying the framework, and what that means for the implementation of RL algorithms.
In particular, I am going to discuss my implementations of two simple algorithms, SARSA and Q learning, with the goal to play the game Blackjack. Once you understand what actors and agents are and how they are designed in Acme, I have no doubt that you will quickly get the idea of how to implement (probably) any reinforcement learning algorithm you can think of.
Let’s dive into an actual example. As already mentioned, we want our agent to play Blackjack.
Acme agents are not designed to interact with Gym environments. Instead, DeepMind has their own RL environment API. The difference is mainly in how the timesteps are represented.
Fortunately, however, you can still make use of Gym environments too since Acme’s developers have provided wrapper functions for this.
env = acme.wrappers.GymWrapper(gym.make('Blackjack-v0'))
There are 32 x 11 x 2 states in Blackjack, although not all those states can actually occur during a game, and there are two actions, “hit” and “stick”. For a description of these dimensions and why not all states can occur, you can check out the environment on GitHub here. I know this is not really obvious at first.
It is critical to understand that there is a distinction between actors, learners, and agents. Actors interact with the environment. That is, they observe states and take actions according to some action-selection policy. The image below illustrates this.
The learner consumes the data collected by the actor in order to learn or improve the policy, usually in an iterative, online fashion. For example, learning might consist of updates to a neural network’s parameters. The new parameters are passed to the actor, who then behaves according to the updated policy.
The agent simply combines both the acting and the learning component, but usually no additional reinforcement learning logic needs to be implemented. The image below includes all three components.
One of the main reasons for this factorization into actors, learners, and agents is to facilitate distributed reinforcement learning. If we do not care about that and if the algorithm is simple enough, however, it can also be sufficient to only implement the actor and simply integrate the learning steps within the actor’s update method. For the sake of simplicity, this also the approach that I am taking here.
For example, the random agent below inherits from the acme.Actor class. The methods that must be implemented by the developer (you) are select_action, observe_first, observe, and update. As just mentioned, the latter is where the learning happens without the additional learner component. Note that this agent would work the same way without subclassing acme.Actor. The base class simply determines the methods that you must overwrite. This also ensures that the agent integrates as intended with other Acme components such as the environment loop that I am going to introduce below.
This agent uses a policy that simply chooses hit or stick randomly, but generally the framework allows you huge flexibility with respect to how you want to implement a policy. Later, you will see an epsilon-greedy policy. In other cases, the policy may include a neural network that you can implement with TensorFlow, PyTorch, or JAX. In that sense, Acme is framework agnostic, so you can combine it with whatever machine learning library you prefer.
In the update method, the actor usually only pulls the most recent parameters from the learner. If you do not use a separate learner, however, the RL logic goes inside the update method (you will see that later).
If you already know a bit about reinforcement learning and have already implemented an RL algorithm, the following loop will be very familiar to you. Each episode consists of four steps, which are repeated until a terminal state is reached.
observe a statetake an action according to the behavior policyobserve a rewardupdate the policy
observe a state
take an action according to the behavior policy
observe a reward
update the policy
In code ...
Sometimes you may have to implement a loop like that, especially if you wish to customize it. But mostly, this loop will always be exactly the same.
Conveniently, there is a shortcut in Acme: the EnvironmentLoop, which performs pretty much exactly the steps seen above. You just have to pass your environment and agent instances and then you can run either a single episode or as many as you want with a single line of code. There are also a number of loggers available that track important metrics such as the number of steps taken in each episode and the collected rewards.
# init Acme's environment looploop = EnvironmentLoop(env, agent, logger=InMemoryLogger())# run a single episodeloop.run_episode()# or run multiple episodesloop.run(10)
Of course, a random agent is not very useful. My promise was to show how to implement some actual reinforcement learning algorithms. So here we go.
By the way, if you are not familiar with RL at all, check out the book Reinforcement Learning: An Introduction by Sutton and Barto (2018). The two algorithms you usually learn first — either in the book or if you take an RL course at university — are SARSA and Q learning.
By now, you already know how Acme agents (or actors) are designed. Let’s look at how you can implement the SARSA algorithm in Acme.
SARSA is an on-policy algorithm whose updates depend on the state, action, reward, next state, and next action (hence the name). Since this is not meant to be a theoretical RL tutorial, I do not go into the details of the algorithm itself here.
First things first, inside the agent’s __init__ method, we initialize Q, the state-action value matrix, and the behavior policy, which is an epsilon greedy policy here. Note also that this agent must always store its last timestep, action, and next timestep as they are needed in the update step. So we also initialize them.
In observe, there is usually not much you have to do. In this case, we simply store the observed timestep and the action that was taken. This is not always necessary, however. Sometimes, for example, you may want to store timesteps (and whole trajectories) in a dataset or a replay buffer instead. Acme also provides dataset and adder components for this purpose. In fact, there is also a library which was developed by DeepMind too. It’s called Reverb (see GitHub here).
The transform_state method above is only a helper function to put states into the right format in order to correctly index the Q matrix.
Finally, to train SARSA on the environment for 500,000 episodes, simply run
agent = SarsaAgent()loop = EnvironmentLoop(env, agent, logger=InMemoryLogger())loop.run(500000)
The Q learning agent below is very similar to the SARSA agent. They only differ in the way how updates to the Q matrix are made. This is because Q learning is an off-policy algorithm.
To train the Q learning agent on the environment for 500,000 episodes, run
agent = QLearningAgent()loop = EnvironmentLoop(env, agent, logger=InMemoryLogger())loop.run(500000)
I think Acme is a really great framework for reinforcement learning because you don’t have to develop your algorithms from scratch. So instead of figuring out by yourself how to write readable and reproducible RL code, you can just rely on the undoubtedly smart researchers and developers at DeepMind, who have done that for you.
Acme allows you to implement any reinforcement learning algorithm, and you can combine it with any other machine learning framework, including TensorFlow, PyTorch and JAX.
If you want to learn more about Acme, you can read DeepMind’s research paper and check out their GitHub repository.
There you will also find the implementation of some common algorithms such as Deep Q-Networks (DQN), Deep Deterministic Policy Gradient (DDPG), Monte Carlo Tree Search (MCTS), Behavior Cloning (BC), IMPALA, and so on.
No matter if you are a advanced researcher or only a beginner interested in reinforcement learning, I encourage you to try it out.
Thank you for reading. If you have any questions, let me know.
My Jupyter Notebook containing the code can be found here.
If you are interested, you can also check out my reinforcement learning course project. Besides SARSA and Q-learning, I also implemented dyna-Q, prioritized sweeping and a Mote Carlo Tree Search agent.
Hoffman et al (2020): Acme: A Research Framework for Distributed Reinforcement Learning. ArXiv.
Sutton and Barto (2018): Reinforcement Learning: An Introduction.