tf-agent/2_environments_tutorial_cn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Copyright 2018 The TF-Agents Authors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Started\n",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/agents/blob/master/tf_agents/colabs/2_environments_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/tensorflow/agents/blob/master/tf_agents/colabs/2_environments_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "</table>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: If you haven't installed tf-agents or gym yet, run:\n",
    "!pip install tf-agents-nightly\n",
    "!pip install gym\n",
    "try:\n",
    "      %%tensorflow_version 2.x\n",
    "except:\n",
    "      pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 导入包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import abc\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "from tf_agents.environments import py_environment\n",
    "from tf_agents.environments import tf_environment\n",
    "from tf_agents.environments import tf_py_environment\n",
    "from tf_agents.environments import utils\n",
    "from tf_agents.specs import array_spec\n",
    "from tf_agents.environments import wrappers\n",
    "from tf_agents.environments import suite_gym\n",
    "from tf_agents.trajectories import time_step as ts\n",
    "\n",
    "tf.compat.v1.enable_v2_behavior()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 简介\n",
    "\n",
    "强化学习（RL）的目标就是设计一个代理（Agent），利用代理和环境（Environment）之间的交互进行学习。在标准的RL设置中，代理在每步接收一个观察值并选择一个动作。这个动作被应用到环境中，环境会返回一个奖励和一个新的观察结果。代理需要训练一种策略（Policy）来选择行动以使回报最大化，也称为回报。\n",
    "\n",
    "在TF-Agents中，环境可以使用Python或者TensorFlow来实现，Python实现的环境通常易于实现、理解和调试，但是TensorFlow实现的环境会更高效并且允许并行。最常见的工作方式是用Python实现一个环境，并使用TF-Agents的一个包装器将其自动转换为TensorFlow。\n",
    "\n",
    "让我们先来看看Python环境。TensorFlow环境遵循非常类似的API。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python环境\n",
    "\n",
    "Python实现的环境有`step(action) -> next_time_step`方法可以对环境实施某种动作(action) ，可以返回下一个步骤的如下信息：\n",
    "1. `observation`: 这是代理可以观察的环境状态的一部分，以便在下一步中选择其操作。\n",
    "2. `reward`: 代理可以学习如何在多个步骤中最大化这些奖励的总和。\n",
    "3. `step_type`: 与环境的交互通常是一个序列（sequence）/事件（episode）的一部分。例如:国际象棋中的多步。step_type可以是`FIRST`、`MID`或`LAST`来表示这个时间步是序列中的第一个、中间步骤还是最后一个步骤。\n",
    "4. `discount`: 这是一个浮点数，表示下一个时间步的奖励相对于当前时间步的奖励的权重。\n",
    "\n",
    "它们被分组到一个名为`TimeStep(step_type, reward, discount, observation)`的元组。\n",
    "\n",
    "所有python环境必须实现的接口都在`environments/py_environment.PyEnvironment`里面，主要的方法为："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PyEnvironment(object):\n",
    "    def reset(self):\n",
    "        \"\"\"Return initial_time_step.\"\"\"\n",
    "        self._current_time_step = self._reset()\n",
    "        return self._current_time_step\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"Apply action and return new time_step.\"\"\"\n",
    "        if self._current_time_step is None:\n",
    "            return self.reset()\n",
    "        self._current_time_step = self._step(action)\n",
    "        return self._current_time_step\n",
    "\n",
    "    def current_time_step(self):\n",
    "        return self._current_time_step\n",
    "\n",
    "    def time_step_spec(self):\n",
    "        \"\"\"Return time_step_spec.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def observation_spec(self):\n",
    "        \"\"\"Return observation_spec.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def action_spec(self):\n",
    "        \"\"\"Return action_spec.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def _reset(self):\n",
    "        \"\"\"Return initial_time_step.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def _step(self, action):\n",
    "        \"\"\"Apply action and return new time_step.\"\"\"\n",
    "        self._current_time_step = self._step(action)\n",
    "        return self._current_time_step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "除了`step()`方法之外，环境也提供 `reset()` 方法，他可以启动一个新序列并提供一个初始`TimeStep`。没有必要显式地调用`reset`方法。我们假设环境是自动重置的，当它们到达某个事件的末尾或者第一次调用step()时都是如此。\n",
    "\n",
    "注意，子类并不直接实现`step()` 或`reset()`，而是覆盖了`_step()`和`_reset()` 方法。从这些方法返回的时间步将被缓存并通过`current_time_step()`公开。\n",
    "\n",
    "方法`observation_spec` 和`action_spec` 返回一组`(Bounded)ArraySpecs` ，分别描述观察值和操作的名称、形状、数据类型和范围。\n",
    "\n",
    "在TF-Agents中，我们反复引用套接字，套接字被定义为由列表、元组、命名元组或字典组成的任何树状结构。这些可以任意组合，以维护观察和操作的结构。我们发现这对于有许多观察和操作且更复杂的环境非常有用。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用标准的环境\n",
    "\n",
    "TF Agents 有许多标准环境的内置包装，比如OpenAI Gym、DeepMind-control和Atari，因此它们可以遵循我们的`py_environment.PyEnvironment` 接口。可以使用我们的环境套件轻松地加载这些包装环境。让我们从OpenAI Gym加载CartPole环境，并查看动作和time_step_spec。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action_spec: BoundedArraySpec(shape=(), dtype=dtype('int64'), name='action', minimum=0, maximum=1)\n",
      "time_step_spec.observation: BoundedArraySpec(shape=(4,), dtype=dtype('float32'), name='observation', minimum=[-4.8000002e+00 -3.4028235e+38 -4.1887903e-01 -3.4028235e+38], maximum=[4.8000002e+00 3.4028235e+38 4.1887903e-01 3.4028235e+38])\n",
      "time_step_spec.step_type: ArraySpec(shape=(), dtype=dtype('int32'), name='step_type')\n",
      "time_step_spec.discount: BoundedArraySpec(shape=(), dtype=dtype('float32'), name='discount', minimum=0.0, maximum=1.0)\n",
      "time_step_spec.reward: ArraySpec(shape=(), dtype=dtype('float32'), name='reward')\n"
     ]
    }
   ],
   "source": [
    "environment = suite_gym.load('CartPole-v0')\n",
    "print('action_spec:', environment.action_spec())\n",
    "print('time_step_spec.observation:', environment.time_step_spec().observation)\n",
    "print('time_step_spec.step_type:', environment.time_step_spec().step_type)\n",
    "print('time_step_spec.discount:', environment.time_step_spec().discount)\n",
    "print('time_step_spec.reward:', environment.time_step_spec().reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，我们看到环境期望在中执行的动作为[0,1]之间类型为`int64`数值，并返回`TimeSteps` ，其中观察值是长度为4的' float32 '向量，衰减因子是[0.0,1.0]中 `float32`的浮点数。现在，让我们尝试为整个事件采取一个固定的动作`(1,)`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TimeStep(step_type=array(0, dtype=int32), reward=array(0., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.01572075, -0.00113396,  0.0223387 , -0.02902855], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.01569807,  0.19366063,  0.02175813, -0.3145805 ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.01957129,  0.388466  ,  0.01546652, -0.600323  ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.02734061,  0.5833682 ,  0.00346006, -0.88809437], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.03900797,  0.77844304, -0.01430183, -1.1796876 ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.05457683,  0.97374773, -0.03789558, -1.4768193 ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.07405178,  1.1693115 , -0.06743196, -1.7810929 ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.09743802,  1.365124  , -0.10305382, -2.093954  ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.1247405 ,  1.5611216 , -0.14493291, -2.416635  ], dtype=float32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(1., dtype=float32), discount=array(1., dtype=float32), observation=array([ 0.15596293,  1.757171  , -0.1932656 , -2.7500873 ], dtype=float32))\n",
      "TimeStep(step_type=array(2, dtype=int32), reward=array(1., dtype=float32), discount=array(0., dtype=float32), observation=array([ 0.19110635,  1.9530503 , -0.24826735, -3.094903  ], dtype=float32))\n"
     ]
    }
   ],
   "source": [
    "action = 1\n",
    "time_step = environment.reset()\n",
    "print(time_step)\n",
    "while not time_step.is_last():\n",
    "    time_step = environment.step(action)\n",
    "    print(time_step)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建自己的python环境\n",
    "\n",
    "对于许多用户来说，一个常见的用例是将TF-Agents中的一个标准代理(参见agents/)应用于他们的问题。为了做到这一点，他们必须把他们的问题框定成一个环境。因此，让我们看看如何在Python中实现环境。\n",
    "\n",
    "假设我们想要训练一名代理玩下面的纸牌游戏(受到Black Jack启发):\n",
    "1. 这款游戏使用无限的一副牌，编号为1…10。\n",
    "2. 每轮代理可以做两件事:获得一张新的随机纸牌，或者停止当前的回合。\n",
    "3. 目标是在这一轮结束时，你的牌的编号总数尽可能接近21张，而不超过21张。\n",
    "\n",
    "代表游戏的环境可以是这样的:\n",
    "1. 动作（Actions）:我们有两个动作。动作0:获得新卡，动作1:终止当前回合。\n",
    "2. 观察（Observations）: 当前回合中所有牌的和。\n",
    "3. 回报（Reward）: 我们的目标是在不超越的情况下尽可能接近21，所以我们可以在这一轮结束时使用以下奖励来实现这一目标:\n",
    "    sum_of_cards - 21 if sum_of_cards <= 21, else -21"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CardGameEnv(py_environment.PyEnvironment):\n",
    "    def __init__(self):\n",
    "        self._action_spec = array_spec.BoundedArraySpec(shape=(),\n",
    "                                                        dtype=np.int32,\n",
    "                                                        minimum=0,\n",
    "                                                        maximum=1,\n",
    "                                                        name='action')\n",
    "        self._observation_spec = array_spec.BoundedArraySpec(\n",
    "            shape=(1, ), dtype=np.int32, minimum=0, name='observation')\n",
    "        self._state = 0\n",
    "        self._episode_ended = False\n",
    "\n",
    "    def action_spec(self):\n",
    "        return self._action_spec\n",
    "\n",
    "    def observation_spec(self):\n",
    "        return self._observation_spec\n",
    "\n",
    "    def _reset(self):\n",
    "        self._state = 0\n",
    "        self._episode_ended = False\n",
    "        return ts.restart(np.array([self._state], dtype=np.int32))\n",
    "\n",
    "    def _step(self, action):\n",
    "\n",
    "        if self._episode_ended:\n",
    "            # The last action ended the episode. Ignore the current action and start\n",
    "            # a new episode.\n",
    "            return self.reset()\n",
    "\n",
    "        # Make sure episodes don't go on forever.\n",
    "        if action == 1:\n",
    "            self._episode_ended = True\n",
    "        elif action == 0:\n",
    "            new_card = np.random.randint(1, 11)\n",
    "            self._state += new_card\n",
    "        else:\n",
    "            raise ValueError('`action` should be 0 or 1.')\n",
    "\n",
    "        if self._episode_ended or self._state >= 21:\n",
    "            reward = self._state - 21 if self._state <= 21 else -21\n",
    "            return ts.termination(np.array([self._state], dtype=np.int32),\n",
    "                                  reward)\n",
    "        else:\n",
    "            return ts.transition(np.array([self._state], dtype=np.int32),\n",
    "                                 reward=0.0,\n",
    "                                 discount=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "让我们确保正确地定义了上面的环境。在创建自己的环境时，必须确保生成的观察结果和time_steps符合规范中定义的正确形状和类型。这些是用来生成TensorFlow图的，因此如果我们做错了，就会产生难以调试的问题。\n",
    "\n",
    "为了验证我们的环境，我们将使用一个随机策略来生成动作，并且我们将迭代5个事件以确保事情按照预期进行。如果接收到的time_step不符合环境规范，则会引发错误。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "environment = CardGameEnv()\n",
    "utils.validate_py_environment(environment, episodes=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们知道环境正在正常工作，让我们使用一个固定的策略来运行这个环境:请求3张牌，然后结束这一轮。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TimeStep(step_type=array(0, dtype=int32), reward=array(0., dtype=float32), discount=array(1., dtype=float32), observation=array([0], dtype=int32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(0., dtype=float32), discount=array(1., dtype=float32), observation=array([3], dtype=int32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(0., dtype=float32), discount=array(1., dtype=float32), observation=array([4], dtype=int32))\n",
      "TimeStep(step_type=array(1, dtype=int32), reward=array(0., dtype=float32), discount=array(1., dtype=float32), observation=array([12], dtype=int32))\n",
      "TimeStep(step_type=array(2, dtype=int32), reward=array(-9., dtype=float32), discount=array(0., dtype=float32), observation=array([12], dtype=int32))\n",
      "Final Reward =  -9.0\n"
     ]
    }
   ],
   "source": [
    "get_new_card_action = 0\n",
    "end_round_action = 1\n",
    "\n",
    "environment = CardGameEnv()\n",
    "time_step = environment.reset()\n",
    "print(time_step)\n",
    "cumulative_reward = time_step.reward\n",
    "\n",
    "for _ in range(3):\n",
    "    time_step = environment.step(get_new_card_action)\n",
    "    print(time_step)\n",
    "    cumulative_reward += time_step.reward\n",
    "\n",
    "time_step = environment.step(end_round_action)\n",
    "print(time_step)\n",
    "cumulative_reward += time_step.reward\n",
    "print('Final Reward = ', cumulative_reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 环境包装\n",
    "\n",
    "环境包装器接受python环境并返回环境的修改版本。原始环境和修改后的环境都是`py_environment.PyEnvironment`的实例。而且多个包装器可以链接在一起。\n",
    "\n",
    "一些常见的包装器可以在`environments/wrappers.py`中找到。例如:\n",
    "1. `ActionDiscretizeWrapper`: 将连续的动作空间转换为离散的动作空间。\n",
    "2. `RunStats`: 获取环境的运行统计数据，如所采取的步骤数、完成的事件数等。\n",
    "3. `TimeLimit`: 在固定数量的步骤之后终止该事件。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 示例 1: Action Discretize Wrapper\n",
    "\n",
    "倒立摆是一个PyBullet环境，它接受`[-1, 1]`范围内的连续动作。如果我们想要在这种环境下训练一个离散的动作代理，比如DQN，我们必须对动作空间进行离散化(量化)。这正是`ActionDiscretizeWrapper`所做的。比较`action_spec` 之前和之后的包装:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Action Spec: BoundedArraySpec(shape=(1,), dtype=dtype('float32'), name='action', minimum=-2.0, maximum=2.0)\n",
      "Discretized Action Spec: BoundedArraySpec(shape=(1,), dtype=dtype('int32'), name='action', minimum=0, maximum=[4])\n"
     ]
    }
   ],
   "source": [
    "env = suite_gym.load('Pendulum-v0')\n",
    "print('Action Spec:', env.action_spec())\n",
    "\n",
    "discrete_action_env = wrappers.ActionDiscretizeWrapper(env, num_actions=5)\n",
    "print('Discretized Action Spec:', discrete_action_env.action_spec())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "包装的 `discrete_action_env`是 `py_environment.PyEnvironment` 的一个实例，可以像普通的python环境一样对待。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensorflow的环境\n",
    "\n",
    "TF环境的接口在`environments/tf_environment.TFEnvironment` 中定义，它看起来与Python环境非常相似。TF环境与python envs在以下几个方面不同:\n",
    "\n",
    "* 它生成的是张量对象而不是数组\n",
    "* TF环境为生成的张量添加一个批处理维度(与specs相比)。\n",
    "\n",
    "将python环境转换为TFEnvs允许tensorflow对操作进行优化。例如，可以定义一个`collect_experience_op` ，它从环境中收集数据并添加到`replay_buffer`,中，以及一个`train_op` ，它从`replay_buffer` 中读取数据并训练代理，然后在TensorFlow中自然地并行运行它们。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TFEnvironment(object):\n",
    "    def time_step_spec(self):\n",
    "        \"\"\"Describes the `TimeStep` tensors returned by `step()`.\"\"\"\n",
    "    def observation_spec(self):\n",
    "        \"\"\"Defines the `TensorSpec` of observations provided by the environment.\"\"\"\n",
    "    def action_spec(self):\n",
    "        \"\"\"Describes the TensorSpecs of the action expected by `step(action)`.\"\"\"\n",
    "    def reset(self):\n",
    "        \"\"\"Returns the current `TimeStep` after resetting the Environment.\"\"\"\n",
    "        return self._reset()\n",
    "\n",
    "    def current_time_step(self):\n",
    "        \"\"\"Returns the current `TimeStep`.\"\"\"\n",
    "        return self._current_time_step()\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"Applies the action and returns the new `TimeStep`.\"\"\"\n",
    "        return self._step(action)\n",
    "\n",
    "    @abc.abstractmethod\n",
    "    def _reset(self):\n",
    "        \"\"\"Returns the current `TimeStep` after resetting the Environment.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def _current_time_step(self):\n",
    "        \"\"\"Returns the current `TimeStep`.\"\"\"\n",
    "    @abc.abstractmethod\n",
    "    def _step(self, action):\n",
    "        \"\"\"Applies the action and returns the new `TimeStep`.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`current_time_step()`方法返回当前的time_step并在需要时初始化环境。\n",
    "\n",
    "`reset()` 方法强制环境中的重置并返回current_step。\n",
    "\n",
    "如果“动作”不依赖于前面的 `time_step`，在“Graph”模式中需要 `tf.control_dependency` 。\n",
    "\n",
    "现在，让我们看看如何创建`TFEnvironments` 。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建你自己的 TensorFlow 环境\n",
    "\n",
    "这比在Python中创建环境要复杂得多，所以我们不会在这个主题中讨论它。[例子](https://github.com/tensorflow/agents/blob/master/tf_agents/environments/tf_environment_test.py)。\n",
    "更常见的用例是用Python实现环境，并使用`TFPyEnvironment` 包装器(参见下面)将其封装在TensorFlow中。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 将Python环境包装为Tensorflow\n",
    "\n",
    "我们可以使用`TFPyEnvironment` 包装器轻松地将任何Python环境包装到TensorFlow环境中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "TimeStep Specs: TimeStep(step_type=TensorSpec(shape=(), dtype=tf.int32, name='step_type'), reward=TensorSpec(shape=(), dtype=tf.float32, name='reward'), discount=BoundedTensorSpec(shape=(), dtype=tf.float32, name='discount', minimum=array(0., dtype=float32), maximum=array(1., dtype=float32)), observation=BoundedTensorSpec(shape=(4,), dtype=tf.float32, name='observation', minimum=array([-4.8000002e+00, -3.4028235e+38, -4.1887903e-01, -3.4028235e+38],\n",
      "      dtype=float32), maximum=array([4.8000002e+00, 3.4028235e+38, 4.1887903e-01, 3.4028235e+38],\n",
      "      dtype=float32)))\n",
      "Action Specs: BoundedTensorSpec(shape=(), dtype=tf.int64, name='action', minimum=array(0), maximum=array(1))\n"
     ]
    }
   ],
   "source": [
    "env = suite_gym.load('CartPole-v0')\n",
    "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
    "\n",
    "print(isinstance(tf_env, tf_environment.TFEnvironment))\n",
    "print(\"TimeStep Specs:\", tf_env.time_step_spec())\n",
    "print(\"Action Specs:\", tf_env.action_spec())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "注意，现在的配置类型是:`(Bounded)TensorSpec`。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 实用例子"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 简单的例子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[TimeStep(step_type=array([0], dtype=int32), reward=array([0.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.02460271,  0.02203766, -0.01728856, -0.01298076]],\n",
      "      dtype=float32)), array([0], dtype=int32), TimeStep(step_type=array([1], dtype=int32), reward=array([1.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.02416196, -0.17283215, -0.01754818,  0.27419767]],\n",
      "      dtype=float32))]\n",
      "[TimeStep(step_type=array([1], dtype=int32), reward=array([1.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.02416196, -0.17283215, -0.01754818,  0.27419767]],\n",
      "      dtype=float32)), array([1], dtype=int32), TimeStep(step_type=array([1], dtype=int32), reward=array([1.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.0276186 ,  0.02253573, -0.01206422, -0.02396793]],\n",
      "      dtype=float32))]\n",
      "[TimeStep(step_type=array([1], dtype=int32), reward=array([1.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.0276186 ,  0.02253573, -0.01206422, -0.02396793]],\n",
      "      dtype=float32)), array([0], dtype=int32), TimeStep(step_type=array([1], dtype=int32), reward=array([1.], dtype=float32), discount=array([1.], dtype=float32), observation=array([[-0.02716789, -0.17241114, -0.01254358,  0.26488432]],\n",
      "      dtype=float32))]\n",
      "Total reward: [3.]\n"
     ]
    }
   ],
   "source": [
    "env = suite_gym.load('CartPole-v0')\n",
    "\n",
    "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
    "# reset() creates the initial time_step after resetting the environment.\n",
    "time_step = tf_env.reset()\n",
    "num_steps = 3\n",
    "transitions = []\n",
    "reward = 0\n",
    "for i in range(num_steps):\n",
    "    action = tf.constant([i % 2])\n",
    "    # applies the action and returns the new TimeStep.\n",
    "    next_time_step = tf_env.step(action)\n",
    "    transitions.append([time_step, action, next_time_step])\n",
    "    reward += next_time_step.reward\n",
    "    time_step = next_time_step\n",
    "\n",
    "np_transitions = tf.nest.map_structure(lambda x: x.numpy(), transitions)\n",
    "print('\\n'.join(map(str, np_transitions)))\n",
    "print('Total reward:', reward.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 整个事件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_episodes: 5 num_steps: 91\n",
      "avg_length 18.2 avg_reward: 18.2\n"
     ]
    }
   ],
   "source": [
    "env = suite_gym.load('CartPole-v0')\n",
    "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
    "\n",
    "time_step = tf_env.reset()\n",
    "rewards = []\n",
    "steps = []\n",
    "num_episodes = 5\n",
    "\n",
    "for _ in range(num_episodes):\n",
    "    episode_reward = 0\n",
    "    episode_steps = 0\n",
    "    while not time_step.is_last():\n",
    "        action = tf.random.uniform([1], 0, 2, dtype=tf.int32)\n",
    "        time_step = tf_env.step(action)\n",
    "        episode_steps += 1\n",
    "        episode_reward += time_step.reward.numpy()\n",
    "    rewards.append(episode_reward)\n",
    "    steps.append(episode_steps)\n",
    "    time_step = tf_env.reset()\n",
    "\n",
    "num_steps = np.sum(steps)\n",
    "avg_length = np.mean(steps)\n",
    "avg_reward = np.mean(rewards)\n",
    "\n",
    "print('num_episodes:', num_episodes, 'num_steps:', num_steps)\n",
    "print('avg_length', avg_length, 'avg_reward:', avg_reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": true,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}