PyTorch vs JAX for RL — which do you actually use in 2026?

[RewardGuard]

PyTorch vs JAX for RL training — which do you actually use?
There's a lot of debate lately about whether JAX is taking over RL research from PyTorch. Curious what the RewardGuard community is actually using in production and research.

🔵 PyTorch — ecosystem is unbeatable, SB3/CleanRL support, just works
🟡 JAX — XLA compilation, vmap/jit are game changers for RL at scale
🟠 TensorFlow — still here, still valid
⚫ Other (MXNet, custom, etc.)

Drop your reasons below — especially curious if anyone has switched from PyTorch to JAX for RL specifically and whether it was worth the migration pain.