Geometry of Uncertainty: Learning Metric Spaces for Multimodal State Estimation in RL

TL;DR

The paper tackles robust state estimation in RL from high-dimensional multimodal observations by learning a latent metric space in which distances reflect the minimum action distance between states. It introduces MetricMM, with per-modality encoders, a latent transition model, and an inverse-distance fusion scheme, trained via three losses—contrastive temporal distance, latent transition, and multimodal invariance—plus a cross-modal aggregation that does not rely on explicit noise models. Empirical results on MuJoCo and Fetch show that MetricMM achieves higher robustness to unseen sensor corruptions and improves RL performance without noise-augmentation during training. The approach offers a scalable geometric interpretation of uncertainty in sequential decision-making, enabling robust and efficient state estimation across diverse tasks and modalities.

Abstract

Estimating the state of an environment from high-dimensional, multimodal, and noisy observations is a fundamental challenge in reinforcement learning (RL). Traditional approaches rely on probabilistic models to account for the uncertainty, but often require explicit noise assumptions, in turn limiting generalization. In this work, we contribute a novel method to learn a structured latent representation, in which distances between states directly correlate with the minimum number of actions required to transition between them. The proposed metric space formulation provides a geometric interpretation of uncertainty without the need for explicit probabilistic modeling. To achieve this, we introduce a multimodal latent transition model and a sensor fusion mechanism based on inverse distance weighting, allowing for the adaptive integration of multiple sensor modalities without prior knowledge of noise distributions. We empirically validate the approach on a range of multimodal RL tasks, demonstrating improved robustness to sensor noise and superior state estimation compared to baseline methods. Our experiments show enhanced performance of an RL agent via the learned representation, eliminating the need of explicit noise augmentation. The presented results suggest that leveraging transition-aware metric spaces provides a principled and scalable solution for robust state estimation in sequential decision-making.

Figures (9)

• Figure 1: We propose MetricMM, a novel state estimation model from noisy multimodal observations. Each observation ($o_1, o_2$ on the left) is mapped into a joint metric space ($Z$ on the right) where distances from the latent dynamics prediction ($\hat{z}_t = \varphi_T(z_{t-1}, a_{t-1})$) are correlated with their uncertainty.
• Figure 2: Mean and standard deviation over 5 seeds of the training return for a SAC agent on the Hopper-v5 environment. The policies are trained with different state estimation modules (LinearComb, ConCat, CURL, GMC) and different amounts of Gaussian noise on the observations. With an increase in noise, the expected average return sensibly decreases for all the estimators. Epochs are in thousands.
• Figure 3: Mean and standard deviation over 5 seeds and 50 trajectories of the testing return for a SAC agent on the Mujoco suite. The policies are tested with different state estimation modules and different amounts of noise (perturbations of one modality at a time). MetricMM is the only estimator that allows for a consistent return with high-frequency perturbations.
• Figure 4: Mean and standard deviation over 5 seeds and 50 trajectories of the testing return for a SAC agent on the Fetch suite under time-persistent sensor failure. The policies are tested with different state estimation modules and different amounts of noise. Each time the noise is applied, it persists for either 3 or 10 consecutive frames ($K$) on that specific sensor modality. MetricMM is the only estimator that allows for a consistent return with high-frequency perturbations.
• Figure 5: Average return for a SAC agent on the one-dimensional pendulum environment and increasing level of Gaussian noise. A MetricMM exhibits robust performances, both of the ablation variants degrade with an increase in noise.