Use of JuliaPOMDP with cognitive modelling package

[PTWaade]

Hey developers of JuliaPOMP !

I am working on a package for cognitive modelling (i.e. the field of reverse-engineering which models how people use to solve tasks), and in that context, I wanted to see if I could rely on your extensive ecosystem for POMDP-related models.
To test that, I just need an example with a simple updater/solver for a simple POMDP (notably, I do not necesarily need to rely on your POMDP tasks or environments).
However, I need to be able to run the action-perception loop manually (because I need to intervene in it), so I can’t use your otherwise very useful simulation functions.
I have been looking around, but have found it a little difficult to figure out how to do this naturally.
I wondered if you would be kind enough to provide me with a brief dummy code example looking something like this:

using ...

#create pomdp

#create updater

#create solver

#initialise beliefs

for i in 1:n_timesteps
   #update belief based on observation (and previous action)
   #choose action based on belief
end

Notably, it would be a great advantage if the loop were fully differentiable (preferably with ForwardDiff.jl, although that requires support for Dual Numbers, so alternatively Mooncake.jl or Enzyme.jl).

Additionally, I would much prefer it if it were possible to return not just an action, but a probability distribution over actions (or something from which I could generate such a distribution). Then the action could be sampled from that afterwards.

And finally, I would also be happy if I could construct a simple interface for setting parameters of the POMDP solver and belief updater (since those are the ones I would try to estimate for a given experimental subject, based on their behaviour in an experiment).

Let me know if there is anything I can do - and lots of admiration for your incredible work!

[PTWaade]

I am likely to be particularly interested in temporal difference models and SARSA models, which I think can be found in the https://github.com/JuliaPOMDP/TabularTDLearning.jl repository :)

[dylan-asmar]

Hello @PTWaade! The process to make your own simulation loop isn't too difficult. The main idea is to sample from the distributions that return from your functions, or create a generative function and use the @gen macro.

However, I don't know if you could make it differentiable.

For your POMDP, are you defining it explicitly (defining the functions) or using a more generative approach and targeting online methods?

Here is the basic flow:

using POMDPs
using POMDPTools
using POMDPModels # To get TigerPOMDP
using POMDPSolve # To get a policy

using Random

pomdp = TigerPOMDP()
solver = POMDPSolveSolver()
policy = solve(solver, pomdp)
up = DiscreteUpdater(pomdp)

seed = 42
rng = MersenneTwister(seed)
max_steps = 10


# Using the `@gen` macro
b = initialstate(pomdp)
s = rand(rng, b)

history_vec = []
for t in 1:max_steps
    a = action(policy, b)
    
    (sp, r, o) = @gen(:sp, :r, :o)(pomdp, s, a, rng)
    
    bp = update(up, b, a, o)
    
    push!(history_vec, (s, a, r, sp, o))
    
    # If you do any modification of the state or belief, 
    # you might want to use `deepcopy` here instead
    s = sp 
    b = bp
    
    if isterminal(pomdp, sp)
        break
    end
end


# Sampling from the distributions
b = initialstate(pomdp)
s = rand(rng, b)

history_vec = []
for t in 1:max_steps
    a = action(policy, b) # Get action
    r = reward(pomdp, s, a) # Get reward
    transition_dist = transition(pomdp, s, a) # Get transition distribution
    sp = rand(rng, transition_dist) # Sample next state
    observation_dist = observation(pomdp, s, a, sp) # Get observation distribution
    o = rand(rng, observation_dist) # Sample observation

    bp = update(up, b, a, o) # Update belief

    push!(history_vec, (s, a, r, sp, o))

    s = sp
    b = bp

    if isterminal(pomdp, sp)
        break
    end
end


# If you want a policy to return a distribution over actions, you can define your own
# action function. Here is an example that takes a policy and returns a distribution over
# the softmax of the action-values.
# Note: the POMDPs.jl ecosystem is designed for the `action` function to return a single
# action and not a distribution over actions. Though if you aren't using tools that
# depend on `action`, then it shouldn't be a problem.
struct CustomPolicy <: Policy
    policy::AlphaVectorPolicy
    λ::Float64
end

function POMDPs.action(p::CustomPolicy, b)
    q = actionvalues(p.policy, b)
    q_softmax = exp.(p.λ .* q) ./ sum(exp.(p.λ .* q))
    # Return a SparseCat distribution
    return SparseCat(1:length(q_softmax), q_softmax)
end

# Here is what the simulation loop might look like now
b = initialstate(pomdp)
s = rand(rng, b)

custom_policy = CustomPolicy(policy, 1.0)

