[Feature] Add the DreamerV3 block RSSM core

[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 the grouped block-linear recurrent core used by DreamerV3
• add RMS-normalized prior and posterior predictor options
• preserve the existing GRU implementation as the default for checkpoint compatibility

Rationale:
The maintained DreamerV3 implementation used a conventional GRUCell, while the
reference architecture projects deterministic state, stochastic state, and action
separately before applying grouped recurrent gates. Providing the block core as an
explicit mode enables architecture-parity experiments without invalidating existing
TorchRL models.

Test plan:

• pytest test/modules/test_dreamer_components.py -k DreamerV3Components -q
• pytest test/objectives/test_dreamer_v3.py -q
• verify the frozen numerical fixture, normalized-action invariance, gradients,
  RMS-normalized posterior shapes, and torch.compile full-graph execution

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1ea75bd 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 c58247eb9921. 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 current PR head a1dc57e7fd08d6d74867770fa45f0eed334a8a78 and the latest CI/bot report. The requested rebase onto freshly fetched origin/main could not complete: replaying the lower-stack base-update commit 0b8b08d65 conflicts in sota-implementations/dreamer_v3/dreamer_v3.py. I left that rebase in progress and did not force-push.

On the PR-specific diff, I checked the grouped block projections/gating, action normalization, RMS-normalized prior/posterior paths, the frozen numerical fixture, gradients, and full-graph compile coverage. pytest -q test/modules/test_dreamer_components.py -k DreamerV3Components passes (4 passed). I did not find an additional functional issue at this head, but the stack conflict must be resolved before it can be refreshed.

The current source/config CI failure is actionable but outside these two changed files: µfmt wants the SymExpTwoHot import reordered in torchrl/modules/__init__.py; the reported bulk LBFGSConfig parity failures are present on trunk/merge-base.

[vmoens]

Addressed the review pass and resubmitted this commit after rebasing the stack onto current main. No layer-specific code change was needed; block-RSSM validation, dtype preservation, action normalization, and compile coverage are preserved. The full all-files lint suite passes at this PR boundary.