Problem Running Vectorized Gymnasium Environments in PyTorch

jpmartin42 commented 2 months ago

Description

When running examples, I was having issues running vectorized Gynasium environments.

For instance, running the following script worked fine:

import gymnasium as gym
from skrl.envs.wrappers.torch import wrap_env

env = gym.make("CartPole-v1")

env = wrap_env(env)

but the following script:

import gymnasium as gym
from skrl.envs.wrappers.torch import wrap_env

env = gym.vector.make("CartPole-v1", num_envs=4, asynchronous=False)
env = wrap_env(env)

... returned the following error:

Traceback (most recent call last):
  File "torch_gymnasium_pendulum_vector_ddpg.py", line 60, in <module>
    env = wrap_env(env)
  File "/home/josh/anaconda3/envs/raisim_python3/lib/python3.8/site-packages/skrl/envs/wrappers/torch/__init__.py", line 154, in wrap_env
    raise ValueError(f"Unknown wrapper type: {wrapper}")
ValueError: Unknown wrapper type: ['gymnasium.vector.vector_env.VectorEnv']

Note that, if I switch over to using gym.vec_make() instead of gym.vector.make(), then I get the same error as above except that the unknown wrapper type becomes ['gymnasium.experimental.vector.vector_env.VectorEnv']

I traced this issue back to __init__.py in wrappers/torch: when checking for Gymnasium environment types (line 100), it only checks for the single environment and not the vectorized environment version. Adding to the line to become the following fixed my issue:

elif _in("gymnasium.core.Env", base_classes) or _in("gymnasium.core.Wrapper", base_classes) or _in("gymnasium.vector.vector_env.VectorEnv", base_classes) or _in("gymnasium.experimental.vector.vector_env.VectorEnv"):

Following this change, I'm able to run the script above and run the training examples with vectorized Gymnasium environments.

Not sure if I just made a simple mistake to give myself this error, but the fix was quick and I thought I'd pass it along just in case!