Deep Reinforcement Learning with TensorFlow 2.0

In this tutorial, I will showcase the upcoming TensorFlow 2.0 features through the lens of deep reinforcement learning (DRL) by implementing an advantage actor-critic (A2C) agent to solve the classic CartPole-v0 environment. While the goal is to showcase TensorFlow 2.0, I will do my best to make the DRL aspect approachable as well, including a brief overview of the field.

In fact, since the main focus of the 2.0 release is making developers’ lives easier, it’s a great time to get into DRL with TensorFlow — our full agent source is under 150 lines! The code is available as a notebook here and online on Google Colab here.

Setup

As TensorFlow 2.0 is still in an experimental stage, I recommend installing it in a separate (virtual) environment. I prefer Anaconda, so I’ll illustrate with it:

> conda create -n tf2 python=3.6 
> source activate tf2
> pip install tf-nightly-2.0-preview

Let’s quickly verify that everything works as expected:

>>> import tensorflow as tf 
>>> print(tf.__version__)
2.0.0-dev20190129
>>> print(tf.executing_eagerly())
True

Note that we’re now in eager mode by default!

>>> print("1 + 2 + 3 + 4 + 5 =", tf.reduce_sum([1, 2, 3, 4, 5])) 
1 + 2 + 3 + 4 + 5 = tf.Tensor(15, shape=(), dtype=int32)

If you’re not yet familiar with the eager mode, then, in essence, it means that computation is executed at runtime, rather than through a pre-compiled graph. You can find a good overview in the TensorFlow documentation.

Deep Reinforcement Learning

Generally speaking, reinforcement learning is a high-level framework for solving sequential decision-making problems. An RL agent navigates an environment by taking actions based on some observations, receiving rewards as a result. Most RL algorithms work by maximizing sum of rewards an agent collects in a trajectory, e.g., during one in-game round.

The output of an RL based algorithm is typically a policy - a function that maps states to actions. A valid policy can be as simple as a hard-coded no-op action. A stochastic policy is represented as a conditional probability distribution of actions, given some state.

Actor-Critic Methods

RL algorithms are often grouped based on the objective function they are optimized with. Value-based methods, such as DQN, work by reducing the error of the expected state-action values.

Policy Gradients methods directly optimize the policy itself by adjusting its parameters, typically via gradient descent. Calculating gradients fully is usually intractable, so instead, they are often estimated via monte-carlo methods.

The most popular approach is a hybrid of the two: actor-critic methods, where agents policy is optimized through policy gradients, while value based method is used as a bootstrap for the expected value estimates.

Deep Actor-Critic Methods

While much of the fundamental RL theory was developed on the tabular cases, modern RL is almost exclusively done with function approximators, such as artificial neural networks. Specifically, an RL algorithm is considered “deep” if the policy and value functions are approximated with deep neural networks.

(Asynchronous) Advantage Actor-Critic

Over the years, a number of improvements have been added to address sample efficiency and stability of the learning process.

First, gradients are weighted with returns: discounted future rewards, which somewhat alleviates the credit assignment problem, and resolves theoretical issues with infinite timesteps.

Second, an advantage function is used instead of raw returns. Advantage is formed as the difference between returns and some baseline (e.g., state-action estimate) and can be thought of as a measure of how good a given action is compared to some average.

Third, an additional entropy maximization term is used in the objective function to ensure agent sufficiently explores various policies. In essence, entropy measures how random a probability distribution is, maximized with uniform distribution.

Finally, multiple workers are used in parallel to speed up sample gathering while helping decorrelate them during training.

Incorporating all of these changes with deep neural networks we arrive at the two of the most popular modern algorithms: (asynchronous) advantage actor critic, or A3C/A2C for short. The difference between the two is more technical than theoretical: as the name suggests, it boils down to how the parallel workers estimate their gradients and propagate them to the model.

With this I will wrap up our tour of DRL methods as the focus of the blog post is more on the TensorFlow 2.0 features. Don’t worry if you’re still unsure about the subject; things should become clearer with code examples. If you want to learn more, then one good resource to get started is Spinning Up in Deep RL.

Advantage Actor-Critic with TensorFlow 2.0

Now that we’re more or less on the same page let’s see what it takes to implement the basis of many modern DRL algorithms: an actor-critic agent, described in the previous section. For simplicity, we won’t implement parallel workers, though most of the code will have support for it. An interested reader could then use this as an exercise opportunity.

As a testbed, we will use the CartPole-v0 environment. Somewhat simplistic, it’s still a great option to get started with. I always rely on it as a sanity check when implementing RL algorithms.

Policy & Value via Keras Model API

First, let’s create the policy and value estimate NNs under a single model class:

