feat(tasks): multi-image trials, memory traits, fused extractors and implied reaction times

.github/workflows/umi-v2.yml

@@ -11,9 +11,9 @@ jobs:
   coordinated:
     name: UMI v2
     # Update the workflow and runner revision together.
-    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@fd3bec0605731b9c72900dfd7e85137a3aa56312
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@346b90ab5093d520d4c49c10752acce0afa6fb39
     with:
       candidate: core
       candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
       candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }}
-      unified_revision: fd3bec0605731b9c72900dfd7e85137a3aa56312
+      unified_revision: 346b90ab5093d520d4c49c10752acce0afa6fb39

brainscore_core/_evaluation/compatibility.py

@@ -160,6 +160,14 @@ def check_channel_compatibility(subject: Subject, benchmark) -> CompatibilityMat
             f"'{benchmark_id}' sends {_format_channels(bench_inputs)}."
         )
 
+    required_traits = set(getattr(benchmark, 'required_traits', None) or ())
+    missing_traits = required_traits - set(subject.traits) if required_traits else set()
+    if missing_traits:
+        raise CompatibilityError(
+            f"Benchmark '{benchmark_id}' requires subject traits {sorted(missing_traits)}; "
+            f"subject '{subject_id}' declares {sorted(subject.traits) or 'none'}."
+        )
+
     if getattr(benchmark, "uses_session", False):
         probe = getattr(subject, "check_session_support", None)
         if callable(probe):

brainscore_core/_model/dispatch.py

@@ -69,6 +69,16 @@ def _is_clip(stimuli) -> bool:
     return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS)
 
 
+def fused_for(owner, detected):
+    """The fused extractor taking every detected input type, if one is registered."""
+    if len(detected) < 2:
+        return None
+    for types, extractor in owner._fused_extractors:
+        if set(detected) <= types:
+            return extractor
+    return None
+
+
 def extractor_for(owner, modality: str, stimuli):
     """The extractor for ``modality``; with a still and a clip extractor, the one
     matching the stimulus's form. None when the model has no such extractor."""
@@ -93,7 +103,18 @@ def normalize_extractors(owner, preprocessors, activations_model):
     import inspect
 
     if isinstance(activations_model, Mapping):
-        extractors = dict(activations_model)
+        # A fused extractor is registered under every input type it takes at
+        # once, {('vision', 'text'): fused}. It also serves each type alone
+        # unless a single-type extractor is registered for it.
+        fused = [(frozenset(canonical_modality(m) for m in key), value)
+                 for key, value in activations_model.items() if isinstance(key, tuple)]
+        extractors = {}
+        for types, value in fused:
+            for modality in types:
+                extractors.setdefault(modality, value)
+        extractors.update({key: value for key, value in activations_model.items()
+                           if not isinstance(key, tuple)})
+        owner._fused_extractors = fused
     elif activations_model is not None:
         if not preprocessors:
             raise ValueError(
@@ -283,6 +304,13 @@ def process_multi_modality(self, stimuli, detected: Set[str],
         import xarray as xr
 
         owner = self.owner
+        fused = fused_for(owner, detected)
+        if fused is not None:
+            # One fused extractor takes every input type: one forward pass.
+            assembly = fused(stimuli, layers=layers)
+            if owner._recorder.is_multi_region:
+                assembly = owner._tag_neuroids_with_regions(assembly)
+            return assembly
         ordered = sorted(
             detected,
             key=lambda m: (

brainscore_core/_sessions/drivers.py

@@ -394,10 +394,6 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     task_context = getattr(session, "task_context", None)
     if task_context is None:
         raise ValueError("A behavior session needs session.task_context.")
-    if getattr(session, "streaming", False):
-        raise NotImplementedError(
-            "Recording during a task presents the stimulus set at once; "
-            "use streaming=False.")
     time_bins = getattr(session, "time_bins", None)
     if time_bins is not None and len(time_bins) > 1:
         raise NotImplementedError(
@@ -407,24 +403,30 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     answers = [c for c in requested if parse_channel(c)[0] == "behavior"]
     _refuse_variance(model, session)
     model.start_task(task_context)
+    model._behavioral.history = []     # memory spans one session, not the next
     stimuli = _behavior_stimuli_from_session(session, task_context)
     if stimuli is None:
         raise ValueError("Recording during a task needs stimuli to present.")
-    if task_context.response is not None:
-        behavior, neural = model._behavioral.respond(stimuli, answers, regions)
-    else:
-        if set(answers) - {"behavior"}:
-            raise ValueError(f"{answers} need TaskContext(response=...); task_type "
-                             f"{task_context.task_type!r} answers on 'behavior'.")
-        plain, neural = model._behavioral.respond_and_record(stimuli, regions)
-        behavior = {"behavior": plain} if answers else {}
-    for channel in answers:
-        session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=0.0,
-                                 meta={"driver": driver, "task_type": task_context.task_type}))
-    for region in regions:
-        session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0,
-                                 meta={"driver": driver, "record": region,
-                                       "during": "behavior"}))
+    if task_context.response is None and set(answers) - {"behavior"}:
+        raise ValueError(f"{answers} need TaskContext(response=...); task_type "
+                         f"{task_context.task_type!r} answers on 'behavior'.")
+    # A streaming session gets each trial's outputs as soon as that trial ends;
+    # otherwise the whole set is answered and sent at once.
+    streaming = getattr(session, "streaming", False)
+    batches = ([(index, _preserve_stimulus_set_type(stimuli, stimuli.iloc[[index]]))
+                for index in range(len(stimuli))] if streaming else [(0, stimuli)])
+    for index, batch in batches:
+        if task_context.response is not None:
+            behavior, neural = model._behavioral.respond(batch, answers, regions)
+        else:
+            plain, neural = model._behavioral.respond_and_record(batch, regions)
+            behavior = {"behavior": plain} if answers else {}
+        for channel in answers:
+            session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=float(index),
+                                     meta={"driver": driver, "task_type": task_context.task_type}))
+        for region in regions:
+            session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=float(index),
+                                     meta={"driver": driver, "record": region, "during": "behavior"}))
 
 
 _drive_neural_session_via_process = drive_neural_session

