ED2: Environment Dynamics Decomposition World Models for Continuous Control

TL;DR

ED2 addresses the core challenge of model error in model-based RL by explicitly decomposing environment dynamics into sub-dynamics. It introduces SD2 to automatically partition the action space and D2P to build a decomposed, end-to-end trainable world model that aggregates multiple sub-dynamics. Empirical results show ED2 reduces model errors, boosts sample efficiency, and improves asymptotic performance across diverse continuous-control tasks when combined with Dreamer, MBPO, or TDMPC, with clustering-based SD2 often delivering the strongest gains. This framework offers a practical, architecture-agnostic backbone for advancing MBRL by aligning world-model structure with the intrinsic dynamics and action causes of the environment.

Abstract

Model-based reinforcement learning (MBRL) achieves significant sample efficiency in practice in comparison to model-free RL, but its performance is often limited by the existence of model prediction error. To reduce the model error, standard MBRL approaches train a single well-designed network to fit the entire environment dynamics, but this wastes rich information on multiple sub-dynamics which can be modeled separately, allowing us to construct the world model more accurately. In this paper, we propose the Environment Dynamics Decomposition (ED2), a novel world model construction framework that models the environment in a decomposing manner. ED2 contains two key components: sub-dynamics discovery (SD2) and dynamics decomposition prediction (D2P). SD2 discovers the sub-dynamics in an environment automatically and then D2P constructs the decomposed world model following the sub-dynamics. ED2 can be easily combined with existing MBRL algorithms and empirical results show that ED2 significantly reduces the model error, increases the sample efficiency, and achieves higher asymptotic performance when combined with the state-of-the-art MBRL algorithms on various continuous control tasks. Our code is open source and available at https://github.com/ED2-source-code/ED2.

Figures (13)

• Figure 1: Motivation Example. The Walker task has six action dimensions, including $\{\textit{hip}, \textit{knee}, \textit{ankle}\}$ for the $\{\textit{left}, \textit{right}\}$ limbs. Decomposability: The dynamics can be decomposed into multiple sub-dynamics in various ways, each sub-dynamics is described with different background colors. Each circle corresponds to a joint including single or multiple action dimensions. Traceability: Dynamics can be traced to the impact caused by the action, and for each sub-dynamics, we show the meanings of the action dimensions it traced to. Model Error Curve: The model error comparison on the Walker task of original Dreamer and D2P-Dreamer methods with different decomposition show that modeling each decomposed sub-dynamics separately can significantly reduce the model error.
• Figure 2: Overview of the ED2 Framework. ED2 contains two components: SD2 and D2P. SD2 decomposes the dynamics by generating a partition $\mathcal{G}$ on action dimensions. D2P decomposes the action $a_{t-1}$ into multiple sub-actions according to $\mathcal{G}$ and makes decomposing predictions based on $s_t$ and each sub-action $a^{G_k}$. The prediction $h_t$ is the combined output of all sub-dynamics models, from which the next state $s_{t}$ and reward $r_{t}$ are generated.
• Figure 3: Extension of RSSM with D2P.
• Figure 4: Example image of continuous control benchmark tasks: Deepmind Control (DMC), DMControl Remastered (DMCR), and Mujoco (left to right).
• Figure 5: Comparisons of ED2-Dreamer vs. Dreamer. The x- and y-axis represent the environment steps and the performance. The line and shaded area denote the mean value and standard deviation.

Theorems & Definitions (1)

• Definition 1