Using envpool's xla step function

[francelico]

Hi @vwxyzjn

I'm using cleanba as a basis when implementing RL algorithms. I was wondering why you're not using envpool's jittable step function (i.e. as in https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py). Is the speed-up not that significant? Or is there something about the cleanba implementation that makes it impractical to use? Thanks :)