Feat/jumanji support

[felixchalumeau]

Related issues: #83

This PR aims at easing the use of Jumanji environments in QDax. Jumanji is an open-source Jax package for games and combinatorial optimization problems written as MDPs.

This PR introduces:

• utils to score a policy in a Jumanji environment (play_step, rollout, scoring)
• an example of how to use a Genetic Algorithm or MAP-Elites on Jumanji-Snake
• a small modification of the MakeFile
• I had to pin jinja to fix the docs

Checks

• a clear description of the PR has been added
• sufficient tests have been written
• relevant section added to the documentation
• example notebook added to the repo
• clean docstrings and comments have been written
• if any issue/observation has been discovered, a new issue has been opened

Future improvements

The way the dependencies are handled in readthedocs is quite annoying at the moment, because the installation of qdax overwrites the requirements. We need to udpate our readthedocs file to fix this.

[codecov-commenter]

Codecov Report

Merging #136 (e561918) into develop (8fd641b) will increase coverage by 0.10%.
The diff coverage is 98.96%.

@@             Coverage Diff             @@
##           develop     #136      +/-   ##
===========================================
+ Coverage    92.31%   92.41%   +0.10%     
===========================================
  Files          103      105       +2     
  Lines         5816     5910      +94     
===========================================
+ Hits          5369     5462      +93     
- Misses         447      448       +1     

Impacted Files	Coverage Δ
tests/default_tasks_test/jumanji_envs_test.py	98.14% <98.14%> (ø)
qdax/tasks/brax_envs.py	96.66% <100.00%> (ø)
qdax/tasks/jumanji_envs.py	100.00% <100.00%> (ø)
qdax/types.py	96.66% <100.00%> (+0.37%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[felixchalumeau]

Thanks @Lookatator for the comments. I agree with them. I'll update accordingly.