🤖 AI

Toward Learning POMDPs Beyond Full-Rank Actions and State Observability

This paper presents a method for learning the explicit transition and observation matrices of a Partially Observable Markov Decision Process (POMDP) from sequential data by leveraging spectral approaches and tensor decomposition under mild rank assumptions, enabling agents to plan for diverse goals while proving that learning beyond a specific state partition is theoretically impossible.

Original authors: Seiji Shaw, Travis Manderson, Chad Kessens, Nicholas Roy

Published 2026-07-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Seiji Shaw, Travis Manderson, Chad Kessens, Nicholas Roy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to navigate a spooky, pitch-black mansion. You can't see the rooms, but you can hear creaks, feel drafts, and bump into furniture. Every time the robot takes a step, it gets a tiny clue about where it might be. This is the world of "Partially Observable" systems: the agent (the robot) knows what actions it can take and what it senses, but it doesn't know the true map of the house or exactly where it is standing at any given moment. To act smartly, the robot needs to build a mental model of this hidden world. For decades, scientists have developed ways to guess this model, but they often hit a wall: the models they built were like "black boxes." They could predict what would happen next, but they couldn't explain why or easily change their plans if the goal changed (like switching from "find the cat" to "find the treasure").

This paper tackles that black box problem. The researchers, working at MIT and the Army Research Laboratory, wanted to teach agents to learn the actual "rules of the game"—the specific probabilities of moving from one hidden room to another and the specific clues each room gives off. They focus on a special kind of mathematical puzzle called a POMDP (Partially Observable Markov Decision Process). Think of a POMDP as a board game where you can't see the board, only the dice rolls and the pieces you land on. The goal is to figure out the board's layout just by playing the game over and over again. The big question they ask is: Can we learn the exact layout of the board, even when some rooms look exactly the same from the outside?

The team's main discovery is a clever new way to learn these hidden maps, but with a very important catch. They found that if the robot's actions are "full-rank" (a fancy way of saying the actions are diverse enough to shake up the system and not get stuck in loops), the robot can learn the true probabilities of moving between groups of rooms. However, they proved that if two different rooms give off the exact same clues for every possible action, the robot can never tell them apart just by walking around. It's like trying to tell the difference between two identical twins who wear the same clothes and speak in the same voice; no matter how many times you ask them questions, you can't be sure which one is which.

So, what did they actually do? They combined two existing mathematical tricks. The first trick, called Predictive State Representations (PSR), is great at guessing the future but bad at explaining the past. The second trick, involving "tensor decomposition," is like a high-tech detective that can separate mixed-up signals to find the original sources. The authors realized that if they used the PSR method first to get a rough sketch of the world, they could then use the tensor method to "rotate" that sketch until it matched the real world's layout.

Here is the magic part: In many real-world scenarios, like a robot arm trying to pick up a cup, the robot might fail sometimes (slipping) or succeed other times. These "messy" actions are actually helpful because they create enough variety in the data to let the math work. The authors showed that by using these full-rank actions, their method can learn the transition and observation matrices (the rules of the game) up to a "partition." This means the robot learns that "Room A and Room B are in the same group because they look identical," and it learns the rules for moving between these groups. If every single room has a unique fingerprint, the robot learns the entire map perfectly. If some rooms are twins, it learns the map of the groups.

The paper also explicitly rules out the idea that we can ever learn the exact difference between those "twin" rooms using just a single line of walking data. They constructed a mathematical proof showing that two completely different house layouts could produce the exact same sequence of creaks and bumps. If the robot can't tell the difference in the clues, it can't tell the difference in the map. This isn't a failure of their algorithm; it's a fundamental limit of the universe.

In their experiments, they tested this on several digital worlds, including a "Tiger" game (where a tiger is behind one of two doors) and a "Sense-Float-Reset" domain (a robot moving on a line). They found that their method successfully learned the correct probabilities for moving between the "groups" of states. More importantly, because they learned the actual rules of the game rather than just a black-box predictor, they could change the robot's goal after the learning was done. For example, they could tell the robot, "Now, go to the noisy middle room," and the robot could figure out how to get there using the map it had already built. Previous methods would have needed to start all over again to learn a new goal.

The authors suggest that this approach is a significant step forward for robots that need to be flexible. It allows them to learn a model of the world that is interpretable and reusable. However, they are careful to note that this works best when the robot has a diverse set of actions to try. If the robot is too clumsy or the environment is too static, the math might not be able to separate the hidden states. They also point out that their current method works best for systems with a small number of hidden states, as the math gets very heavy for huge, complex worlds. But for now, they have shown that we can peek behind the curtain of hidden states, as long as we accept that some twins will always remain indistinguishable.

Technical Summary: Toward Learning POMDPs Beyond Full-Rank Actions and State Observability

Problem Statement
The paper addresses the challenge of enabling autonomous agents to learn and reason about systems with hidden states, specifically modeling these systems as discrete Partially Observable Markov Decision Processes (POMDPs). The agent is provided with the action and observation spaces but lacks prior knowledge of the state space, transition dynamics, or observation models. The goal is to estimate the number of hidden states and the explicit transition and observation matrices from a sequence of actions and observations collected via random exploration.

