具有 Actor-Learner API 的 SAC Minitaur

Copyright 2023 The TF-Agents Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

具有 Actor-Learner API 的 SAC Minitaur#

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

简介#

本例介绍如何在 Minitaur 环境中训练 Soft Actor Critic 代理。

如果您使用过 DQN Colab，应该会对这个环境很熟悉。一些明显的变化包括：

  • 将代理从 DQN 改成了 SAC。

  • 在 Minitaur 中训练，与 CartPole 相比，此环境要复杂得多。Minitaur 环境旨在训练一个四足机器人向前移动。

  • 使用 TF-Agents Actor-Learner API 进行分布式强化学习。

该 API 既支持使用经验回放缓冲区和可变容器（参数服务器）的分布式数据收集，也支持跨多个设备的分布训练，其设计非常简单，并且是模块化的。对于回放缓冲区和可变容器，我们都利用 Reverb，而对于 GPU 和 TPU 上的分布训练，我们都利用 TF DistributionStrategy API。

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

!sudo apt-get update
!sudo apt-get install -y xvfb ffmpeg
!pip install 'imageio==2.4.0'
!pip install matplotlib
!pip install tf-agents[reverb]
!pip install pybullet

设置#

首先，我们将导入所需的不同工具。

import base64
import imageio
import IPython
import matplotlib.pyplot as plt
import os
import reverb
import tempfile
import PIL.Image

import tensorflow as tf

from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import sac_agent
from tf_agents.agents.sac import tanh_normal_projection_network
from tf_agents.environments import suite_pybullet
from tf_agents.metrics import py_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import greedy_policy
from tf_agents.policies import py_tf_eager_policy
from tf_agents.policies import random_py_policy
from tf_agents.replay_buffers import reverb_replay_buffer
from tf_agents.replay_buffers import reverb_utils
from tf_agents.train import actor
from tf_agents.train import learner
from tf_agents.train import triggers
from tf_agents.train.utils import spec_utils
from tf_agents.train.utils import strategy_utils
from tf_agents.train.utils import train_utils

tempdir = tempfile.gettempdir()

超参数#

env_name = "MinitaurBulletEnv-v0" # @param {type:"string"}

# Use "num_iterations = 1e6" for better results (2 hrs)
# 1e5 is just so this doesn't take too long (1 hr)
num_iterations = 100000 # @param {type:"integer"}

initial_collect_steps = 10000 # @param {type:"integer"}
collect_steps_per_iteration = 1 # @param {type:"integer"}
replay_buffer_capacity = 10000 # @param {type:"integer"}

batch_size = 256 # @param {type:"integer"}

critic_learning_rate = 3e-4 # @param {type:"number"}
actor_learning_rate = 3e-4 # @param {type:"number"}
alpha_learning_rate = 3e-4 # @param {type:"number"}
target_update_tau = 0.005 # @param {type:"number"}
target_update_period = 1 # @param {type:"number"}
gamma = 0.99 # @param {type:"number"}
reward_scale_factor = 1.0 # @param {type:"number"}

actor_fc_layer_params = (256, 256)
critic_joint_fc_layer_params = (256, 256)

log_interval = 5000 # @param {type:"integer"}

num_eval_episodes = 20 # @param {type:"integer"}
eval_interval = 10000 # @param {type:"integer"}

policy_save_interval = 5000 # @param {type:"integer"}

环境#

在强化学习 (RL) 中，环境代表要解决的任务或问题。在 TF-Agents 中，使用 suites 可以轻松创建标准环境。我们提供了不同的 suites，只需提供一个字符串环境名称，即可从 OpenAI Gym、Atari、DM Control 等来源加载环境。

现在，我们从 Pybullet 套件加载 Minituar 环境。

env = suite_pybullet.load(env_name)
env.reset()
PIL.Image.fromarray(env.render())

在该环境中，代理的目标是训练一个控制 Minitaur 机器人的策略，让机器人以尽可能快的速度向前移动。片段会持续 1000 个步骤，回报是整个片段的奖励总和。

我们看看该环境提供的 observation 信息，该策略将使用这些信息生成 actions。

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

该观测值非常复杂。我们收到了 28 个值，分别代表所有电机的角度、速度和扭矩。作为回应，该环境希望获得操作介于 [-1, 1] 之间的 8 个值。这些值是需要的电机角度。

通常，我们创建两个环境：一个用于在训练过程中收集数据，另一个用于评估。这些环境使用纯 Python 语言编写，并且使用 NumPy 数组（由 Actor Learner API 直接使用）。

collect_env = suite_pybullet.load(env_name)
eval_env = suite_pybullet.load(env_name)

分布策略#

我们使用 DistributionStrategy API 来支持跨多个设备（如使用数据并行的多个 GPU 或 TPU）运行训练步骤计算。该训练步骤：

  • 接收一批训练数据

  • 将数据分发给设备

  • 计算前向步骤

  • 聚合并计算损失平均值

  • 计算后向步骤并执行梯度变量更新

