Jumanji: Reinforcement Learning Environments in JAX

GITHUB PAPER

Authors: Clément Bonnet et al.

Jumanji is a diverse suite of scalable reinforcement learning environments written in JAX. It features 22 environments, ranging from simple games to NP-hard combinatorial problems, enabling faster iteration and large-scale experimentation in RL research.

Contribution: I contributed to the development of the Sokoban environment, one of the routing challenges featured in the Jumanji suite.

Flashbax: Accelerated Replay Buffers in JAX

GITHUB DOCUMENTATION

Authors: Edan Toledo et al.

Flashbax is a library designed to streamline the use of experience replay buffers within reinforcement learning. It provides efficient buffer implementations like Flat, Trajectory, and Prioritized Buffers, tailored for JAX-based environments.

Contribution: I contributed to the development and design of the buffer library during my internship at InstaDeep.