While spectral methods like Predictive State Representations (PSRs) can learn sufficient statistics to predict future outcomes, they function as black-box models. They do not yield explicit transition and observation likelihoods, making it difficult to manipulate the model for new planning tasks (e.g., changing reward functions) without relearning the entire system. Conversely, existing tensor decomposition methods for recovering full POMDP models typically rely on strict assumptions: that transition matrices are full-rank and that every state has a unique observation distribution for every action. These assumptions fail in many real-world scenarios involving state aliasing (where distinct states share identical observation distributions) or actions that may fail probabilistically.

Methodology
The authors propose a method that combines spectral decomposition with tensor decomposition to learn POMDP parameters under milder assumptions. The approach proceeds in three main stages:

Hankel Matrix Construction and PSR Learning:
The method begins by constructing a Hankel matrix from the observed action-observation trajectories. This matrix encodes the joint likelihoods of past and future observations. By performing a rank factorization (via Singular Value Decomposition) on this matrix, the algorithm derives a linear Predictive State Representation (PSR). This yields a set of linear update matrices ( $M_{ao}$ ), an initial vector ( $m_0$ ), and a final vector ( $m_\infty$ ). These components represent the system dynamics up to an unknown similarity transform $P$ .
Recovering the Similarity Transform:
To recover the explicit transition ( $T^a$ ) and observation ( $O_{ao}$ ) matrices, the algorithm must estimate the similarity transform $P$ . The authors reformulate tensor decomposition techniques to estimate $P$ .
- First, they marginalize observations to obtain transition matrices up to the transform: $\sum_o M_{ao} = P T^a P^{-1}$ .
- For actions with full-rank transition matrices, they compute the product $M_{ao} M_a^{-1} = P O_{ao} P^{-1}$ .
- Since $O_{ao}$ are diagonal, the eigenvalues of these products correspond to the diagonal entries. To resolve ambiguity caused by repeated eigenvalues (which occur when states share observation distributions), the method employs a joint diagonalization of a weighted sum of these matrices. Random weights are sampled to ensure that eigenvalues are distinct almost surely for states with unique observation distributions across all full-rank actions.
Learning up to Observability Partitions:
The paper introduces the concept of a Full-Rank Observability Partition. If multiple states share the exact same observation distribution for all actions with full-rank transitions, they cannot be distinguished from sequential data alone. The algorithm recovers the similarity transform $P$ up to this partition. Consequently, the learned model provides explicit transition and observation matrices between partitions of states rather than individual states. For states within a single partition, the model learns the aggregate dynamics.

Key Contributions

Relaxed Assumptions: The method extends tensor decomposition approaches to learn POMDPs where states may be aliased (sharing observation distributions) and where transition matrices are not necessarily full-rank for every action, provided the product of forward and backward matrices is full-rank.
Explicit Model Recovery: Unlike PSRs, the method recovers explicit transition and observation matrices (up to a partition), enabling the specification of new reward functions and goals after the initial learning phase without retraining.
Theoretical Limits: The authors prove that learning transition and observation dynamics beyond the observability partition is impossible from a single sequential data trajectory. They construct a counterexample demonstrating two distinct POMDPs with different internal dynamics that yield identical observation distributions for all action sequences, establishing the theoretical limit of identifiability.

Experimental Results
The authors evaluate their method on several domains, including the "Sense-Float-Reset" domain (featuring state aliasing), the Tiger problem, and T-Maze.

Convergence: The method successfully recovers the underlying observation models and partition-level transition matrices, converging to ground truth values as the number of exploration interactions increases. In contrast, Expectation-Maximization (EM) often converges to local minima that predict observations correctly but yield incorrect model parameters.
Planning Performance: When used with a standard planner (PO-UCT), the learned partition-level POMDPs achieve planning performance comparable to ground-truth models and PSRs.
Reward Specification: In novel "Noisy Hallway" and "Directional Hallway" domains, the authors demonstrate that the explicit transition and observation likelihoods allow an agent to be directed toward specific latent states (e.g., a middle hallway state with ambiguous observations) by assigning rewards to states rather than just observation-action pairs. This capability is not available with standard PSRs or EM-based models that fail to recover the correct state structure.

Significance and Claims
The paper claims that its method bridges the gap between the predictive power of PSRs and the interpretability of explicit POMDP models. By learning up to an observability partition, the method captures a broader class of real-world systems (particularly in robotic manipulation where actions may fail or states may be aliased) than previous tensor-based methods.

The authors modestly assert that while their method cannot recover the exact internal structure of aliased states from sequential data, it recovers the maximal amount of information theoretically possible. The ability to manipulate the learned model for new tasks (via explicit reward specification) is highlighted as a primary advantage over black-box predictive models. The work suggests that for many practical applications, learning the dynamics between observability partitions is sufficient for effective planning and reasoning.

Technical Summary: Toward Learning POMDPs Beyond Full-Rank Actions and State Observability

More like this