DRiFT: Differentiable Grid-Based Rigid-Fluid Coupling for Learning and Control

Intelligent agents, interacting with physical environments, require an accurate understanding of the consequences of their actions for efficient learning. Such agents are often trained inside simulated environments to alleviate over dependence on data, and gradients from such a simulation can help to train the agent. For this purpose, we present an end-to-end differentiable grid-based fluid simulation with strong two-way coupling with rigid bodies. In the forward pass, the solid-fluid boundary conditions are solved monolithically in a pressure linear system for stability. For backpropagation, we introduce a novel method of calculating and propagating gradients of combined fluid-solid state using adjoint method, which runs as fast as the forward solve. Our method, customized for coupling rigid bodies with inviscid fluids, is more suitable over general purpose methods like automatic differentiation, for use cases where performance is key for analyzing overall flow patterns and learning fluid properties. We demonstrate how our simulator can help to find an optimal force sequence to guide a pill through a complex fluid-filled blood vessel network, and how it can further generalize to arbitrary target states by helping to train a neural network controller. Additionally, we show the effectiveness of our differentiable simulator in isolation, in optimizing initial states which naturally evolve into desired target states. Finally, we analyze the accuracy, robustness and most importantly, the efficiency of our gradient computation method.