Cpaniaguam/rlssm simplified interface by cpaniaguam · Pull Request #955 · lnccbrown/HSSM

cpaniaguam · 2026-05-06T18:51:40Z

flowchart TD
    subgraph SSM_Registry
      A1["_SSM_REGISTRY (dict)"]
      A2["_SSM_LOGP_CACHE (dict)"]
    end
    subgraph RLSSM_Registry
      B1["_RLSSM_REGISTRY (dict)"]
    end
    subgraph User_API
      C1[register_ssm]
      C2[register_rlssm_model]
      C3[get_rlssm_model_config]
    end
    C1-->|adds entry|A1
    C1-->|adds entry|A2
    C2-->|adds entry|B1
    C3-->|reads entry|B1
    C3-->|reads entry|A1
    C3-->|calls _get_ssm_logp|A2
    A2-->|lazy build if needed|A1
    C3-->|returns RLSSMConfig|D1["RLSSMConfig"]
    style D1 fill:#1e7a1e,stroke:#fff,stroke-width:2px,color:#fff
    style A1 fill:#1e3a7a,stroke:#fff,stroke-width:2px,color:#fff
    style A2 fill:#1e3a7a,stroke:#fff,stroke-width:2px,color:#fff
    style B1 fill:#7a1e1e,stroke:#fff,stroke-width:2px,color:#fff
    style C1 fill:#7a7a1e,stroke:#fff,stroke-width:2px,color:#fff
    style C2 fill:#7a7a1e,stroke:#fff,stroke-width:2px,color:#fff
    style C3 fill:#7a7a1e,stroke:#fff,stroke-width:2px,color:#fff

Copilot

Pull request overview

This PR introduces a simplified public interface for RLSSM by adding a named-model registry and wrapping the existing RLSSM implementation behind a friendlier constructor that can build an RLSSMConfig from a model name (defaulting to "rldm").

Changes:

Split the existing RLSSM implementation into an internal base class (_RLSSM) and a public wrapper (RLSSM) with a simplified constructor.
Added an RLSSM/SSM registry and factory (get_rlssm_model_config, register_rlssm_model, register_ssm) to construct configs from named models.
Expanded the RLSSM test suite to cover the simplified interface, registry behavior, and the new wrapper semantics.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`tests/test_rlssm.py`	Adds coverage for the new simplified `RLSSM` constructor and registry-based config creation.
`src/hssm/rl/rlssm.py`	Renames the prior implementation to `_RLSSM` and adds the public `RLSSM` wrapper + blocked-attribute behavior.
`src/hssm/rl/registry.py`	New registry/factory module for named RLSSM models and SSM base logp functions.
`src/hssm/rl/__init__.py`	Exposes `_RLSSM` and registry helpers in the `hssm.rl` public API.
`src/hssm/__init__.py`	Exposes `register_rlssm_model` at the top-level `hssm` API.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…caching mechanism

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

- Added KeyError for missing SSM in the registry. - Improved ValueError messages for non-callable logp functions. - Implemented runtime checks for list_params and loglik in RLSSM. - Ensured defensive copying of mutable parameters in registry functions.

…rs for RL parameters, bounds, and defaults

…y .computed attribute

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

…consistency

…l registration process

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

AlexanderFengler

Left a few comments, mostly looks good already.
@krishnbera can you take a look too?

AlexanderFengler · 2026-05-08T00:51:43Z

+- ``RLSSMConfig``: the config class for RL + SSM models in :mod:`hssm.rl.config`.
+- ``get_rlssm_model_config``: factory that builds a config from a named model.
+- ``register_rlssm_model``: register a custom named RLSSM model.
+- ``register_ssm``: register a custom SSM base logp function.


would we want/need a register_learning_process equivalent here? ( @krishnbera )

yes we should have a way of registering the learning process separately.

…SSM model registration

…od function handling

cpaniaguam · 2026-05-08T17:59:58Z

