#@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
"""


!pip install tf-agents

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

import abc
import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np

from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.networks import network

from tf_agents.policies import py_policy
from tf_agents.policies import random_py_policy
from tf_agents.policies import scripted_py_policy

from tf_agents.policies import tf_policy
from tf_agents.policies import random_tf_policy
from tf_agents.policies import actor_policy
from tf_agents.policies import q_policy
from tf_agents.policies import greedy_policy

from tf_agents.trajectories import time_step as ts

"""
Explanation: 策略
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://tensorflow.google.cn/agents/tutorials/3_policies_tutorial"><img src="https://tensorflow.google.cn/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/zh-cn/agents/tutorials/3_policies_tutorial.ipynb"><img src="https://tensorflow.google.cn/images/colab_logo_32px.png">在 Google Colab 运行</a>
</td>
  <td>     <a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/agents/tutorials/3_policies_tutorial.ipynb"><img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png">在 Github 上查看源代码</a>
</td>
  <td>     <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/agents/tutorials/3_policies_tutorial.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png">下载笔记本</a>   </td>
</table>

简介
在强化学习中，“策略”一词是指将环境的观测值映射到操作或操作分布。在 TF-Agents 中，环境的观测值包含在命名元组 TimeStep('step_type', 'discount', 'reward', 'observation') 内，而策略会将时间步骤映射到操作或操作分布。大多数策略使用 timestep.observation，某些策略会使用 timestep.step_type（例如，在有状态策略中，在片段开始时重置状态），但 timestep.discount 和 timestep.reward 通常会被忽略。
策略与 TF-Agents 中其他组成部分的相关性如下。大多数策略都通过神经网络来计算时间步骤的操作和/或操作分布。代理可以包含一个或多个用于不同目的的策略，例如采用一个主要策略进行训练用以部署，并采用一个噪声策略进行数据收集。策略可以保存/恢复，并且可以在不依赖于代理的状态下用于数据收集、评估等工作。
一些策略在 Tensorflow 中更容易编写（例如神经网络策略），而其他策略在 Python 中更容易编写（例如基于操作脚本的策略）。 因此，在 TF-Agent 中，Python 和 Tensorflow 策略均支持使用。此外，有时可能需要在 Python 环境中使用通过 TensorFlow 编写的策略，反之亦然（例如使用 TensorFlow 策略进行训练，但之后需要在生产 python 环境中进行部署）。为了简化这一步骤，我们提供了用于在 python 和 TensorFlow 策略之间进行转换的包装器。
与策略有关的另一项有趣内容是策略包装器，它们能够以某种方式修改给定策略，例如添加特定类型的噪声、制作随机策略的贪心或 ϵ 贪心版本、随机混合多个策略等。 
设置
如果尚未安装 TF-Agents，请运行以下命令：
End of explanation
"""


class Base(object):

  @abc.abstractmethod
  def __init__(self, time_step_spec, action_spec, policy_state_spec=()):
    self._time_step_spec = time_step_spec
    self._action_spec = action_spec
    self._policy_state_spec = policy_state_spec

  @abc.abstractmethod
  def reset(self, policy_state=()):
    # return initial_policy_state.
    pass

  @abc.abstractmethod
  def action(self, time_step, policy_state=()):
    # return a PolicyStep(action, state, info) named tuple.
    pass

  @abc.abstractmethod
  def distribution(self, time_step, policy_state=()):
    # Not implemented in python, only for TF policies.
    pass

  @abc.abstractmethod
  def update(self, policy):
    # update self to be similar to the input `policy`.
    pass

  @property
  def time_step_spec(self):
    return self._time_step_spec

  @property
  def action_spec(self):
    return self._action_spec

  @property
  def policy_state_spec(self):
    return self._policy_state_spec

"""
Explanation: Python 策略
Python 策略的接口在 policies/py_policy.PyPolicy 中进行定义。主要方法为：
End of explanation
"""


action_spec = array_spec.BoundedArraySpec((2,), np.int32, -10, 10)
my_random_py_policy = random_py_policy.RandomPyPolicy(time_step_spec=None,
    action_spec=action_spec)
time_step = None
action_step = my_random_py_policy.action(time_step)
print(action_step)
action_step = my_random_py_policy.action(time_step)
print(action_step)

"""
Explanation: 最重要的方法为 action(time_step)，该方法可将包含环境观测值的 time_step 映射到包含以下特性的 PolicyStep 命名元组：

action：应用于环境的操作。
state：将被馈入下一个操作调用的策略状态（例如 RNN 状态）。
info：可选辅助信息，例如操作日志概率。

