Leveraging rotational equivariance for reinforcement learning in tractography

Brain tractography involves mapping diffusion-weighted images (DWI) onto streamlines representing neural fibre bundles. Recent research avenues have framed tractography into a reinforcement learning (RL) framework with actor-critic models. However, previous RL-based methods may compromise geometrical relations between the input (DWI) and output (tractogram). More specifically, 3D rotations applied to the input of RL-based tractography are not adequately reflected in the output, indicating a lack of SO(3) equivariance. This study aims to restore the equivariance present in previous non-learning-based methods (e.g., iFOD2 from MRtrix3) to RL-based tractography. To achieve this, we introduce SO(3) equivariant and invariant components for the actors (direction prediction model) and critics (Q-value prediction model), respectively. We employ an SE(3)-equivariant transformer as the next direction prediction function. The fact that both the input DWI and the output directional update can be represented as spherical signals that transform under representations of SO(3) makes this formulation a natural fit for the present problem. The contribution of this work is twofold. First, we discuss rotational equivariance in streamline tractography on a theoretical level. Second, we propose a method that combines RL-based tractography with a rotationally equivariant model. We evaluate the equivariance of the proposed method both locally and globally with phantom and in vivo data. The results show that the proposed method restores the equivariance of Track-to-Learn, which is the state-of-the-art for RL-based tractography. Our code is available at https://github.com/minnelab/SO3TrackToLearn.