HSSM Seems to Return Priors as Posterior

[theonlydvr]

I've been working on implementing an RLDDM in HSSM following some advice from @krishnbera which so far has worked great! One of the issues I've run into though has been speed so I've tried working on setting up a pytensor version of the code so I can take advantage of the faster samplers and a GPU. To test this approach out, I've started by working with a simple RL-only case rather than the full RLDDM. However, whenever I try sampling this model, I seem to just get the priors back. I've included my code and some other model details below. Any help/advice would be greatly appreciated!

def update_Q(choice, reward, alpha, reset, Qs):
    if reset:
        Qs = pt.set_subtensor(Qs[0], 0.5)
        Qs = pt.set_subtensor(Qs[1], 0.5)
    else:
        pe = reward - Qs[(choice + 1) // 2]
        Qs = pt.set_subtensor(Qs[(choice + 1) // 2], Qs[(choice + 1) // 2] + alpha * pe)
    return Qs

def pytensor_llik(data, rl_alpha, rl_beta, response_s, feedback_s, split_by):
    # Convert to tensors
    responses_ = data[:,1].astype('int32')
    feedback_s_ = pt.as_tensor_variable(feedback_s, dtype="int32")
    response_s_ = pt.as_tensor_variable(response_s, dtype="int32")
    split_by_ = pt.as_tensor_variable(split_by, dtype="int32")
    
    # Compute Q-values
    reset = pt.concatenate([pt.zeros(1), pt.neq(pt.diff(split_by_), 0)], axis=0)
    Qs = pt.as_tensor_variable([0.5, 0.5], dtype="float64")
    Qs, _ = pytensor.scan(fn=update_Q, sequences=[response_s_, feedback_s_, rl_alpha, reset], outputs_info=[Qs])
    
    # Compute trial-wise log-likelihood
    Qs = Qs * pt.shape_padright(rl_beta,1)
    logp_actions = Qs - pt.logsumexp(Qs, axis=1, keepdims=True)
    return logp_actions[pt.arange(responses_.shape[0]), (responses_ + 1) // 2]

model_config_combined = {
    'list_params': ['rl.alpha', 'rl.beta'],
    'bounds': {
        'rl.alpha': (0, 1.0),
        'rl.beta': (0, np.inf)
    },
    'backend': 'pytensor',
    'extra_fields': ["response_s", "feedback_s", "split_by"],
    'choices': [-1, 1]
}

model = hssm.HSSM(
                data,
                model='2AB_softmax',
                model_config=model_config_combined,
                loglik=pytensor_llik,
                loglik_kind="analytical",
                # link_settings="log_logit",
                p_outlier=None,
                lapse=None,
                )

res = model.sample(draws=1200, tune=1200, chains=1, cores=4, idata_kwargs = {'log_likelihood': True}, discard_tuned_samples=False)

[AlexanderFengler]

Hi @theonlydvr,

I don't have an immediate answer looking at this, need to actually look at the likelihood.
I did see something like this before, and in that case there was a mix-up of columns so that the connection between parameters and output didn't respect the expected relation at all --> parameters didn't really affect anything, so you would get back the prior.

However can't confirm this is the case here from just looking at it. Will get back.

[theonlydvr]

That seems like it could be the source. Would that sort of thing happen if the rt column had no connection to the parameters (as would be the case in this pure RL test model)? If it would be helpful for checking things, I can also send over the test data.