[Feature] Implements DreamerV3 (Mastering Diverse Domains in World Models, Hafn…

[theap06]

Fixes #3613

Implements DreamerV3 (Mastering Diverse Domains in World Models, Hafner et al. 2023) as a set of composable loss modules and latent model components on top of the existing DreamerV1 infrastructure.

New files:

torchrl/objectives/dreamer_v3.py: DreamerV3ModelLoss, DreamerV3ActorLoss, DreamerV3ValueLoss; symlog/symexp transforms; two_hot_encode/two_hot_decode; categorical_kl_balanced with KL balancing (α=0.8) and free bits
torchrl/modules/models/model_based_v3.py: RSSMPriorV3, RSSMPosteriorV3, RSSMRolloutV3 with discrete categorical latent and straight-through gradient estimator
Updated:

torchrl/objectives/init.py: exports V3 losses and utility functions
torchrl/modules/models/init.py: exports V3 RSSM classes
test/test_objectives.py: TestDreamerV3 class with 30 tests (all passing)
docs/source/reference/objectives_other.rst: DreamerV3 API reference section
docs/source/reference/modules_models.rst: V3 RSSM classes in world models section
Motivation and Context
TorchRL supports DreamerV1/V2 but not DreamerV3, which introduces key improvements: discrete categorical latent states, KL balancing between prior and posterior, symlog-compressed targets, and two-hot categorical value distributions. These changes significantly improve stability and sample efficiency across diverse domains. This PR adds V3 as a self-contained extension — V1 is untouched and existing users are unaffected.

I have raised an issue to propose this change (required for new features and bug fixes)
Types of changes
New feature (non-breaking change which adds core functionality)
Documentation (update in the documentation)
Checklist
I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull workflow for PyTorch trunk commits

❌ 13 New Failures, 1 Cancelled Job

As of commit 182811b with merge base be89a87 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  RuntimeError: Command docker exec -t 113bcdb47cbac7f496f11ee7ef15c0dfc85a63d134f8ab5c224de331934f9616 /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  RuntimeError: Command docker exec -t 0e35f065306ea4f5d0f1f922a943f766e62bb54b239df1d166953b4117eb8d00 /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  RuntimeError: Command docker exec -t 5fb101aed2106b3747e848259e1e32d874babd967fb759c1e9599e6080e7823e /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.13) / linux-job (gh)
  RuntimeError: Command docker exec -t 249a94df27e7aa9b5da3bff806b7f31824e5e3d3efeebcac1d055b20f2ccbd3a /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.14) / linux-job (gh)
  RuntimeError: Command docker exec -t 009178223be151186fa6a1c3d69a38539d66ca86282001c4c3085e3e361ef27d /exec failed with exit code 2
• Unit-tests on Linux / tests-gpu (3.12, 13.0, 3) / linux-job (gh)
  RuntimeError: Command docker exec -t 51789f3a5f27d5296668a37e67e23b1460b77ffc4685bab7ce18b84a1122bd00 /exec failed with exit code 1
• Unit-tests on Linux / tests-gpu-distributed (3.12, 13.0) / linux-job (gh)
  test/test_distributed.py::TestRayTrajsPerBatch::test_ray_trajs_per_batch_replay_buffer
• Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 35bbb5f48bc5c72e68970478a472831dcc99935c68ac184db00af536e0212dcf /exec failed with exit code 2
• Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh)
  RuntimeError: Command docker exec -t 21321934cb9e94d5347bc95e1731c012513570ca9823b5280e71fdb0371db907 /exec failed with exit code 2
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 3) / linux-job (gh)
  RuntimeError: Command docker exec -t bb4a3e04c53c39f1cd216757555f8646e89fd1c9182ada5762d1c973f093c554 /exec failed with exit code 1
• Unit-tests on Linux / tests-stable-gpu-distributed (3.12, 13.0) / linux-job (gh)
  test/test_distributed.py::TestRayTrajsPerBatch::test_ray_trajs_per_batch_replay_buffer
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)
  Process completed with exit code 2.

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

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration

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

[github-actions]

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Implements DreamerV3 (Mastering Diverse Domains in World Models, Hafn…

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Quality]	Quality	[Quality] Fix typos and add codespell

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).