[Feature] Complete DreamerV3 continuation training semantics

[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:

• share a trained continuation head with imagination and use it in lambda returns
• propagate cumulative continuation weights to actor and value objectives
• add episode-reset masking, symlog vector encoding, and RMS-normalized DreamerV3 heads
• align the maintained optimizer, warmup, initialization, and horizon settings

Rationale:
A continuation predictor is only useful when its probabilities govern imagined
returns and objective weighting. Previously the optional auxiliary loss did not
affect imagination, recurrent carries were not explicitly reset at episode
boundaries, and the maintained setup used materially different preprocessing and
optimization settings. These gaps prevented meaningful reference-curve comparison.

Test plan:

• pytest test/objectives/test_dreamer_v3.py -q
• pytest test/modules/test_dreamer_components.py -k DreamerV3Components -q
• pytest --doctest-modules torchrl/modules/models/model_based_v3.py torchrl/objectives/dreamer_v3.py -q
• run the maintained DreamerV3 SOTA smoke with 400 environment steps
• verify continuation lambda fixtures, cumulative weights, shared parameters,
  terminal target scaling, recurrent resets, symlog inputs, AGC, and zero-scale heads

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 049ac18 with merge base d7659c7 ():
💚 Looks good so far! There are no failures yet. 💚

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

[unography]

@vmoens actor seems to be missing observation?
jax version uses obs + belief + state

[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 5eed361d224f. 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.

I attempted the requested rebase of current head 416937e85f onto the freshly fetched origin/main; it stops at commit 66463c675 with a content conflict in sota-implementations/dreamer_v3/dreamer_v3.py, so I left the rebase in progress and did not push a partial rewrite.

I reviewed the existing head and the latest inline comment meanwhile. The cited DreamerV3 implementation feeds the policy feat2tensor(feat), which concatenates the deterministic and flattened stochastic RSSM features; those correspond to this implementation's belief and state, not the raw observation. So the current actor in_keys=["state", "belief"] look correct and I do not think an observation input is missing. I also ran the two focused suites on this head: 266 passed (test/objectives/test_dreamer_v3.py and test/modules/test_dreamer_components.py). The remaining blocker is resolving the rebase conflict and rerunning the maintained SOTA smoke on the rebased result.