Behrens' Bayesian Learner Model : How to share parameters between steps?

Questions
27

You inspired my code like this:

# Create the test data
p = np.repeat([.75, .75, .75, .25, .75, .25], 5)
env= stats.binom.rvs(1, p)
choose = list(map(bool, env))

observed_data = pd.DataFrame({"choose": choose, 
                              "index" : range(len(env))})

with pm.Model() as bayesian_lerner_model:
    k = pm.GaussianRandomWalk("k", mu = 5, sigma = 100, shape = len(env))
    k_ = pm.Deterministic('k_cap', pm.math.exp(k))
    v = pm.GaussianRandomWalk("v", mu = 0.7, sigma = k_, testval = 0.05, shape = len(env))
    v_ = pm.Deterministic('v_cap', pm.math.exp(v))

    r = []
    for ti in range(len(env)):
        if ti == 0:
            # Testvals are used to prevent -inf initial probability
            r.append(pm.Beta(f'r{ti}', mu=0.5, sigma=1))
        else: 
            r.append(pm.Beta(f'r{ti}', mu=r[ti-1], sigma=v_[ti - 1]))

    r = pm.Deterministic('r', pm.math.stack(r))
    y = pm.Bernoulli("y", p = r, observed = env)
        
    map_estimate = pm.find_MAP()

Unfortunately, I still meet the error about Initial evaluation of model.

Here is my notebook.