3D support with Warp

[knyazer]

Is your feature request related to a problem? Please describe. So, would be nice to have a support for 3D if we want to publish this thingy. This is of course a low priority. But, we can use nvidia warp, which is super-experimental, but relatively working, and it was created precisely with intention of making fast simulator in python. I looked through it, and it seems like it might be possible to adapt it to be continuous-time - in the examples it uses semi explicit euler, which looks like it has direct python (jitted of course) evaluation instead of cuda-one, so might be possible.

Describe the solution you'd like Warp integration in the overall framework, for example having cotix.2d and cotix.3d would be nice

Describe alternatives you've considered We can of course implement it by hands (brax inspired, etc), but we have finite man-power, so we generally don't have a choice. And, while Warp does not work for cpu/tpu - that's not that big of a deal, since most people use gpu's anyway, and if somebody (like us) wants to use TPUs, we can use the 2d one, since it is anyway going to be faster. And less prone to bugs, probably ?

Additional context Of course this is not for now, but for the not-so-far-away future (probably February or something)

[knyazer]

Or with Mujoco JAX? Idk what is better, but probably mjx