Source-linked AI summary
Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, Martin Riedmiller
TL;DR
Control from raw images requires jointly identifying useful state representations and nonlinear dynamics despite high-dimensional observations. E2C learns a variational-autoencoder-based latent model whose dynamics are locally linear, enabling stochastic optimal control and long-term image prediction. Across four challenging benchmarks, E2C found embeddings supporting control close to that achievable with the real system model.
Problem
Raw-image control is difficult because observations are high-dimensional and nonlinearly related to the underlying system dynamics, while the system model is unknown.
Method
E2C learns an unsupervised deep generative model with a low-dimensional latent representation and locally linear latent dynamics derived from stochastic optimal control.
Results
Across four challenging benchmarks, E2C found embeddings on which control was easy and achieved performance close to optimal control using the real system model.
Takeaways & Limitations
Constraining latent transitions to be locally linear makes locally optimal control applicable to image-based nonlinear systems while supporting long-term planning.
Takeaways & Limitations
The formulation assumes visual observations are fully observed depictions of the system state, although this assumption is later relaxed.
Abstract
from arXiv · showhide
We introduce Embed to Control (E2C), a method for model learning and control of non-linear dynamical systems from raw pixel images. E2C consists of a deep generative model, belonging to the family of variational autoencoders, that learns to generate image trajectories from a latent space in which the dynamics is constrained to be locally linear. Our model is derived directly from an optimal control formulation in latent space, supports long-term prediction of image sequences and exhibits strong performance on a variety of complex control problems.
1 Introduction
E2C addresses control from raw images by learning a low-dimensional latent space with locally linear dynamics, avoiding direct optimization in high-dimensional nonlinear sensory spaces.
- Raw-image stochastic optimal control is difficult because images are high-dimensional and their content depends nonlinearly on underlying dynamics.
- E2C learns a low-dimensional latent state space in which locally optimal control can be performed robustly and easily.
- E2C is an unsupervised variational-autoencoder-based generative model that represents viable sensory-space trajectories and supports long-term latent-space planning.
- The method is evaluated on four challenging raw-image control tasks and compared with unsupervised representation-learning methods.
2 The Embed to Control (E2C) model
E2C combines inference, locally linear latent transitions, and image generation so stochastic optimal control can operate on a learned low-dimensional representation.
- 2.1 Problem Formulation: The problem is to control an unknown stochastic dynamical system when only visual depictions of its state are available.
- 2.1 Problem Formulation: E2C seeks a low-dimensional mapping from images xt to latent states zt, with nz ≪ nx, while retaining information needed for control.
- 2.2 Stochastic locally optimal control in latent spaces: The latent dynamics are locally linearized around reference trajectories, producing matrices A and B, an offset, and Gaussian transition noise.
- 2.2 Stochastic locally optimal control in latent spaces: Quadratic latent-state and control costs yield a local linear-quadratic-Gaussian problem solvable by iLQR or AICO.
- 2.3 A locally linear latent state space model for dynamical systems: The inference network represents latent states with a diagonal Gaussian distribution and uses the reparameterization trick for gradient-based learning.
- 2.3 A locally linear latent state space model for dynamical systems: The generative model samples latent transitions in the required linear form and reconstructs images through a Bernoulli decoder.
- 2.3 A locally linear latent state space model for dynamical systems: The transformation network predicts local transition matrices and offsets, using a rank-one perturbation of identity to reduce parameters for A.
- 2.3 A locally linear latent state space model for dynamical systems: E2C enforces similarity between predicted next-state samples and encoded observations so multi-step latent predictions remain valid image trajectories.
3 Experimental Results
E2C was evaluated across four visual control tasks and compared with model variants and baselines. Its latent dynamics constraints produced useful embeddings, accurate long-term predictions, and stable control across the tested systems.
- Experimental setup: E2C was evaluated on planar navigation, visual inverted-pendulum swing-up, cart-pole balancing, and three-link robot-arm control.The experiments used both fully connected and deep convolutional encoder-decoder architectures for different image sizes.
- Model comparisons: The study compared E2C with globally linear and nonlinear latent-dynamics variants, plus standard variational and deep autoencoder baselines.The variants differ in whether local transformation parameters are globally linear or derived by linearizing a nonlinear latent transition model.
- Planar control: Including latent dynamics constraints yielded planar embeddings approaching the optimal state-space structure, unlike separately trained autoencoders.The autoencoders produced visually pleasing reconstructions but failed to recover the underlying state-space structure needed for dynamics estimation and distance-based costs.
- Planar control: Latent-space plans from globally or locally linear E2C matched real-system planar performance, whereas other models failed on long-term prediction.The evaluation accumulated latent and real trajectory costs after executing 40 pre-computed actions from random starting positions.
- Inverted pendulum: E2C was the only compared model to produce stable swing-up and balance behavior for the visual inverted pendulum.Although E2C was not best on reconstruction performance, it recovered a meaningful latent structure separating positions and velocities.
- Cart-pole and robot arm: The cart-pole and robot-arm trajectory costs were 11.13 and 85.12, versus 7.28 and 60.74 for AICO using real system dynamics.These E2C trajectories were described as only slightly worse than the corresponding real-dynamics trajectories.
4 Comparison to recent work
The paper situates E2C among representation-learning and image-control methods. It distinguishes E2C by constraining latent dynamics for long-term prediction and local linearizability.
- Related approaches: Prior work includes deep autoencoders for control, model-free deep Q-learning for Atari, kernel and deep policy learning, and nonlinear latent dynamics models.These approaches span representation learning, model-free control, and learned forward-dynamics methods for image-based control.
- Comparison to recent work: Compared with a closely related autoencoder approach, E2C does not require PCA preprocessing and explicitly targets non-diverging, linearizable latent predictions.The cited alternative learns a nonlinear forward-dynamics model jointly with its autoencoder.
- Connections to latent representation learning: E2C is also related to variational autoencoders and prior methods that enforce desired transformations in latent space during learning.The paper places its model within the VAE family while emphasizing latent-space structure for modeling.
5 Conclusion
E2C learns a latent dynamics model constrained to locally linear state transitions, enabling stochastic optimal control from high-dimensional image streams. Its derivation also addresses efficient KL-divergence computation for the transition model.
- Conclusion: E2C extracts a latent dynamics model with locally linear state transitions for stochastic optimal control on high-dimensional image streams.On four challenging benchmarks, control performance was close to that achievable with optimal control on the real system model.
- State transition matrix factorization: The full transition matrix A_t requires n_z × n_z predicted parameters and matrix inversion for the KL-divergence penalty.These requirements make symbolic differentiation and backpropagation inconvenient.
- State transition matrix factorization: Representing A_t as I + v_t r_t^T reduces the transition network’s output burden to estimating v_t and r_t.This reformulation avoids estimating every element of the full matrix.
- KL divergence: The KL-divergence derivation is designed to make partial derivatives efficiently computable in symbolic form and convenient for batched processing.The resulting expressions avoid numerical trace or determinant computations and excessive intermediate tensor products.
- KL divergence: The simplified final KL expression is easy to implement and requires summation only over the non-batch dimension.The difference-of-means term can be derived with the same summation scheme.
B Supplementary to the experimental setup
The supplementary setup uses convolutional inference and up-convolutional decoder networks for high-resolution image tasks. A slowness-regularized VAE slightly improves latent coherence, but slowness alone does not make locally linear control feasible.
- Network architecture: Up-convolutional decoder networks replace fully connected networks for generating high-dimensional images.The decoder mirrors the encoder with up-sampling and convolution steps.
- Network architecture: The decoder architecture mirrors each encoder convolution and pooling stage with a corresponding up-sampling and convolution stage.The reported upsampling strategy is simple perforated upsampling.
- Slowness variant: A slowness VAE variant adds a KL-divergence term that enforces similarity between encodings of temporally close images.Temporal slowness is motivated as a proxy for representation learning in reinforcement learning and videos.
- Slowness variant: A slowness term alone slightly improves coherence of similar latent states but does not sufficiently structure the latent space for locally linear prediction and control.The limitation is reported from the experiments.
B.3 Evaluation criteria
The evaluation separates reconstruction, next-state decoding, latent trajectory costs, and real trajectory costs. Real trajectory cost is the decisive optimality criterion, with true-system control providing an upper bound.
- Evaluation criteria: The evaluation measures autoencoding, next-state decoding, latent trajectory costs, and real trajectory costs.These criteria assess reconstruction, predictive consistency, latent planning, and executed-system performance.
- Evaluation criteria: Reconstruction and next-state decoding are assessed using mean log loss for Bernoulli image distributions.For Bernoulli distributions, mean log loss is the cross-entropy error function.
- Evaluation criteria: Real trajectory cost is the optimality criterion because it tests whether latent action sequences reflect true state trajectory costs.The same evaluation cost matrices are used across models, though they may differ from optimization matrices.
- Evaluation criteria: Latent and real trajectory costs are accumulated from sampled starting states according to the stochastic optimal control method.The real cost evaluates the executed system states using quadratic costs on deviation from a goal state.
- Evaluation criteria: Applying iLQR/AICO to the known simulated system dynamics provides an upper bound on control performance achievable by the learned models.All experiments were performed in simulation, making the true dynamics model available.
B.6.1 Experimental Setup
The experiments use preconstructed transition datasets with separate training, validation, and test splits. E2C trains directly on transitions, while image-only baselines receive extracted image data and separately learned latent dynamics.
- Dataset construction: All datasets are created in advance as transition tuples and divided into training, validation, and test splits.The transition dataset is D = {(x_1, u_1, x_2), ..., (x_T−1, u_T−1, x_T)}.
- Dataset construction: E2C models train on the full transition dataset D, whereas autoencoder and VAE baselines train on extracted images.The slowness VAE trains on consecutive image pairs.
- Dynamics training: For image-only autoencoders, latent representations are combined with actions to form dynamics-training tuples.The dynamics MLPs are trained on these low-dimensional representations so all methods use the same underlying data.
B.6.2 Implementation details
The experiments use stabilized neural-network training and latent-space visualizations to assess multi-step predictions across planar, pendulum, and inverted-pendulum systems.
- Training: Orthogonal initialization and Adam were important for stabilizing training and achieving good reconstructions across methods.The authors also report that both techniques reduced the hyperparameter search required.
- Trajectory prediction: Multi-step predictions are generated by encoding a starting state, applying actions to its latent representation, decoding each predicted latent state, and visualizing the resulting images.This procedure is used for the planar system and inverted pendulum.
- Pendulum experiments: For passive pendulum dynamics, the models predict trajectories from an initial angle θ = −π without applied force.The passive setting requires the models to predict the dynamics with no force input.
- Pendulum experiments: For controlled pendulum dynamics, full force is applied rightward for 6 timesteps and leftward for 4 timesteps from an initial angle θ = π.The trajectory depiction evaluates predictions under this prescribed force sequence.
- Latent representations: The inverted-pendulum models are compared through their learned latent spaces, with the pendulum depictions encoded into a 3-dimensional representation.Figure 8 presents latent spaces for the baseline models and E2C variants.
- Cart-pole predictions: For cart-pole, dreamed trajectories include uncontrolled and controlled dynamics, with long-term predictions generated by following latent-space dynamics without additional input.The uncontrolled case uses u = 0 for all steps.
C.4 Exemplary trajectory taken for three-link arm task
Figure 10 presents a 10-step controlled trajectory executed by E2C with model predictive control, alongside one-step-ahead predictions.
- Trajectory execution: E2C executes a controlled three-link-arm trajectory while combined with model predictive control.The figure shows an executed trajectory segment rather than a long-term prediction.
- Trajectory execution: The visualization covers 10 trajectory steps and includes one-step-ahead predictions.Additional visualizations and controlled trajectories are provided in the supplementary video.
C.5 Comparison of different models for cart-pole and robot arm
The cart-pole and three-link-arm comparison evaluates real trajectory cost and task success across models, while additional experiments compare AICO with iLQR.
- Evaluation setup: Table 2 averages real trajectory cost and task success over 30 starting states with a fixed goal state.The cart-pole begins near its goal with Gaussian noise, while the three-link arm begins from a random configuration.
- Model comparison: Only E2C and its non-linear variant successfully perform the three-link-arm task, with a large performance gap between them.The authors attribute the weaker alternatives to accumulated error from linearizing nonlinear dynamics after training.
- Trajectory cases: The cart-pole evaluation includes uncontrolled and controlled trajectories, with the red image marking the initial configuration and columns depicting different motion cases.The controlled cases include moving left and moving right.
- Model comparison: The standard autoencoder, variational autoencoder, and Global E2C are omitted from Table 2 because they failed on the task.Their performance was similar to the VAE with slowness.
- Trajectory optimizers: AICO and iLQR perform about the same on the cart-pole and three-link-arm tasks.This indicates that the covariance matrices estimated by E2C do not hurt planning, but considering them does not improve performance either.