Speed difference of MJX with impl='warp' vs. mjwarp code

[nico-bohlinger]

Hey everyone,

I was wondering what the speed difference is when using MJX with impl='warp' vs. writing mjwarp code directly and using it with JAX later.
I usually write MJX code so just adding impl='warp' is convenient for me but if writing everything in mjwarp (like it's done on mjlab) is significantly faster it could be worth to switch.

Any insights on this are appreciated!