def get_dataset_paths(dataset_name): if dataset_name == "push_box_hardware": episodes_root = os.path.join(get_data_ssd_root(), "dataset/push_box_hardware") episodes_config = load_yaml( os.path.join( get_project_root(), 'experiments/exp_22_push_box_hardware/push_box_hardware_episodes_config.yaml' )) transporter_model_chkpt = None dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth" return { 'dataset_name': dataset_name, 'dataset_root': episodes_root, 'episodes_config': episodes_config, 'main_camera_name': 'd415_01', 'dense_descriptor_camera_list': ['d415_01', 'd415_02'], 'transporter_model_chkpt': transporter_model_chkpt, 'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt, } elif dataset_name == "push_box_string_pull": episodes_root = os.path.join(get_data_ssd_root(), "dataset/push_box_string_pull") episodes_config = load_yaml( os.path.join( get_project_root(), 'experiments/exp_22_push_box_hardware/push_box_string_pull_episodes_config.yaml' )) transporter_model_chkpt = None dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth" return { 'dataset_name': dataset_name, 'dataset_root': episodes_root, 'episodes_config': episodes_config, 'main_camera_name': 'd415_01', 'dense_descriptor_camera_list': ['d415_01', 'd415_02'], 'transporter_model_chkpt': transporter_model_chkpt, 'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt, } else: raise ValueError("unknown dataset:", 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
def test_pusher_slider_dataset(): # dataset, config = create_pusher_slider_dataset() project_root = get_project_root() config_file = os.path.join(project_root, "experiments/01/config.yaml") config = load_yaml(config_file) # new dataset loading approach episodes = load_episodes_from_config(config) action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.function_from_config( config) dataset = MultiEpisodeDataset(config, action_function=action_function, observation_function=observation_function, episodes=episodes, phase="train") data = dataset[0] # test the getitem print("type(data)", type(data)) print("list(data)", list(data)) print(type(data["observations"])) print("observations.shape", data["observations"].shape) print("actions.shape", data["actions"].shape) print("observations", data["observations"]) print("actions", data["actions"]) stats = dataset.compute_dataset_statistics() print("stats", stats)
def sample_random_mug(): sdf_dir = os.path.join(get_data_root(), "stable/sim_assets/anzu_mugs") # sdf_file = random.choice(SDFHelper.get_sdf_list(sdf_dir)) mug_list = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/mugs.yaml')) sdf_file = random.choice(mug_list['corelle_mug-small']) sdf_file = os.path.join(sdf_dir, sdf_file) return sdf_file
def load_default_config(): """ Loads the experiments/05/config.yaml :return: :rtype: """ config_file = os.path.join(get_project_root(), 'experiments/05/config.yaml') config = load_yaml(config_file) return config
def f(q_tmp): config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) config['dataset']['num_episodes'] = num_episodes_per_thread out = collect_episodes(config, output_dir=OUTPUT_DIR, visualize=False, debug=False, run_from_thread=True) q_tmp.put(out)
def test_pusher_slider_keypoint_dataset(): project_root = get_project_root() config_file = os.path.join(project_root, "experiments/02/config.yaml") config = load_yaml(config_file) config["n_history"] = 1 config["n_roll"] = 0 # new dataset loading approach episodes = load_episodes_from_config(config) action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.function_from_config( config) dataset = MultiEpisodeDataset(config, action_function=action_function, observation_function=observation_function, episodes=episodes, phase="train") # dataset, config = create_pusher_slider_keypoint_dataset(config=config) episode_names = dataset.get_episode_names() episode_names.sort() episode_name = episode_names[0] episode = dataset.episode_dict[episode_name] obs_raw = episode.get_observation(0) obs_raw['slider']['angle'] = 0 dataset.observation_function(obs_raw) print("20 degrees\n\n\n\n") obs_raw['slider']['angle'] = np.deg2rad(90) dataset.observation_function(obs_raw) quit() data = dataset[0] # test the getitem print("type(data)", type(data)) print("data.keys()", data.keys()) print(type(data["observations"])) print("observations.shape", data["observations"].shape) print("actions.shape", data["actions"].shape) print("observations", data["observations"]) print("actions", data["actions"])
def load_episodes_from_config(config): """ Loads episodes using the path specified in the config :param config: :type config: :return: :rtype: """ data_path = config["dataset"]["data_path"] if not os.path.isabs(data_path): data_path = os.path.join(get_project_root(), data_path) # load the data print("loading data from disk . . . ") raw_data = load_pickle(data_path) print("finished loading data") episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data) return episodes
def debug(): save_dir = os.path.join(get_project_root(), 'sandbox/mpc/push_right_box_horizontal') save_dir = "/home/manuelli/data/key_dynam/sandbox/2020-07-07-20-09-54_push_right_box_horizontal" plan_msg = load_pickle(os.path.join(save_dir, 'plan_msg.p'), encoding='latin1') plan_msg = zmq_utils.convert(plan_msg) compute_control_action_msg = load_pickle(os.path.join( save_dir, 'compute_control_action_msg.p'), encoding='latin1') compute_control_action_msg = zmq_utils.convert(compute_control_action_msg) K_matrix = None T_world_camera = None if 'K_matrix' in plan_msg: K_matrix = plan_msg['K_matrix'] if 'T_world_camera' in plan_msg: T_world_camera = plan_msg['T_world_camera'] d = load_model_and_data( K_matrix=K_matrix, T_world_camera=T_world_camera, ) controller = ZMQController( config=d['config'], model_dy=d['model_dy'], action_function=d['action_function'], observation_function=d['observation_function'], visual_observation_function=d['visual_observation_function'], planner=d['planner'], debug=True, zmq_enabled=False, camera_info=d['camera_info'], ) controller._on_plan_msg(plan_msg) controller._on_compute_control_action(compute_control_action_msg)
def construct_dataset_from_config( config, # dict for global config phase="train", # str: either "train" or "valid" episodes=None, # optional dict of episodes to use instead of loading data ): """ Construct dataset based on global config :param config: :type config: :return: :rtype: """ assert phase in ["train", "valid"] data_path = config["dataset"]["data_path"] if not os.path.isabs(data_path): data_path = os.path.join(get_project_root(), data_path) if episodes is None: # load the data if not passed in episodes = load_episodes_from_config(config) action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.function_from_config( config) dataset = MultiEpisodeDataset(config, action_function=action_function, observation_function=observation_function, episodes=episodes, phase=phase) return_data = { "dataset": dataset, "action_function": action_function, "observation_function": observation_function, "episodes": episodes, } return return_data
def main(): start_time = time.time() config = load_yaml(os.path.join(get_project_root(), 'experiments/drake_pusher_slider/env_config.yaml')) config['dataset']['num_episodes'] = 1000 # half for train, half for valid set_seed(500) # just randomly chosen num_episodes = config['dataset']['num_episodes'] DATASET_NAME = "box_push_%d" %(num_episodes) OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME) if not os.path.exists(OUTPUT_DIR): os.makedirs(OUTPUT_DIR) collect_episodes( config, output_dir=OUTPUT_DIR, visualize=False, debug=False) elapsed = time.time() - start_time print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def main(): start_time = time.time() config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) config['dataset']['num_episodes'] = 10 set_seed(500) # just randomly chosen DATASET_NAME = "mugs_%d" % (config['dataset']['num_episodes']) OUTPUT_DIR = os.path.join(get_data_root(), 'sandbox', DATASET_NAME) print("OUTPUT_DIR:", OUTPUT_DIR) collect_episodes(config, output_dir=OUTPUT_DIR, visualize=False, debug=False) elapsed = time.time() - start_time print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def create_pusher_slider_keypoint_dataset(config=None): # load some previously generated data project_root = get_project_root() if config is None: config_file = os.path.join(project_root, "experiments/02/config.yaml") config = load_yaml(config_file) action_function = ActionFunctionFactory.pusher_velocity obs_function = ObservationFunctionFactory.pusher_pose_slider_keypoints( config) DATA_PATH = os.path.join( project_root, "test_data/pusher_slider_10_episodes/2019-10-22-21-30-02-536750.p") raw_data = load_pickle(DATA_PATH) episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data) # create MultiEpisodeDataset dataset = MultiEpisodeDataset(config, action_function=action_function, observation_function=obs_function, episodes=episodes) episode = dataset.get_random_episode() data_0 = episode.get_observation(0) data_1 = episode.get_observation(1) print("time 0", data_0["sim_time"]) print("time 1", data_1["sim_time"]) # episode_name = episodes.keys()[0] # episode = episodes[episode_name] # data = episode.data # print("episode.data.keys()", episode.data.keys()) # print("test ", type(data["trajectory"][0].keys())) # print("test ", data["trajectory"][0].keys()) return dataset, config
def run_interactive_circle_slider(): """ Launch interactive environment where you can move the pusher around with the arrow keys :return: :rtype: """ config_file = os.path.join(get_project_root(), 'experiments/03/config.yaml') config = load_yaml(config_file) env = PusherSlider(config=config) env.reset() while env._running: action = env.process_events() env.step(action) obs, reward, done, info = env.step(action) env.render(mode='human') if True: print("\n\n\n") print("slider position", obs['slider']['position']) print("pusher position", obs['pusher']['position'])
def main(): start_time = time.time() config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_18_box_on_side/config.yaml')) # config['dataset']['num_episodes'] = 500 # half for train, half for valid config['dataset']['num_episodes'] = 600 # half for train, half for valid set_seed(500) # just randomly chosen DATASET_NAME = "dps_box_on_side_%d" % (config['dataset']['num_episodes']) OUTPUT_DIR = os.path.join(get_data_root(), "dev/experiments/18/data", DATASET_NAME) print("OUTPUT_DIR:", OUTPUT_DIR) collect_episodes(config, output_dir=OUTPUT_DIR, visualize=False, debug=False) elapsed = time.time() - start_time print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def load_model_and_data( K_matrix=None, T_world_camera=None, ): dataset_name = "push_box_hardware" model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug" train_dir = os.path.join( get_data_root(), "dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics") # train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics" train_dir = os.path.join(train_dir, model_name) ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth") train_config = load_yaml(os.path.join(train_dir, 'config.yaml')) state_dict = torch.load(ckpt_file) # build dynamics model model_dy = build_dynamics_model(train_config) # print("state_dict.keys()", state_dict.keys()) model_dy.load_state_dict(state_dict) model_dy = model_dy.eval() model_dy = model_dy.cuda() # load the dataset dataset_paths = get_dataset_paths(dataset_name) dataset_root = dataset_paths['dataset_root'] episodes_config = dataset_paths['episodes_config'] spatial_descriptor_data = load_pickle( os.path.join(train_dir, 'spatial_descriptors.p')) metadata = load_pickle(os.path.join(train_dir, 'metadata.p')) ref_descriptors = spatial_descriptor_data['spatial_descriptors'] ref_descriptors = torch_utils.cast_to_torch(ref_descriptors).cuda() # dense descriptor model model_dd_file = metadata['model_file'] model_dd = torch.load(model_dd_file) model_dd = model_dd.eval() model_dd = model_dd.cuda() camera_name = train_config['dataset']['visual_observation_function'][ 'camera_name'] camera_info = None if (T_world_camera is not None) and (K_matrix is not None): camera_info = { "K": K_matrix, 'T_world_camera': T_world_camera, } else: camera_info = get_spartan_camera_info(camera_name) camera_info['camera_name'] = camera_name visual_observation_function = \ VisualObservationFunctionFactory.descriptor_keypoints_3D(config=train_config, camera_name=camera_name, model_dd=model_dd, ref_descriptors=ref_descriptors, K_matrix=camera_info['K'], T_world_camera=camera_info['T_world_camera'], ) action_function = ActionFunctionFactory.function_from_config(train_config) observation_function = ObservationFunctionFactory.function_from_config( train_config) #### PLANNER ####### planner = None # make a planner config planner_config = copy.copy(train_config) config_tmp = load_yaml( os.path.join(get_project_root(), 'experiments/exp_22_push_box_hardware/config_DD_3D.yaml')) planner_config['mpc'] = config_tmp['mpc'] if PLANNER_TYPE == "random_shooting": planner = RandomShootingPlanner(planner_config) elif PLANNER_TYPE == "mppi": planner = PlannerMPPI(planner_config) else: raise ValueError("unknown planner type: %s" % (PLANNER_TYPE)) return { "model_dy": model_dy, 'config': train_config, 'spatial_descriptor_data': spatial_descriptor_data, 'action_function': action_function, 'observation_function': observation_function, 'visual_observation_function': visual_observation_function, 'planner': planner, 'camera_info': camera_info, }
def test_gym_API(): """ Collects two episodes and saves them to disk Currently no images are being saved :return: :rtype: """ config_file = os.path.join(get_project_root(), 'experiments/01/config.yaml') config = load_yaml(config_file) env = PusherSlider(config=config) env.setup_environment() episode_container = None # for namespacing purposes multi_episode_container = MultiEpisodeContainer() print("fps", env._fps) num_episodes = 1 # num_timesteps_per_episode = 40 for episode_idx in range(num_episodes): env.reset() obs_prev = env.get_observation() episode_container = EpisodeContainer() episode_container.set_config(env.config) for i in range(env.config['dataset']['num_timesteps']): env.render(mode='human') if i == 50: # save this image out in local directory img = env.render(mode='rgb_array') if DEBUG_PRINTS: print("saving figure to file") print("type(img)", type(img)) print(img.shape) print(img.dtype) pil_img = numpy_to_PIL(img) pil_img.save('test_PIL.png') action = np.array([100, 0]) obs, reward, done, info = env.step(action) episode_container.add_obs_action(obs_prev, action) obs_prev = obs if DEBUG_PRINTS: print("\n\nsim_time = ", env._sim_time) print("obs sim_time = ", obs['sim_time']) multi_episode_container.add_episode(episode_container) if True: save_file = os.path.join(os.getcwd(), "single_episode_log.p") episode_container.save_to_file(save_file) multi_episode_save_file = os.path.join(os.getcwd(), "multi_episode_log.p") multi_episode_container.save_to_file(multi_episode_save_file)
def load_config(): return load_yaml( os.path.join( get_project_root(), 'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
def multiprocess_main(num_episodes=1000, num_threads=4): set_seed(500) # just randomly chosen start_time = time.time() config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) num_episodes_per_thread = math.ceil(num_episodes / num_threads) num_episodes = num_threads * num_episodes_per_thread # DATASET_NAME = "mugs_random_colors_%d" % (num_episodes) # DATASET_NAME = "single_mug_%d" # DATASET_NAME = "correlle_mug-small_single_color_%d" %(num_episodes) # DATASET_NAME = "single_corelle_mug_%d" %(num_episodes) # DATASET_NAME = "correlle_mug-small_many_colors_%d" %(num_episodes) DATASET_NAME = "correlle_mug-small_many_colors_random_%d" % (num_episodes) # OUTPUT_DIR = os.path.join(get_data_root(), 'sandbox', DATASET_NAME) OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME) print("OUTPUT_DIR:", OUTPUT_DIR) output_dir = OUTPUT_DIR if not os.path.exists(output_dir): os.makedirs(output_dir) def f(q_tmp): config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) config['dataset']['num_episodes'] = num_episodes_per_thread out = collect_episodes(config, output_dir=OUTPUT_DIR, visualize=False, debug=False, run_from_thread=True) q_tmp.put(out) q = Queue() process_list = [] for i in range(num_threads): p = Process(target=f, args=(q, )) p.start() process_list.append(p) metadata = {'episodes': {}} for p in process_list: while p.is_alive(): p.join(timeout=1) # empty out the queue while not q.empty(): out = q.get() metadata['episodes'].update(out['metadata']['episodes']) # double check for p in process_list: p.join() time.sleep(1.0) print("All threads joined") elapsed = time.time() - start_time # collect the metadata.yaml files while not q.empty(): out = q.get() metadata['episodes'].update(out['metadata']['episodes']) save_yaml(metadata, os.path.join(OUTPUT_DIR, 'metadata.yaml')) print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def load_simple_config(): config_file = os.path.join(get_project_root(), 'config/simple_config.yaml') config = load_yaml(config_file) return config
def run_gym_env(): """ Runs the gym env :return: :rtype: """ DEBUG_PRINTS = True USE_PYGAME = True try: # setup pygame thing if USE_PYGAME: pygame.init() screen = pygame.display.set_mode( (640, 480)) # needed for grabbing focus clock = pygame.time.Clock() default_action = np.zeros(2) velocity = 0.2 if USE_PYGAME: action = process_pygame_events() else: action = default_action config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) env = DrakeMugsEnv(config) env.reset() context = env.get_mutable_context() pos = np.array([0, 0, 0.1]) # quat = transforms3d.euler.euler2quat(np.deg2rad(90), 0, 0) quat = np.array([1, 0, 0, 0]) q = np.concatenate((quat, pos)) env.set_object_position(context=context, q=q) # set the box pose # context = env.get_mutable_context() # pos = np.array([1.56907481e-04, 1.11390697e-06, 5.11972761e-02]) # quat = np.array([ 7.13518047e-01, -6.69765583e-07, -7.00636851e-01, -6.82079212e-07]) # q_slider = np.concatenate((quat, pos)) # env.set_slider_position(context, q=q_slider) env.simulator.set_target_realtime_rate(1.0) # move box araound # context = env.get_mutable_context() # q_slider = [-0.05, 0, 0.03] num_model_instances = env.diagram_wrapper.mbp.num_model_instances() print("num_model_instances", num_model_instances) print("num_positions", env.diagram_wrapper.mbp.num_positions()) label_db = env.diagram_wrapper.get_label_db() print("label db:", label_db.all()) mask_db = env.diagram_wrapper.get_labels_to_mask() print("mask db:", mask_db.all()) # context = env.get_mutable_context() # # # set the position of pusher # q_pusher = [0.2,0.2] # mbp = env.diagram_wrapper.mbp # mbp_context = env.diagram.GetMutableSubsystemContext(mbp, context) # mbp.SetPositions(mbp_context, env.diagram_wrapper.models['pusher'], q_pusher) camera_names = list( config['env']['rgbd_sensors']['sensor_list'].keys()) camera_names.sort() image_vis = ImageVisualizer(len(camera_names), 1) print("running sim") while True: if USE_PYGAME: action = velocity * process_pygame_events() else: action = default_action # print("action:", action) obs, reward, done, info = env.step(action) # print("obs\n", obs) # visualize RGB images in matplotlib for idx, camera_name in enumerate(camera_names): rgb_image = obs['images'][camera_name]['rgb'] image_vis.draw_image(idx, 0, rgb_image) image_vis.visualize_interactive() # # print unique depth values # depth_32F = obs['images']['camera_0']['depth_32F'] # print("unique depth_32F vals", np.unique(depth_32F)) # # print unique depth values # depth_16U = obs['images']['camera_0']['depth_16U'] # print("unique depth_16U vals", np.unique(depth_16U)) # build simulator except KeyboardInterrupt: pygame.quit() plt.close()
def test_dynanet_mlp(): # just try doing a single forward pass dataset, config = create_pusher_slider_dataset() stats = dataset.compute_dataset_statistics() n_history = config["train"]["n_history"] # obs_mean_repeat = stats['observations']['mean'].repeat(n_history, 1) # obs_std_repeat = stats['observations']['std'].repeat(n_history, 1) obs_mean_repeat = stats['observations']['mean'] obs_std_repeat = stats['observations']['std'] observations_normalizer = DataNormalizer(obs_mean_repeat, obs_std_repeat) # action_mean_repeat = stats['actions']['mean'].repeat(n_history, 1) # action_std_repeat = stats['actions']['std'].repeat(n_history, 1) action_mean_repeat = stats['actions']['mean'] action_std_repeat = stats['actions']['std'] actions_normalizer = DataNormalizer(action_mean_repeat, action_std_repeat) config["dataset"]["state_dim"] = 5 config["dataset"]["action_dim"] = 2 model = DynaNetMLP(config) # print summary of model before adding new modules print("\n\n -----summary of model BEFORE adding normalization modules") print("num trainable parameters", count_trainable_parameters(model)) print("num non-trainable parameters ", count_non_trainable_parameters(model)) print("\n\n") # summary of model after adding new params model.set_action_normalizer(actions_normalizer) model.set_state_normalizer(observations_normalizer) print("\n\n -----summary of model AFTER adding normalization modules") print("num trainable parameters", count_trainable_parameters(model)) print("num non-trainable parameters ", count_non_trainable_parameters(model)) print("\n\n") # unsqueeze to mimic dataloader with batch size of 1 data = dataset[0] # test the getitem observations = data['observations'].unsqueeze(0) actions = data['actions'].unsqueeze(0) obs_slice = observations[:, :n_history, :] action_slice = actions[:, :n_history, :] print("action_slice.shape", action_slice.shape) print("obs_slice.shape", obs_slice.shape) # run the model forwards one timestep output = model.forward(obs_slice, action_slice) print("output.shape", output.shape) # save the model with torch.save and torch.load save_dir = os.path.join(get_project_root(), 'sandbox') model_save_file = os.path.join(save_dir, "model.pth") torch.save(model, model_save_file) # load the model model_load = torch.load(model_save_file) print("\n\n -----summary of model LOADED from disk") print("num trainable parameters", count_trainable_parameters(model_load)) print("num non-trainable parameters ", count_non_trainable_parameters(model_load)) print("\n\n") # now try doing the same but with the state dict # my hunch is that this won't work . . . params_save_file = os.path.join(save_dir, "model_params.pth") torch.save(model.state_dict(), params_save_file) # load the model model_load = DynaNetMLP(config) state_dict = torch.load(params_save_file) for param_tensor in state_dict: print(param_tensor, "\t", state_dict[param_tensor].size()) # try creating some dummy DataNormalizer objects # model_load.set_state_normalizer(DataNormalizer(0.0,1.0)) # model_load.set_action_normalizer(DataNormalizer(0.0,1.0)) model_load.load_state_dict(state_dict) print("\n\n -----summary of model LOADED from disk with state_dict method") print("num trainable parameters", count_trainable_parameters(model_load)) print("num non-trainable parameters ", count_non_trainable_parameters(model_load)) print("\n\n") print("model_load._action_normalizer._mean", model_load.action_normalizer._mean) print("model._action_normalizer._mean", model.action_normalizer._mean)
def main(): # load dynamics model model_dict = load_autoencoder_model() model = model_dict['model_dy'] model_ae = model_dict['model_ae'] visual_observation_function = model_dict['visual_observation_function'] config = model.config env_config = load_yaml( os.path.join(get_project_root(), 'experiments/exp_18_box_on_side/config.yaml')) env_config['env']['observation']['depth_int16'] = True n_history = config['train']['n_history'] # create the environment # create the environment env = DrakePusherSliderEnv(env_config) env.reset() # create another environment for doing rollouts env2 = DrakePusherSliderEnv(env_config, visualize=False) env2.reset() action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.drake_pusher_position_3D( config) episode = OnlineEpisodeReader() mpc_input_builder = DynamicsModelInputBuilder( observation_function=observation_function, visual_observation_function=visual_observation_function, action_function=action_function, episode=episode) vis = meshcat_utils.make_default_visualizer_object() vis.delete() initial_cond = get_initial_state() reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) obs_init = env.get_observation() # visualize starting position of the object print("obs_init.keys()", obs_init.keys()) print("obs_init['slider']['position']", obs_init['slider']['position']) T = DrakePusherSliderEnv.object_position_from_observation(obs_init) vis['start_pose'].set_object(triad(scale=0.1)) vis['state_pose'].set_transform(T) #### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############ reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) # add just some large number of these episode.clear() for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) #### ROLLOUT THE ACTION SEQUENCE USING THE SIMULATOR ########## # rollout single action sequence using the simulator gt_rollout_data = env_utils.rollout_action_sequence( env, initial_cond['action_sequence'].cpu().numpy()) env_obs_rollout_gt = gt_rollout_data['observations'] gt_rollout_episode = gt_rollout_data['episode_reader'] for i, env_obs in enumerate(gt_rollout_data['observations']): T = DrakePusherSliderEnv.object_position_from_observation(env_obs) vis_name = "GT_trajectory/%d" % (i) vis[vis_name].set_object(triad(scale=0.1)) vis[vis_name].set_transform(T) action_state_gt = mpc_input_builder.get_action_state_tensors( start_idx=0, num_timesteps=N, episode=gt_rollout_episode) state_rollout_gt = action_state_gt['states'] action_rollout_gt = action_state_gt['actions'] z_object_rollout_gt = model.compute_z_state( state_rollout_gt)['z_object_flat'] print('state_rollout_gt.shape', state_rollout_gt.shape) print("z_object_rollout_gt.shape", z_object_rollout_gt.shape) def goal_func(obs_tmp): state_tmp = mpc_input_builder.get_state_input_single_timestep( {'observation': obs_tmp})['state'] return model.compute_z_state( state_tmp.unsqueeze(0))['z_object'].flatten() # using the vision model to get "goal" keypoints z_object_goal = goal_func(env_obs_rollout_gt[-1]) z_object_goal_np = torch_utils.cast_to_numpy(z_object_goal) # input("press Enter to continue") #### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############ reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) # add just some large number of these episode.clear() for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) # [n_history, state_dim] idx = episode.get_latest_idx() dyna_net_input = mpc_input_builder.get_dynamics_model_input( idx, n_history=n_history) state_init = dyna_net_input['states'].cuda() # [n_history, state_dim] action_init = dyna_net_input['actions'] # [n_history, action_dim] print("state_init.shape", state_init.shape) print("action_init.shape", action_init.shape) print("n_history", n_history) action_seq_gt_torch = initial_cond['action_sequence'] action_input = torch.cat( (action_init[:(n_history - 1)], action_seq_gt_torch), dim=0).cuda() print("action_input.shape", action_input.shape) # rollout using the ground truth actions and learned model # need to add the batch dim to do that z_init = model.compute_z_state(state_init)['z'] rollout_pred = rollout_model(state_init=z_init.unsqueeze(0), action_seq=action_input.unsqueeze(0), dynamics_net=model, compute_debug_data=True) state_pred_rollout = rollout_pred['state_pred'].squeeze(0) print("state_pred_rollout.shape", state_pred_rollout.shape) # input("press Enter to continue") # check L2 distance between predicted and actual # basically comparing state_pred_rollout and state_rollout_gt print("state_rollout_gt[-1]\n", state_rollout_gt[-1]) print("state_pred_rollout[-1]\n", state_pred_rollout[-1]) index_dict = get_object_and_robot_state_indices(config) object_indices = index_dict['object_indices'] # reset the environment and use the MPC controller to stabilize this # now setup the MPC to try to stabilize this . . . . reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) episode.clear() # add just some large number of these for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) # input("press Enter to continue") # make a planner config planner_config = copy.copy(config) config_tmp = load_yaml( os.path.join(get_project_root(), 'experiments/drake_pusher_slider/eval_config.yaml')) planner_config['mpc'] = config_tmp['mpc'] planner_config['mpc']['mppi']['terminal_cost_only'] = TERMINAL_COST_ONLY planner_config['mpc']['random_shooting'][ 'terminal_cost_only'] = TERMINAL_COST_ONLY planner = None if PLANNER_TYPE == "random_shooting": planner = RandomShootingPlanner(planner_config) elif PLANNER_TYPE == "mppi": planner = PlannerMPPI(planner_config) else: raise ValueError("unknown planner type: %s" % (PLANNER_TYPE)) mpc_out = None action_seq_mpc = None state_pred_mpc = None counter = -1 while True: counter += 1 print("\n\n-----Running MPC Optimization: Counter (%d)-------" % (counter)) obs_cur = env.get_observation() episode.add_observation_only(obs_cur) if counter == 0 or REPLAN: print("replanning") ####### Run the MPC ########## # [1, state_dim] n_look_ahead = N - counter if USE_FIXED_MPC_HORIZON: n_look_ahead = MPC_HORIZON elif USE_SHORT_HORIZON_MPC: n_look_ahead = min(MPC_HORIZON, N - counter) if n_look_ahead == 0: break start_idx = counter end_idx = counter + n_look_ahead print("start_idx:", start_idx) print("end_idx:", end_idx) print("n_look_ahead", n_look_ahead) # start_time = time.time() # idx of current observation idx = episode.get_latest_idx() mpc_start_time = time.time() mpc_input_data = mpc_input_builder.get_dynamics_model_input( idx, n_history=n_history) state_cur = mpc_input_data['states'] action_his = mpc_input_data['actions'] if SEED_WITH_NOMINAL_ACTIONS: action_seq_rollout_init = action_seq_gt_torch[ start_idx:end_idx] else: if mpc_out is not None: action_seq_rollout_init = mpc_out['action_seq'][1:] print("action_seq_rollout_init.shape", action_seq_rollout_init.shape) if action_seq_rollout_init.shape[0] < n_look_ahead: num_steps = n_look_ahead - action_seq_rollout_init.shape[ 0] action_seq_zero = torch.zeros([num_steps, 2]) action_seq_rollout_init = torch.cat( (action_seq_rollout_init, action_seq_zero), dim=0) print("action_seq_rollout_init.shape", action_seq_rollout_init.shape) else: action_seq_rollout_init = None # run MPPI z_cur = None with torch.no_grad(): z_cur = model.compute_z_state( state_cur.unsqueeze(0).cuda())['z'].squeeze(0) if action_seq_rollout_init is not None: n_look_ahead = action_seq_rollout_init.shape[0] obs_goal = None print("z_object_rollout_gt.shape", z_object_rollout_gt.shape) if TRAJECTORY_GOAL: obs_goal = z_object_rollout_gt[start_idx:end_idx] print("n_look_ahead", n_look_ahead) print("obs_goal.shape", obs_goal.shape) # add the final goal state on as needed if obs_goal.shape[0] < n_look_ahead: obs_goal_final = z_object_rollout_gt[-1].unsqueeze(0) num_repeat = n_look_ahead - obs_goal.shape[0] obs_goal_final_expand = obs_goal_final.expand( [num_repeat, -1]) obs_goal = torch.cat((obs_goal, obs_goal_final_expand), dim=0) else: obs_goal = z_object_rollout_gt[-1] obs_goal = torch_utils.cast_to_numpy(obs_goal) print("obs_goal.shape", obs_goal.shape) seed = 1 set_seed(seed) mpc_out = planner.trajectory_optimization( state_cur=z_cur, action_his=action_his, obs_goal=obs_goal, model_dy=model, action_seq_rollout_init=action_seq_rollout_init, n_look_ahead=n_look_ahead, eval_indices=object_indices, rollout_best_action_sequence=True, verbose=True, add_grid_action_samples=True, ) print("MPC step took %.4f seconds" % (time.time() - mpc_start_time)) action_seq_mpc = torch_utils.cast_to_numpy(mpc_out['action_seq']) state_pred_mpc = torch_utils.cast_to_numpy(mpc_out['state_pred']) # Rollout with ground truth simulator dynamics env2.set_simulator_state_from_observation_dict( env2.get_mutable_context(), obs_cur) obs_mpc_gt = env_utils.rollout_action_sequence( env2, action_seq_mpc)['observations'] vis['mpc_3D'].delete() vis['mpc_GT_3D'].delete() L = len(obs_mpc_gt) print("L", L) if L == 0: break for i in range(L): # ground truth rollout of the MPC action_seq name = "mpc_GT_3D/%d" % (i) T_W_obj = DrakePusherSliderEnv.object_position_from_observation( obs_mpc_gt[i]) vis[name].set_object(triad(scale=0.1)) vis[name].set_transform(T_W_obj) action_cur = action_seq_mpc[0] print("action_cur", action_cur) print("action_GT", initial_cond['action']) input("press Enter to continue") # add observation actions to the episode obs_cur = env.get_observation() episode.replace_observation_action(obs_cur, action_cur) # step the simulator env.step(action_cur) # update the trajectories, in case we aren't replanning action_seq_mpc = action_seq_mpc[1:] state_pred_mpc = state_pred_mpc[1:] pose_error = compute_pose_error(env_obs_rollout_gt[-1], obs_cur) print("position_error: %.3f" % (pose_error['position_error'])) print("angle error degrees: %.3f" % (pose_error['angle_error_degrees'])) obs_final = env.get_observation() pose_error = compute_pose_error(env_obs_rollout_gt[-1], obs_final) print("position_error: %.3f" % (pose_error['position_error'])) print("angle error degrees: %.3f" % (pose_error['angle_error_degrees']))
import os from key_dynam.utils.utils import get_project_root, get_data_root SIM_ASSETS_ROOT = os.path.join(get_project_root(), 'sim_assets') block_push = os.path.join(SIM_ASSETS_ROOT, 'block_push.urdf') extra_heavy_duty_table = os.path.join(SIM_ASSETS_ROOT, "extra_heavy_duty_table_surface_only_collision.sdf") xy_slide = os.path.join(SIM_ASSETS_ROOT, "xy_slide.urdf") ycb_model_paths = dict({ 'cracker_box': os.path.join(SIM_ASSETS_ROOT, "cracker_box/003_cracker_box.sdf"), 'sugar_box': os.path.join(SIM_ASSETS_ROOT, "sugar_box/004_sugar_box.sdf"), 'tomato_soup_can': os.path.join(SIM_ASSETS_ROOT, "tomato_soup_can/005_tomato_soup_can.sdf"), 'mustard_bottle': os.path.join(SIM_ASSETS_ROOT, "mustard_bottle/006_mustard_bottle.sdf"), 'gelatin_box': os.path.join(SIM_ASSETS_ROOT, "gelatin_box/009_gelatin_box.sdf"), 'potted_meat_can': os.path.join(SIM_ASSETS_ROOT, "potted_meat_can/010_potted_meat_can.sdf") }) ycb_model_baselink_names = dict({ 'cracker_box': 'base_link_cracker', 'sugar_box': 'base_link_sugar', 'tomato_soup_can': 'base_link_soup', 'mustard_bottle': 'base_link_mustard', 'gelatin_box': 'base_link_gelatin', 'potted_meat_can': 'base_link_meat' }) LARGE_SIM_ASSETS_ROOT = os.path.join(get_data_root(), 'stable/sim_assets')
from key_dynam.utils import utils, torch_utils from key_dynam.controller.zmq_utils import ZMQServer from key_dynam.dataset.online_episode_reader import OnlineEpisodeReader from key_dynam.dataset.mpc_dataset import DynamicsModelInputBuilder from key_dynam.controller.plan_container import PlanContainer from key_dynam.dynamics.models_dy import rollout_model, rollout_action_sequences, get_object_and_robot_state_indices from key_dynam.planner import utils as planner_utils from key_dynam.utils import meshcat_utils from key_dynam.experiments.exp_22_push_box_hardware.visualize_dynamics_model import \ visualize_model_prediction_single_timestep # dense correspondence from dense_correspondence_manipulation.utils import meshcat_utils as pdc_meshcat_utils from dense_correspondence_manipulation.utils.constants import DEPTH_IM_SCALE PLAN_MSG_FILE = os.path.join(utils.get_project_root(), 'sandbox', 'plan_msg.p') COMPUTE_CONTROL_ACTION_MSG_FILE = os.path.join(utils.get_project_root(), 'sandbox', 'compute_control_action_msg.p') SAVE_MESSAGES = True class ControllerState(Enum): STOPPED = 0 PLAN_READY = 1 RUNNING = 2 class Controller: """
def get_env_config(): return load_yaml( os.path.join(get_project_root(), 'experiments/exp_29_mugs/config.yaml'))
CUDA_VISIBLE_DEVICES = [0] set_cuda_visible_devices(CUDA_VISIBLE_DEVICES) config = dict() config['train'] = dict() config['train']['train_valid_ratio'] = 0.9 DATASET_NAME = "top_down_rotated" # DATASET_NAME = "2019-12-05-15-58-48-462834_top_down_rotated_250" # # MODEL_NAME = "2019-12-04-01-32-12-010393_top_down_rotated_sigma_5" # MODEL_NAME = "2019-12-26-20-49-35-944397_dataset_2019-12-05-15-58-48-462834_top_down_rotated_250" MODEL_NAME = "2019-12-27-22-29-23-359200_dataset_2019-12-05-15-58-48-462834_top_down_rotated_250" 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())
import os import pydrake import time from dense_correspondence_manipulation.utils.utils import set_cuda_visible_devices set_cuda_visible_devices([1]) from key_dynam.utils.utils import get_project_root, save_yaml, load_yaml, get_current_YYYY_MM_DD_hh_mm_ss_ms, get_data_root from key_dynam.transporter.train_transporter import train_transporter from key_dynam.dense_correspondence.dc_drake_sim_episode_reader import DCDrakeSimEpisodeReader # project root PROJECT_ROOT = get_project_root() # load config config = load_yaml( os.path.join(get_project_root(), "experiments/exp_18_box_on_side/config.yaml")) # specify the dataset dataset_name = "dps_box_on_side_600" config['perception']['dataset_name'] = dataset_name dataset_name = config['perception']['dataset_name'] config['dataset']['set_epoch_size_to_num_images'] = False config['dataset']['epoch_size'] = {'train': 6000, 'valid': 300} # camera_name = "camera_angled" camera_name = "camera_1_top_down"
def main(): # load dynamics model model_dict = load_model_state_dict() model = model_dict['model_dy'] model_dd = model_dict['model_dd'] config = model.config env_config = load_yaml(os.path.join(get_project_root(), 'experiments/exp_20_mugs/config.yaml')) env_config['env']['observation']['depth_int16'] = True n_history = config['train']['n_history'] initial_cond = generate_initial_condition(env_config, push_length=PUSH_LENGTH) env_config = initial_cond['config'] # enable the right observations camera_name = model_dict['metadata']['camera_name'] spatial_descriptor_data = model_dict['spatial_descriptor_data'] ref_descriptors = spatial_descriptor_data['spatial_descriptors'] K = ref_descriptors.shape[0] ref_descriptors = torch.Tensor(ref_descriptors).cuda() # put them on the GPU print("ref_descriptors\n", ref_descriptors) print("ref_descriptors.shape", ref_descriptors.shape) # create the environment # create the environment env = DrakeMugsEnv(env_config) env.reset() T_world_camera = env.camera_pose(camera_name) camera_K_matrix = env.camera_K_matrix(camera_name) # create another environment for doing rollouts env2 = DrakeMugsEnv(env_config, visualize=False) env2.reset() action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.drake_pusher_position_3D(config) visual_observation_function = \ VisualObservationFunctionFactory.descriptor_keypoints_3D(config=config, camera_name=camera_name, model_dd=model_dd, ref_descriptors=ref_descriptors, K_matrix=camera_K_matrix, T_world_camera=T_world_camera, ) episode = OnlineEpisodeReader() mpc_input_builder = DynamicsModelInputBuilder(observation_function=observation_function, visual_observation_function=visual_observation_function, action_function=action_function, episode=episode) vis = meshcat_utils.make_default_visualizer_object() vis.delete() reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) obs_init = env.get_observation() #### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############ reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) # add just some large number of these episode.clear() for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) def goal_func(obs_tmp): state_tmp = mpc_input_builder.get_state_input_single_timestep({'observation': obs_tmp})['state'] return model.compute_z_state(state_tmp.unsqueeze(0))['z_object'].flatten() # idx = episode.get_latest_idx() obs_raw = episode.get_observation(idx) z_object_goal = goal_func(obs_raw) z_keypoints_init_W = keypoints_3D_from_dynamics_model_output(z_object_goal, K) z_keypoints_init_W = torch_utils.cast_to_numpy(z_keypoints_init_W) z_keypoints_obj = keypoints_world_frame_to_object_frame(z_keypoints_init_W, T_W_obj=slider_pose_from_observation(obs_init)) color = [1, 0, 0] meshcat_utils.visualize_points(vis=vis, name="keypoints_W", pts=z_keypoints_init_W, color=color, size=0.02, ) # input("press Enter to continue") # rollout single action sequence using the simulator action_sequence_np = torch_utils.cast_to_numpy(initial_cond['action_sequence']) N = action_sequence_np.shape[0] obs_rollout_gt = env_utils.rollout_action_sequence(env, action_sequence_np)[ 'observations'] # using the vision model to get "goal" keypoints z_object_goal = goal_func(obs_rollout_gt[-1]) z_object_goal_np = torch_utils.cast_to_numpy(z_object_goal) z_keypoints_goal = keypoints_3D_from_dynamics_model_output(z_object_goal, K) z_keypoints_goal = torch_utils.cast_to_numpy(z_keypoints_goal) # visualize goal keypoints color = [0, 1, 0] meshcat_utils.visualize_points(vis=vis, name="goal_keypoints", pts=z_keypoints_goal, color=color, size=0.02, ) # input("press Enter to continue") #### ROLLOUT USING LEARNED MODEL + GROUND TRUTH ACTIONS ############ reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) # add just some large number of these episode.clear() for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) # [n_history, state_dim] idx = episode.get_latest_idx() dyna_net_input = mpc_input_builder.get_dynamics_model_input(idx, n_history=n_history) state_init = dyna_net_input['states'].cuda() # [n_history, state_dim] action_init = dyna_net_input['actions'] # [n_history, action_dim] print("state_init.shape", state_init.shape) print("action_init.shape", action_init.shape) action_seq_gt_torch = torch_utils.cast_to_torch(initial_cond['action_sequence']) action_input = torch.cat((action_init[:(n_history-1)], action_seq_gt_torch), dim=0).cuda() print("action_input.shape", action_input.shape) # rollout using the ground truth actions and learned model # need to add the batch dim to do that z_init = model.compute_z_state(state_init)['z'] rollout_pred = rollout_model(state_init=z_init.unsqueeze(0), action_seq=action_input.unsqueeze(0), dynamics_net=model, compute_debug_data=True) state_pred_rollout = rollout_pred['state_pred'] print("state_pred_rollout.shape", state_pred_rollout.shape) for i in range(N): # vis GT for now name = "GT_3D/%d" % (i) T_W_obj = slider_pose_from_observation(obs_rollout_gt[i]) # print("T_W_obj", T_W_obj) # green color = np.array([0, 1, 0]) * get_color_intensity(i, N) meshcat_utils.visualize_points(vis=vis, name=name, pts=z_keypoints_obj, color=color, size=0.01, T=T_W_obj) # red color = np.array([0, 0, 1]) * get_color_intensity(i, N) state_pred = state_pred_rollout[:, i, :] pts_pred = keypoints_3D_from_dynamics_model_output(state_pred, K).squeeze() pts_pred = pts_pred.detach().cpu().numpy() name = "pred_3D/%d" % (i) meshcat_utils.visualize_points(vis=vis, name=name, pts=pts_pred, color=color, size=0.01, ) # input("finished visualizing GT rollout\npress Enter to continue") index_dict = get_object_and_robot_state_indices(config) object_indices = index_dict['object_indices'] # reset the environment and use the MPC controller to stabilize this # now setup the MPC to try to stabilize this . . . . reset_environment(env, initial_cond['q_pusher'], initial_cond['q_slider']) episode.clear() # add just some large number of these for i in range(n_history): action_zero = np.zeros(2) obs_tmp = env.get_observation() episode.add_observation_action(obs_tmp, action_zero) # input("press Enter to continue") # make a planner config planner_config = copy.copy(config) config_tmp = load_yaml(os.path.join(get_project_root(), 'experiments/drake_pusher_slider/eval_config.yaml')) planner_config['mpc'] = config_tmp['mpc'] planner = None if PLANNER_TYPE == "random_shooting": planner = RandomShootingPlanner(planner_config) elif PLANNER_TYPE == "mppi": planner = PlannerMPPI(planner_config) else: raise ValueError("unknown planner type: %s" % (PLANNER_TYPE)) mpc_out = None action_seq_mpc = None state_pred_mpc = None counter = -1 while True: counter += 1 print("\n\n-----Running MPC Optimization: Counter (%d)-------" % (counter)) obs_cur = env.get_observation() episode.add_observation_only(obs_cur) if counter == 0 or REPLAN: print("replanning") ####### Run the MPC ########## # [1, state_dim] n_look_ahead = N - counter if USE_FIXED_MPC_HORIZON: n_look_ahead = MPC_HORIZON if n_look_ahead == 0: break # start_time = time.time() # idx of current observation idx = episode.get_latest_idx() mpc_start_time = time.time() mpc_input_data = mpc_input_builder.get_dynamics_model_input(idx, n_history=n_history) state_cur = mpc_input_data['states'] action_his = mpc_input_data['actions'] if mpc_out is not None: action_seq_rollout_init = mpc_out['action_seq'][1:] else: action_seq_rollout_init = None # run MPPI z_cur = None with torch.no_grad(): z_cur = model.compute_z_state(state_cur.unsqueeze(0).cuda())['z'].squeeze(0) mpc_out = planner.trajectory_optimization(state_cur=z_cur, action_his=action_his, obs_goal=z_object_goal_np, model_dy=model, action_seq_rollout_init=action_seq_rollout_init, n_look_ahead=n_look_ahead, eval_indices=object_indices, rollout_best_action_sequence=True, verbose=True, ) print("MPC step took %.4f seconds" %(time.time() - mpc_start_time)) action_seq_mpc = mpc_out['action_seq'].cpu().numpy() # Rollout with ground truth simulator dynamics action_seq_mpc = torch_utils.cast_to_numpy(mpc_out['action_seq']) env2.set_simulator_state_from_observation_dict(env2.get_mutable_context(), obs_cur) obs_mpc_gt = env_utils.rollout_action_sequence(env2, action_seq_mpc)['observations'] state_pred_mpc = torch_utils.cast_to_numpy(mpc_out['state_pred']) vis['mpc_3D'].delete() vis['mpc_GT_3D'].delete() L = len(obs_mpc_gt) print("L", L) if L == 0: break for i in range(L): # red color = np.array([1, 0, 0]) * get_color_intensity(i, L) state_pred = state_pred_mpc[i, :] state_pred = np.expand_dims(state_pred, 0) # may need to expand dims here pts_pred = keypoints_3D_from_dynamics_model_output(state_pred, K).squeeze() name = "mpc_3D/%d" % (i) meshcat_utils.visualize_points(vis=vis, name=name, pts=pts_pred, color=color, size=0.01, ) # ground truth rollout of the MPC action_seq name = "mpc_GT_3D/%d" % (i) T_W_obj = slider_pose_from_observation(obs_mpc_gt[i]) # green color = np.array([1, 1, 0]) * get_color_intensity(i, L) meshcat_utils.visualize_points(vis=vis, name=name, pts=z_keypoints_obj, color=color, size=0.01, T=T_W_obj) action_cur = action_seq_mpc[0] print("action_cur", action_cur) # print("action_GT", initial_cond['action']) input("press Enter to continue") # add observation actions to the episode obs_cur = env.get_observation() episode.replace_observation_action(obs_cur, action_cur) # step the simulator env.step(action_cur) # visualize current keypoint positions obs_cur = env.get_observation() T_W_obj = slider_pose_from_observation(obs_cur) # yellow color = np.array([1, 1, 0]) meshcat_utils.visualize_points(vis=vis, name="keypoint_cur", pts=z_keypoints_obj, color=color, size=0.02, T=T_W_obj) action_seq_mpc = action_seq_mpc[1:] state_pred_mpc = state_pred_mpc[1:] obs_final = env.get_observation() pose_error = compute_pose_error(obs_rollout_gt[-1], obs_final) print("position_error: %.3f" %(pose_error['position_error'])) print("angle error degrees: %.3f" %(pose_error['angle_error_degrees']))
def load_model_and_data(): dataset_name = "push_box_hardware" # model_name = "DD_2D/2020-06-24-22-22-58-234812_DD_3D_n_his_2" # this model is actually 3D # model_name = "DD_3D/2020-06-25-00-49-29-679042_DD_3D_n_his_2_T_aug" # model_name = "DD_3D/2020-06-25-00-39-29-020621_DD_3D_n_his_2" model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug" train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics" train_dir = os.path.join(train_dir, model_name) ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth") config = load_yaml(os.path.join(train_dir, 'config.yaml')) state_dict = torch.load(ckpt_file) # build dynamics model model_dy = build_dynamics_model(config) # print("state_dict.keys()", state_dict.keys()) model_dy.load_state_dict(state_dict) model_dy = model_dy.eval() model_dy = model_dy.cuda() spatial_descriptor_data = load_pickle( os.path.join(train_dir, 'spatial_descriptors.p')) metadata = load_pickle(os.path.join(train_dir, 'metadata.p')) # build dense-descriptor model model_dd_file = metadata['model_file'] model_dd = torch.load(model_dd_file) model_dd = model_dd.eval() model_dd = model_dd.cuda() # load the dataset 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') 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) visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config( config, keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx']) action_function = ActionFunctionFactory.function_from_config(config) observation_function = ObservationFunctionFactory.function_from_config( config) dataset = MultiEpisodeDataset( config, action_function=action_function, observation_function=observation_function, episodes=multi_episode_dict, visual_observation_function=visual_observation_function, phase="valid", # this means no data augmentation ) #### PLANNER ####### planner = None # make a planner config planner_config = copy.copy(model_dy.config) config_tmp = load_yaml( os.path.join(get_project_root(), 'experiments/exp_22_push_box_hardware/config_DD_3D.yaml')) planner_config['mpc'] = config_tmp['mpc'] if PLANNER_TYPE == "random_shooting": planner = RandomShootingPlanner(planner_config) elif PLANNER_TYPE == "mppi": planner = PlannerMPPI(planner_config) else: raise ValueError("unknown planner type: %s" % (PLANNER_TYPE)) return { "model_dy": model_dy, 'model_dd': model_dd, 'dataset': dataset, 'config': config, "multi_episode_dict": multi_episode_dict, 'spatial_descriptor_data': spatial_descriptor_data, 'planner': planner, 'observation_function': observation_function, 'action_function': action_function, }