Learning What To Do by Simulating the Past

1 INTRODUCTION . As deep learning has become popular , many parts of AI systems that were previously designed by hand have been replaced with learned components . Neural architecture search has automated architecture design ( Zoph & Le , 2017 ; Elsken et al. , 2019 ) , population-based training has automated hyperparameter tuning ( Jaderberg et al. , 2017 ) , and self-supervised learning has led to impressive results in language modeling ( Devlin et al. , 2019 ; Radford et al. , 2019 ; Clark et al. , 2020 ) and reduced the need for labels in image classification ( Oord et al. , 2018 ; He et al. , 2020 ; Chen et al. , 2020 ) . However , in reinforcement learning , one component continues to be designed by humans : the task specification . Handcoded reward functions are notoriously difficult to specify ( Clark & Amodei , 2016 ; Krakovna , 2018 ) , and learning from demonstrations ( Ng et al. , 2000 ; Fu et al. , 2018 ) or preferences ( Wirth et al. , 2017 ; Christiano et al. , 2017 ) requires a lot of human input . Is there a way that we can automate even the specification of what must be done ? It turns out that we can learn part of what the user wants simply by looking at the state of the environment : after all , the user will already have optimized the state towards their own preferences ( Shah et al. , 2019 ) . For example , when a robot is deployed in a room containing an intact vase , it can reason that if its user wanted the vase to be broken , it would already have been broken ; thus she probably wants the vase to remain intact . However , we must ensure that the agent distinguishes between aspects of the state that the user couldn ’ t control from aspects that the user deliberately designed . This requires us to simulate what the user must have done to lead to the observed state : anything that the user put effort into in the past is probably something the agent should do as well . As illustrated in Figure 1 , if we observe a Cheetah balancing on its front leg , we can infer how it must have launched itself into that position . Unfortunately , it is unclear how to simulate these past trajectories that lead to the observed state . So far , this has only been done in gridworlds , where all possible trajectories can be considered using dynamic programming ( Shah et al. , 2019 ) . Our key insight is that we can sample such trajectories by starting at the observed state and simulating backwards in time . To enable this , we derive a gradient that is amenable to estimation through backwards simulation , and learn an inverse policy and inverse dynamics model using supervised ∗Work done at the Center for Human-Compatible AI , UC Berkeley . learning to perform the backwards simulation . Then , the only remaining challenge is finding a reward representation that can be meaningfully updated from a single state observation . To that end , rather than defining the reward directly on the raw input space , we represent it as a linear combination of features learned through self-supervised representation learning . Putting these components together , we propose the Deep Reward Learning by Simulating the Past ( Deep RLSP ) algorithm . We evaluate Deep RLSP on MuJoCo environments and show that it can recover fairly good performance on the task reward given access to a small number of states sampled from a policy optimized for that reward . We also use Deep RLSP to imitate skills generated using a skill discovery algorithm ( Sharma et al. , 2020 ) , in some cases given just a single state sampled from the policy for that skill . Information from the environment state can not completely replace reward supervision . For example , it would be hard to infer how clean Bob would ideally want his room to be , if the room is currently messy because Bob is too busy to clean it . Nonetheless , we are optimistic that information from the environment state can be used to significantly reduce the burden of human supervision required to train useful , capable agents . 2 METHOD . In this section , we describe how Deep RLSP can learn a reward function for high dimensional environments given access only to a simulator and the observed state s0 . Notation . A finite-horizon Markov Decision Process ( MDP ) M = 〈S , A , T , r , P , T 〉 contains a set of states S and a set of actions A . The transition function T : S × A × S 7→ [ 0 , 1 ] determines the distribution over next states given a state and an action , and P is a prior distribution over initial states . The reward function r : S 7→ R determines the agent ’ s objective . T ∈ Z+ is a finite planning horizon . A policy π : S ×A 7→ [ 0 , 1 ] specifies how to choose actions given a state . Given an initial state distribution , a policy and the transition function , we can sample a trajectory τ by sampling the first state from P , every subsequent action from π , and every subsequent state from T . We denote the probability distribution over trajectories as 〈P , π , T 〉 and write τ ∼ 〈P , π , T 〉 for the sampling step . We will sometimes write a single state s instead of a distribution P if the initial state is deterministic . The goal of reinforcement learning ( RL ) is to find a policy π∗ that maximizes the expected cumulative reward Eτ∼〈P , π , T 〉 [ ∑T t=1 r ( st ) ] . We use φ : S → Rn to denote a feature function ( whether handcoded or learned ) that produces a feature vector of length n for every state . The reward function r is linear over φ if it can be expressed in the form r ( s ) = θTφ ( s ) for some θ ∈ Rn . We assume that some past trajectory τ−T :0 = s−Ta−T . . . a−1s0 produced the observed state s0 . 2.1 IDEALIZED ALGORITHM . We first explain what we would ideally do , if we had a handcoded a feature function φ and an enumerable ( small ) state space S that affords dynamic programming . This is a recap of Reward Learning by Simulating the Past ( RLSP ; Shah et al. , 2019 ) . We assume the human follows a Boltzmann-rational policy πt ( a | s , θ ) ∝ exp ( Qt ( s , a ; θ ) ) , where the Q values are computed using soft value iteration . Marginalizing over past trajectories , yields a distribution over the observed state p ( s0 | θ ) = ∑ s−T ... a−1 p ( τ = s−Ta−T . . . a−1s0 | θ ) . We compute the maximum likelihood estimate , argmaxθ ln p ( s0 | θ ) , via gradient ascent , by expressing the gradient of the observed state as a weighted combination of gradients of consistent trajectories ( Shah et al. , 2019 , Appendix B ) : ∇θ ln p ( s0 | θ ) = E τ−T : −1 ∼ p ( τ−T : −1|s0 , θ ) [ ∇θ ln p ( τ | θ ) ] ( 1 ) ∇θ ln p ( τ | θ ) is a gradient for inverse reinforcement learning . Since we assume a Boltzmann-rational human , this is the gradient for Maximum Causal Entropy Inverse Reinforcement Learning ( MCEIRL ; Ziebart et al. , 2010 ) . However , we still need to compute an expectation over all trajectories that end in s0 , which is in general intractable . Shah et al . ( 2019 ) use dynamic programming to compute this gradient in tabular settings . 2.2 GRADIENT AS BACKWARDS-FORWARDS CONSISTENCY . Approximating the expectation . For higher-dimensional environments , we must approximate the expectation over past trajectories p ( τ−T : −1 | s0 , θ ) . We would like to sample from the distribution , but it is not clear how to sample the past conditioned on the present . Our key idea is that just as we can sample the future by rolling out forwards in time , we should be able to sample the past by rolling out backwards in time . Note that by the Markov property we have : p ( τ−T : −1 | s0 , θ ) = −1∏ t=−T p ( st | at , st+1 , . . . s0 , θ ) p ( at | st+1 , at+1 , . . . s0 , θ ) = −1∏ t=−T p ( st | at , st+1 , θ ) p ( at | st+1 , θ ) Thus , given the inverse policy π−1t ( at | st+1 , θ ) , the inverse dynamics T −1t ( st | at , st+1 , θ ) , and the observed state s0 , we can sample a past trajectory τ−T : −1 ∼ p ( τ−T : −1 | s0 , θ ) by iteratively applying π−1 and T −1 , starting from s0 . Analogous to forward trajectories , we express the sampling as τ−T : −1 ∼ 〈s0 , π−1 , T −1〉 . Thus , we can write the gradient in Equation 1 as Eτ−T : −1 ∼ 〈s0 , π−1 , T −1〉 [ ∇θ ln p ( τ | θ ) ] . Learning π , π−1 and T −1 . In order to learn π−1 , we must first know π . We assumed that the human was Boltzmann-rational , which corresponds to the maximum entropy reinforcement learning objective ( Levine , 2018 ) . We use the Soft Actor-Critic algorithm ( SAC ; Haarnoja et al. , 2018 ) to estimate the policy π ( a | s , θ ) , since it explicitly optimizes the maximum entropy RL objective . Given the forward policy π ( a | s , θ ) and simulator T , we can construct a dataset of sampled forward trajectories , and learn the inverse policy π−1 and the inverse dynamics T −1 using supervised learning . Given these , we can then sample τ−T : −1 , allowing us to approximate the expectation in the gradient . In general , both π−1 and T −1 could be stochastic and time-dependent . Estimating the gradient for a trajectory . We now turn to the term within the expectation , which is the inverse reinforcement learning gradient given a demonstration trajectory τ = s−Ta−T . . . s0 . Assuming that the user is Boltzmann-rational , this is the MCEIRL gradient ( Ziebart et al. , 2010 ) , which can be written as ( Shah et al. , 2019 , Appendix A ) : ∇θ ln p ( τ | θ ) = ( 0∑ t=−T φ ( st ) ) −F−T ( s−T ) + −1∑ t=−T ( E s′t+1∼T ( ·|st , at ) [ Ft+1 ( s′t+1 ) ] −Ft+1 ( st+1 ) ) ( 2 ) F is the expected feature count under π , that is , F−t ( s−t ) , Eτ−t:0 ∼ 〈s−t , π , T 〉 [ ∑0 t′=−t φ ( st′ ) ] . The first term computes the feature counts of the demonstrated trajectory τ , while the second term computes the feature counts obtained by the policy for the current reward function θ ( starting from the initial state s−T ) . Since r ( s ) = θTφ ( s ) , these terms increase the reward of features present in the demonstration τ and decrease the reward of features under the current policy . Thus , the gradient incentivizes consistency between the demonstration and rollouts from the learned policy . The last term is essentially a correction for the observed dynamics : if we see that st , at led to st+1 , it corrects for the fact that we “ could have ” seen some other state s′t+1 . Since this correction is zero in expectation ( and expensive to compute ) , we drop it for our estimator . Gradient estimator . After dropping the last term in Equation 2 , expanding the definition of F , and substituting in to Equation 1 , our final gradient estimator is : ∇θ ln p ( s0 | θ ) = E τ−T : −1 ∼ 〈s0 , π−1 , T −1〉 [ ( 0∑ t=−T φ ( st ) ) − E τ ′ ∼ 〈s−T , π , T 〉 [ ( 0∑ t=−T φ ( s′t ) ) ] ] ( 3 ) Thus , given s0 , θ , π , T , π−1 , and T −1 , computing the gradient consists of three steps : 1 . Simulate backwards from s0 , and compute the feature counts of the resulting trajectories . 2 . Simulate forwards from s−T of these trajectories , and compute their feature counts . 3 . Take the difference between these two quantities . This again incentivizes consistency , this time between the backwards and forwards trajectories : the gradient leads to movement towards “ what the human must have done ” and away from “ what the human would do if they had this reward ” . The gradient becomes zero when they are identical . It may seem like the backwards and forwards trajectories should always be consistent with each other , since π−1 and T −1 are inverses of π and T . The key difference is that s0 imposes constraints on the backwards trajectories , but not on the forward trajectories . For example , suppose we observe s0 in which a vase is unbroken , and our current hypothesis is that the user wants to break the vase . When we simulate backwards , our trajectory will contain an unbroken vase , but when we simulate forwards from s−T , π will break the vase . The gradient would then reduce the reward for a broken vase and increase the reward for an unbroken vase .

Summary: This paper introduces an algorithm, called deep reward learning by simulating the past (deep RLSP), that seeks to infer a reward function by looking at states in demonstration data. An example of this described in the paper is an environment with a vase: if demonstration data shows an intact vase in the presence of an embodied agent then breaking the vase is unlikely to be the intended behavior. Otherwise the vase would already be broken in the demo.