diff --git a/schema/aind_behavior_dynamic_foraging.json b/schema/aind_behavior_dynamic_foraging.json
index 52049e3..a68120a 100644
--- a/schema/aind_behavior_dynamic_foraging.json
+++ b/schema/aind_behavior_dynamic_foraging.json
@@ -2355,6 +2355,27 @@
],
"x-sgen-typename": "AllenNeuralDynamics.AindManipulator.MicrostepResolution"
},
+ "MinimumProbabilityPerseverationConditions": {
+ "description": "Defines the conditions under which a block is extended due to perseveration on minimum probability side.",
+ "properties": {
+ "choice_streak": {
+ "default": 4,
+ "description": "Number of consecutive choices on the minimum-probability side required to extend both block lengths.",
+ "minimum": 1,
+ "title": "Choice Streak",
+ "type": "integer"
+ },
+ "block_extension": {
+ "default": 4,
+ "description": "Number of trials to extend both block lengths when the choice streak condition is met.",
+ "minimum": 1,
+ "title": "Block Extension",
+ "type": "integer"
+ }
+ },
+ "title": "MinimumProbabilityPerseverationConditions",
+ "type": "object"
+ },
"MotorOperationMode": {
"description": "Motor operation mode",
"enum": [
@@ -3876,6 +3897,16 @@
"description": "Maximum number of consecutive blocks a side can have the higher probability.",
"title": "Maximum Dominance Streak",
"type": "integer"
+ },
+ "minimum_probability_perseveration_conditions": {
+ "oneOf": [
+ {
+ "$ref": "#/$defs/MinimumProbabilityPerseverationConditions"
+ },
+ {
+ "type": "null"
+ }
+ ]
}
},
"title": "UncoupledTrialGeneratorSpec",
diff --git a/schema/uncoupled.json b/schema/uncoupled.json
index d4d4cf1..1bb4e6b 100644
--- a/schema/uncoupled.json
+++ b/schema/uncoupled.json
@@ -468,7 +468,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
@@ -554,7 +558,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
@@ -640,7 +648,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
diff --git a/schema/uncoupled_baiting.json b/schema/uncoupled_baiting.json
index 2939ee2..435c65b 100644
--- a/schema/uncoupled_baiting.json
+++ b/schema/uncoupled_baiting.json
@@ -468,7 +468,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
@@ -554,7 +558,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
@@ -640,7 +648,11 @@
0.4,
0.7
],
- "maximum_dominance_streak": 3
+ "maximum_dominance_streak": 3,
+ "minimum_probability_perseveration_conditions": {
+ "choice_streak": 4,
+ "block_extension": 4
+ }
}
},
"version": "0.0.2",
diff --git a/src/Extensions/AindBehaviorDynamicForaging.Generated.cs b/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
index ec914b9..176fc4b 100644
--- a/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
+++ b/src/Extensions/AindBehaviorDynamicForaging.Generated.cs
@@ -4442,6 +4442,101 @@ public override string ToString()
}
+ ///
+ /// Defines the conditions under which a block is extended due to perseveration on minimum probability side.
+ ///
+ [System.CodeDom.Compiler.GeneratedCodeAttribute("Bonsai.Sgen", "0.9.0.0 (Newtonsoft.Json v13.0.0.0)")]
+ [System.ComponentModel.DescriptionAttribute("Defines the conditions under which a block is extended due to perseveration on mi" +
+ "nimum probability side.")]
+ [Bonsai.WorkflowElementCategoryAttribute(Bonsai.ElementCategory.Source)]
+ [Bonsai.CombinatorAttribute(MethodName="Generate")]
+ public partial class MinimumProbabilityPerseverationConditions
+ {
+
+ private int _choiceStreak;
+
+ private int _blockExtension;
+
+ public MinimumProbabilityPerseverationConditions()
+ {
+ _choiceStreak = 4;
+ _blockExtension = 4;
+ }
+
+ protected MinimumProbabilityPerseverationConditions(MinimumProbabilityPerseverationConditions other)
+ {
+ _choiceStreak = other._choiceStreak;
+ _blockExtension = other._blockExtension;
+ }
+
+ ///
+ /// Number of consecutive choices on the minimum-probability side required to extend both block lengths.
+ ///
+ [Newtonsoft.Json.JsonPropertyAttribute("choice_streak")]
+ [System.ComponentModel.DescriptionAttribute("Number of consecutive choices on the minimum-probability side required to extend " +
+ "both block lengths.")]
+ public int ChoiceStreak
+ {
+ get
+ {
+ return _choiceStreak;
+ }
+ set
+ {
+ _choiceStreak = value;
+ }
+ }
+
+ ///
+ /// Number of trials to extend both block lengths when the choice streak condition is met.
+ ///
+ [Newtonsoft.Json.JsonPropertyAttribute("block_extension")]
+ [System.ComponentModel.DescriptionAttribute("Number of trials to extend both block lengths when the choice streak condition is" +
+ " met.")]
+ public int BlockExtension
+ {
+ get
+ {
+ return _blockExtension;
+ }
+ set
+ {
+ _blockExtension = value;
+ }
+ }
+
+ public System.IObservable Generate()
+ {
+ return System.Reactive.Linq.Observable.Defer(() => System.Reactive.Linq.Observable.Return(new MinimumProbabilityPerseverationConditions(this)));
+ }
+
+ public System.IObservable Generate(System.IObservable source)
+ {
+ return System.Reactive.Linq.Observable.Select(source, _ => new MinimumProbabilityPerseverationConditions(this));
+ }
+
+ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
+ {
+ stringBuilder.Append("ChoiceStreak = " + _choiceStreak + ", ");
+ stringBuilder.Append("BlockExtension = " + _blockExtension);
+ 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();
+ }
+ }
+
+
///
/// Settings for the quick retract feature.
///
@@ -7017,6 +7112,8 @@ public partial class UncoupledTrialGeneratorSpec : TrialGeneratorSpec
private int _maximumDominanceStreak;
+ private MinimumProbabilityPerseverationConditions _minimumProbabilityPerseverationConditions;
+
public UncoupledTrialGeneratorSpec()
{
_rewardSize = new RewardSize();
@@ -7048,6 +7145,7 @@ protected UncoupledTrialGeneratorSpec(UncoupledTrialGeneratorSpec other) :
_trialGenerationEndParameters = other._trialGenerationEndParameters;
_rewardProbabilities = other._rewardProbabilities;
_maximumDominanceStreak = other._maximumDominanceStreak;
+ _minimumProbabilityPerseverationConditions = other._minimumProbabilityPerseverationConditions;
}
///
@@ -7266,6 +7364,20 @@ public int MaximumDominanceStreak
}
}
+ [System.Xml.Serialization.XmlIgnoreAttribute()]
+ [Newtonsoft.Json.JsonPropertyAttribute("minimum_probability_perseveration_conditions")]
+ public MinimumProbabilityPerseverationConditions MinimumProbabilityPerseverationConditions
+ {
+ get
+ {
+ return _minimumProbabilityPerseverationConditions;
+ }
+ set
+ {
+ _minimumProbabilityPerseverationConditions = value;
+ }
+ }
+
public System.IObservable Generate()
{
return System.Reactive.Linq.Observable.Defer(() => System.Reactive.Linq.Observable.Return(new UncoupledTrialGeneratorSpec(this)));
@@ -7293,7 +7405,8 @@ protected override bool PrintMembers(System.Text.StringBuilder stringBuilder)
stringBuilder.Append("IsBaiting = " + _isBaiting + ", ");
stringBuilder.Append("TrialGenerationEndParameters = " + _trialGenerationEndParameters + ", ");
stringBuilder.Append("RewardProbabilities = " + _rewardProbabilities + ", ");
- stringBuilder.Append("MaximumDominanceStreak = " + _maximumDominanceStreak);
+ stringBuilder.Append("MaximumDominanceStreak = " + _maximumDominanceStreak + ", ");
+ stringBuilder.Append("MinimumProbabilityPerseverationConditions = " + _minimumProbabilityPerseverationConditions);
return true;
}
}
@@ -8555,6 +8668,11 @@ public System.IObservable Process(System.IObservable source)
return Process(source);
}
+ public System.IObservable Process(System.IObservable source)
+ {
+ return Process(source);
+ }
+
public System.IObservable Process(System.IObservable source)
{
return Process(source);
@@ -8702,6 +8820,7 @@ public System.IObservable Process(System.IObservable source)
[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/uncoupled_trial_gnerator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/uncoupled_trial_gnerator.py
index 189448a..a8f1732 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
@@ -1,6 +1,6 @@
import logging
from datetime import datetime, timedelta
-from typing import Literal
+from typing import Literal, Optional
import numpy as np
from aind_behavior_services.task.distributions import (
@@ -38,6 +38,21 @@ class UncoupledTrialGenerationEndConditions(BaseModel):
min_time: float = Field(default=30 * 60, ge=0, description="Minimum session duration (sec)")
+class MinimumProbabilityPerseverationConditions(BaseModel):
+ """Defines the conditions under which a block is extended due to perseveration on minimum probability side."""
+
+ choice_streak: int = Field(
+ default=4,
+ ge=1,
+ description="Number of consecutive choices on the minimum-probability side required to extend both block lengths.",
+ )
+ block_extension: int = Field(
+ default=4,
+ ge=1,
+ description="Number of trials to extend both block lengths when the choice streak condition is met.",
+ )
+
+
class UncoupledTrialGeneratorSpec(BlockBasedTrialGeneratorSpec):
"""
Configuration specification for the Uncoupled Trial Generator.
@@ -78,6 +93,9 @@ class UncoupledTrialGeneratorSpec(BlockBasedTrialGeneratorSpec):
default=3,
description="Maximum number of consecutive blocks a side can have the higher probability.",
)
+ minimum_probability_perseveration_conditions: Optional[MinimumProbabilityPerseverationConditions] = Field(
+ default_factory=MinimumProbabilityPerseverationConditions,
+ )
block_length: UniformDistribution = Field(
default=UniformDistribution(
@@ -117,14 +135,20 @@ def __init__(self, spec: UncoupledTrialGeneratorSpec) -> None:
block_length_min = spec.block_length.distribution_parameters.min
block_length_max = spec.block_length.distribution_parameters.max
self.block_length_stagger = np.floor(
- round((block_length_max - block_length_min - 0.5) / 2 + block_length_min) / 2
+ round(((block_length_max - 1) - block_length_min - 0.5) / 2 + block_length_min) / 2
)
self.block = self._generate_first_block()
+
+ # right counters
self.trials_in_right_block = 0
self.right_dominance_streak = 0
+ self.right_perseveration_streak = 0
+
+ # left counters
self.trials_in_left_block = 0
self.left_dominance_streak = 0
+ self.left_perseveration_streak = 0
def _add_extra_metadata(self, extra_metadata: BlockBasedTrialMetadata) -> BlockBasedTrialMetadata:
"""Adds time remaining metadata to the trial.
@@ -189,12 +213,17 @@ def update(self, outcome: TrialOutcome | str) -> None:
outcome: The TrialOutcome from the most recently completed trial.
"""
+ if isinstance(outcome, str):
+ outcome = TrialOutcome.model_validate_json(outcome)
+
super().update(outcome)
self.trials_in_left_block += 1
self.trials_in_right_block += 1
- switches = []
+ self._update_perseveration_streaks(outcome)
+
+ switches: list[bool] = []
if left_switching := self._is_block_switch_allowed(self.trials_in_left_block, self.block.left_length):
switches.append(False)
if right_switching := self._is_block_switch_allowed(self.trials_in_right_block, self.block.right_length):
@@ -213,16 +242,67 @@ def update(self, outcome: TrialOutcome | str) -> None:
block_stagger=self.block_length_stagger,
block=self.block,
)
- # reset the counter for any side whose probability changed
- if new_block.p_right_reward != self.block.p_right_reward:
- self.trials_in_right_block = 0
- if new_block.p_left_reward != self.block.p_left_reward:
- self.trials_in_left_block = 0
- self.block = new_block
- logger.info(
- "New block generated: p_right_reward=%s, p_left_reward=%s, right_length=%s, left_length=%s."
- % (self.block.p_right_reward, self.block.p_left_reward, self.block.right_length, self.block.left_length)
- )
+ # reset the counter for any side whose probability changed
+ if new_block.p_right_reward != self.block.p_right_reward:
+ self.trials_in_right_block = 0
+ if self.right_dominance_streak >= self.spec.maximum_dominance_streak:
+ self._reset_dominance_streaks()
+
+ if new_block.p_left_reward != self.block.p_left_reward:
+ self.trials_in_left_block = 0
+ if self.left_dominance_streak >= self.spec.maximum_dominance_streak:
+ self._reset_dominance_streaks()
+
+ self.block = new_block
+ logger.info(
+ "New block generated: p_right_reward=%s, p_left_reward=%s, right_length=%s, left_length=%s."
+ % (
+ self.block.p_right_reward,
+ self.block.p_left_reward,
+ self.block.right_length,
+ self.block.left_length,
+ )
+ )
+ if self._is_extend_block_allowed():
+ self._extend_block_lengths()
+
+ def _extend_block_lengths(self) -> None:
+ """Extends both block lengths by the specified extension amount."""
+ self.block.right_length += self.spec.minimum_probability_perseveration_conditions.block_extension
+ self.block.left_length += self.spec.minimum_probability_perseveration_conditions.block_extension
+ self.right_perseveration_streak = 0
+ self.left_perseveration_streak = 0
+ logger.info(
+ "Block lengths extended: right_length=%s, left_length=%s."
+ % (self.block.right_length, self.block.left_length)
+ )
+
+ def _is_extend_block_allowed(self) -> bool:
+ """Return True if either side has exceeded the minimum probability perseveration choice streak."""
+
+ if self.spec.minimum_probability_perseveration_conditions is None:
+ logger.debug("Minimum probability perseveration conditions are not set. Block extension is disabled.")
+ return False
+ return (
+ self.right_perseveration_streak >= self.spec.minimum_probability_perseveration_conditions.choice_streak
+ or self.left_perseveration_streak >= self.spec.minimum_probability_perseveration_conditions.choice_streak
+ )
+
+ def _update_perseveration_streaks(self, trial_outcome: TrialOutcome) -> None:
+ """Update the per-side perseveration streak counters based on the current trial outcome."""
+
+ if trial_outcome.is_right_choice is None:
+ return
+
+ if trial_outcome.is_right_choice:
+ self.left_perseveration_streak = 0
+ if self.block.p_right_reward == min(self.spec.reward_probabilities):
+ self.right_perseveration_streak += 1
+
+ if not trial_outcome.is_right_choice:
+ self.right_perseveration_streak = 0
+ if self.block.p_left_reward == min(self.spec.reward_probabilities):
+ self.left_perseveration_streak += 1
def _update_dominance_streak(self) -> None:
"""Update the per-side dominance streak counters based on the current block.
@@ -245,9 +325,14 @@ def _update_dominance_streak(self) -> None:
% (self.right_dominance_streak, self.left_dominance_streak)
)
+ def _reset_dominance_streaks(self) -> None:
+ """Reset the per-side dominance streak counters to zero."""
+ self.right_dominance_streak = 0
+ self.left_dominance_streak = 0
+
def _is_block_switch_allowed(self, trials_in_block: int, block_length: int) -> bool:
"""Return True if the trial counter has exceeded the current block length."""
- return trials_in_block > block_length
+ return trials_in_block >= block_length
def _generate_first_block(self) -> Block:
"""Generate the initial block for both sides, ensuring neither side starts at minimum
@@ -260,8 +345,9 @@ def _generate_first_block(self) -> Block:
logger.info("Generating first block.")
p_left_reward = np.random.choice(self.spec.reward_probabilities)
p_right_reward = np.random.choice(self.spec.reward_probabilities)
- right_length = np.floor(draw_sample(self.spec.block_length))
left_length = np.floor(draw_sample(self.spec.block_length))
+ right_length = np.floor(draw_sample(self.spec.block_length))
+
while p_right_reward == p_left_reward == min(self.spec.reward_probabilities):
if np.random.choice([True, False]):
logger.debug("Right and left reward are both equal to min. Redrawing right probability.")
@@ -291,8 +377,8 @@ def _generate_first_block(self) -> Block:
left_length=left_length,
)
- @staticmethod
def _generate_next_block(
+ self,
right_switching: bool,
right_dominance_streak: int,
left_dominance_streak: int,
@@ -338,8 +424,12 @@ def _generate_next_block(
if right_switching:
logger.info("Generating right block.")
- r_available = [x for x in reward_probabilities if x != p_right_reward]
- p_right_reward = np.random.choice(r_available) if right_dominance_streak < max_dominance_streak else p_min
+ if right_dominance_streak < max_dominance_streak:
+ p_right_reward = self._draw_reward_probability(p_right_reward, reward_probabilities)
+ else:
+ logger.info("Right dominance streak exceeded max. Forcing right reward probability to minimum.")
+ p_right_reward = p_min
+
right_length = np.floor(draw_sample(block_length))
if p_right_reward == p_left_reward == p_min:
@@ -347,12 +437,16 @@ def _generate_next_block(
"Right and left reward are both equal to min. Staggering right block length and generating new left block."
)
right_length -= block_stagger
- p_left_reward = np.random.choice([x for x in reward_probabilities if x != p_min])
+ p_left_reward = self._draw_reward_probability(p_left_reward, reward_probabilities)
left_length = np.floor(draw_sample(block_length))
else:
logger.info("Generating left block.")
- l_available = [x for x in reward_probabilities if x != p_left_reward]
- p_left_reward = np.random.choice(l_available) if left_dominance_streak < max_dominance_streak else p_min
+ if left_dominance_streak < max_dominance_streak:
+ p_left_reward = self._draw_reward_probability(p_left_reward, reward_probabilities)
+ else:
+ logger.info("Left dominance streak exceeded max. Forcing left reward probability to minimum.")
+ p_left_reward = p_min
+
left_length = np.floor(draw_sample(block_length))
if p_right_reward == p_left_reward == p_min:
@@ -360,7 +454,7 @@ def _generate_next_block(
"Right and left reward are both equal to min. Staggering left block length and generating new right block."
)
left_length -= block_stagger
- p_right_reward = np.random.choice([x for x in reward_probabilities if x != p_min])
+ p_right_reward = self._draw_reward_probability(p_right_reward, reward_probabilities)
right_length = np.floor(draw_sample(block_length))
return Block(
@@ -369,3 +463,16 @@ def _generate_next_block(
right_length=right_length,
left_length=left_length,
)
+
+ @staticmethod
+ def _draw_reward_probability(previous_probability: float, reward_probabilities: list[float]) -> float:
+ """Draw a new reward probability from the available probabilities, excluding the previous probability.
+
+ Args:
+ previous_probability: The reward probability from the previous block.
+ reward_probabilities: List of candidate probabilities to sample from."""
+
+ p_reward = np.random.choice(reward_probabilities)
+ while p_reward == previous_probability:
+ p_reward = np.random.choice(reward_probabilities)
+ return p_reward