flowchart TD
    USER["User: RLSSM(model='rldm')"]

    RLSSM_REG["_RLSSM_REGISTRY\nnamed models\ne.g. 'rldm'"]

    SSM_REG["_SSM_REGISTRY\ncustom SSMs only\n(empty by default)"]

    MODELCONFIG["hssm.modelconfig\nbuilt-in SSMs\nddm · angle · weibull · ornstein · …"]

    CACHE["_SSM_LOGP_CACHE\nlazy ONNX → JAX fn"]

    OUTPUT["RLSSMConfig"]

    USER --> RLSSM_REG
    RLSSM_REG -- decision_process name --> LOOKUP

    LOOKUP{"registered\ncustom SSM?"}
    LOOKUP -- yes --> SSM_REG
    LOOKUP -- no  --> MODELCONFIG

    SSM_REG & MODELCONFIG --> CACHE
    CACHE -- annotated JAX fn --> OUTPUT
    RLSSM_REG -- rl params / bounds / learning process --> OUTPUT

…M model config

…_models_config_structure

Copilot

Pull request overview

Copilot reviewed 8 out of 11 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 8 out of 11 changed files in this pull request and generated no new comments.

… organizing test cases

…odelconfig

…onals

…g lp directly

…ching logic

krishnbera

looks good overall. added minor comments.

Copilot

Pull request overview

Copilot reviewed 9 out of 12 changed files in this pull request and generated 1 comment.

+    Examples
+    --------
+    >>> import hssm
+    >>> hssm.rl.list_models()


cpaniaguam added 4 commits May 6, 2026 14:50

Update init files

b152169

Add registry

c3a09be

Update rlssm module

e4fb4ac

Add tests for simplified RLSSM interface and model registration

9dc18ac

cpaniaguam marked this pull request as draft May 6, 2026 18:52

cpaniaguam requested a review from Copilot May 6, 2026 18:52

Copilot started reviewing on behalf of cpaniaguam May 6, 2026 18:53 View session

Copilot AI reviewed May 6, 2026

View reviewed changes

Comment thread src/hssm/rl/registry.py Outdated

Comment thread src/hssm/rl/rlssm.py Outdated

cpaniaguam added 2 commits May 6, 2026 15:01

Fix condition in RLSSM class to handle model configuration correctly

413fe36

Implement lazy loading for SSM base log-likelihood functions and add …

19a39aa

…caching mechanism

cpaniaguam requested a review from Copilot May 6, 2026 19:07

Copilot started reviewing on behalf of cpaniaguam May 6, 2026 19:08 View session

Copilot AI reviewed May 6, 2026

View reviewed changes

Comment thread src/hssm/rl/registry.py

Comment thread src/hssm/rl/registry.py

Comment thread src/hssm/rl/registry.py Outdated

cpaniaguam added 3 commits May 6, 2026 15:54

Refactor RLSSM model configuration to respect explicit empty containe…

b23c8c2

…rs for RL parameters, bounds, and defaults

Add validation to prevent registration of SSM functions with non-empt…

78bd5c5

…y .computed attribute

cpaniaguam requested a review from Copilot May 6, 2026 20:14

Copilot started reviewing on behalf of cpaniaguam May 6, 2026 20:15 View session

Copilot AI reviewed May 6, 2026

View reviewed changes

Comment thread src/hssm/rl/registry.py Outdated

Comment thread src/hssm/rl/registry.py Outdated

cpaniaguam added 2 commits May 7, 2026 11:31

Refactor SSM registry to convert parameters to appropriate types for …

00638ca

…consistency

Update error message in get_rlssm_model_config to clarify custom mode…

101444e

…l registration process

cpaniaguam requested a review from Copilot May 7, 2026 15:50

Copilot started reviewing on behalf of cpaniaguam May 7, 2026 15:51 View session

Copilot AI reviewed May 7, 2026

View reviewed changes

