Registering custom RLSSM models in HSSM¶

This is the shortest tutorial in the RLSSM suite. It teaches the HSSM-native registry path for building a reinforcement-learning sequential-sampling model (RLSSM): you write a learning rule, hand it to register_rlssm_model(), and get a named, reusable model configuration you can pull out later with get_rlssm_model_config().

Two ways to build an RLSSM — pick based on where you want the model to live:

	Define in `ssms.rl`, bridge with `from_ssms_model`	Register HSSM-side with `register_rlssm_model` (this notebook)
Learning rule lives in	an `ssms.rl` model/learner you author	a function you write in your own code
You get	a full simulator and a fittable config	a fittable config under a named key
Best when	you also need to simulate the model (data, PPC)	you already have data and just want a named, reusable fitting config
Shown in	Advanced, Restless learner	here

The important thing to notice: either way, you bring your own learning rule. The registry does not hand you a learning function — it wraps your function together with HSSM's decision-process likelihood into one tidy RLSSMConfig.

New to all of this? Start with the Basic tutorial, which walks through the RLSSM idea, hierarchical priors, recovery, and posterior predictive checks in full. Here we assume that background and keep the demonstration light.

1. Setup¶

In [1]:

Copied!





import logging
import os
import warnings

import arviz as az
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from jax.lax import scan
from ssms import rl

import hssm
from hssm.rl import get_rlssm_model_config, list_models, register_rlssm_model
from hssm.utils import annotate_function  # attaches input/output metadata to a function

warnings.filterwarnings("ignore")
logging.getLogger("jax._src.xla_bridge").setLevel(logging.ERROR)

hssm.set_floatX("float32", update_jax=True)
RANDOM_SEED = 20260707
import logging
import os
import warnings

import arviz as az
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from jax.lax import scan
from ssms import rl

import hssm
from hssm.rl import get_rlssm_model_config, list_models, register_rlssm_model
from hssm.utils import annotate_function  # attaches input/output metadata to a function

warnings.filterwarnings("ignore")
logging.getLogger("jax._src.xla_bridge").setLevel(logging.ERROR)

hssm.set_floatX("float32", update_jax=True)
RANDOM_SEED = 20260707

Setting PyTensor floatX type to float32.

Setting "jax_enable_x64" to False. If this is not intended, please set `jax` to False.

Simulation scale¶

In [2]:

Copied!





# One notebook, two scales. CI executes at the small "doc" scale; the committed
# outputs come from a richer FULL_RUN=1 pass. This notebook is deliberately light,
# so even the full run is smaller than the Basic/Advanced tutorials.
FULL_RUN = os.environ.get("FULL_RUN", "0") == "1"

N_PARTICIPANTS = 12 if FULL_RUN else 4
N_TRIALS = 150 if FULL_RUN else 60
N_CHAINS = 2
N_TUNE = 800 if FULL_RUN else 300
N_DRAWS = 400 if FULL_RUN else 300

# The second, dual-learning-rate fit estimates only two parameters, so it uses a
# smaller dataset and a single chain. It is an integration demo, not a recovery study.
DUAL_N_PARTICIPANTS = 4 if FULL_RUN else 2
DUAL_N_TRIALS = 100 if FULL_RUN else 40
DUAL_N_TUNE = 300 if FULL_RUN else 100
DUAL_N_DRAWS = 300 if FULL_RUN else 100

print(
    f"FULL_RUN={FULL_RUN} | participants={N_PARTICIPANTS} trials={N_TRIALS} "
    f"tune={N_TUNE} draws={N_DRAWS}"
)
print(
    f"dual-alpha demo | participants={DUAL_N_PARTICIPANTS} trials={DUAL_N_TRIALS} "
    f"tune={DUAL_N_TUNE} draws={DUAL_N_DRAWS}"
)
# One notebook, two scales. CI executes at the small "doc" scale; the committed
# outputs come from a richer FULL_RUN=1 pass. This notebook is deliberately light,
# so even the full run is smaller than the Basic/Advanced tutorials.
FULL_RUN = os.environ.get("FULL_RUN", "0") == "1"

N_PARTICIPANTS = 12 if FULL_RUN else 4
N_TRIALS = 150 if FULL_RUN else 60
N_CHAINS = 2
N_TUNE = 800 if FULL_RUN else 300
N_DRAWS = 400 if FULL_RUN else 300

# The second, dual-learning-rate fit estimates only two parameters, so it uses a
# smaller dataset and a single chain. It is an integration demo, not a recovery study.
DUAL_N_PARTICIPANTS = 4 if FULL_RUN else 2
DUAL_N_TRIALS = 100 if FULL_RUN else 40
DUAL_N_TUNE = 300 if FULL_RUN else 100
DUAL_N_DRAWS = 300 if FULL_RUN else 100

print(
    f"FULL_RUN={FULL_RUN} | participants={N_PARTICIPANTS} trials={N_TRIALS} "
    f"tune={N_TUNE} draws={N_DRAWS}"
)
print(
    f"dual-alpha demo | participants={DUAL_N_PARTICIPANTS} trials={DUAL_N_TRIALS} "
    f"tune={DUAL_N_TUNE} draws={DUAL_N_DRAWS}"
)

FULL_RUN=True | participants=12 trials=150 tune=800 draws=400
dual-alpha demo | participants=4 trials=100 tune=300 draws=300

2. What is already registered?¶

HSSM discovers public RLSSM presets from ssms.rl.preset at runtime and merges any HSSM-side custom registrations into the same list_models() discovery surface. get_rlssm_model_config(name) materializes any listed model into an RLSSMConfig.

In [3]:

Copied!





for name, description in list_models().items():
    print(f"{name:28s} {description}")

# Peek at one built-in template.
builtin = get_rlssm_model_config("2AB_RW_Angle")
print("\ndecision_process :", builtin.decision_process)
print("computed by learner:", set(builtin.ssm_logp_func.computed))  # {'v'}
print("sampled parameters :", builtin.list_params)
for name, description in list_models().items():
    print(f"{name:28s} {description}")

