Object-Aware Regularization for Addressing Causal Confusion in Imitation Learning

1 Introduction . Imitation learning ( IL ) holds the promise of learning skills or behaviors directly from expert demonstrations , effectively reducing the need for costly and dangerous environment interaction [ 21 , 45 ] . Its simplest and effective form is behavioral cloning ( BC ) , which learns a policy by solving a supervised learning problem over state-action pairs from expert demonstrations . While being simple , BC has been successful in a wide range of tasks [ 4 , 7 , 31 , 33 ] with careful designs . However , it has been recently evidenced that BC often suffers from the causal confusion problem , where the policy relies on nuisance variables strongly correlated with expert actions , instead of the true causes [ 11 , 12 , 54 ] . For example , when we train a BC policy on the Atari Pong environment ( see Figure 1a ) , we observe that a policy relies on nuisance variables in images ( i.e. , scores ) for predicting expert actions , instead of learning the underlying fundamental rule of the environment that experts would have used for making decisions . In particular , Table 1c shows that the policy trained using images with scores struggles to generalize to images with scores masked out ( see Figure 1b ) . However , the policy trained with masked images could generalize to original images with scores , which shows that it successfully learned the rule of the environment . This implies that learning the policy that can identify the true cause of expert actions is important for stable performance at deployment time , where nuisance correlates usually do not hold as in expert demonstrations . ∗Equal contribution , in alphabetical order . { jongjin.park , younggyo.seo } @ kaist.ac.kr †This work was done while the author was an intern at Microsoft Research Asia 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In order to address this causal confusion problem , one can consider causal discovery approaches to deduce the cause-effect relationships from observational data [ 26 , 48 ] . However , it is difficult to apply these approaches to domains with high-dimensional inputs , as ( i ) causal discovery from observational data is impossible in general without certain conditions3 [ 38 ] , and ( ii ) these domains usually do not satisfy the assumption that inputs are structured into random variables connected by a causal graph , e.g. , objects in images [ 29 , 46 ] . To address these limitations , de Haan et al . [ 12 ] recently proposed a method that learns a policy on top of disentangled representations from a β-VAE encoder [ 19 ] with random masking , and infers an optimal causal mask during the environment interaction by querying interactive experts [ 43 ] or environment returns . However , given that environment interaction could be dangerous and incur additional costs , we argue that it is important to develop a method for learning the policy robust to causal confusion problem without such a costly environment interaction . In this paper , we present OREO : Object-aware REgularizatiOn , a new regularization technique that addresses the causal confusion problem in imitation learning without environment interaction . The key idea of our method is to regularize a policy to attend uniformly to all semantic objects in images , in order to prevent the policy from exploiting nuisance correlates for predicting expert actions . To this end , we propose to extract semantic objects from raw images by utilizing vector-quantized variational autoencoder ( VQ-VAE ) [ 35 ] . In our experiments , we discover that the units of a feature map corresponding to the objects with similar semantics , e.g. , backgrounds , scores , and characters , are mapped into the same or similar discrete codes ( see Figure 3 ) . Based upon this observation , we propose to regularize the policy by randomly dropping units that share the same discrete code together throughout training . Namely , our method randomly masks out semantically similar objects , which allows object-aware regularization of the policy . We highlight the main contributions of this paper below : • We present OREO , a simple and effective regularization method for addressing the causal confusion problem , and support the effectiveness of OREO with extensive experiments . • We show that OREO significantly improves the performance of behavioral cloning on confounded Atari environments [ 5 , 12 ] , outperforming various other regularization methods [ 13 , 15 , 55 , 50 ] and causality-based methods [ 12 , 47 ] . • We show that OREO even outperforms inverse reinforcement learning methods trained with a considerable amount of environment interaction [ 8 , 20 ] . 2 Related work . Imitation learning . Imitation learning ( IL ) aims to solve complex tasks where learning a policy from scratch is difficult or even impossible , by learning useful skills or behaviors from expert demonstrations [ 2 , 18 , 25 , 37 , 39 , 40 , 56 ] . There are two main approaches for IL : inverse reinforcement learning ( IRL ) methods that find a cost function under which the expert is uniquely optimal [ 8 , 20 , 34 , 44 , 59 ] , and behavioral cloning methods that formulate the IL problem as a supervised learning problem that predicts expert actions from states [ 3 , 4 , 7 , 31 , 33 , 40 ] . Our work employs 3de Haan et al . [ 12 ] showed that causal discovery methods that depend on faithfulness condition [ 38 ] are not applicable to imitation learning setup , as the condition does not hold in environments with nuisance correlates . behavioral cloning as it exhibits the benefit of avoiding costly and dangerous environment interaction , which is crucial for applying imitation learning to real-world scenarios . Distributional shift and the causal confusion problem . Despite its simplicity , BC is known to suffer from the distributional shift , where the state distribution induced by a policy gets different from the training distribution on which the policy was trained . Several approaches have been proposed for learning the policy robust to distributional shift , including interactive IL methods that query experts [ 42 , 43 , 51 ] , and regularization techniques [ 4 , 7 ] . Recently , it has been evidenced that distributional shift leads to the causal confusion problem [ 4 , 12 , 54 ] where a policy exploits the nuisance correlates in states for predicting expert actions . To address this problem , Bansal et al . [ 4 ] proposed to randomly drop previous samples from a sequence of samples , and Wen et al . [ 54 ] proposed an adversarial training scheme of removing information related to previous actions . The work closest to ours is de Haan et al . [ 12 ] , which learns a policy with randomly masked disentangled representations and infers the best mask through during environment interaction . Our approach differs in that we regularize the policy to be robust to the causal confusion problem , without any environment interaction . Causal discovery from observational data . Causal discovery aims to discover causal relations among variables by utilizing observational data [ 38 ] . Most prior approaches assume that inputs are structured as disentangled variables [ 6 , 16 , 26 , 36 , 48 , 47 ] , which often does not hold in domains with high-dimensional inputs , i.e. , images . While Lopez-Paz et al . [ 29 ] demonstrated the possibility of observational causal discovery from high-dimensional images , combining causal models and representation learning in such domains still remains an open problem [ 46 ] . Hence , we instead explore the approach of regularizing a policy that operates on high-dimensional states . 3 Method . 3.1 Preliminaries . We consider the standard imitation learning ( IL ) framework where an agent learns to solve a target task from expert demonstrations . Specifically , IL is typically defined in the context of a discretetime Markov decision process ( MDP ) [ 52 ] without an explicitly-defined reward function , which is defined as a tuple ( S , A , p , γ , ρ0 ) . Here , S is the state space , A is the action space , p ( s′|s , a ) is the transition dynamics , ρ0 is the initial state distribution , and γ ∈ [ 0 , 1 ) is the discount factor . The goal of IL is to learn a policy π , mapping from states to actions , using a set of expert demonstrations D = { ( si , ai ) } Ni=1 . In our problem setup , an agent can not interact with the environment , hence it should learn the policy by using only expert demonstrations . Behavioral cloning . Behavioral cloning ( BC ) reduces an imitation learning problem to the supervised learning problem of training a policy that imitates expert actions . Specifically , we introduce a policy π that maps a state st to an action at , and a convolutional encoder f that maps a state st to a low-dimensional feature map . Then π and f are learned by minimizing the negative log-likelihood of expert actions from demonstrations as follows : LBC ( st , at ) = − log π ( at|f ( st ) ) , ( 1 ) where π is modeled as a multinomial distribution over actions to handle discrete action spaces . Vector quantized variational autoencoder . The VQ-VAE [ 35 ] model consists of an encoder g that compresses images into discrete latent representations , and a decoder d that reconstructs images from these discrete representations . Both encoder and decoder share a codebook C of prototype vectors which are also learned throughout training . Formally , given a state st , the encoder g encodes st into a feature map ht ∈ RL×D that consists of a series of L latent vectors ht , i ∈ RD , i ∈ { 1 , 2 , ... , L } . Then ht = g ( st ) is quantized to discrete representations e ∈ RL×D based on the distance of latent vectors ht , i to the prototype vectors in the codebook C = { ek } Kk=1 as follows : et = ( eq ( t,1 ) , eq ( t,2 ) , · · · , eq ( t , L ) ) , where q ( t , i ) = argmin k∈ [ K ] ‖ht , i − ek‖2 , ( 2 ) where [ K ] is the set { 1 , · · · , K } . Then the decoder d learns to reconstruct st from discrete representations et . The VQ-VAE is trained by minimizing the following objective : LVQVAE ( st ) = ‖st − d ( et ) ‖22︸ ︷︷ ︸ Lrecon + ‖sg [ ht ] − et‖22︸ ︷︷ ︸ Lcodebook +β · ‖sg [ et ] − ht‖22︸ ︷︷ ︸ Lcommit , ( 3 ) where the operator sg refers to a stop-gradient operator , Lrecon is a reconstruction loss for learning representations useful for reconstructing images , Lcodebook is a codebook loss to bring codebook representations closer to corresponding encoder outputs h , and Lcommit is a commitment loss weighted by β to prevent encoder outputs from fluctuating frequently between different representations . 3.2 OREO : Object-aware regularization for behavioral cloning . In this section , we present OREO : Object-aware REgularizatiOn that regularizes a policy in an objectaware manner to address the causal confusion problem . Our main idea is to encourage the policy to uniformly attend to all semantic objects in images , in order to prevent the policy from exploiting nuisance variables strongly correlated with expert actions . To this end , we introduce a two-stage approach : we first train a VQ-VAE model that encodes images into discrete codes , then learn the policy with our regularization scheme of randomly dropping units that share the same discrete codes ( see Figure 2 and Algorithm 1 for the overview and pseudocode of OREO , respectively ) . Algorithm 1 Object-aware regularization ( OREO ) Initialize parameters of encoder g , decoder d , codebook C , policy π. while not converged do // VQ-VAE TRAINING Sample { si } Bi=1 ∼ D. Update parameters of g , d , C by minimizing ∑B i=1 LVQVAE ( si ) end while Initialize encoder f with parameters of g. while not converged do // UPDATE POLICY VIA BEHAVIORAL CLONING Sample { si , ai } Bi=1 ∼ D Get random masks { Mi } Bi=1 in ( 4 ) Update parameters of f , π by minimizing ∑B i=1 LOREO ( si , ai , Mi ) in ( 5 ) end while Extracting semantic objects . To regularize a policy in an object-aware manner , we propose to utilize discrete representations from a VQ-VAE model trained by optimizing the objective in ( 3 ) with images from expert demonstrations D. Our motivation comes from the observation that the units of a feature map corresponding to similar objects are mapped into similar discrete codes ( see Figure 3 ) . Then , in order to extract semantic objects from images and utilize them for regularizing the policy , we propose to randomly drop the units of a feature map that share the same discrete code together throughout training . Formally , for each state st , we sample K binary random variables mk ∈ { 0 , 1 } , k = 1 , 2 , ... , K from a Bernoulli distribution with probability 1− p , where p is the drop probability . Then , we construct a mask Mt by utilizing the discrete representations et in ( 2 ) as follows : Mt = ( mq ( t,1 ) , mq ( t,2 ) , · · · , mq ( t , L ) ) , where q ( t , i ) = argmin k∈ [ K ] ‖ht , i − ek‖2 . ( 4 ) By considering units of a feature map with the same discrete code , we remark that our method can effectively extract semantic objects from high-dimensional images . Behavioral cloning with OREO . Now we propose to utilize our object-aware masking scheme for the regularization of a policy . To this end , we first initialize a convolutional encoder f with the parameters of a VQ-VAE encoder g. We empirically find that employing f as our backbone encoder for π instead of a fixed encoder g is more effective , as it allows an encoder to learn useful information for predicting actions . Then , we train the policy π by minimizing the following objective : LOREO ( st , at , Mt ) = − log π ( at|f ( st ) Mt ) , ( 5 ) where denotes elementwise product , and the mask Mt is shared across all channels in a feature map f ( st ) . Here , our intuition is that our object-aware regularization scheme should be useful for enforcing the policy not to exploit specific objects strongly correlated with expert actions , as the policy should utilize all semantic objects throughout training . Additionally , following Srivastava et al . [ 50 ] , we scale the masked features by a factor of 1/ ( 1− p ) during the training to ensure the scale of the expected output with masked features to match the scale of outputs at test time .

Summary: The paper considers a recently identified problem in imitation learning, that the trained policy learns relationships to nuisance variables, which do not have a casual relationship to the main goal of the action.   For instance, when learning pong, the learned policy becomes dependent on the score display and the learned policy is difficult to transfer to an environment where the score display does not exist. On the other hand, a policy learned on an environment where the score display was masked out transfers better to an environment where it is present.   The authors are proposing a technique to regularize the imitation learning process by masking out certain objects from the environment. They achieve this by training a VQ-VAE, and assuming that units of a feature map corresponding to the same objects are mapped to similar discrete codes. The authors then randomly drop the units of a feature map that share the same discrete code during training.   The authors test the proposed approach on a number of games from the Atari environment using the "confounded input model", where the action from the previous time step is encoded into the current input. Previous work had shown that such a confounded input, although it contains more information than the original form of the input, is in practice harder to learn.