Copied!





import logging
import os
import time
import warnings

# This how-to does not need a GPU. Keep Molab and headless documentation
# builds away from host CUDA plugins before importing HSSM/JAX.
os.environ["JAX_PLATFORMS"] = "cpu"
os.environ["JAX_SKIP_CUDA_CONSTRAINTS_CHECK"] = "1"
# Silence only PyMC's notice that numpyro ignores the `init` HSSM passes,
# numpyro's notice that one CPU device runs the chains in sequence, and Bambi
# deprecation notices raised inside HSSM itself; every other warning,
# including sampling diagnostics, stays visible.
warnings.filterwarnings(
    "ignore", message=r"`init='adapt_diag'` is ignored", category=UserWarning
)
warnings.filterwarnings(
    "ignore", message=r"There are not enough devices", category=UserWarning
)
warnings.filterwarnings(
    "ignore", message=r".*is deprecated; use", category=FutureWarning
)
logging.getLogger("jax._src.xla_bridge").setLevel(logging.CRITICAL)

import arviz as az
import marimo as mo
import numpy as np
import pytensor
import pytensor.tensor as pt
from ssms import rl

import hssm
from hssm.likelihoods import logp_ddm, logp_ddm_jax

hssm.set_floatX("float64")

FULL_RUN = os.environ.get("FULL_RUN") == "1"
import logging
import os
import time
import warnings

# This how-to does not need a GPU. Keep Molab and headless documentation
# builds away from host CUDA plugins before importing HSSM/JAX.
os.environ["JAX_PLATFORMS"] = "cpu"
os.environ["JAX_SKIP_CUDA_CONSTRAINTS_CHECK"] = "1"
# Silence only PyMC's notice that numpyro ignores the `init` HSSM passes,
# numpyro's notice that one CPU device runs the chains in sequence, and Bambi
# deprecation notices raised inside HSSM itself; every other warning,
# including sampling diagnostics, stays visible.
warnings.filterwarnings(
    "ignore", message=r"`init='adapt_diag'` is ignored", category=UserWarning
)
warnings.filterwarnings(
    "ignore", message=r"There are not enough devices", category=UserWarning
)
warnings.filterwarnings(
    "ignore", message=r".*is deprecated; use", category=FutureWarning
)
logging.getLogger("jax._src.xla_bridge").setLevel(logging.CRITICAL)

import arviz as az
import marimo as mo
import numpy as np
import pytensor
import pytensor.tensor as pt
from ssms import rl

import hssm
from hssm.likelihoods import logp_ddm, logp_ddm_jax

hssm.set_floatX("float64")

FULL_RUN = os.environ.get("FULL_RUN") == "1"

Setting PyTensor floatX type to float64.
Setting "jax_enable_x64" to True. If this is not intended, please set `jax` to False.

Use the analytical DDM as an RLSSM decision process¶

The built-in RLSSM presets evaluate their decision process with a likelihood approximation network (LAN) by default. For the DDM you can ask for HSSM's analytical likelihood (Navarro & Fuss, 2009) instead:

hssm.rl.RLSSM(data, model="2AB_RW_DDM", decision_process_loglik_kind="analytical")

This keeps the preset's Rescorla–Wagner learner and replaces only the decision likelihood. It is the same likelihood that hssm.HSSM(model="ddm", loglik_kind="analytical") uses, in the JAX form hssm.likelihoods.logp_ddm_jax that the RL likelihood needs. The LAN is not loaded at all. This guide checks that likelihood, fits a simulated two-armed bandit with it, and compares the fit with the LAN.

Requires HSSM 0.5.1 or later.

The JAX likelihood¶

logp_ddm_jax is a port of logp_ddm with the same parameterization, response coding and lower bound. On random parameters and response times from 1 ms to 10 s after t, the two agree to rounding error:

Copied!





_rng = np.random.default_rng(0)
_n = 500
_v = _rng.uniform(-3.0, 3.0, _n)
_a = _rng.uniform(0.3, 2.5, _n)
_z = _rng.uniform(0.1, 0.9, _n)
_t = _rng.uniform(0.0, 1.0, _n)
_obs = np.column_stack([_t + 10 ** _rng.uniform(-3.0, 1.0, _n), np.sign(_v)])

