Move envpool to examples

Author: kingjuno

examples/envpool/README.md

@@ -0,0 +1,20 @@
+## Installation
+
+ 
+Install envpool with:
+
+``` shell
+pip install envpool
+```
+
+Note 1: envpool only supports Linux operating system.
+
+## Usage
+
+You can use `OpenRL` to train Cartpole (envpool) via:
+
+``` shell
+PYTHON_PATH train_ppo.py
+```
+
+You can also add custom wrappers in `envpool_wrapper.py`. Currently we have `VecAdapter` and `VecMonitor` wrappers.

openrl/envs/wrappers/envpool_wrappers.py examples/envpool/envpool_wrappers.pyopenrl/envs/wrappers/envpool_wrappers.py renamed to examples/envpool/envpool_wrappers.py

@@ -0,0 +1,128 @@
+import copy
+import inspect
+from typing import Callable, Iterable, List, Optional, Union
+
+import envpool
+from gymnasium import Env
+
+
+from openrl.envs.vec_env import (AsyncVectorEnv, RewardWrapper,
+                                 SyncVectorEnv, VecMonitorWrapper)
+from openrl.envs.vec_env.vec_info import VecInfoFactory
+from openrl.envs.wrappers.base_wrapper import BaseWrapper
+from openrl.rewards import RewardFactory
+
+
+def build_envs(
+    make,
+    id: str,
+    env_num: int = 1,
+    wrappers: Optional[Union[Callable[[Env], Env], List[Callable[[Env], Env]]]] = None,
+    need_env_id: bool = False,
+    **kwargs,
+) -> List[Callable[[], Env]]:
+    cfg = kwargs.get("cfg", None)
+
+    def create_env(env_id: int, env_num: int, need_env_id: bool) -> Callable[[], Env]:
+        def _make_env() -> Env:
+            new_kwargs = copy.deepcopy(kwargs)
+            if need_env_id:
+                new_kwargs["env_id"] = env_id
+                new_kwargs["env_num"] = env_num
+            if "envpool" in new_kwargs:
+                # for now envpool doesnt support any render mode
+                # envpool also doesnt stores the id anywhere
+                new_kwargs.pop("envpool")
+                env = make(
+                    id,
+                    **new_kwargs,
+                )
+                env.unwrapped.spec.id = id
+
+            if wrappers is not None:
+                if callable(wrappers):
+                    if issubclass(wrappers, BaseWrapper):
+                        env = wrappers(env, cfg=cfg)
+                    else:
+                        env = wrappers(env)
+                elif isinstance(wrappers, Iterable) and all(
+                    [callable(w) for w in wrappers]
+                ):
+                    for wrapper in wrappers:
+                        if (
+                            issubclass(wrapper, BaseWrapper)
+                            and "cfg" in inspect.signature(wrapper.__init__).parameters
+                        ):
+                            env = wrapper(env, cfg=cfg)
+                        else:
+                            env = wrapper(env)
+                else:
+                    raise NotImplementedError
+
+            return env
+
+        return _make_env
+
+    env_fns = [create_env(env_id, env_num, need_env_id) for env_id in range(env_num)]
+    return env_fns
+
+
+def make_envpool_envs(
+    id: str,
+    env_num: int = 1,
+    **kwargs,
+):
+    assert "env_type" in kwargs
+    assert kwargs.get("env_type") in ["gym", "dm", "gymnasium"]
+    kwargs["envpool"] = True
+
+    if 'env_wrappers' in kwargs:
+        env_wrappers = kwargs.pop("env_wrappers")
+    else:
+        env_wrappers = []
+    env_fns = build_envs(
+        make=envpool.make,
+        id=id,
+        env_num=env_num,
+        wrappers=env_wrappers,
+        **kwargs,
+    )
+    return env_fns
+
+
+def make(
+    id: str,
+    env_num: int = 1,
+    asynchronous: bool = False,
+    add_monitor: bool = True,
+    render_mode: Optional[str] = None,
+    auto_reset: bool = True,
+    **kwargs,
+):
+    cfg = kwargs.get("cfg", None)
+    if id in envpool.registration.list_all_envs():
+        env_fns = make_envpool_envs(
+            id=id.split(":")[-1],
+            env_num=env_num,
+            **kwargs,
+        )
+        if asynchronous:
+            env = AsyncVectorEnv(
+                env_fns, render_mode=render_mode, auto_reset=auto_reset
+            )
+        else:
+            env = SyncVectorEnv(env_fns, render_mode=render_mode, auto_reset=auto_reset)
+
+        reward_class = cfg.reward_class if cfg else None
+        reward_class = RewardFactory.get_reward_class(reward_class, env)
+
+        env = RewardWrapper(env, reward_class)
+
+        if add_monitor:
+            vec_info_class = cfg.vec_info_class if cfg else None
+            vec_info_class = VecInfoFactory.get_vec_info_class(vec_info_class, env)
+            env = VecMonitorWrapper(vec_info_class, env)
+
+        return env
+    else:
+        raise NotImplementedError(f"env {id} is not supported")

