Planning from Pixels using Inverse Dynamics Models

1 INTRODUCTION . Deep reinforcement learning has proven to be a powerful and effective framework for solving a diversity of challenging decision-making problems ( Silver et al. , 2017a ; Berner et al. , 2019 ) . However these algorithms are typically trained to maximize a single reward function , ignoring information that is not directly relevant to the associated task at hand . This way of learning is in stark contrast to how humans learn ( Tenenbaum , 2018 ) . Without being prompted by a specific task , humans can still explore their environment , practice achieving imaginary goals , and in so doing learn about the dynamics of the environment . When subsequently presented with a novel task , humans can utilize this learned knowledge to bootstrap learning — a property we would like our artificial agents to have . In this work , we investigate one way to bridge this gap by learning world models ( Ha & Schmidhuber , 2018 ) that enable the realization of previously unseen tasks . By modeling the task-agnostic dynamics of an environment , an agent can make predictions about how its own actions may affect the environment state without the need for additional samples from the environment . Prior work has shown that by using powerful function approximators to model environment dynamics , training an agent entirely within its own world models can result in large gains in sample efficiency ( Ha & Schmidhuber , 2018 ) . However , learning world models that are both accurate and general has largely remained elusive , with these models experiencing many performance issues in the multi-task setting . The main reason for poor performance is the so-called planning horizon dilemma ( Wang et al. , 2019 ) : accurately modeling dynamics over a long horizon is necessary to accurately estimate rewards , but performance is often poor when planning over long sequences due to the accumulation of errors . These modeling errors are especially prevalent in high-dimensional observation spaces where loss functions that operate on pixels may focus model capacity on task-irrelevant features ( Kaiser et al. , 2020 ) . Recent work ( Hafner et al. , 2020 ; Schrittwieser et al. , 2019 ) has attempted to side-step these issues by learning a world model in a latent space and propagating gradients over multiple time-steps . While these methods are able to learn accurate world models and achieve high performance on benchmark tasks , their representations are usually trained with task-specific information such as rewards , encouraging the model to focus on tracking task-relevant features but compromising their ability to generalize to new tasks . In this work , we propose to learn powerful , latent world models that can predict environment dynamics when planning for a distribution of tasks . The main contributions of our paper are three-fold : we propose to learn a latent world model conditioned on a goal ; we train our latent representation to model inverse dynamics — sequences of actions that take the agent from one state to another , rather than training it to capture information about reward ; and we show that by combining our inverse dynamics model and a prior over action sequences , we can quickly construct plans that maximize the probability of reaching a goal state . We evaluate our world model on a diverse distribution of challenging visual goals in Atari games and the Deepmind Control Suite ( Tassa et al. , 2018 ) to assess both its accuracy and sample efficiency . We find that when planning in our latent world model , our agent outperforms prior , model-free methods across most tasks , while providing an order of magnitude better sample efficiency on some tasks . 2 RELATED WORK . Model-based RL has typically focused on learning powerful forward dynamics models , which are trained to predict the next state given the current state and action . In works such as ( Kaiser et al. , 2020 ) , these models are trained to predict the next state in observation space - often by minimizing L2 distance . While the performance of these algorithms in the low data regime is often strong , they can struggle to reach the asymptotic performance of model-free methods ( Hafner et al. , 2020 ) . An alternative approach is to learn a forward model in a latent space , which may be able to avoid modeling irrelevant features and better optimize for long-term consistency . These latent spaces can be trained to maximize mutual information with the observations ( Hafner et al. , 2020 ; 2019 ) or even task-specific quantities like the reward , value , or policy ( Schrittwieser et al. , 2019 ) . Using a learned forward model , there are several ways that an agent could create a policy . While forward dynamics models map a state and action to the next state , an inverse dynamics model maps two subsequent states to an action . Inverse dynamics models have been used in various ways in sequential decision making . In exploration , inverse dynamics serves as a way to learn representations of the controllable aspects of the state ( Pathak et al. , 2017 ) . In imitation learning , inverse dynamics models can be used to map a sequence of states to the actions needed to imitate the trajectory ( Pavse et al. , 2019 ) . Christiano et al . ( 2016 ) use inverse dynamics models to translate actions taken in a simulated environment to the real world . Recently , there has been an emergence of work ( e.g. , Ghosh et al. , 2020 ; Schmidhuber , 2019 ; Srivastava et al. , 2019 ) highlighting the relationship between imitation learning and reinforcement learning . Specifically , rather than learn to map states and actions to reward , as is typical in reinforcement learning , Srivastava et al . ( 2019 ) train a model to predict actions given a state and an outcome , which could be the amount of reward the agent is to collect within a certain amount of time . Ghosh et al . ( 2020 ) use a similar idea , predicting actions conditioned on an initial state , a goal state , and the amount of time left to achieve the goal . As explored in Appendix A.1 , these methods are perhaps the nearest neighbors to our algorithm . In our paper , we tackle a visual goal-completion task due to its generality and ability to generate tasks with no domain knowledge . Reinforcement learning with multiple goals has been studied since Kaelbling ( 1993 ) . Most agents that are trained to achieve multiple goals are trained with offpolicy reinforcement learning combined with a form of hindsight relabeling ( Andrychowicz et al. , 2017 ) , where trajectories that do not achieve the desired goal are relabeled as a successful trajectory that achieves the goal that was actually reached . Andrychowicz et al . ( 2017 ) uses value-based reinforcement learning with a reward based on the euclidean distance between physical objects , which is only possible with access to an object-oriented representation of the state . In environments with high-dimensional observation spaces , goal-achievement rewards are more difficult to design . Nair et al . ( 2018 ) use a VAE ( Kingma & Welling , 2014 ) trained on observations to construct a latent space and uses distances in the latent space for a reward . These distances , however , may contain features that are uncontrollable or irrelevant . Warde-Farley et al . ( 2019 ) attempt to solve this issue by framing the goal-achievement task as maximizing the mutual information between the goal and achieved state I ( sg , sT ) . Our method differs from these approaches since we aim simply to maximize an indicator reward 1 ( sT = sg ) and do not explicitly learn a value or Q-function . 3 METHOD . 3.1 PROBLEM FORMULATION . Reinforcement learning is a framework in which an agent acts in an unknown environment and adapts based on its experience . We model the problem with an MDP , defined as the tuple ( S , A , T , R , γ ) . S is a set of states ; A is a set of actions ; the transition probabilities T : S×A×S → [ 0 , 1 ] defines the probability of the environment transitioning from state s to s′ given that the agent acts with action a ; the reward function R : S × A× S → R maps a state-action transition to a real number ; and 0 ≤ γ ≤ 1 is the discount factor , which controls how much an agent should prefer rewards sooner rather than later . An agent acts in the MDP with a policy π : S ×A→ [ 0 , 1 ] , which determines the probability of the agent taking action a while in state s. The expected return of a policy is denoted : JRL ( π ) = Eτ∼P ( τ |π ) [ ∑ t γtR ( st , at , s ′ t ) ] , ( 1 ) that is the averaged discounted future rewards for trajectories τ = { ( st , at ) } Tt=1 of states and actions sampled from the policy . A reinforcement learning agent ’ s objective is to find the optimal policy π∗ = argmaxπ JRL ( π ) that maximizes the expected return . In goal-conditioned reinforcement learning , an agent ’ s objective is to find a policy that maximizes this return over the distribution of goals g ∼ p ( g ) when acting with a policy that is now also conditioned on g. In our work , g ∈ S and we consider goal achievement rewards of the form Rg ( s ) = 1 ( s = g ) . Additionally , we consider a trajectory to be complete when any Rg ( st ) = 1 and denote this time-step t = T . With these rewards , an optimal goal-achieving agent maximizes : J ( π ) = Eg∼p ( g ) [ EsT∼p ( sT |πg ) [ γ TRg ( sT ) ] ] . ( 2 ) Note that unlike prior works , we consider both the probability of goal achievement as well as the length of the trajectory T in our objective . 3.2 PLANNING . We consider the problem of finding an optimal action sequence a1 , . . . , ak−1 to maximize expected return J ( s , g , a1 , . . . , ak−1 ) : J ( s , g , a1 , . . . , ak−1 ) = Esk∼p ( sk|s , a1 , ... , ak−1 ) [ γ krg ( sk ) ] = γ T p ( sk = g|s , a1 , . . . , ak−1 ) ( 3 ) Thus , the optimal action sequence is found by solving the following optimization problem : a∗1 , . . . , a ∗ k−1 = argmax a1 , ... , ak−1 γkp ( sk = g|s1 , a1 , . . . , ak−1 ) ( 4 ) Even with access to a perfect model of p ( sk = g|s1 , a1 , . . . , ak−1 ) , solving this optimization may be difficult . In many environments , the number of action sequences that reach the goal are vastly outnumbered by the action sequences that do not . Without a heuristic or reward-shaping , there is little hope of solving this problem in a reasonable amount of time . 3.3 GLAMOR : GOAL-CONDITIONED LATENT ACTION MODELS FOR RL . Inspired by sequence modeling in NLP , we propose to rewrite Equation 4 in a way that permits factoring across the actions in the action sequence . By factoring , planning in our model can use the heuristic search algorithms that enable sampling high quality language sequences that are hundreds of tokens long . First , note that12 : p ( sk = g|s1 , a1 , . . . , ak−1 ) ∝ k−1∏ i=1 p ( sk = g|s1 , a < i , ai ) p ( sk = g|s1 , a < i ) ( 5 ) Let z ( s1 , g , a < i , ai ) , p ( sk=g|s1 , a < i , ai ) p ( sk=g|s1 , a < i ) . Intuitively , these terms are equal to the relative gain in probability of reaching state g conditioned on taking action ai versus the marginal probability of reaching the goal without conditioning on that action . These terms provide useful information3 that can guide search towards high scoring action sequences when constructing a plan . To learn the values of the z ( s1 , g , a < i , ai ) , we use Bayes ’ rule to show that we can equivalently learn two auto-regressive behavioral models : p ( sk = g|s1 , a≤i ) p ( sk = g|s1 , a≤i−1 ) = p ( ai|sk = g , s1 , a≤i−1 ) p ( sk = g , s1 , a≤i−1 ) p ( s1 , a≤i−1 ) p ( s1 , a≤i ) p ( sk = g , s1 , a≤i−1 ) ( 6 ) = p ( ai|s1 , sk = g , a≤i−1 ) p ( ai|s1 , a≤i−1 ) ( 7 ) We refer to p ( a1 , . . . , ak|s1 , sk = g ) as the inverse dynamics model and p ( a1 , . . . , ak|s1 ) as the action prior . Using these models , we can find an optimal plan by optimizing the following objective : a∗1 , . . . , a ∗ k−1 = argmax a1 , ... , ak−1 γk p ( a1 , . . . , ak−1|s1 , sk = g ) p ( a1 , . . . , ak−1|s1 ) ( 8 )

Summary: The author proposes Goal-Conditioned Latent Action Models for RL (GLAMOR) a novel approach to learn latent world models by modeling inverse dynamics. The proposed approach learns to track task-relevant dynamics for a diverse distribution of tasks and provide a strong heuristic that enables efficient planning. GLAMOR demonstrates good performance against its baselines in terms of achieving accurately goals, sample efficiency and effective planning.