import numpy as np import torch import os import h5py from torch.utils.data import TensorDataset, DataLoader import IPython e = IPython.embed class EpisodicDataset(torch.utils.data.Dataset): def __init__(self, episode_ids, dataset_dir, camera_names, norm_stats, max_episode_len=None, stateless=False): super(EpisodicDataset).__init__() self.episode_ids = episode_ids self.dataset_dir = dataset_dir self.camera_names = camera_names self.norm_stats = norm_stats self.stateless = stateless # image-only ACT: zero proprio post-normalization # Pad all actions/is_pad to this fixed length so batch collate works # across variable-length episodes. Determined automatically by scanning # max episode length if not provided. if max_episode_len is None: max_episode_len = 0 for ep in episode_ids: p = os.path.join(dataset_dir, f'episode_{ep}.hdf5') with h5py.File(p, 'r') as rt: max_episode_len = max(max_episode_len, rt['/action'].shape[0]) self.max_episode_len = max_episode_len self.is_sim = None self.__getitem__(0) # initialize self.is_sim def __len__(self): return len(self.episode_ids) def __getitem__(self, index): sample_full_episode = False # hardcode episode_id = self.episode_ids[index] 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'] original_action_shape = root['/action'].shape episode_len = original_action_shape[0] if sample_full_episode: start_ts = 0 else: start_ts = np.random.choice(episode_len) # get observation at start_ts only qpos = root['/observations/qpos'][start_ts] qvel = root['/observations/qvel'][start_ts] image_dict = dict() 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: action = root['/action'][start_ts:] action_len = episode_len - start_ts else: action = root['/action'][max(0, start_ts - 1):] # hack, to make timesteps more aligned action_len = episode_len - max(0, start_ts - 1) # hack, to make timesteps more aligned self.is_sim = is_sim # Pad to max_episode_len (constant across episodes) so batch collate works. pad_to = self.max_episode_len padded_action = np.zeros((pad_to, original_action_shape[1]), dtype=np.float32) n_copy = min(action_len, pad_to) padded_action[:n_copy] = action[:n_copy] is_pad = np.ones(pad_to, dtype=bool) is_pad[:n_copy] = False # new axis for different cameras all_cam_images = [] 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) # construct observations image_data = torch.from_numpy(all_cam_images) qpos_data = torch.from_numpy(qpos).float() action_data = torch.from_numpy(padded_action).float() is_pad = torch.from_numpy(is_pad).bool() # channel last image_data = torch.einsum('k h w c -> k c h w', image_data) # normalize image and change dtype to float image_data = image_data / 255.0 action_data = (action_data - self.norm_stats["action_mean"]) / self.norm_stats["action_std"] qpos_data = (qpos_data - self.norm_stats["qpos_mean"]) / self.norm_stats["qpos_std"] if self.stateless: qpos_data = torch.zeros_like(qpos_data) # image-only: policy relies purely on RGB return image_data, qpos_data, action_data, is_pad def get_norm_stats(dataset_dir, num_episodes): # Variable-length episodes — concatenate along time instead of stacking. all_qpos_data = [] all_action_data = [] for episode_idx in range(num_episodes): dataset_path = os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5') with h5py.File(dataset_path, 'r') as root: qpos = root['/observations/qpos'][()] qvel = root['/observations/qvel'][()] action = root['/action'][()] all_qpos_data.append(torch.from_numpy(qpos)) all_action_data.append(torch.from_numpy(action)) all_qpos_data = torch.cat(all_qpos_data, dim=0) # (sum_T, qpos_dim) all_action_data = torch.cat(all_action_data, dim=0) # (sum_T, action_dim) # normalize action data — reduce over the (sum_T,) axis action_mean = all_action_data.mean(dim=0, keepdim=True) action_std = all_action_data.std(dim=0, keepdim=True) action_std = torch.clip(action_std, 1e-2, np.inf) # clipping # normalize qpos data qpos_mean = all_qpos_data.mean(dim=0, keepdim=True) qpos_std = all_qpos_data.std(dim=0, keepdim=True) qpos_std = torch.clip(qpos_std, 1e-2, np.inf) # clipping stats = {"action_mean": action_mean.numpy().squeeze(), "action_std": action_std.numpy().squeeze(), "qpos_mean": qpos_mean.numpy().squeeze(), "qpos_std": qpos_std.numpy().squeeze(), "example_qpos": qpos} return stats def load_data(dataset_dir, num_episodes, camera_names, batch_size_train, batch_size_val, stateless=False): print(f'\nData from: {dataset_dir}\n') # obtain train test split train_ratio = float(os.environ.get('TRAIN_RATIO', '0.8')) # env-var override (e.g. 0.95); default 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):] # obtain normalization stats for qpos and action norm_stats = get_norm_stats(dataset_dir, num_episodes) # construct dataset and dataloader train_dataset = EpisodicDataset(train_indices, dataset_dir, camera_names, norm_stats, stateless=stateless) val_dataset = EpisodicDataset(val_indices, dataset_dir, camera_names, norm_stats, stateless=stateless) 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) return train_dataloader, val_dataloader, norm_stats, train_dataset.is_sim ### env utils def sample_box_pose(): x_range = [0.0, 0.2] y_range = [0.4, 0.6] z_range = [0.05, 0.05] ranges = np.vstack([x_range, y_range, z_range]) cube_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) cube_quat = np.array([1, 0, 0, 0]) return np.concatenate([cube_position, cube_quat]) def sample_insertion_pose(): # Peg x_range = [0.1, 0.2] y_range = [0.4, 0.6] z_range = [0.05, 0.05] ranges = np.vstack([x_range, y_range, z_range]) peg_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) peg_quat = np.array([1, 0, 0, 0]) peg_pose = np.concatenate([peg_position, peg_quat]) # Socket x_range = [-0.2, -0.1] y_range = [0.4, 0.6] z_range = [0.05, 0.05] ranges = np.vstack([x_range, y_range, z_range]) socket_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) socket_quat = np.array([1, 0, 0, 0]) socket_pose = np.concatenate([socket_position, socket_quat]) return peg_pose, socket_pose ### helper functions def compute_dict_mean(epoch_dicts): result = {k: None for k in epoch_dicts[0]} num_items = len(epoch_dicts) for k in result: value_sum = 0 for epoch_dict in epoch_dicts: value_sum += epoch_dict[k] result[k] = value_sum / num_items return result def detach_dict(d): new_d = dict() for k, v in d.items(): new_d[k] = v.detach() return new_d def set_seed(seed): torch.manual_seed(seed) np.random.seed(seed)