PPO + JAX + EnvPool + Atari

[vwxyzjn]

Description

Types of changes

• Bug fix
• New feature
• New algorithm
• Documentation

Checklist:

• I've read the CONTRIBUTION guide (required).
• I have ensured pre-commit run --all-files passes (required).
• I have updated the documentation and previewed the changes via mkdocs serve.
• I have updated the tests accordingly (if applicable).

If you are adding new algorithms or your change could result in performance difference, you may need to (re-)run tracked experiments. See #137 as an example PR.

• I have contacted @vwxyzjn to obtain access to the openrlbenchmark W&B team (required).
• I have tracked applicable experiments in openrlbenchmark/cleanrl with --capture-video flag toggled on (required).
• I have added additional documentation and previewed the changes via mkdocs serve.
  • I have explained note-worthy implementation details.
  • I have explained the logged metrics.
  • I have added links to the original paper and related papers (if applicable).
  • I have added links to the PR related to the algorithm.
  • I have created a table comparing my results against those from reputable sources (i.e., the original paper or other reference implementation).
  • I have added the learning curves (in PNG format with width=500 and height=300).
  • I have added links to the tracked experiments.
• I have updated the tests accordingly (if applicable).

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Updated
cleanrl	✅ Ready (Inspect)	Visit Preview	Oct 6, 2022 at 1:23AM (UTC)

[vwxyzjn]

The parameters count match

[vwxyzjn]

Preliminary docs are available at https://cleanrl-git-jax-ppo-envpool-atari-vwxyzjn.vercel.app/rl-algorithms/ppo/#ppo_atari_envpool_xla_jaxpy.

[vwxyzjn]

Hi @yooceii @kinalmehta, I have addressed most of your concerns. Please let me know if additional tweaks are needed.

[kinalmehta]

All the concerns seem to be addressed.

LGTM

Great Work!!

concerns addressed