多进程训练Super Mario Bros RL模型时GrayScaleObservation与DummyVecEnv报错求助

解决方案 #

一、报错修复 #

1. 修复CustomRewardWrapper核心错误

CustomRewardWrapper中调用self.env.step([action])是错误的——单个环境的step方法仅接受单个action值，而非列表。这是引发后续所有维度不匹配错误的根源。
修改代码：

def step(self, action):
    # 移除action外的列表包裹，直接传入单个action
    obs, reward, done, info = self.env.step(action)
    custom_reward = 0
    if done and info["life"] == 0:
        custom_reward = -100
    elif self.currentInfo["life"] != info["life"]:
        custom_reward = -50 if self.currentInfo["life"] > info["life"] else 50
    elif self.currentInfo["world"] < info["world"]:
        custom_reward = (info["world"] - self.currentInfo["world"]) * 50
    elif self.currentInfo["stage"] < info["stage"]:
        custom_reward = (info["stage"] - self.currentInfo["stage"]) * 10
    else: 
        custom_reward = (info["x_pos_screen"] - self.currentInfo["x_pos_screen"]) * 2 - (info["time"] - self.currentInfo["time"]) * 1
    
    self.currentInfo = info
    return obs, custom_reward, done, info

2. 修复GrayScaleObservation报错

在多进程环境中，GrayScaleObservation需在每个子进程的单个环境中单独包装，无需嵌套DummyVecEnv。修改make_env函数：

def make_env(rank: int, seed: int = 0):
    def _init():
        env = gym_super_mario_bros.make('SuperMarioBros-v0')
        env.reset(seed=seed + rank)
        env = JoypadSpace(env, [
            ['NOOP'],
            ['right'],
            ['right', 'A'],
            ['right', 'B'],
            ['right', 'A', 'B'],
            ['A'],
            ['down'],
            ['left'],
        ])
        env = CustomRewardWrapper(env)
        # 启用灰度处理，keep_dim=True保留通道维度，适配CnnPolicy的输入要求
        env = GrayScaleObservation(env, keep_dim=True)
        return env
    set_random_seed(seed)
    return _init

3. 修复DummyVecEnv报错

禁止在SubprocVecEnv的子进程中嵌套DummyVecEnv——SubprocVecEnv本身就是多进程向量环境，每个子进程应返回单个非向量环境。若需使用单进程向量环境，直接替换为DummyVecEnv即可，无需混用两种向量环境。

二、内存占用过高问题解决 #

设置的n_steps=10000是内存爆炸的关键：该参数表示每个环境收集的步数，总收集步数为n_steps × num_envs（4个环境时为40000步），导致模型更新时的经验池远超内存容量。
修改方案：
将n_steps调整为合理值（如1024或2048，根据内存/显存大小调整）：

model = PPO('CnnPolicy', env, verbose=1, tensorboard_log=LOG_DIR, learning_rate=0.00003, n_steps=2048)

三、训练加速建议 #

1. 合理设置进程数：进程数不宜超过CPU核心数的80%，避免进程调度开销过大：

num_of_cpus = os.cpu_count()
if num_of_cpus is None:
    num_of_cpus = 4
else:
    num_of_cpus = max(1, num_of_cpus - 2)

2. 启用帧堆叠：在向量环境外添加VecFrameStack，提升模型对时间信息的捕捉能力，且计算量增加有限：

env = SubprocVecEnv([make_env(i) for i in range(num_of_cpus)])
env = VecMonitor(env)
# 堆叠4帧，适配CnnPolicy对时序信息的需求
env = VecFrameStack(env, n_stack=4)

3. GPU加速：确保安装支持GPU的PyTorch版本，Stable Baselines3会自动将模型和数据迁移到GPU，大幅提升训练速度。
4. 调整batch_size：可将batch_size设置为n_steps × num_envs // 4，让模型更新更频繁，同时平衡内存压力。
5. 定期清理内存：导入gc模块，在训练回调中定期调用gc.collect()，清理无用内存对象（避免过于频繁调用）。

内容的提问来源于stack exchange，提问作者Mr. Blue