Comment thread tests/rl/test_rlssm.py Outdated

Comment thread tests/rl/test_rlssm.py Outdated

Comment thread tests/rl/test_rlssm.py Outdated

Comment thread tests/rl/test_rlssm.py Outdated

Comment thread tests/rl/test_rlssm.py Outdated

Use more transparent terminology

91b2572

AlexanderFengler reviewed May 8, 2026

View reviewed changes

AlexanderFengler mentioned this pull request May 8, 2026

Cpaniaguam/add rlssm model registry #956

Closed

cpaniaguam added 2 commits May 8, 2026 11:52

Implement defensive copying for decision process specifications in RL…

c82db43

…SSM model registration

Refactor SSM registry documentation and enhance SSM base log-likeliho…

9310e88

…od function handling

cpaniaguam added 2 commits May 14, 2026 15:10

fix: ensure defensive copying of learning process and choices in RLSS…

3626f4d

…M model config

fix: update choices assertion to use tuple format in test_builtin_2ab…

f3a94d6

…_models_config_structure

cpaniaguam requested a review from Copilot May 14, 2026 19:30

Copilot started reviewing on behalf of cpaniaguam May 14, 2026 19:30 View session

Copilot AI reviewed May 14, 2026

View reviewed changes

cpaniaguam requested a review from Copilot May 15, 2026 15:16

Copilot started reviewing on behalf of cpaniaguam May 15, 2026 15:17 View session

Copilot AI reviewed May 15, 2026

View reviewed changes

cpaniaguam requested review from AlexanderFengler, digicosmos86 and krishnbera May 15, 2026 15:46

cpaniaguam marked this pull request as ready for review May 15, 2026 18:03

cpaniaguam added 8 commits May 15, 2026 14:29

chore: remove unnecessary comments

54cebfe

refactor: clean up test_rlssm.py by removing unnecessary comments and…

08ba693

… organizing test cases

Improve model_config handling in RLSSM class

1e987d5

fix: update logger initialization to use module name for better context

80d3c98

fix: simplify parameter default calculation in _build_ssm_spec_from_m…

abd86dc

…odelconfig

fix: optimize SSM logp retrieval by caching and restructuring conditi…

aac4f6f

…onals

fix: streamline parameter filtering in get_rlssm_model_config by usin…

1042c70

…g lp directly

fix: enhance RLSSM registry tests with improved error handling and ca…

5320621

…ching logic

krishnbera previously approved these changes May 21, 2026

View reviewed changes

Comment thread src/hssm/rl/registry.py

Comment thread src/hssm/rl/registry.py Outdated

Comment thread src/hssm/rl/registry.py Outdated

Comment thread src/hssm/rl/registry.py Outdated

Comment thread src/hssm/rl/registry.py Outdated

fix: rename _derive_rl_params to _derive_lp_params

bc9c88f

cpaniaguam dismissed krishnbera’s stale review via bc9c88f May 21, 2026 18:03

cpaniaguam added 4 commits May 21, 2026 14:26

Rename rl_params to learning_process_params

b47d9db

Fix line too long

4973e8a

Rename rl_bounds to learning_process_bounds

52d5fd6

Rename rl_params_default to learning_process_params_default

fe0be81

cpaniaguam requested a review from Copilot May 22, 2026 20:42

Copilot started reviewing on behalf of cpaniaguam May 22, 2026 20:42 View session

Copilot AI reviewed May 22, 2026

View reviewed changes

Comment thread src/hssm/rl/registry.py

Examples

--------

>>> import hssm

>>> hssm.rl.list_models()

Conversation

cpaniaguam commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AlexanderFengler left a comment

Choose a reason for hiding this comment

Uh oh!

AlexanderFengler May 8, 2026

Choose a reason for hiding this comment

Uh oh!

krishnbera May 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

cpaniaguam commented May 8, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

krishnbera left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

cpaniaguam commented May 6, 2026 •

edited

Loading