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,152 @@ 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
+            ready_obs = deep_merge_dicts({i: o for i, o in zip(env_id, obs)}, ready_obs)
+            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/example/dqn_nstep_envpool.py

@@ -0,0 +1,118 @@
+import datetime
+from easydict import EasyDict
+from ditk import logging
+from ding.model import DQN
+from ding.policy import DQNFastPolicy
+from ding.envs.env_manager.envpool_env_manager import PoolEnvManagerV2
+from ding.data import DequeBuffer
+from ding.config import compile_config
+from ding.framework import task, ding_init
+from ding.framework.context import OnlineRLContext
+from ding.framework.middleware import envpool_evaluator, data_pusher, \
+    eps_greedy_handler, CkptSaver, ContextExchanger, ModelExchanger, online_logger, \
+    termination_checker, wandb_online_logger, epoch_timer, EnvpoolStepCollector, EnvpoolOffPolicyLearner
+from ding.utils import set_pkg_seed
+from dizoo.atari.config.serial import pong_dqn_envpool_config
+
+
+def main(cfg):
+    logging.getLogger().setLevel(logging.INFO)
+    cfg.exp_name = 'Pong-v5-DQN-envpool-' + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+
+    collector_env_cfg = EasyDict(
+        {
+            'env_id': cfg.env.env_id,
+            'env_num': cfg.env.collector_env_num,
+            'batch_size': cfg.env.collector_batch_size,
+            # env wrappers
+            'episodic_life': True,  # collector: True
+            'reward_clip': False,  # collector: True
+            'gray_scale': cfg.env.get('gray_scale', True),
+            'stack_num': cfg.env.get('stack_num', 4),
+        }
+    )
+    cfg.env["collector_env_cfg"] = collector_env_cfg
+    evaluator_env_cfg = EasyDict(
+        {
+            'env_id': cfg.env.env_id,
+            'env_num': cfg.env.evaluator_env_num,
+            'batch_size': cfg.env.evaluator_batch_size,
+            # env wrappers
+            'episodic_life': False,  # evaluator: False
+            'reward_clip': False,  # evaluator: False
+            'gray_scale': cfg.env.get('gray_scale', True),
+            'stack_num': cfg.env.get('stack_num', 4),
+        }
+    )
+    cfg.env["evaluator_env_cfg"] = evaluator_env_cfg
+    cfg = compile_config(cfg, PoolEnvManagerV2, DQNFastPolicy, save_cfg=task.router.node_id == 0)
+    ding_init(cfg)
+    with task.start(async_mode=False, ctx=OnlineRLContext()):
+        collector_env = PoolEnvManagerV2(cfg.env.collector_env_cfg)
+        evaluator_env = PoolEnvManagerV2(cfg.env.evaluator_env_cfg)
+        collector_env.seed(cfg.seed)
+        evaluator_env.seed(cfg.seed)
+        set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
+
+        model = DQN(**cfg.policy.model)
+        buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size)
+        policy = DQNFastPolicy(cfg.policy, model=model)
+
+        # Consider the case with multiple processes
+        if task.router.is_active:
+            # You can use labels to distinguish between workers with different roles,
+            # here we use node_id to distinguish.
+            if task.router.node_id == 0:
+                task.add_role(task.role.LEARNER)
+            elif task.router.node_id == 1:
+                task.add_role(task.role.EVALUATOR)
+            else:
+                task.add_role(task.role.COLLECTOR)
+
+            # Sync their context and model between each worker.
+            task.use(ContextExchanger(skip_n_iter=1))
+            task.use(ModelExchanger(model))
+        task.use(epoch_timer())
+        task.use(envpool_evaluator(cfg, policy.eval_mode, evaluator_env))
+        task.use(eps_greedy_handler(cfg))
+        task.use(
+            EnvpoolStepCollector(
+                cfg,
+                policy.collect_mode,
+                collector_env,
+                random_collect_size=cfg.policy.random_collect_size \
+                       if hasattr(cfg.policy, 'random_collect_size') else 0,
+                    )
+                )
+        task.use(data_pusher(cfg, buffer_))
+        task.use(EnvpoolOffPolicyLearner(cfg, policy, buffer_))
+        task.use(online_logger(train_show_freq=10))
+        task.use(
+            wandb_online_logger(
+                metric_list=policy._monitor_vars_learn(),
+                model=policy._model,
+                exp_config=cfg,
+                anonymous=True,
+                project_name=cfg.exp_name,
+                wandb_sweep=False,
+            )
+        )
+        #task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000))
+        task.use(termination_checker(max_env_step=10000000))
+        task.run()
+
+
+if __name__ == "__main__":
+
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--collector_env_num", type=int, default=8, help="collector env number")
+    parser.add_argument("--collector_batch_size", type=int, default=8, help="collector batch size")
+    arg = parser.parse_args()
+
+    pong_dqn_envpool_config.env.collector_env_num = arg.collector_env_num
+    pong_dqn_envpool_config.env.collector_batch_size = arg.collector_batch_size
+    pong_dqn_envpool_config.seed = arg.seed
+
+    main(pong_dqn_envpool_config)