定义奖励函数
==============

奖励函数是强化学习中非常重要的一部分。它定义了策略在环境中采取动作所得到的奖励，这个奖励可以是正数、负数或零。
强化学习的目标就是最大化累积奖励，因此奖励函数定义了策略的优化目标。

在强化学习中，智能体通过不断地与环境交互来寻找最优策略。策略会根据当前状态选择一个行动，并接收相应的奖励以及下一个状态。
这样的过程会不断地进行，直到达到终止状态为止。如果策略在一系列动作和状态转换中获得高额的奖励，
那么它就找到了一种较好的策略。因此，奖励函数的设计对于强化学习策略的优化至关重要。一个好的奖励函数应该能够指导策略向着预期的方向进行学习。

REVIVE SDK支持以python源文件的方式定义奖励函数。奖励函数定义了策略的优化目标。奖励函数的输入是单步决策流的数据，
奖励函数的输出当前步策略获得的奖励值。下面给出一些奖励函数的示例：


.. note:: 奖励函数的名称必须是 ``get_reward`` 。

假设我们要教一个智能体玩一款跳跃游戏。该智能体应该根据屏幕上的信息，学习如何跳过障碍。奖励函数可以设计为，每当智能体跳过一个障碍，
就给予一定的正面奖励；每次掉入障碍时，就给予一定的负面奖励。对应的奖励函数示例如下：

.. code:: python

  import torch
  from typing import Dict

  def get_reward(data: Dict[str, torch.Tensor]) -> torch.Tensor:
      """奖励函数"""

      # 获得是否跳过障碍的状态
      is_passed = torch.gt(data['obs'][..., :1], 0.5)

      # 通过障碍给予100的奖励，不通过则给予-10的奖励(惩罚)
      passed_reward = 100
      unpassed_reward = -10

      # 根据条件计算每一个时间步的奖励
      reward = torch.where(is_passed, passed_reward, unpassed_reward)

      return reward

如果我们面对一个机器人控制任务，策略的目标是将将机器人从一个位置移动到另一个位置的情况，
奖励函数会根据机器人与目标位置的距离的缩小值来计算。那么对应的奖励函数示例如下：

.. code:: python

  import torch
  from typing import Dict

  def get_reward(data: Dict[str, torch.Tensor], target_pos: torch.Tensor) -> torch.Tensor:
      """机器人控制任务奖励函数"""

      # 获取当前机器人在做动作之前的位置
      pos = data['obs'][..., :2]
      # 获取当前机器人在做动作之后的位置
      next_pos = data['next_obs'][..., :2]
      # 获取目标位置
      target_pos = data['target_pos'][..., :2]

      # 计算机器人与目标位置的距离的缩小值作为reward
      dist1 = torch.norm(pos - target_pos)
      dist2 = torch.norm(next_pos - target_pos)
      reward = (dist1 - dist2) / dist1

      return reward

需要注意，在使用奖励函数对数据进行处理时，通常会将多个数据按批量（batch）组织起来一次性进行处理。这种方式可以提高代码的运行效率。
因此，在编写奖励函数时，需要注意保证函数能够处理与输入张量形状相对应的多维数据。此外，在计算奖励时，
我们通常会关注最后一维的特征维度。为方便处理，奖励函数的计算维度通常都设在了最后一维。因此，在使用数据时需要使用切片（ ``[..., n:m ]`` ）的方式获取数据的最后一维的特征，
并对特征进行计算。例如，取出obs数据的最后一维的前两个特征，可以使用以下方式：

.. code:: python

  obs_features = data['obs'][..., :2]  # 取最后一维的第一列和第二列


而对应返回的reward应该是一个对应的Pytorch Tensor，batch维度保持和输入数据一致， 最后一维特征的维度应该是1。

.. note:: 只有在训练策略模型时才需要并使用定义的奖励函数。训练虚拟环境模型时不需要定义奖励函数。因此REVIVE SDK支持在训练策略时使用不同的奖励函数训练多个策略模型。