Reinforcement learning: a sequential family from scratch

frmtmb has no reinforcement-learning family. This page writes one and fits it hierarchically through the ordinary formula grammar. The whole family is under 200 lines, the taped likelihood is 23 of them, and about a third of the rest is the text of the messages it owes a user when it refuses something.

The point is not the model. The point is the seam. A likelihood that reads a whole sequence rather than one row at a time reaches the core through frmtmb_structure(), and everything else keeps working: random effects, links, standard errors, prediction.

Read vignette("frmtmb") first for the grammar, and ?frmtmb_structure for the reference description of the protocol used here.

The model

The task is a two-armed bandit. On each trial a subject picks one of two arms and receives a payoff. The Rescorla-Wagner delta rule says the subject keeps a value estimate Q for each arm, starts both at zero, and after each choice moves the chosen arm’s estimate toward what it just paid:

Q[chosen] <- Q[chosen] + alpha * (reward - Q[chosen])

The choice itself is a softmax over the two values, which for two arms is a logistic function of their difference:

P(arm 1) <- plogis(beta * (Q1 - Q2))

alpha is the learning rate, on (0, 1). beta is the inverse temperature, on (0, Inf): a large beta means the subject almost always takes the arm it currently values more. This is the model hBayesDM calls bandit2arm_delta (Ahn, Haines and Zhang 2017, Computational Psychiatry 1, 24). None of hBayesDM’s code is used here; the equations above are the published model, and the family below is written from them.

Each of the two parameters becomes a distributional parameter with its own linear predictor, so a learning rate can carry a condition effect and a subject-level random intercept exactly the way a mean would.

Why this needs the structured protocol

The likelihood does factorize over trials: given the parameters, Q is a deterministic function of the subject’s own earlier choices and payoffs, so each trial contributes one Bernoulli factor.

What it is not is ROWWISE. A row’s factor cannot be computed from that row’s parameters and response alone, because Q carries the whole history. frmtmb_family(lpdf =) is handed one row’s worth of everything, so it cannot express this. frmtmb_structure(loglik =) is handed the whole response and returns one number, so it can.

The family

The data a trial carries

# reward() carries what each arm WOULD have paid on this trial, in arm
# order. The likelihood only ever reads the chosen arm's entry, so data
# that records the received outcome alone passes it twice. The second
# column is what makes the simulator coherent, because a simulated
# choice needs the payoff of the arm the subject did not take in the
# data. vreal(pay1, pay2) would carry the same two columns. The registry
# is what lets a family name its data in the words of its own
# literature, and `reward` is the word this literature uses.
frmtmb_register_aterm("reward", arity = 2L)

The block: from a long data frame to subject by trial

The recursion is sequential in trials and independent across subjects, so it should walk trials once and update every subject at each step. That needs the rows arranged as a subject-by-trial matrix, which is what the frame_block slot is for. It runs once, at frame assembly, and everything it returns is plain data that gets saved inside the fit.

# The frame block is DATA, resolved once, and it is what turns a long
# data frame into the subject-by-trial layout the recursion walks. Every
# subject gets a row of `idx` holding its row numbers in trial order,
# padded on the right to the longest subject with a repeat of its own
# first row. `mask` is 1 on a real trial and 0 on a pad. It multiplies
# both the log-likelihood and the learning rate, so a padded cell
# contributes nothing and updates nothing.
rw_block <- function(resp, spec, av, mf, y, n) {
  rw <- resp$family[["rw"]]
  gv <- eval(rw[["subject_expr"]], mf, resp$formula_env)
  if (anyNA(gv)) {
    frm_stop("rw_delta(): every row needs a subject, and ",
             deparse1(rw[["subject_expr"]]), " has ", sum(is.na(gv)),
             " missing value(s)", call. = FALSE)
  }
  gv <- factor(gv)
  tv <- if (is.null(rw[["trial_expr"]])) {
    # no trial column: the data frame's own row order within a subject
    out <- integer(n)
    for (g in split(seq_len(n), gv)) out[g] <- seq_along(g)
    out
  } else {
    v <- eval(rw[["trial_expr"]], mf, resp$formula_env)
    if (anyNA(v)) {
      frm_stop("rw_delta(): the trial variable '", deparse1(rw[["trial_expr"]]),
               "' has missing values, so the order of the recursion is ",
               "undefined at those rows", call. = FALSE)
    }
    as.numeric(v)
  }
  rows <- lapply(split(seq_len(n), gv), function(r) r[order(tv[r])])
  key <- paste(as.integer(gv), tv, sep = "|")
  if (anyDuplicated(key)) {
    dup <- key[duplicated(key)][1L]
    frm_stop("rw_delta(): trial numbers must be unique within a subject; ",
             sum(key == dup), " rows share one. A delta rule updates once ",
             "per trial, so two rows at one trial have no order to learn in",
             call. = FALSE)
  }
  len <- lengths(rows)
  nt <- max(len)
  # The matrix() before each transpose is load-bearing: vapply() drops to
  # a plain vector when nt == 1, and t() of a vector is 1-by-n_subj, so
  # the block would come back with every subject read as one more trial
  # of a single learner and Q carried across subject boundaries.
  idx <- t(matrix(vapply(rows, function(r) c(r, rep(r[1L], nt - length(r))),
                         integer(nt)), nrow = nt))
  mask <- t(matrix(vapply(len, function(l) as.numeric(seq_len(nt) <= l),
                          numeric(nt)), nrow = nt))
  # `group` is a reserved name: it is what the per-subject
  # log-likelihood returns one value per, in level order, and what the
  # core checks a model's own grouping factor against before it lets
  # frm(importance =) sum this family's pieces into its groups.
  list(idx = idx, mask = mask, len = len, n_subj = length(rows),
       n_trial = nt, group = gv, subject = gv, trial = tv,
       levels = levels(gv), n = n)
}

