MissingInputError in PyMC model with Theano/PyTensor

I’m trying to write a simple TD-learning model with theano/pytensor:

import numpy as np
import pymc as pm

import pytensor
import pytensor.tensor as T

def sig(x):
    """Just a sigmoid function"""
    return 1/(1 + np.exp(x))

def Q_update(opt1, opt2, Q, lr):
    # Q is 1x4 vector
    # choice is scalar (which element of Q should be updated)
    # lr = learning rate
    # outcome : reward (1) or no reward (0)
    
    ch_opt1 = T.cast(opt1, 'int64')
    ch_opt2 = T.cast(opt2, 'int64')

    p1 = pm.Deterministic("p1", sig(Q[ch_opt1]-Q[ch_opt2])) # prob of choosing 1st response option (=better response option)
    choice = pm.Binomial("choice", n=1, p = p1) # 1 if opt1 was chosen, 0 if opt2 was chosen
    
    outcome = pm.Binomial("outcome", n=1, p = 0.2+choice*0.6) # prob is 0.8 if opt1 was chosen and 0.2 if opt2 was chosen

    chidx = ch_opt1*choice + ch_opt2*(1-choice) # Either ch_opt1 (if opt1 was chosen) or ch_opt2 (otherwise)

    Q = T.set_subtensor(Q[chidx], Q[chidx] + lr*(outcome-Q[chidx]))
        
    return Q

with pm.Model() as vbm_model:
    lr = pm.Beta('lr', alpha = 1, beta = 1) # learning rate

    Q = np.array([0,0,0,0]).astype(pytensor.config.floatX)
    # Get opt1_val and opt2_val from experiment, for now generate randomly
    opt1_values = np.random.randint(0,4,100).astype('int32')
    opt2_values = np.random.randint(0,4,100).astype('int32')
    
    Q = pytensor.shared(Q, "Q")
    opt1 = pytensor.shared(opt1_values, "opt1")
    opt2 = pytensor.shared(opt2_values, "opt2")
            
    output, updates = pytensor.scan(fn=Q_update, sequences=[opt1, opt2], outputs_info=[Q], non_sequences=[lr])

with vbm_model:        
    trace = pm.sample(2000, tune=3000, target_accept=0.9, return_inferencedata=False)

I keep getting this error:

MissingInputError: Input 0 (opt2[t]) of the graph (indices start from 0), used to compute Elemwise{Cast{int64}}(opt2[t]), was not provided and not given a value. Use the PyTensor flag exception_verbosity='high', for more information on this error.

I think it has to do with the line choice = pm.Binomial("choice", n=1, p = p1), because when I replace p1 with a number between 0 & 1, it works. Also, this error seems to occur when values of random variables depend on values that were defined elsewhere in the script (like on p1), see here or here.I still don’t know how to remedy this, tho. I’m using pymc 5.0.1, pytensor 2.8.11, python 3.9.13, numpy 1.21.5.

Any help is greatly appreciated!

You shouldn’t create PyMC variables inside a Scan.

We have a RL example that may help: Fitting a Reinforcement Learning Model to Behavioral Data with PyMC — PyMC example gallery

Also PyMC 5.0.1 is quite old already, I suggest upgrading if you can

Thanks. I would like my pymc model to be able to simulate data (like choices and outcomes) when running pm.sample() without the observed= argument. At each step, the choice depends on the previous choices and outcomes, so that’s why I put the pymc variables for choices and outcomes inside the scan. In the notebook you referenced, data is generated using numpy, and pymc is only used for inference, so essentially the model has to be coded twice which is error-prone (especially for large models), and that’s what I’d like to avoid. But in order to achieve this, I’d have to define outcome and choice as pymc variables.

Maybe this gist is of help then? ARMA-GARCH in PyMC · GitHub

It shows how a timeseries can be defined using Scan and registering the outputs as model variables. Note that you have to use .dist for variables inside the Scan. You can’t register those variables directly as model variables, only the whole scan sequences.

This relies on PyMC being able to infer the logp of each Scan step.

If this works for you, in a future release of PyMC you’ll be able to put Scans inside CustomDists like this: Support Scans in `CustomDist` by ricardoV94 · Pull Request #6696 · pymc-devs/pymc · GitHub

There’s also an example with a Discrete Markov chain floating around, which may be even more relevant for your application.

Discrete Markov Chain example is here. Note that it was written in PyMC4, so you’ll need to change aesara to pytensor in PyMC5. You could also try the discrete markov chain distribution in pymc-experimental. Here’s an example notebook showing how that works.