PaperlayerView exact revision on arXiv

Source-linked AI summary

Offline Reinforcement Learning from Images with Latent Space Models

Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, Chelsea Finn

arXiv:2012.11547v1cs.LGcs.AIcs.RO

TL;DR

Offline RL needs to learn from static data, but extending it to image observations is difficult because model uncertainty is hard to quantify in pixel space. LOMPO learns latent-state dynamics and penalizes uncertainty in the latent space. Across simulated visual control and a real-robot drawer task, it outperforms or matches prior methods, while relying on heuristic uncertainty estimates in practice.

  • Problem

    Offline RL has mostly studied compact states, leaving the challenge of learning safely from high-dimensional image observations where model uncertainty is difficult to quantify.

  • Method

    LOMPO learns a variational image-based latent dynamics model and trains policies in an uncertainty-penalized latent MDP using ensemble disagreement as the uncertainty estimate.

  • Results

    LOMPO empirically outperforms or matches prior model-based and model-free methods across four simulated tasks and one real-world robotic manipulation task.

  • Takeaways & Limitations

    Latent-space uncertainty makes pessimistic model-based offline RL tractable for image-based locomotion and manipulation, including a real-robot drawer-closing task.

  • Takeaways & Limitations

    The practical method cannot access an uncertainty oracle and instead estimates latent model uncertainty with heuristics such as ensemble disagreement.

Abstract

from arXiv · show

Offline reinforcement learning (RL) refers to the problem of learning policies from a static dataset of environment interactions. Offline RL enables extensive use and re-use of historical datasets, while also alleviating safety concerns associated with online exploration, thereby expanding the real-world applicability of RL. Most prior work in offline RL has focused on tasks with compact state representations. However, the ability to learn directly from rich observation spaces like images is critical for real-world applications such as robotics. In this work, we build on recent advances in model-based algorithms for offline RL, and extend them to high-dimensional visual observation spaces. Model-based offline RL algorithms have achieved state of the art results in state based tasks and have strong theoretical guarantees. However, they rely crucially on the ability to quantify uncertainty in the model predictions, which is particularly challenging with image observations. To overcome this challenge, we propose to learn a latent-state dynamics model, and represent the uncertainty in the latent space. Our approach is both tractable in practice and corresponds to maximizing a lower bound of the ELBO in the unknown POMDP. In experiments on a range of challenging image-based locomotion and manipulation tasks, we find that our algorithm significantly outperforms previous offline model-free RL methods as well as state-of-the-art online visual model-based RL methods. Moreover, we also find that our approach excels on an image-based drawer closing task on a real robot using a pre-existing dataset. All results including videos can be found online at https://sites.google.com/view/lompo/ .

1. Introduction

LOMPO targets offline visual control because online learning from images is costly and unsafe, while prior offline RL largely assumes compact states. It learns visuomotor policies from static image datasets using latent dynamics and uncertainty-aware policy optimization.

  • Interactive learning of control policies from vision has prohibitively high sample complexity, while real-world physical exploration raises safety challenges.These constraints motivate approaches that learn from existing offline datasets without additional environment interaction.
  • Offline RL uses static pre-collected data to reduce online-exploration safety concerns, and existing datasets often contain image or video observations.Prior offline RL work has focused mainly on compact state representations rather than pixel-based control.
  • LOMPO enables learning visuomotor policies from offline datasets across four simulated tasks and one real-world robotic manipulation task.Its uncertainty-penalized latent MDP is designed to regularize policy learning against distributional shift and model exploitation.
  • LOMPO learns an image encoder, decoder, and ensemble of latent dynamics models, then penalizes latent-state uncertainty estimated from ensemble disagreement.The resulting policy is optimized in the learned latent space using the offline dataset.

2. Related Work

Related work spans offline RL methods that regularize several algorithmic families and pixel-based control methods that learn from images either end-to-end or through visual representations. These lines of work motivate combining offline learning with high-dimensional visual control.

  • The paper sits at the intersection of offline RL and control from high-dimensional inputs such as images.
  • Offline RL: Offline RL research has developed regularized importance-sampling, actor-critic, approximate dynamic programming, and model-based algorithms.
  • Control from Pixels: Pixel-based control methods learn policies directly from images or on top of unsupervised visual representations, making high-dimensional control more practical for robotics.

3. Preliminaries

