[Test] Test state_dict and loss modules

[vmoens]

depends on pytorch/tensordict#528

[matteobettini]

LGTM

[github-actions]

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

Total Benchmarks: 89. Improved: $\large\color{#35bf28}5$. Worsened: $\large\color{#d91a1a}10$.

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_single	74.1498ms	73.4676ms	13.6114 Ops/s	13.6853 Ops/s	$\color{#d91a1a}-0.54\%$
test_sync	0.1174s	43.2681ms	23.1117 Ops/s	22.9656 Ops/s	$\color{#35bf28}+0.64\%$
test_async	0.1233s	40.4018ms	24.7514 Ops/s	25.0717 Ops/s	$\color{#d91a1a}-1.28\%$
test_simple	0.6908s	0.6193s	1.6146 Ops/s	1.5961 Ops/s	$\color{#35bf28}+1.16\%$
test_transformed	0.8357s	0.7898s	1.2661 Ops/s	1.2479 Ops/s	$\color{#35bf28}+1.46\%$
test_serial	1.7763s	1.7151s	0.5830 Ops/s	0.5751 Ops/s	$\color{#35bf28}+1.37\%$
test_parallel	1.6320s	1.5450s	0.6473 Ops/s	0.6705 Ops/s	$\color{#d91a1a}-3.47\%$
test_step_mdp_speed[True-True-True-True-True]	0.1908ms	44.4822μs	22.4809 KOps/s	22.2425 KOps/s	$\color{#35bf28}+1.07\%$
test_step_mdp_speed[True-True-True-True-False]	63.2000μs	25.5911μs	39.0762 KOps/s	39.5944 KOps/s	$\color{#d91a1a}-1.31\%$
test_step_mdp_speed[True-True-True-False-True]	57.9010μs	31.4197μs	31.8272 KOps/s	31.6973 KOps/s	$\color{#35bf28}+0.41\%$
test_step_mdp_speed[True-True-True-False-False]	45.5010μs	17.5600μs	56.9476 KOps/s	56.9217 KOps/s	$\color{#35bf28}+0.05\%$
test_step_mdp_speed[True-True-False-True-True]	73.8010μs	46.0908μs	21.6963 KOps/s	21.6462 KOps/s	$\color{#35bf28}+0.23\%$
test_step_mdp_speed[True-True-False-True-False]	71.5010μs	27.1506μs	36.8316 KOps/s	36.8478 KOps/s	$\color{#d91a1a}-0.04\%$
test_step_mdp_speed[True-True-False-False-True]	58.1010μs	33.3646μs	29.9719 KOps/s	29.6487 KOps/s	$\color{#35bf28}+1.09\%$
test_step_mdp_speed[True-True-False-False-False]	63.5010μs	19.6049μs	51.0078 KOps/s	51.2607 KOps/s	$\color{#d91a1a}-0.49\%$
test_step_mdp_speed[True-False-True-True-True]	72.6010μs	48.0380μs	20.8169 KOps/s	20.6361 KOps/s	$\color{#35bf28}+0.88\%$
test_step_mdp_speed[True-False-True-True-False]	96.7010μs	29.1145μs	34.3472 KOps/s	34.1523 KOps/s	$\color{#35bf28}+0.57\%$
test_step_mdp_speed[True-False-True-False-True]	57.6010μs	33.3059μs	30.0247 KOps/s	29.8111 KOps/s	$\color{#35bf28}+0.72\%$
test_step_mdp_speed[True-False-True-False-False]	0.2061ms	19.7896μs	50.5317 KOps/s	51.4702 KOps/s	$\color{#d91a1a}-1.82\%$
test_step_mdp_speed[True-False-False-True-True]	0.1250ms	49.4990μs	20.2024 KOps/s	20.1621 KOps/s	$\color{#35bf28}+0.20\%$
test_step_mdp_speed[True-False-False-True-False]	92.3010μs	30.8482μs	32.4168 KOps/s	32.3106 KOps/s	$\color{#35bf28}+0.33\%$
test_step_mdp_speed[True-False-False-False-True]	54.9010μs	34.9752μs	28.5917 KOps/s	28.8471 KOps/s	$\color{#d91a1a}-0.89\%$
test_step_mdp_speed[True-False-False-False-False]	54.9010μs	21.4403μs	46.6412 KOps/s	47.0897 KOps/s	$\color{#d91a1a}-0.95\%$
test_step_mdp_speed[False-True-True-True-True]	75.0010μs	47.8169μs	20.9131 KOps/s	20.7795 KOps/s	$\color{#35bf28}+0.64\%$
test_step_mdp_speed[False-True-True-True-False]	96.3010μs	29.3084μs	34.1199 KOps/s	34.0197 KOps/s	$\color{#35bf28}+0.29\%$
test_step_mdp_speed[False-True-True-False-True]	0.1130ms	37.1135μs	26.9444 KOps/s	26.8822 KOps/s	$\color{#35bf28}+0.23\%$
test_step_mdp_speed[False-True-True-False-False]	3.3618ms	21.6943μs	46.0951 KOps/s	45.4859 KOps/s	$\color{#35bf28}+1.34\%$
test_step_mdp_speed[False-True-False-True-True]	78.5010μs	49.4293μs	20.2309 KOps/s	20.1209 KOps/s	$\color{#35bf28}+0.55\%$
test_step_mdp_speed[False-True-False-True-False]	52.9010μs	30.7134μs	32.5591 KOps/s	32.1567 KOps/s	$\color{#35bf28}+1.25\%$
test_step_mdp_speed[False-True-False-False-True]	89.1010μs	38.6143μs	25.8972 KOps/s	25.8216 KOps/s	$\color{#35bf28}+0.29\%$
test_step_mdp_speed[False-True-False-False-False]	0.3775ms	23.3459μs	42.8341 KOps/s	42.2889 KOps/s	$\color{#35bf28}+1.29\%$
test_step_mdp_speed[False-False-True-True-True]	98.9010μs	52.0999μs	19.1939 KOps/s	19.4556 KOps/s	$\color{#d91a1a}-1.35\%$
test_step_mdp_speed[False-False-True-True-False]	67.6010μs	32.5604μs	30.7122 KOps/s	30.3836 KOps/s	$\color{#35bf28}+1.08\%$
test_step_mdp_speed[False-False-True-False-True]	89.8010μs	39.1806μs	25.5229 KOps/s	25.9457 KOps/s	$\color{#d91a1a}-1.63\%$
test_step_mdp_speed[False-False-True-False-False]	74.4010μs	23.3706μs	42.7888 KOps/s	42.7245 KOps/s	$\color{#35bf28}+0.15\%$
test_step_mdp_speed[False-False-False-True-True]	86.6010μs	52.8386μs	18.9256 KOps/s	18.9603 KOps/s	$\color{#d91a1a}-0.18\%$
test_step_mdp_speed[False-False-False-True-False]	63.9010μs	33.9816μs	29.4277 KOps/s	29.0462 KOps/s	$\color{#35bf28}+1.31\%$
test_step_mdp_speed[False-False-False-False-True]	83.0010μs	40.3752μs	24.7677 KOps/s	24.7818 KOps/s	$\color{#d91a1a}-0.06\%$
test_step_mdp_speed[False-False-False-False-False]	47.8010μs	24.8591μs	40.2268 KOps/s	40.3326 KOps/s	$\color{#d91a1a}-0.26\%$
test_values[generalized_advantage_estimate-True-True]	15.6447ms	14.0503ms	71.1731 Ops/s	71.1997 Ops/s	$\color{#d91a1a}-0.04\%$
test_values[vec_generalized_advantage_estimate-True-True]	46.0222ms	40.7492ms	24.5404 Ops/s	23.7624 Ops/s	$\color{#35bf28}+3.27\%$
test_values[td0_return_estimate-False-False]	0.4491ms	0.2010ms	4.9754 KOps/s	5.1559 KOps/s	$\color{#d91a1a}-3.50\%$
test_values[td1_return_estimate-False-False]	13.9960ms	13.6740ms	73.1315 Ops/s	74.5302 Ops/s	$\color{#d91a1a}-1.88\%$
test_values[vec_td1_return_estimate-False-False]	46.1648ms	40.9839ms	24.3998 Ops/s	24.1025 Ops/s	$\color{#35bf28}+1.23\%$
test_values[td_lambda_return_estimate-True-False]	37.8114ms	32.9748ms	30.3262 Ops/s	30.7027 Ops/s	$\color{#d91a1a}-1.23\%$
test_values[vec_td_lambda_return_estimate-True-False]	46.5087ms	41.6838ms	23.9902 Ops/s	23.7383 Ops/s	$\color{#35bf28}+1.06\%$
test_gae_speed[generalized_advantage_estimate-False-1-512]	15.8057ms	12.2723ms	81.4843 Ops/s	80.9560 Ops/s	$\color{#35bf28}+0.65\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512]	5.5875ms	3.3654ms	297.1444 Ops/s	303.8898 Ops/s	$\color{#d91a1a}-2.22\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512]	6.6550ms	0.4853ms	2.0605 KOps/s	2.1701 KOps/s	$\textbf{\color{#d91a1a}-5.05\%}$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512]	59.7372ms	55.4786ms	18.0250 Ops/s	17.3819 Ops/s	$\color{#35bf28}+3.70\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512]	3.5001ms	2.8966ms	345.2270 Ops/s	359.6085 Ops/s	$\color{#d91a1a}-4.00\%$
test_dqn_speed	7.8476ms	1.8621ms	537.0316 Ops/s	486.5256 Ops/s	$\textbf{\color{#35bf28}+10.38\%}$
test_ddpg_speed	8.4730ms	2.8220ms	354.3603 Ops/s	359.9958 Ops/s	$\color{#d91a1a}-1.57\%$
test_sac_speed	16.5088ms	8.5636ms	116.7728 Ops/s	121.0420 Ops/s	$\color{#d91a1a}-3.53\%$
test_redq_speed	28.8888ms	17.3439ms	57.6573 Ops/s	61.1037 Ops/s	$\textbf{\color{#d91a1a}-5.64\%}$
test_redq_deprec_speed	18.9700ms	13.6569ms	73.2230 Ops/s	77.2015 Ops/s	$\textbf{\color{#d91a1a}-5.15\%}$
test_td3_speed	11.8437ms	10.3545ms	96.5766 Ops/s	98.6611 Ops/s	$\color{#d91a1a}-2.11\%$
test_cql_speed	37.0919ms	28.2992ms	35.3366 Ops/s	37.4596 Ops/s	$\textbf{\color{#d91a1a}-5.67\%}$
test_a2c_speed	11.9366ms	5.4535ms	183.3677 Ops/s	189.8489 Ops/s	$\color{#d91a1a}-3.41\%$
test_ppo_speed	11.6516ms	5.8491ms	170.9668 Ops/s	180.7691 Ops/s	$\textbf{\color{#d91a1a}-5.42\%}$
test_reinforce_speed	9.4274ms	4.2027ms	237.9426 Ops/s	239.3049 Ops/s	$\color{#d91a1a}-0.57\%$
test_iql_speed	30.1413ms	22.1366ms	45.1740 Ops/s	46.5045 Ops/s	$\color{#d91a1a}-2.86\%$
test_sample_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	3.3789ms	2.6509ms	377.2323 Ops/s	374.9553 Ops/s	$\color{#35bf28}+0.61\%$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	3.8653ms	2.7856ms	358.9872 Ops/s	360.9474 Ops/s	$\color{#d91a1a}-0.54\%$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	0.1517s	3.2885ms	304.0934 Ops/s	353.3019 Ops/s	$\textbf{\color{#d91a1a}-13.93\%}$
test_sample_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	2.9589ms	2.6690ms	374.6734 Ops/s	382.1700 Ops/s	$\color{#d91a1a}-1.96\%$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	4.4310ms	2.8561ms	350.1279 Ops/s	347.6927 Ops/s	$\color{#35bf28}+0.70\%$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	6.0098ms	2.8999ms	344.8423 Ops/s	350.6989 Ops/s	$\color{#d91a1a}-1.67\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	3.7727ms	2.7486ms	363.8221 Ops/s	381.6543 Ops/s	$\color{#d91a1a}-4.67\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	5.1897ms	2.8717ms	348.2278 Ops/s	356.8805 Ops/s	$\color{#d91a1a}-2.42\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	0.1473s	3.3069ms	302.4003 Ops/s	357.6752 Ops/s	$\textbf{\color{#d91a1a}-15.45\%}$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	3.8248ms	2.7456ms	364.2131 Ops/s	375.6836 Ops/s	$\color{#d91a1a}-3.05\%$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	4.3147ms	2.8942ms	345.5244 Ops/s	341.8177 Ops/s	$\color{#35bf28}+1.08\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	4.8153ms	2.8879ms	346.2750 Ops/s	356.9951 Ops/s	$\color{#d91a1a}-3.00\%$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	3.7288ms	2.6871ms	372.1431 Ops/s	290.4150 Ops/s	$\textbf{\color{#35bf28}+28.14\%}$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	4.6629ms	2.7792ms	359.8165 Ops/s	356.4535 Ops/s	$\color{#35bf28}+0.94\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	6.0237ms	2.9028ms	344.4938 Ops/s	348.5298 Ops/s	$\color{#d91a1a}-1.16\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	3.5560ms	2.7074ms	369.3535 Ops/s	378.8372 Ops/s	$\color{#d91a1a}-2.50\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	5.1763ms	2.8764ms	347.6529 Ops/s	351.6251 Ops/s	$\color{#d91a1a}-1.13\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	4.4463ms	2.8579ms	349.9027 Ops/s	361.8547 Ops/s	$\color{#d91a1a}-3.30\%$
test_populate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-400]	0.2364s	31.7021ms	31.5436 Ops/s	33.8884 Ops/s	$\textbf{\color{#d91a1a}-6.92\%}$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]	0.1464s	27.2321ms	36.7214 Ops/s	34.6699 Ops/s	$\textbf{\color{#35bf28}+5.92\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]	0.1474s	29.7935ms	33.5644 Ops/s	36.6506 Ops/s	$\textbf{\color{#d91a1a}-8.42\%}$
test_populate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]	0.1447s	26.9986ms	37.0390 Ops/s	33.9626 Ops/s	$\textbf{\color{#35bf28}+9.06\%}$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]	0.1386s	29.0161ms	34.4637 Ops/s	37.0771 Ops/s	$\textbf{\color{#d91a1a}-7.05\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]	0.1351s	28.7294ms	34.8076 Ops/s	33.9527 Ops/s	$\color{#35bf28}+2.52\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]	0.1352s	26.3812ms	37.9058 Ops/s	37.8172 Ops/s	$\color{#35bf28}+0.23\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]	0.1366s	29.4973ms	33.9014 Ops/s	34.4095 Ops/s	$\color{#d91a1a}-1.48\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]	0.1359s	26.3565ms	37.9413 Ops/s	34.1369 Ops/s	$\textbf{\color{#35bf28}+11.14\%}$