diff --git a/ssms/rl/assembled.py b/ssms/rl/assembled.py index 761f1f8e..6403af7e 100644 --- a/ssms/rl/assembled.py +++ b/ssms/rl/assembled.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from enum import StrEnum -from typing import Any, Literal +from typing import Any, ClassVar, Literal import numpy as np @@ -111,6 +111,24 @@ class AssembledModel: computed_params: list[str] response_to_choice: dict[int, int] + # 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", + self.config.get_observation_metadata(), + ) + @classmethod def from_config( cls, @@ -137,6 +155,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/ssms/rl/config.py b/ssms/rl/config.py index 89593734..78fbf5b7 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,13 @@ class ModelConfig: default_factory=lambda: {"delta_t": 0.001, "max_t": 20.0} ) + # 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.""" # Convert TaskConfig -> TaskEnvironment @@ -584,6 +597,81 @@ 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, + ) + + 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}" + ) + 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( + { + "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(). 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") == [ diff --git a/tests/rl/test_observation_metadata.py b/tests/rl/test_observation_metadata.py new file mode 100644 index 00000000..b280cc24 --- /dev/null +++ b/tests/rl/test_observation_metadata.py @@ -0,0 +1,258 @@ +"""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)), + {"name": "latency", "kind": "continuous"}, + ), + "names and order must exactly match response", + ), + ( + ( + {"name": "latency", "kind": "continuous"}, + _response_field((1, -1)), + ), + "values equal to choices in order", + ), + ( + ( + {"name": "latency", "kind": "continuous"}, + {"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"), + response=["latency", "response"], + observation_schema=schema, + ) + + with pytest.raises(ValueError, match=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}, + ), + ) + assembled = config.assemble(backend="python") + + raw_values[0] = 99 + first = get_observation_metadata(assembled) + first["observation_schema"][2]["values"] = (99, 100) + second = get_observation_metadata(assembled) + + 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] + + +@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() + + 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)