<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2021 The TF-Agents Authors.
Step1:  DQN C51/Rainbow
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https
Step2:  ハイパーパラメータ
Step3:  環境
前回のように環境を読み込みます。1つはトレーニング用で、もう1つは評価用です。ここでは、最大報酬が200ではなく500であるCartPole-v1（DQNチュートリアルではCartPole-v0を使用）を使用します。
Step4:  エージェント
C51は、DQNに基づくQ学習アルゴリズムです。 DQNと同様に、個別の行動空間がある任意の環境で使用できます。
C51とDQNの主な違いは、各状態と行動のペアのQ値を単に予測するのではなく、C51はQ値の確率分布のヒストグラムモデルを予測することです。
単なる推定値ではなく分布を学習することで、アルゴリズムはトレーニング時に安定性を維持でき、最終的なパフォーマンスが向上します。これは、単一の平均では正確な画像が得られない、バイモーダルまたはマルチモーダルの値の分布がある場合に特に有用です。
C51は値ではなく確率分布でトレーニングするために、損失関数を計算するために複雑な分布計算を実行する必要がありますが、これらはすべてTF-Agentで処理されます。
C51 Agentを作成するには、まずCategoricalQNetworkを作成する必要があります。CategoricalQNetworkのAPIは、QNetworkのAPIと同じですが、引数num_atomsが追加されます。これは、確率分布の推定におけるサポートポイントの数を表します。（上の画像には10個のサポートポイントが含まれており、それぞれが青い縦棒で表されています。）名前からわかるように、デフォルトのatom数は51です。
Step5:  また、先ほど作成したネットワークをトレーニングするためのoptimizerと、ネットワークが更新された回数を追跡するためのtrain_step_counter変数も必要です。
簡単なDqnAgentとのもう1つの重要な違いは、引数としてmin_q_valueとmax_q_valueを指定する必要があることです。これらは、サポートの最も極端な値（51のatomの最も極端な値）を指定します。特定の環境に合わせてこれらを適切に選択してください。 ここでは、-20と20を使用します。
Step6:  最後に注意すべき点は、$n$ = 2のnステップ更新を使用する引数も追加されていることです。１ステップのQ学習（$n$ = 1）では、（ベルマン最適化方程式に基づいて）１ステップの戻り値を使用して、その時点のタイムステップと次のタイムステップでの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 agentの中核となっていますが、優先再生の実装による測定可能な改善は見られませんでした。さらに、C51エージェントをnステップの更新のみと組み合わせた場合、エージェントは、テストしたAtari環境のサンプルで他のRainbowエージェントと同様に機能することが明らかになっています。
指標と評価
ポリシーの評価に使用される最も一般的なメトリックは、平均利得です。利得は、エピソードの環境でポリシーを実行中に取得した報酬の総計です。通常、数エピソードが実行され、平均利得が生成されます。
Step7:  データ収集
DQNチュートリアルと同様に、ランダムポリシーを使用して再生バッファーと初期データ収集を設定します。
Step8:  エージェントのトレーニング
トレーニングループには、環境からのデータ収集とエージェントのネットワークの最適化の両方が含まれます。途中で、エージェントのポリシーを時々評価して、状況を確認します。
以下の実行には7分ほどかかります。
Step9:  可視化
プロット
利得とグローバルステップをプロットして、エージェントのパフォーマンスを確認できます。Cartpole-v1では、棒が立ったままでいるタイムステップごとに環境は+1の報酬を提供します。最大ステップ数は500であるため、可能な最大利得値も500です。
Step11:  動画
各ステップで環境をレンダリングすると、エージェントのパフォーマンスを可視化できます。その前に、このColabに動画を埋め込む関数を作成しましょう。
Step12:  次のコードは、いくつかのエピソードに渡るエージェントのポリシーを可視化します。
<ASSISTANT_TASK:>
Python Code:
#@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.
Explanation: Copyright 2021 The TF-Agents Authors.
End of explanation
!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
!pip install pyglet
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
# Set up a virtual display for rendering OpenAI gym environments.
display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()
Explanation: DQN C51/Rainbow
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/agents/tutorials/9_c51_tutorial">     <img src="https://www.tensorflow.org/images/tf_logo_32px.png">     TensorFlow.org で表示</a>
</td>
  <td>     <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/agents/tutorials/9_c51_tutorial.ipynb">     <img src="https://www.tensorflow.org/images/colab_logo_32px.png">     Google Colab で実行</a>
</td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ja/agents/tutorials/9_c51_tutorial.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png">GitHub でソースを表示</a></td>
  <td>     <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/agents/tutorials/9_c51_tutorial.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a>   </td>
