Closed tiffanyona closed 1 day ago
def LookUpDepletion(
amount_drop: int = 5,
available_water: int = 50,
max_p=0.9,
stop_duration: float = 0.5,
delay_mean: float = 0.5,
option: str = "single",
):
if option == "delayed":
lut_values = [0,0,max_p, max_p, max_p] + list(np.repeat(0, 45))
elif option == "single":
lut_values = [max_p] + list(np.repeat(0, 49))
elif option == "no_reward":
lut_values = list(np.repeat(0, 50))
return vr_task_logic.RewardSpecification(
operant_logic=OperantLogicHelper(stop_duration=stop_duration, is_operant=False),
delay=NormalDistributionHelper(delay_mean, 0.15, 0.0, 1),
reward_function=vr_task_logic.PatchRewardFunction(
amount=vr_task_logic.RewardFunction(vr_task_logic.ConstantFunction(value=amount_drop)),
probability=vr_task_logic.RewardFunction(vr_task_logic.LookupTableFunction(
lut_keys=list(np.arange(0,50)), lut_values=lut_values
),
),
available=vr_task_logic.RewardFunction(vr_task_logic.ConstantFunction(value=available_water)),
depletion_rule=vr_task_logic.DepletionRule.ON_CHOICE,
),
)