Two things to notice. idx holds row NUMBERS, not values, so the recursion can index any per-row quantity with it: the response, the payoffs, and the linear predictors of alpha and beta. And subjects with different trial counts are padded to the longest one, with mask marking the padding. The mask is what makes the padding free of branches on the tape: it multiplies the log-likelihood, so a padded cell contributes nothing, and it multiplies the learning rate, so a padded cell learns nothing.

Padding is not the only way. frmtmb.latent::hmm() sorts its sequences by decreasing length instead, which makes the sequences still running at step s a prefix of the order and needs no padded cells at all. That is faster and harder to read. Measured on 40 subjects whose trial counts run from 30 to 100, the padded design carries 41 percent empty cells and builds its tape in 0.09 seconds against the balanced design’s 0.08, so for a tutorial the mask wins.

The likelihood

# Broadcast a dpar the core handed back as a single value. rep(v, each =)
# strips the advector class (see the RTMB notes in dev/); multiplying by
# a vector of ones does not.
rw_bcast <- function(v, n) if (length(v) == 1L) v * rep(1, n) else v

# The subject-by-trial layout, repeated once per replicate of the
# design. The importance correction evaluates the whole design once per
# draw, stacked, so the recursion has to run over subject-crossed-with-
# draw rather than over subject. Because the loop below is already
# vectorized across subjects, that is the same loop over a longer
# vector: `nrep` blocks of subjects, each block's row numbers shifted by
# a whole design.
rw_stack <- function(block, nrep) {
  idx <- block[["idx"]]
  msk <- block[["mask"]]
  if (nrep == 1L) return(list(idx = idx, mask = msk))
  take <- rep(seq_len(nrow(idx)), nrep)
  list(idx = idx[take, , drop = FALSE] +
         rep((seq_len(nrep) - 1L) * block[["n"]], each = nrow(idx)),
       mask = msk[take, , drop = FALSE])
}

# The trial-by-trial factors of the likelihood, as one vector per trial
# over all (subject, replicate) pairs. Every quantity the family reports
# is built from these, so there is one recursion and no second
# definition of the same arithmetic to drift from it.
#
# One iteration per TRIAL, each one a handful of vector operations over
# all subjects at once. Writing it the other way round, one iteration
# per row with Q[s] <- ... inside, is the shape RTMB punishes: an
# element-wise assignment into a taped vector costs about three orders
# of magnitude more than the vector operation it replaces.
rw_terms <- function(y, dpars, aterms, block) {
  n <- block[["n"]]
  nrep <- NROW(y) %/% n
  sk <- rw_stack(block, nrep)
  idx <- sk[["idx"]]
  msk <- sk[["mask"]]
  nn <- n * nrep
  alpha <- rw_bcast(dpars[["alpha"]], nn)
  beta <- rw_bcast(dpars[["beta"]], nn)
  pay1 <- aterms[["reward1"]]
  pay2 <- aterms[["reward2"]]
  q1 <- rep(0, nrow(idx))
  q2 <- q1
  out <- vector("list", ncol(idx))
  for (t in seq_len(ncol(idx))) {
    i <- idx[, t]
    m <- msk[, t]
    c1 <- y[i]
    eta <- beta[i] * (q1 - q2)
    # log P(choice) = c1 * eta - log(1 + exp(eta)), folded with
    # logspace_add so a decisive subject does not overflow
    out[[t]] <- m * (c1 * eta - RTMB::logspace_add(0 * eta, eta))
    # the mask enters through the learning rate, so a padded cell learns
    # nothing and no branch reaches the tape
    a <- alpha[i] * m
    q1 <- q1 + (a * c1) * (pay1[i] - q1)
    q2 <- q2 + (a * (1 - c1)) * (pay2[i] - q2)
  }
  list(terms = out, idx = idx, mask = msk)
}

# One value per subject, which is this family's coarsest honest piece
# and the one frm(importance =) resamples. A subject's trials are a
# product, so its log-likelihood is the elementwise sum of the per-trial
# vectors: the piece is free, given the recursion.
#
# weights() is refused in check_spec, so `weights` is 1 in all three
# slots. A trial's factor cannot be reweighted on its own: its value
# depends on every earlier trial of the same subject.
rw_loglik_group <- function(y, dpars, aterms, weights, block, extra) {
  Reduce(`+`, rw_terms(y, dpars, aterms, block)[["terms"]])
}

