def collect_episodes(config, output_dir=None, visualize=True): if output_dir is None: output_dir = os.path.join(os.getcwd(), 'data') if not os.path.exists(output_dir): os.makedirs(output_dir) # save the config config_save_file = os.path.join(output_dir, 'config.yaml') save_yaml(config, config_save_file) save_file = os.path.join(output_dir, "%s.p" % (get_current_YYYY_MM_DD_hh_mm_ss_ms())) # initialize config for DataCollector dc = PusherSliderDataCollector(config) num_episodes = config['dataset']['num_episodes'] multi_episode_container = MultiEpisodeContainer() for i in range(num_episodes): print("collecting episode %d of %d" % (i + 1, num_episodes)) name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i) episode = dc.collect_single_episode(visualize, episode_name=name) multi_episode_container.add_episode(episode) print("saving data to %s" % save_file) multi_episode_container.save_to_file(save_file)
def train_dense_descriptor_vision(dataset_name): def load_config(): config = load_yaml( os.path.join( get_project_root(), "experiments/exp_22_push_box_hardware/integral_heatmap_3d.yaml" )) return config dataset_paths = get_dataset_paths(dataset_name) episodes_config = dataset_paths['episodes_config'] episodes_root = dataset_paths['dataset_root'] # data augmentation if True: model_name = "data_aug_%s" % (get_current_YYYY_MM_DD_hh_mm_ss_ms()) output_dir = os.path.join( get_experiment_save_root(dataset_name), 'trained_models/perception/dense_descriptors', model_name) dataset_paths = get_dataset_paths(dataset_name) dataset_root = dataset_paths['dataset_root'] multi_episode_dict = DCDynamicSpartanEpisodeReader.load_dataset( config=episodes_config, episodes_root=episodes_root) config = load_config() config['dataset']['data_augmentation'] = True config['dataset']['camera_names'] = dataset_paths[ 'dense_descriptor_camera_list'] train_dense_descriptors(config, train_dir=output_dir, multi_episode_dict=multi_episode_dict, verbose=False)
def save_data( self, save_dir=None, ): """ Saves data from the - PlanContainer - OnlineEpisode """ if save_dir is None: save_dir = os.path.join( utils.get_data_root(), 'hardware_experiments/closed_loop_rollouts/sandbox', utils.get_current_YYYY_MM_DD_hh_mm_ss_ms()) if not os.path.exists(save_dir): os.makedirs(save_dir) print("saving MPC rollout data at: %s" % (save_dir)) save_data = { 'episode': self._state_dict['episode'].get_save_data(), 'plan': self._state_dict['plan'].get_save_data(), } save_file = os.path.join(save_dir, "data.p") utils.save_pickle(save_data, save_file) print("done saving data")
def __init__(self, name=None): if name is None: name = get_current_YYYY_MM_DD_hh_mm_ss_ms() self._data = dict() self._data['name'] = name self._data['config'] = None self._data['trajectory'] = []
def collect_episodes(config, output_dir=None, visualize=True, use_threads=False): if output_dir is None: output_dir = os.path.join(os.getcwd(), 'data') if not os.path.exists(output_dir): os.makedirs(output_dir) # save the config config_save_file = os.path.join(output_dir, 'config.yaml') save_yaml(config, config_save_file) # initialize config for DataCollector dc = DrakePusherSliderEpisodeCollector(config) num_episodes = config['dataset']['num_episodes'] # record some metadata metadata = dict() metadata['episodes'] = dict() for i in range(num_episodes): print("\n") start_time = time.time() print("collecting episode %d of %d" % (i + 1, num_episodes)) name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i) episode = dc.collect_single_episode(visualize, episode_name=name) print("saving to disk") metadata['episodes'][name] = dict() image_data_file = episode.save_images_to_hdf5(output_dir) non_image_data_file = episode.save_non_image_data_to_pickle(output_dir) print("non_image_data.keys()", episode.non_image_data.keys()) metadata['episodes'][name]['non_image_data_file'] = non_image_data_file metadata['episodes'][name]['image_data_file'] = image_data_file print("done saving to disk") elapsed = time.time() - start_time print("single episode took: %.2f seconds" % (elapsed)) save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
def collect_single_episode( env, # gym env action_seq, # [N, action_dim] episode_name=None): """ Rollout the action sequence using the simulator Record the observations """ if episode_name is None: episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() episode_container = EpisodeContainer() episode_container.set_config(env.config) episode_container.set_name(episode_name) episode_container.set_metadata(env.get_metadata()) obs_list = [] N = action_seq.shape[0] obs_prev = env.get_observation() for i in range(N): action = action_seq[i] obs, reward, done, info = env.step(action) episode_container.add_obs_action(obs_prev, action) obs_prev = obs # terminate if outside boundary if not env.slider_within_boundary(): print('slider outside boundary, terminating episode') break if not env.pusher_within_boundary(): print('pusher ouside boundary, terminating episode') break return {'episode_container': episode_container}
def collect_episodes(config, output_dir=None, visualize=True, debug=False): if output_dir is None: output_dir = os.path.join(os.getcwd(), 'data') if not os.path.exists(output_dir): os.makedirs(output_dir) # save the config config_save_file = os.path.join(output_dir, 'config.yaml') save_yaml(config, config_save_file) # initialize config for DataCollector dc = DrakePusherSliderEpisodeCollector(config, visualize=visualize) num_episodes = config['dataset']['num_episodes'] # record some metadata metadata = dict() metadata['episodes'] = dict() while (len(metadata['episodes']) < num_episodes): i = len(metadata['episodes']) if debug: input("Press Enter to continue...") print("\n") start_time = time.time() print("collecting episode %d of %d" % (i + 1, num_episodes)) name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i) ic = generate_initial_condition(config=config, T_aug_enabled=True) if debug: print("initial condition\n", ic) episode = dc.collect_single_episode(visualize=visualize, episode_name=name, q_pusher=ic['q_pusher'], q_slider=ic['q_slider'], # v_pusher=ic['v_pusher'], action_seq=ic['action_seq'], ) # potentially discard it if the object didn't move during the data collection if len(episode._data['trajectory']) < 10: print("trajectory was too short, skipping") continue obs_start = episode._data['trajectory'][5]['observation'] obs_end = episode._data['trajectory'][-1]['observation'] q_slider_start = obs_start['slider']['position']['translation'] q_slider_end = obs_end['slider']['position']['translation'] dq_slider = obs_start['slider']['position']['translation'] - obs_end['slider']['position']['translation'] if debug: print("len(episode._data['trajectory'])", len(episode._data['trajectory'])) print("q_slider_start", q_slider_start) print("q_slider_end", q_slider_end) print("dq_slider", dq_slider) print("np.linalg.norm(dq_slider)", np.linalg.norm(dq_slider)) # if slider didn't move by at least 1 mm then discard this episode if (np.linalg.norm(dq_slider) < 0.001): # one mm print("discarding episode since slider didn't move") continue print("saving to disk") metadata['episodes'][name] = dict() image_data_file = episode.save_images_to_hdf5(output_dir) non_image_data_file = episode.save_non_image_data_to_pickle(output_dir) print("output_dir:", output_dir) print("non_image_data.keys()", episode.non_image_data.keys()) metadata['episodes'][name]['non_image_data_file'] = non_image_data_file metadata['episodes'][name]['image_data_file'] = image_data_file print("done saving to disk") elapsed = time.time() - start_time print("single episode took: %.2f seconds" % (elapsed)) save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
def main(): dataset_name, multi_episode_dict = load_episodes() ## Load Model model_name, model_file = get_DD_model_file() model = torch.load(model_file) model = model.cuda() model = model.eval() # make this unique output_dir = os.path.join(get_data_root(), "dev/experiments/09/precomputed_vision_data", "dataset_%s" % (dataset_name), "model_name_%s" % (model_name), get_current_YYYY_MM_DD_hh_mm_ss_ms()) camera_name = "camera_1_top_down" episode_name = "2020-05-13-21-55-01-487901_idx_33" episode_idx = 22 # compute descriptor confidence scores if True: print("\n\n---------Computing Descriptor Confidence Scores-----------") metadata_file = os.path.join(output_dir, 'metadata.p') if os.path.isfile(metadata_file): answer = input( "metadata.p file already exists, do you want to overwrite it? y/n" ) if answer == "y": os.rmdir(output_dir) print("removing existing file and continuing") else: print("aborting") quit() compute_descriptor_confidences( multi_episode_dict, model, output_dir, batch_size=10, num_workers=20, model_file=model_file, camera_name=camera_name, num_ref_descriptors=50, num_batches=10, episode_name_arg=episode_name, episode_idx=episode_idx, ) if True: metadata_file = os.path.join(output_dir, 'metadata.p') metadata = load_pickle(metadata_file) # metadata_file = "/media/hdd/data/key_dynam/dev/experiments/09/precomputed_vision_data/dataset_2020-03-25-19-57-26-556093_constant_velocity_500/model_name_2020-04-07-14-31-35-804270_T_aug_dataset/2020-04-09-20-51-50-624799/metadata.p" # metadata = load_pickle(metadata_file) print("\n\n---------Precomputing Descriptor Keypoints-----------") descriptor_keypoints_output_dir = os.path.join(output_dir, "descriptor_keypoints") precompute_descriptor_keypoints( multi_episode_dict, model, descriptor_keypoints_output_dir, ref_descriptors_metadata=metadata, batch_size=10, num_workers=20, ) if True: confidence_score_data_file = os.path.join(output_dir, 'data.p') confidence_score_data = load_pickle(confidence_score_data_file) print( "\n\n---------Selecting Spatially Separated Keypoints-----------") score_and_select_spatially_separated_keypoints( metadata, confidence_score_data=confidence_score_data, K=5, position_diff_threshold=25, output_dir=output_dir, ) print("Data saved at: ", output_dir) print("Finished Normally")
def collect_single_episode( self, visualize=True, episode_name=None, q_pusher=None, q_slider=None, v_pusher=None, action_sampler=None, # functional action_seq=None, # sequence of actions [N, action_dim] ): """ This collects a single episode by performing the following steps - sample initial conditions for environment - sample action for environment - collect the interaction :return: EpisodeContainer :rtype: """ if episode_name is None: episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() env = self._env env.reset() if visualize: print("setting target realtime rate 1.0") env.simulator.set_target_realtime_rate(1.0) # it takes a while for this to take effect, so it might initially # look like the sim is running really fast, but this is # just a cosmetic issue # sample if they weren't passed in if q_pusher is None or q_slider is None: q_pusher, q_slider = self.get_pusher_slider_initial_positions() use_constant_action = (v_pusher is None) if action_sampler is None: action_sampler = self.sample_pusher_velocity def get_action(counter): if v_pusher is not None: return v_pusher elif action_seq is not None: return action_seq[counter] elif action_sampler is not None: return action_sampler() else: return self.sample_pusher_velocity() context = env.get_mutable_context() env.set_pusher_position(context, q_pusher) env.set_slider_position(context, q_slider) action_zero = np.zeros(2) # if visualize: # env.step(action_zero) # print("sleeping for several seconds to let meshcat visualize") # time.sleep(5.0) # log data into EpisodeContainer object episode_container = EpisodeContainer() episode_container.set_config(env.config) episode_container.set_name(episode_name) episode_container.set_metadata(env.get_metadata()) # step for 2 seconds to let object fall env.step(action_zero, dt=2.0) obs_prev = env.get_observation() start_time = time.time() num_timesteps = self._config['dataset'][ 'num_timesteps'] + self._config['train']['n_history'] counter = 0 for i in range(num_timesteps): start_time_1 = time.time() # stand still for n_history timesteps if i < self._config['train']['n_history']: action = action_zero # just sit still # this makes it easier to train our dynamics model # since we need to get a history of observations else: action = get_action(counter) counter += 1 # exploration_type = self._config['dataset']['data_generation']['exploration_type'] # if exploration_type == 'random': # action = self.sample_pusher_velocity() # elif exploration_type == 'constant': # action = action_actual # single sim time obs, reward, done, info = env.step(action) episode_container.add_obs_action(obs_prev, action) obs_prev = obs if visualize: # note that sim wil elapsed = time.time() - start_time_1 hz = 1.0 / elapsed # print("sim rate (hz)", hz) # print("pusher velocity", obs['pusher']['velocity']) # terminate if outside boundary if not env.slider_within_boundary(): print('slider outside boundary, terminating episode') break if not env.pusher_within_boundary(): print('pusher ouside boundary, terminating episode') break elapsed = time.time() - start_time if DEBUG: print("elapsed wall clock", elapsed) print("elapsed sim time", obs_prev['sim_time']) return episode_container
def DD_3D_dynamics(dataset_name): from key_dynam.training.train_dynamics_pusher_slider_precomputed_keypoints import train_dynamics from key_dynam.experiments.drake_pusher_slider import DD_utils def load_config(): return load_yaml( os.path.join( get_project_root(), 'experiments/exp_22_push_box_hardware/config_DD_3D.yaml')) dataset_paths = get_dataset_paths(dataset_name) dataset_root = dataset_paths['dataset_root'] episodes_config = dataset_paths['episodes_config'] precomputed_vision_data_root = DD_utils.get_precomputed_data_root( dataset_name)['precomputed_data_root'] # descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints') descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints') config = load_config() multi_episode_dict = DynamicSpartanEpisodeReader.load_dataset( config=config, episodes_config=episodes_config, episodes_root=dataset_paths['dataset_root'], load_image_episode=True, precomputed_data_root=descriptor_keypoints_root, max_num_episodes=None) experiment_save_root = get_experiment_save_root(dataset_name) # experiment_save_root = os.path.join(get_data_root(), 'sandbox') # standard if True: TRAIN = True spatial_descriptor_data = load_pickle( os.path.join(precomputed_vision_data_root, 'spatial_descriptors.p')) metadata = load_pickle( os.path.join(precomputed_vision_data_root, 'metadata.p')) train_dir = None config = load_config() config['train']['n_history'] = 2 config['train']['random_seed'] = 1 config['dataset']['visual_observation_function'][ 'camera_name'] = dataset_paths['main_camera_name'] config['dataset'][ 'precomputed_data_root'] = precomputed_vision_data_root suffix = "_DD_3D_n_his_2" model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix if TRAIN: train_dir = os.path.join(experiment_save_root, 'trained_models/dynamics/DD_3D/', model_name) os.makedirs(train_dir) visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config( config, keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx'] ) train_dynamics( config=config, train_dir=train_dir, multi_episode_dict=multi_episode_dict, visual_observation_function=visual_observation_function, metadata=metadata, spatial_descriptors_data=spatial_descriptor_data) # data_aug if False: TRAIN = True spatial_descriptor_data = load_pickle( os.path.join(precomputed_vision_data_root, 'spatial_descriptors.p')) metadata = load_pickle( os.path.join(precomputed_vision_data_root, 'metadata.p')) train_dir = None config = load_config() config['train']['n_history'] = 2 config['train']['random_seed'] = 1 config['dataset']['visual_observation_function'][ 'camera_name'] = dataset_paths['main_camera_name'] config['dataset']['data_augmentation']['enabled'] = True config['dataset'][ 'precomputed_data_root'] = precomputed_vision_data_root suffix = "_DD_3D_n_his_2_T_aug" model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix if TRAIN: train_dir = os.path.join(experiment_save_root, 'trained_models/dynamics/DD_3D/', model_name) os.makedirs(train_dir) visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config( config, keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx'] ) train_dynamics( config=config, train_dir=train_dir, multi_episode_dict=multi_episode_dict, visual_observation_function=visual_observation_function, metadata=metadata, spatial_descriptors_data=spatial_descriptor_data)
# load config # specify the dataset dataset_name = "dps_box_on_side_600" # prepare folders data_dir = os.path.join(get_data_root(), 'dev/experiments/18/data', dataset_name) start_time = time.time() multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset( dataset_root=data_dir) # placeholder for now config = dev_utils.load_integral_heatmap_3d_config() config['dataset']['name'] = dataset_name config['dataset']['camera_names'] = ['camera_angled', 'camera_angled_rotated'] config['dataset']['train_valid_ratio'] = 0.83 model_name = "3D_loss_%s_%s" % ("camera_angled", get_current_YYYY_MM_DD_hh_mm_ss_ms()) train_dir = os.path.join( get_data_root(), 'dev/experiments/18/trained_models/perception/dense_descriptor', model_name) train_dense_descriptors(config, train_dir=train_dir, multi_episode_dict=multi_episode_dict, verbose=False)
# camera_name = "camera_angled" camera_name = "camera_1_top_down" config['perception']['camera_name'] = camera_name config['train'] = config['train_transporter'] config['train']['train_valid_ratio'] = 0.83 print(config) # prepare folders data_dir = os.path.join(get_data_root(), 'dev/experiments/18/data', dataset_name) train_dir = os.path.join( get_data_root(), 'dev/experiments/18/trained_models/perception', "transporter_%s_%s" % (camera_name, get_current_YYYY_MM_DD_hh_mm_ss_ms())) ckp_dir = os.path.join( train_dir, 'train_nKp%d_invStd%.1f' % (config['perception']['n_kp'], config['perception']['inv_std'])) start_time = time.time() multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset( dataset_root=data_dir) print("loading dataset took %d seconds", time.time() - start_time) train_transporter(config=config, train_dir=train_dir, ckp_dir=ckp_dir, multi_episode_dict=multi_episode_dict)
dataset_root = os.path.join(get_project_root(), "data/dev/experiments/05/data", DATASET_NAME) multi_episode_dict = DrakeSimEpisodeReader.load_dataset(dataset_root) network_folder = os.path.join(get_project_root(), "data/dev/experiments/05/trained_models", MODEL_NAME) epoch = 55 iter = 0 model_file = os.path.join(network_folder, "net_dy_epoch_%d_iter_%d_model.pth" %(epoch, iter)) model = torch.load(model_file) model.cuda() model = model.eval() # make this unique output_dir = os.path.join(get_project_root(), "data/dev/experiments/05/precomputed_descriptor_keypoints", DATASET_NAME, MODEL_NAME, get_current_YYYY_MM_DD_hh_mm_ss_ms()) camera_name = "camera_1_top_down" precompute_descriptor_keypoints(multi_episode_dict, model, output_dir, batch_size=10, num_workers=20, model_file=model_file, camera_name=camera_name, num_ref_descriptors=32, )
def collect_episodes(config, metadata, data_collector, n_episode=1, data_dir=None, visualize=False, episode_name=None, exploration_type='random', model_dy=None): """ This collects a single episode by performing the following steps - sample initial conditions for environment - sample action from controller for environment - collect the interaction """ if data_dir is None: data_dir = os.path.join(os.getcwd(), 'data') if not os.path.exists(data_dir): os.makedirs(data_dir) for idx_episode in range(n_episode): start_time = time.time() print("collecting episode %d of %d" % (idx_episode + 1, n_episode)) episode_name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), idx_episode) env = data_collector._env env.reset() if visualize: print("setting target realtime rate 1.0") env.simulator.set_target_realtime_rate(1.0) q_pusher, q_slider = data_collector.get_pusher_slider_initial_positions( ) context = env.get_mutable_context() env.set_pusher_position(context, q_pusher) env.set_slider_position(context, q_slider) action_zero = np.zeros(2) env.step(action_zero, dt=1.0) # to get the box to drop down # log data into EpisodeContainer object episode_container = EpisodeContainer() episode_container.set_config(env.config) episode_container.set_name(episode_name) episode_container.set_metadata(env.get_metadata()) obs_prev = env.get_observation() num_timesteps = config['mpc']['num_timesteps'] + config['train'][ 'n_history'] for i in range(num_timesteps): # stand still for n_history timesteps if i < config['train']['n_history']: action = action_zero # this makes it easier to train our dynamics model # since we need to get a history of observations else: if exploration_type == 'random': action = data_collector.sample_pusher_velocity() elif exploration_type == 'mppi': action = data_collector.sample_pusher_velocity() # single sim time obs, reward, done, info = env.step(action) episode_container.add_obs_action(obs_prev, action) obs_prev = obs # terminate if outside boundary if not env.slider_within_boundary(): print('slider outside boundary, terminating episode') break if not env.pusher_within_boundary(): print('pusher ouside boundary, terminating episode') break # print("saving to disk") metadata['episodes'][episode_name] = dict() image_data_file = episode_container.save_images_to_hdf5(data_dir) non_image_data_file = episode_container.save_non_image_data_to_pickle( data_dir) metadata['episodes'][episode_name][ 'non_image_data_file'] = non_image_data_file metadata['episodes'][episode_name]['image_data_file'] = image_data_file # print("done saving to disk") elapsed = time.time() - start_time print("single episode took: %.2f seconds" % (elapsed))
def collect_single_episode(self, visualize=True, episode_name=None): """ This collects a single episode by performing the following steps - sample initial conditions for environment - sample action for environment - collect the interaction :return: EpisodeContainer :rtype: """ if episode_name is None: episode_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() env = self._env env.reset() # set pusher position # this is in Body frame pusher_position_B = self.pusher_initial_position_space_from_config( env.config).sample() # this is in world frame pusher_position_W = env.slider_body.local_to_world(pusher_position_B) # set pusher position env.pusher_body.position = pusher_position_W # step so we can get an observation that has correct image env.step([0, 0]) # get action # don't worry about converting it to world frame for now pusher_velocity_B = self.sample_pusher_velocity() action_actual = pusher_velocity_B action_zero = np.zeros(2) episode_container = EpisodeContainer() episode_container.set_config(env.config) episode_container.set_name(episode_name) obs_prev = env.get_observation() num_timesteps = self._config['dataset'][ 'num_timesteps'] + self._config['train']['n_history'] for i in xrange(num_timesteps): if visualize: env.render(mode='human') # stand still for n_history timesteps if i < self._config['train']['n_history']: action = action_zero # just sit still # this makes it easier to train our dynamics model # since we need to get a history of observations else: exploration_type = self._config['dataset']['data_generation'][ 'exploration_type'] if exploration_type == 'random': action = self.sample_pusher_velocity() elif exploration_type == 'constant': action = action_actual obs, reward, done, info = env.step(action) episode_container.add_obs_action(obs_prev, action) obs_prev = obs if DEBUG: print("\n\nsim_time = ", env._sim_time) print("obs sim_time = ", obs['sim_time']) return episode_container
def collect_episodes(config, output_dir=None, visualize=True, debug=False, run_from_thread=False, seed=None): # gets a random seed for each thread/process independently if seed is None: seed = np.random.RandomState().randint(0, 10000) set_seed(seed) if output_dir is None: output_dir = os.path.join(os.getcwd(), 'data') if not os.path.exists(output_dir): os.makedirs(output_dir) # save the config config_save_file = os.path.join(output_dir, 'config.yaml') save_yaml(config, config_save_file) # initialize config for DataCollector num_episodes = config['dataset']['num_episodes'] # record some metadata metadata = dict() metadata['episodes'] = dict() while (len(metadata['episodes']) < num_episodes): i = len(metadata['episodes']) if debug: input("Press Enter to continue...") print("\n") start_time = time.time() print("collecting episode %d of %d" % (i + 1, num_episodes)) name = "%s_idx_%d" % (get_current_YYYY_MM_DD_hh_mm_ss_ms(), i) n_his = config['train_dynamics']['n_history'] ic = generate_initial_condition( config=config, T_aug_enabled=True, n_his=n_his, randomize_velocity=True, randomize_sdf=True, randomize_color=True, ) env = DrakeMugsEnv(ic['config'], visualize=visualize) if debug: print("initial condition\n", ic) # set initial condition on environment if visualize: print("setting target realtime rate 1.0") env.simulator.set_target_realtime_rate(1.0) env.reset() context = env.get_mutable_context() env.set_object_position(context, ic['q_slider']) env.set_pusher_position(context, ic['q_pusher']) print("ic['action_sequence'].shape", ic['action_sequence'].shape) # simulate for 10 seconds to let the mug stabilize action_zero = env.get_zero_action() env.step(action_zero, dt=10.0) episode = collect_single_episode( env, action_seq=ic['action_sequence'])['episode_container'] # potentially discard it if the object didn't move during the data collection if len(episode._data['trajectory']) < 10: print("trajectory was too short, skipping") continue obs_start = episode._data['trajectory'][0]['observation'] obs_end = episode._data['trajectory'][-1]['observation'] q_slider_start = obs_start['slider']['position']['translation'] q_slider_end = obs_end['slider']['position']['translation'] dq_slider = obs_start['slider']['position']['translation'] - obs_end[ 'slider']['position']['translation'] if debug: print("len(episode._data['trajectory'])", len(episode._data['trajectory'])) print("q_slider_start", q_slider_start) print("q_slider_end", q_slider_end) print("dq_slider", dq_slider) print("np.linalg.norm(dq_slider)", np.linalg.norm(dq_slider)) pose_error = compute_pose_error(obs_start, obs_end) # if slider didn't move by at least 1 mm then discard this episode if (pose_error['position_error'] < 0.01) and (pose_error['angle_error_degrees'] < 10): print( "discarding episode since slider didn't move sufficiently far") continue print("saving to disk") metadata['episodes'][name] = dict() image_data_file = episode.save_images_to_hdf5(output_dir) non_image_data_file = episode.save_non_image_data_to_pickle(output_dir) print("output_dir:", output_dir) print("non_image_data.keys()", episode.non_image_data.keys()) metadata['episodes'][name]['non_image_data_file'] = non_image_data_file metadata['episodes'][name]['image_data_file'] = image_data_file print("done saving to disk") elapsed = time.time() - start_time print("single episode took: %.2f seconds" % (elapsed)) if not run_from_thread: save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml')) print("Finished collecting episodes") return {'metadata': metadata}