Mitigating Partial Observability in Decision Processes via the Lambda Discrepancy

Authors: cameron  allen
u c berkeley
aaron  kirtland
yu  ruo
  tao
sam  lobel
daniel  scott
nicholas  petrocelli
omer  gottesman
ronald  parr
michael l littman
george  konidaris
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:50:33

Abstract: Reinforcement learning algorithms typically rely on the assumption that the environment dynamics and value function can be expressed in terms of a Markovian state representation. However, when state information is only partially observable, how can an agent learn such a state representation, and how can it detect when it has found one? We introduce a metric that can accomplish both objectives, without requiring access to---or knowledge of---an underlying, unobservable state space. Our metric, the λ-discrepancy, is the difference between two distinct temporal difference (TD) value estimates, each computed using TD(λ) with a different value of λ. Since TD(λ=0) makes an implicit Markov assumption and TD(λ=1) does not, a discrepancy between these estimates is a potential indicator of a non-Markovian state representation. Indeed, we prove that the λ-discrepancy is exactly zero for all Markov decision processes and almost always non-zero for a broad class of partially observable environments. We also demonstrate empirically that, once detected, minimizing the λ-discrepancy can help with learning a memory function to mitigate the corresponding partial observability. We then train a reinforcement learning agent that simultaneously constructs two recurrent value networks with different λ parameters and minimizes the difference between them as an auxiliary loss. The approach scales to challenging partially observable domains, where the resulting agent frequently performs significantly better (and never performs worse) than a baseline recurrent agent with only a single value network.

Main paper: 
 introduction: 
 The dominant modeling frameworks for reinforcement learning [Sutton and Barto, 2018] define environments in terms of an underlying Markovian state representation. This modeling choice, called the Markov assumption, is nearly ubiquitous in reinforcement learning, because it allows environment dynamics, rewards, value functions, and policies all to be expressed as functions that are independent of the past given the most recent state. In principle, an environment can be modeled as either a Markov decision process (MDP) [Puterman, 1994] , or its partially observable counterpart, a POMDP [Kaelbling et al., 1998] , as long as the underlying state representation contains enough information to satisfy the Markov assumption. The POMDP framework is more general, but POMDPs are typically much harder to solve than MDPs [Zhang et al., 2012] , so it is important to know when it is appropriate to use the simpler MDP framework. Ideally, a system designer can ensure that their reinforcement learning agent is configured to use the appropriate problem model. If an environment is partially observable, the designer may manually add relevant decision-making features to the agent's state representation, either by concatenating observations together [Mnih et al., 2015] or by using other types of feature engineering [Bellemare et al., 2020 , Galataud, 2021 , Tao et al., 2023] . Alternatively, the designer can manually specify a set of possible environment states over which the agent can maintain a Markovian belief distribution [Kaelbling et al., 1998 ]. The challenge is that it is not always obvious when the designer has supplied sufficient information to satisfy the Markov assumption. These approaches require the person deploying the agent to know details about the very problem the agent is supposed to help them solve.An alternative approach-that we explore in this paper-is to let the agent learn a good state representation for the problem it is solving. The conventional deep-learning-era wisdom is that the best problem representations come from training a system to solve a task, rather than from human designers. When faced with a potentially partially observable environment, if we provide the agent with a large enough recurrent neural network (RNN), it can perhaps discover an internal state representation "endto-end" that maximizes reward via gradient descent [Bakker, 2001 , Hausknecht and Stone, 2015 , Ni et al., 2022 , Dong et al., 2022] . Indeed, end-to-end RNN architectures work well for many problems and are both general-purpose and scalable. However, these techniques only implicitly address the problem of learning a Markovian state representation; in this paper, we show that we can achieve much better learning performance when we explicitly tackle the problem of partial observability.In our approach, the agent learns a state representation by directly minimizing a metric that assesses whether the environment is fully or partially observable. We call our metric the λ-discrepancy, and define it as the difference between two distinct value functions estimated using temporal difference learning (TD) [Sutton, 1988] . Specifically, the TD(λ) method defines a smooth trade-off between one-step TD (i.e., λ = 0), which makes an implicit Markov assumption, and Monte Carlo (MC) estimation (λ = 1), which does not, and intermediate λ values interpolate between these extremes. By comparing value estimates for two distinct values of λ, we can check that the agent's observations support Markovian value prediction, and augment them with memory if we find they are incomplete.Our main contributions4 are as follows:1. We introduce and formally define the λ-discrepancy and prove that it is exactly zero for MDPs and almost always non-zero for a broad class of POMDPs that we characterize. This analysis tells us that our metric reliably detects partial observability.2. We then consider a tabular, proof-of-concept experiment that adjusts the parameters of a memory function to minimize the λ-discrepancy via gradient descent. For this experiment, we compute the λ-discrepancy in closed form. Our results demonstrate that minimizing the λ-discrepancy is a viable path to mitigating partial observability.3. Finally, we integrate our approach into a deep reinforcement learning agent and evaluate on a set of large and challenging POMDP benchmarks. We find that minimizing the λdiscrepancy between two learned value functions, as an auxiliary loss alongside traditional reinforcement learning, is often significantly more effective (and never worse) than training the same agent with just a single value function.Overall, we find that the λ-discrepancy is a simple yet powerful metric for detecting and mitigating partial observability. The metric is also practical, since it can be computed directly from value functions, which are ubiquitous in reinforcement learning. Furthermore, such value functions need only consider observable parts of the environment, so the λ-discrepancy remains applicable even without the common assumption that the agent knows the full set of possible POMDP states. 
background: 
 We consider two frameworks for modeling sequential decision processes: MDPs and POMDPs. The MDP framework [Puterman, 1994] consists of a state space S, action space A, reward function R : S×A → R, transition function T : S×A → ∆S mapping to a distribution over states, discount factor γ ∈ [0, 1], and initial state distribution p 0 ∈ ∆S. The agent's goal is to find a policy π S : S → ∆A that selects actions to maximize return, g t , the discounted sum of future rewards starting from timestep t: g π S t = ∞ i=0 γ i r t+i , where r i is the reward at timestep i. We denote the expectation of these returns as value functionsV π S (s) = E π S [g t | s t = s] and Q π S (s, a) = E π S [g t | s t = s, a t = a].The POMDP framework [Kaelbling et al., 1998] additionally includes a set of observations Ω and an observation function Φ : S → ∆Ω that describes the probability Φ(ω|s) of seeing observation ω in latent state s. POMDPs are a more general model of the world, since they contain MDPs as a special case: namely, when observations have a one-to-one correspondence with states. Similarly, POMDPs where states correspond to disjoint sets of observations are called block MDPs [Du et al., 2019] . However, such examples are rare; in typical POMDPs, a single observation ω does not contain enough information to fully resolve the state s. While agents need not fully resolve the underlying state to behave optimally, they must retain at least enough information across timesteps that the optimal policy becomes clear.We are interested in the learning setting, where the agent has no knowledge of the underlying state s nor even the set of possible states S (let alone the transition, reward, and observation functions). It receives an observation ω t ∈ Ω at each timestep and must find a way to maximize expected return. One way to do this is to construct a state representation, perhaps using some form of memory, on which it can condition its behavior. A state representation z ∈ Z is Markovian if at any timestep t, the representation z t and action a t together are a sufficient statistic for predicting the reward r t and next representation z t+1 , instead of requiring the agent's whole history:Pr(z t+1 , r t |z t , a t ) = Pr(z t+1 , r t |z t , a t , . . . , z 0 , a 0 ).(1)The definition can be applied to states (s), observations (ω), and even memory-augmented observations (defined in Sec. 4), by replacing z with the relevant quantity. States and observations are equivalent in MDPs, and this property is satisfied for both by definition, but in POMDPs it typically only holds for the underlying, unobserved state s-not the observations.Markovian state representations have several desirable implications. First, if the Markov property holds then so does the Bellman equation: V π S (s) = a∈A π S (a | s) R(s, a) + γ s ′ ∈S T (s ′ | s, a)V π S (s ′ ) . The Bellman equation allows agents to estimate the value of a policy, from experiences, and without knowing T or R, via a recurrence relation over one-step returns,EQUATIONwhich converges to the unique fixed point V π S . A second implication is that the transition and reward functions, and consequently the value functions V π S and Q π S , have fixed-sized inputs and are therefore easy to parameterize, learn, and reuse. Finally, it follows from the Markov property that the optimal policy π * S can be expressed deterministically and does not require memory [Puterman, 1994] . We can unroll the Bellman equation over multiple timesteps to obtain a similar estimator that uses n-step returns:V π S (s) = E π S [g t:t+n | s t = s], where g t:t+n := r t + γr t+1 + γ 2 r t+2 + • • • + γ n V π S (s t+n ), with g t:t+n := g t if the episode terminates before t + n has been reached. The same process works for weighted combinations of such returns, including the exponential average:V λ π S (s) = E π S (1 -λ) ∞ n=1 λ n-1 g t:t+n s t = s , (3) with V λ=1 π S (s) = E π S [g t | s t = s]as a special case. Equation (3) defines the TD(λ) value function as an expectation over the so-called λ-return [Sutton, 1988] . Low values of λ give more weight to shorter returns where the value term V π S (s t+n ) is discounted less (and thus has greater significance), whereas larger λ values put more weight on longer returns where the value term is more heavily discounted. Given an MDP and a fixed policy, the recurrence relations for all TD(λ) value functions share the same fixed point for any λ ∈ [0, 1]; however, if the Markov property does not hold, different λ may have different TD(λ) fixed points. In this work, we seek to characterize this phenomenon and leverage it for detecting and mitigating partial observability. Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We describe limitations in Appendix A. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
