Hey there op. I don't know what your sensors are measuring (distance to a point maybe? Or angle from a Valve lighthouse for inside-out tracking?)

But here's my "why didn't you just"

Since you have a forward simulation function (pose to measurements), why didn't you use an iterative solver to reverse it? Coordinate descent is easy to code and if you have a constrained range of poses you can probably just use multiple starting points to avoid getting stuck with a local minimum. Then use the last solution as a starting point for the next one to save iterations.

Sure it's not closed-form like an NN and it can still have pathological cases, but the code is a little more transparent