Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Closed
wants to merge 24 commits into from
Closed

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

wants to merge 24 commits into from

Conversation

sven1977
Copy link
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 Outdated
@@ -84,8 +84,9 @@ def _value(self, obs):
return self.model.value_function()[0]


A3CTorchPolicy = build_torch_policy(
A3CTorchPolicy = build_policy_class(
Copy link
Contributor

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
Contributor Author

Choose a reason for hiding this comment

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

done

Copy link
Contributor Author

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
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
Contributor Author

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

@bveeramani
Copy link
Member

‼️ 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

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

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