feat(2048): environment performance improvements

[aar65537]

This PR improves the performance of the Game2048 environment. The improvements include

• Minimizing logic inside of jax.lax.cond and jax.lax.switch
• Using jax.vmap over jax.lax.scan where possible
• A new move implementation
• A can_move implementation that validates an action without mutating the board

	no vmap	vmap 103	vmap 106
cpu	64.36%	201.80%	392.29%
cuda	900.12%	1923.08%	706.87%

The above figure shows the total performance improvement measured as percent increase in steps/sec. For more detailed benchmarking, see here.

[clement-bonnet]

Hi @aar65537, thanks a lot for your suggestions for speed improvement! We will look into them, check that the environment's behavior has not changed, and get back to you shortly.

[clement-bonnet]

Thank you very much for the code and speed improvements! I have left a few comments for which I don't feel very strongly about.

[clement-bonnet]

I have checked that this version is equivalent (in terms of environment behavior) to the current version. I obtained the same learning curves.
Moreover, on a TPU-v4 with 8 cores, I got the performances below (orange is 2048 in main, pink is this updated version):

Environment steps per second (random agent): higher is better

x12 improvement when randomly rolling out.

Train epoch time (a2c agent): lower is better

x12 improvement when training.

The other curves (learning metrics or episode returns) are completely equivalent in both versions.

[clement-bonnet]

If that's okay, I will resolve the comments by applying suggestions and will merge.