TF-Agents 0.4 Tutorials : DQN C51/Rainbow (翻訳/解説)

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

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

DQN C51/Rainbow

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

DQN C51/Rainbow

イントロダクション

このサンプルは TF-Agents ライブラリでカートポール環境上カテゴリカル DQN (C51) エージェントをどのように訓練するかを示します。

必要条件として DQN チュートリアルに目を通していることを確実にしてください。このチュートリアルは DQN チュートリアルに精通していることを仮定しています ; それは DQN と C51 の間の違いに主として焦点を当てます。

セットアップ

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

!sudo apt-get install -y xvfb ffmpeg
!pip install 'gym==0.10.11'
!pip install 'imageio==2.4.0'
!pip install PILLOW
!pip install 'pyglet==1.3.2'
!pip install pyvirtualdisplay
!pip install --upgrade tensorflow-probability
!pip install tf-agents

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 pyvirtualdisplay

import tensorflow as tf

from tf_agents.agents.categorical_dqn import categorical_dqn_agent
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_gym
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 categorical_q_network
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

tf.compat.v1.enable_v2_behavior()


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

ハイパーパラメータ

env_name = "CartPole-v1" # @param {type:"string"}
num_iterations = 15000 # @param {type:"integer"}

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

fc_layer_params = (100,)

batch_size = 64  # @param {type:"integer"}
learning_rate = 1e-3  # @param {type:"number"}
gamma = 0.99
log_interval = 200  # @param {type:"integer"}

num_atoms = 51  # @param {type:"integer"}
min_q_value = -20  # @param {type:"integer"}
max_q_value = 20  # @param {type:"integer"}
n_step_update = 2  # @param {type:"integer"}

num_eval_episodes = 10  # @param {type:"integer"}
eval_interval = 1000  # @param {type:"integer"}

環境

前のように環境をロードします、一つは訓練のためそして一つは評価のためです。ここでは CartPole-v1 (vs. CartPole-v0 in DQN チュートリアル) を使用します、これは 200 よりも大きい 500 の大きな最大報酬を持ちます。

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)

エージェント

C51 は DQN に基づく Q-学習アルゴリズムです。DQN のように、それは離散アクション空間を持つ任意の環境上で利用できます。

C51 と DQN の間の主要な違いは各状態-アクションペアのための Q-値を単純に予測するのではなく、C51 は Q-値の確率分布のためのヒストグラム・モデルを予測します :

単純な期待値よりも分布を学習することにより、アルゴリズムは訓練の間よりステーブルのままでいることが可能で、改良された最終的なパフォーマンスに繋がります。これは特に二峰性分布の状況あるいは多峰性値分布でさえ特に真で、そこでは単一平均は正確な描写を提供しません。

値の上ではなく確率分布上で訓練するため、C51 はその損失関数を計算するためにある複雑な分布計算を遂行しなければなりません。しかし心配しないでください、これの総ては貴方のために TF-Agents でケアされます！

C51 エージェントを作成するため、最初に CategoricalQNetwork を作成する必要があります。CategoricalQNetwork の API は QNetwork のそれと同じです、追加引数 num_atoms があることを除いて。これは確率分布推定のサポート点の数を表します。(上の画像は 10 サポート点を含みます、各々は垂直の青色棒により表されます。) 名前から分かるように、アトムのデフォルト数は 51 です。

categorical_q_net = categorical_q_network.CategoricalQNetwork(
    train_env.observation_spec(),
    train_env.action_spec(),
    num_atoms=num_atoms,
    fc_layer_params=fc_layer_params)

ちょうど作成したネットワークを訓練するために optimizer、そしてネットワークが何回更新されたかを追跡するために train_step_counter 変数も必要です。

vanilla DqnAgent との一つの他の重要な違いは今は引数として min_q_value と max_q_value を指定する必要があることです。これらはサポートの最も極値を指定します (換言すれば、両側の 51 アトムの最も極値)。貴方の特定の環境のためにこれらを適切に選択することを確実にしてください。ここでは -20 と 20 を使用します。

optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)

train_step_counter = tf.compat.v2.Variable(0)

agent = categorical_dqn_agent.CategoricalDqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    categorical_q_network=categorical_q_net,
    optimizer=optimizer,
    min_q_value=min_q_value,
    max_q_value=max_q_value,
    n_step_update=n_step_update,
    td_errors_loss_fn=common.element_wise_squared_loss,
    gamma=gamma,
    train_step_counter=train_step_counter)
