Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models

World models learn the consequences of actions in vision-based interactive systems. However, in practical scenarios such as autonomous driving, there commonly exists noncontrollable dynamics independent of the action signals, making it difficult to learn effective world models. To tackle this problem, we present a novel reinforcement learning approach named Iso-Dream, which improves the Dream-to-Control framework in two aspects. First, by optimizing the inverse dynamics, we encourage the world model to learn controllable and noncontrollable sources of spatiotemporal changes on isolated state transition branches. Second, we optimize the behavior of the agent on the decoupled latent imaginations of the world model. Specifically, to estimate state values, we roll-out the noncontrollable states into the future and associate them with the current controllable state. In this way, the isolation of dynamics sources can greatly benefit long-horizon decision-making of the agent, such as a self-driving car that can avoid potential risks by anticipating the movement of other vehicles. Experiments show that Iso-Dream is effective in decoupling the mixed dynamics and remarkably outperforms existing approaches in a wide range of visual control and prediction domains.


Introduction
Humans can infer and predict real-world dynamics by simply observing and interacting with the environment. Inspired by this, many cutting-edge AI agents use self-supervised learning [38,20,12] or reinforcement learning [39,22,42] techniques to acquire knowledge from their surroundings. Among them, world models [20] have received widespread attention in the field of robot visual control, and led the recent progress in model-based reinforcement learning (MBRL) [22,42,24,30]. A typical approach [22] is to use the trajectories of observations and control signals collected by an RL agent to learn a differentiable simulator of the environment, namely the world model, and then update the RL agent by optimizing the behaviors on the latent imaginations of the world model. However, since the observation sequence is high-dimensional, non-stationary, and often driven by multiple sources of physical dynamics, how to learn effective world models in complex visual scenes remains an open problem. In realistic scenarios such as autonomous driving, we can generally divide spatiotemporal dynamics in the system into controllable parts that perfectly respond to action signals, and parts beyond the control of the agent, such as the movement of other vehicles and other external changes. The isolation of controllable and noncontrollable states can improve MBRL in two aspects: • Modular representation improves the generalization of the agent to non-stationary environments with noises, such as the time-varying background in our modified DeepMind Control Suite. • More importantly, it improves long-horizon RL tasks that can greatly benefit from decisions based on predictions of future noncontrollable dynamics. For example, in autonomous driving, potential risks can be better avoided by predicting the movement of other vehicles.
We present Iso-Dream, a novel MBRL framework that learns to decouple and leverage the controllable and noncontrollable state transitions. Accordingly, it improves the original Dreamer [22] from two perspectives: (i) a new form of world model representations and (ii) a new actor-critic algorithm to derive the behavior from the world model. As shown in Figure 1, the foundation of decoupling the world model is to separate the mixed latent states into an action-conditioned branch and an action-free branch, which can individually transit different sources of visual dynamics. The components are jointly trained to maximize the variational lower bounds. To further isolate the controllable states, the action-conditioned branch is also optimized with inverse dynamics, that is, to reason about the actions that have driven the state transitions between adjacent time steps.
Another contribution of Iso-Dream is to find that disentangling physical dynamics can greatly benefit the downstream decision-making tasks by more accurately foreseeing the inherent changes in the environment. Intuitively, humans can decide how to interact with the environment at each moment based on their anticipation of future changes in the surroundings. To make more forward-looking decisions, as shown by the blue arrows in Figure 1, the policy network integrates the current controllable state and multiple steps of predicted noncontrollable states through an attention mechanism. It enables the agent to thoroughly consider possible future interactions with the environment.
We evaluate Iso-Dream in the following domains: The modified DeepMind Control Suite with noisy video background; The CARLA autonomous driving environment in which other vehicles can be naturally viewed as noncontrollable components; The real-world BAIR robot dataset and the RoboNet dataset that are helpful to validate the effectiveness of the world model for disentanglement. On all benchmarks, Iso-Dream remarkably outperforms the existing approaches by large margins.