The preliminaries formulate visual control as offline RL in a POMDP and connect model-based pessimism to lower-bound optimization. They establish why uncertainty-penalized models are useful when no further environment interaction is available.

  • POMDPs: A visual POMDP separates observations, latent states, actions, transitions, rewards, initialization, and discounting, with the policy maximizing discounted expected return.
  • Control as inference: Control as inference represents optimality with binary variables and relates maximum-entropy RL to approximate inference over trajectories.In the POMDP extension, the ELBO combines inference, latent dynamics, and policy terms, with rewards represented through optimality likelihoods.
  • Offline RL: Offline RL learns a policy from a fixed dataset of trajectories containing high-dimensional observations, actions, and rewards, with no additional environment interaction.The dataset is generated by a behavior policy that may mix multiple policies.
  • Model-based offline RL: Model-based offline RL constructs an uncertainty-penalized MDP from learned dynamics and rewards, and optimizing it corresponds to a lower bound on return in the true MDP.This pessimistic approach is effective in low-dimensional settings but is difficult to scale to image observations.

4. LOMPO: Latent Offline Model-Based Policy Optimization

LOMPO addresses unreliable visual model rollouts by learning latent-state dynamics and penalizing uncertainty in a compact latent space. It trains the latent model and policy through an uncertainty-penalized POMDP objective with practical ensemble-based uncertainty estimation.

  • 4. LOMPO: Latent Offline Model-Based Policy Optimization: LOMPO learns a latent dynamics model and estimates model uncertainty in compact latent space to support offline visual model-based RL.This avoids directly applying computationally expensive visual dynamics ensembles while addressing model exploitation from out-of-distribution rollouts.
  • 4.1. Quantifying Model Uncertainty in the Latent Space: The method constructs an uncertainty-penalized latent MDP with modified reward er(s, a) = r(s, a) − λu(s, a), then induces a corresponding penalized POMDP.The uncertainty term is intended to reduce the impact of unreliable model predictions during offline policy optimization.
  • 4.1. Quantifying Model Uncertainty in the Latent Space: The uncertainty-penalized latent-MDP return is a lower bound on the true latent-MDP return, with the learned-policy gap depending on latent dynamics-model error at the optimal policy.The guarantee follows under the assumption that the latent state space forms an MDP.
  • 4.2. Latent Model Training and Policy Optimization with Uncertainty-Penalized ELBO: LOMPO maximizes an ELBO in the uncertainty-penalized POMDP that lower-bounds the ELBO of the original POMDP.Latent model training uses reconstruction and consistency terms, while policy optimization uses an uncertainty-penalized reward analogous to MOPO.
  • 4.3. Practical Implementation of LOMPO: The practical implementation trains an ensemble of latent transition models, samples one forward model per rollout step, and uses ensemble disagreement as the uncertainty penalty.The ensemble is grounded in a shared inference-induced latent space, and disagreement is estimated from the variance of ensemble log-likelihoods.

5. Experiments

Experiments evaluate LOMPO on simulated image-based robotics tasks and a real drawer-closing task using offline data. LOMPO generally outperforms model-free and model-based baselines, including on the real robot.

  • Experimental Scope: The experiments target offline visual control questions spanning realistic robot dynamics, baseline comparisons, dataset effects, and real-robot deployment.The simulated benchmark was designed because existing offline RL benchmarks largely lack image-based robot-control problems.
  • Simulated Experiments: LOMPO consistently outperforms LMBRL, offline SLAC, CQL, and behavioral cloning in almost all simulated Door Open settings.Scores are undiscounted average returns normalized to roughly 0–100, with 0 representing a random policy and 100 an expert.
  • Simulated Experiments: Across four high-fidelity simulated environments, LOMPO achieves high scores from raw observations and is the only model-based method succeeding on several environments.It outperforms other model-based methods across the board and generally exceeds CQL and behavioral cloning, except on the expert Door Open dataset.
  • Real Robot Experiments: On the real drawer-closing task, LOMPO fully closes the drawer in 19 of 25 rollouts, for a 76% success rate.It reaches the handle and pushes the drawer in 24 of 25 rollouts; the authors attribute some failures to classifier rewards incorrectly predicting closure when the drawer remains slightly open.
  • Real Robot Experiments: LMBRL, Offline SLAC, and visual foresight achieve 0% success on the real task because they do not navigate successfully to the correct handle location.Visual foresight achieves 65% on an easier shorter-horizon version of the task, whereas LOMPO succeeds in the more difficult setting.

6. Conclusion

The paper concludes that LOMPO makes offline model-based RL practical for high-dimensional visual observations by modeling dynamics and uncertainty in latent space. It reports strong performance across simulated and real manipulation tasks while identifying multitask learning as future work.

  • Conclusion: LOMPO handles high-dimensional observations by combining latent dynamics models with uncertainty quantification.It penalizes latent states according to disagreement among an ensemble of latent models.
  • Conclusion: LOMPO empirically outperforms prior latent model-based and model-free methods across four simulated tasks and one real-world robotic manipulation task.
  • Conclusion: Future work includes applying LOMPO to multitask settings, where shared models may improve vision and dynamics modeling, sample efficiency, and generalization.

