#@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[reverb]
!pip install pyglet

from __future__ import absolute_import, division, print_function

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

import tensorflow as tf

from tf_agents.agents.dqn import dqn_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.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import sequential
from tf_agents.policies import py_tf_eager_policy
from tf_agents.policies import random_tf_policy
from tf_agents.replay_buffers import reverb_replay_buffer
from tf_agents.replay_buffers import reverb_utils
from tf_agents.trajectories import trajectory
from tf_agents.specs import tensor_spec
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()

tf.version.VERSION

"""
Explanation: TF-Agent を使用した Deep Q Network のトレーニング
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/agents/tutorials/1_dqn_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/1_dqn_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/1_dqn_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/1_dqn_tutorial.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a></td>
</table>

はじめに
この例は、Cartpole環境でTF-Agentsライブラリを使用して DQN（Deep Q Networks）エージェントをトレーニングする方法を示しています。

ここでは、トレーニング、評価、データ収集のための強化学習（RL）パイプラインのすべてのコンポーネントについて説明します。
このコードをライブで実行するには、上の [Google Colabで実行] リンクをクリックしてください。
セットアップ
以下の依存関係をインストールしていない場合は、実行します。
End of explanation
"""


num_iterations = 20000 # @param {type:"integer"}

initial_collect_steps = 100  # @param {type:"integer"}
collect_steps_per_iteration =   1# @param {type:"integer"}
replay_buffer_max_length = 100000  # @param {type:"integer"}

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

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

"""
Explanation: ハイパーパラメータ
End of explanation
"""


env_name = 'CartPole-v0'
env = suite_gym.load(env_name)

"""
Explanation: 環境
強化学習（RL）では、環境はタスクまたは解決すべき問題を表します。標準環境は、tf_agents.environmentsスイートを使用して TF-Agent で作成できます。TF-Agent には、OpenAI Gym、Atari、DM Control などのソースから環境を読み込むためのスイートがあります。
OpenAI Gym スイートから CartPole 環境を読み込みます。
End of explanation
"""


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

"""
Explanation: この環境をレンダリングして、どのように見えるかを確認できます。台車の上に自由に振り動く棒を立て、その棒が倒れないように台車を左右に動かすことが目標です。
End of explanation
"""


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

print('Reward Spec:')
print(env.time_step_spec().reward)

"""
Explanation: environment.stepメソッドは、環境内でactionを取り、環境の次の観測と行動の報酬を含むTimeStepタプルを返します。
time_step_spec()メソッドは、TimeStepタプルの仕様を返します。そのobservation属性は、観測の形状、データ型、および許容値の範囲を示します。reward属性は、報酬についての同じ詳細を示します。
End of explanation
"""


print('Action Spec:')
print(env.action_spec())

"""
Explanation: action_spec()メソッドは、有効な行動の形状、データタイプ、および許容値を返します。
End of explanation
"""


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)

"""
Explanation: CartPole 環境では

observationは以下の 4 つの float 型の配列です。
台車の位置と速度
棒の角度位置と速度


rewardはスカラーの浮動小数点値です
actionは可能な値が 2 つだけのスカラー整数です。
0 — 「左に移動」
1 — 「右に移動」
End of explanation
"""


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

"""
Explanation: 通常、2 つの環境（トレーニング用、評価用）のみがインスタンス化されます。
End of explanation
"""


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

"""
Explanation: CartPole 環境は、ほとんどの環境と同様に純粋な Python で記述されています。これは、TFPyEnvironmentラッパーを使用して TensorFlow に変換されます。
従来の環境の API は Numpy 配列を使用します。TFPyEnvironmentは、これらをTensorsに変換して、Tensorflow エージェントおよびポリシーとの互換性を確保します。
End of explanation
"""


fc_layer_params = (100, 50)
action_tensor_spec = tensor_spec.from_spec(env.action_spec())
num_actions = action_tensor_spec.maximum - action_tensor_spec.minimum + 1

# Define a helper function to create Dense layers configured with the right
# activation and kernel initializer.
def dense_layer(num_units):
  return tf.keras.layers.Dense(
      num_units,
      activation=tf.keras.activations.relu,
      kernel_initializer=tf.keras.initializers.VarianceScaling(
          scale=2.0, mode='fan_in', distribution='truncated_normal'))

# QNetwork consists of a sequence of Dense layers followed by a dense layer
# with `num_actions` units to generate one q_value per available action as
# its output.
dense_layers = [dense_layer(num_units) for num_units in fc_layer_params]
q_values_layer = tf.keras.layers.Dense(
    num_actions,
    activation=None,
    kernel_initializer=tf.keras.initializers.RandomUniform(
        minval=-0.03, maxval=0.03),
    bias_initializer=tf.keras.initializers.Constant(-0.2))
