diff --git a/schema/aind_behavior_dynamic_foraging.json b/schema/aind_behavior_dynamic_foraging.json
index deec6f0..711b0e6 100644
--- a/schema/aind_behavior_dynamic_foraging.json
+++ b/schema/aind_behavior_dynamic_foraging.json
@@ -433,6 +433,34 @@
"title": "AuditorySecondaryReinforcer",
"type": "object"
},
+ "AutoWaterParameters": {
+ "properties": {
+ "min_ignored_trials": {
+ "default": 3,
+ "description": "Minimum consecutive ignored trials before auto water is triggered.",
+ "minimum": 0,
+ "title": "Min Ignored Trials",
+ "type": "integer"
+ },
+ "min_unrewarded_trials": {
+ "default": 3,
+ "description": "Minimum consecutive unrewarded trials before auto water is triggered.",
+ "minimum": 0,
+ "title": "Min Unrewarded Trials",
+ "type": "integer"
+ },
+ "reward_fraction": {
+ "default": 0.8,
+ "description": "Fraction of full reward volume delivered during auto water (0=none, 1=full).",
+ "maximum": 1,
+ "minimum": 0,
+ "title": "Reward Fraction",
+ "type": "number"
+ }
+ },
+ "title": "AutoWaterParameters",
+ "type": "object"
+ },
"Axis": {
"description": "Motor axis available",
"enum": [
@@ -958,6 +986,22 @@
},
"description": "Parameters defining the reward probability structure."
},
+ "autowater_parameters": {
+ "default": {
+ "min_ignored_trials": 3,
+ "min_unrewarded_trials": 3,
+ "reward_fraction": 0.8
+ },
+ "description": "Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.",
+ "oneOf": [
+ {
+ "$ref": "#/$defs/AutoWaterParameters"
+ },
+ {
+ "type": "null"
+ }
+ ]
+ },
"is_baiting": {
"default": false,
"description": "Whether uncollected rewards carry over to the next trial.",
diff --git a/src/Extensions/AindBehaviorDynamicForaging.Generated.cs b/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
index 3d3ef8c..572bfd5 100644
--- a/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
+++ b/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
@@ -1060,6 +1060,116 @@ public override string ToString()
}
+ [System.CodeDom.Compiler.GeneratedCodeAttribute("Bonsai.Sgen", "0.7.2.0 (Newtonsoft.Json v13.0.0.0)")]
+ [Bonsai.WorkflowElementCategoryAttribute(Bonsai.ElementCategory.Source)]
+ [Bonsai.CombinatorAttribute(MethodName="Generate")]
+ public partial class AutoWaterParameters
+ {
+
+ private int _minIgnoredTrials;
+
+ private int _minUnrewardedTrials;
+
+ private double _rewardFraction;
+
+ public AutoWaterParameters()
+ {
+ _minIgnoredTrials = 3;
+ _minUnrewardedTrials = 3;
+ _rewardFraction = 0.8D;
+ }
+
+ protected AutoWaterParameters(AutoWaterParameters other)
+ {
+ _minIgnoredTrials = other._minIgnoredTrials;
+ _minUnrewardedTrials = other._minUnrewardedTrials;
+ _rewardFraction = other._rewardFraction;
+ }
+
+ ///
+ /// Minimum consecutive ignored trials before auto water is triggered.
+ ///
+ [Newtonsoft.Json.JsonPropertyAttribute("min_ignored_trials")]
+ [System.ComponentModel.DescriptionAttribute("Minimum consecutive ignored trials before auto water is triggered.")]
+ public int MinIgnoredTrials
+ {
+ get
+ {
+ return _minIgnoredTrials;
+ }
+ set
+ {
+ _minIgnoredTrials = value;
+ }
+ }
+
+ ///
+ /// Minimum consecutive unrewarded trials before auto water is triggered.
+ ///
+ [Newtonsoft.Json.JsonPropertyAttribute("min_unrewarded_trials")]
+ [System.ComponentModel.DescriptionAttribute("Minimum consecutive unrewarded trials before auto water is triggered.")]
+ public int MinUnrewardedTrials
+ {
+ get
+ {
+ return _minUnrewardedTrials;
+ }
+ set
+ {
+ _minUnrewardedTrials = value;
+ }
+ }
+
+ ///
+ /// Fraction of full reward volume delivered during auto water (0=none, 1=full).
+ ///
+ [Newtonsoft.Json.JsonPropertyAttribute("reward_fraction")]
+ [System.ComponentModel.DescriptionAttribute("Fraction of full reward volume delivered during auto water (0=none, 1=full).")]
+ public double RewardFraction
+ {
+ get
+ {
+ return _rewardFraction;
+ }
+ set
+ {
+ _rewardFraction = value;
+ }
+ }
+
+ public System.IObservable Generate()
+ {
+ return System.Reactive.Linq.Observable.Defer(() => System.Reactive.Linq.Observable.Return(new AutoWaterParameters(this)));
+ }
+
+ public System.IObservable Generate(System.IObservable source)
+ {
+ return System.Reactive.Linq.Observable.Select(source, _ => new AutoWaterParameters(this));
+ }
+
+ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
+ {
+ stringBuilder.Append("MinIgnoredTrials = " + _minIgnoredTrials + ", ");
+ stringBuilder.Append("MinUnrewardedTrials = " + _minUnrewardedTrials + ", ");
+ stringBuilder.Append("RewardFraction = " + _rewardFraction);
+ return true;
+ }
+
+ public override string ToString()
+ {
+ System.Text.StringBuilder stringBuilder = new System.Text.StringBuilder();
+ stringBuilder.Append(GetType().Name);
+ stringBuilder.Append(" { ");
+ if (PrintMembers(stringBuilder))
+ {
+ stringBuilder.Append(" ");
+ }
+ stringBuilder.Append("}");
+ return stringBuilder.ToString();
+ }
+ }
+
+
///
/// Motor axis available
///
@@ -2416,6 +2526,8 @@ public partial class CoupledTrialGeneratorSpec : TrialGeneratorSpec
private RewardProbabilityParameters _rewardProbabilityParameters;
+ private AutoWaterParameters _autowaterParameters;
+
private bool _isBaiting;
private CoupledTrialGenerationEndConditions _trialGenerationEndParameters;
@@ -2434,6 +2546,7 @@ public CoupledTrialGeneratorSpec()
_minBlockReward = 1;
_kernelSize = 2;
_rewardProbabilityParameters = new RewardProbabilityParameters();
+ _autowaterParameters = new AutoWaterParameters();
_isBaiting = false;
_trialGenerationEndParameters = new CoupledTrialGenerationEndConditions();
_behaviorStabilityParameters = new BehaviorStabilityParameters();
@@ -2451,6 +2564,7 @@ protected CoupledTrialGeneratorSpec(CoupledTrialGeneratorSpec other) :
_minBlockReward = other._minBlockReward;
_kernelSize = other._kernelSize;
_rewardProbabilityParameters = other._rewardProbabilityParameters;
+ _autowaterParameters = other._autowaterParameters;
_isBaiting = other._isBaiting;
_trialGenerationEndParameters = other._trialGenerationEndParameters;
_behaviorStabilityParameters = other._behaviorStabilityParameters;
@@ -2594,6 +2708,25 @@ public RewardProbabilityParameters RewardProbabilityParameters
}
}
+ ///
+ /// Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.
+ ///
+ [System.Xml.Serialization.XmlIgnoreAttribute()]
+ [Newtonsoft.Json.JsonPropertyAttribute("autowater_parameters")]
+ [System.ComponentModel.DescriptionAttribute("Auto water settings. If set, free water is delivered when the animal exceeds the " +
+ "ignored or unrewarded trial thresholds.")]
+ public AutoWaterParameters AutowaterParameters
+ {
+ get
+ {
+ return _autowaterParameters;
+ }
+ set
+ {
+ _autowaterParameters = value;
+ }
+ }
+
///
/// Whether uncollected rewards carry over to the next trial.
///
@@ -2690,6 +2823,7 @@ protected override bool PrintMembers(System.Text.StringBuilder stringBuilder)
stringBuilder.Append("MinBlockReward = " + _minBlockReward + ", ");
stringBuilder.Append("KernelSize = " + _kernelSize + ", ");
stringBuilder.Append("RewardProbabilityParameters = " + _rewardProbabilityParameters + ", ");
+ stringBuilder.Append("AutowaterParameters = " + _autowaterParameters + ", ");
stringBuilder.Append("IsBaiting = " + _isBaiting + ", ");
stringBuilder.Append("TrialGenerationEndParameters = " + _trialGenerationEndParameters + ", ");
stringBuilder.Append("BehaviorStabilityParameters = " + _behaviorStabilityParameters + ", ");
@@ -9508,6 +9642,11 @@ public System.IObservable Process(System.IObservable(source);
}
+ public System.IObservable Process(System.IObservable source)
+ {
+ return Process(source);
+ }
+
public System.IObservable Process(System.IObservable source)
{
return Process(source);
@@ -9833,6 +9972,7 @@ public System.IObservable Process(System.IObservable
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
+ [System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
[System.Xml.Serialization.XmlIncludeAttribute(typeof(Bonsai.Expressions.TypeMapping))]
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 1010f42..ed19686 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
@@ -20,6 +20,21 @@
logger = logging.getLogger(__name__)
+class AutoWaterParameters(BaseModel):
+ min_ignored_trials: int = Field(
+ default=3, ge=0, description="Minimum consecutive ignored trials before auto water is triggered."
+ )
+ min_unrewarded_trials: int = Field(
+ default=3, ge=0, description="Minimum consecutive unrewarded trials before auto water is triggered."
+ )
+ reward_fraction: float = Field(
+ default=0.8,
+ ge=0,
+ le=1,
+ description="Fraction of full reward volume delivered during auto water (0=none, 1=full).",
+ )
+
+
class RewardProbabilityParameters(BaseModel):
"""Defines the reward probability structure for a dynamic foraging task.
@@ -91,6 +106,12 @@ class BlockBasedTrialGeneratorSpec(BaseTrialGeneratorSpecModel):
default=RewardProbabilityParameters(), description="Parameters defining the reward probability structure."
)
+ autowater_parameters: Optional[AutoWaterParameters] = Field(
+ default=AutoWaterParameters(),
+ validate_default=True,
+ description="Auto water settings. If set, free water is delivered when the animal exceeds the ignored or unrewarded trial thresholds.",
+ )
+
is_baiting: bool = Field(default=False, description="Whether uncollected rewards carry over to the next trial.")
def create_generator(self) -> "BlockBasedTrialGenerator":
@@ -155,20 +176,24 @@ def next(self) -> Trial | None:
iti = draw_sample(self.spec.inter_trial_interval_duration)
quiescent = draw_sample(self.spec.quiescent_duration)
- p_reward_left = self.block.p_left_reward
- p_reward_right = self.block.p_right_reward
+ # determine baiting
+ random_numbers = np.random.random(2)
+ is_left_baited = self.block.p_left_reward > random_numbers[0]
+ is_right_baited = self.block.p_right_reward > random_numbers[1]
if self.spec.is_baiting:
- random_numbers = np.random.random(2)
-
- is_left_baited = self.block.p_left_reward > random_numbers[0] or self.is_left_baited
- logger.debug(f"Left baited: {is_left_baited}")
- p_reward_left = 1 if is_left_baited else p_reward_left
-
- is_right_baited = self.block.p_right_reward > random_numbers[1] or self.is_right_baited
- logger.debug(f"Right baited: {is_left_baited}")
- p_reward_right = 1 if is_right_baited else p_reward_right
-
+ is_left_baited = is_left_baited or self.is_left_baited
+ is_right_baited = is_right_baited or self.is_right_baited
+ logger.debug(f"Left baited: {is_left_baited}, Right baited: {is_right_baited}")
+
+ # determine autowater
+ is_autowater_trial = self._are_autowater_conditions_met()
+ is_left_autowater = is_left_baited and is_autowater_trial
+ is_right_autowater = is_right_baited and is_autowater_trial
+
+ p_reward_left = 1 if (is_left_baited or is_left_autowater) else self.block.p_left_reward
+ p_reward_right = 1 if (is_right_baited or is_right_autowater) else self.block.p_right_reward
+
return Trial(
p_reward_left=p_reward_left,
p_reward_right=p_reward_right,
@@ -176,8 +201,32 @@ def next(self) -> Trial | None:
response_deadline_duration=self.spec.response_duration,
quiescence_period_duration=quiescent,
inter_trial_interval_duration=iti,
+ is_auto_response_right=is_right_autowater,
)
+ def _are_autowater_conditions_met(self) -> bool:
+ """Checks whether autowater should be given.
+
+ Returns:
+ True if autowater conditions are met, False otherwise.
+ """
+
+ if self.spec.autowater_parameters is None: # autowater disabled
+ return False
+
+ min_ignore = self.spec.autowater_parameters.min_ignored_trials
+ min_unreward = self.spec.autowater_parameters.min_unrewarded_trials
+
+ is_ignored = [choice is None for choice in self.is_right_choice_history]
+ if all(is_ignored[-min_ignore:]):
+ return True
+
+ is_unrewarded = [not reward for reward in self.reward_history]
+ if all(is_unrewarded[-min_unreward:]):
+ return True
+
+ return False
+
@abstractmethod
def _are_end_conditions_met(self) -> bool:
"""Checks whether the session should end.
@@ -193,7 +242,7 @@ def _generate_next_block(
reward_pairs: list[list[float, float]],
base_reward_sum: float,
block_len: Union[UniformDistribution, ExponentialDistribution],
- current_block: Optional[None] = None,
+ current_block: Optional[Block] = None,
) -> Block:
"""Generates the next block, avoiding repeating the current block's side bias.