PyTorch implementation of Belief Propagation (BP) algorithms. This code accompanies the paper Stein variational belief propagation for multi-robot coordination. Please cite us if you use this code for your research.
The following algorithms are implemented (see the full citations below):
- Discrete Belief Propagation: This is the standard belief propagation algorithm (Wainwright & Jordan, 2008).
- Stein Variational Belief Propagation (SVBP): A particle-based belief propagation algorithm based on Stein Variational Gradient Descent (Pavlasek et al., 2024).
- Particle Belief Propagation (PBP): Importance sampling-based belief propagation (Ihler & McAllester, 2009).
- Gaussian Belief Propagation (GaBP): Based on the implementation described by Davison & Ortiz (2019). See also the excellent interactive article by the same authors.
The first three algorithms (discrete BP, SVBP, PBP) share a common base implementation, since the belief of each node can be represented by a batch of PyTorch tensors, and the standard belief propagation equations can be implemented using PyTorch operations. Gaussian BP is also implemented in PyTorch, but uses operations on Gaussian distributions which are separately implemented.
It is recommended to install in a Python virtual environment.
Install requirements:
- Python >=3.8
- PyTorch >=2.0 (tested with up to v2.2 with CUDA 12.1)
Then install this package:
pip install -e .To run the example notebooks, you will also need to install matplotlib (>=3.7) and Jupyter Notebook.
Example usage can be found in the provided Jupyter notebooks.
TODO: Document how to define factors.
- M. J. Wainwright, M. I. Jordan et al., “Graphical models, exponential families, and variational inference,” Foundations and Trends in Machine Learning, vol. 1, no. 1–2, pp. 1–305, 2008. (link)
- J. Pavlasek, J. J. Z. Mah, R. Xu, O. C. Jenkins, and F. Ramos, "Stein variational belief propagation for multi-robot coordination," in Robotics and Automation Letters (RA-L), 2024. (link)
- A. Ihler and D. McAllester, “Particle belief propagation,” in Artificial Intelligence and Statistics, 2009, pp. 256–263. (link)
- A. J. Davison and J. Ortiz, “FutureMapping 2: Gaussian belief propagation for spatial AI,” arXiv preprint arXiv:1910.14139, 2019. (link)
This code accompanies the paper Stein variational belief propagation for multi-robot coordination (Robotics and Automation Letters, 2024). If you use it in your research, please cite:
@inproceedings{pavlasek2024stein,
title={Stein Variational Belief Propagation for Multi-Robot Coordination},
author={Pavlasek, Jana and Mah, Joshua Jing Zhi and Xu, Ruihan and Jenkins, Odest Chadwicke and Ramos, Fabio},
booktitle={Robotics and Automation Letters (RA-L)},
year={2024}
}