ALE/Breakout-v5环境在Jupyter Notebook中Render可视化失效问题

问题根源 #

Jupyter Notebook的图形显示机制和本地Python脚本不同，Stable Baselines3的VecEnv（向量环境）默认不支持直接弹出原生渲染窗口，加上VecFrameStack包装后，evaluate_policy的render=True参数只会触发环境内部的渲染逻辑，但无法在Notebook中显示画面。

具体解决方案 #

方案1：创建单环境手动渲染（最直观）

放弃向量环境，直接创建单实例Atari环境，手动循环执行评估步骤，用Matplotlib逐帧显示游戏画面：

import os
import matplotlib.pyplot as plt
from stable_baselines3 import A2C
from gymnasium import make
from gymnasium.wrappers import AtariPreprocessing, FrameStack

# 初始化单环境，指定rgb_array模式用于获取帧
env = make('ALE/Breakout-v5', render_mode='rgb_array')
# 添加Atari必要的预处理（裁剪、灰度化、帧跳过等）
env = AtariPreprocessing(env)
# 帧堆叠，和训练时的设置保持一致
env = FrameStack(env, num_stack=4)

# 加载预训练模型
a2c_path = os.path.join('Training', 'Saved Models', 'A2C_Breakout_Model')
model = A2C.load(a2c_path, env)

# 手动执行评估并渲染
episodes_to_eval = 10
total_rewards = []

for _ in range(episodes_to_eval):
    obs, _ = env.reset()
    done = False
    episode_reward = 0
    
    while not done:
        # 模型预测动作
        action, _ = model.predict(obs, deterministic=True)
        # 执行动作
        obs, reward, terminated, truncated, _ = env.step(action)
        episode_reward += reward
        
        # 获取当前游戏帧并显示
        frame = env.render()
        plt.imshow(frame)
        plt.axis('off')  # 隐藏坐标轴
        plt.show(block=False)  # 非阻塞显示
        plt.pause(0.01)  # 控制播放速度，数值越小越快
        plt.clf()  # 清除当前帧，避免画面重叠
        
        done = terminated or truncated
    
    total_rewards.append(episode_reward)
    print(f"单局奖励: {episode_reward}")

print(f"平均奖励: {sum(total_rewards)/len(total_rewards):.2f}")

方案2：适配向量环境的渲染

如果必须使用向量环境（比如训练时用了多环境），可以启用Jupyter的交互绘图模式，手动处理向量环境的渲染输出：

import os
import matplotlib.pyplot as plt
from stable_baselines3.common.vec_env import make_atari_env, VecFrameStack
from stable_baselines3 import A2C

# 启用Jupyter交互绘图模式
%matplotlib notebook

# 创建向量环境时指定rgb_array渲染模式
env = make_atari_env('ALE/Breakout-v5', n_envs=1, seed=0, render_mode='rgb_array')
env = VecFrameStack(env, n_stack=4)

# 加载模型
a2c_path = os.path.join('Training', 'Saved Models', 'A2C_Breakout_Model')
model = A2C.load(a2c_path, env)

# 自定义评估函数，加入渲染逻辑
def custom_evaluate(model, env, eval_episodes=10):
    total_rewards = []
    for _ in range(eval_episodes):
        obs = env.reset()
        done = [False]
        episode_reward = 0
        
        while not all(done):
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, done, _ = env.step(action)
            episode_reward += reward[0]  # 向量环境返回的奖励是数组，取第一个环境的值
            
            # 渲染第一个环境的画面
            frames = env.render()
            plt.imshow(frames[0])
            plt.axis('off')
            plt.show(block=False)
            plt.pause(0.01)
            plt.clf()
        
        total_rewards.append(episode_reward)
        print(f"单局奖励: {episode_reward}")
    
    avg_reward = sum(total_rewards)/eval_episodes
    print(f"平均奖励: {avg_reward:.2f}")
    return avg_reward

# 执行评估
custom_evaluate(model, env, eval_episodes=10)

方案3：无桌面环境下的虚拟渲染（服务器场景）

如果你的Jupyter运行在无图形界面的服务器上，需要用虚拟显示器来模拟桌面环境：

from pyvirtualdisplay import Display
import os
import matplotlib.pyplot as plt
from stable_baselines3 import A2C
from gymnasium import make
from gymnasium.wrappers import AtariPreprocessing, FrameStack

# 创建虚拟显示器
display = Display(visible=0, size=(1280, 720))
display.start()

# 后续代码和方案1/2一致
env = make('ALE/Breakout-v5', render_mode='rgb_array')
env = AtariPreprocessing(env)
env = FrameStack(env, num_stack=4)

a2c_path = os.path.join('Training', 'Saved Models', 'A2C_Breakout_Model')
model = A2C.load(a2c_path, env)

# ... 执行评估和渲染代码 ...

# 结束后关闭虚拟显示器
display.stop()

注意：需要先安装依赖：pip install pyvirtualdisplay，Ubuntu服务器还需安装sudo apt-get install xvfb。

注意事项 #

• 不要用n_eval_episodes=100来渲染，大量局数会导致Notebook卡顿，测试时先用1-10局验证效果。
• Jupyter不支持render_mode='human'的原生窗口弹出，必须用rgb_array模式结合Matplotlib显示。
• 确保训练时的环境预处理（帧堆叠、AtariPreprocessing）和评估时完全一致，避免模型预测出错。

内容的提问来源于stack exchange，提问作者Gamigj