# Peek at one built-in template.
builtin = get_rlssm_model_config("2AB_RW_Angle")
print("\ndecision_process :", builtin.decision_process)
print("computed by learner:", set(builtin.ssm_logp_func.computed))  # {'v'}
print("sampled parameters :", builtin.list_params)

2AB_RW_DDM                         Two-armed bandit with a Rescorla-Wagner delta-rule learner and a DDM decision process.
2AB_RW_Angle                       Two-armed bandit with a Rescorla-Wagner delta-rule learner and an angle decision process.
2AB_RW_Weibull                     Two-armed bandit with a Rescorla-Wagner delta-rule learner and a Weibull-bound DDM decision process.
2AB_RW_DualAlpha_Angle             Two-armed bandit with a dual-alpha Rescorla-Wagner learner and an angle decision process.
2AB_RW_InvTempSoftmax              Two-armed bandit with a Rescorla-Wagner delta-rule learner and a choice-only inverse-temperature softmax decision process.
2AB_RW_DualAlpha_InvTempSoftmax    Two-armed bandit with a dual-alpha Rescorla-Wagner learner and a choice-only inverse-temperature softmax decision process.
3AB_RW_InvTempSoftmax              Three-armed bandit with a Rescorla-Wagner delta-rule learner and a choice-only inverse-temperature softmax decision process.
4AB_RW_InvTempSoftmax              Four-armed bandit with a Rescorla-Wagner delta-rule learner and a choice-only inverse-temperature softmax decision process.

decision_process : angle
computed by learner: {'v'}
sampled parameters : ['rl_alpha', 'scaler', 'a', 'z', 't', 'theta']

3. Write your own learning rule¶

To register a model we need a learning function — the rule that turns trial-by-trial experience into a decision parameter. Here we implement the classic Rescorla–Wagner (RW) update for a two-armed bandit and use it to drive the diffusion drift rate v:

Each arm has an expected value q. Before the choice, the drift is the value difference scaled by a gain: v = (q_high - q_low) * scaler.
After feedback r, the chosen arm's value moves toward the outcome by the learning rate: q_chosen += rl_alpha * (r - q_chosen).

Two contract details matter for the function to plug into HSSM's registry:

Signature. A learning function receives one array of shape (n_trials, 4) whose columns are, in order, the four inputs we declare below: ["rl_alpha", "scaler", "response", "feedback"]. It returns the computed v for every trial. We use jax.lax.scan so the whole thing is JAX-differentiable (HSSM needs gradients for NUTS).
Response coding. The 2AB_RW_Angle simulator codes responses as -1/+1 (which diffusion boundary was hit). The learner needs a 0-based arm index to look up and update q, so we map response > 0 → arm 1, else arm 0. This mapping is on you: HSSM passes the raw response column straight through to your function.

In [4]:

Copied!





def rw_trial(q, row):
    """Run one RW trial: emit the pre-update drift, then update the chosen arm."""
    rl_alpha, scaler, response, feedback = row
    arm = jnp.astype(response > 0, jnp.int32)  # -1 -> arm 0, +1 -> arm 1
    v = (q[1] - q[0]) * scaler  # drift BEFORE this trial's update
    delta = feedback - q[arm]  # reward prediction error
    q = q.at[arm].set(q[arm] + rl_alpha * delta)  # RW value update
    return q, v


@annotate_function(inputs=["rl_alpha", "scaler", "response", "feedback"], outputs=["v"])
def compute_v(subject_trials):
    """Compute the RW drift for one participant (n_trials, 4) -> (n_trials,)."""
    q0 = jnp.ones(2) * 0.5  # both arms start at value 0.5
    _, v = scan(rw_trial, q0, subject_trials)
    return v


# `annotate_function` only attaches metadata (.inputs / .outputs) that the
# registry reads to wire columns to arguments -- it does not change how the
# function is called.
print("inputs :", compute_v.inputs)
print("outputs:", compute_v.outputs)
def rw_trial(q, row):
    """Run one RW trial: emit the pre-update drift, then update the chosen arm."""
    rl_alpha, scaler, response, feedback = row
    arm = jnp.astype(response > 0, jnp.int32)  # -1 -> arm 0, +1 -> arm 1
    v = (q[1] - q[0]) * scaler  # drift BEFORE this trial's update
    delta = feedback - q[arm]  # reward prediction error
    q = q.at[arm].set(q[arm] + rl_alpha * delta)  # RW value update
    return q, v


@annotate_function(inputs=["rl_alpha", "scaler", "response", "feedback"], outputs=["v"])
def compute_v(subject_trials):
    """Compute the RW drift for one participant (n_trials, 4) -> (n_trials,)."""
    q0 = jnp.ones(2) * 0.5  # both arms start at value 0.5
    _, v = scan(rw_trial, q0, subject_trials)
    return v


# `annotate_function` only attaches metadata (.inputs / .outputs) that the
# registry reads to wire columns to arguments -- it does not change how the
# function is called.
print("inputs :", compute_v.inputs)
print("outputs:", compute_v.outputs)

inputs : ['rl_alpha', 'scaler', 'response', 'feedback']
outputs: ['v']

4. Register the model (the recommended HSSM-side path)¶

register_rlssm_model() stores everything the fitter needs under one name. Register once; retrieve anywhere with get_rlssm_model_config(). Every argument is commented below.

In [5]:

Copied!





MODEL_NAME = "tutorial_2AB_RW_Angle"

if MODEL_NAME not in list_models():
    register_rlssm_model(
        name=MODEL_NAME,
        # built-in HSSM SSM likelihood (collapsing-bound DDM)
        decision_process="angle",
        # our function computes the drift `v`
        learning_process={"v": compute_v},
        # the RL params we will fit
        learning_process_params=["rl_alpha", "scaler"],
        # supported ranges for those RL params
        learning_process_bounds={
            "rl_alpha": (0.0, 1.0),
            "scaler": (0.0, 10.0),
        },
        # aligned with learning_process_params
        learning_process_params_default=[0.2, 2.0],
        # data column the learner needs beyond rt/response
        extra_fields=["feedback"],
        # response coding for the DECISION side; the angle default is [0, 1]
        # and would mislabel our data
        choices=[-1, 1],
        description="Tutorial RW + angle RLSSM registered entirely HSSM-side.",
    )

