Merge pull request #13 from tianyu-z/main

Fix some typo and make English Translation

.gitignore

+*.pyc
+results
\ No newline at end of file

README.md

@@ -2,43 +2,41 @@
 
 Lightweight version of MAPPO to help you quickly migrate to your local environment.
 
-轻量版MAPPO，帮助你快速移植到本地环境。
-
-- [视频解析](https://www.bilibili.com/video/BV1bd4y1L73N/?spm_id_from=333.999.0.0&vd_source=d8ab7686ea514acb6635faa5d2227d61)
-
+- [Video (in Chinese)](https://www.bilibili.com/video/BV1bd4y1L73N)  
+This is a translated English version. Please click [here](README_CN.md) for the orginal Chinese readme.
 
 ## Table of Contents
 
-- [背景](#背景)
-- [安装](#安装)
-- [用法](#用法)
+- [Background](#Background)
+- [Installation](#Installation)
+- [Usage](#Usage)
 
-## 背景
+## Background
 
-MAPPO原版代码对于环境的封装过于复杂，本项目直接将环境封装抽取出来。更加方便将MAPPO代码移植到自己的项目上。
+The original MAPPO code was too complex in terms of environment encapsulation, so this project directly extracts and encapsulates the environment. This makes it easier to transfer the MAPPO code to your own project.
 
-## 安装
+## Installation
 
-直接将代码下载下来，创建一个Conda环境，然后运行代码，缺啥补啥包。具体什么包以后再添加。
+Simply download the code, create a Conda environment, and then run the code, adding packages as needed. Specific packages will be added later.
 
-## 用法
+## Usage
 
-- 环境部分是一个空的的实现，文件`light_mappo/envs/env_core.py`里面环境部分的实现：[Code](https://github.com/tinyzqh/light_mappo/blob/main/envs/env_core.py)
+- The environment part is an empty implementation, and the implementation of the environment part in the light_mappo/envs/env_core.py file is: [Code] (https://github.com/tinyzqh/light_mappo/blob/main/envs/env_core.py)
 
 ```python
 import numpy as np
 class EnvCore(object):
     """
-    # 环境中的智能体
+    # Environment Agent
     """
     def __init__(self):
-        self.agent_num = 2  # 设置智能体(小飞机)的个数，这里设置为两个
-        self.obs_dim = 14  # 设置智能体的观测纬度
-        self.action_dim = 5  # 设置智能体的动作纬度，这里假定为一个五个纬度的
+        self.agent_num = 2 # set the number of agents(aircrafts), here set to two
+        self.obs_dim = 14 # set the observation dimension of agents
+        self.action_dim = 5 # set the action dimension of agents, here set to a five-dimensional
 
     def reset(self):
         """
-        # self.agent_num设定为2个智能体时，返回值为一个list，每个list里面为一个shape = (self.obs_dim, )的观测数据
+        # When self.agent_num is set to 2 agents, the return value is a list, and each list contains observation data of shape = (self.obs_dim,)
         """
         sub_agent_obs = []
         for i in range(self.agent_num):
@@ -48,8 +46,8 @@ class EnvCore(object):
 
     def step(self, actions):
         """
-        # self.agent_num设定为2个智能体时，actions的输入为一个2纬的list，每个list里面为一个shape = (self.action_dim, )的动作数据
-        # 默认参数情况下，输入为一个list，里面含有两个元素，因为动作纬度为5，所里每个元素shape = (5, )
+        # When self.agent_num is set to 2 agents, the input of actions is a two-dimensional list, and each list contains action data of shape = (self.action_dim,).
+        # By default, the input is a list containing two elements, because the action dimension is 5, so each element has a shape of (5,)
         """
         sub_agent_obs = []
         sub_agent_reward = []
@@ -65,9 +63,9 @@ class EnvCore(object):
 ```
 
 
-只需要编写这一部分的代码，就可以无缝衔接MAPPO。在env_core.py之后，单独提出来了两个文件env_discrete.py和env_continuous.py这两个文件用于封装处理动作空间和离散动作空间。在algorithms/utils/act.py中elif self.continuous_action:这个判断逻辑也是用来处理连续动作空间的。和runner/shared/env_runner.py部分的# TODO 这里改造成自己环境需要的形式即可都是用来处理连续动作空间的。
+Just write this part of the code, and you can seamlessly connect with MAPPO. After env_core.py, two files, env_discrete.py and env_continuous.py, were separately extracted to encapsulate the action space and discrete action space. In elif self.continuous_action: in algorithms/utils/act.py, this judgment logic is also used to handle continuous action spaces. The # TODO here in runner/shared/env_runner.py is also used to handle continuous action spaces.
 
-在train.py文件里面，选择注释连续环境，或者离散环境进行demo环境的切换。
+In the train.py file, choose to comment out continuous environment or discrete environment to switch the demo environment.
 
 ## Related Efforts
 
@@ -77,6 +75,9 @@ class EnvCore(object):
 
 [@tinyzqh](https://github.com/tinyzqh).
 
+## Translator
+[@tianyu-z](https://github.com/tianyu-z)
+
 ## License
 
 [MIT](LICENSE) © tinyzqh

README_CN.md

+# light_mappo
+
+Lightweight version of MAPPO to help you quickly migrate to your local environment.
+
+轻量版MAPPO，帮助你快速移植到本地环境。
+
+- [视频解析](https://www.bilibili.com/video/BV1bd4y1L73N/?spm_id_from=333.999.0.0&vd_source=d8ab7686ea514acb6635faa5d2227d61)  
+
+英文翻译版readme，请点击[这里](README.md)
+
+## Table of Contents
+
+- [背景](#背景)
+- [安装](#安装)
+- [用法](#用法)
+
+## 背景
+
+MAPPO原版代码对于环境的封装过于复杂，本项目直接将环境封装抽取出来。更加方便将MAPPO代码移植到自己的项目上。
+
+## 安装
+
+直接将代码下载下来，创建一个Conda环境，然后运行代码，缺啥补啥包。具体什么包以后再添加。
+
+## 用法
+
+- 环境部分是一个空的的实现，文件`light_mappo/envs/env_core.py`里面环境部分的实现：[Code](https://github.com/tinyzqh/light_mappo/blob/main/envs/env_core.py)
+
+```python
+import numpy as np
+class EnvCore(object):
+    """
+    # 环境中的智能体
+    """
+    def __init__(self):
+        self.agent_num = 2  # 设置智能体(小飞机)的个数，这里设置为两个
+        self.obs_dim = 14  # 设置智能体的观测维度
+        self.action_dim = 5  # 设置智能体的动作维度，这里假定为一个五个维度的
+
+    def reset(self):
+        """
+        # self.agent_num设定为2个智能体时，返回值为一个list，每个list里面为一个shape = (self.obs_dim, )的观测数据
+        """
+        sub_agent_obs = []
+        for i in range(self.agent_num):
+            sub_obs = np.random.random(size=(14, ))
+            sub_agent_obs.append(sub_obs)
+        return sub_agent_obs
+
+    def step(self, actions):
+        """
+        # self.agent_num设定为2个智能体时，actions的输入为一个2纬的list，每个list里面为一个shape = (self.action_dim, )的动作数据
+        # 默认参数情况下，输入为一个list，里面含有两个元素，因为动作维度为5，所里每个元素shape = (5, )
+        """
+        sub_agent_obs = []
+        sub_agent_reward = []
+        sub_agent_done = []
+        sub_agent_info = []
+        for i in range(self.agent_num):
+            sub_agent_obs.append(np.random.random(size=(14,)))
+            sub_agent_reward.append([np.random.rand()])
+            sub_agent_done.append(False)
+            sub_agent_info.append({})
+
+        return [sub_agent_obs, sub_agent_reward, sub_agent_done, sub_agent_info]
+```
+
+
+只需要编写这一部分的代码，就可以无缝衔接MAPPO。在env_core.py之后，单独提出来了两个文件env_discrete.py和env_continuous.py这两个文件用于封装处理动作空间和离散动作空间。在algorithms/utils/act.py中elif self.continuous_action:这个判断逻辑也是用来处理连续动作空间的。和runner/shared/env_runner.py部分的# TODO 这里改造成自己环境需要的形式即可都是用来处理连续动作空间的。
+
+在train.py文件里面，选择注释连续环境，或者离散环境进行demo环境的切换。
+
+## Related Efforts
+
+- [on-policy](https://github.com/marlbenchmark/on-policy) - 💌 Learn the author implementation of MAPPO.
+
+## Maintainers
+
+[@tinyzqh](https://github.com/tinyzqh).
+
+## License
+
+[MIT](LICENSE) © tinyzqh
+

envs/env_continuous.py

@@ -5,7 +5,11 @@ from envs.env_core import EnvCore
 
 
 class ContinuousActionEnv(object):
-    """对于连续动作环境的封装"""
+    """
+    对于连续动作环境的封装
+    Wrapper for continuous action environment.
+    """
+
     def __init__(self):
         self.env = EnvCore()
         self.num_agent = self.env.agent_num
@@ -27,7 +31,12 @@ class ContinuousActionEnv(object):
         total_action_space = []
         for agent in range(self.num_agent):
             # physical action space
-            u_action_space = spaces.Box(low=-np.inf, high=+np.inf, shape=(self.signal_action_dim,), dtype=np.float32)
+            u_action_space = spaces.Box(
+                low=-np.inf,
+                high=+np.inf,
+                shape=(self.signal_action_dim,),
+                dtype=np.float32,
+            )
 
             if self.movable:
                 total_action_space.append(u_action_space)
@@ -37,17 +46,31 @@ class ContinuousActionEnv(object):
 
             # observation space
             share_obs_dim += self.signal_obs_dim
-            self.observation_space.append(spaces.Box(low=-np.inf, high=+np.inf, shape=(self.signal_obs_dim,),
-                                                     dtype=np.float32))  # [-inf,inf]
-
-        self.share_observation_space = [spaces.Box(low=-np.inf, high=+np.inf, shape=(share_obs_dim,),
-                                                   dtype=np.float32) for _ in range(self.num_agent)]
+            self.observation_space.append(
+                spaces.Box(
+                    low=-np.inf,
+                    high=+np.inf,
+                    shape=(self.signal_obs_dim,),
+                    dtype=np.float32,
+                )
+            )  # [-inf,inf]
+
+        self.share_observation_space = [
+            spaces.Box(
+                low=-np.inf, high=+np.inf, shape=(share_obs_dim,), dtype=np.float32
+            )
+            for _ in range(self.num_agent)
+        ]
 
     def step(self, actions):
         """
-        输入actions纬度假设：
+        输入actions维度假设：
         # actions shape = (5, 2, 5)
         # 5个线程的环境，里面有2个智能体，每个智能体的动作是一个one_hot的5维编码
+
+        Input actions dimension assumption:
+        # actions shape = (5, 2, 5)
+        # 5 threads of environment, there are 2 agents inside, and each agent's action is a 5-dimensional one_hot encoding
         """
 
         results = self.env.step(actions)
@@ -65,4 +88,4 @@ class ContinuousActionEnv(object):
         pass
 
     def seed(self, seed):
-        pass
\ No newline at end of file
+        pass

envs/env_core.py

@@ -5,25 +5,29 @@ class EnvCore(object):
     """
     # 环境中的智能体
     """
+
     def __init__(self):
-        self.agent_num = 2  # 设置智能体(小飞机)的个数，这里设置为两个
-        self.obs_dim = 14  # 设置智能体的观测纬度
-        self.action_dim = 5  # 设置智能体的动作纬度，这里假定为一个五个纬度的
+        self.agent_num = 2  # 设置智能体(小飞机)的个数，这里设置为两个 # set the number of agents(aircrafts), here set to two
+        self.obs_dim = 14  # 设置智能体的观测维度 # set the observation dimension of agents
+        self.action_dim = 5  # 设置智能体的动作维度，这里假定为一个五个维度的 # set the action dimension of agents, here set to a five-dimensional
 
     def reset(self):
         """
         # self.agent_num设定为2个智能体时，返回值为一个list，每个list里面为一个shape = (self.obs_dim, )的观测数据
+        # When self.agent_num is set to 2 agents, the return value is a list, each list contains a shape = (self.obs_dim, ) observation data
         """
         sub_agent_obs = []
         for i in range(self.agent_num):
-            sub_obs = np.random.random(size=(14, ))
+            sub_obs = np.random.random(size=(14,))
             sub_agent_obs.append(sub_obs)
         return sub_agent_obs
 
     def step(self, actions):
         """
         # self.agent_num设定为2个智能体时，actions的输入为一个2纬的list，每个list里面为一个shape = (self.action_dim, )的动作数据
-        # 默认参数情况下，输入为一个list，里面含有两个元素，因为动作纬度为5，所里每个元素shape = (5, )
+        # 默认参数情况下，输入为一个list，里面含有两个元素，因为动作维度为5，所里每个元素shape = (5, )
+        # When self.agent_num is set to 2 agents, the input of actions is a 2-dimensional list, each list contains a shape = (self.action_dim, ) action data
+        # The default parameter situation is to input a list with two elements, because the action dimension is 5, so each element shape = (5, )
         """
         sub_agent_obs = []
         sub_agent_reward = []
@@ -35,4 +39,4 @@ class EnvCore(object):
             sub_agent_done.append(False)
             sub_agent_info.append({})
 
-        return [sub_agent_obs, sub_agent_reward, sub_agent_done, sub_agent_info]
\ No newline at end of file
+        return [sub_agent_obs, sub_agent_reward, sub_agent_done, sub_agent_info]

envs/env_discrete.py

@@ -12,7 +12,11 @@ from envs.env_core import EnvCore
 
 
 class DiscreteActionEnv(object):
-    """对于离散动作环境的封装"""
+    """
+    对于离散动作环境的封装
+    Wrapper for discrete action environment.
+    """
+
     def __init__(self):
         self.env = EnvCore()
         self.num_agent = self.env.agent_num
@@ -42,8 +46,15 @@ class DiscreteActionEnv(object):
             # total action space
             if len(total_action_space) > 1:
                 # all action spaces are discrete, so simplify to MultiDiscrete action space
-                if all([isinstance(act_space, spaces.Discrete) for act_space in total_action_space]):
-                    act_space = MultiDiscrete([[0, act_space.n - 1] for act_space in total_action_space])
+                if all(
+                    [
+                        isinstance(act_space, spaces.Discrete)
+                        for act_space in total_action_space
+                    ]
+                ):
+                    act_space = MultiDiscrete(
+                        [[0, act_space.n - 1] for act_space in total_action_space]
+                    )
                 else:
                     act_space = spaces.Tuple(total_action_space)
                 self.action_space.append(act_space)
@@ -52,17 +63,30 @@ class DiscreteActionEnv(object):
 
             # observation space
             share_obs_dim += self.signal_obs_dim
-            self.observation_space.append(spaces.Box(low=-np.inf, high=+np.inf, shape=(self.signal_obs_dim,),
-                                                     dtype=np.float32))  # [-inf,inf]
-
-        self.share_observation_space = [spaces.Box(low=-np.inf, high=+np.inf, shape=(share_obs_dim,),
-                                                   dtype=np.float32) for _ in range(self.num_agent)]
+            self.observation_space.append(
+                spaces.Box(
+                    low=-np.inf,
+                    high=+np.inf,
+                    shape=(self.signal_obs_dim,),
+                    dtype=np.float32,
+                )
+            )  # [-inf,inf]
+
+        self.share_observation_space = [
+            spaces.Box(
+                low=-np.inf, high=+np.inf, shape=(share_obs_dim,), dtype=np.float32
+            )
+            for _ in range(self.num_agent)
+        ]
 
     def step(self, actions):
         """
-        输入actions纬度假设：
+        输入actions维度假设：
         # actions shape = (5, 2, 5)
         # 5个线程的环境，里面有2个智能体，每个智能体的动作是一个one_hot的5维编码
+        Input actions dimension assumption:
+        # actions shape = (5, 2, 5)
+        # 5 threads of the environment, with 2 intelligent agents inside, and each intelligent agent's action is a 5-dimensional one_hot encoding
         """
 
         results = self.env.step(actions)
@@ -82,7 +106,8 @@ class DiscreteActionEnv(object):
     def seed(self, seed):
         pass
 
-class MultiDiscrete():
+
+class MultiDiscrete:
     """
     - The multi-discrete action space consists of a series of discrete action spaces with different parameters
     - It can be adapted to both a Discrete action space or a continuous (Box) action space
@@ -107,14 +132,22 @@ class MultiDiscrete():
         self.n = np.sum(self.high) + 2
 
     def sample(self):
-        """ Returns a array with one sample from each discrete action space """
+        """Returns a array with one sample from each discrete action space"""
         # For each row: round(random .* (max - min) + min, 0)
         random_array = np.random.rand(self.num_discrete_space)
-        return [int(x) for x in np.floor(np.multiply((self.high - self.low + 1.), random_array) + self.low)]
+        return [
+            int(x)
+            for x in np.floor(
+                np.multiply((self.high - self.low + 1.0), random_array) + self.low
+            )
+        ]
 
     def contains(self, x):
-        return len(x) == self.num_discrete_space and (np.array(x) >= self.low).all() and (
-                    np.array(x) <= self.high).all()
+        return (
+            len(x) == self.num_discrete_space
+            and (np.array(x) >= self.low).all()
+            and (np.array(x) <= self.high).all()
+        )
 
     @property
     def shape(self):
@@ -124,8 +157,10 @@ class MultiDiscrete():
         return "MultiDiscrete" + str(self.num_discrete_space)
 
     def __eq__(self, other):
-        return np.array_equal(self.low, other.low) and np.array_equal(self.high, other.high)
+        return np.array_equal(self.low, other.low) and np.array_equal(
+            self.high, other.high
+        )
 
 
 if __name__ == "__main__":
-    DiscreteActionEnv().step(actions=None)
\ No newline at end of file
+    DiscreteActionEnv().step(actions=None)

runner/separated/env_runner.py

@@ -29,7 +29,9 @@ class EnvRunner(Runner):
         self.warmup()
 
         start = time.time()
-        episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads
+        episodes = (
+            int(self.num_env_steps) // self.episode_length // self.n_rollout_threads
+        )
 
         for episode in range(episodes):
             if self.use_linear_lr_decay:
@@ -38,12 +40,29 @@ class EnvRunner(Runner):
 
             for step in range(self.episode_length):
                 # Sample actions
-                values, actions, action_log_probs, rnn_states, rnn_states_critic, actions_env = self.collect(step)
+                (
+                    values,
+                    actions,
+                    action_log_probs,
+                    rnn_states,
+                    rnn_states_critic,
+                    actions_env,
+                ) = self.collect(step)
 
                 # Obser reward and next obs
                 obs, rewards, dones, infos = self.envs.step(actions_env)
 
-                data = obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic
+                data = (
+                    obs,
+                    rewards,
+                    dones,
+                    infos,
+                    values,
+                    actions,
+                    action_log_probs,
+                    rnn_states,
+                    rnn_states_critic,
+                )
 
                 # insert data into buffer
                 self.insert(data)
@@ -53,34 +72,47 @@ class EnvRunner(Runner):
             train_infos = self.train()
 
             # post process
-            total_num_steps = (episode + 1) * self.episode_length * self.n_rollout_threads
+            total_num_steps = (
+                (episode + 1) * self.episode_length * self.n_rollout_threads
+            )
 
             # save model
-            if (episode % self.save_interval == 0 or episode == episodes - 1):
+            if episode % self.save_interval == 0 or episode == episodes - 1:
                 self.save()
 
             # log information
             if episode % self.log_interval == 0:
                 end = time.time()
-                print("\n Scenario {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
-                      .format(self.all_args.scenario_name,
-                              self.algorithm_name,
-                              self.experiment_name,
-                              episode,
-                              episodes,
-                              total_num_steps,
-                              self.num_env_steps,
-                              int(total_num_steps / (end - start))))
+                print(
+                    "\n Scenario {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n".format(
+                        self.all_args.scenario_name,
+                        self.algorithm_name,
+                        self.experiment_name,
+                        episode,
+                        episodes,
+                        total_num_steps,
+                        self.num_env_steps,
+                        int(total_num_steps / (end - start)),
+                    )
+                )
 
                 if self.env_name == "MPE":
                     for agent_id in range(self.num_agents):
                         idv_rews = []
                         for info in infos:
-                            if 'individual_reward' in info[agent_id].keys():
-                                idv_rews.append(info[agent_id]['individual_reward'])
-                        train_infos[agent_id].update({'individual_rewards': np.mean(idv_rews)})
+                            if "individual_reward" in info[agent_id].keys():
+                                idv_rews.append(info[agent_id]["individual_reward"])
                         train_infos[agent_id].update(
-                            {"average_episode_rewards": np.mean(self.buffer[agent_id].rewards) * self.episode_length})
+                            {"individual_rewards": np.mean(idv_rews)}
+                        )
+                        train_infos[agent_id].update(
+                            {
+                                "average_episode_rewards": np.mean(
+                                    self.buffer[agent_id].rewards
+                                )
+                                * self.episode_length
+                            }
+                        )
                 self.log_train(train_infos, total_num_steps)
 
             # eval
@@ -113,27 +145,35 @@ class EnvRunner(Runner):
 
         for agent_id in range(self.num_agents):
             self.trainer[agent_id].prep_rollout()
-            value, action, action_log_prob, rnn_state, rnn_state_critic \
-                = self.trainer[agent_id].policy.get_actions(self.buffer[agent_id].share_obs[step],
-                                                            self.buffer[agent_id].obs[step],
-                                                            self.buffer[agent_id].rnn_states[step],
-                                                            self.buffer[agent_id].rnn_states_critic[step],
-                                                            self.buffer[agent_id].masks[step])
+            value, action, action_log_prob, rnn_state, rnn_state_critic = self.trainer[
+                agent_id
+            ].policy.get_actions(
+                self.buffer[agent_id].share_obs[step],
+                self.buffer[agent_id].obs[step],
+                self.buffer[agent_id].rnn_states[step],
+                self.buffer[agent_id].rnn_states_critic[step],
+                self.buffer[agent_id].masks[step],
+            )
             # [agents, envs, dim]
             values.append(_t2n(value))
             action = _t2n(action)
             # rearrange action
-            if self.envs.action_space[agent_id].__class__.__name__ == 'MultiDiscrete':
+            if self.envs.action_space[agent_id].__class__.__name__ == "MultiDiscrete":
                 for i in range(self.envs.action_space[agent_id].shape):
-                    uc_action_env = np.eye(self.envs.action_space[agent_id].high[i] + 1)[action[:, i]]
+                    uc_action_env = np.eye(
+                        self.envs.action_space[agent_id].high[i] + 1
+                    )[action[:, i]]
                     if i == 0:
                         action_env = uc_action_env
                     else:
                         action_env = np.concatenate((action_env, uc_action_env), axis=1)
-            elif self.envs.action_space[agent_id].__class__.__name__ == 'Discrete':
-                action_env = np.squeeze(np.eye(self.envs.action_space[agent_id].n)[action], 1)
+            elif self.envs.action_space[agent_id].__class__.__name__ == "Discrete":
+                action_env = np.squeeze(
+                    np.eye(self.envs.action_space[agent_id].n)[action], 1
+                )
             else:
                 # TODO 这里改造成自己环境需要的形式即可
+                # TODO Here, you can change the action_env to the form you need
                 action_env = actions
                 # raise NotImplementedError
 
@@ -157,15 +197,36 @@ class EnvRunner(Runner):
         rnn_states = np.array(rnn_states).transpose(1, 0, 2, 3)
         rnn_states_critic = np.array(rnn_states_critic).transpose(1, 0, 2, 3)
 
-        return values, actions, action_log_probs, rnn_states, rnn_states_critic, actions_env
+        return (
+            values,
+            actions,
+            action_log_probs,
+            rnn_states,
+            rnn_states_critic,
+            actions_env,
+        )
 
     def insert(self, data):
-        obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic = data
-
-        rnn_states[dones == True] = np.zeros(((dones == True).sum(), self.recurrent_N, self.hidden_size),
-                                             dtype=np.float32)
-        rnn_states_critic[dones == True] = np.zeros(((dones == True).sum(), self.recurrent_N, self.hidden_size),
-                                                    dtype=np.float32)
+        (
+            obs,
+            rewards,
+            dones,
+            infos,
+            values,
+            actions,
+            action_log_probs,
+            rnn_states,
+            rnn_states_critic,
+        ) = data
+
+        rnn_states[dones == True] = np.zeros(
+            ((dones == True).sum(), self.recurrent_N, self.hidden_size),
+            dtype=np.float32,
+        )
+        rnn_states_critic[dones == True] = np.zeros(
+            ((dones == True).sum(), self.recurrent_N, self.hidden_size),
+            dtype=np.float32,
+        )
         masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
         masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)
 
@@ -178,46 +239,70 @@ class EnvRunner(Runner):
             if not self.use_centralized_V:
                 share_obs = np.array(list(obs[:, agent_id]))
 
-            self.buffer[agent_id].insert(share_obs,
-                                         np.array(list(obs[:, agent_id])),
-                                         rnn_states[:, agent_id],
-                                         rnn_states_critic[:, agent_id],
-                                         actions[:, agent_id],
-                                         action_log_probs[:, agent_id],
-                                         values[:, agent_id],
-                                         rewards[:, agent_id],
-                                         masks[:, agent_id])
+            self.buffer[agent_id].insert(
+                share_obs,
+                np.array(list(obs[:, agent_id])),
+                rnn_states[:, agent_id],
+                rnn_states_critic[:, agent_id],
+                actions[:, agent_id],
+                action_log_probs[:, agent_id],
+                values[:, agent_id],
+                rewards[:, agent_id],
+                masks[:, agent_id],
+            )
 
     @torch.no_grad()
     def eval(self, total_num_steps):
         eval_episode_rewards = []
         eval_obs = self.eval_envs.reset()
 
-        eval_rnn_states = np.zeros((self.n_eval_rollout_threads, self.num_agents, self.recurrent_N, self.hidden_size),
-                                   dtype=np.float32)
-        eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)
+        eval_rnn_states = np.zeros(
+            (
+                self.n_eval_rollout_threads,
+                self.num_agents,
+                self.recurrent_N,
+                self.hidden_size,
+            ),
+            dtype=np.float32,
+        )
+        eval_masks = np.ones(
+            (self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32
+        )
 
         for eval_step in range(self.episode_length):
             eval_temp_actions_env = []
             for agent_id in range(self.num_agents):
                 self.trainer[agent_id].prep_rollout()
-                eval_action, eval_rnn_state = self.trainer[agent_id].policy.act(np.array(list(eval_obs[:, agent_id])),
-                                                                                eval_rnn_states[:, agent_id],
-                                                                                eval_masks[:, agent_id],
-                                                                                deterministic=True)
+                eval_action, eval_rnn_state = self.trainer[agent_id].policy.act(
+                    np.array(list(eval_obs[:, agent_id])),
+                    eval_rnn_states[:, agent_id],
+                    eval_masks[:, agent_id],
+                    deterministic=True,
+                )
 
                 eval_action = eval_action.detach().cpu().numpy()
                 # rearrange action
-                if self.eval_envs.action_space[agent_id].__class__.__name__ == 'MultiDiscrete':
+                if (
+                    self.eval_envs.action_space[agent_id].__class__.__name__
+                    == "MultiDiscrete"
+                ):
                     for i in range(self.eval_envs.action_space[agent_id].shape):
-                        eval_uc_action_env = np.eye(self.eval_envs.action_space[agent_id].high[i] + 1)[
-                            eval_action[:, i]]
+                        eval_uc_action_env = np.eye(
+                            self.eval_envs.action_space[agent_id].high[i] + 1
+                        )[eval_action[:, i]]
                         if i == 0:
                             eval_action_env = eval_uc_action_env
                         else:
-                            eval_action_env = np.concatenate((eval_action_env, eval_uc_action_env), axis=1)
-                elif self.eval_envs.action_space[agent_id].__class__.__name__ == 'Discrete':
-                    eval_action_env = np.squeeze(np.eye(self.eval_envs.action_space[agent_id].n)[eval_action], 1)
+                            eval_action_env = np.concatenate(
+                                (eval_action_env, eval_uc_action_env), axis=1
+                            )
+                elif (
+                    self.eval_envs.action_space[agent_id].__class__.__name__
+                    == "Discrete"
+                ):
+                    eval_action_env = np.squeeze(
+                        np.eye(self.eval_envs.action_space[agent_id].n)[eval_action], 1
+                    )
                 else:
                     raise NotImplementedError
 
@@ -233,21 +318,36 @@ class EnvRunner(Runner):
                 eval_actions_env.append(eval_one_hot_action_env)
 
             # Obser reward and next obs
-            eval_obs, eval_rewards, eval_dones, eval_infos = self.eval_envs.step(eval_actions_env)
+            eval_obs, eval_rewards, eval_dones, eval_infos = self.eval_envs.step(
+                eval_actions_env
+            )
             eval_episode_rewards.append(eval_rewards)
 
             eval_rnn_states[eval_dones == True] = np.zeros(
-                ((eval_dones == True).sum(), self.recurrent_N, self.hidden_size), dtype=np.float32)
-            eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)
-            eval_masks[eval_dones == True] = np.zeros(((eval_dones == True).sum(), 1), dtype=np.float32)
+                ((eval_dones == True).sum(), self.recurrent_N, self.hidden_size),
+                dtype=np.float32,
+            )
+            eval_masks = np.ones(
+                (self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32
+            )
+            eval_masks[eval_dones == True] = np.zeros(
+                ((eval_dones == True).sum(), 1), dtype=np.float32
+            )
 
         eval_episode_rewards = np.array(eval_episode_rewards)
 
         eval_train_infos = []
         for agent_id in range(self.num_agents):
-            eval_average_episode_rewards = np.mean(np.sum(eval_episode_rewards[:, :, agent_id], axis=0))
-            eval_train_infos.append({'eval_average_episode_rewards': eval_average_episode_rewards})
-            print("eval average episode rewards of agent%i: " % agent_id + str(eval_average_episode_rewards))
+            eval_average_episode_rewards = np.mean(
+                np.sum(eval_episode_rewards[:, :, agent_id], axis=0)
+            )
+            eval_train_infos.append(
+                {"eval_average_episode_rewards": eval_average_episode_rewards}
+            )
+            print(
+                "eval average episode rewards of agent%i: " % agent_id
+                + str(eval_average_episode_rewards)
+            )
 
         self.log_train(eval_train_infos, total_num_steps)
 
@@ -258,12 +358,21 @@ class EnvRunner(Runner):
             episode_rewards = []
             obs = self.envs.reset()
             if self.all_args.save_gifs:
-                image = self.envs.render('rgb_array')[0][0]
+                image = self.envs.render("rgb_array")[0][0]
                 all_frames.append(image)
 
-            rnn_states = np.zeros((self.n_rollout_threads, self.num_agents, self.recurrent_N, self.hidden_size),
-                                  dtype=np.float32)
-            masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
+            rnn_states = np.zeros(
+                (
+                    self.n_rollout_threads,
+                    self.num_agents,
+                    self.recurrent_N,
+                    self.hidden_size,
+                ),
+                dtype=np.float32,
+            )
+            masks = np.ones(
+                (self.n_rollout_threads, self.num_agents, 1), dtype=np.float32
+            )
 
             for step in range(self.episode_length):
                 calc_start = time.time()
@@ -273,22 +382,36 @@ class EnvRunner(Runner):
                     if not self.use_centralized_V:
                         share_obs = np.array(list(obs[:, agent_id]))
                     self.trainer[agent_id].prep_rollout()
-                    action, rnn_state = self.trainer[agent_id].policy.act(np.array(list(obs[:, agent_id])),
-                                                                          rnn_states[:, agent_id],
-                                                                          masks[:, agent_id],
-                                                                          deterministic=True)
+                    action, rnn_state = self.trainer[agent_id].policy.act(
+                        np.array(list(obs[:, agent_id])),
+                        rnn_states[:, agent_id],
+                        masks[:, agent_id],
+                        deterministic=True,
+                    )
 
                     action = action.detach().cpu().numpy()
                     # rearrange action
-                    if self.envs.action_space[agent_id].__class__.__name__ == 'MultiDiscrete':
+                    if (
+                        self.envs.action_space[agent_id].__class__.__name__
+                        == "MultiDiscrete"
+                    ):
                         for i in range(self.envs.action_space[agent_id].shape):
-                            uc_action_env = np.eye(self.envs.action_space[agent_id].high[i] + 1)[action[:, i]]
+                            uc_action_env = np.eye(
+                                self.envs.action_space[agent_id].high[i] + 1
+                            )[action[:, i]]
                             if i == 0:
                                 action_env = uc_action_env
                             else:
-                                action_env = np.concatenate((action_env, uc_action_env), axis=1)
-                    elif self.envs.action_space[agent_id].__class__.__name__ == 'Discrete':
-                        action_env = np.squeeze(np.eye(self.envs.action_space[agent_id].n)[action], 1)
+                                action_env = np.concatenate(
+                                    (action_env, uc_action_env), axis=1
+                                )
+                    elif (
+                        self.envs.action_space[agent_id].__class__.__name__
+                        == "Discrete"
+                    ):
+                        action_env = np.squeeze(
+                            np.eye(self.envs.action_space[agent_id].n)[action], 1
+                        )
                     else:
                         raise NotImplementedError
 
@@ -307,13 +430,19 @@ class EnvRunner(Runner):
                 obs, rewards, dones, infos = self.envs.step(actions_env)
                 episode_rewards.append(rewards)
 
-                rnn_states[dones == True] = np.zeros(((dones == True).sum(), self.recurrent_N, self.hidden_size),
-                                                     dtype=np.float32)
-                masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
-                masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)
+                rnn_states[dones == True] = np.zeros(
+                    ((dones == True).sum(), self.recurrent_N, self.hidden_size),
+                    dtype=np.float32,
+                )
+                masks = np.ones(
+                    (self.n_rollout_threads, self.num_agents, 1), dtype=np.float32
+                )
+                masks[dones == True] = np.zeros(
+                    ((dones == True).sum(), 1), dtype=np.float32
+                )
 
                 if self.all_args.save_gifs:
-                    image = self.envs.render('rgb_array')[0][0]
+                    image = self.envs.render("rgb_array")[0][0]
                     all_frames.append(image)
                     calc_end = time.time()
                     elapsed = calc_end - calc_start
@@ -322,8 +451,17 @@ class EnvRunner(Runner):
 
             episode_rewards = np.array(episode_rewards)
             for agent_id in range(self.num_agents):
-                average_episode_rewards = np.mean(np.sum(episode_rewards[:, :, agent_id], axis=0))
-                print("eval average episode rewards of agent%i: " % agent_id + str(average_episode_rewards))
+                average_episode_rewards = np.mean(
+                    np.sum(episode_rewards[:, :, agent_id], axis=0)
+                )
+                print(
+                    "eval average episode rewards of agent%i: " % agent_id
+                    + str(average_episode_rewards)
+                )
 
         if self.all_args.save_gifs:
-            imageio.mimsave(str(self.gif_dir) + '/render.gif', all_frames, duration=self.all_args.ifi)
\ No newline at end of file
+            imageio.mimsave(
+                str(self.gif_dir) + "/render.gif",
+                all_frames,
+                duration=self.all_args.ifi,
+            )

runner/shared/env_runner.py

@@ -9,6 +9,7 @@ import time
 import numpy as np
 import torch
 from runner.shared.base_runner import Runner
+
 # import imageio
 
 
@@ -26,7 +27,9 @@ class EnvRunner(Runner):
         self.warmup()
 
         start = time.time()
-        episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads
+        episodes = (
+            int(self.num_env_steps) // self.episode_length // self.n_rollout_threads
+        )
 
         for episode in range(episodes):
             if self.use_linear_lr_decay:
@@ -34,12 +37,29 @@ class EnvRunner(Runner):
 
             for step in range(self.episode_length):
                 # Sample actions
-                values, actions, action_log_probs, rnn_states, rnn_states_critic, actions_env = self.collect(step)
+                (
+                    values,
+                    actions,
+                    action_log_probs,
+                    rnn_states,
+                    rnn_states_critic,
+                    actions_env,
+                ) = self.collect(step)
 
                 # Obser reward and next obs
                 obs, rewards, dones, infos = self.envs.step(actions_env)
 
-                data = obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic
+                data = (
+                    obs,
+                    rewards,
+                    dones,
+                    infos,
+                    values,
+                    actions,
+                    action_log_probs,
+                    rnn_states,
+                    rnn_states_critic,
+                )
 
                 # insert data into buffer
                 self.insert(data)
@@ -49,24 +69,29 @@ class EnvRunner(Runner):
             train_infos = self.train()
 
             # post process
-            total_num_steps = (episode + 1) * self.episode_length * self.n_rollout_threads
+            total_num_steps = (
+                (episode + 1) * self.episode_length * self.n_rollout_threads
+            )
 
             # save model
-            if (episode % self.save_interval == 0 or episode == episodes - 1):
+            if episode % self.save_interval == 0 or episode == episodes - 1:
                 self.save()
 
             # log information
             if episode % self.log_interval == 0:
                 end = time.time()
-                print("\n Scenario {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
-                      .format(self.all_args.scenario_name,
-                              self.algorithm_name,
-                              self.experiment_name,
-                              episode,
-                              episodes,
-                              total_num_steps,
-                              self.num_env_steps,
-                              int(total_num_steps / (end - start))))
+                print(
+                    "\n Scenario {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n".format(
+                        self.all_args.scenario_name,
+                        self.algorithm_name,
+                        self.experiment_name,
+                        episode,
+                        episodes,
+                        total_num_steps,
+                        self.num_env_steps,
+                        int(total_num_steps / (end - start)),
+                    )
+                )
 
                 # if self.env_name == "MPE":
                 #     env_infos = {}
@@ -78,8 +103,14 @@ class EnvRunner(Runner):
                 #         agent_k = 'agent%i/individual_rewards' % agent_id
                 #         env_infos[agent_k] = idv_rews
 
-                train_infos["average_episode_rewards"] = np.mean(self.buffer.rewards) * self.episode_length
-                print("average episode rewards is {}".format(train_infos["average_episode_rewards"]))
+                train_infos["average_episode_rewards"] = (
+                    np.mean(self.buffer.rewards) * self.episode_length
+                )
+                print(
+                    "average episode rewards is {}".format(
+                        train_infos["average_episode_rewards"]
+                    )
+                )
                 self.log_train(train_infos, total_num_steps)
                 # self.log_env(env_infos, total_num_steps)
 
@@ -94,7 +125,9 @@ class EnvRunner(Runner):
         # replay buffer
         if self.use_centralized_V:
             share_obs = obs.reshape(self.n_rollout_threads, -1)  # shape = (5, 28)
-            share_obs = np.expand_dims(share_obs, 1).repeat(self.num_agents, axis=1)  # shape = (5, 2, 28)
+            share_obs = np.expand_dims(share_obs, 1).repeat(
+                self.num_agents, axis=1
+            )  # shape = (5, 2, 28)
         else:
             share_obs = obs
 
@@ -104,43 +137,78 @@ class EnvRunner(Runner):
     @torch.no_grad()
     def collect(self, step):
         self.trainer.prep_rollout()
-        value, action, action_log_prob, rnn_states, rnn_states_critic \
-            = self.trainer.policy.get_actions(np.concatenate(self.buffer.share_obs[step]),
-                                              np.concatenate(self.buffer.obs[step]),
-                                              np.concatenate(self.buffer.rnn_states[step]),
-                                              np.concatenate(self.buffer.rnn_states_critic[step]),
-                                              np.concatenate(self.buffer.masks[step]))
+        (
+            value,
+            action,
+            action_log_prob,
+            rnn_states,
+            rnn_states_critic,
+        ) = self.trainer.policy.get_actions(
+            np.concatenate(self.buffer.share_obs[step]),
+            np.concatenate(self.buffer.obs[step]),
+            np.concatenate(self.buffer.rnn_states[step]),
+            np.concatenate(self.buffer.rnn_states_critic[step]),
+            np.concatenate(self.buffer.masks[step]),
+        )
         # [self.envs, agents, dim]
         values = np.array(np.split(_t2n(value), self.n_rollout_threads))
         actions = np.array(np.split(_t2n(action), self.n_rollout_threads))
-        action_log_probs = np.array(np.split(_t2n(action_log_prob), self.n_rollout_threads))
+        action_log_probs = np.array(
+            np.split(_t2n(action_log_prob), self.n_rollout_threads)
+        )
         rnn_states = np.array(np.split(_t2n(rnn_states), self.n_rollout_threads))
-        rnn_states_critic = np.array(np.split(_t2n(rnn_states_critic), self.n_rollout_threads))
+        rnn_states_critic = np.array(
+            np.split(_t2n(rnn_states_critic), self.n_rollout_threads)
+        )
         # rearrange action
-        if self.envs.action_space[0].__class__.__name__ == 'MultiDiscrete':
+        if self.envs.action_space[0].__class__.__name__ == "MultiDiscrete":
             for i in range(self.envs.action_space[0].shape):
-                uc_actions_env = np.eye(self.envs.action_space[0].high[i] + 1)[actions[:, :, i]]
+                uc_actions_env = np.eye(self.envs.action_space[0].high[i] + 1)[
+                    actions[:, :, i]
+                ]
                 if i == 0:
                     actions_env = uc_actions_env
                 else:
                     actions_env = np.concatenate((actions_env, uc_actions_env), axis=2)
-        elif self.envs.action_space[0].__class__.__name__ == 'Discrete':
+        elif self.envs.action_space[0].__class__.__name__ == "Discrete":
             # actions  --> actions_env : shape:[10, 1] --> [5, 2, 5]
             actions_env = np.squeeze(np.eye(self.envs.action_space[0].n)[actions], 2)
         else:
             # TODO 这里改造成自己环境需要的形式即可
+            # TODO Here, you can change the shape of actions_env to fit your environment
             actions_env = actions
             # raise NotImplementedError
 
-        return values, actions, action_log_probs, rnn_states, rnn_states_critic, actions_env
+        return (
+            values,
+            actions,
+            action_log_probs,
+            rnn_states,
+            rnn_states_critic,
+            actions_env,
+        )
 
     def insert(self, data):
-        obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic = data
-
-        rnn_states[dones == True] = np.zeros(((dones == True).sum(), self.recurrent_N, self.hidden_size),
-                                             dtype=np.float32)
-        rnn_states_critic[dones == True] = np.zeros(((dones == True).sum(), *self.buffer.rnn_states_critic.shape[3:]),
-                                                    dtype=np.float32)
+        (
+            obs,
+            rewards,
+            dones,
+            infos,
+            values,
+            actions,
+            action_log_probs,
+            rnn_states,
+            rnn_states_critic,
+        ) = data
+
+        rnn_states[dones == True] = np.zeros(
+            ((dones == True).sum(), self.recurrent_N, self.hidden_size),
+            dtype=np.float32,
+        )
+        rnn_states_critic[dones == True] = np.zeros(
+            ((dones == True).sum(), *self.buffer.rnn_states_critic.shape[3:]),
+            dtype=np.float32,
+        )
         masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
         masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)
 
@@ -150,52 +218,93 @@ class EnvRunner(Runner):
         else:
             share_obs = obs
 
-        self.buffer.insert(share_obs, obs, rnn_states, rnn_states_critic, actions, action_log_probs, values, rewards,
-                           masks)
+        self.buffer.insert(
+            share_obs,
+            obs,
+            rnn_states,
+            rnn_states_critic,
+            actions,
+            action_log_probs,
+            values,
+            rewards,
+            masks,
+        )
 
     @torch.no_grad()
     def eval(self, total_num_steps):
         eval_episode_rewards = []
         eval_obs = self.eval_envs.reset()
 
-        eval_rnn_states = np.zeros((self.n_eval_rollout_threads, *self.buffer.rnn_states.shape[2:]), dtype=np.float32)
-        eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)
+        eval_rnn_states = np.zeros(
+            (self.n_eval_rollout_threads, *self.buffer.rnn_states.shape[2:]),
+            dtype=np.float32,
+        )
+        eval_masks = np.ones(
+            (self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32
+        )
 
         for eval_step in range(self.episode_length):
             self.trainer.prep_rollout()
-            eval_action, eval_rnn_states = self.trainer.policy.act(np.concatenate(eval_obs),
-                                                                   np.concatenate(eval_rnn_states),
-                                                                   np.concatenate(eval_masks),
-                                                                   deterministic=True)
-            eval_actions = np.array(np.split(_t2n(eval_action), self.n_eval_rollout_threads))
-            eval_rnn_states = np.array(np.split(_t2n(eval_rnn_states), self.n_eval_rollout_threads))
-
-            if self.eval_envs.action_space[0].__class__.__name__ == 'MultiDiscrete':
+            eval_action, eval_rnn_states = self.trainer.policy.act(
+                np.concatenate(eval_obs),
+                np.concatenate(eval_rnn_states),
+                np.concatenate(eval_masks),
+                deterministic=True,
+            )
+            eval_actions = np.array(
+                np.split(_t2n(eval_action), self.n_eval_rollout_threads)
+            )
+            eval_rnn_states = np.array(
+                np.split(_t2n(eval_rnn_states), self.n_eval_rollout_threads)
+            )
+
+            if self.eval_envs.action_space[0].__class__.__name__ == "MultiDiscrete":
                 for i in range(self.eval_envs.action_space[0].shape):
-                    eval_uc_actions_env = np.eye(self.eval_envs.action_space[0].high[i] + 1)[eval_actions[:, :, i]]
+                    eval_uc_actions_env = np.eye(
+                        self.eval_envs.action_space[0].high[i] + 1
+                    )[eval_actions[:, :, i]]
                     if i == 0:
                         eval_actions_env = eval_uc_actions_env
                     else:
-                        eval_actions_env = np.concatenate((eval_actions_env, eval_uc_actions_env), axis=2)
-            elif self.eval_envs.action_space[0].__class__.__name__ == 'Discrete':
-                eval_actions_env = np.squeeze(np.eye(self.eval_envs.action_space[0].n)[eval_actions], 2)
+                        eval_actions_env = np.concatenate(
+                            (eval_actions_env, eval_uc_actions_env), axis=2
+                        )
+            elif self.eval_envs.action_space[0].__class__.__name__ == "Discrete":
+                eval_actions_env = np.squeeze(
+                    np.eye(self.eval_envs.action_space[0].n)[eval_actions], 2
+                )
             else:
                 raise NotImplementedError
 
             # Obser reward and next obs
-            eval_obs, eval_rewards, eval_dones, eval_infos = self.eval_envs.step(eval_actions_env)
+            eval_obs, eval_rewards, eval_dones, eval_infos = self.eval_envs.step(
+                eval_actions_env
+            )
             eval_episode_rewards.append(eval_rewards)
 
             eval_rnn_states[eval_dones == True] = np.zeros(
-                ((eval_dones == True).sum(), self.recurrent_N, self.hidden_size), dtype=np.float32)
-            eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)
-            eval_masks[eval_dones == True] = np.zeros(((eval_dones == True).sum(), 1), dtype=np.float32)
+                ((eval_dones == True).sum(), self.recurrent_N, self.hidden_size),
+                dtype=np.float32,
+            )
+            eval_masks = np.ones(
+                (self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32
+            )
+            eval_masks[eval_dones == True] = np.zeros(
+                ((eval_dones == True).sum(), 1), dtype=np.float32
+            )
 
         eval_episode_rewards = np.array(eval_episode_rewards)
         eval_env_infos = {}
-        eval_env_infos['eval_average_episode_rewards'] = np.sum(np.array(eval_episode_rewards), axis=0)
-        eval_average_episode_rewards = np.mean(eval_env_infos['eval_average_episode_rewards'])
-        print("eval average episode rewards of agent: " + str(eval_average_episode_rewards))
+        eval_env_infos["eval_average_episode_rewards"] = np.sum(
+            np.array(eval_episode_rewards), axis=0
+        )
+        eval_average_episode_rewards = np.mean(
+            eval_env_infos["eval_average_episode_rewards"]
+        )
+        print(
+            "eval average episode rewards of agent: "
+            + str(eval_average_episode_rewards)
+        )
         self.log_env(eval_env_infos, total_num_steps)
 
     @torch.no_grad()
@@ -207,14 +316,23 @@ class EnvRunner(Runner):
         for episode in range(self.all_args.render_episodes):
             obs = envs.reset()
             if self.all_args.save_gifs:
-                image = envs.render('rgb_array')[0][0]
+                image = envs.render("rgb_array")[0][0]
                 all_frames.append(image)
             else:
-                envs.render('human')
-
-            rnn_states = np.zeros((self.n_rollout_threads, self.num_agents, self.recurrent_N, self.hidden_size),
-                                  dtype=np.float32)
-            masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
+                envs.render("human")
+
+            rnn_states = np.zeros(
+                (
+                    self.n_rollout_threads,
+                    self.num_agents,
+                    self.recurrent_N,
+                    self.hidden_size,
+                ),
+                dtype=np.float32,
+            )
+            masks = np.ones(
+                (self.n_rollout_threads, self.num_agents, 1), dtype=np.float32
+            )
 
             episode_rewards = []
 
@@ -222,21 +340,29 @@ class EnvRunner(Runner):
                 calc_start = time.time()
 
                 self.trainer.prep_rollout()
-                action, rnn_states = self.trainer.policy.act(np.concatenate(obs),
-                                                             np.concatenate(rnn_states),
-                                                             np.concatenate(masks),
-                                                             deterministic=True)
+                action, rnn_states = self.trainer.policy.act(
+                    np.concatenate(obs),
+                    np.concatenate(rnn_states),
+                    np.concatenate(masks),
+                    deterministic=True,
+                )
                 actions = np.array(np.split(_t2n(action), self.n_rollout_threads))
-                rnn_states = np.array(np.split(_t2n(rnn_states), self.n_rollout_threads))
+                rnn_states = np.array(
+                    np.split(_t2n(rnn_states), self.n_rollout_threads)
+                )
 
-                if envs.action_space[0].__class__.__name__ == 'MultiDiscrete':
+                if envs.action_space[0].__class__.__name__ == "MultiDiscrete":
                     for i in range(envs.action_space[0].shape):
-                        uc_actions_env = np.eye(envs.action_space[0].high[i] + 1)[actions[:, :, i]]
+                        uc_actions_env = np.eye(envs.action_space[0].high[i] + 1)[
+                            actions[:, :, i]
+                        ]
                         if i == 0:
                             actions_env = uc_actions_env
                         else:
-                            actions_env = np.concatenate((actions_env, uc_actions_env), axis=2)
-                elif envs.action_space[0].__class__.__name__ == 'Discrete':
+                            actions_env = np.concatenate(
+                                (actions_env, uc_actions_env), axis=2
+                            )
+                elif envs.action_space[0].__class__.__name__ == "Discrete":
                     actions_env = np.squeeze(np.eye(envs.action_space[0].n)[actions], 2)
                 else:
                     raise NotImplementedError
@@ -245,22 +371,31 @@ class EnvRunner(Runner):
                 obs, rewards, dones, infos = envs.step(actions_env)
                 episode_rewards.append(rewards)
 
-                rnn_states[dones == True] = np.zeros(((dones == True).sum(), self.recurrent_N, self.hidden_size),
-                                                     dtype=np.float32)
-                masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)
-                masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)
+                rnn_states[dones == True] = np.zeros(
+                    ((dones == True).sum(), self.recurrent_N, self.hidden_size),
+                    dtype=np.float32,
+                )
+                masks = np.ones(
+                    (self.n_rollout_threads, self.num_agents, 1), dtype=np.float32
+                )
+                masks[dones == True] = np.zeros(
+                    ((dones == True).sum(), 1), dtype=np.float32
+                )
 
                 if self.all_args.save_gifs:
-                    image = envs.render('rgb_array')[0][0]
+                    image = envs.render("rgb_array")[0][0]
                     all_frames.append(image)
                     calc_end = time.time()
                     elapsed = calc_end - calc_start
                     if elapsed < self.all_args.ifi:
                         time.sleep(self.all_args.ifi - elapsed)
                 else:
-                    envs.render('human')
+                    envs.render("human")
 
-            print("average episode rewards is: " + str(np.mean(np.sum(np.array(episode_rewards), axis=0))))
+            print(
+                "average episode rewards is: "
+                + str(np.mean(np.sum(np.array(episode_rewards), axis=0)))
+            )
 
         # if self.all_args.save_gifs:
         #     imageio.mimsave(str(self.gif_dir) + '/render.gif', all_frames, duration=self.all_args.ifi)

train/train.py

@@ -13,6 +13,13 @@ import setproctitle
 import numpy as np
 from pathlib import Path
 import torch
+
+# Get the parent directory of the current file
+parent_dir = os.path.abspath(os.path.join(os.getcwd(), "."))
+
+# Append the parent directory to sys.path, otherwise the following import will fail
+sys.path.append(parent_dir)
+
 from config import get_config
 from envs.env_wrappers import DummyVecEnv
 
@@ -23,13 +30,17 @@ def make_train_env(all_args):
     def get_env_fn(rank):
         def init_env():
             # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
+            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.
             from envs.env_continuous import ContinuousActionEnv
+
             env = ContinuousActionEnv()
             # from envs.env_discrete import DiscreteActionEnv
             # env = DiscreteActionEnv()
             env.seed(all_args.seed + rank * 1000)
             return env
+
         return init_env
+
     return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)])
 
 
@@ -37,20 +48,26 @@ def make_eval_env(all_args):
     def get_env_fn(rank):
         def init_env():
             # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
+            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.
             from envs.env_continuous import ContinuousActionEnv
+
             env = ContinuousActionEnv()
             # from envs.env_discrete import DiscreteActionEnv
             # env = DiscreteActionEnv()
             env.seed(all_args.seed + rank * 1000)
             return env
+
         return init_env
+
     return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)])
 
 
 def parse_args(args, parser):
-    parser.add_argument('--scenario_name', type=str, default='MyEnv', help="Which scenario to run on")
+    parser.add_argument(
+        "--scenario_name", type=str, default="MyEnv", help="Which scenario to run on"
+    )
     parser.add_argument("--num_landmarks", type=int, default=3)
-    parser.add_argument('--num_agents', type=int, default=2, help="number of players")
+    parser.add_argument("--num_agents", type=int, default=2, help="number of players")
 
     all_args = parser.parse_known_args(args)[0]
 
@@ -62,15 +79,21 @@ def main(args):
     all_args = parse_args(args, parser)
 
     if all_args.algorithm_name == "rmappo":
-        assert (all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy), ("check recurrent policy!")
+        assert (
+            all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy
+        ), "check recurrent policy!"
     elif all_args.algorithm_name == "mappo":
-        assert (all_args.use_recurrent_policy == False and all_args.use_naive_recurrent_policy == False), (
-            "check recurrent policy!")
+        assert (
+            all_args.use_recurrent_policy == False
+            and all_args.use_naive_recurrent_policy == False
+        ), "check recurrent policy!"
     else:
         raise NotImplementedError
 
-    assert (all_args.share_policy == True and all_args.scenario_name == 'simple_speaker_listener') == False, (
-        "The simple_speaker_listener scenario can not use shared policy. Please check the config.py.")
+    assert (
+        all_args.share_policy == True
+        and all_args.scenario_name == "simple_speaker_listener"
+    ) == False, "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."
 
     # cuda
     if all_args.cuda and torch.cuda.is_available():
@@ -86,27 +109,41 @@ def main(args):
         torch.set_num_threads(all_args.n_training_threads)
 
     # run dir
-    run_dir = Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[
-                       0] + "/results") / all_args.env_name / all_args.scenario_name / all_args.algorithm_name / all_args.experiment_name
+    run_dir = (
+        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
+        / all_args.env_name
+        / all_args.scenario_name
+        / all_args.algorithm_name
+        / all_args.experiment_name
+    )
     if not run_dir.exists():
         os.makedirs(str(run_dir))
 
     if not run_dir.exists():
-        curr_run = 'run1'
+        curr_run = "run1"
     else:
-        exst_run_nums = [int(str(folder.name).split('run')[1]) for folder in run_dir.iterdir() if
-                         str(folder.name).startswith('run')]
+        exst_run_nums = [
+            int(str(folder.name).split("run")[1])
+            for folder in run_dir.iterdir()
+            if str(folder.name).startswith("run")
+        ]
         if len(exst_run_nums) == 0:
-            curr_run = 'run1'
+            curr_run = "run1"
         else:
-            curr_run = 'run%i' % (max(exst_run_nums) + 1)
+            curr_run = "run%i" % (max(exst_run_nums) + 1)
     run_dir = run_dir / curr_run
     if not run_dir.exists():
         os.makedirs(str(run_dir))
 
-    setproctitle.setproctitle(str(all_args.algorithm_name) + "-" + \
-                              str(all_args.env_name) + "-" + str(all_args.experiment_name) + "@" + str(
-        all_args.user_name))
+    setproctitle.setproctitle(
+        str(all_args.algorithm_name)
+        + "-"
+        + str(all_args.env_name)
+        + "-"
+        + str(all_args.experiment_name)
+        + "@"
+        + str(all_args.user_name)
+    )
 
     # seed
     torch.manual_seed(all_args.seed)
@@ -124,7 +161,7 @@ def main(args):
         "eval_envs": eval_envs,
         "num_agents": num_agents,
         "device": device,
-        "run_dir": run_dir
+        "run_dir": run_dir,
     }
 
     # run experiments
@@ -141,7 +178,7 @@ def main(args):
     if all_args.use_eval and eval_envs is not envs:
         eval_envs.close()
 
-    runner.writter.export_scalars_to_json(str(runner.log_dir + '/summary.json'))
+    runner.writter.export_scalars_to_json(str(runner.log_dir + "/summary.json"))
     runner.writter.close()