别再死磕贝尔曼方程了!用Python可视化理解MDP中的“状态访问分布”与“占用度量”

别再死磕贝尔曼方程了!用Python可视化理解MDP中的“状态访问分布”与“占用度量” 用Python动态可视化强化学习中的状态访问分布与占用度量当你第一次接触强化学习时那些晦涩的数学公式是否让你望而却步贝尔曼方程、状态访问分布、占用度量...这些概念听起来就像天书。但今天我们将用Python和可视化工具把这些抽象的理论变成直观的动画和图表。告别枯燥的数学推导让我们一起动手构建一个简单的网格世界亲眼看看不同策略如何影响智能体的行为轨迹。1. 准备工作搭建网格世界环境在开始之前我们需要准备一个简单的实验环境。这里我们使用Python的Gymnasium库原OpenAI Gym的分支来创建一个5x5的网格世界。这个环境将作为我们探索状态访问分布和占用度量的 playground。import numpy as np import matplotlib.pyplot as plt import matplotlib.animation as animation from matplotlib.colors import ListedColormap import gymnasium as gym from gymnasium import spaces class GridWorldEnv(gym.Env): def __init__(self, size5): self.size size # 网格大小 self.action_space spaces.Discrete(4) # 上、下、左、右 self.observation_space spaces.Discrete(size * size) # 状态空间 # 定义特殊格子障碍物和终点 self.obstacles [(1, 1), (2, 2), (3, 3)] self.goal (4, 4) # 可视化相关设置 self.cmap ListedColormap([white, black, green]) def reset(self): self.state (0, 0) # 总是从左上角开始 return self.state def step(self, action): x, y self.state # 根据动作更新位置 if action 0: y max(0, y-1) # 上 elif action 1: y min(self.size-1, y1) # 下 elif action 2: x max(0, x-1) # 左 elif action 3: x min(self.size-1, x1) # 右 # 检查是否碰到障碍物或到达终点 if (x, y) in self.obstacles: return self.state, -1, False, {} elif (x, y) self.goal: return (x, y), 10, True, {} self.state (x, y) return self.state, 0, False, {} def render(self): grid np.zeros((self.size, self.size)) for obs in self.obstacles: grid[obs[1], obs[0]] 1 # 障碍物标记为1 grid[self.goal[1], self.goal[0]] 2 # 终点标记为2 grid[self.state[1], self.state[0]] 0.5 # 当前位置 plt.imshow(grid, cmapself.cmap) plt.xticks([]); plt.yticks([]) plt.show()这个简单的网格世界环境包含5x5的网格空间三个障碍物黑色格子一个终点绿色格子智能体从左上角(0,0)出发四个基本动作上、下、左、右2. 理解状态访问分布策略如何影响状态访问频率状态访问分布v^π(s)描述了在策略π下智能体访问各个状态s的频率。这个概念之所以重要是因为它直接影响策略的价值函数——智能体更常访问的状态对策略评估的贡献更大。让我们实现两个不同的策略并可视化它们的状态访问分布def random_policy(state): 随机策略等概率选择所有可用动作 return np.random.randint(0, 4) def goal_oriented_policy(state): 目标导向策略倾向于向终点(4,4)移动 x, y state if x 4 and y 4: return np.random.choice([1, 3], p[0.7, 0.3]) # 70%概率向下30%向右 elif x 4: return 3 # 向右 elif y 4: return 1 # 向下 else: return np.random.randint(0, 4) # 在终点随机选择现在我们编写一个函数来模拟多个episode并统计状态访问频率def compute_state_visitation(env, policy, episodes1000, gamma0.9): visitation np.zeros((env.size, env.size)) for _ in range(episodes): state env.reset() done False t 0 while not done: action policy(state) next_state, _, done, _ env.step(action) # 更新状态访问计数考虑折扣因子 visitation[state[1], state[0]] (gamma ** t) state next_state t 1 # 归一化处理 visitation / np.sum(visitation) return visitation让我们可视化两种策略的状态访问分布env GridWorldEnv() # 计算两种策略的状态访问分布 random_visitation compute_state_visitation(env, random_policy) goal_visitation compute_state_visitation(env, goal_oriented_policy) # 绘制对比图 fig, (ax1, ax2) plt.subplots(1, 2, figsize(12, 5)) im1 ax1.imshow(random_visitation, cmapviridis) ax1.set_title(随机策略的状态访问分布) plt.colorbar(im1, axax1) im2 ax2.imshow(goal_visitation, cmapviridis) ax2.set_title(目标导向策略的状态访问分布) plt.colorbar(im2, axax2) plt.tight_layout() plt.show()从可视化结果中我们可以清晰地看到随机策略访问分布相对均匀但障碍物周围的状态访问频率略低因为碰到障碍物会结束episode目标导向策略访问分布明显集中在通往终点的路径上特别是右下角区域3. 深入占用度量状态-动作对的访问概率占用度量ρ^π(s,a)比状态访问分布更进一步它记录了在策略π下特定状态-动作对被访问的概率。这个概念在策略梯度方法和模仿学习中尤为重要。让我们扩展之前的函数来计算占用度量def compute_occupancy_measure(env, policy, episodes1000, gamma0.9): 计算状态-动作对的占用度量 occupancy np.zeros((env.size, env.size, 4)) # (x, y, action) for _ in range(episodes): state env.reset() done False t 0 while not done: action policy(state) next_state, _, done, _ env.step(action) # 更新占用度量考虑折扣因子 x, y state occupancy[y, x, action] (gamma ** t) state next_state t 1 # 归一化处理 occupancy / np.sum(occupancy) return occupancy为了更直观地理解占用度量我们可以为每个状态绘制动作分布def plot_occupancy_measure(occupancy, title): 可视化占用度量 fig, ax plt.subplots(figsize(8, 8)) # 绘制网格 for i in range(env.size 1): ax.axhline(i, colorblack, lw1) ax.axvline(i, colorblack, lw1) # 绘制每个状态的动作分布 for y in range(env.size): for x in range(env.size): total np.sum(occupancy[y, x, :]) if total 0: actions occupancy[y, x, :] / total # 绘制箭头表示动作偏好 arrow_params {head_width: 0.1, head_length: 0.1, fc: red, ec: red} if actions[0] 0.1: # 上 ax.arrow(x 0.5, y 0.7, 0, -0.4 * actions[0], **arrow_params) if actions[1] 0.1: # 下 ax.arrow(x 0.5, y 0.3, 0, 0.4 * actions[1], **arrow_params) if actions[2] 0.1: # 左 ax.arrow(x 0.7, y 0.5, -0.4 * actions[2], 0, **arrow_params) if actions[3] 0.1: # 右 ax.arrow(x 0.3, y 0.5, 0.4 * actions[3], 0, **arrow_params) ax.set_xlim(0, env.size) ax.set_ylim(0, env.size) ax.set_title(title) ax.set_xticks([]) ax.set_yticks([]) plt.show() # 计算并可视化两种策略的占用度量 random_occupancy compute_occupancy_measure(env, random_policy) goal_occupancy compute_occupancy_measure(env, goal_oriented_policy) plot_occupancy_measure(random_occupancy, 随机策略的占用度量) plot_occupancy_measure(goal_occupancy, 目标导向策略的占用度量)从占用度量的可视化中我们可以观察到随机策略每个状态下的动作分布相对均匀箭头大小相近目标导向策略明显倾向于选择朝向终点的动作特别是在靠近终点的区域4. 动态演示策略如何随时间改变访问分布为了更生动地展示策略对状态访问分布的影响我们可以创建一个动画展示智能体在多个episode中的轨迹如何逐渐形成特定的访问模式。def animate_learning_process(env, policy, episodes50, interval200): 创建学习过程的动画 fig, ax plt.subplots(figsize(8, 8)) visitation np.zeros((env.size, env.size)) def update(frame): ax.clear() # 运行一个episode并更新访问分布 state env.reset() done False t 0 while not done: action policy(state) next_state, _, done, _ env.step(action) visitation[state[1], state[0]] (0.9 ** t) state next_state t 1 # 绘制当前访问分布 im ax.imshow(visitation, cmapviridis, vmin0, vmaxnp.max(visitation)) ax.set_title(fEpisode {frame 1}) # 绘制网格线和特殊格子 for i in range(env.size 1): ax.axhline(i - 0.5, colorwhite, lw1) ax.axvline(i - 0.5, colorwhite, lw1) for obs in env.obstacles: ax.add_patch(plt.Rectangle((obs[0] - 0.5, obs[1] - 0.5), 1, 1, fillTrue, colorblack)) ax.add_patch(plt.Rectangle((env.goal[0] - 0.5, env.goal[1] - 0.5), 1, 1, fillTrue, colorgreen)) return [im] ani animation.FuncAnimation(fig, update, framesepisodes, intervalinterval, blitTrue) plt.close() return ani # 创建并保存动画 random_ani animate_learning_process(env, random_policy) goal_ani animate_learning_process(env, goal_oriented_policy) # 需要保存动画时可以取消下面的注释 # random_ani.save(random_policy.gif, writerpillow) # goal_ani.save(goal_policy.gif, writerpillow)观看这些动画你会清晰地看到随机策略如何逐渐填满整个状态空间目标导向策略如何快速聚焦于通往终点的路径障碍物周围如何形成阴影区访问频率较低5. 实际应用从占用度量反推策略一个有趣的事实是给定一个合法的占用度量我们可以唯一地确定产生它的策略。这在模仿学习中特别有用——通过观察专家的状态-动作分布我们可以尝试恢复其策略。让我们实现这个逆向过程def occupancy_to_policy(occupancy): 从占用度量推导策略 policy np.zeros((env.size, env.size, 4)) for y in range(env.size): for x in range(env.size): state_occupancy occupancy[y, x, :] total np.sum(state_occupancy) if total 0: policy[y, x, :] state_occupancy / total else: policy[y, x, :] 0.25 # 未访问的状态使用均匀分布 return policy # 从目标导向策略的占用度量反推策略 recovered_policy occupancy_to_policy(goal_occupancy) # 比较原始策略和恢复的策略在特定状态下的动作分布 sample_state (2, 3) print(原始策略在状态(2,3)的动作分布:, goal_oriented_policy(sample_state)) print(恢复策略在状态(2,3)的动作分布:, recovered_policy[3, 2, :])这个例子展示了占用度量如何完整地编码策略信息。在实际应用中这种方法可以用于模仿学习从专家演示中学习策略逆向强化学习推断专家的奖励函数策略蒸馏将复杂策略压缩为更简单的形式6. 高级可视化3D状态访问分布为了进一步提升直观理解我们可以将状态访问分布以3D形式呈现from mpl_toolkits.mplot3d import Axes3D def plot_3d_visitation(visitation, title): 3D可视化状态访问分布 fig plt.figure(figsize(10, 8)) ax fig.add_subplot(111, projection3d) x np.arange(env.size) y np.arange(env.size) X, Y np.meshgrid(x, y) # 创建3D柱状图 dx dy 0.5 dz visitation.flatten() ax.bar3d(X.flatten(), Y.flatten(), np.zeros_like(dz), dx, dy, dz, shadeTrue) ax.set_title(title) ax.set_xlabel(X坐标) ax.set_ylabel(Y坐标) ax.set_zlabel(访问频率) plt.show() plot_3d_visitation(random_visitation, 随机策略的3D状态访问分布) plot_3d_visitation(goal_visitation, 目标导向策略的3D状态访问分布)3D可视化提供了另一种视角柱状图高度直观显示访问频率差异可以清晰比较不同区域的访问密度特别适合展示策略间的对比7. 扩展到连续状态空间虽然我们使用了离散的网格世界作为示例但这些概念同样适用于连续状态空间。在连续空间中我们通常使用密度函数来表示状态访问分布和占用度量。以下是连续空间中的近似计算方法def continuous_visitation(policy, episodes1000, gamma0.9): 连续空间中的状态访问分布近似 from scipy.stats import gaussian_kde states [] weights [] for _ in range(episodes): state env.reset() # 假设env现在输出连续状态 done False t 0 while not done: action policy(state) next_state, _, done, _ env.step(action) states.append(state) weights.append(gamma ** t) state next_state t 1 # 使用核密度估计 states np.array(states) kde gaussian_kde(states.T, weightsweights) return kde在连续空间中可视化通常采用等高线图热力图3D曲面图散点图带透明度这些技术可以帮助我们理解更复杂的强化学习系统如机器人控制或自动驾驶等应用场景。