[Feature] VecNormV2

[vmoens]

Stack from ghstack (oldest at bottom):

• -> [Feature] VecNormV2 #2867

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 14 New Failures, 8 Unrelated Failures

As of commit f1f373f with merge base 7df8317 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / build-wheel-py3_9-cpu (gh)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_9-cuda11_8 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu118_
• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 30afb0895d343c3995e1c90e607ea3da3019b618e762ee35afd68c27253bbccb /exec failed with exit code 1
• LLM Tests on Linux / unittests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 371ae66446cdd27193207fb2d24ad128def27f960f9628a51fd2f85b30caf09e /exec failed with exit code 247
• SOTA Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t ba726d0b1b7276615411808adabbd470ceddd1a884758c09c3df669eebc556e0 /exec failed with exit code 1
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  test/test_tensordictmodules.py::TestTDSequence::test_stateful_probabilistic[True-bounded-False]
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  test/test_modules.py::TestMultiAgent::test_multiagent_mlp_lazy
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_tensordictmodules.py::TestTDSequence::test_stateful[True-bounded-False]
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  test/test_tensordictmodules.py::TestTDModule::test_stateful_probabilistic[deterministic-True-out_keys0-bounded-True]
• Unit-tests on Linux / tests-gpu (3.11, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 05a9d789557effd48742b9f5f1d7cfbfa382ecf706632318c06069b31768dbf5 /exec failed with exit code 250
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 4227a89a7066bad9f2a412a4843111fd8ec27ba6f152f45aca5bc9a39f4457c6 /exec failed with exit code 1
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh)
  test/test_tensordictmodules.py::TestTDSequence::test_stateful_probabilistic[True-bounded-False]

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

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

• Build Windows Wheels / pytorch/rl / build-wheel-py3_9-cuda11_8 (gh) (trunk failure)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/rl / build-wheel-py3_9-cuda12_6 (gh) (trunk failure)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/rl / build-wheel-py3_9-cuda12_8 (gh) (trunk failure)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_9-cpu (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_9-cuda12_6 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu126_
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_9-cuda12_8 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu128_
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh) (trunk failure)
  test/test_transforms.py::TestTimer::test_transform_model
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh) (trunk failure)
  test/test_transforms.py::TestTimer::test_transform_env

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

[vmoens]

There is a noticeable difference with the first version which is that here we first normalize then update, whereas the other version first updates then normalizes.
I think it's better to use an unbiased estimate but i'm open to align to the previous version.

List of TODOs:

• Better coverage of stateless
• warning about deprec of v1
• docstrings
• add to doc
• test collision between reward and obs specs, esp if 2 vecnorm are stacked

I guess long-term we should think about collectives for distributed but i'll leave this for enthusiastic young minds