detecting partial observability: 
 Before introducing our partial observability metric, let us first consider the T-maze example of Figure 1 under our two candidate decision-making frameworks: MDPs and POMDPs. In the T-maze, the agent only observes the color of its current square, and must remember the start square color (sampled uniformly from BLUE/RED) to select the action at the junction that leads to the positive reward.To model the T-maze as a POMDP, the state must include the agent's position and the goal location. The transitions are deterministic; each action moves the agent in the corresponding direction. The goal location is sampled at the start of an episode, and defines both the reward function and the observation for the initial square. The observation function uniquely identifies these starting states, but all corridor states are aliased together (they map to the same observation), as are the two junction states.We can convert our POMDP model into an MDP model by using observations Ω as the MDP's "states". This requires new transition and reward functions, T Ω and R Ω , which we define as:T Ω (ω ′ | ω, a) := s,s ′ ∈S Φ(ω ′ | s ′ )T S (s ′ | s, a) Pr(s | ω) and R Ω (ω, a) := s∈S R S (s, a) Pr(s | ω), where Pr(s | ω) is policy-dependent and describes how each hidden state s i ∈ S contributes to the overall environment behavior when we see observation ω. While there are several sensible choices of weighting scheme for this type of state aggregation (uniform over states, relative frequency, etc.), only one of these choices coincides with the distribution Pr(s | ω). That particular weighting scheme is the one that averages the time-dependent Pr(s t | ω t ) over all timesteps, weighted by visitation probability under the agent's policy, and discounted by γ. 5 We explain how to construct this expression in Appendix C. We call the MDP defined in this way the effective MDP model for a given POMDP.The effective MDP model marginalizes over histories (and POMDP hidden states). For example, the model predicts that going UP from the junction will reach the goal exactly half the time, because the transition dynamics marginalize over (equally likely) starting observations. Note that the POMDP model does no such averaging: if the agent initially observes BLUE, then going UP from the junction will always reach the goal, but DOWN never will. Thus, we see that, for the T-maze, there is a mismatch between the POMDP model and the effective MDP model. In other words, the POMDP's hidden states S are a Markovian representation, but its observations Ω are not, despite the fact that the effective MDP model treats them as Markovian "states".In principle, an agent could measure partial observability by simultaneously modeling its environment as both a POMDP and an MDP, and comparing the models' predictions. Unfortunately, since the agent lacks any information about the unobserved state space S, the POMDP model would require variable-length history inputs at each time step, and would come with computational and memory costs that grow exponentially with history length. Instead, we propose a model-free approach, using value functions conditioned only on observations, to approximate this comparison in a tractable way.The Bellman equation and its sample-based recurrence relations (2) and (3) are defined for Markovian states. If we apply them to the observations of a POMDP, we are actually working with the effective MDP model of that POMDP, instead of the POMDP itself. To see this, consider one-step TD (λ = 0), where we use the same recurrence relation (2) but now our expectation is sampling from the POMDP:EQUATIONwhere we have suppressed V Ω 's dependence on the observation-based policy π for ease of notation.We see from Equation (4) that the value function TD computes for a POMDP is the fixed point of the Bellman operator for the effective MDP model. 6 By contrast, the Monte Carlo value estimator (λ = 1) does not exploit the Bellman equation at all; it simply averages over returns:V λ=1 π (s) = E π [g t | s t = s].Translating the Monte Carlo estimator to POMDPs merely requires one additional expectation to convert from states to observations:V λ=1 Ω (ω) = E π [g t | ω t = ω] = s∈S Pr(s | ω) E π [g t | s t = s] = s∈S Pr(s | ω)V λ=1 S (s).(EQUATIONThis means MC (λ = 1) effectively takes the hidden-state value function V S of the POMDP and projects it into observation space, whereas TD (λ = 0) directly computes the value function for the projected model as an MDP, by treating observations as states. Interpolating λ between 0 and 1 smoothly varies the value estimate's reliance on the effective MDP model.For generality, we derive an expression (see Appendix C) for Q λ π -values in terms of a given λ parameter (expressed in tensor notation for compactness) to reveal how λ trades off between the projected state value function and the value function of the projected model:EQUATIONwhere the tensor product . . contracts two indices instead of one,7 with tensors defined as follows: Let us take a moment to parse this equation. First, Π S and ΦW Π are mappings from states to stateaction pairs. We call the former the MC policy spread, Π S s,s ′ ,a = 1[s = s ′ ] ω Φ s,ω π ω,a , which maps states to the expected policy under their observation distribution. We call the latter the TD policy spread, (ΦW Π ) s,s ′ ,a = ω Φ s,ω π ω,a W ω,s ′ , which reallocates the policy probabilities for a given observation ω across all states s ′ that produce that observation, weighted by W . K λ π is a convex combination of these two policy spread tensors, parameterized by λ. Multiplying a policy spread tensor on the left by T produces an (S ×A×S ×A) transition-policy tensor, describing the probability of each state-action transition (s, a) → ∆(s ′ , a ′ ) under the policy. Intuitively, Equation (6) says that TD(λ) computes Q-values for a POMDP as though state-action pairs are evolving according toQ λ π (Ω × A) is a matrix of Q-values; W (Ω × S)T K λ π = λT Π S + (1 -λ)T ΦW Π ,which is a mixture of the policy dynamics under two transition models: the MC transition-policy (T Π S ) and the TD transition-policy (T ΦW Π ). The expression (I -γT K λ π ) -1 . . R computes the state-space Q-values for this hybrid transition model. Finally, these are projected through W to compute the observation-space Q-values. (See Appendix C for more details.)We have shown that, under partial observability, Q λ π value functions may differ for different λ parameters, due to varying reliance on the effective MDP model in the TD(λ) estimator. We call this difference the λ-discrepancy, and we propose to use it as a measure of partial observability. Φ = (1 -p) • Φ Perfect + p • Φ Aliased .Definition 1 For a given POMDP model P and policy π, the λ-discrepancy Λ λ1,λ2 P,π is the weighted norm of the difference between the Q-functions estimated by TD(λ) for λ ∈ {λ 1 , λ 2 }:Λ λ1,λ2 P,π := Q λ1 π -Q λ2 π = W I -γT K λ1 π -1 -I -γT K λ2 π -1 . . R SA .The choice of norm can be arbitrary, as can the norm weighting scheme, as long as it assigns positive weight to all reachable observation-action pairs. We discuss choices of weighted norm in Appendix G.2. For brevity, we suppress the P subscript when the POMDP model is clear from context.A useful property (that we will prove in Theorem 2) is that if the POMDP has Markovian observations, the λ-discrepancy is exactly zero for all policies. However, for it to be a useful measure of partial observability, we must also show that the λ-discrepancy is reliably non-zero when observations are non-Markovian. For this, we have the following theorem.Theorem 1 Given a POMDP model P and distinct λ ̸ = λ ′ , if there exists a policy π : Ω → ∆A such that Λ λ,λ ′ P,π ̸ = 0, then Λ λ,λ ′ P,π ̸ = 0 for all policies except at most a set of measure zero.Proof sketch: We formulate the λ-discrepancy of a given policy as the norm of an analytic function, then use the fact that analytic functions are zero everywhere or almost nowhere, along with the fact that norms preserve this property. The full proof is given in Appendix D.Intuitively, Theorem 1 says that the λ-discrepancy can detect non-Markovian observations. If it is possible to reveal that a POMDP's observation-space value function does not match that of its effective MDP model, then almost all policies will do so. The theorem further suggests that even if a particular policy has zero λ-discrepancy, small perturbations to that policy will almost surely detect partial observability if it is present.We also find that increasing amounts of partial observability lead to larger λ-discrepancy. In Figure 3 we interpolate between Markovian and aliased T-Maze environments, and consider three types of state aliasing. Here we use a fixed policy that goes RIGHT until the junction, then UP with probability 2/3 and DOWN w.p. 1/3. To avoid artifacts due to interactions between the discount factor, weighted norm, and changing observation function, we set γ ≈ 1 and use the (policy-weighted) maximum norm (see Appendix G.2). 8 The results confirm that the λ-discrepancy is a useful indicator of partial observability, provided that it is non-zero for at least one policy.What conditions cause the λ-discrepancy to be zero?We can characterize the cases in which the λ-discrepancy is zero for all policies by inspecting Definition 1. Because norms are positive definite, it suffices to consider the expression inside the norm. The only ways for this expression to equal zero are either when the difference between policy spread tensors K λ1 π -K λ2 π is zero (which we will show implies Markovian observations), or it is non-zero but is projected away by T , R SA , and/or W . We first consider when the policy spread tensors-which are the only terms that depend on λ-are equal, i.e. K λ1 π = K λ2 π .Theorem 2 For any POMDP P and any λ 1 ̸ = λ 2 , K λ1 π = K λ2 π if and only if P is a block MDP.Proof sketch: Using Eq. ( 6), K λ1 π = K λ2 π can be simplified to Π S = ΦW Π , which is satisfied if and only if ΦW = I, i.e. each observation is generated by exactly one hidden state. Proof in Appendix E.Thus, when the policy spread tensors are equal, it means there is no state aliasing and P is a block MDP. Now consider the cases where the difference between K λ π is projected away by T , R SA , and/or W , starting from the policy spread tensors in Definition 1 and working our way outwards.If the policy spread tensors differ, the transition tensor may project these differences away (T K λ1 π = T K λ2 π ), for example, if the transition probabilities leading to the aliased states always occur in some fixed proportion. In such cases, we can collapse the aliased states to non-aliased ones without any loss of information by reformulating the incoming transition probabilities to the aliased states as stochastic transition dynamics (and possibly rewards) after the non-aliased states. In other words, P is equivalent to an MDP. We prove this claim and provide an example of such an environment in Appendix F. Finally, differences between transition-policies may be projected away by the state-blending weights W , the reward function R SA , or both. W (I -γT K λ1 π ) -1 . . R SA = W (I -γT K λ2 π ) -1 . . R SA . In this case, the effective MDP may be lossy, which is a limitation of the λ-discrepancy metric. Fortunately, such environments appear to be rare, and they are easy to handle, as we show in the example below.Parity Check Environment. The Parity Check POMDP of Figure 4 has four equally likely starting states, each associated with a pair of colors that the agent will see during the first two timesteps. At the subsequent (white) junction state, the agent must decide based on these colors whether to go UP (black arrow) or DOWN (white arrow). The rewards are defined such that UP is optimal if and only if the color family matches (i.e. RED → PINK or BLUE → CYAN).The transition-policies for the first eight states differ with and without aliasing, but the differences are perfectly symmetric with respect to rewards and observations.9 Every observation has zero expected return, and so the λ-discrepancy is zero for all policies. Our analysis in Appendix H.5 suggests such edge-cases are the exception and not the rule. Even so, such examples are easy to handle: we can simply add a small amount of memory. If the environment is truly an MDP, adding memory should have no effect, but if it is a POMDP, a randomly initialized memory function (defined in the next section) may be enough to break symmetry and produce a λ-discrepancy (as seen in Figure 4 , center).Taken together, Theorems 1 and 2 suggest that the λ-discrepancy could be a useful measure for detecting partial observability. In the next section, we demonstrate the efficacy of using the λdiscrepancy to learn memory functions that mitigate partial observability. 
memory learning with the λ-discrepancy: 
 The λ-discrepancy can identify that we need memory, but it cannot yet tell us what to remember. For that, we must replace the POMDP P in Definition 1 with a memory-augmented version. In general, Performance is calculated as the expected start-state value, and is normalized between a random policy (y = 0) and the optimal belief state policy (y = 1) found with a POMDP solver [Cassandra, 2003] . Error bars are 95% confidence intervals over 30 seeds. a memory can be any mapping from agent histories (ω 0 , a 0 , ..., a t-1 , r t-1 , ω t ) to internal memory states m within some set of memory states M. For practical reasons, we restrict our attention to recurrent memories that update an internal representation incrementally from fixed-size inputs.We define a memory function µ : Ω × A × M → M as a mapping from an observation, action, and memory state, (ω, a, m), to a next memory state m ′ = µ(ω, a, m). Given a POMDP P, a memory function µ induces a memory-augmented POMDP P µ , with extended states S M = S × M, actions A M = A × M, and observations Ω M = Ω × M. The augmented transition dynamics T M preserve the original transition dynamics T for states S while allowing the agent full observability and control over memory states M. Memory functions naturally lead to memory-augmented policies π µ : Ω M → ∆A M and value functions V π,µ : Ω M → R and Q π,µ : Ω M × A M → R that reflect the expected return under such policies. For details, see Appendix G.1.The λ-discrepancy (Definition 1) applies equally well to memory-augmented POMDPs, and can thus be used to determine whether a particular memory function µ induces a POMDP P µ with a Markovian observation space Ω M . We can also use it as a training objective for learning such a memory function. We conduct a proof-of-concept experiment on several classic POMDPs for which we can obtain closed-form gradients of the λ-discrepancy with respect to a parametrized memory function. In each domain, we randomly generate a set of stochastic policies, select the one with maximal λdiscrepancy Λ 0,1 P , and adjust the parameters of a memory function µ to minimize Λ 0,1 P µ via gradient descent. Figure 5 shows the improvement in policy gradient performance due to the resulting memory function for increasing memory sizes. The details of this experiment are provided in Appendix H. We also run the same experiment on the Parity Check example, and provide results in Figure 4 (right).We see that the λ-discrepancy can help mitigate partial observability, but it is somewhat inconveniently defined in terms of the closed-form value function fixed-points. Fortunately, with the appropriate choice of weighted norm, we can estimate the λ-discrepancy purely from sampled observation-action pairs (or, in the case of memory, observation-memory-action tuples):EQUATION)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. 