# One value per trial: the choice probability CONDITIONAL on everything
# the subject saw before it. These are the factors themselves, scattered
# back to the rows they came from, and a Bernoulli factor's saturated
# value is zero, which is what makes a deviance residual definable here.
rw_loglik_row <- function(y, dpars, aterms, weights, block, extra) {
  "c" <- RTMB::ADoverload("c")
  "[<-" <- RTMB::ADoverload("[<-")
  tm <- rw_terms(y, dpars, aterms, block)
  # ONE sub-assignment, not one per trial: `[<-` copies the whole vector
  # each time it is called on a taped one. The padded cells are dropped
  # rather than written, because a pad repeats its subject's FIRST row
  # number and would otherwise overwrite that trial's own factor with a
  # masked zero.
  keep <- as.vector(tm[["mask"]]) == 1
  out <- rep(0, length(y))
  out[as.vector(tm[["idx"]])[keep]] <-
    do.call(c, tm[["terms"]])[keep]
  # off the tape only: the numeric path is the one residuals() reads,
  # and an advector carries no attributes
  if (!inherits(out, "advector")) {
    attr(out, "saturated") <- rep(0, length(out))
  }
  out
}

# The whole response's log-likelihood, defined as the sum of the pieces
# rather than accumulated separately, so the three slots cannot disagree.
rw_loglik <- function(y, dpars, aterms, weights, block, extra) {
  sum(rw_loglik_group(y, dpars, aterms, weights, block, extra))
}

The loop body is the delta rule, written for every subject at once. q1 and q2 are vectors with one entry per subject, and one iteration of the loop advances all of them. There is no conditional anywhere: the chosen arm is selected by multiplying by the response, which is data, and the padding is selected by multiplying by the mask, which is also data.

rw_terms() stops one step short of a log-likelihood: it returns the per-trial factors themselves, and the three slots below are three ways of adding them up. A subject’s log-likelihood is the elementwise sum across trials, a trial’s is the factor scattered back to its own row, and the total is the sum of either. Writing them as one recursion with three tails is not tidiness, it is the only way to be sure the three numbers agree; the suite asserts that they do.

Why three at all. loglik is what the objective needs. loglik_group is what frm(importance = ) resamples, one value per subject. loglik_row is what a deviance residual compares against its saturated fit, which for a Bernoulli factor is zero, and that zero is attached to the row values so the core does not have to guess it.

The nrep line is the stacking contract. The importance correction evaluates the whole design once per draw, stacked end to end, so the recursion has to run over subject-crossed-with-draw. Because the loop is already vectorized across subjects, that is the same loop over a longer vector rather than a new algorithm.

weights arrives and is ignored, which the family is allowed to do only because it refuses weights() outright further down. A row weight has no meaning here. A trial’s factor depends on every earlier trial of the same subject, so there is nothing separable to reweight.

Fitted values and draws

The same recursion in plain numbers answers two more questions.

# The same recursion off the tape, in plain numbers, at the estimates.
# It answers two questions: the per-trial choice probability, which is
# the fitted value, and a forward draw, which is the simulator. They
# share one function because a draw is a replay whose choices come from
# rbinom() instead of from the data.
rw_replay <- function(dpars, aterms, block, y, draw = FALSE) {
  idx <- block[["idx"]]
  msk <- block[["mask"]]
  n <- block[["n"]]
  alpha <- rep(as.numeric(dpars[["alpha"]]), length.out = n)
  beta <- rep(as.numeric(dpars[["beta"]]), length.out = n)
  pay1 <- aterms[["reward1"]]
  pay2 <- aterms[["reward2"]]
  p <- rep(NA_real_, n)
  q1 <- rep(0, nrow(idx))
  q2 <- q1
  for (t in seq_len(ncol(idx))) {
    i <- idx[, t]
    m <- msk[, t]
    keep <- m == 1
    pt <- stats::plogis(beta[i] * (q1 - q2))
    p[i[keep]] <- pt[keep]
    c1 <- if (draw) stats::rbinom(length(i), 1L, pt) else y[i]
    if (draw) y[i[keep]] <- c1[keep]
    a <- alpha[i] * m
    q1 <- q1 + (a * c1) * (pay1[i] - q1)
    q2 <- q2 + (a * (1 - c1)) * (pay2[i] - q2)
  }
  list(p = p, y = y)
}

# The three post-fit slots, each one line of work on top of the replay.
rw_parts <- function(fit) {
  rspec <- single_response(fit, "a rw_delta() fit")
  list(dpars = eval_dpars(fit)[[rspec$resp_name]],
       aterms = fit$frame[["aterm_values"]][[rspec$resp_name]],
       y = fit$frame[["y"]][[rspec$resp_name]])
}

rw_fitted <- function(fit, block) {
  pp <- rw_parts(fit)
  rw_replay(pp$dpars, pp$aterms, block, pp$y)$p
}

The per-trial choice probability is the natural fitted value for this family, and its binomial variance is what pearson residuals divide by. A forward draw is the same walk with rbinom() in place of the observed choice, which is why one function serves both.

# The structured simulator. simulate(), posterior_predict() and
# frm_simulate() all reach it through the same context, so the draw is
# written once. It walks each subject forward, drawing a choice from the
# current Q pair and learning from the payoff that choice earns.
rw_sim <- function(ctx) {
  blk <- ctx[["block"]]
  rw_replay(ctx[["dpars"]], ctx[["aterms"]], blk,
            y = rep(0, blk[["n"]]), draw = TRUE)$y
}

The family object

# The family. Two distributional parameters, each with its own linear
# predictor: alpha on (0, 1) through a logit link, beta on (0, Inf)
# through a log link. `primary_dpars = "alpha"` sends the main
# right-hand side of the formula to the learning rate.
rw_delta <- function(subject, trial = NULL) {
  subject_expr <- substitute(subject)
  trial_expr <- substitute(trial)
  if (is.null(subject_expr)) {
    frm_stop("rw_delta(subject =) names the column that separates one ",
             "learner's trial sequence from the next", call. = FALSE)
  }
  fam <- frmtmb_family(
    "rw_delta",
    dpars = c("alpha", "beta"),
    links = list(alpha = "logit", beta = "log"),
    primary_dpars = "alpha",
    type = "discrete",
    required_aterms = c("reward1", "reward2"),
    # The likelihood factorizes over trials but is not ROWWISE: a
    # trial's factor reads the whole history of its subject, which this
    # signature cannot see. Returning something anyway is the silent lie
    # the structured protocol exists to remove.
    lpdf = function(y, dpars, aterms, extra = NULL) {
      frm_stop("A rw_delta() trial's probability depends on every earlier ",
               "trial of the same subject, so the family has no row-wise ",
               "log-density. Use logLik() for the total, or fitted() for ",
               "the per-trial choice probabilities", call. = FALSE)
    },
    valid_y = function(y, aterms) {
      if (!all(y %in% c(0, 1))) {
        frm_stop("rw_delta(): the response is the arm chosen on each trial, ",
                 "coded 1 for the first arm and 0 for the second",
                 call. = FALSE)
      }
    },
    init_dpars = list(alpha = function(y, aterms) 0.3,
                      beta = function(y, aterms) 1),
    structure = rw_structure()
  )
  fam[["rw"]] <- list(subject_expr = subject_expr, trial_expr = trial_expr)
  fam
}

# The structure is the family's half of the protocol: where its data
# block comes from, what its likelihood is, what it can answer
# afterwards, and what it refuses. Every capability flag starts FALSE
# for a likelihood that does not factorize over rows, so the refusals
# below are explanations, not switches.
rw_structure <- function() {
  frmtmb_structure(
    frame_vars = function(fam) {
      list(fam[["rw"]][["subject_expr"]], fam[["rw"]][["trial_expr"]])
    },
    # An unanswered trial produces no prediction error, so dropping its
    # row IS the right recursion. This is where a delta rule differs
    # from a hidden Markov chain, which keeps the row because the state
    # still moves; see frmtmb.latent::hmm().
    keep_na = FALSE,
    check_spec = function(resp, spec, av) {
      if (length(spec$responses) > 1L || isTRUE(spec$rescor)) {
        frm_stop("rw_delta() supports univariate models only: the recursion ",
                 "is a likelihood over one response's trial sequences",
                 call. = FALSE)
      }
      bad <- intersect(c("weights", "cens", "trunc_lb", "trunc_ub", "se"),
                       names(av))
      if (length(bad)) {
        frm_stop("rw_delta() cannot be combined with ", bad[1L], "(): that ",
                 "term reshapes a per-row likelihood contribution, and a ",
                 "trial's contribution here is conditional on every earlier ",
                 "trial of the same subject", call. = FALSE)
      }
    },
    frame_block = rw_block,
    loglik = rw_loglik,
    # The two factorizations. `unit` is a different question and keeps
    # its own answer: the finest factorization here is the TRIAL, while
    # the smallest thing that can honestly be left out is a whole
    # subject, because dropping one trial changes every later trial's Q.
    loglik_row = rw_loglik_row,
    loglik_group = rw_loglik_group,
    unit = "one subject's trial sequence",
    fitted_mean = rw_fitted,
    fitted_var = function(fit, block) {
      p <- rw_fitted(fit, block)
      p * (1 - p)
    },
    sim_ctx = rw_sim,
    # The one capability a factorized likelihood buys back here. Every
    # other flag stays FALSE: a per-trial factor is not a per-trial
    # MODEL, so newdata, conditional effects and one-step-ahead
    # residuals are refused for the reasons below and not for want of
    # the factors.
    supports = list(deviance = TRUE),
    refusals = list(
      newdata_response = paste0(
        "predict(type = 'response') on a rw_delta() fit is not ",
        "available for newdata: a trial's choice probability is ",
        "conditional on the payoffs and choices of every earlier trial ",
        "of the same subject, and newdata carries no block to replay. ",
        "Use fitted() on the training data, or predict(dpar = 'alpha') ",
        "for a learning rate"),
      conditional_effects = paste0(
        "conditional_effects() is not available for a rw_delta() fit: ",
        "the expected response is a choice probability that depends on ",
        "a whole trial history, which the synthetic grid this function ",
        "builds does not have. Plot the learning rate itself with ",
        "predict(dpar = 'alpha')"),
      osa = paste0(
        "residuals(type = 'osa') is not available for a rw_delta() fit: ",
        "the tape holds the recursion over each whole sequence with no ",
        "registered observation vector. Use type = 'pearson', which ",
        "divides by the binomial variance of the choice probability")
    )
  )
}