_obs_var = pt.matrix("obs", dtype="float64")
_param_vars = [pt.vector(name, dtype="float64") for name in "vazt"]
_pytensor_logp = pytensor.function(
    [_obs_var, *_param_vars], logp_ddm(_obs_var, *_param_vars)
)
_difference = np.abs(
    np.asarray(logp_ddm_jax(_obs, _v, _a, _z, _t))
    - _pytensor_logp(_obs, _v, _a, _z, _t)
)
mo.md(f"Largest absolute difference over {_n} trials: `{_difference.max():.1e}`")
_rng = np.random.default_rng(0)
_n = 500
_v = _rng.uniform(-3.0, 3.0, _n)
_a = _rng.uniform(0.3, 2.5, _n)
_z = _rng.uniform(0.1, 0.9, _n)
_t = _rng.uniform(0.0, 1.0, _n)
_obs = np.column_stack([_t + 10 ** _rng.uniform(-3.0, 1.0, _n), np.sign(_v)])

_obs_var = pt.matrix("obs", dtype="float64")
_param_vars = [pt.vector(name, dtype="float64") for name in "vazt"]
_pytensor_logp = pytensor.function(
    [_obs_var, *_param_vars], logp_ddm(_obs_var, *_param_vars)
)
_difference = np.abs(
    np.asarray(logp_ddm_jax(_obs, _v, _a, _z, _t))
    - _pytensor_logp(_obs, _v, _a, _z, _t)
)
mo.md(f"Largest absolute difference over {_n} trials: `{_difference.max():.1e}`")

Largest absolute difference over 500 trials: 4.3e-14

Simulate a two-armed bandit¶

ssms.rl simulates the same preset, so the data use the columns and response coding the model expects: participant_id, rt, response (-1 or 1) and feedback, grouped by participant in trial order, with the same number of trials per participant.

ssms simulates the DDM with Euler steps of 1 ms by default. That widens the effective boundary slightly, and an exact likelihood then recovers a about 0.02 above its generating value. A step of 0.1 ms shrinks that offset to about 0.005–0.01, well inside the posterior uncertainty, so we set it through the preset's ssm_kwargs.

Copied!





true_params = {"rl_alpha": 0.3, "scaler": 3.0, "a": 1.2, "z": 0.5, "t": 0.3}
preset = rl.preset.get("2AB_RW_DDM")
preset.ssm_kwargs = {"delta_t": 1e-4, "max_t": 20.0}
simulator = rl.Simulator(preset)
data = simulator.simulate(
    theta=true_params, n_trials=200, n_participants=5, random_state=1
)
data.head()
true_params = {"rl_alpha": 0.3, "scaler": 3.0, "a": 1.2, "z": 0.5, "t": 0.3}
preset = rl.preset.get("2AB_RW_DDM")
preset.ssm_kwargs = {"delta_t": 1e-4, "max_t": 20.0}
simulator = rl.Simulator(preset)
data = simulator.simulate(
    theta=true_params, n_trials=200, n_participants=5, random_state=1
)
data.head()

	trial_id	rt	response	feedback
0	0	0.554569	1	1.0
1	1	1.563886	1	0.0
2	2	2.479979	1	1.0
3	3	1.688493	1	0.0
4	4	0.838366	-1	0.0

Fit with the analytical and the LAN likelihood¶

The two models differ only in decision_process_loglik_kind. The simulated data contain no lapses, so both switch off HSSM's default 5% lapse component with p_outlier=None; with real data, keep it unless you have a reason not to. Both use numpyro, the default sampler for RLSSM models.

Copied!





_draws = 1_000 if FULL_RUN else 300
fits = {}
for _kind in ["analytical", "approx_differentiable"]:
    _model = hssm.rl.RLSSM(
        data,
        model="2AB_RW_DDM",
        decision_process_loglik_kind=_kind,
        p_outlier=None,
    )
    _start = time.perf_counter()
    _idata = _model.sample(
        chains=4, draws=_draws, tune=_draws, random_seed=1, progressbar=False
    )
    fits[_kind] = {"idata": _idata, "seconds": time.perf_counter() - _start}
_draws = 1_000 if FULL_RUN else 300
fits = {}
for _kind in ["analytical", "approx_differentiable"]:
    _model = hssm.rl.RLSSM(
        data,
        model="2AB_RW_DDM",
        decision_process_loglik_kind=_kind,
        p_outlier=None,
    )
    _start = time.perf_counter()
    _idata = _model.sample(
        chains=4, draws=_draws, tune=_draws, random_seed=1, progressbar=False
    )
    fits[_kind] = {"idata": _idata, "seconds": time.perf_counter() - _start}

