Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models

Requirements | Training | Results | Contributing

This repository is the official implementation of Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models.

Requirements

To install requirements:

pip install -r requirements.txt

We build our code using JAX. The code of the algorithm is in the folder dgm.

Training

Examples how to train the model from the paper are in the folder examples. To train e.g. Double Pendulum model on multiple trajectories, go to folder examples/non_parameteric/multiple_trajectories and run:

python double_pendulum.py 

Results

Our model achieves the following Log-Likelihood score on different datasets:

Dataset	Single Trajectory [LL]	Multiple Trajectories [LL]
Lotka Volterra	1.96 ± 0.21	1.81 ± 0.08
Lorenz	−0.57 ± 0.11	−2.18 ± 0.76
Double Pendulum	2.13 ± 0.14	1.86 ± 0.05
Quadrocopter	0.64 ± 0.07	-0.54 ± 0.36

Contributing

If you would like to contribute to the project please reach out to Lenart Treven or Philippe Wenk. If you found this library useful in your research, please consider citing.

@article{treven2021distributional,
  title={Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models},
  author={Treven, Lenart and Wenk, Philippe and Dörfler, Florian and Krause, Andreas},
  journal={Advances in Neural Information Processing Systems},
  year={2021}
}