Support for Envpool XLA (full-blown atari games)

[e-zorzi]

Dear maintainers,
thank you for the very cool repo!
I'm an RL researcher, and I love the fact that with principled Jax implementations I can run experiments much faster than before! This is especially cool for complex environments. I know that your library supports BRAX, which is cool for Mujoco-like environments, but what about full-blown ATARI games?

It seems to me that the easiest way to handle (efficiently) ATARI games in Jax is by using Envpool XLA interface, like CleanRL does (whose PPO implementation is not that different to yours). Does this library support this possibility? Have you tried it before?
If not, do you think it would be easy to adapt some of your implementations to handle this possibility, if it's not currently supported?

Thanks for your work!