ModelTrain/module/utils.py

import numpy as np, torch, os, h5py, fnmatch, cv2
from torch.utils.data import TensorDataset, DataLoader
import torchvision.transforms as transforms
import IPython
e = IPython.embed

def flatten_list(l):
    return [item for sublist in l for item in iter(sublist)]


class EpisodicDataset(torch.utils.data.Dataset):

    def __init__(self, dataset_path_list, camera_names, norm_stats, episode_ids, episode_len, chunk_size, policy_class, use_vitg=False, tactile_camera_names=None):
        super(EpisodicDataset).__init__()
        self.episode_ids = episode_ids
        self.dataset_path_list = dataset_path_list
        self.camera_names = camera_names
        self.tactile_camera_names = tactile_camera_names if tactile_camera_names else []
        self.norm_stats = norm_stats
        self.episode_len = episode_len
        self.chunk_size = chunk_size
        self.cumulative_len = np.cumsum(self.episode_len)
        self.max_episode_len = max(episode_len)
        self.policy_class = policy_class
        self.use_vitg = use_vitg
        if self.policy_class == "Diffusion":
            self.augment_images = True
        else:
            self.augment_images = False
        self.transformations = None
        self.__getitem__(0)
        self.is_sim = False

    def __len__(self):
        return len(self.episode_ids)

    def _locate_transition(self, index):
        assert index < self.cumulative_len[-1]
        episode_index = np.argmax(self.cumulative_len > index)
        start_ts = index - (self.cumulative_len[episode_index] - self.episode_len[episode_index])
        episode_id = self.episode_ids[episode_index]
        return (episode_id, start_ts)

    def __getitem__(self, index):
        episode_id, start_ts = self._locate_transition(index)
        dataset_path = self.dataset_path_list[episode_id]
        with h5py.File(dataset_path, "r") as root:
            try:
                is_sim = root.attrs["sim"]
            except:
                is_sim = False

            compressed = root.attrs.get("compress", False)
            if "/base_action" in root:
                base_action = root["/base_action"][()]
                base_action = preprocess_base_action(base_action)
                action = np.concatenate([root["/action"][()], base_action], axis=(-1))
            else:
                action = root["/action"][()]
                dummy_base_action = np.zeros([action.shape[0], 2])
                action = np.concatenate([action, dummy_base_action], axis=(-1))
            original_action_shape = action.shape
            episode_len = original_action_shape[0]
            qpos = root["/observations/qpos"][start_ts]
            qvel = root["/observations/qvel"][start_ts]
            image_dict = dict()
            for cam_name in self.camera_names:
                # Try /observations/images/{cam_name} first (RGB cameras)
                # Then try /observations/{cam_name} (tactile sensors)
                if f"/observations/images/{cam_name}" in root:
                    image_dict[cam_name] = root[f"/observations/images/{cam_name}"][start_ts]
                elif f"/observations/{cam_name}" in root:
                    image_dict[cam_name] = root[f"/observations/{cam_name}"][start_ts]
                else:
                    raise KeyError(f"Cannot find {cam_name} in /observations/images/ or /observations/")

            if compressed:
                for cam_name in image_dict.keys():
                    decompressed_image = cv2.imdecode(image_dict[cam_name], 1)
                    image_dict[cam_name] = np.array(decompressed_image)

            if is_sim:
                action = action[start_ts:]
                action_len = episode_len - start_ts
            else:
                action = action[max(0, start_ts - 1):]
                action_len = episode_len - max(0, start_ts - 1)
        padded_action = np.zeros((self.max_episode_len, original_action_shape[1]), dtype=(np.float32))
        padded_action[:action_len] = action
        is_pad = np.zeros(self.max_episode_len)
        is_pad[action_len:] = 1
        padded_action = padded_action[:self.chunk_size]
        is_pad = is_pad[:self.chunk_size]
        
        # Separate RGB cameras from tactile sensors
        rgb_cameras = [cam for cam in self.camera_names if cam not in self.tactile_camera_names]
        
        # Process RGB camera images (stack together)
        rgb_images = []
        for cam_name in rgb_cameras:
            img = image_dict[cam_name]  # (H, W, C)
            rgb_images.append(img)
        
        if rgb_images:
            # Stack RGB images and convert to tensor
            rgb_stacked = np.stack(rgb_images, axis=0)  # (num_rgb, H, W, C)
            rgb_tensor = torch.from_numpy(rgb_stacked)
            rgb_tensor = rgb_tensor.permute(0, 3, 1, 2)  # (num_rgb, C, H, W)
            rgb_tensor = rgb_tensor / 255.0  # Normalize
        else:
            rgb_tensor = None
        
        # Process tactile images (keep separate, resize for ViTG)
        tactile_images = []
        for cam_name in self.tactile_camera_names:
            img = image_dict[cam_name]  # (H, W, C)
            
            # Convert to tensor
            img_tensor = torch.from_numpy(img).float()
            img_tensor = img_tensor.permute(2, 0, 1)  # (C, H, W)
            img_tensor = img_tensor / 255.0  # Normalize
            
            # Resize to 224x224 for ViTG
            if self.use_vitg:
                resize_transform = transforms.Resize((224, 224), antialias=True)
                img_tensor = resize_transform(img_tensor)
            
            tactile_images.append(img_tensor)
        
        # Prepare image_data (RGB only, already stacked)
        if rgb_tensor is not None:
            image_data = rgb_tensor  # (num_rgb, C, H, W)
        else:
            # No RGB cameras, create empty tensor
            image_data = torch.empty(0, 3, 480, 640)
        
        qpos_data = torch.from_numpy(qpos).float()
        action_data = torch.from_numpy(padded_action).float()
        is_pad = torch.from_numpy(is_pad).bool()
        
        # Apply augmentation if needed (only to RGB, not tactile)
        if self.augment_images and image_data.shape[0] > 0:
            if self.transformations is None:
                print("Initializing transformations")
                self.transformations = [
                    transforms.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08)]
            
            # Apply to RGB images
            for transform in self.transformations:
                image_data = transform(image_data)

        if self.policy_class == "Diffusion":
            action_data = (action_data - self.norm_stats["action_min"]) / (self.norm_stats["action_max"] - self.norm_stats["action_min"]) * 2 - 1
        else:
            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"]

        return (image_data, tactile_images, qpos_data, action_data, is_pad)


