HMC with Bernoulli variable?

[tejasramdas]

I'm trying to run HMC on a fairly complicated generative model for the sake of efficiency. Importance resampling/MH work but they don't do particularly well on reconstructing the prior distribution (i.e. inferring the right parameters for the forward simulation). Here's a minimal example of the core issue. I have a generative function that flips a coin and determine the support for another variable a.

@gen function flip_and_sample()
    flip ~ bernoulli(0.5)
    if flip
        a ~ uniform(0,1)
    else
        a ~ uniform(2,3)
    end
end

In terms of inference, I want to infer the value for flip based on the observation for a.

observation=Gen.choicemap()
observation[:a] = 2.5
(inferred_trace,_) = Gen.<inference_algorithm>(flip_and_sample,(),observation,...)

I run into an error if I try to do a single step of HMC for both flip and a

observation=Gen.choicemap()
observation[:a] = 2.5
(inferred_trace,_)=Gen.generate(flip_and_sample,(),observation)

# works
(inferred_trace,_)=hmc(inferred_trace,select(:a)

# does not work
(inferred_trace,_)=hmc(inferred_trace,select(:flip,:a)

I understand that the Bernoulli variable doesn't have a gradient but I'm wondering if there's a workaround. Does it make sense to do mh for flip and hmc for the rest? Also, is there a tutorial for HMC on Gen.dev?

[ztangent]

You unfortunately can't run HMC on discrete variables! And yes, it makes sense to instead write a MCMC algorithm which applies HMC to some variables and regular MH (or another MCMC kernel) to the discrete variables.

We don't have a HMC tutorial right now I believe -- partly because the HMC implementation in Gen.jl is far from optimized -- but as a starting point, you can try modifying the inference algorithm over the Bayesian linear regression with outliers model in this tutorial: https://www.gen.dev/tutorials/iterative-inference/tutorial#mcmc-2

In part 5 of that tutorial, there is a compound MCMC kernel called gaussian_drift_update that iterates between MH with a Gaussian proposal on some of the continuous variables, and block resimulation MH on the discrete variables (e.g. the Boolean is_outlier variables). You should be able to replace the Gaussian drift MH kernel with a call to HMC on the continuous variables.

For your example model (which we call an "undisciplined generative function" in the Gen docs), the support of a changes entirely depending on flip, so you have to be careful to resimulate both flip and a jointly (e.g. via block resimulation, or a custom joint proposal). Otherwise, if you only propose a change to flip but not to a with mh, then I believe MH will always reject (since the proposed trace ends up having negative infinite probability).