Skip to content

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:

  1. synthetic data generation in ssms.rl.Simulator;
  2. RLSSM likelihood construction in HSSM through hssm.rl.RLSSMConfig.from_ssms_model(...);
  3. 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 the LearningProcess protocol
  • task_environment — bandit or other task environment (or TaskConfig shorthand)
  • response_to_choice — map SSM response labels to zero-based learning choices
  • learning_backend / gradient — backend policy for simulation and HSSM export
  • context_fields — observable per-trial context columns such as "feedback", "condition", "block", or "stimulus_id"
  • include_choice — optionally include the derived zero-based choice column 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:

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

Required columns are derived from the model config:

  • participant_id
  • every name in config.response (default rt, response)
  • every name in config.context_fields (default feedback for 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:

config = rl.ModelConfig(
    ...,
    learning_process=choice_only_learning,
    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_kind
  • participant_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 LearningProcess protocol.

  • task_environment (TaskEnvironment | TaskConfig) –

    Task environment instance or a TaskConfig to auto-build one. If TaskConfig, 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 by task_environment.response_labels order.

  • 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 choice column. 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__

__post_init__()

Auto-build task environment and derive missing fields.

ssms.rl.config.ModelConfig.assemble

assemble(backend='auto')

Return a validated executable assembled model.

ssms.rl.config.ModelConfig.participant_contract

participant_contract(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the derived participant input layout for this config.

ssms.rl.config.ModelConfig.required_params property

required_params

Parameters that simulation requires from theta.

ssms.rl.config.ModelConfig.resolved_response_to_choice property

resolved_response_to_choice

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

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()

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

ssms.rl.config.ModelConfig.validate_data

validate_data(data)

Validate trial-level data against this model's RLSSM contract.

Returns a report with readable print() output and raise_for_errors() for fail-fast usage.

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

from_config(config, backend=AUTO)

Build an assembled model from a structural model config.

ssms.rl.assembled.AssembledModel.get_participant_input_fields

get_participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the default participant input columns derived from the config.

ssms.rl.assembled.AssembledModel.participant_input_fields

participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Backward-compatible alias for :meth:get_participant_input_fields.

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.

Resolve a preset name or validate an existing RLSSM model config.

Get a named RLSSM preset config. Returns a fresh instance.

List available RLSSM preset names.

Return readable metadata for a named RLSSM preset.

Register a named RLSSM preset.

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:

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

available_backends

Learning backends implemented by this process.

ssms.rl.learning.LearningProcess.compute_python

compute_python(state, params, context)

Compute SSM parameters from explicit Python/NumPy state.

ssms.rl.learning.LearningProcess.compute_ssm_params

compute_ssm_params(trial_params)

Compute SSM parameters from current learning state.

ssms.rl.learning.LearningProcess.computed_params property

computed_params

SSM parameter names this process computes (e.g., ['v']).

ssms.rl.learning.LearningProcess.default_params property

default_params

Default values for each free param.

ssms.rl.learning.LearningProcess.free_params property

free_params

RL parameter names this process requires from theta.

ssms.rl.learning.LearningProcess.init_state

init_state()

Return an explicit initial learning state for one participant.

ssms.rl.learning.LearningProcess.param_bounds property

param_bounds

Bounds for each free param.

ssms.rl.learning.LearningProcess.required_context_fields property

required_context_fields

Context keys this process needs for compute/update.

ssms.rl.learning.LearningProcess.reset

reset(**kwargs)

Reset internal state for a new participant.

ssms.rl.learning.LearningProcess.supports_gradient property

supports_gradient

Whether the differentiable backend supports gradient-based inference.

ssms.rl.learning.LearningProcess.update

update(action, reward, trial_params)

Update learning state given the choice outcome.

ssms.rl.learning.LearningProcess.update_python

update_python(state, params, context)

Return the next explicit Python/NumPy state.

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_ssm_params(trial_params)

Compute pre-update SSM parameters from current learning state.

ssms.rl.learning.RescorlaWagnerDeltaRule.q_values property

q_values

Current Q-values. None if reset() has not been called.

ssms.rl.learning.RescorlaWagnerDeltaRule.update

update(action, reward, trial_params)

Update Q[action] from the observed outcome.

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 learner emitting pre-update Q-values q0..qN.

Bases: RescorlaWagnerDeltaRule

Rescorla-Wagner learning core with separate learning rates.

Positive prediction errors use rl_alpha and negative prediction errors use rl_alpha_neg.

ssms.rl.learning.RescorlaWagnerDualAlphaRule.update_python

update_python(state, params, context)

Update Q[action] with sign-dependent learning rates.

Bases: RescorlaWagnerDualAlphaRule

Dual-alpha Rescorla-Wagner learner emitting two-action drift v.

Bases: RescorlaWagnerDualAlphaRule

Dual-alpha Rescorla-Wagner learner emitting Q-values q0..qN.

See also the full package reference on the ssms API page.