Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models
Authors: Minting Pan
Xiangming Zhu
Yunbo Wang
Xiaokang Yang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement learning
World model
Visual dynamics

Abstract: 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. Naturally, therefore, we need to enable the world models to decouple the controllable and noncontrollable dynamics from the entangled spatiotemporal data. To this end, we present a reinforcement learning approach named Iso-Dream, which expands the Dream-to-Control framework in two aspects. First, the world model contains a three-branch neural architecture. By solving the inverse dynamics problem, it learns to factorize latent representations according to the responses to action signals. Second, in the process of behavior learning, we estimate the state values by rolling-out a sequence of noncontrollable states (less related to the actions) into the future and associate the current controllable state with them. In this way, the isolation of mixed dynamics can greatly facilitate long-horizon decision-making tasks in realistic scenes, such as avoiding potential future risks by predicting the movement of other vehicles in autonomous driving. 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.

Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models
Minting Pan" MoE Key Lab of Artificial Intelligence, Al Institute, Shanghai Jiao Tong University {panmt53, xmzhu76, yunbow, xkyang) tu. edu.
Xiangming Zhu"
Yunbo Wang Xiaokang Yang
Abstract
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 -se-Dremm, which improves the Tream-to-Conine framework [22] in two aspects. First, by optimizing the inverse dynamics. we encourage the world model learn controllable and noncontrollable sources of spatiotemporal changes 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- onng-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 wide range of visual control and prediction domains.
isolated state transition branches. Second,
Introduction
Humans 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 [37 20 [2]
techniques acquire knowledge from their surroundings. Among them, world models 20 have received widespread attention in the field of robot visual
reinforcement learning 38 22
control, and led the recent progress model-based reinforcement learning (MBRL) 22 24 29 A typical approach [22) is to use the trajectories of observations and control signals collected by an RL agent learn 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 digheiimessional, non-stationary, and often driven by multiple sources physical dynamics, how to learn effective world models in complex visual scenes remains 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 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 stationay environments with noises, such as the time- varying background in our modified DeepMind Control Suite.
'Equal contribution.
'Corresponding author: Yunbo Wang.
Code available
httphhaww...mmaaaamaaeee com/panntIIso. Dream
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Inverse dynammics State inference Action, dependemoy Controllable state / Now.cowtvollable state
$1
$2
53
O1
a1
22 2
Q2
Ø1
22
23
24
Figure Probabilistic graph of Aso-Dream. It learns to decouple complex visual dynamics into controllable states (s,) and noncontrollable states (2z) by optimizing the inverse dynamics (Red dashed arrows). On top of the disentangled states, it performs model-based reinforcement learning by explicitly considering the predicted noncontrollable component of future dynamics (Blue arrows).
More importantly, it improves long- horizon RL tasks that greatly benefit from decisions based on predictions 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) new actor-critic algorithm to derive the behavior from the world model. As shown in Figure T the foundation of decoupling the world model is to separate the mixed latent states into an action- conii-nond branch and an action-free 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 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 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 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 [] the policy network integrates the current control- lable state and multiple steps of predicted noncontrollable states through an attention mechanism. It enables the agent thoroughly consider possible future interactions with the environment.
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.
We evaluate Iso-Dream in the following domains:
2 Related Work
Action-conditioned video prediction. A straightforward deep learning solution to visual control problems to learn action-conditioned video prediction models [37] g 52 and then perform Monte-Carde importance sampling and optimization algorithms, such as the cross-entropy methods, over available behaviors [[5] 28). Hot topics in video prediction mainly includes long-term and
high-fidelity future frames generation 43
42
50
5 51
49
53 40
39
35
55
27
2 dynamics
uncertainty modeling [| 10
66
54
object-centric scene decomposition
46
26
18
57 3)
and space-time disentanglement 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
48
26
19
The corresponding technical improvements mainly

---

model in Iso- They commonly separate visual dynamics into content and motion vectors, or long-term and short-term states. In contrast, Iso-Dream is designed 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 33 56
31
32 2 and model-based methods [15] 38 20
23
22 29
[58 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, propose 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.
3 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.)p 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). Finally, we discuss how Iso-Dream is deployed to interact with the environment (Section 3.4
3.1 Basic Assumptions of Iso-Dream
As shown in Figure 1 when the agent receives sequence of visual observations 01:7, the underlying spatiotemporal dynamics can be defined as u1.7. Our goal is to understand the inner relationships of different dynamies by decoupling u1.r into controllable latent states 81:7 and noncontrollable latent states 21. that vary in spacetime, such that:
U1.T ~ (s, 2)1:T, S1+1 ~ p(5 St, az),
2++1 ~ p(2(+1 2t):
(1)
where ag is the action signal. To achieve long-term prediction, we isolate 8, and z to each other and model their state transitions of p(se+1 4)) and p(z++1 =x) 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 Taa | st). 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
at TTae) St, 21:47);
(2)
where we calculate the relationships between 8, and the imagined noncontrollable states over time horizon 7. It assumes that, in specific long-horizon tasks, the agent can greatly benefit from predicting the consequences of external noncontrollable forces.
3.2 Representation Learning of Controllable and Noncontrollable Dynamies
showing that modular structures are effective for disentan- glement learning, we leverage three-branch architecture to decouple in into controllable dynamics state st, noncontrollable dynamics state 21, and iim-invariant representation of the background. As
Inspired by previous approaches 36 17
3