</table>
はじめに
ここでは、Cartpole環境でTF-Agentライブラリを使用してCategorical DQN (C51)エージェントをトレーニングする方法を紹介します。
はじめる前に、 DQNチュートリアルをご覧ください。このチュートリアルは、DQNチュートリアルに精通されていることを前提とし、主にDQNとC51の違いについて見ていきます。
セットアップ
TF-agentをまだインストールしていない場合は、次を実行します。
End of explanation
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"}
Explanation: ハイパーパラメータ
End of explanation
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)
Explanation: 環境
前回のように環境を読み込みます。1つはトレーニング用で、もう1つは評価用です。ここでは、最大報酬が200ではなく500であるCartPole-v1（DQNチュートリアルではCartPole-v0を使用）を使用します。
End of explanation
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)
Explanation: エージェント
C51は、DQNに基づくQ学習アルゴリズムです。 DQNと同様に、個別の行動空間がある任意の環境で使用できます。
C51とDQNの主な違いは、各状態と行動のペアのQ値を単に予測するのではなく、C51はQ値の確率分布のヒストグラムモデルを予測することです。
単なる推定値ではなく分布を学習することで、アルゴリズムはトレーニング時に安定性を維持でき、最終的なパフォーマンスが向上します。これは、単一の平均では正確な画像が得られない、バイモーダルまたはマルチモーダルの値の分布がある場合に特に有用です。
C51は値ではなく確率分布でトレーニングするために、損失関数を計算するために複雑な分布計算を実行する必要がありますが、これらはすべてTF-Agentで処理されます。
C51 Agentを作成するには、まずCategoricalQNetworkを作成する必要があります。CategoricalQNetworkのAPIは、QNetworkのAPIと同じですが、引数num_atomsが追加されます。これは、確率分布の推定におけるサポートポイントの数を表します。（上の画像には10個のサポートポイントが含まれており、それぞれが青い縦棒で表されています。）名前からわかるように、デフォルトのatom数は51です。
End of explanation
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)
train_step_counter = tf.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()
Explanation: また、先ほど作成したネットワークをトレーニングするためのoptimizerと、ネットワークが更新された回数を追跡するためのtrain_step_counter変数も必要です。
簡単なDqnAgentとのもう1つの重要な違いは、引数としてmin_q_valueとmax_q_valueを指定する必要があることです。これらは、サポートの最も極端な値（51のatomの最も極端な値）を指定します。特定の環境に合わせてこれらを適切に選択してください。 ここでは、-20と20を使用します。
End of explanation
#@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]
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
# metrics.
Explanation: 最後に注意すべき点は、$n$ = 2のnステップ更新を使用する引数も追加されていることです。１ステップのQ学習（$n$ = 1）では、（ベルマン最適化方程式に基づいて）１ステップの戻り値を使用して、その時点のタイムステップと次のタイムステップでの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 agentの中核となっていますが、優先再生の実装による測定可能な改善は見られませんでした。さらに、C51エージェントをnステップの更新のみと組み合わせた場合、エージェントは、テストしたAtari環境のサンプルで他のRainbowエージェントと同様に機能することが明らかになっています。
指標と評価
ポリシーの評価に使用される最も一般的なメトリックは、平均利得です。利得は、エピソードの環境でポリシーを実行中に取得した報酬の総計です。通常、数エピソードが実行され、平均利得が生成されます。
End of explanation
#@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)
Explanation: データ収集
DQNチュートリアルと同様に、ランダムポリシーを使用して再生バッファーと初期データ収集を設定します。
End of explanation
#@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)
Explanation: エージェントのトレーニング
トレーニングループには、環境からのデータ収集とエージェントのネットワークの最適化の両方が含まれます。途中で、エージェントのポリシーを時々評価して、状況を確認します。
以下の実行には7分ほどかかります。
End of explanation
#@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)
Explanation: 可視化
プロット
利得とグローバルステップをプロットして、エージェントのパフォーマンスを確認できます。Cartpole-v1では、棒が立ったままでいるタイムステップごとに環境は+1の報酬を提供します。最大ステップ数は500であるため、可能な最大利得値も500です。
End of explanation
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)
Explanation: 動画
各ステップで環境をレンダリングすると、エージェントのパフォーマンスを可視化できます。その前に、このColabに動画を埋め込む関数を作成しましょう。
End of explanation
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)
Explanation: 次のコードは、いくつかのエピソードに渡るエージェントのポリシーを可視化します。
End of explanation
<END_TASK>