q_net = sequential.Sequential(dense_layers + [q_values_layer])

"""
Explanation: エージェント
RL問題の解決に使用されるアルゴリズムは、Agentで表されます。TF-Agent は、以下を含むさまざまなAgentsの標準実装を提供します。

DQN（本チュートリアルで使用）
REINFORCE
DDPG
TD3
PPO
SAC

DQN エージェントは、個別の行動領域がある任意の環境で使用できます。
DQN エージェントの中心にあるのがQNetworkです。これは、環境から観察を得て、すべての行動のQValues（予期される戻り値）を予測する方法を学習するニューラルネットワークモデルです。
tf_agents.networks. を使用して QNetwork を作成します。ネットワークは一連の tf.keras.layers.Dense レイヤーを構成し、最後のレイヤーは、可能なアクションごとに 1 つの出力を持ちます。
End of explanation
"""


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

train_step_counter = tf.Variable(0)

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    optimizer=optimizer,
    td_errors_loss_fn=common.element_wise_squared_loss,
    train_step_counter=train_step_counter)

agent.initialize()

"""
Explanation: tf_agents.agents.dqn.dqn_agentを使用してDqnAgentをインスタンス化します。time_step_spec、action_spec、および QNetwork に加えて、エージェントコンストラクタにもオプティマイザ（この場合、AdamOptimizer）、損失関数、および整数ステップカウンタが必要です。
End of explanation
"""


eval_policy = agent.policy
collect_policy = agent.collect_policy

"""
Explanation: ポリシー
ポリシーは、エージェントが環境で行動する方法を定義します。通常、強化学習の目標は、ポリシーが望ましい結果を生成するまで、基礎となるモデルをトレーニングすることです。
このチュートリアルでは

望ましい結果は、台車の上の棒が倒れないようにバランスを保つことです。
ポリシーは、time_step観測ごとに行動（左または右）を返します。

エージェントには 2 つのポリシーが含まれています。

agent.policy — 評価とデプロイに使用される主なポリシー。
agent.collect_policy — データ収集に使用される補助的なポリシー。
End of explanation
"""


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

"""
Explanation: ポリシーはエージェントとは無関係に作成できます。たとえば、tf_agents.policies.random_tf_policyを使用して、各time_stepの行動をランダムに選択するポリシーを作成できます。
End of explanation
"""


example_environment = tf_py_environment.TFPyEnvironment(
    suite_gym.load('CartPole-v0'))

time_step = example_environment.reset()

random_policy.action(time_step)

"""
Explanation: ポリシーから行動を取得するには、policy.action(time_step)メソッドを呼び出します。time_stepには、環境からの観測が含まれています。このメソッドは、3 つのコンポーネントを持つ名前付きタプルであるPolicyStepを返します。

action — 実行する行動（ここでは0または1)
state — ステートフルポリシー（RNNベース）に使用
info — 行動のログ確率などの補助データ
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]


# See also the metrics module for standard implementations of different metrics.
# https://github.com/tensorflow/agents/tree/master/tf_agents/metrics

"""
Explanation: 指標と評価
ポリシーの評価に使用される最も一般的な指標は、平均リターンです。リターンは、エピソードの環境でポリシーを実行中に取得した報酬の合計です。エピソードは何回か実行され、平均リターンが生成されます。
次の関数は、ポリシー、環境、およびエピソードの数を指定して、ポリシーの平均リターンを計算します。
End of explanation
"""


compute_avg_return(eval_env, random_policy, num_eval_episodes)

"""
Explanation: この計算をrandom_policyで実行すると、環境のベースラインパフォーマンスが示されます。
End of explanation
"""


table_name = 'uniform_table'
replay_buffer_signature = tensor_spec.from_spec(
      agent.collect_data_spec)
replay_buffer_signature = tensor_spec.add_outer_dim(
    replay_buffer_signature)

table = reverb.Table(
    table_name,
    max_size=replay_buffer_max_length,
    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(
    agent.collect_data_spec,
    table_name=table_name,
    sequence_length=2,
    local_server=reverb_server)

rb_observer = reverb_utils.ReverbAddTrajectoryObserver(
  replay_buffer.py_client,
  table_name,
  sequence_length=2)

"""
Explanation: 再生バッファ
環境から収集されたデータを追跡するために、効率的かつ拡張可能で使いやすい Reverb という Deepmind の再生システムを使用します。軌跡を収集する際に経験データを保存し、トレーニング中に消費されます。
再生バッファは、格納されるテンソルを記述する仕様を使用して構築されます。これは、agent.collect_data_spec を使用してエージェントから取得できます。
End of explanation
"""


agent.collect_data_spec

agent.collect_data_spec._fields

"""
Explanation: ほとんどのエージェントの場合、collect_data_specは、Trajectoryと呼ばれる名前付きタプルであり、観測、行動、報酬、およびその他の要素の仕様が含まれています。
End of explanation
"""


#@test {"skip": true}
py_driver.PyDriver(
    env,
    py_tf_eager_policy.PyTFEagerPolicy(
      random_policy, use_tf_function=True),
    [rb_observer],
    max_steps=initial_collect_steps).run(train_py_env.reset())

"""
Explanation: データ収集
次に、環境でランダムポリシーを数ステップ実行し、データを再生バッファに記録します。
ここでは、'PyDriver' を使用して、経験を収集するループを実行しています。TF Agents ドライバーについては、ドライバーのチュートリアルをご覧ください。
End of explanation
"""


# For the curious:
# Uncomment to peel one of these off and inspect it.
# iter(replay_buffer.as_dataset()).next()

"""
Explanation: 再生バッファは Trajectory のコレクションではありません。
End of explanation
"""


# 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)

dataset

iterator = iter(dataset)
print(iterator)

# For the curious:
# Uncomment to see what the dataset iterator is feeding to the agent.
# Compare this representation of replay data 
# to the collection of individual trajectories shown earlier.

# iterator.next()

"""
Explanation: エージェントは再生バッファにアクセスする必要があります。これは、エージェントにデータをフィードするイテレーション可能なtf.data.Datasetパイプラインを作成することにより提供されます。
再生バッファの各行には、1 つの観測ステップのみが格納されます。しかし、DQN エージェントは損失を計算するためにその時点の観測と次の観測を両方必要とするので、データセットパイプラインは、バッチ内の要素ごとに2つの隣接する行をサンプリングします(num_steps=2)。
また、このデータセットは、並列呼び出しを実行してデータをプリフェッチすることにより最適化されます。
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]

