Impulse vs XPBD; and assorted thoughts

Hi, just found this project. Don't really have a concrete issue, but some design questions id like to discuss.

Ive recently started work on a similar project in JAX, that tries to use [XPBD](https://matthias-research.github.io/pages/publications/PBDBodies.pdf) rather than impulse based joints. Its too early to tell how thats going to work out exactly; keeping in mind things like the time-of-impact issues raised in the difftaichi paper. Not that this is a concern for my specific application I have in mind, but I think it does point out a more general issue, that the details of the simulator matter to the effective differentiability. So we will see how that works out; hopefully I will know in a week or two.

The reason I went with this particular approach rather than spring or impulse based is that it should provide quite high stiffness yet unconditionally stable simulations, while at the same time not being restricted to infinitely stiff joints but allowing for compliance where required, and at the same time consisting of a rather simple implementation using an explicit integrator that maps well to JAX; and to differentiability in general I think. The only tension is that the XPBD authors recommend gauss-seidel relaxation as converging faster than jacobi, whereas in JAX id rather have the increased parallelism of a jacobi iteration for solving the constraints. I suppose some kind of graph-partitioned block-gauss-seidel would be optimal on an accelerator. Though for simulations with a handful of constraints it might not matter. One thing to optimize along the way I suppose.

Is XPBD a simulation paradigm you have considered, and do you have any thoughts on it? Perhaps you have already tried and found it wanting in some way, in which case I can save myself the trouble. Just noticed the second author of the above paper is also the author of the recent [GradSim](https://gradsim.github.io/). They make no mention of the XPBD paradigm though; given that they write their simulator as handwritten cuda kernels wrapped in pytorch its also a very different approach from Brax in general. They seem to treat collision resolution in a pure springlike manner, fully resolving the contact dynamics, which is guess isnt great for timestep size but if you are biting the bullet of using an implicit solver anyway I suppose it is very clean in terms of differentiability. No code available yet so kinda hard to tell what is going on exactly.

Given that there is no shortage of physics simulation paradigms, it might be useful to generalize the physics integrator; make it swappable within brax, and be able to compare their merits for various applications in an easy manner. Come to think of it it probably makes sense to port my code to be a fork of brax to save myself from reinventing a few wheels (overall design is similar anyway), and perhaps if I manage to pull off that level of generality in a backwards-compatible manner, I could make a PR here, if there is any interest in that?

You mention in your paper that a simple optimization differentiating through the simulation does not result in trainable walking policies. I wonder what the crucial difference with difftaichi's walker is, or the one in GradSim. Could it be some subtlety relating to time-of-impact collision dynamics like in difftaichi, or some other subtlety of the simulation? Again, I think a swappable physics backend would be useful here I think; implementing a simple pure mass-spring explicitly integrated non-rigid-body engine should be trivial (havnt done anything with sparse matrices in JAX/TPU before; perhaps the scatter/gather approach you use for joints would also scale fine to a large collection of springs? should be fine for academic purposes in any case); although the interfaces need to be sufficiently generic given that youd have a somewhat different datastructure. But do you think that level of generality is feasible, and we can just plug in a mass-spring backend, while retaining a reasonable degree of backwards compatibility? Havnt looked at the Brax code enough yet to tell.