a scalable, online learning objective: 
 So far, we have shown that the λ-discrepancy can detect partial observability in theory and can mitigate it under certain idealized conditions. Now we demonstrate how to integrate our metric into samplebased deep reinforcement learning to solve problems requiring large, complex memory functions.To minimize the λ-discrepancy, we augment a recurrent version of the proximal policy optimization (PPO) algorithm [Schulman et al., 2017] with an auxiliary loss. We use recurrent PPO as our base algorithm due to its strong performance in many POMDPs [Ni et al., 2022] , and since the λdiscrepancy is a natural extension of generalized advantage estimation [Schulman et al., 2016] , which is used in PPO. In this algorithm, a recurrent neural network [Amari, 1972] (specifically a gated recurrent unit, or GRU [Cho et al., 2014] ) is used as the memory function µ that returns a latent state representation given previous latent state and an observation. This latent representation is used as input to an actor network to return a distribution over actions, as well as a critic. The critic is usually a value function network which learns a truncated TD(λ) value estimate for its advantage estimate.To estimate the λ-discrepancy, we learn two TD(λ) value function networks with different λ, parameterized by θ V,1 and θ V,2 respectively, and minimize their mean squared difference as an auxiliary loss:EQUATIONwhere z t = µ θRNN (ω t , z t-1 ) is the latent state output by the GRU, and θ represents all parameters.We train this neural network end-to-end with the standard PPO actor-critic losses and backpropagation through time [Mozer, 1995] . Full details of the algorithm are provided in Appendix I.1.Any algorithm that uses value functions could, in principle, also leverage the λ-discrepancy. However, the theoretical results in Section 3 only hold for value function fixed points and not necessarily for estimated values prior to convergence. Inaccurate value functions may lead to an inaccurate λdiscrepancy. Fortunately, value functions often provide a useful learning signal before convergence, and we observe in the following sections that the λ-discrepancy is also useful during learning.We evaluate our approach on a suite of four hard partially observable environments that require complex memory functions. Battleship [Silver and Veness, 2010] requires reasoning about unknown ship positions and remembering previous shots. Partially observable PacMan [Silver and Veness, 2010] requires localization within the map while tracking dot status and avoiding ghosts with only short-range sensors. RockSample (11, 11) and RockSample (15, 15) [Smith and Simmons, 2004] have stochastic sensors and require remembering which rocks have been sampled. While these environments were originally used to evaluate partially observable planning algorithms in large-scale POMDPs [Silver and Veness, 2010] , we use them to test our sample-based learning algorithms due to the complexity of their required memory functions.10 See Appendix I.2 for more details.We conduct experiments with regular PPO, recurrent PPO, and our λ-discrepancy-augmented recurrent PPO in Figure 6 . We also visualize the agent's memory for the partially observable PacMan environment by reconstructing the dot locations from the RNN latent state to show where the agent "thinks" it has been (see Appendix I.5). We performed a hyperparameter sweep for each method and report learning curves for undiscounted return (see Appendix I.3 for discounted learning curves and Appendix I.4 for additional experimental details).The λ-discrepancy objective leads to significantly better final performance and learning rate versus recurrent and memoryless PPO in all environments. In RockSample (15, 15), the baseline agents quickly learn to exit immediately (for +10 reward), but never improve on this. By contrast, minimizing λ-discrepancy helps the agent learn the missing features that allow it to express better policies. We also run experiments on the classic POMDPs from Section 4, but due to the size of these problems, both baseline and our proposed approach performed almost optimally (see Appendix I.6).We also observe that the best performing hyperparameters from our sweep offer further evidence for the theory developed in Section 3. Our theory suggests the λ parameters should be well separated in order to see the largest λ-discrepancy. Indeed, we find that a large difference between λ 1 and λ 2 was beneficial, with optimal λs close to either 0 or 1. In addition to this, our hyperparameter sweep also includes a PPO variant with two different TD(λ) value functions but without the λ-discrepancy auxiliary loss described in Section 5.1. These agents were never selected in the sweep; using the loss in Equation 8seems to only help with performance in tested environments.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: As stated in Section 1, our code is available at https://github.com/ brownirl/lambda_discrepancy. This repository both has detailed instructions on how to conduct the sweeps, as well as the best runs over the 30 seeds mentioned in both experiments. All experiments were done by setting seeds for the pseudo-random number generator, so results are deterministic with respect to the seed and exactly reproducible. Data from our runs are also provided in the repo.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
