JC6123/aloha
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add task configs to constant.py to reduce command line arguments

Browse Source
This commit is contained in:
Tony Zhao
2023-03-05 16:52:47 -08:00
parent 092735ddb9
commit 5a33ee8db0
11 changed files with 131 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
utils.py
View File

@@ -3,16 +3,16 @@ import torch
import os
import h5py
from torch.utils.data import TensorDataset, DataLoader
from constants import SIM_CAMERA_NAMES, CAMERA_NAMES
import IPython
e = IPython.embed
class EpisodicDataset(torch.utils.data.Dataset):
def __init__(self, episode_ids, dataset_dir, norm_stats):
def __init__(self, episode_ids, dataset_dir, camera_names, norm_stats):
super(EpisodicDataset).__init__()
self.episode_ids = episode_ids
self.dataset_dir = dataset_dir
self.camera_names = camera_names
self.norm_stats = norm_stats
self.is_sim = None
self.__getitem__(0) # initialize self.is_sim
@@ -27,10 +27,6 @@ class EpisodicDataset(torch.utils.data.Dataset):
dataset_path = os.path.join(self.dataset_dir, f'episode_{episode_id}.hdf5')
with h5py.File(dataset_path, 'r') as root:
is_sim = root.attrs['sim']
if is_sim:
camera_names = SIM_CAMERA_NAMES
else:
camera_names = CAMERA_NAMES
original_action_shape = root['/action'].shape
episode_len = original_action_shape[0]
if sample_full_episode:
@@ -41,7 +37,7 @@ class EpisodicDataset(torch.utils.data.Dataset):
qpos = root['/observations/qpos'][start_ts]
qvel = root['/observations/qvel'][start_ts]
image_dict = dict()
for cam_name in camera_names:
for cam_name in self.camera_names:
image_dict[cam_name] = root[f'/observations/images/{cam_name}'][start_ts]
# get all actions after and including start_ts
if is_sim:
@@ -59,7 +55,7 @@ class EpisodicDataset(torch.utils.data.Dataset):
# new axis for different cameras
all_cam_images = []
for cam_name in camera_names:
for cam_name in self.camera_names:
all_cam_images.append(image_dict[cam_name])
all_cam_images = np.stack(all_cam_images, axis=0)
@@ -112,9 +108,9 @@ def get_norm_stats(dataset_dir, num_episodes):
return stats
def load_data(dataset_dir, num_episodes, batch_size_train, batch_size_val):
def load_data(dataset_dir, num_episodes, camera_names, batch_size_train, batch_size_val):
# obtain train test split
train_ratio = 0.8 # TODO
train_ratio = 0.8
shuffled_indices = np.random.permutation(num_episodes)
train_indices = shuffled_indices[:int(train_ratio * num_episodes)]
val_indices = shuffled_indices[int(train_ratio * num_episodes):]
@@ -123,8 +119,8 @@ def load_data(dataset_dir, num_episodes, batch_size_train, batch_size_val):
norm_stats = get_norm_stats(dataset_dir, num_episodes)
# construct dataset and dataloader
train_dataset = EpisodicDataset(train_indices, dataset_dir, norm_stats)
val_dataset = EpisodicDataset(val_indices, dataset_dir, norm_stats)
train_dataset = EpisodicDataset(train_indices, dataset_dir, camera_names, norm_stats)
val_dataset = EpisodicDataset(val_indices, dataset_dir, camera_names, norm_stats)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size_train, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size_val, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)