hssm.rl¶

The hssm.rl module provides reinforcement-learning extensions for HSSM, integrating reinforcement-learning update rules with sequential-sampling decision models (SSMs).

RLSSM¶

Use the hssm.rl.RLSSM class to construct a reinforcement-learning sequential sampling model from a named model string.

hssm.rl.RLSSM ¶

RLSSM(
    data: DataFrame,
    model: str | None = DEFAULT_RLSSM_MODEL,
    choices: list[int] | None = None,
    include: list[dict[str, Any] | Any] | None = None,
    model_config: RLSSMConfig | None = None,
    learning_process: dict[str, Any] | None = None,
    decision_process: str | None = None,
    participant_col: str = "participant_id",
    p_outlier: float | dict | Prior | None = 0.05,
    lapse: float | dict | Prior | None = None,
    link_settings: Literal["log_logit"] | None = None,
    prior_settings: Literal["safe"] | None = "safe",
    extra_namespace: dict[str, Any] | None = None,
    process_initvals: bool = True,
    initval_jitter: float = INITVAL_JITTER_SETTINGS["jitter_epsilon"],
    **kwargs: Any,
)

Bases: _RLSSM

Fit reinforcement-learning sequential sampling models from trial data.

RLSSM combines a reinforcement-learning process with a sequential-sampling decision model in a single likelihood. In the common case, you choose a named RLSSM model with model and optionally override its learning_process, decision_process, or choices settings. Use RLSSM.list_models to inspect the named models available in HSSM.

If you already have a fully built RLSSMConfig, you can pass it as model_config instead of selecting a named model.

RLSSM currently requires balanced panel data and does not support missing_data, deadline, or loglik_missing_data handling.

Parameters:

data (DataFrame) –

Trial-level data (balanced panel required).
model (str | None, default: DEFAULT_RLSSM_MODEL ) –

Name of an ssms.rl preset or custom registered RLSSM model. Defaults to "2AB_RW_DDM".
choices (list[int] | None, default: None ) –

Override the choice values in the registry. None uses the registry default.
include (list | None, default: None ) –

Parameter specifications forwarded to hssm.base.HSSMBase.
model_config (RLSSMConfig | None, default: None ) –

Fully built config. When provided, model, learning_process, decision_process, and choices are ignored (a warning is emitted if they are non-default).
learning_process (dict | None, default: None ) –

Override the learning-process dict in the registry. None uses the registry default.
decision_process (str | None, default: None ) –

Override the SSM name in the registry. None uses the registry default.
participant_col (str, default: 'participant_id' ) –

Column identifying participants. Defaults to "participant_id".
p_outlier (float | dict | Prior | None, default: 0.05 ) –

Lapse probability. Defaults to 0.05.
lapse (dict | Prior | None, default: None ) –

Lapse distribution. Defaults to None.
link_settings (Literal['log_logit'] | None, default: None ) –

Link-function preset. Defaults to None.
prior_settings (Literal['safe'] | None, default: 'safe' ) –

Prior preset. Defaults to "safe".
extra_namespace (dict | None, default: None ) –

Extra variables for formula evaluation. Defaults to None.
process_initvals (bool, default: True ) –

Whether to post-process initial values. Defaults to True.
initval_jitter (float, default: INITVAL_JITTER_SETTINGS['jitter_epsilon'] ) –

Jitter magnitude for initial values.
**kwargs (Any, default: {} ) –

Additional keyword arguments forwarded to bmb.Model.

Methods:

list_models –

All registered RLSSM models and their descriptions.

hssm.rl.RLSSM.list_models ¶

list_models() -> dict[str, str | None]

All registered RLSSM models and their descriptions.

This is the recommended entry point for newcomers to discover which models are available without constructing a full model instance.

Returns:

dict[str, str | None] –

Mapping of model name → description (None if not provided).

Examples:

>>> from hssm.rl import RLSSM
>>> RLSSM.list_models
{'2AB_RW_DDM': 'Two-armed bandit with ...', ...}

RLSSMConfig¶

Configuration object for RLSSM models.

hssm.rl.RLSSMConfig `dataclass` ¶

RLSSMConfig(
    model_name: str,
    description: str | None = None,
    response: list[str] | None = DEFAULT_SSM_OBSERVED_DATA.copy(),
    choices: tuple[int, ...] | None = DEFAULT_SSM_CHOICES,
    list_params: list[str] | None = None,
    bounds: dict[str, tuple[float, float]] = dict(),
    loglik: LogLik | None = None,
    loglik_kind: LoglikKind | None = None,
    backend: Literal["jax", "pytensor"] | None = None,
    extra_fields: list[str] | None = None,
    rv: Any | None = None,
    *,
    decision_process_loglik_kind: str,
    learning_process_kind: str,
    params_default: list[float],
    decision_process: str | "ModelConfig",
    learning_process: dict[str, Any],
    ssm_logp_func: Any = None,
)