related work: 
 There is an interesting connection between state abstraction [Li et al., 2006] , which selectively removes information from state, and partial observability mitigation, which aims to recover state from incomplete observations. Allen et al. [2021] investigated the state abstraction perspective and characterized the properties under which abstract state representations of MDPs either do or do not preserve the Markov property. Several other approaches characterize and measure partial observability and POMDPs. While POMDPs have been shown to be computationally intractable in general [Papadimitriou and Tsitsiklis, 1987] , various works have studied complexity measures [Zhang et al., 2012] and defined subclasses with tractable solutions [Littman, 1993 , Liu et al., 2022] .The most common strategies for mitigating partial observability are memory-based approaches that summarize history. Early approaches relied on discrete representations of history, including tree representations [McCallum, 1996] or finite-state controllers [Meuleau et al., 1999] . Modern approaches mostly use RNNs [Amari, 1972] trained via backpropagation through time (BPTT) [Mozer, 1995] to tackle non-Markovian decision processes [Schmidhuber, 1990] . Various approaches use recurrent function approximation to learn better state representations. One successful approach is learning a recurrent value function [Lin and Mitchell, 1993 , Bakker, 2001 , Hausknecht and Stone, 2015 ] that uses TD error as a learning signal for memory. Policy gradient methods, including PPO (which we compare to), have also been used with recurrent learning to mitigate partial observability [Wierstra et al., 2007 , Heess et al., 2015] . Model-based methods can learn a recurrent dynamics model [Hafner et al., 2020] to facilitate planning alongside reinforcement learning. These approaches learn their representations implicitly to improve prediction error, rather than explicitly to mitigate partial observability.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: These are present in Appendix H and Appendix I. Specifically, in I.4.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 We introduce the λ-discrepancy: an observable and differentiable measure of non-Markovianity suitable for mitigating partial observability. The λ-discrepancy is the difference between two distinct value functions estimated using TD(λ), for two different λ values. We characterize the λ-discrepancy and prove that it reliably distinguishes MDPs from POMDPs. We then use it as a memory learning objective and demonstrate that minimizing it in closed-form helps learn useful memory functions in small-scale POMDPs. Finally, we propose a deep reinforcement learning algorithm which leverages the λ-discrepancy as an auxiliary loss, and show that it significantly improves the performance of a baseline recurrent PPO agent on a set of large and challenging partially observable tasks.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: Throughout the work, all error bars are included for 95% confidence intervals over 30 runs, as described in Sections 4 and 5. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Details of compute usage, hardware and time taken for experiments are discussed in Appendix I.4. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 9. Code Of Ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: We followed the Code of Ethics in its entirety. No human research subjects or participants were utilized. We used no datasets that contained personally identifiable information, required participant consent, or were deprecated. We violated no copyright or fair use licenses, and representative evaluation was not a concern.Our work is of a fundamental nature that has no specific negative impact for safety as compared to other fundamental reinforcement learning works. It contains no particular threats towards security, discrimination, surveillance, deception and harassment, the environment, human rights, or bias and fairness.The POMDP environments used and the models using them are described in Section 4 and Section 5. There is no sensitive data requiring strong security measures. Standard security measures are in place as provided through GitHub. Our models do not have a high risk for misuse. Information required to understand research artifacts and reproduce all results in this paper are provided at the code repository. Legal compliance has been ensured. An Apache License v2.0 is provided for the code at https://github.com/brownirl/ lambda_discrepancy. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discuss the broader impacts of this paper in Appendix B. Our paper, as a work on the foundations of reinforcement learning, is not tied to any particular applications. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our work focuses on the foundations of reinforcement learning and does not entail a high risk for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The main theoretical limitations of our work come from the class of POMDPs that the λ-discrepancy is relevant to. The λ-discrepancy is capable of detecting non-Markovian observations in a "large" class of POMDPs, which we have described using an analytic almost-all characterization. However, there do exist POMDPs in which the fact that observations are non-Markovian cannot be detected by the λ-discrepancy, and we provide one such example in Section 3.2. The examples of such POMDPs we have found seem to exist on a "knife's edge", meaning that slightly perturbing the parameters of the POMDP results in a non-zero λ-discrepancy. Furthermore, incorporating memory into the value function seems to reliably produce a λ-discrepancy, without having to modify the underlying decision process.The recurrent deep reinforcement learning algorithm we develop is also limited in that it only approximates the λ-discrepancy. The auxiliary loss in Section 5 is only an approximation of the λdiscrepancy. It minimizes the difference in value function estimates as opposed to the difference in fixed points that we consider in Section 3. The loss function is therefore a departure from our theoretical analysis, though the empirical results suggest that the value functions need not have fully converged for the discrepancy to be useful. Future work ought to investigate why this works so well in practice.As a work studying the foundations of partially observable reinforcement learning, this paper does not have any particular negative societal applications. Application domains could include robotics and tasks such as navigation. The risk for misuse of this work are no different than any other works aimed towards algorithmic improvements for reinforcement leanring. Negative impact mitigation strategies are likewise the same as other similar works.Here we derive the fixed point of the TD(λ) action-value update rule in a POMDP, following the Markov version by Sutton [1988] . First, define the expected return given initial observation ω 0 and initial action a 0 asE π (G n |ω 0 , a 0 ) = s0 Pr(s 0 |ω 0 ) r0 Pr(r 0 |s 0 , a 0 )r 0 + γ s0 Pr(s 0 |ω 0 ) s1 Pr(s 1 |s 0 , a 0 ) ω1 a1Pr(ω 1 |s 1 )Pr(a1 |ω 1 ) r1 Pr(r 1 |s 1 , a 1 )r 1 + γ 2 s0 Pr(s 0 |ω 0 ) s1 Pr(s 1 |s 0 , a 0 ) ω1 a1 s2Pr(ω 1 |s 1 )Pr(a1 |ω 1 )Pr(s 2 |s 1 , a 1 ) * ω2 a2 Pr(ω 2 |s 2 )Pr(a 2 |ω 2 ) r2 Pr(r 2 |s 2 , a 2 )r 2 + . . .We can define the n-step bootstrapped update rule from this given a value matrix Q π by replacing part of the term with coefficient γ n with a Q π value, e.g. for the n = 2 case, we getQ (i+1) π (ω 0 , a 0 ) = s0 Pr(s 0 |ω 0 ) r0 Pr(r 0 |s 0 , a 0 )r 0 + γ s0 Pr(s 0 |ω 0 ) s1 Pr(s 1 |s 0 , a 0 ) ω1 a1 Pr(ω 1 |s 1 )Pr(a 1 |ω 1 ) * r1 Pr(r 1 |s 1 , a 1 )r 1 + γ 2 s0 Pr(s 0 |ω 0 ) s1 Pr(s 1 |s 0 , a 0 ) ω1 a1 Pr(ω 1 |s 1 )Pr(a 1 |ω 1 ) * s2 Pr(s 2 |s 1 , a 1 ) ω2 a2 Pr(ω 2 |s 2 )Pr(a 2 |ω 2 )Q (i) π (ω 2 , a 2 ).Translating these expressions into matrix notation, we haveE π (G n |ω 0 , a 0 ) = s0 W ω0,s0 R s0,a0 + γ s0 W ω0,s0 s1 T s0,a0,s1 ω1 a1 Φ s1,ω1 π ω1,a1 R s1,a1 + γ 2 s0 W ω0,s0 s1 T s0,a0,s1 ω1 a1 s2 Φ s1,ω1 π ω1,a1 T s1,a1,s2 * ω2 a2 Φ s2,ω2 π ω2,a2 R s2,a2 + . . . ,where the terms W , T , R, and Φ are as in Equation 6, and π is the Ω → ∆A policy written as an Ω × A tensor with entries in [0, 1]. In particular, W ω,s = Pr(s|ω), which averages Pr(s t |ω t ) over all timesteps, weighted by visitation probability and discounted by γ. This is a well-defined stationary quantity, and it can be computed as follows. First solve the system Cx = b to find the discounted state occupancy counts x = c(s), where b = p 0 is the initial state distribution over S, and C = (I -γ(T π ) ⊤ ) accounts for the policy-dependent state-to-state transition dynamics T π defined by T π s,s ′ = ω a Pr(ω|s) Pr(a|ω) Pr(s ′ |s, a). Then Pr(s|ω) ∝ c(s) * Φ(ω|s), so we can just multiply these terms together and renormalize. For the 2-step bootstrapped update rule above, we have:Q (i+1) π (ω 0 , a 0 ) = s0 W ω0,s0 R s0,a0 + γ s0 W ω0,s0 s1 T s0,a0,s1 ω1 a1 Φ s1,ω1 π ω1,a1 R s1,a1 + γ 2 s0 W ω0,s0 s1 T s0,a0,s1 ω1 a1 s2 Φ s1,ω1 π ω1,a1 T s1,a1,s2 * ω2 Φ s2,ω2 a2 Π ω2,a2 Q (i) π (ω 2 , a 2 ).Let Π be an Ω×Ω×A representation of the Ω×A policy π with Π ω,ω ′ ,a = 1[ω = ω ′ ]π ω,a . Likewise, let Π S be the effective policy over latent states, an S×S×A representation of the matrix Φπ, i.e.Π S s,s ′ ,a = 1[s = s ′ ](Φπ) s,a . Given a matrix of action-values Q (i)π , the n-step update rule is:Q (i+1) π = Q π,n (Q (i) π ) := W n-1 k=0 (γT Π S ) k . . R SA + γ(γT Π S ) n-1 . . T ΦΠ . . Q (i) π ,where R SA is the matrix of reward values as described in Equation 6, and we use . . to denote double tensor contraction, while all other tensor products contract a single index.We also have the standard definition of the TD(λ) update rule as:Q (i+1) π = (1 -λ) ∞ n=1 λ n-1 Q n (Q (i) π ).We are concerned with the fixed point of this update rule, which we refer to as Q λ π :Q λ π = (1 -λ) ∞ n=1 λ n-1 W n-1 k=0 (γT Π S ) k . . R SA + γ(γT Π S ) n-1 . . T ΦΠ . . Q λ π .Separating this into a reward part with factor R SA and a value part with factor Q λ π , we find that the value part is(1 -λ)W ∞ n=1 (λγT Π S ) n-1 . . γT ΦΠ . . Q λ π = (1 -λ)W (I -λγT Π S ) -1 . . γT ΦΠ . . Q λ π ,and for the reward part, we have the coefficients of R SA in the table below for values of n and kk = 0 1 2 . . . n = 1 1 . . . 2 λ λγT Π S . . . 3 λ 2 λ 2 γT Π S λ 2 (γT Π S ) 2 . . . . . . . . . . . . . . . . . .where each term is multiplied by (1 -λ)W in front. We can then see by summing over rows before columns that the reward part is:(1 -λ)W ∞ k=0 1 1 -λ (λγT Π S ) k . . R SA = W (I -λγT Π S ) -1 . . R SA .So we rewrite Q λ π as follows:Q λ π = W I SA -λγT Π S -1 . . R SA + (1 -λ)γT ΦΠ . . Q λ π . Now let W A = W ⊗I A , which is Ω × A × S × A, and W Π = Π . . W A , which is Ω × S × A.Here, ⊗ means the Kronecker product. This essentially repeats the W matrix |A| times to incorporate actions into the tensor. Note that for any S×A tensor G,(W A . . G) = W G. This is because (W A . . G) ωa = s,a ′ W A ωasa ′ G sa ′, and the only nonzero terms in this sum are those such that a = a ′ . For these indices,W A ωasa ′ = W ωs , so s,a ′ W A ωasa ′ G sa ′ = s W ωs G sa = (W G) ωa .Also, let F = I SA -λγT Π S -1 , which is S × A × S × A. Then we find:Q λ π = W A . . I -λγT Π S -1 . . R SA + (1 -λ)γT ΦΠ . . Q λ π = W A . . F . . R SA + (1 -λ)γT ΦΠ . . Q λ π = W A . . F . . R SA + W A . . F . . (1 -λ)γT ΦΠ . . Q λ π .At this point, we can subtract the second term on the right hand side from both sides, factor out Q λ π on the right, and multiply by I ΩA -(1 -λ)γW A . . F . . T ΦΠ -1 on the left of both sides to obtain:Q λ π = I ΩA -(1 -λ)γW A . . F . . T ΦΠ -1 W A . . F . . R SA = W A . . I SA -(1 -λ)γF . . T ΦΠ . . W A -1 . . F . . R SA = W I SA -(1 -λ)γF . . T ΦW Π -1 . . F . . R SA = W F + (1 -λ)γF . . T ΦW Π . . F + . . . + (1 -λ) k γ k (F . . T ΦW Π ) k . . F + . . . . . R SA ,where the last equality follows from expanding the geometric series. Now we use the identity(A -B) -1 = ∞ k=0 (A -1 B) k A -1 to find: Q λ π = W F -1 -(1 -λ)γT ΦW Π -1 . . R SA = W I -γT λΠ S + (1 -λ)ΦW Π -1 . . R SA .To recap our previous definitions, W is an Ω×S tensor, I is an S×A×S×A tensor, T is an S×A×S tensor, Π S is an S×S×A tensor, Φ is an S×Ω tensor, W Π is an Ω×S×A tensor, and R SA is an S×A tensor.Lastly, we briefly note that one can get the V λ π values by replacing W in the above equation with W Π and using a double contraction:V λ π = W Π . . I -γT λΠ S + (1 -λ)ΦW Π -1 . . R SA .We can confirm that V λ π = a π ω,a Q ω,a , by rewriting the expression on the right as follows:a π ω,a Q ω,a = a π ω,a s,a ′ W A ω,a,s,a ′ B s,a ′ = a π ω,a s W A ω,a,s,a B s,a = a π ω,a s W ω,s B s,a = s,a π ω,a W ω,s W Π ω,s,a B s,a = V λ π ,whereB = I -γT λΠ S + (1 -λ)ΦW Π -1 . . R SA is an S × A tensor.In this section we prove Theorem 1, that there is either a λ-discrepancy for almost all policies or for no policies. Fix λ and λ ′ . Recall that we define the λ-discrepancy as follows:Λ π := Q λ π -Q λ ′ π 2,π = W (I -γT K λ π ) -1 -(I -γT K λ ′ π ) -1 . . R SA ,where K λ π = λΠ S + (1 -λ)ΦW Π . Let Y be the largest open set in the space of stochastic Ω × A matrices, considered as a subset of R |Ω|(|A|-1) . Now consider the λ-discrepancy as a function of the policy π. In other words, we defineΛ : Y → R, π → Q λ π -Q λ ′ π .Let X be an open subset of R n . We say that a function f : X → R is real analytic on X if for all x ∈ X, f can be written as a convergent power series in some neighborhood of x. For this proof, we will utilize the following facts:1. the composition of analytic functions is analytic [Krantz and Parks, 2002] , 2. the quotient of two analytic functions is analytic where the denominator is nonzero, 3. a real analytic function on a domain X is either identically zero or only zero on a set of measure zero [Mityagin, 2020] .We will also use the fact that for an invertible matrix A, each entry A -1 ij is analytic in the entries of A where the entries of A yield a nonzero determinant. We can prove this by first writingA -1 = det(A) -1 adj(A) = det(A) -1 cof(A) ⊤where adj A is the adjugate of A and cof(A) is the cofactor matrix of A. Each entry of the cofactor matrix is a cofactor that is polynomial in the entries of A, and is therefore analytic in them. Therefore, each entry of A -1 is the quotient of two analytic functions and is therefore analytic except where det A = 0.Next, we will show that Λ is an analytic function. Note that the variable terms in the equation are W , W Π , R SA , T , Π S , and Φ. Of these, R SA , T , and Φ are constant with respect to π. Π S ilj = k 1[i = l]Φ ik π kj , so each entry of Π S is analytic on Y in the entries of π. Likewise, P ij = k,a Φ ik π ka T iaj is analytic on Y . Therefore, the state-occupancy counts c = p 0 + γP ⊤ p 0 + γ 2 (P ⊤ ) 2 p 0 + • • • = (I -γP ⊤ ) -1 p 0 , where p 0 contains the initial state probabilities, are the composition of analytic functions and thus analytic on Y . W ωs = Φsωcs s ′ Φ s ′ ω c s ′ is analytic on Y for the same reason, and the denominator of W ωs , s ′ Φ s ′ ω c s ′ , is nonzero for all observations able to be observed with nonzero probability. Finally, Λ is then a composition of analytic functions on Y and thus analytic itself.To handle the norm weighting, we note that w π is analytic in π as w π = (1, π(a|ω)), and the dot product of w π with Λ is also analytic. Now, we use the fact mentioned above that the zero set of a nontrivial analytic function is of measure zero. Therefore, the zero set of Λ • w π is either zero for all policies or zero only on a set of measure zero. To finish, we note that because norms are positive definite, Λ π = 0 if and only if Λ • w π = 0, so this result extends to the normed λ-discrepancy as well.In this section, we prove Theorem 2 concerning when the system is a block MDP. Recall that in Equation 6 we defineK λ π = λΠ S + (1 -λ)ΦW Π . Suppose K λ π = K λ ′ π . We can rewrite this as (λ - λ ′ )Π S -(λ-λ ′ )ΦW Π = (λ-λ ′ )(Π S -ΦW Π ) = 0. This implies that either λ = λ ′ or Π S = ΦW Π .Recall the definition of Π S and W Π in Section C. Expanding the equation Π S = ΦW Π using these definitions, we have thatΠ S s,s ′ ,a = 1[s = s ′ ](Φπ) s,a = 1[s = s ′ ] ωPr(ω|s)Pr(a|ω);(ΦW Π ) s,s ′ ,a = ω Φ s,ω W Π ω,s ′ ,a = ω Pr(ω|s)Pr(a|ω)Pr(s ′ |ω).Then by setting the rightmost expression in each equation equal and simplifying with the indicator function, we have that for all i, j, k,ω Pr(ω|s i )Pr(a k |ω)Pr(s j |ω) = ω Pr(ω|s i )Pr(a k |ω) i = j, 0 i ̸ = j.We will first consider the case where i ̸ = j. We have that for all i ̸ = j and all k, ω Pr(ω|s i ) Pr(a k |ω)Pr(s j |ω) = 0. Because each term in the sum is nonnegative, this is equivalent to the statement that for all i ̸ = j, all k, and all ω, Pr(ω|s i )Pr(a k |ω)Pr(s j |ω) = 0. Now note that for all observations ω, there exists some k ′ such that Pr(a k ′ |ω) > 0. Therefore, we have that for all i ̸ = j and all ω, there exists a k ′ such that Pr(a k ′ |ω) > 0. This implies that for all i ̸ = j and all ω, Pr(ω|s i ) Pr(a k ′ |ω)Pr(s j |ω) = 0 and thus Pr(ω|s i )Pr(s j |ω) = 0. This means that if state s i produces an observation ω, then ω cannot be produced by any other reachable state s j ̸ = s i , where two states are said to be reachable if there exists a sequence of actions sampled from the policy that enable the agent to reach state s i from s j with nonzero probability. In other words, each observation uniquely identifies the hidden state, and the POMDP is a block MDP with corresponding Markovian observations. Next, we consider the case where i = j. We have that for all i, k, ω Pr(ω|s i )Pr(a k |ω)(1 -Pr(s i |ω)) = 0. Because each term is nonnegative, this is equivalent to Pr(ω|s i )Pr(a k |ω)(Pr(s i |ω)-1) = 0. Because we again have that for all observations there exists an action a k ′ with nonzero probability, this means we can choose k = k ′ to find Pr(ω|s i ) = 0 or Pr(s i |ω) = 1 for all ω, s i . This means that either the state s i does not produce an observation ω, or the observation ω uniquely determines which state the agent is in. For all ω and i = j, either Pr(ω|s i ) = 0 or Pr(s j |ω) = 1, so if we restrict our focus to the set O(s i ) := {ω ∈ Ω : Pr(ω|s i ) > 0}, we see that ω∈O(si) Pr(ω|s i ) Pr(s j |ω) = ω∈O(si) Pr(ω|s i ) • 1 = 1. For the remaining ω / ∈ O(s i ), Pr(ω|s i ) = 0.To recap, the first case tells us that for all ω and i ̸ = j, ω∈Ω Pr(ω|s i ) Pr(s j |ω) = 0. The second case tells us that for all ω and i = j, ω∈Ω Pr(ω|s i ) Pr(s j |ω) = 1. Combining both these cases, we see that ω Pr(ω|s i ) Pr(s j |ω) = 1[s i = s j ], which we can write more succinctly as: ΦW = I. We call ΦW the state confusion matrix. For block MDPs, observations cause no confusion about which state the agent is in.Lastly, by going backwards through the proof, we see that the converse is also true. If the system is a block MDP, then Π S = ΦW Π and soK λ π = K λ ′ π .F Proof that matching transition-policies is equivalent to an MDP.Assume that T K = T K ′ . We can reduce this as follows:0 = T K -T K ′ = T (λΠ S + (1 -λ)ΦW Π ) -T (λ ′ Π S + (1 -λ ′ )ΦW Π ) = (λ -λ ′ )T (Π S -ΦW Π )which implies either λ = λ ′ or T Π S = T ΦW Π .Claim: if P is a POMDP such that T Π S = T ΦW Π for any policy, then P is equivalent to an MDP.Recall that the transition function of the effective MDP model is W T Φ, and the rewards are W R SA . Intuitively, if T and W T Φ are equivalent transition functions, it means they induce the same visitation after any number of steps. Note that due to the premise, the following quantities are equal:W I SA = I ΩA . . W A ; W (T Π S ) = W (T ΦW Π ) = W T ΦΠ . . W A = (W T ΦΠ) . . W A ; W (T Π S ) 2 = W (T Π S . . T Π S ) = W (T ΦW Π . . T ΦW Π ) = W (T ΦΠ . . W A . . T ΦΠ . . W A ) = W (T ΦΠ . . W T ΦΠ . . W A ) = (W T ΦΠ) . . (W T ΦΠ) . . W A = (W T ΦΠ) 2 . . W A ; W (T Π S ) 3 = W (T Π S . . T Π S . . T Π S ) = W (T ΦW Π . . T ΦW Π . . T ΦW Π ) = W (T ΦΠ . . W A . . T ΦΠ . . W A . . T ΦΠ . . W A ) = W (T ΦΠ . . W T ΦΠ . . W T ΦΠ . . W A ) = (W T ΦΠ) . . (W T ΦΠ) . . (W T ΦΠ) . . W A = (W T ΦΠ) 3 . . W A ; and so on.Multiplying each successive equation by γ i and summing the results gives:W (I SA -γT Π S ) -1 = (I ΩA -γ(W T Φ)Π) -1 . . W A .For any given observation, the POMDP model and the effective MDP model make the same predictions about the expected discounted visitation to future state-action pairs. In other words, they have the same successor representation.Multiplying both sides on the right by an arbitrary reward function R SA ,W (I SA -γT Π S ) -1 . . R SA = (I ΩA -γ(W T Φ)Π) -1 . . W A . . R SA = (I ΩA -γ(W T Φ)Π) -1 . . (W R SA ),we see that the MC value according to the POMDP model (left hand side) is equal to the TD value under the effective MDP model (right hand side). Both models make the same value predictions for any reward function, therefore the models are equivalent.We show an example of such a POMDP and its equivalent MDP in Figure 7 . In this section, we clarify theoretical details in augmenting a POMDP with memory and the λdiscrepancy norm.As referenced in Section 4, here we will explain how to define a memory-augmented POMDP from a base POMDP (S, A, T, R, Ω, Φ, γ). Given a set of memory states M, we will augment the POMDP as follows:S M := S × M A M := A × M Ω M := Ω × M T M : S M × A M → ∆S M ; (s 0 , m 0 , a 1 , m 1 ) → T (•|s 0 , a 1 ) × δ m1 (•) R S M A M : S M × A M → R; (s 0 , m 0 , a 1 , m 1 ) → R(s 0 , a 1 ) Φ M : S M → ∆Ω M ; (s 0 , m 0 ) → Φ(• | s 0 ) × δ m0 (•) γ M := γ,where δ x0 (•) is the discrete point distribution centered at x 0 , and × between two distributions denotes the product distribution. To demonstrate what this means, the notation above is equivalent to seeing T M instead as a function mapping to [0, 1] and definingT M : S M × A M × S M → [0, 1], (s 0 , m 0 , a 1 , m 1 , s 2 , m 2 ) → T (s 2 |s 0 , a 1 )1[m 1 = m 2 ],where a memory-augmented action (a 1 , m 1 ) means to take action a 1 and set the memory state to m 1 . This augmentation scheme uses the memory states M in three ways: as augmentations of states, actions, and observations. The state augmentation concatenates the environment state S with the agent's internal memory state M. Meanwhile, the action augmentation A M provides the agent with a means of managing its internal memory state. Together, these augmentations allow writing the augmented transition dynamics T M , which are defined so as to preserve the underlying statetransition dynamics T while allowing the agent full control to select its desired next memory state. The observation augmentation Ω M provides the agent with additional context with which to make policy decisions, and the observation function Φ M preserves the original behavior of the observation function Φ while giving the agent complete information about the internal memory state.We define an augmented policy π M and its tensor counterparts Π M and Π S M over observations and states, respectively, as follows:π M : Ω M → ∆A M ; Π M is of size (Ω M × Ω M × A M ); Π S M is of size (S M × S M × A M ); where Π M (ω 0 , m 0 , ω 1 , m 1 , a 2 , m 2 ) = 1[(ω 0 , m 0 ) = (ω 1 , m 1 )]π M (a 2 , m 2 |ω 1 , m 1 ); Π S M (s 0 , m 0 , s 1 , m 1 , a 2 , m 2 ) = 1[(s 0 , m 0 ) = (s 1 , m 1 )] (ω,m)∈Ω M Φ M (ω, m|s 0 , m 0 )π M (a 2 , m 2 |ω, m).We extend the Pr(s|ω) tensors to S M and Ω M as follows:W M is of size (Ω M × S M ); W Π M is of size(Ω M × S M × A M ); where W M (ω 0 , m 0 , s 1 , m 1 ) = 1[m 0 = m 1 ] Pr(s 1 |ω 0 ); W Π M (ω 0 , m 0 , s 1 , m 1 , a 2 , m 2 ) = 1[m 0 = m 1 ] Pr(s 1 |ω 0 )π M (a 2 , m 2 |ω 0 , m 0 ).we assign the (ω, a) entry of C the weight (1, π(a|ω)). In the occupancy-weighted case, we assign the (ω, a) entry of C the weight (Pr(ω), π(a|ω)), where Pr(ω) is proportional to the discounted observation occupancy.In Figure 4 and Figure 5 , we use the policy-weighted L 2 norm. We use this norm due to issues with performing closed-form gradient descent on the T-maze POMDP. In T-maze, states are aliased in two ways: aliasing between the two hallways (goal up or down), and aliasing between the hallway states. Occupancy weighting puts more weight on discriminating between the hallway states, since the hallway observations appear more frequently for most policies. Uniformly weighting the λdiscrepancy puts more weight on the initial observations, which results in memory functions that place more emphasis on resolving λ-discrepancy due to the starting observations. Closed-form experiments with occupancy weighting yielded less consistent results for T-maze. This seems to only be an issue for closed-form calculation of λ-discrepancy, as we show in Appendix I.6 that occupancy weighting in the sample-based setting does not detract from performance in these environments.In Figure 3 , we use the occupancy-weighted max norm to avoid artifacts due the changing observation function (see text in Section 3.2).In our sample-based settings in Figures 6, 11, 12, and 14, and Table 2 , we use the occupancy-weighted L 2 norm, since it is the simplest choice when using samples as it does not require re-weighting or importance sampling.Regardless of which norm we choose, we can estimate gradients as in Equation 7because:Λ λ1,λ2 P,π := Q λ1 π -Q λ2 π = ω,a Pr(ω) Pr(a|ω)(Q λ1 π (ω, a) -Q λ2 π (ω, a)) 2 1/2 = ω,a Pr(ω, a)(Q λ1 π (ω, a) -Q λ2 π (ω, a)) 2 1/2 = E (ω,a)∼π Q λ1 π (ω, a) -Q λ2 π (ω, a) 2 1/2 .In practice, when optimizing the λ-discrepancy, we omit the outer square root and simply compute the mean squared difference between Q-values, since the square-root function is a monotonic transformation and its argument is non-negative.We now describe the proof-of-concept experiments from Section 4 involving closed-form memory optimization on a range of small-scale classic POMDPs. The environments are: T-maze [Bakker, 2001] , the Tiger problem [Cassandra et al., 1994] , Paint [Kushmerick et al., 1995] , Cheese Maze, Network, Shuttle [Chrisman, 1992] , and the partially observable version of the 4 × 3 maze [Parr and Russell, 1995] .We first begin by detailing the small-scale POMDPs we use. Then we describe the algorithm used to calculate closed-form λ-discrepancy gradients, as well as procedures for policy improvement after learning a memory function. Apart from the T-maze, all other POMDPs used in Section 4 were taken from pre-defined POMDP definitions [Cassandra, 2003] .We made one slight modification to the Tiger environment that preserves the original environment behavior but adapts the domain specification to match our formalism such that observations are only a function of state. The original Tiger domain used a hand-coded initial belief distribution that was uniform over the two states L/R, and did not emit an observation until after the first action was selected. Thereafter, the observation function was action-dependent, with state-action pair (L, listen) emitting observations left and right with probability 0.85 and 0.15 respectively, and other actions (L, *) emitting uniform observations and returning to the initial belief distribution. Since our agent does not have access to the set of states, it cannot use an initial belief distribution.To achieve the same behavior, we modified the domain by splitting each state L/R into an initial state L 1 /R 1 that always emits an initial observation, and a post-listening state L 2 /R 2 that uses the 0.85/0.15 probabilities. We visualize these changes in Figure 8 . This type of modification is always possible for finite POMDPs and does not change the underlying dynamics. We use T-maze with corridor length 5 as an instructive example (see Figure 1 ). The environment has 18 underlying MDP states: one initial state for each reward configuration (reward is either up or down), five for each corridor, one for each junction, and finally one for each terminal state.11 There are 5 observations in this environment -one for each of the initial states, a corridor observation shared by all corridor states, a junction observation shared by both junction states, and a terminal observation. The action space is defined by movement in the cardinal directions. If the agent tries to move into a wall, it remains in the current state. From the junction state, the agent receives a reward of +4 for going north, and -0.1 for going south in the first reward configuration. The rewards are flipped for the second configuration. The environment has a discount rate of γ = 0.9. This environment makes it easy to see the differences between MC and TD approaches to value function estimation. We visualize these differences in Figure 9 . MC computes the average value for each observation by averaging the return over all trajectories starting from that observation. By contrast, TD averages over the 1-step observation transition dynamics and rewards, and bootstraps off the value of the next observation. For a policy that always goes directly down the corridor and north at the junction, this leads to an average (undiscounted) return for the blue observation of +4 with MC and (4 -0.1)/2 = 1.95 with TD.Algorithm 1 describes our memory optimization procedure, which reduces the λ-discrepancy to learn a memory function, then learns an optimal memory-augmented policy.Input: Randomly initialized policy parameters θ π , where Π = softmax(θ π ), randomly initialized memory parameters θ µ , POMDP parameters P := (S, A, T, R SA , Ω, Φ, p 0 , γ), number of memory improvement steps n steps,µ , number of policy iteration steps n steps,π , learning rate α ∈ [0, 1], and number of initial random policies n.Output: Optimized memory-augmented policy parameters θ πµ and memory parameters θ µ .// Initialize random policies and select policy to fix for memory learning. {θ 0 , . . . , θ n-1 } ← randomly init n policies(n) θ π ← select argmax lambda discrepancy(θ 0 , . . . , θ n-1 }) // Repeat policy over all memory states. θ πµ ← repeat(θ π , |M|) // Improve memory function. θ µ ← memory improvement(θ µ , θ πµ , P, n steps,µ ) // Augment POMDP with memory function. P µ ← expand over memory(P, θ µ ) // Improve memory-augmented policy over memory-augmented POMDP θ πµ ← policy improvement(θ πµ , P µ , n steps,π ) return θ πµ , θ µIn Algorithm 1, the policy improvement function can be any function which improves a parameterized policy. We use an analytical version of the policy gradient [Sutton et al., 1999] algorithm. The randomly init n policies function returns n randomly initialized policies, and the select argmax lambda discrepancy function picks the one with the largest λ-discrepancy. The memory improvement and expand over memory functions are defined in Appendices H.2.1 and H.3, respectively.We have noticed that a larger λ-discrepancy tends to provide a better signal for learning memory functions. Sampling random policies for memory improvement is highly likely to reveal a λdiscrepancy, as suggested by the theory in Section 3.2. For this reason, we consider many random policies (n = 100), then use the policy which had maximum λ-discrepancy as the basis for memory optimization.Note that einsum is the Einstein summation, and kron is the Kronecker product. The augmented initial state distribution is simply the same distribution as p 0 , except with 0 probability mass over all non-zero memory states, since the memory state always initializes to memory state 0.For all experiments in Section 4, we run memory optimization on the suite of POMDPs with the following hyperparameters. We optimize memory for n steps,µ = 20K steps and run policy iteration for n steps,π = 10K steps. For all gradient-based experiments, we use the Adam optimizer [Kingma and Ba, 2015] .For the belief-state baselines, solutions were calculated using a POMDP solver from the pomdp-solve package [Cassandra, 2003] . The performance of the belief-state optimal policy was calculated by taking the dot product between the initial belief state and the maximal alpha vector for that belief state. This returns a metric comparable to the initial state distribution weighted value function norm, which we use as a performance metric for our memory-augmented agents.The belief-state solution for the 4 × 3 maze was solved using an epsilon parameter of ϵ = 0.01, due to convergence issues with the environment when utilizing the POMDP solver.In Section 3.3 we introduced the Parity Check environment as an example of a POMDP with zero λ-discrepancy for all policies. There we showed that in such environments, we can still detect partial observability by using randomly initialized memory functions. But the question remains: how prevalent are such environments?We find that these sorts of environments appear to be quite rare. The symmetry disappears if we modify the example even slightly, such as by changing the start state probabilities or by introducing a "stay-in-place" action at any one color. These modifications lead to the λ-discrepancies in Figure 10 , which offer some reassurance that such edge cases are the exception and not the rule.Besides modifications to the POMDP itself, adding memory can induce a non-zero λ-discrepancy. In particular, random memory functions reveal a λ-discrepancy in the system that can be minimized to learn memory, as demonstrated in Figure 4 . We now describe this memory optimization procedure.Memory Learning Memory optimization for the Parity Check experiments in Figure 4 (right) followed almost the same procedure as Algorithm 1 except for selecting the initial policy. In this case, any randomly initialized memoryless policy has zero λ-discrepancy, but augmenting with a randomly initialized memory function reveals a non-zero measure. Thus, for this experiment, we augment our POMDP with a random memory function before choosing the initial policy that maximizes the λdiscrepancy. All other algorithmic details remained the same. I Scaling the λ-Discrepancy with PPO and RNNsWe build our λ-discrepancy minimizing reinforcement learning algorithm on top of online recurrent PPO. Normally, PPO would have two losses: one for the actor and one for the critic.EQUATIONHere, L CLIP is the clipped surrogate loss for the actor, defined by:EQUATIONwhere ρ t (θ) is the ratio between policy probabilities of the old policy and new policy, Ât is the estimated advantage at timestep t, and c V and L V are the value coefficient and value function loss, respectively. L V is the mean-squared TD error, with targets calculated using a truncated version of λreturn [Sutton and Barto, 2018] . Finally, c Ent and H[π µ (• | ω t , z t-1 )] are the entropy coefficient and entropy of the policy π µ that is parameterized by θ RNN and θ Actor . This term is used for exploration through entropy maximization [Mnih et al., 2016] . For our memoryless baseline, we use a nonrecurrent architecture that is a function of only the current observation (and optionally the most recent action).Our proposed algorithm replaces L V with L V,λ1,λ2 , which has three components: two value losses for two separate value heads (each estimating λ-returns for different λs), and our λ-discrepancy auxiliary loss.L V,λ1,λ2 (θ Critic , θ RNN ) = βL Λ (θ) + (1 -β) (L V,λ1 (θ Critic1 , θ RNN ) + L V,λ2 (θ Critic2 , θ RNN )) ,where L Λ (θ) is our λ-discrepancy loss, as defined in Equation 8, and L V,λ is the λ-return-as-target TD error for a given λ. Here, β ∈ [0, 1] is a hyperparameter trading off between the λ-discrepancy loss and the value losses. Finally, the agent only uses V λ1 as its value estimate for its advantage function.The computational and memory requirements for PPO with and without the λ-discrepancy auxiliary loss are similar. While the λ-discrepancy loss requires training a second value head, this only adds two additional layers to the critic network and results in a small number of additional parameters (about 12% overhead). As for computation, the addition of this auxiliary loss only adds a small constant factor to the backpropagation of gradients, since we learn everything end-to-end. Overall, the wall-clock time of runs between the baseline algorithm and our algorithm is comparable.All four of the environments used for evaluation in Section 5 were all re-implementations of environments used in Silver and Veness [2010] , in JAX [Bradbury et al., 2018] , allowing for massive hardware acceleration. We now give details of these environments.Partially observable Battleship, a one-player, limited observation variant of the popular two-player board game, was first introduced as a benchmark for partially observable Monte-Carlo planning [Silver and Veness, 2010] . This variant is particularly challenging since the agent has to reason about the unknown positions of the ship and keep track of past shots. The observation space is Ω = {0, 1}; at every step, the agent receives a binary observation: 0 if the last shot missed a ship, and 1 if the last shot hit a ship. The state space in this game contains all possible board states, which is astronomically large. The action space is A = {1, . . . , 10} × {1, . . . , 10}, or all possible grid locations. The agent is only allowed to take valid actions at every step (no position can be selected twice), which is achieved through action masking of the actor. This makes the problem a finite horizon problem, with a horizon at most 10 × 10 = 100. The environment terminates when all positions on the grid with a ship are hit. The rewards at every step are -1, with a positive reward of 10 × 10 = 100 when all ships are hit/the environment terminates. The discount factor γ is set to 1 here, since we are in the finite horizon setting.At every environment reset, 4 ships of length (5, 4, 3, 2) are uniformly randomly placed on a 10 × 10 grid. At every step, to allow for easier learning of all agents, we concatenate the last action to the observation, so that our RNN conditions on the entire history, as opposed to only observations.Partially observable PacMan (a.k.a. PocMan), a variant of the popular arcade game, was also introduced to test partially observable Monte-Carlo planning with a simulator [Silver and Veness, 2010] . In this version, the agent can only observe indicators of its surroundings, including walls in its cardinal directions, whether there is a ghost in its line of sight, a power pellet nearby, or food nearby. These highly obscured agent-centric observations require the agent to localize within the map, seek food and power pellets, and avoid ghosts with only local sensor information.Concretely, the observation space is an 11-dimensional binary vector. The first four values are 1 if there is a wall in each of the four cardinal directions of the agent. The next element is 1 if there is food Manhattan distance 2 or less away from the agent. The next value is 1 if there is a ghost a Manhattan distance 2 or less away. The next four elements are 1 if there is at least one ghost in the line of sight in each of the four cardinal directions. The last element is if the agent currently has a power pellet.As with Battleship, the state space in PocMan is also very large. The map is a 19 × 21 sized map defined in the original instantiation of the environment [Silver and Veness, 2010] , and the state consists of the agent and ghost locations, as well as the binary status of every dot and power pill. The action space has four actions, corresponding to moving in the four cardinal directions. The agent receives a reward of +200 if it eats a ghost, +10 if it collects a pellet, and +20 if it collects a power pellet. The discount factor is set to γ = 0.95.The episode terminates when the agent is killed by a ghost, the agent collects all pellets, or the timer for the environment elapses. At the beginning of each episode, all ghost locations are reset, and the agent is placed in the same fixed starting position. In this setting, we also append the agent's most recent action to its observation.In our work, we implement this domain by adding the observation function on top of the PacMan environment in the Jumanji reinforcement learning framework [Bonnet et al., 2024] .RockSample [Smith and Simmons, 2004] , a well-known partially observable benchmark in POMDP planning literature, is a rock-collecting problem simulating a rover scanning potential rock samples and collecting them if they are desirable. RockSample (n, k) is a problem with a grid of size n × n, and k rocks distributed randomly throughout the environment. The two variants we test on have a grid size of 11 × 11 and 15 × 15, as well as 11 and 15 rocks in each environment, respectively. When instantiated, each RockSample environment for each seed randomly samples rock positions. At the reset of each environment, each of these rocks are uniformly randomly assigned to be either a good or bad rock. The action space in this environment is 4 + 1 + k = 5 + k. The first 4 actions correspond to moving in the cardinal directions. Action 5 corresponds to sampling in the current position. The last k actions correspond to checking each of the k rocks. Checking a rock will probabilistically tell the agent the correct parity of the rock, depending on the half-efficiency distance d hed and the l 2 distance d i (s) between the agent's position and the rock being checked, given the current state s:EQUATIONIn this work, we simply set d hed to be the maximum possible distance between any two points in the grid. Traversing to and then sampling a good rock gives a reward of +10. Sampling a bad rock gives a reward of -10. Exiting to the east border of the environment will terminate the environment and result in a reward of +10. After sampling a good rock, sampling the same rock again will result in a reward of -10. This environment has a discount rate of γ = 0.99The state space of this environment is all the possible combinations of rock positions, as well as agent positions and rock parities. The agent receives observations of the form of a vector of size 2n + k binary values, where the first 2n elements are a two-hot representation of the xy-coordinates of the agent. The final k observations are set after choosing a check action, and set a 1 at the position of the checked rock if it appears to be good, and a 0 if it appears to be bad. This observation function is a slight departure from the previous RockSample problem definition [Smith and Simmons, 2004, Silver and Veness, 2010] in two aspects, based on an implementation of RockSample that is used to test RNNs [Tao et al., 2023] , as opposed to model-based planning algorithms. The first is that after an agent checks a rock and the agent receives a positive sensor reading, the observation bit corresponding to this rock is set to 1, and remains 1 until sampled. This makes the memory learning portion of the problem easier, but the agent is still required to remember which rocks it has sampled before, and also deal with the stochasticity of checking. The second change is to not have an explicit negative sensor reading. Instead, after checking a rock, the observation bit corresponding to the sensor reading of the rock remains 0. This makes the problem harder, since the agent has to infer the negative sensor reading, and disambiguate it from a null sensor reading, when the agent takes non-check actions.In the main text (Section 5), we follow the common practice of reporting results for undiscounted returns. However, since our optimization objective considers discounted returns, we also include learning curves for discounted returns in Figure 11 . For ease of comparison, we also duplicate the undiscounted learning curves (from Figure 6 ) in Figure 12 . Note that Battleship is unchanged, since that domain is finite-horizon and undiscounted. Our base PPO algorithm is an online PPO algorithm that trains over a vectorized environment, all parallelized using JAX [Bradbury et al., 2018] and the PureJaxRL batch experimentation library [Lu et al., 2022] with hardware acceleration. The hyperparameter sweep was performed on a cluster of NVIDIA 3090 GPUs, and the best seeds presented were run on one GPU for each algorithm, running for 1 to 12 hours, depending on the domain. Our environment is vectorized over 4 copies, and we use truncated-backpropogation through time [Jaeger, 2002] as our gradient descent algorithm, with a truncation length of 128. Our L CLIP clipping ϵ is set to 0.2. The value loss coefficient is set to c V = 0.5. We also anneal our learning rate over all training steps, and clip gradients when the norm is larger than 0.5 by their global norm [Pascanu et al., 2013] .We now detail the hyperparameters swept across all environments and all algorithms. We do so in Table 1 .Step size [2.5 × 10 -3 , 2.5 × 10 -4 , 2.5 × 10 -5 , 2.5 × 10 -6 ] λ 1 [0.1, 0.5, 0.7, 0.9, 0.95] λ 2 (λ-discrepancy) [0.1, 0.5, 0.7, 0.9, 0.95] β (λ-discrepancy)[0, 0.125, 0.25, 0.5]Table 1 : Hyperparameters swept across all algorithms. Rows labelled with λ-discrepancy are hyperparameters swept specific to our algorithm.We conduct this sweep across all environments for 5 seeds each, and use the highest area under the learning curve (AUC) score in each environment to select the best hyperparameters for each environment. Note, we swept over 10 seeds for recurrent PPO and λ-discrepancy-augmented recurrent PPO for the PocMan environment, due to high variance in returns. After hyperparameter selection, we re-run all algorithms on all environments with the best selected hyperparameters for 30 different seeds to produce the learning curves in Figure 6 . We present the best hyperparameters found for both algorithms in Figure 2 .Step size λ 1 λ 2 β Battleship 2.5 × 10 -4 0.1 0.95 0.5 PocMan 2.5 × 10 -4 0.95 0.5 0.5 RockSample (11, 11) 2.5 × 10 -4 0.1 0.95 0.5 RockSample (15, 15) 2.5 × 10 -4 0.1 0.5 0.25 (a) Best hyperparameters found for λ-discrepancy-augmented recurrent PPO.Step size λ 1 Battleship 2.5 × 10 -5 0.7 PocMan 2.5 × 10 -5 0.7 RockSample (11, 11) 2.5 × 10 -4 0.7 RockSample (15, 15) 2.5 × 10 -5 0.5 (b) Best hyperparameters found for the recurrent PPO baseline.Step size λ 1 Battleship 2.5 × 10 -5 0.1 PocMan 2.5 × 10 -4 0.7 RockSample (11, 11) 2.5 × 10 -4 0.1 RockSample (15, 15) 2.5 × 10 -5 0.1 (c) Best hyperparameters found for the memoryless PPO baseline.Table 2 : Best hyperparameters for each environment and each algorithm. Hyperparameters were found using 5 seeds, and taking the maximum AUC.Our network architectures are standard multi-layer perceptions (MLPs). Both actor and critic networks are two-layer MLPs with ReLU [Nair and Hinton, 2010] activations between layers. The actor network applies a softmax function to its output logits. Our recurrent neural network is a dense layer, followed by ReLU activation, then the GRU cell, then another dense layer. For Battleship, to better condition on the hit-or-miss observation, we add an additional dense layer after the first dense layer of the RNN (which has 2× latent size for this environment), that takes as input the outputs of the previous layer concatenated with the hit-or-miss bit (a residual connection). For our memoryless baseline, we replace the recurrent neural network with a 3-layer MLP with ReLU activations between each layer. Put together, our actor-critic network takes in an input observation (and optionally, the previous action) and passes it, as well as the previous latent state, through the GRU for a new latent state. For the memoryless baseline, the 3-layer MLP simply encodes the observation (and potentially action) into the latent state, with no recurrence. It uses this new latent state as inputs to both the actor and critic.We now detail environment-specific hyperparameters in Table 3 Training was done over 2M time steps, where 1M steps were collected from each of these agents.After collection, all trajectories were run through both agents to collect 2M time steps to collect both RNN + PPO hidden states and the RNN + PPO + LD hidden states. With this dataset, the each probe was trained with the corresponding RNN hidden states as inputs, with the pellet occupancy of all potential pellet positions as the target for the prediction. Both probes are 3-layer neural networks with ReLU activations [Nair and Hinton, 2010] between each layer, and a sigmoid function over the final logits to map outputs to 0 and 1. We use a binary cross entropy loss between these predictions and the pellet occupancy targets. The hidden size of the network was 1024, with a step size of 0.0001. At every step of training, a batch of 32 time steps are uniformly randomly sampled from the dataset, and the binary cross entropy loss was minimized. The agent performs 10M steps of training to reach the performance visualized.We also run both recurrent PPO and our λ-discrepancy-augmented PPO algorithm on the small-scale POMDPs evaluated in the analytical experiments in Section 4. For these experiments, all agents had a latent state size of 32, with c Ent = 0.05, and the same hyperparameters swept as in Appendix I.4. We show results in Figure 14 .These results imply that the baseline algorithm is already sufficient for solving these tasks, and any additional auxiliary losses cannot help with performance, since performance is already near optimal. We would like to note that, as with the results in Figure 6 , adding our λ-discrepancy auxiliary loss never harms the performance of PPO.Figure 14 : The λ-discrepancy performance over recurrent PPO on small-scale POMDPs. Learning curves show mean and 95% confidence interval over 30 seeds.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The theoretical claims posed in the abstract and introduction such as "the λdiscrepancy is exactly zero for all Markov decision processes and almost always non-zero for a broad class of partially observable environments" are presented with linked proofs in Section 3. Our analytical results that show mitigation of partial observability are presented in Section 4. Our RNN-based results in Section 5 show that minimizing the λ-discrepancy between two learned value functions as an auxiliary loss alongside traditional reinforcement learning, is often significantly more effective (and never worse) than training the same agent with just a single value function.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Code: https://github.com/brownirl/lambda_discrepancy Videos: https://lambda-discrepancy.github.ioOne intuitive way to think about Pr(s|ω) is that it reflects the relative frequency of trajectories matching both s and ω versus those just matching ω:Eτ [ ∞ t=0 γ t 1[s t =s,ω t =ω]] Eτ [ ∞ t=0 γ t 1[ω t =ω]].This equivalence justifies our choice of Pr(s|ω) when defining TΩ and RΩ for the effective MDP model, since states appear in precisely this proportion when we generate experiences in the POMDP.For 3-dimensional tensors A and B, (AB) ijlm = k A ijk B klm , and (A . . B) il = jk A ijk B jkl .Such artifacts are not a concern in normal environments, since the observation function is fixed.This environment requires projection by both W and R SA to zero the differences; neither alone is sufficient.These domains have extremely large state spaces: Battleship has ∼2 100 states, without accounting for possible ship locations, and PacMan has ∼5 400 states, even when ignoring food dot configurations.Technically, we group the four terminal states into a single state for conciseness, which is functionally equivalent. 