def get_norm_stats(dataset_path_list):
    all_qpos_data = []
    all_action_data = []
    all_episode_len = []
    for dataset_path in dataset_path_list:
        try:
            with h5py.File(dataset_path, "r") as root:
                qpos = root["/observations/qpos"][()]
                qvel = root["/observations/qvel"][()]
                if "/base_action" in root:
                    base_action = root["/base_action"][()]
                    base_action = preprocess_base_action(base_action)
                    action = np.concatenate([root["/action"][()], base_action], axis=(-1))
                else:
                    action = root["/action"][()]
                    dummy_base_action = np.zeros([action.shape[0], 2])
                    action = np.concatenate([action, dummy_base_action], axis=(-1))
        except Exception as e:
            try:
                print(f"Error loading {dataset_path} in get_norm_stats")
                print(e)
                quit()
            finally:
                e = None
                del e

        else:
            all_qpos_data.append(torch.from_numpy(qpos))
            all_action_data.append(torch.from_numpy(action))
            all_episode_len.append(len(qpos))
    else:
        all_qpos_data = torch.cat(all_qpos_data, dim=0)
        all_action_data = torch.cat(all_action_data, dim=0)
        action_mean = all_action_data.mean(dim=[0]).float()
        action_std = all_action_data.std(dim=[0]).float()
        action_std = torch.clip(action_std, 0.01, np.inf)
        qpos_mean = all_qpos_data.mean(dim=[0]).float()
        qpos_std = all_qpos_data.std(dim=[0]).float()
        qpos_std = torch.clip(qpos_std, 0.01, np.inf)
        action_min = all_action_data.min(dim=0).values.float()
        action_max = all_action_data.max(dim=0).values.float()
        eps = 0.0001
        stats = {'action_mean':(action_mean.numpy)(),  'action_std':(action_std.numpy)(),  'action_min':(action_min.numpy()) - eps, 
         'action_max':(action_max.numpy()) + eps,  'qpos_mean':(qpos_mean.numpy)(), 
         'qpos_std':(qpos_std.numpy)(),  'example_qpos':qpos}
        return (
         stats, all_episode_len)


def find_all_hdf5(dataset_dir, skip_mirrored_data):
    hdf5_files = []
    for root, dirs, files in os.walk(dataset_dir):
        for filename in fnmatch.filter(files, "*.hdf5"):
            if "features" in filename:
                pass
            elif skip_mirrored_data and "mirror" in filename:
                pass
            else:
                hdf5_files.append(os.path.join(root, filename))
        else:
            print(f"Found {len(hdf5_files)} hdf5 files")
            return hdf5_files


def BatchSampler(batch_size, episode_len_l, sample_weights):
    sample_probs = np.array(sample_weights) / np.sum(sample_weights) if sample_weights is not None else None
    sum_dataset_len_l = np.cumsum([0] + [np.sum(episode_len) for episode_len in episode_len_l])
    while True:
        batch = []
        for _ in range(batch_size):
            episode_idx = np.random.choice((len(episode_len_l)), p=sample_probs)
            step_idx = np.random.randint(sum_dataset_len_l[episode_idx], sum_dataset_len_l[episode_idx + 1])
            batch.append(step_idx)
        else:
            yield batch