Bases: BaseModelConfig

Config for reinforcement learning + sequential sampling models.

Extends BaseModelConfig with the fields required by the RLSSM likelihood pipeline. The key extra fields are:

ssm_logp_func: the annotated JAX SSM log-likelihood function (see below) whose computed dict drives per-parameter RL computations.
learning_process: a mapping that declares how each computed parameter is specified (see below).
decision_process: the name (string) or ModelConfig instance that identifies the SSM decision process (e.g. "ddm", "angle").
decision_process_loglik_kind / learning_process_kind: string tags that record which kind of likelihood and which kind of learning rule are used (e.g. "approx_differentiable" / "blackbox").

ssm_logp_func: A JAX function decorated with @annotate_function. It must carry:

- ``.inputs`` — ordered list of all parameter names the function
  expects (e.g. ``["v", "a", "z", "t", "theta", "rt", "response"]``).
- ``.outputs`` — list of output names (e.g. ``["logp"]``).
- ``.computed`` — dict mapping each *computed* parameter name to the
  annotated function that produces it.  For example::

      {"v": compute_v_annotated}

  where ``compute_v_annotated`` is itself decorated with
  ``@annotate_function`` and carries ``.inputs`` / ``.outputs``.

``make_rl_logp_op`` inspects ``ssm_logp_func.computed`` to resolve
which parameters come from data / sampled posteriors and which must
be computed by the RL learning rule at each gradient step.

learning_process: A dict keyed by the name of each computed parameter (matching the keys in ssm_logp_func.computed). Values record how that parameter is specified. The dict is intentionally permissive — current supported value forms are:

- **callable** — an annotated function (or plain function) used to
  compute the parameter.  The actual computation at runtime is driven
  by ``ssm_logp_func.computed``; this entry serves as declarative
  documentation and for config serialisation / round-trip::

      learning_process = {"v": compute_v_annotated}

- **string** — a symbolic identifier for declarative / YAML-based
  configs that can be resolved to a callable by the caller::

      learning_process = {"v": "subject_wise_function"}

An empty dict ``{}`` is valid when the SSM has no computed parameters
(i.e. ``ssm_logp_func.computed`` is also empty).

.. note::
    The dict is *not* directly consumed by ``make_rl_logp_op``.
    The actual compute functions used at runtime come from
    ``ssm_logp_func.computed``.  ``learning_process`` therefore acts
    as a config-level record of intent and is useful for inspection,
    serialisation, and future higher-level tooling.

Methods:

from_ssms_model –

Build an HSSM RLSSM config from a canonical ssms.rl model.

hssm.rl.RLSSMConfig.from_ssms_model `classmethod` ¶

from_ssms_model(
    model: str | Any,
    *,
    backend: Literal["auto", "jax"] = "jax",
    decision_process_loglik_kind: str = "approx_differentiable",
) -> "RLSSMConfig"

Build an HSSM RLSSM config from a canonical ssms.rl model.

Thin bridge between ssm-simulators' ssms.rl package — which owns the RLSSM model registry and the learning kernel — and HSSM, which owns the decision-process SSM log-likelihood and inference. The learning recursion (and the response_to_choice mapping) is taken from the ssms assembled participant function; the decision-process SSM logp is built HSSM-side from hssm.modelconfig + ONNX via :mod:hssm.rl.registry.

Parameters:

model (str | Any) –

A registered ssms.rl preset name (e.g. "2AB_RW_Angle") or an ssms.rl.ModelConfig instance.
backend (Literal['auto', 'jax'], default: 'jax' ) –

Learning backend requested from ssms; must yield gradient support. Defaults to "jax".
decision_process_loglik_kind (str, default: 'approx_differentiable' ) –

Loglik kind for the HSSM decision process. Defaults to "approx_differentiable".

Returns:

RLSSMConfig –

An HSSM-ready config, usable as hssm.RLSSM(data, model_config=config).

Registry functions¶

Helpers for discovering, building, and registering named RLSSM models and custom SSM base log-likelihood functions.

hssm.rl.get_rlssm_model_config ¶

get_rlssm_model_config(
    model: str = DEFAULT_RLSSM_MODEL,
    choices: list[int] | None = None,
    learning_process: dict[str, Any] | None = None,
    decision_process: str | None = None,
) -> RLSSMConfig

Build an RLSSMConfig from a named model.

Parameters:

model (str, default: DEFAULT_RLSSM_MODEL ) –

Name of an ssms.rl preset (e.g. "2AB_RW_DDM") or a custom HSSM-side model registered with register_rlssm_model.
choices (list[int] | None, default: None ) –

Override the response choice values stored in the registry.
learning_process (dict[str, Any] | None, default: None ) –

Override the learning process dict stored in the registry.
decision_process (str | None, default: None ) –

Override the SSM name stored in the registry.

Returns:

