Reinforcement Learning for Phylogenetic Trees

Reproducing and extending an RL approach to building evolutionary trees.

For the Advanced Deep Reinforcement Learning seminar at Leiden University, my partner and I reproduced and extended a recent scientific paper: treat the search for the best evolutionary tree as a reinforcement learning problem. Every epoch, the agent starts from a random starting tree and repeatedly rearranges it using Subtree Prune and Regraft (SPR) moves. The reward is defined as how much better the new tree describes the relationship between nodes (DNA sequences) over the starting tree. We began by re-implementing and optimizing the method of Azouri and colleagues, and followed this up with our own extensions on top of it, which we documented in the form of a scientific paper.

The SPR move the agent makes at each step. It prunes a branch out of the tree, shown in red, and regrafts it back somewhere else, shown in blue, then earns a reward based on how faithfully the new tree describes the genetic data when compared to the starting tree.

Links · 💻 Code · 📄 Paper

What we added

We first reproduced their deep Q network baseline, then tried to address the gaps in their approach with our own extensions. The main addition is a Soft Q Learning agent with entropy regularisation, which learned more reliably and with less variance than the baseline. We extended the hand crafted features the agent sees, added a faster feature extraction pipeline, and swept the discount factor. One revealing insight we gained is that a near sighted agent did just as well as a far sighted one, which raises a fair question about where the original method’s gains really come from, and whether the reinforcement learning approach even carries merit for this task.

As a proof of concept we also tried a graph neural network policy that reads the tree structure directly instead of relying on hand crafted features. It learned, but was slower and did not outperform the feature based agent, so we kept it as an experiment rather than the headline result.

The environment. Each alignment sets up a space of possible trees, and the agent moves through it from a starting tree toward the one that best fits the data.

This was a two person seminar project and a reproduction of published work rather than something we invented from scratch. The full write up, extensions and their limits are in the paper.