in ss_baselines/savi/models/rollout_storage.py [0:0]
def recurrent_generator(self, advantages, num_mini_batch):
num_processes = self.rewards.size(1)
assert num_processes >= num_mini_batch, (
"Trainer requires the number of processes ({}) "
"to be greater than or equal to the number of "
"trainer mini batches ({}).".format(num_processes, num_mini_batch)
)
num_envs_per_batch = num_processes // num_mini_batch
perm = torch.randperm(num_processes)
for start_ind in range(0, num_processes, num_envs_per_batch):
observations_batch = defaultdict(list)
recurrent_hidden_states_batch = []
actions_batch = []
prev_actions_batch = []
value_preds_batch = []
return_batch = []
masks_batch = []
old_action_log_probs_batch = []
adv_targ = []
if self.use_external_memory:
em_store_batch = []
em_masks_batch = []
else:
em_store_batch = None
em_masks_batch = None
for offset in range(num_envs_per_batch):
ind = perm[start_ind + offset]
for sensor in self.observations:
observations_batch[sensor].append(
self.observations[sensor][: self.step, ind]
)
recurrent_hidden_states_batch.append(
self.recurrent_hidden_states[0, :, ind]
)
actions_batch.append(self.actions[: self.step, ind])
prev_actions_batch.append(self.prev_actions[: self.step, ind])
value_preds_batch.append(self.value_preds[: self.step, ind])
return_batch.append(self.returns[: self.step, ind])
masks_batch.append(self.masks[: self.step, ind])
old_action_log_probs_batch.append(
self.action_log_probs[: self.step, ind]
)
adv_targ.append(advantages[: self.step, ind])
if self.use_external_memory:
em_store_batch.append(self.em.memory[:, : self.step, ind])
em_masks_batch.append(self.em_masks[: self.step, ind])
T, N = self.step, num_envs_per_batch
# These are all tensors of size (T, N, -1)
for sensor in observations_batch:
observations_batch[sensor] = torch.stack(
observations_batch[sensor], 1
)
actions_batch = torch.stack(actions_batch, 1)
prev_actions_batch = torch.stack(prev_actions_batch, 1)
value_preds_batch = torch.stack(value_preds_batch, 1)
return_batch = torch.stack(return_batch, 1)
masks_batch = torch.stack(masks_batch, 1)
old_action_log_probs_batch = torch.stack(
old_action_log_probs_batch, 1
)
adv_targ = torch.stack(adv_targ, 1)
if self.use_external_memory:
# This is a (em_size, num_steps, bs, em_dim) tensor
em_store_batch = torch.stack(em_store_batch, 2)
# This is a (num_steps, bs, em_size) tensor
em_masks_batch = torch.stack(em_masks_batch, 1)
# States is just a (num_recurrent_layers, N, -1) tensor
recurrent_hidden_states_batch = torch.stack(
recurrent_hidden_states_batch, 1
)
# Flatten the (T, N, ...) tensors to (T * N, ...)
for sensor in observations_batch:
observations_batch[sensor] = self._flatten_helper(
T, N, observations_batch[sensor]
)
actions_batch = self._flatten_helper(T, N, actions_batch)
prev_actions_batch = self._flatten_helper(T, N, prev_actions_batch)
value_preds_batch = self._flatten_helper(T, N, value_preds_batch)
return_batch = self._flatten_helper(T, N, return_batch)
masks_batch = self._flatten_helper(T, N, masks_batch)
old_action_log_probs_batch = self._flatten_helper(
T, N, old_action_log_probs_batch
)
adv_targ = self._flatten_helper(T, N, adv_targ)
if self.use_external_memory:
em_store_batch = em_store_batch.view(-1, T * N, self.em_dim)
em_masks_batch = self._flatten_helper(T, N, em_masks_batch)
yield (
observations_batch,
recurrent_hidden_states_batch,
actions_batch,
prev_actions_batch,
value_preds_batch,
return_batch,
masks_batch,
old_action_log_probs_batch,
adv_targ,
em_store_batch,
em_masks_batch,
)