RLSSM Simulation and HSSM Handoff¶

Audience:

• Users who can already simulate a basic SSM and now want trial-wise learning models.
• HSSM users who want to understand what ssms contributes before inference.

Prerequisites:

• ssm-simulators installed with compiled simulator extensions.
• Familiarity with response-time/choice data and basic reinforcement learning terms.

By the end you will be able to:

• Explain the three moving parts of an RLSSM.
• Simulate synthetic RT + choice RLSSM data with a built-in preset.
• Validate the panel that will be passed to HSSM.
• Understand how the same ssms learning rule supports simulation, likelihood construction, and posterior predictive checks.

Outline¶

1. Load a built-in RLSSM preset.
2. Inspect the learning process, decision process, and task environment.
3. Generate a small synthetic panel.
4. Validate the data contract used by HSSM.
5. Run observed-history-conditioned posterior predictive simulation.

In [1]:

from __future__ import annotations

import pandas as pd

import ssms.rl as rl

pd.set_option("display.max_columns", 20)
pd.set_option("display.width", 120)

SEED = 37

from __future__ import annotations import pandas as pd import ssms.rl as rl pd.set_option("display.max_columns", 20) pd.set_option("display.width", 120) SEED = 37

1. What is an RLSSM?¶

A reinforcement-learning sequential sampling model combines two loops.

• Across trials, a learning rule updates latent state, such as Q-values, from choices and feedback.
• Within each trial, a decision process maps the current learned state to a response, and for RT + choice models, a response time.

In ssms this is explicit: a ModelConfig combines a learning process, a task environment, and a decision process. HSSM can later consume the same structural config for inference through RLSSMConfig.from_ssms_model(...).

In [2]:

config = rl.preset.get("2AB_RW_Angle")
print(rl.preset.info("2AB_RW_Angle"))

config = rl.preset.get("2AB_RW_Angle") print(rl.preset.info("2AB_RW_Angle"))

Preset: 2AB_RW_Angle
Description: Two-armed bandit with a Rescorla-Wagner delta-rule learner and an angle decision process.
Task: two-armed Bernoulli bandit
Learning process: RescorlaWagnerDrift
Decision process: angle
Required parameters: rl_alpha, scaler, a, z, t, theta
Default parameters: rl_alpha=0.2, scaler=2, a=1, z=0.5, t=0.001, theta=0
Response labels: (-1, 1)
Response to choice: {-1: 0, 1: 1}
Context fields: ['feedback']
Learning backend: jax
Gradient support: available
HSSM participant contract: yes

2. Inspect the model anatomy¶

2AB_RW_Angle is a two-armed Bernoulli bandit. The Rescorla-Wagner learner computes trial-wise drift v; the angle decision process supplies the RT + choice likelihood family.

In [3]:

model_summary = pd.DataFrame(
    {
        "component": [
            "model name",
            "learning process",
            "decision process",
            "response columns",
            "response labels",
            "context fields",
            "required parameters",
        ],
        "value": [
            config.model_name,
            type(config.learning_process).__name__,
            config.decision_process,
            config.response,
            config.choices,
            config.context_fields,
            config.required_params,
        ],
    }
)
model_summary

model_summary = pd.DataFrame( { "component": [ "model name", "learning process", "decision process", "response columns", "response labels", "context fields", "required parameters", ], "value": [ config.model_name, type(config.learning_process).__name__, config.decision_process, config.response, config.choices, config.context_fields, config.required_params, ], } ) model_summary

Out[3]:

	component	value
0	model name	2AB_RW_Angle
1	learning process	RescorlaWagnerDrift
2	decision process	angle
3	response columns	[rt, response]
4	response labels	(-1, 1)
5	context fields	[feedback]
6	required parameters	[rl_alpha, scaler, a, z, t, theta]

3. Simulate synthetic participants¶

theta contains both learning parameters (rl_alpha, scaler) and fixed decision-process parameters (a, z, t, theta). The learning process computes drift v on every trial from the current Q-values.

In [4]:

theta = {
    "rl_alpha": 0.20,
    "scaler": 2.0,
    "a": 1.5,
    "z": 0.5,
    "t": 0.3,
    "theta": 0.2,
}

sim = rl.Simulator(config)
data = sim.simulate(
    theta=theta,
    n_trials=80,
    n_participants=4,
    random_state=SEED,
)

print(data.shape)
data.head()

theta = { "rl_alpha": 0.20, "scaler": 2.0, "a": 1.5, "z": 0.5, "t": 0.3, "theta": 0.2, } sim = rl.Simulator(config) data = sim.simulate( theta=theta, n_trials=80, n_participants=4, random_state=SEED, ) print(data.shape) data.head()

(320, 5)

Out[4]:

	participant_id	trial_id	rt	response	feedback
0	0	0	1.314903	-1	1.0
1	0	1	1.943819	1	0.0
2	0	2	1.070650	-1	1.0
3	0	3	3.246168	1	1.0
4	0	4	1.358311	1	1.0

In [5]:

summary = (
    data.groupby("participant_id")
    .agg(
        n_trials=("trial_id", "count"),
        mean_rt=("rt", "mean"),
        p_upper_response=("response", lambda x: (x == 1).mean()),
        mean_feedback=("feedback", "mean"),
    )
    .round(3)
)
summary

