feat(experiments): one-presentation recording, task responses and native input forms

.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@4c76a4757fe6bbc1751638ce6df7e23b7ffdcf2b
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@fd3bec0605731b9c72900dfd7e85137a3aa56312
     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: 4c76a4757fe6bbc1751638ce6df7e23b7ffdcf2b
+      unified_revision: fd3bec0605731b9c72900dfd7e85137a3aa56312

brainscore_core/_model/behavior.py

@@ -24,6 +24,28 @@
 )
 
 
+class RatingRegressor:
+    """Ridge regression readout for numeric answers (ratings).
+
+    Scales features, fits ridge regression on numeric labels, predicts one
+    value per presentation. The rating counterpart of ProbabilitiesClassifier.
+    """
+
+    def __init__(self, alpha: float = 1.0):
+        import sklearn.linear_model
+        import sklearn.preprocessing
+        self._regressor = sklearn.linear_model.Ridge(alpha=alpha)
+        self._scaler = sklearn.preprocessing.StandardScaler()
+
+    def fit(self, features, labels: Iterable) -> 'RatingRegressor':
+        X = self._scaler.fit_transform(np.asarray(features))
+        self._regressor.fit(X, np.asarray(list(labels), dtype=float))
+        return self
+
+    def predict(self, features) -> np.ndarray:
+        return self._regressor.predict(self._scaler.transform(np.asarray(features)))
+
+
 class ProbabilitiesClassifier:
     """Logistic regression readout that produces label probabilities.
 

brainscore_core/_model/dispatch.py

@@ -61,6 +61,25 @@ def register_input_handler(cls, event_type: type, handler_name: str,
     cls._INPUT_HANDLERS.append(entry)
 
 
+CLIP_COLUMNS = ('video_path', 'video_file_name')
+
+
+def _is_clip(stimuli) -> bool:
+    """Whether a stimulus table presents clips (video) rather than stills."""
+    return any(column in getattr(stimuli, 'columns', ()) for column in CLIP_COLUMNS)
+
+
+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."""
+    modality = canonical_modality(modality)
+    matching = [key for key in owner._extractors if canonical_modality(key) == modality]
+    if len(matching) > 1 and modality == 'vision':
+        wanted = 'video' if _is_clip(stimuli) else 'vision'
+        matching = [key for key in matching if key == wanted] or matching
+    return owner._extractors[matching[0]] if matching else None
+
+
 def normalize_extractors(owner, preprocessors, activations_model):
     """Adapt registration forms once to ``extractor(stimuli, layers=...)``.
 
@@ -248,10 +267,10 @@ def extract_for_modality(self, stimuli, modality: str, layers: List[str]):
 
     def _extract_for_modality_inner(self, stimuli, modality: str, layers: List[str]):
         owner = self.owner
+        extractor = extractor_for(owner, modality, stimuli)
+        if extractor is not None:
+            return extractor(stimuli, layers=layers)
         modality = canonical_modality(modality)
-        for key, extractor in owner._extractors.items():
-            if canonical_modality(key) == modality:
-                return extractor(stimuli, layers=layers)
         raise NotImplementedError(
             f"Model '{owner.identifier}' declares {modality!r} input but has no "
             'activation extractor for it. Pass extractor={modality: extractor} to '

brainscore_core/_model/perturbation.py

@@ -1,5 +1,6 @@
 """State-change and perturbation lifecycle support for BrainScoreModel."""
 
+from contextlib import contextmanager
 from dataclasses import replace
 from typing import Callable, Dict
 
@@ -13,6 +14,23 @@ class PerturbationManager:
     def __init__(self, owner) -> None:
         self.owner = owner
         self.active_perturbations: Dict[str, Callable[[], None]] = {}
+        self._changes: Dict[str, object] = {}   # handle -> resolved StateChange
+
+    @contextmanager
+    def intact(self):
+        """Undo every active change for the block, then apply each again.
+
+        Handles keep their identity, so a later reset by handle still works.
+        """
+        handles = list(self.active_perturbations)
+        for handle in reversed(handles):
+            self.active_perturbations[handle]()
+        try:
+            yield
+        finally:
+            for handle in handles:
+                _, cleanup = self.owner._state_change_fn(self._changes[handle])
+                self.active_perturbations[handle] = cleanup
 
     def reset(self) -> None:
         """Undo every active perturbation.
@@ -32,6 +50,7 @@ def reset(self) -> None:
                 failures.append((handle_id, error))
             else:
                 self.active_perturbations.pop(handle_id, None)
+                self._changes.pop(handle_id, None)
         if failures:
             import warnings
             detail = '; '.join(f'{handle!r}: {type(err).__name__}: {err}'
@@ -79,6 +98,7 @@ def dispatch_state_change(self, state_change):
             # model perturbed with no way to retry.
             cleanup()
             self.active_perturbations.pop(handle_id, None)
+            self._changes.pop(handle_id, None)
             return None
 
         if owner._state_change_fn is None:
@@ -107,6 +127,7 @@ def dispatch_state_change(self, state_change):
                 f"state_change_fn must produce unique handle_ids per call."
             )
         self.active_perturbations[applied.handle_id] = cleanup
