search

DLR-RM / stable-baselines3

PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms.
https://stable-baselines3.readthedocs.io
MIT License
8.38k stars 1.61k forks source link

[Question] SubprocVecEnv doesn't work with registered custom environments #1869

Closed marcusfechner closed 3 months ago

marcusfechner commented 3 months ago

🐛 Bug

Hi,

I found it very strange, that when I register custom environments they are not found when using SubprocVecEnv. I guess it has to do with how things are handles with multiprocessing. A temporary, but unsave hotfix would be to set SubprocVecEnv(..., start_method="fork").

I would highly appreciate any help or feedback, because this is driving me crazy. Furthermore I looked into using the gymnasium.vector.SyncVectorEnv, which works upon creation, but is not compatible with SB3.

Thanks for your help, Marcus

Package install:

pip install stable-baselines3==2.2.1 gymnasium==0.29.1 minigrid==2.3.1

Code example

from typing import Callable

import gymnasium as gym
from gymnasium.envs.registration import register
from minigrid.wrappers import ImgObsWrapper
from stable_baselines3.common.utils import set_random_seed
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv

def register_minigrid_envs():
    """Register custom Minigrid environments.
    """
    register(
        id="MiniGrid-Dynamic-Obstacles-Random-16x16-v0",
        entry_point="minigrid.envs:DynamicObstaclesEnv",
        kwargs={"size": 16, "agent_start_pos": None, "n_obstacles": 5},
    )

def make_env(env_id: str, rank: int, seed: int = 0) -> Callable:
    """
    Utility function for multiprocessed env.

    :param env_id: (str) the environment ID
    :param num_env: (int) the number of environment you wish to have in subprocesses
    :param seed: (int) the inital seed for RNG
    :param rank: (int) index of the subprocess
    :return: (Callable)
    """

    def _init() -> gym.Env:
        env = gym.make(env_id)
        env = ImgObsWrapper(env)
        env.reset(seed=seed + rank)
        return env

    set_random_seed(seed)
    return _init

if __name__ == "__main__":
    register_minigrid_envs()
    env_name = "MiniGrid-Dynamic-Obstacles-Random-16x16-v0"

    dummy_vec_env = DummyVecEnv([make_env(env_name, rank=i) for i in range(2)])
    print("Dummy vector environments work with custom registered environments. :)")

    subproc_vec_env = SubprocVecEnv([make_env(env_name, rank=i) for i in range(2)])
    print("Multiprocessing vector environments don't work with custom registered environments. :(")

Relevant log output / Error message

python test.py
pygame 2.5.2 (SDL 2.28.3, Python 3.10.13)
Hello from the pygame community. https://www.pygame.org/contribute.html
Dummy vector environments work with custom registered environments. :)
Process ForkServerProcess-1:
Traceback (most recent call last):
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py", line 29, in _worker
    env = _patch_env(env_fn_wrapper.var())
  File "/Users/marcusfechner/Projects/imitation-from-observation-on-toy-environments/test.py", line 31, in _init
    env = gym.make(env_id)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 741, in make
    env_spec = _find_spec(id)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 527, in _find_spec
    _check_version_exists(ns, name, version)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 393, in _check_version_exists
    _check_name_exists(ns, name)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/gymnasium/envs/registration.py", line 370, in _check_name_exists
    raise error.NameNotFound(
gymnasium.error.NameNotFound: Environment `MiniGrid-Dynamic-Obstacles-Random-16x16` doesn't exist. Did you mean: `MiniGrid-Dynamic-Obstacles-Random-6x6`?
Traceback (most recent call last):
  File "/Users/marcusfechner/Projects/imitation-from-observation-on-toy-environments/test.py", line 46, in <module>
    subproc_vec_env = SubprocVecEnv([make_env(env_name, rank=i) for i in range(2)])
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py", line 119, in __init__
    observation_space, action_space = self.remotes[0].recv()
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/multiprocessing/connection.py", line 250, in recv
    buf = self._recv_bytes()
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/multiprocessing/connection.py", line 414, in _recv_bytes
    buf = self._recv(4)
  File "/Users/marcusfechner/miniconda3/envs/imitation/lib/python3.10/multiprocessing/connection.py", line 383, in _recv
    raise EOFError
EOFError

System Info

pygame 2.5.2 (SDL 2.28.3, Python 3.10.13) Hello from the pygame community. https://www.pygame.org/contribute.html

Checklist

araffin commented 3 months ago

Hello, I think it is somehow an issue with gym. The workaround should be the one present in the RL Zoo: https://github.com/DLR-RM/rl-baselines3-zoo/blob/aa3814530ba8c5263267979bf14499734bfef86d/rl_zoo3/exp_manager.py#L612-L616

See https://github.com/HumanCompatibleAI/imitation/pull/160

marcusfechner commented 3 months ago

Thank you so much @araffin. Works like a charm. Here is the modified code for all other people.

from typing import Callable

import gymnasium as gym
from gymnasium.envs.registration import register
from minigrid.wrappers import ImgObsWrapper
from stable_baselines3.common.utils import set_random_seed
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv

def register_minigrid_envs():
    """Register custom Minigrid environments.
    """
    register(
        id="MiniGrid-Dynamic-Obstacles-Random-16x16-v0",
        entry_point="minigrid.envs:DynamicObstaclesEnv",
        kwargs={"size": 16, "agent_start_pos": None, "n_obstacles": 5},
    )

def make_env(env_id: str, rank: int, seed: int = 0) -> Callable:
    """
    Utility function for multiprocessed env.

    :param env_id: (str) the environment ID
    :param num_env: (int) the number of environment you wish to have in subprocesses
    :param seed: (int) the inital seed for RNG
    :param rank: (int) index of the subprocess
    :return: (Callable)
    """
    spec = gym.spec(env_id)
    # Define make_env here, so it works with subprocesses
    # when the registry was modified with `--gym-packages`
    # See https://github.com/HumanCompatibleAI/imitation/pull/160
    def make_env(**kwargs) -> gym.Env:
        return spec.make(**kwargs)

    def _init() -> gym.Env:
        env = make_env()
        env = ImgObsWrapper(env)
        env.reset(seed=seed + rank)
        return env

    set_random_seed(seed)
    return _init

if __name__ == "__main__":
    register_minigrid_envs()
    env_name = "MiniGrid-Dynamic-Obstacles-Random-16x16-v0"

    dummy_vec_env = DummyVecEnv([make_env(env_name, rank=i) for i in range(2)])
    print("Dummy vector environments work with custom registered environments. :)")

    subproc_vec_env = SubprocVecEnv([make_env(env_name, rank=i) for i in range(2)])
    print("Multiprocessing vector environments now work with custom registered environments. :)")
araffin commented 3 months ago

good to hear =) another solution might be to use the "package:env-id" syntax when registering the env.