[Feature] CompressedStorage

[vmoens]

@AdrianOrenstein feedback on this?

[pytorch-bot]

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

NEW FAILURES - The following jobs have failed:

• Build M1 Wheels / pytorch/rl / build-wheel-py3_9-cpu (gh)
  Process completed with exit code 127.
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  RuntimeError: PassManager::run failed
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 451bca60cf7e529adbe8616fd07f542f3434d7341739ef688fbdb73f36d96517 /exec failed with exit code 1
• Lint / python-source-and-configs / linux-job (gh)
  RuntimeError: Command docker exec -t f54bc5c56a0c51d19e8e924398e1f6c7ea15b272d45799a86529675bb8446b67 /exec failed with exit code 1
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression
• Unit-tests on Linux / tests-gpu (3.11, 12.8) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
  test/test_rb.py::TestCompressedStorage::test_compressed_storage_custom_compression

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

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

• Build M1 Wheels / pytorch/rl / upload / upload-wheel-py3_9-cpu (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cpu_
• Continuous Benchmark / CPU Pytest benchmark (gh) (trunk failure)
• Continuous Benchmark / GPU Pytest benchmark (gh) (trunk failure)
  RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor []] is at version 4; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh) (trunk failure)
  test/test_libs.py::TestGym::test_gym_fake_td[True-False-3-HalfCheetah-v2]
• LLM Tests on Linux / unittests (3.9, 12.8) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• Unit-tests on Linux / tests-olddeps (3.9, 11.6) / linux-job (gh) (trunk failure)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[pendulum-SamplingStrategy.RANDOM-False-True]

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

[AdrianOrenstein]

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.

[vmoens]

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!

[AdrianOrenstein]

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.

[vmoens]

Closing in favor of #3062