[BugFix] Fix KL penalty

[vmoens]

Stack from ghstack (oldest at bottom):

• [Quality] Better device checks #2909
• -> [BugFix] Fix KL penalty #2908

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 5 New Failures, 1 Cancelled Job, 9 Unrelated Failures

As of commit a6d2c96 with merge base f5f3ae4 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Generate documentation / build-docs (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 22ad63395be1d6973c82da9f9e3b296a082091c6af8b2eb2346d0c1025e59ab3 /exec failed with exit code 2
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 2800599bcb13cdf9d2281a5826431f162fa523b4520d27242e09d806c8b1d7b5 /exec failed with exit code 1
• SOTA Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 2ee802fa215601ed832f2a448f9e2186337c0a71e3366a31b175b98eaa534b48 /exec failed with exit code 1
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 3c73f1dba05b64b95bf1a9c81e9ce22bf77140e5b836b976c1b94cd315dd59de /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh) (similar failure)
  test/test_env.py::TestParallel::test_multitask

BROKEN TRUNK - The following jobs failed but was 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 Linux / tests-cpu (3.10) / linux-job (gh) (trunk failure)
  test/test_env.py::TestParallel::test_multitask
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh) (trunk failure)
  test/test_env.py::TestParallel::test_multitask
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh) (trunk failure)
  test/test_env.py::TestNonTensorEnv::test_parallel[False-False]
• Unit-tests on Linux / tests-gpu (3.11, 12.8) / linux-job (gh) (trunk failure)
  test/test_env.py::TestParallel::test_multi_task_serial_parallel[True-stand,stand,stand]
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh) (trunk failure)
  test/test_collector.py::TestPolicyFactory::test_weight_update
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh) (trunk failure)
  test/test_env.py::TestParallel::test_multi_task_serial_parallel[True-stand,stand,stand]
• 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.