import numpy as np
import tensorflow as tf
import tensorflow.keras.layers as kl
class ProbabilityDistribution(tf.keras.Model):
def call(self, logits):
# sample a random categorical action from given logits
return tf.squeeze(tf.random.categorical(logits, 1), axis=-1)
class Model(tf.keras.Model):
def __init__(self, num_actions):
super().__init__('mlp_policy')
# no tf.get_variable(), just simple Keras API
self.hidden1 = kl.Dense(128, activation='relu')
self.hidden2 = kl.Dense(128, activation='relu')
self.value = kl.Dense(1, name='value')
# logits are unnormalized log probabilities
self.logits = kl.Dense(num_actions, name='policy_logits')
self.dist = ProbabilityDistribution()
def call(self, inputs):
# inputs is a numpy array, convert to Tensor
x = tf.convert_to_tensor(inputs, dtype=tf.float32)
# separate hidden layers from the same input tensor
hidden_logs = self.hidden1(x)
hidden_vals = self.hidden2(x)
return self.logits(hidden_logs), self.value(hidden_vals)
def action_value(self, obs):
# executes call() under the hood
logits, value = self.predict(obs)
action = self.dist.predict(logits)
# simpler option, will be clear below why we don't use it
# action = tf.random.categorical(logits, 1)
return np.squeeze(action, axis=-1), np.squeeze(value, axis=-1)

And let’s verify the model works as expected:

import gym env = gym.make('CartPole-v0')
model = Model(num_actions=env.action_space.n)
obs = env.reset()
# no feed_dict or tf.Session() needed at all
action, value = model.action_value(obs[None, :])
print(action, value) # [1] [-0.00145713]

Things to note here:

  • Model layers and execution path are defined separately
  • There is no “input” layer; model will accept raw numpy arrays
  • Two computation paths can be defined in one model via functional API
  • A model can contain helper methods such as action sampling
  • In eager mode, everything works from raw numpy arrays

Random Agent

Now we can move on to the fun stuff — the A2CAgent class. First, let’s add a test method that runs through a full episode and returns sum of rewards.

class A2CAgent:
def __init__(self, model):
self.model = model
def test(self, env, render=True):
obs, done, ep_reward = env.reset(), False, 0
while not done:
action, _ = self.model.action_value(obs[None, :])
obs, reward, done, _ = env.step(action)
ep_reward += reward
if render:
env.render()
return ep_reward

Let’s see how much our model scores with randomly initialized weights:

agent = A2CAgent(model)
rewards_sum = agent.test(env)
print("%d out of 200" % rewards_sum)
# 18 out of 200

Not even close to optimal, time to get to the training part!

Loss / Objective Function

As I’ve described in the DRL overview section, an agent improves its policy through gradient descent based on some loss (objective) function. In actor-critic we train on three objectives: improving policy with advantage weighted gradients plus entropy maximization, and minimizing value estimate errors.

import tensorflow.keras.losses as kls
import tensorflow.keras.optimizers as ko
class A2CAgent:
def __init__(self, model):
# hyperparameters for loss terms
self.params = {'value': 0.5, 'entropy': 0.0001}
self.model = model
self.model.compile(
optimizer=ko.RMSprop(lr=0.0007),
# define separate losses for policy logits and value estimate
loss=[self._logits_loss, self._value_loss]
)
def test(self, env, render=True):
# unchanged from previous section
...
def _value_loss(self, returns, value):
# value loss is typically MSE between value est and returns
loss = kls.mean_squared_error(returns, value)
return self.params['value']*loss
def _logits_loss(self, acts_and_advs, logits):
# a trick to input actions and advantages through same API
actions, advs = tf.split(acts_and_advs, 2, axis=-1)
# sparse categorical CE loss obj, supports sample_weight arg
# from_logits arg ensures normalized probabilities
wght_sp_ce = kls.SparseCategoricalCrossentropy(from_logits=True)
# policy loss: policy gradients, weighted by advantages
# note: we only calculate the loss on the actions we've taken
actions = tf.cast(actions, tf.int32)
policy_loss = wght_sp_ce(actions, logits, sample_weight=advs)
# entropy loss can be calculated via CE over itself
entropy_loss = \
kls.categorical_crossentropy(logits, logits, from_logits=True)
# here signs are flipped because optimizer minimizes
return policy_loss - self.params['entropy']*entropy_loss

Agent Training Loop

Finally, there’s the train loop itself. It’s relatively long but fairly straightforward: collect samples, calculate returns and advantages, and train the model on them.

class A2CAgent:
def __init__(self, model):
# hyperparameters for loss terms
self.params = {'value': 0.5, 'entropy': 0.0001, 'gamma': 0.99}
# unchanged from previous section
...

def train(self, env, batch_sz=32, updates=1000):
# storage helpers for a single batch of data
actions = np.empty((batch_sz,), dtype=np.int32)
rewards, dones, values = np.empty((3, batch_sz))
obs = np.empty((batch_sz,) + env.observation_space.shape)
# training loop: collect samples, send to optimizer, train
ep_rews = [0.0]
next_obs = env.reset()
for update in range(updates):
for step in range(batch_sz):
obs[step] = next_obs.copy()
actions[step], values[step] = \
self.model.action_value(next_obs[None, :])
next_obs, rewards[step], dones[step], _ = \
env.step(actions[step])
ep_rews[-1] += rewards[step]
if dones[step]:
ep_rews.append(0.0)
next_obs = env.reset()
_, next_val = self.model.action_value(next_obs[None, :])
rets, advs = \
self._returns_advantages(rewards, dones, values, next_val)