+        self._changes[applied.handle_id] = state_change
         return applied
 
 

brainscore_core/_sessions/drivers.py

@@ -381,6 +381,52 @@ def drive_behavior_session(
     ))
 
 
+def drive_behavior_and_neural_session(model, session, driver: str = "interact") -> None:
+    """Record ``neural:<region>`` while the subject does the session's task.
+
+    ``behavior`` is sent too when requested.
+
+    Each stimulus is shown once: the regions are recorded during the same
+    forward pass that produced the answer (see
+    ``BehavioralReadout.respond_and_record``). Only ``BrainScoreModel`` runs
+    this; it is reached through its session capability.
+    """
+    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(
+            "Recording during a task returns one time bin; request at most one.")
+    requested = _requested_output_channels(session)
+    regions = [parse_channel(c)[1] for c in requested if parse_channel(c)[0] == "neural"]
+    answers = [c for c in requested if parse_channel(c)[0] == "behavior"]
+    _refuse_variance(model, session)
+    model.start_task(task_context)
+    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"}))
+
+
 _drive_neural_session_via_process = drive_neural_session
 
 

brainscore_core/_sessions/helpers.py

@@ -70,6 +70,7 @@
     _drive_via_process,
     _refuse_variance,
     _start_recording,
+    drive_behavior_and_neural_session,
     drive_behavior_session,
     drive_neural_session,
     drive_typed_session,

brainscore_core/brainscore_model.py

