shared_gat_buffer.py
Dict on-policy 数据与 GAE
原文件 · 593 行 · 32258 字节
utils/shared_gat_buffer.py
SHA-256: 6ae61dc38042189dc2bfd236e3c7c019b1eddf73fe31d82e6c67f3cdaced5b1f
1import torch2import numpy as np3from utils.util import get_shape_from_obs_space, get_shape_from_act_space4from utils.obs_utils import POLICY_OBS_KEYS567def _flatten(T, N, x):8 return x.reshape(T * N, *x.shape[2:])91011def _cast(x):12 return x.transpose(1, 2, 0, *range(3, x.ndim)).reshape(-1, *x.shape[3:])131415def _normalize_shape(shape):16 if type(shape[-1]) == list:17 return shape[:1]18 return shape192021def _main_obs_shape(obs_shape):22 if isinstance(obs_shape, dict):23 return (int(sum(np.prod(obs_shape[k]) for k in POLICY_OBS_KEYS)),)24 return _normalize_shape(obs_shape)252627class SharedGatReplayBuffer(object):28 """29 Buffer to store training data.30 :param args: (argparse.Namespace) arguments containing relevant model, policy, and env information.31 :param num_agents: (int) number of agents in the env.32 :param obs_space: (gym.Space) observation space of agents.33 :param cent_obs_space: (gym.Space) centralized observation space of agents.34 :param act_space: (gym.Space) action space for agents.35 """3637 def __init__(self, args, num_agents, obs_space, cent_obs_space, act_space):38 self.episode_length = args.episode_length39 self.n_rollout_threads = args.n_rollout_threads40 self.hidden_size = args.hidden_size41 self.recurrent_N = args.recurrent_N42 self.gamma = args.gamma43 self.gae_lambda = args.gae_lambda44 self._use_gae = args.use_gae45 self._use_popart = args.use_popart46 self._use_valuenorm = args.use_valuenorm47 self._use_proper_time_limits = args.use_proper_time_limits4849 obs_shape_dict = get_shape_from_obs_space(obs_space) # dict for Dict space50 share_obs_shape = get_shape_from_obs_space(cent_obs_space)5152 if isinstance(share_obs_shape, dict):53 self.share_obs = {54 k: np.zeros(55 (self.episode_length + 1, self.n_rollout_threads, num_agents, *_normalize_shape(v)),56 dtype=np.float32,57 )58 for k, v in share_obs_shape.items()59 }60 else:61 share_obs_shape = _main_obs_shape(share_obs_shape)62 self.share_obs = np.zeros(63 (self.episode_length + 1, self.n_rollout_threads, num_agents, *share_obs_shape),64 dtype=np.float32,65 )66 if isinstance(obs_shape_dict, dict):67 self.obs = {68 k: np.zeros((self.episode_length + 1, self.n_rollout_threads, num_agents, *_normalize_shape(v)),69 dtype=np.float32)70 for k, v in obs_shape_dict.items()71 }72 else:73 obs_shape = _normalize_shape(obs_shape_dict)74 self.obs = {75 "obs": np.zeros((self.episode_length + 1, self.n_rollout_threads, num_agents, *obs_shape),76 dtype=np.float32)77 }7879 self.rnn_states = np.zeros(80 (self.episode_length + 1, self.n_rollout_threads, num_agents, self.recurrent_N, self.hidden_size),81 dtype=np.float32)82 self.rnn_states_critic = np.zeros_like(self.rnn_states)8384 self.value_preds = np.zeros(85 (self.episode_length + 1, self.n_rollout_threads, num_agents, 1), dtype=np.float32)86 self.returns = np.zeros_like(self.value_preds)8788 # 为啥只有 Discrete 才有 available_actions ？给你注释了89 # if act_space.__class__.__name__ == 'Discrete':90 self.available_actions = np.ones((self.episode_length + 1, self.n_rollout_threads, num_agents, act_space.n),91 dtype=np.float32)92 # else:93 # self.available_actions = None9495 act_shape = get_shape_from_act_space(act_space)9697 self.actions = np.zeros(98 (self.episode_length, self.n_rollout_threads, num_agents, act_shape), dtype=np.float32)99 self.action_log_probs = np.zeros(100 (self.episode_length, self.n_rollout_threads, num_agents, act_shape), dtype=np.float32)101 # hierarchical reconfiguration gate 的 z 动作必须像 task action 一样存入102 # rollout buffer；PPO update 会用它重放 joint logprob。103 self.gate_actions = np.zeros(104 (self.episode_length, self.n_rollout_threads, num_agents, 1), dtype=np.float32)105 self.rewards = np.zeros(106 (self.episode_length, self.n_rollout_threads, num_agents, 1), dtype=np.float32)107108 self.masks = np.ones((self.episode_length + 1, self.n_rollout_threads, num_agents, 1), dtype=np.float32)109 self.bad_masks = np.ones_like(self.masks)110 self.active_masks = np.ones_like(self.masks)111112 self.step = 0113114 def insert(self, share_obs, obs, rnn_states_actor, rnn_states_critic, actions, action_log_probs,115 gate_actions,116 value_preds, rewards, masks, bad_masks=None, active_masks=None, available_actions=None):117 """118 Insert data into the buffer.119 :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.120 :param obs: (np.ndarray) local agent observations.121 :param rnn_states_actor: (np.ndarray) RNN states for actor network.122 :param rnn_states_critic: (np.ndarray) RNN states for critic network.123 :param actions:(np.ndarray) actions taken by agents.124 :param action_log_probs:(np.ndarray) log probs of actions taken by agents125 :param value_preds: (np.ndarray) value function prediction at each step.126 :param rewards: (np.ndarray) reward collected at each step.127 :param masks: (np.ndarray) denotes whether the environment has terminated or not.128 :param bad_masks: (np.ndarray) action space for agents.129 :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.130 :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.131 """132 if isinstance(self.share_obs, dict):133 for k in self.share_obs:134 self.share_obs[k][self.step + 1] = share_obs[k].copy()135 else:136 self.share_obs[self.step + 1] = share_obs.copy()137 for k in self.obs:138 self.obs[k][self.step + 1] = obs[k].copy()139 self.rnn_states[self.step + 1] = rnn_states_actor.copy()140 self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()141 self.actions[self.step] = actions.copy()142 self.action_log_probs[self.step] = action_log_probs.copy()143 # 非 hierarchical 运行写入零占位，保证 generator 的 sample tuple 不分叉。144 self.gate_actions[self.step] = gate_actions.copy()145 self.value_preds[self.step] = value_preds.copy()146 self.rewards[self.step] = rewards.copy()147 self.masks[self.step + 1] = masks.copy()148 if bad_masks is not None:149 self.bad_masks[self.step + 1] = bad_masks.copy()150 if active_masks is not None:151 self.active_masks[self.step + 1] = active_masks.copy()152 if available_actions is not None:153 self.available_actions[self.step + 1] = available_actions.copy()154155 self.step = (self.step + 1) % self.episode_length156157 def chooseinsert(self, share_obs, obs, rnn_states, rnn_states_critic, actions, action_log_probs,158 value_preds, rewards, masks, bad_masks=None, active_masks=None, available_actions=None):159 """160 Insert data into the buffer. This insert function is used specifically for Hanabi, which is turn based.161 :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.162 :param obs: (np.ndarray) local agent observations.163 :param rnn_states_actor: (np.ndarray) RNN states for actor network.164 :param rnn_states_critic: (np.ndarray) RNN states for critic network.165 :param actions:(np.ndarray) actions taken by agents.166 :param action_log_probs:(np.ndarray) log probs of actions taken by agents167 :param value_preds: (np.ndarray) value function prediction at each step.168 :param rewards: (np.ndarray) reward collected at each step.169 :param masks: (np.ndarray) denotes whether the environment has terminated or not.170 :param bad_masks: (np.ndarray) denotes indicate whether whether true terminal state or due to episode limit171 :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.172 :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.173 """174 if isinstance(self.share_obs, dict):175 for k in self.share_obs:176 self.share_obs[k][self.step] = share_obs[k].copy()177 else:178 self.share_obs[self.step] = share_obs.copy()179 for k in self.obs:180 self.obs[k][self.step] = obs[k].copy()181 self.rnn_states[self.step + 1] = rnn_states.copy()182 self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()183 self.actions[self.step] = actions.copy()184 self.action_log_probs[self.step] = action_log_probs.copy()185 self.value_preds[self.step] = value_preds.copy()186 self.rewards[self.step] = rewards.copy()187 self.masks[self.step + 1] = masks.copy()188 if bad_masks is not None:189 self.bad_masks[self.step + 1] = bad_masks.copy()190 if active_masks is not None:191 self.active_masks[self.step] = active_masks.copy()192 if available_actions is not None:193 self.available_actions[self.step] = available_actions.copy()194195 self.step = (self.step + 1) % self.episode_length196197 def after_update(self):198 """Copy last timestep data to first index. Called after update to model."""199 if isinstance(self.share_obs, dict):200 for k in self.share_obs:201 self.share_obs[k][0] = self.share_obs[k][-1].copy()202 else:203 self.share_obs[0] = self.share_obs[-1].copy()204 for k in self.obs:205 self.obs[k][0] = self.obs[k][-1].copy()206 self.rnn_states[0] = self.rnn_states[-1].copy()207 self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()208 self.masks[0] = self.masks[-1].copy()209 self.bad_masks[0] = self.bad_masks[-1].copy()210 self.active_masks[0] = self.active_masks[-1].copy()211 if self.available_actions is not None:212 self.available_actions[0] = self.available_actions[-1].copy()213214 def chooseafter_update(self):215 """Copy last timestep data to first index. This method is used for Hanabi."""216 self.rnn_states[0] = self.rnn_states[-1].copy()217 self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()218 self.masks[0] = self.masks[-1].copy()219 self.bad_masks[0] = self.bad_masks[-1].copy()220221 def compute_returns(self, next_value, value_normalizer=None):222 """223 Compute returns either as discounted sum of rewards, or using GAE.224 :param next_value: (np.ndarray) value predictions for the step after the last episode step.225 :param value_normalizer: (PopArt) If not None, PopArt value normalizer instance.226 """227 if self._use_proper_time_limits:228 if self._use_gae:229 self.value_preds[-1] = next_value230 gae = 0231 for step in reversed(range(self.rewards.shape[0])):232 if self._use_popart or self._use_valuenorm:233 # step + 1234 delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(235 self.value_preds[step + 1]) * self.masks[step + 1] \236 - value_normalizer.denormalize(self.value_preds[step])237 gae = delta + self.gamma * self.gae_lambda * gae * self.masks[step + 1]238 gae = gae * self.bad_masks[step + 1]239 self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])240 else:241 delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \242 self.value_preds[step]243 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae244 gae = gae * self.bad_masks[step + 1]245 self.returns[step] = gae + self.value_preds[step]246 else:247 self.returns[-1] = next_value248 for step in reversed(range(self.rewards.shape[0])):249 if self._use_popart or self._use_valuenorm:250 self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[251 step]) * self.bad_masks[step + 1] \252 + (1 - self.bad_masks[step + 1]) * value_normalizer.denormalize(253 self.value_preds[step])254 else:255 self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[256 step]) * self.bad_masks[step + 1] \257 + (1 - self.bad_masks[step + 1]) * self.value_preds[step]258 else:259 if self._use_gae:260 self.value_preds[-1] = next_value261 gae = 0262 for step in reversed(range(self.rewards.shape[0])):263 if self._use_popart or self._use_valuenorm:264 delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(265 self.value_preds[step + 1]) * self.masks[step + 1] \266 - value_normalizer.denormalize(self.value_preds[step])267 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae268 self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])269 else:270 delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \271 self.value_preds[step]272 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae273 self.returns[step] = gae + self.value_preds[step]274 else:275 self.returns[-1] = next_value276 for step in reversed(range(self.rewards.shape[0])):277 self.returns[step] = self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[step]278279 def feed_forward_generator(self, advantages, num_mini_batch=None, mini_batch_size=None):280 """281 Yield training data for MLP policies.282 :param advantages: (np.ndarray) advantage estimates.283 :param num_mini_batch: (int) number of minibatches to split the batch into.284 :param mini_batch_size: (int) number of samples in each minibatch.285 """286 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]287 batch_size = n_rollout_threads * episode_length # * num_agents288289 if mini_batch_size is None:290 assert batch_size >= num_mini_batch, (291 "PPO requires the number of processes ({}) "292 "* number of steps ({}) * number of agents ({}) = {} "293 "to be greater than or equal to the number of PPO mini batches ({})."294 "".format(n_rollout_threads, episode_length, num_agents,295 n_rollout_threads * episode_length * num_agents,296 num_mini_batch))297 mini_batch_size = batch_size // num_mini_batch298299 rand = torch.randperm(batch_size).numpy()300 sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]301302 if isinstance(self.share_obs, dict):303 share_obs = {304 k: v[:-1].reshape(-1, *v.shape[2:])305 for k, v in self.share_obs.items()306 }307 else:308 share_obs = self.share_obs[:-1].reshape(-1, *self.share_obs.shape[2:])309 obs = {310 k: v[:-1].reshape(-1, *v.shape[2:])311 for k, v in self.obs.items()312 }313 rnn_states = self.rnn_states[:-1].reshape(-1, *self.rnn_states.shape[2:])314 rnn_states_critic = self.rnn_states_critic[:-1].reshape(-1, *self.rnn_states_critic.shape[2:])315 actions = self.actions.reshape(-1, *self.actions.shape[2:])316 if self.available_actions is not None:317 available_actions = self.available_actions[:-1].reshape(-1, *self.available_actions.shape[2:])318 value_preds = self.value_preds[:-1].reshape(-1, *self.value_preds.shape[2:])319 returns = self.returns[:-1].reshape(-1, *self.returns.shape[2:])320 masks = self.masks[:-1].reshape(-1, *self.masks.shape[2:])321 active_masks = self.active_masks[:-1].reshape(-1, *self.active_masks.shape[2:])322 action_log_probs = self.action_log_probs.reshape(-1, *self.action_log_probs.shape[2:])323 gate_actions = self.gate_actions.reshape(-1, *self.gate_actions.shape[2:])324 advantages = advantages.reshape(-1, *advantages.shape[2:])325326 for indices in sampler:327 # obs size [T+1 N M Dim]-->[T N M Dim]-->[T*N*M,Dim]-->[index,Dim]328 share_obs_batch = (329 {k: v[indices] for k, v in share_obs.items()}330 if isinstance(share_obs, dict)331 else share_obs[indices]332 )333 obs_batch = {k: v[indices] for k, v in obs.items()}334335 rnn_states_batch = rnn_states[indices].reshape(-1, rnn_states.shape[-1])336 rnn_states_critic_batch = rnn_states_critic[indices].reshape(-1, rnn_states_critic.shape[-1])337 actions_batch = actions[indices].reshape(-1, actions.shape[-1])338 if self.available_actions is not None:339 available_actions_batch = available_actions[indices].reshape(-1, available_actions.shape[-1])340 else:341 available_actions_batch = None342 value_preds_batch = value_preds[indices].reshape(-1, value_preds.shape[-1])343 return_batch = returns[indices].reshape(-1, returns.shape[-1])344 masks_batch = masks[indices].reshape(-1, masks.shape[-1])345 active_masks_batch = active_masks[indices].reshape(-1, active_masks.shape[-1])346 old_action_log_probs_batch = action_log_probs[indices].reshape(-1, action_log_probs.shape[-1])347 gate_actions_batch = gate_actions[indices].reshape(-1, gate_actions.shape[-1])348 if advantages is None:349 adv_targ = None350 else:351 adv_targ = advantages[indices].reshape(-1, advantages.shape[-1])352353 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\354 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\355 adv_targ, available_actions_batch, gate_actions_batch356357 def naive_recurrent_generator(self, advantages, num_mini_batch):358 """359 Yield training data for non-chunked RNN training.360 :param advantages: (np.ndarray) advantage estimates.361 :param num_mini_batch: (int) number of minibatches to split the batch into.362 """363 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]364 batch_size = n_rollout_threads * num_agents365 assert n_rollout_threads * num_agents >= num_mini_batch, (366 "PPO requires the number of processes ({})* number of agents ({}) "367 "to be greater than or equal to the number of "368 "PPO mini batches ({}).".format(n_rollout_threads, num_agents, num_mini_batch))369 num_envs_per_batch = batch_size // num_mini_batch370 perm = torch.randperm(batch_size).numpy()371372 if isinstance(self.share_obs, dict):373 share_obs = {374 k: v.reshape(-1, batch_size, *v.shape[3:])375 for k, v in self.share_obs.items()376 }377 else:378 share_obs = self.share_obs.reshape(-1, batch_size, *self.share_obs.shape[3:])379 obs = {380 k: v.reshape(-1, batch_size, *v.shape[3:])381 for k, v in self.obs.items()382 }383 rnn_states = self.rnn_states.reshape(-1, batch_size, *self.rnn_states.shape[3:])384 rnn_states_critic = self.rnn_states_critic.reshape(-1, batch_size, *self.rnn_states_critic.shape[3:])385 actions = self.actions.reshape(-1, batch_size, self.actions.shape[-1])386 if self.available_actions is not None:387 available_actions = self.available_actions.reshape(-1, batch_size, self.available_actions.shape[-1])388 value_preds = self.value_preds.reshape(-1, batch_size, 1)389 returns = self.returns.reshape(-1, batch_size, 1)390 masks = self.masks.reshape(-1, batch_size, 1)391 active_masks = self.active_masks.reshape(-1, batch_size, 1)392 action_log_probs = self.action_log_probs.reshape(-1, batch_size, self.action_log_probs.shape[-1])393 gate_actions = self.gate_actions.reshape(-1, batch_size, 1)394 advantages = advantages.reshape(-1, batch_size, 1)395396 for start_ind in range(0, batch_size, num_envs_per_batch):397 share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []398 obs_batch = {k: [] for k in self.obs}399 rnn_states_batch = []400 rnn_states_critic_batch = []401 actions_batch = []402 available_actions_batch = []403 value_preds_batch = []404 return_batch = []405 masks_batch = []406 active_masks_batch = []407 old_action_log_probs_batch = []408 gate_actions_batch = []409 adv_targ = []410411 for offset in range(num_envs_per_batch):412 ind = perm[start_ind + offset]413 if isinstance(self.share_obs, dict):414 for k in share_obs_batch:415 share_obs_batch[k].append(share_obs[k][:-1, ind])416 else:417 share_obs_batch.append(share_obs[:-1, ind])418 for k in obs_batch:419 obs_batch[k].append(obs[k][:-1, ind])420 rnn_states_batch.append(rnn_states[0:1, ind])421 rnn_states_critic_batch.append(rnn_states_critic[0:1, ind])422 actions_batch.append(actions[:, ind])423 if self.available_actions is not None:424 available_actions_batch.append(available_actions[:-1, ind])425 value_preds_batch.append(value_preds[:-1, ind])426 return_batch.append(returns[:-1, ind])427 masks_batch.append(masks[:-1, ind])428 active_masks_batch.append(active_masks[:-1, ind])429 old_action_log_probs_batch.append(action_log_probs[:, ind])430 gate_actions_batch.append(gate_actions[:, ind])431 adv_targ.append(advantages[:, ind])432433 # [N[T, dim]]434 T, N = self.episode_length, num_envs_per_batch435 # These are all from_numpys of size (T, N, -1)436 if isinstance(share_obs_batch, dict):437 share_obs_batch = {k: np.stack(v, 1) for k, v in share_obs_batch.items()}438 else:439 share_obs_batch = np.stack(share_obs_batch, 1)440 obs_batch = {k: np.stack(v, 1) for k, v in obs_batch.items()}441 actions_batch = np.stack(actions_batch, 1)442 if self.available_actions is not None:443 available_actions_batch = np.stack(available_actions_batch, 1)444 value_preds_batch = np.stack(value_preds_batch, 1)445 return_batch = np.stack(return_batch, 1)446 masks_batch = np.stack(masks_batch, 1)447 active_masks_batch = np.stack(active_masks_batch, 1)448 old_action_log_probs_batch = np.stack(old_action_log_probs_batch, 1)449 gate_actions_batch = np.stack(gate_actions_batch, 1)450 adv_targ = np.stack(adv_targ, 1)451452 # States is just a (N, dim) from_numpy [N[1,dim]]453 rnn_states_batch = np.stack(rnn_states_batch).reshape(N, *self.rnn_states.shape[3:])454 rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(N, *self.rnn_states_critic.shape[3:])455456 # Flatten the (T, N, ...) from_numpys to (T * N, ...)457 if isinstance(share_obs_batch, dict):458 share_obs_batch = {k: _flatten(T, N, v) for k, v in share_obs_batch.items()}459 else:460 share_obs_batch = _flatten(T, N, share_obs_batch)461 obs_batch = {k: _flatten(T, N, v) for k, v in obs_batch.items()}462 actions_batch = _flatten(T, N, actions_batch)463 if self.available_actions is not None:464 available_actions_batch = _flatten(T, N, available_actions_batch)465 else:466 available_actions_batch = None467 value_preds_batch = _flatten(T, N, value_preds_batch)468 return_batch = _flatten(T, N, return_batch)469 masks_batch = _flatten(T, N, masks_batch)470 active_masks_batch = _flatten(T, N, active_masks_batch)471 old_action_log_probs_batch = _flatten(T, N, old_action_log_probs_batch)472 gate_actions_batch = _flatten(T, N, gate_actions_batch)473 adv_targ = _flatten(T, N, adv_targ)474475 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\476 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\477 adv_targ, available_actions_batch, gate_actions_batch478479 def recurrent_generator(self, advantages, num_mini_batch, data_chunk_length):480 """481 Yield training data for chunked RNN training.482 :param advantages: (np.ndarray) advantage estimates.483 :param num_mini_batch: (int) number of minibatches to split the batch into.484 :param data_chunk_length: (int) length of sequence chunks with which to train RNN.485 """486 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]487 chunks_per_env = episode_length // data_chunk_length488 data_chunks = n_rollout_threads * chunks_per_env489 mini_batch_size = data_chunks // num_mini_batch490491 assert mini_batch_size > 0, (492 "PPO requires rollout_threads ({}) * chunks_per_env ({}) to be "493 "greater than or equal to num_mini_batch ({}).".format(494 n_rollout_threads, chunks_per_env, num_mini_batch495 )496 )497498 rand = torch.randperm(data_chunks).numpy()499 sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]500501 for indices in sampler:502 share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []503 obs_batch = {k: [] for k in self.obs}504 rnn_states_batch = []505 rnn_states_critic_batch = []506 actions_batch = []507 available_actions_batch = []508 value_preds_batch = []509 return_batch = []510 masks_batch = []511 active_masks_batch = []512 old_action_log_probs_batch = []513 gate_actions_batch = []514 adv_targ = []515516 for index in indices:517 env_ind = index // chunks_per_env518 chunk_ind = index % chunks_per_env519 start = chunk_ind * data_chunk_length520 end = start + data_chunk_length521522 if isinstance(self.share_obs, dict):523 for k in share_obs_batch:524 share_obs_batch[k].append(self.share_obs[k][start:end, env_ind])525 else:526 share_obs_batch.append(self.share_obs[start:end, env_ind])527 for k in obs_batch:528 obs_batch[k].append(self.obs[k][start:end, env_ind])529 actions_batch.append(self.actions[start:end, env_ind])530 if self.available_actions is not None:531 available_actions_batch.append(self.available_actions[start:end, env_ind])532 value_preds_batch.append(self.value_preds[start:end, env_ind])533 return_batch.append(self.returns[start:end, env_ind])534 masks_batch.append(self.masks[start:end, env_ind])535 active_masks_batch.append(self.active_masks[start:end, env_ind])536 old_action_log_probs_batch.append(self.action_log_probs[start:end, env_ind])537 gate_actions_batch.append(self.gate_actions[start:end, env_ind])538 adv_targ.append(advantages[start:end, env_ind])539 rnn_states_batch.append(self.rnn_states[start, env_ind])540 rnn_states_critic_batch.append(self.rnn_states_critic[start, env_ind])541542 L, N = data_chunk_length, mini_batch_size543544 # Keep the agent axis grouped for graph encoders: [L, N, A, ...].545 if isinstance(share_obs_batch, dict):546 share_obs_batch = {k: np.stack(v, axis=1) for k, v in share_obs_batch.items()}547 else:548 share_obs_batch = np.stack(share_obs_batch, axis=1)549 obs_batch = {k: np.stack(v, axis=1) for k, v in obs_batch.items()}550551 actions_batch = np.stack(actions_batch, axis=1)552 if self.available_actions is not None:553 available_actions_batch = np.stack(available_actions_batch, axis=1)554 value_preds_batch = np.stack(value_preds_batch, axis=1)555 return_batch = np.stack(return_batch, axis=1)556 masks_batch = np.stack(masks_batch, axis=1)557 active_masks_batch = np.stack(active_masks_batch, axis=1)558 old_action_log_probs_batch = np.stack(old_action_log_probs_batch, axis=1)559 gate_actions_batch = np.stack(gate_actions_batch, axis=1)560 adv_targ = np.stack(adv_targ, axis=1)561562 # One initial RNN state per sampled sequence and per agent.563 rnn_states_batch = np.stack(rnn_states_batch).reshape(N * num_agents, *self.rnn_states.shape[3:])564 rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(565 N * num_agents, *self.rnn_states_critic.shape[3:]566 )567568 # obs/share_obs stay grouped as [L * N, A, ...]; PPO targets are flat [L * N * A, ...].569 if isinstance(share_obs_batch, dict):570 share_obs_batch = {k: _flatten(L, N, v) for k, v in share_obs_batch.items()}571 else:572 share_obs_batch = _flatten(L, N, share_obs_batch)573 obs_batch = {k: _flatten(L, N, v) for k, v in obs_batch.items()}574 actions_batch = _flatten(L, N, actions_batch).reshape(-1, actions_batch.shape[-1])575 if self.available_actions is not None:576 available_actions_batch = _flatten(L, N, available_actions_batch).reshape(577 -1, available_actions_batch.shape[-1]578 )579 else:580 available_actions_batch = None581 value_preds_batch = _flatten(L, N, value_preds_batch).reshape(-1, value_preds_batch.shape[-1])582 return_batch = _flatten(L, N, return_batch).reshape(-1, return_batch.shape[-1])583 masks_batch = _flatten(L, N, masks_batch).reshape(-1, masks_batch.shape[-1])584 active_masks_batch = _flatten(L, N, active_masks_batch).reshape(-1, active_masks_batch.shape[-1])585 old_action_log_probs_batch = _flatten(L, N, old_action_log_probs_batch).reshape(586 -1, old_action_log_probs_batch.shape[-1]587 )588 gate_actions_batch = _flatten(L, N, gate_actions_batch).reshape(-1, 1)589 adv_targ = _flatten(L, N, adv_targ).reshape(-1, adv_targ.shape[-1])590591 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\592 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\593 adv_targ, available_actions_batch, gate_actions_batch