From 02dffb22c99de9d9a46bfebfb49f7f0b05a9f674 Mon Sep 17 00:00:00 2001
From: Micah Woodard <micah.woodard@alleninstitute.org>
Date: Tue, 11 Aug 2026 11:45:18 -0700
Subject: [PATCH 1/2] adds extra metadata to trial generator

---
 .../block_based_trial_generator.py            | 70 +++++++++++++++++-
 .../coupled_trial_generator.py                | 22 ++++--
 .../uncoupled_trial_gnerator.py               | 24 ++++++-
 .../task_logic/utils/__init__.py              |  3 +-
 .../utils/calculate_foraging_efficiency.py    | 69 ++++++++++++++++++
 .../metrics.py                                | 71 ++-----------------
 6 files changed, 182 insertions(+), 77 deletions(-)
 create mode 100644 src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py

diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
index 417a9442..e1069119 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
@@ -1,3 +1,4 @@
+import datetime
 import logging
 from abc import ABC, abstractmethod
 from typing import Literal, Optional
@@ -16,7 +17,7 @@
     BiasIntervention,
     BiasInterventionParameters,
 )
-from aind_behavior_dynamic_foraging.task_logic.utils.calculate_bias import calculate_bias
+from aind_behavior_dynamic_foraging.task_logic.utils import calculate_bias, calculate_foraging_efficiency
 
 from ..trial_models import Metadata, RewardSize, Trial, TrialMetrics
 from ._base import BaseTrialGeneratorSpecModel, ITrialGenerator, TrialOutcome
@@ -27,6 +28,20 @@
 class BlockBasedTrialMetadata(BaseModel):
     """Metadata for block based trial. These fields will NOT be used by the task engine."""
 
+    time_elapsed: Optional[float] = Field(
+        default=None, description="Time elapsed in session at start of trial (in minutes)."
+    )
+    time_remaining: Optional[float] = Field(
+        default=None, description="Time remaining in session at start of trial (in minutes)."
+    )
+    current_trial: Optional[int] = Field(default=None, description="Current trial number in session.")
+    responses: Optional[int] = Field(default=None, description="Number of responses made in session.")
+    ignored: Optional[int] = Field(default=None, description="Number of ignored trials in session.")
+    earned_water: Optional[float] = Field(default=None, description="Total water earned in session (in mL).")
+    total_water: Optional[float] = Field(default=None, description="Total water delivered in session (in mL).")
+    foraging_efficiency: Optional[float] = Field(
+        default=None, description="Foraging efficiency in session (earned water / total water)."
+    )
     is_autowater: bool = Field(default=False, description="Flag indicating if autowater is given for trial.")
     is_bias_water_intervention: bool = Field(
         default=False, description="Flag indicating if bias water intervention is given for trial."
@@ -138,6 +153,7 @@ def __init__(self, spec: BlockBasedTrialGeneratorSpec) -> None:
         """
 
         self.spec = spec
+        self.start_time = datetime.datetime.now()
         self.outcome_history: list[TrialOutcome] = []
         self.is_right_choice_history: list[bool | None] = []
         self.reward_history: list[bool] = []
@@ -228,7 +244,7 @@ def next(self) -> Trial | None:
                 % (is_auto_reward_right, lickspout_offset_delta)
             )
 
-        return Trial(
+        trial = Trial(
             p_reward_left=1 if (self.is_left_baited or is_auto_reward_right is False) else self.block.p_left_reward,
             p_reward_right=1 if (self.is_right_baited or is_auto_reward_right) else self.block.p_right_reward,
             reward_consumption_duration=self.spec.reward_consumption_duration,
@@ -250,6 +266,56 @@ def next(self) -> Trial | None:
                 ),
             ),
         )
+        trial = self.add_extra_metadata(trial)
+        return trial
+
+    def add_extra_metadata(self, trial: Trial) -> Trial:
+        """Adds extra metadata to the trial.
+
+        Args:
+            trial: The Trial to add metadata to.
+
+        Returns:
+            The Trial with extra metadata added.
+        """
+
+        if trial.metadata is None:
+            trial.metadata = Metadata()
+
+        trial.metadata.extra = BlockBasedTrialMetadata(
+            time_elapsed=(datetime.datetime.now() - self.start_time).total_seconds() / 60,
+            current_trial=len(self.outcome_history),
+            responses=sum([1 for choice in self.is_right_choice_history if choice is not None]),
+            ignored=sum([1 for choice in self.is_right_choice_history if choice is None]),
+            earned_water=sum(
+                [
+                    oc.trial.reward_size.left
+                    for oc in self.outcome_history
+                    if oc.is_rewarded and not oc.is_right_choice and oc.trial.is_auto_reward_right is None
+                ]
+                + [
+                    oc.trial.reward_size.right
+                    for oc in self.outcome_history
+                    if oc.is_rewarded and oc.is_right_choice and oc.trial.is_auto_reward_right is None
+                ]
+            ),
+            total_water=sum(
+                [oc.trial.reward_size.left for oc in self.outcome_history if oc.is_rewarded and not oc.is_right_choice]
+                + [oc.trial.reward_size.right for oc in self.outcome_history if oc.is_rewarded and oc.is_right_choice]
+            ),
+            foraging_efficiency=calculate_foraging_efficiency(
+                is_baiting=self.spec.is_baiting,
+                is_rewarded=self.reward_history,
+                p_left_reward=[oc.trial.metadata.p_reward_left for oc in self.outcome_history],
+                p_right_reward=[oc.trial.metadata.p_reward_right for oc in self.outcome_history],
+            ),
+            is_autowater=self._are_autowater_conditions_met(),
+            is_bias_water_intervention=self.bias_intervention.are_antibias_conditions_met(self.bias)
+            and trial.is_auto_reward_right is not None,
+            is_bias_stage_intervention=self.bias_intervention.are_antibias_conditions_met(self.bias)
+            and trial.lickspout_offset_delta != 0,
+        )
+        return trial
 
     def get_metrics(self) -> TrialMetrics:
         """Return metrics at current state of the trial generator."""
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
index 897e5d20..571065f1 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
@@ -5,6 +5,7 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
+from ...trial_models import Trial
 from .base_coupled_trial_generator import (
     BaseCoupledTrialGenerator,
     BaseCoupledTrialGeneratorSpec,
@@ -93,15 +94,26 @@ class CoupledTrialGenerator(BaseCoupledTrialGenerator):
 
     spec: CoupledTrialGeneratorSpec
 
-    def __init__(self, spec: CoupledTrialGeneratorSpec) -> None:
-        """Initializes the generator and records the session start time.
+    def add_extra_metadata(self, trial: Trial) -> Trial:
+        """Adds time remaining metadata to the trial.
 
         Args:
-            spec: The CoupledTrialGeneratorSpec defining task parameters.
+            trial: The Trial object to which metadata will be added.
+
+        Returns:
+            The Trial object with additional metadata.
         """
 
-        super().__init__(spec)
-        self.start_time = datetime.now()
+        trial = super().add_extra_metadata(trial)
+        trial.metadata.extra.time_remaining = (
+            (
+                timedelta(seconds=self.spec.trial_generation_end_parameters.max_time)
+                - (datetime.now() - self.start_time)
+            ).total_seconds()
+            / 60,
+        )
+
+        return trial
 
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
index 61ebf6e8..2d99a209 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
@@ -10,7 +10,7 @@
 from aind_behavior_services.task.distributions_utils import draw_sample
 from pydantic import BaseModel, Field
 
-from ..trial_models import TrialOutcome
+from ..trial_models import TrialOutcome, Trial
 from .block_based_trial_generator import (
     Block,
     BlockBasedTrialGenerator,
@@ -112,7 +112,6 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
         """
 
         super().__init__(spec)
-        self.start_time = datetime.now()
 
         block_length_min = spec.block_length.distribution_parameters.min
         block_length_max = spec.block_length.distribution_parameters.max
@@ -126,6 +125,27 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
         self.trials_in_left_block = 0
         self.left_dominance_streak = 0
 
+    def add_extra_metadata(self, trial: Trial) -> Trial:
+        """Adds time remaining metadata to the trial.
+
+        Args:
+            trial: The Trial object to which metadata will be added.
+
+        Returns:
+            The Trial object with additional metadata.
+        """
+
+        trial = super().add_extra_metadata(trial)
+        trial.metadata.extra.time_remaining = (
+            (
+                timedelta(seconds=self.spec.trial_generation_end_parameters.max_time)
+                - (datetime.now() - self.start_time)
+            ).total_seconds()
+            / 60,
+        )
+
+        return trial
+
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
 
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py b/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py
index ed9492e4..7fea0474 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/utils/__init__.py
@@ -1,3 +1,4 @@
 from .calculate_bias import calculate_bias
+from .calculate_foraging_efficiency import calculate_foraging_efficiency
 
-__all__ = ["calculate_bias"]
+__all__ = ["calculate_bias", "calculate_foraging_efficiency"]
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py b/src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py
new file mode 100644
index 00000000..99ac12a7
--- /dev/null
+++ b/src/aind_behavior_dynamic_foraging/task_logic/utils/calculate_foraging_efficiency.py
@@ -0,0 +1,69 @@
+import logging
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_foraging_efficiency(
+    is_baiting: bool, is_rewarded: list[bool], p_right_reward: list[float], p_left_reward: list[float]
+) -> Optional[float]:
+    """
+    Compute foraging efficiency for a two-arm bandit task.
+
+    This function calculates the ratio of actual rewards obtained to the
+    optimal expected rewards for a session. The implementation is adapted from the Allen Institute dynamic foraging
+    analysis codebase.
+
+    Args:
+        is_baiting (bool):
+            Whether the task uses a baiting schedule. If True, rewards can
+            accumulate on unchosen options; if False, rewards are independent
+            per trial.
+
+        is_rewarded (list[bool | None]):
+            List indicating whether each trial resulted in a reward. `True`
+            indicates a rewarded trial, `False` indicates no reward.
+
+        p_right_reward (list[float]):
+            Probability of reward for the right option on each trial.
+
+        p_left_reward (list[float]):
+            Probability of reward for the left option on each trial.
+
+    Returns:
+        float:
+            Foraging efficiency, defined as the ratio of the number of
+            rewarded trials to the optimal expected number of rewards for
+            the session.
+
+    Raises:
+        ValueError:
+            If input lists have mismatched lengths.
+
+    Notes:
+        Adapted from:
+        https://github.com/AllenNeuralDynamics/aind-dynamic-foraging-basic-analysis/blob/main/src/aind_dynamic_foraging_basic_analysis/metrics/foraging_efficiency.py
+    """
+
+    if not is_baiting:
+        logger.debug("Calculated non baiting foraging efficiency.")
+        optimal_rewards_per_session = np.nanmean(np.max([p_right_reward, p_left_reward], axis=0)) * len(p_left_reward)
+    else:
+        logger.debug("Calculated baiting foraging efficiency.")
+        p_max = np.maximum(p_left_reward, p_right_reward)
+        p_min = np.minimum(p_left_reward, p_right_reward)
+
+        with np.errstate(divide="ignore", invalid="ignore"):
+            optimal_visit_ratio = np.floor(np.log(1 - p_max) / np.log(1 - p_min))
+            optimal_general_reward_rates = p_max + (1 - (1 - p_min) ** (optimal_visit_ratio + 1) - p_max**2) / (
+                optimal_visit_ratio + 1
+            )
+
+        simple_case = (p_min == 0) | (p_max >= 1)
+        optimal_reward_per_trial = np.where(simple_case, p_max, optimal_general_reward_rates)
+
+        optimal_rewards_per_session = np.nanmean(optimal_reward_per_trial) * len(p_left_reward)
+    foraging_efficiency = float(is_rewarded.count(True) / optimal_rewards_per_session)
+    return round(foraging_efficiency, 3)
diff --git a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
index d49731cc..10ba03bb 100644
--- a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
+++ b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
@@ -2,9 +2,9 @@
 import os
 from typing import Annotated, List, Literal, Optional
 
-import numpy as np
 from aind_behavior_curriculum import Metrics
 from aind_behavior_dynamic_foraging.data_contract import dataset as df_foraging_dataset
+from aind_behavior_dynamic_foraging.task_logic.utils.calculate_foraging_efficiency import calculate_foraging_efficiency
 from pydantic import BeforeValidator, Field
 
 STAGE_NAMES = Literal["stage_1_warmup", "stage_1", "stage_2", "stage_3", "final", "graduated"]
@@ -72,9 +72,9 @@ def metrics_from_dataset(
     ]
     is_right_choice = [to["is_right_choice"] for to in filtered]
     is_rewarded = [to["is_rewarded"] for to in filtered]
-    p_right_reward = [to["trial"]["p_reward_right"] for to in filtered]
-    p_left_reward = [to["trial"]["p_reward_left"] for to in filtered]
-    foraging_efficiency = compute_foraging_efficiency(
+    p_right_reward = [to["trial"]["metadata"]["p_reward_right"] for to in filtered]
+    p_left_reward = [to["trial"]["metadata"]["p_reward_left"] for to in filtered]
+    foraging_efficiency = calculate_foraging_efficiency(
         is_baiting=is_baiting, is_rewarded=is_rewarded, p_left_reward=p_left_reward, p_right_reward=p_right_reward
     )
     logger.debug(f"Calculated foraging efficiency as {foraging_efficiency}")
@@ -102,66 +102,3 @@ def metrics_from_dataset(
         consecutive_sessions_at_current_stage=consecutive_sessions_at_current_stage + 1,
         stage_name=stage_name,
     )
-
-
-def compute_foraging_efficiency(
-    is_baiting: bool, is_rewarded: list[bool], p_right_reward: list[float], p_left_reward: list[float]
-) -> Optional[float]:
-    """
-    Compute foraging efficiency for a two-arm bandit task.
-
-    This function calculates the ratio of actual rewards obtained to the
-    optimal expected rewards for a session. The implementation is adapted from the Allen Institute dynamic foraging
-    analysis codebase.
-
-    Args:
-        is_baiting (bool):
-            Whether the task uses a baiting schedule. If True, rewards can
-            accumulate on unchosen options; if False, rewards are independent
-            per trial.
-
-        is_rewarded (list[bool | None]):
-            List indicating whether each trial resulted in a reward. `True`
-            indicates a rewarded trial, `False` indicates no reward.
-
-        p_right_reward (list[float]):
-            Probability of reward for the right option on each trial.
-
-        p_left_reward (list[float]):
-            Probability of reward for the left option on each trial.
-
-    Returns:
-        float:
-            Foraging efficiency, defined as the ratio of the number of
-            rewarded trials to the optimal expected number of rewards for
-            the session.
-
-    Raises:
-        ValueError:
-            If input lists have mismatched lengths.
-
-    Notes:
-        Adapted from:
-        https://github.com/AllenNeuralDynamics/aind-dynamic-foraging-basic-analysis/blob/main/src/aind_dynamic_foraging_basic_analysis/metrics/foraging_efficiency.py
-    """
-
-    if not is_baiting:
-        logger.debug("Calculated non baiting foraging efficiency.")
-        optimal_rewards_per_session = np.nanmean(np.max([p_right_reward, p_left_reward], axis=0)) * len(p_left_reward)
-    else:
-        logger.debug("Calculated baiting foraging efficiency.")
-        p_max = np.maximum(p_left_reward, p_right_reward)
-        p_min = np.minimum(p_left_reward, p_right_reward)
-
-        with np.errstate(divide="ignore", invalid="ignore"):
-            optimal_visit_ratio = np.floor(np.log(1 - p_max) / np.log(1 - p_min))
-            optimal_general_reward_rates = p_max + (1 - (1 - p_min) ** (optimal_visit_ratio + 1) - p_max**2) / (
-                optimal_visit_ratio + 1
-            )
-
-        simple_case = (p_min == 0) | (p_max >= 1)
-        optimal_reward_per_trial = np.where(simple_case, p_max, optimal_general_reward_rates)
-
-        optimal_rewards_per_session = np.nanmean(optimal_reward_per_trial) * len(p_left_reward)
-    foraging_efficiency = float(is_rewarded.count(True) / optimal_rewards_per_session)
-    return round(foraging_efficiency, 3)

From 0a9d48d15092e07af1b66a8eeae2af0b509b090a Mon Sep 17 00:00:00 2001
From: Micah Woodard <micah.woodard@alleninstitute.org>
Date: Tue, 11 Aug 2026 15:52:04 -0700
Subject: [PATCH 2/2] adds extra metadata to trial

---
 .../block_based_trial_generator.py            | 82 ++++++++-----------
 .../coupled_trial_generator.py                | 25 +++---
 .../uncoupled_trial_gnerator.py               | 24 +++---
 .../test_coupled_trial_generator.py           |  3 +-
 .../test_uncoupled_trial_generator.py         |  3 +-
 .../test_warmup_trial_generator.py            |  3 +-
 tests/trial_generators/util.py                | 10 ++-
 .../metrics.py                                |  4 +
 8 files changed, 75 insertions(+), 79 deletions(-)

diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
index e1069119..934d90a0 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/block_based_trial_generator.py
@@ -259,63 +259,53 @@ def next(self) -> Trial | None:
             metadata=Metadata(
                 p_reward_left=self.block.p_left_reward,
                 p_reward_right=self.block.p_right_reward,
-                extra=BlockBasedTrialMetadata(
-                    is_autowater=is_autowater and not is_bias_intervention,
-                    is_bias_water_intervention=is_bias_intervention and is_auto_reward_right is not None,
-                    is_bias_stage_intervention=is_bias_intervention and lickspout_offset_delta != 0,
-                ),
             ),
         )
-        trial = self.add_extra_metadata(trial)
+        extra_metadata = BlockBasedTrialMetadata(
+            is_autowater=is_autowater and not is_bias_intervention,
+            is_bias_water_intervention=is_bias_intervention and is_auto_reward_right is not None,
+            is_bias_stage_intervention=is_bias_intervention and lickspout_offset_delta != 0,
+        )
+        trial.metadata.extra = self._add_extra_metadata(extra_metadata)
         return trial
 
-    def add_extra_metadata(self, trial: Trial) -> Trial:
-        """Adds extra metadata to the trial.
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
+        """Adds extra metadata.
 
         Args:
-            trial: The Trial to add metadata to.
+            extra_metadata: The extra metadata to add to.
 
         Returns:
-            The Trial with extra metadata added.
+            Extra metadata added.
         """
 
-        if trial.metadata is None:
-            trial.metadata = Metadata()
-
-        trial.metadata.extra = BlockBasedTrialMetadata(
-            time_elapsed=(datetime.datetime.now() - self.start_time).total_seconds() / 60,
-            current_trial=len(self.outcome_history),
-            responses=sum([1 for choice in self.is_right_choice_history if choice is not None]),
-            ignored=sum([1 for choice in self.is_right_choice_history if choice is None]),
-            earned_water=sum(
-                [
-                    oc.trial.reward_size.left
-                    for oc in self.outcome_history
-                    if oc.is_rewarded and not oc.is_right_choice and oc.trial.is_auto_reward_right is None
-                ]
-                + [
-                    oc.trial.reward_size.right
-                    for oc in self.outcome_history
-                    if oc.is_rewarded and oc.is_right_choice and oc.trial.is_auto_reward_right is None
-                ]
-            ),
-            total_water=sum(
-                [oc.trial.reward_size.left for oc in self.outcome_history if oc.is_rewarded and not oc.is_right_choice]
-                + [oc.trial.reward_size.right for oc in self.outcome_history if oc.is_rewarded and oc.is_right_choice]
-            ),
-            foraging_efficiency=calculate_foraging_efficiency(
-                is_baiting=self.spec.is_baiting,
-                is_rewarded=self.reward_history,
-                p_left_reward=[oc.trial.metadata.p_reward_left for oc in self.outcome_history],
-                p_right_reward=[oc.trial.metadata.p_reward_right for oc in self.outcome_history],
-            ),
-            is_autowater=self._are_autowater_conditions_met(),
-            is_bias_water_intervention=self.bias_intervention.are_antibias_conditions_met(self.bias)
-            and trial.is_auto_reward_right is not None,
-            is_bias_stage_intervention=self.bias_intervention.are_antibias_conditions_met(self.bias)
-            and trial.lickspout_offset_delta != 0,
+        extra_metadata.time_elapsed = (datetime.datetime.now() - self.start_time).total_seconds() / 60
+        extra_metadata.current_trial = len(self.outcome_history)
+        extra_metadata.responses = sum([1 for choice in self.is_right_choice_history if choice is not None])
+        extra_metadata.ignored = sum([1 for choice in self.is_right_choice_history if choice is None])
+        extra_metadata.earned_water = sum(
+            [
+                oc.trial.reward_size.left
+                for oc in self.outcome_history
+                if oc.is_rewarded and not oc.is_right_choice and oc.trial.is_auto_reward_right is None
+            ]
+            + [
+                oc.trial.reward_size.right
+                for oc in self.outcome_history
+                if oc.is_rewarded and oc.is_right_choice and oc.trial.is_auto_reward_right is None
+            ]
         )
-        return trial
+        extra_metadata.total_water = sum(
+            [oc.trial.reward_size.left for oc in self.outcome_history if oc.is_rewarded and not oc.is_right_choice]
+            + [oc.trial.reward_size.right for oc in self.outcome_history if oc.is_rewarded and oc.is_right_choice]
+        )
+        extra_metadata.foraging_efficiency = calculate_foraging_efficiency(
+            is_baiting=self.spec.is_baiting,
+            is_rewarded=self.reward_history,
+            p_left_reward=[oc.trial.metadata.p_reward_left for oc in self.outcome_history],
+            p_right_reward=[oc.trial.metadata.p_reward_right for oc in self.outcome_history],
+        )
+        return extra_metadata
 
     def get_metrics(self) -> TrialMetrics:
         """Return metrics at current state of the trial generator."""
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
index 571065f1..11d75d36 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_trial_generator.py
@@ -5,7 +5,9 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
-from ...trial_models import Trial
+from ..block_based_trial_generator import (
+    BlockBasedTrialMetadata,
+)
 from .base_coupled_trial_generator import (
     BaseCoupledTrialGenerator,
     BaseCoupledTrialGeneratorSpec,
@@ -94,26 +96,21 @@ class CoupledTrialGenerator(BaseCoupledTrialGenerator):
 
     spec: CoupledTrialGeneratorSpec
 
-    def add_extra_metadata(self, trial: Trial) -> Trial:
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
         """Adds time remaining metadata to the trial.
 
         Args:
-            trial: The Trial object to which metadata will be added.
+            extra_metadata: The extra metadata to which additional metadata will be added.
 
         Returns:
-            The Trial object with additional metadata.
+            The extra metadata with additional metadata.
         """
 
-        trial = super().add_extra_metadata(trial)
-        trial.metadata.extra.time_remaining = (
-            (
-                timedelta(seconds=self.spec.trial_generation_end_parameters.max_time)
-                - (datetime.now() - self.start_time)
-            ).total_seconds()
-            / 60,
-        )
-
-        return trial
+        extra_metadata = super()._add_extra_metadata(extra_metadata)
+        extra_metadata.time_remaining = (
+            timedelta(seconds=self.spec.trial_generation_end_parameters.max_time) - (datetime.now() - self.start_time)
+        ).total_seconds() / 60
+        return extra_metadata
 
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
index 2d99a209..189448a5 100644
--- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
+++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
@@ -10,11 +10,12 @@
 from aind_behavior_services.task.distributions_utils import draw_sample
 from pydantic import BaseModel, Field
 
-from ..trial_models import TrialOutcome, Trial
+from ..trial_models import TrialOutcome
 from .block_based_trial_generator import (
     Block,
     BlockBasedTrialGenerator,
     BlockBasedTrialGeneratorSpec,
+    BlockBasedTrialMetadata,
 )
 
 logger = logging.getLogger(__name__)
@@ -125,26 +126,21 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
         self.trials_in_left_block = 0
         self.left_dominance_streak = 0
 
-    def add_extra_metadata(self, trial: Trial) -> Trial:
+    def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
         """Adds time remaining metadata to the trial.
 
         Args:
-            trial: The Trial object to which metadata will be added.
+            extra_metadata: The extra metadata to which additional metadata will be added.
 
         Returns:
-            The Trial object with additional metadata.
+            The extra metadata with additional metadata.
         """
 
-        trial = super().add_extra_metadata(trial)
-        trial.metadata.extra.time_remaining = (
-            (
-                timedelta(seconds=self.spec.trial_generation_end_parameters.max_time)
-                - (datetime.now() - self.start_time)
-            ).total_seconds()
-            / 60,
-        )
-
-        return trial
+        extra_metadata = super()._add_extra_metadata(extra_metadata)
+        extra_metadata.time_remaining = (
+            timedelta(seconds=self.spec.trial_generation_end_parameters.max_time) - (datetime.now() - self.start_time)
+        ).total_seconds() / 60
+        return extra_metadata
 
     def _are_end_conditions_met(self) -> bool:
         """Checks whether the session should end.
diff --git a/tests/trial_generators/test_coupled_trial_generator.py b/tests/trial_generators/test_coupled_trial_generator.py
index dca8efdf..2d259e7e 100644
--- a/tests/trial_generators/test_coupled_trial_generator.py
+++ b/tests/trial_generators/test_coupled_trial_generator.py
@@ -19,7 +19,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -33,6 +33,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:
diff --git a/tests/trial_generators/test_uncoupled_trial_generator.py b/tests/trial_generators/test_uncoupled_trial_generator.py
index d9521724..b5905884 100644
--- a/tests/trial_generators/test_uncoupled_trial_generator.py
+++ b/tests/trial_generators/test_uncoupled_trial_generator.py
@@ -26,7 +26,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -40,6 +40,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:
diff --git a/tests/trial_generators/test_warmup_trial_generator.py b/tests/trial_generators/test_warmup_trial_generator.py
index 80b4fc50..291d03b8 100644
--- a/tests/trial_generators/test_warmup_trial_generator.py
+++ b/tests/trial_generators/test_warmup_trial_generator.py
@@ -21,7 +21,7 @@ def setUp(self):
     def test_session(self):
         """Simulates a full experimental session to verify generator stability."""
 
-        trial = Trial()
+        trial = self.generator.next()
         outcome = TrialOutcome(
             trial=trial,
             is_right_choice=np.random.choice([True, False, None]),
@@ -35,6 +35,7 @@ def test_session(self):
                 previous_choice=outcome.is_right_choice,
                 previous_left_bait=False,
                 previous_right_bait=False,
+                trial=trial,
             )
 
         if not trial:
diff --git a/tests/trial_generators/util.py b/tests/trial_generators/util.py
index 9ae56250..55ecae34 100644
--- a/tests/trial_generators/util.py
+++ b/tests/trial_generators/util.py
@@ -1,10 +1,16 @@
+from typing import Optional
+
 import numpy as np
 
 from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome
 
 
 def simulate_response(
-    previous_reward: bool, previous_choice: bool | None, previous_left_bait: bool, previous_right_bait: bool
+    previous_reward: bool,
+    previous_choice: bool | None,
+    previous_left_bait: bool,
+    previous_right_bait: bool,
+    trial: Optional[Trial] = None,
 ) -> TrialOutcome:
 
     np.random.seed(42)
@@ -23,4 +29,4 @@ def simulate_response(
     else:
         is_rewarded = previous_right_bait if is_right_choice else previous_left_bait
 
-    return TrialOutcome(trial=Trial(), is_right_choice=is_right_choice, is_rewarded=is_rewarded)
+    return TrialOutcome(trial=trial or Trial(), is_right_choice=is_right_choice, is_rewarded=is_rewarded)
diff --git a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
index 10ba03bb..9010a4a8 100644
--- a/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
+++ b/workspace/aind_behavior_dynamic_foraging_curricula/src/aind_behavior_dynamic_foraging_curricula/metrics.py
@@ -102,3 +102,7 @@ def metrics_from_dataset(
         consecutive_sessions_at_current_stage=consecutive_sessions_at_current_stage + 1,
         stage_name=stage_name,
     )
+
+
+if __name__ == "__main__":
+    print(metrics_from_dataset(r"C:\Users\micah.woodard\Downloads\864253_2026-07-24T194251Z").model_dump_json(indent=4))