@@ -94,6 +94,8 @@ def __init__(
         answer: Optional[Callable] = None,
         policy: Optional[Callable] = None,
         intervention: Optional[Callable] = None,
+        reply: Optional[Callable] = None,
+        likelihood: Optional[Callable] = None,
         behavioral_readout_layer: Optional[str] = None,
         generation_fn: Optional[Callable] = None,
         action_fn: Optional[Callable] = None,
@@ -130,6 +132,14 @@ def __init__(
         :param answer: ``(stimulus_row, instruction, label_set) -> label``;
             generates and parses a behavioral answer for instruction-following
             tasks.
+        :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.
         :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an
             action for closed-loop / embodied evaluation.
         :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``;
@@ -223,6 +233,7 @@ def __init__(
         self._behavioral_readout_layer = capability_values.get(
             'behavioral_readout_layer')
         self._generation_fn = capability_values.get('generation_fn')
+        self._reply_fn, self._likelihood_fn = reply, likelihood
         self._action_fn = capability_values.get('action_fn')
         self._state_change_fn = capability_values.get('state_change_fn')
 
@@ -270,9 +281,14 @@ def _declared_capabilities(self):
 
     @property
     def in_channels(self):
-        return super().in_channels | set().union(
+        return super().in_channels | self._visual_forms() | set().union(
             *(c.declared_input_channels(self) for c in self._declared_capabilities()))
 
+    def _visual_forms(self):
+        """The visual forms taken natively: a ``vision`` key takes stills, a ``video`` key clips."""
+        forms = {'vision': 'vision:still', 'video': 'vision:clip'}
+        return {forms[key] for key in self._preprocessors if key in forms}
+
     @property
     def out_channels(self):
         # Every output comes from a registered part; nothing is found by name.
@@ -613,7 +629,8 @@ def reset(self) -> None:
                 failures.append(error)
         # Callable capabilities and extractors may own episode/history state.
         # Reset each provider once, including bound-method registrations.
-        providers = [self._action_fn, self._generation_fn, self._activations_model,
+        providers = [self._action_fn, self._generation_fn, self._reply_fn, self._likelihood_fn,
+                     self._activations_model,
                      *self._preprocessors.values(), *self._extractors.values()]
         seen = set()
         from .capabilities import capability_registry

brainscore_core/capabilities/sessions.py

@@ -41,30 +41,59 @@ def supports_session(self, model, channels) -> bool:
         return bool(channels) and all(c.startswith('neural:') for c in channels)
 
     def participate(self, model, session) -> None:
-        if getattr(session, 'streaming', False):
-            _drivers._drive_neural_session_streaming(model, session)
-        else:
-            _drivers.drive_neural_session(model, session, allow_multi_modality=True)
+        task = getattr(session, 'task_context', None)
+        if task is not None and task.task_type != 'passive':
+            # The session's task decides what the subject does (rule 4.1):
+            # record during the task; the answers are not sent.
+            _drivers.drive_behavior_and_neural_session(model, session)
+            return
+        # A passive session records activity even after a task was configured.
+        with model._behavioral.suspended():
+            if getattr(session, 'streaming', False):
+                _drivers._drive_neural_session_streaming(model, session)
+            else:
+                _drivers.drive_neural_session(model, session, allow_multi_modality=True)
 
 
 class BehaviorSessions(_SessionCapability):
     identifier = 'behavior-sessions'
 
     def enabled_for(self, model) -> bool:
         return (model._behavioral_readout_layer is not None
-                or model._generation_fn is not None)
+                or model._generation_fn is not None
+                or model._reply_fn is not None or model._likelihood_fn is not None)
 
     def declared_input_channels(self, model):
-        return {'instruction'} if model._generation_fn is not None else set()
+        replies = (model._generation_fn, model._reply_fn, model._likelihood_fn)
+        return {'instruction'} if any(part is not None for part in replies) else set()
 
     def declared_output_channels(self, model):
-        return {'behavior'}
+        # 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'}
+        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'}
+        return channels
 
     def supports_session(self, model, channels) -> bool:
-        return set(channels) == {'behavior'}
+        # Behavior alone, or behavior plus regions recorded during the same answers.
+        behavior = {c for c in channels if c == 'behavior' or c.startswith('behavior:')}
+        recorded = set(channels) - behavior
+        return bool(behavior) and behavior <= self.declared_output_channels(model) and recorded <= {
+            f'neural:{region}' for region in model.region_layer_map}
 
     def participate(self, model, session) -> None:
-        if getattr(session, 'streaming', False):
+        requested = set(_drivers._requested_output_channels(session))
+        task = getattr(session, 'task_context', None)
+        if requested != {'behavior'} or (task is not None and task.response is not None):
+            _drivers.drive_behavior_and_neural_session(model, session)
+        elif getattr(session, 'streaming', False):
             _drivers._drive_behavior_session_streaming(model, session)
         else:
             _drivers.drive_behavior_session(model, session)

brainscore_core/contract.py

@@ -38,13 +38,53 @@ class TaskContext:
     ``'probabilities'``, ``'classification'``, ``'label'``. Put the stimuli to
     answer (for example a held-out test set) in ``metadata['stimulus_set']``
     when using ``behavioral_response``.
+
+    ``response`` is the form of the answer (``brainscore_core.responses``):
+    ``Choice``, ``Scale`` or ``Text``; ``task_type='passive'`` means no answer.
+    With a response, outputs are requested on ``behavior:<kind>`` channels.
+    The older ``task_type`` values ``'label'``, ``'classification'`` and
+    ``'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
+    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
+    (a lesion, a stimulation) is active: ``'current'`` (default) fits on the
+    changed network; ``'intact'`` fits on the network with every change undone,
+    then tests the changed network with that readout, as an animal trained
+    before a lesion is tested after it.
     """
-    task_type: str
+    task_type: Optional[str] = None
     label_set: Optional[List[str]] = None
     fitting_stimuli: Optional[Any] = None
     instruction: Optional[str] = None
     prefer_path: Optional[str] = None  # 'auto' | 'generation' | 'readout'
     metadata: Dict[str, Any] = field(default_factory=dict)
+    readout_fit: str = 'current'  # 'current' | 'intact'
+    response: Any = None  # a brainscore_core.responses Choice, Scale or Text
+    present: Optional[Tuple[str, ...]] = None  # order of a trial's parts: 'text', 'image', 'video', 'audio'
+
+    def __post_init__(self):
+        # The response, when given, decides the task type; label lists follow it.
+        kinds = {'Choice': 'choice', 'Scale': 'rating', 'Text': 'text'}
+        if self.response is not None:
+            kind = kinds.get(type(self.response).__name__)
+            if kind is None:
+                raise TypeError('response must be a Choice, Scale or Text from brainscore_core.responses')
+            if self.task_type is None:
+                self.task_type = kind
+            elif self.task_type != kind:
+                raise ValueError(f'task_type {self.task_type!r} contradicts a {type(self.response).__name__} response')
+            if kind == 'choice' and self.label_set is None:
+                self.label_set = list(self.response.labels)
+        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'}
+            if unknown or isinstance(self.present, str):
+                raise ValueError(f"present lists parts from 'text', 'image', 'video', 'audio'; got {self.present!r}")
+            self.present = tuple(self.present)
 
 
 

brainscore_core/io_catalog.py

@@ -296,6 +296,9 @@ def validate(
                  "(H, W, 3) uint8 image or (T, H, W, 3) uint8 clip, or a path to one",
                  "image/video preprocessor + a vision model harness (wrapper)",
                  meta_keys=("stimulus_id",), expected_ndim=(3, 4), expected_dtype="uint8",
+                 addressing=("address is the form a subject takes natively: 'still' "
+                             "(images) or 'clip' (video); a subject declares each form it "
+                             "has, and a benchmark adapts a stimulus only when a form is missing"),
                  materializers=("StimulusSet",)),
     CatalogEntry("text", INPUT, "StimulusSet column",
                  "a Python str",
@@ -359,6 +362,8 @@ def validate(
                  "a label, a probability vector, or generated text",
                  "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"),
                  materializers=("start_task", "BehavioralAssembly")),
     # --- outputs: motor (carried by EnvironmentResponse, harness-defined) ---
     CatalogEntry("motor", OUTPUT, "EnvironmentResponse.action",