[Feature] Support lazy tensordict inputs in KL reward transform

[vmoens]

Stack from ghstack (oldest at bottom):

• -> [Feature] Support lazy tensordict inputs in KL reward transform #2884

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 11 New Failures, 2 Unrelated Failures

As of commit b88c700 with merge base e1d3fd4 ():

NEW FAILURES - The following jobs have failed:

• 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 67b19bd07567b037e4411894b0493f233eb678d4793ec8825624011b519c0e33 /exec failed with exit code 1
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-gpu (3.11, 12.8) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 894d80e01cdb51343fec6d58cba8769f5b5876f2fbf93481637ae5d74bad078f /exec failed with exit code 1
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm

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

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

• LLM Tests on Linux / unittests (3.9, 12.8) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• 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.