[Paper Review] Learning Latent Dynamics for Planning from Pixels

Danijar Hafner, Timothy Lillicrap+5 more|arXiv (Cornell University)|Nov 12, 2018

Multimodal Machine Learning Applications367 citations

TL;DR

PlaNet is a model-based agent that learns latent environment dynamics from pixel observations and plans in latent space using online model-predictive control, achieving strong performance with far fewer environment interactions compared to model-free methods.

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.

Motivation & Objective

• Motivate planning in unknown environments by learning dynamics from pixel observations.
• Develop a latent dynamics model with both deterministic and stochastic components for robust long-horizon predictions.
• Introduce latent overshooting to improve multi-step predictive accuracy in latent space.
• Demonstrate planning in latent space to achieve strong performance on image-based continuous control tasks.
• Show data efficiency gains over model-free baselines while maintaining competitive computation time.

Proposed method

• Use a recurrent state-space model (RSSM) that combines deterministic and stochastic latent transitions to model environment dynamics from pixels.
• Train an encoder to infer latent beliefs over states from past observations and actions, enabling filtering in a POMDP setting.
• Apply a model-predictive control (MPC) planner based on cross-entropy method (CEM) to search over action sequences in latent space without rendering images.
• Optimize a variational bound extended with latent overshooting to encourage accurate multi-step predictions in latent space (aggregate over distances 1..D).
• Collect data online by planning with the current model and executing the first action, then re-plan at each step (online learning loop).
• Evaluate PlaNet on six pixel-based continuous control tasks from the DeepMind Control Suite, comparing against model-free baselines and ablations.

Experimental results

Research questions

• RQ1 Can planning in a learned latent space from pixel observations achieve competitive performance on continuous control tasks with partial observability?
• RQ2 Does a latent dynamics model with both deterministic and stochastic components improve planning performance compared to purely deterministic or purely stochastic models?
• RQ3 Does latent overshooting improve multi-step predictive accuracy and final task performance under planning?
• RQ4 How data-efficient is PlaNet relative to state-of-the-art model-free methods when trained from pixels?

Key findings

• PlaNet achieves strong final performance on six image-based continuous control tasks, sometimes surpassing strong model-free methods.
• PlaNet with RSSM (deterministic + stochastic transitions) outperforms purely deterministic and purely stochastic variants across tasks.
• Latent overshooting improves long-term predictions and planning effectiveness, though RSSM alone provides substantial gains on many tasks.
• PlaNet reaches final performance close to or higher than top model-free methods with about 200x fewer environment interactions in the reported comparisons.
• Online planning with MPC (CEM) in latent space provides robust planning performance and benefits from iterative refinement over planning horizons.
• PLA Net with latent Dynamics trained from pixels can solve tasks with partial observability and sparse rewards that were previously challenging for planning-based approaches.