Copyright 2023 The TF-Agents Authors.

REINFORCE 代理#

在 TensorFlow.org 上查看 在 Google Colab 运行 在 Github 上查看源代码 下载笔记本

简介#

本例介绍如何使用 TF-Agents 库在 Cartpole 环境中训练 REINFORCE 代理，与 DQN 教程比较相似。

Cartpole environment

我们会引导您完成强化学习 (RL) 流水线中关于训练、评估和数据收集的所有部分。

设置#

如果尚未安装以下依赖项，请运行以下命令：

!sudo apt-get update
!sudo apt-get install -y xvfb ffmpeg freeglut3-dev
!pip install 'imageio==2.4.0'
!pip install pyvirtualdisplay
!pip install tf-agents[reverb]
!pip install pyglet xvfbwrapper

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import base64
import imageio
import IPython
import matplotlib.pyplot as plt
import numpy as np
import PIL.Image
import pyvirtualdisplay
import reverb

import tensorflow as tf

from tf_agents.agents.reinforce import reinforce_agent
from tf_agents.drivers import py_driver
from tf_agents.environments import suite_gym
from tf_agents.environments import tf_py_environment
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import py_tf_eager_policy
from tf_agents.replay_buffers import reverb_replay_buffer
from tf_agents.replay_buffers import reverb_utils
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import trajectory
from tf_agents.utils import common

# Set up a virtual display for rendering OpenAI gym environments.
display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()

超参数#

env_name = "CartPole-v0" # @param {type:"string"}
num_iterations = 250 # @param {type:"integer"}
collect_episodes_per_iteration = 2 # @param {type:"integer"}
replay_buffer_capacity = 2000 # @param {type:"integer"}

fc_layer_params = (100,)

learning_rate = 1e-3 # @param {type:"number"}
log_interval = 25 # @param {type:"integer"}
num_eval_episodes = 10 # @param {type:"integer"}
eval_interval = 50 # @param {type:"integer"}

环境#

RL 环境用于描述要解决的任务或问题。在 TF-Agents 中，使用 suites 可以轻松创建标准环境。我们提供了不同的 suites，只需提供一个字符串环境名称，即可帮助您从来源加载环境，如 OpenAI Gym、Atari、DM Control 等。

现在，我们试试从 OpenAI Gym 套件加载 CartPole 环境。

env = suite_gym.load(env_name)

我们可以渲染此环境以查看其形式：小车上连接一条自由摆动的长杆。目标是向右或向左移动小车，使长杆保持朝上。

#@test {"skip": true}
env.reset()
PIL.Image.fromarray(env.render())

在该环境中，time_step = environment.step(action) 语句用于执行 action。返回的 TimeStep 元组包含该操作在环境中的下一个观测值和奖励。环境中的 time_step_spec() 和 action_spec() 方法分别返回 time_step 和 action 的规范（类型、形状、边界）。

print('Observation Spec:')
print(env.time_step_spec().observation)
print('Action Spec:')
print(env.action_spec())

我们可以看到，该观测值是一个包含 4 个浮点数的数组：小车的位置和速度，长杆的角度位置和速度。由于只有两个操作（向左或向右移动），因此，action_spec 是一个标量，其中 0 表示“向左移动”，1 表示“向右移动”。

time_step = env.reset()
print('Time step:')
print(time_step)

action = np.array(1, dtype=np.int32)

next_time_step = env.step(action)
print('Next time step:')
print(next_time_step)

通常，我们会创建两个环境：一个用于训练，另一个用于评估。大部分环境都是使用纯 Python 语言编写的，但是使用 TFPyEnvironment 包装器可轻松将其转换至 TensorFlow 环境。原始环境的 API 使用 NumPy 数组，但凭借 TFPyEnvironment，这些数组可以与 Tensors 相互转换，从而更轻松地与 TensorFlow 策略和代理交互。

train_py_env = suite_gym.load(env_name)
eval_py_env = suite_gym.load(env_name)

