[Feature] Normalize DreamerV3 REINFORCE return scale

[vmoens]

Stack from ghstack (oldest at bottom):

• [Algorithm] Add a reproducible DreamerV3 DMC Walker benchmark #4075
• [Feature] Complete DreamerV3 continuation training semantics #4074
• [BugFix] Match DreamerV3 categorical and KL semantics #4073
• [Feature] Add the DreamerV3 block RSSM core #4072
• [Feature] Add DreamerV3 slow-critic regularization #4068
• -> [Feature] Normalize DreamerV3 REINFORCE return scale #4067
• [BugFix] Decode DreamerV3 rewards and values end to end #4066
• [BugFix] Correct DreamerV3 symexp two-hot semantics #4065

Summary:

• checkpoint EMA fifth- and ninety-fifth-percentile return statistics in the actor loss
• update them at rate 0.01 without debiasing and clamp the span to at least 1.0
• divide detached REINFORCE advantages by the uncentered span
• expose low, high, and scale diagnostics plus normalization configuration
• enable the maintained example with the paper defaults

Rationale:
DreamerV3 normalizes the scale of REINFORCE advantages so reward magnitude does not directly determine actor-gradient magnitude. Persisting the EMA state keeps resumed training numerically continuous while evaluation remains read-only.

Test plan:

• uv run pytest -q test/objectives/test_dreamer_v3.py
• run the maintained DreamerV3 SOTA smoke with 11-bin CI overrides
• verify the normalized REINFORCE equation, EMA updates, eval stability, torch.compile, gradients, and state-dict round trip
• flake8 --config=setup.cfg on the changed Python files
• ufmt check on the changed Python files
• git diff --check

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 3 Unrelated Failures

As of commit 194e534 with merge base 2a77e48 ():

NEW FAILURES - The following jobs have failed:

• Lint / lint-done (gh)
• Lint / python-source-and-configs / linux-job (gh)
  RuntimeError: Command docker exec -t 79cf88944f2604596bc59e5883df9a2ac6f61790ff66df40d4f8f82e158868fc /exec failed with exit code 1

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

• Unit-tests on Linux / tests-cpu (3.14, bulk) / linux-job (gh) (similar failure)
  test/test_configs.py::TestConfigClassParity::test_wrapped_class_kwargs_have_config_fields[LBFGSConfig]

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

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

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh) (trunk failure)
  test/test_configs.py::TestConfigClassParity::test_wrapped_class_kwargs_have_config_fields[LBFGSConfig]
• Unit-tests on Linux / tests-cpu (3.13, bulk) / linux-job (gh) (trunk failure)
  test/test_configs.py::TestConfigClassParity::test_wrapped_class_kwargs_have_config_fields[LBFGSConfig]

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

[vmoens]

Reviewed published head 194e534a1a57, including the complete layer diff, latest commit/check state, and the full PR conversation (no human review threads are present). The requested rebase onto freshly fetched origin/main stopped on a content conflict in sota-implementations/dreamer_v3/dreamer_v3.py; I preserved the in-progress rebase and did not push anything partial.

The EMA percentile-span normalization is applied only to detached REINFORCE advantages, its buffers are checkpointed, validation covers rates/quantiles/minimum scale, and the reparameterization path remains unscaled. I found no additional code defect in this layer. Both new focused tests (normalization behavior and checkpoint restoration) pass, and git diff --check is clean.

Residual risk: the published head has not been tested against current main because of the conflict, and torch.quantile remains the main compile/performance-sensitive operation in the actor hot path. The bot lint failure on the published head is a runner/container failure, not a reported lint diagnostic. Please resolve the script conflict, then rerun the focused tests and compile/SOTA coverage on the rebased head.