Discovering and Achieving Goals via World Models

• Traditional RL requires human effort and environment interaction for each new task, motivating unsupervised goal-conditioned learning from which agents can later solve goal-image tasks without further training [2].Source: §1 · p. 1
• Earlier exploration methods revisit rare goals [14] [18] [56] or sample goals from generative models [33] [38], keeping goals within or near the frontier of previously visited states.Source: §1 · p. 2
• LEXA trains separate explorer and achiever policies in a shared world model: the explorer plans imagined actions for novel, high-information states, and the achiever practices reaching discovered states using on-policy imagined rollouts.The achiever is trained without experience relabeling, using foresight rather than hindsight.Source: Figure 1 · p. 1, §1 · pp. 2–3
• LEXA outperforms prior methods, is the first to show success in Kitchen, and achieves goals requiring multiple objects to be moved across a new benchmark of 40 goal images in four robot environments.The benchmark covers locomotion and manipulation tasks.Source: §1 · p. 2

2.1 World Model

• An RSSM predicts compact recurrent model states from image inputs, enabling efficient parallel planning; its states combine deterministic recurrent and stochastic components.The recurrent component uses a GRU, while the stochastic component has a diagonal-covariance Gaussian distribution.Source: §2.1 · p. 3

2.2 Explorer

• Ensemble disagreement supplies an exploration reward, and Dreamer-trained imagined rollouts seek informative states before the explorer is deployed to add trajectories to replay.The reward is the variance of ensemble transition predictions across model-state dimensions, which approximates expected information gain.Source: §2.2 · p. 4

2.3 Achiever

• The achiever reuses collected experience across goals by training on imagined on-policy trajectories, with replay-buffer images as goals and unsupervised practice trials in the environment.Practice trials are interleaved equally with exploration episodes, helping correct model inaccuracies along goal-reaching trajectories.Source: §2.3 · pp. 4–5

2.4 Latent Distances

• Temporal distance rewards focus on steps needed to reach a goal, avoiding cosine-distance bias toward controllable robot positions when goals also involve objects.LEXA trains the distance predictor on imagined trajectories and includes states from different trajectories as maximum-distance examples.Source: §2.4 · pp. 5–6

3.1 Experimental setup

• The new benchmark covers RoboYoga poses, RoboBins reaching and manipulation tasks, and RoboKitchen goals involving one or two objects.For baseline comparisons, the authors implemented methods using the same world model and policy optimization where possible, while retaining SkewFit’s original model-free implementation.Source: §3.1 · pp. 6–7

3.2 Performance on New Benchmark

• LEXA with temporal distance outperforms prior methods on RoboBins and learns RoboKitchen tasks, including sequences involving two objects; it is the only approach progressing on every RoboBins goal.On RoboYoga, cosine-distance LEXA outperforms prior approaches, while temporal distances improve on more goals despite average success similar to DDL. In RoboKitchen, prior methods make little progress, and cosine-distance LEXA also makes very limited progress.Source: §3.2 · pp. 7–8, Table 1 · p. 7, Figure 6 · p. 7

Single Agent Across All Environments

• A single LEXA agent with learned temporal distance makes progress across RoboKitchen, RoboBins reaching and pick-and-place, and Walker tasks.The best prior single-environment method, DDL, mainly solves Walker and RoboBins reaching.Source: Single Agent Across All Environments · p. 8

3.3 Performance on Prior benchmarks

• On prior benchmarks, LEXA significantly outperforms previous methods on SkewFit tasks, solves many reproduced DISCERN tasks, and beats Plan2Explore on most tasks while matching oracle agents comparably.For Plan2Explore tasks, LEXA receives goal images but not reward functions; the compared oracle agents use true task rewards during training.Source: §3.3 · pp. 8–9

3.4 Analysis

• Ablations show that a separate explorer is crucial, negative sampling speeds learning, and imagination data performs better than real data on the hardest tasks.During exploration, LEXA encounters tasks involving multiple objects more often; without a separate explorer, it does not discover the more interesting tasks.Source: §3.4 · p. 9, Figure 8 · p. 9

Learning to Achieve Goals

