Reinforcement Learning Based Tractography With SO(3) Equivariant Agents

Fabian Leander Sinzinger, Rodrigo Moreno

Keywords: Tractography, Streamlines, Reinforcement Learning, Spherical CNN, SO(3) Equivariance

TL;DR: We combine reinforcement learning with rotational equivariant (SO3) agents for reliable diffusion-based tractography.

Abstract: Diffusion MRI-based tractography is a promising noninvasive brain connectivity and neural pathway approximation method. While traditional tractography methods relied on deterministic/probabilistic tracking algorithms, recent methods implement data-driven machine learning approaches. The successive prediction of streamlined directions can be biased towards directions frequently occurring in the training data. We propose restoring the rotational equivariance between predicted directions and underlying diffusion MRI data. Therefore, we combine reinforcement learning-based tractography with spherical harmonics-based convolutional neural networks. Our experiments suggest that the proposed method preserves the equivariance and increases explainability.