# Reset the environment.
time_step = train_py_env.reset()

# Create a driver to collect experience.
collect_driver = py_driver.PyDriver(
    env,
    py_tf_eager_policy.PyTFEagerPolicy(
      agent.collect_policy, use_tf_function=True),
    [rb_observer],
    max_steps=collect_steps_per_iteration)

for _ in range(num_iterations):

  # Collect a few steps and save to the replay buffer.
  time_step, _ = collect_driver.run(time_step)

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

  step = agent.train_step_counter.numpy()

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

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

"""
Explanation: エージェントのトレーニング
トレーニングループ時には、以下の 2 つが行われる必要があります。

環境からデータを収集する
そのデータを使用してエージェントのニューラルネットワークをトレーニングする

この例では、定期的にポリシーを評価し、その時点のスコアを出力します。
以下の実行には 5 分ほどかかります。
End of explanation
"""


#@test {"skip": true}

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

"""
Explanation: 可視化
プロット
matplotlib.pyplotを使用して、トレーニング中にポリシーがどのように改善されたかをグラフ化します。
Cartpole-v0の 1 回のイテレーションには、200のタイムステップがあります。環境は、棒が立ったままでいる各ステップに対して+1の報酬を与えるので、1 つのエピソードの最大リターンは 200 です。グラフは、トレーニング中に評価されるたびに、最大値に向かって増加するリターンを示しています。（多少不安定になり、毎回単調に増加しないこともあります。）
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: 動画
グラフは便利ですが、エージェントが環境内で実際にタスクを実行している様子を動画で視覚化するとさらに分かりやすくなります。
まず、ノートブックに動画を埋め込む関数を作成します。
End of explanation
"""


def create_policy_eval_video(policy, filename, num_episodes=5, fps=30):
  filename = filename + ".mp4"
  with imageio.get_writer(filename, fps=fps) 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 = policy.action(time_step)
        time_step = eval_env.step(action_step.action)
        video.append_data(eval_py_env.render())
  return embed_mp4(filename)

create_policy_eval_video(agent.policy, "trained-agent")

"""
Explanation: CartPole ゲームのいくつかのエピソードをエージェントで繰り返します。基礎となる Python 環境（TensorFlow 環境ラッパーの「内部」のもの）は、環境の状態のイメージを出力するrender()メソッドを提供します。これらは動画に収集できます。
End of explanation
"""


create_policy_eval_video(random_policy, "random-agent")

"""
Explanation: トレーニングされたエージェント（上記）をランダムに移動するエージェントと比較してみてください。（ランダムに移動するエージェントは、トレーニングされたエージェントと比べて上手くできません。）
End of explanation
"""


