How is the distance getting estimated? Like another comment said, if a good distance estimator is provided this simplifies the task. Is there a baseline that uses distance in its input as well?