---

Action-free
4-1
GRU,
8
Action-free brnnch
R
Dec,
End
EnGy
Actioe-conitioeed branch
dp. 9
op.
Sp-1
GRU,
Dec
S
Enn
EnCoy
9-11
Rgg
GRU,
GRU,
Inverse Cell
shared
EnCe
d/11
parameters
Static branch
0]
Actio--codiloonal branch
011x
EnGy
attention mechanism
concatenation
(a)
World model of (so-Dream
(b) Learn behavior in imagination
Figure 2: The overall architecture the world model and the behavior learning algorithm in Iso- Dream. (a) World model with three branches to explicitly disentangle controllable, noncontrollable, and static components from visual data, where the netion-conditioned branch learns controllable state transitions by modeling inverse dynamics. (b) The agent optimizes the behaviors in imaginations of the world model through a future state attention mechanism.
shown in Figure ((a) the action-condition branch models p(sr41 8t, dz). It follows the RSSM architecture from PlaNet [23] to use a recurrent neural network GRU, the deterministic hidden state he, and the stochastic state s, to form the transition model, where the GRU keeps the historical information of the controllable dynamiic. The action-free branch models p(2(11 24) with similar network structures. The transition models with separate parameters can be written as follows:
p(5e 8et, (2t)) = p( he), where h = GRU, (14 8t
8t-1, (t-1),
p (22))= = p(z k(), where Ki = GRU,
We here use & and & to denote the prior representations. We optimize the transition models with posterior representations that are derived from s, q($t h, o,) and 2, q(2, M4, o,). We learn the posteriors from the observation at current time step 0g E x by a shared encoder Enc, and subsequent branch-specific encoders Enc,, and Enc,2--
To enhance the disentanglement representation learning corresponding to the control signals, we introduce the training objective of inverse dynamics. Accordingly, design Inverse Cell of 2-layer MLP to infer the actions that lead to certain transitions of the controllable states:
Inverse dynamics: ā--1 MLP(sr-1,
(4)
where the inputs are the posterior representations in the action- condiition branch. By learning to regress the true behavior d,-1, the Inverse Cell facilitates the action-conditioned branch to isolate the representation of the controllable dynamics. To avoid the training collapse where the action- conditioned branch captures most of the useful information, while the action-free branch learns almost nothing, in the process of image reconstruction, state 2₁ to generate the controllable visual component oj H with mask M 1x H x and the noncontrollable component 0j E 3x H with M; E RxxHxW By further integrating the information extracted from the first K frames, we have
respectively use the prior state &, and the posterior
0, M ① of + M O of + - Mi Mf) @ o where 0² = Decps (Ence, dy (o1:k))) (5)
For reward modeling, we have two options with the action-free branch. In one case, the noncontrol- lable dynamics can be considered as noises that are not related to the task, and therefore is no longer useful during imagination. In other words, 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 leamn alternative reward models p(re st) or p(r: st, 2) in forms of MLPs.

---

Algorithm 1: Iso-Dream (Highlight: Our modifications to behavior learning & policy deployment)
Hyperparameters: L: Imagination horizon; Window size for future attention Initialize the replay buffer with random episodes. while not converged do
for updute step = Co
Draw data sequences ((or, ar)))) ~
// Representan
learning
Compute world model loss using Eq. (6) and update model parameters.
// Behavior learning
Roll-out the noncontrollable states im141 from through the action-free branch alone. for time step j = i... + do
10
Compute latent state Ej Attention(Ej, Bgg+r) using Eq.
Imagine action aj = [aal(e())
13
Predict the next controllable state 55-1 p(5y, a,) using the action-conditioned branch alone. end
Update the policy and value models in Eq. using estimated rewards and values.
16
end
17
// Environment interaction enver))
19
for time step 1...T do
20
Calculate the posterior representation 8, q (st. ls, 0t): 2z 4 (z 14, or). Roll-out the noncontrollable states ēt+1d0+ from through the action-free branch alone.
21
Generate a, (aa 8t, R(414r)) using future state attention in Eq. 7
23
Jt, 0++1 - env. step(a))
24
end
Add experience to the replay buffer B +- B (1
25
26 end
For a sequence (ot, dz, rT sampled from the replay buffer during training, the world model can be optimized using the following loss functions, where a, 31, and B2 are hyper-parameters:
C B
Inp hes St, M1g ze) In p(re hes St, M/y 24) In p(7) his St, higs 2x.)
image log loss
reward log
discount log loss
(6)
+ (at, a.)
B KLL[[g(s: he, ot) pp(s hh)] 32KT,[((2+ M1 o.) p(ze kt]l.
action loss
KL divergence
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 maintain the semanties of ě, Otherwise, if the action-free features only used prevent noisy distractions from affecting the training process of Iso-Dream, rather than being used for behavior learning, simply train the action-free branch with the reconstruction loss alone.
3.3 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 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- conniiiooodi branch, and 2) exploits the future future information of noncontrollable dynamics to make more forward-looking decisions.
Suppose we are making decisions at time step from the original Dreamer method is to learn an action model and value model based on the isolated controllable state & 1xddd However, we notice that by employing an attention mechanism, we can explicitly calculate its relations to sequence of future noncontrollable states 2,44 gd where
imagination period. A straightforward solution
5

---

T is the length of a sliding window from now on.
Future state attention:
e = softmax ((+++) 2eetr + &t.
(7)
In this way, & evolves to a more "visionary" representation e E Ixddd We update the action model and the value model in Dreamer [22] as follows:
I+L
Action model: ae ~ π (a ex), Value model: vg(e)) = Be(-le))
(8)
where L is the imagination time horizon. As shown in Alg. T during imagination, we first use the cetion-free transition model to obtain sequences of noncontrollable states of length L - т, denoted by {&] )+ At each time step in the imagination period, the agent draws an action aj from the visionary state E,, which is derived from Eq. Z The action-conditioned branch uses the action a, in latent imagination and predicts the next controllable state 5j+1- We follow DreamerV2 [24] to train the action model to maximize the X-return [4], and train the value model to regress the A-return
3.4 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 Lines 21-22 in Alg. [T] the action- -ree branch consecutively predicts the next T - 1 noncontrollable states 2t+1:t++ starting from the current posterior state 2, Similar Eq. future state attention network to adaptively integrate 8t, 2t and & Based on the integrated feature et, the Iso-Dream agent draws a from the action model to interact with the environment.
the process of behavior learning, we here use the learned
4 Experiments
4.1 Experimental Setup
We quantirativey and qualitatively evaluate Iso-Dream on two reinforcement learning environments, i.e., DeepMind Control Suite [45] and CARLA (Л], and two real-world datasets for action-conditionn video prediction, i..., BAIR robot pushing [13] and RoboNet [9]. The video prediction experiments can provide more intuitive visualizations of disentanglement learning.
Benchmarks.
Compared methods. For the visual control tasks, we compare our approach with five baselines, including both model-based and model-free methods, i.e.. DreamerV2 [24]. CURL [33, SVEA (25]
SAC (21), and DBC [58). For action-conditionss video prediction, we mainly compare our decoupled world model with three approaches, i.e., SVG 10), SA-COnvLSTM 34 and PhyDNet [9
4.2 DeepMind Control Suite
Implementation details. In order to verify the enhancement of Iso-Dream by disentangling differ- ent components under complex visual dynamics, evaluate Iso-Dream on environments from DMC Generalization Benchmark. Instead of training on original DeepMind Control Suite environments, agents 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 non- controllable motion of the background learning of agents. Therefore, to obtain 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, simply train the action-free branch with only reconstruction loss and discard in imagination and policy deployment. We evaluate our model with baselines in 4 tasks from four different domains. The number of environment steps limited 500k.
affect the procedure of dynamics learning and behavior
Details of the loss functions can be found Eq. (5-6) in the paper of DreamerV2 24

---

Table 1: Performance of visual control tasks in the DMC Suite. The agents trained and evaluated in environments with video _ass dynamic background. We report the mean and std of final performance over seeds and 5 trajectories. *We use different setup from that in the paper of DBC.
TASK
SVEA
CURL
DBC*
DREAMERV2
ISO-DREAM
WALKER
826 + 65
443 206
32 7
655
911 + 50
CHEETAH RUN
178 +
269 + 24
1
475 159
659 +
FINGER SPIN
562 + 22
280 50
1
755
800 1 59
HOPPER STAND
6 8
451 + 250
5
260 + 366
746 +: 312
5
1-
3-
50
5
15
t-
True frame
Prediction
Action-cond RGB
Action-cond mask
Action-frce RGB
Action-free mask
40
Figure 3: Video prediction results the DMC (left) and CARLA (right) benchmarks of Iso- Dream. For each sequence, we use the first 5 images as context frames. Iso- So-Dream successfully disentangles controllable and noncontrollable components.
Quantitative results. To evaluate the performance, we train and test the agents in environments with video backgrounds, As shown in Table Iso-Dream exceeds the performance of DreamerV2 and other baselines in all tasks, indicating that the three-branch structure effectively learn task-related visual representations and alleviate complex background interference visual data.
Qualitative results. We leverage Iso-Dream to complete video prediction tasks in video rideo_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, visualize the masks and visual decoupled components from the action-coo ciiiionnonditionee 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 predict long-term sequence and disentangle controllable and noncontrollable dynamics from images in video_ _easy_____y__o environ- ments. As shown the third and fourth row action-conditioned branch output Figure 3 the controllable representation has been successfully isolated and matches its mask. Besides, in this visualization, the -free component in this background video the motion of sea waves, which is captured by the fifth and sixth row of action-firt free branch outputs.

---

CURL
60
DBC
Dreamer/22
SAC
Episode Retum
40
SVEA
Episode Return
20
Iso Uream wio Inverse Cell
wio zolling
bo Dream wis linc-itvarien mudklng
-40
le5
3c5 Environment Steps
265
4e5
les
2c5 Environment Steps
3e5
4-5
Se5
(a) Comparison with state-of-the-arts
(b) Ablation study of Iso-Dream.
Figure Performance with 3 seeds on the CARLA driving task. (a) Comparison of existing methods, in which Iso-Dream outperforms DreamerV2 by large margin. (b) Ablation studies that can show the respective impact optimizing the inverse dynamics (orange), rolling out noncontrollable states (green), and modeling the time-invariant information with separate network branch (red).
4.3 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 [58], in order to encourage highway progression and penalise collisions, the reward is formulated oT di At & - &2 |steer], where Vego is the velocity vector of the ego-vehicle, projected onto the highway's 's unit vector ūn and multiplied by time discretization At = 0.05 to measure highway progression in meters. Impulse II I is caused by collisions, and steering penalty sleer & [-1, facilitates lane-keeping. The hyper-parameters ε and & are set to 10-4 and respectively. We use 81 B2 and 1 in Eq. 66 and 5 in Eq.
Quantitative results. As shown in Figure ((a) Sso-Dream has significant advantages compared to other basclines and outperforms DreamerV by large margin. Furthermore, we conduct ablation studies to confirm the validity of inverse dynamics and the rolling-ous strategy 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 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 we observe that rolling-out noncontrollable states in the action-free free branch can significantly improve the agent's decision-makir by about 15% in the absence of separate network branch that captures the static information.
input. Comparing the blue curve and green curve,
results. The red curve shows that the performance of Iso- -se-Dream degrades
Qualitative results. Reconstruction results of predictions in CARLA environment shown in Figure Bright column). In CARLA, we observe that the agent actions potentially affect all pixel values the observation, the camera the main (i.e., the agent) moves. Therefore, view the visual dynamics of other vehicles as combination of controllable and noncontrollable states. Accordingly, model determine which component is dominant by learning attention masks (values between 0 and 1) across the action 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 vehieles. We include more showcases with different numbers of vehicles in the supplementary materials.
and action-free branches. The "action- free
4.4 BAIR & RoboNet for Action-Conditionel Video Prediction
Implementation details. order to evaluate the effectiveness our world model a more complex environment, test the video prediction ability of the proposed structure BAIR and

---

Table 2: Video prediction results BAIR and RoboNet datasets with bouncing balls. We the first 2 frames as input to predict the next 28 frames on BAIR and the next 18 frames on RoboNet.
BAIR
ROBONET
MODEL
PSNR
SSIM 1
PSNR
SSIM f
SVG TO
18.12
0.712
19.86
0.708
SA-CONVLSTM
18.28
0.677
19.30
0.638
PHYDNET 19
18.91
0.743
20.89
0.727
ISO- ISO-DREAM
19.51
0.768
21.71
0.769
Training sel. BAIR
Teating sot: BAIR
Training set: BAIR+boll Testing sct: BAIR
PSNR
SSIM
22
0.85
21
0.80
20
0.75
19
0.70
17
0.60
SVG
SA-ComWLSTM
PhyDNet
Is0- Jao-Dremm
SVG
SA-CemvLSTM
PhyDNet
Iso- Iso-Dreum
Figure 5: The results of models trained on BAIR (blue) and BAIR bouncing balls (red), and tested on BAIR. We use the first frames as input to predict the next frames. The horizontal axis represents the different models, and the vertical axes represent test results of PSNR and SSIM.
RoboNet dataset. Moreover, we add predictable visual dynamics unrelated to the control signals the raw observations, i.e., bouncing balls of the same size and speed. In the training phase, 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 imputs for training and testing are resized to 64 x 64. Considering the simplicity and predictability of bouncing balls, in the action-free branch, we use similar structure as the DMC experiment. Moreover, we replace the GRU cell with two layers of ST-LSTM unit [51] 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.
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% BAIR and 9.3% in RoboNet. Compared with PhyDNet, which also disentangles features in two branches, Iso- Iso-Dream achieves better performance in both PSNR and SSIM. shows that our Iso-Dream has stronger ability of disentanglement learning achieve long-term prediction. Moreover, Figure 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 Fired bars). In this training setup, it performs best on the standard set without balls. Iso-Dream is built on more efficient architecture than the baseline models. It provides general framework that can be easily extended to other backbones.
Qualitative results. We visualize sequence of predicted frames on BAIR with bouncing balls in F 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. shows the fact that the action-free branch learns noncontrollable dynamics, while the action-conditionnd branch learns controllable dynamics related to input action.

---

I-15
I-18
True frame
PhyDNet
Isc- Dreom
Statie mask
Action-fres
Action free mask
<7
Accion-condizioned branch
Actio--coddtioond mask
Showcases of video prediction results on the BAIR robot pushing dataset. We display every 3 frames the prediction horizon. The generated masks show that each branch Iso-Dream captures coarse localisation of controllable representations and noncontrollable representations.
Figure 6:
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 current behavior. Iso-Dream achieves competitive results on the CARLA autonomous driving task, where other vehicles can be naturally viewed noncontrollable components, indicating that with the help of decoupled latent states, the agent can make more forward-lookin decisions by previewing possible future states in the action-free network branch. Besides, Iso-Dream shown to effectively improve the visual control task 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, requires longer training time per episode due to more intensive state transitions in behavior learning. But fortunately, from Figure ((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 the network structure, which found should be dependent on our prior knowledge of the environments.
Acknowledgements
This work was supported by the Natural Science Foundation of China (U19B2035, 62106144), Shanghai Municipal Science and Technology Major Project (2011SHZDZZX0022,, and Shanghai Sailing Program (212510202133).
References
Mohammad Babaeizadeh. Chelsea Finn, Dumitru Erhan, Roy H Campbell, and Sergey Levine. Stochastic variational video prediction. ICLR, 2018.
[] Nadine Behrmann, Jurgen Gall, and Mehdi Noroozi. Unsupervised video representation learning by bidirectional feature prediction. In WACV, pages 1670-1679, 2021.
[3 Xinzhu Bei, Yanchao Yang, and Stefano Soatto. Learning semantic-awee aware dynamics for video prediction. In CVPR, pages 902 222-912,20 2021.
10

---

[4] Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. prioritization through empowerment in visual model-based RL. In ICLR, 2022.
Information
[5] Prateep Bhattacharjee and Sukhendu Das. Temporal coherency based criteria for predicting video frames using deep multi-stage generative adversarial networks. NewrIPS, pages 4271 17780, 2017.
[6] Navaneeth Bodla, Gaurav Shrivastava, Rama Chellappa, and Abhinav Shrivastava. Hierarchical vidco prediction using relational layouts for human-object interactions. In CVPR, 2021.
[7] Lluis Castrejon, Nicolas Ballas, and Aaron Courville. Improved conditional VRNNs for video prediction. In ICCV, pages 7608-7617, 2019.
[8] Silvia Chiappa, Sébastien Racanicre, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. In ICLR, 2017.
[9] Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bernadette Bucher, Karl Schmeckpeper, Siddharth Singh, Sergey Levine, and Chelsea Finn. Robonet: Large-scale multi-robot learning. arXiv preprint arXiv:1910.12215. 2019.
[10] Emily Denton Rob Fergus. Stochastic video generation with learned prior. ICML. pages 1174-1183. PMLR, 2018.
[11] Alexey Dosovitskiy, Germán Ros, Felipe Codevilla, Antonio M. López, and Vladlen Koltun. CARLA: an open urban driving simulator. In CoRl., volume 78, pages 1-16. PMLR, 2017.
[12] Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual fore- sight: bodel-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:1812.00568, 2018.
[13] Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. CoRL, pages 456, 2017.
[14] Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through vidco prediction. In NeurlPS, pages 64-72, 2016.
[15] Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In ICRA, pages 2786-2793 IEEE, 2017.
[16] Jcan-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. Stochas- tic latent residual video prediction. In ICML, pages 3233-3246 2020.
[17] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In ICI.R, 2021.
[18] Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In JCML, pages 2424-2433, 2019.
[19] Vincent Le Guen and Nicolas Thome. Disentangling physical dynamics from unknown factors for unsupervised video prediction. In CVPR, pages 11474-114444 11484, 2020.
[20] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NewrIPS, 2018.
[21] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry 7hu, Abhishek Gupta, Pieter Abheel, et al. Soft actor-critic algorithms and applications. arXiv preprint wtx.xx:://22055553
2018.
[22] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. JCLR, 2020.
[23] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, pages 2555-2563 PMLR, 2019.
[24] Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.
[25] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In NewrIPS, 2021.

---

[26] Jun-Ting Hsieh, Bingbin Liu, De-An Huang, Li F Fei-Fei, and Juan Carlos Niebles. Learning to decompose and disentangle representations for video prediction. In NeurlPS, pages 517-526, 2018.
[27] Beibei Jin, Yu Hu, Qiankun Tang. Jingyu Niu, Zhiping Shi, Yinhe Han, and Xiaowei Li. Exploring video prediction. In
spatial-temporal multi frequenyyyyy analysis high-fidelitg and CVPR, pages 4554-4563, 2020.
[28] Minju Jung, Takazumi Matsumoto, and Jun Tani. Goal-directed behavior under variational predictive coding: Dynamic organization of visual attention and working memory. In IROS, pages 1040-1047. IEEE, 2019.
[29] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski. Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based rein- forcement learning for Atari. ICLR, 2020.
[30] Taesup Kim, Sungjin Ahn, and Yoshua Bengio. Variational temporal abstraction. NeurIPS, volume 32, pages 11570- 1150-,179,,0000 2019.
[31] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[32] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforce- ment learning with augmented data, arXiv preprint arXiv:2004. 149900, 2020.
[33] Michael Laskin. Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representations for reinforcement learning. In ICMI, volume 199 of Proceedings of Machine Learning Research, pages 5639-5650. PMLR, 2020.
[34] Zhihui Lin, Maomao Li, Zhuobin Zheng, Yangyang Cheng, and Chun Yuan. Self-attention convlstm for spatiotemporal prediction. AAAI, volume 34, pages 11531-1 1538, 2020.
[35] Wenqian Liu, Abhishek Sharma, Octavia Camps, and Mario Sznaier. Dyan: dynamical atoms-based network for video prediction. In ECCV, pages 170-185, 2018
[36] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. objject-cnttre learning with slot attention. Advances in Neural Information Processing Systems, 33: 1325 11538, 2020.
[37] Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard Lewis, and Satinder Singh. Action-conditional vidco prediction using deep networks in atari games. arXiv preprint 107.v01507.08750. 2015.
[38] Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In NeurlPS, 2017.
[39] Marc Oliu, Javier Selva, and Sergio Escalera. Folded recurrent neural networks for future video prediction. In ECCV, pages 716-731, 2018.
[40] Fitsum Reda, Guilin Liu, Kevin J Shih, Robert Kirby, Jon Barker, David Tarjan, Andrew Tao, and Bryan Catanzaro. Sde-net: Video prediction using spatially-displace convolution. In ECCV, pages 718 733, 2018.
[41] Ramanan Sekar, Olch Rybkin, Kostas Daniilidis, Picter Abheel, Danijar Hafner, and Decpak Pathak. Planning to explore via self-supervised world models. In ICML, pages 8583-8599 2020.
[42] Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung. Wai-Kin Kin Wong, and Wang-chun Woo. Convo- lutional LSTM network: A machine learning approach for precipitation nowcasting. In NeurIPS, pages 802 112-810.22 2015.
[43] Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video representa- tions using Istms. In ICML, pages 442- PMLR, 2015.
[44] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018
[45] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li. Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 1801,0000000 2018
[46] Sjoerd van Steenkiste. Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. In ICLR, 2018.
12

---

[47] Ruben Villegas, Arkanath Pathak, Harini Kannan, Dumitru Erhan, Quoe V Le, and Honglak Lee. High fidelity video prediction with large stochastic recurrent neural networks. In NeurIPS, pages 2019.
[48] Ruben Villegas, Jimei Yang. Seunghoon Hong. Xunyu Lin, and Honglak Lee. Decomposing motion and content for natural video sequence prediction. In ICLR, 2017.
[49] Ruben Villegas, Jimei Yang, Yuliang Zou, Sungryull Sohn, Xunyu Lin, and Honglak Lec. Learning generate long-term future via hierarchical prediction. /CML, pages 3560-3569, 2017.
[50] Carl Vondrick, Hamed Pirsiavash, Antonio Torralba. Generating videos with scene dynamics. In NewrIPS, pages (11-621,220 2016.
[51] Yunbo Wang, Mingsheng Long, Jiammin Wang, Zhifeng Gao, and Philip S Yu. Predrnn: Recurrent neural networks for prodictive learning using spatiotemporal Istms. In NeurlPS, pages 879-888, 2017.
[52] Yunbo Wang. Haixu Wu. Jianjin Zhang, Zhifeng Gao, Jianmin Wang, Philip S Yu, and Mingsheng Long. Predrnn: A recurrent neural network for spatiotemporal predictive learning. arXiv preprint arXiv:213.009504 2021.
[53] Nevan Wichers, Ruben Villegas, Dumitru Erhan, and Honglak Lee. Hierarchical long-term video prediction without supervision. ICML, pages 6038-6046, 2018.
[54] Bohan Wu. Suraj Nair, Roberto Martin-Martin. Fei-Fei, and Chelsea Finn. Greedy hierarchical variational autoencoders for large-scale video prediction. In CVPR, pages 2318-2211 2318-2328, 2021.
[55] Jingwei Xu, Bingbing Ni, Zefan Li, Shuo Cheng, and Xiaokang Yang. Structure preserving video prediction. In CVPR, pages 1460-1469, 2018.
[56] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-frre reinforcement learning from images. In AAAI, pages 10674-10681, 2021.
[57] Polina Zablotskaia, Edoardo A Dominici, Leonid Sigal, and Andreas M Lehrmann. Unsupervised video dccomposition using spatio-temporal iterative inference. arXiv preprint arXiv:2006. 14727, 2020.
[58] Amy Zhang, Rowan McAllister, Roherto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In ICLR, 2021.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes] See Abstract and Section 1.
(b) Did you describe the limitations of your work? [Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work
is only for academic research purpose.
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? (Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] We report the experimental results in comparison with most related works to verify the effectiveness
of our method.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or URL)? [No] Our code is
proprictary, but we will release the code once the paper is accepted.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)? [Yes] See Section 4.
13

---

(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Our model trained on Ubuntu system, 8 V100 GPUs with 32G.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See the Reference part.
(b) Did you mention the license of the assets? [No] The data used in our work are publicly
released.
(c) Did you include any new assets either in the supplemental material or as URL.? [No]
(d) Did you discuss whether and how consent obtained from people whose data you' 're using/curating? [Yes] The data and models used our work are publicly released. (c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] The data used in our work does not contain personally identifiable information or offensive content.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A] Our work does not involve human subjects. And the following items are the same.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "isbFq0NmQmQ", "reviewer": "Reviewer_CMPp", "summary": "This paper expands on the Dreamer framework to explicitly model controllable vs non-controllable dynamics, leveraging the fact that controllable dynamics can be explained away using an inverse dynamics model, while non-controllable dynamics ought to be predictable with no action change.", "questions": "- Have the authors looked into how the factorization deals with 'partially controllable' dynamics? When driving, other cars have their own (uncontrollable) intents, but are also affected by one's driving decisions. I suspect the model would still largely treat them as entirely uncontrollable due to the assumption that inverse dynamics ought to be able to resolve them if they were controllable, but I don't have a strong intuition for the actual outcome. The question may be relevant to leveraging this model to build models of other agents in the scene, besides the main one, and using those for behavior prediction.\n- I understand the used of GRUs and LSTMs follows previous papers on the topic, but since everyone seems to have moved on to Transformer-based architectures, it begs the question of whether replacing the backbones with that kind of architecture would further improve the model's efficiency?\n- I may have missed this from the discussion, but is it well-established that the background branch is necessary now that controllable and non-controllable effects are modeled separately? I didn't see an ablation for that.\n", "limitations": "Limitations section is terse and is a non-limitation in disguise.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- strong academic contribution: the approach is novel, well-motivated, and well supported by evidence.\n- the problem of modeling controllable vs. non-controllable effects is at the center of the problem of learning visual dynamics model, and this general area of research is highly relevant to a number of robotics applications.\n\nWeaknesses:\n- the limitations section is short. It would have been interesting to expand on things attempted but that didn't work for instance.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "W7IWEjSNXNH", "reviewer": "Reviewer_MP54", "summary": "This paper proposes a world model that can decompose controlled and uncontrolled factors in the environment. Concretely, their world model is composed of a stream that takes in actions as input and its state is encouraged to contain mostly controlled factors in the video by an inverse dynamics loss on the learned states; the stream that models uncontrolled factors does not take actions as input so that the environment is predicted purely on observations. The two streams as separated in the observation space by masks generated from each stream. In experiments the authors show that the proposed method outperforms other world models in a modified version of DMC and the CARLA driving environment. Finally, they also show the proposed method is able to perform action conditioned video prediction in the BAIR robot push and RobotNet datasets.", "questions": "Please see the weaknesses section.", "limitations": "There is no negative societal impact discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n+ A world model that decomposes uncontrolled and controlled factors in the input observation.\n+ Outperforms recent world models when the environment involves mostly controlled factors.\n\n\nWeaknesses:\n- Performance comparisons in similar battleground.\nThe authors clearly show an advantage of their model over the baselines in the setting the proposed model was designed for. For completeness, and to see whether the proposed model performs similar to the baselines, the authors may want to provide experiments in a setting where the baselines were tested. For example, the authors can just test their method on the non-modified DMC environment, and can provide results of this to show that proposed model is able to improve on the baselines on the settings in this paper, but doesn't make them much worse in the original settings the baselines were tested.\n\n\n- Special treatment for different environments.\nAs far as I understand, the authors choose which version of the proposed world model to use depending on whether the uncontrolled dynamics serve a purpose for the environment (Lines 143-147). Should the model learn this without the need of being told/trained depending on the environment. To me, this takes from the generality of the model and could imply that the model is simply good in specific scenarios. Can the authors comment on this? Or provide results without special treatment?\n\n\n\n\n- Video prediction setup, results and other baselines:\nThe video prediction setup in this paper is unusual. The authors deliberately put bouncing balls on top of the frames in order to show the proposed network has an advantage over other networks in the task of video prediction. I would suggest the authors directly test on the clean data and show the performance there. The model should be able to handle both the original and the modified version of the data.\n\n\nThere are also more up to date baselines missing from this comparison such as FitVid (https://arxiv.org/pdf/2106.13195.pdf)\n\n\n\nSuggestions:\n\n- Modeling relationships between current controlled state and future uncontrolled states.\nIn equation 7, the authors model relationships between the uncontrolled and controlled steps where the uncontrolled states are unrolled into the future to compute a representation that is used to make decisions on what action to take next. While this is sensible to provide the model a look into the future, the future also should consider what actions the agent may take. The uncontrolled environment also depends on what actions the controlled part of the environment take. For example, if a controlled car turns left, we should expect parts of the environment that are not seemingly controlled, to react to the left turn, and this may cause a chain reaction with all cars in the environment. Therefore, it seems more sensible to unroll both the controlled and uncontrolled states, and then equation 7 should be based on the result of this. If time permits, could the authors unroll both states together and see if this results in improvements? This is more of a suggestion out of curiosity and not a weakness of the model.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ESqvxu5aezq", "reviewer": "Reviewer_CjTg", "summary": "This paper presents Iso-dream, which uses a variant of RSSM as a world model that explicitly represents controllable and non-controllable aspects of the world state.  Because these factors are modeled independently, Iso-dream allows its policy to condition on samples of future non-controlled state rollouts from the model.  This method is experimentally evaluated on several benchmarks, some to evaluate the policy and world model together, and some to evaluate the performance of dynamics predictions.  In all benchmarks, Iso-dream outperforms examined baselines.\n\n##Update after responses\n\nI'd like to thank the authors for the clear responses, both through discussion and experimentation, to my questions and concerns.  I have raised my score.", "questions": "Some issues to help with my understanding, or for the authors to address:\n- Comparison to prior work as mentioned above\n- In the CARLA environment, agent actions have the potential to control all pixels (turning, or accelerating, would change the positions of other cars in the observations as well as the background).  Can you clarify what controllable and non-controllable aspects of the state would be in this scenario?  What training signal makes the non-controllable mask represent other cars, and is that signal brittle?\n- When rolling out the state for conditioning the policy, is there a reason this cannot be done with controllable state as well?  I wonder how a baseline that used undifferentiated world state like RSSM, but conditioned its policy on imagined future rollouts would perform (and also how well Iso-dream would work without this future rollout conditioning)  I am curious how much the controllable vs non-controllable split contributes to the improvement in results compared to the conditioning and would love to see these additional experiments in the ablations in Fig 4b.  Since I see this as the main contribution of the paper, I'd like to see this explored a bit more.\n- For the BAIR and RoboNet datasets, I'm curious whether the improved results for Iso-dream are attributable more to the action-free or action-dependent parts of the input.  Is Iso-dream doing better reconstructing the 'distractions' than other approaches but performing similarly for the 'relevant' parts of the observations? \n\nMy score depends mostly on the fact that I think this is a solid, incremental contribution that is well evaluated but could be better.  My rating would be higher if my questions either result from some misunderstanding on my part or can be addressed in a response.", "limitations": "Computational limitations are mentioned, societal impact is not discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Originality\nThe approach, to the best of my knowledge, is novel.  Some aspects of the separability of controllable and non-controllable states have been explored (InfoPower, as cited in the paper), and use of inverse dynamics losses are not new on their own.  However, the conditioning of the policy on future rollouts of the model is an interesting idea that I have not seen in this form before.  It is somewhat reminiscent of methods that incorporate planning into policy learning, and I would like to see the authors discuss the differences between this aspect of their approach and work like [Wang and Ba, Exploring Model-based Planning with Policy Networks, ICLR 2020].\n\n### Quality & Clarity\nThe paper is well-written, technically sound, and mostly clear.  The diagrams are informative and help scaffold a quick understanding of Iso-dream. I'll list a few suggestions and questions that may aid in clarifying a few points, or which may be typos\n- In line 130, the paper mentions a training collapse, but this collapse is not described\n- Figure 5 contains a label that reads 'Truth futher'\n- The image results in Figure 3 are somewhat small and difficult to make out\n\n### Significance\nThe changes proposed in Iso-dream are incremental improvements to both the world model and policy learning of a prominent model-based RL approach.  These changes are evaluated in a thorough set of experiments.  This paper is a solid contribution.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept