Heteroscedastic Uncertainty for Robust Generative Latent Dynamics

TL;DR

This paper introduces a method for learning robust latent dynamics models that capture heteroscedastic uncertainty, improving prediction and control in noisy, real-world robotic tasks.

Contribution

The paper presents a novel approach to jointly learn latent representations and dynamics that incorporate input-specific uncertainty for better robustness.

Findings

Enhanced prediction accuracy under noisy conditions

Improved control performance with heteroscedastic uncertainty modeling

Effective in both simulated and real-world robotic tasks

Abstract

Learning or identifying dynamics from a sequence of high-dimensional observations is a difficult challenge in many domains, including reinforcement learning and control. The problem has recently been studied from a generative perspective through latent dynamics: high-dimensional observations are embedded into a lower-dimensional space in which the dynamics can be learned. Despite some successes, latent dynamics models have not yet been applied to real-world robotic systems where learned representations must be robust to a variety of perceptual confounds and noise sources not seen during training. In this paper, we present a method to jointly learn a latent state representation and the associated dynamics that is amenable for long-term planning and closed-loop control under perceptually difficult conditions. As our main contribution, we describe how our representation is able to capture a notion of heteroscedastic or input-specific uncertainty at test time by detecting novel or out-of-distribution (OOD) inputs. We present results from prediction and control experiments on two image-based tasks: a simulated pendulum balancing task and a real-world robotic manipulator reaching task. We demonstrate that our model produces significantly more accurate predictions and exhibits improved control performance, compared to a model that assumes homoscedastic uncertainty only, in the presence of varying degrees of input degradation.