examples/envpool/make_env.py

@@ -18,8 +18,8 @@
 import numpy as np
 
 from openrl.configs.config import create_config_parser
-from openrl.envs.common import make
-from openrl.envs.wrappers.envpool_wrappers import VecAdapter, VecMonitor
+from make_env import make
+from examples.envpool.envpool_wrappers import VecAdapter, VecMonitor
 from openrl.modules.common import PPONet as Net
 from openrl.modules.common.ppo_net import PPONet as Net
 from openrl.runners.common import PPOAgent as Agent
@@ -32,7 +32,7 @@ def train():
 
     # create environment, set environment parallelism to 9
     env = make(
-        "envpool:CartPole-v1",
+        "CartPole-v1",
         render_mode=None,
         env_num=9,
         asynchronous=False,
@@ -45,7 +45,7 @@ def train():
         cfg=cfg,
     )
     # initialize the trainer
-    agent = Agent(net, use_wandb=False, project_name="envpool:CartPole-v1")
+    agent = Agent(net, use_wandb=False, project_name="CartPole-v1")
     # start training, set total number of training steps to 20000
     agent.train(total_time_steps=20000)
 
@@ -58,7 +58,14 @@ def evaluation(agent):
     # Create an environment for testing and set the number of environments to interact with to 9. Set rendering mode to group_human.
     render_mode = "group_human"
     render_mode = None
-    env = make("CartPole-v1", render_mode=render_mode, env_num=9, asynchronous=True)
+    env = make(
+        "CartPole-v1",
+        env_wrappers=[VecAdapter, VecMonitor],
+        render_mode=render_mode,
+        env_num=9,
+        asynchronous=True,
+        env_type="gym",
+    )
     # The trained agent sets up the interactive environment it needs.
     agent.set_env(env)
     # Initialize the environment and get initial observations and environmental information.

examples/envpool/train_ppo.py

@@ -17,7 +17,6 @@
 """"""
 from typing import Callable, Optional
 
-import envpool
 import gymnasium as gym
 
 import openrl
@@ -155,18 +154,6 @@ def make(
             env_fns = make_PettingZoo_envs(
                 id=id, env_num=env_num, render_mode=convert_render_mode, **kwargs
             )
-        elif (
-            "envpool:" in id
-            and id.split(":")[-1] in envpool.registration.list_all_envs()
-        ):
-            from openrl.envs.envpool import make_envpool_envs
-
-            env_fns = make_envpool_envs(
-                id=id.split(":")[-1],
-                env_num=env_num,
-                render_mode=convert_render_mode,
-                **kwargs,
-            )
         else:
             raise NotImplementedError(f"env {id} is not supported.")
 

openrl/envs/common/registration.py

@@ -76,7 +76,6 @@ def get_extra_requires() -> dict:
             "async_timeout",
             "pettingzoo[classic]",
             "trueskill",
-            "envpool",
         ],
         "selfplay_test": [
             "ray[default]>=2.7",
@@ -85,7 +84,6 @@ def get_extra_requires() -> dict:
             "fastapi",
             "pettingzoo[mpe]",
             "pettingzoo[butterfly]",
-            "envpool",
         ],
         "retro": ["gym-retro"],
         "super_mario": ["gym-super-mario-bros"],