RLSSM simulation (ssms.rl)
The ssms.rl namespace provides a compositional framework for simulating
reinforcement-learning sequential sampling models (RLSSMs). Combine a learning
process, SSM decision process, and task environment; simulate balanced panels
compatible with HSSM inference.
An RLSSM links two time scales:
- within a trial, a decision process generates a response, and sometimes an RT;
- across trials, a learning process updates latent states such as Q-values from choices, feedback, rewards, or other task context.
The same ssms-defined learning rule can therefore serve three workflows:
- synthetic data generation in
ssms.rl.Simulator; - RLSSM likelihood construction in HSSM through
hssm.rl.RLSSMConfig.from_ssms_model(...); - posterior predictive simulation by conditioning learning on observed trial histories and resimulating responses with posterior parameter draws.
Quick start
import ssms.rl as rl
print(rl.preset.info("2AB_RW_Angle"))
config = rl.preset.get("2AB_RW_Angle")
sim = rl.Simulator(config)
data = sim.simulate(
theta={
"rl_alpha": 0.2,
"scaler": 2.0,
"a": 1.5,
"z": 0.5,
"t": 0.3,
"theta": 0.2,
},
n_trials=200,
n_participants=20,
random_state=42,
)
See the RLSSM tutorial for presets, building a model, simulating participants, validation, and plots. For focused examples, see RLSSM simulation and HSSM handoff and choice-only RL models.
Public API
| Export | Role |
|---|---|
ModelConfig |
Structural model specification (no concrete theta values) |
Simulator |
Trial-wise generative simulation loop |
AssembledModel |
Validated executable form of a config (inference-oriented) |
resolve_model |
Resolve preset name or validate a ModelConfig |
env |
Task environments (Bandit, TaskConfig, …) |
learning |
Learning processes (RescorlaWagnerDrift, RescorlaWagnerSoftmax, …) |
preset |
Preset registry (get, list, info, register) |
validate_rlssm_data |
Standalone validation helper used by ModelConfig.validate_data() |
Import style: import ssms.rl as rl.
Most users should start from rl.preset.get(...) and rl.Simulator(...). Use
ModelConfig directly when you need a custom task environment, response mapping,
or learning process. Use assemble(backend="jax") only when integrating with an
inference package or inspecting the participant-wise computed-parameter contract.
Model configuration
ModelConfig describes model structure, not parameter values. Pass concrete
values as theta to Simulator.simulate().
Important fields:
decision_process— SSM name ("angle","ddm", …)learning_process— instance satisfying theLearningProcessprotocoltask_environment— bandit or other task environment (orTaskConfigshorthand)response_to_choice— map SSM response labels to zero-based learning choiceslearning_backend/gradient— backend policy for simulation and HSSM exportcontext_fields— observable per-trial context columns such as"feedback","condition","block", or"stimulus_id"include_choice— optionally include the derived zero-basedchoicecolumn in simulator output
Derived decision-process config (_ssm_config)
ModelConfig builds an internal decision-process configuration in
__post_init__ via ModelConfigBuilder.from_model(decision_process). Users
never construct or pass this layer directly.
It supplies SSM parameter names, default bounds, default values, and choice
labels used to validate choices, derive list_params / bounds, and resolve
which SSM parameters are computed by the learning process versus fixed in
simulator theta. The assembled model and HSSM bridge consume the derived
public fields (list_params, computed_params, response_to_choice, …), not
_ssm_config itself.
Built-in Rescorla-Wagner learning processes
The built-in Rescorla-Wagner classes separate the Q-value update rule from the decision-process parameters emitted on each trial:
| Class | Emits | Actions | Use case |
|---|---|---|---|
RescorlaWagnerDeltaRule |
none | n_actions >= 2 |
Core single-alpha Q-value state/update class for custom adapters |
RescorlaWagnerDrift |
v |
n_actions == 2 |
Two-action SSMs that need trial-wise drift, such as angle |
RescorlaWagnerSoftmax |
q0, q1, ... |
n_actions >= 2 |
Choice-only inverse-temperature softmax decision processes |
RescorlaWagnerDualAlphaRule |
none | n_actions >= 2 |
Core dual-alpha Q-value state/update class |
RescorlaWagnerDualAlphaDrift |
v |
n_actions == 2 |
Two-action drift models with separate positive/negative learning rates |
RescorlaWagnerDualAlphaSoftmax |
q0, q1, ... |
n_actions >= 2 |
Choice-only softmax models with separate positive/negative learning rates |
Use RescorlaWagnerDeltaRule and RescorlaWagnerDualAlphaRule when you need
the update rule but want to write a custom decision-facing adapter. Use the
concrete drift or softmax classes directly for standard presets and HSSM handoff.
For drift models, the learner computes v = (Q[1] - Q[0]) * scaler, so scaler
is a free learning-process parameter. For softmax models, the learner emits the
raw q0..qN values and the decision process uses the fixed SSM parameter beta
as the inverse temperature.
Task environment protocols
TaskEnvironment is the base protocol for per-trial context and post-decision
signals. Models that map SSM response labels to learning choices require a
DiscreteChoiceEnvironment (adds n_choices and response_labels). Built-in
bandits implement DiscreteChoiceEnvironment; Bandit.n_arms is an alias for
n_choices.
Participant-wise parameters
Simulator.simulate() accepts scalar theta values shared by all participants
and one-dimensional participant-wise values. When any theta value is
participant-wise, all participant-wise values must have the same length. If
n_participants is omitted, that length is used as the participant count:
data = sim.simulate(
theta={
"rl_alpha": [0.15, 0.25, 0.35],
"scaler": 2.0,
"a": [1.1, 1.4, 1.7],
"z": 0.5,
"t": 0.3,
"theta": 0.2,
},
n_trials=200,
random_state=42,
)
Passing n_participants explicitly is allowed, but it must match the
participant-wise theta length.
Simulation modes
Simulator.simulate() supports two modes:
mode="generative"— the default unconstrained simulation loop. The simulator samples responses, task context, and learning updates end to end.mode="ppc"— observed-history-conditioned posterior predictive simulation. Learning state is conditioned on observed trial history; RT/response are resimulated and observed context fields are copied into output. After HSSM inference, posterior draws can be routed through the same simulator contract to check whether inferred learning and decision parameters reproduce behavior.
PPC mode uses the same data contract as inference validation (see below). The
observed panel must include participant_id, all config.response columns
(default rt and response), and every configured context field (default
feedback for the built-in bandit):
Data validation
ModelConfig.validate_data() validates external trial panels — empirical
data or simulated panels you plan to pass to PPC mode or HSSM. Generative
simulation does not self-validate its output; only mode="ppc" validates
user-supplied observed_data before conditioning on it.
Validate empirical or simulated panels before PPC or HSSM handoff:
Required columns are derived from the model config:
participant_id- every name in
config.response(defaultrt,response) - every name in
config.context_fields(defaultfeedbackfor the built-in bandit)
The validator checks balanced panels, contiguous participant blocks, response labels
compatible with config.choices and response_to_choice, missing values, and
omission sentinels. Within each participant, rows are processed in their existing
order. trial_id is an ordinary data/context column, not a reserved reset or ordering
field. Errors include repair hints, for example renaming a reward column or adding it
to ModelConfig(context_fields=[...]).
PPC mode example (observed data must satisfy the same contract):
observed = sim.simulate(
theta={
"rl_alpha": 0.2,
"scaler": 2.0,
"a": 1.5,
"z": 0.5,
"t": 0.3,
"theta": 0.2,
},
n_trials=200,
n_participants=20,
random_state=1,
)
ppc = sim.simulate(
theta={
"rl_alpha": 0.2,
"scaler": 2.0,
"a": 1.5,
"z": 0.5,
"t": 0.3,
"theta": 0.2,
},
mode="ppc",
observed_data=observed,
random_state=2,
)
The observed response history is used only to condition learning state; PPC output responses are newly simulated.
Choice-only inverse-temperature softmax presets
2AB_RW_InvTempSoftmax and 3AB_RW_InvTempSoftmax are response-only RL presets.
They use RescorlaWagnerSoftmax to emit q0..qN, and the
inv_temp_softmax_N decision process uses beta as the inverse temperature for
choice probabilities.
These presets declare response=["response"] because the softmax decision
process has no response-time likelihood. The low-level softmax simulator still
returns an rt array for compatibility with the generic simulator interface,
but every value is -1.0 and should be treated only as a non-omission
placeholder. It is not a response time, and it is intentionally distinct from
OMISSION_SENTINEL == -999.0, which ssms and HSSM use for omissions,
deadline/no-response trials, and missing RT handling.
For validation and HSSM handoff, omit the placeholder column:
config = rl.preset.get("2AB_RW_InvTempSoftmax")
data = rl.Simulator(config).simulate(theta=theta, n_trials=200)
report = config.validate_data(data.drop(columns=["rt"]))
report.raise_for_errors()
Choice-only PPC uses the same response-only contract. Empirical observed_data
must not contain an rt column for these presets, and PPC output omits rt:
response_only = data.drop(columns=["rt"])
ppc = rl.Simulator(config).simulate(
theta={"rl_alpha": 0.2, "beta": 2.0},
mode="ppc",
observed_data=response_only,
random_state=13,
)
The lower-level inv_temp_softmax_4 decision process is also available for
four-choice softmax simulation. Built-in RL presets currently cover the two- and
three-choice bandit cases; custom ModelConfig objects can pair
RescorlaWagnerSoftmax(n_actions=4) with decision_process="inv_temp_softmax_4".
Context fields
Outcome-like values are ordinary context fields. By default, built-in bandits emit a
"feedback" column and built-in Rescorla-Wagner learners require context["feedback"]
for updates. Use a custom feedback field by configuring both the learner and the model
context:
config = rl.ModelConfig(
...,
learning_process=rl.learning.RescorlaWagnerDrift(feedback_field="reward"),
context_fields=["reward"],
)
For learning processes that update from choices only, declare required_context_fields
with runtime fields such as "choice" and use no observable context fields:
Assembled model (inference integration)
Assemble a config when you need validated metadata or participant-wise computed parameter functions for downstream packages (for example HSSM):
assembled = config.assemble(backend="jax")
# Derived from config — no manual field lists for standard models
fields = assembled.get_participant_input_fields()
compute_params = assembled.assemble_participant_fn()
assemble_participant_fn() accepts optional overrides (input_fields,
response_field) for non-standard layouts. Runtime context fields such as choice
are derived internally from response_to_choice; observable context fields such as
feedback come from config.context_fields.
Advanced resolution:
config = rl.resolve_model("2AB_RW_Angle") # str or ModelConfig
assembled = config.assemble(backend="auto")
HSSM bridge
The active HSSM handoff path is HSSM's bridge factory:
import hssm
import ssms.rl as rl
ssms_config = rl.preset.get("2AB_RW_Angle")
hssm_config = hssm.rl.RLSSMConfig.from_ssms_model(ssms_config)
model = hssm.RLSSM(data=data, model_config=hssm_config)
RLSSMConfig.from_ssms_model(...) resolves the ssms.rl model, assembles it
with the JAX backend, checks gradient support, and wraps
AssembledModel.assemble_participant_fn(output="dict") for HSSM's annotated
computed-parameter contract.
This bridge is what lets HSSM evaluate RLSSM likelihoods while keeping the
learning rule, response-to-choice mapping, and task context source of truth in
ssms. For choice-only RL models, pass response-only data to HSSM, not the
generative simulator's placeholder rt column.
Note: HSSM's bridge factory still calls the pre-refactor compile() API
until the separate hssm-rlssm-api task updates it to assemble().
ModelConfig.to_hssm_config_dict() remains useful for structural inspection
and compatibility with lower-level HSSM config workflows. It exports shared
structural fields, plus:
learning_backend,gradient,learning_process_kindparticipant_contract— derived trial input layout (trial_params,response_field,context_fields,input_fields). Users never construct this directly; it is exported for bridge metadata and debugging.
Inference-only placeholders in to_hssm_config_dict() (ssm_logp_func,
learning_process) are not a complete model by themselves. A higher-level
hssm.RLSSM(data, model=...) wrapper that consumes ssms.rl directly is
planned separately in HSSM.
Module reference
RLSSM model configuration for ssm-simulators.
Describes the structural specification of an RLSSM model:
which learning process, which decision process (SSM), and which task
environment. Concrete parameter values are NOT stored here — they
are passed as theta to Simulator.simulate().
Parameters:
-
model_name(str) –Unique identifier for this RLSSM model (e.g., "rlssm_angle_rw").
-
description(str) –Human-readable model description.
-
decision_process(str) –SSM model name in ssm-simulators registry (e.g., "angle", "ddm"). Must be resolvable via
ModelConfigBuilder.from_model(). -
learning_process(LearningProcess) –Instance of a class satisfying the
LearningProcessprotocol. -
task_environment(TaskEnvironment | TaskConfig) –Task environment instance or a
TaskConfigto auto-build one. IfTaskConfig,build_environment()is called in__post_init__. -
list_params(list[str] | None, default:None) –All free parameter names (RL + fixed SSM), in order. If None, auto-derived:
learning_process.free_params+ fixed SSM params. -
bounds(dict[str, tuple[float, float]] | None, default:None) –Parameter bounds. If None, auto-derived from learning_process.param_bounds + SSM model config param_bounds.
-
params_default(list[float] | None, default:None) –Default values in same order as list_params. If None, auto-derived.
-
choices(tuple[int, ...] | None, default:None) –SSM response labels (e.g., (-1, 1)). If None, taken from task_environment.
-
response(list[str], default:(lambda: ['rt', 'response'])()) –Response column names. Default ["rt", "response"].
-
response_to_choice(Literal['auto'] | dict[int, int], default:'auto') –Mapping from SSM response labels to zero-based learning choices.
"auto"maps labels bytask_environment.response_labelsorder. -
learning_backend(Literal['auto', 'python', 'jax'], default:'auto') –Learning-process backend used for simulation and exported HSSM metadata.
"auto"selects JAX when the process implements it and JAX is installed; otherwise it selects Python. -
gradient(Literal['auto', 'available', 'unavailable'], default:'auto') –Gradient-support policy for HSSM integration metadata.
-
include_choice(bool, default:False) –Whether simulator output includes the derived zero-based
choicecolumn. Default False. -
context_fields(list[str] | None, default:None) –Data/context columns beyond response required by the environment or learning process. Default derives a union from component declarations.
-
computed_param_mapping(dict[str, str] | None, default:None) –Optional override for non-name-matching handshakes. Maps learning process output name -> SSM param name. E.g., {"drift": "v"} if learning process outputs "drift" but SSM expects "v". Default: None (same-name linking).
-
ssm_kwargs(dict, default:(lambda: {'delta_t': 0.001, 'max_t': 20.0})()) –Default kwargs for the underlying SSM simulator call. Default: {"delta_t": 0.001, "max_t": 20.0}.
ssms.rl.config.ModelConfig.__post_init__
Auto-build task environment and derive missing fields.
ssms.rl.config.ModelConfig.assemble
Return a validated executable assembled model.
ssms.rl.config.ModelConfig.participant_contract
Return the derived participant input layout for this config.
ssms.rl.config.ModelConfig.required_params
property
Parameters that simulation requires from theta.
ssms.rl.config.ModelConfig.resolved_response_to_choice
property
Concrete response-label -> choice-index map.
__post_init__ normalizes response_to_choice (including the
"auto" default) into a plain dict; this accessor exposes that
post-init invariant with a narrowed type.
ssms.rl.config.ModelConfig.to_hssm_config_dict
Produce a dict compatible with HSSM's RLSSMConfig.from_rlssm_dict().
The output contains all fields from _HSSM_SHARED_FIELDS plus placeholder values for inference-only fields that the user must fill in on the HSSM side.
Returns:
-
dict[str, Any]–Dict ready for
RLSSMConfig.from_rlssm_dict(result)after user fills in inference-only fields.
ssms.rl.config.ModelConfig.validate
Validate config consistency. Called by Simulator.init().
Checks: 1. decision_process exists in ssm-simulators registry 2. Handshake: computed + fixed params cover all SSM params exactly once 3. No param is both computed and fixed 4. list_params length matches params_default length 5. All list_params have bounds
Validated executable form of an RLSSM ModelConfig.
The assembled model exposes package-neutral metadata and pure Python/JAX
computed-parameter functions that downstream packages can wrap without
importing HSSM or PyTensor in ssm-simulators.
ssms.rl.assembled.AssembledModel.assemble_participant_fn
assemble_participant_fn(
input_fields=None,
*,
response_field=DEFAULT_RESPONSE_FIELD,
output=ARRAY
)
Assemble a participant-wise computed-parameter function.
By default, input_fields are derived from the model config. Pass
explicit values only for non-standard layouts.
The returned function accepts a (n_trials, n_fields) array whose
columns match input_fields. It computes SSM parameters before each
learning update, maps response labels to zero-based action indices, and
updates learning state from the response and optional outcome.
ssms.rl.assembled.AssembledModel.from_config
classmethod
Build an assembled model from a structural model config.
ssms.rl.assembled.AssembledModel.get_participant_input_fields
Return the default participant input columns derived from the config.
RLSSM simulator composing a learning process with an SSM decision process.
Runs the interleaved trial-by-trial loop: compute SSM params -> simulate SSM -> observe choice -> generate reward -> update learning.
Reuses the existing ssm-simulators simulator() function with n_samples=1
for each trial. No Cython modifications needed — all 40+ SSM models work as
decision processes out of the box.
Parameters:
-
config(ModelConfig) –Structural model configuration. Validated on construction.
ssms.rl.simulator.Simulator.simulate
simulate(
theta,
n_trials=200,
n_participants=None,
random_state=None,
mode="generative",
observed_data=None,
)
Run full RLSSM simulation.
Parameters:
-
theta(dict[str, Any]) –Concrete parameter values. Must contain all params required by the learning process and fixed SSM parameters. Each value can be a scalar shared by all participants or a one-dimensional list/array with one value per participant.
-
n_trials(int, default:200) –Number of trials per participant. Default 200.
-
n_participants(int | None, default:None) –Number of participants to simulate. If None, inferred from participant-wise theta values when present; otherwise defaults to 20.
-
random_state(int | None, default:None) –Seed for reproducibility. If None, non-deterministic.
-
mode(('generative', 'ppc'), default:"generative") –Simulation mode.
"generative"runs the unconstrained simulator loop."ppc"runs observed-history-conditioned posterior predictive simulation. -
observed_data(DataFrame | None, default:None) –Observed participant history required for
mode="ppc".
Returns:
-
DataFrame–Balanced panel with participant_id, trial_id, configured response columns, configured context fields, and optional derived choice.
Validate a data panel against the RLSSM model contract.
Parameters:
-
config(ModelConfig) –Structural RLSSM configuration. Should already pass
config.validate(). -
data(DataFrame) –Empirical or simulated trial-level panel.
Returns:
-
DataValidationReport–Validation findings. Call
raise_for_errors()to fail fast on errors.
Bases: Protocol
Protocol for RLSSM learning processes.
A learning process maintains internal state (e.g., Q-values) and computes SSM parameters (e.g., drift rate) from that state on each trial. After each trial's decision and reward, the state is updated.
The computed_params property is the formal handshake between the learning
process and the decision process: it declares which SSM parameters the
learning process produces. The simulator validates that these, together with
fixed SSM params provided by the user, cover all parameters required by the
decision process model.
ssms.rl.learning.LearningProcess.available_backends
property
Learning backends implemented by this process.
ssms.rl.learning.LearningProcess.compute_python
Compute SSM parameters from explicit Python/NumPy state.
ssms.rl.learning.LearningProcess.compute_ssm_params
Compute SSM parameters from current learning state.
ssms.rl.learning.LearningProcess.computed_params
property
SSM parameter names this process computes (e.g., ['v']).
ssms.rl.learning.LearningProcess.default_params
property
Default values for each free param.
ssms.rl.learning.LearningProcess.free_params
property
RL parameter names this process requires from theta.
ssms.rl.learning.LearningProcess.init_state
Return an explicit initial learning state for one participant.
ssms.rl.learning.LearningProcess.required_context_fields
property
Context keys this process needs for compute/update.
ssms.rl.learning.LearningProcess.supports_gradient
property
Whether the differentiable backend supports gradient-based inference.
ssms.rl.learning.LearningProcess.update
Update learning state given the choice outcome.
Rescorla-Wagner delta learning core.
Updates Q-values via Q[action] += alpha * (reward - Q[action]). This
class owns Q-value state and replay/update behavior but emits no SSM
parameters by itself. Use RescorlaWagnerDrift for two-action drift
models and RescorlaWagnerSoftmax for inverse-temperature softmax models.
ssms.rl.learning.RescorlaWagnerDeltaRule.compute_ssm_params
Compute pre-update SSM parameters from current learning state.
ssms.rl.learning.RescorlaWagnerDeltaRule.q_values
property
Current Q-values. None if reset() has not been called.
Bases: RescorlaWagnerDeltaRule
Rescorla-Wagner learner emitting two-action drift v.
Computes drift rate as scaled Q-value difference:
v = (Q[1] - Q[0]) * scaler.
Bases: RescorlaWagnerDeltaRule
Rescorla-Wagner learning core with separate learning rates.
Positive prediction errors use rl_alpha and negative prediction errors
use rl_alpha_neg.
See also the full package reference on the ssms API page.