使用 TF-Agents Learner API 和 DistributionStrategy API，从在 GPU 上运行训练步骤（使用 MirroredStrategy）切换到 TPU（使用 TPUStrategy）非常容易，无需更改下面的任何训练逻辑。

启用 GPU#

如果您希望在 GPU 上运行，首先需要为笔记本启用 GPU：

  • 导航至 Edit→Notebook Settings

  • 从 Hardware Accelerator 下拉列表中选择 GPU

选择策略#

使用 strategy_utils 生成策略。在后台传递参数：

  • use_gpu = False 返回 tf.distribute.get_strategy()，该策略使用 CPU

  • use_gpu = True 返回 tf.distribute.MirroredStrategy()，该策略使用一台计算机上 TensorFlow 可见的所有 GPU。

use_gpu = True #@param {type:"boolean"}

strategy = strategy_utils.get_strategy(tpu=False, use_gpu=use_gpu)

您需要在 strategy.scope() 下创建所有变量和代理，如下所示。

代理#

要创建 SAC 代理，首先要创建通过该代理训练的网络。SAC 是一个 Actor-Critic 代理，所以我们需要两个网络。

Critic 会提供 Q(s,a) 的值估算结果。也就是说，它会接收一个观测值和一个操作作为输入，同时提供该操作在指定状态下的表现的估算结果。

observation_spec, action_spec, time_step_spec = (
      spec_utils.get_tensor_specs(collect_env))

with strategy.scope():
  critic_net = critic_network.CriticNetwork(
        (observation_spec, action_spec),
        observation_fc_layer_params=None,
        action_fc_layer_params=None,
        joint_fc_layer_params=critic_joint_fc_layer_params,
        kernel_initializer='glorot_uniform',
        last_kernel_initializer='glorot_uniform')

我们将使用该 Critic 来训练 actor 网络，通过该网络，我们可以利用提供的观测值生成操作。

ActorNetwork 将预测 tanh-squashed MultivariateNormalDiag 分布的参数。随后，只要我们需要生成操作，就可以对该分布进行采样，根据当前观测值建立条件。

with strategy.scope():
  actor_net = actor_distribution_network.ActorDistributionNetwork(
      observation_spec,
      action_spec,
      fc_layer_params=actor_fc_layer_params,
      continuous_projection_net=(
          tanh_normal_projection_network.TanhNormalProjectionNetwork))

凭借这些已有的网络，我们现在可以实例化代理。

with strategy.scope():
  train_step = train_utils.create_train_step()

  tf_agent = sac_agent.SacAgent(
        time_step_spec,
        action_spec,
        actor_network=actor_net,
        critic_network=critic_net,
        actor_optimizer=tf.keras.optimizers.Adam(
            learning_rate=actor_learning_rate),
        critic_optimizer=tf.keras.optimizers.Adam(
            learning_rate=critic_learning_rate),
        alpha_optimizer=tf.keras.optimizers.Adam(
            learning_rate=alpha_learning_rate),
        target_update_tau=target_update_tau,
        target_update_period=target_update_period,
        td_errors_loss_fn=tf.math.squared_difference,
        gamma=gamma,
        reward_scale_factor=reward_scale_factor,
        train_step_counter=train_step)

  tf_agent.initialize()

回放缓冲区#

为了跟踪从环境收集的数据，我们将使用 Reverb——Deepmind 出品的一款高效、可扩展且易于使用的回放系统。它会存储 Actor 收集的经验数据，供 Learner 在训练时使用。

在本教程中，其重要性不如 max_size。但是，在使用异步收集和训练的分布设置中，您可能希望使用位置在 2 到 1000 之间的 samples_per_insert 来尝试执行 rate_limiters.SampleToInsertRatio。例如：

rate_limiter=reverb.rate_limiters.SampleToInsertRatio(samples_per_insert=3.0, min_size_to_sample=3, error_buffer=3.0)

table_name = 'uniform_table'
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))

reverb_server = reverb.Server([table])

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

由于 SAC 代理同时需要当前和下一个观察值才能计算损失，因此，我们设置 sequence_length=2。

reverb_replay = reverb_replay_buffer.ReverbReplayBuffer(
    tf_agent.collect_data_spec,
    sequence_length=2,
    table_name=table_name,
    local_server=reverb_server)

现在，我们创建一个驱动器来积累经验，利用这些经验设置回放缓冲区的种子。驱动器提供了一种简单的方法，让我们可以使用特定策略在环境中收集 n 个步骤或片段的数据。

dataset = reverb_replay.as_dataset(
      sample_batch_size=batch_size, num_steps=2).prefetch(50)
experience_dataset_fn = lambda: dataset

策略#

在 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 — 用于数据收集的第二策略。

tf_eval_policy = tf_agent.policy
eval_policy = py_tf_eager_policy.PyTFEagerPolicy(
  tf_eval_policy, use_tf_function=True)

tf_collect_policy = tf_agent.collect_policy
collect_policy = py_tf_eager_policy.PyTFEagerPolicy(
  tf_collect_policy, use_tf_function=True)

策略可以独立于代理进行创建。例如，使用 tf_agents.policies.random_py_policy 创建策略，将为每个 time_step 随机选择一项操作。

