OpenAI Gym 中的多重处理与绳降

Question

我正在使用 abseil 库在 OpenAI Gym 中进行多处理。基本上，gym.make 似乎有效。但是，我正在尝试使用不起作用的 gym-super-mario-bros。下面是一个最小的工作示例：

from absl import app
import os
os.environ['OMP_NUM_THREADS'] = '1'

import gym
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT
from nes_py.wrappers import JoypadSpace
import multiprocessing as mp
import torch
from torch import nn
import time


def get_env():
    env = JoypadSpace(gym_super_mario_bros.make('SuperMarioBros-1-1-v0'), SIMPLE_MOVEMENT)
    # env = gym.make('LunarLander-v2') # other environment such as this one and others works well
    return env

def do_something(env, net1, net2):
    print('inside do_something')
    obs = env.reset()
    print(f'after reset {obs.shape}')
    net2.load_state_dict(net1.state_dict())
    print('after load_state_dict')

def main(args):
    del args
    env = get_env()

    net1 = nn.Sequential(nn.Conv2d(1, 20, 5), nn.ReLU())
    net2 = nn.Sequential(nn.Conv2d(1, 20, 5), nn.ReLU())

    net1.share_memory()
    net2.share_memory()

    device = torch.device('cuda')
    net1 = net1.to(device)
    net2 = net2.to(device)

    p = mp.Process(target=do_something, args=(env, net1, net2,))
    p.start()

    time.sleep(4.0) # wait for the above process to execute print statements
    env.close()

if __name__ == '__main__':
    mp.set_start_method('spawn')
    app.run(main)

运行上面的代码后，打印

inside do_something

后似乎卡住了。它还在如下所示的终端中报告

CUDA warning

：

$ python mwe.py 
inside do_something
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)

请注意，在

CUDA warning

的情况下，它可以在没有任何

gym.make

的情况下顺利运行。

版本信息

我使用的是 Ubuntu 22.04.2 LTS 操作系统。以下是版本信息：

图书馆	版本
absl-py	1.3.0
cuda	11.7
健身房	0.17.2
健身房超级马里奥兄弟	7.4.0
nes-py	8.2.1
numpy	1.21.0
蟒蛇	3.9.16
手电筒	1.13.1

是否有使马里奥在多处理环境中工作的解决方法？