Related Work
Action-conditioned video prediction. A straightforward deep learning solution to visual control problems is to learn action-conditioned video prediction models [38,14,8,53] and then perform Monte-Carlo importance sampling and optimization algorithms, such as the cross-entropy methods, over available behaviors [15,12,29]. Hot topics in video prediction mainly includes long-term and high-fidelity future frames generation [44,43,51,5,52,50,54,41,40,36,56,28,2], dynamics uncertainty modeling [1,10,48,31,7,16,55], object-centric scene decomposition [47,27,18,58,3], and space-time disentanglement [49,27,19,6]. The corresponding technical improvements mainly involve the use of more effective neural architectures, novel probabilistic modeling methods, and specific forms of video representation. The disentanglement methods are closely related to the world model in Iso-Dream. They commonly separate visual dynamics into content and motion vectors, or long-term and short-term states. In contrast, Iso-Dream is designed to learns a decoupled world model based on controllability, which contributes more to the downstream behavior learning process.
Visual MBRL. In visual control tasks, the agents have to learn the action policy directly from high-dimensional observations. They can be roughly grouped into two categories, that is, model-free methods [34,57,32,33,25] and model-based methods [15,39,20,23,22,30,42,59,4]. Among them, the MBRL approaches explicitly model the state transitions and generally yield higher sample efficiency than the model-free methods. Ha and Schmidhuber [20] proposed the World Models that first learn compressed latent states of the environment in a self-supervised manner, and then train the agent on the latent states generated by the world model. Following the two-stage training procedure, PlaNet [23] uses an action-conditioned, recurrent state-space model (RSSM) as the world model, and optimizes the action policy on the recurrent states with the cross-entropy methods. In Dreamer [22] and DreamerV2 [24], agents learn behaviors by optimizing the expected values over the predicted latent states in RSSM. InfoPower [4] prioritizes functional-related information from visual observations to obtain a more robust representation for MBRL. Notably, Iso-Dream is very different from InfoPower in two aspects. First, we explicitly model the state transitions of controllable and noncontrollable dynamics, so that it is possible to choose whether to take the noncontrollable states into behavior learning according to the prior knowledge of a specific domain. Second, we propose a new behavior learning method that greatly benefits from the decoupled world model, so that we can preview possible future states of noncontrollable patterns before making decisions at this moment.

Method
In this section, we first present basic assumptions and the general framework of Iso-Dream for decoupling and leveraging controllable and noncontrollable dynamics for visual control (Section 3.1).
For representation learning, we introduce the three-branch world model and its training objectives of inverse dynamics (Section 3.2). For behavior learning, we present an actor-critic method that is trained on the imaginations of the decoupled world model latent states, so that the agent may consider possible future states of noncontrollable dynamics (Section 3.3). Finally, we discuss how Iso-Dream is deployed to interact with the environment (Section 3.4).

Basic Assumptions of Iso-Dream
As shown in Figure 1, when the agent receives a sequence of visual observations o 1:T , the underlying spatiotemporal dynamics can be defined as u 1:T . Our goal is to understand the inner relationships of different dynamics by decoupling u 1:T into controllable latent states s 1:T and noncontrollable latent states z 1:T that vary in spacetime, such that: where a t is the action signal. To achieve long-term prediction, we isolate s t and z t to each other and model their state transitions of p(s t+1 | s t , a t ) and p(z t+1 | z t ) respectively.
According to our prior knowledge of the environment, we can optionally choose whether to roll out the noncontrollable states and consider them during behavior learning. For tasks where the noncontrollable components can be viewed as time-varying noise, we simply derive the action policy by a t ∼ π(a t | s t ). The isolation of controllable states improves the generalization of the agent to non-stationary systems. For tasks like autonomous driving, the behaviors are derived by a t ∼ π(a t | s t , z t:t+τ ), where we calculate the relationships between s t and the imagined noncontrollable states over time horizon τ . It assumes that, in specific long-horizon tasks, the agent can greatly benefit from predicting the consequences of external noncontrollable forces.

Representation Learning of Controllable and Noncontrollable Dynamics
Inspired by previous approaches [37,17] showing that modular structures are effective for disentanglement learning, we leverage a three-branch architecture to decouple u t into controllable dynamics state s t , noncontrollable dynamics state z t , and time-invariant representation of the background. As shown in Figure 2(a), the action-conditioned branch models p(s t+1 | s t , a t ). It follows the RSSM architecture from PlaNet [23] to use a recurrent neural network GRU s (·), the deterministic hidden state h t , and the stochastic state s t to form the transition model, where the GRU keeps the historical information of the controllable dynamics. The action-free branch models p(z t+1 | z t ) with similar network structures. The transition models with separate parameters can be written as follows: Action-free branch Static branch Action-conditioned branch Action-free branch We here uses t andz t to denote the prior representations. We optimize the transition models with posterior representations that are derived from We learn the posteriors from the observation at current time step o t ∈ R 3×H×W by a shared encoder Enc θ and subsequent branch-specific encoders Enc φ1 and Enc φ2 .
To enhance the disentanglement representation learning corresponding to the control signals, we introduce the training objective of inverse dynamics. Accordingly, we design an Inverse Cell of a 2-layer MLP to infer the actions that lead to certain transitions of the controllable states: where the inputs are the posterior representations in the action-conditioned branch. By learning to regress the true behavior a t−1 , the Inverse Cell facilitates the action-conditioned branch to isolate the representation of the controllable dynamics. To avoid the training collapse where the actionconditioned branch captures most of the useful information, while the action-free branch learns almost nothing, in the process of image reconstruction, we respectively use the prior states t and the posterior state z t to generate the controllable visual componentô s t ∈ R 3×H×W with mask M s t ∈ R 1×H×W and the noncontrollable componentô z t ∈ R 3×H×W with M z t ∈ R 1×H×W . By further integrating the time-invariant information extracted from the first K frames, we havê For reward modeling, we have two options with the action-free branch. In one case, the noncontrollable dynamics can be considered as noises that are not related to the task, and therefore z t is no longer useful during imagination. In other words, the policy and the predicted reward are only related to the controllable states. In the other case, future noncontrollable states would affect how the agent makes decisions, and we consider the action-free components during behavior learning. For this, we learn alternative reward models p(r t | s t ) or p(r t | s t , z t ) in forms of MLPs.
For a sequence of (o t , a t , r t ) T t=1 sampled from the replay buffer during training, the world model can be optimized using the following loss functions, where α, β 1 , and β 2 are hyper-parameters: Roll-out the noncontrollable states {zi} t+L+τ i=t+1 from zt through the action-free branch alone. 10 for time step j = i . . . i + L do 11 Compute latent state ej ∼ Attention(sj,zj:j+τ ) using Eq. (7). 12 Imagine an action aj ∼ π(aj|ej). 13 Predict the next controllable statesj+1 ∼ p(sj, aj) using the action-conditioned branch alone.
Roll-out the noncontrollable stateszt+1:t+τ from zt through the action-free branch alone. 22 Generate at ∼ π(at | st, zt,zt+1:t+τ ) using future state attention in Eq. (7). The world model training approach can be partly customized for different environments. In situations where noncontrollable states are indeed involved in behavior learning, minimizing the ELBO objective can maintain the semantics ofz t . Otherwise, if the action-free features are only used to prevent noisy distractions from affecting the training process of Iso-Dream, rather than being used for behavior learning, we can simply train the action-free branch with the reconstruction loss alone.

Behavior Learning in Decoupled Imaginations
Thanks to the decoupled world model, we can optimize the agent behaviors to adaptively consider the relations between available actions and possible future states of the noncontrollable dynamics. A practical example is autonomous driving, where the motion of other vehicles can be naturally viewed as noncontrollable but predictable components. As shown in Figure 2(b), we here propose an improved actor-critic learning algorithm that 1) allows the action-free branch to foresee the future ahead of the action-conditioned branch, and 2) exploits the predicted future information of noncontrollable dynamics to make more forward-looking decisions.
Suppose we are making decisions at time step t in the imagination period. A straightforward solution from the original Dreamer method is to learn an action model and a value model based on the isolated controllable states t ∈ R 1×d . However, we notice that by employing an attention mechanism, we can explicitly calculate its relations to a sequence of future noncontrollable statesz t:t+τ ∈ R τ ×d , where τ is the length of a sliding window from now on.
In this way,s t evolves to a more "visionary" representation e t ∈ R 1×d . We update the action model and the value model in Dreamer [22] as follows: Action model: a t ∼ π(a t | e t ), where L is the imagination time horizon. As shown in Alg. 1, during imagination, we first use the action-free transition model to obtain sequences of noncontrollable states of length L + τ , denoted . At each time step in the imagination period, the agent draws an action a j from the visionary state e j , which is derived from Eq. (7). The action-conditioned branch uses the action a j in latent imagination and predicts the next controllable state s j+1 . We follow DreamerV2 [24] to train the action model to maximize the λ-return [45], and train the value model to regress the λ-return 3 .

