ED2: An Environment Dynamics Decomposition Framework for World Model Construction

1 INTRODUCTION . Reinforcement Learning ( RL ) is a general learning framework for solving sequential decision-making problems and has made significant progress in many fields ( Mnih et al. , 2015 ; Silver et al. , 2016 ; Vinyals et al. , 2019 ; Schrittwieser et al. , 2019 ) . 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 et 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. , 2020 ; Yu et al. , 2020 ) ; the second category leverages the differentiability of the world model , and generates differentiable trajectories for policy optimization ( Deisenroth & Rasmussen , 2011 ; 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 et al. , 2018 ; Feinberg et al. , 2018 ) ; 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 et 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 1Our code is open source and available at https : //github.com/ED2-source-code/ED2 prediction ( Luo et al. , 2019 ; Kurutach et al. , 2018 ) . Model ensemble is also widely used in model construction for uncertainty estimation , which provides a more reliable prediction ( Janner et al. , 2019 ; Pan et al. , 2020 ) . To reduce the model error in long trajectory generation , optimizing the multi-step prediction errors is also an effective technique ( Hafner et al. , 2019 ) . However , these techniques improve the environment modeling in a black-box way , which ignores the inner decomposed structure of environment dynamics . For example , Figure 1 ( a ) shows the Cheetah task from DeepMindControl ( DMC ) Suite tasks , where the dynamics can be decomposed in various ways . Figure 1 ( b ) shows various decomposition on the dynamics : according to the role of sub-dynamics , we can decompose it into : { thigh , shin , foot } ; alternatively , according to the position of sub-dynamics , we can decompose it into : { back , front } . Figure 1 ( d ) shows that no matter whether we decompose the Cheetah task according to role or position , modeling each decomposed sub-dynamics separately can significantly reduce the model error of the existing MBRL algorithm ( e.g . Dreamer ( Hafner et al. , 2020 ) ) . Inspired by the above example , we propose environment dynamics decomposition ( ED2 ) , a 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 prediction ( 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-based method . D2P is proposed to construct the world model from the decomposed dynamics in SD2 , which models each sub-dynamics separately in an 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 MBRL algorithms . 2 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 , γ , O , Ω , T ) , where S ∈ Rn is the state space , and A ∈ Rm is the action space , R : S × A → R denotes the reward function , P : S × A → S denotes the environment dynamics , γ is the discount factor . The agent receives an observation o ∈ Ω , which contain partial information about the state s ∈ S. O is the observation function , which mapping states to probability distributions over observations . The decision process length is denoted as T . Let η denote the expected return of a policy π over the initial state distribution ρ0 . The goal of an RL agent is to find the optimal policy π∗ which maximizes the expected return : π∗ = arg max π η [ π ] = arg max π Eπ [ T∑ t=0 γtR ( st , at ) ] , where s0 ∼ ρ0 , ot ∼ O ( ·|st ) , at ∼ π ( ·|ot ) , st+1 ∼ P ( ·|st , at ) . If the environment is fully observable , i.e. , Ω = S and O is an identity function , POMDP is equivalent to the MDP : ( S , A , R , P , γ , T ) . 2.2 REPRESENTATIVE WORLD MODELS IN MBRL . The world model is a key component of MBRL that directly impacts policy training . World models are usually formulated with latent dynamics ( Janner et al. , 2019 ; Hafner et al. , 2020 ) , and the general form of the latent dynamics model can be summarized as follows : Latent transition kernel : ht = f ( s≤t−1 , a≤t−1 ) Stochastic state function : p ( st|ht ) Reward function : p ( rt|ht ) The latent transition kernel ( shorthand as kernel ) predicts the latent state ht with input s≤t−1 and a≤t−1 . Based on latent state ht , the stochastic state function and reward function decode the state st and reward rt . For the partially observable environment , two additional functions are required : Observation function : p ( ot|st ) Representation function : p ( st|ht , ot ) 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 : ht = { f ( st−1 , at−1 ) With non-recurrent kernel f ( ht−1 , st−1 , at−1 ) With recurrent kernel Non-recurrent kernel are relatively basic kernel for modeling , which are often implemented as Fully-Connected Networks . Non-recurrent kernel takes the current state st−1 and action at−1 as input , outputs the latent state prediction ht . Compare to non-recurrent kernel , recurrent kernel is implemented as RNN and takes the additional input ht−1 , which performs better under POMDP setting . For both kernels , the st and rt can be generated from the latent prediction ht . 3 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 1 . However , these techniques improve the environment modeling in a 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 : 1 ) Decomposability : The environment dynamics can be decomposed into multiple subdynamics in various ways and the decomposed sub-dynamics can be combined to reconstruct the entire dynamics . 2 ) Traceability : The environment dynamics can be traced to the action ’ s impact on the environment , and each sub-dynamics can be traced to the impact caused by a part of the action . For example in the Cheetah task , Figure 1 ( b ) demonstrates the decomposability : we can decompose the dynamics into { thigh , shin , foot } 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 . Figure 1 ( c ) explains the traceability : each sub-dynamics can be 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-dynamics , which allows us to model the dynamics separately , while the traceability investigates the causes of the dynamics and guides us to decompose the dynamics at its root ( i.e . the action ) . 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 to discover the latent sub-dynamics by analyzing the action ( SD2 , the blue part in Figure 2 ) ; by considering the decomposability , we propose to construct the world model in a decomposing manner ( D2P , the green part in Figure 2 ) . Our framework can be used as a 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 ⊂ Rm , the index of each action dimension constitutes a set Λ = { 1 , 2 , · · · , m } , any disjoint partition G = { G1 , . . . , Gk } over Λ corresponds to a particular way of decomposing action space . For each action dimension i in Λ , we define the action space as Ai , which satisfied A = A1 × · · · ×Am . The action space decomposition under partition G is defined as AG = { AG1 , · · · , AGk } , where sub-action space AGj = ∏ x∈Gj A x . Based on above definitions , we define the dynamics decomposition for P under partition G as follows : Definition 1 Given a partition G , the decomposition for P : S ×A→ S can be defined as : P ( s , a ) = fc ( 1 k k∑ i=1 Pi ( s , a Gi ) ) , ∀s , a ∈ S ×A , ( 1 ) with a set of sub-dynamics functions { P1 , ... , Pk } that Pi : S ×AGi → H , and a decoding function fc : H → S. Note H is a latent space and aGi ∈ AGi 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 . In this section , we mainly focus on dynamics modeling , and we will introduce how to derive the partition G by using SD2 in section 3.3 . To implement D2P , we use model M iφi parameterized by φi ( i.e. , neural network parameters ) to approximate each sub-dynamics Pi . As illustrated in Figure 2 , given a partition G , an action a is divided into multiple sub-actions { aG1 , · · · , aGk } , each model M iφi takes state s and the sub-action aGi as input and output a latent prediction hi ∈ H . The separate latent predictions { h1 , · · · , hk } are aggregated and then decoded for the generation of state s′ and reward r. For each kernel described in Section 2.2 , we provide the formal description here when combine with D2P : ht = { 1 k ∑k i=1 f ( st−1 , a Gi t−1 ) For non-recurrent kernel 1 k ∑k i=1 f ( ht−1 , st−1 , a Gi t−1 ) For recurrent kernel We propose a set of kernels , where each kernel models a specific sub-dynamics with the input of current state s , corresponding sub-action aGi and hidden state ht−1 ( ignored when applying on non-recurrent kernel ) . The output of all kernels is averaged to get the final output ht . The prediction of reward rt and state st is generated from the output ht . Specifically , we provide an example when combining with the kernel of Recurrent State-Space Model ( RSSM ) ( Hafner et al. , 2019 ) in Figure 3 , which is a representative recurrent kernel-based world model . The original single kernel implemented as GRU are replace by multiple kernels with different action input .

Summary: 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.  We can divide the method into two parts: * The one for partitioning the action coordinates, and * The NN architecture, one for each partition and then combining the outcomes ($h^i_t$) of partitions into the final one ($h_t$). This paper also shows a clustering based on the correlation, human-provided partition, and full/random partitioning algorithms.   The experiment shows the results on Mujoco-based environments that have orthogonal coordinates per position and velocity, so the action coordinate partitioning can be readily applicable.