{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef80017f",
   "metadata": {},
   "source": [
    "# C10 演员评论家方法\n",
    "\n",
    "## 最基础的AC算法：QAC\n",
    "\n",
    "Actor网络：策略梯度上升\n",
    "\n",
    "直接更新策略参数，以采取更优的动作。\n",
    "\n",
    "我们希望最优化一个标量指标$J(\\theta)$，他的梯度包含一个期望值：\n",
    "\n",
    "$$\\nabla_{\\theta}J(\\theta)=\\mathbb{E}_{S\\sim\\eta,A\\sim\\pi}[\\nabla_{\\theta}\\ln \\pi(A|S,\\theta_{t})q_{\\pi}(S,A)]$$\n",
    "\n",
    "由于实际工程中很难求出这个数学期望的解析解，因此采用随机的采样方法进行梯度上升的近似更新：\n",
    "\n",
    "$$\\theta_{t+1}=\\theta_{t}+\\alpha_{\\theta}\\nabla_{\\theta}\\ln \\pi(a_{t}|s_{t},\\theta_{t})q(s_{t},a_{t},w_{t})$$\n",
    "\n",
    "如果某个动作在状态 $s_t$ 下得到了很高的 $q$ 值，梯度更新就会使得网络参数 $\\theta$ 发生改变，从而提高未来在相同状态下采取动作 $a_t$ 的概率 $\\pi$。这就是 Actor 学习的过程。\n",
    "\n",
    "Critic网络：价值的时序差分更新\n",
    "\n",
    "在 Actor 的更新公式中，我们需要知道动作价值 $q_t$ 。如果采用蒙特卡洛方法（走完一整个回合再算），那就是 REINFORCE 算法 。但为了实现更高效的单步在线学习，我们使用时序差分（TD）学习来估计这个值，这就是 Actor-Critic 的精髓。\n",
    "\n",
    "Critic 的任务是通过评估动作的价值来“批评” Actor 。在最基础的 QAC 算法中，它使用类似 Sarsa 的方式更新自己的参数 $w$：\n",
    "\n",
    "$$w_{t+1}=w_{t}+\\alpha_{w}[r_{t+1}+\\gamma q(s_{t+1},a_{t+1},w_{t})-q(s_{t},a_{t},w_{t})]\\nabla_{w}q(s_{t},a_{t},w_{t})$$。\n",
    "\n",
    "公式方括号内的部分 $[r_{t+1}+\\gamma q(s_{t+1},a_{t+1},w_{t})-q(s_{t},a_{t},w_{t})]$ 就是著名的 TD 误差。它衡量了“当前预测的 $q$ 值”与“实际得到的奖励加上下一步预测的 $q$ 值”之间的差距。\n",
    "\n",
    "算法流程：\n",
    "\n",
    "1. 策略执行（Actor 正向计算）： 系统当前处于状态 $s_t$，Actor 依据当前的概率分布 $\\pi(a|s_{t},\\theta_{t})$ 采样出一个动作$a_t$并执行，获得真实物理系统的反馈奖励 $r_{t+1}$ 和下一个状态 $s_{t+1}$ 。\n",
    "2. 价值评估（Critic 反向更新）： 系统依据策略再预演一步动作 $a_{t+1}$，以此计算 TD 误差，并使用梯度下降法更新打分器参数 $w_{t+1}$。\n",
    "3. 策略优化（Actor 反向更新）： 拿着刚刚打出的分数 $q(s_{t},a_{t},w_{t})$，沿着梯度上升的方向更新策略参数 $\\theta_{t+1}$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c2f49de7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch 版本: 2.5.1+cu121\n",
      "GPU 是否可用: True\n",
      "当前显卡: NVIDIA GeForce RTX 5060 Ti\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# 打印 PyTorch 版本\n",
    "print(f\"PyTorch 版本: {torch.__version__}\")\n",
    "# 检查 GPU 是否可用 (如果输出 True，说明大功告成！)\n",
    "print(f\"GPU 是否可用: {torch.cuda.is_available()}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    # 打印当前使用的显卡型号\n",
    "    print(f\"当前显卡: {torch.cuda.get_device_name(0)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "79226b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torch.distributions as distributions\n",
    "\n",
    "# 定义超参数\n",
    "# 学习率\n",
    "ACTOR_LR = 0.001\n",
    "CRITIC_LR = 0.002\n",
    "# 折扣因子\n",
    "GAMMA = 0.99\n",
    "# 训练迭代次数\n",
    "NUM_EPISODES = 1000\n",
    "\n",
    "# 1. 定义 Actor 网络 (演员：输出策略)\n",
    "class Actor(nn.Module):\n",
    "    def __init__(self, state_dim, action_dim):\n",
    "        super(Actor, self).__init__()\n",
    "        # 定义两层全连接网络提取状态特征\n",
    "        self.fc1 = nn.Linear(state_dim, 128)\n",
    "        self.fc2 = nn.Linear(128, action_dim)\n",
    "        \n",
    "    def forward(self, state):\n",
    "        # 使用 ReLU 激活函数\n",
    "        x = F.relu(self.fc1(state))\n",
    "        # 使用 Softmax 输出离散动作的概率分布\n",
    "        action_probs = F.softmax(self.fc2(x), dim=-1)\n",
    "        return action_probs\n",
    "\n",
    "# 2. 定义 Critic 网络 (评论家：输出动作价值 Q 值)\n",
    "class Critic(nn.Module):\n",
    "    def __init__(self, state_dim, action_dim):\n",
    "        super(Critic, self).__init__()\n",
    "        # 定义全连接网络，输入是状态，输出是各个动作的 Q 值\n",
    "        self.fc1 = nn.Linear(state_dim, 128)\n",
    "        self.fc2 = nn.Linear(128, action_dim)\n",
    "        \n",
    "    def forward(self, state):\n",
    "        x = F.relu(self.fc1(state))\n",
    "        q_values = self.fc2(x)\n",
    "        return q_values\n",
    "\n",
    "\n",
    "# 3. 核心训练逻辑 (对应 Algorithm 10.1)\n",
    "def train_step(actor, critic, actor_optimizer, critic_optimizer, \n",
    "               state, action, reward, next_state, next_action, done):\n",
    "    \"\"\"\n",
    "    执行一步 QAC 算法的参数更新\n",
    "    \"\"\"\n",
    "    # 转换数据格式为 Tensor\n",
    "    state = torch.FloatTensor(state)\n",
    "    next_state = torch.FloatTensor(next_state)\n",
    "    reward = torch.FloatTensor([reward])\n",
    "    \n",
    "    # -----------------------------------------\n",
    "    # Critic 更新 (价值更新)\n",
    "    # -----------------------------------------\n",
    "    # 计算当前状态动作的 Q(s_t, a_t, w_t)\n",
    "    q_values = critic(state)        # 选取某个state对应的所有动作的Q值\n",
    "    current_q = q_values[action]    # 选取当前动作对应的Q值\n",
    "    \n",
    "    # 计算下一状态动作的 Q(s_{t+1}, a_{t+1}, w_t)\n",
    "    # 使用 .detach() 截断梯度，因为目标值不需要传递梯度回网络\n",
    "    next_q_values = critic(next_state).detach()\n",
    "    next_q = next_q_values[next_action]\n",
    "    \n",
    "    # 计算 TD 目标：r_{t+1} + gamma * Q(s_{t+1}, a_{t+1}) (如果是终止状态则没有下一步的Q)\n",
    "    td_target = reward + GAMMA * next_q * (1 - int(done))\n",
    "    \n",
    "    # 计算 TD 误差并更新 Critic 参数\n",
    "    # 对应公式: w_{t+1} = w_t + alpha_w * TD_Error * grad(Q)\n",
    "    critic_loss = F.mse_loss(current_q, td_target)\n",
    "    \n",
    "    critic_optimizer.zero_grad()\n",
    "    critic_loss.backward()\n",
    "    critic_optimizer.step()\n",
    "    \n",
    "    # -----------------------------------------\n",
    "    # Actor 更新 (策略更新)\n",
    "    # -----------------------------------------\n",
    "    # 获取当前状态下所有动作的概率分布\n",
    "    action_probs = actor(state)\n",
    "    # 构建概率分布对象，方便计算对数概率\n",
    "    dist = distributions.Categorical(action_probs)\n",
    "    \n",
    "    # 计算 ln(pi(a_t | s_t, theta_t))\n",
    "    log_prob = dist.log_prob(torch.tensor(action))\n",
    "    \n",
    "    # Actor 梯度上升目标：ln(pi) * Q(s, a)\n",
    "    # 在 PyTorch 中优化器默认执行梯度下降，所以加个负号变成最小化损失\n",
    "    # 注意这里使用的是刚刚更新前算出的 current_q，为了阻断梯度传到 Critic，使用 .detach()\n",
    "    actor_loss = -log_prob * current_q.detach()\n",
    "    \n",
    "    actor_optimizer.zero_grad()\n",
    "    actor_loss.backward()\n",
    "    actor_optimizer.step()\n",
    "\n",
    "    return actor_loss.item(), critic_loss.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81a39722",
   "metadata": {},
   "source": [
    "\n",
    "## 控制任务\n",
    "\n",
    "本处使用OpenAI的Gymnasium环境中的CartPole-v1作为测试环境。这个环境的目标是通过控制一个小车来保持杆子竖直。状态空间是连续的，包含了小车的位置、速度以及杆子的角度和角速度；动作空间是离散的，只有两个动作：向左或向右移动小车。\n",
    "\n",
    "状态空间: 一个四维的连续向量\n",
    "- $s_0$: 小车的位置\n",
    "- $s_1$: 小车的速度\n",
    "- $s_2$: 杆子的角度\n",
    "- $s_3$: 杆子的角速度\n",
    "\n",
    "动作空间: 两个离散动作\n",
    "- $a_0$: 向左移动小车\n",
    "- $a_1$: 向右移动小车\n",
    "\n",
    "奖励函数: 每个时间步获得的奖励为 +1，直到杆子倒下或小车移出边界。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cec5bd6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cpu\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_77052/2107052713.py:73: UserWarning: Using a target size (torch.Size([1])) that is different to the input size (torch.Size([])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
      "  critic_loss = F.mse_loss(current_q, td_target)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 100/1000 completed. 近100回合平均reward: 15.25\n",
      "Episode 200/1000 completed. 近100回合平均reward: 33.00\n",
      "Episode 300/1000 completed. 近100回合平均reward: 53.27\n",
      "Episode 400/1000 completed. 近100回合平均reward: 34.57\n",
      "Episode 500/1000 completed. 近100回合平均reward: 59.50\n",
      "Episode 600/1000 completed. 近100回合平均reward: 51.37\n",
      "Episode 700/1000 completed. 近100回合平均reward: 77.24\n",
      "Episode 800/1000 completed. 近100回合平均reward: 77.66\n",
      "Episode 900/1000 completed. 近100回合平均reward: 78.07\n",
      "Episode 1000/1000 completed. 近100回合平均reward: 51.60\n",
      "\n",
      "=== QAC 训练统计 ===\n",
      "总回合数: 1000\n",
      "平均Reward: 53.15\n",
      "最高Reward: 336.00\n",
      "最后100回合平均Reward: 51.60\n",
      "图片已保存到: qac_training_results.png\n"
     ]
    }
   ],
   "source": [
    "# 设置设备 - 使用 CPU\n",
    "device = torch.device(\"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "\n",
    "# 训练的主循环\n",
    "def main():\n",
    "    # 初始化环境\n",
    "    env = gym.make('CartPole-v1')\n",
    "    state_dim = env.observation_space.shape[0]  # type: ignore\n",
    "    action_dim = int(env.action_space.n) # type: ignore\n",
    "\n",
    "    # 实例化 Actor 和 Critic 网络\n",
    "    actor = Actor(state_dim, action_dim).to(device)\n",
    "    critic = Critic(state_dim, action_dim).to(device)\n",
    "    actor_optimizer = optim.Adam(actor.parameters(), lr=ACTOR_LR)\n",
    "    critic_optimizer = optim.Adam(critic.parameters(), lr=CRITIC_LR)\n",
    "    \n",
    "    # 记录训练数据\n",
    "    episode_rewards = []\n",
    "    actor_losses = []\n",
    "    critic_losses = []\n",
    "\n",
    "    for episode in range(NUM_EPISODES):\n",
    "        state, _ = env.reset()\n",
    "        episode_reward = 0\n",
    "        episode_actor_loss = 0\n",
    "        episode_critic_loss = 0\n",
    "        steps = 0\n",
    "\n",
    "        # 预先生成一个初始动作 a_0\n",
    "        state_tensor = torch.FloatTensor(state).to(device)\n",
    "        action_probs = actor(state_tensor)\n",
    "        action = distributions.Categorical(action_probs).sample().item()\n",
    "\n",
    "        while True:\n",
    "            # 与环境交互，获取当前状态\n",
    "            next_state, reward, terminated, truncated, _ = env.step(action)\n",
    "            done = terminated or truncated\n",
    "\n",
    "            # 预先生成下一个动作 a_{t+1}，供 Critic 更新使用\n",
    "            next_state_tensor = torch.FloatTensor(next_state).to(device)\n",
    "            next_action_probs = actor(next_state_tensor)\n",
    "            next_action = distributions.Categorical(next_action_probs).sample().item()\n",
    "\n",
    "            # 执行单步的学习\n",
    "            actor_loss, critic_loss = train_step(actor, critic, actor_optimizer, critic_optimizer,\n",
    "                          state, action, reward, next_state, next_action, done)\n",
    "            \n",
    "            # 记录loss\n",
    "            episode_actor_loss += actor_loss\n",
    "            episode_critic_loss += critic_loss\n",
    "            steps += 1\n",
    "            \n",
    "            # 更新状态和动作\n",
    "            state = next_state\n",
    "            action = next_action\n",
    "            episode_reward += reward  # type: ignore\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        # 记录本回合数据\n",
    "        episode_rewards.append(episode_reward)\n",
    "        actor_losses.append(episode_actor_loss / steps)\n",
    "        critic_losses.append(episode_critic_loss / steps)\n",
    "\n",
    "        # 打印训练进度\n",
    "        if (episode + 1) % 100 == 0:\n",
    "            avg_reward = np.mean(episode_rewards[-100:])\n",
    "            print(f\"Episode {episode + 1}/{NUM_EPISODES} completed. 近100回合平均reward: {avg_reward:.2f}\")\n",
    "    \n",
    "    env.close()\n",
    "    \n",
    "    # 绘制训练曲线\n",
    "    plot_qac_results(episode_rewards, actor_losses, critic_losses)\n",
    "\n",
    "def plot_qac_results(rewards, actor_losses, critic_losses):\n",
    "    \"\"\"绘制QAC训练结果（保存到文件）\"\"\"\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
    "    \n",
    "    # 绘制 reward 曲线\n",
    "    axes[0].plot(rewards, alpha=0.6, label='Episode Reward')\n",
    "    # 添加移动平均线\n",
    "    window = min(50, len(rewards) // 10)\n",
    "    if window > 1:\n",
    "        moving_avg = np.convolve(rewards, np.ones(window)/window, mode='valid')\n",
    "        axes[0].plot(np.arange(window-1, len(rewards)), moving_avg, 'r-', label=f'{window}-Episode Moving Avg')\n",
    "    axes[0].set_xlabel('Episode')\n",
    "    axes[0].set_ylabel('Reward')\n",
    "    axes[0].set_title('QAC Training Rewards')\n",
    "    axes[0].legend()\n",
    "    axes[0].grid(True, alpha=0.3)\n",
    "    \n",
    "    # 绘制 Actor loss\n",
    "    axes[1].plot(actor_losses, alpha=0.6, color='orange')\n",
    "    axes[1].set_xlabel('Episode')\n",
    "    axes[1].set_ylabel('Loss')\n",
    "    axes[1].set_title('QAC Actor Loss')\n",
    "    axes[1].grid(True, alpha=0.3)\n",
    "    \n",
    "    # 绘制 Critic loss\n",
    "    axes[2].plot(critic_losses, alpha=0.6, color='green')\n",
    "    axes[2].set_xlabel('Episode')\n",
    "    axes[2].set_ylabel('Loss')\n",
    "    axes[2].set_title('QAC Critic Loss')\n",
    "    axes[2].grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    # 保存图片到文件\n",
    "    plt.savefig('qac_training_results.png', dpi=150)\n",
    "    plt.close()\n",
    "    \n",
    "    # 打印统计信息\n",
    "    print(f\"\\n=== QAC 训练统计 ===\")\n",
    "    print(f\"总回合数: {len(rewards)}\")\n",
    "    print(f\"平均Reward: {np.mean(rewards):.2f}\")\n",
    "    print(f\"最高Reward: {np.max(rewards):.2f}\")\n",
    "    print(f\"最后100回合平均Reward: {np.mean(rewards[-100:]):.2f}\")\n",
    "    print(f\"图片已保存到: qac_training_results.png\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cc0f13e",
   "metadata": {},
   "source": [
    "## A2C算法\n",
    "\n",
    "通过上一个部分，我们看到其实最基础的QAC算法在这个控制算例中效果并不好，原因是他的Actor网络的更新时直接采用了Critic网络的输出作为动作价值的估计，这个估计可能非常不准确，导致Actor网络的更新方向错误，从而无法有效地学习到好的策略。另外，由于QAC是一个On-Policy的算法，一般产生一条轨迹，更新一次参数，而之后这些参数就不再被使用了，这样就浪费了很多数据。我们先解决第一个问题，即方差过大的问题，引入一个基线函数来降低方差，这就是 Advantage Actor-Critic (A2C) 算法。A2C 的核心思想是引入一个基线函数 $b(s)$ 来减去动作价值 $q(s,a)$ 中的平均水平，从而得到优势函数 $A(s,a) = q(s,a) - b(s)$。这个优势函数可以更准确地反映某个动作相对于平均水平的好坏，从而降低了梯度估计的方差。\n",
    "\n",
    "从控制系统设计的角度来看，最基础的 QAC 就像是一个只看绝对误差、没有稳态基准的开环打分器，系统噪声和方差极其容易被放大，导致策略输出剧烈震荡。\n",
    "\n",
    "A2C引入状态价值 $v_{\\pi}(s)$ 作为基线（Baseline），这就如同在闭环控制中引入了一个动态的参考基准。Critic 不再评估具体的动作有多好，而是评估当前状态的平均预期。Actor 更新的依据变成了优势函数（Advantage）：\n",
    "$$\\delta_t = r_{t+1} + \\gamma v(s_{t+1}) - v(s_t)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4809b58c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cpu\n",
      "开始 A2C 训练...\n",
      "回合 100/1000 完成. 近 100 回合平均 Reward: 13.11\n",
      "回合 200/1000 完成. 近 100 回合平均 Reward: 65.78\n",
      "回合 300/1000 完成. 近 100 回合平均 Reward: 59.20\n",
      "回合 400/1000 完成. 近 100 回合平均 Reward: 9.48\n",
      "回合 500/1000 完成. 近 100 回合平均 Reward: 65.32\n",
      "回合 600/1000 完成. 近 100 回合平均 Reward: 67.78\n",
      "回合 700/1000 完成. 近 100 回合平均 Reward: 9.36\n",
      "回合 800/1000 完成. 近 100 回合平均 Reward: 9.26\n",
      "回合 900/1000 完成. 近 100 回合平均 Reward: 9.68\n",
      "回合 1000/1000 完成. 近 100 回合平均 Reward: 51.43\n",
      "\n",
      "=== A2C 训练统计 ===\n",
      "总回合数: 1000\n",
      "平均Reward: 36.04\n",
      "最高Reward: 500.00\n",
      "最后100回合平均Reward: 51.43\n",
      "图片已保存到: a2c_training_results.png\n"
     ]
    }
   ],
   "source": [
    "# 设置设备 - 使用 CPU\n",
    "device = torch.device(\"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "# 定义超参数\n",
    "ACTOR_LR = 0.001\n",
    "CRITIC_LR = 0.002\n",
    "GAMMA = 0.99\n",
    "NUM_EPISODES = 1000\n",
    "\n",
    "# 1. 定义 Actor 网络 (演员：输出离散动作的概率分布)\n",
    "class ActorA2C(nn.Module):\n",
    "    def __init__(self, state_dim, action_dim):\n",
    "        super(ActorA2C, self).__init__()\n",
    "        self.fc1 = nn.Linear(state_dim, 128)\n",
    "        self.fc2 = nn.Linear(128, action_dim)\n",
    "        \n",
    "    def forward(self, state):\n",
    "        x = F.relu(self.fc1(state))\n",
    "        action_probs = F.softmax(self.fc2(x), dim=-1)\n",
    "        return action_probs\n",
    "    \n",
    "# 2. 定义 Critic 网络 (评论家：输出标量 V 值)\n",
    "class CriticA2C(nn.Module):\n",
    "    def __init__(self, state_dim):\n",
    "        super(CriticA2C, self).__init__()\n",
    "        self.fc1 = nn.Linear(state_dim, 128)\n",
    "        self.fc2 = nn.Linear(128, 1)\n",
    "        \n",
    "    def forward(self, state):\n",
    "        x = F.relu(self.fc1(state))\n",
    "        v_value = self.fc2(x)\n",
    "        return v_value\n",
    "\n",
    "# 3. 核心训练逻辑：基于优势函数 (Advantage)\n",
    "def train_step_a2c(actor, critic, actor_optimizer, critic_optimizer, \n",
    "                   state, action, reward, next_state, done):\n",
    "    state = torch.FloatTensor(state).unsqueeze(0).to(device)\n",
    "    next_state = torch.FloatTensor(next_state).unsqueeze(0).to(device)\n",
    "    reward = torch.FloatTensor([reward]).unsqueeze(0).to(device)\n",
    "    \n",
    "    # Critic 更新\n",
    "    v_value = critic(state)\n",
    "    next_v_value = critic(next_state).detach()\n",
    "    td_target = reward + GAMMA * next_v_value * (1 - int(done))\n",
    "    advantage = td_target - v_value\n",
    "    critic_loss = F.mse_loss(v_value, td_target)\n",
    "    \n",
    "    critic_optimizer.zero_grad()\n",
    "    critic_loss.backward()\n",
    "    critic_optimizer.step()\n",
    "    \n",
    "    # Actor 更新\n",
    "    action_probs = actor(state)\n",
    "    dist = distributions.Categorical(action_probs)\n",
    "    log_prob = dist.log_prob(torch.tensor([action]).to(device))\n",
    "    actor_loss = -(log_prob * advantage.detach()).mean()\n",
    "    \n",
    "    actor_optimizer.zero_grad()\n",
    "    actor_loss.backward()\n",
    "    actor_optimizer.step()\n",
    "\n",
    "    return actor_loss.item(), critic_loss.item()\n",
    "\n",
    "\n",
    "def plot_a2c_results(rewards, actor_losses, critic_losses):\n",
    "    \"\"\"绘制A2C训练结果（保存到文件）\"\"\"\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
    "    \n",
    "    axes[0].plot(rewards, alpha=0.6, label='Episode Reward')\n",
    "    window = min(50, len(rewards) // 10)\n",
    "    if window > 1:\n",
    "        moving_avg = np.convolve(rewards, np.ones(window)/window, mode='valid')\n",
    "        axes[0].plot(np.arange(window-1, len(rewards)), moving_avg, 'r-', label=f'{window}-Episode Moving Avg')\n",
    "    axes[0].set_xlabel('Episode')\n",
    "    axes[0].set_ylabel('Reward')\n",
    "    axes[0].set_title('A2C Training Rewards')\n",
    "    axes[0].legend()\n",
    "    axes[0].grid(True, alpha=0.3)\n",
    "    \n",
    "    axes[1].plot(actor_losses, alpha=0.6, color='orange')\n",
    "    axes[1].set_xlabel('Episode')\n",
    "    axes[1].set_ylabel('Loss')\n",
    "    axes[1].set_title('A2C Actor Loss')\n",
    "    axes[1].grid(True, alpha=0.3)\n",
    "    \n",
    "    axes[2].plot(critic_losses, alpha=0.6, color='green')\n",
    "    axes[2].set_xlabel('Episode')\n",
    "    axes[2].set_ylabel('Loss')\n",
    "    axes[2].set_title('A2C Critic Loss')\n",
    "    axes[2].grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig('a2c_training_results.png', dpi=150)\n",
    "    plt.close()\n",
    "    \n",
    "    print(f\"\\n=== A2C 训练统计 ===\")\n",
    "    print(f\"总回合数: {len(rewards)}\")\n",
    "    print(f\"平均Reward: {np.mean(rewards):.2f}\")\n",
    "    print(f\"最高Reward: {np.max(rewards):.2f}\")\n",
    "    print(f\"最后100回合平均Reward: {np.mean(rewards[-100:]):.2f}\")\n",
    "    print(f\"图片已保存到: a2c_training_results.png\")\n",
    "\n",
    "\n",
    "# 4. 主循环函数\n",
    "def main_a2c():\n",
    "    env = gym.make('CartPole-v1')\n",
    "    state_dim = env.observation_space.shape[0]   # type: ignore\n",
    "    action_dim = int(env.action_space.n)         # type: ignore\n",
    "\n",
    "    actor = ActorA2C(state_dim, action_dim).to(device)\n",
    "    critic = CriticA2C(state_dim).to(device)\n",
    "    actor_optimizer = optim.Adam(actor.parameters(), lr=ACTOR_LR)\n",
    "    critic_optimizer = optim.Adam(critic.parameters(), lr=CRITIC_LR)\n",
    "    \n",
    "    episode_rewards = []\n",
    "    actor_losses = []\n",
    "    critic_losses = []\n",
    "    \n",
    "    print(\"开始 A2C 训练...\")\n",
    "    for episode in range(NUM_EPISODES):\n",
    "        state, _ = env.reset()\n",
    "        episode_reward = 0\n",
    "        episode_actor_loss = 0\n",
    "        episode_critic_loss = 0\n",
    "        steps = 0\n",
    "        \n",
    "        while True:\n",
    "            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)\n",
    "            action_probs = actor(state_tensor)\n",
    "            action = distributions.Categorical(action_probs).sample().item()\n",
    "            \n",
    "            next_state, reward, terminated, truncated, _ = env.step(action)\n",
    "            done = terminated or truncated\n",
    "            \n",
    "            actor_loss, critic_loss = train_step_a2c(actor, critic, actor_optimizer, critic_optimizer, \n",
    "                           state, action, reward, next_state, done)\n",
    "            \n",
    "            episode_actor_loss += actor_loss\n",
    "            episode_critic_loss += critic_loss\n",
    "            steps += 1\n",
    "            \n",
    "            state = next_state\n",
    "            episode_reward += reward\n",
    "            \n",
    "            if done:\n",
    "                break\n",
    "        \n",
    "        episode_rewards.append(episode_reward)\n",
    "        actor_losses.append(episode_actor_loss / steps)\n",
    "        critic_losses.append(episode_critic_loss / steps)\n",
    "        \n",
    "        if (episode + 1) % 100 == 0:\n",
    "            avg_reward = np.mean(episode_rewards[-100:])\n",
    "            print(f\"回合 {episode + 1}/{NUM_EPISODES} 完成. 近 100 回合平均 Reward: {avg_reward:.2f}\")\n",
    "\n",
    "    env.close()\n",
    "    plot_a2c_results(episode_rewards, actor_losses, critic_losses)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main_a2c()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "RL_Env",
   "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.11.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}