Poor Inference results using FFBSStep

[jcatulo]

Hi!

I am trying to estimate the parameters of a Hidden Markov Model (HMM) with gaussian emissions. For that, I am using the model below.

N_states = 7
with pm.Model() as model:

    observations = pm.Data('data', observations)

    Pt = pm.Dirichlet("p_transition", np.ones( (N_states, N_states) ), shape=(N_states, N_states))
    P0 = pm.Dirichlet("p_init", np.ones((N_states,)), shape=(N_states,))

    mu1 = pm.Normal('mu1', mu=-28, sigma=4 )
    mu2 = pm.Normal('mu2', mu=-20, sigma=2 )
    mu3 = pm.Normal('mu3', mu=-16, sigma=2 )
    mu4 = pm.Normal('mu4', mu=-12, sigma=2 )
    mu5 = pm.Normal('mu5', mu=-8, sigma=2 )
    mu6 = pm.Normal('mu6', mu=-6, sigma=2 )
    mu7 = pm.Normal('mu7', mu=-4, sigma=2 )

    mu = tt.stack( [mu1,mu2,mu3,mu4,mu5,mu6,mu7])

    sigma = pm.HalfNormal("sigma", sigma=2, shape=(N_states,) )

    comp_dists = [ pm.TruncatedNormal.dist(mu = mu[i], sigma = sigma[i], lower=-30, upper=0) for i in range(0, N_states) ]

    Z_rv = DiscreteMarkovChain("Z_t", tt.shape_padleft(Pt), P0, shape = (22,) )
    X_rv = SwitchingProcess("X_t", comp_dists, Z_rv, observed = observations)

When I sample from the posterior without defining any step method (Metropolis sampler is automatically assigned to the DiscreteMarkovChain random variables and NUTS is assigned to the rest) the results converge. The code and some of the results are represented below.

• Code:

with model:
    posterior_trace = pm.sample( 
        return_inferencedata=True,
        chains=2,
        progressbar=True
    )

• Results:
  

However, when I define the sampling steps for the hidden states Z_t to be performed using the FFBSStep, my samples do not converge. You can find the code and some of the results below.

• Code:

with model:
    states_step = FFBSStep([model.Z_t])
    mu_step = pm.NUTS( [model.mu1, model.mu2, model.mu3, model.mu4, model.mu5, model.mu6, model.mu7, model.sigma, model.p_transition, model.p_init] )

    posterior_trace = pm.sample(
        step=[states_step, mu_step],
        return_inferencedata=True,
        chains=2,
        progressbar=True,
    )

• Results:
  

Why does this happen? Am I doing someting wrong?