Learning Latent Dynamics for Planning from Pixels
TL;DR PlaNet learns a latent dynamics model from 64x64 images and picks every action by planning in that latent space, with no policy or value network. On six DeepMind control suite tasks it comes close to D4PG's final scores, and beats D4PG by 26% on cheetah, while using 1,000 episodes instead of 100,000.
What they do
Learned latent models had only planned simple pixel tasks, such as balancing a cartpole from dense rewards. PlaNet is a purely model-based agent. It learns the dynamics from images, actions and rewards, then chooses each action by planning in the model’s latent space with model-predictive control. It has no policy or value network.
How it works
- Recurrent state-space model (RSSM). The latent state has a deterministic part (a 200-unit GRU) and a stochastic part (a 30-dimensional diagonal Gaussian). It is trained like a sequential VAE: an encoder infers states from past images, and decoders reconstruct the image and the reward.
- Latent planning. At every step, the cross-entropy method searches 12-step action sequences: 1000 candidates, refit to the best 100, for 10 iterations. Rewards are predicted from latent states, so planning never generates images.
- Data collection. After 5 random episodes, the agent adds one planned episode with action noise every 100 updates.
- Latent overshooting. A generalized variational bound trains multi-step predictions in latent space. It helps several models, but the final RSSM agent does not need it.
Results
- On six control-suite tasks seen only through 64x64 images, PlaNet scores 821/832/662/700/930/951 (cartpole, reacher, cheetah, finger, cup, walker) after 1,000 episodes. D4PG scores 862/967/524/985/980/968 after 100,000 episodes.
- Within 100 episodes it beats A3C trained on proprioceptive states for 100,000 episodes, on every task. Training takes 10-20 h on one V100.
- Both model parts matter on every task, and without the stochastic part the agent does not learn. Online data collection is necessary for cartpole, finger and walker, and CEM beats picking the best of 1000 random sequences.
- A single agent trained on all six tasks, without being told which task it faces, solves all of them but learns more slowly.
Limitations
Results are in simulation only, and the action repeat (2-8) is tuned per task. Finger spin stays well below D4PG (700 vs 985). There is no value function, so rewards beyond the planning horizon are ignored, and the representation relies on image reconstruction. The authors suggest dropping reconstruction for visually diverse tasks.
Figures
Abstract
Planning has been very successful for control tasks with known environment dynamics. To leverage planning in unknown environments, the agent needs to learn the dynamics from interactions with the world. However, learning dynamics models that are accurate enough for planning has been a long-standing challenge, especially in image-based domains. We propose the Deep Planning Network (PlaNet), a purely model-based agent that learns the environment dynamics from images and chooses actions through fast online planning in latent space. To achieve high performance, the dynamics model must accurately predict the rewards ahead for multiple time steps. We approach this using a latent dynamics model with both deterministic and stochastic transition components. Moreover, we propose a multi-step variational inference objective that we name latent overshooting. Using only pixel observations, our agent solves continuous control tasks with contact dynamics, partial observability, and sparse rewards, which exceed the difficulty of tasks that were previously solved by planning with learned models. PlaNet uses substantially fewer episodes and reaches final performance close to and sometimes higher than strong model-free algorithms.
Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.