polish ppof rewardclip and add atari config

Author: karroyan

ding/bonus/config.py

@@ -77,6 +77,45 @@ def get_instance_config(env: str) -> EasyDict:
             critic_head_hidden_size=256,
             actor_head_hidden_size=256,
         )
+    elif env == 'qbert':
+        cfg.n_sample = 1024
+        cfg.batch_size = 128
+        cfg.epoch_per_collect = 10
+        cfg.learning_rate = 0.0001
+        cfg.model = dict(
+            obs_shape=[4, 84, 84],
+            action_shape=6,
+            encoder_hidden_size_list=[32, 64, 64, 128],
+            actor_head_hidden_size=128,
+            critic_head_hidden_size=128,
+            critic_head_layer_num=2,
+        )
+    elif env == 'kangaroo':
+        cfg.n_sample = 1024
+        cfg.batch_size = 128
+        cfg.epoch_per_collect = 10
+        cfg.learning_rate = 0.0001
+        cfg.model = dict(
+            obs_shape=[4, 84, 84],
+            action_shape=18,
+            encoder_hidden_size_list=[32, 64, 64, 128],
+            actor_head_hidden_size=128,
+            critic_head_hidden_size=128,
+            critic_head_layer_num=2,
+        )
+    elif env == 'bowling':
+        cfg.n_sample = 1024
+        cfg.batch_size = 128
+        cfg.epoch_per_collect = 10
+        cfg.learning_rate = 0.0001
+        cfg.model = dict(
+            obs_shape=[4, 84, 84],
+            action_shape=6,
+            encoder_hidden_size_list=[32, 64, 64, 128],
+            actor_head_hidden_size=128,
+            critic_head_hidden_size=128,
+            critic_head_layer_num=2,
+        )
     else:
         raise KeyError("not supported env type: {}".format(env))
     return cfg
@@ -152,6 +191,36 @@ def get_instance_env(env: str) -> BaseEnv:
             },
             seed_api=False,
         )
+    elif env == 'qbert':
+        from dizoo.atari.envs.atari_env import AtariEnv
+        cfg = EasyDict({
+            'env_id': 'QbertNoFrameskip-v4',
+            'env_wrapper': 'atari_default',
+        })
+        ding_env_atari = DingEnvWrapper(gym.make('QbertNoFrameskip-v4'), cfg=cfg)
+        #ding_env_atari.enable_save_replay('atari_log/')
+        obs = ding_env_atari.reset()
+        return ding_env_atari
+    elif env == 'kangaroo':
+        from dizoo.atari.envs.atari_env import AtariEnv
+        cfg = EasyDict({
+            'env_id': 'KangarooNoFrameskip-v4',
+            'env_wrapper': 'atari_default',
+        })
+        ding_env_atari = DingEnvWrapper(gym.make('KangarooNoFrameskip-v4'), cfg=cfg)
+        #ding_env_atari.enable_save_replay('atari_log/')
+        obs = ding_env_atari.reset()
+        return ding_env_atari
+    elif env == 'bowling':
+        from dizoo.atari.envs.atari_env import AtariEnv
+        cfg = EasyDict({
+            'env_id': 'BowlingNoFrameskip-v4',
+            'env_wrapper': 'atari_default',
+        })
+        ding_env_atari = DingEnvWrapper(gym.make('BowlingNoFrameskip-v4'), cfg=cfg)
+        #ding_env_atari.enable_save_replay('atari_log/')
+        obs = ding_env_atari.reset()
+        return ding_env_atari
     else:
         raise KeyError("not supported env type: {}".format(env))
 

ding/bonus/ppof.py

@@ -1,6 +1,7 @@
 from typing import Optional, Union
 from ditk import logging
 from easydict import EasyDict
+from functools import partial
 import os
 import gym
 import torch
@@ -30,6 +31,10 @@ class PPOF:
         'mario',
         'di_sheep',
         'procgen_bigfish',
