TD3 jax fix

[joaogui1]

Description

JAX version of #281

Types of changes

• [x] Bug fix
• [ ] New feature
• [ ] New algorithm
• [ ] Documentation

Checklist:

• [x] I've read the CONTRIBUTION guide (required).
• [x I have ensured pre-commit run --all-files passes (required).
• [ ] 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 https://github.com/vwxyzjn/cleanrl/pull/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[bot]]

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

Name	Status	Preview	Updated
cleanrl	✅ Ready (Inspect)	Visit Preview	Oct 20, 2022 at 9:34PM (UTC)

[vwxyzjn]

Find no sign of regression

Regression report: https://wandb.ai/openrlbenchmark/cleanrl-cache/reports/-285-MuJoCo-CleanRL-s-TD3-JAX--VmlldzoyODI2ODIy

[vwxyzjn]

Thanks @joaogui1!