Paper: Resolving Partial Observability in Decision Processes via the Lambda Discrepancy (ICLR-2024)

Reviewer comment (Writing):
The notations are very confusing, and many parts of the technical presentations are very difficult to follow. The products of tensors are presented as regular matrix multiplications without specifying which dimensions are used in the multiplications. Many definitions are not precise, like effective policy over latent state, and some augment representations of some tensors. In the proofs, it is very difficult to follow when the indices are meaningless $i, j, k, l, ...$ without specifying their domains.

Author response:
$ $

-----
### Mathematical presentation

$ $

> - notations are very confusing
> - The products of tensors are presented as regular matrix multiplications without specifying which dimensions are used in the multiplications
> - many definitions are not precise, like effective policy over latent state
> - indices are meaningless i,j,k,l,...

$ $

We have fixed the mathematical presentation in the paper to make it easier to follow. Specifically,

1. We changed the indices from $i,j,k,l,...$ to indices that indicate their domain, such as $a_1, o_1,$ and $s_1$.
2. We have clarified the contractions used in the proof of Appendix A. We kept the matrix multiplications written as-is because they make sense as long as one follows the following guideline, which we included in the original paper: in the main paper equations, each tensor contraction is a single contraction except in the products involving $R^{SA}$ and $T$ on the right hand side, which are double contractions. We also tried writing the contractions with dots, but decided this tensor notation offered the best compromise between clarity and succinctness.
3. The effective policy over latent state $\Pi^S$ is an $S \times S \times A$ representation of the matrix $\phi \pi$, where $\phi$ is the $S \times \Omega$ observation function, and $\pi$ is the $\Omega \times A$ observation policy. Formally, $\Pi^S_{s,s’,a} = \delta_{s,s’} \phi_{s,o} \pi_{o,a}$.

Reviewer comment (Theory):
There are issues in many definitions in Section 3.1. As mentioned in the background section, the value functions of partial observable problems generally depend on the agent's entire history. But in Section 3.1, the Q-function is given as a function of the observation and action without any justification. This observation-action Q-function may be some useful approximation, but it is used throughout the paper as the correct value estimation without any discussion. Several variables of conditional probabilities like Pr$(s|\omega)$ are presented as given constants, but they are not provided by the observation model like Pr$(\omega|s) = O(\omega|s)$; the probability Pr$(s_t|\omega_t)$ depends on the policy and the time instant $t$ in general.

Author response:
$ $

-----
### Value definitions

$ $

> The value functions of partial observable problems generally depend on the agent's entire history. But in Section 3.1, the Q-function is given as a function of the observation and action without any justification. This observation-action Q-function may be some useful approximation, but it is used throughout the paper as the correct value estimation without any discussion.

$ $

The MC Q-function is indeed well defined, but it is distinct from the Q-function conditioned on the agent’s full history. We define $Q(\omega,a)$ in the usual way: the expected discounted return starting from observation $\omega$, taking action a, and following the current policy thereafter. This expectation averages over all possible agent histories consistent with reaching observation $\omega$. We do not claim that this is equivalent to either the state value function $Q(s,a)$ or the history value function $Q(h,a)$, but it is nevertheless “correct” in that MC is an unbiased estimator of $Q(\omega,a)$, whereas TD($\lambda$) is not, for any $\lambda < 1$.

$ $

-----
### Clarifying conditionals

$ $

> Several variables of conditional probabilities like $Pr(s|\omega)$ are presented as given constants, but they are not provided by the observation model like $Pr(\omega|s)=O(\omega|s)$; the probability $Pr(s_t|\omega_t)$ depends on the policy and the time instant $t$ in general.

$ $

We now describe in Appendix A that $P(s|\omega)$ reflects the average of $P(s_t|\omega_t)$ over all timesteps, weighted by visitation probability and discounted by $\gamma$. This is a well-defined stationary quantity, and it can be computed as follows. First solve the system $Ax = b$ to find the discounted state occupancy counts $x = c(s)$, where $A = (I - \gamma (T^\pi)^\top)$ accounts for the policy-dependent state-to-state transition dynamics $T^\pi$, and $b = P_0$ is the initial state distribution over $s$. Then $P(s|\omega) \propto c(s) * P(\omega|s)$, so we can just multiply these terms together and renormalize.

We will release the code alongside the paper so that these and other such computations will be clear.

Reviewer comment (Theory):
The theoretical analysis in Appendix A seems incorrect. Below the first big equation in Appendix A, it is claimed that one can bootstrap by replacing the last line with a term with $Q(\omega_2, a_2)$ when $n=2$, but the claim means that 
$$\sum_{w_2, a_2} P(\omega_2 |s_2) P(a_2|\omega_2) Q(\omega_2, a_2) = \sum_{w_2, a_2} P(\omega_2 |s_2) P(a_2|\omega_2)
\sum_{s_3, r_2} P(s_3|s_2, a_2)P(r_2|s_2, a_2, s_3) r_2$$
This equation does not look correct. Even if what the authors mean is to include all the later terms in $...$, the bootstrap does not seem to hold given the time-homogenous definition of the Q-function unless with further justification.

Author response:
$ $

-----
### Clarifying bootstrapping

$ $

> Below the first big equation in Appendix A, it is claimed that one can bootstrap by replacing the last line with a term with Q(ω_2,a_2) when n=2... the bootstrap does not seem to hold given the time-homogenous definition of the Q-function unless with further justification.

$ $

Yes, we did mean for all the latter terms in "$\dots$" to be included. We rewrote the equations to make this clearer. We agree that using the time-homogenous Q-function bootstrap as is given here results in a quantity not equal to $\mathbb{E}[G_n|ω_0,a_0]$; that corresponds to the fact that the Q-value estimates are biased. The Q-values are commonly used this way, and doing so does not invalidate the results.

Reviewer comment (Theory):
There are some other likely errors in Appendix A. There is a missing summation over $l$ in the second big equation in Appendix A. In the equation for $Q^\lambda$, $Q_n$ is replaced by the equation derived above, but the inner $Q_n$ in that equation is directly replaced by $Q^\lambda$ without any argument.

Author response:
$ $

-----
### Fixed inaccuracies

$ $

> There is a missing summation over l in the second big equation in Appendix A. In the equation for Q^λ, Q_n is replaced by the equation derived above, but the inner Q_n in that equation is directly replaced by Q^λ without any argument.

$ $

Thank you for catching that. Our original end result is correct, though our definition was not. The equation that defines $Q_n$ should take as an argument a matrix for bootstrapping $Q$. Then, the equation that previously defined $Q^\lambda$ should instead define the $Q^\lambda$ update rule as $Q\gets (1-\lambda)\sum_{n=1}^{\infty}Q_n(Q)$. We define the matrix $Q^\lambda$ as the fixed point of the $Q^\lambda$ update rule, and then we proceed to compute it. See Appendix A in the updated paper.

$ $

Reviewer comment (Theory):
Definition 1 defines $\lambda$-discrepancy by an unspecified norm. It is revealed in Appendix E.2 that the norm is l2, but this is not consistent with the proof in Appendix B where the $\lambda$-discrepancy is a weighted difference between the two Q-functions without squares.

Author response:
$ $

-----
### Norm

$ $

> Definition 1 defines λ-discrepancy by an unspecified norm. It is revealed in Appendix E.2 that the norm is l2, but this is not consistent with the proof in Appendix B where the λ-discrepancy is a weighted difference between the two Q-functions without squares.

$ $

Thank you for noticing this. We define the norm to be a weighted $L^2$ norm, with the weighting dependent on the policy. We have corrected the notation throughout the paper to be consistent with this. The weighted $L^2$ norm is equivalent to taking the dot product between a policy-dependent weight matrix and the difference of the value functions, then taking the $L^2$ norm of that. Because the weight matrix is also an analytic function of $\pi$ and the $L^2$ norm, as a norm, is positive-definite, the normed $\lambda$-discrepancy is 0 if and only if the value function difference is 0; thus the result in Appendix B still holds.

Reviewer comment (Experiments):
Even if some of the analysis is true, they only provide properties when $\lambda$-discrepancy is zero. They still cannot explain whether decreasing $\lambda$-discrepancy can reduce the issue of partial observability in some sense. This kind of property may be explored via numerical experiments, but it is not done in the paper.

Author response:
$ $

-----
### Correlation with value error

$ $

> Even if the some of the analysis is true, they only provide properties when λ-discrepancy is zero. They still cannot explain whether decreasing λ-discrepancy can reduce the issue of partial observability in some sense. This kind of property may be explored via numerical experiments, but it’s not done in the paper.

$ $

Minimizing $\lambda$-discrepancy reduces the value error due to partial observability. To provide additional evidence for this claim, we show in Appendix F that $\lambda$-discrepancy is positively correlated with value error in every domain we considered.

Reviewer comment (Evaluation):
In the numerical experiments, the proposed method achieves the optimal performance only in the cheese environment. Since increasing the memory size might lead to optimal policy in principle, the inability to achieve optimal performance with more memory in many environments makes the ability of the proposed methods questionable.

Author response:
$ $

-----
### Comparison with optimal

$ $

> In the numerical experiments, the proposed method achieves the optimal performance only in the cheese environment. Since increasing the memory size might lead to optimal policy in principle, the inability to achieve optimal performance with more memory in many environments makes the ability of the proposed methods questionable.

$ $

The optimal performance is computed by a belief-state planner, which means the agent has access to a perfect environment model *and* enough memory bits to distinguish between all possible belief states. By contrast, our method only requires observable quantities, and works even for agents with limited-memory. Nevertheless, we do see that increasing the memory size improves performance.

Reviewer comment (Experiments):
This paper mentions that the POMDP problem is inherently complex. While memory augmentation is shown to improve performance, it doesn't directly address the computational complexity of solving POMDPs. The efficiency of the proposed approach in more challenging POMDP scenarios is not fully explored.

Author response:
The computational and sample complexity of computing the $\lambda$-discrepancy are asymptotically the same as computing either of the two Q functions. The algorithm we present is only intended to provide an initial demonstration of the concept; it is not intended to be optimal. Our paper is focused solely on showing (both theoretically and empirically) that minimizing the $\lambda$-discrepancy leads to useful memory functions. The approach is applicable to any POMDP for which we can obtain an accurate observation-based Q function, but obtaining such a Q function is the subject of a large body of work and is outside the scope of this paper.

Reviewer comment (Experiments):
It is still not clear to me why we need to reduce the $\lambda$-discrepancy, and at least an improved result should be shown by using this technique compared to using a single $\lambda$.

Author response:
We are not aware of any method that can learn a memory function from a fixed TD($\lambda$) value function. TD($\lambda$) implicitly makes a Markov assumption for all $\lambda<1$, but there is no way to determine from a single value function whether the Markov assumption is a reasonable one. This is precisely why we need the second $\lambda$. The discrepancy between value functions reliably reveals partial observability and allows us to find a memory function that corrects for it.

Reviewer comment (Theory):
The condition in Lemma 2 is a sufficient condition; can a necessary condition be provided? That may help in designing the memory function.

Author response:
Yes, we have the following necessary condition: Given a fixed policy and $λ_1$, $λ_2$ values, for almost all discount factor values $\gamma$, the converse holds, namely the condition in Lemma 2 implies that there is zero $\lambda$-discrepancy. We have added this to the updated paper, and included the proof in Appendix D.

Reviewer comment (Theory):
It seems that you approximate Equation (7) with Equation (10), what is the justification for that?

Author response:
Thank you for bringing this up. The approximation was not central to any of the claims in the paper, so we have since removed Equation (10) from the work and replaced all experimental results that used the approximation with results computed exactly, following Equation (7). This only affected the results for the gradient-based optimization in Section 4.2, and the new results match or exceed the previous performance.

Reviewer comment (Theory):
If I am understanding this correctly, the proposed metric is flawed in that errors in the Q function (not the state estimator) will result in a $\lambda$-discrepancy. For example, given a perfect state estimator $M^*(o_1, o_2, ..., o_n) = s_n$ and two imperfect Q functions $Q_\theta, Q_\phi$: $ \lVert Q_{\theta}^{\lambda_1}(s_n) - Q_{\phi}^{\lambda_2}(s_n) \rVert > 0$. In this case, I do not see how $\lambda$-discrepancy is a better metric of Markovness than the traditional TD(0) Q learning error.

Author response:
In POMDPs, value error comes from two places: estimation error (as in MDPs) and partial observability. Any method that estimates a value function will have some estimation error, and good value estimation approaches will hopefully drive this error to zero over time. The results we present here show that even if we achieve zero estimation error (which is the case for our value function oracle), there would still be error due to partial observability. Whereas Bellman error is currently the main method we use to reduce the former, it does nothing to reduce the latter. Using the $\lambda$-discrepancy, we are able to reduce the error due to partial observability.

To provide further evidence of this, we show in Appendix F that $\lambda$-discrepancy is positively correlated with value error in every domain we considered.