Training a machine learning model for underwater chemical source localization in simulated turbulent flows

Abstract

Underwater chemical source localization is challenging due to the dynamic and chaotic processes involved. Averaged across long time scales, the geometry of the chemical plume is determined by the mean water flow. At shorter time scales, the turbulence of water tends to create swirls, eddies, and vortices, preventing the observation of a smooth gradient of chemical concentration. Instead, the chemical concentration in the plume downstream from a source is intermittent with mostly low-level concentrations interspersed with short high-concentration segments. Various underwater platforms could deploy chemical sensors to sample the chemical concentration and measure the water flow as they move. By traveling with a predefined trajectory, the sensors can collect observations at different positions. However, these observations may consist of only a few non-zero chemical concentration measurements along the path through the turbulent plume. It is non-trivial to process these measurements to recognize the geometry of the chemical plume and predict the chemical source’s location. In order to predict the location of any chemical sources, we train recurrent neural networks to process the input time series jointly consisting of the chemical concentration observations, water flow measurements, and sensor platform movements. From there, the neural network model constructs a heatmap that represents the probability that the chemical source is located at different locations around the sensing platform. This heatmap is trained based on simulations where a sensor platform moves along different trajectories across numerous scenarios of various source locations, water flows, and turbulence characteristics. In each simulated trajectory, the heatmap at each time step (after an initial non-zero chemical concentration measurement) is compared to the true source location using a Wasserstein distance metric as the loss function. This encourages the heatmap to minimize the expected distance given the source localization predictions and the true source location, which is known during simulation. Since Wasserstein distance keeps the geometries of distributions in consideration and it does not require the support of distributions to be the same, it provides an additional advantage in comparison to the traditional cross-entropy-based loss functions. Thus, when the source is out of the prediction range, the heatmap can still be useful to predict the direction of the chemical source location which respects the sensor platform’s current location. Additionally, we show that the expected Wasserstein distance for cases where no chemical is detected leads to a regularization term that shrinks the variance of source localization predictions. In order to train and test our methodology, we created a particle-based turbulence simulation based on prior work. The simulation models the Spatio-temporal variation in water flow along with the diffusion of the dissolved chemical. In every simulation episode, the source location is randomized radially symmetric around the sensor platform. At each time step, the sensor platform moves at a fixed speed for a predefined number of steps. To assess performance, we measure the resolution-accuracy tradeoff of the heatmap prediction under various water flow characteristics. The results indicate the potential for predicting chemical source locations from chemical sensor readings from limited observations in turbulent environments.