Source-linked AI summary
Discovering and Achieving Goals via World Models
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, Deepak Pathak
TL;DR
Agents must discover and achieve many diverse goals in complex visual environments without supervision. LEXA learns a world model from images and trains explorer and achiever policies through imagined rollouts, with the explorer seeking unseen surprising states. LEXA outperforms prior unsupervised goal-reaching methods across prior and new benchmarks and achieves challenging sequential multi-object goals, including with one agent trained across four environments.
Problem
Unsupervised agents need to discover diverse goals and learn to achieve image-specified tasks without rewards or demonstrations.
Method
LEXA learns a world model from images and uses imagined rollouts to train a forward-looking explorer and a goal-conditioned achiever.
Results
LEXA outperforms prior approaches on prior and new benchmarks, achieves sequential multi-object goals, and makes progress across four environments with one general agent.
Takeaways & Limitations
Forward-looking exploration in imagination supports zero-shot image-goal achievement in more diverse visual environments than prior unsupervised methods.
Abstract
from arXiv · showhide
How can artificial agents learn to solve many diverse tasks in complex visual environments in the absence of any supervision? We decompose this question into two problems: discovering new goals and learning to reliably achieve them. We introduce Latent Explorer Achiever (LEXA), a unified solution to these that learns a world model from image inputs and uses it to train an explorer and an achiever policy from imagined rollouts. Unlike prior methods that explore by reaching previously visited states, the explorer plans to discover unseen surprising states through foresight, which are then used as diverse targets for the achiever to practice. After the unsupervised phase, LEXA solves tasks specified as goal images zero-shot without any additional learning. LEXA substantially outperforms previous approaches to unsupervised goal-reaching, both on prior benchmarks and on a new challenging benchmark with a total of 40 test tasks spanning across four standard robotic manipulation and locomotion domains. LEXA further achieves goals that require interacting with multiple objects in sequence. Finally, to demonstrate the scalability and generality of LEXA, we train a single general agent across four distinct environments. Code and videos at https://orybkin.github.io/lexa/
1 Introduction
LEXA addresses unsupervised learning of diverse visual tasks by separating goal discovery from goal achievement. It uses a learned world model and imagination to explore novel goals, train an achiever, and reach user-specified image goals zero-shot.
- Motivation: Unsupervised goal-conditioned RL seeks diverse skills without task rewards or demonstrations, reducing the human effort required to specify and train each new task.Traditional RL requires domain knowledge, task rewards, and environment interaction for every new task.
- Evaluation: LEXA is evaluated across four visual control environments, including complex locomotion, multi-object manipulation, and kitchen tasks requiring strategies such as opening a cabinet.The new benchmark contains 40 diverse goal images across four robot locomotion and manipulation environments.
- Challenge: Prior exploration methods revisit rare states or generate goals from replay-buffer data, limiting goals to or near the agent’s experience frontier.LEXA instead targets goals beyond that frontier through forward-looking exploration.
- Approach: The explorer plans action sequences in the imagined world model to discover novel states rather than merely revisiting previously visited states.The achiever is trained on these discovered goals and on replay-buffer images using a latent distance objective.
- Results: LEXA outperforms prior methods, succeeds for the first time in the Kitchen environment, and reaches goals requiring multiple objects to be moved.After unsupervised training, the achiever handles user-specified goals zero-shot without further test-time learning.
- Approach: LEXA learns a world model from images and uses imagined rollouts to train separate explorer and achiever policies.The explorer discovers new images, while the achiever learns to reach them reliably.
2 Latent Explorer Achiever (LEXA)
LEXA uses a learned world model to train separate explorer and achiever policies in imagination. The explorer seeks informative unseen states, while the achiever practices reaching replay-buffer goals and can handle multi-object and temporally extended tasks.
- Overview: LEXA trains separate explorer and achiever policies using imagined rollouts from a learned world model.The world model is trained from image inputs and supports both exploration and goal-achievement learning.
- 2.2 Explorer: The explorer plans through imagined latent trajectories to seek surprising states rather than revisiting previously discovered states.An ensemble estimates epistemic uncertainty through prediction disagreement, which defines the exploration reward.
- 2.2 Explorer: The explorer is periodically deployed to collect novel environment trajectories, expanding replay data for world-model and achiever training.LEXA alternates imagined explorer training with environment exploration and achiever practice using replay-buffer goals.
- 2.3 Achiever: The achiever receives a goal embedding and is trained in imagination and practice trials to reach diverse images encountered during exploration.Goals are sampled from replay-buffer images, allowing trajectories collected for one goal to support learning toward others.
- 2.3 Achiever: LEXA supports multi-object and temporally extended goal-reaching, including manipulating three RoboKitchen objects and picking and placing two RoboBins objects sequentially.These behaviors are illustrated as successful trajectories from test-set goal images.
- 2.4 Latent Distances: Temporal distance training uses imagined achiever rollouts and maximum-distance labels for goals sampled from different trajectories.The approach encourages preference for goals observed within the same trajectory while incorporating learning signals for far-away goals.
3 Experiments
LEXA is evaluated on prior benchmarks and a new 40-task benchmark spanning four robotic manipulation and locomotion domains. It achieves higher success than prior approaches across challenging tasks, while ablations and cross-environment experiments examine which components support performance and scalability.
- Experimental setup: The evaluation covers prior SkewFit, DISCERN, and Plan2Explore benchmarks plus a new benchmark spanning four domains.The new benchmark reports goal success rates averaged over test goals within each environment.
- New benchmark: LEXA achieves higher average success than all prior approaches on the new benchmark.The comparison evaluates a single image-trained agent on goal images without rewards.
- RoboBins: LEXA with temporal distance is the only approach making progress on all RoboBins goals, including tasks involving two blocks.Prior methods mainly solve reaching and simple pushing, whereas the hardest tasks involve manipulating two blocks.
- RoboKitchen: On RoboKitchen, temporal-distance LEXA learns multiple tasks, including some requiring sequential completion of two tasks, while prior methods barely progress.The gap between temporal and cosine distance is larger in RoboKitchen than in RoboBins.
- Single agent across environments: A single LEXA agent makes progress across RoboKitchen, RoboBins Reaching, RoboBins Pick & Place, and Walker, whereas DDL mainly solves Walker and RoboBin reaching.This experiment tests generalization across four visually distinct environments.
- Prior benchmarks: LEXA significantly outperforms prior work on previously used benchmarks and is often close to the optimal policy.On Plan2Explore tasks, it outperforms Plan2Explore on most tasks and performs comparably to oracle agents trained with true rewards.
- Ablations: A separate explorer is crucial for most tasks, while negative sampling and imagined training improve learning on the hardest tasks.Training temporal distance on negative samples speeds learning, and real-data training converges to slightly lower success than imagination-data training.
4 Related Work
Prior work learns goal-conditioned policies and generates training goals, but visual-goal methods struggle with meaningful exploration and novel-state discovery. Other approaches avoid or mitigate this challenge using behavior datasets or previously seen states.
- Model-free methods commonly learn a single goal-conditioned policy for reaching diverse goals.
- Visual-goal methods use contrastive or reconstructive representations but struggle to explore meaningfully without a clear reward signal.
- Some methods avoid exploration challenges with large datasets of interesting behaviors, while others generate goals similar to previously observed states.
5 Conclusion
LEXA is presented as an unsupervised agent that explores environments, learns to achieve discovered goals, and solves image-based tasks zero-shot. Its novelty-driven imagined exploration discovers meaningful behaviors across more diverse environments than prior work.
- LEXA explores environments and learns to achieve the goals it discovers without supervision.
- At test time, LEXA solves downstream tasks specified as images in a zero-shot way.
- Searching for novelty in imagination enables LEXA to discover meaningful behaviors in more diverse environments than prior work.
A Experimental Details
The experiments evaluate LEXA across robotic manipulation and locomotion benchmarks, including a new four-environment benchmark and prior goal-reaching tasks. Results indicate that exploration and distance choices matter across task types, with temporal distance helping on multi-object and fine-grained manipulation tasks.
- Benchmarks: The benchmark includes 12 Walker tasks, 12 Quadruped tasks, 8 RoboBin tasks, and 12 RoboKitchen tasks.
- Benchmark results: LEXA substantially outperforms competing agents on the RoboYoga Walker and Quadruped benchmarks, solving a majority of tasks.
- Distance functions: Temporal distance outperforms cosine distance on RoboBin tasks requiring manipulation of several blocks.
- Benchmark results: LEXA makes progress on four of six RoboKitchen base tasks, where most prior agents fail all tasks.
- Generalization: A single LEXA agent with temporal distance makes progress on tasks from Kitchen, RoboBin, and Walker, unlike LEXA+cosine and DDL on Kitchen tasks.
- Prior benchmarks: The evaluation also covers prior SkewFit, DISCERN, and DeepMind Control tasks, with comparisons to their reported baselines.