Policy Deployment by Rolling-out Noncontrollable Dynamics
As discussed above, in the cases that noncontrollable dynamics are irrelevant to the control task, when interacting with the environment, we only use the state of controllable dynamics to generate the policy at each time step t. However, for the situation where noncontrollable dynamics should be closely related to the behavior of the agent, as shown in Lines 21-22 in Alg. 1, the action-free branch consecutively predicts the next τ − 1 noncontrollable statesz t+1:t+τ starting from the current posterior state z t . Similar to Eq. (7) in the process of behavior learning, we here use the learned future state attention network to adaptively integrate s t , z t andz t+1:t+τ . Based on the integrated feature e t , the Iso-Dream agent draws a t from the action model to interact with the environment.

Experimental Setup
Benchmarks. We quantitatively and qualitatively evaluate Iso-Dream on two reinforcement learning environments, i.e., DeepMind Control Suite [46] and CARLA [11], and two real-world datasets for action-conditioned video prediction, i.e., BAIR robot pushing [13] and RoboNet [9]. The video prediction experiments can provide more intuitive visualizations of disentanglement learning.

DeepMind Control Suite
Implementation details. In order to verify the enhancement of Iso-Dream by disentangling different components under complex visual dynamics, we evaluate Iso-Dream on environments from DMC Generalization Benchmark. Instead of training on original DeepMind Control Suite environments, agents are trained and tested both with natural video backgrounds (i.e. video_easy environments). In this environment, since the background is randomly replaced by a real-world video, the noncontrollable motion of the background can affect the procedure of dynamics learning and behavior learning of agents. Therefore, to obtain a better decision policy and avoid the disruption from noisy backgrounds, the agent may decouple noncontrollable representation (i.e., dynamic background) and controllable representation in spacetime, and only use controllable representation for control. To this end, we simply train the action-free branch with only reconstruction loss and discard it in imagination and policy deployment. We evaluate our model with baselines in 4 tasks from four different domains. The number of environment steps is limited to 500k.
Quantitative results. To evaluate the performance, we train and test the agents in environments with video backgrounds. As shown in Table 1, Iso-Dream exceeds the performance of DreamerV2 and other baselines in all tasks, indicating that the three-branch structure can effectively learn task-related visual representations and alleviate complex background interference in visual data. For each sequence, we use the first 5 images as context frames. Iso-Dream successfully disentangles controllable and noncontrollable components.
Qualitative results. We leverage Iso-Dream to complete video prediction tasks in video_easy environments. The sequence of frames and actions are randomly collected during test episodes. The first 5 frames are given to the model and the next 45 frames are predicted only based on action inputs. To show the qualitative results, we visualize the masks and visual decoupled components from the action-conditioned and action-free branches. The overall visualization is shown in Figure 3(left). From this prediction result, we can find that Iso-Dream has the ability to predict long-term sequence and disentangle controllable and noncontrollable dynamics from images in video_easy environments. As shown in the third and fourth row of action-conditioned branch output in Figure 3, the controllable representation has been successfully isolated and matches its mask. Besides, in this visualization, the action-free component in this background video is the motion of sea waves, which is captured by the fifth and sixth row of action-free branch outputs.

CARLA Autonomous Driving Environment
Implementation details. In the autonomous driving task, We use a camera with 60 degree view on the roof of the ego-vehicle, which obtains images of 64 × 64 pixels. Following the setting in the DBC [59], in order to encourage highway progression and penalise collisions, the reward is formulated as: where v ego is the velocity vector of the ego-vehicle, projected onto the highway's unit vectorû h , and multiplied by time discretization ∆t = 0.05 to measure highway progression in meters. Impulse I ∈ R + is caused by collisions, and a steering penalty steer ∈ [−1, 1] facilitates lane-keeping. The hyper-parameters ξ 1 and ξ 2 are set to 10 −4 and 1, respectively. We use β 1 = 1, β 2 = 1 and α = 1 in Eq. (6) and τ = 5 in Eq. (7).
Quantitative results. As shown in Figure 4(a), Iso-Dream has significant advantages compared to other baselines and outperforms DreamerV2 by a large margin. Furthermore, we conduct ablation studies to confirm the validity of inverse dynamics and the rolling-out strategy of noncontrollable states. Figure 4(b) shows that the performance drops when Inverse Cell is removed, indicating the importance of modeling inverse dynamics to isolate controllable and noncontrollable components from the whole dynamics. In order to verify the effectiveness of the proposed attention mechanism, we conduct experiments to evaluate Iso-Dream where policy networks directly concatenate the current controllable state and the noncontrollable state as input. Comparing the blue curve and green curve,   we observe that rolling-out noncontrollable states in the action-free branch can significantly improve the agent's decision-making results. The red curve shows that the performance of Iso-Dream degrades by about 15% in the absence of a separate network branch that captures the static information.
Qualitative results. Reconstruction results of predictions in CARLA environment are shown in Figure 3(right column). In CARLA, we observe that the agent actions potentially affect all pixel values in the observation, as the camera on the main car (i.e., the agent) moves. Therefore, we view the visual dynamics of other vehicles as a combination of controllable and noncontrollable states. Accordingly, our model can determine which component is dominant by learning attention masks (values between 0 and 1) across the action-conditioned and action-free branches. The "action-free masks" present hot spots around other vehicles, while the attention values in corresponding areas on the "action-cond masks" are still greater than 0. The agent can avoid collisions by rolling-out noncontrollable components to preview possible future states of other vehicles. We include more showcases with different numbers of vehicles in the supplementary materials.

