Causal Curiosity: RL Agents Discovering Self-supervised Experiments for Causal Representation Learning

1 INTRODUCTION . Discovering causation in environments an agent might encounter remains an open and challenging problem for causal reinforcement learning ( Schölkopf ( 2015 ) , Bengio et al . ( 2013 ) , Schölkopf ( 2019 ) ) . Most approaches take the form of BAMDPs ( Bayes Adaptive Markov Decision Processes ) ( Zintgraf et al . ( 2019 ) ) or Hi-Param MDP ( Hidden Parameter MDPs ) ( Doshi-Velez & Konidaris ( 2016 ) ; Yao et al . ( 2018 ) ; Killian et al . ( 2017 ) ; Perez et al . ( 2020 ) ) which condition the transition p ( st+1|st , at ; H ) and/or reward function R ( rt+1|st , at , st+1 ; H ) of each environment on hidden parameters ( also referred to as causal factors in some of the above studies ) . Let s ∈ S , a ∈ A , r ∈ R , H ∈ H where S , A , R , and H are the set of states , actions , rewards and feasible hidden parameters . In the physical world and in the case of mechanical systems , examples of the parameter hj ∈ H include gravity , coefficients of friction , masses and sizes of objects . Typically , H is treated as a latent variable for which an embedding is learned during training , using variational methods ( Kingma et al . ( 2014 ) ; Ilse et al . ( 2019 ) ) . Let s0 : T be the entire state trajectory of length T . Similarly , a0 : T is the sequence of actions applied during that trajectory by the agent that results in s0 : T . In an environment parameterized by these causal factors , these latent variable approaches define a probability distribution over the entire sequence of ( rewards , states , actions ) conditioned on a latent z as p ( r0 : T , s0 : T , a0 : T−1 ; z ) that factorizes as T−1∏ i=1 p ( rt+1|st , at , st+1 , z ) p ( st+1|st , at , z ) p ( at|st , z ) ( 1 ) due to the Markov assumption . At test time , the agent infers the causal factor associated with its environment by observing the trajectories produced by its initial actions that can be issued by any policy such as model-based reinforcement learning . In practice , however , discovering causal factors in a physical environment is prone to various challenges that are caused by the disjointed nature of the influence of these factors on the produced trajectories . More specifically , at each time step , the transition function is affected by a subset of global causal factors . This subset is implicitly defined on the basis of the current state and the action taken . For example , if a body in an environment loses contact with the ground , the coefficient of friction between the body and the ground no longer affects the outcome of any action that is taken . Likewise , the outcome of an upward force applied by the agent to a body on the ground is unaffected by the friction coefficient . We can therefore take advantage of this natural discontinuity to discern causal factors . Without knowledge of how independent causal mechanisms affect the outcome of a particular action in a given state in an environment , it becomes impossible for the agent to conclude where the variation it encountered came from . Unsurprisingly , Hi-Param and BAMDP approaches fail to learn a disentangled embedding for the causal factors , making their behaviors uninterpretable ( Perez et al . ( 2020 ) ) . For example , if , in an environment , a body remains stationary under a particular force , the Hi-Param or BAMDP agent may apply a higher force to achieve its goal of perhaps moving the body , but will be unable to conclude whether the `` un-movability '' was caused by high friction or high mass of the body . Additionally , these approaches require human-supervised reward engineering , making it difficult to apply them outside of the simulated environments they are tested in . Our goal is , instead of focusing on maximizing reward for some particular task , to allow agents to discover causal processes through exploratory interaction . During training , our agents discover self-supervised experimental behaviors which they apply to a set of training environments . These behaviors allow them to learn about the various causal mechanisms that govern the transitions in each environment . During inference in a novel environment , they perform these discovered behaviors sequentially and use the outcome of each behavior to infer the embedding for a single causal factor ( Figure 1 ) . The main challenge while learning a disentangled representation for the causal factors of the world is that several causal factors may affect the outcome of behaviors in each environment . For example , when pushing a body on the ground , the outcome , i.e. , whether the body moves , or how far the body is pushed , depends on several factors , e.g. , mass , shape and size , frictional coefficients , etc . However , if , instead of pushing on the ground , the agent executes a perfect grasp-and-lift behavior , only mass will affect whether the body is lifted off the ground or not . Thus , it is clear that not all experimental behaviors are created equal and that the outcomes of some behaviors are caused by fewer causal factors than others . Our agents learn these behaviors without supervision using causal curiosity , an intrinsic reward . The outcome of a single such experimental behavior is then used to infer a binary quantized embedding describing the single isolated causal factor . Even though causal factors of variation in a physical world are easily identifiable to humans , a concrete definition is required to back up our proposed method . We conjecture that the causality of a factor of variation depends on the available actions to the agent . If the set of actions that an agent can take is very limited , there is no way for it to discern a diverse set of causal factors in the environment . Definition 1 ( Causal factors ) . Consider the POMDP ( O , S , A , p , r ) with observation space O , state space S , action spaceA , the transition function p , and the reward function r. Let o0 : T ∈ OT denotes a trajectory of observations and T be the length of such trajectories . Let d ( · , · ) : OT ×OT → R+ be a distance function defined on the space of trajectories of length T . The set H = { h1 , h2 , . . . , hk } is called a set of −causal factors if for every hj ∈ H , there exists a unique sequence of actions a0 : T that clusters the state trajectories into two sets S and S′ such that min { d ( o0 : T , o′0 : T ) : o0 : T ∈ O , o′0 : T ∈ O′ } > ( 2 ) and that hj is the cause of the trajectory of states obtained i.e. , p ( o0 : T |do ( hj = k ) , a0 : T ) 6= p ( o0 : T |do ( hj = k′ ) , a0 : T ) ∀k 6= k′ ( 3 ) Intuitively , a factor of variation affecting a set of environments is called causal if there exists a sequence of actions available to the agent where the resultant trajectories are clustered into two or more sets ( for simplicity here we assume binary clusters ) . This is analogous to the human ability to conclude whether objects are heavy or light , big or small . For a gentle introduction to the intuition about this definition , we refer the reader to Appendix D. According to Def . 1 , a causal factor is a parameter in the environment whose value , when intervened on ( i.e . varied ) over a set of values , results in trajectories of states that are divisible into disjoint clusters under a particular sequence of actions . These clusters represent the quantized values of the causal factor . For example , mass , which is a causal factor of a body , under an action sequence of a grasping and lifting motion , results in 2 clusters , liftable ( low mass ) and not liftable ( high mass ) . However , such an action sequence is not known in advance . Therefore , discovering a causal factor in the environment boils down to finding a sequence of actions that makes the effect of that factor prominent by producing clustered trajectories for different values of that environmental factor . Using the above , we propose an intrinsic reward , which allows our agents to discover experimental behaviors which are semantically meaningful and can be used to re-train for downstream tasks , resulting in high sample efficiency . Our work , therefore , forms an important link between structured representation learning and skill discovery , two largely disjoint fields in RL , which stand to benefit from each other . The contributions of the work are as follows : • We equip agents with the ability to perform experiments and behave meaningfully in a set of environments in an unsupervised manner . These behaviors can expose or obfuscate specific independent causal mechanisms that occur in the world of the agent , allowing the agent to learn about each in the absence of the others , an important human behavioral trait . • We introduce an intrinsic reward , causal curiosity , which allows our agents to discover these behaviors without human-engineered rewards . The outcomes of the experiments are used to learn a disentangled quantized binary representation for the causal factors of the environment , analogous to the human ability to conclude whether objects are light/heavy , big/small etc . • Through extensive experiments , we conclude that knowledge of the causal factors aids sample efficiency in two ways - first , that the knowledge of the causal factors aids transfer learning across multiple environments , and , second , that the experimental behaviors acquired can be repurposed for downstream tasks . 2 METHOD . Consider a set of N environments E with e ( i ) ∈ E where e ( i ) denotes the ith environment . The letter H is overloaded . While H is a set of global causal factors ( as defined in Def . 1 ) such that hj ∈ H , each causal factor hj is itself a random variable which assumes a particular value for every instantiation of an environment . Thus every environment e ( i ) is represented by a set of causal factors { h ( i ) j ∀j } . For each environment e ( i ) , ( z ( i ) ( 0 ) , z ( i ) ( 1 ) ... z ( i ) ( K−1 ) ) represents the disentangled embedding vector , such that z ( i ) ( j ) encodes h ( i ) j . Algorithm 1 Training Scheme 1 : Initialize j = 0 2 : Initialize training environment set Envs 3 : for iteration m to M do . Experiment Planner Training Loop 4 : Sample experimental behavior a0 : T ∼ CEM ( · ) 5 : for ith env in Envs do 6 : Apply a0 : T to env 7 : Collect S ( i ) = O ( i ) 0 : T 8 : Reset env 9 : Calculate −L ( S|M ) given that M is bimodal clustering model . Calculate Curiosity 10 : Update CEM ( · ) distribution with highest reward trajectories 11 : Use learnt qM ( z|S ) for cluster assignment of each env in Envs i.e . z ( i ) j = qM ( z|S ( i ) ) 12 : Update j = j + 1 13 : Repeat from step 2 , first setting Envs = { e ( i ) : z ( i ) j−1 = 0 } and then , setting Envs = { e ( i ) : z ( i ) j−1 = 1 }

Summary: This paper develops an intrinsic reward to help identify factors of variation within a family of MDPs. This intrinsic reward takes a form of curiosity and is used to develop initial behaviors to identify the causes of the latent variation in the environment dynamics. The experiments are used to validate the proposed intrinsic reward across several analyses used to identify its utility and effectiveness.