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 call(self, example, training, mask=None): actions = {k: example[k] for k in self.action_keys} input_sequence_length = actions[self.action_keys[0]].shape[1] s_0 = {k: example[k][:, 0] for k in self.state_keys} pred_states = [s_0] for t in range(input_sequence_length): s_t = pred_states[-1] action_t = {k: example[k][:, t] for k in self.action_keys} local_action_t = self.scenario.put_action_local_frame( s_t, action_t) s_t_local = self.scenario.put_state_local_frame(s_t) states_and_actions = list(s_t_local.values()) + list( local_action_t.values()) # concat into one big state-action vector z_t = self.concat(states_and_actions) for dense_layer in self.dense_layers: z_t = dense_layer(z_t) delta_s_t = vector_to_dict(self.state_description, z_t) s_t_plus_1 = self.scenario.integrate_dynamics(s_t, delta_s_t) pred_states.append(s_t_plus_1) pred_states_dict = sequence_of_dicts_to_dict_of_tensors(pred_states, axis=1) return pred_states_dict
def from_example(self, example, training: bool = False): """ This is the function that all other functions eventually call """ if 'batch_size' not in example: example['batch_size'] = self.get_batch_size(example) outputs = [ net(net.preprocess_no_gradient(example, training), training=training) for net in self.nets ] outputs_dict = sequence_of_dicts_to_dict_of_tensors(outputs) outputs_dict = { k: flatten_after(outputs_dict[k], axis=self.get_num_batch_axes()) for k in self.get_output_keys() } # axis 0 is the different networks mean = { k: tf.math.reduce_mean(outputs_dict[k], axis=0) for k in self.get_output_keys() } stdev = { k: tf.math.reduce_std(outputs_dict[k], axis=0) for k in self.get_output_keys() } # each output variable has its own vector of variances, # and here we sum all the elements of all the vectors to get a single scalar all_stdevs = tf.concat(list(stdev.values()), axis=-1) mean['stdev'] = tf.reduce_sum(all_stdevs, axis=-1, keepdims=True) return mean, stdev
def check_constraint_tf(self, environment: Dict, states_sequence: List[Dict], actions: List[Dict]): environment = add_batch(environment) states_sequence_dict = sequence_of_dicts_to_dict_of_tensors(states_sequence) states_sequence_dict = add_batch(states_sequence_dict) state_sequence_length = len(states_sequence) actions_dict = sequence_of_dicts_to_dict_of_tensors(actions) actions_dict = add_batch(actions_dict) mean_probabilities, stdev_probabilities = self.check_constraint_batched_tf(environment=environment, predictions=states_sequence_dict, actions=actions_dict, batch_size=1, state_sequence_length=state_sequence_length) mean_probabilities = remove_batch(mean_probabilities) stdev_probabilities = remove_batch(stdev_probabilities) return mean_probabilities, stdev_probabilities
def batch_stateless_sample_action(scenario: ExperimentScenario, environment: Dict, state: Dict, batch_size: int, n_action_samples: int, n_actions: int, action_params: Dict, action_rng: np.random.RandomState): # TODO: make the lowest level sample_action operate on batched state dictionaries action_sequences = scenario.sample_action_sequences( environment=environment, state=state, action_params=action_params, n_action_sequences=n_action_samples, action_sequence_length=n_actions, validate=False, action_rng=action_rng) action_sequences = [ sequence_of_dicts_to_dict_of_tensors(a) for a in action_sequences ] action_sequences = sequence_of_dicts_to_dict_of_tensors(action_sequences) return { k: tf.tile(v[tf.newaxis], [batch_size, 1, 1, 1]) for k, v in action_sequences.items() }
def propagate_differentiable(self, environment: Dict, start_state: Dict, actions: List[Dict]) -> Tuple[ List[Dict], List[Dict]]: # add time dimension of size 1 net_inputs = {k: tf.expand_dims(start_state[k], axis=0) for k in self.state_keys} net_inputs.update(environment) net_inputs.update(sequence_of_dicts_to_dict_of_tensors(actions)) net_inputs = add_batch(net_inputs) net_inputs = make_dict_tf_float32(net_inputs) # the network returns a dictionary where each value is [T, n_state] # which is what you'd want for training, but for planning and execution and everything else # it is easier to deal with a list of states where each state is a dictionary mean_predictions, stdev_predictions = self.from_example(net_inputs, training=False) mean_predictions = remove_batch(mean_predictions) stdev_predictions = remove_batch(stdev_predictions) mean_predictions = dict_of_sequences_to_sequence_of_dicts_tf(mean_predictions) stdev_predictions = dict_of_sequences_to_sequence_of_dicts_tf(stdev_predictions) return mean_predictions, stdev_predictions
def predict(fwd_model: BaseDynamicsFunction, environment: Dict, start_states: List[Dict], actions: List[List[List[Dict]]], n_actions: int, n_actions_sampled: int, n_start_states: int) -> List[List[List[Dict]]]: # reformat the inputs to be efficiently batched actions_batched = batch_of_many_of_actions_sequences_to_dict(actions, n_actions_sampled, n_start_states, n_actions) start_states = sequence_of_dicts_to_dict_of_tensors(start_states) start_states = make_dict_tf_float32({k: tf.expand_dims(v, axis=1) for k, v in start_states.items()}) # copy from start states for each random action start_states_tiled = {k: tf.concat([v] * n_actions_sampled, axis=0) for k, v in start_states.items()} # Actually do the predictions predictions_dict, _ = fwd_model.propagate_differentiable_batched(environment=environment, state=start_states_tiled, actions=actions_batched) # break out the num actions and num start states n_states = n_actions + 1 predictions_dict = {k: tf.reshape(v, [n_start_states, n_actions_sampled, n_states, -1]) for k, v in predictions_dict.items()} # invert structure to List[List[List[Dict]]] and return predictions_list = [] for i in range(n_start_states): predictions_i = [] for j in range(n_actions_sampled): predictions_ij = [] for t in range(n_actions + 1): prediction_ijt = {k: predictions_dict[k][i, j, t] for k, v in predictions_dict.items()} predictions_ij.append(prediction_ijt) predictions_i.append(predictions_ij) predictions_list.append(predictions_i) return predictions_list
def call(self, example, training, mask=None): batch_size = example['batch_size'] time = example[self.action_keys[0]].shape[1] + 1 s_0 = {k: example[k][:, 0] for k in self.state_keys} # Construct a [b, h, w, c, 3] grid of the indices which make up the local environment pixel_row_indices = tf.range(0, self.local_env_h_rows, dtype=tf.float32) pixel_col_indices = tf.range(0, self.local_env_w_cols, dtype=tf.float32) pixel_channel_indices = tf.range(0, self.local_env_c_channels, dtype=tf.float32) x_indices, y_indices, z_indices = tf.meshgrid(pixel_col_indices, pixel_row_indices, pixel_channel_indices) # Make batched versions for creating the local environment batch_y_indices = tf.cast( tf.tile(tf.expand_dims(y_indices, axis=0), [batch_size, 1, 1, 1]), tf.int64) batch_x_indices = tf.cast( tf.tile(tf.expand_dims(x_indices, axis=0), [batch_size, 1, 1, 1]), tf.int64) batch_z_indices = tf.cast( tf.tile(tf.expand_dims(z_indices, axis=0), [batch_size, 1, 1, 1]), tf.int64) # Convert for rastering state pixel_indices = tf.stack([y_indices, x_indices, z_indices], axis=3) pixel_indices = tf.expand_dims(pixel_indices, axis=0) pixel_indices = tf.tile(pixel_indices, [batch_size, 1, 1, 1, 1]) # # DEBUG # # plot the occupancy grid # time_steps = np.arange(time) # anim = RvizAnimationController(time_steps) # b = 0 # full_env_dict = { # 'env': example['env'][b], # 'origin': example['origin'][b], # 'res': example['res'][b], # 'extent': example['extent'][b], # } # self.scenario.plot_environment_rviz(full_env_dict) # # END DEBUG pred_states = [s_0] for t in range(time - 1): s_t = pred_states[-1] s_t_local = self.scenario.put_state_local_frame(s_t) action_t = {k: example[k][:, t] for k in self.action_keys} local_action_t = self.scenario.put_action_local_frame( s_t, action_t) local_env_center_t = self.scenario.local_environment_center_differentiable( s_t) # by converting too and from the frame of the full environment, we ensure the grids are aligned indices = batch_point_to_idx_tf_3d_in_batched_envs( local_env_center_t, example) local_env_center_t = batch_idx_to_point_3d_in_env_tf( *indices, example) local_env_t, local_env_origin_t = get_local_env( center_point=local_env_center_t, full_env=example['env'], full_env_origin=example['origin'], res=example['res'], local_h_rows=self.local_env_h_rows, local_w_cols=self.local_env_w_cols, local_c_channels=self.local_env_c_channels, batch_x_indices=batch_x_indices, batch_y_indices=batch_y_indices, batch_z_indices=batch_z_indices, batch_size=batch_size) local_voxel_grid_t_array = tf.TensorArray(tf.float32, size=0, dynamic_size=True) local_voxel_grid_t_array = local_voxel_grid_t_array.write( 0, local_env_t) for i, state_component_t in enumerate(s_t.values()): state_component_voxel_grid = raster_3d( state=state_component_t, pixel_indices=pixel_indices, res=example['res'], origin=local_env_origin_t, h=self.local_env_h_rows, w=self.local_env_w_cols, c=self.local_env_c_channels, k=self.rope_image_k, batch_size=batch_size) local_voxel_grid_t_array = local_voxel_grid_t_array.write( i + 1, state_component_voxel_grid) local_voxel_grid_t = tf.transpose(local_voxel_grid_t_array.stack(), [1, 2, 3, 4, 0]) # add channel dimension information because tf.function erases it somehow... local_voxel_grid_t.set_shape( [None, None, None, None, len(self.state_keys) + 1]) # # DEBUG # local_env_dict = { # 'env': tf.clip_by_value(tf.reduce_sum(local_voxel_grid_t[b], axis=-1), 0, 1), # # 'env': local_voxel_grid_t[b], # 'origin': local_env_origin_t[b].numpy(), # 'res': example['res'][b].numpy(), # } # msg = environment_to_occupancy_msg(local_env_dict, frame='local_occupancy') # link_bot_sdf_utils.send_occupancy_tf(self.scenario.broadcaster, local_env_dict, frame='local_occupancy') # self.debug_pub.publish(msg) # self.scenario.plot_state_rviz(numpify(index_dict_of_batched_vectors_tf(s_t, b)), label='pred') # local_extent = compute_extent_3d(*local_voxel_grid_t[b].shape[:3], resolution=example['res'][b].numpy()) # depth, width, height = extent_to_env_size(local_extent) # bbox_msg = BoundingBox() # bbox_msg.header.frame_id = 'local_occupancy' # bbox_msg.pose.position.x = width / 2 # bbox_msg.pose.position.y = depth / 2 # bbox_msg.pose.position.z = height / 2 # bbox_msg.dimensions.x = width # bbox_msg.dimensions.y = depth # bbox_msg.dimensions.z = height # self.local_env_bbox_pub.publish(bbox_msg) # anim.step() # # END DEBUG # CNN z_t = local_voxel_grid_t for conv_layer, pool_layer in zip(self.conv_layers, self.pool_layers): z_t = conv_layer(z_t) z_t = pool_layer(z_t) conv_z_t = self.flatten_conv_output(z_t) for dense_layer in self.conv_only_dense_layers: conv_z_t = dense_layer(conv_z_t) # concat into one big state-action vector states_and_actions = list(s_t_local.values()) + list( local_action_t.values()) state_action_t = self.concat2(states_and_actions) for dense_layer in self.state_action_dense_layers: state_action_t = dense_layer(state_action_t) # combine state/action vector with output of CNN full_z_t = self.combine_state_action_and_conv_output( [state_action_t, conv_z_t]) # dense layers for combined vector for dense_layer in self.final_dense_layers: full_z_t = dense_layer(full_z_t) delta_s_t = self.vector_to_state_dict(full_z_t) s_t_plus_1 = self.scenario.integrate_dynamics(s_t, delta_s_t) pred_states.append(s_t_plus_1) pred_states_dict = sequence_of_dicts_to_dict_of_tensors(pred_states, axis=1) return pred_states_dict
def predict_and_classify(fwd_model: Ensemble, classifier: NNClassifierWrapper, environment: Dict, start_states: List[Dict], actions: List[List[List[Dict]]], n_actions: int, n_actions_sampled: int, n_start_states: int): state_sequence_length = n_actions + 1 # reformat the inputs to be efficiently batched actions_dict = {} for actions_for_start_state in actions: for actions in actions_for_start_state: for action in actions: for k, v in action.items(): if k not in actions_dict: actions_dict[k] = [] actions_dict[k].append(v) environment_batched = {k: tf.stack([v] * n_actions_sampled * n_start_states, axis=0) for k, v in environment.items()} actions_batched = {k: tf.reshape(v, [n_actions_sampled * n_start_states, n_actions, -1]) for k, v in actions_dict.items()} start_states = sequence_of_dicts_to_dict_of_tensors(start_states) start_states = make_dict_tf_float32({k: tf.expand_dims(v, axis=1) for k, v in start_states.items()}) # copy from start states for each random action start_states_tiled = {k: tf.concat([v] * n_actions_sampled, axis=0) for k, v in start_states.items()} # Actually do the predictions predictions_dict = fwd_model.propagate_differentiable_batched(start_states=start_states_tiled, actions=actions_batched) # Run classifier accept_probabilities = classifier.check_constraint_batched_tf(environment=environment_batched, predictions=predictions_dict, actions=actions_batched, state_sequence_length=state_sequence_length, batch_size=n_start_states * n_actions_sampled) accept_probabilities = tf.reshape( accept_probabilities, [n_start_states, n_actions_sampled, state_sequence_length - 1, -1]) # break out the num actions and num start states state_sequence_length = n_actions + 1 predictions_dict = {k: tf.reshape(v, [n_start_states, n_actions_sampled, state_sequence_length, -1]) for k, v in predictions_dict.items()} # invert structure to List[List[List[Dict]]] and return predictions_list = [] accept_probabilities_list = [] for i in range(n_start_states): predictions_i = [] accept_probabilities_i = [] for j in range(n_actions_sampled): predictions_ij = [] accept_probabilities_ij = [] for t in range(n_actions + 1): prediction_ijt = {k: predictions_dict[k][i, j, t] for k, v in predictions_dict.items()} predictions_ij.append(prediction_ijt) if t < n_actions: accept_probabilities_ij.append(accept_probabilities[i, j, t]) predictions_i.append(predictions_ij) accept_probabilities_i.append(accept_probabilities_ij) predictions_list.append(predictions_i) accept_probabilities_list.append(accept_probabilities_i) return predictions_list, accept_probabilities_list
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)
def check_constraint_from_example(self, example: Dict, training: Optional[bool] = False): predictions = [net(net.preprocess_no_gradient(example, training), training=training) for net in self.nets] predictions_dict = sequence_of_dicts_to_dict_of_tensors(predictions) mean_predictions = {k: tf.math.reduce_mean(v, axis=0) for k, v in predictions_dict.items()} stdev_predictions = {k: tf.math.reduce_std(v, axis=0) for k, v in predictions_dict.items()} return mean_predictions, stdev_predictions