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 proposes to jointly train two agents in a differentiable world model: an *explorer* that is rewarded to explore uncertain states and an *achiever* whose task is to achieve a given goal state. After training, the *achiever* can be deployed as a goal-reaching policy without further tuning. Authors also spent time investigating how to design the loss function for the *achiever*. In the experiments, three new benchmarks are introduced and the proposed method is shown to outperform baselines.