history_vec = []
for t in 1:max_steps
    action_dist = action(custom_policy, b) # Get action distribution
    a = rand(rng, action_dist) # Sample action
    
    # The rest is the same as before. You could use the `@gen` macro here as well.
    r = reward(pomdp, s, a) 
    transition_dist = transition(pomdp, s, a)
    sp = rand(rng, transition_dist) 
    observation_dist = observation(pomdp, s, a, sp)
    o = rand(rng, observation_dist)

    bp = update(up, b, a, o)

    push!(history_vec, (s, a, r, sp, o))

    s = sp
    b = bp

    if isterminal(pomdp, sp)
        break
    end
end

[PTWaade]

Dear @dylan-asmar

Thank your for this incredibly useful response! It is very appreciated.

I have been able set up two simple example models which fits your example policy, as well as a Q-learning policy from TabularTDLearning.jl, to behavioural data.
There are a few things - well, probably just one more thing - that I would need to make this fully functional. If you would look at it - whenever you have time - I would be very grateful. If we can make this work, I can make JuliaPOMDPs extensive ecosystem available to cognitive modellers - and that would be quite useful.
Note that, because this is to be used with behavioural experiments where we have pre-recorded the actions and observations of the participant, we don't need to simulate the POMDP environment here. Simulating the environment is useful for other things (such as parameter recovery), but in that case, I would probably package reward, transition and observe inside a single function.

After the two model examples, I'll list the last things I'd need for this to be directly applicable. Most important, however, is the point that cognitive modelling often targets the learning process itself, and aim to estimate parameters such as the learning rate of the Q-learning policy. This means that the learning loop that I presume that place inside solve would have to be unpacked, which I would need help with.

See below :)

Functional Turing model with your example policy

For this toy example, imagine a participant in a psychological experiment has played the Tiger POMDP game as a task. We therefore already know the observations and rewards the participant received, and the actions they took. This means we do not need to simulate the POMDP-environment's transitions and observation-generation (this was done during the execution of the experiment). What we are interested in, however, is reconstructing the parameters of the participant's strategy for solving the task.

using POMDPs
using POMDPTools
using POMDPModels      # To get TigerPOMDP
using POMDPSolve         # For standard POMDP solver

using LogExpFunctions  #For softmax function

using Turing                   #For fitting to behavioural data with MCMC etc
using Distributions         #For distributions comaptible with Turing


## Setup ##

#Create pomdp object (Tiger pomdp)
pomdp = TigerPOMDP()
#Create an updater object for the pomdp
updater = DiscreteUpdater(pomdp)
#Instantiate a solver type
solver = POMDPSolveSolver()
#Apply the solver to the pomdp to get a trained policy
policy = solve(solver, pomdp)


## Turing model ##

#Simple model for fitting to one set of behavioural data
@model function pomdp_model(
    policy,
    updater;
    actions = [1, 1, 1],                            # Empirical actions
    observations = [true, false, false], # Empirical observations
    rewards = [10.0, 10.0, 10.0],           # Empirical rewards
)

    ## Sample parameters ##
    #Inverse temperature parameter for softmax (i.e. action noise)
    β ~ Exponential(0.1)
    #Initial belief about environment
    initial_belief ~ LogitNormal(0.0, 0.1)

    ## Simulate behaviour ##
    #Initialise belief about environment
    belief = BoolDistribution(initial_belief)

    #For each trial in the data
    for t in 1:length(actions)

        #Get expected values for each possible action, given belief and policy
        q = actionvalues(policy, belief)

        #Transform action values into a probability distribution
        action_probabilities = softmax(q * β)

        #Create a Categorical distribution from the action probabilities
        action_dist = Categorical(action_probabilities)

        #Condition the model on the empirically observed action
        actions[t] ~ action_dist

        #Update belief based on the empirical observation and action
        belief = update(updater, belief, actions[t], observations[t])
    end
end


## Fit model to data ##

#Sample model with metropolis-hastings sampler
chains = sample(pomdp_model(policy, updater), MH(), 1000)

#See results
describe(chains)

#Plot results
using StatsPlots
plot(chains)

As you anticipated, the model isn't differentiable (FiniteDifferences works, but is slow and runs into numerical issues).
I might make new issues with Enzyme and Mooncake to ask what it would take to make the model differentiable - it would be great to be able to use gradient-based samplers (but I can also make do without for now).

#Loading Enzyme, Mooncake and FiniteDifferences as AD backends
using ADTypes: AutoEnzyme, AutoMooncake, AutoFiniteDifferences
import Enzyme: set_runtime_activity, Forward, Reverse
import Mooncake
import FiniteDifferences: central_fdm