# a trick to input actions and advantages through same API
acts_and_advs = \
np.concatenate([actions[:, None], advs[:, None]], axis=-1)
# performs a full training step on the collected batch
# note: no need to mess with gradients, Keras API handles it

_ = self.model.train_on_batch(obs, [acts_and_advs, rets])
return ep_rews
def _returns_advantages(self, rewards, dones, values, next_value):
# next_value is the bootstrap val est of a future state
rets = np.append(np.zeros_like(rewards), next_value, axis=-1)
# returns are calculated as discounted sum of future rewards
for t in reversed(range(rewards.shape[0])):
rets[t] = \
rewards[t] + self.params['gamma'] * rets[t+1] * (1-dones[t])
rets = rets[:-1]
# advantages are returns - baseline, value estimates in our case
advantages = returns - values
return returns, advantages
def test(self, env, render=True):
# unchanged from previous section
...
def _value_loss(self, returns, value):
# unchanged from previous section
...
def _logits_loss(self, acts_and_advs, logits):
# unchanged from previous section
...

Training & Results

We’re now all set to train our single-worker A2C agent on CartPole-v0! Training process shouldn’t take longer than a couple of minutes. After training is complete, you should see an agent successfully achieve the target 200 out of 200 score.

rewards_history = agent.train(env)
print("Finished training, testing...")
print("%d out of 200" % agent.test(env))
# 200 out of 200

In the source code, I include some additional helpers that print out running episode rewards and losses, along with basic plotter for the rewards_history.

Static Computational Graph

With all of this eager mode excitement, you might wonder if static graph execution is even possible anymore. Of course, it is! Moreover, it takes just one additional line to enable it!

with tf.Graph().as_default():
print(tf.executing_eagerly()) # False

model = Model(num_actions=env.action_space.n)
agent = A2CAgent(model)
rewards_history = agent.train(env)
print("Finished training, testing...")
print("%d out of 200" % agent.test(env))
# 200 out of 200

There’s one caveat that during static graph execution we can’t just have Tensors laying around, which is why we needed that trick with CategoricalDistribution during model definition. In fact, while I was looking for a way to execute in static mode, I discovered one interesting low-level detail about models built through the Keras API…

One More Thing…

Remember when I said TensorFlow runs in eager mode by default, even proving it with a code snippet? Well, I lied! Kind of.

If you use Keras API to build and manage your models, then it will attempt to compile them as static graphs under the hood. So what you end up getting is the performance of static computational graphs with the flexibility of eager execution.

You can check the status of your model via the model.run_eagerly flag. You can also force eager mode by setting this flag to True, though most of the times you probably don’t need to - if Keras detects that there’s no way around eager mode, it will back off on its own.

To illustrate that it’s indeed running as a static graph here’s a simple benchmark:

# create a 100000 samples batch
env = gym.make('CartPole-v0')
obs = np.repeat(env.reset()[None, :], 100000, axis=0)

Eager Benchmark

%%time
model = Model(env.action_space.n)
model.run_eagerly = True
print("Eager Execution:  ", tf.executing_eagerly())
print("Eager Keras Model:", model.run_eagerly)
_ = model(obs)
######## Results #######
Eager Execution:   True
Eager Keras Model: True
CPU times: user 639 ms, sys: 736 ms, total: 1.38 s

Static Benchmark

%%time
with tf.Graph().as_default():
model = Model(env.action_space.n)
print("Eager Execution:  ", tf.executing_eagerly())
print("Eager Keras Model:", model.run_eagerly)
_ = model.predict(obs)
######## Results #######
Eager Execution:   False
Eager Keras Model: False
CPU times: user 793 ms, sys: 79.7 ms, total: 873 ms

Default Benchmark

%%time
model = Model(env.action_space.n)
print("Eager Execution:  ", tf.executing_eagerly())
print("Eager Keras Model:", model.run_eagerly)
_ = model.predict(obs)
######## Results #######
Eager Execution:   True
Eager Keras Model: False
CPU times: user 994 ms, sys: 23.1 ms, total: 1.02 s

As you can see eager mode is behind static mode, and by default, our model was indeed executing statically, more or less matching explicit static graph execution.

Conclusion

Hopefully, this has been an illustrative tour of both DRL and the things to come in TensorFlow 2.0. Note that this is still just a nightly preview build, not even a release candidate. Everything is subject to change and if there’s something about TensorFlow you especially dislike (or like :) ), let the developers know!

A lingering question people might have is if TensorFlow is better than PyTorch? Maybe. Maybe not. Both are great libraries, so it is hard to say one way or the other. If you’re familiar with PyTorch, you probably noticed that TensorFlow 2.0 not only caught up but also avoided some of the PyTorch API pitfalls.

In either case, what is clear is that this competition has resulted in a net positive outcome for both camps and I am excited to see what will become of the frameworks in the future.


Originally published at inoryy.com on January 20, 2019.