The refusals are the interesting part. A structured family starts with every capability turned off, and each one it leaves off should say why in its own words. conditional_effects() is refused not because it is hard but because it is undefined: the expected response here is a choice probability that depends on a whole trial history, and the synthetic grid that function builds has no history.

deviance is the one flag turned back on, and it is on because of a slot rather than in spite of one: with loglik_row declared the core has each trial’s log-density and its saturated value, which is what a unit deviance is made of. unit is a separate declaration and keeps its own answer. The finest factorization here is the TRIAL; the smallest thing that can honestly be left out is a whole SUBJECT, because dropping one trial changes every later trial’s value store.

Fitting it

# Scaffolding for the worked example, not part of the family. It builds
# a two-armed bandit with the payoff schedule of BOTH arms fixed in
# advance, one row per subject and trial, and a placeholder response the
# simulator overwrites. Arm 1 pays with probability p1 and arm 2 with
# probability p2, so the learnable contrast is p1 - p2.
rl_bandit_design <- function(n_subj, n_trial, p1 = 0.7, p2 = 0.3) {
  dd <- expand.grid(trial = seq_len(n_trial), id = seq_len(n_subj))
  dd$condition <- factor(rep(0:1, length.out = n_subj)[dd$id], 0:1,
                         c("ctl", "trt"))
  dd$id <- factor(dd$id)
  dd$pay1 <- stats::rbinom(nrow(dd), 1L, p1)
  dd$pay2 <- stats::rbinom(nrow(dd), 1L, p2)
  dd$choice <- 0
  # expand.grid() leaves an out.attrs attribute that str() prints at
  # length and nothing reads
  attr(dd, "out.attrs") <- NULL
  dd
}

# The learning rate carries the condition effect; the softmax
# temperature is a per-subject intercept correlated with it.
rl_bform <- bf(choice | reward(pay1, pay2) ~ condition + (1 | p | id),
               beta ~ 1 + (1 | p | id))

rl_family <- function() rw_delta(subject = id, trial = trial)

# Parameter values in frm_simulate()'s spelling, brms's names: the same
# strings variables() returns for the fit.
rl_truth <- list(
  b_alpha_Intercept = stats::qlogis(0.35),
  b_alpha_conditiontrt = 0.8,
  b_beta_Intercept = log(3),
  sd_id__alpha_Intercept = 0.5,
  sd_id__beta_Intercept = 0.4,
  cor_id__alpha_Intercept__beta_Intercept = 0.3
)

# Draws come from the family's own simulator: frm_simulate() draws a new
# set of subject effects from `truth`, then walks each subject forward.
rl_simulate <- function(design, truth = rl_truth, nsim = 1, seed = NULL) {
  sims <- frm_simulate(rl_bform, data = design, family = rl_family(),
                       newparams = truth, nsim = nsim, seed = seed)
  lapply(seq_len(nsim), function(s) {
    d <- design
    d$choice <- sims[[s]]
    d
  })
}

set.seed(2024)
design <- rl_bandit_design(n_subj = 30, n_trial = 80)
dd <- rl_simulate(design, nsim = 1, seed = 2024)[[1]]
str(dd[1:3, ])
#> 'data.frame':    3 obs. of  6 variables:
#>  $ trial    : int  1 2 3
#>  $ id       : Factor w/ 30 levels "1","2","3","4",..: 1 1 1
#>  $ condition: Factor w/ 2 levels "ctl","trt": 1 1 1
#>  $ pay1     : int  0 1 1
#>  $ pay2     : int  0 0 1
#>  $ choice   : num  1 0 0

rl_simulate() goes through frm_simulate(), which draws a fresh set of subject effects from rl_truth and then calls the family’s own simulator. The choices in dd therefore come from the same code path the fit is about to score.

Now the fit. This is an ordinary frm() call:

fit <- frm(bf(choice | reward(pay1, pay2) ~ condition + (1 | p | id),
              beta ~ 1 + (1 | p | id)),
           family = rw_delta(subject = id, trial = trial),
           data = dd)
fixef(fit)
#>                      Estimate  Est.Error       Q2.5     Q97.5
#> alpha_Intercept    -0.7919201 0.21949895 -1.2221301 -0.361710
#> beta_Intercept      1.0586371 0.08594413  0.8901897  1.227084
#> alpha_conditiontrt  1.1220594 0.30016589  0.5337451  1.710374

The main right-hand side is the LEARNING RATE’s predictor, because the family declared primary_dpars = "alpha". The | p | in both terms correlates each subject’s learning rate with its own temperature, which is the usual hierarchical parameterization for these models.

