Spaces:
Runtime error
Runtime error
| """Image dataset.""" | |
| import os | |
| import pickle | |
| import warnings | |
| import numpy as np | |
| from torch.utils.data import Dataset | |
| from cliport import tasks | |
| from cliport.tasks import cameras | |
| from cliport.utils import utils | |
| import traceback | |
| # See transporter.py, regression.py, dummy.py, task.py, etc. | |
| PIXEL_SIZE = 0.003125 | |
| CAMERA_CONFIG = cameras.RealSenseD415.CONFIG | |
| BOUNDS = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]]) | |
| # Names as strings, REVERSE-sorted so longer (more specific) names are first. | |
| TASK_NAMES = (tasks.names).keys() | |
| TASK_NAMES = sorted(TASK_NAMES)[::-1] | |
| class RavensDataset(Dataset): | |
| """A simple image dataset class.""" | |
| def __init__(self, path, cfg, n_demos=0, augment=False): | |
| """A simple RGB-D image dataset.""" | |
| self._path = path | |
| self.cfg = cfg | |
| self.sample_set = [] | |
| self.max_seed = -1 | |
| self.n_episodes = 0 | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = self.cfg['dataset']['cache'] | |
| self.n_demos = n_demos | |
| self.augment = augment | |
| self.aug_theta_sigma = self.cfg['dataset']['augment']['theta_sigma'] if 'augment' in self.cfg['dataset'] else 60 # legacy code issue: theta_sigma was newly added | |
| self.pix_size = 0.003125 | |
| self.in_shape = (320, 160, 6) | |
| self.cam_config = cameras.RealSenseD415.CONFIG | |
| self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]]) | |
| # Track existing dataset if it exists. | |
| color_path = os.path.join(self._path, 'action') | |
| if os.path.exists(color_path): | |
| for fname in sorted(os.listdir(color_path)): | |
| if '.pkl' in fname: | |
| seed = int(fname[(fname.find('-') + 1):-4]) | |
| self.n_episodes += 1 | |
| self.max_seed = max(self.max_seed, seed) | |
| self._cache = {} | |
| if self.n_demos > 0: | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = self.cfg['dataset']['cache'] | |
| # Check if there sufficient demos in the dataset | |
| if self.n_demos > self.n_episodes: | |
| # raise Exception(f"Requested training on {self.n_demos} demos, but only {self.n_episodes} demos exist in the dataset path: {self._path}.") | |
| print(f"Requested training on {self.n_demos} demos, but only {self.n_episodes} demos exist in the dataset path: {self._path}.") | |
| self.n_demos = self.n_episodes | |
| episodes = np.random.choice(range(self.n_episodes), self.n_demos, False) | |
| self.set(episodes) | |
| def add(self, seed, episode): | |
| """Add an episode to the dataset. | |
| Args: | |
| seed: random seed used to initialize the episode. | |
| episode: list of (obs, act, reward, info) tuples. | |
| """ | |
| color, depth, action, reward, info = [], [], [], [], [] | |
| for obs, act, r, i in episode: | |
| color.append(obs['color']) | |
| depth.append(obs['depth']) | |
| action.append(act) | |
| reward.append(r) | |
| info.append(i) | |
| color = np.uint8(color) | |
| depth = np.float32(depth) | |
| def dump(data, field): | |
| field_path = os.path.join(self._path, field) | |
| if not os.path.exists(field_path): | |
| os.makedirs(field_path) | |
| fname = f'{self.n_episodes:06d}-{seed}.pkl' # -{len(episode):06d} | |
| with open(os.path.join(field_path, fname), 'wb') as f: | |
| pickle.dump(data, f) | |
| dump(color, 'color') | |
| dump(depth, 'depth') | |
| dump(action, 'action') | |
| dump(reward, 'reward') | |
| dump(info, 'info') | |
| self.n_episodes += 1 | |
| self.max_seed = max(self.max_seed, seed) | |
| def set(self, episodes): | |
| """Limit random samples to specific fixed set.""" | |
| self.sample_set = episodes | |
| def load(self, episode_id, images=True, cache=False): | |
| # TODO(lirui): consider loading into memory | |
| def load_field(episode_id, field, fname): | |
| # Check if sample is in cache. | |
| if cache: | |
| if episode_id in self._cache: | |
| if field in self._cache[episode_id]: | |
| return self._cache[episode_id][field] | |
| else: | |
| self._cache[episode_id] = {} | |
| # Load sample from files. | |
| path = os.path.join(self._path, field) | |
| data = pickle.load(open(os.path.join(path, fname), 'rb')) | |
| if cache: | |
| self._cache[episode_id][field] = data | |
| return data | |
| # Get filename and random seed used to initialize episode. | |
| seed = None | |
| path = os.path.join(self._path, 'action') | |
| for fname in sorted(os.listdir(path)): | |
| if f'{episode_id:06d}' in fname: | |
| seed = int(fname[(fname.find('-') + 1):-4]) | |
| # Load data. | |
| color = load_field(episode_id, 'color', fname) | |
| depth = load_field(episode_id, 'depth', fname) | |
| action = load_field(episode_id, 'action', fname) | |
| reward = load_field(episode_id, 'reward', fname) | |
| info = load_field(episode_id, 'info', fname) | |
| # Reconstruct episode. | |
| episode = [] | |
| for i in range(len(action)): | |
| obs = {'color': color[i], 'depth': depth[i]} if images else {} | |
| episode.append((obs, action[i], reward[i], info[i])) | |
| return episode, seed | |
| print(f'{episode_id:06d} not in ', path) | |
| def get_image(self, obs, cam_config=None): | |
| """Stack color and height images image.""" | |
| # if self.use_goal_image: | |
| # colormap_g, heightmap_g = utils.get_fused_heightmap(goal, configs) | |
| # goal_image = self.concatenate_c_h(colormap_g, heightmap_g) | |
| # input_image = np.concatenate((input_image, goal_image), axis=2) | |
| # assert input_image.shape[2] == 12, input_image.shape | |
| if cam_config is None: | |
| cam_config = self.cam_config | |
| # Get color and height maps from RGB-D images. | |
| cmap, hmap = utils.get_fused_heightmap( | |
| obs, cam_config, self.bounds, self.pix_size) | |
| img = np.concatenate((cmap, | |
| hmap[Ellipsis, None], | |
| hmap[Ellipsis, None], | |
| hmap[Ellipsis, None]), axis=2) | |
| assert img.shape == self.in_shape, img.shape | |
| return img | |
| def process_sample(self, datum, augment=True): | |
| # Get training labels from data sample. | |
| (obs, act, _, info) = datum | |
| img = self.get_image(obs) | |
| # p0, p1 = None, None | |
| # p0_theta, p1_theta = None, None | |
| # perturb_params = None | |
| p0, p1 = np.zeros(1), np.zeros(1) | |
| p0_theta, p1_theta = np.zeros(1), np.zeros(1) | |
| perturb_params = np.zeros(5) | |
| if act: | |
| p0_xyz, p0_xyzw = act['pose0'] | |
| p1_xyz, p1_xyzw = act['pose1'] | |
| p0 = utils.xyz_to_pix(p0_xyz, self.bounds, self.pix_size) | |
| p0_theta = -np.float32(utils.quatXYZW_to_eulerXYZ(p0_xyzw)[2]) | |
| p1 = utils.xyz_to_pix(p1_xyz, self.bounds, self.pix_size) | |
| p1_theta = -np.float32(utils.quatXYZW_to_eulerXYZ(p1_xyzw)[2]) | |
| p1_theta = p1_theta - p0_theta | |
| p0_theta = 0 | |
| # Data augmentation. | |
| if augment: | |
| img, _, (p0, p1), perturb_params = utils.perturb(img, [p0, p1], theta_sigma=self.aug_theta_sigma) | |
| # print("sample", p0,p1,p0_theta,p1_theta,perturb_params) | |
| sample = { | |
| 'img': img.copy(), | |
| 'p0': np.array(p0).copy(), 'p0_theta': np.array(p0_theta).copy(), | |
| 'p1': np.array(p1).copy(), 'p1_theta': np.array(p1_theta).copy() , | |
| 'perturb_params': np.array(perturb_params).copy() | |
| } | |
| # Add language goal if available. | |
| if 'lang_goal' not in info: | |
| warnings.warn("No language goal. Defaulting to 'task completed.'") | |
| if info and 'lang_goal' in info: | |
| sample['lang_goal'] = info['lang_goal'] | |
| else: | |
| sample['lang_goal'] = "task completed." | |
| return sample | |
| def process_goal(self, goal, perturb_params): | |
| # Get goal sample. | |
| (obs, act, _, info) = goal | |
| img = self.get_image(obs) | |
| # p0, p1 = None, None | |
| # p0_theta, p1_theta = None, None | |
| p0, p1 = np.zeros(1), np.zeros(1) | |
| p0_theta, p1_theta = np.zeros(1), np.zeros(1) | |
| # Data augmentation with specific params. | |
| # try: | |
| if perturb_params is not None and len(perturb_params) > 1: | |
| img = utils.apply_perturbation(img, perturb_params) | |
| sample = { | |
| 'img': img.copy(), | |
| 'p0': p0 , 'p0_theta': np.array(p0_theta).copy(), | |
| 'p1': p1, 'p1_theta': np.array(p1_theta).copy(), | |
| 'perturb_params': np.array(perturb_params).copy() | |
| } | |
| # Add language goal if available. | |
| if 'lang_goal' not in info: | |
| warnings.warn("No language goal. Defaulting to 'task completed.'") | |
| # print("goal",p0,p1,p0_theta,p1_theta,perturb_params) | |
| if info and 'lang_goal' in info: | |
| sample['lang_goal'] = info['lang_goal'] | |
| else: | |
| sample['lang_goal'] = "task completed." | |
| return sample | |
| def __len__(self): | |
| return len(self.sample_set) | |
| def __getitem__(self, idx): | |
| # Choose random episode. | |
| # if len(self.sample_set) > 0: | |
| # episode_id = np.random.choice(self.sample_set) | |
| # else: | |
| # episode_id = np.random.choice(range(self.n_episodes)) | |
| episode_id = self.sample_set[idx] | |
| res = self.load(episode_id, self.images, self.cache) | |
| if res is None: | |
| print("in get item", episode_id, self._path) | |
| print("load sample return None. Reload") | |
| print("Exception:", str(traceback.format_exc())) | |
| return self[0] # | |
| episode, _ = res | |
| # Is the task sequential like stack-block-pyramid-seq? | |
| is_sequential_task = '-seq' in self._path.split("/")[-1] | |
| # Return random observation action pair (and goal) from episode. | |
| i = np.random.choice(range(len(episode)-1)) | |
| g = i+1 if is_sequential_task else -1 | |
| sample, goal = episode[i], episode[g] | |
| # Process sample. | |
| sample = self.process_sample(sample, augment=self.augment) | |
| goal = self.process_goal(goal, perturb_params=sample['perturb_params']) | |
| return sample, goal | |
| class RavensMultiTaskDataset(RavensDataset): | |
| def __init__(self, path, cfg, group='multi-all', | |
| mode='train', n_demos=100, augment=False): | |
| """A multi-task dataset.""" | |
| self.root_path = path | |
| self.mode = mode | |
| if group not in self.MULTI_TASKS: | |
| # generate the groups on the fly | |
| self.tasks = list(set(group)) # .split(" ") | |
| else: | |
| self.tasks = self.MULTI_TASKS[group][mode] | |
| print("self.tasks:", self.tasks) | |
| self.attr_train_task = self.MULTI_TASKS[group]['attr_train_task'] if group in self.MULTI_TASKS and 'attr_train_task' in self.MULTI_TASKS[group] else None | |
| self.cfg = cfg | |
| self.sample_set = {} | |
| self.max_seed = -1 | |
| self.n_episodes = 0 | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = self.cfg['dataset']['cache'] | |
| self.n_demos = n_demos | |
| self.augment = augment | |
| self.aug_theta_sigma = self.cfg['dataset']['augment']['theta_sigma'] if 'augment' in self.cfg['dataset'] else 60 # legacy code issue: theta_sigma was newly added | |
| self.pix_size = 0.003125 | |
| self.in_shape = (320, 160, 6) | |
| self.cam_config = cameras.RealSenseD415.CONFIG | |
| self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]]) | |
| self.n_episodes = {} | |
| episodes = {} | |
| for task in self.tasks: | |
| task_path = os.path.join(self.root_path, f'{task}-{mode}') | |
| action_path = os.path.join(task_path, 'action') | |
| n_episodes = 0 | |
| if os.path.exists(action_path): | |
| for fname in sorted(os.listdir(action_path)): | |
| if '.pkl' in fname: | |
| n_episodes += 1 | |
| self.n_episodes[task] = n_episodes | |
| if n_episodes == 0: | |
| raise Exception(f"{task}-{mode} has 0 episodes. Remove it from the list in dataset.py") | |
| # Select random episode depending on the size of the dataset. | |
| episodes[task] = np.random.choice(range(n_episodes), min(self.n_demos, n_episodes), False) | |
| if self.n_demos > 0: | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = False # TODO(mohit): fix caching for multi-task dataset | |
| self.set(episodes) | |
| self._path = None | |
| self._task = None | |
| def __len__(self): | |
| # Average number of episodes across all tasks | |
| total_episodes = 0 | |
| for _, episode_ids in self.sample_set.items(): | |
| total_episodes += len(episode_ids) | |
| avg_episodes = total_episodes # // len(self.sample_set) | |
| return avg_episodes | |
| def __getitem__(self, idx): | |
| # Choose random task. | |
| self._task = self.tasks[idx % len(self.tasks)] # np.random.choice(self.tasks) | |
| self._path = os.path.join(self.root_path, f'{self._task}') | |
| # Choose random episode. | |
| if len(self.sample_set[self._task]) > 0: | |
| episode_id = np.random.choice(self.sample_set[self._task]) | |
| else: | |
| episode_id = np.random.choice(range(self.n_episodes[self._task])) | |
| res = self.load(episode_id, self.images, self.cache) | |
| if res is None: | |
| print("failed in get item", episode_id, self._task, self._path) | |
| print("Exception:", str(traceback.format_exc())) | |
| return self[np.random.randint(len(self))] # | |
| episode, _ = res | |
| # Is the task sequential like stack-block-pyramid-seq? | |
| is_sequential_task = '-seq' in self._path.split("/")[-1] | |
| # Return observation action pair (and goal) from episode. | |
| if len(episode) > 1: | |
| i = np.random.choice(range(len(episode)-1)) | |
| g = i+1 if is_sequential_task else -1 | |
| sample, goal = episode[i], episode[g] | |
| else: | |
| sample, goal = episode[0], episode[0] | |
| # Process sample | |
| sample = self.process_sample(sample, augment=self.augment) | |
| goal = self.process_goal(goal, perturb_params=sample['perturb_params']) | |
| return sample, goal | |
| def add(self, seed, episode): | |
| raise Exception("Adding tasks not supported with multi-task dataset") | |
| def load(self, episode_id, images=True, cache=False): | |
| # if self.attr_train_task is None or self.mode in ['val', 'test']: | |
| # self._task = np.random.choice(self.tasks) | |
| # else: | |
| # all_other_tasks = list(self.tasks) | |
| # all_other_tasks.remove(self.attr_train_task) | |
| # all_tasks = [self.attr_train_task] + all_other_tasks # add seen task in the front | |
| # # 50% chance of sampling the main seen task and 50% chance of sampling any other seen-unseen task | |
| # mult_attr_seen_sample_prob = 0.5 | |
| # sampling_probs = [(1-mult_attr_seen_sample_prob) / (len(all_tasks)-1)] * len(all_tasks) | |
| # sampling_probs[0] = mult_attr_seen_sample_prob | |
| # self._task = np.random.choice(all_tasks, p=sampling_probs) | |
| self._path = os.path.join(self.root_path, f'{self._task}-{self.mode}') | |
| return super().load(episode_id, images, cache) | |
| def get_curr_task(self): | |
| return self._task | |
| MULTI_TASKS = { | |
| # new expeeriments | |
| 'multi-gpt-test': { | |
| 'train': ['align-box-corner', 'rainbow-stack'], | |
| 'val': ['align-box-corner', 'rainbow-stack'], | |
| 'test': ['align-box-corner', 'rainbow-stack'] | |
| }, | |
| # all tasks | |
| 'multi-all': { | |
| 'train': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| 'align-rope', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'val': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| 'align-rope', | |
| 'assembling-kits-seq-seen-colors', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'packing-boxes-pairs-seen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-seen-colors', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-seen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-seen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-seen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| 'align-rope', | |
| 'assembling-kits-seq-seen-colors', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'packing-boxes-pairs-seen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-seen-colors', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-seen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-seen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-seen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| }, | |
| # demo-conditioned tasks | |
| 'multi-demo-conditioned': { | |
| 'train': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| ], | |
| 'val': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| ], | |
| 'test': [ | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'block-insertion', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| ], | |
| }, | |
| # goal-conditioned tasks | |
| 'multi-language-conditioned': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-unseen-colors', # unseen here refers to training only seen splits to be consitent with single-task setting | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'val': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-seen-colors', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'packing-boxes-pairs-seen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-seen-colors', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-seen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-seen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-seen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-seen-colors', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'packing-boxes-pairs-seen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl-seen-colors', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'stack-block-pyramid-seq-seen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'separating-piles-seen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'towers-of-hanoi-seq-seen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| }, | |
| ##### multi-attr tasks | |
| 'multi-attr-align-rope': { | |
| 'train': [ | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'align-rope', | |
| ], | |
| 'test': [ | |
| 'align-rope', | |
| ], | |
| 'attr_train_task': None, | |
| }, | |
| 'multi-attr-packing-shapes': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'packing-shapes', | |
| ], | |
| 'test': [ | |
| 'packing-shapes', | |
| ], | |
| 'attr_train_task': None, | |
| }, | |
| 'multi-attr-assembling-kits-seq-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-seen-colors', # seen only | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'assembling-kits-seq-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'assembling-kits-seq-unseen-colors', | |
| ], | |
| 'attr_train_task': 'assembling-kits-seq-seen-colors', | |
| }, | |
| 'multi-attr-packing-boxes-pairs-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-seen-colors', # seen only | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'packing-boxes-pairs-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'packing-boxes-pairs-unseen-colors', | |
| ], | |
| 'attr_train_task': 'packing-boxes-pairs-seen-colors', | |
| }, | |
| 'multi-attr-packing-unseen-google-objects-seq': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'packing-unseen-google-objects-seq', | |
| ], | |
| 'test': [ | |
| 'packing-unseen-google-objects-seq', | |
| ], | |
| 'attr_train_task': 'packing-seen-google-objects-group', | |
| }, | |
| 'multi-attr-packing-unseen-google-objects-group': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'packing-unseen-google-objects-group', | |
| ], | |
| 'test': [ | |
| 'packing-unseen-google-objects-group', | |
| ], | |
| 'attr_train_task': 'packing-seen-google-objects-seq', | |
| }, | |
| 'multi-attr-put-block-in-bowl-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-seen-colors', # seen only | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'put-block-in-bowl-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'put-block-in-bowl-unseen-colors', | |
| ], | |
| 'attr_train_task': 'put-block-in-bowl-seen-colors', | |
| }, | |
| 'multi-attr-stack-block-pyramid-seq-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-seen-colors', # seen only | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| ], | |
| 'attr_train_task': 'stack-block-pyramid-seq-seen-colors', | |
| }, | |
| 'multi-attr-separating-piles-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-seen-colors', # seen only | |
| 'towers-of-hanoi-seq-full', | |
| ], | |
| 'val': [ | |
| 'separating-piles-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'separating-piles-unseen-colors', | |
| ], | |
| 'attr_train_task': 'separating-piles-seen-colors', | |
| }, | |
| 'multi-attr-towers-of-hanoi-seq-unseen-colors': { | |
| 'train': [ | |
| 'align-rope', | |
| 'assembling-kits-seq-full', | |
| 'packing-boxes-pairs-full', | |
| 'packing-shapes', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq-seen-colors', # seen only | |
| ], | |
| 'val': [ | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'test': [ | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| ], | |
| 'attr_train_task': 'towers-of-hanoi-seq-seen-colors', | |
| }, | |
| } | |
| class RavenMultiTaskDatasetBalance(RavensMultiTaskDataset): | |
| def __init__(self, path, cfg, group='multi-all', | |
| mode='train', n_demos=100, augment=False, balance_weight=0.1): | |
| """A multi-task dataset for balancing data.""" | |
| self.root_path = path | |
| self.mode = mode | |
| if group not in self.MULTI_TASKS: | |
| # generate the groups on the fly | |
| self.tasks = group# .split(" ") | |
| else: | |
| self.tasks = self.MULTI_TASKS[group][mode] | |
| print("self.tasks:", self.tasks) | |
| self.attr_train_task = self.MULTI_TASKS[group]['attr_train_task'] if group in self.MULTI_TASKS and 'attr_train_task' in self.MULTI_TASKS[group] else None | |
| self.cfg = cfg | |
| self.sample_set = {} | |
| self.max_seed = -1 | |
| self.n_episodes = 0 | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = self.cfg['dataset']['cache'] | |
| self.n_demos = n_demos | |
| self.augment = augment | |
| self.aug_theta_sigma = self.cfg['dataset']['augment']['theta_sigma'] if 'augment' in self.cfg['dataset'] else 60 # legacy code issue: theta_sigma was newly added | |
| self.pix_size = 0.003125 | |
| self.in_shape = (320, 160, 6) | |
| self.cam_config = cameras.RealSenseD415.CONFIG | |
| self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]]) | |
| self.n_episodes = {} | |
| episodes = {} | |
| for task in self.tasks: | |
| task_path = os.path.join(self.root_path, f'{task}-{mode}') | |
| action_path = os.path.join(task_path, 'action') | |
| n_episodes = 0 | |
| if os.path.exists(action_path): | |
| for fname in sorted(os.listdir(action_path)): | |
| if '.pkl' in fname: | |
| n_episodes += 1 | |
| self.n_episodes[task] = n_episodes | |
| if n_episodes == 0: | |
| raise Exception(f"{task}-{mode} has 0 episodes. Remove it from the list in dataset.py") | |
| # Select random episode depending on the size of the dataset. | |
| if task in self.ORIGINAL_NAMES and self.mode == 'train': | |
| assert self.n_demos < 200 # otherwise, we need to change the code below | |
| episodes[task] = np.random.choice(range(n_episodes), min(int(self.n_demos*balance_weight), n_episodes), False) | |
| else: | |
| episodes[task] = np.random.choice(range(n_episodes), min(self.n_demos, n_episodes), False) | |
| if self.n_demos > 0: | |
| self.images = self.cfg['dataset']['images'] | |
| self.cache = False | |
| self.set(episodes) | |
| self._path = None | |
| self._task = None | |
| ORIGINAL_NAMES = [ | |
| # demo conditioned | |
| 'align-box-corner', | |
| 'assembling-kits', | |
| 'assembling-kits-easy', | |
| 'block-insertion', | |
| 'block-insertion-easy', | |
| 'block-insertion-nofixture', | |
| 'block-insertion-sixdof', | |
| 'block-insertion-translation', | |
| 'manipulating-rope', | |
| 'packing-boxes', | |
| 'palletizing-boxes', | |
| 'place-red-in-green', | |
| 'stack-block-pyramid', | |
| 'sweeping-piles', | |
| 'towers-of-hanoi', | |
| 'gen-task', | |
| # goal conditioned | |
| 'align-rope', | |
| 'assembling-kits-seq', | |
| 'assembling-kits-seq-seen-colors', | |
| 'assembling-kits-seq-unseen-colors', | |
| 'assembling-kits-seq-full', | |
| 'packing-shapes', | |
| 'packing-boxes-pairs', | |
| 'packing-boxes-pairs-seen-colors', | |
| 'packing-boxes-pairs-unseen-colors', | |
| 'packing-boxes-pairs-full', | |
| 'packing-seen-google-objects-seq', | |
| 'packing-unseen-google-objects-seq', | |
| 'packing-seen-google-objects-group', | |
| 'packing-unseen-google-objects-group', | |
| 'put-block-in-bowl', | |
| 'put-block-in-bowl-seen-colors', | |
| 'put-block-in-bowl-unseen-colors', | |
| 'put-block-in-bowl-full', | |
| 'stack-block-pyramid-seq', | |
| 'stack-block-pyramid-seq-seen-colors', | |
| 'stack-block-pyramid-seq-unseen-colors', | |
| 'stack-block-pyramid-seq-full', | |
| 'separating-piles', | |
| 'separating-piles-seen-colors', | |
| 'separating-piles-unseen-colors', | |
| 'separating-piles-full', | |
| 'towers-of-hanoi-seq', | |
| 'towers-of-hanoi-seq-seen-colors', | |
| 'towers-of-hanoi-seq-unseen-colors', | |
| 'towers-of-hanoi-seq-full', | |
| ] |