Mu of pm.Normal() depends on output of RandomVariable on previous time step. How do I implement this?

Questions v5
,
7

It turns out Jesse was right - I can just use the normal parameterised by des and sig. I tried this ages ago but the posterior distributions indicated that it wasn’t working, and there were lots of divergences, so I assumed that it just wasn’t doing what I expected. It turns out I just needed to sample a lot more than I was.

Thanks to Jesse and Ricardo for your input!

Thusly, for anyone that might find this useful, my working code is:

def lik(self,rot,nSteps,obs):
    with pm.Model() as model:
        
        def step(*args):
            desired,sig,idx,executionNoise,exploratoryNoise,lr,rot,desList = args
            prevAim = pm.Normal.dist(mu=desired,sigma=sig)
            prevRew = T.switch(T.le(T.abs(prevAim - rot),3),1,0)
            des = T.switch(T.eq(prevRew,1),desired + ((prevAim - desired) * lr) % 360,desired)
            des = T.switch(T.ge(des,180),des -360,des)
            sig = T.switch(T.eq(prevRew,1),executionNoise,executionNoise+exploratoryNoise)
            aimUpdate = collect_default_updates(inputs=args,outputs=[des,sig,prevAim])
            T.set_subtensor(desList[idx], des, inplace=True)
            idx+=1
            return (des,sig,prevAim,idx), aimUpdate 
        
        def generatePredictions(desiredCurr,sigma,idx,exectutionNoise,exploratoryNoise,lr,rot,desList,nSteps,size):
            outputs_info,_ = pt.scan(fn=step,
                                                    outputs_info=[desiredCurr,sigma,None,idx],
                                                    non_sequences=[executionNoise,exploratoryNoise,lr,rot,desList],
                                                    strict=True,
                                                    n_steps=nSteps)
            
            desired,sig,sampledAims,_ = outputs_info
            return sampledAims
        
        nSteps = pm.intX(pm.ConstantData("nSteps", nSteps))
        rot = T.as_tensor_variable(rot)
        
        executionNoise = pm.Uniform(name='executionNoise',lower=0.01,upper=100.0)
        exploratoryNoise = pm.Uniform(name='exploratoryNoise',lower=0.0,upper=10000.0)
        lr = pm.Uniform(name='lr',lower=0.0,upper=1.0)
        
        observations = pm.Data("observations", obs)
        desiredCurr = np.array([0.]).astype('float64')[0]
        sigma = np.array([0.01]).astype('float64')[0]
        usedAim = np.array([0.]).astype('float64')[0]
        outcome = np.array([0]).astype('int8')[0]
        idx = np.array([0]).astype('int8')[0]
        desList = T.empty(nSteps)
        
        lik = pm.CustomDist("lik",desiredCurr,sigma,idx,executionNoise,exploratoryNoise,lr,rot,desList,nSteps,
                                    dist=generatePredictions,observed = observations)
        
        states = pm.Deterministic("states", desList)
        return model

So the posterior parameter distributions are now correct:
download
Notice in the code and the plot how I have attempted to record the state (the desired aim) of the model on each timestep, but clearly I’ve failed to do that correctly as when I plot the mean against the observations, I get:
download
or even this, for a different participant, which I do not understand at all:
download

Nonetheless, the parameter distributions are roughly what is expected, given the data, so I’ll just trust that the model is working as intended and hope that subsequent model comparisons with my other two models will indicate whether or not this is the case.

If anyone reading this does know a simple way to plot the state as I intend to, please let me know. If I just plot the result of my CustomDist (“lik”), then it perfectly matches the observations, so that’s also not what I want.