Source-linked AI summary
Planning to Explore via Self-Supervised World Models
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, Deepak Pathak
TL;DR
Task-specific reinforcement learning requires costly interaction and often does not generalize efficiently to unseen tasks. Plan2Explore uses reward-free planning through a learned world model to seek expected future novelty, then adapts to downstream tasks from high-dimensional images. It achieves state-of-the-art zero-shot performance, with few-shot performance eventually matching or outperforming supervised alternatives and zero-shot performance nearly matching a reward-informed oracle.
Problem
Task-specific reinforcement learning requires substantial task-specific interaction, creating a need for efficient zero-shot or few-shot generalization to unseen tasks.
Method
Plan2Explore learns a world model through reward-free self-supervised exploration, plans toward expected future novelty using ensemble disagreement, and trains downstream policies in imagination.
Results
Plan2Explore achieves state-of-the-art zero-shot performance across challenging visual control tasks and nearly matches the reward-informed supervised oracle, with few-shot performance eventually matching or outperforming it on some tasks.
Takeaways & Limitations
A single reward-free exploration phase can support solving multiple downstream tasks with zero or few task-specific interactions.
Takeaways & Limitations
The exploration objective is evaluated in challenging visual control settings and depends on learning an accurate world model from high-dimensional inputs.
Abstract
from arXiv · showhide
Reinforcement learning allows solving complex tasks, however, the learning tends to be task-specific and the sample efficiency remains a challenge. We present Plan2Explore, a self-supervised reinforcement learning agent that tackles both these challenges through a new approach to self-supervised exploration and fast adaptation to new tasks, which need not be known during exploration. During exploration, unlike prior methods which retrospectively compute the novelty of observations after the agent has already reached them, our agent acts efficiently by leveraging planning to seek out expected future novelty. After exploration, the agent quickly adapts to multiple downstream tasks in a zero or a few-shot manner. We evaluate on challenging control tasks from high-dimensional image inputs. Without any training supervision or task-specific interaction, Plan2Explore outperforms prior self-supervised exploration methods, and in fact, almost matches the performances oracle which has access to rewards. Videos and code at https://ramanans1.github.io/plan2explore/
1. Introduction
Plan2Explore addresses task-specific learning and inefficient exploration by learning a world model through reward-free planning, then adapting to unseen downstream tasks with zero or few task-specific interactions.
- Task-specific training commonly requires substantial environment interaction for each task, motivating rapid generalization to unseen tasks.
- Plan2Explore collects diverse experience once through self-supervised exploration, then uses downstream rewards to solve multiple tasks with limited or no further interaction.
- Prior model-free exploration methods compute novelty retrospectively after reaching observations, which can make exploration inefficient.
- Plan2Explore instead learns a world model and plans in latent space toward expected novelty of future situations.
- The method targets high-dimensional visual control and studies transfer, adaptation, generalization, and expected-future novelty against supervised and unsupervised alternatives.
- Its exploration objective approximates expected information gain with disagreement among an ensemble of one-step predictive models, emphasizing learnable uncertainty over stochastic uncertainty.
2. Control with Latent Dynamics Models
The control framework uses compact latent dynamics to model image-based environments and trains policies in imagination by propagating predicted long-term value through the learned world model.
- World models compress past experience into representations that support prediction of imagined future sequences.
- The latent model combines an image decoder with shared recurrent posterior and prior dynamics, while jointly trained components provide learning signals for the dynamics.
- Dreamer learns a latent-state value network and actor so the policy can maximize predicted long-term rewards inside the world model.
- The actor is trained on imagined latent sequences initialized from replay-buffer images, with gradients propagated through the neural-network dynamics model.
3. Planning to Explore
Plan2Explore learns a task-agnostic world model and uses imagined latent disagreement to plan toward informative, expected future states. The resulting model supports zero- or few-shot adaptation by optimizing downstream policies in imagination.
- Self-supervised exploration: During exploration, Plan2Explore trains a global world model and an exploration policy that seeks novel states using ensemble disagreement in imagined latent trajectories.The collected data expands the dataset and improves the world model iteratively.
- Task adaptation: After exploration, a single task-agnostic model can optimize policies for multiple downstream reward functions with no or limited additional interaction.Task adaptation distills the supplied reward into a latent reward predictor and trains the task policy in imagination.
- Latent disagreement: Ensemble disagreement estimates epistemic uncertainty by measuring variance across one-step models’ predicted next image embeddings.The ensemble members are trained alongside the world model and become more similar as visited states enter the training data.
- Zero-shot performance: On zero-shot evaluation, Plan2Explore achieves state-of-the-art performance across a range of tasks and remains competitive with supervised Dreamer.The evaluation uses task policies trained from exploration snapshots after rewards are specified at test time.
- Expected information gain: The exploration objective approximates expected information gain, targeting trajectories that efficiently reduce model uncertainty.This objective distinguishes reducible uncertainty from stochastic data noise.
4. Experimental Setup
The experiments evaluate visual-only continuous-control learning from 64 × 64 × 3 pixel observations across multiple exploration and supervised baselines. Few-shot adaptation is measured after 1M reward-free exploration steps and limited task data.
- Environment and protocol: Experiments use the DM Control Suite with visual observations of 64 × 64 × 3 pixels, 1000-step episodes, action repeat R = 2, and three random seeds.Shaded regions show standard deviation.
- Few-shot evaluation: Few-shot adaptation evaluates behavior after a 1M-step reward-free exploration phase followed by a small amount of task-specific data.The figure reports adaptation after up to 1000 exploration episodes and limited subsequent task interaction.
- Baselines: Plan2Explore is compared with supervised Dreamer, unsupervised Curiosity and MAX, and random data collection using a common base agent for fairness.Curiosity and MAX are adapted for the model-based visual setting.
5. Results and Analysis
Plan2Explore learns a self-supervised world model that supports zero-shot transfer and few-shot adaptation across tasks, while planning for expected novelty improves exploration.
- Zero-shot transfer: Plan2Explore trains a downstream agent in imagination using a self-supervised world model and no new environment interaction.The task is specified through a reward function used to train a reward predictor.
- Zero-shot transfer: Plan2Explore generally outperforms prior pixel-based exploration strategies and can be competitive with Dreamer in zero-shot performance.It even outperforms Dreamer on the hopper hop task.
- Zero-shot transfer: Plan2Explore almost matches Dreamer despite exploring without task rewards, whereas Dreamer receives task rewards during exploration.Dreamer is treated as an oracle because it has access to task rewards.
- Few-shot adaptation: During exploration, unsupervised methods may achieve task rewards coincidentally; after task information arrives, they switch to task-oriented behavior.The reported transition occurs after 1000 exploration episodes, followed by 150 episodes of task-oriented behavior.
- Generalization: On four cheetah tasks, Dreamer solves running forward but struggles on the other tasks, while Plan2Explore performs well across all tasks.The comparison is based on raw-pixel zero-shot performance without state-space input.
- Expected novelty: Planning expected future novelty gives Plan2Explore superior performance over a one-step retrospective-disagreement approach.Plan2Explore imagines future states beyond those already represented in the replay buffer.
6. Related Work
Related work spans intrinsic-motivation exploration, model-based control, self-supervised visual learning, and information-gain methods, while identifying gaps in scalable integrated visual exploration.
- Exploration: Traditional exploration methods based on visitation counts or Bayesian approaches are difficult to generalize to high-dimensional inputs such as images.These methods are effective in tabular settings but face scaling challenges for visual observations.
- Intrinsic motivation: Intrinsic-motivation methods explore without rewards by seeking prediction errors, influence over inputs, or rare states.The passage frames these as examples of exploration objectives for complex environments.
- Model-based control: Model-based reinforcement learning has improved data efficiency, especially when low-dimensional state information is available.Recent latent-dynamics work has also achieved competitive performance with model-free agents.
- Positioning: Plan2Explore connects these lines by using model-based self-supervised exploration with visual observations and few-shot task solving.The related-work context positions this combination as the paper’s distinctive setting.
- Visual model-based learning: Prior model-based exploration was not shown to scale to complex visual observations, while other visual approaches separated exploration from model learning.Some methods used only proprioceptive information or performed exploration and learning in stages.
- Information gain: Information gain provides a longstanding foundation for actively selecting informative data, including model-based reinforcement learning objectives.Earlier work optimized information gain in Bayesian and tabular settings.
7. Discussion
Plan2Explore learns an environment world model through unsupervised exploration and uses it for zero-shot or few-shot task solving, achieving strong control performance and data efficiency.
- Discussion: Plan2Explore learns a world model through unsupervised exploration and uses it to solve downstream tasks in zero-shot or few-shot settings.The method is framed as self-supervised reinforcement learning.
- Discussion: Plan2Explore achieves state-of-the-art zero-shot task performance on the DeepMind Control Suite.Its zero-shot performance is competitive with Dreamer on some tasks.
- Discussion: Few-shot performance eventually matches or outperforms the supervised Dreamer agent.The paper connects the method to expected information gain as a principled exploration objective.
- Discussion: The authors position the approach as a step toward scalable, data-efficient reinforcement learning for real-world systems.This is presented as the broader potential of learning behavior for many tasks from shared exploration.
A. Appendix
The appendix details the visual-control benchmark, implementation choices, and zero-shot/adaptation evaluations used to assess Plan2Explore.
- Generalization: The new Cheetah evaluation defines running backward, flipping forward, and flipping backward as downstream generalization tasks.These tasks test whether a model trained during task-agnostic exploration transfers beyond forward running.
- Environment: Experiments use DeepMind Control Suite tasks with only 64 × 64 visual observations, 1000-step episodes, and action repeat R = 2.The benchmark includes tasks with sparse rewards, high-dimensional action spaces, unstable equilibria, or long planning horizons.
- Implementation: Plan2Explore uses an ensemble of 5 one-step prediction models with two hidden-layer MLPs for latent disagreement.The models take the RSSM RNN-state and action as inputs and predict encoder features; disagreement is scaled by 10,000 for the final intrinsic reward.
- Evaluation: At 3.5 million environment steps, Table 1 reports zero-shot performance for Plan2Explore, Curiosity, Random, MAX, Retrospective, and Dreamer.Scores average the final 20 episodes and use the mode of the actor without action noise; highest no-task-reward performance is highlighted per task.
- Evaluation: Table 2 reports adaptation performance after 1M task-agnostic and 150K task-specific environment steps across the same agent comparisons.Scores average the final 20 episodes before 1.15M steps, with highest self-supervised performance highlighted for each task.
- Evaluation: Figure 6 evaluates zero-shot performance from raw pixels across all DM Control Suite tasks and reports state-of-the-art performance for Plan2Explore on a range of tasks.The figure also reports competitive performance with supervised Dreamer, using snapshots taken throughout exploration to train downstream task policies.