File size: 1,269 Bytes
987ed1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import torch
def preprocess_obs_batch(obs_dict):
B, OH, OD = obs_dict['obs'].shape
obs_dict_batch = dict()
for key in obs_dict.keys():
obs_dict_batch[key] = obs_dict[key].unsqueeze(1).repeat(1, num_sample, 1, 1).reshape(B * num_sample, OH, OD)
return obs_dict_batch
def postprocess_action_batch(action_dict_batch):
AH, PH, AD = action_dict_batch['action'].shape[1], action_dict_batch['action_pred'].shape[1], action_dict_batch['action_pred'].shape[2]
action_dict_batch['action'] = action_dict_batch['action'].reshape(B, num_sample, AH, AD)
action_dict_batch['action_pred'] = action_dict_batch['action_pred'].reshape(B, num_sample, PH, AD)
if 'action_obs_pred' in action_dict_batch:
action_dict_batch['action_obs_pred'] = action_dict_batch['action_obs_pred'].reshape(B, num_sample, AH, OD)
if 'obs_pred' in action_dict_batch:
action_dict_batch['obs_pred'] = action_dict_batch['obs_pred'].reshape(B, num_sample, PH, OD)
return action_dict_batch
def slice_action_batch(action_dict_batch):
action_dict = dict()
range_tensor = torch.arange(B, device=index.device)
for key in action_dict_batch.keys():
action_dict[key] = action_dict_batch[key][range_tensor, index]
return action_dict
|