From Optimal Actions to World Models: Identifiability of Transition Kernels in Discounted MDPs

Technical Summary: Identifiability of Transition Kernels in Discounted MDPs

Problem Statement
This paper investigates the inverse problem of recovering the transition probabilities of a finite, discounted Markov Decision Process (MDP) given only the set of optimal actions for every possible reward function within a specific class. While the transition model dictates both the optimal policy and the dynamics of the environment, two distinct transition kernels may recommend identical optimal actions. The central question is: under what conditions do two kernels $P$ and $Q$ become indistinguishable solely based on the optimal actions they induce?

The study compares three distinct classes of rewards, which test the dynamics in different ways:

1. State rewards: $r(s)$, depending only on the current state.
2. State–action rewards: $r(s, a)$, depending on the current state and action.
3. Transition-dependent rewards: $r(s, a, s')$, depending on the state, action, and the resulting next state.

The paper defines three notions of equivalence: $P \equiv_{tr} Q$ (indistinguishable under transition rewards), $P \equiv_{sa} Q$ (indistinguishable under state–action rewards), and $P \equiv_{s} Q$ (indistinguishable under state rewards).

Methodology
The analysis relies on the structural properties of the Bellman optimality equations and the relationship between value functions and transition kernels.

• State–Action Rewards: The author constructs a specific family of rewards, $r^P_v(s, a) = v(s) - \gamma P_{s,a}v$, which render every action optimal for a given value vector $v$. By comparing the policy evaluation equations under kernels $P$ and $Q$ for these rewards, the paper derives a necessary and sufficient condition for equivalence. This involves identifying a unique invertible matrix $L$ that relates the value functions of the two kernels.
• Transition-Dependent Rewards: The methodology exploits the ability to isolate specific transition rows. If a state has multiple actions, a transition-dependent reward can be designed to have zero expected value under one transition row and non-zero under another, effectively distinguishing them. For states with a single action, the analysis relies on whether the specific transition row affects the optimality of actions elsewhere in the state space.
• State Rewards: The analysis characterizes equivalence by examining the "policy regions" $R_P(\pi)$, defined as the set of state rewards for which a deterministic policy $\pi$ is optimal. Two kernels are equivalent if and only if these regions are identical for all deterministic policies.

Key Results

1. State–Action Equivalence ($P \equiv_{sa} Q$):
   Two kernels $P$ and $Q$ yield the same optimal actions for every state–action reward if and only if there exists a unique invertible matrix $L$ (satisfying $L\mathbf{1} = \mathbf{1}$) such that:
   $$Q_{s,a} = \left( P_{s,a} + \frac{1}{\gamma} e_s^T (L - I) \right) L^{-1}$$
   This implies that the advantage of any action relative to a policy is preserved under the transformation. Near a kernel with strictly positive entries, the set of indistinguishable kernels forms a smooth manifold of dimension $n(n-1)$, where $n$ is the number of states. This is a strict subset of the ambient kernel space, meaning that while the transition kernel is not uniquely determined, the ambiguity is constrained to a specific lower-dimensional family. The result holds even if the data is restricted to rewards with a unique optimal action at every state.

2. Transition-Dependent Equivalence ($P \equiv_{tr} Q$):
   Transition-dependent rewards provide significantly more information.

   • At any state with at least two actions, the transition row is uniquely identified.
   • At states with a single action, the row is identified unless the specific transition probabilities never influence the optimality of actions at other states (a condition that holds only for a measure-zero set of strictly positive kernels).
   • Consequently, if the MDP contains at least one state with a genuine choice (multiple actions), almost every strictly positive transition kernel is uniquely identified by transition-dependent rewards.

3. State Reward Equivalence ($P \equiv_{s} Q$):
   State rewards provide the least information. Two kernels are equivalent if and only if they share the same set of optimal deterministic policies for every state reward. This condition is equivalent to the equality of the policy regions $R_P(\pi) = R_Q(\pi)$ for all $\pi$. This equivalence is strictly weaker than state–action equivalence; distinct kernels can be indistinguishable under state rewards while being distinguishable under state–action rewards.

4. Hierarchy of Identifiability:
   The paper establishes a strict hierarchy of equivalence:
   $$P \equiv_{tr} Q \implies P \equiv_{sa} Q \implies P \equiv_{s} Q$$
   Neither implication is reversible. Allowing the reward to depend on more components of the transition (specifically the next state) strictly increases the ability to distinguish between different transition models.

Significance and Claims
The paper claims to resolve the question of how much of the transition dynamics can be recovered from optimal behavior alone. It demonstrates that:

• Optimal actions are insufficient for full recovery: Even with knowledge of optimal actions for all rewards, the transition kernel is generally not uniquely determined when rewards depend only on states or state–actions.
• Structural constraints exist: For state–action rewards, the ambiguity is not arbitrary; it is precisely characterized by a matrix transformation that preserves action comparisons (advantages) across all policies.
• Reward structure dictates recoverability: The form of the reward is critical. Transition-dependent rewards can recover the kernel almost entirely, whereas state rewards may leave substantial structural differences hidden.
• Distinction between decision-making and prediction: The results highlight that a model may be sufficient for reproducing optimal decisions (policy reproduction) without being sufficient for predicting transitions, simulating counterfactuals, or transferring to unseen reward classes.

The work extends previous literature (e.g., Letcher et al.) by shifting the focus from recovering kernels from numerical Q-values to recovering them from the discrete set of optimal actions, while simultaneously expanding the scope of the reward class considered. The paper concludes that while optimal actions reveal significant structure about the dynamics, they do not generally reveal the dynamics themselves unless the reward function is sufficiently rich (i.e., depends on the next state).