Decision: This paper tackles reinforcement learning problems under partial observability. The paper proposes a new metric to build state representation based on the TD($\lambda$), which makes the novelty of the method is clear. While the empirical analysis still seems limited, missing a comparison with state-of-the-art RL methods for POMDPs, the paper provides sufficient theoretical support for the proposed approach. Considering all reviewers are positive about the paper, I recommend its acceptance.

Reviews: review: paper_summary: paper_summary: The Authors address the problem of missingness or partial observability in Markov Decision Problems (MDPs), where the Markov assumption does always hold. They introduce a measure, the Lambda Discrepancy (LD), which serves as an indicator for how closely a Markovian state representation is given.

The LD learns two different value estimates, ideally one making an implicit Markov assumption and one that does not. It is also shown that using the LD it is often clear whether a problem can be defined with a Markovian state representation or not.

Further, the authors use the LD to define an auxiliary loss and included it into a Proximal Policy Optimization algorithm. By using two value estimators for, one for each of two distinct lambdas, the resulting agent outperformed agents that were trained without LD on multiple tasks.
strength_weakness: strength_weakness: - In section “Related Work”, it is said that the “dominant paradigm for resolving partial observability are memory-based approaches”. 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.
    
- The paper shows experimental results for POMDPs, but there is no example of an environment that is a pure MDP for comparison. Ideally, it would be most interesting seeing results for environments with Markovian state representation that gradually shift to POMDPs that are not Markovian.