train_env = tf_py_environment.TFPyEnvironment(train_py_env)
eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)

代理#

我们用于解决 RL 问题的算法以 Agent 形式表示。除了 REINFORCE 代理，TF-Agents 还为各种 Agents 提供了标准实现，如 DQN、DDPG、TD3、PPO 和 SAC。

要创建 REINFORCE 代理，首先需要有一个通过环境提供的观测值，学会预测操作的 Actor Network。

使用观测值和操作的规范，我们可以轻松创建 Actor Network。我们也可以在网络中指定层，本例中是设置为 ints 元祖（表示每个隐藏层的大小）的 fc_layer_params 参数（请参阅上面的“超参数”部分）。

actor_net = actor_distribution_network.ActorDistributionNetwork(
    train_env.observation_spec(),
    train_env.action_spec(),
    fc_layer_params=fc_layer_params)

我们还需要一个 optimizer 来训练刚才创建的网络，以及一个跟踪网络更新次数的 train_step_counter 变量。

optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

train_step_counter = tf.Variable(0)

tf_agent = reinforce_agent.ReinforceAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    actor_network=actor_net,
    optimizer=optimizer,
    normalize_returns=True,
    train_step_counter=train_step_counter)
tf_agent.initialize()

策略#

在 TF-Agents 中，策略是 RL 中的标准策略概念：给订 time_step 来产生操作或操作的分布。主要方法是 policy_step = policy.step(time_step)，其中 policy_step 是命名元祖 PolicyStep(action, state, info)。policy_step.action 是要应用到环境的 action，state 表示有状态 (RNN) 策略的状态，而 info 可能包含辅助信息（如操作的对数几率）。

代理包含两个策略：一个是用于评估/部署的主要策略 (agent.policy)，另一个是用于数据收集的策略 (agent.collect_policy)。

eval_policy = tf_agent.policy
collect_policy = tf_agent.collect_policy

指标和评估#

用于评估策略的最常用指标是平均回报。回报就是在环境中运行策略时，某个片段获得的奖励总和，我们通常会计算几个片段的平均值。计算平均回报指标的代码如下。

#@test {"skip": true}
def compute_avg_return(environment, policy, num_episodes=10):

  total_return = 0.0
  for _ in range(num_episodes):

    time_step = environment.reset()
    episode_return = 0.0

    while not time_step.is_last():
      action_step = policy.action(time_step)
      time_step = environment.step(action_step.action)
      episode_return += time_step.reward
    total_return += episode_return

  avg_return = total_return / num_episodes
  return avg_return.numpy()[0]


# Please also see the metrics module for standard implementations of different
# metrics.

回放缓冲区#

为了跟踪从环境收集的数据，我们将使用 Reverb，这是 Deepmind 打造的一款高效、可扩展且易于使用的回放系统。它会在我们收集轨迹时存储经验数据，并在训练期间使用。

此回放缓冲区使用描述要存储的张量的规范进行构造，可以使用 tf_agent.collect_data_spec 从代理获取这些张量。

table_name = 'uniform_table'
replay_buffer_signature = tensor_spec.from_spec(
      tf_agent.collect_data_spec)
replay_buffer_signature = tensor_spec.add_outer_dim(
      replay_buffer_signature)
table = reverb.Table(
    table_name,
    max_size=replay_buffer_capacity,
    sampler=reverb.selectors.Uniform(),
    remover=reverb.selectors.Fifo(),
    rate_limiter=reverb.rate_limiters.MinSize(1),
    signature=replay_buffer_signature)

reverb_server = reverb.Server([table])

replay_buffer = reverb_replay_buffer.ReverbReplayBuffer(
    tf_agent.collect_data_spec,
    table_name=table_name,
    sequence_length=None,
    local_server=reverb_server)

rb_observer = reverb_utils.ReverbAddEpisodeObserver(
    replay_buffer.py_client,
    table_name,
    replay_buffer_capacity
)

