L EARNING T ASK I NFORMED A BSTRACTIONS
A BSTRACT
Current model-based reinforcement learning methods struggle when operating from complex visual scenes due to their inability to prioritize task-relevant features.
To mitigate this problem, we propose learning Task Informed Abstractions (TIA) that separates reward-correlated visual features from background distractions.
For learning TIA, we introduce the formalism of Task Informed MDP (TiMDP), which is realized by training two models that learn visual features via cooperative reconstruction, but one model is adversarially dissociated from the reward signal.
Empirical evaluation shows that TIA leads to signiﬁcant performance gains over state-of-the-art methods on many visual control tasks where natural and unconstrained visual distractions pose a formidable challenge.
I NTRODUCTION
[Figure]
Consider results of a simple experiment reported in Figure 1.
We train state-of-the-art model-based reinforcement learning agents (Hafner et al., 2020) operating from visual inputs to solve two versions of the
Cheetah Run task (Tassa et al., 2018): one with a simple and the other with a visually complex background (Zhang et al., 2021).
For each version we train three model variants that only differ in in the capacity of their world model, which contain 0 .
5 × (small), 1 × (medium) and 2 × (large) of the parameters of the original implementation.
Not surprisingly, the performance with the simple background is only marginally affected by model capacity, which shows that even the smallest model is sufﬁcient for learning the features necessary for solving the task.
Performance is much worse with the complex background, but it increases monotonically with the model size.
Given that the task-relevant information (i.e., joint information of the cheetah ) is the same in both, the performance improvements with model size indicates that the increase in representational capacity is used to encode the complex background.
The background conveys no information about the task and therefore interferes with the learning of task-relevant information by consuming model capacity.
Here “relevant” refers to those features that are needed to predict the optimal actions, whereas “irrelevant” refers to everything else that makes up the observation.
There are two main components to a model-based learner: (i) a forward dynamics model that predicts future events resulting from executing a sequence of actions from the current state and (ii) a reward predictor for evaluating possible future states.
The policy performance critically depends on the prediction accuracy of this model, which is intimately tied to the feature space in which the future ∗ Equal contribution.
is predicted.
Similar to the complex background version of Cheetah run , the real world contains observations that are full of irrelevant content.
Therefore to learn directly in the real-world, a compact feature space that only captures task-relevant information could make the learning problem much easier.
Without this bias towards “task-relevant features”, spurious features that confounds with the task would unnecessarily increase the data requirement or lead to training issues.
Larger models will be also be needed for complex domains, despite that the dynamics of the core control problem is simple.
A popular choice for feature learning is to reconstruct the raw observations (Kingma & Welling, 2014;
Kingma et al., 2014; Watter et al., 2015; Hafner et al., 2020).
Often these features are encouraged to be disentagled (Bengio, 2013; Higgins et al., 2016; 2017) to identify distinct factors of variation.
Since disentanglement simply re-formats the input space, the disentangled feature space would still contain irrelevant information and does not address the core issue of learning task-relevant features.
Because the model-based agent also predicts rewards from the feature space, one might expect the combination of disentanglement and reward prediction sufﬁcient to incentivize learning of task-relevant features.
However, rewards provide insufﬁcient supervision for feature learning (Yarats et al., 2019).
For instance, just knowing the center of mass of a humanoid moving forward is sufﬁcient to predict the reward, whereas it would require the full pose to come up with the optimal action.
In a nutshell reconstruction captures too much information , whereas reward-prediction captures too little .
Several works attempt to combine these two training signals (Hafner et al., 2020;
Oh et al., 2017) but still struggle to learn in complex visual scenarios.
Since the goal of the agent is to maximize the expected return, predicting the value function instead of one-step reward may aid in learning all the relevant information (Silver et al., 2017;
Oh et al., 2017).
However, because the value function is learned simultaneously with the model, it is not stationary and may not provide a stable training signal.
These challenges inspired several works to investigate feature learning methods that neither rely on reconstruction nor solely depend on rewards.
One line of work biases the learned features to only capture controllable parts of the environment using an inverse model that predicts actions from a pair of states (Agrawal et al., 2015; Jayaraman & Grauman, 2015; Agrawal et al., 2016; Pathak et al., 2017), or using metrics such as empowerment (Klyubin et al., 2005; Gregor et al., 2016).
To understand their shortcoming, consider the scenario of the arm pushing an object.
Here both the arm and the object are controllable.
While it is easy to capture the part that is directly controllable (e.g., the arm), capturing all controllable features (i.e., arm and the object) without imposing a reconstruction loss is non-trivial.
Another idea that has shown promise is the bisimulation metric (Ferns et al., 2011;
Zhang et al., 2021).
Because supervision in bisimulation comes solely from rewards, it is subject to the same insufﬁciency mentioned earlier.
Another possibility is to use contrastive learning (Chen et al., 2020;
Oord et al., 2018), but without additional constraints, these methods do not distinguish between relevant and irrelevant features.
The ongoing discussion illustrates the fundamental challenge in learning task-relevant features: some objectives (e.g., reconstruction) capture too much information, whereas others (e.g., rewards, inverse models, empowerment) capture too little.
Using a weighted loss function that combines these objectives has been empirically found not to learn task-relevant features (see Figure1).
In this work, we revisit feature learning by using reconstruction and the reward but propose to explicitly “explain away” irrelevant features by constructing a co-operative two-player game between two models.
These models, dubbed as task and distractor, learn task-relevant ( s + t ) and irrelevant features ( s − t ) of the observation ( o t ) respectively.
Similar to prior work, we force the task model to learn task-relevant features ( s + t ) by predicting the reward.
But unlike past work, we also force the distractor model to learn task-irrelevant features ( s − t ) via adversarial dissociation with the reward signal.
However, both models cooperate to reconstruct o t by maximizing p ( o t | s + t , s − t ) .
Our method models a Markov decision process (MDP) of a speciﬁc factored structure, which we call Task Informed MDP (TiMDP) (see Figure 2b).
It is worth noting that TiMDP is structurally similar to the relaxed block MDP (Zhang et al., 2020) formulation in partitioning the state-space into two separate components.
However, Zhang et al. 2020 neither proposes a practical method for segregating relevant information nor does it provide any experimental validation in the context of learning from complex visual inputs.
We evaluate our method on a custom ManyWorld environment, a suite of control tasks that speciﬁcally test the robustness of learning to visual distractions (Zhang et al., 2021) and ATARI games.
The results convincingly demonstrate that our method, which we call
T ask I nformed A bstractions (TIA), successfully learns relevant features and outperforms existing state-of-the-art methods.
P RELIMINARIES
A Markov Decision Process is represented as the tuple (cid:104)S , O , A , T, r, γ, ρ 0 (cid:105) where O is a highdimensional observation space.
A is the space of actions.
S is the state space.
ρ 0 is the initial state distribution.
r : S (cid:55)→ R is the scalar reward.
The goal of RL is to learn a policy π ∗ ( a | s ) that maximizes cumulative reward J π = arg max π E (cid:80) t γ t − 1 r t discounted by γ .
Our primary contribution is in learning a feature space for forward dynamics and it is agnostic to the speciﬁc choice of the model-based algorithm.
We choose to build upon the state-of-the-art method known as Dreamer (Hafner et al., 2020).
The main components of this model are:
Representation model:
[Equation]
Model Learning Dreamer (Hafner et al., 2019) makes future predictions in a feature space that is supervised by three signals: (a) image reconstruction (cid:2) J t O .
=
ln q ( o t | s t ) (cid:3) , (b) reward prediction (cid:2) J t R .
=
ln q ( r t | s t ) (cid:3) and (c) dynamics regularization (cid:104)
J t D .
= − β KL (cid:0) p ( s t | s t − 1 , a t − 1 , o t ) (cid:13)(cid:13)
q ( s t | s t − 1 , a t − 1 ) (cid:1)(cid:105) .
The overall objective is:
[Equation]
optimized over the agent’s experience τ .
To achieve competitive performance on ATARI, a few modiﬁcations that are incorporated in this variant DreamerV2 described in (Hafner et al., 2021).
Policy Learning Dreamer uses the learned forward dynamics model to train a policy using an actor-critic formulation described below:
[Equation]
The action model is trained to maximize cumulative rewards over a ﬁxed horizon H .
Both the action and value models are learned using imagined rollouts from the forward dynamics model.
We refer the reader to (Hafner et al., 2020) for more details.
L EARNING T ASK I NFORMED A BSTRACTIONS
Task Informed MDP
The state space of an MDP is often not directly observed in the real world, where high-dimensional observations are produced by latent states.
We present the graphical model of this common scenario in Figure 2a.
As discussed earlier the raw visual observations typically contain both task-relevant and irrelevant features.
In order to explicitly segregate these factors, we propose to factor the state space S into two components: a task-relevant component S + and a task-irrelevant component S − .
We assume that the reward is fully determined by the task-relevant component r : S
+ (cid:55)→ R , and the task-irrelevant component contains no information about the reward:
MI( r t ; s − t )
=
at each time step t .
In the most general case, s − t +1 can depend on s + t
and s + t +1 can depend on s − t .
However, in many realistic scenarios the task vs distractor distinction often follows factored forward dynamics (Guestrin et al., 2003; Pitis et al., 2020) which greatly simpliﬁes the learning model.
For this reason we further incorporate this factored structure into our formulation through the assumption: p ( s t +1 | s t , a t ) =
p ( s + t
+
| s + t , a t ) p ( s − t +1 | s − t , a t ) .
Our method involves learning two models: one model captures the task-relevant state component s + t , which we call the task model .
The other model captures the task-irrelevant state component
[Figure]
rr
[Figure]
oos
[Figure]
s
[Figure]
aa
[Figure]
rrrooos + s + s + s −
[Figure]
−
[Figure]
[Figure] (cid:30)
[Figure]
(cid:29) (cid:26)(cid:27)
[Figure]
(cid:24) (cid:23) (cid:27) (cid:23) (cid:27)(cid:26)(cid:25)
J R adv (cid:23) (cid:27)
Task Model
Distractor Model
Reconstruction (cid:31)
J Os
J Oj (a) Learning Task Informed World Models (cid:31) (cid:31) (cid:30) (cid:29) (cid:30) (cid:28) (cid:27) (cid:29) (cid:27) (cid:26) (cid:27) (cid:28) (cid:25) (cid:29) (cid:25) (cid:26) (cid:25) (cid:28) (cid:24) (cid:29) (cid:24) (cid:26) (cid:24) (cid:28) (cid:23) (cid:23) (cid:23) (cid:24) (cid:29) (cid:22) (cid:24) (cid:21) (cid:22) (cid:24) (cid:21) (cid:23) (cid:30) (cid:26) Task Model
Distractor Model (not used) (b) Policy Learning only unrolls in S + .
s − t , which we call the distractor model .
The learning objective for these two models are denoted by J P and J S (task and distractor), and expanded in Equation (4).
A visual illustration is provided in Figure 3a.
[Equation]
We will explain each component in the following section.
Reward Dissociation for the distractor model is accomplished via the adversarial objective J t Radv .
This is a minimax setup where we interleave optimizing the distractor model’s reward prediction head (for multiple iterations/training step) with training the distractor model.
While the reward prediction head is trained to minimize the reward prediction loss − ln q ( r t | s − t ) , the distractor model maximizes this objective, so as to exclude reward-correlated information from its learned features (Ganin &
Lempitsky, 2015).
The reward prediction loss is computed using ln N ( r t ; ˆ r t , 1) , where N ( · ; µ, 1) is the unit Gaussian, and ˆ r t is the predicted reward.
[Figure]
Cooperative Reconstruction By jointly reconstructing the image, the distractor model that’s biased towards capturing task-irrelevant information will enable the task model to focus on task-relevant features.
We implement joint reconstruction through the objective J t Oj .
Starting with a sequence of observation and actions { o [ <t ] , a [ <t ] } , we ﬁrst pass this sequence through the two separate recurrent state space model (RSSM, Hafner et al. 2019) to produce the states s + t and s − t , which are then used to decode two images ˆ o + t and ˆ o − t .
Given the observation o t , the joint reconstruction is achieved through a learned mixing where each model additionally decodes a 64 × 64 × 3 tensor.
These two tensors are concatenated channel-wise before being passed through a 1 × 1 convolution layer followed by sigmoid activation to obtain a 64 × 64 × 1 mask M t with the value between (0 , 1) .
The ﬁnal reconstruction is obtained through ˆ o t
= ˆ
o + t (cid:12) M t
+ ˆ o − t (cid:12) (1 − M t ) , where (cid:12) denotes element-wise product.
The reconstruction objective is computed as ln N ( o t ; ˆ o t , 1) .
Distractor-model-only Reconstruction A degenerate case exists where a distractor model that captures no information at all can still satisfy the two objectives above, letting the task model reconstruct the entire observation by itself.
To avoid such degeneracy, we add an additional image decoder to encourage the distractor model to capture as much information from the observation as possible.
This is implemented via the objective J tOs .
Policy Learning is similar to Dreamer, except that we replace the world model with the task model.
This way, the forward unroll only involves the S + subspace.
[Equation]
An illustration of the policy learning stage is in Figure 3b.
E XPERIMENTS
Our empirical evaluation aims to answer if our method outperforms existing methods when learning in environments that contain irrelevant information in the form of distractor objects or visually complex backgrounds.
For this purpose we make use of three environments that are described in Section 4.1.
Baseline wise, we compare our method against several alternatives described in Section 4.2.
4. E NVIRONMENTS
[Figure]
ManyWorld (Figure 4a)
We want a test environment where one can vary the level of distraction in a controlled manner.
For this purpose, we introduce ManyWorld, a physics domain where one can control the number of objects and their dynamics.
The task is to move a target block (in blue) to a location indicated visually by a translucent red sphere.
Other objects act as distractors.
We turn off collision between objects so that they do not physically interfere with the target object, but occlusion does occur.
The visual similarity between the objects introduces confusion, requiring additional effort to resolve when learning a world model.
Kinematic Control with Natural Video Distraction
(Figure 4b)
We consider the DeepMind Control (DMC) suite with natural video background from the
Kinetics dataset (Kay et al., 2017) used in prior work (Zhang et al., 2021), which was introduced speciﬁcally to test learning under natural visual distraction.
These control tasks involve different types of challenges, such as long planning horizon (Hopper), contact and collision (Walker), and larger state/action space (Cheetah).
The natural video backgrounds in this test suite contain a large number of factors of variation, adding additional burden to world-model learning and negatively affecting policy optimization downstream.
It is a challenging domain for both model-based and model-free algorithms to master.
The video backgrounds are from the class “driving car” and are grayscale as in Zhang et al.
Atari Learning Environments (ALE) (Figure 6) are a standard benchmark for vision-based control.
The visuals of these games naturally contain many distractor objects that are irrelevant to the game objective.
Our limited compute resources only allowed us to experiment on ﬁve games.
Each seed takes a week on a V100 Volta GPU.
We present results on games where state-of-the-art modelbased algorithms perform signiﬁcantly worse than model-free algorithms or human performance in the hope of closing this gap.
4. B ASELINE M ETHODS
We include both model-based and model-free baselines, a few proposed speciﬁcally to tackle learning in the presence of task-irrelevant distractions.
In particular we compare against Dreamer (Hafner et al., 2020) which is a state of the art model-based algorithm on DMC.
On ALE, we compare against an improved variant Dreamer(V2)
(Hafner et al., 2021).
We compare against a strong model-free method, Deep Bisimulation for Control (DBC) (Zhang et al., 2021), which uses the bisimulation metric and is developed speciﬁcally to tackle task-irrelevant distractions.
Finally, we also include
DeepMDP (Gelada et al., 2019) which learns a forward model with the sole purpose of acquiring a
[Figure]
Representation learning through an inverse model Inverse model take the observations o t and o t +1 as input and predict the intervening action a t (Agrawal et al., 2015; Jayaraman & Grauman, 2015).
To investigate if features learned by inverse model sufﬁce for learning task-relevant features, we constructed the Dreamer-Inverse model.
In this model, the learning objective becomes the following: J t Inverse =
J t inv + J tR + β J tD where J t inv .
= ln q ( a t | s t , s t +1 ) is the inverse model objective and J t D is the dynamics regularizer described in Section 2.
4. C AN TIA DISASSOCIATE TASK - IRRELEVANT INFORMATION ?
We ﬁrst evaluate our method on the ManyWorld domain, where relevant information comprises the agent (blue block) and the goal (red sphere).
Figure 4a provides a qualitative comparison of the information represented by TIA and the baseline method of Dreamer .
In many cases, Dreamer mistakes the distractor (yellow) for the agent (blue).
On the other hand, the task model of TIA isolates the agent (blue) and the goal (red), while the distractor model of TIA captures the distractors (yellow and green).
This demonstrates that indeed s + t captures task-relevant information and successfully ignores the rest.
Quantitative evaluation depicted in Figure 5 (a) and (b) demonstrates that TIA outperforms the baselines, and the performance gap increases with the number of distractors.
Next, we considered the DMC domains with natural video distractions.
We evaluate on three environments, Cheetah Run , Walker Run and Hopper Stand since they are visually different and cover a variety of learning challenges described in Section 4.1.
The image reconstruction results
Figure 4b and Figure 4c show that Dreamer performs poorly in capturing the full state of the agents and is distracted by the background.
In contrast, the task model of our TIA method accurately recovers the relevant part of the raw observation, which happens to be the agent’s body in these examples.
Quantitative performance reported in Figure 5 (c,d,e) clearly shows that TIA outperforms strong baseline methods described in Section 4.2.
Overall, these results suggest that TIA is the new state-of-the-art in learning from cluttered observations.
4. R ESULTS ON ATARI
In previous test environments, the irrelevant factors were manually injected.
Due to human intervention, it is possible these tasks were biased, which our method exploited.
To investigate performance in more natural settings, we experimented on ATARI games that innately contain signiﬁcant visual distractions.
E.g., images of robotank game contain several visual signatures that change during the game: the number of enemy tanks destroyed, the rotating radar scanner, the green sprites, and so on.
We evaluated performance on ﬁve ATARI games that are known to be challenging for model-based methods (Hafner et al., 2020) without any hyperparameter tuning (i.e., just a single value chosen based on the intuition described in Section 4.5).
The results reported in Figure 7 demonstrate that we substantially outperform the strong baseline of
DreamerV2.
Furthermore, for the games of Chopper Command , Demon Attack and Yars Revenge we either match or outperform strong model-free baselines of DQN/Rainbow trained for 200M steps, while our method is only trained for 50M steps.
These results convincingly demonstrate the superiority of our method.
Figure 6 shows the image reconstruction of TIA’s task model in Robotank .
4. H YPERPARAMETER S ELECTION
[Figure]
We plot the negative log-likelihood loss of the reward prediction − ln p ( r ) , (lower bounded by 0 . 92 ) of a mean predictor, the reward prediction module of the primary model, and the reward prediction module of the secondary model.
The features from the primary model contain sufﬁcient information for reward prediction.
The performance of the reward predictor for the secondary model follows the same trend as the mean predictor, indicating that the features learned by the secondary model are reward-independent.
The two important hyperparameters are λ Radv and λ Os .
One particular mode of failure is when the distractor model takes over the reconstruction.
It strips the task model of task-relevant information, thus preventing the policy from learning any meaningful behavior.
Our reward dissociation scheme relies on informative reward signals to work.
Yet, at the beginning of training, the reward collected by a random policy tend to be sparse and noisy, making λ Radv less effective at preventing a dominant distractor model.
This scenario suggests using a large λ Radv at the beginning of training and slowly increasing the weight λ
Os for the distractor reconstruction loss.
The other extreme is for the distractor model to collapse into degeneracy, where it fails to capture any information from the observations.
TIA degenerates into Dreamer in this case.
We can increase λ
Os so that the distractor model is encouraged to capture more.
To gain more insight into the reward dissociation process, we want to know how much information of the reward does the distractor model capture during learning.
We use errors in predicting the reward, measured with the log-likelihood (as unit Gaussian) in Figure 8.
We estimate the upper bound of this prediction error using the trailing average of the reward.
This corresponds to an uninformed reward predictor (the “Mean Predictor”) that always guesses the average.
The reward prediction loss of the distractor model trails behind this upper bound during learning (see Figure 8, in red), whereas the reward prediction loss of the main model remains small (in blue), slightly above the loss lower bound, which equals to − ln N (0; 0 , 1) = − ln 1 √ 2 π ≈ 0 .
C ONCLUSION AND D ISCUSSION
In this work, we have shown that the TiMDP formulation that explains away task-irrelevant information can successfully learn from cluttered visual inputs.
Our approach of learning T ask I nformed
A bstractions (TIA) outperforms previous state-of-the-art model-based RL methods on multiple standard benchmarks.
An issue worth mentioning is that while one set of hyperparameters worked well across ATARI games, in the DMC suite, the choice of hyperparameters λ Radv and λ Os is domain-dependent.
We discussed good practices for choosing these hyperparameters in Section 4.5.
Based on these practices, it might be possible to automatically tune the hyperparameters by considering the reconstruction and reward-prediction loss of the two models that constitute TIA.
Our goal in this work is to build agents that operate from complex visual imagery.
While we outperform previous methods, all of our evaluation is on simulated data.
We plan to test our method on real-world data in the future.
Another area of potential investigation is to characterize performance as a function of the sparsity of reward signals.
We hypothesize that a low-noise reward signal is a key factor for the robustness of TIA, and the performance might drop in scenarios with sparse rewards.
Developing methods that can overcome this “potential” challenge is another avenue for future work.
A CKNOWLEDGEMENT
This work is supported by the MIT-IBM grant on adversarial learning of multi-modal and structured data, the DARPA Machine Common Sense Program and the National Science Foundation under
Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artiﬁcial Intelligence and Fundamental Interactions, http://iaiﬁ.org/).
The authors also acknowledge the MIT SuperCloud and the
Lincoln Laboratory Supercomputing Center for providing interactive HPC resources.
A O THER R ELATED W ORKS
State-of-the-art deep reinforcement learning algorithms often jointly optimize the the discounted return together with an auxiliary representation learning objective such as image reconstruction (Watter et al., 2015; Wahlstr¨om et al., 2015) or contrastive learnpredictive coding (CPC, Sermanet et al.
Oord et al.
2018; Yan et al. 2020; Lee et al.
2020; Mazoure et al. 2020).
In model-based reinforcement learning methods, reward and value prediction (Oh et al., 2017; Racani`ere et al., 2017;
Silver et al., 2017; Tamar et al., 2016; Feinberg et al., 2018; Schrittwieser et al., 2020) are also shown to improve performance, along with learning a world model.
More recently, data augmentations are found to improve sample complexity (Laskin et al., 2020a;
Yarats et al., 2021; Raileanu et al., 2020;
Laskin et al., 2020b; Stooke et al., 2020) by taking advantage of domain knowledge of symmetry transformations in the space of data.
A recent benchmark (Stone et al., 2021) however, shows that data augmentation helps but still performs poorly in the presence of complex visual distractions.
We do ﬁnd data augmentation to be an orthogonal and complementary approach to our proposal, which focuses on the way to inform representation learning for RL, which feature is more useful, and therefore should be learned ﬁrst.
A more principled approach exploits additional structure in the real world to learn state-abstractions.
The most representative are block MDPs (Du et al., 2019)
, bisimulation (Givan et al., 2003; Dean et al., 1997) and bisimulation metrics (Ferns et al., 2004; 2011; Ferns & Precup, 2014), including recent impressive empirical gains on natural scene domains (Zhang et al., 2021; Gelada et al., 2019;
Agarwal et al., 2021).
The work presented here fall under the broad umbrella of learning state and temporal abstractions and is the most closely related to utile distinction (McCallum, 1997) which is limited to ﬁnite-state machines.
Like utile distinction , TIA is distinct from DBC and bisimulation embeddings in that we rely on the task speciﬁcation for feature separation, and we are not concerned about generalization across a set of MDPs related by dynamics.
The spiritual nearest neighbor to our work is the idea of using a primary-bias duo for model debiasing in supervised learning (Clark et al., 2019; Cadene et al., 2019;
Wang et al., 2019; He et al., 2019;
Bahng et al., 2020; Clark et al., 2020).
These methods remove dataset bias by imposing independence constraints between the primary model and a model that’s biased by design.
Our approach focuses on
RL and uses two models for removing task-irrelevant features for policy learning.
Our implementation extends recent work in model-based RL from pixels (Watter et al., 2015; Finn & Levine, 2017;
Banijamali et al., 2018;
Ha & Schmidhuber, 2018; Kaiser et al., 2020; Hafner et al., 2020; 2021;
Srinivas et al., 2018; Pertsch et al., 2020).
B P ERFORMANCE G AP W HEN L EARNING I N T HE P RESENCE OF V ISUAL D ISTRACTIONS
When complex visual distractors are present in the observations, state of the art model-based agents struggle to maintain their original sample efﬁciency and asymptotic performance.
We produce this gap in details with Figure 9, where we include three additional (six in total) domains from the
DeepMind Control suite.
We arrange the domains in the order of difﬁculty levels, from easy to hard for the model-based approach.
This gap is smaller on the two easiest domains, Walker stand and Walker walk ; and is much more pronounced in Walker run , Hopper stand , and Cheetah run .
On the most challenging task, Finger spin , where the model-based algorithm dreamer Hafner et al.
(2020) performs much worse than model-free approaches, the gap almost disappears – making the domain ill-suited for testing our proposal.
Note that the maximum episodic return on these tasks is calibrated to 1000 Tassa et al.
We additionally provide the ratio for the number of pixels that are replaced by the background video for each domain, as a proxy for how much visual distraction is introduced: Walker 63%, Hopper 77%, Cheetah: 83%, and Finger 92%.
Our intent with this paper and the proposal, learning Task Informed Abstractions , is to close this gap, such that model-based agents can retain its performance even when learning in the presence of complex visual distractions.
[Figure]
C B RIDGING THE G AP BY L EARNING T ASK I NFORMED A BSTRACTIONS
We produce the complete result on DeepMind Control Suite including the 3 additional tasks from above where we only expect marginal improvements.
We setup the experiments by matching the total number of parameters, so that each one of TIA’s two models are only half in size as the single world model from dreamer.
This comparison is disadvantage to our method, because the smaller task model runs the risk of being too small to capture the set of task-relevant features in its entirety.
Despite of this disadvantage in model capacity, task informed abstraction is able to reduce the gap on Walker run while making the gap signiﬁcantly smaller on Hopper stand and Cheetah run (see Figure 9).
On Walker stand and Walker walk the gap is small to begin with, therefore our method cannot bring much beneﬁt.
Finally, despite being a poor choice for testing our method, we include results on Finger spin for the sake of completeness.
D U NDERSTANDING F AILURE C ASES
In Section 4 .
5 we provide the principled approach to tuning the hyperapameters λ
Radv and λ Os , where we balance these two terms to avoid either one of the two models taking over the entire reconstruction.
We label these two extreme cases, where either the distractor model, or the task model takes over, as type I, and type II.
We provide detailed renderings of these failure cases below Figure 11, in comparison to a successfully learned world model.
The ﬁrst column in Figure 11a shows a successfully trained agent whose task model is able to perfectly reconstruct the walker agent.
In type I failure mode (Figure 11b)
the distractor model would take over the entire reconstruction, causing the task model to lose its grasp on the task-relevant features.
This would prevent the policy from learning any useful behavior, which makes the collected reward practically random.
With random rewards, the term J Radv is ineffective at dissociating task-relevant features, causing training to fail.
The policy performance under this scenario is usually close to that of a random policy.
When this happens, we want to increase λ Radv to dissociate the distractor model from the reward, or decrease λ
Os that weighs the distractor reconstruction term.
In the type II failure case (see Figure 11c), the task model takes over the reconstruction.
In this case the model degenerates into dreamer without separating out the task-irrelevant features, and the performance is close to a dreamer agent with a smaller model.
In this case we would increase the weight λ
Os on the distractor reconstruction to encourage it to learn more features.
By tuning these two parameters λ Radv and λ
Os , we were able to avoid these two failure modes, and stabilize training on a wide variety of domains.
A future direction would be to tune these to parameters automatically using the signals mentioned in Section 4 . 5 .
[Figure]
E M ODEL S IZES A ND A RCHITECTURE D ETAILS
For fair comparison, on DMC and ManyWorld we match the total number of parameters of our two world models combined with that of a single, large Dreamer model.
The total number of parameters is 10 million on DMC, and 1 million on ManyWorld.
We divide the two models equally in size, with an additional image reconstruction head for the distractor model.
Dreamer-Inverse has the same size for all model components as the Dreamer model except it replaces the deconvolution head with an action prediction head for learning the inverse dynamics.
We scale the model size by changing the width of each layer in the networks, without changing the overall architecture or the depth of the networks.
All other hyperparameters such as learning rates are kept the same as Hafner et al.
On the ATARI Learning Environments we compare against the state-of-the-art on this domain,
DreamerV2 Hafner et al.
(2021), which uses a world model of 20 million trainable parameters.
We made the task model the same size as the original implementation while adding a smaller distractor model which contains 12 million parameters.
A key difference between Dreamer and DreamerV2 is that the latter has an additional prediction head for the discount factor γ t besides the standard reward prediction head.
This discount factor head plays an instrumental role in allowing DreamerV2 to solve
ATARI games, therefore we additionally dissociate the distractor model from information about the discount factor, by adding an additional adversarial prediction loss.
We use the same scale for the discount factor γ as that for reward: λ γ adv = λ
Radv .
All other hyperparameters such as learning rates are kept the same as Hafner et al.
We use an input size of 32 × 32 × 3 in ManyWorld, 64 × 64 × 3 in DeepMind Control Suite, and 64 × 64 ×in ATARI games.
We use grayscale for the natural video backgrounds, the same as previous work
Zhang et al.
F H YPERPARAMETERS
For fair comparison, we did not tweak existing hyperparameters from Dreamer and used identical settings as Hafner et al.
(2020) and Hafner et al.
Our reward-dissociation scheme introduces two new hyperparameters λ Radv and λ Os .
We scale the reward dissociation loss via λ
Radv such that the term matches reconstruction losses in magnitude.
For this reason, the differences in scale in Table 1 mostly reﬂect the differences in input image sizes.
We tweaked λ
Os to stabilize training.
Detailed settings for each domain are in Table 1.
[Table]
G T RANSFER TO N OVEL D ISTRACTIONS
The task informed abstraction we introduce in this paper improves learning when distractions are present.
To adapt to out-of-distribution scenarios unseen during training, additional architectural changes that reject distracting image features on the ﬂy may be required.
To provide a baseline and intuitions for this future direction, we evaluate how well existing agents perform under this type of domain shift.
In Table 2 we take agents that are trained (1) without video background, (2) with background videos from the driving car class or (3) with white noise backgrounds [Table] and evaluate against background videos from a different class, walking the dog (labeled as transfer , see Table 2) [Table]
[Table]
The Dreamer agent trained with no background distraction fails to transfer its performance when background videos are introduced at test time, which is expected.
In the second experiment we train both dreamer and TIA with video background
[Table]
but test using videos from a different category.
Dreamer did not learn as well as TIA as indicated by its poor performance in the training environments, but both methods retain their training performance post-transfer, unaffected by the change in the background video.
As control, we also train both methods using white noise as the background.
The training and transfer performance are both identical between the two methods, and the transfer performance is worse than performance on white noise.
We additionally evaluate ManyWorld agents that are trained with (1) no distractor, (2) one distractor, or (3) two distractors, with an additional distractor (three distractors, see Table 3) [Table]
[Table]
Both results show that while TIA learns better from cluttered scenes, mechanism to reject unseen backgrounds at decision time is required to transfer successfully.
This points to the incorporation of attention as a great avenue for future work.