JAX, M.D.: A Framework for Differentiable Physics
We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional code. Since the simulations themselves are differentiable functions, entire trajectories can be differentiated to perform meta-optimization. These features are built on primitive operations, such as spatial partitioning, that allow simulations to scale to hundreds-of-thousands of particles on a single GPU. These primitives are flexible enough that they can be used to scale up workloads outside of molecular dynamics. We present several examples that highlight the features of JAX MD including: integration of graph neural networks into traditional simulations, meta-optimization through minimization of particle packings, and a multi-agent flocking simulation. JAX MD is available at www.github.com/google/jax-md.
Code (1)
Tasks
Drug DiscoveryGPUSimilar Papers 제목 키워드 기반
Differentiable Physics Simulation
Differentiable physics simulation is a powerful family of new techniques that applies gradient-based methods to learning and control of physical systems. It enables optimization for control, and can also be integrated in…
Scalable Differentiable Physics for Learning and Control
Differentiable physics is a powerful approach to learning and control problems that involve physical objects and environments. While notable progress has been made, the capabilities of differentiable physics solvers rema…
Dynamic Visual Reasoning by Learning Differentiable Physics Models from Video and Language
In this work, we propose a unified framework, called Visual Reasoning with Differ-entiable Physics (VRDP), that can jointly learn visual concepts and infer physics models of objects and their interactions from videos and…
counterfactualVisual ReasoningComplex Locomotion Skill Learning via Differentiable Physics
Differentiable physics enables efficient gradient-based optimizations of neural network (NN) controllers. However, existing work typically only delivers NN controllers with limited capability and generalizability. We pre…
DiversityDifferentiable Chemistry in PINNs for Solving Parameterized and Stiff Reaction Systems
From neural ODEs to continuous-time machine learning, differentiable solvers allow physics, optimization, and simulation to become trainable components within deep learning systems. This has opened the path to a new gene…