- The “Limitations” section should be part of the main corpus.
    
- 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.
    
- The paper explains how to interpret the extremes of $\lambda$ being 0 or 1. However the explanation of “Intermediate values of λ smoothly interpolate between the two” ([pdf](zotero://open-pdf/library/items/7RNSWK98?page=5&annotation=ESN6JY42)) (l. 163) values is not intuitive. An example at $\lambda=0.5$ would be helpful, where not only the previous state is taken into account (as would be with $\lambda=0$ in an MDP setting) but neither all of the previous observation (as would be $\lambda=1$), but a balance in between.
    
- While it is stated that the “approach scales to challenging partially observable domains” the experiments remain on a low complexity levels and no comment is made about how the scaling the complexity of environments affects the performance using the LD method.
questions: questions: - Q1: Could you please clarify the setup regarding the $\text{L}_\Lambda$ loss. You train two value function models for two distinct $\lambda$ values in parallel? Hence, the models share a single loss function? This means after training is complete we end up with two models. How are the results of both networks used/combined?
    
- Q2: How does the LD method benefit when there is existing domain knowledge about an environment? How would knowledge whether the environment should be modeled as a (PO)MDP including how many prev. states are required to describe the next state be taken into account? In that case, what are the expected advantages of using the LD method?
    
- Q3: How is the LD method affected when switching between online and offline-RL settings?
limitations: limitations: - Using LD for more complex environments results in adding another value estimator to the optimization algorithm. For complex environments requiring large models this can become a limiting factor.
    
- Introducing the LD loss introduces multiple hyperparameters $\lambda_1, \lambda_2$ and $\beta$. While it is clear from empirical research that the $\lambda$s should be on separate ends of the interval $[0,1]$ , it is less so for $\beta$. Performing hyperparameter optimization for these three additional parameters is a further addition to computational requirements using LD.
paper_summary: paper_summary: The authors consider a problem of reinforcement learning in partially observable systems. Instead of constructing an information state using the entire trajectory history with some peripheral objective, the authors propose using discrepancy in the values of TD$(\lambda)$ estimates for different values of $\lambda$. Their argument is that TD(0) makes a Markovian assumption, whereas TD(1) (Monte Carlo) does not, therefore any discrepancy in their values can be used to construct a metric called the $\lambda$-discrepancy, which measures the non-Markovian nature or partial observability of the currently used state representation. The authors then use this $\lambda$-discrepancy to construct an RL algorithm that learns a memory function to mitigate partial observability. They do this by traiing two different recurrent networks to predict TD$(\lambda)$-values fpr two different values of $\lambda$ and add an auxiliary loss that minimizes the discrepancy between these twp. Using numerical experiments, they show that their method performs better or at least as good as conventional history based methods for POMDPs.
strength_weakness: strength_weakness: The metric defined by the authors requires operating over the entire state, observation and action space for its computation as given by the authors in their theoretical derivation. The authors do not explain how the per state/observation, action value determined in practice relates to this metric?
questions: questions: 1. The authors state in Line 146 that: "Without state, a model-based comparison would require variable-length history inputs to ensure a Markovian representation, and the computational and memory requirements would scale exponentially with history length. Instead, we propose a model-free approach, using value functions defined only over observations, to approximate this comparison in a tractable way." However, the authors also use two RNNs in their proposed approach that uses the entire history. So how does this argument hold?
2. The authors have given one example (Parity check environment) where equation (7) is satisfied and the $\lambda$-discrepancy metric fails. The authors also show that this can be addressed by perturbing the initial state distribution or action space. 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.
3. How will the convergence analysis of PPO change after the addition of the auxiliary loss given by equation (8)? The authors have proven that $\lambda$-discrepancy reflects partial observability only when the corresponding Q functions are learnt. When these functions are being learnt (still have errors) the $\lambda$-discrepancy will also have an error. How do the authors address this?
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?
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?
6. How does the memory learnt by this approach compare with representations learnt by POMDP based approaches such as PSRs etc.?
7. Can $\lambda$-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 equation (8) relates to the one in Definition 1?
limitations: limitations: The authors address the limitations of their work in a separate section in the Appendix.
paper_summary: paper_summary: This paper studies policy planning in a partially observable Markov decision process (POMDP). It has been acknowledged that planning in a POMDP is computationally challenging in general. This paper attempts to address this challenge by proposing novel planning heuristics First, the authors formalize the $\lambda$-discrepency, which measures the distance of $\lambda$-returns using different $\lambda$ parameters. The learner then obtains a policy by minimizing the $\lambda$-discrepency. Simulations results demonstrate that the proposed method could outperform standard PPO using recurrent NN in a select set of POMDP environments.
strength_weakness: strength_weakness: - As for Theorems 1 and 2, do the policies $\pi$ 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. Could the authors make the connection with the weakly revealing POMDPs?
- The authors could further elaborate on the training procedures for minimizing the loss function in Eq. 8. Please see my questions in the next section.
- Additional discussion on why minimizing $\lambda$-discrepency could improve the planning performance would be appreciated. Could the authors provide some intuitions?
questions: questions: - What does LD stand for in Figure 4?
- How to select parameters $\lambda_1$ and $\lambda_2$ when minimizing the loss function in Eq. 8? Does one have to try multiple pairs of such $\lambda$s?
- For policies $\pi$ 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.
limitations: limitations: The authors have adequately addressed the limitations of the paper.
paper_summary: paper_summary: The authors propose the lambda discrepancy, a measure of partial observability based on TD-lambda. Importantly, this metric can be estimated from samples of POMDP observations (and a bootstrapped value function estimator) without relying on a transition model, access to POMDP states, or even knowledge of what information the states contain. It is shown theoretically that the lambda discrepancy is zero under a Markovian state representation and almost always non-zero otherwise. Experimentally, it is shown that optimizing the lambda discrepancy can improve representation learning and downstream performance both on smaller POMDP planning domains as well as larger problems solved by PPO.
strength_weakness: strength_weakness: While my overall impression of the idea is highly positive, unfortunately I believe that it has not been investigated sufficiently in this work. I don't find the theory/discussion/experiments entirely convincing that the lambda discrepancy is a reliable and broadly applicable metric across general POMDP problems. Furthermore, the manuscript has some serious clarity issues and perhaps tries to convey too much material in too little space (and relegating a lot of important parts to the appendix). 

**Quality**: The discussion of some important points -- both theory-wise and experiment-wise -- is missing. In (rougly) chronological order:

(Theoretical:) 
- While I appreciate that there is a discussion of when the lambda discrepancy can be 0 aside from when the environment is an MDP (which explores its limitation as a measure of "non-Markovian"-ness), I'm not sure what the conclusion is. Based on Eq 7 and the subsequent discussion, it seems that there are many ways this can occur and no good way of characterizing them. (I couldn't really understand the text nor the example in this section, but if the authors disagree with this point, please bring it up in the rebuttal.)

- There should be some theoretical discussion on how to choose $\lambda_1, \lambda_2$ in the discrepancy. For example, it could be mentioned that for large values of $\lambda \approx 1$, the TD error estimator has very high variance and may not be practical for predicting the TD error.

(Experimental:)
- I would've appreciated some experiments of an "observational" nature to verify the lambda discrepancy actually measures how Markovian the state representations are. For example, maybe you could formulate a class of environments that become more "partially observable" according to a parameter $\epsilon$, and show how the lambda discrepancy changes with $\epsilon$. Another interesting experiment is to train a state representation from the history using an RNN, and compare lambda discrepancy under the state representation with some other known measures of "Markovian"-ness like bisimulation distance [1].

- I appreciate the experiments on larger POMDP planning benchmarks, since they require complex memory that seem non-trivial to learn. However, the transition/observation system is presumably simple enough for someone to specify. It would be interesting to apply the method on domains where the observation/transition function is more complex, such as image-based continuous control.

**Clarity**: Overall, I believe this is the weakest point of the paper. I had a hard time reading the paper since it glossed over important points, with a lot of material placed in the appendix. I understand the struggle to fit everything within the page limit, but there were some paragraphs that I felt did not need to be in the main text as well. A non-exhaustive list of things that could be potentially confusing to a reader:

- The definition of a *Markovian* state representation in text (line 95-97) and mathematically (Eq 1) do not describe the same thing. At the very least, Eq 1 needs to include observations $\omega$. 
- The definition of empirical MDPs (lines 128-135) is confusing for a few reasons. First, it's not a real MDP, and the way it's presented makes it seem like an existing concept rather than a new concept by the authors. Second, terms like $Pr(s | \omega)$ are not well defined without further context. I understand that it's somewhat clarified in the next sentence, but I think a precise definition from the start would help someone reading things in order. I suggest it be defined using expectations over trajectories under the policy $\pi$.
- In equation 6, is the introduction of tensor notation necessary? Is it not possible to derive this using probability notation directly? Also, what is the relevance of using Q-functions vs value functions in the definition of lambda discrepancy? 
- I cannot make sense of the explanation for when Eq 7 can be zero, even after several reads. It's also unclear to me what about the Parity Check example makes it one of the degenerate cases for Eq 7. 

**Significance**: No weaknesses. I believe this work has great potential, but the other weaknesses affect how this work would be received by the broader community. 

[1] Zhang, Amy, et al. "Learning Invariant Representations for Reinforcement Learning without Reconstruction." International Conference on Learning Representations.
questions: questions: Please see weaknesses above.
limitations: limitations: Yes.

Comments: comment: > Re: "empirical MDPs", I'm skeptical that the transition function $T(ω'|a,ω)$ is actually Markovian. Can't we get a better prediction of $ω'$ given not just $ω$ but the entire sequence of observations $ω_1, …, ω_t$ up to the current timestep?

