def plot_sampled_minibatch(sampled_output, experts=None, figsize=(8, 20), partial_write_np_image_to_tb=None, without_samples=False, plot_bev_kwargs={}, tensorstr='', fig=None, axes=None, eager=False): fig, axes = get_figure(fig=fig, axes=axes, figsize=figsize) B = sampled_output.phi_metadata.B A = sampled_output.rollout.event_shape[0] limit = tensoru.size(sampled_output.phi.overhead_features, 1) # Make the expert color different from the model if there's only one agent. expert_kwargs = {'color': COLORS[1]} if A == 1 else {} for b in range(min(B, 10)): ax = axes.ravel()[b] magic_reset_axis(ax) if not without_samples: plot_single_sampled_output(sampled_output, batch_index=b, render=False, fig=fig, ax=ax) plot_past(sampled_output.phi.S_past_grid_frame, b=b, fig=fig, ax=ax, limit=limit) plot_rollout(sampled_output.rollout, fig=fig, ax=ax, b=b) plot_bev(sampled_output, batch_index=b, ax=ax, **plot_bev_kwargs) if experts is not None: plot_expert(experts, batch_index=b, fig=fig, ax=ax, render=False, limit=tensoru.size( sampled_output.phi.overhead_features, 1), **expert_kwargs) # TODO make work with eager mode. plot_figure('sampled_minibatch' + tensorstr, fig, partial_write_np_image_to_tb=partial_write_np_image_to_tb) if not eager: figsclose() return fig
def get_map_feats(feature_map, batch_shape, batch_size, last_agent_positions_grid, A, phi=None): """Interpolate the positions into the feature map. :param feature_map: (B, H, W, F), feature map :param batch_shape: first several dimensions of positions. could be inferred but we require it for sanity checks. :param batch_size: size of batch. could be inferred but we require it for sanity checks. :param last_agent_positions_grid: batch_shape + (A, D) :param A: number of agents. could be inferred but we require it for sanity checks. :param F: feature map dimension. could be inferred but we require it for sanity checks. :param phi: DEPRECATED :returns: :rtype: """ assert (tensoru.rank(feature_map) == 4) assert (tensoru.rank(last_agent_positions_grid) >= 3) F = tensoru.size(feature_map, -1) # (B, batch_shape[1:]*A, d) last_agent_positions_grid_r_xy = tf.reshape(last_agent_positions_grid, (batch_shape[0], -1, 2)) last_agent_positions_grid_r_ij = last_agent_positions_grid_r_xy[..., ::-1] # (B, batch_shape[1:]*A, F) # N.B. the indexing order! Our points are stored in xy-format, with x corresponding to the W dimension of the feature grid. map_feat = tfu.interpolate_bilinear(feature_map, last_agent_positions_grid_r_ij, indexing='ij') # (BSize, A, F) map_feats = tf.reshape(map_feat, (batch_size, A, F)) return map_feats
def _static_prepare_samples(self, global_step): for t in self.input_singleton.sample_placeholders: self.model_collections.add_sample_input(t) for t in self.input_singleton.placeholders: self.model_collections.add_infer_input(t) # Track some intermediate tensors that we'll want to reconstruct when we load the model. self.model_collections.add_intermediate_input( self.input_singleton.phi.S_past_car_frames) self.model_collections.add_intermediate_input( self.input_singleton.phi.S_past_grid_frame) self.model_collections.add_intermediate_input( self.input_singleton.phi_m.agent_counts) self.model_collections.add_intermediate_label( self.input_singleton.experts.S_future_car_frames) self.model_collections.add_intermediate_label( self.input_singleton.experts.S_future_grid_frame) log.info("Computing static-mode samples.") self.sampled_output = self.model_distribution.sample( phi=self.input_singleton.phi, phi_metadata=self.input_singleton.phi_m, T=self.dataset.T) # dbg self.model_distribution.bijection.check_gradients( self.sampled_output.base_and_log_q.Z_sample) # Z input is a sample input self.model_collections.add_sample_input( self.sampled_output.base_and_log_q.Z_sample) # Record outputs for sampling. for t in self.sampled_output.rollout.rollout_outputs: self.model_collections.add_sample_output(t) # Log q is a sample output. self.model_collections.add_sample_output( self.sampled_output.base_and_log_q.log_q_samples) # Functions to write a numpy image into tensorboard. The function takes a single input (np.ndarray) self.partial_np2tb_smbs = [] for c in range( tensoru.size(self.input_singleton.phi.overhead_features, -1)): self.partial_np2tb_smbs.append( plot.bind_write_np_image_to_tb( sess=self.sess, writer=self.writer, global_step=global_step, key='sampled_minibatch_bev_{}'.format(c))) self.partial_np2tb_smb_joint = plot.bind_write_np_image_to_tb( sess=self.sess, writer=self.writer, global_step=global_step, key='sampled_minibatch_bev_joint') self.partial_np2tb_cnn0 = plot.bind_write_np_image_to_tb( sess=self.sess, writer=self.writer, global_step=global_step, key='CNN_0') log.info("Initializing variables...") self.sess.run(tfv1.global_variables_initializer())
def plot_whiskers(rollout, batch_index=0, k_slice=slice(0, 1), a_slice=slice(0, 5), period=5, fig=None, ax=None, render=True): A, T, d = rollout.event_shape whiskers = [ _['whiskers_grid'][batch_index, k_slice, a_slice] for _ in rollout.metadata_list[::period] ] frames = [_['local2grid'] for _ in rollout.metadata_list] origins = tf.stack([_.t[batch_index, k_slice, a_slice] for _ in frames], axis=-2) limit = tensoru.size(rollout.phi.overhead_features, 1) # plot_joint_trajectory(whiskers[0], key='whiskers0', render=render, fig=fig, ax=ax, marker='d', zorder=2, alpha=0.5, limit=limit) plot_joint_trajectory(origins, key='origins', render=render, fig=fig, ax=ax, marker='+', zorder=2, alpha=0.5, limit=limit) # plot_joint_trajectory(whiskers[-1], key='whiskers1', render=render, fig=fig, ax=ax, marker='d', zorder=2, alpha=0.5, limit=limit) return fig, ax
def get_whisker_map_feats(feature_map, batch_shape, batch_size, cars2grid, template_cars, A, phi=None): """Interpolate the positions into the feature map. :param feature_map: (B, H, W, F), feature map :param batch_shape: first several dimensions of positions. could be inferred but we require it for sanity checks. :param batch_size: size of batch. could be inferred but we require it for sanity checks. :param last_agent_positions_cars: batch_shape + (A, D) :param cars2grid: SimilarityTransform from car frames to grid :param template: (N, D) template of positions in local frame at which to interpolate :param A: number of agents. could be inferred but we require it for sanity checks. :param F: feature map dimension. could be inferred but we require it for sanity checks. :param phi: DEPRECATED :returns: :rtype: """ assert (tensoru.rank(feature_map) == 4) F = tensoru.size(feature_map, -1) if len(batch_shape) == 2: points_ein = 'bkaNj' template_cars = template_cars[None] elif len(batch_shape) == 1: points_ein = 'baNj' else: raise ValueError n_whiskers = tensoru.size(template_cars, -2) whiskers_grid = cars2grid.apply(template_cars, points_ein=points_ein) whiskers_grid_r_xy = tf.reshape(whiskers_grid, (batch_shape[0], -1, 2)) # (B, batch_shape[1:]*A, F) # N.B. the indexing order! Our points are stored in xy-format, with x corresponding to the W dimension of the feature grid. map_feat = tfu.interpolate_bilinear(feature_map, whiskers_grid_r_xy, indexing='xy') # (B, ..., A, n_whiskers, F) map_feat_r = tf.reshape(map_feat, batch_shape + (A, n_whiskers, F)) # (B*..., A, n_whiskers*F) map_feats = tf.reshape(map_feat_r, (batch_size, A, n_whiskers * F)) return map_feats, whiskers_grid
def __init__(self, S_past_world_frame, yaws, overhead_features, agent_presence, feature_pixels_per_meter, is_training, light_strings=None, yaws_in_degrees=True, past_perturb_epsilon=5e-2, name=False): """ :param S_past_world_frame: (B, A, T, D) :param yaws: (B, A) :param overhead_features: (B, H, W, C) :param agent_presence: (B, A) :param yaws_in_degrees: bool """ assert (feature_pixels_per_meter >= 1) if light_strings is None: light_strings = np.array(['NONE'] * tensoru.size(yaws, 0), dtype=np.unicode_) # Overwrite these members with tensorized versions of them. self.tensor_init(S_past_world_frame, yaws, overhead_features, agent_presence, light_strings, is_training) self._frames_init() past_noise_world_frame = tf.random.normal( mean=0., stddev=self.past_perturb_epsilon, shape=tensoru.shape(self.S_past_world_frame), dtype=tf.float64) # Always create an alternative noisy-past (learning may not use it). self.S_past_world_frame_noisy = self.S_past_world_frame + past_noise_world_frame # (B, A, T, D) self.S_past_car_frames = self.world2local.apply( self.S_past_world_frame) self.S_past_car_frames_noisy = self.world2local.apply( self.S_past_world_frame_noisy) # (B, A, T, D) self.S_past_grid_frame = self.world2grid.apply(self.S_past_world_frame, dtype=tf.float64) self.S_past_grid_frame_noisy = self.world2grid.apply( self.S_past_world_frame_noisy, dtype=tf.float64) self.light_features = one_hotify_light_strings(self.light_strings) if name: # Name some intermediate computations (not placeholders) self.S_past_grid_frame = tf.identity(self.S_past_grid_frame, name='S_past_grid_frame') self.S_past_car_frames = tf.identity(self.S_past_car_frames, name='S_past_car_frames') self.light_features = tf.identity(self.light_features, name='light_features')
def plot_rollout(rollout, b=None, fig=None, ax=None): assert (b is not None) limit = tensoru.size(rollout.phi.overhead_features, 1) plot_joint_trajectory(joint_traj=rollout.S_grid_frame[b], key='rollout_future', render=False, fig=fig, ax=ax, marker='o', zorder=3, alpha=0.5, limit=limit)
def plot_sample(sampled_output, expert=None, b=0, figsize=(4, 4), partial_write_np_image_to_tb=None, bev_kwargs={}): fig, ax = plt.subplots(1, 1, figsize=figsize) magic_reset_axis(ax) plot_single_sampled_output(sampled_output, batch_index=b, render=False, fig=fig, ax=ax) plot_bev(sampled_output, batch_index=b, ax=ax, **bev_kwargs) A = tensoru.size(sampled_output.rollout.S_car_frames, 2) expert_kwargs = {'color': COLORS[1]} if A == 1 else {} limit = tensoru.size(sampled_output.phi.overhead_features, 1) plot_past(sampled_output.phi.S_past_grid_frame, b=b, fig=fig, ax=ax, limit=limit) if expert is not None: plot_expert(expert, batch_index=b, fig=fig, ax=ax, render=False, limit=tensoru.size(sampled_output.phi.overhead_features, 1), **expert_kwargs) res = plot_figure( 'sampled_minibatch', fig, partial_write_np_image_to_tb=partial_write_np_image_to_tb) plt.close('all') return res
def __init__(self, phi, metadata_list, name=False): name = None if not name else 'agent_counts' assert (isinstance(self.metadata_list, MetadataList)) assert (all([isinstance(_, MetadataItem) for _ in metadata_list])) self.agent_counts = tf.reduce_sum(tf.cast(phi.agent_presence, tf.float64), axis=1, name=name) B0 = self.phi.S_past_car_frames.shape[0] B1 = self.phi.yaws.shape[0] B2 = self.phi.overhead_features.shape[0] B3 = self.phi.agent_presence.shape[0] assert (B0 == B1 == B2 == B3) self.B = B0.value self.H = tensoru.size(self.phi.overhead_features, 1) if len(self.metadata_list): self.tensor_init(**self.metadata_list.to_dict())
def _static_plot(self, minibatch): log.debug("Static plotting") sessrun = functools.partial(self.sess.run, feed_dict=minibatch) # Convert data to numpy to prepare for plotting. sampled_output_np = self.sampled_output.to_numpy(sessrun) experts_np = self.input_singleton.experts.to_numpy(sessrun) # Plot things over every channel of the BEV. for c in range( tensoru.size(self.input_singleton.phi.overhead_features, -1)): plot_bev_kwargs = { 'onechannel': True, 'channel_idx': c, 'allchannel': False } plot.plot_sampled_minibatch( sampled_output=sampled_output_np, experts=experts_np, partial_write_np_image_to_tb=self.partial_np2tb_smbs[c], figsize=self.figsize, without_samples=self.plot_without_samples, plot_bev_kwargs=plot_bev_kwargs, tensorstr='_bev-{}'.format(c)) plot_bev_kwargs = {'onechannel': False, 'allchannel': False} plot.plot_sampled_minibatch( sampled_output=sampled_output_np, experts=experts_np, partial_write_np_image_to_tb=self.partial_np2tb_smb_joint, figsize=self.figsize, without_samples=self.plot_without_samples, plot_bev_kwargs=plot_bev_kwargs, tensorstr='_bev-joint') try: feature_map = sessrun( self.model_distribution.bijection.feature_map) feature_map_fig_w = int(np.ceil(np.sqrt(feature_map.shape[-1]))) plot.plot_feature_map( feature_map=feature_map, partial_write_np_image_to_tb=self.partial_np2tb_cnn0, nrows=feature_map_fig_w, ncols=feature_map_fig_w) except AttributeError as e: log.error(e)
def plot_single_sampled_output(sampled_output, batch_index=0, fig=None, ax=None, render=True): S = sampled_output.rollout.S_grid_frame S_past = sampled_output.phi.S_past_grid_frame limit = tensoru.size(sampled_output.phi.overhead_features, 1) # Plot future. fig, ax = plot_joint_trajectory(S[batch_index], key='sampled_trajectories', render=False, marker='o', zorder=1, alpha=.4, fig=fig, ax=ax, limit=limit) return fig, ax
def plot_joint_trajectory(joint_traj, key='joint_trajectories', fig=None, ax=None, render=True, limit=100, **kwargs): """ :param joint_traj: K x A x T x d :param key: :param fig: :param ax: :param render: :param kwargs: :returns: :rtype: """ assert (tensoru.rank(joint_traj) == 4) A = tensoru.size(joint_traj, 1) for a in range(A): render_a = (a == A - 1) and render color = kwargs.get('color', COLORS[a]) if isinstance(joint_traj, tf.Tensor): single_traj = joint_traj[:, a].numpy() else: single_traj = joint_traj[:, a] kwargs.pop('color', None) fig, ax = plot_trajectory(key, single_traj, color=color, render=render_a, fig=fig, ax=ax, axis=[0, limit, limit, 0], **kwargs) assert (fig is not None) return fig, ax
def plot_joint_trajectory(joint_traj, limit, scale=1, agents=None, fig=None, ax=None, **kwargs): assert (tensoru.rank(joint_traj) == 4) if agents is None: agents = range(tensoru.size(joint_traj, 1)) for a in agents: if "color" in kwargs: color = kwargs.pop("color") else: color = cm.get_cmap("tab10").colors[a] if isinstance(joint_traj, tf.Tensor): single_traj = joint_traj[:, a].numpy().copy() else: single_traj = joint_traj[:, a].copy() if scale > 0: single_traj *= scale single_traj[..., -1] *= -1 fig, ax = plot_trajectory(single_traj, color=color, fig=fig, ax=ax, axis=limit, **kwargs) assert (fig is not None) return fig, ax
def _prepare(self, batch_shape, phi): self.t = 0 ldb = lambda x: log.info("Step {}, ".format(self.t) + x) self.step_generate_record.append({}) assert (isinstance(batch_shape, tuple)) # (B, ...), e.g. (B,); (B,K), etc. self.batch_shape = batch_shape self.batch_size = functools.reduce(operator.mul, self.batch_shape) self.S_past_car_frames = tf.cond( tf.logical_and(phi.is_training, tf.convert_to_tensor(self.rnnconf.past_perturb)), lambda: phi.S_past_car_frames_noisy, lambda: phi.S_past_car_frames) # To represent tensors packed to shape (flattened_batch * A_agents, ...), e.g. (B*K*A, ...) self.A_batch_size = self.batch_size * self.A self.rnn_state = self.rnn.zero_state(self.A_batch_size, dtype=tf.float64) self.mu_shape = self.batch_shape + (self.A, self.D) self.sigma_shape = self.batch_shape + (self.A, self.D, self.D) if len(self.batch_shape) == 2: self.batch_str = 'bk' elif len(self.batch_shape) == 1: self.batch_str = 'b' else: raise ValueError("Unhandled batch size") if len(self.batch_shape) == 2: # (B, K, A) yaws_batch = tensoru.expand_and_tile_axis(phi.yaws, axis=1, N=self.batch_shape[1]) # (BKA, 1) self.yaws_A_batch = tf.reshape(yaws_batch, (self.A_batch_size, 1)) elif len(self.batch_shape) == 1: # (BA, 1) self.yaws_A_batch = tf.reshape(phi.yaws, (self.A_batch_size, 1)) else: raise ValueError("Unhandled batch size") # Create past encodings once. if self.past_encodings is None: if self.rnnconf.past_do_preconv: ldb("Doing preconv") # 1D-convolve over the past states before plugging them into the RNN. self.preconv_W = tf.Variable(tf.ones( (self.rnnconf.preconv_horizon, self.D, self.rnnconf.past_gru_units), dtype=tf.float64), name="W_preconv") self.preconv_b = tf.Variable(1e-5 * tf.ones( (self.rnnconf.past_gru_units, ), dtype=tf.float64), name="b_preconv") past_rnn_inputs = [ tf.nn.conv1d(self.S_past_car_frames[:, a], self.preconv_W, stride=1, padding="SAME") + self.preconv_b for a in range(self.A) ] else: ldb("Not doing preconv") past_rnn_inputs = [ self.S_past_car_frames[:, a] for a in range(self.A) ] # For every agent, run the RNN cell and retrieve the state at the last time step. self.past_encodings = [ tf.nn.dynamic_rnn(cell=self.past_rnn, inputs=past_rnn_inputs[a], initial_state=None, time_major=False, dtype=tf.float64)[1] for a in range(self.A) ] # Create feature map once. if self.feature_map is None: if self.cnnconf.create_residual_connections: ldb("Creating a CNN with residual connections") if self.cnnconf.do_batchnorm: self.feature_map = tensoru.convnet_with_residuals_and_batchnorm( self.convnet, self.batchnorms, phi.overhead_features, skip_indices=set([0, self.cnnconf.n_conv_layers - 1]), is_training=phi.is_training) else: self.feature_map = tensoru.convnet_with_residuals( self.convnet, phi.overhead_features, skip_indices=set([0, self.cnnconf.n_conv_layers - 1])) else: ldb("Creating a CNN without residual connections") self.feature_map = tensoru.convnet(self.convnet, phi.overhead_features) if self.cnnconf.append_cnn_input_to_cnn_output: ldb("append_cnn_input_to_cnn_output=True") self.feature_map = tf.concat( (self.feature_map, phi.overhead_features), axis=-1) if self.cnnconf.create_overhead_feature: # Global avg, max, min. self.overhead_feature = tf.concat( (tf.reduce_mean(self.feature_map, axis=(-3, -2)), tf.reduce_max(self.feature_map, axis=(-3, -2)), tf.reduce_min(self.feature_map, axis=(-3, -2))), axis=-1) assert (len(batch_shape) > 1) tilex = functools.reduce(operator.mul, batch_shape[1:]) self.overhead_feature = tf.tile(self.overhead_feature, (tilex, 1)) self.feature_map_C = tensoru.size(self.feature_map, -1) self.radii_feat_size = self.feature_map_C * len( self.whiskerconf.radii) * self.whiskerconf.n_samples # Tile the past encodings for the current batch_shape assert (tensoru.size(phi.light_features, 0) == self.batch_shape[0]) if len(batch_shape) > 1: assert (len(batch_shape) == 2) #[(BK, F) ... ] self.past_encodings_batch = [ tensoru.expand_and_tile_and_pack(e, 1, 0, N=batch_shape[1]) for e in self.past_encodings ] # (BK, F) light_features = tf.keras.backend.repeat_elements( phi.light_features, self.lightconf.lightrep, axis=-1) light_features_batch = tensoru.expand_and_tile_and_pack( light_features, 1, 0, N=batch_shape[1]) # (BKA, F) self.light_features_A_batch = tensoru.expand_and_tile_and_pack( light_features_batch, 1, 0, N=self.A) else: # [(B, F) ... ] self.past_encodings_batch = self.past_encodings # (BA, F) light_features_batch = tf.keras.backend.repeat_elements( phi.light_features, self.lightconf.lightrep, axis=-1) self.light_features_A_batch = tensoru.expand_and_tile_and_pack( light_features_batch, 1, 0, N=self.A) self.past_encodings_joint = [] if self.A > 1: for a in range(self.A): others_feat = tf.reduce_sum(self.past_encodings_batch[:a] + self.past_encodings_batch[a + 1:], axis=0) # Create a feature for each agent that depends on its own encoding and the sum of the other agent's encodings. self.past_encodings_joint.append( tf.concat((self.past_encodings_batch[a], others_feat), axis=-1)) past_encodings_joint_shape = 2 * self.rnnconf.past_gru_units self.social_map_feat_size = 2 * self.feature_map_C else: self.past_encodings_joint = [self.past_encodings_batch[0]] past_encodings_joint_shape = self.rnnconf.past_gru_units self.social_map_feat_size = self.feature_map_C self.past_encodings_A_batch = tf.reshape( tf.stack(self.past_encodings_joint, axis=-2), (self.A_batch_size, past_encodings_joint_shape))