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

[BugFix] Fix Brax #2233

Merged
merged 20 commits into from
Jun 18, 2024
Merged

[BugFix] Fix Brax #2233

merged 20 commits into from
Jun 18, 2024

Conversation

vmoens
Copy link
Contributor

@vmoens vmoens commented Jun 15, 2024

No description provided.

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jun 15, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2233

Note: Links to docs will display an error until the docs builds have been completed.

❌ 12 New Failures, 1 Unrelated Failure

As of commit 4e91e53 with merge base 35df59e (image):

NEW FAILURES - The following jobs have failed:

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 15, 2024
@vmoens vmoens linked an issue Jun 15, 2024 that may be closed by this pull request
[BUG] BRAX: batch size issue #2184
Closed
3 tasks
@github-actions GitHub Actions
Copy link

github-actions bot commented Jun 15, 2024 •
edited
Loading

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 91. Improved: $\large\color{#35bf28}9$. Worsened: $\large\color{#d91a1a}7$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_single 0.1208s 61.2415ms 16.3288 Ops/s 17.3478 Ops/s $\textbf{\color{#d91a1a}-5.87\%}$
test_sync 44.5593ms 32.9545ms 30.3449 Ops/s 32.6524 Ops/s $\textbf{\color{#d91a1a}-7.07\%}$
test_async 58.3457ms 29.4297ms 33.9793 Ops/s 35.2069 Ops/s $\color{#d91a1a}-3.49\%$
test_simple 0.3948s 0.3939s 2.5390 Ops/s 2.6000 Ops/s $\color{#d91a1a}-2.35\%$
test_transformed 0.5593s 0.5551s 1.8015 Ops/s 1.8375 Ops/s $\color{#d91a1a}-1.96\%$
test_serial 1.3337s 1.2820s 0.7801 Ops/s 0.7763 Ops/s $\color{#35bf28}+0.48\%$
test_parallel 1.1465s 1.0808s 0.9252 Ops/s 0.9163 Ops/s $\color{#35bf28}+0.98\%$
test_step_mdp_speed[True-True-True-True-True] 0.1814ms 22.0389μs 45.3743 KOps/s 46.5350 KOps/s $\color{#d91a1a}-2.49\%$
test_step_mdp_speed[True-True-True-True-False] 39.4040μs 13.4731μs 74.2218 KOps/s 76.9803 KOps/s $\color{#d91a1a}-3.58\%$
test_step_mdp_speed[True-True-True-False-True] 55.2100μs 12.9549μs 77.1911 KOps/s 78.5572 KOps/s $\color{#d91a1a}-1.74\%$
test_step_mdp_speed[True-True-True-False-False] 98.5340μs 7.8531μs 127.3383 KOps/s 124.6592 KOps/s $\color{#35bf28}+2.15\%$
test_step_mdp_speed[True-True-False-True-True] 0.1217ms 23.3367μs 42.8509 KOps/s 43.8262 KOps/s $\color{#d91a1a}-2.23\%$
test_step_mdp_speed[True-True-False-True-False] 56.8860μs 14.7310μs 67.8838 KOps/s 70.1521 KOps/s $\color{#d91a1a}-3.23\%$
test_step_mdp_speed[True-True-False-False-True] 46.8570μs 14.3709μs 69.5850 KOps/s 72.3203 KOps/s $\color{#d91a1a}-3.78\%$
test_step_mdp_speed[True-True-False-False-False] 43.5720μs 9.2294μs 108.3497 KOps/s 112.5670 KOps/s $\color{#d91a1a}-3.75\%$
test_step_mdp_speed[True-False-True-True-True] 56.8670μs 24.8302μs 40.2735 KOps/s 41.4181 KOps/s $\color{#d91a1a}-2.76\%$
test_step_mdp_speed[True-False-True-True-False] 42.6700μs 16.3304μs 61.2353 KOps/s 63.9451 KOps/s $\color{#d91a1a}-4.24\%$
test_step_mdp_speed[True-False-True-False-True] 64.7010μs 14.3061μs 69.9003 KOps/s 71.4902 KOps/s $\color{#d91a1a}-2.22\%$
test_step_mdp_speed[True-False-True-False-False] 43.7920μs 9.2122μs 108.5523 KOps/s 112.1665 KOps/s $\color{#d91a1a}-3.22\%$
test_step_mdp_speed[True-False-False-True-True] 65.7730μs 25.9806μs 38.4902 KOps/s 38.8857 KOps/s $\color{#d91a1a}-1.02\%$
test_step_mdp_speed[True-False-False-True-False] 44.5330μs 17.4351μs 57.3554 KOps/s 59.2486 KOps/s $\color{#d91a1a}-3.20\%$
test_step_mdp_speed[True-False-False-False-True] 52.0370μs 15.4804μs 64.5980 KOps/s 65.6932 KOps/s $\color{#d91a1a}-1.67\%$
test_step_mdp_speed[True-False-False-False-False] 42.7400μs 10.4500μs 95.6939 KOps/s 98.6729 KOps/s $\color{#d91a1a}-3.02\%$
test_step_mdp_speed[False-True-True-True-True] 52.1470μs 25.1497μs 39.7618 KOps/s 41.4181 KOps/s $\color{#d91a1a}-4.00\%$
test_step_mdp_speed[False-True-True-True-False] 43.8410μs 16.3483μs 61.1686 KOps/s 63.6955 KOps/s $\color{#d91a1a}-3.97\%$
test_step_mdp_speed[False-True-True-False-True] 59.6220μs 16.5324μs 60.4872 KOps/s 61.5494 KOps/s $\color{#d91a1a}-1.73\%$
test_step_mdp_speed[False-True-True-False-False] 51.9670μs 10.4632μs 95.5732 KOps/s 97.6693 KOps/s $\color{#d91a1a}-2.15\%$
test_step_mdp_speed[False-True-False-True-True] 53.4600μs 26.1933μs 38.1777 KOps/s 39.5422 KOps/s $\color{#d91a1a}-3.45\%$
test_step_mdp_speed[False-True-False-True-False] 44.9730μs 17.3359μs 57.6839 KOps/s 58.8661 KOps/s $\color{#d91a1a}-2.01\%$
test_step_mdp_speed[False-True-False-False-True] 54.6320μs 18.0509μs 55.3988 KOps/s 57.8520 KOps/s $\color{#d91a1a}-4.24\%$
test_step_mdp_speed[False-True-False-False-False] 38.1210μs 11.8310μs 84.5237 KOps/s 88.3219 KOps/s $\color{#d91a1a}-4.30\%$
test_step_mdp_speed[False-False-True-True-True] 94.6700μs 27.6291μs 36.1937 KOps/s 37.7420 KOps/s $\color{#d91a1a}-4.10\%$
test_step_mdp_speed[False-False-True-True-False] 58.0680μs 18.9532μs 52.7614 KOps/s 55.6274 KOps/s $\textbf{\color{#d91a1a}-5.15\%}$
test_step_mdp_speed[False-False-True-False-True] 49.1620μs 17.9145μs 55.8208 KOps/s 58.1955 KOps/s $\color{#d91a1a}-4.08\%$
test_step_mdp_speed[False-False-True-False-False] 37.4700μs 11.8001μs 84.7449 KOps/s 88.3674 KOps/s $\color{#d91a1a}-4.10\%$
test_step_mdp_speed[False-False-False-True-True] 90.1080μs 29.3766μs 34.0407 KOps/s 35.6768 KOps/s $\color{#d91a1a}-4.59\%$
test_step_mdp_speed[False-False-False-True-False] 65.4820μs 20.0103μs 49.9743 KOps/s 53.2425 KOps/s $\textbf{\color{#d91a1a}-6.14\%}$
test_step_mdp_speed[False-False-False-False-True] 63.6690μs 18.9535μs 52.7607 KOps/s 55.2961 KOps/s $\color{#d91a1a}-4.59\%$
test_step_mdp_speed[False-False-False-False-False] 40.6960μs 12.8567μs 77.7802 KOps/s 81.6779 KOps/s $\color{#d91a1a}-4.77\%$
test_values[generalized_advantage_estimate-True-True] 10.8125ms 9.4824ms 105.4586 Ops/s 103.6959 Ops/s $\color{#35bf28}+1.70\%$
test_values[vec_generalized_advantage_estimate-True-True] 35.3908ms 33.3382ms 29.9956 Ops/s 28.3871 Ops/s $\textbf{\color{#35bf28}+5.67\%}$
test_values[td0_return_estimate-False-False] 0.2335ms 0.1713ms 5.8379 KOps/s 5.2958 KOps/s $\textbf{\color{#35bf28}+10.24\%}$
test_values[td1_return_estimate-False-False] 27.3805ms 24.4720ms 40.8631 Ops/s 41.7520 Ops/s $\color{#d91a1a}-2.13\%$
test_values[vec_td1_return_estimate-False-False] 34.8556ms 33.4332ms 29.9103 Ops/s 28.2884 Ops/s $\textbf{\color{#35bf28}+5.73\%}$
test_values[td_lambda_return_estimate-True-False] 35.7163ms 35.2620ms 28.3591 Ops/s 29.0844 Ops/s $\color{#d91a1a}-2.49\%$
test_values[vec_td_lambda_return_estimate-True-False] 34.0309ms 33.3835ms 29.9550 Ops/s 28.2427 Ops/s $\textbf{\color{#35bf28}+6.06\%}$
test_gae_speed[generalized_advantage_estimate-False-1-512] 8.4363ms 8.3255ms 120.1125 Ops/s 121.0846 Ops/s $\color{#d91a1a}-0.80\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 2.2890ms 1.8679ms 535.3532 Ops/s 499.0042 Ops/s $\textbf{\color{#35bf28}+7.28\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.5599ms 0.3629ms 2.7558 KOps/s 2.7966 KOps/s $\color{#d91a1a}-1.46\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 41.4826ms 39.7537ms 25.1549 Ops/s 21.6349 Ops/s $\textbf{\color{#35bf28}+16.27\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 3.6657ms 3.0302ms 330.0122 Ops/s 328.6545 Ops/s $\color{#35bf28}+0.41\%$
test_dqn_speed 6.8898ms 1.3669ms 731.5912 Ops/s 733.8489 Ops/s $\color{#d91a1a}-0.31\%$
test_ddpg_speed 3.2142ms 2.8821ms 346.9694 Ops/s 346.8161 Ops/s $\color{#35bf28}+0.04\%$
test_sac_speed 8.9444ms 8.4497ms 118.3469 Ops/s 116.1816 Ops/s $\color{#35bf28}+1.86\%$
test_redq_speed 14.4980ms 13.3618ms 74.8401 Ops/s 72.6233 Ops/s $\color{#35bf28}+3.05\%$
test_redq_deprec_speed 97.6210ms 15.0233ms 66.5634 Ops/s 70.3567 Ops/s $\textbf{\color{#d91a1a}-5.39\%}$
test_td3_speed 9.4412ms 8.4809ms 117.9114 Ops/s 114.4313 Ops/s $\color{#35bf28}+3.04\%$
test_cql_speed 37.3315ms 36.7218ms 27.2318 Ops/s 26.8867 Ops/s $\color{#35bf28}+1.28\%$
test_a2c_speed 9.2594ms 7.6764ms 130.2692 Ops/s 126.9589 Ops/s $\color{#35bf28}+2.61\%$
test_ppo_speed 9.2663ms 7.8421ms 127.5166 Ops/s 122.5795 Ops/s $\color{#35bf28}+4.03\%$
test_reinforce_speed 8.0190ms 6.7407ms 148.3515 Ops/s 144.4586 Ops/s $\color{#35bf28}+2.69\%$
test_iql_speed 35.2524ms 33.4674ms 29.8798 Ops/s 29.9446 Ops/s $\color{#d91a1a}-0.22\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 4.0188ms 3.6051ms 277.3858 Ops/s 271.7521 Ops/s $\color{#35bf28}+2.07\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 99.7646ms 0.5718ms 1.7490 KOps/s 1.9697 KOps/s $\textbf{\color{#d91a1a}-11.21\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.9304ms 0.4767ms 2.0978 KOps/s 2.0676 KOps/s $\color{#35bf28}+1.46\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 3.7739ms 3.4587ms 289.1241 Ops/s 274.9542 Ops/s $\textbf{\color{#35bf28}+5.15\%}$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 0.8198ms 0.4956ms 2.0179 KOps/s 1.9655 KOps/s $\color{#35bf28}+2.67\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3.4804ms 0.4741ms 2.1093 KOps/s 2.0950 KOps/s $\color{#35bf28}+0.68\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 2.5278ms 1.7566ms 569.2971 Ops/s 578.9065 Ops/s $\color{#d91a1a}-1.66\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 4.8091ms 1.6663ms 600.1171 Ops/s 606.2467 Ops/s $\color{#d91a1a}-1.01\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 5.0541ms 3.8509ms 259.6777 Ops/s 270.4754 Ops/s $\color{#d91a1a}-3.99\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1.5628ms 0.6403ms 1.5618 KOps/s 1.3793 KOps/s $\textbf{\color{#35bf28}+13.23\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 3.6500ms 0.6014ms 1.6627 KOps/s 1.6901 KOps/s $\color{#d91a1a}-1.62\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 5.4655ms 3.6900ms 270.9998 Ops/s 275.8994 Ops/s $\color{#d91a1a}-1.78\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.6041ms 0.4971ms 2.0115 KOps/s 1.9844 KOps/s $\color{#35bf28}+1.37\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3.5846ms 0.4790ms 2.0877 KOps/s 2.0907 KOps/s $\color{#d91a1a}-0.14\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 4.3055ms 3.5826ms 279.1283 Ops/s 271.8278 Ops/s $\color{#35bf28}+2.69\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.1658ms 0.4941ms 2.0237 KOps/s 1.9798 KOps/s $\color{#35bf28}+2.22\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.7518ms 0.4720ms 2.1186 KOps/s 2.0689 KOps/s $\color{#35bf28}+2.40\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 5.3764ms 3.6449ms 274.3540 Ops/s 267.7801 Ops/s $\color{#35bf28}+2.45\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1.1935ms 0.6205ms 1.6115 KOps/s 1.5996 KOps/s $\color{#35bf28}+0.74\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 3.7356ms 0.5973ms 1.6741 KOps/s 1.6850 KOps/s $\color{#d91a1a}-0.65\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.1153s 7.9220ms 126.2303 Ops/s 161.6687 Ops/s $\textbf{\color{#d91a1a}-21.92\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 14.5298ms 12.5303ms 79.8063 Ops/s 80.4049 Ops/s $\color{#d91a1a}-0.74\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 1.5131ms 1.0526ms 950.0195 Ops/s 948.0677 Ops/s $\color{#35bf28}+0.21\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 0.1004s 5.6843ms 175.9245 Ops/s 123.3011 Ops/s $\textbf{\color{#35bf28}+42.68\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 14.7812ms 12.5917ms 79.4172 Ops/s 80.0518 Ops/s $\color{#d91a1a}-0.79\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 3.6570ms 1.1210ms 892.0795 Ops/s 889.5786 Ops/s $\color{#35bf28}+0.28\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 0.1151s 6.1384ms 162.9099 Ops/s 165.5429 Ops/s $\color{#d91a1a}-1.59\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 15.2180ms 12.7620ms 78.3578 Ops/s 78.3119 Ops/s $\color{#35bf28}+0.06\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 1.8918ms 1.2186ms 820.6123 Ops/s 824.1003 Ops/s $\color{#d91a1a}-0.42\%$

@github-actions GitHub Actions
Copy link

github-actions bot commented Jun 15, 2024 •
edited
Loading

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 94. Improved: $\large\color{#35bf28}11$. Worsened: $\large\color{#d91a1a}4$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_single 0.1688s 0.1155s 8.6610 Ops/s 9.0735 Ops/s $\color{#d91a1a}-4.55\%$
test_sync 0.1044s 0.1013s 9.8694 Ops/s 9.7381 Ops/s $\color{#35bf28}+1.35\%$
test_async 0.1873s 94.7095ms 10.5586 Ops/s 10.6807 Ops/s $\color{#d91a1a}-1.14\%$
test_single_pixels 0.1182s 0.1162s 8.6088 Ops/s 8.5895 Ops/s $\color{#35bf28}+0.22\%$
test_sync_pixels 82.3177ms 79.1536ms 12.6337 Ops/s 12.4434 Ops/s $\color{#35bf28}+1.53\%$
test_async_pixels 0.1525s 66.2254ms 15.0999 Ops/s 15.3857 Ops/s $\color{#d91a1a}-1.86\%$
test_simple 0.8325s 0.7789s 1.2839 Ops/s 1.2977 Ops/s $\color{#d91a1a}-1.06\%$
test_transformed 1.0835s 1.0194s 0.9810 Ops/s 0.9798 Ops/s $\color{#35bf28}+0.12\%$
test_serial 2.4486s 2.3998s 0.4167 Ops/s 0.4272 Ops/s $\color{#d91a1a}-2.46\%$
test_parallel 2.3637s 2.3283s 0.4295 Ops/s 0.4316 Ops/s $\color{#d91a1a}-0.49\%$
test_step_mdp_speed[True-True-True-True-True] 0.2577ms 31.6267μs 31.6188 KOps/s 32.4791 KOps/s $\color{#d91a1a}-2.65\%$
test_step_mdp_speed[True-True-True-True-False] 0.1423ms 18.7183μs 53.4236 KOps/s 53.7501 KOps/s $\color{#d91a1a}-0.61\%$
test_step_mdp_speed[True-True-True-False-True] 0.1175ms 18.0284μs 55.4680 KOps/s 56.7429 KOps/s $\color{#d91a1a}-2.25\%$
test_step_mdp_speed[True-True-True-False-False] 78.1110μs 10.9090μs 91.6677 KOps/s 94.4813 KOps/s $\color{#d91a1a}-2.98\%$
test_step_mdp_speed[True-True-False-True-True] 91.7810μs 33.2722μs 30.0551 KOps/s 30.4764 KOps/s $\color{#d91a1a}-1.38\%$
test_step_mdp_speed[True-True-False-True-False] 52.9610μs 20.3724μs 49.0861 KOps/s 49.8203 KOps/s $\color{#d91a1a}-1.47\%$
test_step_mdp_speed[True-True-False-False-True] 66.2910μs 19.5787μs 51.0760 KOps/s 51.0189 KOps/s $\color{#35bf28}+0.11\%$
test_step_mdp_speed[True-True-False-False-False] 41.0500μs 12.7142μs 78.6520 KOps/s 79.9820 KOps/s $\color{#d91a1a}-1.66\%$
test_step_mdp_speed[True-False-True-True-True] 74.0420μs 35.0913μs 28.4971 KOps/s 28.9694 KOps/s $\color{#d91a1a}-1.63\%$
test_step_mdp_speed[True-False-True-True-False] 98.2920μs 22.6203μs 44.2081 KOps/s 45.2191 KOps/s $\color{#d91a1a}-2.24\%$
test_step_mdp_speed[True-False-True-False-True] 58.3310μs 19.5787μs 51.0759 KOps/s 52.5329 KOps/s $\color{#d91a1a}-2.77\%$
test_step_mdp_speed[True-False-True-False-False] 50.2110μs 12.6235μs 79.2174 KOps/s 80.7677 KOps/s $\color{#d91a1a}-1.92\%$
test_step_mdp_speed[True-False-False-True-True] 84.5710μs 36.3999μs 27.4726 KOps/s 27.8849 KOps/s $\color{#d91a1a}-1.48\%$
test_step_mdp_speed[True-False-False-True-False] 54.0210μs 23.9067μs 41.8292 KOps/s 42.1947 KOps/s $\color{#d91a1a}-0.87\%$
test_step_mdp_speed[True-False-False-False-True] 51.7810μs 21.2718μs 47.0107 KOps/s 47.7615 KOps/s $\color{#d91a1a}-1.57\%$
test_step_mdp_speed[True-False-False-False-False] 40.5010μs 14.1878μs 70.4830 KOps/s 71.4262 KOps/s $\color{#d91a1a}-1.32\%$
test_step_mdp_speed[False-True-True-True-True] 74.7920μs 35.1698μs 28.4334 KOps/s 28.8284 KOps/s $\color{#d91a1a}-1.37\%$
test_step_mdp_speed[False-True-True-True-False] 61.4110μs 22.2274μs 44.9895 KOps/s 45.6948 KOps/s $\color{#d91a1a}-1.54\%$
test_step_mdp_speed[False-True-True-False-True] 59.7310μs 23.4077μs 42.7210 KOps/s 43.8775 KOps/s $\color{#d91a1a}-2.64\%$
test_step_mdp_speed[False-True-True-False-False] 82.6820μs 14.1792μs 70.5260 KOps/s 70.3560 KOps/s $\color{#35bf28}+0.24\%$
test_step_mdp_speed[False-True-False-True-True] 72.5110μs 36.6089μs 27.3158 KOps/s 27.4964 KOps/s $\color{#d91a1a}-0.66\%$
test_step_mdp_speed[False-True-False-True-False] 0.1600ms 23.8582μs 41.9143 KOps/s 41.9628 KOps/s $\color{#d91a1a}-0.12\%$
test_step_mdp_speed[False-True-False-False-True] 60.1310μs 24.8144μs 40.2992 KOps/s 40.2264 KOps/s $\color{#35bf28}+0.18\%$
test_step_mdp_speed[False-True-False-False-False] 45.4210μs 15.7705μs 63.4095 KOps/s 63.6805 KOps/s $\color{#d91a1a}-0.43\%$
test_step_mdp_speed[False-False-True-True-True] 78.7420μs 38.3664μs 26.0645 KOps/s 26.2067 KOps/s $\color{#d91a1a}-0.54\%$
test_step_mdp_speed[False-False-True-True-False] 0.1379ms 25.3596μs 39.4327 KOps/s 39.1973 KOps/s $\color{#35bf28}+0.60\%$
test_step_mdp_speed[False-False-True-False-True] 90.0920μs 24.2362μs 41.2607 KOps/s 41.1835 KOps/s $\color{#35bf28}+0.19\%$
test_step_mdp_speed[False-False-True-False-False] 40.3810μs 15.4394μs 64.7694 KOps/s 63.7252 KOps/s $\color{#35bf28}+1.64\%$
test_step_mdp_speed[False-False-False-True-True] 50.0210μs 39.1400μs 25.5493 KOps/s 24.4698 KOps/s $\color{#35bf28}+4.41\%$
test_step_mdp_speed[False-False-False-True-False] 75.3820μs 26.8058μs 37.3054 KOps/s 36.7748 KOps/s $\color{#35bf28}+1.44\%$
test_step_mdp_speed[False-False-False-False-True] 0.1017ms 25.7042μs 38.9042 KOps/s 39.2986 KOps/s $\color{#d91a1a}-1.00\%$
test_step_mdp_speed[False-False-False-False-False] 0.1596ms 17.3345μs 57.6884 KOps/s 58.3930 KOps/s $\color{#d91a1a}-1.21\%$
test_values[generalized_advantage_estimate-True-True] 25.5392ms 24.5243ms 40.7758 Ops/s 43.4686 Ops/s $\textbf{\color{#d91a1a}-6.19\%}$
test_values[vec_generalized_advantage_estimate-True-True] 0.1013s 2.8926ms 345.7068 Ops/s 353.9171 Ops/s $\color{#d91a1a}-2.32\%$
test_values[td0_return_estimate-False-False] 87.7820μs 62.5483μs 15.9876 KOps/s 15.2669 KOps/s $\color{#35bf28}+4.72\%$
test_values[td1_return_estimate-False-False] 56.8154ms 52.7598ms 18.9538 Ops/s 19.4307 Ops/s $\color{#d91a1a}-2.45\%$
test_values[vec_td1_return_estimate-False-False] 1.3676ms 1.0486ms 953.6937 Ops/s 963.5616 Ops/s $\color{#d91a1a}-1.02\%$
test_values[td_lambda_return_estimate-True-False] 90.4488ms 86.7102ms 11.5327 Ops/s 12.2485 Ops/s $\textbf{\color{#d91a1a}-5.84\%}$
test_values[vec_td_lambda_return_estimate-True-False] 1.4158ms 1.0429ms 958.8914 Ops/s 951.4904 Ops/s $\color{#35bf28}+0.78\%$
test_gae_speed[generalized_advantage_estimate-False-1-512] 23.6609ms 23.2865ms 42.9434 Ops/s 43.2997 Ops/s $\color{#d91a1a}-0.82\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 0.9196ms 0.6875ms 1.4546 KOps/s 1.4617 KOps/s $\color{#d91a1a}-0.49\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.7807ms 0.6310ms 1.5847 KOps/s 1.5918 KOps/s $\color{#d91a1a}-0.44\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 1.6362ms 1.4430ms 692.9830 Ops/s 696.3226 Ops/s $\color{#d91a1a}-0.48\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 0.8071ms 0.6443ms 1.5520 KOps/s 1.5611 KOps/s $\color{#d91a1a}-0.59\%$
test_dqn_speed 1.5219ms 1.3158ms 760.0039 Ops/s 708.2996 Ops/s $\textbf{\color{#35bf28}+7.30\%}$
test_ddpg_speed 2.9612ms 2.7699ms 361.0191 Ops/s 355.2777 Ops/s $\color{#35bf28}+1.62\%$
test_sac_speed 8.7084ms 7.8170ms 127.9271 Ops/s 113.9567 Ops/s $\textbf{\color{#35bf28}+12.26\%}$
test_redq_speed 0.1030s 11.5522ms 86.5635 Ops/s 94.0328 Ops/s $\textbf{\color{#d91a1a}-7.94\%}$
test_redq_deprec_speed 11.6258ms 11.0846ms 90.2150 Ops/s 89.9196 Ops/s $\color{#35bf28}+0.33\%$
test_td3_speed 8.1237ms 7.8920ms 126.7108 Ops/s 127.0892 Ops/s $\color{#d91a1a}-0.30\%$
test_cql_speed 27.6424ms 24.9064ms 40.1503 Ops/s 38.9000 Ops/s $\color{#35bf28}+3.21\%$
test_a2c_speed 5.7723ms 5.4527ms 183.3953 Ops/s 186.1157 Ops/s $\color{#d91a1a}-1.46\%$
test_ppo_speed 6.1255ms 5.7372ms 174.2999 Ops/s 175.6431 Ops/s $\color{#d91a1a}-0.76\%$
test_reinforce_speed 4.7068ms 4.4436ms 225.0444 Ops/s 225.8963 Ops/s $\color{#d91a1a}-0.38\%$
test_iql_speed 20.3687ms 19.5554ms 51.1369 Ops/s 51.4133 Ops/s $\color{#d91a1a}-0.54\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 4.6855ms 4.4582ms 224.3045 Ops/s 223.8371 Ops/s $\color{#35bf28}+0.21\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.1053s 0.3955ms 2.5282 KOps/s 2.2623 KOps/s $\textbf{\color{#35bf28}+11.75\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.6331ms 0.4061ms 2.4625 KOps/s 2.4316 KOps/s $\color{#35bf28}+1.27\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 4.7577ms 4.4380ms 225.3265 Ops/s 225.8467 Ops/s $\color{#d91a1a}-0.23\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.1402ms 0.4028ms 2.4829 KOps/s 2.3241 KOps/s $\textbf{\color{#35bf28}+6.83\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.6592ms 0.3999ms 2.5005 KOps/s 2.6718 KOps/s $\textbf{\color{#d91a1a}-6.41\%}$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 2.1078ms 1.6337ms 612.1238 Ops/s 630.1330 Ops/s $\color{#d91a1a}-2.86\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 8.0343ms 1.5293ms 653.8880 Ops/s 675.3597 Ops/s $\color{#d91a1a}-3.18\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 4.8590ms 4.5754ms 218.5583 Ops/s 217.6529 Ops/s $\color{#35bf28}+0.42\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 0.8045ms 0.5153ms 1.9407 KOps/s 1.8088 KOps/s $\textbf{\color{#35bf28}+7.29\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 9.2517ms 0.5270ms 1.8974 KOps/s 1.8382 KOps/s $\color{#35bf28}+3.22\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 4.7085ms 4.4720ms 223.6144 Ops/s 223.1390 Ops/s $\color{#35bf28}+0.21\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.5182ms 0.3096ms 3.2302 KOps/s 2.1645 KOps/s $\textbf{\color{#35bf28}+49.24\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 9.3939ms 0.2926ms 3.4176 KOps/s 2.2955 KOps/s $\textbf{\color{#35bf28}+48.88\%}$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 4.7170ms 4.4057ms 226.9792 Ops/s 224.9524 Ops/s $\color{#35bf28}+0.90\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.2315ms 0.3542ms 2.8232 KOps/s 2.6594 KOps/s $\textbf{\color{#35bf28}+6.16\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.5039ms 0.3366ms 2.9705 KOps/s 2.6791 KOps/s $\textbf{\color{#35bf28}+10.88\%}$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 4.7652ms 4.6016ms 217.3163 Ops/s 218.0765 Ops/s $\color{#d91a1a}-0.35\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1.1106ms 0.4587ms 2.1801 KOps/s 2.0568 KOps/s $\textbf{\color{#35bf28}+6.00\%}$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 0.6762ms 0.4387ms 2.2796 KOps/s 2.1751 KOps/s $\color{#35bf28}+4.80\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.1591s 7.7756ms 128.6071 Ops/s 130.4754 Ops/s $\color{#d91a1a}-1.43\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 19.7395ms 14.8794ms 67.2071 Ops/s 65.0278 Ops/s $\color{#35bf28}+3.35\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 0.9847ms 0.8823ms 1.1334 KOps/s 1.0808 KOps/s $\color{#35bf28}+4.87\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 0.1249s 9.5173ms 105.0715 Ops/s 103.6303 Ops/s $\color{#35bf28}+1.39\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 19.6792ms 14.9394ms 66.9370 Ops/s 66.6088 Ops/s $\color{#35bf28}+0.49\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 2.0104ms 0.9182ms 1.0891 KOps/s 1.1167 KOps/s $\color{#d91a1a}-2.47\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 0.1267s 7.2643ms 137.6601 Ops/s 136.6844 Ops/s $\color{#35bf28}+0.71\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 19.9586ms 15.1357ms 66.0690 Ops/s 66.5189 Ops/s $\color{#d91a1a}-0.68\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 2.1051ms 1.0779ms 927.7648 Ops/s 841.7706 Ops/s $\textbf{\color{#35bf28}+10.22\%}$

@vmoens vmoens changed the title [WIP] Fix Brax [BugFix] Fix Brax Jun 18, 2024
@vmoens vmoens added bug Something isn't working Environments Adds or modifies an environment wrapper labels Jun 18, 2024
@vmoens vmoens marked this pull request as ready for review June 18, 2024 12:53
@vmoens vmoens merged commit 45ab9de into main Jun 18, 2024
38 of 51 checks passed
@vmoens vmoens deleted the fix-brax branch June 18, 2024 13:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
bug Something isn't working CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Environments Adds or modifies an environment wrapper
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

[BUG] BRAX: batch size issue
2 participants
@vmoens @facebook-github-bot