You supplied a model '2AB_RW_DDM', 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]: [z, t, a, scaler, rl_alpha]
You supplied a model '2AB_RW_DDM', 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]: [z, t, a, scaler, rl_alpha]

Compare the posteriors¶

Copied!





_rows = []
for _name, _value in true_params.items():
    _row = f"| `{_name}` | {_value} |"
    for _kind in ["analytical", "approx_differentiable"]:
        _summary = az.summary(
            fits[_kind]["idata"],
            var_names=[_name],
            ci_kind="hdi",
            ci_prob=0.94,
            round_to="none",
        ).iloc[0]
        _row += (
            f" {_summary['mean']:.3f} [{_summary['hdi94_lb']:.3f}, "
            f"{_summary['hdi94_ub']:.3f}] | {_summary['r_hat']:.3f} |"
        )
    _rows.append(_row)

_stats = []
for _kind in ["analytical", "approx_differentiable"]:
    _divergences = int(fits[_kind]["idata"]["sample_stats"]["diverging"].sum())
    _stats.append(
        f"`{_kind}`: {fits[_kind]['seconds']:.0f} s, {_divergences} divergences"
    )

mo.md(
    "| parameter | truth | analytical: mean [94% HDI] | r_hat "
    "| LAN: mean [94% HDI] | r_hat |\n|---|---|---|---|---|---|\n"
    + "\n".join(_rows)
    + "\n\nSampling time on the machine that built this page, and divergences: "
    + "; ".join(_stats)
    + "."
)
_rows = []
for _name, _value in true_params.items():
    _row = f"| `{_name}` | {_value} |"
    for _kind in ["analytical", "approx_differentiable"]:
        _summary = az.summary(
            fits[_kind]["idata"],
            var_names=[_name],
            ci_kind="hdi",
            ci_prob=0.94,
            round_to="none",
        ).iloc[0]
        _row += (
            f" {_summary['mean']:.3f} [{_summary['hdi94_lb']:.3f}, "
            f"{_summary['hdi94_ub']:.3f}] | {_summary['r_hat']:.3f} |"
        )
    _rows.append(_row)

_stats = []
for _kind in ["analytical", "approx_differentiable"]:
    _divergences = int(fits[_kind]["idata"]["sample_stats"]["diverging"].sum())
    _stats.append(
        f"`{_kind}`: {fits[_kind]['seconds']:.0f} s, {_divergences} divergences"
    )

mo.md(
    "| parameter | truth | analytical: mean [94% HDI] | r_hat "
    "| LAN: mean [94% HDI] | r_hat |\n|---|---|---|---|---|---|\n"
    + "\n".join(_rows)
    + "\n\nSampling time on the machine that built this page, and divergences: "
    + "; ".join(_stats)
    + "."
)

parameter	truth	analytical: mean [94% HDI]	r_hat	LAN: mean [94% HDI]	r_hat
`rl_alpha`	0.3	0.273 [0.237, 0.314]	1.001	0.274 [0.237, 0.316]	1.002
`scaler`	3.0	3.025 [2.759, 3.281]	1.001	3.045 [2.769, 3.310]	1.001
`a`	1.2	1.211 [1.162, 1.265]	1.003	1.188 [1.138, 1.242]	1.002
`z`	0.5	0.493 [0.465, 0.520]	1.002	0.495 [0.465, 0.524]	1.001
`t`	0.3	0.296 [0.275, 0.314]	1.001	0.299 [0.277, 0.319]	1.001

Sampling time on the machine that built this page, and divergences: analytical: 43 s, 0 divergences; approx_differentiable: 153 s, 0 divergences.

Both likelihoods describe the same data, so their posteriors should be close. The analytical likelihood is exact and usually samples faster. It also has no training range. The learner's drift v = scaler * (Q1 - Q0) can reach ±10 within the preset's scaler bounds, while the DDM network was trained on v in (-3, 3) and nothing checks the computed drift against that range.

Check the fit by simulation¶

