Expose observation metadata for RLSSM producers

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,
         *,

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().
 

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") == [