策略梯度算法.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7C0480D97F344544BC2824225ADBF12D",
    "jupyter": {},
    "mdEditEnable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# 策略梯度算法\n",
    "## 简介\n",
    "之前我们介绍的Q-learning和DQN及改进算法都是基于值函数（value-based）的方法，其中Q-learning是处理有限状态的算法，而DQN可以用来解决连续状态的问题。在强化学习中，除了基于值函数的方法，还有一支非常经典的方法，那就是基于策略（policy-based）的方法。对比两者，基于值函数的方法主要是学习值函数，然后根据值函数导出一个策略，此时并不存在一个显式的策略。而基于策略的方法则是直接显式的学习一个目标策略。策略梯度是基于策略的方法的基础，我们将从策略梯度算法说起。\n",
    "## 策略梯度\n",
    "基于策略的方法首先需要参数化策略，我们假设目标策略$\\pi_\\theta$是一个随机性策略，并且处处可微，其中$\\theta$是对应的参数。我们可以用一个线性模型或者神经网络模型来建模这样一个策略函数，输入是某个状态，然后输出一个动作的概率分布。我们的目标是要寻找一个最优策略，来最大化这个策略在环境中的期望回报。我们将策略学习的目标函数定义为\n",
    "$$\n",
    "J(\\theta)= \\mathbb{E}_{s_0}[V^{\\pi_\\theta}(s_0)]\n",
    "$$\n",
    "其中$s_0$表示初始状态。现在有了目标函数，我们将目标函数对策略$\\theta$求导，得到导数后，我们就可以用梯度上升方法来最大化这个目标函数从而得到最优策略。\n",
    "\n",
    "我们之前在MDP章节中学习过在策略$\\pi$下的状态访问分布，我们在此用$\\nu^{\\pi}$表示。然后我们对目标函数求梯度，可以得到如下式子，更详细的推导将在扩展阅读中给出。\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\nabla_{\\theta}J(\\theta) &\\propto \\sum_{s \\in S}\\nu^{\\pi_\\theta}(s)\\sum_{a \\in A}Q^{\\pi_\\theta}(s,a)\\nabla_{\\theta}\\pi_{\\theta}(a|s) \\nonumber \\\\\n",
    "&=\\sum_{s \\in S}\\nu^{\\pi_\\theta}(s)\\sum_{a \\in A}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a)\\frac{\\nabla_{\\theta}\\pi_{\\theta}(a|s)}{\\pi_{\\theta}(a|s)} \\nonumber \\\\\n",
    "&= \\mathbb{E}_{\\pi_\\theta}[Q^{\\pi_\\theta}(s,a)\\nabla_{\\theta}\\log \\pi_{\\theta}(a|s)] \\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "于是，我们就可以用这个梯度来更新策略。需要注意的是，因为上式期望的下标是$\\pi_\\theta$，所以策略梯度算法为在线学习算法，即必须使用策略$\\pi_\\theta$采样得到的数据来计算梯度。更一般地，我们可以把梯度写成下面这个形式：\n",
    "$$\n",
    "g = \\mathbb{E}_{\\pi_\\theta}[\\sum^{\\infty}_{t=0}\\psi_{t}\\nabla_{\\theta}\\log \\pi_{\\theta}(a_{t}|s_{t})]\n",
    "$$\n",
    "其中$\\psi_{t}$可以有很多种形式：\n",
    "$$\n",
    "\\begin{align}\n",
    "&1.\\sum_{t‘=0}^{\\infty}\\gamma^t r_{t’} : 轨迹的总回报  &&4.Q^{\\pi_\\theta}(s_{t},a_{t}) : 动作价值函数 \\\\\n",
    "&2.\\sum_{t'=t}^{\\infty} \\gamma^{t'-t} r_{t'} : 动作 a_{t}之后的回报 &&5.A^{\\pi_\\theta}(s_{t},a_{t}): 优势函数 \\\\\n",
    "&3.\\sum_{t'=t}^{\\infty}r_{t'}-b(s_{t}) : 基线版本的改进 \\quad\\quad \\quad\\quad &&6.r_{t} + V^{\\pi_\\theta}(s_{t+1}) - V^{\\pi_\\theta}(s_t) : 时序差分残差\n",
    "\\end{align}\n",
    "$$\n",
    "至此，我们已经了解了策略梯度的目标与形式，那么我们今天要讲的REINFORCE算法的流程是怎样的呢？\n",
    "## REINFORCE 算法\n",
    "其实REINFORCE就是在上文中当$\\psi=\\sum_{t'=t}^{\\infty}\\gamma^{t'-t}r_{t'}$时的策略梯度算法。他的具体算法流程如下：\n",
    "\n",
    "- 初始化策略参数$\\theta$\n",
    "- 不断进行如下循环（每个循环是一条序列）：\n",
    "    - 用当前策略 $\\pi_\\theta$ 采样轨迹 $\\{s_{1},a_{1},r_{1},s_{2},a_{2},r_{2} ... \\}$\n",
    "\t- 计算当前轨迹每个时刻往后的回报$\\sum_{t'=t}^{\\infty}\\gamma^{t'-t}r_{t'}$作为$\\psi_{t}$\n",
    "\t- 对$\\theta$进行更新 $\\theta = \\theta + \\alpha \\sum_t \\psi_{t}\\nabla_{\\theta}\\log \\pi_{\\theta}(a_{t}|s_{t})$\n",
    "\n",
    "\n",
    "这便是REINFORCE算法的全部了，让我们来用代码实现它看看效果如何吧！\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "880649899BDF4E868E82960D05260E37",
    "jupyter": {},
    "mdEditEnable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## REINFORCE代码实践\n",
    "我们在Cartpole环境中进行REINFORCE算法的实验。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "40B616B26E8B4AA48E98B4258C7988CA",
    "jupyter": {},
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import rl_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7C0E3A3A844549BC905ED109D604EF92",
    "jupyter": {},
    "mdEditEnable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "定义我们的策略网络PolicyNet，输入是状态，输出则是该状态下的动作分布。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "864658B0470F4AB4823322580ED41CDA",
    "jupyter": {},
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class PolicyNet(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super(PolicyNet, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        return  F.softmax(self.fc2(x),dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6A19E6CFCEDC467BA350D1200C7E28B6",
    "jupyter": {},
    "mdEditEnable": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "再定义我们的Reinforce算法。在函数take_action中，我们通过概率对离散的动作进行采样。在更新过程中，我们按照算法，将损失函数写为$-\\sum_t \\psi_{t}\\nabla_{\\theta}\\log \\pi_{\\theta}(a_{t}|s_{t})$，对$\\theta$求导后就可以通过梯度下降更新策略。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "677551348C8D4AAE8AE44C4B46AEFA26",
    "jupyter": {},
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Reinforce:\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma, device):\n",
    "        self.policy_net = PolicyNet(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.optimizer= torch.optim.Adam(self.policy_net.parameters(), lr=learning_rate) # 使用Adam优化器\n",
    "        self.gamma = gamma # 折扣因子\n",
    "\n",
    "    def take_action(self, state): # 根据动作概率分布随机采样\n",
    "        state = torch.tensor([state], dtype=torch.float)\n",
    "        probs = self.policy_net(state)\n",
    "        action_dist = torch.distributions.Categorical(probs)\n",
    "        action = action_dist.sample()\n",
    "        return action.item()\n",
    "\n",
    "    def update(self, transition_dict):\n",
    "        reward_list = transition_dict['rewards']\n",
    "        state_list = transition_dict['states']\n",
    "        action_list = transition_dict['actions']\n",
    "\n",
    "        G = 0\n",
    "        self.optimizer.zero_grad()\n",
    "        for i in reversed(range(len(reward_list))): # 从最后一步算起\n",
    "            reward = reward_list[i]\n",
    "            state = torch.tensor([state_list[i]], dtype=torch.float)\n",
    "            action = torch.tensor([action_list[i]]).view(-1, 1)\n",
    "            log_prob = torch.log(self.policy_net(state).gather(1, action))\n",
    "            G = self.gamma * G + reward\n",
    "            loss = - log_prob * G # 每一步的损失函数\n",
    "            loss.backward() # 反向传播计算梯度\n",
    "        self.optimizer.step() # 梯度下降"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F6DDA156DE31493CA22217FA65E3C43E",
    "jupyter": {},
    "mdEditEnable": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "定义好策略，我们就可以开始实验了，看看Reinforce算法在Cartpole环境上表现如何吧！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "A9993AF344AE4770896E0504E77BAF08",
    "jupyter": {},
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration 0: 100%|██████████| 100/100 [00:04<00:00, 23.88it/s, episode=100, return=55.500]\n",
      "Iteration 1: 100%|██████████| 100/100 [00:08<00:00, 10.45it/s, episode=200, return=75.300]\n",
      "Iteration 2: 100%|██████████| 100/100 [00:16<00:00,  4.75it/s, episode=300, return=178.800]\n",
      "Iteration 3: 100%|██████████| 100/100 [00:20<00:00,  4.90it/s, episode=400, return=164.600]\n",
      "Iteration 4: 100%|██████████| 100/100 [00:21<00:00,  4.58it/s, episode=500, return=156.500]\n",
      "Iteration 5: 100%|██████████| 100/100 [00:21<00:00,  4.73it/s, episode=600, return=187.400]\n",
      "Iteration 6: 100%|██████████| 100/100 [00:22<00:00,  4.40it/s, episode=700, return=194.500]\n",
      "Iteration 7: 100%|██████████| 100/100 [00:23<00:00,  4.24it/s, episode=800, return=200.000]\n",
      "Iteration 8: 100%|██████████| 100/100 [00:23<00:00,  4.33it/s, episode=900, return=200.000]\n",
      "Iteration 9: 100%|██████████| 100/100 [00:22<00:00,  4.14it/s, episode=1000, return=186.100]\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 1e-3\n",
    "num_episodes = 1000\n",
    "hidden_dim = 128\n",
    "gamma = 0.98\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "env_name = \"CartPole-v0\"\n",
    "env = gym.make(env_name)\n",
    "env.seed(0)\n",
    "torch.manual_seed(0)\n",
    "state_dim = env.observation_space.shape[0]\n",
    "action_dim = env.action_space.n\n",
    "agent = Reinforce(state_dim, hidden_dim, action_dim, learning_rate, gamma, device)\n",
    "\n",
    "return_list = []\n",
    "for i in range(10):\n",
    "    with tqdm(total=int(num_episodes/10), desc='Iteration %d' % i) as pbar:\n",
    "        for i_episode in range(int(num_episodes/10)):\n",
    "            episode_return = 0\n",
    "            transition_dict = {'states': [], 'actions': [], 'next_states': [], 'rewards': [], 'dones': []}\n",
    "            state = env.reset()\n",
    "            done = False\n",
    "            while not done:\n",
    "                action = agent.take_action(state)\n",
    "                next_state, reward, done, _ = env.step(action)\n",
    "                transition_dict['states'].append(state)\n",
    "                transition_dict['actions'].append(action)\n",
    "                transition_dict['next_states'].append(next_state)\n",
    "                transition_dict['rewards'].append(reward)\n",
    "                transition_dict['dones'].append(done)\n",
    "                state = next_state\n",
    "                episode_return += reward\n",
    "            return_list.append(episode_return)\n",
    "            agent.update(transition_dict)\n",
    "            if (i_episode+1) % 10 == 0:\n",
    "                pbar.set_postfix({'episode': '%d' % (num_episodes/10 * i + i_episode+1), 'return': '%.3f' % np.mean(return_list[-10:])})\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CB975960A51D4D97A708F800DAEB03EE",
    "jupyter": {},
    "mdEditEnable": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "在CartPole-v0环境中，满分就是200分，我们发现Reinforce算法效果很好，可以达到200分。接下来我们画一下训练过程中每一条轨迹的回报变化图。由于抖动比较大，我们又进行了光滑处理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "D2B1AD392B66419791F4FD6CB800A4A9",
    "jupyter": {},
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://staticcdn.boyuai.com/user-assets/358/oqLvJ1QyqxTnag45S5g49v/qwak896dq3.png!png\">"
      ],
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://staticcdn.boyuai.com/user-assets/358/aA8HRdkJdShwh4v9qWeFBb/qwak89inze.png!png\">"
      ],
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "episodes_list = list(range(len(return_list)))\n",
    "plt.plot(episodes_list,return_list)\n",
    "plt.xlabel('Episodes')   \n",
    "plt.ylabel('Returns')  \n",
    "plt.title('Reinforce on {}'.format(env_name))  \n",
    "plt.show()\n",
    "\n",
    "mv_return = rl_utils.moving_average(return_list, 9)\n",
    "plt.plot(episodes_list, mv_return)\n",
    "plt.xlabel('Episodes')   \n",
    "plt.ylabel('Returns')  \n",
    "plt.title('Reinforce on {}'.format(env_name))  \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C7BF78B87D294A038E9BEA15C75E8940",
    "jupyter": {},
    "mdEditEnable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## 总结\n",
    "REINFORCE算法理论上是能保证局部最优的。REINFORCE算法实际上是借助于蒙特卡洛方法采样轨迹估计动作价值，这种做法的一个优点是梯度是无偏的。但是同样由于因为是蒙特卡洛方法，导致REINFORCE算法梯度估计的方差很大，从而可能会造成一定程度上的不稳定，这也是接下来将介绍的Actor-Critic算法要解决的问题。\n",
    "\n",
    "## 拓展阅读：策略梯度证明\n",
    "我们要证明$\\nabla_{\\theta}J(\\theta) \\propto \\sum_{s \\in S}\\nu^{\\pi}(s)\\sum_{a \\in A}Q^{\\pi_\\theta}(s,a)\\nabla_{\\theta}\\pi_{\\theta}(a|s)$\n",
    "                        \n",
    "先从状态价值函数的推导开始：\n",
    "$$\n",
    "\\begin{align}\n",
    "\\nabla_{\\theta}V^{\\pi_\\theta}(s) &=\\nabla_{\\theta}(\\sum_{a \\in A} \\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a)) \\nonumber\\\\\n",
    "&=\\sum_{a\\in A}(\\nabla_{\\theta}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a) + \\pi_{\\theta}(a|s)\\nabla_{\\theta}Q^{\\pi_\\theta}(s,a))\\nonumber\\\\\n",
    "&=\\sum_{a\\in A}(\\nabla_{\\theta}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a) + \\pi_{\\theta}(a|s)\\nabla_{\\theta}\\sum_{s',r}p(s',r|s,a)(r+V^{\\pi_\\theta}(s'))\\nonumber\\\\\n",
    "&=\\sum_{a\\in A}(\\nabla_{\\theta}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a) + \\pi_{\\theta}(a|s)\\sum_{s',r}p(s',r|s,a)\\nabla_{\\theta}V^{\\pi_\\theta}(s'))\\nonumber\\\\\n",
    "&=\\sum_{a\\in A}(\\nabla_{\\theta}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a) + \\pi_{\\theta}(a|s)\\sum_{s'}p(s'|s,a)\\nabla_{\\theta}V^{\\pi_\\theta}(s'))\\nonumber\\\\\n",
    "\\end{align}\n",
    "$$\n",
    "为了简化表示，我们让$\\phi(s)=\\sum_{a \\in A}\\nabla_{\\theta}\\pi_{\\theta}(a|s)Q^{\\pi_\\theta}(s,a)$, 定义$d^{\\pi_\\theta}(s\\rightarrow x, k)$为策略$\\pi$从状态$s$出发$k$步后到达状态$x$的概率，我们继续推导:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\nabla_{\\theta}V^{\\pi_\\theta}(s) &= \\phi(s) + \\sum_{a}\\pi_{\\theta}(a|s)\\sum_{s'}P(s'|s,a)\\nabla_{\\theta}V^{\\pi_\\theta}(s')\\nonumber\\\\\n",
    "&= \\phi(s) + \\sum_{a}\\sum_{s'}\\pi_{\\theta}(a|s)P(s'|s,a)\\nabla_{\\theta}V^{\\pi_\\theta}(s')\\nonumber\\\\\n",
    "&= \\phi(s) + \\sum_{s'}d^{\\pi_\\theta}(s \\rightarrow s',1)\\nabla_{\\theta}V^{\\pi_\\theta}(s')\\nonumber\\\\\n",
    "&= \\phi(s) + \\sum_{s'}d^{\\pi_\\theta}(s \\rightarrow s',1)[\\phi(s') + \\sum_{s''}d^{\\pi_\\theta}(s' \\rightarrow s'',1)\\nabla_{\\theta}V^{\\pi_\\theta}(s'')]\\nonumber\\\\\n",
    "&= \\phi(s) + \\sum_{s'}d^{\\pi_\\theta}(s \\rightarrow s',1)\\phi(s') + \\sum_{s''}d^{\\pi_\\theta}(s \\rightarrow s'',2)\\nabla_{\\theta}V^{\\pi_\\theta}(s'')\\nonumber\\\\\n",
    "&= \\phi(s) + \\sum_{s'}d^{\\pi_\\theta}(s \\rightarrow s',1)\\phi(s') +\\sum_{s''}d^{\\pi_\\theta}(s' \\rightarrow s'',2)\\phi(s'') + \\sum_{s'''}d^{\\pi_\\theta}(s \\rightarrow s''',3)\\nabla_{\\theta}V^{\\pi_\\theta}(s''')\\nonumber\\\\\n",
    "&= ......\\nonumber\\\\\n",
    "&= \\sum_{x \\in S}\\sum^{\\infty}_{k=0}d^{\\pi_\\theta}(s \\rightarrow x, k)\\phi(x)\\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "我们定义$\\eta(s)= \\mathbb{E}_{s_0}[\\sum^{\\infty}_{k=0}d^{\\pi_\\theta}(s_{0} \\rightarrow s, k)]$。至此我们回到我们的目标函数：\n",
    "$$\n",
    "\\begin{align}\n",
    "\\nabla_{\\theta}J(\\theta) &= \\nabla_{\\theta}\\mathbb{E}_{s_0}[V^{\\pi_\\theta}(s_{0})]\\nonumber\\\\\n",
    "&= \\sum_{s} \\mathbb{E}_{s_0}[ \\sum^{\\infty}_{k=0}d^{\\pi_\\theta}(s_{0} \\rightarrow s, k)]\\phi(s)\\nonumber\\\\\n",
    "&= \\sum_{s}\\eta(s)\\phi(s)\\nonumber\\\\\n",
    "&= (\\sum_{s}\\eta(s))\\sum_{s}\\frac{\\eta(s)}{\\sum_{s}\\eta(s)}\\phi(s)\\nonumber\\\\\n",
    "&\\propto \\sum_{s}\\frac{\\eta(s)}{\\sum_{s}\\eta(s)}\\phi(s)\\nonumber\\\\\n",
    "&= \\sum_{s}\\nu^{\\pi}(s)\\sum_{a}Q^{\\pi_\\theta}(s,a)\\nabla_{\\theta}\\pi_{\\theta}(a|s)\\nonumber\n",
    "\\end{align}\n",
    "$$\n",
    "证明完毕！"
   ]
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}