From d539e4205edd699dc741bc432843530bc936bb37 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 24 Aug 2026 01:57:43 -0400 Subject: [PATCH 1/5] test: define RL observation metadata --- tests/rl/test_observation_metadata.py | 208 ++++++++++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 tests/rl/test_observation_metadata.py diff --git a/tests/rl/test_observation_metadata.py b/tests/rl/test_observation_metadata.py new file mode 100644 index 00000000..95d47de5 --- /dev/null +++ b/tests/rl/test_observation_metadata.py @@ -0,0 +1,208 @@ +"""RLSSM producer observation-metadata contracts.""" + +from dataclasses import replace + +import numpy as np +import pandas as pd +import pytest + +from ssms.basic_simulators import get_observation_metadata +import ssms.rl as rl + + +def _rt_field() -> dict[str, object]: + return { + "name": "rt", + "kind": "continuous", + "lower": 0.0, + "lower_inclusive": False, + } + + +def _response_field(values: tuple[int, ...]) -> dict[str, object]: + return {"name": "response", "kind": "categorical", "values": values} + + +def _descriptor(*fields: dict[str, object]) -> dict[str, object]: + return { + "observation_schema_version": 1, + "observation_schema": fields, + "obs_dim": len(fields), + } + + +def test_rt_response_preset_exposes_exact_metadata_through_public_accessor(): + config = rl.preset.get("2AB_RW_Angle") + assembled = config.assemble(backend="python") + expected = _descriptor(_rt_field(), _response_field((-1, 1))) + + assert get_observation_metadata(config) == expected + assert get_observation_metadata(assembled) == expected + + +@pytest.mark.parametrize( + ("preset_name", "choices"), + [ + ("2AB_RW_InvTempSoftmax", (0, 1)), + ("3AB_RW_InvTempSoftmax", (0, 1, 2)), + ("4AB_RW_InvTempSoftmax", (0, 1, 2, 3)), + ], +) +def test_response_only_presets_expose_exact_metadata(preset_name, choices): + config = rl.preset.get(preset_name) + assembled = config.assemble(backend="python") + expected = _descriptor(_response_field(choices)) + + assert get_observation_metadata(config) == expected + assert get_observation_metadata(assembled) == expected + + +def test_metadata_access_preserves_seeded_choice_only_simulation_and_input_order(): + config = rl.preset.get("2AB_RW_InvTempSoftmax") + assembled = config.assemble(backend="python") + simulator = rl.Simulator(config) + theta = {"rl_alpha": 0.2, "beta": 2.0} + + input_fields_before = assembled.get_participant_input_fields() + simulated_before = simulator.simulate( + theta=theta, + n_trials=8, + n_participants=2, + random_state=37, + ) + + metadata = get_observation_metadata(config) + get_observation_metadata(assembled) + + input_fields_after = assembled.get_participant_input_fields() + simulated_after = simulator.simulate( + theta=theta, + n_trials=8, + n_participants=2, + random_state=37, + ) + + assert ( + input_fields_before + == input_fields_after + == [ + "rl_alpha", + "response", + "feedback", + ] + ) + schema_names = tuple(field["name"] for field in metadata["observation_schema"]) + assert schema_names == ("response",) + assert config.context_fields == ["feedback"] + assert set(config.context_fields).isdisjoint(schema_names) + assert np.all(simulated_before["rt"] == -1.0) + pd.testing.assert_frame_equal(simulated_after, simulated_before) + + +def test_schema_choice_order_does_not_follow_response_to_choice_mapping(): + config = replace( + rl.preset.get("2AB_RW_Angle"), + response_to_choice={1: 0, -1: 1}, + ) + + metadata = get_observation_metadata(config) + + assert tuple(config.resolved_response_to_choice) == (1, -1) + assert config.choices == (-1, 1) + assert metadata["observation_schema"][1]["values"] == (-1, 1) + + +def test_nonstandard_response_layout_requires_an_explicit_v1_ordered_schema(): + response = ["rt", "confidence", "response"] + schema = ( + _rt_field(), + { + "name": "confidence", + "kind": "continuous", + "lower": 0.0, + "upper": 1.0, + }, + _response_field((-1, 1)), + ) + config = replace(rl.preset.get("2AB_RW_Angle"), response=response) + + with pytest.raises(ValueError, match="require an explicit observation_schema"): + get_observation_metadata(config) + + explicit = replace( + config, + observation_schema_version=1, + observation_schema=schema, + ) + assert get_observation_metadata(explicit) == _descriptor(*schema) + + unsupported_version = replace(explicit, observation_schema_version=2) + with pytest.raises(ValueError, match="supported integer version 1"): + get_observation_metadata(unsupported_version) + + +@pytest.mark.parametrize( + ("schema", "match"), + [ + ( + (_response_field((-1, 1)), _rt_field()), + "names and order must exactly match response", + ), + ( + (_rt_field(), _response_field((1, -1))), + "values equal to choices in order", + ), + ( + (_rt_field(), {"name": "response", "kind": "continuous"}), + "must be categorical", + ), + ], + ids=("field-order", "categorical-values", "categorical-kind"), +) +def test_explicit_schema_must_match_the_rl_response_contract(schema, match): + config = replace( + rl.preset.get("2AB_RW_Angle"), + observation_schema=schema, + ) + + with pytest.raises(ValueError, match=match): + get_observation_metadata(config) + + +def test_assembled_metadata_is_an_independent_fresh_snapshot(): + raw_values = [-1, 1] + config = replace( + rl.preset.get("2AB_RW_Angle"), + observation_schema=( + _rt_field(), + {"name": "response", "kind": "categorical", "values": raw_values}, + ), + ) + assembled = config.assemble(backend="python") + + raw_values[0] = 99 + first = get_observation_metadata(assembled) + first["observation_schema"][1]["values"] = (99, 100) + second = get_observation_metadata(assembled) + + assert second == _descriptor(_rt_field(), _response_field((-1, 1))) + assert first is not second + assert first["observation_schema"] is not second["observation_schema"] + assert first["observation_schema"][0] is not second["observation_schema"][0] + + +def test_metadata_access_does_not_change_or_extend_the_hssm_config_dict(): + config = rl.preset.get("2AB_RW_Angle") + before = config.to_hssm_config_dict() + + get_observation_metadata(config) + get_observation_metadata(config.assemble(backend="python")) + + after = config.to_hssm_config_dict() + assert after == before + assert config.observation_schema is None + assert { + "observation_schema_version", + "observation_schema", + "obs_dim", + }.isdisjoint(after) From 88e8a8cd7d1f6ef2f877064b36f4822540bcd83d Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 24 Aug 2026 02:01:10 -0400 Subject: [PATCH 2/5] feat: expose RL observation metadata --- ssms/rl/config.py | 78 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/ssms/rl/config.py b/ssms/rl/config.py index 89593734..8a70157b 100644 --- a/ssms/rl/config.py +++ b/ssms/rl/config.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field import importlib.util from typing import TYPE_CHECKING, Any, Literal, cast @@ -137,6 +138,11 @@ class ModelConfig: ssm_kwargs : dict Default kwargs for the underlying SSM simulator call. Default: {"delta_t": 0.001, "max_t": 20.0}. + observation_schema_version : int + Structured-observation schema version. Version 1 is currently supported. + observation_schema : tuple[Mapping[str, Any], ...] | None + Explicit ordered schema for non-standard response layouts. The existing + ``["rt", "response"]`` and ``["response"]`` layouts are derived directly. """ model_name: str @@ -165,6 +171,11 @@ class ModelConfig: default_factory=lambda: {"delta_t": 0.001, "max_t": 20.0} ) + # Additive observation-metadata contract. Keep these fields last so existing + # positional constructor calls retain their meaning. + observation_schema_version: int = 1 + observation_schema: tuple[Mapping[str, Any], ...] | None = None + def __post_init__(self): """Auto-build task environment and derive missing fields.""" # Convert TaskConfig -> TaskEnvironment @@ -584,6 +595,73 @@ def participant_contract( """Return the derived participant input layout for this config.""" return derive_participant_contract(self, response_field=response_field) + def get_observation_metadata(self) -> dict[str, Any]: + """Return a fresh schema for the configured stochastic response columns.""" + from ssms.basic_simulators.observation_metadata import ( + validate_observation_metadata, + ) + + schema = self.observation_schema + response_fields = tuple(self.response) + if DEFAULT_RESPONSE_FIELD not in response_fields: + raise ValueError( + f"response must include {DEFAULT_RESPONSE_FIELD!r} for observation " + f"metadata; got {response_fields}" + ) + if schema is None: + if response_fields == ("rt", "response"): + schema = ( + { + "name": "rt", + "kind": "continuous", + "lower": 0.0, + "lower_inclusive": False, + }, + self._response_schema_entry(), + ) + elif response_fields == ("response",): + schema = (self._response_schema_entry(),) + else: + raise ValueError( + "RL response layouts other than ['rt', 'response'] and " + "['response'] require an explicit observation_schema" + ) + + descriptor = validate_observation_metadata( + { + "observation_schema_version": self.observation_schema_version, + "observation_schema": schema, + } + ) + field_names = tuple(field["name"] for field in descriptor["observation_schema"]) + if field_names != response_fields: + raise ValueError( + "observation_schema names and order must exactly match response: " + f"expected {response_fields}, got {field_names}" + ) + + response_entry = descriptor["observation_schema"][ + response_fields.index(DEFAULT_RESPONSE_FIELD) + ] + expected_choices = tuple(self.choices or ()) + if ( + response_entry["kind"] != "categorical" + or tuple(response_entry["values"]) != expected_choices + ): + raise ValueError( + "the observation_schema 'response' field must be categorical with " + f"values equal to choices in order: {expected_choices}" + ) + return descriptor + + def _response_schema_entry(self) -> dict[str, Any]: + """Build the categorical schema entry from raw SSM response labels.""" + return { + "name": DEFAULT_RESPONSE_FIELD, + "kind": "categorical", + "values": tuple(self.choices or ()), + } + def to_hssm_config_dict(self) -> dict[str, Any]: """Produce a dict compatible with HSSM's RLSSMConfig.from_rlssm_dict(). From 89500f842b4f3a185445dc5a714c1e46aaa1eb49 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 24 Aug 2026 02:01:38 -0400 Subject: [PATCH 3/5] feat: snapshot assembled observation metadata --- ssms/rl/assembled.py | 19 ++++++++++++++++++- tests/rl/test_assembled_model.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/ssms/rl/assembled.py b/ssms/rl/assembled.py index 761f1f8e..e2f0418d 100644 --- a/ssms/rl/assembled.py +++ b/ssms/rl/assembled.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import StrEnum from typing import Any, Literal @@ -110,6 +110,15 @@ class AssembledModel: context_fields: list[str] computed_params: list[str] response_to_choice: dict[int, int] + _observation_metadata: dict[str, Any] = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + """Snapshot semantic observation metadata without widening the constructor.""" + object.__setattr__( + self, + "_observation_metadata", + self.config.get_observation_metadata(), + ) @classmethod def from_config( @@ -137,6 +146,14 @@ def from_config( response_to_choice=dict(config.resolved_response_to_choice), ) + def get_observation_metadata(self) -> dict[str, Any]: + """Return the observation schema snapshot captured during assembly.""" + from ssms.basic_simulators.observation_metadata import ( + validate_observation_metadata, + ) + + return validate_observation_metadata(self._observation_metadata) + def get_participant_input_fields( self, *, diff --git a/tests/rl/test_assembled_model.py b/tests/rl/test_assembled_model.py index e3ac0833..1d5561b1 100644 --- a/tests/rl/test_assembled_model.py +++ b/tests/rl/test_assembled_model.py @@ -273,6 +273,24 @@ def test_participant_input_fields_can_use_custom_response_field(self): assembled = _make_default_config( learning_backend="python", response=["rt", "response", "choice_response"], + observation_schema=( + { + "name": "rt", + "kind": "continuous", + "lower": 0.0, + "lower_inclusive": False, + }, + { + "name": "response", + "kind": "categorical", + "values": (-1, 1), + }, + { + "name": "choice_response", + "kind": "categorical", + "values": (-1, 1), + }, + ), ).assemble(backend="python") assert assembled.participant_input_fields(response_field="choice_response") == [ From b514af3250c6c1d99a06a8d0e64d452f02e1197a Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 24 Aug 2026 02:03:49 -0400 Subject: [PATCH 4/5] fix: keep standard RL schemas derived --- ssms/rl/config.py | 38 +++++++++++++--------- tests/rl/test_observation_metadata.py | 46 ++++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 20 deletions(-) diff --git a/ssms/rl/config.py b/ssms/rl/config.py index 8a70157b..5ab6aaf4 100644 --- a/ssms/rl/config.py +++ b/ssms/rl/config.py @@ -601,31 +601,39 @@ def get_observation_metadata(self) -> dict[str, Any]: validate_observation_metadata, ) - schema = self.observation_schema response_fields = tuple(self.response) if DEFAULT_RESPONSE_FIELD not in response_fields: raise ValueError( f"response must include {DEFAULT_RESPONSE_FIELD!r} for observation " f"metadata; got {response_fields}" ) - if schema is None: - if response_fields == ("rt", "response"): - schema = ( - { - "name": "rt", - "kind": "continuous", - "lower": 0.0, - "lower_inclusive": False, - }, - self._response_schema_entry(), - ) - elif response_fields == ("response",): - schema = (self._response_schema_entry(),) - else: + standard_layouts = {("rt", "response"), ("response",)} + if self.observation_schema is not None and response_fields in standard_layouts: + raise ValueError( + "observation_schema must be omitted for the standard RL response " + f"layout {list(response_fields)!r}; its schema is derived" + ) + + schema: tuple[Mapping[str, Any], ...] + if response_fields == ("rt", "response"): + schema = ( + { + "name": "rt", + "kind": "continuous", + "lower": 0.0, + "lower_inclusive": False, + }, + self._response_schema_entry(), + ) + elif response_fields == ("response",): + schema = (self._response_schema_entry(),) + else: + if self.observation_schema is None: raise ValueError( "RL response layouts other than ['rt', 'response'] and " "['response'] require an explicit observation_schema" ) + schema = self.observation_schema descriptor = validate_observation_metadata( { diff --git a/tests/rl/test_observation_metadata.py b/tests/rl/test_observation_metadata.py index 95d47de5..c7f4a118 100644 --- a/tests/rl/test_observation_metadata.py +++ b/tests/rl/test_observation_metadata.py @@ -145,15 +145,24 @@ def test_nonstandard_response_layout_requires_an_explicit_v1_ordered_schema(): ("schema", "match"), [ ( - (_response_field((-1, 1)), _rt_field()), + ( + _response_field((-1, 1)), + {"name": "latency", "kind": "continuous"}, + ), "names and order must exactly match response", ), ( - (_rt_field(), _response_field((1, -1))), + ( + {"name": "latency", "kind": "continuous"}, + _response_field((1, -1)), + ), "values equal to choices in order", ), ( - (_rt_field(), {"name": "response", "kind": "continuous"}), + ( + {"name": "latency", "kind": "continuous"}, + {"name": "response", "kind": "continuous"}, + ), "must be categorical", ), ], @@ -162,6 +171,7 @@ def test_nonstandard_response_layout_requires_an_explicit_v1_ordered_schema(): def test_explicit_schema_must_match_the_rl_response_contract(schema, match): config = replace( rl.preset.get("2AB_RW_Angle"), + response=["latency", "response"], observation_schema=schema, ) @@ -169,12 +179,29 @@ def test_explicit_schema_must_match_the_rl_response_contract(schema, match): get_observation_metadata(config) +def test_standard_response_layout_rejects_an_explicit_schema_override(): + config = replace( + rl.preset.get("2AB_RW_Angle"), + observation_schema=(_rt_field(), _response_field((-1, 1))), + ) + + with pytest.raises(ValueError, match="must be omitted for the standard RL"): + get_observation_metadata(config) + + def test_assembled_metadata_is_an_independent_fresh_snapshot(): raw_values = [-1, 1] config = replace( rl.preset.get("2AB_RW_Angle"), + response=["rt", "confidence", "response"], observation_schema=( _rt_field(), + { + "name": "confidence", + "kind": "continuous", + "lower": 0.0, + "upper": 1.0, + }, {"name": "response", "kind": "categorical", "values": raw_values}, ), ) @@ -182,10 +209,19 @@ def test_assembled_metadata_is_an_independent_fresh_snapshot(): raw_values[0] = 99 first = get_observation_metadata(assembled) - first["observation_schema"][1]["values"] = (99, 100) + first["observation_schema"][2]["values"] = (99, 100) second = get_observation_metadata(assembled) - assert second == _descriptor(_rt_field(), _response_field((-1, 1))) + assert second == _descriptor( + _rt_field(), + { + "name": "confidence", + "kind": "continuous", + "lower": 0.0, + "upper": 1.0, + }, + _response_field((-1, 1)), + ) assert first is not second assert first["observation_schema"] is not second["observation_schema"] assert first["observation_schema"][0] is not second["observation_schema"][0] From 0d718e49d53ea454ee68d3a0cc575e36e9673edd Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 24 Aug 2026 02:13:48 -0400 Subject: [PATCH 5/5] fix: preserve RL metadata compatibility --- ssms/rl/assembled.py | 15 ++++++++++++--- ssms/rl/config.py | 10 ++++++---- tests/rl/test_observation_metadata.py | 14 ++++++++++++++ 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/ssms/rl/assembled.py b/ssms/rl/assembled.py index e2f0418d..6403af7e 100644 --- a/ssms/rl/assembled.py +++ b/ssms/rl/assembled.py @@ -3,9 +3,9 @@ from __future__ import annotations from collections.abc import Callable, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from enum import StrEnum -from typing import Any, Literal +from typing import Any, ClassVar, Literal import numpy as np @@ -110,10 +110,19 @@ class AssembledModel: context_fields: list[str] computed_params: list[str] response_to_choice: dict[int, int] - _observation_metadata: dict[str, Any] = field(init=False, repr=False, compare=False) + + # Materialized per instance in __post_init__, while ClassVar keeps the private + # cache outside the public dataclass field/serialization contract. + _observation_metadata: ClassVar[dict[str, Any]] def __post_init__(self) -> None: """Snapshot semantic observation metadata without widening the constructor.""" + if self.response != self.config.response or self.choices != tuple( + self.config.choices or () + ): + raise ValueError( + "AssembledModel response and choices must match its ModelConfig" + ) object.__setattr__( self, "_observation_metadata", diff --git a/ssms/rl/config.py b/ssms/rl/config.py index 5ab6aaf4..78fbf5b7 100644 --- a/ssms/rl/config.py +++ b/ssms/rl/config.py @@ -171,10 +171,12 @@ class ModelConfig: default_factory=lambda: {"delta_t": 0.001, "max_t": 20.0} ) - # Additive observation-metadata contract. Keep these fields last so existing - # positional constructor calls retain their meaning. - observation_schema_version: int = 1 - observation_schema: tuple[Mapping[str, Any], ...] | None = None + # Additive observation-metadata contract. These are keyword-only so the existing + # positional constructor and pattern-matching surfaces remain unchanged. + observation_schema_version: int = field(default=1, kw_only=True) + observation_schema: tuple[Mapping[str, Any], ...] | None = field( + default=None, kw_only=True + ) def __post_init__(self): """Auto-build task environment and derive missing fields.""" diff --git a/tests/rl/test_observation_metadata.py b/tests/rl/test_observation_metadata.py index c7f4a118..b280cc24 100644 --- a/tests/rl/test_observation_metadata.py +++ b/tests/rl/test_observation_metadata.py @@ -227,6 +227,20 @@ def test_assembled_metadata_is_an_independent_fresh_snapshot(): assert first["observation_schema"][0] is not second["observation_schema"][0] +@pytest.mark.parametrize( + "change", + [ + {"response": ["response"]}, + {"choices": (1, -1)}, + ], +) +def test_assembled_response_contract_cannot_diverge_from_config(change): + assembled = rl.preset.get("2AB_RW_Angle").assemble(backend="python") + + with pytest.raises(ValueError, match="must match its ModelConfig"): + replace(assembled, **change) + + def test_metadata_access_does_not_change_or_extend_the_hssm_config_dict(): config = rl.preset.get("2AB_RW_Angle") before = config.to_hssm_config_dict()