Discovering and Achieving Goals via World Models

1 Introduction How can we build an agent that learns to solve hundreds of tasks in complex visual environments , such as rearranging objects with a robot arm or completing chores in a kitchen ? While traditional reinforcement learning ( RL ) has been successful for individual tasks , it requires a substantial amount of human effort for every new task . Specifying task rewards requires domain knowledge , access to object positions , is timeconsuming , and prone to human errors . Moreover , traditional RL would require environment interaction to explore and practice in the environment for every new task . Instead , we approach learning hundreds of tasks through the paradigm of unsupervised goal-conditioned RL , where the agent learns many diverse skills in the environment in the complete absence of supervision , to later solve tasks via user-specified goal images immediately without further training [ 2 , 26 , 40 ] . Challenges Exploring the environment and learning to solve many different tasks is substantially more challenging than traditional RL with a dense reward function or learning from ∗ Equal contribution . Ordering determined at random . Project page : https : //orybkin.github.io/lexa/ 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . expert demonstrations . Existing methods are limited to simple tasks , such as picking or pushing a puck [ 13 , 32 , 37 ] or controlling simple 2D robots [ 50 ] . The key challenge in improving the performance of unsupervised RL is exploration . In particular , previous approaches explore by either revisiting previously seen rare goals [ 14 , 18 , 55 ] or sampling goals from a generative model [ 32 , 37 ] . However , in both these approaches , the policy as well as the generative model are trained on previously visited states from the replay buffer , and hence the sampled goals are either within or near the frontier of agent ’ s experience . Ideally , we would like the agent to discover goals much beyond its frontier for efficient exploration , but how does an agent generate goals that it is yet to encounter ? This is an open question not just for AI but for cognitive science too [ 42 ] . Approach To rectify this issue , we leverage a learned world model to train a separate explorer and achiever policy in imagination . Instead of randomly sampling or generating goals , our explorer policy discovers distant goals by first planning a sequence of actions optimized in imagination of the world model to find novel states with high expected information gain [ 30 , 43 , 44 ] . It then executes those imagined actions in the environment to discover interesting states without the need to generate them . Note these actions are likely to lead the agent to states which are several steps outside the frontier because otherwise the model wouldn ’ t have had high uncertainty or information gain . Finally , these discovered states are used as diverse targets for the achiever to practice . We train the achiever from on-policy imagination rollouts within the world model and without relying on experience relabeling , therefore leveraging foresight over hindsight . After this unsupervised training phase , the achiever solves tasks specified as goal images zero-shot without any additional learning at deployment . Unlike in the conventional RL paradigm [ 31 , 47 ] , our method is trained once and then used to achieve several tasks at test time without any supervision during training or testing . Contributions We introduce Latent Explorer Achiever ( LEXA ) , an unsupervised goal reaching agent that trains an explorer and an achiever within a shared world model . At training , LEXA unlocks diverse data for goal reaching in environments where exploration is nontrivial . At test time , the achiever solves challenging locomotion and manipulation tasks provided as user-specified goal images . Our contributions are summarized as follows : • We propose to learn separate explorer and achiever policies as an approach to overcome the exploration problem of unsupervised goal-conditioned RL . • We show that forward-looking exploration by planning with a learned world model substantially outperforms previous strategies for goal exploration . • To evaluate on challenging tasks , we introduce a new goal reaching benchmark with a total of 40 diverse goal images across 4 different robot locomotion and manipulation environments . • LEXA outperforms prior methods , being the first to show success in the Kitchen robotic manipu- lation environment , and achieves goal images where multiple objects need to be moved . 2 Latent Explorer Achiever ( LEXA ) . Our aim is to build an agent that can achieve arbitrary user-specified goals after learning in the environment without any supervision . This presents two challenges - collecting trajectories that contain diverse goals and learning to achieve these goals when specified as a goal image . We introduce a simple solution based on a world model and imagination training that addresses both challenges . The world model represents the agent ’ s current knowledge about the environment and is used for training two policies , the explorer and the achiever . To explore novel situations , we construct an estimate of which states the world model is still uncertain about . To achieve goals , we train the goal-conditioned achiever in imagination , using the images found so far as unsupervised goals . At test time , the achiever is deployed to reach user-specified goals . The training procedure is in Algorithm 1 . 2.1 World Model . To efficiently predict potential outcomes of future actions in environments with high-dimensional image inputs , we leverage a Recurrent State Space Model ( RSSM ) [ 23 ] that learns to predict forward using compact model states that facilitate planning [ 7 , 51 ] . In contrast to predicting forward in image space , the model states enable efficient parallel planning with a large batch size and can reduce accumulating errors [ 39 ] . The world model consists of the following components : Encoder : et = encφ ( xt ) Posterior : qφ ( st | st−1 , at−1 , et ) Dynamics : pφ ( st | st−1 , at−1 ) Image decoder : pφ ( xt | st ) ( 1 ) The model states st contain a deterministic component ht and a stochastic component zt with diagonalcovariance Gaussian distribution . ht is the recurrent state of a Gated Recurrent Unit ( GRU ) [ 11 ] . The encoder and decoder are convolutional neural networks ( CNNs ) and the remaining components are multi-layer perceptrons ( MLPs ) . The world model is trained end-to-end by optimizing the evidence lower bound ( ELBO ) via stochastic backpropagation [ 28 , 38 ] with the Adam optimizer [ 27 ] . 2.2 Explorer . To efficiently explore , we seek out surprising states imagined by the world model [ 6 , 41 , 43 , 44 , 46 ] , as opposed to retrospectively exploring by revisiting previously novel states [ 4 , 5 , 8 , 34 ] . As the world model can predict model states that correspond to unseen situations in the environment , the imagined trajectories contain more novel goals , compared to model-free exploration that is limited to the replay buffer . To collect informative novel trajectories in the environment , we train an exploration Algorithm 1 : Latent Explorer Achiever ( LEXA ) 1 : initialize : World modelM , Replay buffer D , Explorer πe ( at | zt ) , Achiever πg ( at | zt , g ) 2 : while exploring do 3 : TrainM on D 4 : Train πe in imagination ofM to maximize exploration rewards ∑ t r e t . 5 : Train πg in imagination ofM to maximize ∑ t r g t ( zt , g ) for images g ∼ D. 6 : ( Optional ) Train d ( zi , zj ) to predict distances j − i on the imagination data from last step . 7 : Deploy πe in the environment to explore and grow D. 8 : Deploy πg in the environment to achieve a goal image g ∼ D to grow D. 9 : end while 10 : while evaluating do 11 : given : Evaluation goal g 12 : Deploy πg in the world to reach g. 13 : end while policy πe from the model states st in imagination of the world model to maximize an exploration reward : Explorer : πe ( at | st ) Explorer Value : ve ( st ) ( 2 ) To explore the most informative model states , we estimate the epistemic uncertainty as a disagreement of an ensemble of transition functions . We train an ensemble of 1-step models to predict the next model state from the current model state . The ensemble model is trained alongside the world model on model states produced by the encoder qφ . Because the ensemble models are initialized at random , they will differ , especially for inputs that they have not been trained on [ 29 , 36 ] : Ensemble : f ( st , θk ) = ẑkt+1 for k = 1 .. K ( 3 ) Leveraging the ensemble , we estimate the epistemic uncertainty as the ensemble disagreement . The exploration reward is the variance of the ensemble predictions averaged across dimension of the model state , which approximates the expected information gain [ 3 , 43 ] : ret ( st ) . = 1 N ∑ n Var { k } [ f ( st , θk ) ] n ( 4 ) The explorer πe maximizes the sum of future exploration rewards ret using the Dreamer algorithm [ 24 ] , which considers long-term rewards into the future by maximizing λ-returns under a learned value function . As a result , the explorer is trained to seek out situations are as informative as possible from imagined latent trajectories of the world model , and is periodically deployed in the environment to add novel trajectories to the replay buffer , so the world model and goal achiever policy can improve . 2.3 Achiever . To leverage the knowledge obtained by exploration for learning to reach goals , we train a goal achiever policy πg that receives a model state and a goal as input . Our aim is to train a general policy that is capable of reaching many diverse goals . To achieve this in a data-efficient way , it is crucial that environment trajectories that were collected with one goal in mind are reused to also learn how to reach other goals . While prior work addressed this by goal relabeling which makes off-policy policy optimization a necessity [ 2 ] , we instead leverage past trajectories via the world model trained on them that lets us generate an unlimited amount of new imagined trajectories for training the goal achiever on-policy in imagination . This simplifies policy optimization and can improve stability , while still sharing all collected experience across many goals . Achiever : πg ( at | st , eg ) Achiever Value : vg ( st , eg ) ( 5 ) To train the goal achiever , we sample a goal image xg from the replay buffer and compute its embedding eg = encφ ( xg ) . The achiever aims to maximize an unsupervised goal-reaching reward rg ( st , eg ) . We discuss different choices for this reward in Section 2.4 . We again use the Dreamer algorithm [ 24 ] for training , where now the value function also receives the goal embedding as input . In addition to imagination training , it can also be important to perform practice trials with the goal achiever in the true environment , so that any model inaccuracies along the goal reaching trajectories may be corrected . To perform practice trials , we sample a goal from the replay buffer and execute the goal achiever policy for that goal in the environment . These trials are interleaved with exploration episodes collected by the exploration policy in equal proportion . We note that the goal achiever learning is entirely unsupervised because the practice goals are simply images the agent encountered through exploration or during previous practice trails .

Summary: This paper introduces a framework to train policies for visual goal oriented tasks. The framework comprises two policies, the first, an explorer, which is trained to seek out novel states (where novelty is defined as the level of agreement between an ensemble of future state predictors - an approximate measure of epistemic uncertainty), while the second, an achiever, is trained to reach these states. A model based approach to train the achiever, where the achiever is trained using trajectories generated by a learned dynamics model. The paper also introduces a new benchmark for testing goal oriented polices of this form, comprising a kitchen manipulation task, a block handling task, and a yoga task.