feature(zjow): envpoool env example in new pipeline

ding/entry/utils.py

@@ -1,4 +1,4 @@
-from typing import Optional, Callable, List, Any
+from typing import Optional, Callable, List, Any, Dict
 
 from ding.policy import PolicyFactory
 from ding.worker import IMetric, MetricSerialEvaluator
@@ -46,7 +46,8 @@ def random_collect(
         collector_env: 'BaseEnvManager',  # noqa
         commander: 'BaseSerialCommander',  # noqa
         replay_buffer: 'IBuffer',  # noqa
-        postprocess_data_fn: Optional[Callable] = None
+        postprocess_data_fn: Optional[Callable] = None,
+        collect_kwargs: Optional[Dict] = None,
 ) -> None:  # noqa
     assert policy_cfg.random_collect_size > 0
     if policy_cfg.get('transition_with_policy_data', False):
@@ -55,7 +56,8 @@ def random_collect(
         action_space = collector_env.action_space
         random_policy = PolicyFactory.get_random_policy(policy.collect_mode, action_space=action_space)
         collector.reset_policy(random_policy)
-    collect_kwargs = commander.step()
+    if collect_kwargs is None:
+        collect_kwargs = commander.step()
     if policy_cfg.collect.collector.type == 'episode':
         new_data = collector.collect(n_episode=policy_cfg.random_collect_size, policy_kwargs=collect_kwargs)
     else:

ding/envs/env_manager/envpool_env_manager.py

@@ -2,7 +2,11 @@
 from easydict import EasyDict
 from copy import deepcopy
 import numpy as np
+import torch
+import treetensor.torch as ttorch
+import treetensor.numpy as tnp
 from collections import namedtuple
+import enum
 from typing import Any, Union, List, Tuple, Dict, Callable, Optional
 from ditk import logging
 try:
@@ -17,17 +21,28 @@
 from ding.torch_utils import to_ndarray
 
 
-@ENV_MANAGER_REGISTRY.register('env_pool')
+class EnvState(enum.IntEnum):
+    VOID = 0
+    INIT = 1
+    RUN = 2
+    RESET = 3
+    DONE = 4
+    ERROR = 5
+    NEED_RESET = 6
+
+
+@ENV_MANAGER_REGISTRY.register('envpool')
 class PoolEnvManager:
-    '''
+    """
     Overview:
+        PoolEnvManager supports old pipeline of DI-engine.
         Envpool now supports Atari, Classic Control, Toy Text, ViZDoom.
         Here we list some commonly used env_ids as follows.
         For more examples, you can refer to <https://envpool.readthedocs.io/en/latest/api/atari.html>.
 
         - Atari: "Pong-v5", "SpaceInvaders-v5", "Qbert-v5"
         - Classic Control: "CartPole-v0", "CartPole-v1", "Pendulum-v1"
-    '''
+    """
 
     @classmethod
     def default_config(cls) -> EasyDict:
@@ -39,10 +54,17 @@ def default_config(cls) -> EasyDict:
         # Async mode: batch_size <  env_num
         env_num=8,
         batch_size=8,
+        image_observation=True,
+        episodic_life=False,
+        reward_clip=False,
+        gray_scale=True,
+        stack_num=4,
+        frame_skip=4,
     )
 
     def __init__(self, cfg: EasyDict) -> None:
-        self._cfg = cfg
+        self._cfg = self.default_config()
+        self._cfg.update(cfg)
         self._env_num = cfg.env_num
         self._batch_size = cfg.batch_size
         self._ready_obs = {}
@@ -55,6 +77,7 @@ def launch(self) -> None:
             seed = 0
         else:
             seed = self._seed
+
         self._envs = envpool.make(
             task_id=self._cfg.env_id,
             env_type="gym",
@@ -65,8 +88,10 @@ def launch(self) -> None:
             reward_clip=self._cfg.reward_clip,
             stack_num=self._cfg.stack_num,
             gray_scale=self._cfg.gray_scale,
-            frame_skip=self._cfg.frame_skip
+            frame_skip=self._cfg.frame_skip,
         )
+        self._action_space = self._envs.action_space
+        self._observation_space = self._envs.observation_space
         self._closed = False
         self.reset()
 
@@ -77,6 +102,8 @@ def reset(self) -> None:
             obs, _, _, info = self._envs.recv()
             env_id = info['env_id']
             obs = obs.astype(np.float32)
+            if self._cfg.image_observation:
+                obs /= 255.0
             self._ready_obs = deep_merge_dicts({i: o for i, o in zip(env_id, obs)}, self._ready_obs)
             if len(self._ready_obs) == self._env_num:
                 break
@@ -91,6 +118,8 @@ def step(self, action: dict) -> Dict[int, namedtuple]:
 
         obs, rew, done, info = self._envs.recv()
         obs = obs.astype(np.float32)
+        if self._cfg.image_observation:
+            obs /= 255.0
         rew = rew.astype(np.float32)
         env_id = info['env_id']
         timesteps = {}
@@ -124,3 +153,153 @@ def env_num(self) -> int:
     @property
     def ready_obs(self) -> Dict[int, Any]:
         return self._ready_obs
+
+    @property
+    def observation_space(self) -> 'gym.spaces.Space':  # noqa
+        try:
+            return self._observation_space
+        except AttributeError:
+            self.launch()
+            self.close()
+            return self._observation_space
+
+    @property
+    def action_space(self) -> 'gym.spaces.Space':  # noqa
+        try:
+            return self._action_space
+        except AttributeError:
+            self.launch()
+            self.close()
+            return self._action_space
+
+
+@ENV_MANAGER_REGISTRY.register('envpool_v2')
+class PoolEnvManagerV2:
+    """
+    Overview:
+        PoolEnvManagerV2 supports new pipeline of DI-engine.
+        Envpool now supports Atari, Classic Control, Toy Text, ViZDoom.
+        Here we list some commonly used env_ids as follows.
+        For more examples, you can refer to <https://envpool.readthedocs.io/en/latest/api/atari.html>.
+
+        - Atari: "Pong-v5", "SpaceInvaders-v5", "Qbert-v5"
+        - Classic Control: "CartPole-v0", "CartPole-v1", "Pendulum-v1"
+    """
+
+    @classmethod
+    def default_config(cls) -> EasyDict:
+        return EasyDict(deepcopy(cls.config))
+
+    config = dict(
+        type='envpool_v2',
+        env_num=8,
+        batch_size=8,
+        image_observation=True,
+        episodic_life=False,
+        reward_clip=False,
+        gray_scale=True,
+        stack_num=4,
+        frame_skip=4,
+    )
+
+    def __init__(self, cfg: EasyDict) -> None:
+        super().__init__()
+        self._cfg = self.default_config()
+        self._cfg.update(cfg)
+        self._env_num = cfg.env_num
+        self._batch_size = cfg.batch_size
+
+        self._closed = True
+        self._seed = None
+
+    def launch(self) -> None:
+        assert self._closed, "Please first close the env manager"
+        if self._seed is None:
+            seed = 0
+        else:
+            seed = self._seed
+
+        self._envs = envpool.make(
+            task_id=self._cfg.env_id,
+            env_type="gym",
+            num_envs=self._env_num,
+            batch_size=self._batch_size,
+            seed=seed,
+            episodic_life=self._cfg.episodic_life,
+            reward_clip=self._cfg.reward_clip,
+            stack_num=self._cfg.stack_num,
+            gray_scale=self._cfg.gray_scale,
+            frame_skip=self._cfg.frame_skip,
+        )
+        self._action_space = self._envs.action_space
+        self._observation_space = self._envs.observation_space
+        self._closed = False
+        return self.reset()
+
+    def reset(self) -> None:
+        self._envs.async_reset()
+        ready_obs = {}
+        while True:
+            obs, _, _, info = self._envs.recv()
+            env_id = info['env_id']
+            obs = obs.astype(np.float32)
+            if self._cfg.image_observation:
+                obs /= 255.0
+            for i in range(len(list(env_id))):
+                ready_obs[env_id[i]] = obs[i]
+            if len(ready_obs) == self._env_num:
+                break
+        self._eval_episode_return = [0. for _ in range(self._env_num)]
+
+        return ready_obs
+
+    def send_action(self, action, env_id) -> Dict[int, namedtuple]:
+        self._envs.send(action, env_id)
+
+    def receive_data(self):
+        next_obs, rew, done, info = self._envs.recv()
+        next_obs = next_obs.astype(np.float32)
+        if self._cfg.image_observation:
+            next_obs /= 255.0
+        rew = rew.astype(np.float32)
+
+        return next_obs, rew, done, info
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        # Envpool has no `close` API
+        self._closed = True
+
+    @property
+    def closed(self) -> None:
+        return self._closed
+
+    def seed(self, seed: int, dynamic_seed=False) -> None:
+        # The i-th environment seed in Envpool will be set with i+seed, so we don't do extra transformation here
+        self._seed = seed
+        logging.warning("envpool doesn't support dynamic_seed in different episode")
+
+    @property
+    def env_num(self) -> int:
+        return self._env_num
+
+    @property
+    def observation_space(self) -> 'gym.spaces.Space':  # noqa
+        try:
+            return self._observation_space
+        except AttributeError:
+            self.launch()
+            self.close()
+            self._ready_obs = {}
+            return self._observation_space
+
+    @property
+    def action_space(self) -> 'gym.spaces.Space':  # noqa
+        try:
+            return self._action_space
+        except AttributeError:
+            self.launch()
+            self.close()
+            self._ready_obs = {}
+            return self._action_space

ding/envs/env_manager/tests/test_envpool_env_manager.py

@@ -3,7 +3,7 @@
 import numpy as np
 from easydict import EasyDict
 
-from ..envpool_env_manager import PoolEnvManager
+from ding.envs.env_manager.envpool_env_manager import PoolEnvManager, PoolEnvManagerV2
 
 env_num_args = [[16, 8], [8, 8]]
 
@@ -30,17 +30,51 @@ def test_naive(self, env_num, batch_size):
         env_manager = PoolEnvManager(env_manager_cfg)
         assert env_manager._closed
         env_manager.launch()
-        # Test step
-        start_time = time.time()
-        for count in range(20):
+        for count in range(5):
             env_id = env_manager.ready_obs.keys()
             action = {i: np.random.randint(4) for i in env_id}
             timestep = env_manager.step(action)
             assert len(timestep) == env_manager_cfg.batch_size
-            print('Count {}'.format(count))
-            print([v.info for v in timestep.values()])
-        end_time = time.time()
-        print('total step time: {}'.format(end_time - start_time))
-        # Test close
         env_manager.close()
         assert env_manager._closed
+
+
+@pytest.mark.envpooltest
+@pytest.mark.parametrize('env_num, batch_size', env_num_args)
+class TestPoolEnvManagerV2:
+
+    def test_naive(self, env_num, batch_size):
+        env_manager_cfg = EasyDict(
+            {
+                'env_id': 'Pong-v5',
+                'env_num': env_num,
+                'batch_size': batch_size,
+                'seed': 3,
+                # env wrappers
+                'episodic_life': False,
+                'reward_clip': False,
+                'gray_scale': True,
+                'stack_num': 4,
+                'frame_skip': 4,
+            }
+        )
+        env_manager = PoolEnvManagerV2(env_manager_cfg)
+        assert env_manager._closed
+        ready_obs = env_manager.launch()
+        env_id = list(ready_obs.keys())
+        for count in range(5):
+            action = {i: np.random.randint(4) for i in env_id}
+            action_send = np.array([action[i] for i in action.keys()])
+            env_id_send = np.array(list(action.keys()))
+            env_manager.send_action(action_send, env_id_send)
+            next_obs, rew, done, info = env_manager.receive_data()
+            assert next_obs.shape == (env_manager_cfg.batch_size, 4, 84, 84)
+            assert rew.shape == (env_manager_cfg.batch_size, )
+            assert done.shape == (env_manager_cfg.batch_size, )
+            assert info['env_id'].shape == (env_manager_cfg.batch_size, )
+        env_manager.close()
+        assert env_manager._closed
+
+
+if __name__ == "__main__":
+    TestPoolEnvManagerV2().test_naive(16, 8)