TF-Agents 0.4 : Tutorials : SAC minitaur

Posted on by Sales Information

TF-Agents 0.4 Tutorials : SAC minitaur (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 04/21/2020 (0.4)

* 本ページは、TF Agents の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

SAC minitaur

イントロダクション

このサンプルは TF-Agents ライブラリを使用して Minitaur 環境上 Soft Actor Critic エージェントをどのように訓練するかを示します。

DQN Colab を通して作業したのであれば、これは非常に馴染みやすいはずです。顕著な変更は以下を含みます :

  • エージェントを DQN から SAC に変更する。
  • Minitaur 上で訓練します、これはカートポールよりも遥かに複雑な環境です。Minitaur 環境は四足ロボットを前に進むように訓練することを目的とします。
  • 初期データ収集を遂行するためにランダムポリシーは使用しません。

以下の依存性をまだインストールしてないのであれば、以下を実行します :

!sudo apt-get install -y xvfb ffmpeg
!pip install 'gym==0.10.11'
!pip install 'imageio==2.4.0'
!pip install matplotlib
!pip install PILLOW
!pip install --upgrade tensorflow-probability
!pip install tf-agents
!pip install 'pybullet==2.4.2'

 

セットアップ

最初に必要な各種ツールをインポートして TF-V2 動作を有効にします、何故ならば colab を通して Eager モードで反復することは容易だからです。

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

import base64
import imageio
import IPython
import matplotlib
import matplotlib.pyplot as plt
import PIL.Image

import tensorflow as tf
tf.compat.v1.enable_v2_behavior()

from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import sac_agent
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_pybullet
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import normal_projection_network
from tf_agents.policies import greedy_policy
from tf_agents.policies import random_tf_policy
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.trajectories import trajectory
from tf_agents.utils import common

 

ハイパーパラメータ

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

# use "num_iterations = 1e6" for better results,
# 1e5 is just so this doesn't take too long. 
num_iterations = 100000 # @param {type:"integer"}

initial_collect_steps = 10000 # @param {type:"integer"} 
collect_steps_per_iteration = 1 # @param {type:"integer"}
replay_buffer_capacity = 1000000 # @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"}
gradient_clipping = None # @param

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

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

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

 

環境

RL の環境は解こうとしているタスクや問題を表します。標準的な環境はスーツを使用して TF-Agents で容易に作成できます。文字列環境名が与えられたとき、OpenAI Gym, Atari, DM Control 等のようなソースから環境をロードするために様々なスーツを持ちます。

今は Pybullet スーツから Minituar 環境をロードしましょう。

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

この環境では目標はエージェントが Minitaur ロボットを制御するポリシーを訓練してできる限り速く前に進ませることです。エピソードは 1000 ステップ続いてリターンはエピソードを通した報酬の総計です。環境が観測として提供する情報を見ましょう、これをポリシーはアクションを生成するために利用します。

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

見れるように観測は非常に複雑です。総てのモータのための角度、速度とトルクを表す 28 値を受け取ります。返しでは環境は [-1, 1] の間のアクションのための 8 値を想定します。これらは望まれるモーター角度です。

通常は 2 つの環境を作成します : 一つは訓練のためそして 1 つは評価のためです。殆どの環境は純粋は python で書かれますが、それらは TFPyEnvironment ラッパーを使用して容易に TensorFlow に変換できます。元の環境の API は numpy 配列を使用し、TFPyEnvironment はこれらを TensorFlow ポリシーとエージェントとより容易に相互作用するために貴方のために Tensor に/から変換します。

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

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

 

エージェント

SAC エージェントを作成するため、それが訓練するネットワークを最初に作成する必要があります。SAC は actor-critic エージェントですので、2 つのネットワークを必要とします。

critic (批評家) は Q(s, a) のための値推定を与えます。つまり、それは入力として観測とアクションを受け取り、そしてそれは与えられた状態のためにアクションがどのくらい良いかの推定を与えます。

observation_spec = train_env.observation_spec()
action_spec = train_env.action_spec()
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)

観測が与えられたときアクションを生成することを可能にする actor ネットワークを訓練するためにこの critic を利用します。

ActorNetwork は正規分布のためのパラメータを予測します。それからアクションを生成する必要があるとき、この分布はサンプリングされ、現在の観測で条件付けられます。

def normal_projection_net(action_spec,init_means_output_factor=0.1):
  return normal_projection_network.NormalProjectionNetwork(
      action_spec,
      mean_transform=None,
      state_dependent_std=True,
      init_means_output_factor=init_means_output_factor,
      std_transform=sac_agent.std_clip_transform,
      scale_distribution=True)


actor_net = actor_distribution_network.ActorDistributionNetwork(
    observation_spec,
    action_spec,
    fc_layer_params=actor_fc_layer_params,
    continuous_projection_net=normal_projection_net)

手元のこれらのネットワークで今ではエージェントをインスタンス化できます。

