def viz_func(batch, predictions, test_dataset: DynamicsDatasetLoader): """ we assume batch size of 1 """ test_dataset.scenario.plot_environment_rviz(remove_batch(batch)) anim = RvizAnimationController(np.arange(test_dataset.steps_per_traj)) while not anim.done: t = anim.t() actual_t = numpify( remove_batch( test_dataset.scenario.index_time_batched_predicted(batch, t))) action_t = numpify( remove_batch( test_dataset.scenario.index_time_batched_predicted(batch, t))) test_dataset.scenario.plot_state_rviz(actual_t, label='actual', color='red') test_dataset.scenario.plot_action_rviz(actual_t, action_t, color='gray') prediction_t = remove_batch( test_dataset.scenario.index_time_batched_predicted(predictions, t)) test_dataset.scenario.plot_state_rviz(prediction_t, label='predicted', color='blue') anim.step()
def debug_rviz(self, input_dict: Dict, debug_info_seq: List[Tuple]): import pickle from merrrt_visualization.rviz_animation_controller import RvizSimpleStepper from link_bot_pycommon.bbox_visualization import grid_to_bbox from moonshine.moonshine_utils import numpify from moonshine.indexing import index_dict_of_batched_tensors_tf, index_time_with_metadata from link_bot_pycommon.grid_utils import environment_to_occupancy_msg, send_occupancy_tf import numpy as np from time import time debug_filename = f'debug_{int(time())}.pkl' print(f"Saving debug info to {debug_filename}") with open(debug_filename, "wb") as debug_file: pickle.dump({'input_dict': input_dict, 'debug_info': debug_info_seq}, debug_file) stepper = RvizSimpleStepper() batch_size = input_dict.pop("batch_size").numpy().astype(np.int32) input_dict.pop("time") for b in range(batch_size): example = index_dict_of_batched_tensors_tf(input_dict, b) for t, debug_info_t in enumerate(debug_info_seq): state_t, local_env_origin_t, local_env_t, local_voxel_grid_t = debug_info_t for i, state_component_k_voxel_grid in enumerate(tf.transpose(local_voxel_grid_t, [4, 0, 1, 2, 3])): raster_dict = { 'env': tf.clip_by_value(state_component_k_voxel_grid[b], 0, 1), 'origin': local_env_origin_t[b].numpy(), 'res': input_dict['res'][b].numpy(), } raster_msg = environment_to_occupancy_msg(raster_dict, frame='local_occupancy') self.raster_debug_pubs[i].publish(raster_msg) local_env_dict = { 'env': local_env_t[b], 'origin': local_env_origin_t[b].numpy(), 'res': input_dict['res'][b].numpy(), } send_occupancy_tf(self.scenario.tf.tf_broadcaster, local_env_dict, frame='local_occupancy') pred_t = index_time_with_metadata({}, example, self.pred_state_keys, t) action_t = {k: example[k][t] for k in self.action_keys} self.scenario.plot_state_rviz(numpify(pred_t), label='predicted', color='#0000ffff') if action_t is not None: self.scenario.plot_action_rviz(numpify(pred_t), numpify(action_t), label='action', color='#0000ffff') # # Ground-Truth # true_t = index_time_with_metadata({}, example, self.true_state_keys, t) # self.scenario.plot_state_rviz(numpify(true_t), label='actual', color='#ff0000ff', scale=1.1) # label_t = example['is_close'][1] # self.scenario.plot_is_close(label_t) bbox_msg = grid_to_bbox(rows=self.local_env_h_rows, cols=self.local_env_w_cols, channels=self.local_env_c_channels, resolution=example['res'].numpy()) bbox_msg.header.frame_id = 'local_occupancy' self.local_env_bbox_pub.publish(bbox_msg) stepper.step()
def plan_internal(self, planning_query: PlanningQuery) -> PlanningResult: start_planning_time = perf_counter() action_rng = np.random.RandomState(planning_query.seed) random_actions = self.scenario.sample_action_sequences(environment=planning_query.environment, state=planning_query.start, action_params=self.action_params, n_action_sequences=self.n_samples, action_sequence_length=1, validate=True, action_rng=action_rng) random_actions = [sequence_of_dicts_to_dict_of_tensors(a) for a in random_actions] random_actions = sequence_of_dicts_to_dict_of_tensors(random_actions) environment_batched = {k: tf.stack([v] * self.n_samples, axis=0) for k, v in planning_query.environment.items()} start_state_batched = {k: tf.expand_dims(tf.stack([v] * self.n_samples, axis=0), axis=1) for k, v in planning_query.start.items()} mean_predictions, _ = self.fwd_model.propagate_differentiable_batched(environment=environment_batched, state=start_state_batched, actions=random_actions) final_states = {k: v[:, -1] for k, v in mean_predictions.items()} goal_state_batched = {k: tf.stack([v] * self.n_samples, axis=0) for k, v in planning_query.goal.items()} costs = self.scenario.trajopt_distance_to_goal_differentiable(final_states, goal_state_batched) costs = tf.squeeze(costs) min_idx = tf.math.argmin(costs, axis=0) # print(costs) # print('min idx', min_idx.numpy()) best_indices = tf.argsort(costs) cmap = cm.Blues n_to_show = 5 min_cost = costs[best_indices[0]] max_cost = costs[best_indices[n_to_show]] for j, i in enumerate(best_indices[:n_to_show]): s = numpify({k: v[i] for k, v in start_state_batched.items()}) a = numpify({k: v[i][0] for k, v in random_actions.items()}) c = (costs[i] - min_cost) / (max_cost - min_cost) self.scenario.plot_action_rviz(s, a, label='samples', color=cmap(c), idx1=2 * j, idx2=2 * j + 1) best_actions = {k: v[min_idx] for k, v in random_actions.items()} best_prediction = {k: v[min_idx] for k, v in mean_predictions.items()} best_actions = numpify(dict_of_sequences_to_sequence_of_dicts(best_actions)) best_predictions = numpify(dict_of_sequences_to_sequence_of_dicts(best_prediction)) planning_time = perf_counter() - start_planning_time planner_status = MyPlannerStatus.Solved return PlanningResult(status=planner_status, path=best_predictions, actions=best_actions, time=planning_time, tree={})
def plot_3d(args, dataset: DynamicsDatasetLoader, tf_dataset: tf.data.Dataset): scenario = dataset.scenario image_diff_viz_pub = rospy.Publisher("image_diff_viz", Image, queue_size=10, latch=True) for i, example in enumerate(tf_dataset): if args.start_at is not None and i < args.start_at: continue example = numpify(example) print(f"Example {i}, Trajectory #{int(example['traj_idx'])}") time_steps = example['time_idx'] anim = RvizAnimationController(time_steps) while not anim.done: t = anim.t() scenario.plot_environment_rviz(example) example_t = index_time_np(example, dataset.time_indexed_keys, t) scenario.plot_state_rviz(example_t, label='') scenario.plot_action_rviz_internal(example_t, label='') if t < dataset.steps_per_traj - 1: s_next = index_time_np(example, dataset.time_indexed_keys, t + 1) # diff = s['rgbd'][:, :, :3] - s_next['rgbd'][:, :, :3] # publish_color_image(image_diff_viz_pub, diff) # this will return when either the animation is "playing" or because the user stepped forward anim.step()
def on_trial_complete(self, trial_data: Dict, trial_idx: int): extra_trial_data = { "planner_params": self.planner_params, "scenario": self.planner.scenario.simple_name(), 'current_time': int(time()), 'uuid': uuid.uuid4(), } trial_data.update(extra_trial_data) data_filename = self.outdir / f'{trial_idx}_metrics.json.gz' dummy_proof_write(trial_data, data_filename) if self.record: # TODO: maybe make this happen async? sleep(1) self.service_provider.stop_record_trial() self.bag.close() # compute the current running average success goal = trial_data['planning_queries'][0].goal final_actual_state = numpify(trial_data['end_state']) final_execution_to_goal_error = self.planner.scenario.distance_to_goal(final_actual_state, goal) self.final_execution_to_goal_errors.append(final_execution_to_goal_error) goal_threshold = self.planner_params['goal_params']['threshold'] n = len(self.final_execution_to_goal_errors) success_percentage = np.count_nonzero(np.array(self.final_execution_to_goal_errors) < goal_threshold) / n * 100 update_msg = f"[{self.outdir.name}] Current average success rate {success_percentage:.2f}%" rospy.loginfo(Fore.LIGHTBLUE_EX + update_msg) jobkey = self.jobkey(trial_idx) self.job_chunker.store_result(jobkey, {'data_filename': data_filename})
def plot_recovery(args, scenario, step, metadata): actual_path = step['execution_result']['path'] action = step['recovery_action'] environment = numpify(step['planning_query']['environment']) scenario.plot_environment_rviz(environment) scenario.plot_state_rviz(actual_path[0], idx=1, label='recovery') scenario.plot_action_rviz(actual_path[0], action, label='recovery') scenario.plot_state_rviz(actual_path[1], idx=2, label='recovery')
def viz_main(args): dataset_dirs = args.dataset_dirs checkpoint = args.checkpoint trial_path, params = load_trial(checkpoint.parent.absolute()) dataset = DynamicsDatasetLoader(dataset_dirs) scenario = dataset.scenario tf_dataset = dataset.get_datasets(mode='val') tf_dataset = batch_tf_dataset(tf_dataset, batch_size=1, drop_remainder=True) model = CFM(hparams=params, batch_size=1, scenario=scenario) ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, args.checkpoint, max_to_keep=1) status = ckpt.restore(manager.latest_checkpoint).expect_partial() if manager.latest_checkpoint: print(Fore.CYAN + "Restored from {}".format(manager.latest_checkpoint)) status.assert_existing_objects_matched() else: raise RuntimeError("Failed to restore!!!") for example_idx, example in enumerate(tf_dataset): stepper = RvizAnimationController(n_time_steps=dataset.steps_per_traj) for t in range(dataset.steps_per_traj): output = model( model.preprocess_no_gradient(example, training=False)) actual_t = numpify( remove_batch(scenario.index_time_batched_predicted(example, t))) action_t = numpify( remove_batch(scenario.index_time_batched_predicted(example, t))) scenario.plot_state_rviz(actual_t, label='actual', color='red') scenario.plot_action_rviz(actual_t, action_t, color='gray') prediction_t = remove_batch( scenario.index_time_batched_predicted(output, t)) scenario.plot_state_rviz(prediction_t, label='predicted', color='blue') stepper.step()
def sample(self, environment: Dict, state: Dict): input_dict = environment input_dict.update({add_predicted(k): tf.expand_dims(v, axis=0) for k, v in state.items()}) input_dict = add_batch(input_dict) input_dict = {k: tf.cast(v, tf.float32) for k, v in input_dict.items()} output = self.net.sample(input_dict) output = remove_batch(output) output = numpify(output) return output
def main(): colorama.init(autoreset=True) plt.style.use("slides") np.set_printoptions(precision=3, suppress=True, linewidth=200) parser = argparse.ArgumentParser(formatter_class=my_formatter) parser.add_argument("fwd_model_dir", help="load this saved forward model file", type=pathlib.Path, nargs='+') parser.add_argument("test_config", help="json file describing the test", type=pathlib.Path) parser.add_argument("labeling_params", help='labeling params', type=pathlib.Path) args = parser.parse_args() rospy.init_node('test_model_from_gazebo') test_config = json.load(args.test_config.open("r")) labeling_params = json.load(args.labeling_params.open("r")) labeling_state_key = labeling_params['state_key'] # read actions from config actions = [numpify(a) for a in test_config['actions']] n_actions = len(actions) time_steps = np.arange(n_actions + 1) fwd_model, _ = dynamics_utils.load_generic_model(args.fwd_model_dir) service_provider = GazeboServices() service_provider.setup_env( verbose=0, real_time_rate=0, max_step_size=fwd_model.data_collection_params['max_step_size']) environment = fwd_model.scenario.get_environment( params=fwd_model.data_collection_params) start_state = fwd_model.scenario.get_state() start_state = make_dict_tf_float32(start_state) start_states = [start_state] expanded_actions = [[actions]] predicted_states = predict(fwd_model, environment, start_states, expanded_actions, n_actions, 1, 1) scenario = fwd_model.scenario actual_states_lists = execute(service_provider, scenario, start_states, expanded_actions) visualize(scenario, environment, actual_states_lists, actions, predicted_states, labeling_state_key, time_steps)
def play(self, example: Any): example = numpify(example) for init_func in self.init_funcs: init_func(self.scenario, example) controller = RvizAnimationController(n_time_steps=self.n_time_steps) while not controller.done: t = controller.t() for t_func in self.t_funcs: t_func(self.scenario, example, t) controller.step()
def evaluate(self, batch_size: int, trial_path: pathlib.Path, take: int): print(Fore.GREEN + f"Evaluating {trial_path.as_posix()}") train_metrics = train_test_classifier.eval_main( dataset_dirs=[self.classifier_dataset_dir], checkpoint=trial_path / 'best_checkpoint', mode='train', threshold=self.threshold, trials_directory=self.trials_directory, batch_size=batch_size, use_gt_rope=False, take=take, ) val_metrics = train_test_classifier.eval_main( dataset_dirs=[self.classifier_dataset_dir], checkpoint=trial_path / 'best_checkpoint', mode='val', threshold=self.threshold, trials_directory=self.trials_directory, batch_size=batch_size, use_gt_rope=False, take=take) train_metrics = dict_tools.dict_round(numpify(train_metrics)) val_metrics = dict_tools.dict_round(numpify(val_metrics)) return train_metrics, val_metrics
def viz_func(self, batch, outputs, test_dataset: DynamicsDatasetLoader): """ we assume batch size of 1 """ test_dataset.scenario.plot_environment_rviz(remove_batch(batch)) anim = RvizAnimationController(np.arange(test_dataset.steps_per_traj)) while not anim.done: t = anim.t() if self.first_latent_state is None: self.first_latent_state = outputs['z'][0, 0] pass m = Marker() m.header.frame_id = 'world' m.header.stamp = rospy.Time.now() m.type = Marker.SPHERE m.action = Marker.MODIFY m.scale.x = 0.01 m.scale.y = 0.01 m.scale.z = 0.01 m.color.r = 0.8 m.color.g = 0.2 m.color.b = 0.8 m.color.a = 0.8 m.id = self.idx m.ns = 'latent state' m.pose.position.x = (outputs['z'][0, 0, 0] - self.first_latent_state[0]) * 10 m.pose.position.y = (outputs['z'][0, 0, 1] - self.first_latent_state[1]) * 10 m.pose.position.z = (outputs['z'][0, 0, 2] - self.first_latent_state[2]) * 10 m.pose.orientation.w = 1 self.latent_state_pub.publish(m) e_t = numpify( remove_batch( test_dataset.scenario.index_time_batched_predicted( batch, t))) test_dataset.scenario.plot_state_rviz(e_t, label='actual', color='red') test_dataset.scenario.plot_action_rviz(e_t, e_t, color='gray') self.idx += 1 anim.step()
def plot_plan(args, scenario, step, metadata, fallback_labeing_params: Dict): planner_params = metadata['planner_params'] labeling_params = labeling_params_from_planner_params( planner_params, fallback_labeing_params) goal = step['planning_query']['goal'] environment = numpify(step['planning_query']['environment']) planned_path = step['planning_result']['path'] actual_path = step['execution_result']['path'] planned_actions = step['planning_result']['actions'] scenario.reset_planning_viz() if args.show_tree: def _draw_tree_function(scenario, tree_json): print(f"n vertices {len(tree_json['vertices'])}") for vertex in tree_json['vertices']: scenario.plot_tree_state(vertex, color='#77777722') sleep(0.001) tree_thread = threading.Thread(target=_draw_tree_function, args=( scenario, step['tree_json'], )) tree_thread.start() goal_threshold = get_goal_threshold(planner_params) scenario.animate_evaluation_results(environment=environment, actual_states=actual_path, predicted_states=planned_path, actions=planned_actions, goal=goal, goal_threshold=goal_threshold, labeling_params=labeling_params, accept_probabilities=None, horizon=metadata['horizon'])
def reset_rope(self, action_params: Dict): reset = SetRopeStateRequest() # TODO: rename this to rope endpoints reset positions or something reset.left_gripper.x = numpify( action_params['left_gripper_reset_position'][0]) reset.left_gripper.y = numpify( action_params['left_gripper_reset_position'][1]) reset.left_gripper.z = numpify( action_params['left_gripper_reset_position'][2]) reset.right_gripper.x = numpify( action_params['right_gripper_reset_position'][0]) reset.right_gripper.y = numpify( action_params['right_gripper_reset_position'][1]) reset.right_gripper.z = numpify( action_params['right_gripper_reset_position'][2]) self.set_rope_state_srv(reset)
def plot_steps(args, scenario, datum, metadata, fallback_labeing_params: Dict): planner_params = metadata['planner_params'] goal_threshold = get_goal_threshold(planner_params) labeling_params = labeling_params_from_planner_params( planner_params, fallback_labeing_params) steps = datum['steps'] if len(steps) == 0: q = datum['planning_queries'][0] start = q['start'] goal = q['goal'] environment = q['environment'] anim = RvizAnimationController(n_time_steps=1) scenario.plot_state_rviz(start, label='actual', color='#ff0000aa') scenario.plot_goal_rviz(goal, goal_threshold) scenario.plot_environment_rviz(environment) anim.step() return goal = datum['goal'] first_step = steps[0] environment = numpify(first_step['planning_query']['environment']) all_actual_states = [] types = [] all_predicted_states = [] all_actions = [] for step_idx, step in enumerate(steps): if step['type'] == 'executed_plan': actions = step['planning_result']['actions'] actual_states = step['execution_result']['path'] predicted_states = step['planning_result']['path'] elif step['type'] == 'executed_recovery': actions = [step['recovery_action']] actual_states = step['execution_result']['path'] predicted_states = [None, None] else: raise NotImplementedError(f"invalid step type {step['type']}") actions = numpify(actions) actual_states = numpify(actual_states) predicted_states = numpify(predicted_states) all_actions.extend(actions) types.extend([step['type']] * len(actions)) all_actual_states.extend(actual_states[:-1]) all_predicted_states.extend(predicted_states[:-1]) if args.show_tree and step['type'] == 'executed_plan': def _draw_tree_function(scenario, tree_json): print(f"n vertices {len(tree_json['vertices'])}") for vertex in tree_json['vertices']: scenario.plot_tree_state(vertex, color='#77777722') sleep(0.001) tree_thread = threading.Thread(target=_draw_tree_function, args=( scenario, step['planning_result']['tree'], )) tree_thread.start() # but do add the actual final states all_actual_states.append(actual_states[-1]) all_predicted_states.append(predicted_states[-1]) anim = RvizAnimationController(n_time_steps=len(all_actual_states)) def _type_action_color(type_t: str): if type_t == 'executed_plan': return 'b' elif type_t == 'executed_recovery': return '#ff00ff' scenario.reset_planning_viz() dist_to_goal = np.inf while not anim.done: scenario.plot_environment_rviz(environment) t = anim.t() s_t = all_actual_states[t] s_t_pred = all_predicted_states[t] scenario.plot_state_rviz(s_t, label='actual', color='#ff0000aa') c = '#0000ffaa' if len(all_actions) > 0: if t < anim.max_t: type_t = types[t] action_color = _type_action_color(type_t) scenario.plot_action_rviz(s_t, all_actions[t], color=action_color) else: type_t = types[t - 1] action_color = _type_action_color(type_t) scenario.plot_action_rviz(all_actual_states[t - 1], all_actions[t - 1], color=action_color) if s_t_pred is not None: scenario.plot_state_rviz(s_t_pred, label='predicted', color=c) is_close = scenario.compute_label(s_t, s_t_pred, labeling_params) scenario.plot_is_close(is_close) else: scenario.plot_is_close(None) dist_to_goal = scenario.distance_to_goal(s_t, goal) actually_at_goal = dist_to_goal < goal_threshold scenario.plot_goal_rviz(goal, goal_threshold, actually_at_goal) anim.step() print(f"final dist to goal {dist_to_goal:.3f}")
def __init__(self, planner: MyPlanner, trials: List[int], verbose: int, planner_params: Dict, service_provider: BaseServices, no_execution: bool, test_scenes_dir: Optional[pathlib.Path] = None, save_test_scenes_dir: Optional[pathlib.Path] = None, ): self.planner = planner self.scenario = self.planner.scenario self.scenario.on_before_get_state_or_execute_action() self.trials = trials self.planner_params = planner_params self.verbose = verbose self.service_provider = service_provider self.no_execution = no_execution self.env_rng = np.random.RandomState(0) self.goal_rng = np.random.RandomState(0) self.recovery_rng = np.random.RandomState(0) self.test_scenes_dir = test_scenes_dir self.save_test_scenes_dir = save_test_scenes_dir if self.planner_params['recovery']['use_recovery']: recovery_model_dir = pathlib.Path(self.planner_params['recovery']['recovery_model_dir']) self.recovery_policy = recovery_policy_utils.load_generic_model(model_dir=recovery_model_dir, scenario=self.scenario, rng=self.recovery_rng) else: self.recovery_policy = None self.n_failures = 0 # for saving snapshots of the world self.link_states_listener = Listener("gazebo/link_states", LinkStates) # Debugging if self.verbose >= 2: self.goal_bbox_pub = rospy.Publisher('goal_bbox', BoundingBox, queue_size=10, latch=True) bbox_msg = extent_to_bbox(planner_params['goal_params']['extent']) bbox_msg.header.frame_id = 'world' self.goal_bbox_pub.publish(bbox_msg) goal_params = self.planner_params['goal_params'] if goal_params['type'] == 'fixed': self.goal_generator = lambda e: numpify(goal_params['goal_fixed']) elif goal_params['type'] == 'random': self.goal_generator = lambda e: self.scenario.sample_goal(environment=e, rng=self.goal_rng, planner_params=self.planner_params) elif goal_params['type'] == 'dataset': dataset = DynamicsDatasetLoader([pathlib.Path(goal_params['goals_dataset'])]) tf_dataset = dataset.get_datasets(mode='val') goal_dataset_iterator = iter(tf_dataset) def _gen(e): example = next(goal_dataset_iterator) example_t = dataset.index_time_batched(example_batched=add_batch(example), t=1) goal = remove_batch(example_t) return goal self.goal_generator = _gen else: raise NotImplementedError(f"invalid goal param type {goal_params['type']}")
def propagate(self, environment: Dict, start_states: Dict, actions: List[Dict]) -> Tuple[List[Dict], List[Dict]]: mean_predictions, stdev_predictions = self.propagate_differentiable(environment, start_states, actions) return numpify(mean_predictions), numpify(stdev_predictions)
def index_time_np(e: Dict, time_indexed_keys: List[str], t: int): return numpify(index_time(e, time_indexed_keys, t))
def test_as_inverse_model(filter_model, latent_dynamics_model, test_dataset, test_tf_dataset): scenario = test_dataset.scenario shooting_method = ShootingMethod(fwd_model=latent_dynamics_model, classifier_model=None, scenario=scenario, params={'n_samples': 1000}) trajopt = TrajectoryOptimizer(fwd_model=latent_dynamics_model, classifier_model=None, scenario=scenario, params={ "iters": 100, "length_alpha": 0, "goal_alpha": 1000, "constraints_alpha": 0, "action_alpha": 0, "initial_learning_rate": 0.0001, }) s_color_viz_pub = rospy.Publisher("s_state_color_viz", Image, queue_size=10, latch=True) s_next_color_viz_pub = rospy.Publisher("s_next_state_color_viz", Image, queue_size=10, latch=True) image_diff_viz_pub = rospy.Publisher("image_diff_viz", Image, queue_size=10, latch=True) action_horizon = 1 initial_actions = [] total_errors = [] for example_idx, example in enumerate(test_tf_dataset): stepper = RvizAnimationController( n_time_steps=test_dataset.steps_per_traj) for t in range(test_dataset.steps_per_traj - 1): print(example_idx) environment = {} current_observation = remove_batch( scenario.index_observation_time_batched(add_batch(example), t)) for j in range(action_horizon): left_gripper_position = [0, 0, 0] right_gripper_position = [0, 0, 0] initial_action = { 'left_gripper_position': left_gripper_position, 'right_gripper_position': right_gripper_position, } initial_actions.append(initial_action) goal_observation = { k: example[k][1] for k in filter_model.obs_keys } planning_query = PlanningQuery(start=current_observation, goal=goal_observation, environment=environment, seed=1) planning_result = shooting_method.plan(planning_query) actions = planning_result.actions planned_path = planning_result.latent_path true_action = numpify( {k: example[k][0] for k in latent_dynamics_model.action_keys}) for j in range(action_horizon): optimized_action = actions[j] # optimized_action = { # 'left_gripper_position': current_observation['left_gripper'], # 'right_gripper_position': current_observation['right_gripper'], # } true_action = numpify({ k: example[k][j] for k in latent_dynamics_model.action_keys }) # Visualize s = numpify( remove_batch( scenario.index_observation_time_batched( add_batch(example), 0))) s.update( numpify( remove_batch( scenario.index_observation_features_time_batched( add_batch(example), 0)))) s_next = numpify( remove_batch( scenario.index_observation_time_batched( add_batch(example), 1))) s_next.update( numpify( remove_batch( scenario.index_observation_features_time_batched( add_batch(example), 1)))) scenario.plot_state_rviz(s, label='t', color="#ff000055", id=1) scenario.plot_state_rviz(s_next, label='t+1', color="#aa222255", id=2) # scenario.plot_action_rviz(s, optimized_action, label='inferred', color='#00ff00', id=1) # scenario.plot_action_rviz(s, true_action, label='true', color='#ee770055', id=2) publish_color_image(s_color_viz_pub, s['rgbd'][:, :, :3]) publish_color_image(s_next_color_viz_pub, s_next['rgbd'][:, :, :3]) diff = s['rgbd'][:, :, :3] - s_next['rgbd'][:, :, :3] publish_color_image(image_diff_viz_pub, diff) # Metrics total_error = 0 for v1, v2 in zip(optimized_action.values(), true_action.values()): total_error += -np.dot(v1, v2) total_errors.append(total_error) stepper.step() if example_idx > 100: break print(np.min(total_errors)) print(np.max(total_errors)) print(np.mean(total_errors)) plt.xlabel("total error (meter-ish)") plt.hist(total_errors, bins=np.linspace(0, 2, 20)) plt.show()
def index_time_batched(self, example_batched, t: int): e_t = numpify( remove_batch( index_time_batched(example_batched, self.time_indexed_keys, t))) return e_t
def viz(data_filename, fps, no_plot, save): rospy.init_node("compare_models") # Load the results base_folder = data_filename.parent with gzip.open(data_filename, "rb") as data_file: data_str = data_file.read() saved_data = json.loads(data_str.decode("utf-8")) all_metrics = {} for example_idx, datum in enumerate(saved_data): print(example_idx) # use the first (or any) model data to get the ground truth and dataset_element = numpify(datum.pop("dataset_element")) environment = numpify(datum.pop("environment")) action_keys = datum.pop("action_keys") actions = {k: dataset_element[k] for k in action_keys} models_viz_info = {} n_models = len(datum) time_steps = np.arange(datum.pop('time_steps')) for model_name, data_for_model in datum.items(): scenario = get_scenario(data_for_model['scenario']) # Metrics metrics_for_model = {} predictions = numpify(data_for_model['predictions']) predictions.pop('stdev') metrics = scenario.dynamics_metrics_function(dataset_element, predictions) loss = scenario.dynamics_loss_function(dataset_element, predictions) metrics['loss'] = loss for metric_name, metric_value in metrics.items(): if metric_name not in metrics_for_model: metrics_for_model[metric_name] = [] metrics_for_model[metric_name].append(metric_value.numpy()) for metric_name, metric_values in metrics_for_model.items(): mean_metric_value = float(np.mean(metric_values)) if model_name not in all_metrics: all_metrics[model_name] = {} if metric_name not in all_metrics[model_name]: all_metrics[model_name][metric_name] = [] all_metrics[model_name][metric_name].append(mean_metric_value) models_viz_info[model_name] = (scenario, predictions) if not no_plot and not save: # just use whatever the latest scenario was, it shouldn't matter which we use scenario.plot_environment_rviz(remove_batch(environment)) anim = RvizAnimationController(time_steps) while not anim.done: t = anim.t() actual_t = remove_batch(scenario.index_state_time(dataset_element, t)) action_t = remove_batch(scenario.index_action_time(actions, t)) scenario.plot_state_rviz(actual_t, label='actual', color='#0000ff88') scenario.plot_action_rviz(actual_t, action_t, color='gray') for model_idx, (model_name, viz_info) in enumerate(models_viz_info.items()): scenario_i, predictions = viz_info prediction_t = remove_batch(scenario_i.index_state_time(predictions, t)) color = cm.jet(model_idx / n_models) scenario_i.plot_state_rviz(prediction_t, label=model_name, color=color) anim.step() metrics_by_model = {} for model_name, metrics_for_model in all_metrics.items(): for metric_name, metric_values in metrics_for_model.items(): if metric_name not in metrics_by_model: metrics_by_model[metric_name] = {} metrics_by_model[metric_name][model_name] = metric_values with (base_folder / 'metrics_tables.txt').open("w") as metrics_file: for metric_name, metric_by_model in metrics_by_model.items(): headers = ["Model", "min", "max", "mean", "median", "std"] table_data = [] for model_name, metric_values in metric_by_model.items(): table_data.append([model_name] + row_stats(metric_values)) print('-' * 90) print(Style.BRIGHT + metric_name + Style.NORMAL) table = tabulate(table_data, headers=headers, tablefmt='fancy_grid', floatfmt='6.4f', numalign='center', stralign='left') metrics_file.write(table) print(table) print() print(Style.BRIGHT + f"p-value matrix [{metric_name}]" + Style.NORMAL) print(dict_to_pvalue_table(metric_by_model))
def viz_main(dataset_dirs: List[pathlib.Path], checkpoint: pathlib.Path, mode: str, batch_size: int, only_errors: bool, use_gt_rope: bool, old_compat: bool = False, **kwargs): stdev_pub_ = rospy.Publisher("stdev", Float32, queue_size=10) traj_idx_pub_ = rospy.Publisher("traj_idx_viz", Float32, queue_size=10) ############### # Model ############### trials_directory = pathlib.Path('trials').absolute() trial_path = checkpoint.parent.absolute() _, params = filepath_tools.create_or_load_trial( trial_path=trial_path, trials_directory=trials_directory) model_class = link_bot_classifiers.get_model(params['model_class']) ############### # Dataset ############### dataset = ClassifierDatasetLoader( dataset_dirs, load_true_states=True, use_gt_rope=use_gt_rope, threshold=params['classifier_dataset_hparams']['labeling_params'] ['threshold'], old_compat=old_compat) model = model_class(hparams=params, batch_size=batch_size, scenario=dataset.scenario) tf_dataset = dataset.get_datasets(mode=mode) scenario = dataset.scenario ############### # Evaluate ############### tf_dataset = batch_tf_dataset(tf_dataset, batch_size, drop_remainder=True) model = classifier_utils.load_generic_model([checkpoint]) for batch_idx, example in enumerate( progressbar(tf_dataset, widgets=base_dataset.widgets)): example.update(dataset.batch_metadata) predictions, _ = model.check_constraint_from_example(example, training=False) labels = tf.expand_dims(example['is_close'][:, 1:], axis=2) probabilities = predictions['probabilities'] # Visualization example.pop("time") example.pop("batch_size") decisions = probabilities > 0.5 classifier_is_correct = tf.squeeze(tf.equal(decisions, tf.cast(labels, tf.bool)), axis=-1) for b in range(batch_size): example_b = index_dict_of_batched_tensors_tf(example, b) # if the classifier is correct at all time steps, ignore if only_errors and tf.reduce_all(classifier_is_correct[b]): continue def _custom_viz_t(scenario: Base3DScenario, e: Dict, t: int): if t > 0: accept_probability_t = predictions['probabilities'][ b, t - 1, 0].numpy() else: accept_probability_t = -999 scenario.plot_accept_probability(accept_probability_t) traj_idx_msg = Float32() traj_idx_msg.data = batch_idx * batch_size + b traj_idx_pub_.publish(traj_idx_msg) anim = RvizAnimation(scenario=scenario, n_time_steps=dataset.horizon, init_funcs=[ init_viz_env, dataset.init_viz_action(), ], t_funcs=[ _custom_viz_t, dataset.classifier_transition_viz_t(), ExperimentScenario.plot_stdev_t, ]) with open("debugging.hjson", 'w') as f: my_hdump(numpify(example_b), f) anim.play(example_b)
def filter(self, environment: Dict, state: Optional[Dict], observation: Dict) -> Tuple[Dict, Dict]: mean_state, stdev_state = self.filter_differentiable(environment=environment, state=state, observation=observation) return numpify(mean_state), numpify(stdev_state)
def main(): colorama.init(autoreset=True) parser = argparse.ArgumentParser() parser.add_argument("results_dir", type=pathlib.Path, help='dir containing *_metrics.json.gz') parser.add_argument("plan_idx", type=int) args = parser.parse_args() rospy.init_node("postprocess_result") with (args.results_dir / 'metadata.json').open('r') as metadata_file: metadata_str = metadata_file.read() metadata = json.loads(metadata_str) planner_params = metadata['planner_params'] fwd_model_dirs = [pathlib.Path(p) for p in planner_params['fwd_model_dir']] fwd_model, _ = dynamics_utils.load_generic_model(fwd_model_dirs) scenario = fwd_model.scenario classifier_model_dir = pathlib.Path(planner_params['classifier_model_dir']) classifier = classifier_utils.load_generic_model(classifier_model_dir, scenario=scenario) metrics_filename = args.results_dir / f"{args.plan_idx}_metrics.json.gz" with gzip.open(metrics_filename, 'rb') as f: data_str = f.read() datum = json.loads(data_str.decode("utf-8")) steps = datum['steps'] goal = datum['goal'] first_step = steps[0] environment = numpify(first_step['planning_query']['environment']) all_output_actions = [] for step in steps: if step['type'] == 'executed_plan': actions = postprocess_step(scenario, fwd_model, classifier, environment, step, goal, planner_params) all_output_actions.extend(actions) elif step['type'] == 'executed_recovery': actions = [step['recovery_action']] all_output_actions.extend(actions) # viz the whole thing start_state = steps[0]['planning_result']['path'][0] final_states = fwd_model.propagate(environment, start_state, actions) T = len(final_states) for t, s_t in enumerate(final_states): scenario.plot_state_rviz(s_t, idx=t, label='smoothed', color='#00ff0099') if t < T - 1: scenario.plot_action_rviz(s_t, actions[t], idx=t, color="#ffffff99", label='smoothed') sleep(0.02) # Save the output actions outfilename = args.results_dir / f"{args.plan_idx}_smoothed.json" with outfilename.open("w") as outfile: json.dump(listify(all_output_actions), outfile) print(Fore.GREEN + f"Wrote {outfilename}" + Fore.RESET)
def test_classifier(classifier_model_dir: pathlib.Path, fwd_model_dir: List[pathlib.Path], n_actions: int, saved_state: Optional[pathlib.Path], generate_actions: Callable): fwd_model, _ = dynamics_utils.load_generic_model([pathlib.Path(p) for p in fwd_model_dir]) classifier: NNClassifierWrapper = classifier_utils.load_generic_model([classifier_model_dir]) service_provider = GazeboServices() service_provider.setup_env(verbose=0, real_time_rate=0, max_step_size=fwd_model.data_collection_params['max_step_size'], play=False) if saved_state: service_provider.restore_from_bag(saved_state) scenario = fwd_model.scenario scenario.on_before_get_state_or_execute_action() # NOTE: perhaps it would make sense to have a "fwd_model" have API for get_env, get_state, sample_action, etc # because the fwd_model knows it's scenario, and importantly it also knows it's data_collection_params # which is what we're using here to pass to the scenario methods params = fwd_model.data_collection_params environment = numpify(scenario.get_environment(params)) start_state = numpify(scenario.get_state()) start_state_tiled = repeat(start_state, n_actions, axis=0, new_axis=True) start_states_tiled = add_time_dim(start_state_tiled) actions = generate_actions(environment, start_state_tiled, scenario, params, n_actions) environment_tiled = repeat(environment, n_actions, axis=0, new_axis=True) actions_dict = sequence_of_dicts_to_dict_of_tensors(actions) actions_dict = add_time_dim(actions_dict) predictions, _ = fwd_model.propagate_differentiable_batched(environment=environment_tiled, state=start_states_tiled, actions=actions_dict) # Run classifier state_sequence_length = 2 accept_probabilities, _ = classifier.check_constraint_batched_tf(environment=environment_tiled, predictions=predictions, actions=actions_dict, state_sequence_length=state_sequence_length, batch_size=n_actions) # animate over the sampled actions anim = RvizAnimation(scenario=scenario, n_time_steps=n_actions, init_funcs=[init_viz_env], t_funcs=[ lambda s, e, t: init_viz_env(s, e), viz_transition_for_model_t({}, fwd_model), ExperimentScenario.plot_accept_probability_t, ], ) example = { 'accept_probability': tf.squeeze(accept_probabilities, axis=1), } example.update(environment) example.update(predictions) example.update(actions_dict) anim.play(example)