🔗 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 ():

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

$\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$.

$\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$.

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

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?

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?

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?

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!

Thanks! This works wonders.

do you still need my input : ) ?

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