Repository: pyro-ppl/pyro

Issue:
[Bug] Infinite loop in `HMC._find_reasonable_step_size`
### 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. 

<img width="450" alt="Screen Shot 2022-11-18 at 11 34 59 PM" src="https://user-images.githubusercontent.com/9263852/203143099-00723668-88f4-4e34-b9f2-97e8f7a5ae58.png">

### 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](https://github.com/pyro-ppl/pyro/files/10060315/pyro_debug.pt.zip)

### Quick fix
For a hot-fix, I updated the following line
https://github.com/pyro-ppl/pyro/blob/8b7e5641228b664b3c9b674a079e95a478f3e2a0/pyro/infer/mcmc/hmc.py#L193
with `while direction_new == direction and step_size > 1e-100:` which prevents the infinite loop.

cc @dme65, @balandat



Discussion:
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?

Resolving patch:
diff --git a/pyro/infer/mcmc/hmc.py b/pyro/infer/mcmc/hmc.py
--- a/pyro/infer/mcmc/hmc.py
+++ b/pyro/infer/mcmc/hmc.py
@@ -66,6 +66,8 @@ class HMC(MCMCKernel):
         step size, hence the sampling will be slower and more robust. Default to 0.8.
     :param callable init_strategy: A per-site initialization function.
         See :ref:`autoguide-initialization` section for available functions.
+    :param min_stepsize (float): Lower bound on stepsize in adaptation strategy.
+    :param max_stepsize (float): Upper bound on stepsize in adaptation strategy.
 
     .. note:: Internally, the mass matrix will be ordered according to the order
         of the names of latent variables, not the order of their appearance in
@@ -108,6 +110,9 @@ def __init__(
         ignore_jit_warnings=False,
         target_accept_prob=0.8,
         init_strategy=init_to_uniform,
+        *,
+        min_stepsize: float = 1e-10,
+        max_stepsize: float = 1e10,
     ):
         if not ((model is None) ^ (potential_fn is None)):
             raise ValueError("Only one of `model` or `potential_fn` must be specified.")
@@ -119,6 +124,8 @@ def __init__(
         self._jit_options = jit_options
         self._ignore_jit_warnings = ignore_jit_warnings
         self._init_strategy = init_strategy
+        self._min_stepsize = min_stepsize
+        self._max_stepsize = max_stepsize
 
         self.potential_fn = potential_fn
         if trajectory_length is not None:
@@ -188,9 +195,11 @@ def _find_reasonable_step_size(self, z):
         step_size_scale = 2**direction
         direction_new = direction
         # keep scale step_size until accept_prob crosses its target
-        # TODO: make thresholds for too small step_size or too large step_size
         t = 0
-        while direction_new == direction:
+        while (
+            direction_new == direction
+            and self._min_stepsize < step_size < self._max_stepsize
+        ):
             t += 1
             step_size = step_size_scale * step_size
             r, r_unscaled = self._sample_r(name="r_presample_{}".format(t))
@@ -206,6 +215,8 @@ def _find_reasonable_step_size(self, z):
             energy_new = self._kinetic_energy(r_new_unscaled) + potential_energy_new
             delta_energy = energy_new - energy_current
             direction_new = 1 if self._direction_threshold < -delta_energy else -1
+        step_size = max(step_size, self._min_stepsize)
+        step_size = min(step_size, self._max_stepsize)
         return step_size
 
     def _sample_r(self, name):


Test patch: