format my code

Author: rongkunxue

ding/example/dqn_frozen_lake.py

@@ -12,19 +12,17 @@
 from dizoo.frozen_lake.config.frozen_lake_dqn_config import main_config, create_config
 from dizoo.frozen_lake.envs import FrozenLakeEnv
 
+
 def main():
     logging.getLogger().setLevel(logging.INFO)
-    main_config.exp_name = 'cartpole_dqn_nstep'
     main_config.policy.nstep = 5
     cfg = compile_config(main_config, create_cfg=create_config, auto=True)
     with task.start(async_mode=False, ctx=OnlineRLContext()):
         collector_env = BaseEnvManagerV2(
-            env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.collector_env_num)],
-            cfg=cfg.env.manager
+            env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager
         )
         evaluator_env = BaseEnvManagerV2(
-            env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.evaluator_env_num)],
-            cfg=cfg.env.manager
+            env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager
         )
         set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
 
@@ -44,4 +42,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

dizoo/frozen_lake/config/__init__.py

@@ -1 +1 @@
-from .frozen_lake_dqn_config import main_config, create_config
+from .frozen_lake_dqn_config import main_config, create_config

dizoo/frozen_lake/config/frozen_lake_dqn_config.py

@@ -6,30 +6,29 @@
         collector_env_num=8,
         evaluator_env_num=5,
         n_evaluator_episode=10,
-        env_id = 'FrozenLake-v1',
-        desc = None,
-        map_name = "4x4",
-        is_slippery = False,
-        save_replay_gif = False,
+        env_id='FrozenLake-v1',
+        desc=None,
+        map_name="4x4",
+        is_slippery=False,
+        save_replay_gif=False,
     ),
