def visualize_autoencoder_result(): d = load_model() dataset = d['dataset'] model = d['model'] model = model.train() model = model.cuda() for i in range(10): idx = np.random.randint(0, len(dataset)) data = dataset[idx] input = data['input_tensor'].unsqueeze(0).cuda() out = model(input) target_tensor = data['target_tensor'].squeeze() target_pred = out['output'].squeeze() keypoints_xy = out['expected_xy'].squeeze() print("keypoints_xy[0]", keypoints_xy[0]) print("target_tensor.shape", target_tensor.shape) print("target_pred.shape", target_pred.shape) print("target_pred.dtype", target_pred.dtype) print("target_pred.max()", target_pred.max()) print("target_pred.min()", target_pred.min()) target_tensor_np = convert_float_image_to_uint8(torch_utils.cast_to_numpy(target_tensor)) target_pred_np = convert_float_image_to_uint8(torch_utils.cast_to_numpy(target_pred)) print(target_pred_np.shape) print(target_tensor_np.shape) # draw reticles on input image H, W, _ = data['input'].shape keypoints_uv = torch_utils.convert_xy_to_uv_coordinates(keypoints_xy, H=H, W=W) input_wr = np.copy(data['input']) draw_reticles(input_wr, u_vec=keypoints_uv[:, 0], v_vec=keypoints_uv[:, 1], label_color=[0,255,0]) print("type(input_wr)", type(input_wr)) figsize = 2*np.array([4.8,6.4]) fig = plt.figure(figsize=figsize) ax = fig.subplots(1,3) ax[0].imshow(input_wr) ax[1].imshow(target_tensor_np, cmap='gray', vmin=0, vmax=255) ax[2].imshow(target_pred_np, cmap='gray', vmin=0, vmax=255) plt.show()
def visualize_model_prediction_single_timestep( vis, config, z_pred, # [z_dim] display_idx, name_prefix=None, color=None, display_robot_state=True, ): if color is None: color = [0, 0, 255] idx_dict = get_object_and_robot_state_indices(config) z_object = z_pred[idx_dict['object_indices']].reshape( config['dataset']['object_state_shape']) z_robot = z_pred[idx_dict['robot_indices']].reshape( config['dataset']['robot_state_shape']) name = "z_object/%d" % (display_idx) if name_prefix is not None: name = name_prefix + "/" + name meshcat_utils.visualize_points( vis, name, torch_utils.cast_to_numpy(z_object), color=color, size=0.01, ) if display_robot_state: name = "z_robot/%d" % (display_idx) if name_prefix is not None: name = name_prefix + "/" + name meshcat_utils.visualize_points( vis, name, torch_utils.cast_to_numpy(z_robot), color=color, size=0.01, )
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 visualize_episode_data_single_timestep( vis, dataset, camera_name, episode, episode_idx, display_idx, ): image_episode_idx = episode.image_episode_idx_from_query_idx(episode_idx) image_data = episode.image_episode.get_image_data(camera_name=camera_name, idx=image_episode_idx) depth = image_data['depth_int16'] / DEPTH_IM_SCALE # pointcloud name = "pointclouds/%d" % (display_idx) pdc_meshcat_utils.visualize_pointcloud( vis, name, depth=depth, K=image_data['K'], rgb=image_data['rgb'], T_world_camera=image_data['T_world_camera']) data = dataset._getitem(episode, episode_idx, n_history=1, rollout_length=0) # action position # name = "actions/%d" %(display_idx) # actions = torch_utils.cast_to_numpy(data['actions']) # # meshcat_utils.visualize_points(vis, # name, # actions, # color=[255,0,0], # size=0.01, # ) # observation position name = "observations/%d" % (display_idx) observations = torch_utils.cast_to_numpy(data['observations']) # print("observations.shape", observations.shape) # print("observations", observations) meshcat_utils.visualize_points( vis, name, observations, color=[0, 255, 0], size=0.01, ) # keypoints if True: name = "keypoints/%d" % (display_idx) keypoints = torch_utils.cast_to_numpy( data['visual_observation_func_collated']['keypoints_3d'][0]) # print("keypoints.shape", keypoints.shape) meshcat_utils.visualize_points( vis, name, keypoints, color=[0, 255, 0], size=0.01, )
def evaluate_mpc( model_dy, # dynamics model env, # the environment episode, # OnlineEpisodeReader mpc_input_builder, # DynamicsModelInputBuilder planner, # RandomShooting planner eval_indices=None, goal_func=None, # function that gets goal from observation config=None, wait_for_user_input=False, save_dir=None, model_name="", experiment_name="", generate_initial_condition_func=None, # (optional) function to generate initial condition, takes episode length N as parameter ): if not os.path.exists(save_dir): os.makedirs(save_dir) # must specify initial condition distribution assert generate_initial_condition_func is not None save_yaml(config, os.path.join(save_dir, 'config.yaml')) writer = SummaryWriter(log_dir=save_dir) pandas_data_list = [] for episode_length in config['eval']['episode_length']: counter = 0 seed = 0 while counter < config['eval']['num_episodes']: start_time = time.time() seed += 1 set_seed(seed) # make it repeatable # initial_cond = generate_initial_condition(config, N=episode_length) initial_cond = generate_initial_condition_func(N=episode_length) env.set_initial_condition_from_dict(initial_cond) action_sequence_np = torch_utils.cast_to_numpy( initial_cond['action_sequence']) episode_data = mpc_single_episode( model_dy=model_dy, env=env, action_sequence=action_sequence_np, action_zero=np.zeros(2), episode=episode, mpc_input_builder=mpc_input_builder, planner=planner, eval_indices=eval_indices, goal_func=goal_func, config=config, wait_for_user_input=wait_for_user_input, ) # continue if invalid if not episode_data['valid']: print("invalid episode, skipping") continue pose_error = compute_pose_error( obs=episode_data['obs_mpc_final'], obs_goal=episode_data['obs_goal'], ) object_delta = compute_pose_error( obs=episode_data['obs_init'], obs_goal=episode_data['obs_goal']) print("object_delta\n", object_delta) if wait_for_user_input: print("pose_error\n", pose_error) pandas_data = { 'episode_length': episode_length, 'seed': counter, 'model_name': model_name, 'experiment_name': experiment_name, 'object_pos_delta': object_delta['position_error'], 'object_angle_delta': object_delta['angle_error'], 'object_angle_delta_degrees': object_delta['angle_error_degrees'], } pandas_data.update(pose_error) pandas_data_list.append(pandas_data) # log to tensorboard for key, val in pose_error.items(): plot_name = "%s/episode_len_%d" % (key, episode_length) writer.add_scalar(plot_name, val, counter) writer.flush() print("episode [%d/%d], episode_length %d, duration %.2f" % (counter, config['eval']['num_episodes'], episode_length, time.time() - start_time)) counter += 1 df_tmp = pd.DataFrame(pandas_data_list) keys = ["angle_error_degrees", "position_error"] for key in keys: for i in range(10): mean = df_tmp[key][df_tmp.episode_length == episode_length].mean() median = df_tmp[key][df_tmp.episode_length == episode_length].median() plot_name_mean = "mean/%s/episode_len_%d" % (key, episode_length) writer.add_scalar(plot_name_mean, mean, i) plot_name_median = "median/%s/episode_len_%d" % ( key, episode_length) writer.add_scalar(plot_name_median, median, i) # save some data df = pd.DataFrame(pandas_data_list) df.to_csv(os.path.join(save_dir, "data.csv"))
def _on_compute_control_action( self, msg, visualize=True, ): """ Computes the current control action :param msg: :type msg: :return: :rtype: """ print("\n\n----------") start_time = time.time() assert self._state in [ ControllerState.PLAN_READY, ControllerState.RUNNING ] # allow setting the visualize flag from the message try: visualize = bool(int(msg['debug'])) except KeyError: pass # add data to the OnlineEpisodeReader episode = self._state_dict['episode'] input_builder = self._state_dict['input_builder'] plan = self._state_dict['plan'] observation = msg['data']['observations'] action_dict = msg['data']['actions'] if self._state == ControllerState.PLAN_READY: # this is the first time through the loop # then add a few to populate n_his for i in range(self._n_history): # episode.add_observation_action(copy.deepcopy(observation), copy.deepcopy(action_dict)) episode.add_data(copy.deepcopy(msg['data'])) episode.add_data(msg['data']) # run the planner # seed with previous actions mpc_out = self._state_dict['mpc_out'] n_look_ahead = None # this is the MPC horizon mpc_horizon_type = self._planner.config['mpc']['hardware'][ 'mpc_horizon']["type"] mpc_hardware_config = self._planner.config['mpc']['hardware'] # previous actions to seed with # compute the MPC horizon action_seq_rollout_init = None if mpc_out is not None: # this means it's not our first time through the loop # the plan is already running action_seq_rollout_init = mpc_out['action_seq'][1:] if mpc_horizon_type == "PLAN_LENGTH": n_look_ahead = action_seq_rollout_init.shape[0] # this means we are at the end of the plan # so send the stop message if n_look_ahead == 0: print("Plan finished, sending STOP message") return self.make_stop_message() if mpc_horizon_type == "MIN_HORIZON": n_look_ahead = action_seq_rollout_init.shape[0] H_min = mpc_hardware_config['mpc_horizon']['H_min'] n_look_ahead = max(H_min, n_look_ahead) # maybe set overwrite the config to be traj cost . . . elif mpc_horizon_type == "FIXED": n_look_ahead = mpc_hardware_config['mpc_horizon']["H_fixed"] else: raise ValueError("unknow mpc_horizon_type") # add zeros to the end of the action trajectory as a good # starting point # extend the previous action, either with zeros or by repeating last action if action_seq_rollout_init.shape[0] < n_look_ahead: if mpc_hardware_config['action_extension_type'] == "CONSTANT": num_steps = n_look_ahead - action_seq_rollout_init.shape[0] # [action_dim] action_extend = action_seq_rollout_init[-1] action_extend = action_extend.unsqueeze(0).expand( [num_steps, -1]) # print("action_seq_rollout_init.shape", action_seq_rollout_init.shape) # print("action_extend.shape", action_extend.shape) action_seq_rollout_init = torch.cat( (action_seq_rollout_init, action_extend), dim=0) elif mpc_hardware_config['action_extension_type'] == "ZERO": num_steps = n_look_ahead - action_seq_rollout_init.shape[0] action_seq_zero = torch.zeros([num_steps, 2]).to( action_seq_rollout_init.device) action_seq_rollout_init = torch.cat( (action_seq_rollout_init, action_seq_zero), dim=0) else: if mpc_horizon_type == "FIXED": n_look_ahead = mpc_hardware_config['mpc_horizon']["H_fixed"] else: n_look_ahead = mpc_hardware_config['mpc_horizon']["H_init"] action_seq_rollout_init = None print("n_look_ahead", n_look_ahead) print("plan_counter", plan.counter) start_time_tmp = time.time() idx = episode.get_latest_idx() mpc_input_data = input_builder.get_dynamics_model_input( idx, n_history=self._n_history) print("computing dynamics model input took", time.time() - start_time_tmp) start_time_tmp = time.time() state_cur = mpc_input_data['states'] action_his = mpc_input_data['actions'] current_reward_data = None # run the planner with torch.no_grad(): # z_goal_dict = plan.data[-1]['dynamics_model_input_data']['z'] # obs_goal = plan.data[-1]['dynamics_model_input_data']['z']['z_object_flat'] # convert state_cur to z_cur z_cur_dict = self._model_dy.compute_z_state(state_cur) z_cur = z_cur_dict['z'] z_cur_no_his = z_cur[-1] print("z_cur.shape") # for computing current cost z_batch = None # for now we support just final state rather than trajectory costs z_goal = None if self._planner.config['mpc']['reward'][ "goal_type"] == "TRAJECTORY": z_goal = plan.get_trajectory_goal( counter=self._state_dict['action_counter'], n_look_ahead=n_look_ahead) # [1, 1, z_dim] z_batch = z_cur_no_his.unsqueeze(0).unsqueeze(0) # [1, n_look_aheda, z_dim] z_batch = z_batch.expand([-1, n_look_ahead, -1]) elif self._planner.config['mpc']['reward'][ "goal_type"] == "FINAL_STATE": z_goal = plan.get_final_state_goal() # [1, 1, z_dim] z_batch = z_cur_no_his.unsqueeze(0).unsqueeze(0) else: raise ValueError("unknown goal type") print("z_batch.shape", z_batch.shape) obs_goal = torch.index_select(z_goal, dim=-1, index=self._object_indices) print("obs_goal.shape", obs_goal.shape) # compute the cost of stopping now so we can compare it to the cost # of running the trajectory optimization # note this is on the CPU # [1, 1, z_dim] z_batch = z_cur_no_his.unsqueeze(0).unsqueeze(0) # [goal_dim] obs_goal_final = torch.index_select(plan.get_final_state_goal(), dim=-1, index=self._object_indices) current_reward_data = \ planner_utils.evaluate_model_rollout(state_pred=z_batch, obs_goal=obs_goal_final, eval_indices=self._object_indices, **self._planner.config['mpc']['reward']) print("current reward:", current_reward_data['best_reward']) if current_reward_data['best_reward'] > mpc_hardware_config[ 'goal_reward_threshold']: print("Below goal reward threshold, stopping") return self.make_stop_message() # run the planner mpc_out = self._planner.trajectory_optimization( state_cur=z_cur, action_his=action_his, obs_goal=obs_goal, model_dy=self._model_dy, action_seq_rollout_init=action_seq_rollout_init, n_look_ahead=n_look_ahead, eval_indices=self._object_indices, rollout_best_action_sequence=self._debug, verbose=self._debug, add_grid_action_samples=False, ) # update the action that was actually applied # equivalent to action_seq_mpc = mpc_out['action_seq'] action_dict['ee_setpoint']['setpoint_linear_velocity']['x'] = float( action_seq_mpc[0][0]) action_dict['ee_setpoint']['setpoint_linear_velocity']['x'] = float( action_seq_mpc[0][1]) state_pred = mpc_out['state_pred'] if mpc_out['reward'].cpu() < current_reward_data['best_reward'].cpu( ) + mpc_hardware_config['reward_improvement_tol']: if mpc_hardware_config['terminate_if_no_improvement']: print("Traj opt didn't yield successive improvement, STOPPING") return self.make_stop_message() else: print( "Traj opt didn't yield successive improvement, HOLDING STILL" ) return self.make_zero_action_message() if self._debug: # pass print("action_seq_mpc\n", action_seq_mpc) if visualize: vis = meshcat_utils.make_default_visualizer_object() vis["mpc"].delete() visualize_model_prediction_single_timestep(vis, self._config, z_cur[0], display_idx=0, name_prefix="start", color=[0, 0, 255]) # visualize pointcloud of current position start_data = episode.get_data(0) depth = start_data['observations']['images'][self._camera_info[ 'camera_name']]['depth_int16'] / DEPTH_IM_SCALE rgb = start_data['observations']['images'][ self._camera_info['camera_name']]['rgb'] # pointcloud name = "mpc/pointclouds/start" pdc_meshcat_utils.visualize_pointcloud( vis, name, depth=depth, K=self._camera_info['K'], rgb=rgb, T_world_camera=self._camera_info['T_world_camera']) # # visualize_model_prediction_single_timestep(vis, # self._config, # z_goal_dict['z'], # display_idx=0, # name_prefix="goal", # color=[0, 255, 0]) # # goal_data = plan.data[-1] # depth = goal_data['observations']['images'][self._camera_info['camera_name']][ # 'depth_int16'] / DEPTH_IM_SCALE # rgb = goal_data['observations']['images'][self._camera_info['camera_name']]['rgb'] # # pointcloud # name = "pointclouds/goal" # pdc_meshcat_utils.visualize_pointcloud(vis, # name, # depth=depth, # K=self._camera_info['K'], # rgb=rgb, # T_world_camera=self._camera_info['T_world_camera']) for i in range(state_pred.shape[0]): visualize_model_prediction_single_timestep( vis, self._config, state_pred[i], display_idx=(i + 1), name_prefix="mpc", color=[255, 0, 0], ) # show pointclouds . . . # store the results of this computation self._state_dict['action_counter'] += 1 self._state_dict['timestamp_system'].append(time.time()) self._state_dict['mpc_out'] = mpc_out self._state = ControllerState.RUNNING # or maybe also STOPPED/FINISHED # data we want to save out from the MPC mpc_save_data = { 'n_look_ahead': n_look_ahead, 'action_seq': mpc_out['action_seq'].cpu(), 'reward': mpc_out['reward'].cpu(), 'state_pred': mpc_out['state_pred'].cpu(), 'current_reward': current_reward_data, } # add some data for saving out later idx = episode.get_latest_idx() episode_data = episode.get_data(idx) episode_data['mpc'] = mpc_save_data plan.increment_counter() action_seq_mpc_np = torch_utils.cast_to_numpy(action_seq_mpc) # todo: update the action in the episode reader. Currently # we have an off by one error if False: action = action_seq_mpc[0] idx = episode.get_latest_idx() episode_data = episode.get_data(idx) episode_data['actions']['mpc'] # would be pretty hacky . . . . episode_data['actions']['setpoint_linear_velocity']['x'] = action[ 0] episode_data['actions']['setpoint_linear_velocity']['y'] = action[ 1] resp_data = { 'action': action_seq_mpc_np[0], 'action_seq': action_seq_mpc_np, } resp = {'type': 'CONTROL_ACTION', 'data': resp_data} # we are seeing averages around 0.12 seconds (total) for length 10 plan print("Compute Control Action Took %.3f" % (time.time() - start_time)) return resp
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']))
def localize_transporter_keypoints( model_kp, rgb=None, # np.array [H, W, 3] mask=None, # np.array [H, W] depth=None, # np.array [H, W] in meters K=None, # camera intrinsics matrix [3,3] T_world_camera=None, ): processed_image_dict = process_image(rgb=rgb, config=model_kp.config, mask=mask) rgb_input = processed_image_dict['image'] crop_param = processed_image_dict['crop_param'] H, W, _ = rgb.shape # cast them to torch so we can use them later crop_param_torch = None if crop_param is not None: crop_param_torch = dict() for key, val in crop_param.items(): # cast to torch and add batch dim crop_param_torch[key] = torch.Tensor([val]).unsqueeze(0) rgb_to_tensor = ImageTupleDataset.make_rgb_image_to_tensor_transform() rgb_input_tensor = rgb_to_tensor(rgb_input).unsqueeze( 0).cuda() # make it batch # [1, n_kp, 2] kp_pred = model_kp.predict_keypoint(rgb_input_tensor).cpu() # if it was cropped, then an extra step is needed # kp_pred_full_pixels = None # if crop_param is not None: # kp_pred_full_pixels = map_cropped_pixels_to_full_pixels_torch(kp_pred, # crop_param_torch) # else: # kp_pred_full_pixels = kp_pred # # xy = kp_pred_full_pixels.clone() # xy[:, :, 0] = (xy[:, :, 0]) * 2.0 / W - 1.0 # xy[:, :, 1] = (xy[:, :, 1]) * 2.0 / H - 1.0 # # # get depth values # kp_pred_full_pixels_int = kp_pred_full_pixels.type(torch.LongTensor) full_image_coords = map_transporter_keypoints_to_full_image( kp_pred, crop_params=crop_param_torch, full_image_size=(H, W)) uv = full_image_coords['uv'] uv_int = full_image_coords['uv_int'] xy = full_image_coords['xy'] depth_batch = torch.from_numpy(depth).unsqueeze(0) z = None pts_world_frame = None pts_camera_frame = None if depth is not None: z = pdc_utils.index_into_batch_image_tensor(depth_batch.unsqueeze(1), uv_int.transpose(1, 2)) z = z.squeeze(1) K_inv = np.linalg.inv(K) K_inv_torch = torch.Tensor(K_inv).unsqueeze(0) # add batch dim pts_camera_frame = pdc_torch_utils.pinhole_unprojection( uv, z, K_inv_torch) # print("pts_camera_frame.shape", pts_camera_frame.shape) pts_world_frame = pdc_torch_utils.transform_points_3D( torch.from_numpy(T_world_camera), pts_camera_frame) pts_world_frame_np = torch_utils.cast_to_numpy( pts_world_frame.squeeze()) pts_camera_frame_np = torch_utils.cast_to_numpy( pts_camera_frame.squeeze()) uv_np = torch_utils.cast_to_numpy(uv.squeeze()) uv_int_np = torch_utils.cast_to_numpy(uv_int.squeeze()) return { 'uv': uv_np, 'uv_int': uv_int_np, 'xy': torch_utils.cast_to_numpy(xy.squeeze()), 'z': torch_utils.cast_to_numpy(z.squeeze()), 'pts_world_frame': pts_world_frame_np, 'pts_camera_frame': pts_camera_frame_np, 'pts_W': pts_world_frame_np, # just for backwards compatibility 'kp_pred': kp_pred, 'rgb_input': rgb_input, }
def precompute_transporter_keypoints( multi_episode_dict, model_kp, output_dir, # str batch_size=10, num_workers=10, camera_names=None, model_file=None, ): assert model_file is not None metadata = dict() metadata['model_file'] = model_file save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml')) start_time = time.time() log_freq = 10 device = next(model_kp.parameters()).device model_kp = model_kp.eval() # make sure model is in eval mode image_data_config = { 'rgb': True, 'mask': True, 'depth_int16': True, } # build all the dataset datasets = {} dataloaders = {} for episode_name, episode in multi_episode_dict.items(): single_episode_dict = {episode_name: episode} config = model_kp.config # need to do this since transporter type data sampling only works # with tuple_size = 1 dataset_config = copy.deepcopy(config) dataset_config['dataset']['use_transporter_type_data_sampling'] = False datasets[episode_name] = ImageTupleDataset( dataset_config, single_episode_dict, phase="all", image_data_config=image_data_config, tuple_size=1, compute_K_inv=True, camera_names=camera_names) dataloaders[episode_name] = DataLoader(datasets[episode_name], batch_size=batch_size, num_workers=num_workers, shuffle=False) episode_counter = 0 num_episodes = len(multi_episode_dict) for episode_name, dataset in datasets.items(): episode_counter += 1 print("\n\n") episode = multi_episode_dict[episode_name] hdf5_file = None try: hdf5_file = os.path.basename(episode.image_data_file) except AttributeError: hdf5_file = "%s.h5" % (episode.name) hdf5_file_fullpath = os.path.join(output_dir, hdf5_file) str_split = hdf5_file_fullpath.split(".") assert len(str_split) == 2 pickle_file_fullpath = str_split[0] + ".p" # print("episode_name", episode_name) # print("hdf5_file_fullpath", hdf5_file_fullpath) # print("pickle_file_fullpath", pickle_file_fullpath) if os.path.isfile(hdf5_file_fullpath): os.remove(hdf5_file_fullpath) if os.path.isfile(pickle_file_fullpath): os.remove(pickle_file_fullpath) episode_keypoint_data = dict() episode_start_time = time.time() with h5py.File(hdf5_file_fullpath, 'w') as hf: for i, data in enumerate(dataloaders[episode_name]): data = data[0] rgb_crop_tensor = data['rgb_crop_tensor'].to(device) crop_params = data['crop_param'] depth_int16 = data['depth_int16'] key_tree_joined = data['key_tree_joined'] # print("\n\n i = %d, idx = %d, camera_name = %s" %(i, data['idx'], data['camera_name'])) depth = depth_int16.float() * 1.0 / DEPTH_IM_SCALE if (i % log_freq) == 0: log_msg = "computing [%d/%d][%d/%d]" % ( episode_counter, num_episodes, i + 1, len(dataloaders[episode_name])) print(log_msg) B = rgb_crop_tensor.shape[0] _, H, W, _ = data['rgb'].shape kp_pred = None kp_pred_full_pixels = None with torch.no_grad(): kp_pred = model_kp.predict_keypoint(rgb_crop_tensor) # [B, n_kp, 2] kp_pred_full_pixels = transporter_utils.map_cropped_pixels_to_full_pixels_torch( kp_pred, crop_params) xy = kp_pred_full_pixels.clone() xy[:, :, 0] = (xy[:, :, 0]) * 2.0 / W - 1.0 xy[:, :, 1] = (xy[:, :, 1]) * 2.0 / H - 1.0 # debug # print("xy[0,0]", xy[0,0]) # get depth values kp_pred_full_pixels_int = kp_pred_full_pixels.type( torch.LongTensor) z = pdc_utils.index_into_batch_image_tensor( depth.unsqueeze(1), kp_pred_full_pixels_int.transpose(1, 2)) z = z.squeeze(1) K_inv = data['K_inv'] pts_camera_frame = pdc_torch_utils.pinhole_unprojection( kp_pred_full_pixels, z, K_inv) # print("pts_camera_frame.shape", pts_camera_frame.shape) pts_world_frame = pdc_torch_utils.transform_points_3D( data['T_W_C'], pts_camera_frame) # print("pts_world_frame.shape", pts_world_frame.shape) for j in range(B): keypoint_data = {} # this goes from [-1,1] keypoint_data['xy'] = torch_utils.cast_to_numpy(xy[j]) keypoint_data['uv'] = torch_utils.cast_to_numpy( kp_pred_full_pixels[j]) keypoint_data['uv_int'] = torch_utils.cast_to_numpy( kp_pred_full_pixels_int[j]) keypoint_data['z'] = torch_utils.cast_to_numpy(z[j]) keypoint_data[ 'pos_world_frame'] = torch_utils.cast_to_numpy( pts_world_frame[j]) keypoint_data[ 'pos_camera_frame'] = torch_utils.cast_to_numpy( pts_camera_frame[j]) # save out some data in both hdf5 and pickle format for key, val in keypoint_data.items(): save_key = key_tree_joined[ j] + "/transporter_keypoints/%s" % (key) hf.create_dataset(save_key, data=val) episode_keypoint_data[save_key] = val save_pickle(episode_keypoint_data, pickle_file_fullpath) print("duration: %.3f seconds" % (time.time() - episode_start_time))
def precompute_descriptor_keypoints(multi_episode_dict, model, output_dir, # str ref_descriptors_metadata, batch_size=10, num_workers=10, localization_type="spatial_expectation", # ['spatial_expectation', 'argmax'] compute_3D=True, # in world frame camera_names=None, ): if not os.path.isdir(output_dir): os.makedirs(output_dir) start_time = time.time() log_freq = 10 device = next(model.parameters()).device model = model.eval() # make sure model is in eval mode # build all the dataset datasets = {} dataloaders = {} for episode_name, episode in iteritems(multi_episode_dict): single_episode_dict = {episode_name: episode} config = None datasets[episode_name] = ImageDataset(config, single_episode_dict, phase="all", camera_names=camera_names) dataloaders[episode_name] = DataLoader(datasets[episode_name], batch_size=batch_size, num_workers=num_workers, shuffle=False) # K = num_ref_descriptors metadata = ref_descriptors_metadata ref_descriptors = torch.Tensor(metadata['ref_descriptors']) ref_descriptors = ref_descriptors.cuda() K, _ = ref_descriptors.shape metadata_file = os.path.join(output_dir, 'metadata.p') save_pickle(metadata, metadata_file) episode_counter = 0 num_episodes = len(multi_episode_dict) for episode_name, dataset in iteritems(datasets): episode_counter += 1 print("\n\n") episode = multi_episode_dict[episode_name] hdf5_file = None try: hdf5_file = os.path.basename(episode.image_data_file) except AttributeError: hdf5_file = "%s.h5" % (episode.name) hdf5_file_fullpath = os.path.join(output_dir, hdf5_file) str_split = hdf5_file_fullpath.split(".") assert len(str_split) == 2 pickle_file_fullpath = str_split[0] + ".p" # print("hdf5_file_fullpath", hdf5_file_fullpath) # print("pickle_file_fullpath", pickle_file_fullpath) if os.path.isfile(hdf5_file_fullpath): os.remove(hdf5_file_fullpath) if os.path.isfile(pickle_file_fullpath): os.remove(pickle_file_fullpath) dataloader = dataloaders[episode_name] episode_keypoint_data = dict() episode_start_time = time.time() with h5py.File(hdf5_file_fullpath, 'w') as hf: for i, data in enumerate(dataloaders[episode_name]): rgb_tensor = data['rgb_tensor'].to(device) key_tree_joined = data['key_tree_joined'] if (i % log_freq) == 0: log_msg = "computing [%d/%d][%d/%d]" % (episode_counter, num_episodes, i + 1, len(dataloader)) print(log_msg) # don't use gradients tmp_time = time.time() with torch.no_grad(): out = model.forward(rgb_tensor) # [B, D, H, W] des_img = out['descriptor_image'] B, _, H, W = rgb_tensor.shape # [B, N, 2] batch_indices = None preds_3d = None if localization_type == "spatial_expectation": sigma_descriptor_heatmap = 5 # default try: sigma_descriptor_heatmap = model.config['network']['sigma_descriptor_heatmap'] except: pass # print("ref_descriptors.shape", ref_descriptors.shape) # print("des_img.shape", des_img.shape) d = get_spatial_expectation(ref_descriptors, des_img, sigma=sigma_descriptor_heatmap, type='exp', return_heatmap=True, ) batch_indices = d['uv'] # [B, K, H, W] if compute_3D: # [B*K, H, W] heatmaps_no_batch = d['heatmap_no_batch'] # [B, H, W] depth = data['depth_int16'].to(device) # expand depth images and convert to meters, instead of mm # [B, K, H, W] depth_expand = depth.unsqueeze(1).expand([B, K, H, W]).reshape([B * K, H, W]) depth_expand = depth_expand.type(torch.FloatTensor) / constants.DEPTH_IM_SCALE depth_expand = depth_expand.to(heatmaps_no_batch.device) pred_3d = get_integral_preds_3d(heatmaps_no_batch, depth_images=depth_expand, compute_uv=True) pred_3d['uv'] = pred_3d['uv'].reshape([B, K, 2]) pred_3d['xy'] = pred_3d['xy'].reshape([B, K, 2]) pred_3d['z'] = pred_3d['z'].reshape([B, K]) elif localization_type == "argmax": # localize descriptors best_match_dict = get_argmax_l2(ref_descriptors, des_img) # [B, N, 2] # where N is num_ref_descriptors batch_indices = best_match_dict['indices'] else: raise ValueError("unknown localization type: %s" % (localization_type)) print("computing keypoints took", time.time() - tmp_time) tmp_time = time.time() # iterate over elements in the batch for j in range(B): keypoint_data = {} # dict that stores information to save out # [N,2] # indices = batch_indices[j].cpu().numpy() # key = key_tree_joined[j] + "/descriptor_keypoints" # hf.create_dataset(key, data=indices) # keypoint_indices_dict[key] = indices # stored 3D keypoint locations (in both camera and world frame) if pred_3d is not None: # key_3d_W = key_tree_joined[j] + "/descriptor_keypoints_3d_world_frame" # key_3d_C = key_tree_joined[j] + "/descriptor_keypoints_3d_camera_frame" # T_W_C = data['T_world_camera'][j].cpu().numpy() # K_matrix = data['K'][j].cpu().numpy() T_W_C = torch_utils.cast_to_numpy(data['T_world_camera'][j]) K_matrix = torch_utils.cast_to_numpy(data['K'][j]) uv = torch_utils.cast_to_numpy(pred_3d['uv'][j]) xy = torch_utils.cast_to_numpy(pred_3d['xy'][j]) z = torch_utils.cast_to_numpy(pred_3d['z'][j]) # [K, 3] # this is in camera frame pts_3d_C = pdc_utils.pinhole_unprojection(uv, z, K_matrix) # hf.create_dataset(key_3d_C, data=pts_3d_C) # keypoint_indices_dict[key_3d_C] = pts_3d_C # project into world frame pts_3d_W = transform_utils.transform_points_3D(transform=T_W_C, points=pts_3d_C) # hf.create_dataset(key_3d_W, data=pts_3d_W) # keypoint_indices_dict[key_3d_W] = pts_3d_W keypoint_data['xy'] = torch_utils.cast_to_numpy(xy) keypoint_data['uv'] = torch_utils.cast_to_numpy(uv) keypoint_data['z'] = torch_utils.cast_to_numpy(z) keypoint_data['pos_world_frame'] = torch_utils.cast_to_numpy(pts_3d_W) keypoint_data['pos_camera_frame'] = torch_utils.cast_to_numpy(pts_3d_C) # save out some data in both hdf5 and pickle format for key, val in keypoint_data.items(): save_key = key_tree_joined[j] + "/descriptor_keypoints/%s" % (key) hf.create_dataset(save_key, data=val) episode_keypoint_data[save_key] = val print("saving to disk took", time.time() - tmp_time) # save_pickle(keypoint_indices_dict, pickle_file_fullpath) save_pickle(episode_keypoint_data, pickle_file_fullpath) print("duration: %.3f seconds" % (time.time() - episode_start_time)) print("total time to compute descriptors: %.3f seconds" % (time.time() - start_time))