CUDA Env Supports with concat_vec_envs_v0

[Jarvis-K]

Hi, I am trying to concat my environments with concat_vec_envs_v0. In my environment, the env.step function adopt CUDA to faster gradient calculation.

def step(self,action):
        with torch.enable_grad():
            self.model.zero_grad()
            obj_value = self.obj_function(self.model)
            obj_value.backward()

When I try to concat the environment to multi-process,

env = ss.concat_vec_envs_v0(env, 2, num_cpus=40, base_class='stable_baselines3')

It fails and the error msg is pasted below:

Traceback (most recent call last):
  File "quadraJoint.py", line 79, in <module>
    env = ss.concat_vec_envs_v0(env, 2, num_cpus=40, base_class='stable_baselines3')
  File "/supersuit/vector/vector_constructors.py", line 49, in concat_vec_envs_v0
    vec_env = MakeCPUAsyncConstructor(num_cpus)(*vec_env_args(vec_env, num_vec_envs))
  File "/supersuit/vector/constructors.py", line 38, in constructor
    return ProcConcatVec(cat_env_fns, obs_space, act_space, num_fns * envs_per_env, example_env.metadata)
  File "/supersuit/vector/multiproc_vec.py", line 99, in __init__
    env_nums = self._receive_info()
  File "/supersuit/vector/multiproc_vec.py", line 131, in _receive_info
    raise e
RuntimeError: CUDA error: initialization error

Would you like to help me address this error? Will the CUDA multiprocessing be supported in the supersuit?

[Jarvis-K]

I find a solution. In case someone meets a similar issue, the key is to set a spawn start.

mp.set_start_method('spawn')

Moreover, for users who adopt supersuit==3.2.0, the shared array bug should be noticed and update your version to the latest.

[benblack769]

Sorry about this issue, we have never tested with a GPU accelerated environment before. Honestly, if the GPU is under full load, I don't expect you will see much performance improvement under multiprocessing, but if not, then this is an interesting way of using this wrapper. Kind of cool.

[jkterry1]

@Jarvis-K would you mind letting me see the environment you're working on, even if it's just inviting me since if it's closed source. I would really like to see a hardware accelerated PettingZoo environment,

[Jarvis-K]

Sure. I am working on the Learning to Optimize Environment. This can be much different with the current MARL community. The environment needs to calculate the current gradient information for the agent in the step() function. This can be quickly done with GPU but slow on the CPU.

If you are interested, you can find the environment code here, the lines 341 to 404 are relevant. And the training code is here. But all the things are in progress, sorry for the poor documents and usage.