def __getitem__(self, index):
file_index = index // self.traj_per_file
path = self.filenames[file_index]
try:
with h5py.File(path, 'r') as F:
ex_index = index % self.traj_per_file # get the index
key = 'traj{}'.format(ex_index)
# Fetch data into a dict
if key + '/images' in F.keys():
data_dict = AttrDict(images=(F[key + '/images'].value))
else:
data_dict = AttrDict()
for name in F[key].keys():
if name in ['states', 'actions', 'pad_mask']:
data_dict[name] = F[key + '/' + name].value.astype(np.float32)
# Make length consistent
end_ind = np.argmax(data_dict.pad_mask * np.arange(data_dict.pad_mask.shape[0], dtype=np.float32), 0)
start_ind = np.random.randint(0, end_ind - 1) if self.randomize_start else 0
start_ind, end_ind, data_dict = self.sample_max_len_video(data_dict, start_ind, end_ind)
# Randomize length
if self.randomize_length:
end_ind = self._randomize_length(start_ind, end_ind, data_dict)
# Collect data into the format the model expects
data_dict.end_ind = end_ind
data_dict.start_ind = start_ind
self.process_data_dict(data_dict)
except:
raise ValueError("Problem when loading file from {}".format(path))
return data_dict
def __getitem__(self, index):
if 'one_datum' in self.data_conf and self.data_conf.one_datum:
index = 1
file_index = index // self.traj_per_file
path = self.filenames[file_index]
try:
with h5py.File(path, 'r') as F:
ex_index = index % self.traj_per_file # get the index
key = 'traj{}'.format(ex_index)
# Fetch data into a dict
if key + '/images' in F.keys():
data_dict = AttrDict(images=(F[key + '/images'].value))
else:
data_dict = AttrDict()
for name in F[key].keys():
if name in ['states', 'actions', 'pad_mask']:
data_dict[name] = F[key + '/' + name].value.astype(
np.float32)
# remove spurious states at end of trajectory
if self.filter_repeated_tail:
data_dict = self._filter_tail(data_dict)
# maybe subsample seqs
if self.subsampler is not None:
data_dict = self._subsample_data(data_dict)
if 'robosuite_full_state' in F[key].keys():
data_dict.robosuite_full_state = F[
key + '/robosuite_full_state'].value
if 'regression_state' in F[key].keys():
data_dict.states = F[key +
'/regression_state'].value.astype(
np.float32)
# Make length consistent
end_ind = np.argmax(
data_dict.pad_mask *
np.arange(data_dict.pad_mask.shape[0], dtype=np.float32),
0)
start_ind = np.random.randint(0, end_ind -
1) if self.randomize_start else 0
start_ind, end_ind, data_dict = self.sample_max_len_video(
data_dict, start_ind, end_ind)
# Randomize length
if self.randomize_length:
end_ind = self._randomize_length(start_ind, end_ind,
data_dict)
# repeat last frame until end of sequence
data_dict.norep_end_ind = end_ind
if self.repeat_tail:
data_dict, end_ind = self._repeat_tail(data_dict, end_ind)
# Collect data into the format the model expects
data_dict.end_ind = end_ind
data_dict.start_ind = start_ind
# for roboturk env rendering
if 'robosuite_env_name' in F[key].keys():
data_dict.robosuite_env_name = F[
key + '/robosuite_env_name'].value
if 'robosuite_xml' in F[key].keys():
data_dict.robosuite_xml = F[key + '/robosuite_xml'].value
self.process_data_dict(data_dict)
except: # KeyError:
raise ValueError("Problem when loading file from {}".format(path))
return data_dict
