🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 10 New Failures, 6 Unrelated Failures

As of commit 3da84a9 with merge base 7e8f940 ():

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

I think the PR is great. If I get permissions to modify you could delegate some of the tasks to me if you'd like.

I've left some comments below:

• The conversion of an object to a bytestream is hardcoded. If we move that into a function, it'll be easier to later extend this class to serialise the other objects like numpy.arrays or a pickle dump of some object.

• I don't think we had any tests on numpy objects in the tensordict or arbitrary python objects.

• Is the goal to keep the scope of the PR just for CPU data, CPU storage, CPU compression, and CPU decompression? If so, I might follow up with a PR to include some GPU compression. I like the API and it is in a state where I'd be able to extend it.

• In docs/source/reference/data.rst you mention a 2-10x compression ratio. I was thinking of including pytest-benchmark as a dependancy and then we can get a number on compressing atari frames. This might also be out of scope so I could do it later.

• There is pip install zstd that provides a stateless compression and decompression API which might be more convenient.

• If you're open to adding pytest-benchmark as a dependancy, we could make this a regression test. It might make sense to add an Atari rollout to evaluate the expected compression ratio.

All of these look good! I think it's going to be an amazing feature!

I don't think we want a separate PR for GPU. Not particularly opinionated though, if you feel like splitting this in separate chunks I'm good with it.
Also you can just branch out of this branch and make a new PR with your changes, I'll close this one.
Thanks!

Great, I'll get started shortly on the gpu side of things. Also thanks for getting me started with this pr, it has helped me wrap my head around a lot of the details of the codebase.

Closing in favor of #3062