config = get_rlssm_model_config(MODEL_NAME)
# sampled: [rl_alpha, scaler, a, z, t, theta]
print("sampled parameters:", config.list_params)
print("computed by learner:", set(config.ssm_logp_func.computed))  # {'v'}
config
MODEL_NAME = "tutorial_2AB_RW_Angle"

if MODEL_NAME not in list_models():
    register_rlssm_model(
        name=MODEL_NAME,
        # built-in HSSM SSM likelihood (collapsing-bound DDM)
        decision_process="angle",
        # our function computes the drift `v`
        learning_process={"v": compute_v},
        # the RL params we will fit
        learning_process_params=["rl_alpha", "scaler"],
        # supported ranges for those RL params
        learning_process_bounds={
            "rl_alpha": (0.0, 1.0),
            "scaler": (0.0, 10.0),
        },
        # aligned with learning_process_params
        learning_process_params_default=[0.2, 2.0],
        # data column the learner needs beyond rt/response
        extra_fields=["feedback"],
        # response coding for the DECISION side; the angle default is [0, 1]
        # and would mislabel our data
        choices=[-1, 1],
        description="Tutorial RW + angle RLSSM registered entirely HSSM-side.",
    )

config = get_rlssm_model_config(MODEL_NAME)
# sampled: [rl_alpha, scaler, a, z, t, theta]
print("sampled parameters:", config.list_params)
print("computed by learner:", set(config.ssm_logp_func.computed))  # {'v'}
config

sampled parameters: ['rl_alpha', 'scaler', 'a', 'z', 't', 'theta']
computed by learner: {'v'}

Out[5]:

RLSSMConfig(model_name='tutorial_2AB_RW_Angle', description='Tutorial RW + angle RLSSM registered entirely HSSM-side.', response=['rt', 'response'], choices=(-1, 1), list_params=['rl_alpha', 'scaler', 'a', 'z', 't', 'theta'], bounds={'rl_alpha': (0.0, 1.0), 'scaler': (0.0, 10.0), 'a': (0.3, 3.0), 'z': (0.1, 0.9), 't': (0.001, 2.0), 'theta': (-0.1, 1.3)}, loglik=None, loglik_kind='approx_differentiable', backend=None, extra_fields=['feedback'], rv=None, decision_process_loglik_kind='approx_differentiable', learning_process_kind='blackbox', params_default=[0.2, 2.0, 1.65, 0.5, 1.0005, 0.6], decision_process='angle', learning_process={'v': <function compute_v at 0x123306700>}, ssm_logp_func=<function make_jax_matrix_logp_funcs_from_onnx.<locals>.logp at 0x123306840>)

5. Simulate data to fit¶

The registry gives us a fitting config. It does not simulate — so to get some data (and to run a posterior predictive check later) we still use the ssms.rl simulator for the matching preset. Our registered learner reproduces this preset's Rescorla–Wagner rule exactly, so the two agree.

We reuse the Basic tutorial's ground-truth parameters (small rl_alpha for a gradual, visible learning curve) and draw per-participant values around them.

In [6]:

Copied!





ssms_config = rl.preset.get("2AB_RW_Angle")

GROUP_THETA = {
    "rl_alpha": 0.08,  # learning rate (small -> gradual, visible learning)
    "scaler": 2.5,  # value-difference -> drift gain
    "a": 1.2,  # boundary separation
    "z": 0.5,  # starting-point bias (0.5 = unbiased)
    "t": 0.25,  # non-decision time (s)
    "theta": 0.35,  # boundary-collapse angle
}
SDS = {"rl_alpha": 0.03, "scaler": 0.40, "a": 0.20, "z": 0.06, "t": 0.05, "theta": 0.10}
BOUNDS = {
    "rl_alpha": (0.01, 1.0),
    "scaler": (0.1, 5.0),
    "a": (0.3, 2.5),
    "z": (0.1, 0.9),
    "t": (0.05, 1.0),
    "theta": (0.0, 1.2),
}
LIST_PARAMS = list(GROUP_THETA)

rng = np.random.default_rng(RANDOM_SEED)
theta_arrays = {
    name: np.clip(
        rng.normal(GROUP_THETA[name], SDS[name], N_PARTICIPANTS), *BOUNDS[name]
    )
    for name in LIST_PARAMS
}
true_params = pd.DataFrame(theta_arrays)
true_params.index.name = "participant_id"

data = rl.Simulator(ssms_config).simulate(
    theta=theta_arrays,
    n_trials=N_TRIALS,
    n_participants=N_PARTICIPANTS,
    random_state=RANDOM_SEED,
)
print("rows:", len(data), "| columns:", list(data.columns))
data.head()
ssms_config = rl.preset.get("2AB_RW_Angle")

GROUP_THETA = {
    "rl_alpha": 0.08,  # learning rate (small -> gradual, visible learning)
    "scaler": 2.5,  # value-difference -> drift gain
    "a": 1.2,  # boundary separation
    "z": 0.5,  # starting-point bias (0.5 = unbiased)
    "t": 0.25,  # non-decision time (s)
    "theta": 0.35,  # boundary-collapse angle
}
SDS = {"rl_alpha": 0.03, "scaler": 0.40, "a": 0.20, "z": 0.06, "t": 0.05, "theta": 0.10}
BOUNDS = {
    "rl_alpha": (0.01, 1.0),
    "scaler": (0.1, 5.0),
    "a": (0.3, 2.5),
    "z": (0.1, 0.9),
    "t": (0.05, 1.0),
    "theta": (0.0, 1.2),
}
LIST_PARAMS = list(GROUP_THETA)

rng = np.random.default_rng(RANDOM_SEED)
theta_arrays = {
    name: np.clip(
        rng.normal(GROUP_THETA[name], SDS[name], N_PARTICIPANTS), *BOUNDS[name]
    )
    for name in LIST_PARAMS
}
true_params = pd.DataFrame(theta_arrays)
true_params.index.name = "participant_id"

