本文详解如何将 gymnasium 中的向量型 cartpole 环境封装为返回 rgb 渲染图像的 wrapper，并支持帧堆叠（如 4 帧），使其兼容 stable-baselines3 的 dqn 等基于视觉输入的强化学习算法。
本文详解如何将 gymnasium 中的向量型 cartpole 环境封装为返回 rgb 渲染图像的 wrapper，并支持帧堆叠（如 4 帧），使其兼容 stable-baselines3 的 dqn 等基于视觉输入的强化学习算法。
在使用 Stable-Baselines3 实现 DQN 时，Atari 类环境天然提供图像观测（rgb_array 渲染模式），而经典控制环境（如 CartPole-v1）默认仅输出 4 维状态向量（位置、速度、角度、角速度）。为统一预处理流程（如灰度化、缩放、帧堆叠），需自定义 wrapper 将 CartPole 渲染为图像，并确保其与 VecFrameStack、VecTransposeImage 等向量化工具链兼容。
核心挑战在于：gymnasium.Env.reset() 接口要求支持 seed 参数（用于可复现实验），但 CartPoleEnv 在旧版本或未显式适配时可能不接受该参数，导致 VecEnv 调用失败（报错 TypeError: reset() got an unexpected keyword argument 'seed'）。解决方案是让 wrapper 的 reset() 方法接收任意 **kwargs 并透传给底层环境——现代 gymnasium 已全面支持 seed，因此只需安全忽略不识别的参数即可。
以下是完整、健壮的实现：
import gymnasium as gym
import numpy as np
import cv2
from gymnasium import spaces
from stable_baselines3.common.vec_env import VecFrameStack, VecTransposeImage, DummyVecEnv
from stable_baselines3.common.env_util import make_vec_env
class CartPoleImageWrapper(gym.Wrapper):
def __init__(self, env):
super().__init__(env)
# 定义输出观测空间：84×84 单通道灰度图（uint8）
self.observation_space = spaces.Box(
low=0, high=255,
shape=(84, 84, 1),
dtype=np.uint8
)
def _get_image_observation(self) -> np.ndarray:
"""渲染当前帧 → 缩放 → 灰度化 → 添加通道维度"""
# 注意：render() 返回 ndarray，shape 为 (H, W, 3) 或 (H, W, 4)
img = self.render()
if img is None:
raise RuntimeError("CartPole environment must be created with render_mode='rgb_array'")
# OpenCV 处理：缩放 + 灰度 + 扩展通道
img_resized = cv2.resize(img, (84, 84))
img_gray = cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
return np.expand_dims(img_gray, axis=-1) # shape: (84, 84, 1)
def reset(self, **kwargs):
# 安全调用底层 reset，兼容 seed 等新参数
super().reset(**kwargs)
return self._get_image_observation()
def step(self, action):
obs, reward, terminated, truncated, info = self.env.step(action)
# 注意：gymnasium 的 step 返回 5 元组（含 truncated）
return self._get_image_observation(), reward, terminated, truncated, info
# ✅ 正确构建环境链（关键：render_mode 必须设为 'rgb_array'）
env = CartPoleImageWrapper(
gym.make("CartPole-v1", render_mode="rgb_array")
)
# 使用 DummyVecEnv（非 make_vec_env）避免额外包装干扰
vec_env = DummyVecEnv([lambda: env])
vec_env = VecTransposeImage(vec_env) # (H,W,C) → (C,H,W)
vec_env = VecFrameStack(vec_env, n_stack=4) # 堆叠 4 帧，输出 shape: (1, 4, 84, 84)
obs = vec_env.reset()
print(f"Final observation shape: {obs.shape}") # → (1, 4, 84, 84)
⚠️ 关键注意事项：
- render_mode="rgb_array" 是必需的，否则 render() 返回 None；
- VecTransposeImage 必须在 VecFrameStack 之前应用，否则通道顺序错误；
- 不要使用 make_vec_env(lambda: env, ...)，它会自动添加 Monitor wrapper，可能引发 seed 冲突；推荐 DummyVecEnv 显式构造；
- 若使用 gym（非 gymnasium），需将 step() 返回值改为 4 元组（移除 truncated），并调整 reset() 签名（seed=None）；
- 图像预处理（如归一化）应在网络输入层或自定义 NormalizeObservation wrapper 中完成，而非此处（保持 uint8 原始范围更符合标准流程）。
通过该 wrapper，CartPole 即可无缝接入原 Atari DQN 流程：从 VecFrameStack 获取 4 帧历史图像，直接送入 CNN 主干网络，实现状态表征的一致性与算法复用性。
