rec_buffer.py
off-policy replay
原文件 · 324 行 · 17299 字节
utils/rec_buffer.py
SHA-256: 7819887145daab82a0522d4e44f15a18de47f882ca59851f2fe8939e55ca2c64
1import numpy as np2from utils.util import get_dim_from_space3from utils.segment_tree import SumSegmentTree, MinSegmentTree456def _cast(x):7 return x.transpose(2, 0, 1, 3)8910class RecReplayBuffer(object):11 def __init__(self, policy_info, policy_agents, buffer_size, episode_length, use_same_share_obs, use_avail_acts,12 use_reward_normalization=False):13 """14 Replay buffer class for training RNN policies. Stores entire episodes rather than single transitions.1516 :param policy_info: (dict) maps policy id to a dict containing information about corresponding policy.17 :param policy_agents: (dict) maps policy id to list of agents controled by corresponding policy.18 :param buffer_size: (int) max number of transitions to store in the buffer.19 :param use_same_share_obs: (bool) whether all agents share the same centralized observation.20 :param use_avail_acts: (bool) whether to store what actions are available.21 :param use_reward_normalization: (bool) whether to use reward normalization.22 """23 self.policy_info = policy_info2425 self.policy_buffers = {p_id: RecPolicyBuffer(buffer_size,26 episode_length,27 len(policy_agents[p_id]),28 self.policy_info[p_id]['obs_space'],29 self.policy_info[p_id]['share_obs_space'],30 self.policy_info[p_id]['act_space'],31 use_same_share_obs,32 use_avail_acts,33 use_reward_normalization)34 for p_id in self.policy_info.keys()}3536 def __len__(self):37 return self.policy_buffers['policy_0'].filled_i3839 def insert(self, num_insert_episodes, obs, share_obs, acts, rewards, dones, dones_env, avail_acts):40 """41 Insert a set of episodes into buffer. If the buffer size overflows, old episodes are dropped.4243 :param num_insert_episodes: (int) number of episodes to be added to buffer44 :param obs: (dict) maps policy id to numpy array of observations of agents corresponding to that policy45 :param share_obs: (dict) maps policy id to numpy array of centralized observation corresponding to that policy46 :param acts: (dict) maps policy id to numpy array of actions of agents corresponding to that policy47 :param rewards: (dict) maps policy id to numpy array of rewards of agents corresponding to that policy48 :param dones: (dict) maps policy id to numpy array of terminal status of agents corresponding to that policy49 :param dones_env: (dict) maps policy id to numpy array of terminal status of env50 :param valid_transition: (dict) maps policy id to numpy array of whether the corresponding transition is valid of agents corresponding to that policy51 :param avail_acts: (dict) maps policy id to numpy array of available actions of agents corresponding to that policy5253 :return: (np.ndarray) indexes in which the new transitions were placed.54 """55 for p_id in self.policy_info.keys():56 idx_range = self.policy_buffers[p_id].insert(num_insert_episodes, np.array(obs[p_id]),57 np.array(share_obs[p_id]), np.array(acts[p_id]),58 np.array(rewards[p_id]), np.array(dones[p_id]),59 np.array(dones_env[p_id]), np.array(avail_acts[p_id]))60 return idx_range6162 def sample(self, batch_size):63 """64 Sample a set of episodes from buffer, uniformly at random.65 :param batch_size: (int) number of episodes to sample from buffer.6667 :return: obs: (dict) maps policy id to sampled observations corresponding to that policy68 :return: share_obs: (dict) maps policy id to sampled observations corresponding to that policy69 :return: acts: (dict) maps policy id to sampled actions corresponding to that policy70 :return: rewards: (dict) maps policy id to sampled rewards corresponding to that policy71 :return: dones: (dict) maps policy id to sampled terminal status of agents corresponding to that policy72 :return: dones_env: (dict) maps policy id to sampled environment terminal status corresponding to that policy73 :return: valid_transition: (dict) maps policy_id to whether each sampled transition is valid or not (invalid if corresponding agent is dead)74 :return: avail_acts: (dict) maps policy_id to available actions corresponding to that policy75 """76 inds = np.random.choice(self.__len__(), batch_size)77 obs, share_obs, acts, rewards, dones, dones_env, avail_acts = {}, {}, {}, {}, {}, {}, {}78 for p_id in self.policy_info.keys():79 obs[p_id], share_obs[p_id], acts[p_id], rewards[p_id], dones[p_id], dones_env[p_id], avail_acts[p_id] = \80 self.policy_buffers[p_id].sample_inds(inds)8182 return obs, share_obs, acts, rewards, dones, dones_env, avail_acts, None, None838485class RecPolicyBuffer(object):86 def __init__(self, buffer_size, episode_length, num_agents, obs_space, share_obs_space, act_space,87 use_same_share_obs, use_avail_acts, use_reward_normalization=False):88 """89 Buffer class containing buffer data corresponding to a single policy.9091 :param buffer_size: (int) max number of episodes to store in buffer.92 :param episode_length: (int) max length of an episode.93 :param num_agents: (int) number of agents controlled by the policy.94 :param obs_space: (gym.Space) observation space of the environment.95 :param share_obs_space: (gym.Space) centralized observation space of the environment.96 :param act_space: (gym.Space) action space of the environment.97 :use_same_share_obs: (bool) whether all agents share the same centralized observation.98 :use_avail_acts: (bool) whether to store what actions are available.99 :param use_reward_normalization: (bool) whether to use reward normalization.100 """101 self.buffer_size = buffer_size102 self.episode_length = episode_length103 self.num_agents = num_agents104 self.use_same_share_obs = use_same_share_obs105 self.use_avail_acts = use_avail_acts106 self.use_reward_normalization = use_reward_normalization107 self.filled_i = 0108 self.current_i = 0109110 # obs111 if obs_space.__class__.__name__ == 'Box':112 obs_shape = obs_space.shape113 share_obs_shape = share_obs_space.shape114 elif obs_space.__class__.__name__ == 'list':115 obs_shape = obs_space116 share_obs_shape = share_obs_space117 else:118 raise NotImplementedError119120 self.obs = np.zeros((self.episode_length + 1, self.buffer_size,121 self.num_agents, obs_shape[0]), dtype=np.float32)122123 if self.use_same_share_obs:124 self.share_obs = np.zeros((self.episode_length + 1, self.buffer_size, share_obs_shape[0]), dtype=np.float32)125 else:126 self.share_obs = np.zeros((self.episode_length + 1, self.buffer_size, self.num_agents, share_obs_shape[0]),127 dtype=np.float32)128129 # action130 act_dim = np.sum(get_dim_from_space(act_space))131 self.acts = np.zeros((self.episode_length, self.buffer_size, self.num_agents, act_dim), dtype=np.float32)132 if self.use_avail_acts:133 self.avail_acts = np.ones((self.episode_length + 1, self.buffer_size, self.num_agents, act_dim),134 dtype=np.float32)135136 # rewards137 self.rewards = np.zeros((self.episode_length, self.buffer_size, self.num_agents, 1), dtype=np.float32)138139 # default to done being True140 self.dones = np.ones_like(self.rewards, dtype=np.float32)141 self.dones_env = np.ones((self.episode_length, self.buffer_size, 1), dtype=np.float32)142143 def __len__(self):144 return self.filled_i145146 def insert(self, num_insert_episodes, obs, share_obs, acts, rewards, dones, dones_env, avail_acts=None):147 """148 Insert a set of episodes corresponding to this policy into buffer. If the buffer size overflows, old transitions are dropped.149150 :param num_insert_steps: (int) number of transitions to be added to buffer151 :param obs: (np.ndarray) observations of agents corresponding to this policy.152 :param share_obs: (np.ndarray) centralized observations of agents corresponding to this policy.153 :param acts: (np.ndarray) actions of agents corresponding to this policy.154 :param rewards: (np.ndarray) rewards of agents corresponding to this policy.155 :param dones: (np.ndarray) terminal status of agents corresponding to this policy.156 :param dones_env: (np.ndarray) environment terminal status.157 :param valid_transition: (np.ndarray) whether each transition is valid or not (invalid if agent was dead during transition)158 :param avail_acts: (np.ndarray) available actions of agents corresponding to this policy.159160 :return: (np.ndarray) indexes of the buffer the new transitions were placed in.161 """162163 # obs: [step, episode, agent, dim]164 episode_length = acts.shape[0]165 assert episode_length == self.episode_length, ("different dimension!")166167 if self.current_i + num_insert_episodes <= self.buffer_size:168 idx_range = np.arange(self.current_i, self.current_i + num_insert_episodes)169 else:170 num_left_episodes = self.current_i + num_insert_episodes - self.buffer_size171 idx_range = np.concatenate((np.arange(self.current_i, self.buffer_size), np.arange(num_left_episodes)))172173 if self.use_same_share_obs:174 # remove agent dimension since all agents share centralized observation175 share_obs = share_obs[:, :, 0]176177 self.obs[:, idx_range] = obs.copy()178 self.share_obs[:, idx_range] = share_obs.copy()179 self.acts[:, idx_range] = acts.copy()180 self.rewards[:, idx_range] = rewards.copy()181 self.dones[:, idx_range] = dones.copy()182 self.dones_env[:, idx_range] = dones_env.copy()183184 if self.use_avail_acts:185 self.avail_acts[:, idx_range] = avail_acts.copy()186187 self.current_i = idx_range[-1] + 1188 self.filled_i = min(self.filled_i + len(idx_range), self.buffer_size)189190 return idx_range191192 def sample_inds(self, sample_inds):193 """194 Sample a set of transitions from buffer from the specified indices.195 :param sample_inds: (np.ndarray) indices of samples to return from buffer.196197 :return: obs: (np.ndarray) sampled observations corresponding to that policy198 :return: share_obs: (np.ndarray) sampled observations corresponding to that policy199 :return: acts: (np.ndarray) sampled actions corresponding to that policy200 :return: rewards: (np.ndarray) sampled rewards corresponding to that policy201 :return: dones: (np.ndarray) sampled terminal status of agents corresponding to that policy202 :return: dones_env: (np.ndarray) sampled environment terminal status corresponding to that policy203 :return: valid_transition: (np.ndarray) whether each sampled transition in episodes are valid or not (invalid if corresponding agent is dead)204 :return: avail_acts: (np.ndarray) sampled available actions corresponding to that policy205 """206207 obs = _cast(self.obs[:, sample_inds])208 acts = _cast(self.acts[:, sample_inds])209 if self.use_reward_normalization:210 # mean std211 # [length, envs, agents, 1]212 # [length, envs, 1]213 all_dones_env = np.tile(np.expand_dims(214 self.dones_env[:, :self.filled_i], -1), (1, 1, self.num_agents, 1))215 first_step_dones_env = np.zeros((1, self.filled_i, self.num_agents, 1))216 curr_dones_env = np.concatenate((first_step_dones_env, all_dones_env[:self.episode_length - 1]))217 temp_rewards = self.rewards[:, :self.filled_i].copy()218 temp_rewards[curr_dones_env == 1.0] = np.nan219220 mean_reward = np.nanmean(temp_rewards)221 std_reward = np.nanstd(temp_rewards)222 rewards = _cast(223 (self.rewards[:, sample_inds] - mean_reward) / std_reward)224 else:225 rewards = _cast(self.rewards[:, sample_inds])226227 if self.use_same_share_obs:228 share_obs = self.share_obs[:, sample_inds]229 else:230 share_obs = _cast(self.share_obs[:, sample_inds])231232 dones = _cast(self.dones[:, sample_inds])233 dones_env = self.dones_env[:, sample_inds]234235 if self.use_avail_acts:236 avail_acts = _cast(self.avail_acts[:, sample_inds])237 else:238 avail_acts = None239240 return obs, share_obs, acts, rewards, dones, dones_env, avail_acts241242243class PrioritizedRecReplayBuffer(RecReplayBuffer):244 def __init__(self, alpha, policy_info, policy_agents, buffer_size, episode_length, use_same_share_obs,245 use_avail_acts, use_reward_normalization=False):246 """Prioritized replay buffer class for training RNN policies. See parent class."""247 super(PrioritizedRecReplayBuffer, self).__init__(policy_info, policy_agents, buffer_size,248 episode_length, use_same_share_obs, use_avail_acts,249 use_reward_normalization)250 self.alpha = alpha251 self.policy_info = policy_info252 it_capacity = 1253 while it_capacity < buffer_size:254 it_capacity *= 2255256 self._it_sums = {p_id: SumSegmentTree(257 it_capacity) for p_id in self.policy_info.keys()}258 self._it_mins = {p_id: MinSegmentTree(259 it_capacity) for p_id in self.policy_info.keys()}260 self.max_priorities = {p_id: 1.0 for p_id in self.policy_info.keys()}261262 def insert(self, num_insert_episodes, obs, share_obs, acts, rewards, dones, dones_env, avail_acts=None):263 """See parent class."""264 idx_range = super().insert(num_insert_episodes, obs, share_obs, acts, rewards, dones, dones_env, avail_acts)265 for idx in range(idx_range[0], idx_range[1]):266 for p_id in self.policy_info.keys():267 self._it_sums[p_id][idx] = self.max_priorities[p_id] ** self.alpha268 self._it_mins[p_id][idx] = self.max_priorities[p_id] ** self.alpha269270 return idx_range271272 def _sample_proportional(self, batch_size, p_id=None):273 total = self._it_sums[p_id].sum(0, len(self) - 1)274 mass = np.random.random(size=batch_size) * total275 idx = self._it_sums[p_id].find_prefixsum_idx(mass)276 return idx277278 def sample(self, batch_size, beta=0, p_id=None):279 """280 Sample a set of episodes from buffer; probability of choosing a given episode is proportional to its priority.281 :param batch_size: (int) number of episodes to sample.282 :param beta: (float) controls the amount of prioritization to apply.283 :param p_id: (str) policy which will be updated using the samples.284285 :return: See parent class.286 """287 assert len(288 self) > batch_size, "Cannot sample with no completed episodes in the buffer!"289 assert beta > 0290291 batch_inds = self._sample_proportional(batch_size, p_id)292293 p_min = self._it_mins[p_id].min() / self._it_sums[p_id].sum()294 max_weight = (p_min * len(self)) ** (-beta)295 p_sample = self._it_sums[p_id][batch_inds] / self._it_sums[p_id].sum()296 weights = (p_sample * len(self)) ** (-beta) / max_weight297298 obs, share_obs, acts, rewards, dones, dones_env, avail_acts = {}, {}, {}, {}, {}, {}, {}299 for p_id in self.policy_info.keys():300 p_buffer = self.policy_buffers[p_id]301 obs[p_id], share_obs[p_id], acts[p_id], rewards[p_id], dones[p_id], dones_env[p_id], avail_acts[302 p_id] = p_buffer.sample_inds(batch_inds)303304 return obs, share_obs, acts, rewards, dones, dones_env, avail_acts, weights, batch_inds305306 def update_priorities(self, idxes, priorities, p_id=None):307 """308 Update priorities of sampled transitions.309 sets priority of transition at index idxes[i] in buffer310 to priorities[i].311 :param idxes: ([int]) List of idxes of sampled transitions312 :param priorities: ([float]) List of updated priorities corresponding to transitions at the sampled idxes313 denoted by variable `idxes`.314 """315 assert len(idxes) == len(priorities)316 assert np.min(priorities) > 0317 assert np.min(idxes) >= 0318 assert np.max(idxes) < len(self)319320 self._it_sums[p_id][idxes] = priorities ** self.alpha321 self._it_mins[p_id][idxes] = priorities ** self.alpha322323 self.max_priorities[p_id] = max(324 self.max_priorities[p_id], np.max(priorities))