{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31a4e03b",
   "metadata": {},
   "source": [
    "# 第 2 章：状态价值与贝尔曼方程\n",
    "贝尔曼方程描述了所有状态价值之间的关系。\n",
    "我们将使用**矩阵-向量形式** 来表示贝尔曼方程： $$v_{\\pi} = r_{\\pi} + \\gamma P_{\\pi}v_{\\pi}$$\n",
    "以及使用**迭代法** 来求解它： $$v_{k+1} = r_{\\pi} + \\gamma P_{\\pi}v_{k}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e67777b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- 贝尔曼方程迭代求解过程 ---\n",
      "第 1 次迭代: v(s1)=-0.5000, v(s2)=1.0000, v(s3)=1.0000, v(s4)=1.0000\n",
      "第 2 次迭代: v(s1)=0.4000, v(s2)=1.9000, v(s3)=1.9000, v(s4)=1.9000\n",
      "第 3 次迭代: v(s1)=1.2100, v(s2)=2.7100, v(s3)=2.7100, v(s4)=2.7100\n",
      "第 4 次迭代: v(s1)=1.9390, v(s2)=3.4390, v(s3)=3.4390, v(s4)=3.4390\n",
      "第 5 次迭代: v(s1)=2.5951, v(s2)=4.0951, v(s3)=4.0951, v(s4)=4.0951\n",
      "第 100 次迭代: v(s1)=8.4997, v(s2)=9.9997, v(s3)=9.9997, v(s4)=9.9997\n",
      "\n",
      "最终收敛的状态价值: [ 8.5 10.  10.  10. ]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 1. 设定折扣因子\n",
    "gamma = 0.9\n",
    "\n",
    "# 2. 定义书本图 2.5 中的期望奖励向量 r_pi (Reward vector)\n",
    "# s1 的期望奖励是 0.5*0 + 0.5*(-1) = -0.5\n",
    "r_pi = np.array([-0.5, 1.0, 1.0, 1.0])\n",
    "\n",
    "# 3. 定义状态转移矩阵 P_pi (Transition matrix)\n",
    "# 每一行代表 s1, s2, s3, s4; 每一列代表转移到 s1, s2, s3, s4 的概率\n",
    "P_pi = np.array([\n",
    "    [0.0, 0.5, 0.5, 0.0],  # s1 有 0.5 概率到 s2, 0.5 概率到 s3\n",
    "    [0.0, 0.0, 0.0, 1.0],  # s2 有 1.0 概率到 s4\n",
    "    [0.0, 0.0, 0.0, 1.0],  # s3 有 1.0 概率到 s4\n",
    "    [0.0, 0.0, 0.0, 1.0]   # s4 有 1.0 概率停留在 s4\n",
    "])\n",
    "\n",
    "# 4. 迭代法求解状态价值 (Iterative solution)\n",
    "v_values = np.zeros(4) # 初始时，假设所有状态价值为 0\n",
    "iterations = 100\n",
    "\n",
    "print(\"--- 贝尔曼方程迭代求解过程 ---\")\n",
    "for k in range(iterations):\n",
    "    # 核心公式：v_{k+1} = r_pi + gamma * P_pi * v_k\n",
    "    v_next = r_pi + gamma * P_pi.dot(v_values)\n",
    "    \n",
    "    if k < 5 or k == iterations - 1:\n",
    "        print(f\"第 {k+1} 次迭代: v(s1)={v_next[0]:.4f}, v(s2)={v_next[1]:.4f}, v(s3)={v_next[2]:.4f}, v(s4)={v_next[3]:.4f}\")\n",
    "        \n",
    "    v_values = v_next\n",
    "\n",
    "print(\"\\n最终收敛的状态价值:\", np.round(v_values, 2))\n",
    "# 你可以将这个输出结果与书中 2.5 节手算的 8.5, 10, 10, 10 进行对比！"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2086b1ab",
   "metadata": {},
   "source": [
    "# 2.8 从状态价值到动作价值 (Action Value)\n",
    "动作价值 $q_\\pi(s,a)$ 评估的是在特定状态下采取特定动作的好坏。\n",
    "\n",
    "它由两部分组成：\n",
    "1. **即时奖励的期望**：采取动作 $a$ 后立刻得到的奖励。\n",
    "2. **未来奖励的期望**：进入下一个状态 $s'$ 后，未来的状态价值 $\\gamma v_\\pi(s')$。\n",
    "\n",
    "数学公式为：\n",
    "$$q_{\\pi}(s,a) = \\sum_{r\\in\\mathcal{R}}p(r|s,a)r + \\gamma\\sum_{s'\\in\\mathcal{S}}p(s'|s,a)v_{\\pi}(s')$$\n",
    "\n",
    "状态价值其实就是所有可能动作价值的加权平均：\n",
    "$$v_{\\pi}(s) = \\sum_{a\\in\\mathcal{A}}\\pi(a|s)q_{\\pi}(s,a)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a3e76dd6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- 计算在状态 s1 下，所有可能动作的价值 (Q-values) ---\n",
      "策略内的动作: q(s1, a2向右) = -1 + 0.9 * v(s2) = 8.0\n",
      "策略内的动作: q(s1, a3向下) = 0 + 0.9 * v(s3) = 9.0\n",
      "\n",
      "策略外的动作: q(s1, a1向上) = -1 + 0.9 * v(s1) = 6.65\n",
      "策略外的动作: q(s1, a4向左) = -1 + 0.9 * v(s1) = 6.65\n",
      "策略外的动作: q(s1, a5原地) =  0 + 0.9 * v(s1) = 7.65\n",
      "\n",
      "--- 验证公式：状态价值是动作价值的加权平均 ---\n",
      "根据 Q 值反推的 v(s1) = 0.5 * 8.0 + 0.5 * 9.0 = 8.5\n",
      "这与我们之前迭代出来的 v(s1) = 8.5 完全一致！\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 1. 继承之前的已知数据\n",
    "gamma = 0.9\n",
    "# 假设这是我们在 Cell 8 算出来的最终状态价值\n",
    "v_values = {\"s1\": 8.5, \"s2\": 10.0, \"s3\": 10.0, \"s4\": 10.0}\n",
    "\n",
    "# 2. 定义动作价值计算函数 q(s, a)\n",
    "def calculate_q_value(state, action, reward, next_state, gamma, v_values):\n",
    "    # q(s,a) = 立即奖励 + gamma * 下一个状态的价值\n",
    "    return reward + gamma * v_values[next_state]\n",
    "\n",
    "print(\"--- 计算在状态 s1 下，所有可能动作的价值 (Q-values) ---\")\n",
    "\n",
    "# --- 策略会选择的动作 (向右 a2, 向下 a3) ---\n",
    "# a2: 向右进入 s2，得奖励 -1\n",
    "q_s1_a2 = calculate_q_value(\"s1\", \"a2\", -1, \"s2\", gamma, v_values)\n",
    "print(f\"策略内的动作: q(s1, a2向右) = -1 + 0.9 * v(s2) = {q_s1_a2}\") # 预期为 8.0\n",
    "\n",
    "# a3: 向下进入 s3，得奖励 0\n",
    "q_s1_a3 = calculate_q_value(\"s1\", \"a3\", 0, \"s3\", gamma, v_values)\n",
    "print(f\"策略内的动作: q(s1, a3向下) = 0 + 0.9 * v(s3) = {q_s1_a3}\") # 预期为 9.0\n",
    "\n",
    "\n",
    "# --- 策略不会选择的动作 (向上 a1, 向左 a4, 原地 a5) ---\n",
    "# a1: 向上撞墙，弹回 s1，得奖励 -1\n",
    "q_s1_a1 = calculate_q_value(\"s1\", \"a1\", -1, \"s1\", gamma, v_values)\n",
    "print(f\"\\n策略外的动作: q(s1, a1向上) = -1 + 0.9 * v(s1) = {q_s1_a1:.2f}\") # 预期为 6.65\n",
    "\n",
    "# a4: 向左撞墙，弹回 s1，得奖励 -1\n",
    "q_s1_a4 = calculate_q_value(\"s1\", \"a4\", -1, \"s1\", gamma, v_values)\n",
    "print(f\"策略外的动作: q(s1, a4向左) = -1 + 0.9 * v(s1) = {q_s1_a4:.2f}\") # 预期为 6.65\n",
    "\n",
    "# a5: 原地不动，停在 s1，得奖励 0\n",
    "q_s1_a5 = calculate_q_value(\"s1\", \"a5\", 0, \"s1\", gamma, v_values)\n",
    "print(f\"策略外的动作: q(s1, a5原地) =  0 + 0.9 * v(s1) = {q_s1_a5:.2f}\") # 预期为 7.65\n",
    "\n",
    "\n",
    "print(\"\\n--- 验证公式：状态价值是动作价值的加权平均 ---\")\n",
    "# 我们的策略是: 0.5概率选a2, 0.5概率选a3\n",
    "v_s1_calculated = 0.5 * q_s1_a2 + 0.5 * q_s1_a3\n",
    "print(f\"根据 Q 值反推的 v(s1) = 0.5 * {q_s1_a2} + 0.5 * {q_s1_a3} = {v_s1_calculated}\")\n",
    "print(f\"这与我们之前迭代出来的 v(s1) = {v_values['s1']} 完全一致！\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GymRL",
   "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.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}