ED2: An Environment Dynamics Decomposition Framework for World Model Construction
Authors: Cong Wang
Tianpei Yang
Jianye HAO
YAN ZHENG
Hongyao Tang
Fazl Barez
Jinyi Liu
Jiajie Peng
haiyin piao
Zhixiao Sun
Venue: ICLR 2022 Submitted (2022)
Keywords: model based reinforcement learning

Abstract: Model-based reinforcement learning methods achieve significant sample efficiency in many tasks, but their performance is often limited by the existence of the model error. To reduce the model error, previous works use a single well-designed network to fit the entire environment dynamics, which treats the environment dynamics as a black box. However, these methods lack to consider the environmental decomposed property that the dynamics may contain 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 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 and boosts the performance of the state-of-the-art MBRL algorithms on various continuous control tasks.

Under review as a conference paper at ICLR 2022
ED2: AN ENVIRONMENT DYNAMICS DECOMPOSI- TION FRAMEWORK FOR WORLD MODEL CONSTRUC- TION
Anonymous authors Paper under double-blind review
ABSTRACT
Model-based reinforcement learning methods achieve significant sample efficiency in many tasks, but their performance often limited by the existence of the model error. To reduce the model error, previous works use single well-designed network to fit the entire environment dynamics, which treats the environment dynamics as a black box. However, these methods lack to consider the environmental decomposed property that the dynamics may contain 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 (ED22, a novel world model construction framework that models the environment in decomposing manner. ED2 contains two key components: sub-dynamics discovery (SD2) and dynamics decomposition prediction (D2P). SD2 discovers the sub-dynamies in environment 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 and boosts the performance of the state-of-the-ar MBRL algorithms on various continuous control tasks.
INTRODUCTION
Reinforcement Learning (RL) is general learning framework for solving sequential decision-make problems and has made significant progress in many fields (Mnih et al. 2015 Silver et al. 2016 Vinyals et al. 2019)] Schrittwieser et al.] (019). In general, RL methods can be divided into two categories regarding whether a world model is constructed for the policy deriving: model-free RL (MFRL) and model-based RL (MBRL). MFRL methods train the policy by directly interacting with the environment, which results in good asymptotic performance but low sample efficiency. By contrast, MBRL methods improve the sample efficiency by modeling the environment, but often with limited asymptotic performance and suffer from the model error (Lai al. 2020; Kaiser et al. 2020).
Existing MBRL algorithms can be divided into four categories according to the paradigm they follow: the first category focuses on generating imaginary data by the world model and training the policy with these data via MFRL algorithms (Kidambi et al., (0220: Yu et al. 2020); the second category leverages the diffcrentiability of the world model, and generates differentiable trajectories for policy optimization (Deisenroth & Rasmussen, 201T Levine & Koltun 2013 Zhu et al. 2020): ; the third category aims to obtain an accurate value function by generating imaginations for temporal difference (TD) target calculation (Buckman al. 2018: Feinberg et al. (018):: the last category of works focuses on reducing the computational cost of the policy deriving by combining the optimal control algorithm (e.g. model predictive control) with the learned world models (Chua al. 2018 Okada & Taniguchi 2019 Argenson & Dulac-Arnold] 2020). Regardless of paradigms, the performance of all existing MBRL algorithms depends on the accuracy of the world model. The more accurate the world model is, the more reliable data can be generated, and finally, the better policy performance can be achieved. Therefore, improving the world model accuracy is critical in MBRL.
To this end, various techniques have been proposed to improve the model accuracy. For example, rather than directly predict the next state, some works construct a world model for the state change
Our code is open source and available httph:tp//// htap::///itthbee com/ 6n2 scurce code/ooeeii

---

Under review as a conference paper at ICLR 2022
According role
Traced to action
Model Error
Harch
Douns
D2A Gvernereation
value
Traced to action
Your de
According to position
Step
(a] Cheetah environment
(b) Decomposability
(c) Traceability
(d) Model error
Figure I: (a) The Cheetah task with six action dimensions. (b) The dynamics can be decomposed into multiple sub-dynamics in various ways, each sub-dynamics described with different background colors. (c) Dynamics can be traced to the impact caused by the action, and for each sub-dynamics, we show the meanings of the action dimensions traced to. (d) The model error comparison on the Seaeee-Roooooeleeeeeeeee correspond to decompose the dynamics according to role/position and model each sub-dynamics separately).
Cheetah task of Dreamer and 22--Dreamer methods
prediction (Luo et al., 2019 Kurutach et al. 2018). Model ensemble is also widely used in model construction for uncertainty estimation, which provides more reliable prediction (Janner et al. [0119] Pan et al., 2020). To reduce the model error in long trajectory generation, optimizing the multi-step prediction errors is also effective technique (Hafner et al. 2019). However, these techniques improve the eavironment modeling in a black-box way, which ignores the inner decomposed structure of environment dynamics. For example, Figure (a) (a) shows the Cheetah task from DeepMindControl (DMC) Suite tasks, where the dynamics can be decomposed in various ways. Figure T (b) shows various decomposition
decompose it into: {thigh, (thit, alternatively. according to the position of sub-dynamics, we can decompose into: {back, front.. Figure (d) shows that no matter whether we decompose the Cheetah task according to role position, modeling each decomposed sub-dynamic separately can significantly reduce the model error of the existing MBRL algorithm (e.g. Dreamer (Hafner al. (0200).
the dynamics: according to the role of sub-dynamics,
Inspired by the above example, we propose environment dynamics decomposition (ED2). novel world model construction framework that models the dynamics in a decomposing fashion. ED2 contains two main components: sub-dynamics discovery (SD2) and dynamics decomposition pre- diction (D2P). SD2 is proposed to decompose the dynamics into multiple sub-dynamics, which can be flexibly designed and we also provide three alternative approaches: complete decomposition, human prior, and the clustering-baat based method. D2P proposed to construct the world model from the decomposed dynamics in SD2, which models each sub- dub-dynamics separately in end-to-end training manner. ED2 is orthogonal to existing MBRL algorithms and can be used as a backbone to easily combine with any MBRL algorithm. Experiment shows ED2 improves the model accuracy and boosts the performance significantly when combined with existing MBRI. algorithms.
BACKGROUND
2.1 REINFORCEMENT LEARNING
Given an environment, we can define a finite-horizon partially observable Markov decision process (POMDP) as (S, A, R, P, y, O S2 7), where S E Th" is the state space, and A E IRm is the action space, R x A - IR. denotes the reward function, dynamics, is the discount factor. The agent receives an observation E Ω, which contain partial information about the state s E S. O is the observation function, which mapping states to probability distributions over observations. The decision process length is denoted as T.
S A + S denotes the environment
Let 7 denote the expected return a policy æ over the initial state distribution Po. The goal of an RL agent is to find the optimal policy x² which maximizes the expected return:
arg max n/mT] = arg max E, 7 R(ss, az)];
x
77
2

---

Under review as a conference paper at ICLR 2022
where su Pu. ot O(-(se). ar x(-)o)). 8141 P())s,,a)). If the environment is fully observable,
ie., Q S and O an identity function, POMDP is equivalent to the MDP: R,,
2.2 REPRESENTATIVE WORLD MODELS IN MBRL
The world model a key component MBRL that directly impacts policy training. World models are usually formulated with latent dynamics (Janner et al. [2019] Hafner al. 2020), and the general form of the latent dynamics model can summarized as follows:
Latent transition kernel:
h f($g)-1,
Stochastic state function:
p(selhe)
Reward function:
p(re(he)
The latent transition kernel (shorthand as kernel) predicts the latent state he with input 8ge-1 and aga-1- Based latent state l4, the stochastic state function and reward function decode the state 8t and reward Tr. For the partially observable environment, two additional functions are required:
Observation function:
p(oc)ss))
Representation function: p(s,/he:o,)
In general, world models mainly differ at the implementation of kernel, which can be roughly divided into two categories: with non-recurrent kernel and with recurrent kernel. The formal definition of both kernels are as follows:
hr (
((--) ((h--1,8--100--)
With non-recurrent kernel With recurrent kernel
Non-recurrent kernel relatively basic kernel for modeling, which are often implemented as Colly-Conneeted Networks. Non-recurrent kernel takes the current state $t-1 and action Q,-1 as input, outputs the latent state prediction ht. Compare non-recurrent kernel, recurrent kernel is implemented as RNN and takes the additional input f,,-1, which performs better under POMDP setting. For both kernels, the 8, and r; can be generated from the latent prediction h4-
ENVIRONMENT DYNAMICS DECOMPOSITION
3.1
MOTIVATION
An accurate world model is critical in MBRL policy deriving. To decrease the model error, existing works propose various techniques as introduced in Section [ However, these techniques improve the environment modeling
black-box manner, which ignores the inner properties of environment dynamics, resulting in inaccurate world model construction and poor policy performance. To address this problem, we propose two important environment properties when modeling an environment:
I) Decomposability:
The environment dynamics can be decomposed into multiple sub- dynamics in various ways and the decomposed sub-dynamics be combined reconstruct the entire dynamics.
2) Traceability: The environment dynamics can be traced to the action's 's impact on the en- vironment, and each sub- aub-dynamies can be traced to the impact caused by part of the action.
For example in the Cheetah task, Figure (b) demonstrates the decomposability: we can decompose the dynamics into (thigh, shin,foor) sub-dynamics or (back.front) sub-dynamics, which depends on the different decomposition perspectives and the combination of decomposed sub-dynamics can constitute the entire dynamics. [ig explains the traceability: each sub-dynamics can traced to the corresponding subset of action dimensions: for the thigh dynamics, it can be regarded as the impact caused by the front-thigh and back-thigh action dimensions. The above two properties are closely related to environment modeling: the decomposability reveals the existence of sub-dynam -ob--yymmicss which allows us to model the dynamics separately, while the traceability investigates the causes of the dynamics and guides to decompose the dynamics its root (i.e. the action).
3

---

Under review as a conference paper at ICLR 2022
SD2 Method
ED2 world model
Framework
Action dimensions
Partition &
Sub-a bub-actions
models
predictions
G₁
D2P Framework
G2
SD2 Method
Action desimposing
6
...
Gs
(Partiition
M
Figure 2: Overview of the world model under ED2 Framework. ED2 contains components: SD2 and D2P. SD2 decomposes the dynamics by generating partition G on action dimensions. D2P decomposes action into multiple sub-actions according based on s and each sub-action. The prediction h is the combined output of all sub-dynamics models, from which the next state s' and reward r are generated.
and makes decomposing predictions
To take the above properties into account, we propose the environment dynamics decomposition (ED2) framework (as shown in Figure 2). which contains two key components: sub-dynamics discovery (SD2) and dynamics decomposition prediction (D2P). More specifically, by considering the traceability, we propose discover the latent sub-dynamies by analyzing the action (SD2, the blue part in Figure 2); by considering the decomposability, we propose to construct the world model in decomposing manner (D2P, the green part in Figure 2). Our framework can be used backbone in MBRL and the combination can lead to performance improvements over existing MBRL algorithms.
3.2 DYNAMICS DECOMPOSITION PREDICTION
Given an environment with m-dimensional action space A R the index of each action dimension constitutes set A = {1, 2, ... .m,, any disjoint partition G = [GI,.... Gx over A corresponds to a particular way of decomposing action space. For each action dimension in A, we define the action space as A', which satisfied A = A' x ... x A". The action space decomposition under partition G defined as A (1 ..... A }. where sub-action space A ILLea, A". Based on above definitions, we define the dynamics decomposition for under partition G as follows:
Definition Given a partition G. the decomposition for P: S A + can be defined as:
((s, = fc ( [P((saa))) Sx,
(1)
with set of sub-dynamics functions {P1, Pk} that P. PS:Sx AC. . H. and a decoding function fo H + S. S. Note H is a latent space and ai E A is a sub-action (projection) of action a.
Intuitively, the choice of partition G is significant to the rationality of dynamics decomposition, which should be reasonably derived from the environments. this section, we mainly focus on dynamics modeling, and will introduce how derive the partition & by using SD2 in section 3.3
To implement D2P, we use model M) parameterized by di (i.e., neural network parameters) to approximate each sub-dynamics P.. As illustrated in Figure 2 given a partition G, an action is Ga J. each model M takes state and the sub-action divided into multiple sub-actions (a
G. as input and output a latent prediction h E H. The separate latent predictions (lt,.... , are aggregated and then decoded for the generation of state s' and reward For each kernel described in we provide the formal description here when combine with D2P:
Section 2.2
I
E
181-11,
For non-recurrent kernel
For recurrent kernel
4

---

Under review as a conference paper at ICLR 2022
We propose set of kernels, where each kernel mod- els specific sub-dynamics with the input of current state s, corresponding sub-action a and hidden state - (ignored when applying on non-recurrent kernel). The output of all kernels is averaged get the final output heg The prediction of reward and state 8t generated from the output /ę- Specifi- cally, we provide an example when combining with the kernel of Recurrent State-Space Model (RSSM) (Hafner et al. 2019) in Figure 3 which is repre- sentative recurrent kernel-based world model. The original single kernel implemented as GRU are re- place by multiple kernels with different action input.
Br-1
Me-1
GRUS
GRU
...
GRU
AE
9
Op.7
RSSM
DXPeexetedeed RSSM
Figure 3: Extension of RSSM with D2P.
3.3 SUB- BUB-DYNAMICS DISCOVERY
The traceability of the environment introduced in Section 3.1 provides with basis for dynamics decomposition: the decomposition on dynamics can be converted to the decomposition on the action space. Therefore, we present the SD2 module for the action space decomposition and discuss three im- plementations in this section. With the Cheetah task in Section 3.1 as the example: the straightforward SD2 implementation is the complete decomposition, which regards each action dimension as a sub- dynamics and decomposes the dynamics completely. Specifically, complete decomposition decom- poses the dynamies into six sub-dynamics: {Front, Back) (Thigh, Shin, Foot}. However, com- plete decomposition ignores the inner action dimensions correlations. which limits its performance in Fronl_Shim Frout_Food
many tasks. For example, the three action dimensions (Pronl_Thigg affect the dynamics of the front part together, thus simply separate these action dimensions would affect the prediction accuracy. To include the action dimension correlations, incorporating human prior for the action space decomposition is improved implementation. Based different human prior, we can decompose the dynamics in different ways as introduced in Figure I Nevertheless, although human prior considers the action dimension correlations, it highly subjective and might lead to sub-optimal results due to the limited understanding of tasks (we also provide the correspond- ing experiment in Section 4.2.3). Therefore, human prior not applicable in complex systems which is beyond human understanding.
To better discover the sub-dynamies and eliminate the dependence on human prior, we propose to automatically decompose the action space using the clustering- based method. The clustering- bbased method contains two components: feature extraction and clustering criterion. Feature extraction extracts the properties of action dimension a' into feature vector F*. Then we regard each action dimension as a cluster and aggregate related action dimensions together with the clustering criterion. The effectiveness of the clustering-basee method depends on the quality of feature extraction and the validity of clustering criteria, which may be different in different environments. Therefore, although provide general implementation later, still suggest readers design suitable clustering-basee methods according to task-specific information.
We extract the properties of each action dimension by computing the Pearson correlation coefficient between action dimensions and state dimensions. Specifically, we define where each fij denotes the Pearson correlation the feature vector F (l. (f
Feature Extraction:
coefficient between action dimension i and state dimension j. 5s describes the impact caused by action dimension and (S is calculated by the corresponding action value a' and state value changes As (which is the difference between the next state and the current state):
pij cov(a' As)
where cov denotes the covariance and a denotes the standard deviation.
Clustering Criterion: We define the clustering criterion the relationship between clusters, which can be formalized as follow:
Rela G,) =R(G, G,) R(G,, G-)) W R(G,, G-5) wh- wh--3 + wi.-i
5

---

Under review as a conference paper at ICLR 2022
where = A\ Gi, w = G x |Gy| and R(G,, =
=-
whd EA'EG, Eseeg, 1IF',,'
II IIl measures the distance between vectors under distance function D (we choose the negative cosine similarity as D).
Algorithm presents the overall im- plementation of the clustering-based method. As Algorithm|] Jddseribes. with input task & and clustering threshold 1). we first initialize the cluster set 9 contain- ing clusters (each for single action di- mension), a random policy Trand, and an empty dataset De. Then for T episodes, Trand collects samples from the environ- ment and we calculate F for each action dimension i. After that, for each cluster- ing step. we select the two most relevant clusters from G and cluster them together. The process ends when there is only one cluster, when the correlation of the two most correlated clusters is less than the threshold y. is hyperparameter which assigned with value around 0 and empirically adjusted.
Algorithm 1 Selectable clustering based method.
Input: Task &, clustering threshold
Initialize cluster set 9 = {{1), ... , {m}} according to E, random policy Wrand. dataset D. +
for 1,2,
Collect and store samples in D. with Frand Calculate F' for cach action dimension with De while |g| > do
G Guazy .mary = arg m G, Rela(G,, G₃)
if Rela(Gmaay Gmarg) > 1j then Remove Gmaz; and Grazr from C
Add Gmaz; U Gmazz to G
else
Stop clustering
return
3.4
BD2 FOR MBRL ALGORITIMMS
ED2 is a general framework and can be combined with any existing MBRL al-
Algorithm 2 -00--Drammrr
Input: Task &, clustering threshold
gorithms. Here we provide the practi-
cal combination implementation of ED2 with Dreamer(Hather el al.] 2020) (Al- gorithm 2 and we also combine ED2 with MBPO (Janner et al. 2019) in the appendix. The whole process of ED2 Dreamer contains three phases: 1) SD2 decomposes the environment dynamics of task E, which and can be implemented by three decomposing methods intro- duced in SectionB.3) 2) D2P models each sub-dynamics separately and constructs the ED2-combined world model po by combining all sub-models with decoding network f3 3) The final training phase initializes the policy me and the world model pú. then derive the policy from Dreamer with input ma, Ps and task &.
// Sub-dynamics Discovery (SD2) Phase:
SD2 methods (E,1)
// Dynamics Decomposition Prediction (D2P) Phase:
for = 1,2, ... | do
Build sub-dynamics model: Må, = (hi-1, 85- a
Combining all aob-dynamics models: M. = [9] E Mi.
Combining Me with decoding network for and construct the ED2-combined world model: pa fè (M.) // Training Phase:
Initialize policy ma, model: Ps
Optimize policy with Dreamer: T
Dreamer(E, ma. Po)
4 EXPERIMENTS
Baselines & Benchmarks: For a comparative study, we take two state-of-the-art MBRL methods: MBPO Janner et al. 2019) and Dreamer (Hafner et al. 2020) as the baselines, and extend them with ED2 as ED2-MBPO, ED2-Dreamer To reduce implementation bias, we reuse the code and benchmarks from the prior works: the DMC Suite (Tasaa et al. (018) for Dreamer and Gym- Mujoco (Brockman et al. 2016) for MBPO. Besides, we also provide ablation studies to validate the effectiveness of each component of our ED2.
(1) Although the data required by the clustering bassed method is tiny (less than 1% of the policy training), we include it in the figures for fair comparison. (2) We take the clustering- based method as the main SD2 implementation (denoted as ED2-Methods) and discuss other SD2 methods in Section 4.2.3 and appendix (complete decomposition. human prior denoted CD, HP respectively). (3) Due to the space limit, we leave the result of ABPO/ED2-MBBP in the appendix. (4) All results are averaged over 5 seeds. The hyperparameters setting is left in the appendix.
Clarification:

---

Under review as a conference paper at ICLR 2022
4.1 PERFORMANCE
dmc cheotah run
dinc finger spin
dmc hopperr hop
dmcc humanold stand
600
200
200
52
52
1
Step
dime_ humanoid walk
dmc_reacher
dmc.walker.nog
dmc hopper stand
Elo mnamer
Elo innver
Ele Cnermer
100
Voloo
400
Learme
Stop *10
Stop
Stop
Step
Figure 4: Comparisons of ED2-Dreamer vs. Dreamer. The x- and y-axis represent the training steps and performance. The line and shaded area denotes the mean value and standard deviation.
We evaluate Dreamer and ED2-Dreamer eight DMC tasks with image inputs. As shown in Figure DD2-Dreamer outperforms Dreamer on all tasks. This is because ED2 establishes a more rational and accurate world model, which leads to more efficient policy training. Another finding is that, in tasks like cheetah_ ____ and walker funke ED2-Dreamer achieves lower variance, demonstrating that ED2 can also lead to a more stable training process. Furthermore, Dreamer fails to achieve good performance in difficult tasks such as gumanoid_stane and humanoid_ walk. In contrast, ED2- Dreamer improves the performance significantly, which indicating the superiority of ED2 in complex tasks. In humanoid tasks, the dynamics are too complex for the clustering- based input. Therefore, we use the vector-based state for clustering and keep the policy training on the image-based state (we will further discuss this in Section 5) The performance of MBPO and ED2-MBPO are left in the appendix, which proves that ED2 boost MBPO's performance significantly.
method with image-
4.2
ABLATION STUDIES
4.2.1 THE EFFECTIVENESS OF D2P AND SD2
In this section, we investigate the contribution of each component to performance improvement. We can summarize the improvements into three parts: multiple kernels, decomposing prediction, reasonable partition. There is progressive dependence between these three parts: the decomposing prediction depends on the existence of the multiple kernels and the reasonable partition depends on the decomposing prediction. Therefore, we design an incremental experiment for the validation.
dmcc chostah run
dmnc humanold walk
dmc walker run
dmc por hop
Value
E12
EED Orenae-enniin
3.0
+20
3.0
x19
Figure 5: Performance comparisons of components ablation experiments.
First, we employ the DD2=Dreeamr--nseembeeee D which maintains the multiple kernel structure but without dynamics decomposing (i.e. all kernels input with action a rather than sub-action aC.). We investigate the contribution of multiple kernels by comparing ED2-Dreamer-ensemble with baselines. Second, we employ the ED2-Dreammer-random partition randomly. We investigate the contribution of decomposing prediction by comparing ED2- Dreamer-random with Soo2-Dreamer-ensembee Last, we investigate the contribution of reasonable partition by comparing ED2-Dreamer with ED2-Dreamer-random.
which maintains the D2P structure and obtains
Thhe ensemble refers to kernel ensemble, which is described in detail in the appendix.
7

---

Under review as a conference paper at ICLR 2022
shows that ED2- Dreamee--nnsemte outperform Dreamer on humanoid walk and cheetah run tasks, indicating that multiple kernels help the policy training on some tasks. on humanoid walk task, but not in other tasks (even perform worse in cheetah_ run and hopper_hopp This is due to the different modeling difficulty of tasks: the tasks except humanoid_ walk are relatively simple and can be modeled without D2P directly (hut in sub-optimal way). The modeling process of these tasks can be aided by a reasonable partition but damaged by random partition. The humanoid walk challenging and cannot be modeled directly, therefore decomposing prediction (D2P) is most critical and performance can be boosted even with a random decomposing prediction. Finally, ED2-Dreamer outperforms tasks, which indicates that a reasonable partition (SD2) is critical in dynamics modeling and D2P can not contribute significantly to the modeling process without a reasonable partition.
random
outperforms
102--reamer-ramae
on all
4.2.2
MODEL ERROR
dmc cheetsh
dmc humanoid walk
dmc, wawker.run
dmc per, hop brane
3,015
3.30
2.3
2011
1.5
100g 3.3
Stop
3.00
K10P
16.0
Stup K1/e 3,00
1.0
Figure 6: The model error (KL-Divergence) comparison of DD2-Dreamer and Dreamer.
In this section, we further investigate whether the model error reduced when combined with ED2. We conduct an environment modeling experiment on the same dataset (which is collected in the MBRL training process) and record the model error. Since the policy keeps update in the MBRL training process, the dataset of MBRL also changes with the updated policy. For example, in the MBRL training on the Cheetah task, the model is first trained with the data like Rolling on the ground and finally trained with the data like rumning with high speed. To simulate the MBRL training process, we implement our dataset by slowly expanding it from 0 to all according to the data generation time. This setting can also help to investigate the generalization ability of the model on unfamiliar data (i.e. the data generated by the updated policy). We list parts of the result in Figure and the result of other tasks are shown in the appendix.
F shows that Dreamer. ED2-Dreamer can also achieve more stable world model training (i.e. with low variance) on humanoid walk and walker Bun tasks. We also find that the baseline methods have significantly increasing model error on humanoid walk and walker run tasks, but for ED2-methods, the increase is much smaller. We hypothesize that ED2 produces a reasonable network structure; as the dataset grows,
has a significantly lower model error in all tasks compared with
can generalize to the new data better. This property is significant in MBRL since the stable and accurate model prediction is critical for policy training. We also provide the model error comparison of MBPO and ED2-MBPO in the appendix, which also proves that ED2 can reduce the model errors when combine with MBRL methods.
4.2.3 SD2 COMPARISON
In this section, we compare the performance of three proposed SD2 methods. We list the decom- position obtained by the clustering-basee method and human prior on humanoid walk task for the illustrating purpose and provide the corresponding performance comparison results in Figure 7 More experimental results on other tasks are provided in the appendix.
As shown in Figure human prior can generate different partitions from different task understandings and we average their performance as the final result. Experiment shows that all SD2 methods help method performs best and baseline Dreamer performs
the policy learning. The clustering-basee worst in the comparison. For the complete decomposition, it performs poorly under humanoid, walk, which implies that humanoid walk contains many inner action dimension correlations, and simply complete decomposition heavily breaks this correlation thus hinders the final performance. Compared to complete decomposition, human prior maintains more action dimension correlations by leveraging the human prior knowledge, which leads to better performance. However, the correlations maintained
8

---

Under review as a conference paper at ICLR 2022
dmc humanoid
Omanner
Dreame
EC2 CD Dreamer
100
50
CL-Beedl:
HF-3
Step X207
Figure 7: The performance comparison of Dreamer, ED2-Dreames ED2-CD-Dreamme and ED2-HP- Dreamer. We provide the sub-dynamics visualization the left four figures. Each circle correspond to joint. A joint contains multiple action dimensions when the corresponding circle is separated into multiple parts. We mark the action dimensions the same sub-dynamics with the same color.
by human prior might be false or incomplete due to human limited understanding of tasks. Compare to human prior, the clustering-based method automatically decomposes the action space according the clustering criterion, which decomposes the action space better a mathematical way. For example, human prior aggregates (right_hipp _hip__,_____r right right_hip__n (x,, y, z denote the rotation direc- tion) together and the clustering- based method aggregates {abdomen_ æ, right_ faftt_hhp______ IT, left_hiftt t__ip__a together. Although the action dimensions from human prior sub-dynamics affect the same joint left_hip they rotate in different directions and play different role in the dynamics. In contrast, the sub-dynamics discovered by the clustering based method aggregate the action dimensions that affect the x-direction rotation together. It maintains stronger correlations and helps the world model fitting movement on x-direction better. Therefore, performs better than human prior on this task.
DISCUSSION
In this paper, we regard SD2 a flexible module that can adopt any suitable partition methods considering the task-specific information. Currently, we discuss three kinds of SD2 methods, i.e., human prior, complete decomposition, and the clustering- method is chosen testbeds empirically. We also analyze the reasonable decomposition provided by the clustering- method, which contributes lot to the dynamics modeling process. Nevertheless, the clustering- based method is still faced with extra challenges when solving complex tasks with image- based inputs, such as humanoid tasks. In this paper, use the vector-based state in the clustering stage for humanoid tasks by considering the one-to-one correspondence between vector-based and image- based state representations. How leverage self-supervised learning contrastive learning to learn low-dimension, high-quality state features from raw images to improve the partition discovery effect of SD2 and further apply our ED2 more complex scenarios worthwhile further investigate.
method, and the clustering- main implementation since it outperforms the other two methods on these
Previous work (Doya et al. 2002) also takes the dynamics decomposed prediction into consideration, which achieves better dynamics modeling. decomposes the dynamics from the perspective of state and time. Different from this work, we analyze the cause of dynamics and decompose it from its root: the action space, which makes the modeling of environmental dynamics more reasonable and scalable.
CONCLUSION
In this paper, we propose a novel world model construction framework: Environment Dynamics Decomposition (ED2), which explicitly considers the properties of environment dynamics and models the dynamics in decomposing manner. ED2 contains two components: SD2 and D2P. SD2 decomposes the environment dynamics into several sub-dynamics according to the dynamics-action relation. D2P constructs a decomposing prediction model according to the result of SD2. With combining ED2, the performance of existing MBRL algorithms is significantly boosted. Currently, this work only considers the decomposition the dimension level, and for future work, it is worthwhile investigating how to decompose environment dynamics at the object level, which can further improve the interpretability and generalizability of ED2.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. CoRR, abs/2000.05550 2020.
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, abs/ 100.01140,, 2016.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Advances Neural 2018.
Information Processing Systems 31, pp.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems 31, pp. 1759-47700 2018.
Marc Peter Deisenroth and Carl Edward Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28- July 2, 2011, pp. 465-472. Omnipress, 2011.
Kenji Doya, Kazuyuki Samejima, Ken-ichi Katagiri, and Mitsuo Kawato. Multiple model-based reinforcement learning. Neural Comput., 14(6):1347-136,, 2002.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. CoRR, 1bs//800.00101, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, pp. 2555-2565, 2019.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control:
Learning behaviors by latent imagination. In Proceedings of the 8th International Conference on Learning Representations, 2020.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems 32, pp. 12498-12509, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In Proceedings of the 8th International Conference on Learning Representations, 2020.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. In Proceedings of the 6th International Conference on Learning Representations, 2018
Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model-based policy optimization. In Proceedings of the 37th International Conference on Machine Learning, pp. 5618-5627, 2020.
Sergey Levine and Vladlen Koltun. Guided policy search. In Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013, volume 28 of JMLR Workshop and Conference Proceedings, pp. 1-9. JMLR.org, 2013.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorith- mic framework for model-based deep reinforcement learning with theoretical guarantees. In Proceedings of the 7th International Conference on Learning Representations, 2019.
10

---

Under review as a conference paper at ICLR 2022
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nat., 518(7540):529-555 533, 2015.
Masashi Okada and Tadahiro Taniguchi. Variational inference MPC for bayesian model-based reinforcement learning. In Proceedings of the 3rd Annual Conference on Robot Learning, pp. 258-272 2019.
Shayegan Omidshafiei, Dong-Ki Kim, Jason Pazis, and Jonathan How. Crossmodal attentive skill learner. In Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, pp. 133--46. International Foundation for Autonomous Agents and Multiagent Systems Richland, SC, USA/ / ACM, 2018.
Feiyang Pan, Jia He, Dandan Tu, and Qing He. Trust the model when it is confident: Masked model-based actor-critic. In Advances in Neural Information Processing Systems 33, 2020.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. CoRR, abs/1911.08266. 2019.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George van den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam,, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nat., 529(7587):484-445) 2016.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, abs/ 1001 .00690, 2018.
Oriol Vinyals, Igor Babuschkin, Wojciech M. Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom L. Paine, Çaglar Gülçchre, Ziyu Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in starcraft II using multi-agent reinforcement learning. Nat., 575(7882): 350-354, 2019.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsca Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Guangxiang Zhu, Minghao Zhang, Honglak Lee, and Chongjie Zhang. Bridging imagination and reality for model-based deep reinforcement learning. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
11

---

Under review as a conference paper at ICLR 2022
A EXPERIMENT SETTING
We keep the experiment setting the same with Dreamer and MBPO. For Dreamer, we do the experiment Deep Mind Control environments, with images imput, and each episode length is set to 1000. For MBPO, the experiment is processed on the Gym-Mujoco environment, takes the structured data as input, and sets the episode length 1000. Each experiment is averaged by running five seeds. Our specific computing infrastructure is as shown in Table
Table I: The computing infrastructure of our experiment.
CPU
GPU
MEMORY
XEON(R) SILVER 4214
RTX2080TI
256G
B EXPERIMENT HYPERPARAMETERS
For ED2-Dreamer and ED2-MBPO, we followed the official implementation of Dreamer and MBPO except for the dynamies models. Specifically, for the dynamics models, select the hyperparameters in each environment, as shown in Table
Table 2:
The hidden size and 1 value for each environment. DeepMind denotes the environment that belongs to DeepMind Control Suite, and Gym-Mujoco denotes the environment from Gym-Mujoco.
ENVIRONMENT
HIDDEN SIZE
HOPPER(DEEPMIND)
200
WALKER
200
-0.06
CHEETAH(DEERMINND
200
-0.1
HUMANOID(DDEPPMIND)
200
REACHER(DEEPPINN))
FINGER (INEER(DEEMMIND)
200
HALFCHEETAHI CCCCEEEETAAIGGMMMMMJCCOO
150
0
HOOPER(GYYMMMUI) MUJOCO)
200
-0.3
WALKER(GYM) GEEEEEUOGY-MOOOO
200
-0.2
ANT(GYM- MTT(GYM-MUJOCO)
150
-0.12
For the clustering process, hyperparameter describes the tightness of constraints on inter-group distance and intra- group distance. When = 0, the clustering process stop condition is that the distance between the two most relative clusters is equal to the distance between these two clusters and others (this is general condition in the most environment). In some environments, although 7) is good choice, but the value can be further finetuned to obtain more reasonable clustering results. The 77 value use is in Table
C ED2-MBPO IMPLEMENTATION
The combination of MBPO and ED2 can be described as Algorithm 3 We can also separate it into three parts: the first phase is SD2, which discover the sub-dynamics (partition 9) in the environment by using appropriate SD2 method. Then the D2P phase models each sub-dynamics separately and construct the combineddded world model P.- Finally, we derive the trained policy T by using MBPO method with input task &, initialized policy mo and world model ensemble Pj-
12

---

Under review as a conference paper at ICLR 2022
Algorithm 3 ED2-MBPO
Input: Task E. clustering threshold 1)
// Sub dynamics Discovery (SD2) Phase:
G SD2 methods (E,7) // Dynamics Decomposition Prediction (D2P) Phase:
for = 1,2, ...., G9 do
Construct sub uub-dynamics model: Ma, f( a-11
Combining sub-dynamics models: M, [0] Ei=1 M3.
Combining M. with a decoding network fØg and construct the // Training Phase:
world model: Pa = fōy (Mc)
Initialize policy тa, model ensemble: P: = {po).... ..... .p&.} Optimize policy with MBPO: =6 MBPO no. P2))
D KERNEL ENSEMBLE AND MODEL ENSEMBLE
We retain the multi kernel network structure and all kernel input with the same information: hidden state h (if combined with recurrent kernel), current state s and current action a. Kernel ensemble can be regarded as single model, which average the outputs from all kernels and training in an end-to-end manner for all kernels.
Kernel ensemble is deployed in section 4.2
Model ensemble is widely used MBRL for uncertainty estimation. In MBRL, the model ensemble is generally implemented by setting different initial parameters and sampling different training data from the same dataset. Models in model ensemble are trained separately and no connection between them except training from the same dataset.
Therefore, model ensemble totally different from kernel ensemble. Model ensemble propose to train multiple unrelated models for the same task, which can estimate the uncertainty of prediction. But kernel ensemble propose to use one model (but construct with multiple kernels) for the prediction task. We can also combine kernel ensemble with model ensemble and improve the accuracy of all models (e.g. MBPO).
E MODEL ERROR
In order to verify whether we get more accurate world model, we measure the model error. Figure S the model error curve of Dreamer / EDD-Dreamer. We can see that our framework can significantly reduce the model error.
cheetah Fun Distw
finger spin
dmc dmc_hopper hop
dmc humanoid stand
Daine
335
1.011
3.5
1.004
533
1.1
0.5
3.6
1.
0.5
1.0
dmc humanoid walk
dmc reacher easy
dmc walker run
dmc hopper stand
1.04
3.35
3.30
3.00
2.10
2.15
2.38
3.00
29.3
Stap
Step
Step
Stepp
Figure 8: The model error reduced when combine ED2 with Dreamer.
13

---

Under review as a conference paper at ICLR 2022
F ED2-CD-DREAMER
Here we provide the experiment result of complete decomposition in Figure complete decompo- sition can boost the performance in most environments. But in some complex environments like humanoid and walker, it fails to improve the performance. We analysis that in humanoid and walker, the correlation between action dimensions can't be ignored. Complete decomposition break the correlations between action dimensions and lead to poor performance in these tasks.
dmc cheetah run
dmcc finger spin
dmc hopper hop
dmnc humanold stand
209
300
100
52
Step
Step
Step
Step
humanoid
dmcc reachar easy
dmc walker,
dmc hopper stand
Volue
400
Dianner
Stop
+37
Stop
Stop
1.0
Stop
+27
Figure 9: Comparisons between ED2-CD-Dreamen and Dreamer.
G COMBINE WITH MBPO
Here we provide the experiment results of ED2-MBPO method, which include the performance (Figure TOp, model error evaluation (Figure 11 and performance under complete decomposition
(Figure D2p.
Hopper
HalfCheatah
WalkerZd
8000
4500
Lo0e
2000
2300
+30
+100
1.
12
Figure 10: Performance comparisons between ED2-MBPO and MBPO
Hopper
HalfCheetah
Ant
Walker2d
Ere Mnin
Volue
Volue
0000
6.00
Stop
x20
Stop
5tup
Stup
Figure 11: Model error comparisons between ED2-MBPO and MBPO.
H SD2 CLUSTERING RESULTS
We visualize the final partition result obtained by clustering here with both figure and table form, the figure result of DeepMind Control Suite is shown in Figure [3] and the figure result of Gym-Mujoco is shown in Figure 14 The result shows that the clustering- based method tends to group the relative action dimensions together, and the clustering results are also reasonable from the human point of view.
14

---

Under review as a conference paper at ICLR 2022
Hopper
HalfCheetah
Ant
Walker2d
3900
Step
Step
Step
Step
Figure 12: Performance comparisons between ED2-CD-MBUL and MBPO.
Reacher
Humanoid
Walker
Finger
Hopper
Cheetah
Figure 13: The visualization of final partition of environments in DeepMind Control Suite. Some joints in humanoid is divided into two three parts (e.g. abdomen joint). It indicates that there are multiple action dimensions contained by this joint (e.g. abdomen joint contains abdomen x, abdomen_y and abdomen_2 action dimensions).
As shown in Table 3, each row denotes sub-dynamics discovered by the clustcring-basee method in this environment (expect the final two-row in humanoid environment, they belong the same sub-dynamics. Because of the length of the table, we write as two lines). The sub-dynamics we discovered is very reasonable, and there is obvious connection between the action dimensions in the same sub-dynamics.
I ATARI EXPERIMENTS
We also conducted model error experiments on Atari environment and Atari-like Maze environment (called Minecraft) (Omidshafiei al. (0018). In this experiment, random policy used generate data for dynamics model training. As shown in Figure 15 ED2 could bring more accurate dynamics modeling process.
DREAMER WITH BIGGER HIDDEN SIZE
In this section, we provide the result of Dreamer method under bigger hidden size (which keeps the similar parameter size as ED2-Dreamer). As shown in Figure I6] increasing the size of parameters can not improve the performance.
15

---

Under review as a conference paper at ICLR 2022
ラ Ant
Walker2d
Hopper
HalfCheetah
Figure 14: The visualization of final partition of environments in Gym-Mujoco.
Minecraft
Enduro-vo
Baseline
Baseline
ED2
ED2
0.04
0.0015
Value
ANIPA
0.02
0.0010
0.00
0.0005
2
Step
x10
Step
x10*
Figure 15: Model error experiments Atari and Atari like maze environment.
dmc hopper hop
dmc walker run run
400
Value
Value
200
Dreamer
Dreamer
@heamer_large_heeee hidden
0.0
0.5
1.0
Step
x10
Step
x10
Figure 16: The performance of Dreamer under bigger parameter size.
16

---

Under review as a conference paper at ICLR 2022
Table 3: result)
The meaning of action dimensions in each environment (listed according the clustering
ENVIRONMENT
ACTION DIMENSION MEANNING
RIGHT_ANKLER _AKLL_X
LEFT _EF__ANKL__X
ABDOMEN X, RIGHT _IP_ LEFT HIF
ABDOMEN RIGHT
RIGHT KNEE,RIGHT, ANKLE_
HUAMNOI((DEEPMIND)
LEFT KNBE,LEFT_ANKEE ANKLE_Y
RIGHT HIP Z. LEFT HIP
LEFT SHOULDERL. LEFT ELBOW, RIGHT SHOULDERI, RIGHT ELBOW, RLBOW. HIGH__SHOULDER2, LEFT_ SHOULLEOO,,,, ABDOMEN_Z
LEFT HIP, RIGHT HIP
WALKEN(DEEPMIND)
LEFT KNEE, RIGHT KNEE JEFL_ANKLE, RIGHT_ANKLA
BACK THIGH, BACK SHIN, BACK FOOT
CHEETAH(DEEPMINN)
FRONT THIGH, FRONT SHIN, FRONT FOOT
WAIST, HIP
(DEPMN)
KNEE
ANKLE
SHOULDER
REACHER(DEEPMINN)
WRIST
PROXIMAL
(DEEERDDEEPMIND)
DISTAL
BACK BACK_THIGH, BACK SHIN FRONT_THIGH, FRONT_FSHIN
HAUFCHHEETAM(GGPMMMUUOO))
BACK_FOOT
FRONT_ FOOT
RIGHT THIGH, RIGHT LEG, SEFT_THIGH, LEFT
WALKER20(GYx-MUocc))
RIGHT FOOT
HOPPER(GYM-MUrOCO)
THIGH, LEG
FOOT
LEFT_ SWFT_PRONT_HIP RIGHT RIGHT_FRONT, LEFT_ B_AK_AACKPPPH HIP, RIGHT ______ACC_HPPP_C___
ANT(GYM-MUJOCO)
LEFT_ SEF__FRONT_AKKEEE LEFT SEFT_BACKAANKE
RIGHT FRONT ANKLE, RIGHT BACK ANKLE
17


Reviews:
[{"review_id": "qQGqZ1SV-3m", "reviewer": "Reviewer_6o3p", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "However, I have the following concerns:\n\nQ1. The sub-dynamics model is to output an embedding vector regarding sub-actions and state $s$, but not to predict the next state $s’$ and reward $r$. The sub-dynamics model is therefore not a dynamics model. I don’t think it is precise to claim that ED2 is a dynamics decomposition method. \n\nQ2. As the underlying idea of ED2 is very simple, I believe it is possible to theoretically analyze the advantage of ED2, which can improve the paper’s contribution significantly. \n\nQ3. Doya et al. [1] also leverage the idea of dynamics decomposition to model the environmental dynamics but decompose the dynamics into multiple domains in state space and time. Although their method is different from ED2, I suspect that it is important to discuss the difference between the two studies because both are based on the idea of dynamics decomposition.\n\nReference:\n\n[1] Doya, Kenji, et al. \"Multiple model-based reinforcement learning.\" Neural computation 14.6 (2002): 1347-1369.\n\nQ4. The introduction section mentions many MBRL algorithms, such as ME-TRPO [1] and SimPLe [2], but the authors empirically compare ED2 with only two MBRL methods in the experiment section.\n\nReferences:\n\n[1] Kurutach, Thanard, et al. \"Model-ensemble trust-region policy optimization.\" ICML. 2018.\n\n[2] Kaiser, Łukasz, et al. \"Model Based Reinforcement Learning for Atari.\" ICLR. 2019.\n", "summary_of_the_paper": "The authors propose an environment dynamics decomposition (ED2) framework to decompose the environmental dynamics into multiple sub-dynamics. Empirical results show that ED2 improves the performance of several model-based reinforcement learning (MBRL) algorithms. The paper is well-written and easy to read, and the motivation is clear.", "summary_of_the_review": "Although the paper is well written, the current contribution is not enough for me to recommend acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "W-sXV6GW5rj", "reviewer": "Reviewer_BTNs", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have some issues with the technical approach, raising doubts over how this method will generalize. In particular in light of the small to moderate improvements in the experiments sections and small selection of benchmarks. \n\nTechnical approach: \n- The authors make a strong and seemingly very limiting assumption in that they only partition actions, all sub-models get fed the same state input. Then they average the output of all submodels to get the new state, also a strong assumption that the state can be additively decomposed this way. I assume this means their submodels output the full state vector as well, almost like an ensemble approach for the latent state with an extra clustering step to filter out just the action inputs for each sub-network.\n\n- The clustering approach itself also seems to make a very strong assumption in clustering bon linear correlations, and of the action dimensions vs. the state dimensions. I am not sure why this makes sense because linear correlations strongly depend on how you encode the state and action space. For example, it is signed so if you just negate one action dimension, shouldn't that change which other actions it gets clustered with it, even though it will be functionally the same? Further, as dynamics models, shouldn't actions be clustered based on how they predict the next state, not how they are correlated with the current state? In general, the Feature Extraction and Cluster Criterion sections are crucial to this paper but the theoretical basis of the design choices made there are not clear to me. \n\nExperiments:\n- The Dreamer and MBPO papers you compare against have 20 and 6 benchmarks respectively, it is unfortunate that you only present a subset of 4 from each. This becomes relevant because the improvements aren't large to begin with and some questionable assumptions were seemingly made in the method section. As you allude to in Appendix F, the method also appears to only improve results on problems that are easily decomposable, which seems reasonable.\n\nLanguage: \n- Readable but some bad grammar. I recommend another pass. In particular, there is frequent overuse of \"the\", as in \"by combining the optimal control algorithm\" when you mean \"by combining [with] optical control algorithms\"\n\nMinor:\n- Is it necessary to describe it as a POMDP? The paper by Hafner et al 2020 that this paper builds on does not. Your state space definition is seemingly much more complex than theirs, but as I understand it you use the same benchmarks.", "summary_of_the_paper": "The authors propose a clustering method to decompose world models for model-based RL into sub-models reflecting the dynamics of the task. The method appears to give small to moderate improvements on most of the presented 4+4 benchmarks. \n", "summary_of_the_review": "The problem of learning decomposed world models is interesting and the authors do show a small improvement on most of the presented benchmarks.\n\nUnfortunately, the clustering method, which is the main contribution of the paper, makes a number of strong and unexpected assumptions that make me question how well this approach will generalize. The paper also suffers from poor clarity in the method section on the theoretical motivation for these design decisions. The empirical results are not very strong to begin with, and were only presented for a smaller subset of the benchmarks in the papers they compare against. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6qEXY4tR7WB", "reviewer": "Reviewer_SSLz", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposed to build multiple world models, one per action space partition, and average it to improve the model error.\nThis idea is generic to other MBRL approaches, and the paper implements the method in \"Dreamer\" and \"MBPO\" and shows improvements in the model error.\n\n#### Strengths\n* The idea is simple, and the experiment results showed improvements in model error or value in some domains\n\n#### Weakness\n* If I understood correctly, this action space partitioning can only be applicable when there is a clear decomposable structure in the action space. \n* The result shows some marginal improvements on the five random seeds.\n\n#### Questions\n* Have you experimented with this dynamics decomposition approach to other domains like ATARI?\n\n* Could you provide a side-by-side comparison of NN on ED2-Dreamer and simple Dreamer?\nWill the total number of parameters remain the same or similar? Does the performance changes per the choice of hyperparameters?\n\n* In Figure 6, there is a gap in the model errors between the two approaches. If we run on more sample interactions, will this gap disappear? Or remains the same?\n\n* The main difference is that ED2-Dreamer usually shows lower variance over the five random seeds than Dreamer. Is this due to the averaging effect of the dynamic decomposition? In Figure 5, ED2-Ensemble takes the whole action space, learns multiple models, and averages the results. We can see that the blue and the red curves are less similar in the cheetah and hopper domains. Would it imply that the decomposition doesn't bring improvements in those domains?\n\n* When combining the latent states, the paper proposed to use averaging the outcomes. Have you tested other choices like max or putting an additional layer?\n\n\n", "summary_of_the_paper": "This paper presents a model-based RL algorithm that provides decomposed dynamics models when the state and action spaces are defined as a multi-dimensional Cartesian space. \nWe can divide the method into two parts:\n* The one for partitioning the action coordinates, and\n* The NN architecture, one for each partition and then combining the outcomes ($h^i_t$) of partitions into the final one ($h_t$).\nThis paper also shows a clustering based on the correlation, human-provided partition, and full/random partitioning algorithms.\n \nThe experiment shows the results on Mujoco-based environments that have orthogonal coordinates per position and velocity, so\nthe action coordinate partitioning can be readily applicable.", "summary_of_the_review": "This paper shows MBRL methods that utilized the decomposition of the action space and showed experiment results on the control domains.\nIn some domains, we see improvements in the model error or higher values. But the overall improvement is marginal, and the method can only be applicable to the problems having readily decomposable action space.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: One of the major challenges in model-based RL is the learning of accurate world models. This work proposes a method that can learn to decompose the dynamics of the world into several sub-dynamics, which is postulated to lead to better prediction accuracy (which in turn should lead to better policies/downstream performance). The proposed method clusters the sub-dynamics in latent space, and can be combined with any existing MBRL method. Here it is specifically combined with Dreamer and MBPO and evaluated on the deepmind control suite/mujoco.

**Strengths**
This work addresses an important problem, and on a high-level the problem/approach is well motivated
The proposed algorithm is "simple" (which can be a really good thing) and very general in that it can be combined with seemingly any MBRL method

**Weaknesses**
The manuscript was lacking in clarity on a technical level (partially addressed during rebuttal)
On average the experimental results are not very convincing (yet)

**Rebuttal**
The authors clarified a few misunderstandings the reviewers and also updated the manuscript accordingly

**Summary**
I agree with the reviewers that the experimental results are not fully convincing. When looking at the model error plots, the y-scale is very small, and it looks like there is no significant improvement. Also in the "downstream tasks" only for the humanoid do we seem to see a significant improvement. Overall this seems promising, but the authors should investigate why there seems to be a clearer benefit for the humanoid, and show more such results. Furthermore, there are some concerns/clarity issues with respect to what your approach learns - I would recommend you take a small-ish toy example to introduce the intuition of your approach, and maybe visualize learned sub-dynamics.
Overall, while promising, in it's current state this manuscript is not quite ready yet for publication
Decision: Reject