from typing import Dict, List import torch import numpy as np import h5py from tqdm import tqdm import zarr import os import shutil import copy import json import hashlib import traceback import cv2 from filelock import FileLock from threadpoolctl import threadpool_limits import concurrent.futures import multiprocessing from omegaconf import OmegaConf from diffusion_policy.common.pytorch_util import dict_apply from diffusion_policy.dataset.base_dataset import BaseImageDataset, LinearNormalizer from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer from diffusion_policy.model.common.rotation_transformer import RotationTransformer from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs, Jpeg2k from diffusion_policy.common.replay_buffer import ReplayBuffer from diffusion_policy.common.sampler import SequenceSampler, get_val_mask from diffusion_policy.common.normalize_util import ( robomimic_abs_action_only_normalizer_from_stat, robomimic_abs_action_only_dual_arm_normalizer_from_stat, get_range_normalizer_from_stat, get_image_range_normalizer, get_identity_normalizer_from_stat, array_to_stats ) register_codecs() class AlohaReplayImageDataset(BaseImageDataset): def __init__(self, shape_meta: dict, dataset_path: str, horizon=1, pad_before=0, pad_after=0, n_obs_steps=None, use_legacy_normalizer=False, use_cache=False, seed=42, val_ratio=0.0 ): replay_buffer = None if use_cache: cache_zarr_path = dataset_path + '.zarr.zip' cache_lock_path = cache_zarr_path + '.lock' print('Acquiring lock on cache.') with FileLock(cache_lock_path): if not os.path.exists(cache_zarr_path): # cache does not exists try: print('Cache does not exist. Creating!') # store = zarr.DirectoryStore(cache_zarr_path) replay_buffer = _convert_robomimic_to_replay( store=zarr.MemoryStore(), shape_meta=shape_meta, dataset_path=dataset_path, ) print('Saving cache to disk.') with zarr.ZipStore(cache_zarr_path) as zip_store: replay_buffer.save_to_store( store=zip_store ) except Exception as e: shutil.rmtree(cache_zarr_path) raise e else: print('Loading cached ReplayBuffer from Disk.') print('cache_zarr_path ', cache_zarr_path) with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store: replay_buffer = ReplayBuffer.copy_from_store( src_store=zip_store, store=zarr.MemoryStore()) print('Loaded!') else: replay_buffer = _convert_robomimic_to_replay( store=zarr.MemoryStore(), shape_meta=shape_meta, dataset_path=dataset_path, ) rgb_keys = list() lowdim_keys = list() obs_shape_meta = shape_meta['obs'] for key, attr in obs_shape_meta.items(): type = attr.get('type', 'low_dim') if type == 'rgb': rgb_keys.append(key) elif type == 'low_dim': lowdim_keys.append(key) # for key in rgb_keys: # replay_buffer[key].compressor.numthreads=1 key_first_k = dict() if n_obs_steps is not None: # only take first k obs from images for key in rgb_keys + lowdim_keys: key_first_k[key] = n_obs_steps val_mask = get_val_mask( n_episodes=replay_buffer.n_episodes, val_ratio=val_ratio, seed=seed) train_mask = ~val_mask sampler = SequenceSampler( replay_buffer=replay_buffer, sequence_length=horizon, pad_before=pad_before, pad_after=pad_after, episode_mask=train_mask, key_first_k=key_first_k) self.replay_buffer = replay_buffer self.sampler = sampler self.shape_meta = shape_meta self.rgb_keys = rgb_keys self.lowdim_keys = lowdim_keys self.n_obs_steps = n_obs_steps self.train_mask = train_mask self.horizon = horizon self.pad_before = pad_before self.pad_after = pad_after self.use_legacy_normalizer = use_legacy_normalizer print('episode ends ', replay_buffer.episode_ends[:]) # print('agentview_image', replay_buffer['agentview_image'].shape) # print('robot0_eef_pos ', replay_buffer['robot0_eef_pos'].shape) print('action ', replay_buffer['action'].shape) # print('abs_action ', replay_buffer['abs_action'].shape) def get_validation_dataset(self): val_set = copy.copy(self) val_set.sampler = SequenceSampler( replay_buffer=self.replay_buffer, sequence_length=self.horizon, pad_before=self.pad_before, pad_after=self.pad_after, episode_mask=~self.train_mask ) val_set.train_mask = ~self.train_mask return val_set def get_normalizer(self, **kwargs) -> LinearNormalizer: normalizer = LinearNormalizer() stat = array_to_stats(self.replay_buffer['action']) # already normalized this_normalizer = get_identity_normalizer_from_stat(stat) normalizer['action'] = this_normalizer # obs for key in self.lowdim_keys: stat = array_to_stats(self.replay_buffer[key]) if key.endswith('states'): this_normalizer = get_range_normalizer_from_stat(stat) else: raise RuntimeError('unsupported') normalizer[key] = this_normalizer # image for key in self.rgb_keys: normalizer[key] = get_image_range_normalizer() return normalizer def get_all_actions(self) -> torch.Tensor: return torch.from_numpy(self.replay_buffer['action']) def __len__(self): return len(self.sampler) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: threadpool_limits(1) data = self.sampler.sample_sequence(idx) # to save RAM, only return first n_obs_steps of OBS # since the rest will be discarded anyway. # when self.n_obs_steps is None # this slice does nothing (takes all) T_slice = slice(self.n_obs_steps) obs_dict = dict() for key in self.rgb_keys: # move channel last to channel first # T,H,W,C # convert uint8 image to float32 obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1 ).astype(np.float32) / 255. # T,C,H,W del data[key] for key in self.lowdim_keys: obs_dict[key] = data[key][T_slice].astype(np.float32) del data[key] torch_data = { 'obs': dict_apply(obs_dict, torch.from_numpy), 'action': torch.from_numpy(data['action'].astype(np.float32)) } return torch_data def undo_transform_action(action, rotation_transformer): raw_shape = action.shape if raw_shape[-1] == 20: # dual arm action = action.reshape(-1,2,10) d_rot = action.shape[-1] - 4 pos = action[...,:3] rot = action[...,3:3+d_rot] gripper = action[...,[-1]] rot = rotation_transformer.inverse(rot) uaction = np.concatenate([ pos, rot, gripper ], axis=-1) if raw_shape[-1] == 20: # dual arm uaction = uaction.reshape(*raw_shape[:-1], 14) return uaction def _convert_robomimic_to_replay(store, shape_meta, dataset_path, n_workers=None, max_inflight_tasks=None): if n_workers is None: n_workers = multiprocessing.cpu_count() if max_inflight_tasks is None: max_inflight_tasks = n_workers * 5 # parse shape_meta rgb_keys = list() lowdim_keys = list() # construct compressors and chunks obs_shape_meta = shape_meta['obs'] for key, attr in obs_shape_meta.items(): shape = attr['shape'] type = attr.get('type', 'low_dim') if type == 'rgb': rgb_keys.append(key) elif type == 'low_dim': lowdim_keys.append(key) root = zarr.group(store) data_group = root.require_group('data', overwrite=True) meta_group = root.require_group('meta', overwrite=True) # breakpoint() with h5py.File(dataset_path) as file: # count total steps demos = file['data'] episode_ends = list() prev_end = 0 for i in range(len(demos)): demo = demos[f'demo_{i}'] episode_length = demo['actions'].shape[0] episode_end = prev_end + episode_length prev_end = episode_end episode_ends.append(episode_end) n_steps = episode_ends[-1] episode_starts = [0] + episode_ends[:-1] _ = meta_group.array('episode_ends', episode_ends, dtype=np.int64, compressor=None, overwrite=True) # save lowdim data extra_keys = ['action'] for key in tqdm(lowdim_keys + extra_keys, desc="Loading lowdim data"): data_key = 'obs/' + key if key == 'action': data_key = 'actions' elif key == 'rewards': data_key = 'rewards' this_data = list() for i in range(len(demos)): demo = demos[f'demo_{i}'] this_data.append(demo[data_key][:].astype(np.float32)) this_data = np.concatenate(this_data, axis=0) if key == 'rewards': this_data = this_data[:, None] if key == 'action': assert this_data.shape == (n_steps,) + tuple(shape_meta[key]['shape']) else: print(f"Key: {key}, Shape: {this_data.shape}, Expected: {(n_steps,) + tuple(shape_meta['obs'][key]['shape'])}") assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape']) _ = data_group.array( name=key, data=this_data, shape=this_data.shape, chunks=this_data.shape, compressor=None, dtype=this_data.dtype ) def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx, h, w): # try: img = hdf5_arr[hdf5_idx] # (480, 640, 3) img = cv2.resize(img, (w, h),interpolation=cv2.INTER_AREA) # if h == 640: # print(f"h: {h}, w: {w}") # breakpoint() zarr_arr[zarr_idx] = img # make sure we can successfully decode _ = zarr_arr[zarr_idx] return True # except Exception: # traceback.print_exc() # raise with tqdm(total=n_steps*len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar: # one chunk per thread, therefore no synchronization needed with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor: futures = set() for key in rgb_keys: data_key = 'obs/' + key shape = tuple(shape_meta['obs'][key]['shape']) c,h,w = shape # this_compressor = Jpeg2k(level=20)s this_compressor = None img_arr = data_group.require_dataset( name=key, shape=(n_steps,h,w,c), chunks=(1,h,w,c), compressor=this_compressor, dtype=np.uint8 ) for episode_idx in range(len(demos)): demo = demos[f'demo_{episode_idx}'] hdf5_arr = demo['obs'][key] for hdf5_idx in range(hdf5_arr.shape[0]): if len(futures) >= max_inflight_tasks: # limit number of inflight tasks completed, futures = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED) for f in completed: if not f.result(): raise RuntimeError('Failed to encode image!') pbar.update(len(completed)) zarr_idx = episode_starts[episode_idx] + hdf5_idx futures.add( executor.submit(img_copy, img_arr, zarr_idx, hdf5_arr, hdf5_idx, h,w)) completed, futures = concurrent.futures.wait(futures) for f in completed: if not f.result(): raise RuntimeError('Failed to encode image!') pbar.update(len(completed)) replay_buffer = ReplayBuffer(root) return replay_buffer def normalizer_from_stat(stat): max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max()) scale = np.full_like(stat['max'], fill_value=1/max_abs) offset = np.zeros_like(stat['max']) return SingleFieldLinearNormalizer.create_manual( scale=scale, offset=offset, input_stats_dict=stat )