shared_buffer.py
MAPPO 对照 buffer
原文件 · 588 行 · 31775 字节
utils/shared_buffer.py
SHA-256: c9cf6f5ed3efae47653e39a506f87531a2aaabacd7df6299be71793c9d77e430
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 SharedReplayBuffer(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 self.rewards = np.zeros(102 (self.episode_length, self.n_rollout_threads, num_agents, 1), dtype=np.float32)103104 self.masks = np.ones((self.episode_length + 1, self.n_rollout_threads, num_agents, 1), dtype=np.float32)105 self.bad_masks = np.ones_like(self.masks)106 self.active_masks = np.ones_like(self.masks)107108 self.step = 0109110 def insert(self, share_obs, obs, rnn_states_actor, rnn_states_critic, actions, action_log_probs,111 value_preds, rewards, masks, agent_caps=None, task_reqs=None, bad_masks=None, active_masks=None, available_actions=None):112 """113 Insert data into the buffer.114 :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.115 :param obs: (np.ndarray) local agent observations.116 :param rnn_states_actor: (np.ndarray) RNN states for actor network.117 :param rnn_states_critic: (np.ndarray) RNN states for critic network.118 :param actions:(np.ndarray) actions taken by agents.119 :param action_log_probs:(np.ndarray) log probs of actions taken by agents120 :param value_preds: (np.ndarray) value function prediction at each step.121 :param rewards: (np.ndarray) reward collected at each step.122 :param masks: (np.ndarray) denotes whether the environment has terminated or not.123 :param agent_caps: (np.ndarray) agent capabilities (unused in base buffer, used by GAT buffer).124 :param task_reqs: (np.ndarray) task requirements (unused in base buffer, used by GAT buffer).125 :param bad_masks: (np.ndarray) action space for agents.126 :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.127 :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.128 """129 if isinstance(self.share_obs, dict):130 for k in self.share_obs:131 self.share_obs[k][self.step + 1] = share_obs[k].copy()132 else:133 self.share_obs[self.step + 1] = share_obs.copy()134 for k in self.obs:135 self.obs[k][self.step + 1] = obs[k].copy()136 self.rnn_states[self.step + 1] = rnn_states_actor.copy()137 self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()138 self.actions[self.step] = actions.copy()139 self.action_log_probs[self.step] = action_log_probs.copy()140 self.value_preds[self.step] = value_preds.copy()141 self.rewards[self.step] = rewards.copy()142 self.masks[self.step + 1] = masks.copy()143 if bad_masks is not None:144 self.bad_masks[self.step + 1] = bad_masks.copy()145 if active_masks is not None:146 self.active_masks[self.step + 1] = active_masks.copy()147 if available_actions is not None:148 self.available_actions[self.step + 1] = available_actions.copy()149150 self.step = (self.step + 1) % self.episode_length151152 def chooseinsert(self, share_obs, obs, rnn_states, rnn_states_critic, actions, action_log_probs,153 value_preds, rewards, masks, bad_masks=None, active_masks=None, available_actions=None):154 """155 Insert data into the buffer. This insert function is used specifically for Hanabi, which is turn based.156 :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.157 :param obs: (np.ndarray) local agent observations.158 :param rnn_states_actor: (np.ndarray) RNN states for actor network.159 :param rnn_states_critic: (np.ndarray) RNN states for critic network.160 :param actions:(np.ndarray) actions taken by agents.161 :param action_log_probs:(np.ndarray) log probs of actions taken by agents162 :param value_preds: (np.ndarray) value function prediction at each step.163 :param rewards: (np.ndarray) reward collected at each step.164 :param masks: (np.ndarray) denotes whether the environment has terminated or not.165 :param bad_masks: (np.ndarray) denotes indicate whether whether true terminal state or due to episode limit166 :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.167 :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.168 """169 if isinstance(self.share_obs, dict):170 for k in self.share_obs:171 self.share_obs[k][self.step] = share_obs[k].copy()172 else:173 self.share_obs[self.step] = share_obs.copy()174 for k in self.obs:175 self.obs[k][self.step] = obs[k].copy()176 self.rnn_states[self.step + 1] = rnn_states.copy()177 self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()178 self.actions[self.step] = actions.copy()179 self.action_log_probs[self.step] = action_log_probs.copy()180 self.value_preds[self.step] = value_preds.copy()181 self.rewards[self.step] = rewards.copy()182 self.masks[self.step + 1] = masks.copy()183 if bad_masks is not None:184 self.bad_masks[self.step + 1] = bad_masks.copy()185 if active_masks is not None:186 self.active_masks[self.step] = active_masks.copy()187 if available_actions is not None:188 self.available_actions[self.step] = available_actions.copy()189190 self.step = (self.step + 1) % self.episode_length191192 def after_update(self):193 """Copy last timestep data to first index. Called after update to model."""194 if isinstance(self.share_obs, dict):195 for k in self.share_obs:196 self.share_obs[k][0] = self.share_obs[k][-1].copy()197 else:198 self.share_obs[0] = self.share_obs[-1].copy()199 for k in self.obs:200 self.obs[k][0] = self.obs[k][-1].copy()201 self.rnn_states[0] = self.rnn_states[-1].copy()202 self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()203 self.masks[0] = self.masks[-1].copy()204 self.bad_masks[0] = self.bad_masks[-1].copy()205 self.active_masks[0] = self.active_masks[-1].copy()206 if self.available_actions is not None:207 self.available_actions[0] = self.available_actions[-1].copy()208209 def chooseafter_update(self):210 """Copy last timestep data to first index. This method is used for Hanabi."""211 self.rnn_states[0] = self.rnn_states[-1].copy()212 self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()213 self.masks[0] = self.masks[-1].copy()214 self.bad_masks[0] = self.bad_masks[-1].copy()215216 def compute_returns(self, next_value, value_normalizer=None):217 """218 Compute returns either as discounted sum of rewards, or using GAE.219 :param next_value: (np.ndarray) value predictions for the step after the last episode step.220 :param value_normalizer: (PopArt) If not None, PopArt value normalizer instance.221 """222 if self._use_proper_time_limits:223 if self._use_gae:224 self.value_preds[-1] = next_value225 gae = 0226 for step in reversed(range(self.rewards.shape[0])):227 if self._use_popart or self._use_valuenorm:228 # step + 1229 delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(230 self.value_preds[step + 1]) * self.masks[step + 1] \231 - value_normalizer.denormalize(self.value_preds[step])232 gae = delta + self.gamma * self.gae_lambda * gae * self.masks[step + 1]233 gae = gae * self.bad_masks[step + 1]234 self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])235 else:236 delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \237 self.value_preds[step]238 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae239 gae = gae * self.bad_masks[step + 1]240 self.returns[step] = gae + self.value_preds[step]241 else:242 self.returns[-1] = next_value243 for step in reversed(range(self.rewards.shape[0])):244 if self._use_popart or self._use_valuenorm:245 self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[246 step]) * self.bad_masks[step + 1] \247 + (1 - self.bad_masks[step + 1]) * value_normalizer.denormalize(248 self.value_preds[step])249 else: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]) * self.value_preds[step]253 else:254 if self._use_gae:255 self.value_preds[-1] = next_value256 gae = 0257 for step in reversed(range(self.rewards.shape[0])):258 if self._use_popart or self._use_valuenorm:259 delta = self.rewards[step] + self.gamma * value_normalizer.denormalize(260 self.value_preds[step + 1]) * self.masks[step + 1] \261 - value_normalizer.denormalize(self.value_preds[step])262 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae263 self.returns[step] = gae + value_normalizer.denormalize(self.value_preds[step])264 else:265 delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \266 self.value_preds[step]267 gae = delta + self.gamma * self.gae_lambda * self.masks[step + 1] * gae268 self.returns[step] = gae + self.value_preds[step]269 else:270 self.returns[-1] = next_value271 for step in reversed(range(self.rewards.shape[0])):272 self.returns[step] = self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[step]273274 def feed_forward_generator(self, advantages, num_mini_batch=None, mini_batch_size=None):275 """276 Yield training data for MLP policies.277 :param advantages: (np.ndarray) advantage estimates.278 :param num_mini_batch: (int) number of minibatches to split the batch into.279 :param mini_batch_size: (int) number of samples in each minibatch.280 """281 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]282 batch_size = n_rollout_threads * episode_length * num_agents283284 if mini_batch_size is None:285 assert batch_size >= num_mini_batch, (286 "PPO requires the number of processes ({}) "287 "* number of steps ({}) * number of agents ({}) = {} "288 "to be greater than or equal to the number of PPO mini batches ({})."289 "".format(n_rollout_threads, episode_length, num_agents,290 n_rollout_threads * episode_length * num_agents,291 num_mini_batch))292 mini_batch_size = batch_size // num_mini_batch293294 rand = torch.randperm(batch_size).numpy()295 sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]296297 if isinstance(self.share_obs, dict):298 share_obs = {299 k: v[:-1].reshape(-1, *v.shape[3:])300 for k, v in self.share_obs.items()301 }302 else:303 share_obs = self.share_obs[:-1].reshape(-1, *self.share_obs.shape[3:])304 obs = {305 k: v[:-1].reshape(-1, *v.shape[3:])306 for k, v in self.obs.items()307 }308 rnn_states = self.rnn_states[:-1].reshape(-1, *self.rnn_states.shape[3:])309 rnn_states_critic = self.rnn_states_critic[:-1].reshape(-1, *self.rnn_states_critic.shape[3:])310 actions = self.actions.reshape(-1, self.actions.shape[-1])311 if self.available_actions is not None:312 available_actions = self.available_actions[:-1].reshape(-1, self.available_actions.shape[-1])313 value_preds = self.value_preds[:-1].reshape(-1, 1)314 returns = self.returns[:-1].reshape(-1, 1)315 masks = self.masks[:-1].reshape(-1, 1)316 active_masks = self.active_masks[:-1].reshape(-1, 1)317 action_log_probs = self.action_log_probs.reshape(-1, self.action_log_probs.shape[-1])318 advantages = advantages.reshape(-1, 1)319320 for indices in sampler:321 # obs size [T+1 N M Dim]-->[T N M Dim]-->[T*N*M,Dim]-->[index,Dim]322 share_obs_batch = (323 {k: v[indices] for k, v in share_obs.items()}324 if isinstance(share_obs, dict)325 else share_obs[indices]326 )327 obs_batch = {k: v[indices] for k, v in obs.items()}328 rnn_states_batch = rnn_states[indices]329 rnn_states_critic_batch = rnn_states_critic[indices]330 actions_batch = actions[indices]331 if self.available_actions is not None:332 available_actions_batch = available_actions[indices]333 else:334 available_actions_batch = None335 value_preds_batch = value_preds[indices]336 return_batch = returns[indices]337 masks_batch = masks[indices]338 active_masks_batch = active_masks[indices]339 old_action_log_probs_batch = action_log_probs[indices]340 if advantages is None:341 adv_targ = None342 else:343 adv_targ = advantages[indices]344345 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\346 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\347 adv_targ, available_actions_batch348349 def naive_recurrent_generator(self, advantages, num_mini_batch):350 """351 Yield training data for non-chunked RNN training.352 :param advantages: (np.ndarray) advantage estimates.353 :param num_mini_batch: (int) number of minibatches to split the batch into.354 """355 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]356 batch_size = n_rollout_threads * num_agents357 assert n_rollout_threads * num_agents >= num_mini_batch, (358 "PPO requires the number of processes ({})* number of agents ({}) "359 "to be greater than or equal to the number of "360 "PPO mini batches ({}).".format(n_rollout_threads, num_agents, num_mini_batch))361 num_envs_per_batch = batch_size // num_mini_batch362 perm = torch.randperm(batch_size).numpy()363364 if isinstance(self.share_obs, dict):365 share_obs = {366 k: v.reshape(-1, batch_size, *v.shape[3:])367 for k, v in self.share_obs.items()368 }369 else:370 share_obs = self.share_obs.reshape(-1, batch_size, *self.share_obs.shape[3:])371 obs = {372 k: v.reshape(-1, batch_size, *v.shape[3:])373 for k, v in self.obs.items()374 }375 rnn_states = self.rnn_states.reshape(-1, batch_size, *self.rnn_states.shape[3:])376 rnn_states_critic = self.rnn_states_critic.reshape(-1, batch_size, *self.rnn_states_critic.shape[3:])377 actions = self.actions.reshape(-1, batch_size, self.actions.shape[-1])378 if self.available_actions is not None:379 available_actions = self.available_actions.reshape(-1, batch_size, self.available_actions.shape[-1])380 value_preds = self.value_preds.reshape(-1, batch_size, 1)381 returns = self.returns.reshape(-1, batch_size, 1)382 masks = self.masks.reshape(-1, batch_size, 1)383 active_masks = self.active_masks.reshape(-1, batch_size, 1)384 action_log_probs = self.action_log_probs.reshape(-1, batch_size, self.action_log_probs.shape[-1])385 advantages = advantages.reshape(-1, batch_size, 1)386387 for start_ind in range(0, batch_size, num_envs_per_batch):388 share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []389 obs_batch = {k: [] for k in self.obs}390 rnn_states_batch = []391 rnn_states_critic_batch = []392 actions_batch = []393 available_actions_batch = []394 value_preds_batch = []395 return_batch = []396 masks_batch = []397 active_masks_batch = []398 old_action_log_probs_batch = []399 adv_targ = []400401 for offset in range(num_envs_per_batch):402 ind = perm[start_ind + offset]403 if isinstance(share_obs, dict):404 for k in share_obs_batch:405 share_obs_batch[k].append(share_obs[k][:-1, ind])406 else:407 share_obs_batch.append(share_obs[:-1, ind])408 for k in obs_batch:409 obs_batch[k].append(obs[k][:-1, ind])410 rnn_states_batch.append(rnn_states[0:1, ind])411 rnn_states_critic_batch.append(rnn_states_critic[0:1, ind])412 actions_batch.append(actions[:, ind])413 if self.available_actions is not None:414 available_actions_batch.append(available_actions[:-1, ind])415 value_preds_batch.append(value_preds[:-1, ind])416 return_batch.append(returns[:-1, ind])417 masks_batch.append(masks[:-1, ind])418 active_masks_batch.append(active_masks[:-1, ind])419 old_action_log_probs_batch.append(action_log_probs[:, ind])420 adv_targ.append(advantages[:, ind])421422 # [N[T, dim]]423 T, N = self.episode_length, num_envs_per_batch424 # These are all from_numpys of size (T, N, -1)425 if isinstance(share_obs_batch, dict):426 share_obs_batch = {k: np.stack(v, 1) for k, v in share_obs_batch.items()}427 else:428 share_obs_batch = np.stack(share_obs_batch, 1)429 obs_batch = {k: np.stack(v, 1) for k, v in obs_batch.items()}430 actions_batch = np.stack(actions_batch, 1)431 if self.available_actions is not None:432 available_actions_batch = np.stack(available_actions_batch, 1)433 value_preds_batch = np.stack(value_preds_batch, 1)434 return_batch = np.stack(return_batch, 1)435 masks_batch = np.stack(masks_batch, 1)436 active_masks_batch = np.stack(active_masks_batch, 1)437 old_action_log_probs_batch = np.stack(old_action_log_probs_batch, 1)438 adv_targ = np.stack(adv_targ, 1)439440 # States is just a (N, dim) from_numpy [N[1,dim]]441 rnn_states_batch = np.stack(rnn_states_batch).reshape(N, *self.rnn_states.shape[3:])442 rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(N, *self.rnn_states_critic.shape[3:])443444 # Flatten the (T, N, ...) from_numpys to (T * N, ...)445 if isinstance(share_obs_batch, dict):446 share_obs_batch = {k: _flatten(T, N, v) for k, v in share_obs_batch.items()}447 else:448 share_obs_batch = _flatten(T, N, share_obs_batch)449 obs_batch = {k: _flatten(T, N, v) for k, v in obs_batch.items()}450 actions_batch = _flatten(T, N, actions_batch)451 if self.available_actions is not None:452 available_actions_batch = _flatten(T, N, available_actions_batch)453 else:454 available_actions_batch = None455 value_preds_batch = _flatten(T, N, value_preds_batch)456 return_batch = _flatten(T, N, return_batch)457 masks_batch = _flatten(T, N, masks_batch)458 active_masks_batch = _flatten(T, N, active_masks_batch)459 old_action_log_probs_batch = _flatten(T, N, old_action_log_probs_batch)460 adv_targ = _flatten(T, N, adv_targ)461462 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\463 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\464 adv_targ, available_actions_batch465466 def recurrent_generator(self, advantages, num_mini_batch, data_chunk_length):467 """468 Yield training data for chunked RNN training.469 :param advantages: (np.ndarray) advantage estimates.470 :param num_mini_batch: (int) number of minibatches to split the batch into.471 :param data_chunk_length: (int) length of sequence chunks with which to train RNN.472 """473 episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]474 batch_size = n_rollout_threads * episode_length * num_agents475 data_chunks = batch_size // data_chunk_length # [C=r*T*M/L]476 mini_batch_size = data_chunks // num_mini_batch477478 rand = torch.randperm(data_chunks).numpy()479 sampler = [rand[i * mini_batch_size:(i + 1) * mini_batch_size] for i in range(num_mini_batch)]480481 if isinstance(self.share_obs, dict):482 share_obs = {483 k: v[:-1].transpose(1, 2, 0, *range(3, v.ndim)).reshape(-1, *v.shape[3:])484 for k, v in self.share_obs.items()485 }486 else:487 share_obs = _cast(self.share_obs[:-1])488 obs = {k: _cast(v[:-1]) for k, v in self.obs.items()}489490 actions = _cast(self.actions)491 action_log_probs = _cast(self.action_log_probs)492 advantages = _cast(advantages)493 value_preds = _cast(self.value_preds[:-1])494 returns = _cast(self.returns[:-1])495 masks = _cast(self.masks[:-1])496 active_masks = _cast(self.active_masks[:-1])497 # rnn_states = _cast(self.rnn_states[:-1])498 # rnn_states_critic = _cast(self.rnn_states_critic[:-1])499 rnn_states = self.rnn_states[:-1].transpose(1, 2, 0, 3, 4).reshape(-1, *self.rnn_states.shape[3:])500 rnn_states_critic = self.rnn_states_critic[:-1].transpose(1, 2, 0, 3, 4).reshape(-1,501 *self.rnn_states_critic.shape[502 3:])503504 if self.available_actions is not None:505 available_actions = _cast(self.available_actions[:-1])506507 for indices in sampler:508 share_obs_batch = {k: [] for k in self.share_obs} if isinstance(self.share_obs, dict) else []509 obs_batch = {k: [] for k in self.obs}510 rnn_states_batch = []511 rnn_states_critic_batch = []512 actions_batch = []513 available_actions_batch = []514 value_preds_batch = []515 return_batch = []516 masks_batch = []517 active_masks_batch = []518 old_action_log_probs_batch = []519 adv_targ = []520521 for index in indices:522523 ind = index * data_chunk_length524 # size [T+1 N M Dim]-->[T N M Dim]-->[N,M,T,Dim]-->[N*M*T,Dim]-->[L,Dim]525 if isinstance(share_obs, dict):526 for k in share_obs_batch:527 share_obs_batch[k].append(share_obs[k][ind:ind + data_chunk_length])528 else:529 share_obs_batch.append(share_obs[ind:ind + data_chunk_length])530 for k in obs_batch:531 obs_batch[k].append(obs[k][ind:ind + data_chunk_length])532 actions_batch.append(actions[ind:ind + data_chunk_length])533 if self.available_actions is not None:534 available_actions_batch.append(available_actions[ind:ind + data_chunk_length])535 value_preds_batch.append(value_preds[ind:ind + data_chunk_length])536 return_batch.append(returns[ind:ind + data_chunk_length])537 masks_batch.append(masks[ind:ind + data_chunk_length])538 active_masks_batch.append(active_masks[ind:ind + data_chunk_length])539 old_action_log_probs_batch.append(action_log_probs[ind:ind + data_chunk_length])540 adv_targ.append(advantages[ind:ind + data_chunk_length])541 # size [T+1 N M Dim]-->[T N M Dim]-->[N M T Dim]-->[N*M*T,Dim]-->[1,Dim]542 rnn_states_batch.append(rnn_states[ind])543 rnn_states_critic_batch.append(rnn_states_critic[ind])544545 L, N = data_chunk_length, mini_batch_size546547 # These are all from_numpys of size (L, N, Dim) 548 if isinstance(share_obs_batch, dict):549 share_obs_batch = {k: np.stack(v, axis=1) for k, v in share_obs_batch.items()}550 else:551 share_obs_batch = np.stack(share_obs_batch, axis=1)552 obs_batch = {k: np.stack(v, axis=1) for k, v in obs_batch.items()}553554 actions_batch = np.stack(actions_batch, axis=1)555 if self.available_actions is not None:556 available_actions_batch = np.stack(available_actions_batch, axis=1)557 value_preds_batch = np.stack(value_preds_batch, axis=1)558 return_batch = np.stack(return_batch, axis=1)559 masks_batch = np.stack(masks_batch, axis=1)560 active_masks_batch = np.stack(active_masks_batch, axis=1)561 old_action_log_probs_batch = np.stack(old_action_log_probs_batch, axis=1)562 adv_targ = np.stack(adv_targ, axis=1)563564 # States is just a (N, -1) from_numpy565 rnn_states_batch = np.stack(rnn_states_batch).reshape(N, *self.rnn_states.shape[3:])566 rnn_states_critic_batch = np.stack(rnn_states_critic_batch).reshape(N, *self.rnn_states_critic.shape[3:])567568 # Flatten the (L, N, ...) from_numpys to (L * N, ...)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)575 if self.available_actions is not None:576 available_actions_batch = _flatten(L, N, available_actions_batch)577 else:578 available_actions_batch = None579 value_preds_batch = _flatten(L, N, value_preds_batch)580 return_batch = _flatten(L, N, return_batch)581 masks_batch = _flatten(L, N, masks_batch)582 active_masks_batch = _flatten(L, N, active_masks_batch)583 old_action_log_probs_batch = _flatten(L, N, old_action_log_probs_batch)584 adv_targ = _flatten(L, N, adv_targ)585586 yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\587 value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\588 adv_targ, available_actions_batch