data = rl.Simulator(ssms_config).simulate(
    theta=theta_arrays,
    n_trials=N_TRIALS,
    n_participants=N_PARTICIPANTS,
    random_state=RANDOM_SEED,
)
print("rows:", len(data), "| columns:", list(data.columns))
data.head()

rows: 1800 | columns: ['participant_id', 'trial_id', 'rt', 'response', 'feedback']

Out[6]:

	trial_id	rt	response	feedback
0	0	1.397879	-1	1.0
1	1	0.424477	-1	0.0
2	2	0.498285	-1	1.0
3	3	1.050406	-1	0.0
4	4	1.527007	-1	1.0

6. Fit with the registered config¶

The only new argument versus the Basic tutorial is model_config=config — the config we just pulled from the registry. Everything else is the standard hierarchical RLSSM recipe: identity links, a TruncatedNormal group intercept plus a per-participant random effect for each parameter, and process_initvals=False (mandatory — without it NUTS collapses to the float32 step-size floor and the posterior never leaves the prior).

In [7]:

Copied!





PARTICIPANT_EFFECT_PRIOR = {
    "name": "Normal",
    "mu": 0,
    "sigma": {"name": "HalfNormal", "sigma": 0.5},
}


def hierarchical_param(name, lower, upper, mu, sigma):
    """Build a group intercept (TruncatedNormal) + per-participant random effect."""
    return hssm.Param(
        name,
        formula=f"{name} ~ 1 + (1|participant_id)",
        prior={
            "Intercept": hssm.Prior(
                "TruncatedNormal", lower=lower, upper=upper, mu=mu, sigma=sigma
            ),
            "1|participant_id": PARTICIPANT_EFFECT_PRIOR,
        },
    )


model = hssm.RLSSM(
    data=data,
    model_config=config,  # <- the registered, named config
    p_outlier=0,
    lapse=None,
    process_initvals=False,  # mandatory for RLSSM sampling stability
    include=[
        hierarchical_param("rl_alpha", 0.01, 1.0, 0.15, 0.15),
        hierarchical_param("scaler", 0.1, 5.0, 2.0, 0.8),
        hierarchical_param("a", 0.3, 2.5, 1.1, 0.3),
        hierarchical_param("z", 0.1, 0.9, 0.5, 0.15),
        hierarchical_param("t", 0.05, 1.0, 0.25, 0.1),
        hierarchical_param("theta", 0.0, 1.2, 0.35, 0.15),
    ],
)
print("free parameters:", list(model.params.keys()))
assert "v" not in model.params  # v is computed by our learner, never sampled
PARTICIPANT_EFFECT_PRIOR = {
    "name": "Normal",
    "mu": 0,
    "sigma": {"name": "HalfNormal", "sigma": 0.5},
}


def hierarchical_param(name, lower, upper, mu, sigma):
    """Build a group intercept (TruncatedNormal) + per-participant random effect."""
    return hssm.Param(
        name,
        formula=f"{name} ~ 1 + (1|participant_id)",
        prior={
            "Intercept": hssm.Prior(
                "TruncatedNormal", lower=lower, upper=upper, mu=mu, sigma=sigma
            ),
            "1|participant_id": PARTICIPANT_EFFECT_PRIOR,
        },
    )


model = hssm.RLSSM(
    data=data,
    model_config=config,  # <- the registered, named config
    p_outlier=0,
    lapse=None,
    process_initvals=False,  # mandatory for RLSSM sampling stability
    include=[
        hierarchical_param("rl_alpha", 0.01, 1.0, 0.15, 0.15),
        hierarchical_param("scaler", 0.1, 5.0, 2.0, 0.8),
        hierarchical_param("a", 0.3, 2.5, 1.1, 0.3),
        hierarchical_param("z", 0.1, 0.9, 0.5, 0.15),
        hierarchical_param("t", 0.05, 1.0, 0.25, 0.1),
        hierarchical_param("theta", 0.0, 1.2, 0.35, 0.15),
    ],
)
print("free parameters:", list(model.params.keys()))
assert "v" not in model.params  # v is computed by our learner, never sampled

You supplied a model 'tutorial_2AB_RW_Angle', which is currently not supported in the ssm_simulators package. An error will be thrown when sampling from the random variable or when using any posterior or prior predictive sampling methods.

Model initialized successfully.

free parameters: ['rl_alpha', 'scaler', 'a', 'z', 't', 'theta']

7. Group-level recovery¶

A light check that the fit found the truth: the group intercept (posterior mean + 94% HDI) against each true group mean. For a full per-participant recovery treatment see the Basic and Advanced tutorials — we keep this one short on purpose.

In [9]:

Copied!





def group_recovery(idata, true_group):
    """Group intercept posterior vs. true group mean, per parameter."""
    names = [f"{p}_Intercept" for p in LIST_PARAMS]
    summ = az.summary(
        idata,
        var_names=names,
        kind="stats",
        ci_kind="hdi",
        ci_prob=0.94,
        round_to="none",
    )
    summ.index = LIST_PARAMS
    summ["true"] = [true_group[p] for p in LIST_PARAMS]

    fig, ax = plt.subplots(figsize=(8, 4.5))
    y = np.arange(len(LIST_PARAMS))
    ax.errorbar(
        summ["mean"],
        y,
        xerr=[summ["mean"] - summ["hdi94_lb"], summ["hdi94_ub"] - summ["mean"]],
        fmt="o",
        capsize=4,
        label="posterior (94% HDI)",
    )
    ax.scatter(
        summ["true"], y, color="crimson", marker="D", zorder=5, label="true group mean"
    )
    ax.set_yticks(y)
    ax.set_yticklabels(LIST_PARAMS)
    ax.invert_yaxis()
    ax.set_title("Group-level recovery")
    ax.legend()
    fig.tight_layout()
    plt.show()
    return summ[["mean", "hdi94_lb", "hdi94_ub", "true"]].round(3)


