Support for JAX

[vwxyzjn]

It doesn't work yet... Cache the changes here. The effort does not seem worth it because the bottleneck is in the environment.

ppo_gridnet.py
times, step=1.3973662853240967, inference=1.2995522022247314, get_mask=0.5436160564422607, rollout=3.3615562915802, training=0.68623948097229, update=4.0478010177612305
times, step=1.398798942565918, inference=1.2518138885498047, get_mask=0.5237548351287842, rollout=3.296189546585083, training=0.6800618171691895, update=3.9762566089630127
times, step=1.4957928657531738, inference=1.329256296157837, get_mask=0.5641641616821289, rollout=3.5079431533813477, training=0.6815638542175293, update=4.189512014389038
times, step=1.3465540409088135, inference=1.3321504592895508, get_mask=0.546196699142456, rollout=3.358635187149048, training=0.6794841289520264, update=4.038123369216919
times, step=1.3423476219177246, inference=1.2476074695587158, get_mask=0.5164742469787598, rollout=3.2134127616882324, training=0.6839492321014404, update=3.8973660469055176


ppo_gridnet_jax.py
times, step=1.3827629089355469, inference=0.1793196201324463, get_mask=0.42531347274780273, rollout=1.9898040294647217, training=0.023705720901489258, update=2.0135114192962646
times, step=1.4476242065429688, inference=0.17769885063171387, get_mask=0.3145482540130615, rollout=1.9413518905639648, training=0.023492097854614258, update=1.9648451805114746
times, step=1.4554364681243896, inference=0.1767590045928955, get_mask=0.312000036239624, rollout=1.9453554153442383, training=0.023662328720092773, update=1.969019889831543
times, step=1.553600788116455, inference=0.17900538444519043, get_mask=0.19531893730163574, rollout=1.9289541244506836, training=0.02358078956604004, update=1.9525365829467773
times, step=1.5749571323394775, inference=0.17865371704101562, get_mask=0.2940824031829834, rollout=2.0495758056640625, training=0.023550033569335938, update=2.0731265544891357