告别官方示例：用Gymnasium从零搭建一个‘贪吃蛇’强化学习环境（Python+PyGame）

从零构建贪吃蛇强化学习环境Gymnasium实战指南在强化学习领域标准化的环境库让研究者能专注于算法开发而无需反复造轮子。Gymnasium作为Gym的现代继承者提供了更完善的接口和功能。本文将带你深入一个具体案例——用Gymnasium框架从头实现贪吃蛇游戏环境这比官方网格世界示例更具挑战性和趣味性。1. 环境设计基础架构任何Gymnasium自定义环境都需要继承gymnasium.Env基类并实现几个核心方法。贪吃蛇环境的结构框架如下import numpy as np import pygame import gymnasium as gym from gymnasium import spaces class SnakeEnv(gym.Env): metadata {render_modes: [human, rgb_array], render_fps: 10} def __init__(self, render_modeNone, grid_size10): # 初始化空间定义和游戏参数 pass def reset(self, seedNone, optionsNone): # 重置游戏状态 pass def step(self, action): # 处理动作并返回(observation, reward, terminated, truncated, info) pass def render(self): # 可视化环境状态 pass def close(self): # 清理资源 pass1.1 动作与观察空间设计贪吃蛇的动作空间相对简单——只有四个移动方向self.action_space spaces.Discrete(4) # 0:上, 1:右, 2:下, 3:左观察空间的设计则更有讲究。常见方案有两种完整网格表示将整个游戏板编码为矩阵每个单元格标记为空、蛇身或食物精简特征表示仅包含蛇头位置、方向、食物相对位置等关键信息第一种方案更直观但维度高第二种计算效率更高# 精简特征表示示例 self.observation_space spaces.Dict({ head: spaces.Box(0, grid_size-1, shape(2,), dtypeint), direction: spaces.Discrete(4), food: spaces.Box(0, grid_size-1, shape(2,), dtypeint), body_length: spaces.Discrete(grid_size**2) })2. 游戏逻辑实现2.1 状态初始化与重置reset()方法需要初始化或重置游戏状态def reset(self, seedNone, optionsNone): super().reset(seedseed) # 初始化蛇的位置和方向 self.snake [np.array([self.grid_size//2, self.grid_size//2])] self.direction self.np_random.integers(0, 4) # 随机生成食物位置 self._place_food() # 初始化游戏状态变量 self.score 0 self.steps 0 return self._get_obs(), self._get_info()其中_place_food()需要确保食物不会出现在蛇身上def _place_food(self): while True: self.food self.np_random.integers(0, self.grid_size, size2) if not any(np.array_equal(self.food, segment) for segment in self.snake): break2.2 核心游戏逻辑step()方法包含游戏的核心规则def step(self, action): # 0. 验证动作有效性 if action not in [0,1,2,3]: raise ValueError(Invalid action) # 1. 更新蛇的方向不能直接反向移动 if (action self.direction) % 2 ! 0: # 非相反方向 self.direction action # 2. 移动蛇头 direction_map { 0: np.array([0, -1]), # 上 1: np.array([1, 0]), # 右 2: np.array([0, 1]), # 下 3: np.array([-1, 0]) # 左 } new_head self.snake[0] direction_map[self.direction] # 3. 检查碰撞 terminated ( new_head[0] 0 or new_head[0] self.grid_size or # 撞墙 new_head[1] 0 or new_head[1] self.grid_size or any(np.array_equal(new_head, segment) for segment in self.snake[1:]) # 撞自身 ) # 4. 处理食物和蛇身更新 if np.array_equal(new_head, self.food): # 吃到食物 self.snake.insert(0, new_head) self._place_food() reward 1.0 self.score 1 else: # 正常移动 self.snake.insert(0, new_head) self.snake.pop() reward 0.0 # 5. 检查步数限制防止无限循环 self.steps 1 truncated self.steps self.max_steps return self._get_obs(), reward, terminated, truncated, self._get_info()3. 奖励函数设计奖励函数是强化学习环境设计的核心难点。贪吃蛇常见的奖励设计策略包括奖励类型实现方式优点缺点稀疏奖励吃到食物1死亡-1其他0简单直接学习效率低距离奖励根据与食物距离变化给予奖励提供中间信号可能陷入局部最优生存奖励每存活一步给予小奖励鼓励长期生存可能导致保守策略一个平衡的混合奖励函数示例def _calculate_reward(self, new_head, terminated): if terminated: return -1.0 # 死亡惩罚 if np.array_equal(new_head, self.food): return 1.0 # 吃到食物奖励 # 计算与食物的曼哈顿距离变化 old_dist np.sum(np.abs(self.snake[0] - self.food)) new_dist np.sum(np.abs(new_head - self.food)) distance_reward (old_dist - new_dist) * 0.01 # 小生存奖励 survival_reward 0.001 return distance_reward survival_reward4. 可视化与调试PyGame渲染实现需要注意性能优化def render(self): if self.render_mode rgb_array: return self._render_frame() if self.render_mode human: self._render_frame() def _render_frame(self): if self.window is None and self.render_mode human: pygame.init() pygame.display.init() self.window pygame.display.set_mode((self.window_size, self.window_size)) self.clock pygame.time.Clock() canvas pygame.Surface((self.window_size, self.window_size)) canvas.fill((255, 255, 255)) # 白色背景 # 计算网格单元大小 cell_size self.window_size / self.grid_size # 绘制食物 pygame.draw.rect( canvas, (255, 0, 0), # 红色食物 pygame.Rect( self.food * cell_size, (cell_size, cell_size), ), ) # 绘制蛇身 for i, segment in enumerate(self.snake): color (0, 0, 255) if i 0 else (0, 255, 0) # 蓝色头部绿色身体 pygame.draw.rect( canvas, color, pygame.Rect( segment * cell_size, (cell_size, cell_size), ), ) # 绘制网格线 for x in range(self.grid_size 1): pygame.draw.line( canvas, 0, (0, x * cell_size), (self.window_size, x * cell_size), width1, ) pygame.draw.line( canvas, 0, (x * cell_size, 0), (x * cell_size, self.window_size), width1, ) if self.render_mode human: self.window.blit(canvas, canvas.get_rect()) pygame.event.pump() pygame.display.update() self.clock.tick(self.metadata[render_fps]) else: return np.transpose( np.array(pygame.surfarray.pixels3d(canvas)), axes(1, 0, 2) )5. 环境注册与使用完成环境开发后需要注册才能通过gymnasium.make()使用from gymnasium.envs.registration import register register( idSnake-v1, entry_pointsnake_env:SnakeEnv, max_episode_steps1000, )使用时可以灵活配置参数import gymnasium as gym # 创建10x10的贪吃蛇环境 env gym.make(Snake-v1, grid_size10, render_modehuman) # 典型训练循环 observation, info env.reset() for _ in range(1000): action env.action_space.sample() # 随机策略 observation, reward, terminated, truncated, info env.step(action) if terminated or truncated: observation, info env.reset() env.close()在实际项目中我发现观察空间的设计对训练效果影响极大。精简特征表示虽然效率高但可能需要更复杂的神经网络来学习完整网格表示则可以直接使用CNN处理但训练速度较慢。根据具体需求选择合适的表示方式很关键。