time_step_spec 和 action_spec 分别为输入时间步骤规范和输出操作规范。策略还具有 reset 函数，通常用于重置有状态策略中的状态。update(new_policy) 函数可将 self 朝着 new_policy 更新。
现在，让我们看看 Python 策略的两个示例。
示例 1：随机 Python 策略
RandomPyPolicy 是 PyPolicy 的一个简单示例， 它可以为给定的离散/连续 action_spec 生成随机操作。将忽略输入 time_step。
End of explanation
"""


action_spec = array_spec.BoundedArraySpec((2,), np.int32, -10, 10)
action_script = [(1, np.array([5, 2], dtype=np.int32)), 
                 (0, np.array([0, 0], dtype=np.int32)), # Setting `num_repeats` to 0 will skip this action.
                 (2, np.array([1, 2], dtype=np.int32)), 
                 (1, np.array([3, 4], dtype=np.int32))]

my_scripted_py_policy = scripted_py_policy.ScriptedPyPolicy(
    time_step_spec=None, action_spec=action_spec, action_script=action_script)

policy_state = my_scripted_py_policy.get_initial_state()
time_step = None
print('Executing scripted policy...')
action_step = my_scripted_py_policy.action(time_step, policy_state)
print(action_step)
action_step= my_scripted_py_policy.action(time_step, action_step.state)
print(action_step)
action_step = my_scripted_py_policy.action(time_step, action_step.state)
print(action_step)

print('Resetting my_scripted_py_policy...')
policy_state = my_scripted_py_policy.get_initial_state()
action_step = my_scripted_py_policy.action(time_step, policy_state)
print(action_step)

"""
Explanation: 示例 2：脚本化 Python 策略
脚本化策略可回放以 (num_repeats, action) 元组列表形式呈现的操作脚本。每次调用 action 函数时，它都会返回列表内的下一个操作并使其完成指定的重复次数，然后再继续执行列表中的下一个操作。调用 reset 方法即可从列表开头重新开始执行。
End of explanation
"""


action_spec = tensor_spec.BoundedTensorSpec(
    (2,), tf.float32, minimum=-1, maximum=3)
input_tensor_spec = tensor_spec.TensorSpec((2,), tf.float32)
time_step_spec = ts.time_step_spec(input_tensor_spec)

my_random_tf_policy = random_tf_policy.RandomTFPolicy(
    action_spec=action_spec, time_step_spec=time_step_spec)
observation = tf.ones(time_step_spec.observation.shape)
time_step = ts.restart(observation)
action_step = my_random_tf_policy.action(time_step)

print('Action:')
print(action_step.action)

"""
Explanation: TensorFlow 策略
TensorFlow 策略采用与 Python 策略相同的接口。让我们看几个例子。
示例 1：随机 TF 策略
使用 RandomTFPolicy 可根据给定的离散/连续 action_spec 生成随机操作。将忽略输入 time_step。
End of explanation
"""


class ActionNet(network.Network):

  def __init__(self, input_tensor_spec, output_tensor_spec):
    super(ActionNet, self).__init__(
        input_tensor_spec=input_tensor_spec,
        state_spec=(),
        name='ActionNet')
    self._output_tensor_spec = output_tensor_spec
    self._sub_layers = [
        tf.keras.layers.Dense(
            action_spec.shape.num_elements(), activation=tf.nn.tanh),
    ]

  def call(self, observations, step_type, network_state):
    del step_type

    output = tf.cast(observations, dtype=tf.float32)
    for layer in self._sub_layers:
      output = layer(output)
    actions = tf.reshape(output, [-1] + self._output_tensor_spec.shape.as_list())

    # Scale and shift actions to the correct range if necessary.
    return actions, network_state

"""
Explanation: 示例 2：参与者策略
可以使用将 time_steps 映射到操作的网络或将 time_steps 映射到操作分布的网络创建参与者策略。
使用操作网络
让我们按照以下方式定义网络：
End of explanation
"""


input_tensor_spec = tensor_spec.TensorSpec((4,), tf.float32)
time_step_spec = ts.time_step_spec(input_tensor_spec)
action_spec = tensor_spec.BoundedTensorSpec((3,),
                                            tf.float32,
                                            minimum=-1,
                                            maximum=1)

action_net = ActionNet(input_tensor_spec, action_spec)

my_actor_policy = actor_policy.ActorPolicy(
    time_step_spec=time_step_spec,
    action_spec=action_spec,
    actor_network=action_net)

"""
Explanation: 在 TensorFlow 中，大多数网络层都是针对批量运算而设计的，因此我们希望输入 time_step 得到批处理，网络的输出也得到批处理。另外，网络还负责在给定 action_spec 的正确范围内生成操作。常用方法是对最后一层使用 tanh 激活函数以在 [-1, 1] 区间内生成操作，然后将其缩放并移动到正确的范围作为输入 action_spec（例如，请参阅 tf_agents/agents/ddpg/networks.actor_network()）。
现在，我们可以使用以上网络创建参与者策略。
End of explanation
"""


batch_size = 2
observations = tf.ones([2] + time_step_spec.observation.shape.as_list())

time_step = ts.restart(observations, batch_size)

action_step = my_actor_policy.action(time_step)
print('Action:')
print(action_step.action)

distribution_step = my_actor_policy.distribution(time_step)
print('Action distribution:')
print(distribution_step.action)

"""
Explanation: 我们可以将其应用于遵循 time_step_spec 的任何 time_step 批次：
End of explanation
"""


class ActionDistributionNet(ActionNet):

  def call(self, observations, step_type, network_state):
    action_means, network_state = super(ActionDistributionNet, self).call(
        observations, step_type, network_state)

    action_std = tf.ones_like(action_means)
    return tfp.distributions.MultivariateNormalDiag(action_means, action_std), network_state


action_distribution_net = ActionDistributionNet(input_tensor_spec, action_spec)

my_actor_policy = actor_policy.ActorPolicy(
    time_step_spec=time_step_spec,
    action_spec=action_spec,
    actor_network=action_distribution_net)

action_step = my_actor_policy.action(time_step)
print('Action:')
print(action_step.action)
distribution_step = my_actor_policy.distribution(time_step)
print('Action distribution:')
print(distribution_step.action)

"""
Explanation: 在以上示例中，我们是使用生成操作张量的操作网络来创建的策略。在这种情况下，policy.distribution(time_step) 是围绕 policy.action(time_step) 输出的确定性（增量）分布。生成随机策略的一种方法是在策略包装器内包装参与者策略，为操作增加噪声。另一种方法是使用操作分布网络而非操作网络来创建参与者策略，如下所示。
使用操作分布网络
End of explanation
"""


input_tensor_spec = tensor_spec.TensorSpec((4,), tf.float32)
time_step_spec = ts.time_step_spec(input_tensor_spec)
action_spec = tensor_spec.BoundedTensorSpec((),
                                            tf.int32,
                                            minimum=0,
                                            maximum=2)
num_actions = action_spec.maximum - action_spec.minimum + 1


class QNetwork(network.Network):

  def __init__(self, input_tensor_spec, action_spec, num_actions=num_actions, name=None):
    super(QNetwork, self).__init__(
        input_tensor_spec=input_tensor_spec,
        state_spec=(),
        name=name)
    self._sub_layers = [
        tf.keras.layers.Dense(num_actions),
    ]

  def call(self, inputs, step_type=None, network_state=()):
    del step_type
    inputs = tf.cast(inputs, tf.float32)
    for layer in self._sub_layers:
      inputs = layer(inputs)
    return inputs, network_state


batch_size = 2
observation = tf.ones([batch_size] + time_step_spec.observation.shape.as_list())
time_steps = ts.restart(observation, batch_size=batch_size)

my_q_network = QNetwork(
    input_tensor_spec=input_tensor_spec,
    action_spec=action_spec)
my_q_policy = q_policy.QPolicy(
    time_step_spec, action_spec, q_network=my_q_network)
action_step = my_q_policy.action(time_steps)
distribution_step = my_q_policy.distribution(time_steps)

print('Action:')
print(action_step.action)

print('Action distribution:')
print(distribution_step.action)

"""
Explanation: 请注意，上例中的操作被裁剪到给定操作规范 [-1, 1] 区间内。这是因为 ActorPolicy clip 的构造函数参数的默认值为 True。将其设置为 False 将返回网络生成的未裁剪操作。 
可以使用 GreedyPolicy 包装器将随机策略转换为确定性策略，该包装器选择 stochastic_policy.distribution().mode() 作为其操作，并选择围绕此贪心操作的确定性/增量分布作为其 distribution()。
示例 3：Q 策略
Q 策略适用于 DQN 一类的代理，并基于预测每个离散操作 Q 值的 Q 网络。对于给定的时间步骤，Q 策略中的操作分布是使用 Q 值作为 logits 创建的分类分布。
End of explanation
"""


my_greedy_policy = greedy_policy.GreedyPolicy(my_q_policy)

action_step = my_greedy_policy.action(time_steps)
print('Action:')
print(action_step.action)

distribution_step = my_greedy_policy.distribution(time_steps)
print('Action distribution:')
print(distribution_step.action)

"""
Explanation: 策略包装器
策略包装器可用于包装和修改给定策略，例如增加噪声。策略包装器是策略 (Python/TensorFlow) 的子类，因此可以像其他任何策略一样使用。 
示例：贪心策略
贪心包装器可用于包装实现 distribution() 的任何 TensorFlow 策略。GreedyPolicy.action() 将返回 wrapped_policy.distribution().mode()，而 GreedyPolicy.distribution() 是围绕 GreedyPolicy.action() 的确定性/增量分布：
End of explanation
"""