global_step = tf.compat.v1.train.get_or_create_global_step()
tf_agent = sac_agent.SacAgent(
    train_env.time_step_spec(),
    action_spec,
    actor_network=actor_net,
    critic_network=critic_net,
    actor_optimizer=tf.compat.v1.train.AdamOptimizer(
        learning_rate=actor_learning_rate),
    critic_optimizer=tf.compat.v1.train.AdamOptimizer(
        learning_rate=critic_learning_rate),
    alpha_optimizer=tf.compat.v1.train.AdamOptimizer(
        learning_rate=alpha_learning_rate),
    target_update_tau=target_update_tau,
    target_update_period=target_update_period,
    td_errors_loss_fn=tf.compat.v1.losses.mean_squared_error,
    gamma=gamma,
    reward_scale_factor=reward_scale_factor,
    gradient_clipping=gradient_clipping,
    train_step_counter=global_step)
tf_agent.initialize()

 

ポリシー

TF-Agents では、ポリシーは RL のポリシーの標準的な概念を表します : time_step が与えられたときアクションかアクションに渡る分布を生成します。主要メソッドは policy_step = policy.step(time_step) です、そこでは policy_step は名前付きタプル PolicyStep(action, state, info) です。policy_step.action は環境に適用されるべきアクションで、state はステートフル (RNN) ポリシーのための状態を表して info はアクションの対数確率のような補助的情報を含むかもしれません。

エージェントは 2 つのポリシーを含みます : 主要ポリシー (agent.policy) とデータ収集のために使用される行動 (= behavioral) ポリシー (agent.collect_policy) です。評価/配備については、主要ポリシーを GreedyPolicy() でラップすることにより平均 (= mean) アクションを取ります。

eval_policy = greedy_policy.GreedyPolicy(tf_agent.policy)
collect_policy = tf_agent.collect_policy

 

メトリクスと評価

ポリシーを評価するために使用される最も一般的なメトリックは平均リターンです。リターンはエピソードのための環境でポリシーを実行する間に得られた報酬の総計で、通常はこれを幾つかのエピソードに渡り平均します。次のように平均リターン・メトリックを計算できます。

def compute_avg_return(environment, policy, num_episodes=5):

  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]


compute_avg_return(eval_env, eval_policy, num_eval_episodes)

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

 

再生バッファ

環境から収集されたデータを追跡するために、TFUniformReplayBuffer を利用します。この再生バッファはストアされる tensor を記述する specs を使用して構築されます、これは tf_agent.collect_data_spec を使用してエージェントから得られます。

replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec=tf_agent.collect_data_spec,
    batch_size=train_env.batch_size,
    max_length=replay_buffer_capacity)

殆どのエージェントについて、collect_data_spec は Trajectory と名前付けられたタプルで観測、アクション、報酬等を含みます。

 

データ収集

今はそれで再生バッファを seed するための経験を収集するためにドライバーを作成します。ドライバーは指定されたポリシーを使用して環境上の n ステップかエピソードを収集する単純な方法を提供します。

initial_collect_driver = dynamic_step_driver.DynamicStepDriver(
        train_env,
        collect_policy,
        observers=[replay_buffer.add_batch],
        num_steps=initial_collect_steps)
initial_collect_driver.run()

再生バッファからデータをサンプリングするため、tf.data パイプラインを作成します、これを後で訓練のためにエージェントに供給できます。再生バッファからサンプリングされた項目の数を設定するために sample_batch_size を指定できます。並列呼び出しと先取りを使用してデータパイプラインを最適化することもできます。

スペースを節約するため、再生バッファの各行で現在の観測をストアするだけです。しかし SAC エージェントは損失を計算するために現在と次の観測の両者を必要としますので、num_steps=2 を設定することによりバッチの各項目に対して 2 つの隣接する行を常にサンプリングします。

# Dataset generates trajectories with shape [Bx2x...]
dataset = replay_buffer.as_dataset(
    num_parallel_calls=3, sample_batch_size=batch_size, num_steps=2).prefetch(3)

iterator = iter(dataset)

 

エージェントを訓練する

訓練ループは環境からのデータ収集とエージェントのネットワークの最適化の両者を含みます。道中、どのようにやっているかを見るためにエージェントのポリシーを時々評価します。

collect_driver = dynamic_step_driver.DynamicStepDriver(
    train_env,
    collect_policy,
    observers=[replay_buffer.add_batch],
    num_steps=collect_steps_per_iteration)

#@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)
collect_driver.run = common.function(collect_driver.run)

# 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, eval_policy, num_eval_episodes)
returns = [avg_return]

for _ in range(num_iterations):

  # Collect a few steps using collect_policy and save to the replay buffer.
  for _ in range(collect_steps_per_iteration):
    collect_driver.run()

  # Sample a batch of data from the buffer and update the agent's network.
  experience, unused_info = next(iterator)
  train_loss = tf_agent.train(experience)

  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, eval_policy, num_eval_episodes)
    print('step = {0}: Average Return = {1}'.format(step, avg_return))
    returns.append(avg_return)

 

可視化

プロット

エージェントのパフォーマンスを見るために平均リターン vs グルーバルステップをプロットできます。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 = '''
  '''.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_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)
 

以上