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 strict-length for spanning trajectories #1982

Merged
merged 12 commits into from
Mar 4, 2024
Merged

[BugFix] Fix strict-length for spanning trajectories #1982

merged 12 commits into from
Mar 4, 2024

Conversation

vmoens
Copy link
Contributor

@vmoens vmoens commented Feb 29, 2024 •
edited
Loading

Solves #1969 as discussed in the issue

To test the feature:

from torchrl.data import ReplayBuffer, SliceSampler, LazyTensorStorage
from tensordict import TensorDict
import torch


trajectory0 = torch.tensor([3, 3, 0, 1, 1, 1, 2, 2, 2, 3])
trajectory1 = torch.arange(10)
trajectory = torch.stack([trajectory0, trajectory1], -1)

td = TensorDict({"trajectory": trajectory, "steps": torch.arange(10).view(10, 1).expand(10, 2)}, [10, 2]).transpose(-2, -1).clone()
print(td)

rb = ReplayBuffer(
    sampler=SliceSampler(traj_key="trajectory", num_slices=2),
    storage=LazyTensorStorage(20, ndim=2),
    batch_size=6
)

rb.extend(td)

s = rb.sample()
print(s["steps"])
print(s["trajectory"])

Which can give

tensor([9, 0, 1, 6, 7, 8])
tensor([3, 3, 3, 2, 2, 2])

cc @nicklashansen @dasGringuen @aalmuzairee

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Feb 29, 2024 •
edited
Loading

🔗 Helpful Links

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

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

❌ 3 New Failures

As of commit 87883ee with merge base 9337550 (image):

NEW FAILURES - The following jobs have failed:

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 Feb 29, 2024
@vmoens vmoens added the bug Something isn't working label Feb 29, 2024
@vmoens vmoens linked an issue Feb 29, 2024 that may be closed by this pull request
[BUG] SliceSampler breaks when at capacity #1969
Closed
@github-actions GitHub Actions
Copy link

