From 850c58ca0bef6021403f42393c847eca54ebef6f Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Thu, 1 Oct 2026 14:35:43 -0700 Subject: [PATCH 1/2] fixes warmup baiting --- .../coupled_trial_generators/coupled_warmup_trial_generator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_warmup_trial_generator.py b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_warmup_trial_generator.py index 013e729..fd1cdab 100644 --- a/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_warmup_trial_generator.py +++ b/src/aind_behavior_dynamic_foraging/task_logic/trial_generators/coupled_trial_generators/coupled_warmup_trial_generator.py @@ -50,6 +50,8 @@ def next(self): self.coupled_generator.is_right_choice_history = self.warmup_generator.is_right_choice_history.copy() self.coupled_generator.reward_history = self.warmup_generator.reward_history.copy() self.coupled_generator.start_time = self.warmup_generator.start_time + self.coupled_generator.is_left_baited = self.warmup_generator.is_left_baited + self.coupled_generator.is_right_baited = self.warmup_generator.is_right_baited self._active_generator = self.coupled_generator trial = self._active_generator.next() From 3f7c96e799579cf5f3a981b964522a6f913d2e14 Mon Sep 17 00:00:00 2001 From: Micah Woodard Date: Thu, 1 Oct 2026 14:39:32 -0700 Subject: [PATCH 2/2] adds test --- .../test_coupled_warmup_generator.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/trial_generators/test_coupled_warmup_generator.py b/tests/trial_generators/test_coupled_warmup_generator.py index a046264..ec0c153 100644 --- a/tests/trial_generators/test_coupled_warmup_generator.py +++ b/tests/trial_generators/test_coupled_warmup_generator.py @@ -103,6 +103,36 @@ def test_state_transfer_from_warmup_to_coupled_generator(self): self.assertEqual(coupled.reward_history, warmup.reward_history) self.assertEqual(coupled.start_time, warmup.start_time) + def test_baiting_state_transfer_from_warmup_to_coupled_generator(self): + trial = self.generator.next() + outcome = TrialOutcome( + trial=trial, + is_right_choice=np.random.choice([True, False, None]), + is_rewarded=np.random.choice([True, False]), + ) + + for i in range(50): + trial = self.generator.next() + self.generator.update(outcome) + outcome = simulate_response( + previous_reward=outcome.is_rewarded, + previous_choice=outcome.is_right_choice, + previous_left_bait=False, + previous_right_bait=False, + trial=trial, + ) + + warmup = self.generator.warmup_generator + warmup.is_left_baited = True + warmup.is_right_baited = False + + next_trial = self.generator.next() + + self.assertIsNotNone(next_trial) + self.assertIs(self.generator._active_generator, self.generator.coupled_generator) + self.assertEqual(self.generator.coupled_generator.is_left_baited, warmup.is_left_baited) + self.assertEqual(self.generator.coupled_generator.is_right_baited, warmup.is_right_baited) + if __name__ == "__main__": unittest.main()