Resources
Authors & Affiliations
Valentina Njaradi, Rodrigo Carrasco-Davis, Andrew Saxe
Abstract
Animals constantly decide how to allocate mental resources to thrive in their environments. Managing learning is a key challenge, and recent experiments show that animals often incur attentional costs for the potential of greater future rewards, or disengage when a task appears unlearnable, exhibiting nuanced cognitive control over learning. However, optimizing time-dependent control requires complex computations, including integrating an agent's learning dynamics and predicting how control affects future rewards. Here, we present a solution for optimal control of integrated non-linear learning dynamics. In contrast to prior work, which has relied on numerical simulations, we derive an analytical approximation for optimal learning rate scheduling, interpreted as the agent's level of cognitive engagement during learning. Optimality is ensured by the Hamilton-Jacobi-Bellman equation derived from the learning dynamics, which we solve using the homotopy perturbation method (Nik \& Shirazian 2012). Although we derive our results from a simple learning setting, numerically optimized control signals for more complex models exhibit similar behaviour. Our results align with experimental observations in rodents (Masis et al. 2023), showing that learning rates initially increase despite higher control costs, leading to improved performance later. Additionally, when task noise rises, the learning rate decreases, as exerting control becomes less beneficial. Importantly, the approximated optimal control signal relies on known agent variables at each step, making it suitable for online closed-loop control and enabling practical heuristics that biological agents could use to solve this problem. This work extends beyond learning rate scheduling to broader applications, including the Expected Value of Control theory (Shenhav et al. 2013) and potentially the role of neuromodulators like norepinephrine (Stanley et al. 2023) or acetylcholine (Huang et al. 2022) in learning. Together, these results provide an analytical theory of cognitive control of learning, which may sculpt aspects of animal behaviour in a range of settings.