Appendix A. Environments and Datasets

The appendix evaluates LOMPO on four simulated robotics environments with image observations, varied dynamics, action spaces, rewards, and dataset sizes. Together, the tasks test visual modeling, multimodal state estimation, sparse-reward control, and learning from thin data distributions.

  • Environments: The benchmark includes Walker, D’Claw Screw, Adroit Pen, and Sawyer Door environments with raw image observations and, in two tasks, robot proprioception.Walker uses 64 × 64 images; D’Claw, Adroit, and Sawyer use 128 × 128 images, with proprioception included for D’Claw and Adroit.
  • Environment Challenges: Sawyer Door requires opening a door from 128 × 128 images without robot-state access and uses a sparse reward only when the door is fully opened.The task is described as a difficult multistage robot-arm environment that online RL cannot solve.
  • Environment Challenges: Walker tests learning image-based forward dynamics with a non-stationary camera, while D’Claw and Adroit require combining vision and proprioception under high-dimensional actions and difficult contact dynamics.D’Claw and Adroit have action dimensions 9 and 26, respectively; D’Claw also uses a sparse reward.
  • Datasets: Dataset sizes are 100K, 180K, 1M, and 50K transitions for Walker, D’Claw, Adroit, and Sawyer, respectively, and include data from sub-optimal policies.The datasets are intended to test learning on thin data distributions.

Appendix B. Ablation Studies on Varying Offline Dataset Size

The dataset-size ablation varies the amount of D’Claw Screw data while preserving its source distribution. LOMPO remains effective in the low-data regime relative to regular latent model-based RL and Offline SLAC.

  • Dataset-Size Ablation: LOMPO still performs well in the low-data regime compared with regular latent model-based RL and Offline SLAC.The ablation uses medium-expert D’Claw Screw data and evaluates progressively smaller datasets.
  • Dataset Construction: The ablation creates a 1M-transition medium-expert dataset by mixing data from three policy-training runs, then subsamples it by factors of 5 and 25.All datasets are sampled from the same underlying distribution.
  • Dataset-Size Ablation: Figure 4 reports agent performance as a function of dataset size.

Appendix C. Implementation details

The latent dynamics and observation models use a variational architecture with recurrent deterministic state and stochastic latent transitions. Implementation details specify separate convolutional encoder and decoder designs, ensemble transition models, and fully connected actor-critic networks.

  • Latent dynamics: The latent state combines deterministic and sampled stochastic representations within an RSSM-based dynamics model.The deterministic component is shared across forward and inference models, while stochastic transitions are modeled separately.
  • Transition ensemble: The transition ensemble uses separate MLP stochastic models while sharing a single GRU-based deterministic representation.Each ensemble member models stochastic transitions pθk, while the deterministic model fθ is shared.
  • Visual representation: The image encoder and decoder are convolutional networks whose layer configurations vary across the listed visual-control environments.The encoder uses four or five convolutional layers, while the Walker decoder uses four layers with specified channel sizes, kernels, and strides.
  • Policy optimization: The actor and critic are three-layer fully connected networks with 256 units per layer, trained using Adam at learning rate 3e−4.Both networks share the stated architecture and optimizer setting.

Appendix D. Main Algorithm

The appendix directs readers to Algorithm 1 for the complete LOMPO procedure.

  • Main algorithm: Algorithm 1 presents the full LOMPO algorithm.

Appendix E. Variational Latent Model Samples

The appendix describes how LOMPO trains a variational ensemble model, infers latent states from offline sequences, and generates latent transitions for policy learning. Figure 5 compares ground-truth sequences with posterior samples and ensemble rollouts conditioned on actions.

  • Model training: The algorithm first trains a variational ensemble model on batches of raw observation, action, and reward sequences from the offline dataset.Training uses equation 8 with K models.
  • Real-data replay: Latent states inferred from offline sequences are added to a real replay buffer together with actions, rewards, and successor states.The inference model samples latent states conditioned on observation and action sequences.
  • Latent rollout: The procedure generates a latent replay buffer by sampling a transition model, rolling out actions and next states, computing rewards, and storing the resulting transitions.The latent transition model is selected from the ensemble, with both random-action and policy-action rollout descriptions supplied.
  • Policy update: The resulting real and latent batches are used to update the critic and policy with an off-policy actor-critic algorithm.
  • Model samples: Figure 5 organizes samples into ground-truth sequences, posterior model samples, and ensemble latent-model rollouts conditioned on the action sequence.
Loading 2012.11547v1…