[Bug] Infinite loop in `HMC._find_reasonable_step_size`

[saitcakmak]

Issue Description

In some edge cases, the while loop in HMC._find_reasonable_step_size never exits and becomes an infinite loop. This happens because the direction never changes and the step_size eventually becomes 0. The below printout of t & step_size shows an example of this.

Environment

For any bugs, please provide the following:

• OS and python version: CentOS 8, Python 3.8
• PyTorch version, or if relevant, output of pip freeze. Latest main.
• Pyro version: output of python -c 'import pyro; print pyro.__version__' 1.5.2 (checked that the error is not fixed in latest)

Code Snippet

Unfortunately, I do not have a clean code snippet that reproduces the issue. I have extracted the HMC object and the input z that produces the infinite loop and used torch.save to save them. With the attached output, run

import torch
save_dict = torch.load("/home/saitcakmak/pyro_debug.pt")  # update the file path
save_dict["hmc"]._find_reasonable_step_size(save_dict["z"])

pyro_debug.pt.zip

Quick fix

For a hot-fix, I updated the following line

pyro/pyro/infer/mcmc/hmc.py

Line 193 in 8b7e564

while direction_new == direction:

with while direction_new == direction and step_size > 1e-100: which prevents the infinite loop.

cc @dme65, @Balandat

[fritzo]

Thanks for the detailed report @saitcakmak. What do you think should be ideal behavior here? Hanging is certainly bad, but stepsize=1e-100 seems bad also. Maybe we should raise a ValueError to surface problem early?

[fritzo]

I'm trying to push another release out today or tomorrow, so I'll just add lower and upper bounds say 1e-10 < step_size < 1e10. Please LMK if you want different thresholds in the release.

#3157

[saitcakmak]

Thanks for the quick patch @fritzo! I am not sure what an ideal solution here would look like. Raising an error is also a good option. I don't have a preference either way, as long as it doesn't hang :)

I ran our (BoTorch) unit tests as a quick check and the step size doesn't seem to get below 1e-10. In one other instance, I've seen the step size starting from 4.5e-16 and increasing a bit from there.

This is from the same setup that produced the hanging behavior reported above. FWIW, the covariance matrix here is almost singular (half of the time, it would just raise a LinAlgError in Cholesky). So, that could possibly explain the small step size, but I am just speculating here.