random_policy = random_py_policy.RandomPyPolicy(
  collect_env.time_step_spec(), collect_env.action_spec())

Actor#

Actor 用于管理策略与环境之间的交互。

  • Actor 组件包含环境的一个实例（ 作为 py_environment）和策略变量的一个副本。

  • 给定策略变量的本地值，每个 Actor 工作进程运行一系列数据收集步骤。

  • 在调用 actor.run() 之前，使用训练脚本中的可变容器客户端实例明确完成变量更新。

  • 在每个数据收集步骤中将观察到的经验写入回放缓冲区。

当 Actor 运行数据收集步骤时，它们会将（状态、操作、奖励）的轨迹传递给观察器，而观察器将缓存轨迹并将其写入 Reverb 回放系统。

由于 stride_length=1，因此，我们存储框架 [(t0,t1) (t1,t2) (t2,t3), …]。

rb_observer = reverb_utils.ReverbAddTrajectoryObserver(
  reverb_replay.py_client,
  table_name,
  sequence_length=2,
  stride_length=1)

我们使用随机策略创建一个 Actor，并收集经验以设置回放缓冲区的种子。

initial_collect_actor = actor.Actor(
  collect_env,
  random_policy,
  train_step,
  steps_per_run=initial_collect_steps,
  observers=[rb_observer])
initial_collect_actor.run()

在训练过程中使用收集策略实例化 Actor 以收集更多经验。

env_step_metric = py_metrics.EnvironmentSteps()
collect_actor = actor.Actor(
  collect_env,
  collect_policy,
  train_step,
  steps_per_run=1,
  metrics=actor.collect_metrics(10),
  summary_dir=os.path.join(tempdir, learner.TRAIN_DIR),
  observers=[rb_observer, env_step_metric])

创建一个可用于在训练过程中评估策略的 Actor。我们传入 actor.eval_metrics(num_eval_episodes)，以便随后记录指标。

eval_actor = actor.Actor(
  eval_env,
  eval_policy,
  train_step,
  episodes_per_run=num_eval_episodes,
  metrics=actor.eval_metrics(num_eval_episodes),
  summary_dir=os.path.join(tempdir, 'eval'),
)

Learner#

Learner 组件包含代理，并使用回放缓冲区的经验数据执行策略变量的梯度步骤更新。在经过一个或多个训练步骤后，Learner 可以将一组新的变量值推送到可变容器。

saved_model_dir = os.path.join(tempdir, learner.POLICY_SAVED_MODEL_DIR)

# Triggers to save the agent's policy checkpoints.
learning_triggers = [
    triggers.PolicySavedModelTrigger(
        saved_model_dir,
        tf_agent,
        train_step,
        interval=policy_save_interval),
    triggers.StepPerSecondLogTrigger(train_step, interval=1000),
]

agent_learner = learner.Learner(
  tempdir,
  train_step,
  tf_agent,
  experience_dataset_fn,
  triggers=learning_triggers,
  strategy=strategy)

指标和评估#

我们使用上面的 actor.eval_metrics 实例化评价 Actor，这会创建在策略评估期间最常用的指标：

  • 平均回报。回报是在某一片段的环境中运行策略时获得的回报总和，我们通常会求几个片段的平均值。

  • 平均片段长度。

我们运行 Actor 以生成这些指标。

def get_eval_metrics():
  eval_actor.run()
  results = {}
  for metric in eval_actor.metrics:
    results[metric.name] = metric.result()
  return results

metrics = get_eval_metrics()

def log_eval_metrics(step, metrics):
  eval_results = (', ').join(
      '{} = {:.6f}'.format(name, result) for name, result in metrics.items())
  print('step = {0}: {1}'.format(step, eval_results))

log_eval_metrics(0, metrics)

有关不同指标的其他标准实现，请查看指标模块。

训练代理#

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

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

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

# Evaluate the agent's policy once before training.
avg_return = get_eval_metrics()["AverageReturn"]
returns = [avg_return]

for _ in range(num_iterations):
  # Training.
  collect_actor.run()
  loss_info = agent_learner.run(iterations=1)

  # Evaluating.
  step = agent_learner.train_step_numpy

  if eval_interval and step % eval_interval == 0:
    metrics = get_eval_metrics()
    log_eval_metrics(step, metrics)
    returns.append(metrics["AverageReturn"])

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

rb_observer.close()
reverb_server.stop()

可视化#

绘图#

我们可以通过绘制回报与全局步骤的图形来了解代理的性能。在 Minitaur 中，奖励函数基于 Minitaur 在 1000 个步骤中行走了多长的距离，并扣除能量消耗。

#@test {"skip": true}

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

视频#

渲染每个步骤的环境有助于可视化代理的性能。在此之前，我们先创建一个函数，在该 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 = 'sac_minitaur.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_env.render())
    while not time_step.is_last():
      action_step = eval_actor.policy.action(time_step)
      time_step = eval_env.step(action_step.action)
      video.append_data(eval_env.render())

embed_mp4(video_filename)