What matters in reinforcement learning for tractography

• Published in: Medical image analysis Vol. 93; p. 103085
• Main Authors: Théberge, Antoine, Desrosiers, Christian, Boré, Arnaud, Descoteaux, Maxime, Jodoin, Pierre-Marc
• Format: Journal Article
• Language: English
• Published: Netherlands Elsevier B.V 01.04.2024
• Subjects: Algorithms; Humans; Learning; Machine learning; Reinforcement learning; Reinforcement, Psychology; Reward; Tractography; Tractography; 41A10; Reinforcement learning; Machine learning; 65D05; 65D17; 41A05
• ISSN: 1361-8415; 1361-8423; 1361-8423
• DOI: 10.1016/j.media.2024.103085

Recently, deep reinforcement learning (RL) has been proposed to learn the tractography procedure and train agents to reconstruct the structure of the white matter without manually curated reference streamlines. While the performances reported were competitive, the proposed framework is complex, and little is still known about the role and impact of its multiple parts. In this work, we thoroughly explore the different components of the proposed framework, such as the choice of the RL algorithm, seeding strategy, the input signal and reward function, and shed light on their impact. Approximately 7,400 models were trained for this work, totalling nearly 41,000 h of GPU time. Our goal is to guide researchers eager to explore the possibilities of deep RL for tractography by exposing what works and what does not work with the category of approach. As such, we ultimately propose a series of recommendations concerning the choice of RL algorithm, the input to the agents, the reward function and more to help future work using reinforcement learning for tractography. We also release the open source codebase, trained models, and datasets for users and researchers wanting to explore reinforcement learning for tractography. •Clear set of pitfalls to avoid and recommendations when learning the tractography procedure using reinforcement learning.•Evaluation of multiple input signal, reward function and tracking procedure formulations to guide future research on the subject.•Open-sourced codebase, trained models and datasets to facilitate usage and improvements.