Great question. Yes we can, and that's precisely the point. If we want to predict the *actual*, true-to-the-environment next observation, we are absolutely better off conditioning on the entire history. If we view $T(ω'|a,ω)$ as a marginal over the transition dynamics for all POMDP trajectories leading to ω, then this quantity is indeed non-Markovian, non-stationary, and policy-dependent.

At the same time, we are also free to invent a stochastic tensor of transition probabilities and pretend that is an accurate model of the environment. If we do, accurately predicting the imaginary next state in an MDP governed by this stochastic transition tensor requires no more than a single state and action (by definition). It's just that the predicted MDP state may bear no resemblance to the *actual* next observation. If the probabilities in this transition tensor are those of the empirical MDP (as we define it), then this is exactly what TD(0) is doing implicitly when we apply it to the observations of the POMDP.

So whether $T(ω'|a,ω)$ is Markovian depends on whether we are viewing it as a marginal distribution for the POMDP, or a transition function for a hypothetical MDP.

We hope this clarifies things!
Thank you all for the productive discussions! We appreciated the collaborative spirit of your reviews, and your feedback and suggestions were very helpful. We are happy to hear that we were able to address your main concerns and that you all now recommend acceptance. We will incorporate your latest feedback into the camera ready if the paper is accepted.

You identified several interesting follow-up directions that we are excited to pursue, and we are glad that you see the potential for this work to have such broad impact.

Sincerely,

Submission12090 authors
A thanks to the authors for providing a thorough rebuttal, which lead me to increase my score.
- The global rebuttal as well as the direct one were detailed and addressed many points that were raised.
- As suggested, the authors performed an experiment comparing a gradual shift from a MDP to POMDP setting and will include it in their final submission. 
- Further, the authors will provide information about the choosing of the $\lambda_1$ and $\lambda_2$ hyperparameters.
- Also, the authors arguments on environment complexity were convincing; to demonstrate that the LD is helpful is done by the experiments provided. However, to demonstrate that this is true in practice does require experiments with more complex scenarios. 
- Including the architecture diagram is also very valuable and will help readers to understand the training procedure better. (I would recommend using small labels such as 1, 2 and 3 in the right part of the figure which could be referenced in the caption.)
Thank you for your clarifications. Overall, I found the authors' rebuttal to be both substantial and effective at addressing the majority of my concerns. 

- The authors have clarified my concern re: pathological cases where LD = 0 in POMDPs. The additional experiments showing how it can be practically prevented is great. These pathological cases are reminiscent of saddle points -- theoretically an issue for gradient descent but not a significant concern in practice.

- The authors have promised to include practical insights of choosing $\lambda_1, \lambda_2$ when computing the LD, which I think will be valuable. 

- Per my suggestion, the authors have included an experiment interpolating between MDPs and POMDPs and verifying that LD increases. It's on a toy domain but this is a positive -- it's illustrative, easy to understand, and insightful.

A couple of concerns still stand:

- I still believe the work would be significantly improved by experimenting on more complex POMDPs. This would alleviate concerns about scalability, e.g. TD-$\lambda$ may not be so easy to accurately estimate for high $\lambda$. In fact, despite its popularity in deep RL, there has been evidence that similar estimates like GAE-$\lambda$ do not need to be that accurate for deep RL to work well [1].

- Re: "empirical MDPs", I'm skeptical that the transition function $T(w' | w, a)$ is actually Markovian. Can't we get a better prediction of $w'$ given not just $w$ but the entire sequence of observations $w_1, \ldots, w_t$ up to the current timestep?

In any case, I think the work has significantly improved and I'm pleased to recommend acceptance. Congratulations to the authors on a great idea.  

[1] Ilyas, Andrew, et al. "A Closer Look at Deep Policy Gradients." International Conference on Learning Representations.
I thank the authors for their detailed response. They have addressed my questions. Overall, this paper proposes novel algorithms for policy learning in MDPs with hidden confounders, and will have an impact on the AI community. On the other hand, I also agree with Reviewer rpiT regarding the comparison with other SOTA POMDP algorithms. Due to these reasons, I will keep my current score.
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.