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 pre-print "Leveraging Rotational Equivariance for Reinforcement Learning in Tractography". (Pre-print link: TODO)

Requirements

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

• NVIDIA SE(3) Transformer