别怕数学！用Python手把手带你推导贝尔曼方程（附代码）

用Python代码拆解贝尔曼方程：从数学恐惧到编程实践

1. 为什么我们需要贝尔曼方程？

在强化学习的世界里，贝尔曼方程就像是一张藏宝图，指引着智能体如何在未知环境中做出最优决策。想象你正在玩一个迷宫游戏，每次走到岔路口都需要决定往左还是往右。贝尔曼方程就是那个能告诉你"当前选择对未来奖励影响"的神奇公式。

传统教学中，贝尔曼方程往往以复杂的数学符号呈现：

v_{\pi}(s) = \sum_a \pi(a|s) \sum_{s',r} p(s',r|s,a)[r + \gamma v_\pi(s')]

这让很多开发者望而生畏。但如果我们换种方式，用Python代码一步步构建这个方程，你会发现它其实非常直观。下面是我们将要实现的代码框架：

import numpy as np class BellmanEquation: def __init__(self, states, actions, transition_probs, rewards, gamma=0.9): self.states = states self.actions = actions self.transition_probs = transition_probs # [s, a, s'] self.rewards = rewards # [s, a, s'] self.gamma = gamma # 折扣因子

2. 构建基础环境模型

2.1 定义马尔可夫决策过程(MDP)

任何强化学习问题都始于对环境的建模。我们创建一个简单的网格世界作为示例：

def create_grid_world(size=4): """创建一个size x size的网格世界""" states = [(i, j) for i in range(size) for j in range(size)] actions = ['up', 'down', 'left', 'right'] # 初始化转移概率和奖励 transition_probs = np.zeros((size, size, len(actions), size, size)) rewards = np.zeros((size, size, len(actions), size, size)) # 填充转移规则（实际项目中这里会更复杂） for i in range(size): for j in range(size): for a_idx, action in enumerate(actions): # 简单移动逻辑 next_i, next_j = move((i, j), action, size) transition_probs[i, j, a_idx, next_i, next_j] = 1.0 rewards[i, j, a_idx, next_i, next_j] = -1 # 每步小惩罚 # 设置终点奖励 rewards[size-1, size-1, :, :, :] = 10 return states, actions, transition_probs, rewards def move(state, action, size): """处理移动逻辑""" i, j = state if action == 'up' and i > 0: return i-1, j elif action == 'down' and i < size-1: return i+1, j elif action == 'left' and j > 0: return i, j-1 elif action == 'right' and j < size-1: return i, j+1 return i, j # 碰到边界保持原地

2.2 可视化环境

理解环境结构对调试至关重要。我们可以用matplotlib绘制网格：

import matplotlib.pyplot as plt def plot_grid_world(size, terminal_state=None): fig, ax = plt.subplots(figsize=(size, size)) ax.set_xticks(np.arange(size+1)) ax.set_yticks(np.arange(size+1)) ax.grid(which='both') if terminal_state: rect = plt.Rectangle(terminal_state, 1, 1, facecolor='green', alpha=0.3) ax.add_patch(rect) plt.xlim(0, size) plt.ylim(0, size) plt.gca().invert_yaxis() # 让(0,0)在左上角 plt.show()

3. 实现贝尔曼方程的核心逻辑

3.1 状态值函数计算

贝尔曼方程的核心是递归地评估状态价值。让我们用代码实现这个递归关系：

def calculate_state_value(self, policy, state, current_values): """ 计算给定策略下某状态的价值 :param policy: [s, a] 策略矩阵 :param state: 当前状态 (i,j) :param current_values: 当前各状态的价值估计 [size, size] :return: 该状态的新价值 """ i, j = state total = 0 for a_idx, action in enumerate(self.actions): # 第一部分：即时奖励的期望 immediate_reward = np.sum( self.rewards[i, j, a_idx] * self.transition_probs[i, j, a_idx] ) # 第二部分：未来奖励的期望 future_reward = 0 for next_i in range(len(self.states)): for next_j in range(len(self.states)): prob = self.transition_probs[i, j, a_idx, next_i, next_j] future_reward += prob * current_values[next_i, next_j] # 合并两部分，考虑策略概率 total += policy[i, j, a_idx] * (immediate_reward + self.gamma * future_reward) return total

3.2 策略评估算法

基于贝尔曼方程，我们可以迭代评估策略：

def policy_evaluation(self, policy, threshold=1e-4, max_iter=1000): """策略评估算法""" values = np.zeros((len(self.states), len(self.states))) for _ in range(max_iter): delta = 0 new_values = np.zeros_like(values) for i in range(len(self.states)): for j in range(len(self.states)): state = (i, j) new_v = self.calculate_state_value(policy, state, values) delta = max(delta, abs(new_v - values[i, j])) new_values[i, j] = new_v values = new_values if delta < threshold: break return values

4. 从理论到实践：完整案例解析

4.1 初始化策略和环境

让我们创建一个4x4网格世界并定义随机策略：

# 创建环境 size = 4 states, actions, transition_probs, rewards = create_grid_world(size) # 定义随机策略（每个动作概率均等） random_policy = np.ones((size, size, len(actions))) / len(actions) # 实例化贝尔曼方程求解器 bellman = BellmanEquation(states, actions, transition_probs, rewards, gamma=0.9)

4.2 运行策略评估

执行策略评估并可视化结果：

# 评估随机策略 values = bellman.policy_evaluation(random_policy) # 可视化价值函数 def plot_values(values): plt.figure(figsize=(size, size)) plt.imshow(values, cmap='hot', interpolation='nearest') for i in range(size): for j in range(size): plt.text(j, i, f"{values[i, j]:.1f}", ha='center', va='center', color='blue') plt.colorbar() plt.title("State Values under Random Policy") plt.show() plot_values(values)

4.3 结果分析与优化

观察输出结果，你会发现：

1. 右下角终点状态价值最高（约8-9）
2. 距离终点越远的状态价值越低
3. 边缘状态由于移动受限，价值略低于中心状态

这验证了贝尔曼方程的核心思想：当前状态价值等于即时奖励加上未来奖励的折现期望。我们可以进一步优化策略：

def improve_policy(values, transition_probs, rewards, gamma=0.9): """策略改进：基于当前价值函数选择最优动作""" new_policy = np.zeros_like(random_policy) size = values.shape[0] for i in range(size): for j in range(size): # 计算每个动作的Q值 q_values = [] for a_idx in range(len(actions)): immediate = np.sum(rewards[i, j, a_idx] * transition_probs[i, j, a_idx]) future = gamma * np.sum(transition_probs[i, j, a_idx] * values) q_values.append(immediate + future) # 选择最优动作 best_action = np.argmax(q_values) new_policy[i, j, best_action] = 1.0 return new_policy # 策略迭代过程 optimized_policy = improve_policy(values, transition_probs, rewards) optimized_values = bellman.policy_evaluation(optimized_policy) plot_values(optimized_values)

5. 高级话题与实用技巧

5.1 处理大规模状态空间

当状态空间很大时，直接计算变得不可行。我们可以采用以下优化：

def approximate_policy_evaluation(self, policy, num_samples=1000): """使用采样方法近似计算价值函数""" values = np.zeros((len(self.states), len(self.states))) counts = np.zeros_like(values) for _ in range(num_samples): state = (np.random.randint(size), np.random.randint(size)) total_reward = 0 discount = 1.0 # 模拟一条轨迹 for _ in range(100): # 防止无限循环 # 根据策略选择动作 a_idx = np.random.choice(len(actions), p=policy[state[0], state[1]]) # 根据转移概率得到下一个状态 next_state_probs = self.transition_probs[state[0], state[1], a_idx] next_i, next_j = np.unravel_index( np.random.choice(len(self.states)**2, p=next_state_probs.ravel()), next_state_probs.shape ) # 累积奖励 total_reward += discount * self.rewards[state[0], state[1], a_idx, next_i, next_j] discount *= self.gamma # 更新状态 state = (next_i, next_j) # 如果到达终止状态则结束 if state == (size-1, size-1): break # 更新价值估计 values[state[0], state[1]] += total_reward counts[state[0], state[1]] += 1 # 计算平均值 return np.where(counts > 0, values / counts, 0)

5.2 调试贝尔曼方程实现

常见问题及解决方案：

5.3 扩展应用：Q-Learning算法

贝尔曼方程是许多强化学习算法的基础。以下是Q-Learning的实现片段：

def q_learning(self, episodes=1000, alpha=0.1, epsilon=0.1): """Q-Learning算法实现""" q_table = np.zeros((len(self.states), len(self.states), len(self.actions))) for _ in range(episodes): state = (0, 0) # 起始状态 while state != (size-1, size-1): # 未到达终点 # ε-贪婪策略选择动作 if np.random.random() < epsilon: action_idx = np.random.randint(len(self.actions)) else: action_idx = np.argmax(q_table[state[0], state[1]]) # 执行动作，观察下一个状态和奖励 next_state_probs = self.transition_probs[state[0], state[1], action_idx] next_i, next_j = np.unravel_index( np.random.choice(len(self.states)**2, p=next_state_probs.ravel()), next_state_probs.shape ) reward = self.rewards[state[0], state[1], action_idx, next_i, next_j] # Q值更新（贝尔曼最优方程） best_next_action = np.argmax(q_table[next_i, next_j]) td_target = reward + self.gamma * q_table[next_i, next_j, best_next_action] q_table[state[0], state[1], action_idx] += alpha * ( td_target - q_table[state[0], state[1], action_idx] ) state = (next_i, next_j) return q_table