# Lint as: python3 # pylint: disable=g-bad-file-header # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """Full model implementation accompanying ICML 2020 submission. "Learning to Simulate Complex Physics with Graph Networks" Alvaro Sanchez-Gonzalez*, Jonathan Godwin*, Tobias Pfaff*, Rex Ying, Jure Leskovec, Peter W. Battaglia https://arxiv.org/abs/2002.09405 """ import graph_nets as gn import sonnet as snt import tensorflow.compat.v1 as tf from learning_to_simulate import connectivity_utils from learning_to_simulate import graph_network STD_EPSILON = 1e-8 class LearnedSimulator(snt.AbstractModule): """Learned simulator from https://arxiv.org/pdf/2002.09405.pdf.""" def __init__( self, num_dimensions, connectivity_radius, graph_network_kwargs, boundaries, normalization_stats, num_particle_types, particle_type_embedding_size, name="LearnedSimulator"): """Inits the model. Args: num_dimensions: Dimensionality of the problem. connectivity_radius: Scalar with the radius of connectivity. graph_network_kwargs: Keyword arguments to pass to the learned part of the graph network `model.EncodeProcessDecode`. boundaries: List of 2-tuples, containing the lower and upper boundaries of the cuboid containing the particles along each dimensions, matching the dimensionality of the problem. normalization_stats: Dictionary with statistics with keys "acceleration" and "velocity", containing a named tuple for each with mean and std fields, matching the dimensionality of the problem. num_particle_types: Number of different particle types. particle_type_embedding_size: Embedding size for the particle type. name: Name of the Sonnet module. """ super().__init__(name=name) self._connectivity_radius = connectivity_radius self._num_particle_types = num_particle_types self._boundaries = boundaries self._normalization_stats = normalization_stats with self._enter_variable_scope(): self._graph_network = graph_network.EncodeProcessDecode( output_size=num_dimensions, **graph_network_kwargs) if self._num_particle_types > 1: self._particle_type_embedding = tf.get_variable( "particle_embedding", [self._num_particle_types, particle_type_embedding_size], trainable=True, use_resource=True) def _build(self, position_sequence, n_particles_per_example, global_context=None, particle_types=None): """Produces a model step, outputting the next position for each particle. Args: position_sequence: Sequence of positions for each node in the batch, with shape [num_particles_in_batch, sequence_length, num_dimensions] n_particles_per_example: Number of particles for each graph in the batch with shape [batch_size] global_context: Tensor of shape [batch_size, context_size], with global context. particle_types: Integer tensor of shape [num_particles_in_batch] with the integer types of the particles, from 0 to `num_particle_types - 1`. If None, we assume all particles are the same type. Returns: Next position with shape [num_particles_in_batch, num_dimensions] for one step into the future from the input sequence. """ input_graphs_tuple = self._encoder_preprocessor( position_sequence, n_particles_per_example, global_context, particle_types) normalized_acceleration = self._graph_network(input_graphs_tuple) next_position = self._decoder_postprocessor( normalized_acceleration, position_sequence) return next_position def _encoder_preprocessor( self, position_sequence, n_node, global_context, particle_types): # Extract important features from the position_sequence. most_recent_position = position_sequence[:, -1] velocity_sequence = time_diff(position_sequence) # Finite-difference. # Get connectivity of the graph. (senders, receivers, n_edge ) = connectivity_utils.compute_connectivity_for_batch_pyfunc( most_recent_position, n_node, self._connectivity_radius) # Collect node features. node_features = [] # Normalized velocity sequence, merging spatial an time axis. velocity_stats = self._normalization_stats["velocity"] normalized_velocity_sequence = ( velocity_sequence - velocity_stats.mean) / velocity_stats.std flat_velocity_sequence = snt.MergeDims(start=1, size=2)( normalized_velocity_sequence) node_features.append(flat_velocity_sequence) # Normalized clipped distances to lower and upper boundaries. # boundaries are an array of shape [num_dimensions, 2], where the second # axis, provides the lower/upper boundaries. boundaries = tf.constant(self._boundaries, dtype=tf.float32) distance_to_lower_boundary = ( most_recent_position - tf.expand_dims(boundaries[:, 0], 0)) distance_to_upper_boundary = ( tf.expand_dims(boundaries[:, 1], 0) - most_recent_position) distance_to_boundaries = tf.concat( [distance_to_lower_boundary, distance_to_upper_boundary], axis=1) normalized_clipped_distance_to_boundaries = tf.clip_by_value( distance_to_boundaries / self._connectivity_radius, -1., 1.) node_features.append(normalized_clipped_distance_to_boundaries) # Particle type. if self._num_particle_types > 1: particle_type_embeddings = tf.nn.embedding_lookup( self._particle_type_embedding, particle_types) node_features.append(particle_type_embeddings) # Collect edge features. edge_features = [] # Relative displacement and distances normalized to radius normalized_relative_displacements = ( tf.gather(most_recent_position, senders) - tf.gather(most_recent_position, receivers)) / self._connectivity_radius edge_features.append(normalized_relative_displacements) normalized_relative_distances = tf.norm( normalized_relative_displacements, axis=-1, keepdims=True) edge_features.append(normalized_relative_distances) # Normalize the global context. if global_context is not None: context_stats = self._normalization_stats["context"] # Context in some datasets are all zero, so add an epsilon for numerical # stability. global_context = (global_context - context_stats.mean) / tf.math.maximum( context_stats.std, STD_EPSILON) return gn.graphs.GraphsTuple( nodes=tf.concat(node_features, axis=-1), edges=tf.concat(edge_features, axis=-1), globals=global_context, # self._graph_net will appending this to nodes. n_node=n_node, n_edge=n_edge, senders=senders, receivers=receivers, ) def _decoder_postprocessor(self, normalized_acceleration, position_sequence): # The model produces the output in normalized space so we apply inverse # normalization. acceleration_stats = self._normalization_stats["acceleration"] acceleration = ( normalized_acceleration * acceleration_stats.std ) + acceleration_stats.mean # Use an Euler integrator to go from acceleration to position, assuming # a dt=1 corresponding to the size of the finite difference. most_recent_position = position_sequence[:, -1] most_recent_velocity = most_recent_position - position_sequence[:, -2] new_velocity = most_recent_velocity + acceleration # * dt = 1 new_position = most_recent_position + new_velocity # * dt = 1 return new_position def get_predicted_and_target_normalized_accelerations( self, next_position, position_sequence_noise, position_sequence, n_particles_per_example, global_context=None, particle_types=None): # pylint: disable=g-doc-args """Produces normalized and predicted acceleration targets. Args: next_position: Tensor of shape [num_particles_in_batch, num_dimensions] with the positions the model should output given the inputs. position_sequence_noise: Tensor of the same shape as `position_sequence` with the noise to apply to each particle. position_sequence, n_node, global_context, particle_types: Inputs to the model as defined by `_build`. Returns: Tensors of shape [num_particles_in_batch, num_dimensions] with the predicted and target normalized accelerations. """ # Add noise to the input position sequence. noisy_position_sequence = position_sequence + position_sequence_noise # Perform the forward pass with the noisy position sequence. input_graphs_tuple = self._encoder_preprocessor( noisy_position_sequence, n_particles_per_example, global_context, particle_types) predicted_normalized_acceleration = self._graph_network(input_graphs_tuple) # Calculate the target acceleration, using an `adjusted_next_position `that # is shifted by the noise in the last input position. next_position_adjusted = next_position + position_sequence_noise[:, -1] target_normalized_acceleration = self._inverse_decoder_postprocessor( next_position_adjusted, noisy_position_sequence) # As a result the inverted Euler update in the `_inverse_decoder` produces: # * A target acceleration that does not explicitly correct for the noise in # the input positions, as the `next_position_adjusted` is different # from the true `next_position`. # * A target acceleration that exactly corrects noise in the input velocity # since the target next velocity calculated by the inverse Euler update # as `next_position_adjusted - noisy_position_sequence[:,-1]` # matches the ground truth next velocity (noise cancels out). return predicted_normalized_acceleration, target_normalized_acceleration def _inverse_decoder_postprocessor(self, next_position, position_sequence): """Inverse of `_decoder_postprocessor`.""" previous_position = position_sequence[:, -1] previous_velocity = previous_position - position_sequence[:, -2] next_velocity = next_position - previous_position acceleration = next_velocity - previous_velocity acceleration_stats = self._normalization_stats["acceleration"] normalized_acceleration = ( acceleration - acceleration_stats.mean) / acceleration_stats.std return normalized_acceleration def time_diff(input_sequence): return input_sequence[:, 1:] - input_sequence[:, :-1]