model.sample_posterior_predictive() does not support RLSSM models yet, whichever likelihood they use. Simulate from ssms.rl instead: with mode="ppc", the simulator replays the observed task, so the learner sees the same feedback each participant did. Here we simulate once at the analytical posterior means. This limitation is also why HSSM warned, when it built each model above, that 2AB_RW_DDM is not supported in ssm_simulators: the warning concerns predictive sampling only, not the fit.

Copied!





_posterior = fits["analytical"]["idata"]["posterior"]
_means = {name: float(_posterior[name].mean()) for name in true_params}
_predicted = simulator.simulate(
    theta=_means, mode="ppc", observed_data=data, random_state=2
)
mo.md(
    "| | P(response = 1) | mean RT (s) |\n|---|---|---|\n"
    + "\n".join(
        f"| {_label} | {(_df['response'] == 1).mean():.3f} | {_df['rt'].mean():.3f} |"
        for _label, _df in [("observed", data), ("simulated", _predicted)]
    )
)
_posterior = fits["analytical"]["idata"]["posterior"]
_means = {name: float(_posterior[name].mean()) for name in true_params}
_predicted = simulator.simulate(
    theta=_means, mode="ppc", observed_data=data, random_state=2
)
mo.md(
    "| | P(response = 1) | mean RT (s) |\n|---|---|---|\n"
    + "\n".join(
        f"| {_label} | {(_df['response'] == 1).mean():.3f} | {_df['rt'].mean():.3f} |"
        for _label, _df in [("observed", data), ("simulated", _predicted)]
    )
)

	P(response = 1)	mean RT (s)
observed	0.098	1.135
simulated	0.095	1.139

Notes¶

Which decision processes. "analytical" works for decision processes whose analytical likelihood declares a JAX implementation (jax_loglik in its model configuration): the DDM, and the LBA with 2, 3 or 4 choices (lba2, lba3, lba4; see below). Asking for it on a preset without one, such as "2AB_RW_Angle", raises an error that names what is missing.
Other entry points. hssm.rl.get_rlssm_model_config and hssm.rl.RLSSMConfig.from_ssms_model take the same decision_process_loglik_kind. A model registered with hssm.rl.register_rlssm_model(..., decision_process="ddm", decision_process_loglik_kind="analytical") takes its decision-parameter bounds from the analytical likelihood; the presets keep their ssms bounds whichever likelihood you choose.
Passing a config. With model_config=, the config's decision_process_loglik_kind decides. Passing a different decision_process_loglik_kind alongside it raises an error rather than silently using one of the two.
Response coding. Keep the -1/1 coding of the simulated data: the preset's learner maps it to the two arms itself.
Sampler. Both kinds default to numpyro. If you choose sampler="pymc", use cores=1 (or mp_ctx="forkserver") on Apple-silicon Macs: PyMC's default process start method can hang with HSSM's JAX likelihoods there.

An LBA decision process¶

No ssms preset pairs a learner with the LBA yet, so build the model as an ssms.rl.ModelConfig. RescorlaWagnerRaceDrifts gives arm i the drift v_i = scaler * q_i, and responses are coded 0, ..., n - 1:

import numpy as np
from ssms import rl
from ssms.rl.env import Bandit
from ssms.rl.learning import RescorlaWagnerRaceDrifts

rw_lba3 = rl.ModelConfig(
    model_name="3AB_RW_LBA3",
    description="Three-armed bandit, Rescorla-Wagner learner, LBA decision process.",
    decision_process="lba3",
    learning_process=RescorlaWagnerRaceDrifts(n_actions=3, initial_q=0.5),
    task_environment=Bandit.bernoulli(
        probabilities=[0.7, 0.2, 0.1], response_labels=[0, 1, 2]
    ),
)
# random_state seeds the bandit; the LBA draws use NumPy's global RNG (ssm-simulators#399).
np.random.seed(0)
lba_data = rl.Simulator(rw_lba3).simulate(
    theta={"rl_alpha": 0.3, "scaler": 3.0, "A": 0.3, "b": 0.6},
    n_trials=200,
    n_participants=5,
    random_state=0,
)
config = hssm.rl.RLSSMConfig.from_ssms_model(
    rw_lba3, decision_process_loglik_kind="analytical"
)
model = hssm.rl.RLSSM(lba_data, model_config=config, p_outlier=None)

The free parameters are rl_alpha, scaler, A and b; this LBA has no non-decision time.