Right, that part seems to be based on Schütz et al. 2021 https://arxiv.org/abs/2104.07526