Skip to content

[BugFix] TensorStorage key filtering#3323

Merged
vmoens merged 4 commits into
gh/vmoens/191/basefrom
gh/vmoens/191/head
Jan 13, 2026
Merged

[BugFix] TensorStorage key filtering#3323
vmoens merged 4 commits into
gh/vmoens/191/basefrom
gh/vmoens/191/head

Conversation

@vmoens

@vmoens vmoens commented Jan 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

Implement _storage_keys caching and key filtering in TensorStorage.
This handles cases where the incoming data has extra keys (e.g. from a
policy) that aren't present in the storage. Also improved error messages
for locked storage key mismatches.

@vmoens
Update
c4c1c64 
[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jan 12, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 4 New Failures, 1 Unrelated Failure

As of commit a2c13e0 with merge base 0a98e17 (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.

@meta-cla meta-cla 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 Jan 12, 2026
This was referenced Jan 12, 2026
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
@github-actions

github-actions Bot commented Jan 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

Total Benchmarks: 164. Improved: $\large\color{#35bf28}18$. Worsened: $\large\color{#d91a1a}18$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_tensor_to_bytestream_speed[pickle] 83.2815μs 81.7533μs 12.2319 KOps/s 11.6893 KOps/s $\color{#35bf28}+4.64\%$
test_tensor_to_bytestream_speed[torch.save] 0.1408ms 0.1404ms 7.1234 KOps/s 6.7577 KOps/s $\textbf{\color{#35bf28}+5.41\%}$
test_tensor_to_bytestream_speed[untyped_storage] 0.1258s 0.1256s 7.9648 Ops/s 7.2372 Ops/s $\textbf{\color{#35bf28}+10.05\%}$
test_tensor_to_bytestream_speed[numpy] 2.7882μs 2.7842μs 359.1709 KOps/s 367.1716 KOps/s $\color{#d91a1a}-2.18\%$
test_tensor_to_bytestream_speed[safetensors] 40.0661μs 39.5882μs 25.2600 KOps/s 24.3977 KOps/s $\color{#35bf28}+3.53\%$
test_simple 0.5644s 0.5600s 1.7857 Ops/s 1.7267 Ops/s $\color{#35bf28}+3.42\%$
test_transformed 1.1291s 1.1264s 0.8878 Ops/s 0.8617 Ops/s $\color{#35bf28}+3.02\%$
test_serial 1.6915s 1.6871s 0.5927 Ops/s 0.5814 Ops/s $\color{#35bf28}+1.95\%$
test_parallel 1.1256s 1.1018s 0.9076 Ops/s 0.8616 Ops/s $\textbf{\color{#35bf28}+5.34\%}$
test_step_mdp_speed[True-True-True-True-True] 0.1555ms 44.5507μs 22.4464 KOps/s 22.6079 KOps/s $\color{#d91a1a}-0.71\%$
test_step_mdp_speed[True-True-True-True-False] 61.1930μs 24.5419μs 40.7466 KOps/s 39.9791 KOps/s $\color{#35bf28}+1.92\%$
test_step_mdp_speed[True-True-True-False-True] 61.3430μs 24.8107μs 40.3052 KOps/s 39.4351 KOps/s $\color{#35bf28}+2.21\%$
test_step_mdp_speed[True-True-True-False-False] 48.5120μs 13.6282μs 73.3771 KOps/s 73.2172 KOps/s $\color{#35bf28}+0.22\%$
test_step_mdp_speed[True-True-False-True-True] 96.2450μs 47.2649μs 21.1574 KOps/s 21.1741 KOps/s $\color{#d91a1a}-0.08\%$
test_step_mdp_speed[True-True-False-True-False] 65.4740μs 27.2859μs 36.6490 KOps/s 36.0789 KOps/s $\color{#35bf28}+1.58\%$
test_step_mdp_speed[True-True-False-False-True] 75.6140μs 28.1528μs 35.5204 KOps/s 35.5708 KOps/s $\color{#d91a1a}-0.14\%$
test_step_mdp_speed[True-True-False-False-False] 59.2530μs 16.3623μs 61.1160 KOps/s 61.3935 KOps/s $\color{#d91a1a}-0.45\%$
test_step_mdp_speed[True-False-True-True-True] 90.3240μs 50.5451μs 19.7843 KOps/s 19.8177 KOps/s $\color{#d91a1a}-0.17\%$
test_step_mdp_speed[True-False-True-True-False] 69.5630μs 30.0166μs 33.3149 KOps/s 32.9138 KOps/s $\color{#35bf28}+1.22\%$
test_step_mdp_speed[True-False-True-False-True] 63.6130μs 27.8293μs 35.9334 KOps/s 35.1733 KOps/s $\color{#35bf28}+2.16\%$
test_step_mdp_speed[True-False-True-False-False] 56.1830μs 16.7494μs 59.7037 KOps/s 59.7930 KOps/s $\color{#d91a1a}-0.15\%$
test_step_mdp_speed[True-False-False-True-True] 0.1007ms 52.9383μs 18.8899 KOps/s 18.6513 KOps/s $\color{#35bf28}+1.28\%$
test_step_mdp_speed[True-False-False-True-False] 68.8930μs 33.5119μs 29.8401 KOps/s 30.1461 KOps/s $\color{#d91a1a}-1.02\%$
test_step_mdp_speed[True-False-False-False-True] 64.2630μs 30.1572μs 33.1595 KOps/s 32.2723 KOps/s $\color{#35bf28}+2.75\%$
test_step_mdp_speed[True-False-False-False-False] 57.1720μs 19.3048μs 51.8005 KOps/s 51.8245 KOps/s $\color{#d91a1a}-0.05\%$
test_step_mdp_speed[False-True-True-True-True] 93.0640μs 50.0170μs 19.9932 KOps/s 19.6178 KOps/s $\color{#35bf28}+1.91\%$
test_step_mdp_speed[False-True-True-True-False] 70.9940μs 29.7865μs 33.5723 KOps/s 33.0083 KOps/s $\color{#35bf28}+1.71\%$
test_step_mdp_speed[False-True-True-False-True] 2.3602ms 31.8369μs 31.4101 KOps/s 31.4316 KOps/s $\color{#d91a1a}-0.07\%$
test_step_mdp_speed[False-True-True-False-False] 52.3120μs 18.2274μs 54.8624 KOps/s 54.2007 KOps/s $\color{#35bf28}+1.22\%$
test_step_mdp_speed[False-True-False-True-True] 95.4340μs 52.5364μs 19.0344 KOps/s 18.4978 KOps/s $\color{#35bf28}+2.90\%$
test_step_mdp_speed[False-True-False-True-False] 64.2530μs 32.7880μs 30.4990 KOps/s 29.5370 KOps/s $\color{#35bf28}+3.26\%$
test_step_mdp_speed[False-True-False-False-True] 81.9240μs 34.2581μs 29.1901 KOps/s 29.1499 KOps/s $\color{#35bf28}+0.14\%$
test_step_mdp_speed[False-True-False-False-False] 55.3130μs 20.6693μs 48.3808 KOps/s 48.1494 KOps/s $\color{#35bf28}+0.48\%$
test_step_mdp_speed[False-False-True-True-True] 96.6740μs 56.6355μs 17.6568 KOps/s 17.8917 KOps/s $\color{#d91a1a}-1.31\%$
test_step_mdp_speed[False-False-True-True-False] 73.1940μs 36.2412μs 27.5929 KOps/s 28.0166 KOps/s $\color{#d91a1a}-1.51\%$
test_step_mdp_speed[False-False-True-False-True] 74.2030μs 34.0750μs 29.3470 KOps/s 29.5101 KOps/s $\color{#d91a1a}-0.55\%$
test_step_mdp_speed[False-False-True-False-False] 58.4130μs 20.7987μs 48.0800 KOps/s 47.2135 KOps/s $\color{#35bf28}+1.84\%$
test_step_mdp_speed[False-False-False-True-True] 99.0050μs 57.4627μs 17.4026 KOps/s 17.1520 KOps/s $\color{#35bf28}+1.46\%$
test_step_mdp_speed[False-False-False-True-False] 73.1730μs 38.2845μs 26.1202 KOps/s 26.2490 KOps/s $\color{#d91a1a}-0.49\%$
test_step_mdp_speed[False-False-False-False-True] 77.4740μs 36.2243μs 27.6058 KOps/s 27.4219 KOps/s $\color{#35bf28}+0.67\%$
test_step_mdp_speed[False-False-False-False-False] 58.7030μs 23.0086μs 43.4621 KOps/s 42.8228 KOps/s $\color{#35bf28}+1.49\%$
test_non_tensor_env_rollout_speed[1000-single-True] 0.8766s 0.7738s 1.2922 Ops/s 1.2860 Ops/s $\color{#35bf28}+0.49\%$
test_non_tensor_env_rollout_speed[1000-single-False] 0.7277s 0.6301s 1.5871 Ops/s 1.5641 Ops/s $\color{#35bf28}+1.47\%$
test_non_tensor_env_rollout_speed[1000-serial-no-buffers-True] 1.7609s 1.6770s 0.5963 Ops/s 0.5919 Ops/s $\color{#35bf28}+0.74\%$
test_non_tensor_env_rollout_speed[1000-serial-no-buffers-False] 1.5301s 1.4554s 0.6871 Ops/s 0.6838 Ops/s $\color{#35bf28}+0.49\%$
test_non_tensor_env_rollout_speed[1000-serial-buffers-True] 2.0177s 1.9338s 0.5171 Ops/s 0.5154 Ops/s $\color{#35bf28}+0.34\%$
test_non_tensor_env_rollout_speed[1000-serial-buffers-False] 1.7873s 1.7064s 0.5860 Ops/s 0.5779 Ops/s $\color{#35bf28}+1.40\%$
test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True] 4.7804s 4.7512s 0.2105 Ops/s 0.2121 Ops/s $\color{#d91a1a}-0.77\%$
test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-False] 4.6223s 4.5143s 0.2215 Ops/s 0.2197 Ops/s $\color{#35bf28}+0.84\%$
test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] 2.0299s 1.9440s 0.5144 Ops/s 0.5117 Ops/s $\color{#35bf28}+0.54\%$
test_non_tensor_env_rollout_speed[1000-parallel-buffers-False] 1.8075s 1.6828s 0.5943 Ops/s 0.5833 Ops/s $\color{#35bf28}+1.88\%$
test_values[generalized_advantage_estimate-True-True] 10.7010ms 10.5401ms 94.8759 Ops/s 88.4175 Ops/s $\textbf{\color{#35bf28}+7.30\%}$
test_values[vec_generalized_advantage_estimate-True-True] 19.6939ms 17.6191ms 56.7566 Ops/s 90.7621 Ops/s $\textbf{\color{#d91a1a}-37.47\%}$
test_values[td0_return_estimate-False-False] 0.2294ms 0.1320ms 7.5780 KOps/s 7.8900 KOps/s $\color{#d91a1a}-3.95\%$
test_values[td1_return_estimate-False-False] 29.4833ms 28.6270ms 34.9321 Ops/s 32.8140 Ops/s $\textbf{\color{#35bf28}+6.45\%}$
test_values[vec_td1_return_estimate-False-False] 18.0557ms 17.6622ms 56.6180 Ops/s 90.0413 Ops/s $\textbf{\color{#d91a1a}-37.12\%}$
test_values[td_lambda_return_estimate-True-False] 42.8339ms 42.0778ms 23.7655 Ops/s 22.2590 Ops/s $\textbf{\color{#35bf28}+6.77\%}$
test_values[vec_td_lambda_return_estimate-True-False] 18.5846ms 17.6740ms 56.5803 Ops/s 90.3119 Ops/s $\textbf{\color{#d91a1a}-37.35\%}$
test_gae_speed[generalized_advantage_estimate-False-1-512] 11.8676ms 9.4888ms 105.3876 Ops/s 98.1016 Ops/s $\textbf{\color{#35bf28}+7.43\%}$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 1.7615ms 1.5698ms 637.0387 Ops/s 647.7628 Ops/s $\color{#d91a1a}-1.66\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.5575ms 0.4292ms 2.3297 KOps/s 2.3112 KOps/s $\color{#35bf28}+0.80\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 34.9660ms 34.1497ms 29.2829 Ops/s 39.6375 Ops/s $\textbf{\color{#d91a1a}-26.12\%}$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 1.8513ms 1.7111ms 584.4207 Ops/s 580.0851 Ops/s $\color{#35bf28}+0.75\%$
test_dqn_speed[False-None] 1.5282ms 1.4228ms 702.8395 Ops/s 710.1737 Ops/s $\color{#d91a1a}-1.03\%$
test_dqn_speed[False-backward] 2.0901ms 1.9499ms 512.8404 Ops/s 518.3823 Ops/s $\color{#d91a1a}-1.07\%$
test_dqn_speed[True-None] 0.7468ms 0.5371ms 1.8618 KOps/s 1.7794 KOps/s $\color{#35bf28}+4.63\%$
test_dqn_speed[True-backward] 1.0280ms 0.9819ms 1.0184 KOps/s 992.1439 Ops/s $\color{#35bf28}+2.65\%$
test_dqn_speed[reduce-overhead-None] 0.6546ms 0.5295ms 1.8885 KOps/s 1.8228 KOps/s $\color{#35bf28}+3.61\%$
test_dqn_speed[reduce-overhead-backward] 1.0071ms 0.9654ms 1.0358 KOps/s 1.0118 KOps/s $\color{#35bf28}+2.38\%$
test_ddpg_speed[False-None] 3.2265ms 2.8933ms 345.6268 Ops/s 342.0034 Ops/s $\color{#35bf28}+1.06\%$
test_ddpg_speed[False-backward] 4.2074ms 4.1137ms 243.0917 Ops/s 235.1494 Ops/s $\color{#35bf28}+3.38\%$
test_ddpg_speed[True-None] 1.5010ms 1.3942ms 717.2394 Ops/s 711.5027 Ops/s $\color{#35bf28}+0.81\%$
test_ddpg_speed[True-backward] 2.4361ms 2.3880ms 418.7599 Ops/s 396.2417 Ops/s $\textbf{\color{#35bf28}+5.68\%}$
test_ddpg_speed[reduce-overhead-None] 1.4911ms 1.3873ms 720.8144 Ops/s 709.8456 Ops/s $\color{#35bf28}+1.55\%$
test_ddpg_speed[reduce-overhead-backward] 2.4326ms 2.3859ms 419.1317 Ops/s 412.0448 Ops/s $\color{#35bf28}+1.72\%$
test_sac_speed[False-None] 8.5978ms 8.1130ms 123.2590 Ops/s 123.4230 Ops/s $\color{#d91a1a}-0.13\%$
test_sac_speed[False-backward] 11.8354ms 11.3199ms 88.3399 Ops/s 87.8295 Ops/s $\color{#35bf28}+0.58\%$
test_sac_speed[True-None] 2.8749ms 2.1634ms 462.2418 Ops/s 449.3358 Ops/s $\color{#35bf28}+2.87\%$
test_sac_speed[True-backward] 4.2008ms 4.0642ms 246.0483 Ops/s 214.1380 Ops/s $\textbf{\color{#35bf28}+14.90\%}$
test_sac_speed[reduce-overhead-None] 2.2655ms 2.1343ms 468.5412 Ops/s 454.5062 Ops/s $\color{#35bf28}+3.09\%$
test_sac_speed[reduce-overhead-backward] 4.1054ms 4.0144ms 249.1055 Ops/s 218.8923 Ops/s $\textbf{\color{#35bf28}+13.80\%}$
test_redq_speed[False-None] 14.0925ms 10.5247ms 95.0150 Ops/s 67.6377 Ops/s $\textbf{\color{#35bf28}+40.48\%}$
test_redq_speed[False-backward] 18.6210ms 17.8352ms 56.0690 Ops/s 56.7205 Ops/s $\color{#d91a1a}-1.15\%$
test_redq_speed[True-None] 4.5491ms 4.3565ms 229.5442 Ops/s 226.7621 Ops/s $\color{#35bf28}+1.23\%$
test_redq_speed[True-backward] 9.9613ms 9.5789ms 104.3962 Ops/s 99.7672 Ops/s $\color{#35bf28}+4.64\%$
test_redq_speed[reduce-overhead-None] 4.5381ms 4.3455ms 230.1223 Ops/s 224.2353 Ops/s $\color{#35bf28}+2.63\%$
test_redq_speed[reduce-overhead-backward] 10.0335ms 9.8562ms 101.4588 Ops/s 104.9116 Ops/s $\color{#d91a1a}-3.29\%$
test_redq_deprec_speed[False-None] 11.7070ms 11.2090ms 89.2140 Ops/s 87.8906 Ops/s $\color{#35bf28}+1.51\%$
test_redq_deprec_speed[False-backward] 16.4496ms 16.0754ms 62.2068 Ops/s 61.3915 Ops/s $\color{#35bf28}+1.33\%$
test_redq_deprec_speed[True-None] 4.2320ms 3.6947ms 270.6560 Ops/s 266.7749 Ops/s $\color{#35bf28}+1.45\%$
test_redq_deprec_speed[True-backward] 8.1115ms 7.6421ms 130.8537 Ops/s 130.0933 Ops/s $\color{#35bf28}+0.58\%$
test_redq_deprec_speed[reduce-overhead-None] 3.9954ms 3.6005ms 277.7414 Ops/s 271.5052 Ops/s $\color{#35bf28}+2.30\%$
test_redq_deprec_speed[reduce-overhead-backward] 8.0671ms 7.6351ms 130.9738 Ops/s 125.8881 Ops/s $\color{#35bf28}+4.04\%$
test_td3_speed[False-None] 8.3978ms 8.1489ms 122.7164 Ops/s 122.6937 Ops/s $\color{#35bf28}+0.02\%$
test_td3_speed[False-backward] 11.4240ms 10.9781ms 91.0908 Ops/s 90.9963 Ops/s $\color{#35bf28}+0.10\%$
test_td3_speed[True-None] 1.8753ms 1.8125ms 551.7217 Ops/s 530.6274 Ops/s $\color{#35bf28}+3.98\%$
test_td3_speed[True-backward] 3.7884ms 3.6527ms 273.7672 Ops/s 256.3845 Ops/s $\textbf{\color{#35bf28}+6.78\%}$
test_td3_speed[reduce-overhead-None] 1.8531ms 1.8091ms 552.7716 Ops/s 538.3102 Ops/s $\color{#35bf28}+2.69\%$
test_td3_speed[reduce-overhead-backward] 3.7334ms 3.6427ms 274.5213 Ops/s 245.4486 Ops/s $\textbf{\color{#35bf28}+11.84\%}$
test_cql_speed[False-None] 29.0761ms 26.1642ms 38.2202 Ops/s 39.0415 Ops/s $\color{#d91a1a}-2.10\%$
test_cql_speed[False-backward] 39.0083ms 35.5619ms 28.1200 Ops/s 28.7786 Ops/s $\color{#d91a1a}-2.29\%$
test_cql_speed[True-None] 12.3933ms 12.1553ms 82.2687 Ops/s 87.0335 Ops/s $\textbf{\color{#d91a1a}-5.47\%}$
test_cql_speed[True-backward] 18.5168ms 18.1061ms 55.2301 Ops/s 59.8185 Ops/s $\textbf{\color{#d91a1a}-7.67\%}$
test_cql_speed[reduce-overhead-None] 12.6911ms 12.2813ms 81.4245 Ops/s 78.7635 Ops/s $\color{#35bf28}+3.38\%$
test_cql_speed[reduce-overhead-backward] 18.7022ms 18.1185ms 55.1923 Ops/s 59.2866 Ops/s $\textbf{\color{#d91a1a}-6.91\%}$
test_a2c_speed[False-None] 5.6234ms 5.4778ms 182.5551 Ops/s 191.8925 Ops/s $\color{#d91a1a}-4.87\%$
test_a2c_speed[False-backward] 12.1815ms 11.7915ms 84.8069 Ops/s 85.1557 Ops/s $\color{#d91a1a}-0.41\%$
test_a2c_speed[True-None] 3.8968ms 3.6872ms 271.2108 Ops/s 287.6462 Ops/s $\textbf{\color{#d91a1a}-5.71\%}$
test_a2c_speed[True-backward] 9.1530ms 8.5984ms 116.3005 Ops/s 117.1615 Ops/s $\color{#d91a1a}-0.73\%$
test_a2c_speed[reduce-overhead-None] 4.1977ms 3.7136ms 269.2809 Ops/s 266.3203 Ops/s $\color{#35bf28}+1.11\%$
test_a2c_speed[reduce-overhead-backward] 8.8545ms 8.7376ms 114.4479 Ops/s 112.2462 Ops/s $\color{#35bf28}+1.96\%$
test_ppo_speed[False-None] 6.3982ms 5.9849ms 167.0883 Ops/s 167.4640 Ops/s $\color{#d91a1a}-0.22\%$
test_ppo_speed[False-backward] 12.6991ms 12.4842ms 80.1014 Ops/s 78.0659 Ops/s $\color{#35bf28}+2.61\%$
test_ppo_speed[True-None] 3.9682ms 3.6113ms 276.9053 Ops/s 268.8910 Ops/s $\color{#35bf28}+2.98\%$
test_ppo_speed[True-backward] 8.6250ms 8.3664ms 119.5257 Ops/s 117.8694 Ops/s $\color{#35bf28}+1.41\%$
test_ppo_speed[reduce-overhead-None] 3.7069ms 3.5785ms 279.4474 Ops/s 271.2472 Ops/s $\color{#35bf28}+3.02\%$
test_ppo_speed[reduce-overhead-backward] 8.7900ms 8.6650ms 115.4071 Ops/s 118.1838 Ops/s $\color{#d91a1a}-2.35\%$
test_reinforce_speed[False-None] 4.8007ms 4.5883ms 217.9452 Ops/s 232.6956 Ops/s $\textbf{\color{#d91a1a}-6.34\%}$
test_reinforce_speed[False-backward] 7.5718ms 7.3759ms 135.5765 Ops/s 138.0245 Ops/s $\color{#d91a1a}-1.77\%$
test_reinforce_speed[True-None] 3.1320ms 2.8672ms 348.7684 Ops/s 347.5380 Ops/s $\color{#35bf28}+0.35\%$
test_reinforce_speed[True-backward] 7.9281ms 7.6663ms 130.4406 Ops/s 129.4223 Ops/s $\color{#35bf28}+0.79\%$
test_reinforce_speed[reduce-overhead-None] 3.0172ms 2.8559ms 350.1580 Ops/s 385.7777 Ops/s $\textbf{\color{#d91a1a}-9.23\%}$
test_reinforce_speed[reduce-overhead-backward] 8.0529ms 7.8757ms 126.9735 Ops/s 122.1594 Ops/s $\color{#35bf28}+3.94\%$
test_iql_speed[False-None] 24.8080ms 20.3141ms 49.2269 Ops/s 49.8785 Ops/s $\color{#d91a1a}-1.31\%$
test_iql_speed[False-backward] 35.9260ms 30.4201ms 32.8730 Ops/s 33.3316 Ops/s $\color{#d91a1a}-1.38\%$
test_iql_speed[True-None] 8.9188ms 8.4523ms 118.3113 Ops/s 124.3727 Ops/s $\color{#d91a1a}-4.87\%$
test_iql_speed[True-backward] 16.9828ms 16.6278ms 60.1404 Ops/s 59.3163 Ops/s $\color{#35bf28}+1.39\%$
test_iql_speed[reduce-overhead-None] 8.7073ms 8.4892ms 117.7968 Ops/s 108.2317 Ops/s $\textbf{\color{#35bf28}+8.84\%}$
test_iql_speed[reduce-overhead-backward] 17.4945ms 17.0664ms 58.5947 Ops/s 57.7535 Ops/s $\color{#35bf28}+1.46\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 7.4885ms 6.0998ms 163.9408 Ops/s 163.2845 Ops/s $\color{#35bf28}+0.40\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 0.5463ms 0.3288ms 3.0417 KOps/s 3.2807 KOps/s $\textbf{\color{#d91a1a}-7.28\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.7653ms 0.3246ms 3.0811 KOps/s 3.4321 KOps/s $\textbf{\color{#d91a1a}-10.23\%}$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 6.0688ms 5.8026ms 172.3370 Ops/s 171.0510 Ops/s $\color{#35bf28}+0.75\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.9125ms 0.3092ms 3.2337 KOps/s 3.4426 KOps/s $\textbf{\color{#d91a1a}-6.07\%}$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.7357ms 0.3398ms 2.9429 KOps/s 3.0271 KOps/s $\color{#d91a1a}-2.78\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 1.6752ms 1.3586ms 736.0416 Ops/s 723.3968 Ops/s $\color{#35bf28}+1.75\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 1.5610ms 1.2799ms 781.3252 Ops/s 780.0209 Ops/s $\color{#35bf28}+0.17\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 9.4185ms 6.1085ms 163.7071 Ops/s 166.3287 Ops/s $\color{#d91a1a}-1.58\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2.3809ms 0.4599ms 2.1746 KOps/s 1.9056 KOps/s $\textbf{\color{#35bf28}+14.11\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 0.6870ms 0.4544ms 2.2006 KOps/s 2.1466 KOps/s $\color{#35bf28}+2.52\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 5.9882ms 5.8367ms 171.3290 Ops/s 170.0218 Ops/s $\color{#35bf28}+0.77\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 1.6454ms 0.2857ms 3.5006 KOps/s 2.7539 KOps/s $\textbf{\color{#35bf28}+27.12\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 0.5420ms 0.3649ms 2.7401 KOps/s 3.2583 KOps/s $\textbf{\color{#d91a1a}-15.90\%}$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 6.0495ms 5.7756ms 173.1429 Ops/s 171.1911 Ops/s $\color{#35bf28}+1.14\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1.4900ms 0.3100ms 3.2260 KOps/s 2.9005 KOps/s $\textbf{\color{#35bf28}+11.22\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 0.5041ms 0.2929ms 3.4141 KOps/s 3.8139 KOps/s $\textbf{\color{#d91a1a}-10.48\%}$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 6.3801ms 5.9591ms 167.8117 Ops/s 167.1026 Ops/s $\color{#35bf28}+0.42\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 0.9726ms 0.5191ms 1.9263 KOps/s 1.9737 KOps/s $\color{#d91a1a}-2.40\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 0.7094ms 0.5016ms 1.9938 KOps/s 1.9539 KOps/s $\color{#35bf28}+2.04\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.6493s 17.8917ms 55.8919 Ops/s 49.4486 Ops/s $\textbf{\color{#35bf28}+13.03\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 12.6215ms 1.9478ms 513.4071 Ops/s 540.0465 Ops/s $\color{#d91a1a}-4.93\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 3.3584ms 1.1442ms 873.9902 Ops/s 932.7752 Ops/s $\textbf{\color{#d91a1a}-6.30\%}$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 8.7539ms 5.0313ms 198.7567 Ops/s 195.7159 Ops/s $\color{#35bf28}+1.55\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 3.8480ms 1.6768ms 596.3721 Ops/s 569.6594 Ops/s $\color{#35bf28}+4.69\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 11.8764ms 1.2886ms 776.0324 Ops/s 988.1976 Ops/s $\textbf{\color{#d91a1a}-21.47\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 0.6032s 17.2355ms 58.0199 Ops/s 188.9090 Ops/s $\textbf{\color{#d91a1a}-69.29\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 10.2788ms 2.0400ms 490.2024 Ops/s 475.0794 Ops/s $\color{#35bf28}+3.18\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 9.9290ms 1.3869ms 721.0469 Ops/s 700.2052 Ops/s $\color{#35bf28}+2.98\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] 40.1003ms 37.1478ms 26.9195 Ops/s 27.2263 Ops/s $\color{#d91a1a}-1.13\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 19.9943ms 18.5794ms 53.8230 Ops/s 55.1590 Ops/s $\color{#d91a1a}-2.42\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True] 44.6124ms 39.0292ms 25.6219 Ops/s 26.2588 Ops/s $\color{#d91a1a}-2.43\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] 22.3382ms 18.3114ms 54.6109 Ops/s 53.2500 Ops/s $\color{#35bf28}+2.56\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 40.2179ms 38.5549ms 25.9370 Ops/s 25.1428 Ops/s $\color{#35bf28}+3.16\%$
test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 21.1545ms 19.7554ms 50.6191 Ops/s 48.5817 Ops/s $\color{#35bf28}+4.19\%$

vmoens added 3 commits January 12, 2026 17:37
@vmoens
Update
c9c327b 
[ghstack-poisoned]
@vmoens
Update
ec4fee4 
[ghstack-poisoned]
@vmoens
Update
a2c13e0 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request Jan 13, 2026
@vmoens
[BugFix] TensorStorage key filtering
0667a58 
Implement `_storage_keys` caching and key filtering in `TensorStorage`.
This handles cases where the incoming data has extra keys (e.g. from a
policy) that aren't present in the storage. Also improved error messages
for locked storage key mismatches.

ghstack-source-id: 8b70f22
Pull-Request: #3323
@vmoens vmoens merged commit a2c13e0 into gh/vmoens/191/base Jan 13, 2026
101 of 106 checks passed
@vmoens vmoens deleted the gh/vmoens/191/head branch January 13, 2026 18:50
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

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.

1 participant

@vmoens