Neural Predictive Belief Representations

TL;DR

This work investigates unsupervised learning of belief-state representations for partially observable environments by comparing one-step frame prediction, CPC, and CPC|Action within POMDPs. Using a glass-box evaluation, the authors show that all three methods can encode the agent's state, trajectory, and uncertainty in a compact representation $b_t$, with CPC-based methods—especially CPC with multi-step prediction and action-conditioning—performing best in visually complex environments. The findings suggest that neural belief representations can support planning and exploration under uncertainty, and that incorporating future actions and longer prediction horizons is beneficial. The work proposes a path toward more general, uncertainty-aware representations that can transfer across tasks and modalities.

Abstract

Unsupervised representation learning has succeeded with excellent results in many applications. It is an especially powerful tool to learn a good representation of environments with partial or noisy observations. In partially observable domains it is important for the representation to encode a belief state, a sufficient statistic of the observations seen so far. In this paper, we investigate whether it is possible to learn such a belief representation using modern neural architectures. Specifically, we focus on one-step frame prediction and two variants of contrastive predictive coding (CPC) as the objective functions to learn the representations. To evaluate these learned representations, we test how well they can predict various pieces of information about the underlying state of the environment, e.g., position of the agent in a 3D maze. We show that all three methods are able to learn belief representations of the environment, they encode not only the state information, but also its uncertainty, a crucial aspect of belief states. We also find that for CPC multi-step predictions and action-conditioning are critical for accurate belief representations in visually complex environments. The ability of neural representations to capture the belief information has the potential to spur new advances for learning and planning in partially observable domains, where leveraging uncertainty is essential for optimal decision making.

Figures (7)

• Figure 1: Different architectures used in our experiments.
• Figure 2: Frames from the agent moving at random in a gridworld (outermost cells are walls, see \ref{['fig:gridworld_uncertainty_screenshot:wall']}). In each image, the agent's partial observation is on the left (agent and walls in black, empty spaces in white), the agent's position and orientation are on the centre, and the predicted position and orientation are on the right. The diamond-looking shapes result from flattening the beliefs for each of the four possible orientations in the same cell.
• Figure 3: Examples of agent observations for different environments.
• Figure 4: Example predictions for terrain. In each image, ground truths are on the top, predictions on the bottom, $(x, y, \theta)$ (ground truth and predictions) on the left, and past $(x, y, \theta)$ on the right. \ref{['fig:natlab_position_and_history:00', 'fig:natlab_position_and_history:01', 'fig:natlab_position_and_history:02']} show examples of accurate $(x, y, \theta)$ predictions, \ref{['fig:natlab_position_and_history:10', 'fig:natlab_position_and_history:11', 'fig:natlab_position_and_history:12']} show examples for inaccurate $(x, y, \theta)$ predictions.
• Figure 5: Symmetry of the $(x, y, \theta)$ prediction with the frame predictor in room. Ground truths are on the top, predictions on the bottom, $(x, y, \theta)$ (ground truth and predictions) on the left, and past $(x, y, \theta)$ on the right.