Isolating and Leveraging Controllable and Noncontrollable Visual Dynamics in World Models
Abstract
World models learn the consequences of actions in vision-based interactive systems.
However, in practical scenarios such as autonomous driving, there commonlyexists noncontrollable dynamics independent of the action signals, making itdifficult to learn effective world models.
To tackle this problem, we present anovel reinforcement learning approach named Iso-Dream, which improves the
Dream-to-Control framework [22] in two aspects.
First, by optimizing the inversedynamics , we encourage the world model to learn controllable and noncontrollablesources of spatiotemporal changes on isolated state transition branches.
Second,we optimize the behavior of the agent on the decoupled latent imaginations of theworld model.
Specifically, to estimate state values, we roll-out the noncontrollablestates into the future and associate them with the current controllable state .
In thisway, the isolation of dynamics sources can greatly benefit long-horizon decisionmaking of the agent, such as a self-driving car that can avoid potential risks byanticipating the movement of other vehicles.
Experiments show that Iso-Dream iseffective in decoupling the mixed dynamics and remarkably outperforms existingapproaches 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 theenvironment.
Inspired by this, many cutting-edge AI agents use self-supervised learning [37, 20, 12]or reinforcement learning [38, 22, 41] techniques to acquire knowledge from their surroundings.
Among them, world models [20] have received widespread attention in the field of robot visualcontrol and led the recent progress in model-based reinforcement learning (MBRL) [22, 41, 24, 29].
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 thenupdate 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 bymultiple sources of physical dynamics, how to learn effective world models in complex visual scenesremains an open problem.
In realistic scenarios such as autonomous driving, we can generally dividespatiotemporal 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 externalchanges.
The isolation of controllable and noncontrollable states can improve MBRL in two aspects:
• Modular representation improves the generalization of the agent to non-stationary environmentswith noises, such as the time-varying background in our modified DeepMind Control Suite.
[Equation]
[Figure]
Inverse dynamics
State
[Figure]
• More importantly, it improves long-horizon RL tasks that can greatly benefit from decisions basedon predictions of future noncontrollable dynamics.
For example, in autonomous driving, potentialrisks 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 controllableand noncontrollable state transitions.
Accordingly, it improves the original Dreamer [22] from twoperspectives: (i) a new form of world model representations and (ii) a new actor-critic algorithm toderive the behavior from the world model.
As shown in Figure 1, the foundation of decoupling theworld model is to separate the mixed latent states into an action-conditioned branch and an action-freebranch, which can individually transit different sources of visual dynamics.
The components arejointly 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 theactions that have driven the state transitions between adjacent time steps.
Another contribution of Iso-Dream is to find that disentangling physical dynamics can greatly benefitthe downstream decision-making tasks by more accurately foreseeing the inherent changes in theenvironment.
Intuitively, humans can decide how to interact with the environment at each momentbased on their anticipation of future changes in the surroundings.
To make more forward-lookingdecisions, 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.
Itenables 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 noisyvideo background; The CARLA autonomous driving environment in which other vehicles can benaturally viewed as noncontrollable components; The real-world BAIR robot dataset and the RoboNetdataset that are helpful to validate the effectiveness of the world model for disentanglement.
On allbenchmarks, Iso-Dream remarkably outperforms the existing approaches by large margins.
Related Work
Action-conditioned video prediction.
A straightforward deep learning solution to visual controlproblems is to learn action-conditioned video prediction models [37, 14, 8, 52] and then perform
Monte-Carlo importance sampling and optimization algorithms, such as the cross-entropy methods ,over available behaviors [15, 12, 28].
Hot topics in video prediction mainly includes long-term andhigh-fidelity future frames generation [43, 42, 50, 5, 51, 49, 53, 40, 39, 35, 55, 27, 2], dynamicsuncertainty modeling [1, 10, 47, 30, 7, 16, 54], object-centric scene decomposition [46, 26, 18, 57, 3],and space-time disentanglement [48, 26, 19, 6].
The corresponding technical improvements mainlyinvolve the use of more effective neural architectures, novel probabilistic modeling methods, andspecific forms of video representation.
The disentanglement methods are closely related to the worldmodel 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 worldmodel 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 fromhigh-dimensional observations.
They can be roughly grouped into two categories, that is, model-freemethods [33, 56, 31, 32, 25] and model-based methods [15, 38, 20, 23, 22, 29, 41, 58, 4].
Amongthem, the MBRL approaches explicitly model the state transitions and generally yield higher sampleefficiency than the model-free methods.
Ha and Schmidhuber [20] proposed the World Modelsthat first learn compressed latent states of the environment in a self-supervised manner, and thentrain the agent on the latent states generated by the world model.
Following the two-stage trainingprocedure, PlaNet [23] uses an action-conditioned, recurrent state-space model (RSSM) as the worldmodel, 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 thepredicted latent states in RSSM.
InfoPower [4] prioritizes functional-related information from visualobservations to obtain a more robust representation for MBRL.
Notably, Iso-Dream is very differentfrom InfoPower in two aspects.
First, we explicitly model the state transitions of controllable andnoncontrollable dynamics, so that it is possible to choose whether to take the noncontrollable statesinto behavior learning according to the prior knowledge of a specific domain.
Second, we propose anew behavior learning method that greatly benefits from the decoupled world model, so that we canpreview 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 fordecoupling 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 objectivesof inverse dynamics (Section 3.2).
For behavior learning, we present an actor-critic method that istrained on the imaginations of the decoupled world model latent states, so that the agent may considerpossible future states of noncontrollable dynamics (Section 3.3).
Finally, we discuss how Iso-Dreamis 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 a sequence of visual observations o 1: T , the underlyingspatiotemporal dynamics can be defined as u 1: T .
Our goal is to understand the inner relationships ofdifferent dynamics by decoupling u 1: T into controllable latent states s 1: T and noncontrollable latentstates z 1: T that vary in spacetime, such that:
[Equation]
where a t is the action signal.
To achieve long-term prediction, we isolate s t and z t to each other andmodel 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 rollout the noncontrollable states and consider them during behavior learning.
For tasks where thenoncontrollable components can be viewed as time-varying noise, we simply derive the action policyby a t ∼ π ( a t | s t ) .
The isolation of controllable states improves the generalization of the agent to 106 non-stationary systems.
For tasks like autonomous driving, the behaviors are derived by
[Equation]
where we calculate the relationships between s t and the imagined noncontrollable states over timehorizon τ .
It assumes that, in specific long-horizon tasks, the agent can greatly benefit from predictingthe consequences of external noncontrollable forces.
3.2 Representation Learning of Controllable and Noncontrollable Dynamics
Inspired by previous approaches [36, 17] showing that modular structures are effective for disentanglement learning, we leverage a three-branch architecture to decouple u t into controllable dynamicsstate s t , noncontrollable dynamics state z t , and time-invariant representation of the background.
Asshown in Figure 2(a), the action-conditioned branch models p ( s t +1 | s t , a t ) .
It follows the RSSMarchitecture from PlaNet [23] to use a recurrent neural network GRU s ( · ) , the deterministic hiddenstate h t , and the stochastic state s t to form the transition model, where the GRU keeps the historicalinformation of the controllable dynamics.
The action-free branch models p ( z t +1 | z t ) with similarnetwork structures. The transition models with separate parameters can be written as follows: [Equation] where h t =
GRU s ( h t − 1 , s t − 1 , a t − 1 ) , [Equation] where h ′ t = GRU z ( h ′ t − 1 , z t − 1 ) .
[Equation] Enc θ Enc θ Enc ϕ 2 Enc ϕ Enc ϕ 1 Enc ϕ
Dec φ 2 Dec φ
Dec φ 1 Dec φ
Action [Figure] free branch
Action-conditioned branch
Inverse Cell
GRU s GRU s s t −1 , a t −1 s t −1 , a t − ˜ s t ˜ s t o t o t o 1: K o 1: K s t s t minKL
Enc θ Enc θ o t −1 o t − shared parameters ̂
o t ̂ o t ˜ a t −1 ˜ a t −
GRU z GRU z ˜ z
[Figure]
˜ z [Figure] z t z [Figure]
Enc [Figure] Background [Figure]
st [Figure] st
[Figure]
st [Figure] model [Figure]
h
[Figure]
a
[Figure]
GRU
[Figure]
s [Figure] GRU [Figure]
GRU
[Figure] ˜ [Figure] ˜ z
[Figure]
˜
[Figure]
attention mechanism
[Figure]
̂
[Figure]
̂
[Figure]
a
[Figure]
̂ v
[Figure]
modeling inverse dynamics.
(b) The agent optimizes the behaviors in imaginations of the world model through a future state attention mechanism.
We here use ˜ s t and ˜ z t to denote the prior representations.
We optimize the transition models withposterior representations that are derived from s t ∼ q ( s t | h t , o t ) and z t ∼ q ( z t | h ′ t , o t ) .
We learn 121 the posteriors from the observation at current time step o t ∈ R 3 × H × W by a shared encoder Enc θ andsubsequent branch-specific encoders Enc ϕ
and Enc ϕ 2 .
To enhance the disentanglement representation learning corresponding to the control signals, weintroduce the training objective of inverse dynamics .
Accordingly, we design an Inverse Cell of a2 -layer MLP to infer the actions that lead to certain transitions of the controllable states: Inverse dynamics:
[Equation]
where the inputs are the posterior representations in the action-conditioned branch.
By learning toregress the true behavior a t − 1 , the Inverse Cell facilitates the action-conditioned branch to isolatethe 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 almostnothing, in the process of image reconstruction, we respectively use the prior state ˜ s t and the posteriorstate z t to generate the controllable visual component ˆ o st ∈ R 3 × H × W t with mask M st ∈ R 1 × H × Wand the noncontrollable component ˆ o zt ∈ R 3 × H × W t with M zt ∈ R 1 × H × W .
By further integratingthe time-invariant background extracted from the first K frames, we have
[Equation]
where ˆ o b =
Dec φ 3 ( Enc θ,ϕ [Equation] o 1: K ))) [Equation]
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 nolonger useful during imagination.
In other words, the policy and the predicted reward are only relatedto the controllable states.
In the other case, future noncontrollable states would affect how the agentmakes decisions, and we consider the action-free components during behavior learning.
For this, welearn 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 ) Tt =1 sampled from the replay buffer during training, the world model canbe optimized using the following loss functions, where α , β 1 , and β 2 are hyper-parameters:
[Equation]
The world model training approach can be partly customized for different environments.
In situationswhere noncontrollable states are indeed involved in behavior learning, minimizing the ELBO objectivecan maintain the semantics of ˜ z t .
Otherwise, if the action-free features are only used to prevent noisydistractions from affecting the training process of Iso-Dream, rather than being used for behaviorlearning, we can 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 considerthe 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 naturallyviewed as non-controllable but predictable components.
As shown in Figure 2(b), we here proposean improved actor-critic learning algorithm that 1) allows the action-free branch to foresee thefuture ahead of the action-conditioned branch, and 2) exploits the predicted future information ofnoncontrollable dynamics to make more forward-looking decisions .
Suppose we are making decisions at time step t in the imagination period.
A straightforward solutionfrom the original Dreamer method is to learn an action model and a value model based on the isolatedcontrollable state ˜ s t ∈ R 1 × d .
However, we notice that by employing an attention mechanism, we canexplicitly calculate its relations to a sequence of future noncontrollable states ˜
z t : t + τ ∈ R τ × d , whereτ is the length of a sliding window from now on.
[Equation]
In this way, ˜ s t evolves to a more “ visionary ” representation e t ∈ R 1 × d .
We update the action modeland the value model in Dreamer [22] as follows:
[Equation]
where L is the imagination time horizon.
As shown in Alg. 1, during imagination, we first use theaction-free transition model to obtain sequences of noncontrollable states of length L + τ , denoted
[Table]
At each the period, the agent an action a j thevisionary e j which is derived from Eq action conditioned the latent predicts the next j
DreamerV2 [ maximize the
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 generatethe policy at each time step t .
However, for the situation where noncontrollable dynamics shouldbe closely related to the behavior of the agent, as shown in Lines 21-22 in Alg. 1, the action-freebranch consecutively predicts the next τ − 1 noncontrollable states ˜
z t +1:
t + τ starting from the current 174 posterior state z t .
Similar to Eq.
(7) in the process of behavior learning, we here use the learnedfuture state attention network to adaptively integrate s t , z t and ˜ z t +1: t + τ .
Based on the integratedfeature e t , the Iso-Dream agent draws a t from the action model to interact with the environment.
Experiments
4.1 Experimental Setup
Benchmarks.
We quantitatively and qualitatively evaluate Iso-Dream on two reinforcement learningenvironments, i.e. , DeepMind Control Suite [45] and CARLA [11], and two real-world datasets foraction-conditioned video prediction, i.e. , BAIR robot pushing [13] and RoboNet [9].
The videoprediction experiments can provide more intuitive visualizations of disentanglement learning.
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-conditioned video prediction, we mainly compare our decoupledworld model with three approaches, i.e. , SVG [10], SA-ConvLSTM [34] and PhyDNet [19].
4.2 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 behaviorlearning of agents.
Therefore, to obtain a better decision policy and avoid the disruption from noisybackgrounds, the agent may decouple noncontrollable representation ( i.e. , dynamic background) andcontrollable representation in spacetime, and only use controllable representation for control.
To thisend, we simply train the action-free branch with only reconstruction loss and discard it in imaginationand policy deployment.
We evaluate our model with baselines in 4 tasks from four different domains.
The number of environment steps is limited to 500 k.
Quantitative results.
To evaluate the performance, we train and test the agents in environmentswith video backgrounds.
As shown in Table 1, Iso-Dream exceeds the performance of DreamerV2 andother baselines in all tasks, indicating that the three-branch structure can effectively learn task-relatedvisual representations and alleviate complex background interference in visual data.
t= t= t= t= T r u e fr a m e P r e d i c ti on A c ti on - c ond
R G B
A c ti o n -fr ee
R G B
A c ti on -fr ee m a s k
A c ti on - c ond m a s k t= t= t= t=
Qualitative results.
We leverage Iso-Dream to complete video prediction tasks in video_easyenvironments.
The sequence of frames and actions are randomly collected during test episodes.
Thefirst 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 theaction-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 sequenceand 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, thecontrollable representation has been successfully isolated and matches its mask.
Besides, in thisvisualization, the action-free component in this background video is the motion of sea waves, whichis captured by the fifth and sixth row of action-free branch outputs.
4.3 CARLA Autonomous Driving Environment
Implementation details.
In the autonomous driving task, We use a camera with 60 degree view onthe 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 formulatedas:
[Equation]
1 · I − ξ 2 · | steer | , where v ego is the velocity vector of the ego-vehicle, 220 projected onto the highway’s unit vector ˆ u h , and multiplied by time discretization ∆ t = 0 .
05 tomeasure highway progression in meters.
Impulse I ∈ R + is caused by collisions, and a steeringpenalty steer ∈ [ − 1 , 1] facilitates lane-keeping.
The hyper-parameters ξ 1 and ξ 2 are set to 10 − 4 and1 , respectively.
We use β 1 = 1 , β 2 = 1 and α = 10 − 3 in Eq.
(6) and τ = 5 in Eq.
Quantitative results.
As shown in Figure 4(a), Iso-Dream has significant advantages compared toother baselines and outperforms DreamerV2 by a large margin.
Furthermore, we conduct ablationstudies to confirm the validity of inverse dynamics and the rolling-out strategy of noncontrollablestates.
Figure 4(b) shows that the performance drops when Inverse Cell is removed, indicating theimportance of modeling inverse dynamics to isolate controllable and noncontrollable componentsfrom 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 currentcontrollable state and the noncontrollable state as input.
Comparing the blue curve and green curve,
[Figure] [Figure] [Table] we observe that rolling-out noncontrollable states in the action-free branch can significantly improvethe agent’s decision-making results.
The red curve shows that the performance of Iso-Dream decreasesby 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 pixelvalues in the observation, as the camera on the main car ( i.e. , the agent) moves. Therefore, we canview the dynamics of other vehicles as a combination of controllable and non-controllable 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-freemasks’ present hot spots around other vehicles, while the attention values in corresponding areason the ’action-cond masks’ are still greater than zero.
The agent can avoid collisions by rolling-outnoncontrollable components to preview possible future states of other vehicles.
We include moreshowcases with different numbers of vehicles in the supplementary materials.
4.4 BAIR & RoboNet for Action-Conditioned Video Prediction
Implementation details.
In order to evaluate the effectiveness of our world model in a morecomplex 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 signalsto the raw observations, i.e. , bouncing balls of the same size and speed.
In the training phase, wetrain the model to predict 10 frames into the future from 2 observations.
For testing, we use the first2 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 simplicityand 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 [51] in bothbranches.
The optimization objective consists of image reconstruction loss and action reconstructionloss of Inverse Cell.
SSIM and PSNR are adopted as evaluation metrics.
Quantitative results.
Table 2 gives the quantitative results on BAIR and RoboNet datasets withbouncing balls in the training and testing phase.
Compared with other models, Iso-Dream shows the
[Figure]
Background mask t= t= t= t= t= PhyDNet t= t=
Iso-Dream True frame t= t= t= t= t= t= t=
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 futuredynamics for long-term prediction.
It shows the fact that the action-free branch learns noncontrollabledynamics, 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 thedifficulty of vision-based prediction and control in the presence of complex visual dynamics.
Ourapproach has two novel contributions to world model representation learning and corresponding
MBRL algorithms.
First, it learns to decouple controllable and noncontrollable latent state transitionsvia modular network structures and inverse dynamics.
Further, it makes long-horizon decisions byrolling-out the noncontrollable dynamics into the future and learning their influences on currentbehavior.
Iso-Dream achieves competitive results on the CARLA autonomous driving task, whereother vehicles can be naturally viewed as noncontrollable components, indicating that with the help ofdecoupled latent states, the agent can make more forward-looking decisions by previewing possiblefuture states in the action-free network branch.
Besides, Iso-Dream was shown to effectively improvethe visual control task in a modified DeepMind Control Suite, as well as the visual prediction task onthe BAIR robot pushing dataset and the RoboNet dataset.
The limitations of Iso-Dream are the computational efficiency and special treatment for differentenvironments.
Compared with DreamerV2, it requires longer training time per episode due to moreintensive state transitions in behavior learning.
But fortunately, from Figure 4(a), Iso-Dream is moresample-efficient than existing MBRL methods.
In preliminary experiments, we attempted to usethe same model architecture for all test benchmarks.
However, the experimental results show thatdifferent benchmarks have specific requirements on the network structure, which we found should bedependent on our prior knowledge of the environments.
For video prediction, we adopt PredRNN-V2as the backbone network.
We found that using the RSSM from Dreamer-V2 would decrease thequality of generated future images.
It is because the compact states and rather simple state transitionmodel (in the form of a single-layer GRU) may inevitably result in information loss of video details.
For the control tasks, in DMC, we only disentangle the non-controllable states from the mixed visualdata, but do not roll-out them into the future and use them for behavior learning, as we do in the
CARLA environment.
We believe that whether the non-controllable dynamics should be consideredas noise or as the condition of behavior learning is related to the specific settings of the environment.
Especially in a virtual environment like DMC, it is unreasonable to force the model to learn theinfluence of background noise on the control strategy.