agent.initialize()

注意すべき一つの最後のことは $n$ = 2 で n-step 更新を使用するために引数も追加したことです。私達は (bellman 最適方程式に基づいて) 単一ステップのリターンを使用して現在の時間ステップと次の時間ステップの Q-値の間の誤差を計算するだけです。単一ステップのリターンは次のように定義されます :

$G_t = R_{t + 1} + \gamma V(s_{t + 1})$

ここで $V(s) = \max_a{Q(s, a)}$ を定義します。

N-ステップ更新は標準的な単一ステップの報酬関数を $n$ 回展開を伴います :

$G_t^n = R_{t + 1} + \gamma R_{t + 2} + \gamma^2 R_{t + 3} + \dots + \gamma^n V(s_{t + n})$

N-ステップ更新は更なる将来からブートストラップすることを可能にし、そして $n$ の正しい値で、これはしばしばより高速な学習に繋がります。

C51 と n-ステップ更新はしばしば Rainbow エージェントのコアを形成するために優先付けられた再生と結合されますが、私達は優先付けられた再生の実装から測定可能な改良は見ませんでした。その上、C51 エージェントを n-ステップ更新だけと結合するとき、エージェントはテストしたAtari 環境のサンプル上他の Rainbow エージェントと同じように上手く遂行します。

メトリクスと評価

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

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]


random_policy = random_tf_policy.RandomTFPolicy(train_env.time_step_spec(),
                                                train_env.action_spec())

compute_avg_return(eval_env, random_policy, num_eval_episodes)

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

データ収集

DQN チュートリアル内のように、再生バッファとランダムポリシーで初期データ収集をセットアップします。

#@test {"skip": true}
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec=agent.collect_data_spec,
    batch_size=train_env.batch_size,
    max_length=replay_buffer_capacity)

def collect_step(environment, policy):
  time_step = environment.current_time_step()
  action_step = policy.action(time_step)
  next_time_step = environment.step(action_step.action)
  traj = trajectory.from_transition(time_step, action_step, next_time_step)

  # Add trajectory to the replay buffer
  replay_buffer.add_batch(traj)

for _ in range(initial_collect_steps):
  collect_step(train_env, random_policy)

# This loop is so common in RL, that we provide standard implementations of
# these. For more details see the drivers module.

# Dataset generates trajectories with shape [BxTx...] where
# T = n_step_update + 1.
dataset = replay_buffer.as_dataset(
    num_parallel_calls=3, sample_batch_size=batch_size,
    num_steps=n_step_update + 1).prefetch(3)

iterator = iter(dataset)

エージェントを訓練する

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

以下は実行に ~7 分かかります。

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

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

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

# Evaluate the agent's policy once before training.
avg_return = compute_avg_return(eval_env, agent.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_step(train_env, agent.collect_policy)

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

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

可視化

プロット

エージェントのパフォーマンスを見るためにリターン vs グローバルステップをプロットすることができます。Cartpole-v1 では、ポールが直立し続ける総ての時間ステップのために環境は +1 の報酬を与え、そしてステップのあ最大数は 500 ですから、最大の可能なリターンもまた 500 です。

#@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=550)

ビデオ

各ステップで環境をレンダリングすることによりエージェントのパフォーマンスを可視化することは有用です。それを行なう前に、最初にビデオをこの 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)

以下のコードは 2,3 のエピソードのためにエージェントのポリシーを可視化します :

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 = agent.policy.action(time_step)
      time_step = eval_env.step(action_step.action)
      video.append_data(eval_py_env.render())

embed_mp4(video_filename)

C51 は CartPole-v1 上 DQN よりも僅かに良く行なう傾向がありますが、2 つのエージェントの間の違いはますます複雑な環境で更に更に重要になります。例えば、full Atari 2600 ベンチマーク上、C51 はランダムエージェントに関して正規化後 DQN を越えて 126 % の平均スコア改善を実演します。n-ステップ更新を含むことにより追加の改善が獲得できます。

For a deeper dive into the C51 algorithm, see A Distributional Perspective on Reinforcement Learning (2017).

以上

月	火	水	木	金	土	日
		1	2	3	4	5
6	7	8	9	10	11	12
13	14	15	16	17	18	19
20	21	22	23	24	25	26
27	28	29	30