对于大多数代理，collect_data_spec 是一个 Trajectory 命名元组，其中包含观测值、操作和奖励等。

数据收集#

当 REINFORCE 从全部片段中学习时，我们使用给定数据收集策略定义一个函数来收集片段，并在回放缓冲区中将数据（观测值、操作、奖励等）保存为轨迹。这里我们使用“PyDriver”运行经验收集循环。您可以在我们的 driver 教程中了解到有关 TF Agents driver 的更多信息。

#@test {"skip": true}

def collect_episode(environment, policy, num_episodes):

  driver = py_driver.PyDriver(
    environment,
    py_tf_eager_policy.PyTFEagerPolicy(
      policy, use_tf_function=True),
    [rb_observer],
    max_episodes=num_episodes)
  initial_time_step = environment.reset()
  driver.run(initial_time_step)

训练代理#

训练循环包括从环境收集数据和优化代理的网络。在训练过程中，我们偶尔会评估代理的策略，看看效果如何。

运行下面的代码大约需要 3 分钟。

#@test {"skip": true}
try:
  %%time
except:
  pass

# (Optional) Optimize by wrapping some of the code in a graph using TF function.
tf_agent.train = common.function(tf_agent.train)

# Reset the train step
tf_agent.train_step_counter.assign(0)

# Evaluate the agent's policy once before training.
avg_return = compute_avg_return(eval_env, tf_agent.policy, num_eval_episodes)
returns = [avg_return]

for _ in range(num_iterations):

  # Collect a few episodes using collect_policy and save to the replay buffer.
  collect_episode(
      train_py_env, tf_agent.collect_policy, collect_episodes_per_iteration)

  # Use data from the buffer and update the agent's network.
  iterator = iter(replay_buffer.as_dataset(sample_batch_size=1))
  trajectories, _ = next(iterator)
  train_loss = tf_agent.train(experience=trajectories)  

  replay_buffer.clear()

  step = tf_agent.train_step_counter.numpy()

  if step % log_interval == 0:
    print('step = {0}: loss = {1}'.format(step, train_loss.loss))

  if step % eval_interval == 0:
    avg_return = compute_avg_return(eval_env, tf_agent.policy, num_eval_episodes)
    print('step = {0}: Average Return = {1}'.format(step, avg_return))
    returns.append(avg_return)

可视化#

绘图#

我们可以通过绘制回报与全局步骤的图形来了解代理的性能。在 Cartpole-v0 中，长杆每停留一个时间步骤，环境就会提供一个 +1 的奖励，由于最大步骤数量为 200，所以可以获得的最大回报也是 200。

#@test {"skip": true}

steps = range(0, num_iterations + 1, eval_interval)
plt.plot(steps, returns)
plt.ylabel('Average Return')
plt.xlabel('Step')
plt.ylim(top=250)

视频#

在每个步骤渲染环境有助于可视化代理的性能。在此之前，我们先创建一个函数，在该 Colab 中嵌入视频。

def embed_mp4(filename):
  """Embeds an mp4 file in the notebook."""
  video = open(filename,'rb').read()
  b64 = base64.b64encode(video)
  tag = '''
  <video width="640" height="480" controls>
    <source src="data:video/mp4;base64,{0}" type="video/mp4">
  Your browser does not support the video tag.
  </video>'''.format(b64.decode())

  return IPython.display.HTML(tag)

以下代码用于为代理可视化几个片段的策略：

num_episodes = 3
video_filename = 'imageio.mp4'
with imageio.get_writer(video_filename, fps=60) as video:
  for _ in range(num_episodes):
    time_step = eval_env.reset()
    video.append_data(eval_py_env.render())
    while not time_step.is_last():
      action_step = tf_agent.policy.action(time_step)
      time_step = eval_env.step(action_step.action)
      video.append_data(eval_py_env.render())

embed_mp4(video_filename)