v0.4.0

[krishnbera]

HSSM v0.4.0: Reinforcement Learning (choice-only and joint decision) models, modern Bayesian tooling (PyMC6), and broader model support

We're excited to announce HSSM 0.4.0, a major release that makes HSSM a more capable home for Bayesian computational modeling, from hierarchical response-time models to reinforcement-learning,
attention, multi-choice, and choice-only workflows.

This release adds a first-class RLSSM interface, introduces attentional DDM support, expands choice-only and multi-choice modeling, modernizes HSSM for the PyMC 6 ecosystem, and brings a substantial
refresh of tutorials and documentation. Whether you are modeling learning across trials, fitting gaze-conditioned decisions, comparing multi-alternative choice models, or building custom likelihood
workflows, HSSM 0.4.0 offers a stronger and more extensible foundation.

Install or upgrade directly from PyPI:

pip install -U hssm
# or
uv add hssm

HSSM 0.4.0 supports Python 3.12 through 3.14.

What's new

First-class reinforcement-learning SSMs

The centerpiece of HSSM 0.4.0 is a public RLSSM (https://lnccbrown.github.io/HSSM/api/rl/) interface for fitting reinforcement-learning sequential-sampling models. It connects a trial-by-trial learning
process, such as value updating, with an evidence-accumulation decision process inside HSSM's familiar Bayesian modeling workflow.

The new RLSSM tooling includes:

• Flexible configuration of learning and decision processes
• A registry-based approach for built-in and custom RLSSM specifications
• Support for computed, trial-wise SSM parameters driven by learning dynamics
• Validation and response-label handling designed for ssms.rl-backed models
• Public utilities for extending and registering custom models
• New tutorials covering basic, advanced, restless-learner, and custom-registration workflows

To get started, see the RLSSM basic tutorial (https://lnccbrown.github.io/HSSM/tutorials/rlssm_basic/), then explore custom models with ssms.rl
(https://lnccbrown.github.io/HSSM/tutorials/rlssm_advanced/), restless-learner models (https://lnccbrown.github.io/HSSM/tutorials/rlssm_restless_learner/), and custom model registration
(https://lnccbrown.github.io/HSSM/tutorials/rlssm_hssm_custom_models/).

Choice-only reinforcement-learning models

Many learning experiments record choices without response times. HSSM now supports choice-only RL models, allowing users to model learning-driven choice probabilities directly when RT data are
unavailable or outside the scope of the research question.

This broadens the RLSSM workflow to tasks where choices are the primary outcome, including multi-option learning problems and softmax-style decision rules.

See the choice-only RLSSM tutorial (https://lnccbrown.github.io/HSSM/tutorials/choice_only_rlssm/) and the underlying choice-only models guide
(https://lnccbrown.github.io/HSSM/tutorials/choice_only_models/).

Attentional Drift Diffusion Model

HSSM 0.4.0 adds support for the attentional Drift Diffusion Model (aDDM) through hssm.aDDM and aDDMConfig.

The new aDDM workflow is designed for researchers working with gaze, attention, and value-based choice. It supports:

• Per-trial fixation covariates
• Trial-wise regression and hierarchical priors on core aDDM parameters
• Sampled non-decision time
• Posterior predictive checks conditioned on observed fixations
• A differentiable JAX first-passage-time likelihood
• A dedicated tutorial and parameter-recovery workflow

This makes it easier to connect fixation data to Bayesian decision-modeling workflows without leaving HSSM.

Analytical four-choice LBA model

HSSM 0.4.0 adds analytical support for the four-choice Linear Ballistic Accumulator (LBA4) model.

The LBA4 likelihood has been ported to JAX, enabling JAX/NumPyro-compatible sampling and gradients for the analytical four-choice LBA model. This extends HSSM's support for multi-alternative decisions
and gives users a practical analytical option for four-way choice-and-response-time data.

Per-parameter centered and non-centered parameterization

Hierarchical models often benefit from different parameterizations for different parameters. HSSM now supports per-parameter centered and non-centered parameterization.

For example:

hssm.HSSM(..., noncentered={"v": False, "a": True})

Group-term priors can also define their own noncentered setting. This gives users more control over sampler geometry and can improve model stability in hierarchical workflows.

A modern PyMC 6 foundation

HSSM now aligns with the current Bayesian Python ecosystem:

• PyMC 6+
• Bambi 0.19+
• PyTensor 3+
• ArviZ 1+

This upgrade brings HSSM forward with the tools users increasingly rely on while retaining its focus on accessible hierarchical SSM modeling.

The PyMC 6 migration touches sampling, variational inference, posterior predictive workflows, plotting, tutorials, and internal likelihood infrastructure.

Expanded neural and custom likelihood workflows

For users developing custom models and likelihoods, HSSM continues to support JAX and ONNX-based workflows, with improved integrations across the ecosystem.

This release includes updates around:

• JAX-backed likelihoods
• ONNX likelihood loading
• BayesFlow workflows
• sbi neural-ratio-estimation integration
• safer shape and dtype handling for exported neural likelihoods

See the custom likelihood workflow (https://lnccbrown.github.io/HSSM/tutorials/blackbox_contribution_onnx_example/) and sbi NRE integration tutorial
(https://lnccbrown.github.io/HSSM/tutorials/sbi_nre_integration/).

Expanded tutorials and API documentation

This release adds and refreshes learning material throughout the documentation, including:

• The main HSSM tutorial (https://lnccbrown.github.io/HSSM/tutorials/main_tutorial/)
• RLSSM tutorials and API reference (https://lnccbrown.github.io/HSSM/api/rl/)
• Choice-only RLSSM tutorial (https://lnccbrown.github.io/HSSM/tutorials/choice_only_rlssm/)
• PyMC-level HSSM usage (https://lnccbrown.github.io/HSSM/tutorials/pymc/)
• Custom and ONNX likelihood workflows (https://lnccbrown.github.io/HSSM/tutorials/blackbox_contribution_onnx_example/)
• sbi neural-ratio-estimation integration (https://lnccbrown.github.io/HSSM/tutorials/sbi_nre_integration/)

The goal is to make HSSM easier to approach for new users while giving advanced users clearer paths for custom model development.

Important upgrade notes

HSSM 0.4.0 includes breaking changes. Please review these before upgrading:

• HSSM now requires Python 3.12 or later.
• model.sample() now returns an xarray.DataTree instead of an arviz.InferenceData.
• The default sampling compute backend is now numba.
• model.summary() and model.plot_trace() have been removed. Use arviz.summary() and arviz.plot_trace_dist() instead.
• HSSM is now distributed through PyPI with pip or uv; Conda support for the latest release has ended.
• Dependency floors have been raised, including bambi>=0.19.0 and ssm-simulators>=0.13.1.

See the full changelog (https://lnccbrown.github.io/HSSM/changelog/) and installation guide (https://lnccbrown.github.io/HSSM/getting_started/installation/) for details.

Reliability and compatibility improvements

HSSM 0.4.0 also includes a broad set of fixes and infrastructure improvements:

• Restored controlled initialization jitter for JAX NUTS samplers, preventing unintended additional jitter from NumPyro and BlackJAX
• Fixed Python 3.14 posterior and prior predictive sampling compatibility
• Improved model.vi(..., backend="jax") support for LAN likelihood gradients
• Fixed gradients for choice-only models with all-scalar parameters and outlier modeling
• Cleaned up LAN likelihood operators for long-running workflows
• Migrated internal PyTensor gradient hooks to the newer pullback API
• Improved typing, CI, tests, notebook compatibility, and documentation
• Fixed documentation builds for the marimo aDDM tutorial

Thank you

Thank you to the contributors, reviewers, and users who helped shape this release.

We're looking forward to seeing what you build with HSSM 0.4.0.

---

This discussion was created from the release v0.4.0.