## Sample with NUTS (requires autodiff) ##
# Mooncake errors
chains = sample(
    pomdp_model(policy, updater),
    NUTS(; adtype = AutoMooncake(; config = nothing)),
    1000,
)

# Enzyme errors
chains = sample(
    pomdp_model(policy, updater),
    NUTS(; adtype = AutoEnzyme(; mode = set_runtime_activity(Forward, true))),
    1000,
)
chains = sample(
    pomdp_model(policy, updater),
    NUTS(; adtype = AutoEnzyme(; mode = set_runtime_activity(Reverse, true))),
    1000,
)

# FiniteDifferences works but often runs into numerical issues
chains = sample(
    pomdp_model(policy, updater),
    NUTS(; adtype = AutoFiniteDifferences(; fdm = central_fdm(5, 1))),
    1000,
)

Functional Turing model with TabularTDLearning.jl

It took me a little while to figure out how I had to adapt your original example to the MDP case where state = observation = belief, and where actionvalues takes a vector instead of a probability distribution, and transition and observe and the updater isn't necessary. But it works - and is fully differentiable (!).
For this toy example, we might imagine that a participant in an experiment is playing a kind of GridWorld game as experiment, and we hypothesize that they play this in a way that can be described as doing Q-learning.
The actions here are the moving on the grid, and the observation is (in this case at least) only the position on the grid, and the rewards.
This kind of modelling is very abundant in cognitive modelling, so getting this to work in general would be very exciting indeed!

using TabularTDLearning

## Setup ##

#Create mdp object (simple grid world))
mdp = SimpleGridWorld()

#Set the exploration policy for the learning phase
exppolicy = SoftmaxPolicy(mdp, 1)
#Create the Q-learning solver
solver = QLearningSolver(
    exploration_policy = exppolicy,
    learning_rate = 0.1,
    n_episodes = 5000,
    max_episode_length = 50,
    eval_every = 50,
    n_eval_traj = 100,
)
#Train to get a policy
policy = solve(solver, mdp)


## Turing model ##

#Turing model for fitting to behavioural data
@model function mdp_model(
    policy,
    actions = [1, 3, 1],
    #These would be the positions based on the actions above, assuming that index 1 is row and index 2 is column
    observations = [[2,3], [1,3], [1,4]],
    rewards = [0.0, 0.0, 0.0],
)

    ## Sample parameters ##
    #Inverse temperature parameter for softmax
    β ~ Exponential(0.1)

    ## Simulate behaviour ##
    #Initialse belief about environment (which is also true starting point: [2,2]
    belief = [2,2]

    #For each timestep in the data
    for t = 1:length(actions)

        #Get expected values for each action, given belief and policy
        q = actionvalues(policy, belief)

        #Transform action values into a probability distribution
        action_probabilities = softmax(q * β)

        #Create a Categorical distribution from the action probabilities
        action_dist = Categorical(action_probabilities)

        #Condition on the empirically observed action
        actions[t] ~ action_dist

        #Update belief based on the empirical observation
        belief = observations[t]
    end
end


## Fit model ##

#This works with every relevant AD backend

#Sample with metropolis-hastings (works)
chains = sample(mdp_model(policy), MH(), 1000)
# Mooncake works
chains = sample(
   mdp_model(policy),
    NUTS(; adtype = AutoMooncake(; config = nothing)),
    1000,
)
# Enzyme works
chains = sample(
    mdp_model(policy),
    NUTS(; adtype = AutoEnzyme(; mode = set_runtime_activity(Forward, true))),
    1000,
)
chains = sample(
    mdp_model(policy),
    NUTS(; adtype = AutoEnzyme(; mode = set_runtime_activity(Reverse, true))),
    1000,
)
#ForwardDiff works 
chains = sample(
    mdp_model(policy),
    NUTS(; adtype = AutoForwardDiff()),
    1000,
)
#ReverseDiff works
chains = sample(
    mdp_model(policy),
    NUTS(; adtype = AutoReverseDiff()),
    1000,
)
# FiniteDifferences works but often runs into numerical issues
chains = sample(
    mdp_model(policy),
    NUTS(; adtype = AutoFiniteDifferences(; fdm = central_fdm(5, 1))),
    1000,
)

Outstanding questions

