PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Latent Dynamics for Planning from Pixels

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, James Davidson

arXiv:1811.04551v5cs.LGcs.AIstat.ML

TL;DR

PlaNet tackles the challenge of learning dynamics accurate enough for planning from image observations in unknown environments. It learns a latent model with deterministic and stochastic transition components and uses latent-space online planning, augmented by multi-step latent overshooting. Across difficult pixel-based continuous-control tasks, it achieves performance comparable to strong model-free methods while using 200× fewer episodes and similar or less computation time.

  • Problem

    Learning dynamics models accurate enough for planning remains a long-standing challenge, especially in image-based domains with accumulating multi-step errors and multiple possible futures.

  • Method

    PlaNet learns environment dynamics from pixels in a compact latent space, combines deterministic and stochastic transition components, and uses latent overshooting for multi-step variational prediction.

  • Results

    PlaNet succeeds on difficult image-based continuous-control tasks and reaches performance comparable to strong model-free algorithms while using 200× fewer episodes and similar or less computation time.

  • Takeaways & Limitations

    Planning with learned latent dynamics can support continuous control from pixels across contact dynamics, partial observability, and sparse rewards.

  • Takeaways & Limitations

    Future work is needed for temporal abstraction, longer-horizon value estimation, more efficient planning, and higher visual diversity.

Abstract

from arXiv · show

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.

1. Introduction

PlaNet addresses the difficulty of learning planning-accurate dynamics from images by combining latent-space planning with a dynamics model that captures stochasticity, determinism, and multi-step predictions. It targets challenging continuous-control settings and reports substantially lower interaction requirements than model-free baselines.

  • Learning dynamics accurate enough for planning remains difficult because errors accumulate across predictions, multiple futures may be missed, and models can become overconfident out of distribution.
  • Model-based planning can improve data efficiency, scale performance with search computation, and potentially transfer dynamics across tasks.
  • PlaNet learns environment dynamics from pixels and selects actions through online planning in a compact latent space.
  • PlaNet evaluates image-based control domains involving partial observability, sparse rewards, contact dynamics, larger state spaces, and difficult balance interactions.The experiments include cartpole swingup, reacher, cheetah, finger spinning, cup, and walker tasks.
  • 200× fewer environment interactions are reported than for model-free A3C, with similar computation time and higher final performance in some cases.
  • The latent dynamics model combines deterministic and stochastic components, while latent overshooting extends variational training to multi-step predictions.Latent overshooting is described as a fast latent-space regularizer that can improve long-term predictions and work with latent sequence models.

2. Latent Space Planning

PlaNet learns a latent model from collected episodes and uses online planning to select actions under partial observability. Its planner searches candidate action sequences with the cross entropy method, evaluates predicted latent rewards, and replans after observations.

  • PlaNet iteratively collects experience by planning with a partially trained model and retrains that model on the accumulated episodes.The procedure begins with random seed episodes, then adds episodes at fixed update intervals with exploration noise.
  • The task is modeled as a partially observable Markov decision process because individual image observations generally do not reveal the full environment state.
  • The model learns transition, observation, and reward distributions, plus an encoder that infers the current hidden-state belief from observation and action history.The observation model supplies a rich training signal but is not used directly for planning.
  • The cross entropy method searches for action sequences by sampling candidates, evaluating them under the learned model, and refitting a Gaussian distribution to the top candidates.
  • The planner returns the mean action for the current time step after repeated sampling and refitting, then resets its action-sequence belief after the next observation.
  • Candidate sequences are evaluated by summing predicted mean rewards along a sampled latent-state trajectory, allowing fast batch evaluation without generating images.

3. Recurrent State Space Model

The recurrent state-space model combines stochastic and deterministic latent components to support multi-step prediction from image observations. It uses an encoder to infer hidden states and is designed for efficient latent-space planning.

  • Model design: RSSM predicts forward purely in latent space, enabling efficient evaluation of many action sequences during planning.The model is described as a nonlinear Kalman filter or sequential VAE.
  • Model design: The latent state-space model represents image and reward generation through a hidden state sequence with learned transition, observation, and reward models.The transition, observation, and reward distributions are parameterized by neural networks.
  • Inference: The filtering encoder infers approximate state posteriors from past observations and actions for parameter learning and planning.The encoder uses a diagonal Gaussian whose parameters are produced by convolutional and feed-forward networks.
  • Model design: RSSM splits each state into stochastic and deterministic parts, combining multiple possible futures with reliable information retention across time.The deterministic sequence is implemented through an RNN, while the stochastic component is sampled through the encoder.
  • Training motivation: The standard objective motivates a new training objective because it does not directly train transition chains for multi-step predictions.The paper states that the next section proposes a generalization intended to improve long-term predictions.

