[Feature] WorldModel + WorldModelLoss — general model-based RL abstraction by theap06 · Pull Request #3783 · pytorch/rl

theap06 · 2026-05-20T07:27:11Z

Adds `WorldModel(TensorDictModuleBase)`: a key-driven, architecture-agnostic composition layer for encoder + dynamics + reward/done/decoder heads, with `encode`, `step`, `decode` shortcuts and a `rollout(start_td, policy, horizon)` method whose `[batch, horizon]` output matches `EnvBase.rollout` — making imagined trajectories drop-in compatible with replay buffers, GAE, and loss modules.
Adds `WorldModelLoss(LossModule)`: follows the standard `_AcceptedKeys` / `set_keys()` / `forward() → TensorDict` pattern; supports configurable sub-losses (`reward`, `done`, `reconstruction`, `latent`) with per-loss weights and distance-function choice.
Existing Dreamer components (`WorldModelWrapper`, `DreamerEnv`, `RSSMRollout`, `DreamerModelLoss/ActorLoss/ValueLoss`) are unchanged.

Motivation

TorchRL ships Dreamer-specific model-based components but no general abstraction. Users implementing MBPO, TD-MPC, PlaNet, or any custom world model must hand-wire `TensorDictModule` chains and write bespoke multi-step rollout loops. This adds the missing general layer.

Test plan

`python -m pytest test/test_world_model.py -v` — 19 new tests covering forward, encode/step/decode shortcuts, nested keys, rollout shape and early termination, replay buffer compatibility, all four loss types, `set_keys`, per-loss weights, and gradient flow
`python -c "from torchrl.modules import WorldModel; from torchrl.objectives import WorldModelLoss; print('OK')"` — import smoke test
Existing dreamer-related tests (`pytest test/ -k "dreamer" --ignore=test/llm`) — 409 tests, all passing

pytorch-bot · 2026-05-20T07:27:15Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 6 Pending, 1 Unrelated Failure

As of commit 24f33c9 with merge base 4679d0a ():

NEW FAILURES - The following jobs have failed:

Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
Process completed with exit code 1.
Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
Process completed with exit code 1.
Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh)
test/libs/test_gym.py::TestGym::test_gym_gymnasium_parallel

BROKEN TRUNK - The following job failed but were present on the merge base:

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

Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh) (trunk failure)
##[error]The operation was canceled.

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

github-actions · 2026-05-20T07:27:23Z

⚠️ PR Title Label Error

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

Current title: WorldModel + WorldModelLoss — general model-based RL abstraction

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]).

github-actions · 2026-05-20T07:27:27Z

⚠️ PR Title Label Error

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

Current title: WorldModel + WorldModelLoss — general model-based RL abstraction

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]).

github-actions · 2026-05-20T07:27:30Z

⚠️ PR Title Label Error

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

Current title: WorldModel + WorldModelLoss — general model-based RL abstraction

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]).

github-actions · 2026-05-20T07:27:31Z

⚠️ PR Title Label Error

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

Current title: WorldModel + WorldModelLoss — general model-based RL abstraction

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]).

theap06 · 2026-05-21T05:36:48Z

@vmoens @elin-bdai this was my attempt at creating a World Model abstraction. Lmk if you have any feedback!

theap06 · 2026-05-22T00:06:12Z

@vmoens I believe this issue is tied to the flaky tests from earlier. Lmk if the design seems sound to you.

vmoens

I'm not sure why WorldModel is a not an env - or a module that's used by a model-based env to generate rollouts. That's the usual way to think about this in torchrl (see Dreamer for example). Can you elaborate on this choice>

Per review feedback (#vmoens): move imagined-rollout responsibility off the WorldModel module and onto a new WorldModelEnv(ModelBasedEnvBase), matching the pattern used by DreamerEnv / ImaginedEnv. Eliminates the divergent rollout semantics WorldModel.rollout had vs EnvBase.rollout.

vmoens

Just some minor comments but otherwise LGTM!

@vmoens

Per @vmoens review comments on the WorldModelEnv PR: - WorldModelEnv: - Annotate `world_model` parameter as `WorldModel` (via TYPE_CHECKING to avoid the torchrl.modules <-> torchrl.envs circular import). - Replace `.clone(recurse=False)` with `.copy()`. - Drop unnecessary `{}` in `TensorDict(batch_size=..., device=...)`. - Iterate spec leaves with `include_nested=True, leaves_only=True` so nested observation/reward/done specs survive `_step` end-to-end. - Use `full_*_spec.zero()` (and `full_*_spec[key].zero()`) instead of hand-rolled `torch.zeros(...)` for default reward/done/terminated tensors, keeping shapes / dtypes consistent with the specs. - WorldModelLoss: drop unused `batch_size=[]` from final TensorDict construction. - Tests: consolidate `test/test_world_model.py` (deleted) into: - `test/envs/test_model_based.py` — TestWorldModelForward, TestWorldModelEnv (and a `_SpecOnlyEnv` helper). - `test/objectives/test_dreamer.py` — TestWorldModelLoss. No public-API changes. 22/22 affected tests pass locally.

vmoens

LGTM
do you think we should consider subclassing the dreamer classes with these for a more homogeneous implementation?

theap06 · 2026-06-10T09:51:13Z

LGTM do you think we should consider subclassing the dreamer classes with these for a more homogeneous implementation?

I can open an issue with the refactors for the existing dreamer classes.

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 20, 2026

github-actions Bot added Documentation Improvements or additions to documentation Objectives Modules Integrations/torch_geometric Integrations labels May 20, 2026

theap06 changed the title ~~WorldModel + WorldModelLoss — general model-based RL abstraction~~ [Feature] WorldModel + WorldModelLoss — general model-based RL abstraction May 20, 2026

github-actions Bot added the Feature New feature label May 20, 2026

WorldModel + WorldModelLoss — general model-based RL abstraction

b21a251

theap06 force-pushed the feat/World_Model branch from ddb8e36 to b21a251 Compare May 26, 2026 07:30

vmoens reviewed Jun 8, 2026

View reviewed changes

Comment thread torchrl/modules/tensordict_module/world_models.py

theap06 added 2 commits June 8, 2026 19:09

Merge branch 'main' into feat/World_Model

583ccc0

theap06 requested a review from vmoens June 9, 2026 02:09

vmoens reviewed Jun 9, 2026

View reviewed changes

theap06 requested a review from vmoens June 10, 2026 07:43

vmoens approved these changes Jun 10, 2026

View reviewed changes

vmoens merged commit 78c12fe into pytorch:main Jun 10, 2026
104 of 108 checks passed

Conversation

theap06 commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Test plan

Uh oh!

pytorch-bot Bot commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 3 New Failures, 6 Pending, 1 Unrelated Failure

Uh oh!

github-actions Bot commented May 20, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

github-actions Bot commented May 20, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

github-actions Bot commented May 20, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

github-actions Bot commented May 20, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

theap06 commented May 21, 2026

Uh oh!

theap06 commented May 22, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

theap06 commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

theap06 commented May 20, 2026 •

edited

Loading

pytorch-bot Bot commented May 20, 2026 •

edited

Loading