Clarification:
Firstly, I just want to make sure I've understood the workflow correctly.
Is it understood correctly that the policy = solve(solver, pomdp) call essentially runs the strategy many times on the task and optimizes some kind of parameters so that the agent maximizes the reward? The returned policy is then a pre-trained strategy specifically for the given task, with mappings form (beliefs about) states to expectaed values over actions, is that right? That would be why the policy in the Tiger POMDP example is almost deterministically certain to pick option 1, I assume, which seems to be where the reward is, without needing to do any information gathering: it knows this is where the reward is because it's been trained on the task.
Secondly, in my code here I try to keep the agent and the environment entirely separate (because they are separate in a cognitive modelling workflow). This means that I distinguish between the state of the environment, and the agent's belief about the state of the environment. The b parameter in your oriignal example seems to play both roles (i.e. as a belief about the environment in actionvalues, and as the probability distribution from which to sample the initial state of the environment). Is that correct?

Generality:
I also would like to make sure that the workflow I am creating is general: would the above model work with any pomdp/mdp and any solver, in principle? I realise that there probably has to be two types of models for POMDPs and MDPS, but it would be wonderful if the rest was general.
For example, is actionvalues(policy, belief) pomdp- and solver-general? How about update(updater, belief, actions[t], observations[t])?
I also noticed that using the TabularTDLearning models with the TigerPOMDP caused an error: I presume it is because these methods aren't defined for when there is uncertainty about the environment. It strikes me, however, that I would expect these methods also to work when the environment is only partially observable: do you know if there are Q-learning and SARSA methods that work for POMDPs?

Modelling learning:
Most importantly, as I mentioned, the target of cognitive modelling is often to model how participants learn about their environments, which means that the learning process over multiple attempts at the game is important to include in the model. Here it seems to happen inside solve, however, and so I would need to unpack that call in order to be able to estimate parameters like the Q-learning learning rate or the temperature of the softmax exploration policy.
Is there any chance that that would be possible?
As an alternative (which would also be interesting, but less generally applicable), I could also estimate the results of the learning for a given participant (essentially their state-reward mappings, for example). This would require that I sample these values in the same way that I sample the initial belief in the first example and input them into the policy, somehow.

Modelling reward functions:
Finally, and reminiscent of inverse reinforcement learning, a goal is sometimes in cognitive modelling to infer the loss function of the participant (i.e. how much they value different states of the environment). I am unsure how to do this, because in this framework the reward are often given directly by the environment, but I would expect that I could estimate parameters of the loss function without too much trouble, if the rest works out.

Again - thank you for the help. I appreciate it, and hope to be able to do something very useful for the cognitive modelling community if this ultimately works.
(This would make it very easy to ask questions like whether clinical diagnoses such as schizophrenia affects learning rates or reward functions of participants, or to investigate which type of learning functions people seem to use).

[PTWaade]

Just a soft ping on this issue :)
I know you must have quite limited time resources, so no pressure form my side.
It would be wonderful if I could unpack the learning part in solve, so of all the above questions, that would be what matters most (in case you want to prioritise time).

Otherwise, just the warmest wishes!

[zsunberg]

Hi @PTWaade

Unfortunately, I don't have time to delve in deeply!

I think one thing that is important to understand is our distinction between "state uncertainty" and "model uncertainty". A good reference for how we think about these things is here: https://algorithmsbook.com/decisionmaking/#outline

POMDPs.jl is focused on state uncertainty and NOT on model uncertainty. Most basic reinforcement learning focuses on model uncertainty first and not on state uncertainty.

Basic RL algorithms like td learning, are designed for MDPs with unknown parameters and they gradually learn a policy in the face of model uncertainty by interacting with the environment (Part III in the book). This is why TabularTDLearning works only for MDPs. In contrast, the POMDP paradigm that POMDPs.jl uses (Part IV in the book) assumes a probabilistic model is known, but that the state is not directly observable. Since the probabilistic model is known, we can use a Bayesian updater to update the belief about the state (the update function), and we can find a near optimal policy without repeated interactions with the environment that are used in reinforcement learning. Hence, there is no standard breakdown of the solve function into learning episodes or the like.

Hopefully this helps understand the package. We might be able to add more if we understand your cogsci use case better (I was not able to understand exactly how your Turing model would fit in with the human trial. Otherwise, it is not difficult to implement your own versions of these simple TD algorithms tailored to your use case (see e.g. https://github.com/CU-ADCL/CU-DMU-Materials/blob/master/notebooks/100-SARSA.jl for a pluto notebook).

Another unrelated note: For differentiability, it is sometimes easier to use implicit differentiation for these types of optimizers. Then you do not need to differentiate through the entire solution process, but only through a potential function related to the solution itself. One package for this is https://github.com/byuflowlab/ImplicitAD.jl - we do not have it integrated into POMDPs.jl, but it might be possible.