brainscore_core/brainscore_model.py

@@ -96,6 +96,7 @@ def __init__(
         intervention: Optional[Callable] = None,
         reply: Optional[Callable] = None,
         likelihood: Optional[Callable] = None,
+        memory: Optional[str] = None,
         behavioral_readout_layer: Optional[str] = None,
         generation_fn: Optional[Callable] = None,
         action_fn: Optional[Callable] = None,
@@ -132,14 +133,25 @@ def __init__(
         :param answer: ``(stimulus_row, instruction, label_set) -> label``;
             generates and parses a behavioral answer for instruction-following
             tasks.
+        :param extractor: also accepts a fused extractor under a tuple of input
+            types, ``{('vision', 'text'): fused}``: one forward pass over all of a
+            stimulus's parts. It declares each type, and takes any stimulus that
+            brings several of them at once.
         :param reply: ``(prompt) -> str``; the model's raw reply to one trial.
             ``prompt`` is a ``brainscore_core.responses.Prompt``: the trial's parts
             in order (instruction, stimulus parts, the response's directive).
             Parsing is the task's. Gives choices, ratings, reported confidence
             and written answers.
         :param likelihood: ``(prompt, candidates) -> {candidate: log
             probability}``; scores candidate replies without generating. Gives
-            choices and their implied probabilities.
+            choices and their implied probabilities. Run the prompt on its own
+            before any candidate, so a recording during scoring is the prompt pass.
+        :param memory: ``'session'`` keeps what the model saw and said across the
+            trials of a session, until ``reset()``: each trial's prompt and reply
+            reach ``reply`` and ``likelihood`` as ``prompt.history``. It declares
+            the ``memory:session`` trait. A stateful extractor (a recurrent model)
+            keeps its own state the same way. ``None`` (default): every trial is
+            independent.
         :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an
             action for closed-loop / embodied evaluation.
         :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``;
@@ -220,6 +232,7 @@ def __init__(
         from ._model.dispatch import normalize_extractors
         self._preprocessors: Dict[str, Optional[Callable]] = dict(preprocessors)
         self._activations_model = activations_model
+        self._fused_extractors = []     # (input types, extractor) registered under a tuple
         self._extractors, self._layer_extractors = normalize_extractors(
             self, preprocessors, activations_model)
         for modality in supported_modalities or ():
@@ -234,6 +247,9 @@ def __init__(
             'behavioral_readout_layer')
         self._generation_fn = capability_values.get('generation_fn')
         self._reply_fn, self._likelihood_fn = reply, likelihood
+        if memory not in (None, 'session'):
+            raise ValueError("memory must be None or 'session'")
+        self._memory = memory
         self._action_fn = capability_values.get('action_fn')
         self._state_change_fn = capability_values.get('state_change_fn')
 
@@ -289,6 +305,10 @@ def _visual_forms(self):
         forms = {'vision': 'vision:still', 'video': 'vision:clip'}
         return {forms[key] for key in self._preprocessors if key in forms}
 
+    @property
+    def traits(self):
+        return {'memory:session'} if self._memory == 'session' else set()
+
     @property
     def out_channels(self):
         # Every output comes from a registered part; nothing is found by name.

brainscore_core/capabilities/neural.py

@@ -65,7 +65,8 @@ def process(self, model, stimuli, **kwargs):
                 f"{sorted(model._region_layer_map_dict.keys())}."
             )
 
-        if multi_modality and len(detected) > 1:
+        from .._model.dispatch import fused_for
+        if (multi_modality or fused_for(model, detected) is not None) and len(detected) > 1:
             return model._process_multi_modality(stimuli, detected, layers)
 
         if not multi_modality and len(detected) > 1:

brainscore_core/capabilities/sessions.py

@@ -71,14 +71,14 @@ def declared_output_channels(self, model):
         # Each part declares the kinds of answer it can give (rule 2.4).
         channels = {'behavior'}
         if model._behavioral_readout_layer is not None:
-            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating'}
+            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating', 'behavior:rt'}
         if model._generation_fn is not None:
             channels |= {'behavior:choice'}
         if model._reply_fn is not None:
             channels |= {'behavior:choice', 'behavior:rating', 'behavior:confidence',
                          'behavior:text'}
         if model._likelihood_fn is not None:
-            channels |= {'behavior:choice', 'behavior:probabilities'}
+            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rt'}
         return channels
 
     def supports_session(self, model, channels) -> bool:

brainscore_core/contract.py

@@ -46,7 +46,9 @@ class TaskContext:
     ``'probabilities'`` keep their meaning on the plain ``behavior`` channel.
 
     ``present`` gives the order of a trial's parts when a stimulus has several
-    (e.g. ``('text', 'video')``: a sentence, then a clip). A replying model is
+    (e.g. ``('text', 'video')``: a sentence, then a clip). Name columns to show
+    several parts of one kind in sequence: ``('image:instruction', 'image:word_1',
+    'image:word_2')``. A replying model is
     shown the instruction, those parts in order, then the response's directive.
 
     ``readout_fit`` says which network a readout is fitted on while a change
@@ -81,9 +83,11 @@ def __post_init__(self):
         elif self.task_type is None:
             raise ValueError("Give a response (Choice, Scale, Text) or task_type='passive'.")
         if self.present is not None:
-            unknown = set(self.present) - {'text', 'image', 'video', 'audio'}
+            kinds = {'text', 'image', 'video', 'audio'}
+            unknown = [entry for entry in self.present if entry.split(':', 1)[0] not in kinds]
             if unknown or isinstance(self.present, str):
-                raise ValueError(f"present lists parts from 'text', 'image', 'video', 'audio'; got {self.present!r}")
+                raise ValueError("present lists parts as a kind ('text', 'image', 'video', 'audio') "
+                                 f"or 'kind:column' (e.g. 'image:word_1'); got {self.present!r}")
             self.present = tuple(self.present)
 
 
@@ -123,6 +127,16 @@ def required_channels(self) -> Set[str]:
         """Inputs the subject cannot run without; none by default."""
         return set()
 
+    @property
+    def traits(self) -> Set[str]:
+        """Properties of the subject a benchmark may require; none by default.
+
+        ``'memory:session'``: the subject keeps what it saw and said across the
+        trials of one session, until ``reset()``. A benchmark that needs a trait
+        lists it in ``required_traits``; a subject without it is N/A.
+        """
+        return set()
+
     @abstractmethod
     def participate(self, session: "Session") -> None:
         """Consume session inputs and emit responses until the session ends."""

brainscore_core/io_catalog.py

@@ -363,7 +363,7 @@ def validate(
                  "a fitted readout or the model's generation path",
                  meta_keys=("label_set",),
                  addressing=("address is the kind of answer a TaskContext response asks for: "
-                             "choice, probabilities, confidence, rating or text"),
+                             "choice, probabilities, confidence, rating, rt or text"),
                  materializers=("start_task", "BehavioralAssembly")),
     # --- outputs: motor (carried by EnvironmentResponse, harness-defined) ---
     CatalogEntry("motor", OUTPUT, "EnvironmentResponse.action",

brainscore_core/responses.py

@@ -51,21 +51,32 @@ class Choice:
     ``probabilities=True`` also returns the implied probability of each label
     (from a readout or a language model's likelihoods). ``confidence=Scale(...)``
     also asks the subject to report how sure it is, in the same reply.
+    ``reaction_time`` also returns an implied reaction time from those
+    probabilities, scored by correlation with human RTs as Futrell2018 scores
+    surprisal: ``'evidence_margin'`` (1 minus the gap between the two most
+    probable labels; close calls are slow) or ``'surprisal'`` (-log2 of the
+    chosen label's probability).
     """
     labels: Tuple[str, ...]
     probabilities: bool = False
     confidence: Optional[Scale] = None
+    reaction_time: Optional[str] = None   # 'evidence_margin' or 'surprisal': an implied RT
+
+    REACTION_TIMES = ('evidence_margin', 'surprisal')
 
     def __post_init__(self):
         if isinstance(self.labels, str) or len(self.labels) < 2:
             raise ValueError('Choice needs at least two labels')
+        if self.reaction_time not in (None, *self.REACTION_TIMES):
+            raise ValueError(f'reaction_time is one of {self.REACTION_TIMES}, not {self.reaction_time!r}')
         object.__setattr__(self, 'labels', tuple(self.labels))
 
     @property
     def channels(self) -> Tuple[str, ...]:
         return (('behavior:choice',)
                 + (('behavior:probabilities',) if self.probabilities else ())
-                + (('behavior:confidence',) if self.confidence is not None else ()))
+                + (('behavior:confidence',) if self.confidence is not None else ())
+                + (('behavior:rt',) if self.reaction_time is not None else ()))
 
     def directive(self) -> str:
         text = f'Answer with exactly one of: {", ".join(self.labels)}.'
@@ -126,6 +137,8 @@ class Prompt:
     """
     parts: Tuple[Tuple[str, Any], ...]
     stimulus_id: Optional[str] = None
+    names: Tuple[str, ...] = ()   # one per part: 'instruction', a column name, or 'directive'
+    history: Tuple[Tuple['Prompt', str], ...] = ()   # earlier trials of the session: (prompt, reply)
 
     def text(self) -> str:
         """The text parts joined, for a model that reads text only."""
@@ -145,17 +158,28 @@ def trial_prompt(task, stimulus_row) -> Prompt:
     if task.present is None and len(present) > 1:
         raise ValueError(f'Each stimulus has several parts {sorted(present)}; say the order people '
                          'saw them in with TaskContext(present=(...)), e.g. present=("text", "video").')
-    missing = [kind for kind in order if kind not in present]
-    if missing:
-        raise ValueError(f'TaskContext.present names {missing}, but the stimulus has only {sorted(present)}.')
-    parts = []
+    shown = []          # (kind, value, name)
+    for entry in order:
+        if ':' in entry:
+            # A named part: 'image:word_1' shows the image in column word_1.
+            kind, column = entry.split(':', 1)
+            value = _cell(stimulus_row, column)
+            if value is None:
+                raise ValueError(f'TaskContext.present names column {column!r}, which this stimulus lacks.')
+            shown.append((kind, value, column))
+        elif entry in present:
+            shown.append((entry, present[entry], entry))
+        else:
+            raise ValueError(f'TaskContext.present names {entry!r}, but the stimulus has only {sorted(present)}.')
+    parts, names = [], []
     if (task.instruction or '').strip():
-        parts.append(('text', task.instruction.strip()))
-    parts += [(kind, present[kind]) for kind in order]
+        parts.append(('text', task.instruction.strip())); names.append('instruction')
+    for kind, value, name in shown:
+        parts.append((kind, value)); names.append(name)
     directive = task.response.directive() if task.response is not None else ''
     if directive:
-        parts.append(('text', directive))
-    return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id'))
+        parts.append(('text', directive)); names.append('directive')
+    return Prompt(tuple(parts), stimulus_id=_cell(stimulus_row, 'stimulus_id'), names=tuple(names))
 
 
 def _cell(row, column):

tests/test_compatibility.py

@@ -395,7 +395,7 @@ def _model(self):
     def test_every_part_declares_registered_channels(self):
         model = self._model()
         assert model.out_channels == {'neural:IT', 'behavior', 'behavior:choice',
-                                      'behavior:probabilities', 'behavior:rating',
+                                      'behavior:probabilities', 'behavior:rating', 'behavior:rt',
                                       'motor', 'perturbation'}
         assert {'vision', 'observation', 'lesion', 'stimulation'} <= model.in_channels
         for output in sorted(model.out_channels):