BAIR & RoboNet for Action-Conditioned Video Prediction
Implementation details. In order to evaluate the effectiveness of our world model in a more complex environment, we test the video prediction ability of the proposed structure on BAIR and RoboNet dataset. Moreover, we add predictable visual dynamics unrelated to the control signals to the raw observations, i.e., bouncing balls of the same size and speed. In the training phase, we train the model to predict 10 frames into the future from 2 observations. For testing, we use the first 2 frames as input to predict the next 28 frames in the BAIR dataset, and the next 18 frames in the RoboNet dataset. All inputs for training and testing are resized to 64 × 64. Considering the simplicity and predictability of bouncing balls, in the action-free branch, we use a similar structure as in the DMC experiment. Moreover, we replace the GRU cell with two layers of ST-LSTM unit [52] in both branches. The optimization objective consists of image reconstruction loss and action reconstruction loss of Inverse Cell. SSIM and PSNR are adopted as evaluation metrics.  Figure 5: Showcases of video prediction results on the BAIR robot pushing dataset. We display every 3 frames in the prediction horizon. The generated masks show that each branch of Iso-Dream captures coarse localisation of controllable representations and noncontrollable representations.
Quantitative results. Table 2 gives the quantitative results on BAIR and RoboNet datasets with bouncing balls in the training and testing phase. Compared with other models, Iso-Dream shows the competitive performance in two datasets. For PSNR, Iso-Dream improves SVG by 7.7% in BAIR and 9.3% in RoboNet. Compared with PhyDNet, which also disentangles features in two branches, Iso-Dream achieves better performance in both PSNR and SSIM. It shows that our Iso-Dream has a stronger ability of disentanglement learning to achieve long-term prediction.
Qualitative results. We visualize a sequence of predicted frames on BAIR with bouncing balls in Figure 5. Specifically, the output of two branches and corresponding masks are provided. We can see from these demonstrations that the world model of Iso-Dream is more accurate in modeling future dynamics for long-term prediction. It shows the fact that the action-free branch learns noncontrollable dynamics, while the action-conditioned branch learns controllable dynamics related to input action.

Conclusions
In this paper, we proposed an MBRL framework named Iso-Dream, which mainly tackles the difficulty of vision-based prediction and control in the presence of complex visual dynamics. Our approach has two novel contributions to world model representation learning and corresponding MBRL algorithms. First, it learns to decouple controllable and noncontrollable latent state transitions via modular network structures and inverse dynamics. Further, it makes long-horizon decisions by rolling-out the noncontrollable dynamics into the future and learning their influences on current behavior. Iso-Dream achieves competitive results on the CARLA autonomous driving task, where other vehicles can be naturally viewed as noncontrollable components, indicating that with the help of decoupled latent states, the agent can make more forward-looking decisions by previewing possible future states in the action-free network branch. Besides, Iso-Dream was shown to effectively improve the visual control task in a modified DeepMind Control Suite, as well as the visual prediction task on the BAIR robot pushing dataset and the RoboNet dataset.
One limitation of Iso-Dream is the computational efficiency. Compared with DreamerV2, it requires longer training time per episode due to more intensive state transitions in behavior learning. But fortunately, from Figure 4(a), Iso-Dream is more sample-efficient than existing MBRL methods. Another limitation is the special treatment for different environments. In our preliminary experiments, we attempted to use the same model architecture for all test benchmarks. However, we observed that different benchmarks have specific requirements on the network structure, which we found should be dependent on our prior knowledge of the environments.

