prototype jax with ddpg

[vwxyzjn]

Description

Types of changes

• New algorithm

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 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).

[gitpod-io]

[vercel]

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

Name	Status	Preview	Updated
cleanrl	✅ Ready (Inspect)	Visit Preview	Jul 12, 2022 at 9:17PM (UTC)

[vwxyzjn]

@dosssman @huxiao09 I seem to have gotten CleanRL's DDPG + Jax working: about 5x speed up for free.

[vwxyzjn]

@dosssman and @ikostrikov could you help review this, please? I am unfamiliar with JAX so might be coding up things wrong or have really bad format...

[ikostrikov]

Looks good to me! The only thing I would add is TrainState:
https://flax.readthedocs.io/en/latest/flax.training.html#flax.training.train_state.TrainState

[vwxyzjn]

@dosssman @yooceii could you give a review, please? The changes have been finalized.

[dosssman]

Not very familiar with Jax, so can't really suggest any quality improvement.
Beside that, its relatively easy to understand, and the algorithm logic looks good to me.
Great work as always.

[vwxyzjn]

@dosssman, @yooceii, and @joaogui1 this is ready for review with docs (https://cleanrl-git-jax-ddpg-vwxyzjn.vercel.app/rl-algorithms/ddpg/#ddpg_continuous_action_jaxpy, note some of the links don't work until this PR is merged).

[joaogui1]

This should be x = x * self.action_scale + self.action_bias, no?

[vwxyzjn]

Thanks for this great catch! I am fixing this and will merge after CI passes.