• Model-free methods learn a single goal-conditioned policy [2] and generate goals through asymmetric self-play [34] [46] or intermediate-difficulty sampling [14] [18], but focus on low-dimensional state.These methods can achieve remarkable performance in simulated robotic domains.Source: §4 · p. 10
• Visual-goal methods using contrastive or reconstructive representations [51] [33] [38] struggle to explore without a clear reward; dataset-based methods [10] [50] and seen-goal generation [38] [56] avoid or limit that challenge.The latter generate goals similar to states already seen rather than exploring truly novel states.Source: §4 · p. 10
• Model-based planning and model-regularized goal-reaching methods [13] [17] are limited by short horizons, whereas LEXA uses long-horizon value functions and imagined exploration to find novel states.The paper says this exploration significantly improves exploration and downstream agent capabilities.Source: §4 · p. 10

Learning Distance Functions

• Feature-space distances [9] [33] [51] [52] may poorly reflect goal difficulty, while mutual-information distances [1] [12] [15] [21] have yet to demonstrate scalability to more complex tasks.These are alternative approaches to defining reward or cost for visual goal achievement.Source: §4 · p. 10
• Temporal distance methods use approximate dynamic programming with Q-learning [16] [19]; unlike Hartikainen et al., LEXA trains distance on-policy in imagination and integrates the achiever to practice on discovered goals.Hartikainen et al. learn temporal distance through supervised learning on recent policy experience.Source: §4 · p. 10

• LEXA searches for novelty in imagination to discover meaningful behaviors across more diverse environments than prior work, then learns to achieve discovered goals.Source: §5 · p. 10
• LEXA solves challenging downstream tasks specified as images zero-shot, without supervision from rewards or demonstrations.Source: §5 · p. 10
• The paper proposes a challenging benchmark and presents the first agent to achieve meaningful performance on its tasks, though many remain unsolved.Source: §5 · p. 10
• Future work should improve world models and policy optimization, demonstrate scalability on real-world systems, and enable goal specification through accessible means such as natural language.Source: §5 · p. 10

Evaluation and environments

• Success is measured at each episode’s final step; benchmark experiments use three seeds, while LEXA’s DISCERN and Plan2Explore runs use one seed and are indicative only.Episodes last 150 steps for RoboBin and RoboKitchen and 1000 for RoboYoga; success thresholds vary by task, including 10 cm placement criteria in RoboBin and manually set visual thresholds in RoboKitchen.Source: A Experimental Details · p. 14

Implementation and hyperparameters

• LEXA builds on Dreamer; achiever goals are sampled from replay-buffer trajectories, and real-environment data collection samples goals uniformly from the replay buffer.Initial and goal states for achiever training come from the same sampled batch, creating a mix of easy and hard goals. Training uses Dreamer V2 DM Control defaults, action repeat 2, and one additional hyperparameter: the proportion of negative sampled goals; parameters were manually tuned and shared base settings support fairness.Source: A Experimental Details · p. 14, Table A.5 · p. 15

Baseline details

• DIAYN and GCSL can achieve simple reaching goals but struggle with more complex tasks; both perform best when conditioned on image embedding e rather than latent state s.DIAYN uses a 16-dimensional skill space with a uniform prior and Gaussian reverse predictor; GCSL trains on replay-buffer images and uses imagined rollouts only for explorer training.Source: A Experimental Details · p. 14

Benchmark results

• LEXA makes progress on a majority of RoboYoga Walker and Quadruped tasks; temporal distance supports multi-block RoboBin goals and challenging RoboKitchen tasks.In RoboKitchen, LEXA makes progress on four of six base tasks and can solve combined goal images. The single-agent evaluation also shows progress across Kitchen, RoboBin, and Walker with temporal distance, while cosine-distance LEXA and DDL do not progress on Kitchen.Source: Figure A.2 · p. 16, Figure A.3 · p. 16, Figure A.4 · p. 16

Prior-benchmark comparisons

• On reproduced SkewFit [38], DISCERN [51], and Plan2Explore [44] tasks, LEXA outperforms prior work on SkewFit and the reported DISCERN comparison, and beats Plan2Explore on most tasks.LEXA largely solves visual block pushing and pickup on SkewFit. For DISCERN, the original tasks were unavailable, and reproduced goals may differ; on Plan2Explore tasks, LEXA uses goal images without test-time reward access, unlike Plan2Explore, and results use early stopping.Source: Figure A.7 · p. 18, Figure A.8 · p. 18, Table A.2 · p. 18