summary = ( data.groupby("participant_id") .agg( n_trials=("trial_id", "count"), mean_rt=("rt", "mean"), p_upper_response=("response", lambda x: (x == 1).mean()), mean_feedback=("feedback", "mean"), ) .round(3) ) summary

Out[5]:

	n_trials	mean_rt	p_upper_response	mean_feedback
participant_id
0	80	1.463	0.175	0.600
1	80	1.469	0.212	0.612
2	80	1.586	0.262	0.500
3	80	1.773	0.250	0.625

4. Validate before HSSM handoff¶

The validator checks the columns and response labels required by the ssms/HSSM bridge: participant IDs, configured response columns, context fields, balanced trial panels, and valid response-to-choice mappings.

In [6]:

report = config.validate_data(data)
report.print()
report.raise_for_errors()

report = config.validate_data(data) report.print() report.raise_for_errors()

RLSSM data validation report:
  [WARNING] extra_columns: data has extra columns not required by this model: ['trial_id'].
  panel shape: participants=4, trials_per_participant=80

5. HSSM consumes the ssms model contract¶

HSSM should not duplicate the learning rule. The intended handoff is to pass the ssms model into HSSM's bridge factory:

import hssm

hssm_config = hssm.rl.RLSSMConfig.from_ssms_model(config)  # or "2AB_RW_Angle"
model = hssm.RLSSM(data=data, model_config=hssm_config)
idata = model.sample()

The bridge assembles the ssms learning rule into participant-wise computed parameters used by the HSSM likelihood. This is how the same ssms implementation supports synthetic data generation and RLSSM likelihoods for inference.

6. Posterior predictive simulation¶

mode="ppc" conditions the learning state on an observed trial history, then resimulates responses and RTs. In an inference workflow, the theta values here would come from posterior draws rather than the known generative values used in this small example.

In [7]:

ppc = sim.simulate(
    theta=theta,
    mode="ppc",
    observed_data=data,
    random_state=SEED + 1,
)

comparison = pd.DataFrame(
    {
        "participant_id": data["participant_id"].head(10),
        "trial_id": data["trial_id"].head(10),
        "observed_response": data["response"].head(10),
        "ppc_response": ppc["response"].head(10),
        "observed_rt": data["rt"].head(10).round(3),
        "ppc_rt": ppc["rt"].head(10).round(3),
        "feedback_replayed": ppc["feedback"].head(10),
    }
)
comparison

ppc = sim.simulate( theta=theta, mode="ppc", observed_data=data, random_state=SEED + 1, ) comparison = pd.DataFrame( { "participant_id": data["participant_id"].head(10), "trial_id": data["trial_id"].head(10), "observed_response": data["response"].head(10), "ppc_response": ppc["response"].head(10), "observed_rt": data["rt"].head(10).round(3), "ppc_rt": ppc["rt"].head(10).round(3), "feedback_replayed": ppc["feedback"].head(10), } ) comparison

Out[7]:

	participant_id	trial_id	observed_response	ppc_response	observed_rt	ppc_rt	feedback_replayed
0	0	0	-1	-1	1.315	3.367	1.0
1	0	1	1	1	1.944	0.752	0.0
2	0	2	-1	1	1.071	1.398	1.0
3	0	3	1	-1	3.246	4.523	1.0
4	0	4	1	-1	1.358	2.670	1.0
5	0	5	-1	-1	1.196	0.787	0.0
6	0	6	1	1	1.471	0.957	0.0
7	0	7	-1	1	1.114	0.904	1.0
8	0	8	1	1	2.407	1.599	0.0
9	0	9	1	-1	1.619	1.208	0.0

Exercise¶

Change only rl_alpha and compare the simulated learning curves. Higher values adapt faster to recent feedback; lower values preserve prior Q-values longer.

In [8]:

fast_theta = {**theta, "rl_alpha": 0.55}
fast_data = sim.simulate(
    theta=fast_theta,
    n_trials=80,
    n_participants=4,
    random_state=SEED,
)

pd.DataFrame(
    {
        "baseline_feedback": data.groupby("trial_id")["feedback"].mean(),
        "fast_learning_feedback": fast_data.groupby("trial_id")["feedback"].mean(),
    }
).rolling(10, min_periods=1).mean().tail()

fast_theta = {**theta, "rl_alpha": 0.55} fast_data = sim.simulate( theta=fast_theta, n_trials=80, n_participants=4, random_state=SEED, ) pd.DataFrame( { "baseline_feedback": data.groupby("trial_id")["feedback"].mean(), "fast_learning_feedback": fast_data.groupby("trial_id")["feedback"].mean(), } ).rolling(10, min_periods=1).mean().tail()

Out[8]:

	baseline_feedback	fast_learning_feedback
trial_id
75	0.525	0.525
76	0.500	0.500
77	0.475	0.450
78	0.425	0.400
79	0.425	0.400

Summary¶

• An RLSSM is a learning process plus a task environment plus a decision process.
• ssms simulates the trial-wise learning loop and validates the data contract.
• HSSM consumes the ssms model contract for inference through RLSSMConfig.from_ssms_model(...).
• Posterior predictive checks reuse the same learning logic while conditioning on observed histories.