RLSSMConfig –

Fully populated configuration ready to be passed to _RLSSM.

Raises:

ValueError –

If model or the resolved decision_process is not registered.

hssm.rl.list_models ¶

list_models() -> dict[str, str | None]

Return names and descriptions of available RLSSM models.

This is the recommended starting point for new users who want to discover which models are available. Built-in public presets are discovered from ssms.rl.preset at call time; HSSM-side custom registrations are merged in afterwards and take precedence on name conflicts.

Returns:

dict[str, str | None] –

Mapping of model name → description string (or None if no description was provided at registration time).

Examples:

>>> import hssm
>>> hssm.rl.list_models()
{'2AB_RW_DDM': 'Two-armed bandit with ...', ...}

hssm.rl.register_rlssm_model ¶

register_rlssm_model(
    name: str,
    decision_process: str,
    learning_process: dict[str, Any],
    learning_process_params: list[str],
    learning_process_bounds: dict[str, tuple[float, float]],
    learning_process_params_default: list[float],
    extra_fields: list[str] | None = None,
    choices: list[int] | None = None,
    description: str | None = None,
    decision_process_loglik_kind: str = "approx_differentiable",
    learning_process_kind: str = "blackbox",
) -> None

Register a named RLSSM model in the global registry.

Parameters:

name (str) –

Registry key (e.g. "my_rldm").
decision_process (str) –

Name of the SSM to use. This may be either a custom SSM already registered in the SSM registry via register_ssm, or a built-in HSSM modelconfig SSM name such as "ddm", "angle", or "weibull".
learning_process (dict[str, Any]) –

Dict mapping computed parameter name → annotated learning function.
learning_process_params (list[str]) –

Ordered list of sampled RL parameter names.
learning_process_bounds (dict[str, tuple[float, float]]) –

Parameter bounds for the RL parameters.
learning_process_params_default (list[float]) –

Default values aligned with learning_process_params.
extra_fields (list[str] | None, default: None ) –

Data column names required by the learning process (e.g. ["feedback"]).
choices (list[int] | None, default: None ) –

Response choice values. Defaults to [0, 1].
description (str | None, default: None ) –

Optional human-readable description.
decision_process_loglik_kind (str, default: 'approx_differentiable' ) –

Loglik kind tag. Defaults to "approx_differentiable".
learning_process_kind (str, default: 'blackbox' ) –

Learning process kind tag. Defaults to "blackbox".

hssm.rl.register_ssm ¶

register_ssm(
    name: str,
    ssm_base_logp_func: Any,
    list_params_ssm: list[str],
    bounds_ssm: dict[str, tuple[float, float]],
    params_default_ssm: list[float],
    response: list[str] | None = None,
) -> None

Register an SSM base log-likelihood function in the SSM registry.

Parameters:

name (str) –

Registry key (e.g. "ddm").
ssm_base_logp_func (Any) –

An annotated JAX function (created with @annotate_function) that computes the SSM log-likelihood from a parameter matrix. Must carry .inputs and .outputs attributes but should not have a .computed key — that is injected by the factory at config-build time.
list_params_ssm (list[str]) –

Ordered list of all SSM parameter names (including any that will be computed by the learning process).
bounds_ssm (dict[str, tuple[float, float]]) –

Bounds for the non-computed SSM parameters.
params_default_ssm (list[float]) –

Default values aligned with list_params_ssm.
response (list[str] | None, default: None ) –

Data column names. Defaults to ["rt", "response"].

Utilities¶

hssm.rl.validate_balanced_panel ¶

validate_balanced_panel(
    data: DataFrame, participant_col: str = "participant_id"
) -> tuple[int, int]

Validate that data forms a balanced panel and return its shape.

A balanced panel requires every participant to have exactly the same number of trials (rows in data).

Parameters:

data (DataFrame) –

The DataFrame to validate.
participant_col (str, default: 'participant_id' ) –

Name of the column that identifies participants. Defaults to "participant_id".

Returns:

tuple[int, int] –

(n_participants, n_trials) where n_trials is the number of rows per participant.

Raises:

ValueError –

If participant_col is not present in data, or if the panel is unbalanced (participants have different trial counts).

hssm.rl¶

RLSSM¶

hssm.rl.RLSSM ¶

hssm.rl.RLSSM.list_models ¶

RLSSMConfig¶

hssm.rl.RLSSMConfig dataclass ¶

hssm.rl.RLSSMConfig.from_ssms_model classmethod ¶

Registry functions¶

hssm.rl.get_rlssm_model_config ¶

hssm.rl.list_models ¶

hssm.rl.register_rlssm_model ¶

hssm.rl.register_ssm ¶

Utilities¶

hssm.rl.validate_balanced_panel ¶

hssm.rl.RLSSMConfig `dataclass` ¶

hssm.rl.RLSSMConfig.from_ssms_model `classmethod` ¶