github-actions bot commented Feb 29, 2024 •
edited
Loading

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

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

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_single 61.7763ms 60.8291ms 16.4395 Ops/s 15.8844 Ops/s $\color{#35bf28}+3.49\%$
test_sync 34.7289ms 33.2624ms 30.0640 Ops/s 29.6620 Ops/s $\color{#35bf28}+1.36\%$
test_async 69.7162ms 32.2845ms 30.9746 Ops/s 31.5902 Ops/s $\color{#d91a1a}-1.95\%$
test_simple 0.4794s 0.4250s 2.3528 Ops/s 2.4239 Ops/s $\color{#d91a1a}-2.93\%$
test_transformed 0.6291s 0.5829s 1.7154 Ops/s 1.7308 Ops/s $\color{#d91a1a}-0.89\%$
test_serial 1.4354s 1.3834s 0.7229 Ops/s 0.7233 Ops/s $\color{#d91a1a}-0.05\%$
test_parallel 1.4074s 1.3779s 0.7258 Ops/s 0.7337 Ops/s $\color{#d91a1a}-1.08\%$
test_step_mdp_speed[True-True-True-True-True] 0.1787ms 21.6012μs 46.2937 KOps/s 46.1532 KOps/s $\color{#35bf28}+0.30\%$
test_step_mdp_speed[True-True-True-True-False] 51.0950μs 13.3746μs 74.7685 KOps/s 77.8847 KOps/s $\color{#d91a1a}-4.00\%$
test_step_mdp_speed[True-True-True-False-True] 45.1040μs 12.6683μs 78.9372 KOps/s 79.9987 KOps/s $\color{#d91a1a}-1.33\%$
test_step_mdp_speed[True-True-True-False-False] 47.3890μs 7.6370μs 130.9415 KOps/s 131.6862 KOps/s $\color{#d91a1a}-0.57\%$
test_step_mdp_speed[True-True-False-True-True] 55.0330μs 23.1962μs 43.1105 KOps/s 43.7846 KOps/s $\color{#d91a1a}-1.54\%$
test_step_mdp_speed[True-True-False-True-False] 52.1380μs 14.7089μs 67.9862 KOps/s 69.5996 KOps/s $\color{#d91a1a}-2.32\%$
test_step_mdp_speed[True-True-False-False-True] 39.4440μs 13.8778μs 72.0577 KOps/s 72.8971 KOps/s $\color{#d91a1a}-1.15\%$
test_step_mdp_speed[True-True-False-False-False] 47.7390μs 8.9210μs 112.0949 KOps/s 113.5935 KOps/s $\color{#d91a1a}-1.32\%$
test_step_mdp_speed[True-False-True-True-True] 67.5780μs 24.7176μs 40.4570 KOps/s 41.3482 KOps/s $\color{#d91a1a}-2.16\%$
test_step_mdp_speed[True-False-True-True-False] 88.5450μs 15.9031μs 62.8807 KOps/s 63.9551 KOps/s $\color{#d91a1a}-1.68\%$
test_step_mdp_speed[True-False-True-False-True] 56.3550μs 13.9974μs 71.4418 KOps/s 72.7269 KOps/s $\color{#d91a1a}-1.77\%$
test_step_mdp_speed[True-False-True-False-False] 30.6380μs 9.0941μs 109.9612 KOps/s 114.1353 KOps/s $\color{#d91a1a}-3.66\%$
test_step_mdp_speed[True-False-False-True-True] 49.2320μs 25.7215μs 38.8780 KOps/s 39.2469 KOps/s $\color{#d91a1a}-0.94\%$
test_step_mdp_speed[True-False-False-True-False] 49.1920μs 17.3611μs 57.5999 KOps/s 59.4957 KOps/s $\color{#d91a1a}-3.19\%$
test_step_mdp_speed[True-False-False-False-True] 39.9650μs 15.2793μs 65.4479 KOps/s 65.8953 KOps/s $\color{#d91a1a}-0.68\%$
test_step_mdp_speed[True-False-False-False-False] 38.0710μs 10.3396μs 96.7152 KOps/s 100.2806 KOps/s $\color{#d91a1a}-3.56\%$
test_step_mdp_speed[False-True-True-True-True] 0.1792ms 24.7899μs 40.3390 KOps/s 41.3647 KOps/s $\color{#d91a1a}-2.48\%$
test_step_mdp_speed[False-True-True-True-False] 40.1550μs 15.9864μs 62.5533 KOps/s 63.8482 KOps/s $\color{#d91a1a}-2.03\%$
test_step_mdp_speed[False-True-True-False-True] 45.2650μs 16.5386μs 60.4648 KOps/s 62.5750 KOps/s $\color{#d91a1a}-3.37\%$
test_step_mdp_speed[False-True-True-False-False] 34.8650μs 10.3600μs 96.5247 KOps/s 100.6410 KOps/s $\color{#d91a1a}-4.09\%$
test_step_mdp_speed[False-True-False-True-True] 38.2620μs 26.2010μs 38.1665 KOps/s 38.8912 KOps/s $\color{#d91a1a}-1.86\%$
test_step_mdp_speed[False-True-False-True-False] 54.7320μs 17.3507μs 57.6345 KOps/s 59.5231 KOps/s $\color{#d91a1a}-3.17\%$
test_step_mdp_speed[False-True-False-False-True] 45.8260μs 17.6436μs 56.6777 KOps/s 59.0890 KOps/s $\color{#d91a1a}-4.08\%$
test_step_mdp_speed[False-True-False-False-False] 73.8480μs 11.6900μs 85.5432 KOps/s 89.2058 KOps/s $\color{#d91a1a}-4.11\%$
test_step_mdp_speed[False-False-True-True-True] 65.7930μs 27.2720μs 36.6677 KOps/s 37.5366 KOps/s $\color{#d91a1a}-2.31\%$
test_step_mdp_speed[False-False-True-True-False] 48.6010μs 18.9417μs 52.7935 KOps/s 54.9486 KOps/s $\color{#d91a1a}-3.92\%$
test_step_mdp_speed[False-False-True-False-True] 41.7580μs 17.7038μs 56.4851 KOps/s 58.5821 KOps/s $\color{#d91a1a}-3.58\%$
test_step_mdp_speed[False-False-True-False-False] 33.8540μs 11.6647μs 85.7286 KOps/s 89.5770 KOps/s $\color{#d91a1a}-4.30\%$
test_step_mdp_speed[False-False-False-True-True] 64.7610μs 28.3612μs 35.2594 KOps/s 35.5549 KOps/s $\color{#d91a1a}-0.83\%$
test_step_mdp_speed[False-False-False-True-False] 45.9960μs 19.9776μs 50.0560 KOps/s 52.0839 KOps/s $\color{#d91a1a}-3.89\%$
test_step_mdp_speed[False-False-False-False-True] 59.8410μs 18.5863μs 53.8031 KOps/s 55.0625 KOps/s $\color{#d91a1a}-2.29\%$
test_step_mdp_speed[False-False-False-False-False] 57.3380μs 12.8930μs 77.5613 KOps/s 81.7826 KOps/s $\textbf{\color{#d91a1a}-5.16\%}$
test_values[generalized_advantage_estimate-True-True] 12.2832ms 9.3888ms 106.5097 Ops/s 107.6018 Ops/s $\color{#d91a1a}-1.01\%$
test_values[vec_generalized_advantage_estimate-True-True] 36.5326ms 33.3226ms 30.0096 Ops/s 29.8080 Ops/s $\color{#35bf28}+0.68\%$
test_values[td0_return_estimate-False-False] 0.2275ms 0.1663ms 6.0124 KOps/s 5.9259 KOps/s $\color{#35bf28}+1.46\%$
test_values[td1_return_estimate-False-False] 25.2326ms 22.9215ms 43.6272 Ops/s 43.6317 Ops/s $\color{#d91a1a}-0.01\%$
test_values[vec_td1_return_estimate-False-False] 34.4104ms 33.3473ms 29.9874 Ops/s 29.7476 Ops/s $\color{#35bf28}+0.81\%$
test_values[td_lambda_return_estimate-True-False] 36.2944ms 33.6957ms 29.6774 Ops/s 30.1235 Ops/s $\color{#d91a1a}-1.48\%$
test_values[vec_td_lambda_return_estimate-True-False] 34.6276ms 33.3696ms 29.9674 Ops/s 29.7849 Ops/s $\color{#35bf28}+0.61\%$
test_gae_speed[generalized_advantage_estimate-False-1-512] 11.2937ms 8.1106ms 123.2954 Ops/s 125.6750 Ops/s $\color{#d91a1a}-1.89\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 2.2556ms 2.0166ms 495.8828 Ops/s 526.5559 Ops/s $\textbf{\color{#d91a1a}-5.83\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.4469ms 0.3451ms 2.8977 KOps/s 2.8377 KOps/s $\color{#35bf28}+2.12\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 54.0579ms 42.9868ms 23.2630 Ops/s 24.7822 Ops/s $\textbf{\color{#d91a1a}-6.13\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 4.5980ms 3.0339ms 329.6099 Ops/s 329.8410 Ops/s $\color{#d91a1a}-0.07\%$
test_dqn_speed 6.5513ms 1.3303ms 751.7284 Ops/s 748.7956 Ops/s $\color{#35bf28}+0.39\%$
test_ddpg_speed 3.3143ms 2.6357ms 379.4046 Ops/s 374.1321 Ops/s $\color{#35bf28}+1.41\%$
test_sac_speed 9.9341ms 8.2151ms 121.7265 Ops/s 120.5363 Ops/s $\color{#35bf28}+0.99\%$
test_redq_speed 14.0811ms 13.0632ms 76.5510 Ops/s 76.0597 Ops/s $\color{#35bf28}+0.65\%$
test_redq_deprec_speed 14.5128ms 13.3485ms 74.9145 Ops/s 76.1412 Ops/s $\color{#d91a1a}-1.61\%$
test_td3_speed 14.9264ms 8.1293ms 123.0125 Ops/s 122.8730 Ops/s $\color{#35bf28}+0.11\%$
test_cql_speed 38.5052ms 36.4858ms 27.4079 Ops/s 27.6977 Ops/s $\color{#d91a1a}-1.05\%$
test_a2c_speed 7.9209ms 7.2960ms 137.0618 Ops/s 134.7678 Ops/s $\color{#35bf28}+1.70\%$
test_ppo_speed 9.0499ms 7.6472ms 130.7665 Ops/s 131.9043 Ops/s $\color{#d91a1a}-0.86\%$
test_reinforce_speed 7.4983ms 6.5935ms 151.6639 Ops/s 149.9252 Ops/s $\color{#35bf28}+1.16\%$
test_iql_speed 34.3351ms 32.5191ms 30.7511 Ops/s 29.9846 Ops/s $\color{#35bf28}+2.56\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 3.3065ms 2.2311ms 448.2177 Ops/s 433.9051 Ops/s $\color{#35bf28}+3.30\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 1.0487ms 0.4939ms 2.0245 KOps/s 2.0028 KOps/s $\color{#35bf28}+1.08\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.8075ms 0.4687ms 2.1334 KOps/s 2.1029 KOps/s $\color{#35bf28}+1.45\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 2.6639ms 2.3064ms 433.5849 Ops/s 427.1328 Ops/s $\color{#35bf28}+1.51\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.0280ms 0.4888ms 2.0460 KOps/s 2.0120 KOps/s $\color{#35bf28}+1.69\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.7972ms 0.4727ms 2.1157 KOps/s 2.0968 KOps/s $\color{#35bf28}+0.90\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 1.7309ms 1.2423ms 804.9536 Ops/s 910.8393 Ops/s $\textbf{\color{#d91a1a}-11.63\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 1.7937ms 1.1725ms 852.8734 Ops/s 942.2856 Ops/s $\textbf{\color{#d91a1a}-9.49\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 2.7147ms 2.3797ms 420.2249 Ops/s 413.1225 Ops/s $\color{#35bf28}+1.72\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1.0284ms 0.6066ms 1.6487 KOps/s 1.6097 KOps/s $\color{#35bf28}+2.42\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 0.8808ms 0.5834ms 1.7140 KOps/s 1.6925 KOps/s $\color{#35bf28}+1.27\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 2.6452ms 2.3331ms 428.6204 Ops/s 447.3713 Ops/s $\color{#d91a1a}-4.19\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.9143ms 0.5028ms 1.9889 KOps/s 1.6382 KOps/s $\textbf{\color{#35bf28}+21.41\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.7205ms 0.4757ms 2.1020 KOps/s 2.0982 KOps/s $\color{#35bf28}+0.18\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 2.7424ms 2.3249ms 430.1170 Ops/s 439.8559 Ops/s $\color{#d91a1a}-2.21\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 0.5915ms 0.4955ms 2.0181 KOps/s 2.0213 KOps/s $\color{#d91a1a}-0.16\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3.5948ms 0.4752ms 2.1044 KOps/s 2.1139 KOps/s $\color{#d91a1a}-0.45\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 2.7273ms 2.3678ms 422.3291 Ops/s 417.5306 Ops/s $\color{#35bf28}+1.15\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1.1681ms 0.6158ms 1.6240 KOps/s 1.3147 KOps/s $\textbf{\color{#35bf28}+23.53\%}$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 0.8987ms 0.5913ms 1.6913 KOps/s 1.6680 KOps/s $\color{#35bf28}+1.40\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.1041s 7.6160ms 131.3027 Ops/s 182.9372 Ops/s $\textbf{\color{#d91a1a}-28.23\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 14.4618ms 11.9821ms 83.4578 Ops/s 82.7165 Ops/s $\color{#35bf28}+0.90\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 1.7042ms 1.1048ms 905.1380 Ops/s 969.6981 Ops/s $\textbf{\color{#d91a1a}-6.66\%}$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 0.1041s 7.6071ms 131.4568 Ops/s 136.7237 Ops/s $\color{#d91a1a}-3.85\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 14.5850ms 12.0947ms 82.6810 Ops/s 84.2949 Ops/s $\color{#d91a1a}-1.91\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 3.6485ms 1.1300ms 884.9714 Ops/s 924.3758 Ops/s $\color{#d91a1a}-4.26\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 95.6867ms 5.8573ms 170.7257 Ops/s 133.2171 Ops/s $\textbf{\color{#35bf28}+28.16\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 14.4202ms 12.3489ms 80.9792 Ops/s 81.7523 Ops/s $\color{#d91a1a}-0.95\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 87.5993ms 3.0873ms 323.9065 Ops/s 647.7250 Ops/s $\textbf{\color{#d91a1a}-49.99\%}$

@github-actions GitHub Actions
Copy link

github-actions bot commented Feb 29, 2024 •
edited
Loading

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

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

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_single 0.1108s 0.1098s 9.1109 Ops/s 8.9017 Ops/s $\color{#35bf28}+2.35\%$
test_sync 95.3821ms 95.1831ms 10.5061 Ops/s 10.5402 Ops/s $\color{#d91a1a}-0.32\%$
test_async 0.1798s 91.0513ms 10.9828 Ops/s 10.7635 Ops/s $\color{#35bf28}+2.04\%$
test_single_pixels 0.1207s 0.1206s 8.2943 Ops/s 8.1444 Ops/s $\color{#35bf28}+1.84\%$
test_sync_pixels 83.0314ms 81.3503ms 12.2925 Ops/s 12.4028 Ops/s $\color{#d91a1a}-0.89\%$
test_async_pixels 0.1532s 65.1679ms 15.3450 Ops/s 15.0274 Ops/s $\color{#35bf28}+2.11\%$
test_simple 0.8592s 0.8007s 1.2490 Ops/s 1.2325 Ops/s $\color{#35bf28}+1.34\%$
test_transformed 1.0717s 1.0160s 0.9843 Ops/s 0.9646 Ops/s $\color{#35bf28}+2.04\%$
test_serial 2.3796s 2.3229s 0.4305 Ops/s 0.4195 Ops/s $\color{#35bf28}+2.63\%$
test_parallel 2.0950s 2.0289s 0.4929 Ops/s 0.4778 Ops/s $\color{#35bf28}+3.15\%$
test_step_mdp_speed[True-True-True-True-True] 78.6820μs 32.8739μs 30.4192 KOps/s 29.2275 KOps/s $\color{#35bf28}+4.08\%$
test_step_mdp_speed[True-True-True-True-False] 87.1610μs 19.6432μs 50.9082 KOps/s 49.8784 KOps/s $\color{#35bf28}+2.06\%$
test_step_mdp_speed[True-True-True-False-True] 42.8010μs 18.4023μs 54.3412 KOps/s 52.1801 KOps/s $\color{#35bf28}+4.14\%$
test_step_mdp_speed[True-True-True-False-False] 25.1910μs 11.0996μs 90.0932 KOps/s 87.7226 KOps/s $\color{#35bf28}+2.70\%$
test_step_mdp_speed[True-True-False-True-True] 61.0710μs 34.6520μs 28.8584 KOps/s 27.8188 KOps/s $\color{#35bf28}+3.74\%$
test_step_mdp_speed[True-True-False-True-False] 0.1010ms 21.2773μs 46.9984 KOps/s 45.1503 KOps/s $\color{#35bf28}+4.09\%$
test_step_mdp_speed[True-True-False-False-True] 42.3910μs 20.2928μs 49.2786 KOps/s 47.5096 KOps/s $\color{#35bf28}+3.72\%$
test_step_mdp_speed[True-True-False-False-False] 38.1710μs 12.9590μs 77.1666 KOps/s 75.8213 KOps/s $\color{#35bf28}+1.77\%$
test_step_mdp_speed[True-False-True-True-True] 0.1022ms 36.3643μs 27.4995 KOps/s 26.6601 KOps/s $\color{#35bf28}+3.15\%$
test_step_mdp_speed[True-False-True-True-False] 53.7320μs 23.6780μs 42.2332 KOps/s 41.4471 KOps/s $\color{#35bf28}+1.90\%$
test_step_mdp_speed[True-False-True-False-True] 78.2810μs 19.9640μs 50.0902 KOps/s 47.3988 KOps/s $\textbf{\color{#35bf28}+5.68\%}$
test_step_mdp_speed[True-False-True-False-False] 32.2410μs 12.8482μs 77.8319 KOps/s 75.1817 KOps/s $\color{#35bf28}+3.53\%$
test_step_mdp_speed[True-False-False-True-True] 98.3320μs 38.2182μs 26.1656 KOps/s 25.3310 KOps/s $\color{#35bf28}+3.29\%$
test_step_mdp_speed[True-False-False-True-False] 49.7510μs 25.1371μs 39.7818 KOps/s 38.4522 KOps/s $\color{#35bf28}+3.46\%$
test_step_mdp_speed[True-False-False-False-True] 42.2710μs 22.0146μs 45.4244 KOps/s 43.3996 KOps/s $\color{#35bf28}+4.67\%$
test_step_mdp_speed[True-False-False-False-False] 59.8610μs 15.1280μs 66.1028 KOps/s 65.6548 KOps/s $\color{#35bf28}+0.68\%$
test_step_mdp_speed[False-True-True-True-True] 59.8500μs 36.1870μs 27.6343 KOps/s 26.2713 KOps/s $\textbf{\color{#35bf28}+5.19\%}$
test_step_mdp_speed[False-True-True-True-False] 85.7920μs 23.5119μs 42.5316 KOps/s 41.4550 KOps/s $\color{#35bf28}+2.60\%$
test_step_mdp_speed[False-True-True-False-True] 45.7810μs 23.8176μs 41.9857 KOps/s 40.2298 KOps/s $\color{#35bf28}+4.36\%$
test_step_mdp_speed[False-True-True-False-False] 82.4620μs 14.7497μs 67.7981 KOps/s 64.4164 KOps/s $\textbf{\color{#35bf28}+5.25\%}$
test_step_mdp_speed[False-True-False-True-True] 71.4710μs 38.3370μs 26.0845 KOps/s 25.3259 KOps/s $\color{#35bf28}+3.00\%$
test_step_mdp_speed[False-True-False-True-False] 59.0710μs 25.4246μs 39.3320 KOps/s 37.8655 KOps/s $\color{#35bf28}+3.87\%$
test_step_mdp_speed[False-True-False-False-True] 64.1300μs 26.1339μs 38.2644 KOps/s 37.0030 KOps/s $\color{#35bf28}+3.41\%$
test_step_mdp_speed[False-True-False-False-False] 41.4010μs 16.9644μs 58.9469 KOps/s 57.1687 KOps/s $\color{#35bf28}+3.11\%$
test_step_mdp_speed[False-False-True-True-True] 64.2910μs 39.7299μs 25.1700 KOps/s 23.7900 KOps/s $\textbf{\color{#35bf28}+5.80\%}$
test_step_mdp_speed[False-False-True-True-False] 50.4110μs 27.2890μs 36.6448 KOps/s 35.2273 KOps/s $\color{#35bf28}+4.02\%$
test_step_mdp_speed[False-False-True-False-True] 47.0010μs 26.0333μs 38.4124 KOps/s 36.7594 KOps/s $\color{#35bf28}+4.50\%$
test_step_mdp_speed[False-False-True-False-False] 64.3300μs 16.6860μs 59.9306 KOps/s 57.0068 KOps/s $\textbf{\color{#35bf28}+5.13\%}$
test_step_mdp_speed[False-False-False-True-True] 68.1110μs 41.2334μs 24.2522 KOps/s 23.3870 KOps/s $\color{#35bf28}+3.70\%$
test_step_mdp_speed[False-False-False-True-False] 58.6210μs 28.9421μs 34.5517 KOps/s 33.6870 KOps/s $\color{#35bf28}+2.57\%$
test_step_mdp_speed[False-False-False-False-True] 70.6420μs 27.5436μs 36.3060 KOps/s 35.1185 KOps/s $\color{#35bf28}+3.38\%$
test_step_mdp_speed[False-False-False-False-False] 85.5220μs 18.7249μs 53.4049 KOps/s 51.9513 KOps/s $\color{#35bf28}+2.80\%$
test_values[generalized_advantage_estimate-True-True] 24.1873ms 23.7589ms 42.0895 Ops/s 41.5592 Ops/s $\color{#35bf28}+1.28\%$
test_values[vec_generalized_advantage_estimate-True-True] 92.7487ms 3.4100ms 293.2575 Ops/s 292.2156 Ops/s $\color{#35bf28}+0.36\%$
test_values[td0_return_estimate-False-False] 0.1019ms 62.7879μs 15.9266 KOps/s 15.9059 KOps/s $\color{#35bf28}+0.13\%$
test_values[td1_return_estimate-False-False] 51.5202ms 51.0300ms 19.5963 Ops/s 19.3106 Ops/s $\color{#35bf28}+1.48\%$
test_values[vec_td1_return_estimate-False-False] 2.0871ms 1.7552ms 569.7391 Ops/s 569.2589 Ops/s $\color{#35bf28}+0.08\%$
test_values[td_lambda_return_estimate-True-False] 82.3876ms 81.5317ms 12.2652 Ops/s 12.1217 Ops/s $\color{#35bf28}+1.18\%$
test_values[vec_td_lambda_return_estimate-True-False] 2.1075ms 1.7467ms 572.5013 Ops/s 571.8970 Ops/s $\color{#35bf28}+0.11\%$
test_gae_speed[generalized_advantage_estimate-False-1-512] 22.8109ms 22.2573ms 44.9292 Ops/s 43.4271 Ops/s $\color{#35bf28}+3.46\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 0.8776ms 0.6800ms 1.4706 KOps/s 1.4185 KOps/s $\color{#35bf28}+3.68\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.7659ms 0.6318ms 1.5829 KOps/s 1.5209 KOps/s $\color{#35bf28}+4.07\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 1.5632ms 1.4388ms 695.0381 Ops/s 694.0571 Ops/s $\color{#35bf28}+0.14\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 0.9075ms 0.6543ms 1.5284 KOps/s 1.4967 KOps/s $\color{#35bf28}+2.12\%$
test_dqn_speed 7.9813ms 1.4300ms 699.3155 Ops/s 691.0915 Ops/s $\color{#35bf28}+1.19\%$
test_ddpg_speed 3.0185ms 2.7091ms 369.1320 Ops/s 365.1332 Ops/s $\color{#35bf28}+1.10\%$
test_sac_speed 8.8169ms 7.9433ms 125.8930 Ops/s 124.5979 Ops/s $\color{#35bf28}+1.04\%$
test_redq_speed 12.6450ms 10.1796ms 98.2357 Ops/s 90.7407 Ops/s $\textbf{\color{#35bf28}+8.26\%}$
test_redq_deprec_speed 11.6762ms 11.1218ms 89.9134 Ops/s 91.7028 Ops/s $\color{#d91a1a}-1.95\%$
test_td3_speed 15.3955ms 7.8887ms 126.7642 Ops/s 125.6531 Ops/s $\color{#35bf28}+0.88\%$
test_cql_speed 25.6856ms 24.9482ms 40.0831 Ops/s 40.2163 Ops/s $\color{#d91a1a}-0.33\%$
test_a2c_speed 5.7708ms 5.5347ms 180.6794 Ops/s 181.1693 Ops/s $\color{#d91a1a}-0.27\%$
test_ppo_speed 6.1624ms 5.8913ms 169.7424 Ops/s 170.3590 Ops/s $\color{#d91a1a}-0.36\%$
test_reinforce_speed 5.6109ms 4.5256ms 220.9629 Ops/s 221.9980 Ops/s $\color{#d91a1a}-0.47\%$
test_iql_speed 20.5993ms 19.6334ms 50.9337 Ops/s 50.6851 Ops/s $\color{#35bf28}+0.49\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 3.0928ms 2.8522ms 350.6061 Ops/s 345.7326 Ops/s $\color{#35bf28}+1.41\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 90.0423ms 0.6103ms 1.6386 KOps/s 1.8355 KOps/s $\textbf{\color{#d91a1a}-10.72\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.7082ms 0.5177ms 1.9317 KOps/s 1.9184 KOps/s $\color{#35bf28}+0.70\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 3.1713ms 2.8721ms 348.1757 Ops/s 343.5983 Ops/s $\color{#35bf28}+1.33\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.3479ms 0.5346ms 1.8705 KOps/s 1.8779 KOps/s $\color{#d91a1a}-0.39\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.6781ms 0.5111ms 1.9567 KOps/s 1.9557 KOps/s $\color{#35bf28}+0.05\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 1.6556ms 1.4641ms 683.0064 Ops/s 774.4809 Ops/s $\textbf{\color{#d91a1a}-11.81\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 1.5778ms 1.4043ms 712.1078 Ops/s 801.7071 Ops/s $\textbf{\color{#d91a1a}-11.18\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 3.1730ms 2.9634ms 337.4496 Ops/s 333.7975 Ops/s $\color{#35bf28}+1.09\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 0.8538ms 0.6574ms 1.5211 KOps/s 1.5104 KOps/s $\color{#35bf28}+0.71\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 4.3141ms 0.6342ms 1.5769 KOps/s 1.5596 KOps/s $\color{#35bf28}+1.11\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 3.1788ms 2.8659ms 348.9262 Ops/s 347.8879 Ops/s $\color{#35bf28}+0.30\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.7035ms 0.5374ms 1.8608 KOps/s 1.8443 KOps/s $\color{#35bf28}+0.89\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 4.1726ms 0.5181ms 1.9300 KOps/s 1.9131 KOps/s $\color{#35bf28}+0.88\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 3.0988ms 2.8503ms 350.8433 Ops/s 345.1795 Ops/s $\color{#35bf28}+1.64\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.4100ms 0.5320ms 1.8797 KOps/s 1.8590 KOps/s $\color{#35bf28}+1.12\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.7605ms 0.5182ms 1.9296 KOps/s 1.9193 KOps/s $\color{#35bf28}+0.54\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 3.1802ms 3.0005ms 333.2725 Ops/s 332.7450 Ops/s $\color{#35bf28}+0.16\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 0.9138ms 0.6621ms 1.5102 KOps/s 1.4988 KOps/s $\color{#35bf28}+0.77\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 4.5676ms 0.6424ms 1.5566 KOps/s 1.5573 KOps/s $\color{#d91a1a}-0.05\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.1177s 8.9325ms 111.9502 Ops/s 112.2851 Ops/s $\color{#d91a1a}-0.30\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 16.5396ms 14.2940ms 69.9595 Ops/s 67.8895 Ops/s $\color{#35bf28}+3.05\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 2.2110ms 1.0733ms 931.6858 Ops/s 952.4097 Ops/s $\color{#d91a1a}-2.18\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 0.1015s 8.5943ms 116.3559 Ops/s 115.8142 Ops/s $\color{#35bf28}+0.47\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 16.3305ms 14.3245ms 69.8104 Ops/s 67.5395 Ops/s $\color{#35bf28}+3.36\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1.1220ms 1.0309ms 970.0336 Ops/s 850.4711 Ops/s $\textbf{\color{#35bf28}+14.06\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 0.1010s 7.0125ms 142.6024 Ops/s 141.7408 Ops/s $\color{#35bf28}+0.61\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 17.3522ms 14.8691ms 67.2536 Ops/s 66.2482 Ops/s $\color{#35bf28}+1.52\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 1.5268ms 1.4056ms 711.4154 Ops/s 672.1327 Ops/s $\textbf{\color{#35bf28}+5.84\%}$

@vmoens
Copy link
Contributor Author

vmoens commented Feb 29, 2024

This is much more complex than I thought haha
I'm pretty sure I can get it to work efficiently but that will require some thinking!
Thanks for your patience guys

@nicklashansen
Copy link

This is much more complex than I thought haha I'm pretty sure I can get it to work efficiently but that will require some thinking! Thanks for your patience guys

@vmoens no worries, we really appreciate it! let us know if you need anything from us :-)

@vmoens
Copy link
Contributor Author

vmoens commented Mar 1, 2024

That was fun!
I think I got it right now.
There are various cases that I needed to handle (multiple dims, storage at capacity or not, raising errors if not valid traj can be found).

If you guys want to give it a look I'll leave this PR open for a little longer

@vmoens vmoens marked this pull request as ready for review March 1, 2024 19:29
@vmoens vmoens merged commit 6fb16a2 into main Mar 4, 2024
63 of 66 checks passed
@vmoens vmoens deleted the fix-strict-length branch March 4, 2024 14:09
@nicklashansen
Copy link

@vmoens sorry that it took a little while to get back to you, but this looks great! we'll run some tests on our end and see if it resolves the issue :-)

vmoens added a commit that referenced this pull request Mar 25, 2024
@vmoens
[BugFix] Fix strict-length for spanning trajectories (#1982)
bff93f4 
(cherry picked from commit 6fb16a2)
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.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

[BUG] SliceSampler breaks when at capacity
3 participants
@vmoens @nicklashansen @facebook-github-bot