-
-    policy = dict(
+    policy=dict(
         cuda=True,
         load_path='frozen_lake_seed0/ckpt/ckpt_best.pth.tar',
-        model = dict(
+        model=dict(
             obs_shape=16,
             action_shape=4,
             encoder_hidden_size_list=[128, 128, 64],
             dueling=True,
         ),
-        nstep = 3,
+        nstep=3,
         discount_factor=0.97,
         learn=dict(
             update_per_collect=5,
             batch_size=256,
             learning_rate=0.001,
         ),
-        collect = dict(n_sample=10),
+        collect=dict(n_sample=10),
         eval=dict(evaluator=dict(eval_freq=40, )),
         other=dict(
             eps=dict(
@@ -62,4 +61,4 @@
 if __name__ == "__main__":
     # or you can enter `ding -m serial -c frozen_lake_dqn_config.py -s 0`
     from ding.entry import serial_pipeline
-    serial_pipeline((main_config, create_config), max_env_step=5000,seed=0)
+    serial_pipeline((main_config, create_config), max_env_step=5000, seed=0)

dizoo/frozen_lake/envs/frozen_lake_env.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict,List, Optional
+from typing import Any, Dict, List, Optional
 import imageio
 import os
 import gymnasium as gymn
@@ -7,10 +7,12 @@
 from ding.torch_utils import to_ndarray
 from ding.utils import ENV_REGISTRY
 
+
 @ENV_REGISTRY.register('frozen_lake')
 class FrozenLakeEnv(BaseEnv):
-    def __init__(self,cfg)->None:
-        self._cfg=cfg
+
+    def __init__(self, cfg) -> None:
+        self._cfg = cfg
         assert self._cfg.env_id == "FrozenLake-v1", "yout name is not FrozernLake_v1"
         self._init_flag = False
         self._save_replay_bool = False
@@ -19,31 +21,33 @@ def __init__(self,cfg)->None:
         self._frames = []
         self._replay_path = False
 
-    def reset(self)-> np.ndarray:
+    def reset(self) -> np.ndarray:
         if not self._init_flag:
-            if not self._cfg.desc :#specify maps non-preloaded maps
-                self._env = gymn.make(self._cfg.env_id,
-                                      desc=self._cfg.desc,
-                                      map_name=self._cfg.map_name,
-                                      is_slippery=self._cfg.is_slippery,
-                                      render_mode="rgb_array")
+            if not self._cfg.desc:  #specify maps non-preloaded maps
+                self._env = gymn.make(
+                    self._cfg.env_id,
+                    desc=self._cfg.desc,
+                    map_name=self._cfg.map_name,
+                    is_slippery=self._cfg.is_slippery,
+                    render_mode="rgb_array"
+                )
         self._observation_space = self._env.observation_space
         self._action_space = self._env.action_space
         self._reward_space = gymn.spaces.Box(
-                low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32
-            )
+            low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32
+        )
         self._init_flag = True
         self._eval_episode_return = 0
         if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed:
             np_seed = 100 * np.random.randint(1, 1000)
-            self._env_seed=self._seed + np_seed
+            self._env_seed = self._seed + np_seed
         elif hasattr(self, '_seed'):
-            self._env_seed=self._seed
+            self._env_seed = self._seed
         if hasattr(self, '_seed'):
-            obs,info = self._env.reset(seed=self._env_seed)
+            obs, info = self._env.reset(seed=self._env_seed)
         else:
-            obs,info = self._env.reset()
-        obs = self.onehot_encode(obs)
+            obs, info = self._env.reset()
+        obs = np.eye(16, dtype=np.float32)[obs - 1]
         return obs
 
     def close(self) -> None:
@@ -57,30 +61,30 @@ def seed(self, seed: int, dynamic_seed: bool = True) -> None:
         np.random.seed(self._seed)
 
     def step(self, action: Dict) -> BaseEnvTimestep:
-        obs, rew, terminated, truncated,info = self._env.step(action[0])
+        obs, rew, terminated, truncated, info = self._env.step(action[0])
         self._eval_episode_return += rew
-        obs = self.onehot_encode(obs)
+        obs = np.eye(16, dtype=np.float32)[obs - 1]
         rew = to_ndarray([rew])
         if self._save_replay_bool:
-            picture=self._env.render()
+            picture = self._env.render()
             self._frames.append(picture)
         if terminated or truncated:
             done = True
-        else :
+        else:
             done = False
         if done:
             info['eval_episode_return'] = self._eval_episode_return
             if self._save_replay_bool:
-                assert self._replay_path is not None,"your should have a path"
+                assert self._replay_path is not None, "your should have a path"
                 path = os.path.join(
-                self._replay_path, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count)
-            )
-                self.frames_to_gif(self._frames,path)
+                    self._replay_path, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count)
+                )
+                self.frames_to_gif(self._frames, path)
                 self._frames = []
                 self._save_replay_count += 1
         rew = rew.astype(np.float32)
         return BaseEnvTimestep(obs, rew, done, info)
-    
+
     def random_action(self) -> Dict:
         raw_action = self._env.action_space.sample()
         my_type = type(self._env.action_space)
@@ -109,7 +113,6 @@ def enable_save_replay(self, replay_path: Optional[str] = None) -> None:
         self._save_replay_count = 0
         self._frames = []
 
-
     @staticmethod
     def frames_to_gif(frames: List[imageio.core.util.Array], gif_path: str, duration: float = 0.1) -> None:
         """
@@ -138,9 +141,4 @@ def frames_to_gif(frames: List[imageio.core.util.Array], gif_path: str, duration
         # Clean up temporary image files
         for temp_image_file in temp_image_files:
             os.remove(temp_image_file)
-
         print(f"GIF saved as {gif_path}")
-
-    def onehot_encode(self, x):
-        onehot = np.eye(16, dtype=np.float32)[x - 1]
-        return onehot

dizoo/frozen_lake/envs/test_frozen_lake_env.py

@@ -3,23 +3,26 @@
 from dizoo.frozen_lake.envs import FrozenLakeEnv
 from easydict import EasyDict
 
+
 @pytest.mark.envtest
 class TestGymHybridEnv:
+
     def test_my_lake(self):
         env = FrozenLakeEnv(
-            EasyDict(
-                    {
-                        'env_id': 'FrozenLake-v1',
-                        'desc': None,
-                        'map_name': "4x4",
-                        'is_slippery': False,
-                    }
-                ))
+            EasyDict({
+                'env_id': 'FrozenLake-v1',
+                'desc': None,
+                'map_name': "4x4",
+                'is_slippery': False,
+            })
+        )
         for _ in range(5):
             env.seed(314, dynamic_seed=False)
             assert env._seed == 314
             obs = env.reset()
-            assert obs.shape == (16,), "Considering the one-hot encoding format, your observation should have a dimensionality of 16." 
+            assert obs.shape == (
+                16,
+            ), "Considering the one-hot encoding format, your observation should have a dimensionality of 16."
             for i in range(10):
                 env.enable_save_replay("./video")
                 # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space,