VarCorr(fit)
#> $id
#> $id$sd
#>                  Estimate  Est.Error      Q2.5     Q97.5
#> alpha_Intercept 0.4628031 0.16596184 0.1375239 0.7880824
#> beta_Intercept  0.3793421 0.08175405 0.2191071 0.5395771
#> 
#> $id$cor
#> , , alpha_Intercept
#> 
#>                    Estimate Est.Error       Q2.5     Q97.5
#> alpha_Intercept  1.00000000 0.0000000  1.0000000 1.0000000
#> beta_Intercept  -0.04221552 0.4113754 -0.8484964 0.7640654
#> 
#> , , beta_Intercept
#> 
#>                    Estimate Est.Error       Q2.5     Q97.5
#> alpha_Intercept -0.04221552 0.4113754 -0.8484964 0.7640654
#> beta_Intercept   1.00000000 0.0000000  1.0000000 1.0000000
#> 
#> 
#> $id$cov
#> , , alpha_Intercept
#> 
#>                     Estimate  Est.Error        Q2.5     Q97.5
#> alpha_Intercept  0.214186726 0.15361532 -0.08689376 0.5152672
#> beta_Intercept  -0.007411386 0.07253944 -0.14958608 0.1347633
#> 
#> , , beta_Intercept
#> 
#>                     Estimate  Est.Error        Q2.5     Q97.5
#> alpha_Intercept -0.007411386 0.07253944 -0.14958608 0.1347633
#> beta_Intercept   0.143900424 0.06202550  0.02233267 0.2654682

Compare with the simulated values:

rbind(estimate = c(unlist(fixef_by_dpar(fit)),
                   sd_alpha = VarCorr(fit)$id$sd[1, "Estimate"],
                   sd_beta = VarCorr(fit)$id$sd[2, "Estimate"]),
      truth = c(rl_truth$b_alpha_Intercept, rl_truth$b_alpha_conditiontrt,
                rl_truth$b_beta_Intercept,
                rl_truth$sd_id__alpha_Intercept,
                rl_truth$sd_id__beta_Intercept))
#>          alpha.(Intercept) alpha.conditiontrt beta.(Intercept)  sd_alpha
#> estimate        -0.7919201           1.122059         1.058637 0.4628031
#> truth           -0.6190392           0.800000         1.098612 0.5000000
#>            sd_beta
#> estimate 0.3793421
#> truth    0.4000000

This is one dataset, so read the differences as sampling error and not as bias. The recovery study below is the systematic version.

Everything downstream of the fit works because the structure said it does. The fitted values are per-trial choice probabilities:

p <- fitted(fit)[, "Estimate"]
summary(p)
#>    Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
#> 0.02198 0.59464 0.76165 0.72271 0.89176 0.99444
# the model's mean predicted choice rate against the observed one
c(predicted = mean(p), observed = mean(dd$choice))
#> predicted  observed 
#> 0.7227135 0.7104167

Check 1: an independent reference for the recursion

The taped recursion is vectorized and masked, which is exactly the kind of code that can be wrong in a way that still converges. So write the model out again, the way it is stated, one subject and one trial at a time, over quantities pulled back through the public accessors, and compare.

u <- ranef(fit)[[1]]
sid <- match(as.character(dd$id), rownames(u))
xa <- model.matrix(~ condition, dd)
alpha <- plogis(as.vector(xa %*% unname(fixef_by_dpar(fit)$alpha)) + u[sid, 1])
beta <- exp(unname(fixef_by_dpar(fit)$beta)[[1]] + u[sid, 2])

ll <- 0
for (s in unique(sid)) {
  rows <- which(sid == s)
  rows <- rows[order(dd$trial[rows])]
  q <- c(0, 0)
  for (i in rows) {
    ll <- ll + dbinom(dd$choice[i], 1, plogis(beta[i] * (q[1] - q[2])),
                      log = TRUE)
    k <- if (dd$choice[i] == 1) 1 else 2
    pay <- if (k == 1) dd$pay1[i] else dd$pay2[i]
    q[k] <- q[k] + alpha[i] * (pay - q[k])
  }
}
c(reference = ll,
  from_fitted = sum(dbinom(dd$choice, 1, p, log = TRUE)),
  difference = ll - sum(dbinom(dd$choice, 1, p, log = TRUE)))
#>     reference   from_fitted    difference 
#> -1.149687e+03 -1.149687e+03  2.728484e-12

tests/testthat/test-rl-example.R runs this check on an unbalanced design as well, where the padding is live, and against the full joint density rather than the data part alone.

Check 2: an identity against Stan

The reference above shares R with the thing it checks. A second implementation in another language and another autodiff system does not. tests/testthat/helper-rl.R holds a thirty-line Stan program for the same model, written from the equations at the top of this page, and compares Stan’s log_prob with frmtmb’s joint log density at frmtmb’s own estimates. The Stan program declares exactly the parameters frmtmb estimates, on the scales frmtmb estimates them on, so the map between them is the identity and the comparison carries no Jacobian.

Measured on 30 subjects and 100 trials:

quantity	at the estimates	at a displaced point
Stan log_prob	-1427.8453940	-1460.3042080
frmtmb joint log density	-1427.8453940	-1460.3042080
difference	-4.5e-13	-2.5e-12
largest gradient on the subject effects	1.2e-14	not asserted
largest gradient overall	5.2	not asserted

The difference is zero and is expected to be exactly zero, not merely small: core frm() carries no default prior, the Stan program uses target += with _lpmf and _lpdf so it drops no normalizing constant, and every declared parameter is unbounded so adjust_transform = FALSE has no Jacobian to discard. The second column is a deliberately non-stationary point, so the agreement cannot be an artifact of both sides sitting at an optimum.

The gradient row is the check that catches a wrong map. frmtmb puts the subject effects at their conditional modes, so Stan’s gradient with respect to them must vanish there. The gradient with respect to the fixed effects is NOT zero at that point and is not asserted on: the maximum-likelihood estimate makes the MARGINAL likelihood stationary, not the joint one.

The test is gated the same way the brms comparison tier is, on FRMTMB_BRMS_FIT_TESTS and NOT_CRAN, because it compiles a Stan program.

Check 3: parameter recovery

Simulate 100 datasets at known parameter values through the family’s own simulator, fit each one, and ask whether the estimates land on the truth and whether their intervals cover it. The design is 40 subjects with 100 trials each and a condition effect on the learning rate. The script is described in dev/rl-findings.md.

parameter	truth	bias	Monte Carlo se	sd of estimates	coverage
alpha_(Intercept)	-0.619	+0.009	0.018	0.178	0.94
alpha_conditiontrt	0.800	-0.018	0.025	0.248	0.94
beta_(Intercept)	1.099	+0.002	0.007	0.070	0.96
log sd(alpha)	-0.693	-0.211	0.057	0.573	0.93
log sd(beta)	-0.916	-0.062	0.017	0.167	0.95

The three fixed effects are unbiased to within their Monte Carlo error and their 95 percent Wald intervals cover at the nominal rate. The variance components are the weak point, and the learning-rate one especially: its estimate is biased downward and scattered widely, which is the usual behavior of a variance component estimated from binary data with a small number of levels. Coverage is over the replicates that produced a finite interval, which for log sd(alpha) was 98 of the 100.

The Laplace caveat, measured

frmtmb integrates the subject effects out with a Laplace approximation. The approximation is a Gaussian fitted at the conditional mode, and it is exact only when the conditional log-density really is quadratic. For a binary response it is not, and the fewer trials a subject has, the less quadratic it is.

The way to price that error in frmtmb is frm(importance = ), which reweights draws from the Laplace Gaussian against the exact integrand. It works here, and it works because the family declares how its likelihood factorizes: loglik_group hands the correction one value per subject, which is exactly the piece it resamples. Until that slot existed the correction refused every family with a loglik, whether or not the likelihood factorized.

fit_i <- frm(bf(choice | reward(pay1, pay2) ~ condition + (1 | p | id),
                beta ~ 1 + (1 | p | id)),
             family = rw_delta(subject = id, trial = trial),
             data = dd, importance = 200)
c(laplace = VarCorr(fit)$id$sd[1, "Estimate"],
  corrected = VarCorr(fit_i)$id$sd[1, "Estimate"])
#>   laplace corrected 
#> 0.4628031 0.5406222
fit_i$importance[c("draws", "mcse")]
#> $draws
#> [1] 200
#> 
#> $mcse
#> [1] 0.104886

The chunk is left with error = TRUE on purpose. The correction stops when a draw count is too small for a model, and its message names a count that works; the next section is about what the count buys.

Short sessions, and what the correction says about them

Run the same recovery study at 20 trials per subject instead of 100, and read the bias and the coverage:

parameter	100 trials, bias	coverage	20 trials, bias	coverage
alpha_(Intercept)	+0.009	0.94	+0.064	0.95
alpha_conditiontrt	-0.018	0.94	-0.005	0.95
beta_(Intercept)	+0.002	0.96	-0.000	0.92
log sd(alpha)	-0.211	0.93	-0.628	0.73
log sd(beta)	-0.062	0.95	-0.247	0.89

100 replicates each, 40 subjects, the same true parameter values. Coverage is over the replicates that produced a finite interval: at 20 trials that was 84 of 100 for log sd(alpha) and 96 or more for everything else.

Two readings. The fixed effects survive. At 20 trials per subject their bias is still inside Monte Carlo error and their intervals still cover close to the nominal rate, so a study that wants the condition effect on the learning rate can work with short sessions.

The variance components do not survive. The learning rate’s log standard deviation comes out 0.63 too low, which is a standard deviation about half the true one, and its interval covers 73 times in 100. Sixteen of the hundred short-session fits gave no usable interval for it at all.

Two causes could produce that, and they are not the same thing. A variance component estimated by maximum likelihood from binary data with 40 levels is biased downward whether or not the integral is approximated; and the Laplace approximation is itself worse the fewer trials a subject has. The recovery table cannot separate them. frm(importance = ) can, and now that this family declares its pieces, it can be asked. Six seeds per design (4201 to 4206), the same 40 subjects and the same truth, 100 draws:

	40 x 100 trials	40 x 20 trials
replicates the correction completed	6 of 6	6 of 6
mean shift, alpha_(Intercept)	-0.002	-0.044
mean shift, alpha_conditiontrt	+0.007	+0.027
mean shift, beta_(Intercept)	-0.001	+0.049
mean shift, log sd(alpha)	+0.121	+0.451
largest single shift, log sd(alpha)	0.207	0.914
Monte Carlo standard error	0.11 to 0.15	0.02 to 0.83
smallest effective sample size per draw	0.69	0.03

The fixed effects do not move. That is the recovery table’s first reading confirmed by an independent route: their bias is not Laplace error, because removing the Laplace error leaves them where they were.

The variance component moves up, and the diagnostics say how far to trust it. At 100 trials the shift is about an eighth of a log unit, and it holds as the draws grow: seed 4201 reads -0.8703 under Laplace, and -0.7068, -0.7509 and -0.7464 at 100, 400 and 1600 draws. At 20 trials the shift is larger, and 100 draws are few: on two of the six datasets the smallest effective sample size is below the warning’s threshold, at 0.11 and 0.03, and those numbers need more draws before anyone reads them.

One dataset, named so it can be checked, shows what the draws buy. Take 40 subjects at 20 trials from rl_simulate(rl_bandit_design(40, 20), seed = 4206). Its Laplace log sd(alpha) is -1.2024. At 100 draws the correction returns -0.2879, a standard deviation of 0.75 against a truth of 0.5, with a Monte Carlo standard error of 0.162 and a smallest effective sample size of 0.63 per draw. At 400 draws it returns -0.6501, a standard deviation of 0.52, with a Monte Carlo standard error of 0.036 and an effective sample size of 0.96.

So the honest reading is not that the variance-component bias is Laplace error. At 100 trials some small part of it is. At 20 trials the approximation’s part is large and the correction can measure it if it gets enough draws, and that is a second reason to treat a subject-level standard deviation from 20 binary trials as a lower bound rather than as an estimate. The operational advice is unchanged: report fixed effects from short sessions, and do not report a variance component from them.

Shape and cost

The family loops over TRIALS and vectorizes over subjects. It is worth seeing what the other spelling costs. Here is the same likelihood written the way the model is usually stated, one row at a time, with the value store updated by sub-assignment:

for (t in seq_len(ncol(idx))) {
  for (s in seq_len(nrow(idx))) {          # the extra loop
    i <- idx[s, t]
    eta <- beta[i] * (q1[s] - q2[s])
    ll <- ll + y[i] * eta - RTMB::logspace_add(0 * eta, eta)
    if (y[i] == 1) q1[s] <- q1[s] + alpha[i] * (pay1[i] - q1[s])
    else q2[s] <- q2[s] + alpha[i] * (pay2[i] - q2[s])
  }
}

All three spellings give the same objective value to fourteen significant digits. On 40 subjects and 100 trials, 4000 rows, with the three builds interleaved and each gradient timed over a batch of 100 calls:

loglik shape	tape build	one gradient
vectorized over subjects	0.08 s	7.3 ms
row loop, scalar value store	0.57 s	7.2 ms
row loop, length-n log-density vector	0.65 s	7.2 ms

The whole penalty is paid at TAPE CONSTRUCTION, and none of it per evaluation. That is what it has to be: once the tape exists it is a node list, and the R code that built it is gone. The gradient column is flat to within measurement noise, and an optimization runs hundreds of gradients against one build.

Do not read the build column as a fixed multiplier either. RTMB’s [<- copies the vector it writes into, so the sub-assigning form gets relatively worse as the response grows:

rows	vectorized	length-n sub-assignment	factor
1000	0.03 s	0.16 s	5.3
4000	0.08 s	0.60 s	7.5
10000	0.14 s	1.67 s	11.9
20000	0.34 s	4.33 s	12.7

Values agree to 1e-11 across the whole sweep. The rule to take away is the shape, not a number: vectorize the tape-building loop, and expect the cost of not doing so to grow with the data.

Scaling of the vectorized family, interleaved, frm(dry_run = "objective") with the frame-assembly time subtracted, median of 7 builds and the best of 3 gradient batches:

subjects	trials	rows	tape build	one gradient
10	100	1000	0.04 s	10.0 ms
40	25	1000	0.03 s	8.8 ms
40	100	4000	0.08 s	39.6 ms
40	400	16000	0.28 s	153.4 ms
160	100	16000	0.22 s	160.6 ms
40	800	32000	0.61 s	296.2 ms

Both costs track the ROW count, which is the node count. Read the two 16000-row rows against each other for the part that does not: more trials at a fixed row count costs more to tape, because the R loop runs more iterations, and costs the same per gradient, because the tape is the same size either way. The size of that build gap is not stable enough to quote. Two interleaved runs of this table on one machine put it at 1.27x and 1.70x.

What a built-in family would add

The example above is a complete family. Three things would still change if it moved into a package.

A pointwise log-likelihood for loo(). The family declares its factorization, so frm(importance = ) and deviance residuals work above. loo() and waic() would work too, on a sampled fit: they live in frmtmb.sample and read the same slots, and their columns would be SUBJECTS rather than trials, because unit says a trial cannot honestly be left out. Nothing here exercises that path, because these are maximum-likelihood fits and an elpd is a posterior quantity.

Newdata. frm_linpred(type = "response") on new data would need a block built without a response, and the protocol defers that case on purpose. For this family the meaning is clear enough (replay the new subject’s payoff schedule against the fitted parameters), so a built-in version would be worth the work.

More of the literature. A real package would carry the family of models this one is the simplest member of: separate learning rates for gains and losses, a decay term, more than two arms. All of them share this block and this loop shape, and differ only in the body.

None of that is a limitation of the protocol as a place to put a sequential likelihood. The core knows nothing about reinforcement learning, and the model above fits, predicts, simulates and recovers its parameters through the ordinary grammar.