Importing POMDP code in pymc3 to pymc - Code is running with erroneous results

Questions v5
1

Hi.

I have been trying to replicate the tutorial on POMDP inference (Tutorial-on-POMDP-inference-and-robust-planning/Tutorial_POMDP_robust_planning.ipynb at master · giarcieri/Tutorial-on-POMDP-inference-and-robust-planning · GitHub). The codes were developed using pymc3 and Theano. I have been trying to adapt it to pymc as below:

class HMMstates():
    def __init__(self, p_transition, init_prob,actions, n_states, *args, **kwargs):
        self.p_transition = p_transition
        self.init_prob = init_prob
        self.k = n_states
        self.actions = actions
        self.mode = pt.cast(0, dtype='int64')

    def logp(self, x):
        p_init = self.init_prob
        acts = self.actions[:-1]
        p_tr = self.p_transition[acts, x[:-1]]

        # The logp of the initial state
        initial_state_logp = pm.Categorical.dist(p=p_init).logp(x[0])

        # The logp of the rest of the states
        x_i = x[1:]
        ou_like = pm.Categorical.dist(p=p_tr).logp(x_i)
        transition_logp = pt.sum(ou_like)

        return initial_state_logp + transition_logp


class HMMGaussianEmissions():
    def __init__(self, states, mu, sigma, *args, **kwargs):
        self.states = states
        # self.rate = rate
        self.mu = mu
        self.sigma = sigma

    def logp(self, x):
        """
        x: observations
        """
        states = self.states
        # rate = self.rate[states]  # broadcast the rate across the states.
        mu = self.mu[states]
        sigma = self.sigma[states]
        #return pt.sum(pm.Normal.dist(mu=mu, sigma=sigma).logp(x))
        like = pt.sum(pm.Normal.dist(mu=mu[1:], sigma=sigma[1:]).logp(x[1:]))

        boundary_like = pm.Normal.dist(mu=mu[0], sigma=sigma[0]).logp(x[0])
        return like + boundary_like

n_states = 3
with pm.Model() as model:
    # Priors for transition matrix
    p_transition = pm.Dirichlet("p_trans", a=pt.ones((n_actions, n_states, n_states)))


    init_probs = pm.Dirichlet('init_probs', a = pt.ones((n_states,)), shape=n_states)

    # Prior for mu and sigma
    μ = pm.Normal("mu", mu=[0.5, 0.0, -0.5], sigma=1, shape=(n_states,))
    σ = pm.Exponential("sigma", lam=2, shape=(n_states,))

    # HMM state
    for i in range(number_seq):
        # HMM state
        hmm_states = HMMstates(
        p_transition=p_transition,
        init_prob=init_probs,
        actions=actions_all_seq[i, :],
        n_states=n_states,
        shape=(length_seq)
      )

    # Observed emission likelihood
        obs = HMMGaussianEmissions(
        states=hmm_states,
        mu=μ,
        sigma=σ,
        observed=np.array(emissions_all_seq[i, :]).astype("float")
     )

Based on previous suggestions, I modified the sampler

with model:
    nuts_step1 = CategoricalGibbsMetropolis([model.init_probs, model.p_trans])
    nuts_step2=pm.NUTS([model.mu, model.sigma])
    trace = pm.sample(2000, tune=3000, step=[nuts_step1,nuts_step2], chains=4, cores=4, progressbar=True)

However, I am still not able to get a good inference of the transition matrices I am trying to estimate. I have also tried the custom categorical Distribution suggested in Custom Categorical Distribution - ensure bounded candidates , but could not get a good inference as expected. Does the pymc version affect the results?

1 Like
5

Hey, @DeeJ! If I understood your post correctly, this is something I was looking into a while ago and have a working implementation that you can see here. Not sure if this works better or worse but I guess it wouldn’t hurt to check out.

I can’t recall but I think that there might be a more complete solution in pymc-extras.

1 Like
6

Thank you @Dekermanjian ! Yes, your implementation is a great help. I will look into it more closely, but at this point I want to point out the error in my previous code snippet. Now, I would like to try out the steps you followed (particularly the pm.step_methods.metropolis) and hopefully get the desired inference.

But for the code I posted, the sampler should look more like

   step = pm.CompoundStep([
       CategoricalGibbsMetropolis(hmm_states),
       pm.NUTS(vars=[model.mu, model.sigma, model.init_probs, model.p_transition])
   ])```
1 Like