[Feature] Open-X Embodiement datasets

[vmoens]

The data can be either downloaded or streamed.

Example usage:

import tempfile

from torchrl.data.datasets.openx import OpenXExperienceReplay
# Download the data, and sample 128 elements in each batch out of two trajectories
num_slices = 2
with tempfile.TemporaryDirectory() as root:
    dataset = OpenXExperienceReplay("cmu_stretch", batch_size=128,
        num_slices=num_slices, download=True, streaming=False, root=root)
    for batch in dataset:
        print(batch.reshape(num_slices, -1))
        break
    # Read data from a stream. Deliver entire trajectories when iterating
    dataset = OpenXExperienceReplay("cmu_stretch",
        num_slices=num_slices, download=True, streaming=False, root=root)
    for data in dataset: # data does not have a consistent shape
        break
    # Define batch-size dynamically
    data = dataset.sample(128)  # delivers 2 sub-trajectories of length 64

cc @jxu124 @quanvuong @nicklashansen

• Should we return an exception when shuffle=False and calling sample?

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (8 Unrelated Failures)

As of commit fc2b1ab with merge base 781a5b2 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Build M1 Wheels / pytorch/rl / wheel-py3_8-cpu (gh)
  ImportError: cannot import name 'MemoryMappedTensor' from 'tensordict' (/Users/ec2-user/runner/_work/_temp/conda_environment_7445506770/lib/python3.8/site-packages/tensordict/__init__.py)
• Unit-tests on Linux GPU / tests (3.8, 12.1) / linux-job (gh)
  test/test_libs.py::TestOpenML::test_env[batch_size0-mushroom_num]
• Unit-tests on Linux GPU, latest stable release / tests (3.8, 11.8) / linux-job (gh)
  test/test_libs.py::TestOpenML::test_env[batch_size0-mushroom_num]
• Unit-tests on MacOS CPU / tests (3.8) / macos-job (gh)
  test/test_libs.py::TestOpenML::test_env[batch_size2-magic]
• Unit-tests on Windows / unittests-cpu / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Unit-tests on Windows / unittests-gpu / windows-job (gh)
  ##[error]The operation was canceled.

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

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

• Habitat Tests on Linux / tests (3.9, 11.6) / linux-job (gh)
  test/test_libs.py::TestHabitat::test_habitat_render[False-HabitatPick-v0]
