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.

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

rollout notebook: fixups to benchmarks, plots, text #2456

Merged

Conversation

aftersomemath
Copy link
Contributor

Changes:

  1. MJX benchmarks use jax.lax.scan and block_until_ready
  2. Benchmark parameters updated to match RL use case
  3. humanoid benchmarks are warmstarted by 2 seconds, humanoid100 by 4 seconds
  4. benchmarks use skip_checks and preallocated threadpools
  5. Improved advanced usage plots
  6. Benchmark no longer includes model initialization in timed section
  7. Update notebook text as needed

The benchmark results are qualitatively similar. Except for the N humanoid benchmark where the MJX performance is significantly improved.

It would be helpful if someone can double check the MJX benchmark.

1. MJX benchmarks use jax.lax.scan and block_until_ready
2. Benchmark parameters updated to match RL use case
3. humanoid benchmarks are warmstarted by 2 seconds, humanoid100 by 4 seconds
4. benchmarks use skip_checks and preallocated threadpools
5. Improved advanced usage plots
6. Benchmark no longer includes model initialization in timed section
7. Update notebook text as needed
@copybara-service copybara-service bot merged commit 6b11f2b into google-deepmind:main Feb 25, 2025
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants