Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy

Thank you for the kind words and feedback. We're glad you had such a positive impression of the idea and see potential for very high impact. You cited clarity as your biggest concern, so we'll address that here and put our other comments in the global response.

Markovian Representation

We apologize for the confusion. There was a typo in line 97. It should instead read "next representation z_{t+1}". The definition can be applied to states (s), observations (ω), and even memory-augmented observations (ω, m) by replacing z with the relevant quantity. If equation 1 holds, the corresponding representation is Markovian.

Empirical MDP

You're correct that Pr(s|ω) is an expectation over trajectories under the policy. Assuming all observations have non-zero probability, we can write this in terms of trajectories $\tau=(s_0,ω_0,a_0,r_0, s_1,ω_1,a_1,r_1,...)$:

$Pr(s|ω)=\frac{𝔼\_{\tau\sim π}[\sum_{t=0}^∞ γ^t𝟙[s_t=s,ω_t=ω]]}{𝔼\_{\tau\sim π}[\sum_{t=0}^∞ γ^t 𝟙[ω_t=ω]]}.$

This is equivalent to the definition we use in Appendix C, but we agree it is easier to understand it this way up front.

Your point that the empirical MDP is “not a real MDP” is interesting. MDPs are simply models of environments and can be of arbitrary quality. The empirical MDP is one such model. It is the one practitioners use implicitly whenever they treat POMDP observations as Markovian "states". The empirical MDP model consists of (Markovian) transition and reward functions that condition on observations, which we define in terms of Pr(s|ω), and which, crucially, might make different predictions from the POMDP model. We can thus view Pr(s|ω) as a weighting scheme for aggregating the behavior of the states that produce each observation.

Model selection is not a new idea, but our specific characterization of the empirical MDP for a given POMDP is novel to the best of our knowledge.

Tensor Notation and Q vs. V

Defining LD using probability notation directly would be prohibitively cumbersome and a poor use of space (see Appendix C). Tensor notation succinctly expresses the concept and has precedent in many papers [Lagoudakis & Parr 2003, Tagorti & Scherrer 2015, Gehring et al. 2016]. Q-values are easy to express with tensors that track action and observation indices separately, and they are less likely than V-values to succumb to Parity Check symmetry, since they avoid an additional summation over actions.

LD=0 cases and Parity Check

We found a new way to characterize these edge cases that streamlines the section, removes the need for Eq. 7 entirely, and that we hope is much clearer.

We start with Defn. 1: $Λ_π^{λ_1,λ_2} := \|\| Q_π^{λ_1} - Q_π^{λ_2}\|\| = \Big\|\Big\|W\Big(I-γTK_π^{λ_1})^{-1}-(I-γTK_π^{λ_2})^{-1}\Big):R^{SA}\Big\|\Big\|,$ where $K^λ_π := (λ Π^S + (1-λ) \Phi W^Π)$, and the : operator denotes a tensor product over two indices.

Let us parse this equation by considering one of the two Q-functions in isolation (say $λ_1$).

Policy Tensors. First, note that $Π^S$ and $\Phi W^Π$ are mappings from states to distributions over state-action pairs $s\mapsto\Delta(\tilde s,a)$. We call the former the MC policy tensor, ${Π^S}\_{s,\tilde s,a}=𝟙[s=\tilde s]\sum\_ω \Phi\_{s,ω}π\_{ω,a}$, which maps each state the expected policy under its observation distribution. We call the latter the TD policy tensor, $(\Phi W^Π)_{s,\tilde s,a}=\sum\_ω \Phi\_{s,ω} π\_{ω,a} W\_{ω,\tilde s}$, which spreads the policy probabilities for a given observation $ω$ across all states $\tilde s$ that produce $ω$, weighted by $W=Pr(s|ω)$. $K\_π^{λ_1}$ is a convex combination of these policy tensors, parameterized by ${λ_1}$.

Transition-policy tensors. Left-multiplying a policy tensor by $T$ produces an $(S×A×S×A)$ transition-policy tensor, which describes the probability of each state-action transition $(s,a)\mapsto\Delta(s',a')$ under the policy. Intuitively, TD($λ_1$) computes Q-values for a POMDP as though state-action pairs are evolving according to $TK\_π^{λ_1} = λ_1 TΠ^S + (1-λ_1)T\Phi W^Π$, which is a mixture of the policy dynamics under two models: the MC transition-policy ($TΠ^S$) and the TD transition-policy ($T\Phi W^Π$).

Final projections. The expression $(I - γTK\_π^{λ_1})^{-1}$ computes the state-action successor representation for this hybrid transition model, then multiplication by $R^{SA}$ computes state-space $Q$-values. Finally, these are projected through $W$ to compute observation-space $Q$-values.

Returning to the question of when $Λ=0$ for all policies, there are 3 cases:

1. Policy tensors are equal: $K_{π}^{λ_1}=K_{π}^{λ_2}$. This simplifies to $Π^S=\Phi W^Π$, which is satisfied if and only if $\Phi W = I$, i.e. each observation is generated by exactly one hidden state. This means the POMDP is actually a block MDP.

2. Transition-policy tensors are equal: $TK_{π}^{λ_1}=TK_{π}^{λ_2}$. In this case, the POMDP is equivalent to an MDP. Both transition-policy tensors lead to the same state-action successor representation, so any reward function will lead to the same value under both models.

3. Any differences between transition-policy tensors (and thus the state-action successor representation) are subsequently projected away by either $W$, $R^{SA}$, or both. In this case, LD=0, but the effective MDP may be lossy. Fortunately, such pathological environments appear to be quite rare, and they are easy to handle, as we discuss below.

Parity Check is an example of case 3. One can check (by hand or using our code) that the $λ_1$ and $λ_2$ transition-policies (and thus the successor representations) differ with and without aliasing. However, these differences are perfectly symmetric with respect to observations and rewards, so they are projected away by $W$ and $R^{SA}$. The Q-values are therefore zero at every observation, and so the LD=0 for all policies. Nevertheless, as we cover in our global response, we can still detect and mitigate partial observability for this problem.

Thank you for the feedback and questions.

We will group two of your comments together, and respond to the remainder in order.

As for Theorems 1 and 2, do the policies π have to be Markov policies only taking the current state S as input? If so, the underlying environment is expected to be a block MDP. I am wondering if the result could be more general if we allow non-Markov policies with memory.

Q3: For policies π in Theorems 1 and 2, are they Markov policies only taking the current state $S$ as input? If so, it excludes the non-Markov policy, which could have memories.

There seems to be a minor misconception. The agent’s policy $π(a|ω)$ is conditioned on observations---not hidden state. Rather, the state policy tensor, $Π^S\_{s,s',a}=𝟙[s=s']\sum\_ω \Phi\_{s,ω}π\_{ω,a}$, is the effective policy over states; this tensor maps states to the expected policy under their observation distribution. At no point do we require the agent’s behavior to be conditioned on unobservable quantities. However, our definition does also support policies that condition on both the observation and the memory state, via the memory augmentation process we outline in Sec. 4.

---

Could the authors make the connection with the weakly revealing POMDPs?

Weakly revealing POMDPs (Liu, Chung, et al. 2022) are an interesting class of POMDPs in which an $m$-step future contains enough information to reveal the current state, given sufficient samples, and the paper authors argue that this condition makes partial observability “not scary”. At a high level, the LD also helps to make partial observability “not scary”, by providing a way to detect and mitigate it. So it would make a lot of sense if the weakly revealing literature ultimately turns out to have some connection to this work. We have not been able to prove a formal connection yet, but if you have suggestions on how to do this, please let us know! Ultimately, we hope that the LD will be a useful tool for analyzing many aspects of partially observable reinforcement learning.

