从零定制强化学习环境:GridWorld 悬崖行走 + PPO 训练全解析

前言
本文以 GridWorld 悬崖行走(Cliff Walking)变体 为例,从环境建模、观测空间设计、奖励工程、多模式渲染到 PPO 训练与可视化,逐行拆解一个完整的 Gymnasium 自定义环境,并附训练曲线解读。读完这篇,你应该能独立设计自己的 RL 环境。
完整代码约 200 行,建议边看边跑。
一、任务定义
这是一个 5×5 的网格世界:
- Agent(蓝色圆) :从随机位置出发,需要到达右下角的 Goal(绿色方块 G)
- Cliff(红色方块 C):地图中存在一道悬崖区域,踩上去直接失败
- 边界:撞墙会受惩罚并终止回合
- 动作:上下左右 4 个离散动作
任务示意:
erlang
. . . . .
. . . . .
. . . . .
. C C C G
. . . . A
与标准 Cliff Walking 的核心差异:
| 标准 Cliff Walking | 本文变体 | |
|---|---|---|
| 网格大小 | 4×12 | 5×5 |
| 悬崖位置 | 底行中间 | 第 4 行中间 |
| 起点 | 固定左下角 | 随机(不在悬崖/目标上) |
| 目标 | 固定右下角 | 固定右下角 |
| 观测空间 | 一维索引 | Dict(结构化多模态) |
| 边界处理 | 无(无出界) | 撞墙惩罚 + 终止 |
设计意图:结构化的 Dict 观测空间 + 撞墙惩罚,使这个环境更接近真实机器人场景------需要同时处理坐标信息、障碍物地图和边界感知。
二、依赖安装
bash
pip install gymnasium numpy stable-baselines3 pygame tensorboard
gymnasium:环境标准接口numpy:数值计算stable-baselines3:PPO 算法实现pygame:human 模式下的图形化渲染tensorboard:训练曲线可视化
三、环境代码逐段详解
3.1 类定义与元数据
python
import gymnasium as gym
import numpy as np
from gymnasium import spaces
from typing import Optional, Dict, Any, SupportsFloat
class GridWorldEnv(gym.Env):
metadata = {"render_modes": ["human", "ansi"], "render_fps": 4}
metadata告诉上层框架本环境支持两种渲染模式 :human:使用 Pygame 图形化展示ansi:纯文本输出,适合无 GUI 的服务器环境
render_fps=4:控制 human 模式的刷新率,这里故意设慢以便观察 Agent 行为
3.2 __init__:初始化空间与悬崖配置
python
def __init__(self, size=5, render_mode="ansi") -> None:
super().__init__()
self.size = size
assert render_mode in ("human", "ansi")
self.render_mode = render_mode
self.n_step = 0
self.MAX_STEP = self.size ** 2 # 最大步数 = 25
# 悬崖区域:第4行(x=3),列1~3(非边界)
self.cliff = np.array(
[(3, y) for y in range(1, size - 1)], dtype=np.int32
)
self.cliff_set = {tuple(map(int, p)) for p in self.cliff}
设计要点:
MAX_STEP = size²:最大回合步数等于网格面积。对于 5×5 网格即 25 步,这是一个紧约束------Agent 必须高效移动,不能漫无目的游荡cliff数组:位于第 4 行的中间 3 格(列 1、2、3),排除了左右边界列cliff_set:转为 Python set 用于 O(1) 碰撞检测
3.3 观测空间:Dict 结构化设计(核心亮点)
python
self.observation_space = spaces.Dict({
"agent": spaces.Box(0, 1, shape=(2,), dtype=np.float32),
"target": spaces.Box(0, 1, shape=(2,), dtype=np.float32),
"cliff": spaces.Box(0, 1, shape=(len(self.cliff), 2), dtype=np.float32),
"boundary_flag": spaces.Discrete(5)
})
这是本文环境最值得学习的部分。传统做法是把所有信息展平成一个一维向量,但这里采用了 Dict 观测空间 ,将异构信息按语义拆分:
| 字段 | 空间类型 | 维度 | 语义 |
|---|---|---|---|
agent |
Box | (2,) | Agent 归一化坐标 (x,y),取值范围 0,1 |
target |
Box | (2,) | 目标归一化坐标 |
cliff |
Box | (3, 2) | 三个悬崖格的归一化坐标,相当于局部地图 |
boundary_flag |
Discrete | 5 | 是否贴边:0=无,1=上边,2=下边,3=左边,4=右边 |
为什么要用 Dict?
- 语义清晰 :网络可以分别处理不同类型的信息,例如用卷积处理
cliff地图,用 MLP 处理坐标 - 可扩展:新增传感器(如雷达、图像)只需加字段,不影响已有结构
- 与算法库兼容 :Stable Baselines3 的
MultiInputPolicy可自动处理 Dict 观测,内部会为每个字段创建独立的特征提取器
boundary_flag 的设计意图:
这是一个非常巧妙的设计------Agent 的坐标如果靠近边界,撞墙的概率增大。通过显式给出边界信息,Agent 可以学习"提前减速"或"转弯"。
实际观测生成:
python
def _obs(self):
ax, ay = self._agent
flag = 0
if ax == 0: flag = 1 # 贴上方
elif ax == self.size - 1: flag = 2 # 贴下方
elif ay == 0: flag = 3 # 贴左侧
elif ay == self.size - 1: flag = 4 # 贴右侧
return {
"agent": self._agent.astype(np.float32) / (self.size - 1),
"target": self._target.astype(np.float32) / (self.size - 1),
"cliff": self.cliff.astype(np.float32) / (self.size - 1),
"boundary_flag": flag
}
所有坐标都除以 (size-1) 归一化到 0,1,这是神经网络输入的标准做法,避免不同尺度的数值导致梯度不稳定。
3.4 动作空间
python
self.action_space = spaces.Discrete(4)
self._action_to_dict = {
0: np.array([0, -1]), # 上
1: np.array([0, +1]), # 下
2: np.array([-1, 0]), # 左
3: np.array([+1, 0]) # 右
}
4 个离散动作,对应上下左右移动一格。注意这里的坐标约定:(x,y) 中 x 为行(纵向),y 为列(横向)。
3.5 reset:回合初始化
python
def reset(self, *, seed=None, options=None):
super().reset(seed=seed, options=options)
self.n_step = 0
self._target = np.array(
[self.size - 1, self.np_random.integers(0, self.size)], dtype=np.int32
)
while True:
self._agent = np.array([
self.np_random.integers(0, self.size),
self.np_random.integers(0, self.size)
], dtype=np.int32)
agent_pos = tuple(map(int, self._agent))
target_pos = tuple(map(int, self._target))
if agent_pos not in self.cliff_set and agent_pos != target_pos:
break
return self._obs(), {}
关键细节:
- 目标位置:固定在最后一行,但列在 0~4 中随机选取,增加了任务的多样性
- Agent 起始位置:完全随机(25 个格子中排除悬崖和目标),确保策略不依赖固定起点
- 拒绝采样循环 :不断生成随机位置直到合法。对于 5×5 网格(排除 3 个悬崖 + 1 个目标 = 4 个非法格),期望尝试次数约 25/21≈1.2 次,开销可忽略
- 必须使用
self.np_random而非np.random,否则seed参数不会生效
3.6 step:环境核心逻辑
这是整个环境含金量最高的部分,我们分段解析。
python
def step(self, action):
terminated = False
truncated = False
termination_reason = "running"
action = int(np.asarray(action).item())
direction = self._action_to_dict[action]
old_distance = np.linalg.norm(self._agent - self._target)
candidate = self._agent + direction
action的int(np.asarray(action).item())写法兼容了 numpy 数组和 Python 标量,是从真实工程中总结出的防御性代码old_distance用于后续的势能奖励(见下文)
3.6.1 出界处理
python
reward = -0.1
if (
candidate[0] < 0 or candidate[0] >= self.size
or candidate[1] < 0 or candidate[1] >= self.size
):
reward -= 2.5
self._agent = np.clip(candidate, 0, self.size - 1)
terminated = True
- 每步基础惩罚 :
-0.1,鼓励 Agent 尽快到达目标,避免无意义绕圈 - 撞墙惩罚 :
-2.5(总计 -2.6),Agent 的行为被截断在边界内,回合终止 - 终止原因:撞墙后直接结束,这是一种硬约束设计------撞墙 = 失败
3.6.2 悬崖与目标检测
python
else:
self._agent = candidate
agent_pos = tuple(map(int, self._agent))
if agent_pos == tuple(self._target):
reward += 5.0
terminated = True
termination_reason = "goal"
elif agent_pos in self.cliff_set:
reward -= 5.0
terminated = True
termination_reason = "cliff"
else:
new_distance = np.linalg.norm(self._agent - self._target)
reward += 0.01 * (old_distance - new_distance)
奖励设计总结:
| 事件 | 奖励 | 设计意图 |
|---|---|---|
| 每步基础 | -0.1 | 时间惩罚,鼓励最短路径 |
| 靠近目标 | +0.01 × 距离缩短量 | 势能奖励(Potential-Based),提供稠密信号 |
| 到达目标 | +5.0 | 成功奖励,引导关键行为 |
| 坠落悬崖 | -5.0 | 重大失败惩罚 |
| 撞墙 | -2.5 | 边界违规惩罚 |
势能奖励(Potential-Based Reward Shaping) :
+0.01 * (old_distance - new_distance)是一种经典的奖励塑形技术。它在不改变最优策略的前提下,为每一步提供即时反馈。0.01 的系数很小,确保塑形奖励不会覆盖任务本身的奖励信号。
3.6.3 截断处理
python
self.n_step += 1
if self.n_step >= self.MAX_STEP and not terminated:
truncated = True
reward -= 1.0
termination_reason = "timeout"
truncated=True而非terminated=True:正确区分"任务结束"和"回合被截断"- 超时惩罚
-1.0:略重,意在惩罚效率过低的策略
3.6.4 info 字典
python
info = {
"pos": self._agent.copy(),
"is_success": float(termination_reason == "goal"),
"termination_reason": termination_reason,
}
return self._obs(), float(reward), terminated, truncated, info
info 中包含了调试和评估所需的关键信息,尤其是 termination_reason 可以用于统计 Agent 的失败模式分布(撞墙 vs 坠崖 vs 超时)。
3.7 多模式渲染
3.7.1 ANSI 文本模式
python
def render(self):
if self.render_mode == "ansi":
grid = [["." for _ in range(self.size)] for _ in range(self.size)]
tx, ty = self._target
grid[tx][ty] = "G"
for cx, cy in self.cliff:
grid[cx][cy] = "C"
ax, ay = self._agent
grid[ax][ay] = "A"
return "\n".join(" ".join(row) for row in grid)
输出示例:
erlang
. . . . .
. . . . .
. . . . .
. C C C G
. . . . A
在无 GUI 的远程服务器或终端中可以直接观察 Agent 行为,极大方便了服务器端调试。
3.7.2 Pygame 图形模式
python
def _render_human(self):
import pygame
# ... 初始化窗口和字体 ...
# 绘制方格线
for i in range(self.size + 1):
pygame.draw.line(screen, (80, 80, 80), (i*CELL, 0), (i*CELL, WIN))
pygame.draw.line(screen, (80, 80, 80), (0, i*CELL), (WIN, i*CELL))
# 悬崖:红色方块
for cx, cy in self.cliff:
pygame.draw.rect(screen, (180, 80, 80), ...)
# 目标:绿色方块 + "G" 文字
pygame.draw.rect(screen, (30, 180, 30), ...)
screen.blit(self._font.render("G", True, (255, 255, 255)), ...)
# Agent:蓝色圆形
pygame.draw.circle(screen, (30, 100, 255), (cx, cy), CELL//3)
配色设计:
- 深灰背景(30, 30, 30):专业暗色主题
- 红色悬崖(180, 80, 80):视觉警示,符合人类直觉
- 绿色目标(30, 180, 30):安全信号
- 蓝色 Agent(30, 100, 255):导航主体
四、PPO 训练设置
python
from stable_baselines3 import PPO
from stable_baselines3.common.env_checker import check_env
# 必须先做合规检查
env = GridWorldEnv(size=5, render_mode="ansi")
check_env(env) # 无报错则接口合规
model = PPO(
"MultiInputPolicy", # 适配 Dict 观测空间
env,
normalize_advantage=True, # 优势归一化,提升训练稳定性
verbose=1,
tensorboard_log="./tb/gridworld",
)
model.learn(total_timesteps=60000)
model.save("ppo_gridworld")
关键选择:
MultiInputPolicy:这是 SB3 中专门处理 Dict 观测空间的策略网络。内部会为每个 Dict 字段创建独立的特征提取器(agent和target各走 MLP,cliff走 Flatten 后 MLP,boundary_flag走 Embedding),最终将多路特征融合normalize_advantage=True:对优势函数做批归一化(减均值除以标准差),在小批量训练时能显著提升稳定性和收敛速度- 60,000 步训练量:对于 5×5 小网格,这个规模绰绰有余
五、训练结果分析
训练 60,000 步后,TensorBoard 输出如下曲线:
5.1 平均回合奖励(rollout/ep_rew_mean)
从约 -3 持续上升至约 4.38(S 形学习曲线),说明:
- 初期(0~15k 步):Agent 频繁撞墙、坠崖,奖励为负
- 中期(15k~40k 步):逐渐学会规避悬崖和边界,奖励快速攀升
- 后期(40k~60k 步):奖励趋于平稳在 4+,策略接近收敛
一个完美回合的奖励构成:最短路径约 8 步 × (-0.1) + 距离奖励约 +0.4 + 成功奖励 +5.0 ≈ 4.6,与观测值吻合。
5.2 训练损失(train/loss)
从约 2.2 下降至接近 0,中间在 10k~25k 步出现过剧烈波动(峰值约 2.2),这是 PPO 在学习悬崖规避时策略大幅调整的典型表现。之后损失单调下降,说明策略逐渐收敛。
5.3 价值损失(train/value_loss)
形状与 policy loss 高度同步,这表明本任务中 Actor 和 Critic 的学习难度基本一致,没有出现"Critic 拖后腿"的情况。
六、可视化验证
python
test_env = GridWorldEnv(size=5, render_mode="human")
model = PPO.load("ppo_gridworld")
for ep in range(100):
obs, _ = test_env.reset()
done = False
while not done:
action, _ = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, info = test_env.step(action)
done = terminated or truncated
test_env.render()
运行后你会看到蓝色 Agent 从随机位置出发,绕过红色悬崖区,最终到达绿色目标格。配合 pygame.time.delay(800) 在每个回合间暂停,方便观察轨迹。

七、设计回顾与最佳实践总结
7.1 本环境的特点
| 设计 | 价值 |
|---|---|
| Dict 观测空间 | 结构清晰,兼容 MultiInputPolicy,可扩展性强 |
| 势能奖励塑形 | 在不改变最优策略的前提下提供稠密信号,加速收敛 |
| boundary_flag 字段 | 显式注入边界信息,帮助 Agent 学习撞墙规避 |
| 双渲染模式 | ANSI 用于服务器调试,Pygame 用于演示展示 |
| 细粒度 info 字典 | 支持失败模式分析(区分撞墙/坠崖/超时) |
| 归一化坐标 | 所有坐标 /(size-1) 到 0,1,符合 NN 输入规范 |
7.2 自定义环境的黄金法则
- 先 check_env,再训练 :
check_env能捕获 90% 的接口问题 - 观测归一化:Box 观测永远归一化到小范围,否则梯度爆炸
- 奖励缩放:奖励绝对值建议在 -10, 10 区间,过大会导致值函数不稳定
- 区分 terminated 与 truncated:正确使用 Gymnasium v26+ 的五元返回值
- 用
self.np_random:不要用全局np.random,否则 seed 不生效 - info 要写清 :
termination_reason等信息对调试和评估至关重要 - render 不要引入副作用:渲染应该是纯只读操作,不影响环境状态
八、拓展方向
基于本环境可以进行这些变体实验:
- 连续动作:将动作改为连续二维向量,使用 SAC 或 TD3 训练
- 随机风场:增加随机扰动,测试策略鲁棒性
- 动态悬崖:悬崖位置每回合随机变化,要求更强的泛化能力
- 部分可观测:Agent 只能看到周围 3×3 范围,引入记忆机制(LSTM/Transformer)
- 多目标:多个 Goal,到达任一即可,测试优先级学习
九、完整代码
python
import gymnasium as gym
import numpy as np
from gymnasium import spaces
from typing import Optional, Dict, Any, SupportsFloat
class GridWorldEnv(gym.Env):
metadata = {"render_modes": ["human", "ansi"], "render_fps": 4}
def __init__(self, size=5, render_mode="ansi") -> None:
super().__init__()
self.size = size
assert render_mode in ("human", "ansi")
self.render_mode = render_mode
self.n_step = 0
self.MAX_STEP = self.size ** 2
self.cliff = np.array(
[(3, y) for y in range(1, size - 1)], dtype=np.int32
)
self.cliff_set = {tuple(map(int, position)) for position in self.cliff}
self.observation_space = spaces.Dict({
"agent": spaces.Box(0, 1, shape=(2,), dtype=np.float32),
"target": spaces.Box(0, 1, shape=(2,), dtype=np.float32),
"cliff": spaces.Box(0, 1, shape=(len(self.cliff), 2), dtype=np.float32),
"boundary_flag": spaces.Discrete(5)
})
self.action_space = spaces.Discrete(4)
self._action_to_dict = {
0: np.array([0, -1]),
1: np.array([0, +1]),
2: np.array([-1, 0]),
3: np.array([+1, 0])
}
def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
super().reset(seed=seed, options=options)
self.n_step = 0
self._target = np.array(
[self.size - 1, self.np_random.integers(0, self.size)], dtype=np.int32
)
while True:
self._agent = np.array([
self.np_random.integers(0, self.size),
self.np_random.integers(0, self.size)
], dtype=np.int32)
agent_pos = tuple(map(int, self._agent))
target_pos = tuple(map(int, self._target))
if agent_pos not in self.cliff_set and agent_pos != target_pos:
break
return self._obs(), {}
def step(self, action):
terminated = False
truncated = False
termination_reason = "running"
action = int(np.asarray(action).item())
direction = self._action_to_dict[action]
old_distance = np.linalg.norm(self._agent - self._target)
candidate = self._agent + direction
reward = -0.1
if (
candidate[0] < 0 or candidate[0] >= self.size
or candidate[1] < 0 or candidate[1] >= self.size
):
reward -= 2.5
self._agent = np.clip(candidate, 0, self.size - 1)
terminated = True
else:
self._agent = candidate
agent_pos = tuple(map(int, self._agent))
if agent_pos == tuple(self._target):
reward += 5.0
terminated = True
termination_reason = "goal"
elif agent_pos in self.cliff_set:
reward -= 5.0
terminated = True
termination_reason = "cliff"
else:
new_distance = np.linalg.norm(self._agent - self._target)
reward += 0.01 * (old_distance - new_distance)
self.n_step += 1
if self.n_step >= self.MAX_STEP and not terminated:
truncated = True
reward -= 1.0
termination_reason = "timeout"
info = {
"pos": self._agent.copy(),
"is_success": float(termination_reason == "goal"),
"termination_reason": termination_reason,
}
return self._obs(), float(reward), terminated, truncated, info
def render(self):
if self.render_mode == "ansi":
grid = [["." for _ in range(self.size)] for _ in range(self.size)]
tx, ty = self._target
grid[tx][ty] = "G"
for cx, cy in self.cliff:
grid[cx][cy] = "C"
ax, ay = self._agent
grid[ax][ay] = "A"
return "\n".join(" ".join(row) for row in grid)
elif self.render_mode == "human":
return self._render_human()
def close(self) -> None:
if hasattr(self, "_screen"):
import pygame
pygame.quit()
del self._screen
def _obs(self):
ax, ay = self._agent
flag = 0
if ax == 0:
flag = 1
elif ax == self.size - 1:
flag = 2
elif ay == 0:
flag = 3
elif ay == self.size - 1:
flag = 4
return {
"agent": self._agent.astype(np.float32) / (self.size - 1),
"target": self._target.astype(np.float32) / (self.size - 1),
"cliff": self.cliff.astype(np.float32) / (self.size - 1),
"boundary_flag": flag
}
def _render_human(self):
import pygame
pygame.font.init()
CELL = 60
WIN = self.size * CELL
if not hasattr(self, "_screen"):
pygame.init()
self._screen = pygame.display.set_mode((WIN, WIN))
self._clock = pygame.time.Clock()
self._font = pygame.font.Font(None, 28)
screen = self._screen
screen.fill((30, 30, 30))
for i in range(self.size + 1):
pygame.draw.line(screen, (80, 80, 80), (i*CELL, 0), (i*CELL, WIN))
pygame.draw.line(screen, (80, 80, 80), (0, i*CELL), (WIN, i*CELL))
for cx, cy in self.cliff:
pygame.draw.rect(screen, (180, 80, 80),
(cx*CELL+1, cy*CELL+1, CELL-2, CELL-2))
tx, ty = self._target
pygame.draw.rect(screen, (30, 180, 30),
(tx*CELL+1, ty*CELL+1, CELL-2, CELL-2))
txt = self._font.render("G", True, (255, 255, 255))
screen.blit(txt, (tx*CELL+CELL//2-8, ty*CELL+CELL//2-8))
ax, ay = self._agent
cx = ax*CELL + CELL // 2
cy = ay*CELL + CELL // 2
pygame.draw.circle(screen, (30, 100, 255), (cx, cy), CELL//3)
pygame.display.flip()
self._clock.tick(self.metadata["render_fps"])
if __name__ == "__main__":
from stable_baselines3 import PPO
from stable_baselines3.common.env_checker import check_env
import pygame
print(">>> check_env ...")
env = GridWorldEnv(size=5, render_mode="ansi")
check_env(env)
print(">>> check_env passed")
print(">>> training ...")
model = PPO(
"MultiInputPolicy",
env,
normalize_advantage=True,
verbose=1,
tensorboard_log="./tb/gridworld",
)
model.learn(total_timesteps=60_000)
model.save("ppo_gridworld")
env.close()
print(">>> training done")
test_env = GridWorldEnv(size=5, render_mode="human")
model = PPO.load("ppo_gridworld")
for ep in range(100):
print(f"\n=== Episode {ep} ===")
obs, _ = test_env.reset()
test_env.render()
done = False
steps = 0
step_reward = 0.0
while not done:
for event in pygame.event.get():
if event.type == pygame.QUIT:
test_env.close()
exit()
action, _ = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, info = test_env.step(action)
print(f"Step {steps}: agent={obs}, action={action}")
done = terminated or truncated
step_reward += reward
steps += 1
test_env.render()
print(f" finished in {steps} steps, reward={step_reward}")
pygame.time.delay(800)
print("\n按回车退出窗口...")
input()
test_env.close()
十、总结
从零构建一个强化学习环境并不复杂,核心在于:
- 想清楚任务定义:状态、动作、奖励、终止条件的语义
- 设计好观测空间:结构化的 Dict 观测比扁平化向量更适合复杂场景
- 奖励工程是艺术:势能奖励 + 事件奖励 + 基础惩罚的组合拳
- 双渲染模式:服务器调试用文本,演示展示用图形
- 先检查再训练 :
check_env是最后的防线