Why does running differentiable simulation training for SHAC using backend=mjx easily cause VRAM explosion or take excessive time?

[skywalkerist]

Hello！When using backend=mjx for differentiable simulation training of the GO1 robot dog model with SHAC, even after significantly simplifying the model—retaining only the foot and lower leg colliders and reshaping them into simple geometric colliders—it still frequently runs out of memory (OOM). I think there is no fundamental difference between the Ant environment and the GO1 environment—both are multi-legged robots. Why is there such a severe OOM issue? Are there any solutions? I’ve already minimized the number of environments.
Thanks for your any answers