Warp vs Jax Backend VRAM Usage

[aaprasad]

Hi warp dev team, I've been playing around with using warp as the backend for mujoco playground envs instead of jax and while I've noticed faster compile times and similar performance between both, one glaring difference I've noticed is that warp uses a lot more vram. I was wondering if this is just a result of mujoco-warp still being in development and its a transient issue or if theres something inherent in using warp as the backend that makes it more memory inefficient than jax

[Kenny-Vilella]

Hello,

FYI, there is a similar discussion in #926.
The Warp team has recently fixed an issue with memory allocation during graph capture.
It should improve the situation, if you still see the same issue it may be an issue with JAX <-> warp interop.
In that case, if you could provide a minimal repro, I would forward it to the warp team.

[aaprasad]

oh cool, how can I get those changes? I'm using warp-lang=1.11.0 but I don't have mujoco_warp installed