def load_data(dataset_dir_l, name_filter, camera_names, batch_size_train, batch_size_val, chunk_size, skip_mirrored_data=False, load_pretrain=False, policy_class=None, stats_dir_l=None, sample_weights=None, train_ratio=0.99, use_vitg=False, tactile_camera_names=None):
    if type(dataset_dir_l) == str:
        dataset_dir_l = [
         dataset_dir_l]
    dataset_path_list_list = [find_all_hdf5(dataset_dir, skip_mirrored_data) for dataset_dir in dataset_dir_l]
    num_episodes_0 = len(dataset_path_list_list[0])
    dataset_path_list = flatten_list(dataset_path_list_list)
    dataset_path_list = [n for n in dataset_path_list if name_filter(n)]
    num_episodes_l = [len(dataset_path_list) for dataset_path_list in dataset_path_list_list]
    num_episodes_cumsum = np.cumsum(num_episodes_l)
    shuffled_episode_ids_0 = np.random.permutation(num_episodes_0)
    train_episode_ids_0 = shuffled_episode_ids_0[:int(train_ratio * num_episodes_0)]
    val_episode_ids_0 = shuffled_episode_ids_0[int(train_ratio * num_episodes_0):]
    train_episode_ids_l = [train_episode_ids_0] + [np.arange(num_episodes) + num_episodes_cumsum[idx] for idx, num_episodes in enumerate(num_episodes_l[1:])]
    val_episode_ids_l = [val_episode_ids_0]
    train_episode_ids = np.concatenate(train_episode_ids_l)
    val_episode_ids = np.concatenate(val_episode_ids_l)
    print(f"\n\nData from: {dataset_dir_l}\n- Train on {[len(x) for x in train_episode_ids_l]} episodes\n- Test on {[len(x) for x in val_episode_ids_l]} episodes\n\n")
    _, all_episode_len = get_norm_stats(dataset_path_list)
    train_episode_len_l = [[all_episode_len[i] for i in train_episode_ids] for train_episode_ids in train_episode_ids_l]
    val_episode_len_l = [[all_episode_len[i] for i in val_episode_ids] for val_episode_ids in val_episode_ids_l]
    train_episode_len = flatten_list(train_episode_len_l)
    val_episode_len = flatten_list(val_episode_len_l)
    if stats_dir_l is None:
        stats_dir_l = dataset_dir_l
    else:
        if type(stats_dir_l) == str:
            stats_dir_l = [
             stats_dir_l]
    
    norm_stats, _ = get_norm_stats(flatten_list([find_all_hdf5(stats_dir, skip_mirrored_data) for stats_dir in stats_dir_l]))
    print(f"Norm stats from: {stats_dir_l}")
    if use_vitg:
        print("Dataset configured for ViTG: tactile images will be resized to 224x224")
    batch_sampler_train = BatchSampler(batch_size_train, train_episode_len_l, sample_weights)
    batch_sampler_val = BatchSampler(batch_size_val, val_episode_len_l, None)
    train_dataset = EpisodicDataset(dataset_path_list, camera_names, norm_stats, train_episode_ids, train_episode_len, chunk_size, policy_class, use_vitg=use_vitg, tactile_camera_names=tactile_camera_names)
    val_dataset = EpisodicDataset(dataset_path_list, camera_names, norm_stats, val_episode_ids, val_episode_len, chunk_size, policy_class, use_vitg=use_vitg, tactile_camera_names=tactile_camera_names)
    train_num_workers = (8 if os.getlogin() == "zfu" else 16) if train_dataset.augment_images else 2
    val_num_workers = 8 if train_dataset.augment_images else 2
    print(f"Augment images: {train_dataset.augment_images}, train_num_workers: {train_num_workers}, val_num_workers: {val_num_workers}")
    train_dataloader = DataLoader(train_dataset, batch_sampler=batch_sampler_train, pin_memory=True, num_workers=train_num_workers, prefetch_factor=2)
    val_dataloader = DataLoader(val_dataset, batch_sampler=batch_sampler_val, pin_memory=True, num_workers=val_num_workers, prefetch_factor=2)
    return (
     train_dataloader, val_dataloader, norm_stats, train_dataset.is_sim)


def calibrate_linear_vel(base_action, c=None):
    if c is None:
        c = 0.0
    v = base_action[(Ellipsis, 0)]
    w = base_action[(Ellipsis, 1)]
    base_action = base_action.copy()
    base_action[(Ellipsis, 0)] = v - c * w
    return base_action


def smooth_base_action(base_action):
    return np.stack([np.convolve(base_action[:, i], np.ones(5) / 5, mode="same") for i in range(base_action.shape[1])],
      axis=(-1)).astype(np.float32)


def preprocess_base_action(base_action):
    base_action = smooth_base_action(base_action)
    return base_action


def postprocess_base_action(base_action):
    linear_vel, angular_vel = base_action
    linear_vel *= 1.0
    angular_vel *= 1.0
    return np.array([linear_vel, angular_vel])


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():
    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])
    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)


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]
        else:
            result[k] = value_sum / num_items

    else:
        return result


def detach_dict(d):
    new_d = dict()
    for k, v in d.items():
        new_d[k] = v.detach()
    else:
        return new_d


def set_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)

# okay decompiling utils.pyc