4. Latent Overshooting

The standard variational objective trains latent transitions mainly through one-step predictions, which can be insufficient when planning requires accurate long-horizon forecasts. Latent overshooting generalizes the objective by training multi-step latent predictions toward corresponding posterior beliefs without decoding every predicted observation.

  • Motivation: The standard variational objective trains the stochastic transition path only through one-step KL regularizers, limiting direct learning of multi-step prediction chains.This limitation matters because planning requires accurate multi-step predictions.
  • Latent overshooting: The objective trains multi-step prior beliefs toward corresponding posterior beliefs in latent space, avoiding the expense of decoding every multi-step image prediction.The method is presented as a fast latent-space regularizer that encourages consistency between one-step and multi-step predictions.
  • Scope: The final RSSM-based agent does not require the additional benefit observed for several other dynamics models from latent overshooting.This qualification is reported in the appendix-related discussion.
  • Multi-step prediction: Multi-step predictions are computed by repeatedly applying the transition model while integrating out intermediate latent states.The one-step case, d = 1, recovers the original transition model.
  • Latent overshooting: Latent overshooting generalizes the standard variational bound to train predictions across all distances up to the planning horizon.A fixed-distance objective is insufficient when planning needs accurate predictions for multiple distances.

5. Experiments

PlaNet is evaluated on six image-based continuous-control tasks, where its latent planning, recurrent dynamics design, and online data collection are tested against model-free methods and alternative agent designs.

  • Evaluation: PlaNet evaluates six continuous-control tasks from pixel observations with qualitatively different challenges, including long horizons, partial observability, sparse rewards, and contact dynamics.The tasks use third-person camera images, and the task suite includes cartpole swingup, reacher, finger spinning, and cheetah.
  • Comparison to model-free methods: Within 100 episodes, PlaNet outperforms A3C on all tasks; after 500 episodes, it matches D4PG except on the finger task.A3C was trained from proprioceptive states for 100,000 episodes, while D4PG was trained from images for 100,000 episodes.
  • Comparison to model-free methods: PlaNet surpasses D4PG’s final performance by 26% on cheetah running.The comparison is reported as a relative improvement in final performance.
  • Model designs: The recurrent state-space model’s deterministic and stochastic components are both important, with the stochastic component necessary for learning across the evaluated tasks.The deterministic path supports memory over many time steps, while the experiments report that the agent does not learn without the stochastic component.
  • Agent designs: Online data collection helps on all tasks and is necessary for cartpole, finger, and walker, while iterative CEM search improves performance on every task.These findings compare PlaNet with random collection and random shooting that does not iteratively refine action sequences.
  • One agent all tasks: A single agent solves all six tasks without being told which task it faces, although it learns more slowly than agents trained individually.The agent infers the task from image observations.

6. Related Work

PlaNet is situated among model-based methods that learn dynamics for planning, with an emphasis on pixel-based latent planning and comparisons against alternative models and planning designs.

  • Prior model-based directions: Earlier model-based approaches commonly used low-dimensional state spaces or combined model-based and model-free learning rather than focusing on the same pixel-based latent-planning setting.The related-work discussion lists prior directions including low-dimensional planning, hybrid methods, and model learning without planning.
  • Model comparisons: Figure 4 compares PlaNet’s RSSM with purely deterministic GRU and purely stochastic SSM designs, reporting medians and 5–95 percentile areas over five seeds and ten trajectories.The figure also compares model-free algorithms and notes that sparse rewards produce large shaded areas on two tasks.
  • Benchmark comparisons: Table 1 compares PlaNet with A3C and D4PG using mean final performance over five seeds and ten trajectories, and includes true-dynamics CEM as an estimated performance upper bound.The D4PG episode count needed to reach PlaNet’s final performance is estimated from training curves.
  • Pixel-based planning: Prior pixel-based planning methods often addressed simple tasks, whereas PlaNet targets larger state and action spaces, longer horizons, partial observability, contact dynamics, and sparse rewards.The paper contrasts its simulated latent-planning setting with video-prediction work on simple gripper tasks and local-linear latent models for cartpoles and two-link arms.
  • Agent comparisons: Figure 5 compares PlaNet with random collection and random shooting, using medians and 5–95 percentile areas over five seeds and ten trajectories.Random shooting selects the best of 1000 action sequences without iterative CEM refinement.

7. Discussion

