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: In this paper, the causal confusion problem in imitation learning (IL) is considered. To tackle the confusion that resulted from not knowing the causal relationship, a two-stage approach is proposed. At the first stage, a VQ-VAE model is trained to generate semantic codes for different objects in the state space. At the second stage, a dropout-based approach is adopted to decouple the semantic confusion among the objects. The contributions of the paper are: >- Propose a practical approach to solve the causal confusion problem for IL; >- Conduct extensive experiments under Atari to verify the performance of the proposed approach.