diagnostics = az.summary(idata, kind="diagnostics", round_to="none")
print("max r-hat:", float(diagnostics["r_hat"].max()))
group_recovery(idata, GROUP_THETA)
def group_recovery(idata, true_group):
    """Group intercept posterior vs. true group mean, per parameter."""
    names = [f"{p}_Intercept" for p in LIST_PARAMS]
    summ = az.summary(
        idata,
        var_names=names,
        kind="stats",
        ci_kind="hdi",
        ci_prob=0.94,
        round_to="none",
    )
    summ.index = LIST_PARAMS
    summ["true"] = [true_group[p] for p in LIST_PARAMS]

    fig, ax = plt.subplots(figsize=(8, 4.5))
    y = np.arange(len(LIST_PARAMS))
    ax.errorbar(
        summ["mean"],
        y,
        xerr=[summ["mean"] - summ["hdi94_lb"], summ["hdi94_ub"] - summ["mean"]],
        fmt="o",
        capsize=4,
        label="posterior (94% HDI)",
    )
    ax.scatter(
        summ["true"], y, color="crimson", marker="D", zorder=5, label="true group mean"
    )
    ax.set_yticks(y)
    ax.set_yticklabels(LIST_PARAMS)
    ax.invert_yaxis()
    ax.set_title("Group-level recovery")
    ax.legend()
    fig.tight_layout()
    plt.show()
    return summ[["mean", "hdi94_lb", "hdi94_ub", "true"]].round(3)


diagnostics = az.summary(idata, kind="diagnostics", round_to="none")
print("max r-hat:", float(diagnostics["r_hat"].max()))
group_recovery(idata, GROUP_THETA)

max r-hat: 1.0262178975792724

No description has been provided for this image

Out[9]:

	mean	hdi94_lb	hdi94_ub	true
rl_alpha	0.094	0.062	0.136	0.08
scaler	2.158	1.873	2.483	2.50
a	1.166	1.024	1.287	1.20
z	0.501	0.458	0.542	0.50
t	0.236	0.190	0.286	0.25
theta	0.403	0.312	0.494	0.35

8. A quick posterior predictive check¶

RLSSM-aware PPC re-simulates the learning process conditioned on the observed choices and feedback (mode="ppc"), using posterior-mean parameters. We compare one predicted dataset against the observations along three views: the learning curve (does the fitted model reproduce the climb toward the high-reward arm?), the RT distribution, and the overall response proportions. A fuller multi-draw PPC lives in the other tutorials.

In [10]:

Copied!





# Posterior-mean per-participant parameters (identity links -> natural scale).
post = idata.posterior
ppc_theta = {}
for name in LIST_PARAMS:
    re = post[f"{name}_1|participant_id"]
    pid_dim = [d for d in re.dims if d not in ("chain", "draw")][0]
    per_p = (post[f"{name}_Intercept"] + re).mean(("chain", "draw"))
    ids = [int(v) for v in re[pid_dim].values]
    ppc_theta[name] = (
        pd.Series(np.asarray(per_p.values), index=ids).sort_index().to_numpy()
    )

ppc = rl.Simulator(ssms_config).simulate(
    theta=ppc_theta,
    mode="ppc",
    observed_data=data,
    random_state=RANDOM_SEED + 1,
)


def _clean(df):
    """Drop non-finite RTs and sentinel (missed) responses."""
    return df[np.isfinite(df["rt"]) & (df["rt"] > 0) & (df["response"] > -900)]


