RL functions using blackbox likelihoods in HSSM

[BenJonathanWagner]

Dear all,

i have a question regarding a decision-making task and the HSSM toolbox.
We have a reinforcement learning task where in some trials there is only learning of Q-values and sequence values (repetition of a policy) and in other trials participants have to make a binary choice that is modelled as a linear or non-linear function of Q-values and policy repetition and in some models other variables like e.g. a conflict term. I did fit the model in JAGS, but because of time-pressure in the task, i would like to fit the data with collapsing bounds (which is not possible in JAGS at least if you don´t build your own custom module which is beyond my capabilities).

It would be important that the model does learn different types of values in all trials, and that we can implement our own functions of how these values map on the different DDM parameters when decisions come up every 6 to 10 trials. Further the task has different conditions and days of participation which we all want to put in one hierarchical model.

I want something like:
#Learning in every trial
for(t in NT){
Q <- Q + learning_rate * reward
PolicyRep <- PolicyRep + rep_learning_rate*repetition
}

#Decision only in some trials
for(d in NDecisions){
delta[d] <- thetaQ * delta_Q + thetaRep * delta_Rep + theta_conflict * conflict_term
beta[d] <- thetaRep * delta_Rep
RT ~ wiener+collapsingbounds(alpha,tau, beta[d], delta[d], angle)
}

What would be the key steps to implement such a model in HSSM? That would be great help.
Thank you very much in advance.

Bests,
Ben

Postdoc
Chair of Cognitive Computational Neuroscience
TU Dresden

[krishnbera]

Hi,

A couple of points to implement this --

1. At a high level, you can implement this black box likelihood in Cython or JAX/PyTensor. While JAX/PyTensor implementation will allow the use of gradient-based samplers like NUTS, you will have to rely on Slice sampling if you have a Cython function. While you can also implement a Python function, it will be very inefficient in practice.

2. Black box likelihood function. Let's say we want to have a Cython implementation. The function signature should look like this --

def cython_LL(np.ndarray[double, ndim=1] subject,
                    np.ndarray[double, ndim=1] response,
                    np.ndarray[double, ndim=1] ...,
                    np.ndarray[double, ndim=1] param1_arr,
                    np.ndarray[double, ndim=1] param2_arr,
                    np.ndarray[double, ndim=1] ...,
                    ):

The function should take as inputs, all the required data columns (participant id, trial number, response, feedback, etc.) as well as the parameter columns to compute the likelihood. This function will encapsulate all the necessary likelihood logic. This function returns an array of computed logp values.

3. Since you are interested in using collapsing bound SSMs, you will have to interface your likelihood function with LANs (likelihood approximation networks) to get the amortized likelihood. To do this, you can directly access the LAN files available in onnx or pytorch format and create a simple wrapper over it that performs a forward pass to get the likelihoods. The call to this wrapper function will go inside the cython_LL function.

4. For constructing the HSSM model using this custom likelihood function, you will have to specify "extra_fields" argument in the model_config. "extra_fields" is specified as a list of columns (in the dataframe) that needs to be passed to the likelihood function. By default, HSSM will only pass choice and rt columns to the likelihood function. You can then define the black-box HSSM model as usual.

Here is a template to define the HSSM model with a custom likelihood function -

model_config = {
        "list_params": ["param1", "param2", ...],
        "extra_fields": ["feedback", "response", "participant_id", ...],
    }

model = hssm.HSSM(
                data, 
                model='custom_RL_model', 
                model_config=model_config,
                loglik=cython_logp_fn, 
                loglik_kind="black-box",
                ....
                )

Hope this helps, let me know if you have any questions.

[BenJonathanWagner]

Hi Krishn,

thank you very much for the detailed explanations, i will try to get that work.

Another question thats on my mind. Would it also be advisable to try to translate my model from JAGS to PyMC and then use the DDMCollpasing Bounds function from HSSM for fitting? Are there any downsides to this approach ? I thought that would be cool, because then I would have maximum flexibility in terms of how the individual values (repetition, Q-values) could be related to the DDM parameters, or would that be just as possible with the way you describe?

I was trying to read the tutorial on "HSSM random variables in PyMC" but the link "See the dedicated tutorial in the documentation if you are interested" is not working for me. Is it just me, or are you still working on the tutorial?

Thank you and best wishes
Ben

[krishnbera]

Hi Ben,

Re- "DDMCollapsing Bounds" function- are you referring to the functions in hssm.likelihoods? As of now, HSSM does not have such a directly callable wrapper function. Or are you referring to something else?

In any case, defining a custom likelihood function will give you full flexibility and control over how the RL process is related to the SSM parameters. You can define a custom link function over the individual values (repetition, q-values, etc.) to compute the SSM parameters.

Yes, that link is not rendered properly, we'll look into it. I don't think you'll need to rely on PyMC to build a custom model, HSSM can do the job for you in this case. I would suggest working with a custom likelihood that can be used to construct an HSSM model with loglik_kind="black-box". Re- your original question on translating the model from JAGS, no need to explicitly translate it to PyMC. Once you define a custom likelihood (and prior settings if any), HSSM can take care of the model construction.

[BenJonathanWagner]

Hi Krishn,

Thank you very much for the further detailed explanations.
Regarding my idea with PyMC and HSSM. I had understood the section in the tutorial (with considerable uncertainty) "...we can instead use HSSM to construct valid [PyMC] distributions..." to mean that I could export a distribution/likelihood function from HSSM and then use it in PyMC for fitting. But then I had misunderstood... I do not have much knowledge about software implementation and Python packages, so I was probably on the wrong track.

I will try to get it work via the custom likelihood function!
If its okay I might reach out to you if get stuck in the process?

Thank you for now and best wishes,
Ben

[krishnbera]

Hi Ben,

Sure, happy to help!
Please feel free to reach me at [email protected] if you have specific queries.

Best,
Krishn