A Benchmarks
We quantitatively and qualitatively evaluate Iso-Dream on the following two environments for visual control and two real-world datasets for action-conditioned video prediction.
• DeepMind control suite [46]: A set of stable, well-tested continuous control tasks that are easy to use and modify. For vision-based control, we use a modified version of the DeepMind control suite in DMControl Generalization Benchmark [26] to evaluate Iso-Dream. In this environment, agents are trained to complete different tasks with random natural video as backgrounds, namely video_easy and video_hard benchmarks. We use 4 tasks to test our Iso-Dream, i.e., Finger Spin, Cheetah Run, Walker Walk, Hopper Stand.
• CARLA [11]: An open-source simulator with more complex and realistic visual observations for autonomous driving research. In our experiments, we evaluate Iso-Dream in a first-person highway driving task in "Town04". The agent's goal is to drive as far as possible in 1000 time steps without colliding with the 30 other moving vehicles or barriers.
• BAIR robot pushing [13]: An action-conditioned video prediction dataset composed of hours of self-supervised learning with the robotic arm Sawyer. In each video, a random moving robotic arm pushes a variety of objects on similar tables with a static background. Each video also has recorded actions taken by the robotic arm which correspond to the commanded gripper pose.
• RoboNet [9]: A large-scale dataset contains action-conditioned videos of seven robotic arms interacting with a variety of objects from four different research laboratories, i.e., Berkeley, Google, Penn, and Stanford.

B Compared Methods
For visual MBRL, we compare our method with the following baselines and existing approaches: • DreamerV2 [24]: A model-based RL method that learns directly from latent variables in world models. The latent representation enables agents to imagine thousands of trajectories in parallel.
• CURL [34]: A model-free RL method that extracts high-level features from raw pixels using contrastive learning, maximizing agreement between augmented versions of the same observation.
• SVEA [25]: A framework for data augmentation in deep Q-learning algorithms that improves stability and generalization on off-policy RL.
• SAC [21]: A model-free actor-critic method that optimizes a stochastic policy in an off-policy way.
• DBC [59]: It learns a bisimulation metric representation without reconstruction loss, which are invariant to different task-irrelevant details in the observation.
For video prediction, we compare the proposed world model with the following approaches: • SVG [10]: This model introduces random variables into latent space, which ensures that the future trajectory is inherently random.
• SA-ConvLSTM [35]: Based on the self-attention mechanism, this model uses the self-attention memory to capture long-term spatial dependency.
• PhyDNet [19]: This model uses a two-branch architecture to disentangle PDE dynamics from unknown complementary information.

C Additional Visualization in DMC and CARLA
DeepMind Control suite. In Figure 6, more showcases on the DeepMind Control are presented with different noisy backgrounds. We show the visualization of the masks and decoupled components from three branches of Iso-Dream.
CARLA autonomous driving simulator. In Figure 7, we visualize the video prediction results on the CARLA environment with different numbers of vehicles. We train Iso-Dream with 30 vehicles and test with 10 vehicles and 20 vehicles respectively. D Additional Results on the BAIR Robot Pushing Dataset Figure 8 shows an interesting result of the different training sets (i.e., BAIR, BAIR+bouncing balls) and the same testing set (i.e., BAIR). Iso-Dream is the only approach that achieves improvements when training on noisy data with bouncing balls, as shown in Figure 8(red bars). In this training setup, it performs best on the standard test set without balls. Iso-Dream is built on a more efficient architecture than the baseline models. It provides a general framework that can be easily extended to other backbones.
Ablation study. In Table 3, the first row shows the results of removing the action-free branch in the world model of Iso-Dream. The performance has decreased from 21.43 to 20.47 and from 19.51 to 18.51 in PSNR for predicting the next 18 frames and next 28 frames respectively, indicating that modular network structures are effective for predictive learning by decoupling the controllable and noncontrollable representations. Comparing the second row and third row in the Table 3, we observe that modeling inverse dynamics can improve the performance by learning more deterministic state transitions given particular actions in the action-conditioned branch.

E Network Architectures for Different Environments
The networks and hyper-parameters used for different environments are shown in Table 4. The results of models trained on BAIR (blue) and BAIR + bouncing balls (red), and tested on BAIR. We use the first 2 frames as input to predict the next 18 frames. The horizontal axis represents the different models, and the vertical axes represent test results of PSNR and SSIM.  conv3-3 conv3-3 -