Python code for the paper A Blob Method for Mean Field Control with Terminal Constraints by Katy Craig, Karthik Elamvazhuthi, and Harlin Lee.
BlobOT.py contains the most important class and function definitions for the optimization problem.
-
Modelis (more or less) a dummy class so that we can leverage the autograd functionality of Pytorch. -
KEcalculates the kinetic energy using velocity of the trajectories. -
KE_accuses the acceleration of trajectories. - Use
NLEto calcualte the non-local energy term (without$m_1^2$ , so this may be negative) when the target distribution is an empirical distribution. We assume equal weights per particle, but that should be easily modifiable if you need to. -
NLE_accis for the acceleration control example. It needs special treatment to deal with position and velocity becoming a$d+1$ -dimensional vector together. -
NLE_gaussis the contiuum Gaussian distribution version ofNLE. -
NLE_conscalculates$m_1^2$ , the term that's always positive and doesn't affect the optimization problem. For large$N$ , it would make sense to modify the code such that this is calculated only once. -
NLE_posandNLE_acc_poscome in handy if you want to calculate the non-negative version of the NLE, e.g. for a log-log plot. -
obstaclecalculates how much cost is incurred by the current trajectory with respect to one circular obstacle.
The following notebooks can be run after cloning this repo and installing Pytorch. Every figure, npz and pkl file created by this repo should be in the output folder. If any are missing (such as small_eps_{date}.npz), it's because the file was too large for github.
loss_contour.ipynbdiff_eps.ipynb: This may produce a different output due to randomness, but it should be qualitatively the same.small_eps.ipynb: Run this file afterdiff_eps.ipynbbecause this uses the same source and target distributions as the previous experiment.delta_scaling.ipynbPOT_comparison.ipynb: This may produce a different output due to randomness.continuum_gaussian.ipynb: The empirical distribution version may produce a different output due to randomness.obstacles.ipynbacceleration.ipynb: The shuffled velocity experiment may produce a different output due to randomness.
Helpful tips:
-
blobLossandtraining_loopare written so that they are easily adaptable to other examples. -
zandware the source and target distributions, respectively. -
ydenotes$x(1)$ , i.e. particle positions at$t=1$ . - Even though we use the method
torch.optim.SGD, it's actually GD and not SGD. -
X0andXinside the model is of length$L = 1/\Delta t$ in the time dimension and does not include$z = x(0)$ . They're of the shape$N \times d \times L$ , orparticlesxdimensionxtime steps. - The use of npz vs. pkl is pretty random.
