Integrate designs from the ForgeBench work https://arxiv.org/abs/2504.15185 These are very simple to port over. One can reach out to Andy or Hanqiu for any questions.