TrackToLearn Fork - Spherical Equivariant Agents

This repository is a fork of the TrackToLearn framework, forked at commit 968b28ff300055d4fccc8e024d705310a74b7aee. We highly recommend getting familiar with the original TrackToLearn framework before exploring this adaptation for spherical equivariant models.

This repository accompanies our paper published in Medical Image Analysis: "Leveraging Rotational Equivariance for Reinforcement Learning in Tractography".

https://doi.org/10.1016/j.media.2026.104216

Requirements

To run this version, ensure you have the SE(3) transformer implementation by NVIDIA available in your PYTHONPATH:

• NVIDIA SE(3) Transformer