def learning_curve(df, bin_size=15):
    """P(chose the high-reward arm) over trials (high-reward arm = response == -1)."""
    d = _clean(df).copy()
    d["chose_high"] = (d["response"] == -1).astype(float)
    d["trial_bin"] = (d["trial_id"] // bin_size) * bin_size
    return d.groupby("trial_bin")["chose_high"].mean()


fig, axes = plt.subplots(1, 3, figsize=(15, 4), constrained_layout=True)

# (a) Learning-curve PPC: does the fit reproduce the climb toward the good arm?
obs_curve, ppc_curve = learning_curve(data), learning_curve(ppc)
axes[0].plot(
    obs_curve.index + 7.5, obs_curve.values, "o-", color="black", label="observed"
)
axes[0].plot(
    ppc_curve.index + 7.5, ppc_curve.values, "s--", color="tab:blue", label="PPC"
)
axes[0].axhline(0.5, color="0.7", ls=":", lw=1)
axes[0].set(
    title="Learning-curve PPC",
    xlabel="Trial",
    ylabel="P(chose high-reward arm)",
    ylim=(0, 1),
)
axes[0].legend(frameon=False)

# (b) RT distribution
axes[1].hist(_clean(data)["rt"], bins=25, density=True, alpha=0.5, label="observed")
axes[1].hist(_clean(ppc)["rt"], bins=25, density=True, alpha=0.5, label="PPC")
axes[1].set(title="RT distribution", xlabel="RT (s)", ylabel="density")
axes[1].legend(frameon=False)

# (c) Response proportions
for label, df, off in [("observed", data, -0.2), ("PPC", ppc, 0.2)]:
    props = _clean(df)["response"].value_counts(normalize=True).sort_index()
    axes[2].bar(
        np.array([-1, 1]) + off,
        [props.get(-1, 0), props.get(1, 0)],
        width=0.4,
        label=label,
    )
axes[2].set(title="Response proportions", xticks=[-1, 1], xlabel="response")
axes[2].legend(frameon=False)
plt.show()
# Posterior-mean per-participant parameters (identity links -> natural scale).
post = idata.posterior
ppc_theta = {}
for name in LIST_PARAMS:
    re = post[f"{name}_1|participant_id"]
    pid_dim = [d for d in re.dims if d not in ("chain", "draw")][0]
    per_p = (post[f"{name}_Intercept"] + re).mean(("chain", "draw"))
    ids = [int(v) for v in re[pid_dim].values]
    ppc_theta[name] = (
        pd.Series(np.asarray(per_p.values), index=ids).sort_index().to_numpy()
    )

ppc = rl.Simulator(ssms_config).simulate(
    theta=ppc_theta,
    mode="ppc",
    observed_data=data,
    random_state=RANDOM_SEED + 1,
)


def _clean(df):
    """Drop non-finite RTs and sentinel (missed) responses."""
    return df[np.isfinite(df["rt"]) & (df["rt"] > 0) & (df["response"] > -900)]


def learning_curve(df, bin_size=15):
    """P(chose the high-reward arm) over trials (high-reward arm = response == -1)."""
    d = _clean(df).copy()
    d["chose_high"] = (d["response"] == -1).astype(float)
    d["trial_bin"] = (d["trial_id"] // bin_size) * bin_size
    return d.groupby("trial_bin")["chose_high"].mean()


fig, axes = plt.subplots(1, 3, figsize=(15, 4), constrained_layout=True)

# (a) Learning-curve PPC: does the fit reproduce the climb toward the good arm?
obs_curve, ppc_curve = learning_curve(data), learning_curve(ppc)
axes[0].plot(
    obs_curve.index + 7.5, obs_curve.values, "o-", color="black", label="observed"
)
axes[0].plot(
    ppc_curve.index + 7.5, ppc_curve.values, "s--", color="tab:blue", label="PPC"
)
axes[0].axhline(0.5, color="0.7", ls=":", lw=1)
axes[0].set(
    title="Learning-curve PPC",
    xlabel="Trial",
    ylabel="P(chose high-reward arm)",
    ylim=(0, 1),
)
axes[0].legend(frameon=False)

# (b) RT distribution
axes[1].hist(_clean(data)["rt"], bins=25, density=True, alpha=0.5, label="observed")
axes[1].hist(_clean(ppc)["rt"], bins=25, density=True, alpha=0.5, label="PPC")
axes[1].set(title="RT distribution", xlabel="RT (s)", ylabel="density")
axes[1].legend(frameon=False)

# (c) Response proportions
for label, df, off in [("observed", data, -0.2), ("PPC", ppc, 0.2)]:
    props = _clean(df)["response"].value_counts(normalize=True).sort_index()
    axes[2].bar(
        np.array([-1, 1]) + off,
        [props.get(-1, 0), props.get(1, 0)],
        width=0.4,
        label=label,
    )
axes[2].set(title="Response proportions", xticks=[-1, 1], xlabel="response")
axes[2].legend(frameon=False)
plt.show()

9. Extension: separate learning rates for positive and negative prediction errors¶

The standard RW rule assumes that better-than-expected and worse-than-expected outcomes change values at the same speed. A common extension uses two rates:

rl_alpha when the reward-prediction error is zero or positive;
rl_alpha_neg when the reward-prediction error is negative.

Adding one parameter changes several connected pieces of the recipe. The learner's input signature, registry parameter list, bounds and defaults, simulation values, and fit specification must all include rl_alpha_neg. The example below keeps those changes together so it can serve as a template for other custom learning rules.

In [11]:

Copied!





def dual_alpha_rw_trial(q, row):
    """One asymmetric RW trial with a learning rate selected by the RPE sign."""
    rl_alpha, rl_alpha_neg, scaler, response, feedback = row
    arm = jnp.astype(response > 0, jnp.int32)
    v = (q[1] - q[0]) * scaler
    delta = feedback - q[arm]
    alpha = jnp.where(delta < 0, rl_alpha_neg, rl_alpha)
    q = q.at[arm].set(q[arm] + alpha * delta)
    return q, v


@annotate_function(
    inputs=["rl_alpha", "rl_alpha_neg", "scaler", "response", "feedback"],
    outputs=["v"],
)
def compute_v_dual_alpha(subject_trials):
    """Asymmetric RW drift for one participant."""
    q0 = jnp.ones(2) * 0.5
    _, v = scan(dual_alpha_rw_trial, q0, subject_trials)
    return v
def dual_alpha_rw_trial(q, row):
    """One asymmetric RW trial with a learning rate selected by the RPE sign."""
    rl_alpha, rl_alpha_neg, scaler, response, feedback = row
    arm = jnp.astype(response > 0, jnp.int32)
    v = (q[1] - q[0]) * scaler
    delta = feedback - q[arm]
    alpha = jnp.where(delta < 0, rl_alpha_neg, rl_alpha)
    q = q.at[arm].set(q[arm] + alpha * delta)
    return q, v


@annotate_function(
    inputs=["rl_alpha", "rl_alpha_neg", "scaler", "response", "feedback"],
    outputs=["v"],
)
def compute_v_dual_alpha(subject_trials):
    """Asymmetric RW drift for one participant."""
    q0 = jnp.ones(2) * 0.5
    _, v = scan(dual_alpha_rw_trial, q0, subject_trials)
    return v

Register the extended learner¶

The registration call changes only where the new parameter enters: it is added to the ordered parameter list, given bounds, and assigned a default. The decision process, feedback field, and response coding are unchanged.

In [12]:

Copied!





DUAL_MODEL_NAME = "tutorial_2AB_RW_DualAlpha_Angle"

if DUAL_MODEL_NAME not in list_models():
    register_rlssm_model(
        name=DUAL_MODEL_NAME,
        decision_process="angle",
        learning_process={"v": compute_v_dual_alpha},
        learning_process_params=["rl_alpha", "rl_alpha_neg", "scaler"],
        learning_process_bounds={
            "rl_alpha": (0.0, 1.0),
            "rl_alpha_neg": (0.0, 1.0),
            "scaler": (0.0, 10.0),
        },
        learning_process_params_default=[0.2, 0.2, 2.0],
        extra_fields=["feedback"],
        choices=[-1, 1],
        description="Tutorial asymmetric RW + angle RLSSM.",
    )

dual_config = get_rlssm_model_config(DUAL_MODEL_NAME)
assert {"rl_alpha", "rl_alpha_neg"}.issubset(dual_config.list_params)
assert set(dual_config.ssm_logp_func.computed) == {"v"}
print("sampled parameters:", dual_config.list_params)
DUAL_MODEL_NAME = "tutorial_2AB_RW_DualAlpha_Angle"

if DUAL_MODEL_NAME not in list_models():
    register_rlssm_model(
        name=DUAL_MODEL_NAME,
        decision_process="angle",
        learning_process={"v": compute_v_dual_alpha},
        learning_process_params=["rl_alpha", "rl_alpha_neg", "scaler"],
        learning_process_bounds={
            "rl_alpha": (0.0, 1.0),
            "rl_alpha_neg": (0.0, 1.0),
            "scaler": (0.0, 10.0),
        },
        learning_process_params_default=[0.2, 0.2, 2.0],
        extra_fields=["feedback"],
        choices=[-1, 1],
        description="Tutorial asymmetric RW + angle RLSSM.",
    )

dual_config = get_rlssm_model_config(DUAL_MODEL_NAME)
assert {"rl_alpha", "rl_alpha_neg"}.issubset(dual_config.list_params)
assert set(dual_config.ssm_logp_func.computed) == {"v"}
print("sampled parameters:", dual_config.list_params)

sampled parameters: ['rl_alpha', 'rl_alpha_neg', 'scaler', 'a', 'z', 't', 'theta']

Simulate matched data¶

ssm-simulators already includes RescorlaWagnerDualAlphaRule, so we use that tested implementation to generate a small matched dataset. We still wrote the HSSM-side function independently above because the purpose here is to show how users can introduce different custom learning rules into the registry. If a new rule also needs simulation or posterior predictive checks, implementing its matching ssms.rl learner is the reusable route demonstrated in the Advanced tutorial.

In [13]:

Copied!





dual_ssms_config = rl.ModelConfig(
    model_name="2AB_RW_DualAlpha_Angle_Tutorial",
    description="Two-arm angle RLSSM with separate positive/negative learning rates.",
    decision_process="angle",
    learning_process=rl.learning.RescorlaWagnerDualAlphaRule(
        n_actions=2,
        initial_q=0.5,
    ),
    task_environment=rl.env.Bandit.bernoulli(
        probabilities=[0.7, 0.3],
        response_labels=[-1, 1],
    ),
)

DUAL_THETA = {
    "rl_alpha": 0.35,  # positive RPE learning rate
    "rl_alpha_neg": 0.10,  # negative RPE learning rate
    "scaler": 2.5,
    "a": 1.2,
    "z": 0.5,
    "t": 0.25,
    "theta": 0.35,
}
dual_theta_arrays = {
    name: np.repeat(value, DUAL_N_PARTICIPANTS) for name, value in DUAL_THETA.items()
}
dual_data = rl.Simulator(dual_ssms_config).simulate(
    theta=dual_theta_arrays,
    n_trials=DUAL_N_TRIALS,
    n_participants=DUAL_N_PARTICIPANTS,
    random_state=RANDOM_SEED + 20,
)
dual_ssms_config.validate_data(dual_data)
print("rows:", len(dual_data), "| columns:", list(dual_data.columns))
dual_ssms_config = rl.ModelConfig(
    model_name="2AB_RW_DualAlpha_Angle_Tutorial",
    description="Two-arm angle RLSSM with separate positive/negative learning rates.",
    decision_process="angle",
    learning_process=rl.learning.RescorlaWagnerDualAlphaRule(
        n_actions=2,
        initial_q=0.5,
    ),
    task_environment=rl.env.Bandit.bernoulli(
        probabilities=[0.7, 0.3],
        response_labels=[-1, 1],
    ),
)

DUAL_THETA = {
    "rl_alpha": 0.35,  # positive RPE learning rate
    "rl_alpha_neg": 0.10,  # negative RPE learning rate
    "scaler": 2.5,
    "a": 1.2,
    "z": 0.5,
    "t": 0.25,
    "theta": 0.35,
}
dual_theta_arrays = {
    name: np.repeat(value, DUAL_N_PARTICIPANTS) for name, value in DUAL_THETA.items()
}
dual_data = rl.Simulator(dual_ssms_config).simulate(
    theta=dual_theta_arrays,
    n_trials=DUAL_N_TRIALS,
    n_participants=DUAL_N_PARTICIPANTS,
    random_state=RANDOM_SEED + 20,
)
dual_ssms_config.validate_data(dual_data)
print("rows:", len(dual_data), "| columns:", list(dual_data.columns))

rows: 400 | columns: ['participant_id', 'trial_id', 'rt', 'response', 'feedback']

A deliberately lightweight Bayesian fit¶

To verify the complete registration-to-sampling path without turning this short extension into another recovery tutorial, we estimate only rl_alpha and rl_alpha_neg. We fix scaler, a, z, t, and theta at their generating values. Fixing them is a demonstration-only speed choice: it isolates the newly added rates and keeps sampling fast. In a substantive analysis you would generally estimate scientifically relevant decision parameters too, then study joint identifiability and recovery under the intended design.

The short single-chain run below is an integration check, not a formal recovery analysis. The posterior summary should therefore be read as evidence that both new parameters enter and sample correctly, not as a benchmark for recoverability.

In [14]:

Copied!





dual_model = hssm.RLSSM(
    data=dual_data,
    model_config=dual_config,
    p_outlier=0,
    lapse=None,
    process_initvals=False,
    include=[
        hssm.Param("rl_alpha", prior=hssm.Prior("Beta", alpha=2, beta=2)),
        hssm.Param("rl_alpha_neg", prior=hssm.Prior("Beta", alpha=2, beta=2)),
        hssm.Param("scaler", prior=DUAL_THETA["scaler"]),
        hssm.Param("a", prior=DUAL_THETA["a"]),
        hssm.Param("z", prior=DUAL_THETA["z"]),
        hssm.Param("t", prior=DUAL_THETA["t"]),
        hssm.Param("theta", prior=DUAL_THETA["theta"]),
    ],
)
dual_idata = dual_model.sample(
    sampler="numpyro",
    draws=DUAL_N_DRAWS,
    tune=DUAL_N_TUNE,
    chains=1,
    cores=1,
    target_accept=0.9,
    random_seed=RANDOM_SEED + 20,
)
dual_model = hssm.RLSSM(
    data=dual_data,
    model_config=dual_config,
    p_outlier=0,
    lapse=None,
    process_initvals=False,
    include=[
        hssm.Param("rl_alpha", prior=hssm.Prior("Beta", alpha=2, beta=2)),
        hssm.Param("rl_alpha_neg", prior=hssm.Prior("Beta", alpha=2, beta=2)),
        hssm.Param("scaler", prior=DUAL_THETA["scaler"]),
        hssm.Param("a", prior=DUAL_THETA["a"]),
        hssm.Param("z", prior=DUAL_THETA["z"]),
        hssm.Param("t", prior=DUAL_THETA["t"]),
        hssm.Param("theta", prior=DUAL_THETA["theta"]),
    ],
)
dual_idata = dual_model.sample(
    sampler="numpyro",
    draws=DUAL_N_DRAWS,
    tune=DUAL_N_TUNE,
    chains=1,
    cores=1,
    target_accept=0.9,
    random_seed=RANDOM_SEED + 20,
)

You supplied a model 'tutorial_2AB_RW_DualAlpha_Angle', which is currently not supported in the ssm_simulators package. An error will be thrown when sampling from the random variable or when using any posterior or prior predictive sampling methods.

Model initialized successfully.

Using default initvals.

NUTS[numpyro]: [rl_alpha_neg, rl_alpha]

  0%|                                                  | 0/600 [00:00<?, ?it/s]

warmup:   5%| | 30/600 [00:00<00:09, 63.20it/s, 7 steps of size 1.10e-01. acc.

warmup:  11%| | 68/600 [00:00<00:03, 135.44it/s, 7 steps of size 1.83e-01. acc.

warmup:  17%|▏| 103/600 [00:00<00:02, 188.10it/s, 7 steps of size 5.00e-01. acc

warmup:  23%|▏| 138/600 [00:00<00:02, 228.82it/s, 5 steps of size 1.15e+00. acc

warmup:  31%|▎| 184/600 [00:00<00:01, 289.79it/s, 7 steps of size 9.72e-01. acc

warmup:  37%|▎| 223/600 [00:00<00:01, 315.25it/s, 7 steps of size 7.76e-01. acc

warmup:  45%|▍| 268/600 [00:01<00:00, 351.59it/s, 9 steps of size 4.45e-01. acc

sample:  50%|▌| 302/600 [00:01<00:00, 346.94it/s, 3 steps of size 6.94e-01. acc

sample:  58%|▌| 348/600 [00:01<00:00, 376.67it/s, 7 steps of size 6.94e-01. acc

sample:  66%|▋| 396/600 [00:01<00:00, 402.99it/s, 7 steps of size 6.94e-01. acc

sample:  72%|▋| 433/600 [00:01<00:00, 393.13it/s, 7 steps of size 6.94e-01. acc

sample:  79%|▊| 474/600 [00:01<00:00, 397.28it/s, 3 steps of size 6.94e-01. acc

sample:  85%|▊| 512/600 [00:01<00:00, 391.25it/s, 7 steps of size 6.94e-01. acc

sample:  93%|▉| 557/600 [00:01<00:00, 408.25it/s, 1 steps of size 6.94e-01. acc

sample:  99%|▉| 596/600 [00:01<00:00, 401.09it/s, 3 steps of size 6.94e-01. acc

sample: 100%|█| 600/600 [00:01<00:00, 314.31it/s, 7 steps of size 6.94e-01. acc

Only one chain was sampled, this makes it impossible to run some convergence checks

In [15]:

Copied!





dual_summary = az.summary(
    dual_idata,
    var_names=["rl_alpha", "rl_alpha_neg"],
    kind="stats",
    ci_kind="hdi",
    ci_prob=0.94,
    round_to="none",
)[["mean", "hdi94_lb", "hdi94_ub"]]
dual_summary["true"] = [DUAL_THETA[name] for name in dual_summary.index]
dual_summary.round(3)
dual_summary = az.summary(
    dual_idata,
    var_names=["rl_alpha", "rl_alpha_neg"],
    kind="stats",
    ci_kind="hdi",
    ci_prob=0.94,
    round_to="none",
)[["mean", "hdi94_lb", "hdi94_ub"]]
dual_summary["true"] = [DUAL_THETA[name] for name in dual_summary.index]
dual_summary.round(3)

Out[15]:

	mean	hdi94_lb	hdi94_ub	true
rl_alpha	0.434	0.314	0.578	0.35
rl_alpha_neg	0.107	0.080	0.136	0.10

10. Under the hood: what the registry built¶

A registered model and a bridged model (from_ssms_model) populate the same RLSSMConfig slots — the registry is just a named shortcut. The key fields:

ssm_logp_func — the decision likelihood, with .computed listing learner-driven params.
learning_process — your {computed_param: function} mapping.
extra_fields — data columns the learner needs beyond rt/response.
list_params — every sampled parameter (RL params + non-computed SSM params).

In [16]:

Copied!





bridged = hssm.rl.RLSSMConfig.from_ssms_model(ssms_config)
print("registry list_params:", config.list_params)
print("bridge   list_params:", bridged.list_params)
print("same sampled parameters:", config.list_params == bridged.list_params)
print("registry extra_fields :", config.extra_fields)
print("registry computed     :", set(config.ssm_logp_func.computed))
bridged = hssm.rl.RLSSMConfig.from_ssms_model(ssms_config)
print("registry list_params:", config.list_params)
print("bridge   list_params:", bridged.list_params)
print("same sampled parameters:", config.list_params == bridged.list_params)
print("registry extra_fields :", config.extra_fields)
print("registry computed     :", set(config.ssm_logp_func.computed))

registry list_params: ['rl_alpha', 'scaler', 'a', 'z', 't', 'theta']
bridge   list_params: ['rl_alpha', 'scaler', 'a', 'z', 't', 'theta']
same sampled parameters: True
registry extra_fields : ['feedback']
registry computed     : {'v'}

11. Summary¶

Use register_rlssm_model() when you want a named, reusable RLSSM config built entirely HSSM-side — you supply the learning rule; HSSM supplies the decision likelihood.
You own the learning function. Match its input contract (one ordered column per declared input) and handle the response coding yourself (choices=[-1, 1] fixes the decision side; the arm-index mapping inside the learner fixes the learning side).
Extending a model with rl_alpha_neg requires coordinated changes to the learner signature, registry metadata, simulation parameters, and fit specification. The short sampled example fixes nuisance parameters only to isolate that wiring and stay fast.
To simulate a custom model (not just fit it), define it in ssms.rl and bridge with from_ssms_model instead — see the Advanced and Restless learner tutorials.

HSSM also supports choice-only RL models (choices without response times). Those are documented separately once validated for the current release.