强化学习Atari Breakout可视化视频录制异常问题求助

我正在学习强化学习，选择Atari Breakout环境实践。训练完神经网络后，尝试生成MP4格式可视化视频展示结果，代码如下：

import os

os.environ["OMP_NUM_THREADS"] = "1"
import torch

torch.set_num_threads(1)
torch.set_num_interop_threads(1)

import imageio
from numpy import random

import gymnasium as gym
from gymnasium.wrappers import AtariPreprocessing, FrameStackObservation
import ale_py
gym.register_envs(ale_py)

from AtariBreakout_Refactored import DQN

device = torch.device('cuda')

env = gym.make("BreakoutNoFrameskip-v4", render_mode="rgb_array")  # , render_mode="human")
env = AtariPreprocessing(env)
env = FrameStackObservation(env, 4)


state, _ = env.reset()
state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)

file = "/home/leo/Documents/ReinforcementLearning/projects/PyTorch_tests_0/saved/20250410_001501 20 episodes.pt"
q_net = torch.load(file, weights_only=False)


fps = 30
frames = []

for step in range(3000):  # 3000 frames
   ## Calculate action option 1
   # action = random.randint(4) ## <- If using this line, no issues

   ## Calculate action option 2
   with torch.no_grad():
       q_predicted = q_net(state)
       action = q_predicted.max(1).indices.view(1, 1)  ## <- If using this line, the simulation crashes
       action = int(action) ## Should not be necessary, but just to make sure
   torch.cuda.synchronize()

   state, _, terminated, truncated, _ = env.step(action)
   state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
   torch.cuda.synchronize()

   frames.append(env.render())

   if terminated or truncated:
       state, _ = env.reset()
       terminated, truncated = False, False
       state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
       torch.cuda.synchronize()

env.close()

video_filename = "/home/leo/Documents/ReinforcementLearning/projects/PyTorch_tests_0/saved/breakout_play.mp4"

imageio.mimsave(video_filename, frames, fps=fps, codec="libx264")
print(f"Video saved as {video_filename}") 

问题现象 #

代码执行无报错，但打开MP4文件时出现异常：

• 使用Numpy随机函数生成动作时，视频正常生成，AI表现虽差但画面正常；
• 使用网络输出动作时，循环会执行完3000帧，但2-3局后球会消失，仅能看到控制器循环移动。

已尝试的排查方法 #

• 改用TensorFlow网络替代PyTorch，问题完全相同；
• 强制单线程运行并添加torch.cuda.synchronize()，排查线程竞争；
• 将网络移至CPU运行，排查VRAM冲突；
• 使用未训练网络（行为接近随机），仍会出现该问题；
• 添加动作合法性检查，未发现异常动作：

if action not in [0,1,2,3]:
    print(f"Action is {action}!") ##<-Never prints anything
    break

• 改用env.action_space.sample()生成随机动作；
• 更换视频编码格式；
• 使用gym.wrappers.RecordVideo内置录屏功能。

以上方法均未解决问题。动作仅为0-3的整数，生成方式理论上不应影响环境，但随机动作从未出错，网络生成动作则必现问题，疑似env.reset()未被正确调用，但代码逻辑无明显问题。

补充说明 #

1. 所有视频均可通过VLC正常打开，排除编码问题；
2. 网络生成动作的视频文件更小（117kb vs 199kb），但帧数相同；
3. 球消失后，操控杆有时冻结，有时重复小幅移动，说明视频未冻结。

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