+        # atari
+        'qbert',
+        'kangaroo',
+        'bowling'
     ]
 
     def __init__(
@@ -67,8 +72,11 @@ def __init__(
             action_shape = action_space.shape
         if model is None:
             model = PPOFModel(
-                self.env.observation_space.shape, action_shape, action_space=self.cfg.action_space, **self.cfg.model
+                action_space=self.cfg.action_space, **self.cfg.model
             )
+            # model = PPOFModel(
+            #     self.env.observation_space.shape, action_shape, action_space=self.cfg.action_space, **self.cfg.model
+            # )
         self.policy = PPOFPolicy(self.cfg, model=model)
 
     def train(
@@ -86,7 +94,7 @@ def train(
         logging.debug(self.policy._model)
         # define env and policy
         collector_env = self._setup_env_manager(collector_env_num, context, debug)
-        evaluator_env = self._setup_env_manager(evaluator_env_num, context, debug)
+        evaluator_env = self._setup_env_manager(evaluator_env_num, context, debug, 'evaluator')
 
         with task.start(ctx=OnlineRLContext()):
             task.use(interaction_evaluator_ttorch(self.seed, self.policy, evaluator_env))
@@ -168,7 +176,7 @@ def batch_evaluate(
         if debug:
             logging.getLogger().setLevel(logging.DEBUG)
         # define env and policy
-        env = self._setup_env_manager(env_num, context, debug)
+        env = self._setup_env_manager(env_num, context, debug, 'evaluator')
         if ckpt_path is None:
             ckpt_path = os.path.join(self.exp_name, 'ckpt/eval.pth.tar')
         state_dict = torch.load(ckpt_path, map_location='cpu')
@@ -179,7 +187,7 @@ def batch_evaluate(
             task.use(interaction_evaluator_ttorch(self.seed, self.policy, env, n_evaluator_episode))
             task.run(max_step=1)
 
-    def _setup_env_manager(self, env_num: int, context: Optional[str] = None, debug: bool = False) -> BaseEnvManagerV2:
+    def _setup_env_manager(self, env_num: int, context: Optional[str] = None, debug: bool = False, caller: str = 'collector') -> BaseEnvManagerV2:
         if debug:
             env_cls = BaseEnvManagerV2
             manager_cfg = env_cls.default_config()
@@ -188,4 +196,4 @@ def _setup_env_manager(self, env_num: int, context: Optional[str] = None, debug:
             manager_cfg = env_cls.default_config()
             if context is not None:
                 manager_cfg.context = context
-        return env_cls([self.env.clone for _ in range(env_num)], manager_cfg)
+        return env_cls([partial(self.env.clone, caller) for _ in range(env_num)], manager_cfg)

ding/envs/env/default_wrapper.py

@@ -5,7 +5,8 @@
 eval_episode_return_wrapper = EasyDict(type='eval_episode_return')
 
 
-def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None) -> List[dict]:
+def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None, caller: str = 'collector') -> List[dict]:
+    assert caller == 'collector' or 'evaluator'
     if env_wrapper_name == 'mujoco_default':
         return [
             EasyDict(type='delay_reward', kwargs=dict(delay_reward_step=3)),
@@ -21,7 +22,8 @@ def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None) ->
                 wrapper_list.append(EasyDict(type='fire_reset'))
         wrapper_list.append(EasyDict(type='warp_frame'))
         wrapper_list.append(EasyDict(type='scaled_float_frame'))
-        wrapper_list.append(EasyDict(type='clip_reward'))
+        if caller == 'collector':
+            wrapper_list.append(EasyDict(type='clip_reward'))
         wrapper_list.append(EasyDict(type='frame_stack', kwargs=dict(n_frames=4)))
         wrapper_list.append(copy.deepcopy(eval_episode_return_wrapper))
         return wrapper_list

ding/envs/env/ding_env_wrapper.py

@@ -15,7 +15,7 @@
 
 class DingEnvWrapper(BaseEnv):
 
-    def __init__(self, env: gym.Env = None, cfg: dict = None, seed_api: bool = True) -> None:
+    def __init__(self, env: gym.Env = None, cfg: dict = None, seed_api: bool = True, caller: str = 'collector') -> None:
         """
         You can pass in either an env instance, or a config to create an env instance:
             - An env instance: Parameter `env` must not be `None`, but should be the instance.
@@ -25,6 +25,7 @@ def __init__(self, env: gym.Env = None, cfg: dict = None, seed_api: bool = True)
         self._raw_env = env
         self._cfg = cfg
         self._seed_api = seed_api  # some env may disable `env.seed` api
+        self._caller = caller
         if self._cfg is None:
             self._cfg = dict()
         self._cfg = EasyDict(self._cfg)
@@ -37,7 +38,7 @@ def __init__(self, env: gym.Env = None, cfg: dict = None, seed_api: bool = True)
         if env is not None:
             self._init_flag = True
             self._env = env
-            self._wrap_env()
+            self._wrap_env(caller)
             self._observation_space = self._env.observation_space
             self._action_space = self._env.action_space
             self._action_space.seed(0)  # default seed
@@ -57,7 +58,7 @@ def __init__(self, env: gym.Env = None, cfg: dict = None, seed_api: bool = True)
     def reset(self) -> None:
         if not self._init_flag:
             self._env = gym.make(self._cfg.env_id)
-            self._wrap_env()
+            self._wrap_env(self._caller)
             self._observation_space = self._env.observation_space
             self._action_space = self._env.action_space
             self._reward_space = gym.spaces.Box(
@@ -149,11 +150,11 @@ def random_action(self) -> np.ndarray:
             )
         return random_action
 
-    def _wrap_env(self) -> None:
+    def _wrap_env(self, caller: str = 'collector') -> None:
         # wrapper_cfgs: Union[str, List]
         wrapper_cfgs = self._cfg.env_wrapper
         if isinstance(wrapper_cfgs, str):
-            wrapper_cfgs = get_default_wrappers(wrapper_cfgs, self._cfg.env_id)
+            wrapper_cfgs = get_default_wrappers(wrapper_cfgs, self._cfg.env_id, caller)
         # self._wrapper_cfgs: List[Union[Callable, Dict]]
         self._wrapper_cfgs = wrapper_cfgs
         for wrapper in self._wrapper_cfgs:
@@ -197,12 +198,12 @@ def action_space(self) -> gym.spaces.Space:
     def reward_space(self) -> gym.spaces.Space:
         return self._reward_space
 
-    def clone(self) -> BaseEnv:
+    def clone(self, caller: str = 'collector') -> BaseEnv:
         try:
             spec = copy.deepcopy(self._raw_env.spec)
             raw_env = CloudPickleWrapper(self._raw_env)
             raw_env = copy.deepcopy(raw_env).data
             raw_env.__setattr__('spec', spec)
         except Exception:
             raw_env = self._raw_env
-        return DingEnvWrapper(raw_env, self._cfg, self._seed_api)
+        return DingEnvWrapper(raw_env, self._cfg, self._seed_api, caller)