Skip to content

[RLlib] JAXPolicy (working discrete-actions PPO prototype).#13014

Closed
sven1977 wants to merge 24 commits intoray-project:masterfrom
sven1977:jax_policy
Closed

[RLlib] JAXPolicy (working discrete-actions PPO prototype).#13014
sven1977 wants to merge 24 commits intoray-project:masterfrom
sven1977:jax_policy

Conversation

@sven1977
Copy link
Copy Markdown
Contributor

@sven1977 sven1977 commented Dec 21, 2020 •
edited
Loading

NOTE: Merge this PR here first for separation of concerns: #13091

JAXPolicy (working discrete-actions PPO prototype).

This PR adds:

  • new JAXPolicy class
  • PPOJAXPolicy class and loss
  • Enhancements in JAXModelV2 to enable learning (get/set_weights).
  • Working PPO compilation and LR schedule test case.
  • Learning regression test for CartPole-v0 on PPO.

Why are these changes needed?

Related issue number

Checks

  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

ericl
ericl reviewed Dec 21, 2020
View reviewed changes
rllib/agents/a3c/a3c_torch_policy.py


A3CTorchPolicy = build_torch_policy(
A3CTorchPolicy = build_policy_class(
Copy link
Copy Markdown
Contributor

@ericl ericl Dec 21, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we do these kind of unrelated changes in a separate PR?

Copy link
Copy Markdown
Contributor Author

@sven1977 sven1977 Dec 28, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

Copy link
Copy Markdown
Contributor Author

@sven1977 sven1977 Dec 28, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, review this one here first before proceeding: #13091

@sven1977
WIP.
48d619d
@sven1977
WIP.
9b31872
@sven1977
Fix and LINT.
548b6b5
@sven1977
Merge branch 'master' of https://github.com/ray-project/ray into jax_…
905be68 
…policy
@sven1977
Fix and LINT.
46f483a
@sven1977
Merge branch 'master' of https://github.com/ray-project/ray into jax_…
686b6d3 
…policy

# Conflicts:
#	rllib/BUILD
#	rllib/agents/ddpg/ddpg_torch_model.py
#	rllib/agents/maml/maml_tf_policy.py
#	rllib/policy/policy_template.py
@sven1977 sven1977 mentioned this pull request Dec 28, 2020
Merged
6 tasks
ericl
ericl reviewed Dec 30, 2020
View reviewed changes
Copy link
Copy Markdown
Contributor

@ericl ericl left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be cool to also add JAX icon to the documentation (rllib.rst, rllib-algorithms.rst, rllib-toc.rst)

@sven1977
Copy link
Copy Markdown
Contributor Author

sven1977 commented Dec 31, 2020

A follow up PR will have documentation updates, including algo table with JAX sigils.

@bveeramani
Copy link
Copy Markdown
Member

bveeramani commented Jan 30, 2022

‼️ ACTION REQUIRED ‼️

We've switched our code formatter from YAPF to Black (see #21311).

To prevent issues with merging your code, here's what you'll need to do:

  1. Install Black
pip install -I black==21.12b0
  1. Format changed files with Black
curl -o format-changed.sh https://gist.githubusercontent.com/bveeramani/42ef0e9e387b755a8a735b084af976f2/raw/7631276790765d555c423b8db2b679fd957b984a/format-changed.sh
chmod +x ./format-changed.sh
./format-changed.sh
rm format-changed.sh
  1. Commit your changes.
git add --all
git commit -m "Format Python code with Black"
  1. Merge master into your branch.
git pull upstream master
  1. Resolve merge conflicts (if necessary).

After running these steps, you'll have the updated format.sh.

@stale
Copy link
Copy Markdown

stale bot commented Mar 10, 2022

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

  • If you'd like to keep this open, just leave any comment, and the stale label will be removed.

@stale stale bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Mar 10, 2022
@stale
Copy link
Copy Markdown

stale bot commented Mar 25, 2022

Hi again! The issue will be closed because there has been no more activity in the 14 days since the last message.

Please feel free to reopen or open a new issue if you'd still like it to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for opening the issue!

@stale stale bot closed this Mar 25, 2022
@sven1977 sven1977 deleted the jax_policy branch June 2, 2023 20:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ericl ericl ericl left review comments

@michaelzhiluo michaelzhiluo Awaiting requested review from michaelzhiluo

Assignees

@michaelzhiluo michaelzhiluo

Labels

stale The issue is stale. It will be closed within 7 days unless there are further conversation

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@sven1977 @bveeramani @ericl @michaelzhiluo