• Unit-tests on MacOS CPU / tests (3.11) / macos-job (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 245

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

[github-actions]

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

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

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_single	66.1435ms	65.5893ms	15.2464 Ops/s	15.3780 Ops/s	$\color{#d91a1a}-0.86\%$
test_sync	41.7780ms	35.2085ms	28.4023 Ops/s	28.7720 Ops/s	$\color{#d91a1a}-1.29\%$
test_async	0.1380s	33.6892ms	29.6831 Ops/s	29.6756 Ops/s	$\color{#35bf28}+0.03\%$
test_simple	0.5029s	0.4557s	2.1947 Ops/s	2.2282 Ops/s	$\color{#d91a1a}-1.51\%$
test_transformed	0.6896s	0.6311s	1.5846 Ops/s	1.6406 Ops/s	$\color{#d91a1a}-3.41\%$
test_serial	1.4042s	1.4035s	0.7125 Ops/s	0.7290 Ops/s	$\color{#d91a1a}-2.27\%$
test_parallel	1.4029s	1.3544s	0.7383 Ops/s	0.7444 Ops/s	$\color{#d91a1a}-0.82\%$
test_step_mdp_speed[True-True-True-True-True]	0.1474ms	22.2160μs	45.0126 KOps/s	45.6863 KOps/s	$\color{#d91a1a}-1.47\%$
test_step_mdp_speed[True-True-True-True-False]	57.5080μs	13.2215μs	75.6344 KOps/s	75.1315 KOps/s	$\color{#35bf28}+0.67\%$
test_step_mdp_speed[True-True-True-False-True]	39.4640μs	13.0677μs	76.5245 KOps/s	77.5979 KOps/s	$\color{#d91a1a}-1.38\%$
test_step_mdp_speed[True-True-True-False-False]	73.3170μs	7.8151μs	127.9570 KOps/s	127.1716 KOps/s	$\color{#35bf28}+0.62\%$
test_step_mdp_speed[True-True-False-True-True]	72.0150μs	23.1639μs	43.1707 KOps/s	42.9714 KOps/s	$\color{#35bf28}+0.46\%$
test_step_mdp_speed[True-True-False-True-False]	42.4600μs	14.6679μs	68.1760 KOps/s	67.9130 KOps/s	$\color{#35bf28}+0.39\%$
test_step_mdp_speed[True-True-False-False-True]	42.6390μs	14.3395μs	69.7375 KOps/s	70.4465 KOps/s	$\color{#d91a1a}-1.01\%$
test_step_mdp_speed[True-True-False-False-False]	56.1050μs	9.0534μs	110.4552 KOps/s	108.6424 KOps/s	$\color{#35bf28}+1.67\%$
test_step_mdp_speed[True-False-True-True-True]	76.9680μs	24.5344μs	40.7590 KOps/s	40.5798 KOps/s	$\color{#35bf28}+0.44\%$
test_step_mdp_speed[True-False-True-True-False]	43.6210μs	15.8231μs	63.1987 KOps/s	62.5720 KOps/s	$\color{#35bf28}+1.00\%$
test_step_mdp_speed[True-False-True-False-True]	47.7390μs	14.3678μs	69.5999 KOps/s	71.5648 KOps/s	$\color{#d91a1a}-2.75\%$
test_step_mdp_speed[True-False-True-False-False]	50.8650μs	9.1625μs	109.1407 KOps/s	109.3635 KOps/s	$\color{#d91a1a}-0.20\%$
test_step_mdp_speed[True-False-False-True-True]	67.4630μs	25.6477μs	38.9898 KOps/s	39.1356 KOps/s	$\color{#d91a1a}-0.37\%$
test_step_mdp_speed[True-False-False-True-False]	47.7000μs	17.0182μs	58.7607 KOps/s	58.9699 KOps/s	$\color{#d91a1a}-0.35\%$
test_step_mdp_speed[True-False-False-False-True]	41.0260μs	15.3938μs	64.9613 KOps/s	64.6591 KOps/s	$\color{#35bf28}+0.47\%$
test_step_mdp_speed[True-False-False-False-False]	33.7030μs	10.2478μs	97.5819 KOps/s	97.0410 KOps/s	$\color{#35bf28}+0.56\%$
test_step_mdp_speed[False-True-True-True-True]	71.7640μs	24.8230μs	40.2852 KOps/s	40.9131 KOps/s	$\color{#d91a1a}-1.53\%$
test_step_mdp_speed[False-True-True-True-False]	56.5160μs	16.0290μs	62.3868 KOps/s	62.9904 KOps/s	$\color{#d91a1a}-0.96\%$
test_step_mdp_speed[False-True-True-False-True]	40.8570μs	16.5659μs	60.3649 KOps/s	60.4329 KOps/s	$\color{#d91a1a}-0.11\%$
test_step_mdp_speed[False-True-True-False-False]	37.2400μs	10.3237μs	96.8648 KOps/s	96.2658 KOps/s	$\color{#35bf28}+0.62\%$
test_step_mdp_speed[False-True-False-True-True]	55.9350μs	25.9095μs	38.5959 KOps/s	39.0635 KOps/s	$\color{#d91a1a}-1.20\%$
test_step_mdp_speed[False-True-False-True-False]	57.0070μs	17.0595μs	58.6184 KOps/s	58.4687 KOps/s	$\color{#35bf28}+0.26\%$
test_step_mdp_speed[False-True-False-False-True]	55.7240μs	17.6212μs	56.7497 KOps/s	56.2444 KOps/s	$\color{#35bf28}+0.90\%$
test_step_mdp_speed[False-True-False-False-False]	38.1410μs	11.4655μs	87.2185 KOps/s	87.2489 KOps/s	$\color{#d91a1a}-0.03\%$
test_step_mdp_speed[False-False-True-True-True]	82.8750μs	27.3959μs	36.5017 KOps/s	38.0098 KOps/s	$\color{#d91a1a}-3.97\%$
test_step_mdp_speed[False-False-True-True-False]	77.1440μs	18.2664μs	54.7455 KOps/s	54.5072 KOps/s	$\color{#35bf28}+0.44\%$
test_step_mdp_speed[False-False-True-False-True]	41.5680μs	17.7425μs	56.3617 KOps/s	56.9249 KOps/s	$\color{#d91a1a}-0.99\%$
test_step_mdp_speed[False-False-True-False-False]	31.4690μs	11.5130μs	86.8584 KOps/s	85.6351 KOps/s	$\color{#35bf28}+1.43\%$
test_step_mdp_speed[False-False-False-True-True]	69.5500μs	28.2497μs	35.3987 KOps/s	35.9948 KOps/s	$\color{#d91a1a}-1.66\%$
test_step_mdp_speed[False-False-False-True-False]	53.2900μs	19.4768μs	51.3431 KOps/s	51.8145 KOps/s	$\color{#d91a1a}-0.91\%$
test_step_mdp_speed[False-False-False-False-True]	68.4780μs	18.7083μs	53.4523 KOps/s	54.1036 KOps/s	$\color{#d91a1a}-1.20\%$
test_step_mdp_speed[False-False-False-False-False]	36.8090μs	12.5939μs	79.4038 KOps/s	79.1959 KOps/s	$\color{#35bf28}+0.26\%$
test_values[generalized_advantage_estimate-True-True]	12.8280ms	12.1449ms	82.3390 Ops/s	81.5616 Ops/s	$\color{#35bf28}+0.95\%$
test_values[vec_generalized_advantage_estimate-True-True]	39.3800ms	27.6997ms	36.1015 Ops/s	35.8221 Ops/s	$\color{#35bf28}+0.78\%$
test_values[td0_return_estimate-False-False]	0.2290ms	0.1753ms	5.7051 KOps/s	5.6356 KOps/s	$\color{#35bf28}+1.23\%$
test_values[td1_return_estimate-False-False]	26.4060ms	26.0959ms	38.3202 Ops/s	38.3156 Ops/s	$\color{#35bf28}+0.01\%$
test_values[vec_td1_return_estimate-False-False]	35.7156ms	27.7313ms	36.0604 Ops/s	35.9540 Ops/s	$\color{#35bf28}+0.30\%$
test_values[td_lambda_return_estimate-True-False]	39.3201ms	36.3948ms	27.4765 Ops/s	27.0337 Ops/s	$\color{#35bf28}+1.64\%$
test_values[vec_td_lambda_return_estimate-True-False]	35.7228ms	27.8653ms	35.8869 Ops/s	36.0448 Ops/s	$\color{#d91a1a}-0.44\%$
test_gae_speed[generalized_advantage_estimate-False-1-512]	11.3614ms	8.2958ms	120.5427 Ops/s	122.3831 Ops/s	$\color{#d91a1a}-1.50\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512]	2.2462ms	2.0093ms	497.6746 Ops/s	492.6799 Ops/s	$\color{#35bf28}+1.01\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512]	8.9374ms	0.4429ms	2.2577 KOps/s	2.3218 KOps/s	$\color{#d91a1a}-2.76\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512]	46.8952ms	39.0255ms	25.6243 Ops/s	25.1406 Ops/s	$\color{#35bf28}+1.92\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512]	10.8338ms	2.6698ms	374.5668 Ops/s	371.3894 Ops/s	$\color{#35bf28}+0.86\%$
test_dqn_speed	14.6352ms	7.7637ms	128.8044 Ops/s	121.5552 Ops/s	$\textbf{\color{#35bf28}+5.96\%}$
test_ddpg_speed	22.7205ms	14.8161ms	67.4942 Ops/s	68.8335 Ops/s	$\color{#d91a1a}-1.95\%$
test_sac_speed	37.3970ms	29.7612ms	33.6008 Ops/s	33.6997 Ops/s	$\color{#d91a1a}-0.29\%$
test_redq_speed	0.1095s	37.8732ms	26.4039 Ops/s	27.9250 Ops/s	$\textbf{\color{#d91a1a}-5.45\%}$
test_redq_deprec_speed	30.2391ms	25.8306ms	38.7137 Ops/s	38.8052 Ops/s	$\color{#d91a1a}-0.24\%$
test_td3_speed	28.6539ms	20.5045ms	48.7698 Ops/s	49.1731 Ops/s	$\color{#d91a1a}-0.82\%$
test_cql_speed	97.4798ms	89.2291ms	11.2071 Ops/s	11.4519 Ops/s	$\color{#d91a1a}-2.14\%$
test_a2c_speed	32.1902ms	27.4205ms	36.4691 Ops/s	37.3822 Ops/s	$\color{#d91a1a}-2.44\%$
test_ppo_speed	32.5853ms	27.3432ms	36.5722 Ops/s	36.9261 Ops/s	$\color{#d91a1a}-0.96\%$
test_reinforce_speed	34.9270ms	26.3868ms	37.8978 Ops/s	38.1978 Ops/s	$\color{#d91a1a}-0.79\%$
test_iql_speed	71.5367ms	63.8741ms	15.6558 Ops/s	15.7268 Ops/s	$\color{#d91a1a}-0.45\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	2.1943ms	1.4722ms	679.2610 Ops/s	678.5208 Ops/s	$\color{#35bf28}+0.11\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	8.8403ms	0.5343ms	1.8715 KOps/s	1.8594 KOps/s	$\color{#35bf28}+0.65\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	2.4175ms	0.6301ms	1.5869 KOps/s	1.9115 KOps/s	$\textbf{\color{#d91a1a}-16.98\%}$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	1.8686ms	1.4650ms	682.6136 Ops/s	683.3227 Ops/s	$\color{#d91a1a}-0.10\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	9.7365ms	0.5207ms	1.9206 KOps/s	1.8660 KOps/s	$\color{#35bf28}+2.93\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	0.6830ms	0.4987ms	2.0052 KOps/s	1.9317 KOps/s	$\color{#35bf28}+3.80\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	0.1157s	1.8238ms	548.3002 Ops/s	591.0824 Ops/s	$\textbf{\color{#d91a1a}-7.24\%}$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	4.1680ms	0.6568ms	1.5225 KOps/s	1.4700 KOps/s	$\color{#35bf28}+3.57\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	8.9975ms	0.6483ms	1.5425 KOps/s	1.4973 KOps/s	$\color{#35bf28}+3.02\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	1.8104ms	1.4754ms	677.7695 Ops/s	560.1128 Ops/s	$\textbf{\color{#35bf28}+21.01\%}$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	8.8681ms	0.5299ms	1.8872 KOps/s	1.8800 KOps/s	$\color{#35bf28}+0.38\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	1.6983ms	0.5150ms	1.9416 KOps/s	1.9253 KOps/s	$\color{#35bf28}+0.84\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	1.8608ms	1.4444ms	692.3204 Ops/s	674.6344 Ops/s	$\color{#35bf28}+2.62\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	2.1754ms	0.5135ms	1.9473 KOps/s	1.8697 KOps/s	$\color{#35bf28}+4.15\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	0.6281ms	0.4999ms	2.0004 KOps/s	1.9293 KOps/s	$\color{#35bf28}+3.69\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	2.2360ms	1.6757ms	596.7699 Ops/s	586.7917 Ops/s	$\color{#35bf28}+1.70\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	4.1572ms	0.6594ms	1.5164 KOps/s	1.4615 KOps/s	$\color{#35bf28}+3.76\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	9.0107ms	0.6595ms	1.5162 KOps/s	1.4879 KOps/s	$\color{#35bf28}+1.90\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]	0.1230s	16.6643ms	60.0085 Ops/s	60.1681 Ops/s	$\color{#d91a1a}-0.27\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]	14.2592ms	12.1480ms	82.3178 Ops/s	80.1885 Ops/s	$\color{#35bf28}+2.66\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]	10.0915ms	1.7666ms	566.0496 Ops/s	669.3988 Ops/s	$\textbf{\color{#d91a1a}-15.44\%}$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]	0.1071s	16.1600ms	61.8810 Ops/s	61.3121 Ops/s	$\color{#35bf28}+0.93\%$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]	18.9331ms	12.3214ms	81.1596 Ops/s	81.6679 Ops/s	$\color{#d91a1a}-0.62\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]	5.0857ms	1.5817ms	632.2344 Ops/s	657.1941 Ops/s	$\color{#d91a1a}-3.80\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]	0.1176s	16.7170ms	59.8195 Ops/s	68.7690 Ops/s	$\textbf{\color{#d91a1a}-13.01\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]	14.6687ms	12.3541ms	80.9448 Ops/s	79.9412 Ops/s	$\color{#35bf28}+1.26\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]	2.3964ms	1.6709ms	598.4889 Ops/s	589.2275 Ops/s	$\color{#35bf28}+1.57\%$

[github-actions]

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

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

Expand to view detailed results

Name	Max	Mean	Ops	Ops on Repo HEAD	Change
test_single	0.1285s	0.1265s	7.9067 Ops/s	8.1831 Ops/s	$\color{#d91a1a}-3.38\%$
test_sync	0.1787s	0.1116s	8.9621 Ops/s	9.1187 Ops/s	$\color{#d91a1a}-1.72\%$
test_async	0.2719s	0.1001s	9.9906 Ops/s	10.0976 Ops/s	$\color{#d91a1a}-1.06\%$
test_single_pixels	0.1337s	0.1334s	7.4942 Ops/s	6.8800 Ops/s	$\textbf{\color{#35bf28}+8.93\%}$
test_sync_pixels	0.1008s	97.1117ms	10.2974 Ops/s	10.4030 Ops/s	$\color{#d91a1a}-1.01\%$
test_async_pixels	0.2477s	92.5903ms	10.8003 Ops/s	10.9039 Ops/s	$\color{#d91a1a}-0.95\%$
test_simple	0.9664s	0.8969s	1.1150 Ops/s	1.1216 Ops/s	$\color{#d91a1a}-0.59\%$
test_transformed	1.2110s	1.1469s	0.8719 Ops/s	0.8876 Ops/s	$\color{#d91a1a}-1.76\%$
test_serial	2.5691s	2.5323s	0.3949 Ops/s	0.3942 Ops/s	$\color{#35bf28}+0.18\%$
test_parallel	2.5804s	2.5057s	0.3991 Ops/s	0.3995 Ops/s	$\color{#d91a1a}-0.09\%$
test_step_mdp_speed[True-True-True-True-True]	92.4020μs	32.2779μs	30.9810 KOps/s	30.1805 KOps/s	$\color{#35bf28}+2.65\%$
test_step_mdp_speed[True-True-True-True-False]	56.7500μs	19.9877μs	50.0307 KOps/s	50.7704 KOps/s	$\color{#d91a1a}-1.46\%$
test_step_mdp_speed[True-True-True-False-True]	43.9310μs	18.5651μs	53.8646 KOps/s	53.0372 KOps/s	$\color{#35bf28}+1.56\%$
test_step_mdp_speed[True-True-True-False-False]	41.7010μs	11.3576μs	88.0468 KOps/s	88.6549 KOps/s	$\color{#d91a1a}-0.69\%$
test_step_mdp_speed[True-True-False-True-True]	63.4900μs	33.8840μs	29.5125 KOps/s	29.1412 KOps/s	$\color{#35bf28}+1.27\%$
test_step_mdp_speed[True-True-False-True-False]	46.2110μs	21.1990μs	47.1719 KOps/s	46.7862 KOps/s	$\color{#35bf28}+0.82\%$
test_step_mdp_speed[True-True-False-False-True]	45.2200μs	19.8646μs	50.3407 KOps/s	48.3554 KOps/s	$\color{#35bf28}+4.11\%$
test_step_mdp_speed[True-True-False-False-False]	66.5600μs	13.3018μs	75.1776 KOps/s	76.0733 KOps/s	$\color{#d91a1a}-1.18\%$
test_step_mdp_speed[True-False-True-True-True]	68.1110μs	35.4156μs	28.2362 KOps/s	27.8211 KOps/s	$\color{#35bf28}+1.49\%$
test_step_mdp_speed[True-False-True-True-False]	48.8810μs	23.3647μs	42.7997 KOps/s	42.9028 KOps/s	$\color{#d91a1a}-0.24\%$
test_step_mdp_speed[True-False-True-False-True]	45.3520μs	20.1485μs	49.6315 KOps/s	48.5495 KOps/s	$\color{#35bf28}+2.23\%$
test_step_mdp_speed[True-False-True-False-False]	36.9200μs	13.1582μs	75.9983 KOps/s	76.6530 KOps/s	$\color{#d91a1a}-0.85\%$
test_step_mdp_speed[True-False-False-True-True]	80.0120μs	38.0452μs	26.2845 KOps/s	26.4338 KOps/s	$\color{#d91a1a}-0.56\%$
test_step_mdp_speed[True-False-False-True-False]	55.8800μs	25.5544μs	39.1321 KOps/s	40.1056 KOps/s	$\color{#d91a1a}-2.43\%$
test_step_mdp_speed[True-False-False-False-True]	48.3400μs	21.7739μs	45.9264 KOps/s	45.1771 KOps/s	$\color{#35bf28}+1.66\%$
test_step_mdp_speed[True-False-False-False-False]	36.9710μs	15.0992μs	66.2288 KOps/s	67.3061 KOps/s	$\color{#d91a1a}-1.60\%$
test_step_mdp_speed[False-True-True-True-True]	65.2520μs	36.0496μs	27.7396 KOps/s	27.5708 KOps/s	$\color{#35bf28}+0.61\%$
test_step_mdp_speed[False-True-True-True-False]	50.0100μs	23.1467μs	43.2027 KOps/s	43.3000 KOps/s	$\color{#d91a1a}-0.22\%$
test_step_mdp_speed[False-True-True-False-True]	55.1420μs	24.0103μs	41.6488 KOps/s	39.8304 KOps/s	$\color{#35bf28}+4.57\%$
test_step_mdp_speed[False-True-True-False-False]	41.5610μs	15.2310μs	65.6554 KOps/s	67.3196 KOps/s	$\color{#d91a1a}-2.47\%$
test_step_mdp_speed[False-True-False-True-True]	97.2310μs	38.0292μs	26.2956 KOps/s	26.3566 KOps/s	$\color{#d91a1a}-0.23\%$
test_step_mdp_speed[False-True-False-True-False]	50.9200μs	25.4288μs	39.3255 KOps/s	39.7583 KOps/s	$\color{#d91a1a}-1.09\%$
test_step_mdp_speed[False-True-False-False-True]	51.4810μs	25.9599μs	38.5210 KOps/s	38.0301 KOps/s	$\color{#35bf28}+1.29\%$
test_step_mdp_speed[False-True-False-False-False]	56.3800μs	16.9211μs	59.0978 KOps/s	59.3561 KOps/s	$\color{#d91a1a}-0.44\%$
test_step_mdp_speed[False-False-True-True-True]	83.1510μs	39.7678μs	25.1460 KOps/s	25.0602 KOps/s	$\color{#35bf28}+0.34\%$
test_step_mdp_speed[False-False-True-True-False]	51.1110μs	27.0954μs	36.9067 KOps/s	36.9239 KOps/s	$\color{#d91a1a}-0.05\%$
test_step_mdp_speed[False-False-True-False-True]	60.8710μs	25.7032μs	38.9056 KOps/s	38.0830 KOps/s	$\color{#35bf28}+2.16\%$
test_step_mdp_speed[False-False-True-False-False]	40.2510μs	16.9644μs	58.9471 KOps/s	59.6254 KOps/s	$\color{#d91a1a}-1.14\%$
test_step_mdp_speed[False-False-False-True-True]	69.4710μs	40.5628μs	24.6531 KOps/s	24.1220 KOps/s	$\color{#35bf28}+2.20\%$
test_step_mdp_speed[False-False-False-True-False]	64.3610μs	28.5695μs	35.0024 KOps/s	34.9476 KOps/s	$\color{#35bf28}+0.16\%$
test_step_mdp_speed[False-False-False-False-True]	63.7810μs	26.7088μs	37.4409 KOps/s	36.7322 KOps/s	$\color{#35bf28}+1.93\%$
test_step_mdp_speed[False-False-False-False-False]	42.5910μs	18.3985μs	54.3521 KOps/s	54.8182 KOps/s	$\color{#d91a1a}-0.85\%$
test_values[generalized_advantage_estimate-True-True]	26.8892ms	26.5237ms	37.7021 Ops/s	41.1057 Ops/s	$\textbf{\color{#d91a1a}-8.28\%}$
test_values[vec_generalized_advantage_estimate-True-True]	85.3261ms	3.2809ms	304.7909 Ops/s	294.9295 Ops/s	$\color{#35bf28}+3.34\%$
test_values[td0_return_estimate-False-False]	99.4520μs	63.0587μs	15.8582 KOps/s	16.0484 KOps/s	$\color{#d91a1a}-1.19\%$
test_values[td1_return_estimate-False-False]	57.4010ms	55.6871ms	17.9575 Ops/s	19.1631 Ops/s	$\textbf{\color{#d91a1a}-6.29\%}$
test_values[vec_td1_return_estimate-False-False]	2.0766ms	1.7772ms	562.6814 Ops/s	568.5309 Ops/s	$\color{#d91a1a}-1.03\%$
test_values[td_lambda_return_estimate-True-False]	91.3586ms	88.1068ms	11.3499 Ops/s	11.9733 Ops/s	$\textbf{\color{#d91a1a}-5.21\%}$
test_values[vec_td_lambda_return_estimate-True-False]	2.0330ms	1.7732ms	563.9418 Ops/s	569.0746 Ops/s	$\color{#d91a1a}-0.90\%$
test_gae_speed[generalized_advantage_estimate-False-1-512]	25.0301ms	24.5421ms	40.7462 Ops/s	42.9272 Ops/s	$\textbf{\color{#d91a1a}-5.08\%}$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512]	0.9019ms	0.7195ms	1.3899 KOps/s	1.4168 KOps/s	$\color{#d91a1a}-1.90\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512]	0.7699ms	0.6689ms	1.4951 KOps/s	1.5306 KOps/s	$\color{#d91a1a}-2.32\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512]	1.5244ms	1.4698ms	680.3614 Ops/s	686.8041 Ops/s	$\color{#d91a1a}-0.94\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512]	0.9430ms	0.6894ms	1.4506 KOps/s	1.4817 KOps/s	$\color{#d91a1a}-2.10\%$
test_dqn_speed	14.1001ms	7.5399ms	132.6281 Ops/s	133.6723 Ops/s	$\color{#d91a1a}-0.78\%$
test_ddpg_speed	15.1960ms	14.2131ms	70.3577 Ops/s	68.7426 Ops/s	$\color{#35bf28}+2.35\%$
test_sac_speed	31.5667ms	29.1954ms	34.2519 Ops/s	33.7814 Ops/s	$\color{#35bf28}+1.39\%$
test_redq_speed	35.8688ms	34.8657ms	28.6815 Ops/s	28.3452 Ops/s	$\color{#35bf28}+1.19\%$
test_redq_deprec_speed	0.1115s	26.0251ms	38.4244 Ops/s	41.1647 Ops/s	$\textbf{\color{#d91a1a}-6.66\%}$
test_td3_speed	19.8728ms	19.6891ms	50.7895 Ops/s	49.7667 Ops/s	$\color{#35bf28}+2.06\%$
test_cql_speed	84.3627ms	82.8129ms	12.0754 Ops/s	11.7438 Ops/s	$\color{#35bf28}+2.82\%$
test_a2c_speed	27.3341ms	26.6295ms	37.5523 Ops/s	36.9707 Ops/s	$\color{#35bf28}+1.57\%$
test_ppo_speed	0.1240s	30.6122ms	32.6667 Ops/s	36.4894 Ops/s	$\textbf{\color{#d91a1a}-10.48\%}$
test_reinforce_speed	26.8080ms	25.5459ms	39.1452 Ops/s	38.1420 Ops/s	$\color{#35bf28}+2.63\%$
test_iql_speed	57.9376ms	57.0058ms	17.5421 Ops/s	17.1745 Ops/s	$\color{#35bf28}+2.14\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	0.1043s	2.1183ms	472.0673 Ops/s	516.4896 Ops/s	$\textbf{\color{#d91a1a}-8.60\%}$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	2.1171ms	0.7854ms	1.2732 KOps/s	1.2547 KOps/s	$\color{#35bf28}+1.47\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	0.9759ms	0.7729ms	1.2939 KOps/s	1.2724 KOps/s	$\color{#35bf28}+1.69\%$
test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	2.3997ms	1.8477ms	541.2088 Ops/s	518.9844 Ops/s	$\color{#35bf28}+4.28\%$
test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	0.9050ms	0.7741ms	1.2918 KOps/s	1.2747 KOps/s	$\color{#35bf28}+1.34\%$
test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	4.7704ms	0.7687ms	1.3008 KOps/s	1.2792 KOps/s	$\color{#35bf28}+1.69\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	2.9086ms	2.1539ms	464.2637 Ops/s	450.6544 Ops/s	$\color{#35bf28}+3.02\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	3.4198ms	0.9079ms	1.1014 KOps/s	1.0813 KOps/s	$\color{#35bf28}+1.85\%$
test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	1.0609ms	0.8939ms	1.1187 KOps/s	1.0943 KOps/s	$\color{#35bf28}+2.23\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]	2.6529ms	1.9082ms	524.0633 Ops/s	513.1449 Ops/s	$\color{#35bf28}+2.13\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]	0.9034ms	0.7845ms	1.2747 KOps/s	1.2545 KOps/s	$\color{#35bf28}+1.61\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]	4.9006ms	0.7790ms	1.2836 KOps/s	1.2630 KOps/s	$\color{#35bf28}+1.63\%$
test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]	2.0540ms	1.8829ms	531.0982 Ops/s	524.3690 Ops/s	$\color{#35bf28}+1.28\%$
test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]	2.1838ms	0.7779ms	1.2855 KOps/s	1.2704 KOps/s	$\color{#35bf28}+1.18\%$
test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]	0.9772ms	0.7666ms	1.3044 KOps/s	1.2876 KOps/s	$\color{#35bf28}+1.31\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]	2.7432ms	2.1754ms	459.6852 Ops/s	451.4313 Ops/s	$\color{#35bf28}+1.83\%$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]	0.1488s	1.0871ms	919.8365 Ops/s	1.0789 KOps/s	$\textbf{\color{#d91a1a}-14.75\%}$
test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]	1.0647ms	0.8960ms	1.1161 KOps/s	1.0908 KOps/s	$\color{#35bf28}+2.31\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]	0.1225s	15.4596ms	64.6848 Ops/s	54.8339 Ops/s	$\textbf{\color{#35bf28}+17.96\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]	16.1022ms	13.0119ms	76.8529 Ops/s	79.6047 Ops/s	$\color{#d91a1a}-3.46\%$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]	5.7165ms	1.9132ms	522.6843 Ops/s	531.7753 Ops/s	$\color{#d91a1a}-1.71\%$
test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]	0.1252s	17.7181ms	56.4396 Ops/s	65.0577 Ops/s	$\textbf{\color{#d91a1a}-13.25\%}$
test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]	15.8389ms	13.0402ms	76.6858 Ops/s	67.4876 Ops/s	$\textbf{\color{#35bf28}+13.63\%}$
test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]	2.5414ms	1.8112ms	552.1264 Ops/s	520.4906 Ops/s	$\textbf{\color{#35bf28}+6.08\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]	0.1233s	17.8210ms	56.1137 Ops/s	64.5241 Ops/s	$\textbf{\color{#d91a1a}-13.03\%}$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]	16.4647ms	13.2358ms	75.5526 Ops/s	78.2899 Ops/s	$\color{#d91a1a}-3.50\%$
test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]	2.7093ms	1.9878ms	503.0661 Ops/s	496.2066 Ops/s	$\color{#35bf28}+1.38\%$

[vmoens]

Cc @nicklashansen I think this is more or less ready, if you want to give it a look!

[nicklashansen]

Thanks! Seems to run quite smoothly for me. As one might expect, streaming is extremely slow but practical for debugging; download works like a charm!

One question I have: would there be a computational advantage to only sampling specific keys (provided by the user) rather than sampling the full TensorDict and then only using some of them? E.g. I might want to use only a single image per step, or no images, or leave out the next data field. Is that (a) feasible to implement, and (b) even necessary?

[vmoens]

This should do the trick if you want to exclude ("next", "reward") and "observation" entries:

from torchrl.envs import SelectTransform, ExcludeTransform 
dataset = OpenXExperienceReplay(..., transform=ExcludeTransform(("next", "reward"), "observation"))

SelectTransform has the opposite behaviour.

Would that help?

[nicklashansen]

Yea! I guess my question is whether there is any computational advantage to excluding / selecting data as a transform vs. sampling the full TensorDict and just ignoring that the extra fields are there?

[vmoens]

Yea! I guess my question is whether there is any computational advantage to excluding / selecting data as a transform vs. sampling the full TensorDict and just ignoring that the extra fields are there?

Valid concern.
One solution could be to only store what you need, but (1) it's adding yet another argument to the constructor which is quite loaded already and (2) I don't think the first time you use a dataset should condition what you will do with it later (someone could choose to ignore the reward the first time the dataset is downloaded and use it the next time).

All datasets are thin wrappers around a composed replay buffer, so what I would do to speed up indexing would be

dataset = OpenXExperienceReplay(...)
filtered_storage = TensorStorage(dataset._storage._storage.exclude(*keys)) # you can also use .select(*keys)
dataset = ReplayBuffer(
   storage=filtered_storage,
   sampler=dataset._sampler,
   _collate_fn=dataset._collate_fn,
   writer=dataset._writer,
   transform=dataset._transform,
   batch_size=dataset._batch_size,
)   

Alternatively you can also just use copy.copy and patch the _storage attribute.
A bit verbose but works!

[nicklashansen]

Thanks! This works wonders.

[quanvuong]

do you still need my input : ) ?

[nicklashansen]

@quanvuong We have a working dataloader now, so no need I think! Thanks though :-)