[BugFix] Decode DreamerV3 rewards and values end to end

[vmoens]

Stack from ghstack (oldest at bottom):

• [Doc] Explain DreamerV3 architecture and nomenclature #4095
• [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:

• add configurable reward_logits and value_logits NestedKeys while reserving reward and value keys for decoded scalars
• preserve legacy logits-under-scalar-key inputs with v0.16 deprecation warnings
• expose categorical critic logits and decoded values to lambda returns
• share the trained prior and categorical reward head with imagination
• configure the maintained DreamerV3 example for 255-bin reward and value cross entropy

Rationale:
Training consumes categorical logits, whereas environment rewards, critic baselines, and TD-lambda returns require decoded scalars. Separating these representations prevents logits from leaking into return computation, and sharing the trained heads ensures imagination reflects the optimized world model.

Test plan:

• uv run pytest -q test/objectives/test_dreamer_v3.py
• run the maintained DreamerV3 SOTA smoke with 11-bin CI overrides
• 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/4066

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit 54337d9 with merge base d5a62cb ():

BROKEN TRUNK - The following jobs failed but were 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/rb/test_storages.py::TestStorages::test__rand_given_ndim_recompile
• Unit-tests on Linux / tests-cpu (3.13, bulk) / linux-job (gh) (trunk failure)
  test/rb/test_storages.py::TestStorages::test__rand_given_ndim_recompile

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

[vmoens]

Update after the stack-level ghstack rebase: The full eight-PR Dreamer stack was checked out from #4075, rebased onto current origin/main, linted with pre-commit run --all-files, and submitted successfully as vmoens. This PR is now at ba23945a1465. The core Dreamer suites pass (267 tests) and the two doctest modules pass (15 tests); the DMC-only smoke still requires dm_control, which is not installed locally. Any earlier conflict/head-status note below is superseded; the substantive code-review observations remain applicable.

Reviewed published head 6df1724eea04, including the full layer diff, the latest commit, current checks, and the complete PR conversation (there are no human review threads yet). I also attempted the requested rebase after fetching origin/main, but it stopped on a content conflict in sota-implementations/dreamer_v3/dreamer_v3.py; per the rebase workflow, I left that rebase in progress and did not push a partial result.

Within the published patch, the separation of categorical logits from decoded reward/value keys is coherent: the loss reads the configurable logits keys, the legacy fallback has an explicit v0.16 deprecation, nested-key coverage is present, and imagination shares the trained prior/reward head instead of independent parameters. I did not find an additional code defect in this layer. The four new focused tests (decoded critic value, legacy warning, nested logits keys, and shared imagination parameters) pass, and git diff --check is clean.

Residual risk: this head is still based on the pre-rebase stack, so it has not been validated against current origin/main. The existing bot-reported lint failure is an infrastructure/container failure rather than a source diagnostic. Please resolve the DreamerV3 script conflict and rerun the focused suite plus the SOTA smoke on the rebased result.

[vmoens]

Addressed the review pass and resubmitted this commit after rebasing the stack onto current main. No layer-specific code change was needed; the categorical logits remain separated from decoded reward/value scalars and the shared imagination heads are preserved. The full all-files lint suite passes at this PR boundary.