---

Additional discussion on why minimizing λ-discrepency could improve the planning performance would be appreciated. Could the authors provide some intuitions?

Certainly! Lower LD leads to state representations that are more predictive of the future. Better representations lead to better value function estimation, and this leads to better performance. Why does it help to minimize this LD quantity specifically? Well, to the extent that optimal behavior requires knowledge of the underlying hidden state, zeroing the LD is a reliable way to ensure that the agent’s representation contains a complete picture of that hidden state. On the other hand, if optimal behavior only requires conditioning on the current observation, then LD minimization won’t improve performance, but it can still provide additional confidence that the current observation is indeed sufficient.

To see what these “better” representations look like, check out “Memory visualization” in our global response above. We also find that LD has a minimum at zero and increases with partial observability (see ‘Varying “Markovness”’). We also conjecture that reducing LD also improves value estimation error. We have not yet proved a formal link, but this seems plausible, particularly given that LD is defined as the difference between Q-values, which is essentially a type of value error. Moreover, the empirical results in our paper suggest we are indeed learning better representations, which should allow better value estimation.

---

Questions

Q1: What does LD stand for in Fig. 4?

Sorry for the confusion. LD stands for the Lambda Discrepancy in this Fig. 4. We’ve updated the caption to clarify this.

---

Q2: How to select parameters λ_1 and λ_2 when minimizing the loss function in Eq. 8? Does one have to try multiple pairs of such λ s?

Thanks for the question. Please see “Understanding and selecting λ parameters” in our global response.

---

Q3.

(Addressed above.)

Thank you for your review and feedback. It seems that your main concern is the misconception that LD “requires operating over the entire state, observation and action space for its computation.” We apologize if this was confusing. While the derivation of the metric uses tensors that operate over the entire state, action and observation space, the actual estimation and minimization of the metric only uses estimates of TD($λ_1$) and TD($λ_2$), which are computed directly from sampled (observation, action) pairs. This means that our method does not require the underlying model of the POMDP to mitigate partial observability, only samples.

We realize this may not have been clear in our submission, so we’ve added the following discussion near the end of Sec. 4:

Despite its definition in terms of $S$, we can express the $λ$-discrepancy in a form amenable to estimation from observable quantities only. By selecting the appropriate weighted norm, the $λ$-discrepancy becomes an expectation over observation-action pairs (or, in the case of memory, observation-memory-action tuples): $Λ\_{\mathcal{P},π}^{λ\_1,λ\_2} := \|\|Q\_π^{λ\_1} - Q\_π^{λ\_2}\|\|\_π = 𝔼\_{π}\left[ Q\_π^{λ\_1}(ω, a) - Q\_π^{λ\_2}(ω, a) \right]$ The norm is taken over the on-policy joint observation-action distribution $Pr(ω)π(a|ω)$, which allows us to estimate the metric using samples generated by the agent’s interaction with the environment. We show how to use it as an optimization objective in the following section.

We hope by clarifying our objective's dependence on only observables we've highlighted the broad applicability of this method in partially observable domains.

---

We will now respond to your questions.

1. However, the authors also use two RNNs in their proposed approach that uses the entire history. So how does this argument hold?

   Sorry for the confusion. We meant to say “Without state, a model-based comparison would require variable-length history inputs per time step to ensure a Markovian representation”. So the input size for this model-based comparison would be on the order of O(t), whereas the input size to an RNN is fixed and always O(1). To clarify, we only use a single RNN, with two value heads conditioned on the same internal representation (new Fig 2, left).

2. How can such examples be identified without knowledge of the underlying model and how can they be practically addressed as the changes suggested here are changes to the environment which the agent cannot make.

   Good question! We address this with a new experiment in the global response (“Parity Check”).