PlaNet demonstrates that latent dynamics learned from images can support fast planning on challenging continuous-control tasks, while the paper identifies several directions for extending the approach.

  • Contribution: PlaNet learns latent dynamics from image observations and chooses actions through fast planning in latent space.The design uses both stochastic and deterministic paths to support accurate long-term predictions.
  • Results: PlaNet reaches performance comparable to the best model-free algorithms while using 200× fewer episodes and similar or less computation time.This conclusion summarizes the reported cross-task performance and interaction-efficiency comparison.
  • Conclusion: The results support latent dynamics models for planning in image domains as a promising approach.This is the paper’s stated conclusion about the demonstrated setting.
  • Future work: Future work includes replacing fixed action repeat with temporal abstraction, improving beyond-horizon rewards with a value function, and increasing planning efficiency or visual robustness.The paper also identifies multi-task control through shared dynamics as a starting point for future work.

A. Hyper Parameters

PlaNet uses recurrent latent dynamics, variational image modeling, and cross-entropy-method planning with fixed task-specific action repeats and specified optimization settings.

  • Network architecture: The dynamics model uses a convolutional encoder and decoder, a 200-unit GRU deterministic path, two-layer ReLU networks, and 30-dimensional diagonal-Gaussian latent distributions.The listed fully connected layers have size 200.
  • Optimization: Training uses 5-bit image preprocessing, Adam with learning rate 10^-3, epsilon 10^-4, gradient clipping at 1000, and batches of 50 chunks of length 50.The KL divergence is clipped below 3 free nats rather than relatively scaled to reconstruction terms.
  • Planning: Planning uses CEM with horizon H = 12, I = 10 optimization iterations, J = 1000 candidates, and K = 100 top candidates for refitting.The agent starts with five random-action seed episodes and collects another episode every 100 update steps with action noise.
  • Latent planning with CEM: Latent planning initializes a factorized Gaussian belief over action sequences, evaluates candidates through the current belief, refits to the best candidates, and returns the first action mean.The algorithm’s inputs include horizon, optimization iterations, candidate count, and the number of top candidates used for refitting.

C. Multi-Task Learning

PlaNet can learn multiple tasks with one agent from image observations, although joint training slows learning relative to task-specific agents.

  • Multi-task performance: A single PlaNet agent learns to solve all six tasks while learning more slowly than individually trained agents.Performance is aggregated across six tasks, five seeds, and ten trajectories.
  • Task inference: A task-agnostic agent infers which task it is solving from image observations and solves the tasks with only a moderate learning slowdown.Results use medians and 5th–95th percentile areas over four seeds and ten trajectories.

D. Latent Overshooting

Latent overshooting affects model performance differently across architectures, substantially helping DRNN while slightly reducing RSSM performance; smooth activations benefit purely stochastic models, whereas RSSM is robust to activation choice.

  • Latent overshooting: Latent overshooting substantially improves DRNN performance but slightly reduces performance of the proposed RSSM.The comparison uses medians and 5th–95th percentile areas over five seeds and ten trajectories.
  • Activation functions: Smooth ELU activations improve purely stochastic and deterministic models, while RSSM performance remains robust to the activation function.The deterministic-model result is reported as not shown.

F. Bound Derivations

The derivations establish variational bounds for one-step and multi-step latent predictive distributions, while the accompanying diagnostics examine long-horizon latent and pixel predictions and planning settings.

  • One-step predictive distribution: The one-step latent-dynamics objective combines observation likelihood, transition likelihood, and variational-posterior terms.The bound is estimated without bias by replacing expectations with sample averages.
  • Multi-step predictive distribution: The multi-step predictive distribution replaces one-step transitions with d-step priors whose recursion averages over intermediate latent states.The variational bound contains the observation term, multi-step prior, and posterior term.
  • Relation between bounds: The data-processing inequality gives E[ln p_d(o_1:T)] ≤ E[ln p(o_1:T)], so bounds on multi-step predictions also bound the one-step predictive distribution.This relation is stated for d ≥ 1 under a Markovian latent-state sequence.
  • Prediction diagnostics: Pixel-accurate open-loop predictions extend 50 steps into the future in the cheetah environment, and latent diagnostics predict positions, velocities, and rewards beyond planning horizons.The state diagnostics use frozen dynamics models and learned predictors of simulator quantities.
  • Planning settings: Planning performance ranges from 132 to 837; evaluating more action sequences and iterations helps, while horizons of 6 or much longer perform poorly.The figure attributes long-horizon degradation to the increased search space.
Loading 1811.04551v5…