env_discrete.py
CapCoal 空间封装
原文件 · 192 行 · 7035 字节
envs/capcoal/env_discrete.py
SHA-256: 5ef245ced7f53cc716210a4fac0bfd8a77549b931f0e8d7bae026d8cd97b6cf4
1"""2# @Time : 2021/7/2 5:22 下午3# @Author : hezhiqiang4# @Email : tinyzqh@163.com5# @File : env_discrete.py6"""78import gym9from gym import spaces10import numpy as np11from .env_core import EnvCore121314class DiscreteActionEnv(object):15 """16 对于离散动作环境的封装17 Wrapper for discrete action environment.18 """1920 def __init__(self, config):21 self.env = EnvCore(config)22 self.num_agent = self.env.agent_num2324 self.signal_obs_dim = self.env.obs_dim25 self.signal_action_dim = self.env.action_dim26 self.cap_dim = self.env.cap_dim27 self.task_num = self.env.task_num28 self.per_task_feat = self.env.per_task_feat29 self.ego_dim = self.env.ego_feat_dim30 self.other_agent_feat_dim = self.env.other_agent_feat_dim31 self.global_agent_feat_dim = self.env.global_agent_feat_dim32 self.global_task_feat_dim = self.env.global_task_feat_dim33 self.global_feat_dim = self.env.global_feat_dim3435 # if true, action is a number 0...N, otherwise action is a one-hot N-dimensional vector36 # Edited37 self.discrete_action_input = True3839 self.movable = True4041 # configure spaces42 self.action_space = []43 self.observation_space = []44 self.share_observation_space = []4546 total_action_space = []47 for agent_idx in range(self.num_agent):48 # physical action space49 u_action_space = spaces.Discrete(self.signal_action_dim) # signal_action_dim 个离散的动作5051 total_action_space.append(u_action_space)52 self.action_space.append(total_action_space[agent_idx])5354 # observation space: policy keys plus structured graph features.55 self.observation_space.append(56 spaces.Dict({57 "ego": spaces.Box(58 low=-np.inf,59 high=+np.inf,60 shape=(self.ego_dim,),61 dtype=np.float32,62 ),63 "tasks": spaces.Box(64 low=-np.inf,65 high=+np.inf,66 shape=(self.task_num, self.per_task_feat),67 dtype=np.float32,68 ),69 "others": spaces.Box(70 low=-np.inf,71 high=+np.inf,72 shape=(self.num_agent, self.other_agent_feat_dim),73 dtype=np.float32,74 ),75 })76 )7778 state_space = spaces.Dict({79 "agents": spaces.Box(80 low=-np.inf,81 high=+np.inf,82 shape=(self.num_agent, self.global_agent_feat_dim),83 dtype=np.float32,84 ),85 "tasks": spaces.Box(86 low=-np.inf,87 high=+np.inf,88 shape=(self.task_num, self.global_task_feat_dim),89 dtype=np.float32,90 ),91 "global": spaces.Box(92 low=-np.inf,93 high=+np.inf,94 shape=(self.global_feat_dim,),95 dtype=np.float32,96 ),97 "focal": spaces.Box(98 low=0.0,99 high=1.0,100 shape=(self.num_agent, 1),101 dtype=np.float32,102 ),103 })104 self.share_observation_space = [state_space for _ in range(self.num_agent)]105106 def step(self, actions):107 """108 输入actions维度假设：109 # actions shape = (5, 2, 5)110 # 5个线程的环境，里面有2个智能体，每个智能体的动作是一个one_hot的5维编码111 Input actions dimension assumption:112 # actions shape = (5, 2, 5)113 # 5 threads of the environment, with 2 intelligent agents inside, and each intelligent agent's action is a 5-dimensional one_hot encoding114 """115 obs, reward, dones, info, active_mask, bad_mask, available_actions = self.env.step(actions)116 obs = self._format_obs_dict(obs)117 return obs, reward, dones, info, active_mask, bad_mask, available_actions118119 def reset(self, episode):120 obs, active_mask, available_actions = self.env.reset(episode)121 obs = self._format_obs_dict(obs)122 return obs, active_mask, available_actions123124 def _format_obs_dict(self, obs):125 return {k: v.astype(np.float32) for k, v in obs.items()}126127 def get_global_state(self):128 return self.env.get_global_state()129130 def close(self):131 pass132133 def render(self, mode="rgb_array"):134 pass135136 def set_reset_mode(self, mode):137 self.env.set_reset_mode(mode)138139 def seed(self, seed):140 return self.env.seed(seed)141142143class MultiDiscrete:144 """145 - The multi-discrete action space consists of a series of discrete action spaces with different parameters146 - It can be adapted to both a Discrete action space or a continuous (Box) action space147 - It is useful to represent game controllers or keyboards where each key can be represented as a discrete action space148 - It is parametrized by passing an array of arrays containing [min, max] for each discrete action space149 where the discrete action space can take any integers from `min` to `max` (both inclusive)150 Note: A value of 0 always need to represent the NOOP action.151 e.g. Nintendo Game Controller152 - Can be conceptualized as 3 discrete action spaces:153 1) Arrow Keys: Discrete 5 - NOOP[0], UP[1], RIGHT[2], DOWN[3], LEFT[4] - params: min: 0, max: 4154 2) Button A: Discrete 2 - NOOP[0], Pressed[1] - params: min: 0, max: 1155 3) Button B: Discrete 2 - NOOP[0], Pressed[1] - params: min: 0, max: 1156 - Can be initialized as157 MultiDiscrete([ [0,4], [0,1], [0,1] ])158 """159160 def __init__(self, array_of_param_array):161 super().__init__()162 self.low = np.array([x[0] for x in array_of_param_array])163 self.high = np.array([x[1] for x in array_of_param_array])164 self.num_discrete_space = self.low.shape[0]165 self.n = np.sum(self.high) + 2166167 def sample(self):168 """Returns a array with one sample from each discrete action space"""169 # For each row: round(random .* (max - min) + min, 0)170 random_array = np.random.rand(self.num_discrete_space)171 return [int(x) for x in np.floor(np.multiply((self.high - self.low + 1.0), random_array) + self.low)]172173 def contains(self, x):174 return (175 len(x) == self.num_discrete_space176 and (np.array(x) >= self.low).all()177 and (np.array(x) <= self.high).all()178 )179180 @property181 def shape(self):182 return self.num_discrete_space183184 def __repr__(self):185 return "MultiDiscrete" + str(self.num_discrete_space)186187 def __eq__(self, other):188 return np.array_equal(self.low, other.low) and np.array_equal(self.high, other.high)189190191if __name__ == "__main__":192 DiscreteActionEnv().step(actions=None)