3. How will the convergence analysis of PPO change after the addition of the auxiliary loss given by Eq. (8)? When these functions are being learnt (still have errors) the λ-discrepancy will also have an error. How do the authors address this?

   While we are unaware of any general convergence guarantees for PPO (much less recurrent PPO), our empirical results show that the LD loss frequently performs better (and never performs worse) than the RNN + PPO baseline.

   You are correct that the value functions will be inaccurate during learning, and so too will the LD. Nevertheless, just as value functions are commonly useful for RL even before they have fully converged, so too is the LD, as our experiments demonstrate.

4. The authors compare their approach with a recurrent PPO baseline. Why is there no comparison with state of the art POMDP methods for these problems?

   Our approach is introducing a new metric, the LD, as opposed to advocating for a particular algorithm like PPO. The LD is not limited to only PPO, but could be leveraged in essentially any RL method that uses value functions. We are considering extensions to other approaches for future work, where we believe the LD still has great potential to improve existing algorithms.

   We aren’t sure what you mean by state-of-the-art POMDP methods. If you mean belief-state methods like POMCP, it is unclear what such a comparison would reveal, since those methods assume knowledge of the transition, reward, and observation functions, and the set of hidden states. Our method only assumes access to on-policy observations while acting in the environment, which is a much more challenging setting.

5. The proposed algorithm requires two Q networks, thereby doubling the compute requirements of the critic in actor-critic methods. Is the gain in performance sufficient to justify this?

   The proposed algorithm requires two value function heads (3 layer MLPs), each conditioned on the hidden state of a single RNN. The extra computational overhead is only for this additional value head---about a 12% increase in parameters vs. the RNN baseline, as we discuss in Appendix I. We feel this is a reasonable tradeoff, given the performance gains.

6. How does the memory learnt by this approach compare with representations learnt by POMDP based approaches such as PSRs etc.?

   Good question! We added a new experiment in our global response visualizing these memory functions for the methods in Sec. 5.

   Our approach uses memory functions to incrementally update a summary of past observations. By contrast, PSRs use tests of potential future trajectories, where the main challenge lies in finding these tests. In prior work [Singh, Littman et al., 2003], PSRs were able to learn representations for domains as large as 58 states and 21 observations. By contrast, Battleship (without accounting for possible ship locations) has ~2^100 states. PacMan, even when ignoring food dot configurations, has ~5^400 states. To the best of our knowledge, PSRs have not yet been scaled up to domains of the size we study here.

7. Can λ-discrepancy metric be defined at a state/observation or state/observation x action level? If not, it is not clear how the quantity given in Eq. (8) relates to the one in Definition 1?

   We hope our initial clarification above answered this question!

I thanks the authors for their explanations. Based on these, I have decided to increase my score. Below are my responses to your explanations:

1. Here is one paper that provides convergence analysis for policy gradient algorithms: https://www.jmlr.org/papers/volume22/19-736/19-736.pdf . My question was regarding the impact of your auxiliary loss in such analyses.
2. By state of the art POMDP methods, I mean methods such as predictive state representations, deep recurrent Q learning etc., where the algorithm is designed primarily to handle partial observability. Instead of predicting tests as in PSR, there are also auxiliary tasks such as predicting next observation which have been shown to be useful for information state learning in the literature.

Thank you for your response and feedback! We are glad to hear you thought the writing was “clear and to the point.”

We will respond to your comments in order.

---

The experiments also use a recurrent PPO and compare it to recurrent PPO with lambda discrepancy but it is not explained how they relate to other memory-based methods.

Thank you for the suggestion. While we use RNNs (specifically GRUs) in our experiments due to their popularity in RL, there are other architectures we have not considered. Transformer-based architectures have recently been studied in reinforcement learning settings, and have been shown to help in long-term memory tasks [1]. Seeing as our metric poses no constraints on function approximation or network architecture, an interesting direction for future work would be to see if the LD objective would benefit Transformer-based architectures as well.

As for methods that explicitly consider partial observability with memory, the most dominant approaches have used belief-states. Belief-state methods [2] are a class of solution methods for POMDPs that have seen substantial success in robotics [3], but most approaches are intractable for even small environments where the dynamics are unknown, since most methods rely on some form of state information to form the belief state that acts as memory.

• [1] Ni, Ma, Eysenbach, Bacon. NeurIPS, 2023.
• [2] Kaelbling, Littman, Cassandra. Artificial Intelligence, 1998.
• [3] Thrun, Burgard, Fox. Probabilistic Robotics. MIT Press, 2005.

---

Ideally, it would be most interesting seeing results for environments with Markovian state representation that gradually shift to POMDPs that are not Markovian.

Thanks for the suggestion. We ran this experiment for T-Maze and included the results in our global response.

---

The “Limitations” section should be part of the main corpus.

We had originally moved the limitations section to the appendix due to space constraints, but we will move this section back into the main corpus with the extra page allotted for the camera-ready.

---

Block MDPs are mentioned as a possibility for the LD to be zero, but their relevancy - whether they remain to be an edge case - is not discussed.

We proved block MDPs have LD=0, and this is good, because block MDPs are fully observable; they are MDPs. In fact, classic MDPs can be viewed as block MDPs with a single observation per state. Block MDPs are therefore not an edge case, but an indication that our metric is correctly labeling MDPs as having zero partial observability.

---

However the explanation of “Intermediate values of λ smoothly interpolate between the two” (pdf) (l. 163) values is not intuitive. An example at would be helpful

We apologize for the confusion. You are right that TD(λ=0) and TD(λ=1) don’t provide much intuition on how the intermediate lambda values will behave. For intermediate λ values, TD(λ) is a weighted average of all the n-step returns. The λ parameter controls the weighting scheme: lower values of λ give more weight to the shorter rollouts. When λ=0.5, this defines a weighting scheme where each n-step rollout contributes half as much to the weighted sum as the (n-1)-step rollout. It is possible to view TD(λ) as operating in either a backwards-looking or forwards-looking way, but in the paper, we always take the forward view: the λ parameter does not determine whether previous states are taken into account, but rather how much each n-step future return contributes to the weighted average.

---

The experiments remain on a low complexity levels

Thanks for the feedback. Please see “Environment Complexity" in our global response above.

---

Questions

1. Could you please clarify the setup regarding the LD loss?

We have provided a visualization of the network architecture in new Fig. 2 (left). We train two value heads and a policy head on top of a single shared recurrent representation. Only one of the value heads is needed for the PPO objective, but both are used in computing the LD loss, and the recurrent representation is trained end-to-end. For more details, see Appendix I. We will add the figure to the camera-ready if the paper is accepted.

---

2. How does the LD method benefit when there is existing domain knowledge about an environment?

Great question! Domain knowledge can help a system designer to provide useful (ideally Markovian) features to the decision-making agent. If the system designer wants to check whether these features are in fact Markovian, he or she could use the LD to test for it.

Interestingly, even when the environment is fully observable, the λ-discrepancy may still be useful, because function approximation can sometimes act like partial observability. If the environment has Markovian observations but these observations are sent through a lossy function approximator (like a neural network), then minimizing λ-discrepancy could help the function approximator to learn better features.

---

3. How is the LD method affected when switching between online and offline-RL settings?

Another good question! For offline RL, LD needs to be applied with care when the data is not generated from the policy being evaluated, TD(λ) no longer has the same fixed point for all values of λ under different policies. Our theory relating the Markov assumption to LD would not hold between two value functions under different policies. However, one could use tools from off-policy RL such as importance sampling to adapt our methods to this new setting. Another idea is to minimize the LD for the data-generating policy instead of the evaluation policy. Due to the result in Theorem 1, this should result in a memory function that, in theory, mitigates partial observability for other policies as well. This is similar to what we do in the analytical experiments in Sec. 4. That setup is akin to an offline policy evaluation approach for memory learning, which is an exciting avenue for future work.